diff --git a/go.mod b/go.mod index cea105e99..8b502bf03 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 591255ef9..0122fe526 100644 --- a/go.sum +++ b/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= diff --git a/pkg/worker/lifecycle.go b/pkg/worker/lifecycle.go index 9634f1c3c..c301844d7 100644 --- a/pkg/worker/lifecycle.go +++ b/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} } diff --git a/pkg/worker/storage_manager_test.go b/pkg/worker/storage_manager_test.go index c661a47ce..2e112f4f4 100644 --- a/pkg/worker/storage_manager_test.go +++ b/pkg/worker/storage_manager_test.go @@ -12,7 +12,10 @@ 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 { @@ -20,7 +23,12 @@ func (s *trackedStorage) Unmount(string) error { 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{} diff --git a/pkg/worker/worker.go b/pkg/worker/worker.go index 251a7df7b..23fdbf854 100644 --- a/pkg/worker/worker.go +++ b/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,25 +957,23 @@ 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 { @@ -962,29 +981,43 @@ func (s *Worker) updateContainerStatusOnce(request *types.ContainerRequest) (boo 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 { diff --git a/pkg/worker/worker_events.go b/pkg/worker/worker_events.go index 330e239ad..199394472 100644 --- a/pkg/worker/worker_events.go +++ b/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 }) } diff --git a/pkg/worker/worker_events_test.go b/pkg/worker/worker_events_test.go index 5188ac23a..df43f9cb2 100644 --- a/pkg/worker/worker_events_test.go +++ b/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) +} diff --git a/pkg/worker/worker_test.go b/pkg/worker/worker_test.go index 8df13723b..b167c17e2 100644 --- a/pkg/worker/worker_test.go +++ b/pkg/worker/worker_test.go @@ -12,6 +12,7 @@ import ( "github.com/beam-cloud/beta9/pkg/common" "github.com/beam-cloud/beta9/pkg/runtime" + "github.com/beam-cloud/beta9/pkg/storage" "github.com/beam-cloud/beta9/pkg/types" pb "github.com/beam-cloud/beta9/proto" "github.com/stretchr/testify/require" @@ -318,6 +319,40 @@ func TestWorkspaceOnlyStoppingProtectsRunningSiblings(t *testing.T) { require.True(t, worker.workspaceOnlyStopping("shared")) } +func TestAbortStuckWorkspaceMountWithoutRuntimeState(t *testing.T) { + workspaceName := "shared" + request := &types.ContainerRequest{ + ContainerId: "container-1", + Workspace: types.Workspace{Name: workspaceName}, + } + instances := common.NewSafeMap[*ContainerInstance]() + instances.Set(request.ContainerId, &ContainerInstance{ + ExitCode: -1, + StopReason: types.StopContainerReasonUser, + Request: request, + }) + mount := &trackedStorage{mode: storage.StorageModeGeese} + manager := &WorkspaceStorageManager{ + mounts: common.NewSafeMap[storage.Storage](), + mountLastUsed: common.NewSafeMap[time.Time](), + containerInstances: instances, + mountLocks: make(map[string]*sync.RWMutex), + poolConfig: types.WorkerPoolConfig{StorageMode: storage.StorageModeGeese}, + config: types.StorageConfig{WorkspaceStorage: types.WorkspaceStorageConfig{ + BaseMountPath: t.TempDir(), + }}, + } + manager.mounts.Set(workspaceName, mount) + manager.mountLastUsed.Set(workspaceName, time.Now()) + worker := &Worker{containerInstances: instances, storageManager: manager} + + worker.abortStuckWorkspaceMount(request) + + require.True(t, mount.unmounted) + _, mounted := manager.mounts.Get(workspaceName) + require.False(t, mounted) +} + func TestShutdownWaitDrainsWithoutStoppingActiveContainer(t *testing.T) { worker := &Worker{ containerInstances: common.NewSafeMap[*ContainerInstance](), @@ -457,6 +492,7 @@ func TestFailContainerRequestReportsExitCode(t *testing.T) { } type fakeContainerRepoClient struct { + mu sync.Mutex state *pb.ContainerState getStateCalls int updateStatusCalls int @@ -471,6 +507,8 @@ type fakeContainerRepoClient struct { } func (f *fakeContainerRepoClient) GetContainerState(ctx context.Context, in *pb.GetContainerStateRequest, opts ...grpc.CallOption) (*pb.GetContainerStateResponse, error) { + f.mu.Lock() + defer f.mu.Unlock() f.getStateCalls++ return &pb.GetContainerStateResponse{ Ok: true,