Updatable timeout for envd init

Author: jakubno

packages/orchestrator/internal/sandbox/envd.go

@@ -9,23 +9,31 @@ import (
 	"net/http"
 	"time"
 
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+	"go.opentelemetry.io/otel/trace"
+	"go.uber.org/zap"
+
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
 
 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) {
+	count := int64(0)
 	for {
-		reqCtx, cancel := context.WithTimeout(ctx, requestTimeout)
+		count++
+		start := time.Now()
+		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, count, err
 		}
 
 		// make sure request to already authorized envd will not fail
@@ -38,12 +46,13 @@ func doRequestWithInfiniteRetries(ctx context.Context, method, address string, r
 		cancel()
 
 		if err == nil {
-			return response, nil
+			return response, count, nil
 		}
 
+		zap.L().Error("Error requesting infinite retries", zap.Error(err), logger.WithSandboxID(sandboxID), zap.Int64("elapsed", time.Since(start).Milliseconds()))
 		select {
 		case <-ctx.Done():
-			return nil, fmt.Errorf("%w with cause: %w", ctx.Err(), context.Cause(ctx))
+			return nil, count, fmt.Errorf("%w with cause: %w", ctx.Err(), context.Cause(ctx))
 		case <-time.After(loopDelay):
 		}
 	}
@@ -56,10 +65,12 @@ 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")
+func (s *Sandbox) initEnvd(ctx context.Context, envVars map[string]string, accessToken *string, envdInitRequestTimeout time.Duration) error {
+	childCtx, childSpan := tracer.Start(ctx, "envd-init", trace.WithAttributes(attribute.String("evd_version", s.Config.Envd.Version)))
 	defer childSpan.End()
 
+	attributes := metric.WithAttributes(attribute.String("envd_version", 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)
 	now := time.Now()
@@ -75,7 +86,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(childCtx, "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 +102,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/trace"
 	"go.uber.org/zap"
@@ -31,7 +32,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,
@@ -314,6 +320,7 @@ func ResumeSandbox(
 	devicePool *nbd.DevicePool,
 	useClickhouseMetrics bool,
 	apiConfigToStore *orchestrator.SandboxConfig,
+	envdInitRequestTimeout time.Duration,
 ) (s *Sandbox, e error) {
 	ctx, span, done := startSpan(ctx, "resume sandbox")
 	defer func() { done(e) }()
@@ -522,6 +529,7 @@ func ResumeSandbox(
 	err = sbx.WaitForEnvd(
 		ctx,
 		defaultEnvdTimeout,
+		envdInitRequestTimeout,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to wait for sandbox start: %w", err)
@@ -949,6 +957,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()
@@ -977,7 +986,7 @@ func (s *Sandbox) WaitForEnvd(
 		}
 	}()
 
-	initErr := s.initEnvd(syncCtx, s.Config.Envd.Vars, s.Config.Envd.AccessToken)
+	initErr := s.initEnvd(syncCtx, s.Config.Envd.Vars, s.Config.Envd.AccessToken, envdInitRequestTimeout)
 	if initErr != nil {
 		return fmt.Errorf("failed to init new envd: %w", initErr)
 	}

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/telemetry/meters.go

@@ -24,6 +24,9 @@ const (
 	SandboxCreateMeterName          CounterType = "api.env.instance.started"
 
 	TeamSandboxCreated CounterType = "e2b.team.sandbox.created"
+
+	EnvdInitAttempts CounterType = "orchestrator.envd.init.attempts"
+	EnvdInitSuccess  CounterType = "orchestrator.envd.init.success"
 )
 
 const (
@@ -92,6 +95,8 @@ var counterDesc = map[CounterType]string{
 	BuildResultCounterName:          "Number of template build results",
 	BuildCacheResultCounterName:     "Number of build cache results",
 	TeamSandboxCreated:              "Counter of started sandboxes for the team in the interval",
+	EnvdInitAttempts:                "Number of envd initialization attempts",
+	EnvdInitSuccess:                 "Number of successful envd initializations",
 }
 
 var counterUnits = map[CounterType]string{
@@ -100,6 +105,8 @@ var counterUnits = map[CounterType]string{
 	BuildResultCounterName:          "{build}",
 	BuildCacheResultCounterName:     "{layer}",
 	TeamSandboxCreated:              "{sandbox}",
+	EnvdInitAttempts:                "1",
+	EnvdInitSuccess:                 "1",
 }
 
 var observableCounterDesc = map[ObservableCounterType]string{