Recover containers from stuck FUSE mounts

pkg/storage/geese.go

@@ -543,6 +543,9 @@ func (s *GeeseStorage) Unmount(localPath string) error {
 		case <-time.After(defaultGeeseFSFlushTimeout):
 			errs = errors.Join(errs, fmt.Errorf("timed out waiting for geesefs files to flush after %s", defaultGeeseFSFlushTimeout))
 			log.Warn().Str("local_path", localPath).Dur("timeout", defaultGeeseFSFlushTimeout).Msg("geesefs: flush wait timed out during unmount")
+			if err := abortFuseConnection(localPath); err != nil {
+				errs = errors.Join(errs, fmt.Errorf("abort stuck geesefs connection: %w", err))
+			}
 		}
 	}
 
@@ -560,6 +563,34 @@ func (s *GeeseStorage) Unmount(localPath string) error {
 	return errs
 }
 
+func abortFuseConnection(localPath string) error {
+	mountInfo, err := os.ReadFile("/proc/self/mountinfo")
+	if err != nil {
+		return err
+	}
+	connectionID := fuseConnectionID(string(mountInfo), localPath)
+	if connectionID == "" {
+		return fmt.Errorf("fuse connection not found for %s", localPath)
+	}
+	const connectionsPath = "/sys/fs/fuse/connections"
+	_ = exec.Command("mount", "-t", "fusectl", "fusectl", connectionsPath).Run()
+	return os.WriteFile(connectionsPath+"/"+connectionID+"/abort", []byte("1"), 0200)
+}
+
+func fuseConnectionID(mountInfo, localPath string) string {
+	target := cleanMountInfoPath(localPath)
+	for _, line := range strings.Split(mountInfo, "\n") {
+		fields := strings.Fields(line)
+		if len(fields) < 7 || !strings.Contains(line, " - fuse.") || cleanMountInfoPath(unescapeMountInfoPath(fields[4])) != target {
+			continue
+		}
+		if _, minor, ok := strings.Cut(fields[2], ":"); ok {
+			return minor
+		}
+	}
+	return ""
+}
+
 func unmountGeeseFS(mfs core.MountedFS, localPath string) error {
 	unmounted := make(chan error, 1)
 	go func() {

pkg/worker/criu.go

@@ -1554,7 +1554,8 @@ func (s *Worker) deferStopForCheckpoint(instance *ContainerInstance, kill bool)
 	if instance == nil || instance.Request == nil {
 		return false
 	}
-	if instance.StopReason != types.StopContainerReasonScheduler && instance.StopReason != types.StopContainerReasonTtl {
+	_, stopReason := instance.lifecycleState()
+	if stopReason != types.StopContainerReasonScheduler && stopReason != types.StopContainerReasonTtl {
 		return false
 	}
 
@@ -1570,7 +1571,7 @@ func (s *Worker) deferStopForCheckpoint(instance *ContainerInstance, kill bool)
 
 	log.Info().
 		Str("container_id", event.ContainerId).
-		Str("reason", string(instance.StopReason)).
+		Str("reason", string(stopReason)).
 		Bool("kill", event.Kill).
 		Msg("deferring automatic stop until checkpoint creation finishes")
 	return true

pkg/worker/lifecycle.go

@@ -87,7 +87,7 @@ func (s *Worker) handleStopContainerArgs(stopArgs types.StopContainerArgs, sourc
 
 	if containerInstance, exists := s.containerInstances.Get(stopArgs.ContainerId); exists {
 		log.Info().Str("container_id", stopArgs.ContainerId).Msg("received stop container event")
-		containerInstance.StopReason = reason
+		containerInstance.setStopReason(reason)
 		s.containerInstances.Set(stopArgs.ContainerId, containerInstance)
 		s.recordContainerEvent(context.Background(), containerInstance.Request, types.EventContainerEventSchema{
 			ID:          types.ContainerEventWorkerStopEventReceived,
@@ -172,7 +172,7 @@ func (s *Worker) clearContainer(containerId string, request *types.ContainerRequ
 		if request != nil && request.Stub.Type.Kind() == types.StubTypeSandbox {
 			instance.signalProcessManagerReadiness(false)
 		}
-		instance.ExitCode = exitCode
+		instance.setExitCode(exitCode)
 		s.containerInstances.Set(containerId, instance)
 	}
 
@@ -1200,7 +1200,7 @@ func (s *Worker) spawn(request *types.ContainerRequest, spec *specs.Spec, output
 	containerInstance.BundlePath = opts.BundlePath
 	containerInstance.Overlay = overlay
 	containerInstance.Spec = spec
-	containerInstance.ExitCode = -1
+	containerInstance.setExitCode(-1)
 	containerInstance.OutputWriter = common.NewOutputWriter(func(s string) {
 		outputLogger.Info(s, "done", false, "success", false)
 	})
@@ -1461,8 +1461,11 @@ func (s *Worker) spawn(request *types.ContainerRequest, spec *specs.Spec, output
 
 	stopReason := types.StopContainerReasonUnknown
 	containerInstance, exists = s.containerInstances.Get(containerId)
-	if exists && containerInstance.StopReason != "" {
-		stopReason = types.StopContainerReason(containerInstance.StopReason)
+	if exists {
+		_, instanceStopReason := containerInstance.lifecycleState()
+		if instanceStopReason != "" {
+			stopReason = instanceStopReason
+		}
 	}
 	rawExitCode := exitCode
 

pkg/worker/sandbox.go

@@ -403,7 +403,11 @@ func (s *Worker) dockerStartupCanceled(ctx context.Context, containerId string,
 	if !exists {
 		return true
 	}
-	return instance != nil && instance.StopReason != ""
+	if instance == nil {
+		return false
+	}
+	_, stopReason := instance.lifecycleState()
+	return stopReason != ""
 }
 
 func dockerStartupCanceled(ctx context.Context, err error) bool {

pkg/worker/worker.go

@@ -11,6 +11,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -49,6 +50,7 @@ const (
 	shutdownCleanupReserve         time.Duration = 5 * time.Second
 	workerShutdownRPCTimeout       time.Duration = 5 * time.Second
 	containerStartupTimeout        time.Duration = 15 * time.Minute
+	stuckContainerAbortDelay       time.Duration = 2 * time.Minute
 	gvisorShmemTHPPath                           = "/sys/kernel/mm/transparent_hugepage/shmem_enabled"
 )
 
@@ -147,6 +149,26 @@ type ContainerInstance struct {
 	ContainerAddressMap        map[int32]string
 	Runtime                    runtime.Runtime
 	OOMWatcher                 runtime.OOMWatcher
+	StopEscalationStarted      atomic.Bool
+	stateMu                    sync.RWMutex
+}
+
+func (i *ContainerInstance) lifecycleState() (int, types.StopContainerReason) {
+	i.stateMu.RLock()
+	defer i.stateMu.RUnlock()
+	return i.ExitCode, i.StopReason
+}
+
+func (i *ContainerInstance) setExitCode(exitCode int) {
+	i.stateMu.Lock()
+	defer i.stateMu.Unlock()
+	i.ExitCode = exitCode
+}
+
+func (i *ContainerInstance) setStopReason(reason types.StopContainerReason) {
+	i.stateMu.Lock()
+	defer i.stateMu.Unlock()
+	i.StopReason = reason
 }
 
 func (i *ContainerInstance) setContainerAddressMap(addressMap map[int32]string) {
@@ -817,10 +839,11 @@ func (s *Worker) updateContainerStatusOnce(request *types.ContainerRequest) (boo
 		return true, nil
 	}
 
-	if instance.ExitCode >= 0 {
+	exitCode, _ := instance.lifecycleState()
+	if exitCode >= 0 {
 		log.Debug().
 			Str("container_id", request.ContainerId).
-			Int("exit_code", instance.ExitCode).
+			Int("exit_code", exitCode).
 			Msg("container exited, stopping status heartbeat")
 		return true, nil
 	}
@@ -832,7 +855,7 @@ func (s *Worker) updateContainerStatusOnce(request *types.ContainerRequest) (boo
 	if err != nil {
 		notFoundErr := &types.ErrContainerStateNotFound{}
 		if notFoundErr.From(err) {
-			instance.StopReason = types.StopContainerReasonUnknown
+			instance.setStopReason(types.StopContainerReasonUnknown)
 			s.containerInstances.Set(request.ContainerId, instance)
 			s.stopContainerChan <- stopContainerEvent{ContainerId: request.ContainerId, Kill: true}
 			go s.recordContainerEvent(context.Background(), request, types.EventContainerEventSchema{
@@ -891,6 +914,14 @@ func (s *Worker) updateContainerStatusOnce(request *types.ContainerRequest) (boo
 	// If container is supposed to be stopped, but isn't gone after TerminationGracePeriod seconds
 	// ensure it is killed after that
 	if status == types.ContainerStatusStopping {
+		_, stopReason := instance.lifecycleState()
+		if stopReason == "" {
+			instance.setStopReason(types.StopContainerReasonUnknown)
+			s.containerInstances.Set(request.ContainerId, instance)
+		}
+		if !instance.StopEscalationStarted.CompareAndSwap(false, true) {
+			return false, nil
+		}
 		go func() {
 			time.Sleep(time.Duration(s.config.Worker.TerminationGracePeriod) * time.Second)
 
@@ -910,10 +941,11 @@ func (s *Worker) updateContainerStatusOnce(request *types.ContainerRequest) (boo
 			}
 
 			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(instance.StopReason),
+				Reason:      string(stopReason),
 				Source:      types.EventSourceWorkerStatusHeartbeat.String(),
 				Message:     types.EventMessageStoppingGraceKill.String(),
 				Attrs: map[string]string{
@@ -924,12 +956,54 @@ func (s *Worker) updateContainerStatusOnce(request *types.ContainerRequest) (boo
 				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")
+			}
 		}()
 	}
 
 	return false, nil
 }
 
+func (s *Worker) workspaceOnlyStopping(workspaceName string) bool {
+	safe := workspaceName != ""
+	s.containerInstances.Range(func(_ string, instance *ContainerInstance) bool {
+		if instance == nil || instance.Request == nil || instance.Request.Workspace.Name != workspaceName {
+			return true
+		}
+		exitCode, stopReason := instance.lifecycleState()
+		if exitCode >= 0 {
+			return true
+		}
+		if stopReason == "" {
+			safe = false
+			return false
+		}
+		return true
+	})
+	return safe
+}
+
 func runtimeNeedsGraceKill(ctx context.Context, rt runtime.Runtime, containerID string) bool {
 	if rt == nil {
 		return false
@@ -1211,7 +1285,7 @@ func (s *Worker) stopActiveContainersForShutdown() {
 
 	for _, id := range ids {
 		if instance, exists := s.containerInstances.Get(id); exists {
-			instance.StopReason = types.StopContainerReasonAdmin
+			instance.setStopReason(types.StopContainerReasonAdmin)
 			s.containerInstances.Set(id, instance)
 		}
 		if err := s.stopContainer(id, false); err != nil {

pkg/worker/worker_test.go

@@ -295,6 +295,29 @@ func TestUpdateContainerStatusOnceReconcilesStartedPendingContainer(t *testing.T
 	require.Equal(t, int64(types.ContainerStateTtlS), repoClient.lastUpdateStatus.ExpirySeconds)
 }
 
+func TestWorkspaceOnlyStoppingProtectsRunningSiblings(t *testing.T) {
+	worker := &Worker{containerInstances: common.NewSafeMap[*ContainerInstance]()}
+	request := &types.ContainerRequest{Workspace: types.Workspace{Name: "shared"}}
+	worker.containerInstances.Set("stopping", &ContainerInstance{ExitCode: -1, StopReason: types.StopContainerReasonTtl, Request: request})
+	worker.containerInstances.Set("running", &ContainerInstance{ExitCode: -1, Request: request})
+	require.False(t, worker.workspaceOnlyStopping("shared"))
+
+	instance, _ := worker.containerInstances.Get("running")
+	done := make(chan struct{})
+	go func() {
+		for range 1000 {
+			instance.setStopReason(types.StopContainerReasonUser)
+			instance.setExitCode(-1)
+		}
+		close(done)
+	}()
+	for range 1000 {
+		_ = worker.workspaceOnlyStopping("shared")
+	}
+	<-done
+	require.True(t, worker.workspaceOnlyStopping("shared"))
+}
+
 func TestShutdownWaitDrainsWithoutStoppingActiveContainer(t *testing.T) {
 	worker := &Worker{
 		containerInstances: common.NewSafeMap[*ContainerInstance](),