Updatable timeout for envd init

Author: jakubno

packages/orchestrator/internal/metrics/sandboxes.go

@@ -197,7 +197,7 @@ func (so *SandboxObserver) startObserving() (metric.Registration, error) {
 								logger.WithSandboxID(sbx.Runtime.SandboxID),
 								logger.WithTeamID(sbx.Runtime.TeamID),
 								logger.WithTemplateID(sbx.Runtime.TemplateID),
-								zap.String("envd_version", sbx.Config.Envd.Version),
+								logger.WithEnvdVersion(sbx.Config.Envd.Version),
 								zap.Time("sandbox_start", sbx.StartedAt),
 								zap.Int64("clock_host", hostTm),
 								zap.Int64("clock_sbx", sbxTm),

packages/orchestrator/internal/sandbox/envd.go

@@ -9,23 +9,29 @@ import (
 	"net/http"
 	"time"
 
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
 
 const (
-	requestTimeout = 50 * time.Millisecond
-	loopDelay      = 5 * time.Millisecond
+	loopDelay = 5 * time.Millisecond
 )
 
 // doRequestWithInfiniteRetries does a request with infinite retries until the context is done.
 // The parent context should have a deadline or a timeout.
-func doRequestWithInfiniteRetries(ctx context.Context, method, address string, requestBody []byte, accessToken *string) (*http.Response, error) {
+func doRequestWithInfiniteRetries(ctx context.Context, method, address string, requestBody []byte, accessToken *string, envdInitRequestTimeout time.Duration, sandboxID string) (*http.Response, int64, error) {
+	requestCount := int64(0)
 	for {
-		reqCtx, cancel := context.WithTimeout(ctx, requestTimeout)
+		requestCount++
+		reqCtx, cancel := context.WithTimeout(ctx, envdInitRequestTimeout)
 		request, err := http.NewRequestWithContext(reqCtx, method, address, bytes.NewReader(requestBody))
 		if err != nil {
 			cancel()
-			return nil, err
+			return nil, requestCount, err
 		}
 
 		// make sure request to already authorized envd will not fail
@@ -38,12 +44,12 @@ func doRequestWithInfiniteRetries(ctx context.Context, method, address string, r
 		cancel()
 
 		if err == nil {
-			return response, nil
+			return response, requestCount, nil
 		}
 
 		select {
 		case <-ctx.Done():
-			return nil, fmt.Errorf("%w with cause: %w", ctx.Err(), context.Cause(ctx))
+			return nil, requestCount, fmt.Errorf("%w with cause: %w", ctx.Err(), context.Cause(ctx))
 		case <-time.After(loopDelay):
 		}
 	}
@@ -56,9 +62,11 @@ type PostInitJSONBody struct {
 	Timestamp   *time.Time         `json:"timestamp,omitempty"`
 }
 
-func (s *Sandbox) initEnvd(ctx context.Context, envVars map[string]string, accessToken *string) error {
-	childCtx, childSpan := tracer.Start(ctx, "envd-init")
-	defer childSpan.End()
+func (s *Sandbox) initEnvd(ctx context.Context, envVars map[string]string, accessToken *string, envdInitRequestTimeout time.Duration) error {
+	ctx, span := tracer.Start(ctx, "envd-init", trace.WithAttributes(telemetry.WithEnvdVersion(s.Config.Envd.Version)))
+	defer span.End()
+
+	attributes := metric.WithAttributes(telemetry.WithEnvdVersion(s.Config.Envd.Version), attribute.Int64("timeout_ms", envdInitRequestTimeout.Milliseconds()))
 
 	hyperloopIP := s.Slot.HyperloopIPString()
 	address := fmt.Sprintf("http://%s:%d/init", s.Slot.HostIPString(), consts.DefaultEnvdServerPort)
@@ -75,7 +83,8 @@ func (s *Sandbox) initEnvd(ctx context.Context, envVars map[string]string, acces
 		return err
 	}
 
-	response, err := doRequestWithInfiniteRetries(childCtx, "POST", address, body, accessToken)
+	response, count, err := doRequestWithInfiniteRetries(ctx, "POST", address, body, accessToken, envdInitRequestTimeout, s.Runtime.SandboxID)
+	envdInitAttempts.Add(ctx, count, attributes)
 	if err != nil {
 		return fmt.Errorf("failed to init envd: %w", err)
 	}
@@ -90,5 +99,8 @@ func (s *Sandbox) initEnvd(ctx context.Context, envVars map[string]string, acces
 		return err
 	}
 
+	// Track successful envd init
+	envdInitSuccess.Add(ctx, 1, attributes)
+
 	return nil
 }

packages/orchestrator/internal/sandbox/sandbox.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
@@ -32,7 +33,12 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
-var defaultEnvdTimeout = utils.Must(time.ParseDuration(env.GetEnv("ENVD_TIMEOUT", "10s")))
+var (
+	defaultEnvdTimeout = utils.Must(time.ParseDuration(env.GetEnv("ENVD_TIMEOUT", "10s")))
+	meter              = otel.GetMeterProvider().Meter("github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox")
+	envdInitAttempts   = utils.Must(telemetry.GetCounter(meter, telemetry.EnvdInitAttempts))
+	envdInitSuccess    = utils.Must(telemetry.GetCounter(meter, telemetry.EnvdInitSuccess))
+)
 
 var httpClient = http.Client{
 	Timeout: 10 * time.Second,
@@ -317,6 +323,7 @@ func ResumeSandbox(
 	devicePool *nbd.DevicePool,
 	useClickhouseMetrics bool,
 	apiConfigToStore *orchestrator.SandboxConfig,
+	envdInitRequestTimeout time.Duration,
 ) (s *Sandbox, e error) {
 	ctx, span := tracer.Start(ctx, "resume sandbox")
 	defer func() { endSpan(span, e) }()
@@ -520,6 +527,7 @@ func ResumeSandbox(
 	err = sbx.WaitForEnvd(
 		ctx,
 		defaultEnvdTimeout,
+		envdInitRequestTimeout,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to wait for sandbox start: %w", err)
@@ -931,6 +939,7 @@ func (s *Sandbox) WaitForExit(ctx context.Context) error {
 func (s *Sandbox) WaitForEnvd(
 	ctx context.Context,
 	timeout time.Duration,
+	envdInitRequestTimeout time.Duration,
 ) (e error) {
 	ctx, span := tracer.Start(ctx, "sandbox-wait-for-start")
 	defer span.End()
@@ -959,7 +968,7 @@ func (s *Sandbox) WaitForEnvd(
 		}
 	}()
 
-	if err := s.initEnvd(ctx, s.Config.Envd.Vars, s.Config.Envd.AccessToken); err != nil {
+	if err := s.initEnvd(ctx, s.Config.Envd.Vars, s.Config.Envd.AccessToken, envdInitRequestTimeout); err != nil {
 		return fmt.Errorf("failed to init new envd: %w", err)
 	}
 

packages/orchestrator/internal/server/sandboxes.go

@@ -104,6 +104,13 @@ func (s *server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 		return nil, fmt.Errorf("failed to get template snapshot data: %w", err)
 	}
 
+	// Get timeout from feature flag
+	envdInitRequestTimeoutMs, err := s.featureFlags.IntFlag(ctx, featureflags.EnvdInitTimeoutSeconds)
+	if err != nil {
+		zap.L().Warn("failed to get envd timeout from feature flag, using default", zap.Error(err))
+	}
+	envdInitRequestTimeout := time.Duration(envdInitRequestTimeoutMs) * time.Millisecond
+
 	sbx, err := sandbox.ResumeSandbox(
 		ctx,
 		s.networkPool,
@@ -136,6 +143,7 @@ func (s *server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 		s.devicePool,
 		metricsWriteFlag,
 		req.Sandbox,
+		envdInitRequestTimeout,
 	)
 	if err != nil {
 		err := errors.Join(err, context.Cause(ctx))

packages/orchestrator/internal/template/build/layer/create_sandbox.go

@@ -100,6 +100,7 @@ func (cs *CreateSandbox) Sandbox(
 	err = sbx.WaitForEnvd(
 		ctx,
 		waitEnvdTimeout,
+		defaultEnvdInitRequestTimeout,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("wait for envd: %w", err)

packages/orchestrator/internal/template/build/layer/interfaces.go

@@ -10,7 +10,8 @@ import (
 )
 
 const (
-	waitEnvdTimeout = 60 * time.Second
+	defaultEnvdInitRequestTimeout = 50 * time.Millisecond
+	waitEnvdTimeout               = 60 * time.Second
 )
 
 // SandboxCreator creates sandboxes for layer building

packages/orchestrator/internal/template/build/layer/layer_executor.go

@@ -193,6 +193,7 @@ func (lb *LayerExecutor) updateEnvdInSandbox(
 	err = sbx.WaitForEnvd(
 		ctx,
 		waitEnvdTimeout,
+		defaultEnvdInitRequestTimeout,
 	)
 	if err != nil {
 		return fmt.Errorf("failed to wait for envd initialization after update: %w", err)

packages/orchestrator/internal/template/build/layer/resume_sandbox.go

@@ -46,6 +46,7 @@ func (rs *ResumeSandbox) Sandbox(
 		layerExecutor.devicePool,
 		false,
 		nil,
+		defaultEnvdInitRequestTimeout,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("resume sandbox: %w", err)

packages/shared/pkg/feature-flags/flags.go

@@ -84,4 +84,5 @@ var (
 	BestOfKMaxOvercommit          = newIntFlag("best-of-k-max-overcommit", 400)              // Default R=4 (stored as percentage, max over-commit ratio)
 	BestOfKAlpha                  = newIntFlag("best-of-k-alpha", 50)                        // Default Alpha=0.5 (stored as percentage for int flag, current usage weight)
 	PubsubQueueChannelSize        = newIntFlag("pubsub-queue-channel-size", 8*1024)          // size of the channel buffer used to queue incoming sandbox events
+	EnvdInitTimeoutSeconds        = newIntFlag("envd-init-request-timeout-milliseconds", 50)         // Timeout for envd init request in milliseconds
 )

packages/shared/pkg/logger/fields.go

@@ -32,3 +32,7 @@ func WithClusterID(clusterID uuid.UUID) zap.Field {
 func WithServiceInstanceID(instanceID string) zap.Field {
 	return zap.String("service.instance.id", instanceID)
 }
+
+func WithEnvdVersion(envdVersion string) zap.Field {
+	return zap.String("envd.version", envdVersion)
+}