diff --git a/packages/api/internal/orchestrator/create_instance.go b/packages/api/internal/orchestrator/create_instance.go index 46855aa912..b1b9fb0445 100644 --- a/packages/api/internal/orchestrator/create_instance.go +++ b/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) diff --git a/packages/api/internal/sandbox/storage/memory/operations.go b/packages/api/internal/sandbox/storage/memory/operations.go index 74550992bf..3a4eb2cadd 100644 --- a/packages/api/internal/sandbox/storage/memory/operations.go +++ b/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 { diff --git a/packages/api/internal/sandbox/storage/redis/state_change.go b/packages/api/internal/sandbox/storage/redis/state_change.go index 483642d6f8..d7ac1c7dcc 100644 --- a/packages/api/internal/sandbox/storage/redis/state_change.go +++ b/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))) diff --git a/packages/shared/pkg/middleware/otel/joined/joined.go b/packages/shared/pkg/middleware/otel/joined/joined.go new file mode 100644 index 0000000000..dd7667ab33 --- /dev/null +++ b/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()) +} diff --git a/packages/shared/pkg/middleware/otel/joined/joined_test.go b/packages/shared/pkg/middleware/otel/joined/joined_test.go new file mode 100644 index 0000000000..b695f87b5e --- /dev/null +++ b/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()) +} diff --git a/packages/shared/pkg/middleware/otel/metrics/middleware.go b/packages/shared/pkg/middleware/otel/metrics/middleware.go index 8d02e86363..f4e7da0f23 100644 --- a/packages/shared/pkg/middleware/otel/metrics/middleware.go +++ b/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 { diff --git a/packages/shared/pkg/middleware/otel/tracing/middleware.go b/packages/shared/pkg/middleware/otel/tracing/middleware.go index 0bbe201e4a..da1b6f932a 100644 --- a/packages/shared/pkg/middleware/otel/tracing/middleware.go +++ b/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()))) }