feat(metrics): distinguish joined from regular requests (ENG-4072)

packages/api/internal/orchestrator/create_instance.go

@@ -26,6 +26,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/middleware/otel/joined"
 	sandbox_network "github.com/e2b-dev/infra/packages/shared/pkg/sandbox-network"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 	ut "github.com/e2b-dev/infra/packages/shared/pkg/utils"
@@ -159,6 +160,9 @@ func (o *Orchestrator) CreateSandbox(
 	}
 
 	if waitForStart != nil {
+		// Mark as a joined request for telemetry purposes
+		joined.Mark(ctx)
+
 		logger.L().Info(ctx, "sandbox is already being started, waiting for it to be ready", logger.WithSandboxID(sandboxID))
 
 		sbx, err = waitForStart(ctx)

packages/api/internal/sandbox/storage/memory/operations.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/middleware/otel/joined"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
@@ -183,6 +184,14 @@ func startRemoving(ctx context.Context, sbx *memorySandbox, opts sandbox.RemoveO
 			return false, nil, &sandbox.InvalidStateTransitionError{CurrentState: currentState, TargetState: newState}
 		}
 
+		if currentState == newState {
+			// The caller will inherit the in-flight transition's result
+			// without doing the work itself: this is a joiner. Mark before
+			// waiting so the request stays tagged even if the inherited
+			// transition fails.
+			joined.Mark(ctx)
+		}
+
 		logger.L().Debug(ctx, "State transition already in progress to the same state, waiting", logger.WithSandboxID(sbx.SandboxID()), zap.String("state", string(newState)))
 		err = transition.WaitWithContext(ctx)
 		if err != nil {

packages/api/internal/sandbox/storage/redis/state_change.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/middleware/otel/joined"
 	redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis"
 )
 
@@ -317,7 +318,11 @@ func (s *Storage) handleExistingTransition(
 	transactionID string,
 ) (sandbox.Sandbox, bool, func(context.Context, error), error) {
 	if sbx.State == newState {
-		// Same target state - wait for completion and return alreadyDone=true
+		// Same target state - wait for completion and return alreadyDone=true.
+		// The caller inherits the in-flight transition's result without
+		// doing the work itself: this is a joiner.
+		joined.Mark(ctx)
+
 		logger.L().Debug(ctx, "State transition already in progress to the same state, waiting",
 			logger.WithSandboxID(sbx.SandboxID),
 			zap.String("state", string(newState)))

packages/shared/pkg/middleware/otel/joined/joined.go

@@ -0,0 +1,58 @@
+package joined
+
+import (
+	"context"
+	"sync/atomic"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+)
+
+// AttributeKey is the dotted-lowercase key used for both the histogram
+// attribute and the span attribute.
+const AttributeKey = "request.joined"
+
+type holder struct {
+	joined atomic.Bool
+
+	serverSpan trace.Span
+}
+
+type holderKey struct{}
+
+// WithHolder installs a fresh holder on ctx if one is not already present.
+//
+// Call WithHolder *after* the server span has been started so the holder
+// captures the correct span for Mark to write attributes onto.
+func WithHolder(ctx context.Context) context.Context {
+	if _, ok := ctx.Value(holderKey{}).(*holder); ok {
+		return ctx
+	}
+
+	return context.WithValue(ctx, holderKey{}, &holder{
+		serverSpan: trace.SpanFromContext(ctx),
+	})
+}
+
+// Mark marks the current request as a joiner. First-write-wins
+func Mark(ctx context.Context) {
+	h, ok := ctx.Value(holderKey{}).(*holder)
+	if !ok {
+		return
+	}
+
+	if h.joined.CompareAndSwap(false, true) {
+		h.serverSpan.SetAttributes(
+			attribute.Bool(AttributeKey, true),
+		)
+	}
+}
+
+func Attribute(ctx context.Context) attribute.KeyValue {
+	h, ok := ctx.Value(holderKey{}).(*holder)
+	if !ok {
+		return attribute.Bool(AttributeKey, false)
+	}
+
+	return attribute.Bool(AttributeKey, h.joined.Load())
+}

packages/shared/pkg/middleware/otel/joined/joined_test.go

@@ -0,0 +1,79 @@
+package joined_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/middleware/otel/joined"
+)
+
+// Mark must be safe even when the context carries no holder.
+func TestMark_NoHolder_Noop(t *testing.T) {
+	t.Parallel()
+	joined.Mark(context.Background())
+}
+
+// Attribute must return request.joined=false when no holder is on ctx.
+func TestAttribute_NoHolder_ReturnsFalse(t *testing.T) {
+	t.Parallel()
+
+	a := joined.Attribute(context.Background())
+	assert.Equal(t, joined.AttributeKey, string(a.Key))
+	assert.False(t, a.Value.AsBool())
+}
+
+// Attribute must return request.joined=false on a freshly installed holder
+// before Mark has been called.
+func TestAttribute_FreshHolder_ReturnsFalse(t *testing.T) {
+	t.Parallel()
+
+	ctx := joined.WithHolder(context.Background())
+
+	a := joined.Attribute(ctx)
+	assert.False(t, a.Value.AsBool())
+}
+
+// Mark must flip Attribute to true on the same ctx.
+func TestMark_FlipsAttributeToTrue(t *testing.T) {
+	t.Parallel()
+
+	ctx := joined.WithHolder(context.Background())
+	joined.Mark(ctx)
+
+	a := joined.Attribute(ctx)
+	assert.True(t, a.Value.AsBool())
+}
+
+// WithHolder must be idempotent: calling it twice returns a ctx that shares
+// the same underlying holder (Mark on the first ctx is visible from the
+// second).
+func TestWithHolder_Idempotent(t *testing.T) {
+	t.Parallel()
+
+	ctx1 := joined.WithHolder(context.Background())
+	ctx2 := joined.WithHolder(ctx1)
+
+	joined.Mark(ctx1)
+
+	a := joined.Attribute(ctx2)
+	assert.True(t, a.Value.AsBool(), "second WithHolder must reuse the first holder")
+}
+
+// Mark must be safe when called from a goroutine descended from the
+// request context.
+func TestMark_DescendantGoroutine(t *testing.T) {
+	t.Parallel()
+
+	ctx := joined.WithHolder(context.Background())
+	done := make(chan struct{})
+	go func() {
+		joined.Mark(ctx)
+		close(done)
+	}()
+	<-done
+
+	a := joined.Attribute(ctx)
+	assert.True(t, a.Value.AsBool())
+}

packages/shared/pkg/middleware/otel/metrics/middleware.go

@@ -13,6 +13,7 @@ import (
 	semconv "go.opentelemetry.io/otel/semconv/v1.7.0"
 
 	sharedmiddleware "github.com/e2b-dev/infra/packages/shared/pkg/middleware"
+	"github.com/e2b-dev/infra/packages/shared/pkg/middleware/otel/joined"
 )
 
 const MetricPrefix = "metric."
@@ -57,6 +58,13 @@ func Middleware(meterProvider metric.MeterProvider, service string, options ...O
 	return func(ginCtx *gin.Context) {
 		ctx := ginCtx.Request.Context()
 
+		// Install the request-scoped joined holder so descendant code paths
+		// (orchestrator, storage layer, etc.) can call joined.Mark via
+		// context.Context. Idempotent: if the tracing middleware already
+		// installed the holder, this reuses it.
+		ctx = joined.WithHolder(ctx)
+		ginCtx.Request = ginCtx.Request.WithContext(ctx)
+
 		route := ginCtx.FullPath()
 		if len(route) == 0 {
 			route = "nonconfigured"
@@ -90,6 +98,9 @@ func Middleware(meterProvider metric.MeterProvider, service string, options ...O
 			// Append attributes from ginCtx
 			resAttributes = append(resAttributes, attributesFromGinContext(ginCtx, MetricPrefix)...)
 
+			// Distinguish between regular and joined requests
+			resAttributes = append(resAttributes, joined.Attribute(ctx))
+
 			// Use processing start time if set, otherwise fall back to the middleware start time.
 			effectiveStart := start
 			if processingStart, ok := getProcessingStartTime(ginCtx); ok {

packages/shared/pkg/middleware/otel/tracing/middleware.go

@@ -31,6 +31,7 @@ import (
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	sharedmiddleware "github.com/e2b-dev/infra/packages/shared/pkg/middleware"
+	"github.com/e2b-dev/infra/packages/shared/pkg/middleware/otel/joined"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
@@ -110,6 +111,9 @@ func Middleware(tracerProvider oteltrace.TracerProvider, service string) gin.Han
 		ctx, span := tracer.Start(ctx, spanName, opts...)
 		defer span.End()
 
+		// Install the request-scoped joined holder
+		ctx = joined.WithHolder(ctx)
+
 		// pass the span through the request context
 		c.Request = c.Request.WithContext(ctx)
 
@@ -130,6 +134,9 @@ func Middleware(tracerProvider oteltrace.TracerProvider, service string) gin.Han
 		spanStatus, spanMessage := semconv.SpanStatusFromHTTPStatusCode(status)
 		span.SetStatus(spanStatus, spanMessage)
 
+		// Marks the joined requests for telemetry purposes.
+		span.SetAttributes(joined.Attribute(ctx))
+
 		if len(c.Errors) > 0 {
 			span.SetAttributes(attribute.String("gin.errors", strings.TrimSpace(c.Errors.String())))
 		}