Merge branch 'main' into return-correct-status-code-for-non-existing-sandbox-in-keep-alive

Author: jakubno

iac/provider-gcp/nomad/jobs/template-manager.hcl

@@ -73,7 +73,6 @@ job "template-manager-system" {
         OTEL_COLLECTOR_GRPC_ENDPOINT  = "${otel_collector_grpc_endpoint}"
         LOGS_COLLECTOR_ADDRESS        = "${logs_collector_address}"
         ORCHESTRATOR_SERVICES         = "${orchestrator_services}"
-        ALLOW_SANDBOX_INTERNET        = "${allow_sandbox_internet}"
         SHARED_CHUNK_CACHE_PATH       = "${shared_chunk_cache_path}"
         CLICKHOUSE_CONNECTION_STRING  = "${clickhouse_connection_string}"
         DOCKERHUB_REMOTE_REPOSITORY_URL  = "${dockerhub_remote_repository_url}"

iac/provider-gcp/nomad/main.tf

@@ -527,7 +527,6 @@ resource "nomad_job" "template_manager" {
     otel_collector_grpc_endpoint    = "localhost:${var.otel_collector_grpc_port}"
     logs_collector_address          = "http://localhost:${var.logs_proxy_port.port}"
     orchestrator_services           = "template-manager"
-    allow_sandbox_internet          = var.allow_sandbox_internet
     clickhouse_connection_string    = local.clickhouse_connection_string
     dockerhub_remote_repository_url = var.dockerhub_remote_repository_url
     launch_darkly_api_key           = trimspace(data.google_secret_manager_secret_version.launch_darkly_api_key.secret_data)

packages/api/internal/orchestrator/keep_alive.go

@@ -23,6 +23,10 @@ func (o *Orchestrator) KeepAliveFor(ctx context.Context, sandboxID string, durat
 	now := time.Now()
 
 	updateFunc := func(sbx sandbox.Sandbox) (sandbox.Sandbox, error) {
+		if sbx.State != sandbox.StateRunning {
+			return sbx, &sandbox.NotFoundError{SandboxID: sandboxID}
+		}
+
 		maxAllowedTTL := getMaxAllowedTTL(now, sbx.StartTime, duration, sbx.MaxInstanceLength)
 		newEndTime := now.Add(maxAllowedTTL)
 

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

@@ -120,6 +120,7 @@ func (s *Store) Update(sandboxID string, updateFunc func(sandbox.Sandbox) (sandb
 
 	item.mu.Lock()
 	defer item.mu.Unlock()
+
 	sbx, err := updateFunc(item._data)
 	if err != nil {
 		return sandbox.Sandbox{}, err

tests/integration/internal/tests/api/sandboxes/sandbox_set_timeout_test.go

@@ -0,0 +1,61 @@
+package sandboxes
+
+import (
+	"fmt"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/e2b-dev/infra/tests/integration/internal/api"
+	"github.com/e2b-dev/infra/tests/integration/internal/setup"
+	"github.com/e2b-dev/infra/tests/integration/internal/utils"
+)
+
+func TestSandboxSetTimeoutPausingSandbox(t *testing.T) {
+	c := setup.GetAPIClient()
+
+	t.Run("test set timeout while pausing", func(t *testing.T) {
+		sbx := utils.SetupSandboxWithCleanup(t, c, utils.WithAutoPause(true))
+		sbxId := sbx.SandboxID
+
+		// Pause the sandbox
+		wg := errgroup.Group{}
+		wg.Go(func() error {
+			pauseResp, err := c.PostSandboxesSandboxIDPauseWithResponse(t.Context(), sbxId, setup.WithAPIKey())
+			if err != nil {
+				return err
+			}
+
+			if pauseResp.StatusCode() != http.StatusNoContent {
+				return fmt.Errorf("unexpected status code: %d", pauseResp.StatusCode())
+			}
+
+			return nil
+		})
+
+		for range 5 {
+			time.Sleep(200 * time.Millisecond)
+			wg.Go(func() error {
+				setTimeoutResp, err := c.PostSandboxesSandboxIDTimeoutWithResponse(t.Context(), sbxId, api.PostSandboxesSandboxIDTimeoutJSONRequestBody{
+					Timeout: 15,
+				},
+					setup.WithAPIKey())
+				if err != nil {
+					return err
+				}
+
+				if setTimeoutResp.StatusCode() != http.StatusNotFound {
+					return fmt.Errorf("unexpected status code: %d", setTimeoutResp.StatusCode())
+				}
+
+				return nil
+			})
+		}
+
+		err := wg.Wait()
+		require.NoError(t, err)
+	})
+}