chore: extend the default otelgrpc metric with resume (#2293)

* chore: extend the default otelgrpc metric with resume

* fix: lint

* fix: make is resume metric optional

Author: jakubno

packages/api/internal/orchestrator/nodemanager/metadata.go

@@ -2,11 +2,13 @@ package nodemanager
 
 import (
 	"context"
+	"strconv"
 
 	"google.golang.org/grpc/metadata"
 
 	"github.com/e2b-dev/infra/packages/api/internal/clusters"
 	"github.com/e2b-dev/infra/packages/shared/pkg/edge"
+	grpcshared "github.com/e2b-dev/infra/packages/shared/pkg/grpc"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 )
 
@@ -49,6 +51,10 @@ func (n *Node) GetSandboxCreateCtx(ctx context.Context, req *orchestrator.Sandbo
 		)
 	}
 
+	// Pass snapshot (is_resume) via metadata so the server-side stats handler
+	// can include it in otelgrpc metric attributes during TagRPC.
+	md.Set(grpcshared.IsResumeMetadataKey, strconv.FormatBool(req.GetSandbox().GetSnapshot()))
+
 	// Merge medata from client (auth, routing with service instance id) and event metadata.
 	return n.client, appendMetadataCtx(ctx, md)
 }

packages/orchestrator/pkg/factories/run.go

@@ -602,7 +602,7 @@ func run(config cfg.Config, opts Options) (success bool) {
 	})
 	closers = append(closers, closer{"hyperloop server", hyperloopSrv.Shutdown})
 
-	grpcServer := e2bgrpc.NewGRPCServer(tel)
+	grpcServer := e2bgrpc.NewGRPCServer(tel, e2bgrpc.WithSandboxResumeMetrics())
 	orchestrator.RegisterSandboxServiceServer(grpcServer, orchestratorService)
 	orchestrator.RegisterVolumeServiceServer(grpcServer, volumeService)
 	orchestrator.RegisterChunkServiceServer(grpcServer, orchestratorService)

packages/shared/pkg/grpc/filter.go

@@ -6,6 +6,12 @@ import (
 	"google.golang.org/grpc/stats"
 )
 
+// IsResumeMetadataKey is the gRPC metadata key used to pass the is_resume/snapshot
+// value from the API client to the orchestrator server. This allows the server-side
+// otelgrpc stats handler to include sandbox.resume in metric attributes during TagRPC,
+// before the request payload is available.
+const IsResumeMetadataKey = "x-sandbox-resume"
+
 type noTraceKey struct{}
 
 var noTrace = struct{}{}

packages/shared/pkg/grpc/server.go

@@ -1,21 +1,42 @@
 package grpc
 
 import (
+	"context"
 	"time"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/recovery"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/selector"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
+	"go.opentelemetry.io/otel/attribute"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
-func NewGRPCServer(tel *telemetry.Client) *grpc.Server {
-	opts := []logging.Option{
+// ServerOption configures NewGRPCServer.
+type ServerOption func(*serverOptions)
+
+type serverOptions struct {
+	withSandboxResumeMetrics bool
+}
+
+// WithSandboxResumeMetrics adds sandbox.resume attribute to otelgrpc metrics,
+// read from incoming gRPC metadata.
+func WithSandboxResumeMetrics() ServerOption {
+	return func(o *serverOptions) { o.withSandboxResumeMetrics = true }
+}
+
+func NewGRPCServer(tel *telemetry.Client, opts ...ServerOption) *grpc.Server {
+	var cfg serverOptions
+	for _, o := range opts {
+		o(&cfg)
+	}
+
+	logOpts := []logging.Option{
 		logging.WithLogOnEvents(logging.StartCall, logging.PayloadReceived, logging.PayloadSent, logging.FinishCall),
 		logging.WithLevels(logging.DefaultServerCodeToLevel),
 		logging.WithFieldsFromContext(logging.ExtractFields),
@@ -28,33 +49,57 @@ func NewGRPCServer(tel *telemetry.Client) *grpc.Server {
 		"/InfoService/ServiceInfo",
 	)
 
+	otelOpts := []otelgrpc.Option{
+		otelgrpc.WithTracerProvider(tel.TracerProvider),
+		otelgrpc.WithMeterProvider(tel.MeterProvider),
+	}
+	if cfg.withSandboxResumeMetrics {
+		otelOpts = append(otelOpts, otelgrpc.WithMetricAttributesFn(extractSandboxResumeAttrs))
+	}
+
 	return grpc.NewServer(
 		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
-			MinTime:             5 * time.Second, // Minimum time between pings from client
-			PermitWithoutStream: true,            // Allow pings even when no active streams
+			MinTime:             5 * time.Second,
+			PermitWithoutStream: true,
 		}),
 		grpc.KeepaliveParams(keepalive.ServerParameters{
-			Time:    15 * time.Second, // Server sends keepalive pings every 15s
-			Timeout: 5 * time.Second,  // Wait 5s for response before considering dead
+			Time:    15 * time.Second,
+			Timeout: 5 * time.Second,
 		}),
 		grpc.StatsHandler(
 			NewStatsWrapper(
-				otelgrpc.NewServerHandler(
-					otelgrpc.WithTracerProvider(tel.TracerProvider),
-					otelgrpc.WithMeterProvider(tel.MeterProvider),
-				))),
+				otelgrpc.NewServerHandler(otelOpts...))),
 		grpc.ChainUnaryInterceptor(
 			recovery.UnaryServerInterceptor(),
 			selector.UnaryServerInterceptor(
-				logging.UnaryServerInterceptor(logger.GRPCLogger(logger.L()), opts...),
+				logging.UnaryServerInterceptor(logger.GRPCLogger(logger.L()), logOpts...),
 				ignoredLoggingRoutes,
 			),
 		),
 		grpc.ChainStreamInterceptor(
 			selector.StreamServerInterceptor(
-				logging.StreamServerInterceptor(logger.GRPCLogger(logger.L()), opts...),
+				logging.StreamServerInterceptor(logger.GRPCLogger(logger.L()), logOpts...),
 				ignoredLoggingRoutes,
 			),
 		),
 	)
 }
+
+// extractSandboxResumeAttrs reads sandbox.resume from gRPC metadata set by the
+// API client. Called by otelgrpc during TagRPC — before the request payload is
+// deserialized — so we use metadata instead of the payload.
+func extractSandboxResumeAttrs(ctx context.Context) []attribute.KeyValue {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return nil
+	}
+
+	values := md.Get(IsResumeMetadataKey)
+	if len(values) == 0 {
+		return nil
+	}
+
+	return []attribute.KeyValue{
+		attribute.Bool("sandbox.resume", values[0] == "true"),
+	}
+}