fix: prevent containers from sticking during shutdown

go.mod

@@ -319,7 +319,7 @@ require (
 	tags.cncf.io/container-device-interface/specs-go v0.8.0
 )
 
-replace github.com/yandex-cloud/geesefs => github.com/beam-cloud/geesefs v0.0.0-20260801195151-4d29b1e9baad
+replace github.com/yandex-cloud/geesefs => github.com/beam-cloud/geesefs v0.0.0-20260803222054-a5eacc477b21
 
 replace github.com/aws/aws-sdk-go => github.com/beam-cloud/geesefs/s3ext v0.0.0-20250606164905-2f3593d03f4f
 

go.sum

@@ -133,8 +133,8 @@ github.com/aws/smithy-go v1.22.2/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxY
 github.com/aybabtme/rgbterm v0.0.0-20170906152045-cc83f3b3ce59/go.mod h1:q/89r3U2H7sSsE2t6Kca0lfwTK8JdoNGS/yzM/4iH5I=
 github.com/beam-cloud/clip v0.0.0-20260802023200-1316f48c56cc h1:NNPU35tvmxDopZDtJoxz9ZwCrU+xYqAYm1OOx5DBYSk=
 github.com/beam-cloud/clip v0.0.0-20260802023200-1316f48c56cc/go.mod h1:uymwDNMk1qoxElmDSlPzwYqZDriHDnW5VVyfK3YNuz8=
-github.com/beam-cloud/geesefs v0.0.0-20260801195151-4d29b1e9baad h1:Q2J1/nypbW7BejKuuAO/YiutTEouoAT6CFxhlbc00HE=
-github.com/beam-cloud/geesefs v0.0.0-20260801195151-4d29b1e9baad/go.mod h1:q0qGsu0wiOzf/5m9JGLqZ7WteYDaaLGgBwNmstb3og0=
+github.com/beam-cloud/geesefs v0.0.0-20260803222054-a5eacc477b21 h1:gMxi5aHC7gG21FNDzVtEegbmU9ct8S95Nl/VYyoqoY4=
+github.com/beam-cloud/geesefs v0.0.0-20260803222054-a5eacc477b21/go.mod h1:q0qGsu0wiOzf/5m9JGLqZ7WteYDaaLGgBwNmstb3og0=
 github.com/beam-cloud/geesefs/s3ext v0.0.0-20250606164905-2f3593d03f4f h1:XzHOu+erxeBO6D3fKVbd5DAlipl+PYZ3u+Ywb8m7Ovk=
 github.com/beam-cloud/geesefs/s3ext v0.0.0-20250606164905-2f3593d03f4f/go.mod h1:YT41ScwaZw9hYfM0WbYZ64sQLNhPxWZFOXJOPug7O5M=
 github.com/beam-cloud/go-runc v0.0.0-20250911154456-bb45084abfe1 h1:EUB/gApGrZW0SzPdyk0jQILJk4Q8wUsWevTHGMkWU58=

pkg/worker/lifecycle.go

@@ -99,7 +99,7 @@ func (s *Worker) handleStopContainerArgs(stopArgs types.StopContainerArgs, sourc
 				types.EventAttrForce: fmt.Sprintf("%t", stopArgs.Force),
 			},
 		})
-		s.cancelBuild(stopArgs.ContainerId)
+		s.cancelContainer(stopArgs.ContainerId)
 		s.stopContainerChan <- stopContainerEvent{ContainerId: stopArgs.ContainerId, Kill: stopArgs.Force}
 	}
 

pkg/worker/storage_manager_test.go

@@ -12,15 +12,23 @@ import (
 	"github.com/beam-cloud/beta9/pkg/types"
 )
 
-type trackedStorage struct{ unmounted bool }
+type trackedStorage struct {
+	unmounted bool
+	mode      string
+}
 
 func (*trackedStorage) Mount(string) error { return nil }
 func (s *trackedStorage) Unmount(string) error {
 	s.unmounted = true
 	return nil
 }
 func (*trackedStorage) Format(string) error { return nil }
-func (*trackedStorage) Mode() string        { return storage.StorageModeLocal }
+func (s *trackedStorage) Mode() string {
+	if s.mode != "" {
+		return s.mode
+	}
+	return storage.StorageModeLocal
+}
 
 func TestWorkspaceStorageMountHotPathAllowsConcurrentHealthChecks(t *testing.T) {
 	mount := &trackedStorage{}

pkg/worker/worker.go

@@ -97,7 +97,7 @@ type Worker struct {
 	containerMountManager   *ContainerMountManager
 	imageClient             *ImageClient
 	containerInstances      *common.SafeMap[*ContainerInstance]
-	buildCancels            *common.SafeMap[context.CancelFunc]
+	containerCancels        *common.SafeMap[context.CancelFunc]
 	containerLock           sync.Mutex
 	checkpointCreateLocks   sync.Map
 	containerStartSem       chan struct{}
@@ -431,7 +431,7 @@ func NewWorker() (_ *Worker, err error) {
 		criuManager:             criuManager,
 		podHostName:             podHostName,
 		containerInstances:      containerInstances,
-		buildCancels:            common.NewSafeMap[context.CancelFunc](),
+		containerCancels:        common.NewSafeMap[context.CancelFunc](),
 		containerLock:           sync.Mutex{},
 		containerStartSem:       make(chan struct{}, containerStartLimit),
 		containerStartLimit:     containerStartLimit,
@@ -688,11 +688,9 @@ func (s *Worker) runContainerRequest(request *types.ContainerRequest) {
 	ctx, cancel := context.WithCancel(s.ctx)
 	defer cancel()
 
-	if request.IsBuildRequest() {
-		s.registerBuildCancel(containerId, cancel)
-		defer s.unregisterBuildCancel(containerId)
-		go s.cancelBuildIfAlreadyStopping(cancel, containerId)
-	}
+	s.registerContainerCancel(containerId, cancel)
+	defer s.unregisterContainerCancel(containerId)
+	go s.cancelContainerIfAlreadyStopping(cancel, containerId)
 
 	if err := s.hydrateRuntimeCredentials(ctx, request); err != nil {
 		log.Error().Str("container_id", containerId).Err(err).Msg("unable to hydrate runtime credentials")
@@ -750,16 +748,39 @@ func (s *Worker) failContainerRequest(containerId string, request *types.Contain
 	s.clearContainer(containerId, request, exitCode, false)
 }
 
-// cancelBuildIfAlreadyStopping checks if a build has already been cancelled and cancels the context if it has.
-func (s *Worker) cancelBuildIfAlreadyStopping(cancel context.CancelFunc, containerId string) {
-	containerState, err := handleGRPCResponse(s.containerRepoClient.GetContainerState(context.Background(), &pb.GetContainerStateRequest{ContainerId: containerId}))
+// cancelContainerIfAlreadyStopping closes the startup race where a stop is
+// persisted before this worker registers the in-memory cancellation callback.
+func (s *Worker) cancelContainerIfAlreadyStopping(cancel context.CancelFunc, containerId string) {
+	// A stop event can land after the instance is reserved but before the
+	// startup cancellation callback is registered. Prefer the local lifecycle
+	// state so that race does not depend on another repository round trip.
+	if s.containerInstances != nil {
+		if instance, exists := s.containerInstances.Get(containerId); exists && instance != nil {
+			_, stopReason := instance.lifecycleState()
+			if stopReason != "" {
+				log.Info().Str("container_id", containerId).Msg("container stopped before startup cancellation was registered")
+				cancel()
+				return
+			}
+		}
+	}
+
+	// The worker event stream may have disconnected before delivering the stop
+	// event. Reconcile against persisted state, but keep the lookup bounded so a
+	// repository outage cannot leak one goroutine per startup or reconnect.
+	if s.containerRepoClient == nil {
+		return
+	}
+	stateCtx, stateCancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer stateCancel()
+	containerState, err := handleGRPCResponse(s.containerRepoClient.GetContainerState(stateCtx, &pb.GetContainerStateRequest{ContainerId: containerId}))
 	if err != nil {
 		log.Error().Str("container_id", containerId).Err(err).Msg("failed to get container state")
 		return
 	}
 
 	if types.ContainerStatus(containerState.State.Status) == types.ContainerStatusStopping {
-		log.Info().Str("container_id", containerId).Msg("incoming container state is stopping, cancelling context")
+		log.Info().Str("container_id", containerId).Msg("incoming container state is stopping, cancelling startup context")
 		cancel()
 		return
 	}
@@ -936,55 +957,67 @@ func (s *Worker) updateContainerStatusOnce(request *types.ContainerRequest) (boo
 			stateCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 			needsStop := runtimeNeedsGraceKill(stateCtx, rt, request.ContainerId)
 			cancel()
-			if !needsStop {
-				return
-			}
-
-			log.Info().Str("container_id", request.ContainerId).Int64("grace_period_seconds", s.config.Worker.TerminationGracePeriod).Msg("container still running after stop event")
-			_, stopReason := instance.lifecycleState()
-			s.recordContainerEvent(context.Background(), request, types.EventContainerEventSchema{
-				ID:          types.ContainerEventWorkerStoppingGraceKill,
-				ContainerID: request.ContainerId,
-				Reason:      string(stopReason),
-				Source:      types.EventSourceWorkerStatusHeartbeat.String(),
-				Message:     types.EventMessageStoppingGraceKill.String(),
-				Attrs: map[string]string{
-					types.EventAttrGracePeriodSeconds: fmt.Sprintf("%d", s.config.Worker.TerminationGracePeriod),
-				},
-			})
-			s.stopContainerChan <- stopContainerEvent{
-				ContainerId: request.ContainerId,
-				Kill:        true,
+			if needsStop {
+				log.Info().Str("container_id", request.ContainerId).Int64("grace_period_seconds", s.config.Worker.TerminationGracePeriod).Msg("container still running after stop event")
+				_, stopReason := instance.lifecycleState()
+				s.recordContainerEvent(context.Background(), request, types.EventContainerEventSchema{
+					ID:          types.ContainerEventWorkerStoppingGraceKill,
+					ContainerID: request.ContainerId,
+					Reason:      string(stopReason),
+					Source:      types.EventSourceWorkerStatusHeartbeat.String(),
+					Message:     types.EventMessageStoppingGraceKill.String(),
+					Attrs: map[string]string{
+						types.EventAttrGracePeriodSeconds: fmt.Sprintf("%d", s.config.Worker.TerminationGracePeriod),
+					},
+				})
+				s.stopContainerChan <- stopContainerEvent{
+					ContainerId: request.ContainerId,
+					Kill:        true,
+				}
 			}
 
 			select {
 			case <-time.After(stuckContainerAbortDelay):
 			case <-s.ctx.Done():
 				return
 			}
-			stateCtx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
-			stillRunning := runtimeNeedsGraceKill(stateCtx, rt, request.ContainerId)
-			cancel()
-			if !stillRunning || s.storageManager == nil || s.storageManager.poolConfig.StorageMode != storage.StorageModeGeese {
-				return
-			}
-			unlock := s.storageManager.lockWorkspaceMount(request.Workspace.Name)
-			if !s.workspaceOnlyStopping(request.Workspace.Name) {
-				unlock()
-				return
-			}
-			log.Warn().Str("container_id", request.ContainerId).Str("workspace", request.Workspace.Name).Msg("aborting stuck workspace mount after SIGKILL timeout")
-			err := s.storageManager.unmountLocked(request.Workspace.Name)
-			unlock()
-			if err != nil {
-				log.Warn().Err(err).Str("workspace", request.Workspace.Name).Msg("stuck workspace mount recovery completed with errors")
-			}
+			s.abortStuckWorkspaceMount(request)
 		}()
 	}
 
 	return false, nil
 }
 
+// abortStuckWorkspaceMount recovers the shutdown path even when the runtime
+// was never created (or has already disappeared). The local lifecycle state is
+// authoritative here: a nonterminal instance with a stop reason is still
+// waiting for its cleanup defers, commonly on a wedged FUSE flush.
+func (s *Worker) abortStuckWorkspaceMount(request *types.ContainerRequest) {
+	if request == nil || s.containerInstances == nil || s.storageManager == nil || s.storageManager.poolConfig.StorageMode != storage.StorageModeGeese {
+		return
+	}
+	instance, exists := s.containerInstances.Get(request.ContainerId)
+	if !exists || instance == nil {
+		return
+	}
+	exitCode, stopReason := instance.lifecycleState()
+	if exitCode >= 0 || stopReason == "" {
+		return
+	}
+
+	workspaceName := request.Workspace.Name
+	unlock := s.storageManager.lockWorkspaceMount(workspaceName)
+	defer unlock()
+	if !s.workspaceOnlyStopping(workspaceName) {
+		return
+	}
+
+	log.Warn().Str("container_id", request.ContainerId).Str("workspace", workspaceName).Msg("aborting stuck workspace mount after SIGKILL timeout")
+	if err := s.storageManager.unmountLocked(workspaceName); err != nil {
+		log.Warn().Err(err).Str("workspace", workspaceName).Msg("stuck workspace mount recovery completed with errors")
+	}
+}
+
 func (s *Worker) workspaceOnlyStopping(workspaceName string) bool {
 	safe := workspaceName != ""
 	s.containerInstances.Range(func(_ string, instance *ContainerInstance) bool {

pkg/worker/worker_events.go

@@ -31,7 +31,7 @@ func (s *Worker) listenForWorkerEvents() {
 		}
 
 		delay = workerEventStreamReconnectMin
-		s.cancelStoppingBuilds()
+		s.cancelStoppingContainers()
 
 		for {
 			event, err := stream.Recv()
@@ -101,7 +101,7 @@ func (s *Worker) handleWorkerEvent(event *pb.WorkerEvent) {
 			return
 		}
 
-		s.cancelBuild(e.StopBuild.ContainerId)
+		s.cancelContainer(e.StopBuild.ContainerId)
 	default:
 		log.Warn().Str("event_id", event.EventId).Msg("received unknown worker event")
 	}
@@ -114,44 +114,44 @@ func (s *Worker) storageNodeID() string {
 	return types.StableStorageNodeID(s.machineID, s.workerId)
 }
 
-func (s *Worker) registerBuildCancel(containerID string, cancel context.CancelFunc) {
-	if s.buildCancels == nil {
-		s.buildCancels = common.NewSafeMap[context.CancelFunc]()
+func (s *Worker) registerContainerCancel(containerID string, cancel context.CancelFunc) {
+	if s.containerCancels == nil {
+		s.containerCancels = common.NewSafeMap[context.CancelFunc]()
 	}
 
-	s.buildCancels.Set(containerID, cancel)
+	s.containerCancels.Set(containerID, cancel)
 }
 
-func (s *Worker) unregisterBuildCancel(containerID string) {
-	if s.buildCancels == nil {
+func (s *Worker) unregisterContainerCancel(containerID string) {
+	if s.containerCancels == nil {
 		return
 	}
 
-	s.buildCancels.Delete(containerID)
+	s.containerCancels.Delete(containerID)
 }
 
-func (s *Worker) cancelBuild(containerID string) bool {
-	if s.buildCancels == nil {
+func (s *Worker) cancelContainer(containerID string) bool {
+	if s.containerCancels == nil {
 		return false
 	}
 
-	cancel, ok := s.buildCancels.Get(containerID)
+	cancel, ok := s.containerCancels.Get(containerID)
 	if !ok {
 		return false
 	}
 
-	log.Info().Str("container_id", containerID).Msg("received stop build event")
+	log.Info().Str("container_id", containerID).Msg("cancelling container startup context")
 	cancel()
 	return true
 }
 
-func (s *Worker) cancelStoppingBuilds() {
-	if s.buildCancels == nil {
+func (s *Worker) cancelStoppingContainers() {
+	if s.containerCancels == nil {
 		return
 	}
-	s.buildCancels.Range(func(containerID string, cancel context.CancelFunc) bool {
+	s.containerCancels.Range(func(containerID string, cancel context.CancelFunc) bool {
 		// One slow state lookup must not delay the rest of a burst.
-		go s.cancelBuildIfAlreadyStopping(cancel, containerID)
+		go s.cancelContainerIfAlreadyStopping(cancel, containerID)
 		return true
 	})
 }

pkg/worker/worker_events_test.go

@@ -15,11 +15,11 @@ func TestHandleWorkerEventStopsOwnedContainer(t *testing.T) {
 	worker := &Worker{
 		workerId:           "worker-1",
 		containerInstances: common.NewSafeMap[*ContainerInstance](),
-		buildCancels:       common.NewSafeMap[context.CancelFunc](),
+		containerCancels:   common.NewSafeMap[context.CancelFunc](),
 		stopContainerChan:  make(chan stopContainerEvent, 1),
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	worker.registerBuildCancel("container-1", cancel)
+	worker.registerContainerCancel("container-1", cancel)
 	worker.containerInstances.Set("container-1", &ContainerInstance{
 		Id: "container-1",
 		Request: &types.ContainerRequest{
@@ -94,10 +94,10 @@ func TestHandleWorkerEventIgnoresHeartbeat(t *testing.T) {
 
 func TestHandleWorkerEventCancelsMatchingBuild(t *testing.T) {
 	worker := &Worker{
-		buildCancels: common.NewSafeMap[context.CancelFunc](),
+		containerCancels: common.NewSafeMap[context.CancelFunc](),
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	worker.registerBuildCancel("build-1", cancel)
+	worker.registerContainerCancel("build-1", cancel)
 
 	worker.handleWorkerEvent(&pb.WorkerEvent{
 		EventId: "event-1",
@@ -113,7 +113,7 @@ func TestHandleWorkerEventCancelsMatchingBuild(t *testing.T) {
 	}
 }
 
-func TestReconnectCancelsStoppingBuilds(t *testing.T) {
+func TestReconnectCancelsStoppingContainers(t *testing.T) {
 	repoClient := &fakeContainerRepoClient{
 		state: &pb.ContainerState{
 			ContainerId: "build-1",
@@ -122,20 +122,35 @@ func TestReconnectCancelsStoppingBuilds(t *testing.T) {
 	}
 	worker := &Worker{
 		containerRepoClient: repoClient,
-		buildCancels:        common.NewSafeMap[context.CancelFunc](),
+		containerCancels:    common.NewSafeMap[context.CancelFunc](),
 	}
 	first, cancelFirst := context.WithCancel(context.Background())
 	second, cancelSecond := context.WithCancel(context.Background())
-	worker.registerBuildCancel("build-1", cancelFirst)
-	worker.registerBuildCancel("build-2", cancelSecond)
+	worker.registerContainerCancel("container-1", cancelFirst)
+	worker.registerContainerCancel("container-2", cancelSecond)
 
-	worker.cancelStoppingBuilds()
+	worker.cancelStoppingContainers()
 
 	for _, done := range []<-chan struct{}{first.Done(), second.Done()} {
 		select {
 		case <-done:
 		case <-time.After(time.Second):
-			t.Fatal("expected build context to be cancelled")
+			t.Fatal("expected container context to be cancelled")
 		}
 	}
 }
+
+func TestStartupRegistrationRaceUsesLocalStopState(t *testing.T) {
+	worker := &Worker{
+		containerInstances: common.NewSafeMap[*ContainerInstance](),
+	}
+	worker.containerInstances.Set("container-1", &ContainerInstance{
+		ExitCode:   -1,
+		StopReason: types.StopContainerReasonUser,
+	})
+	ctx, cancel := context.WithCancel(context.Background())
+
+	worker.cancelContainerIfAlreadyStopping(cancel, "container-1")
+
+	require.ErrorIs(t, ctx.Err(), context.Canceled)
+}