From 3d34de96bfebb393a74257fa8344c959cc975f34 Mon Sep 17 00:00:00 2001 From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:14:55 +0800 Subject: [PATCH] refactor(shell): sync.RWMutex for sessions + sync.Map for exited (mirror process pkg) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same shape as the process package refactor: - LocalManager.mu (sync.Mutex) -> sessionsMu (sync.RWMutex). Read paths (Status, Subscribe, Stdin, ExitStatus, Resize) take RLock; write paths (Start, Stop, Restart, List) take Lock. - LocalManager.exited: was a hand-rolled map[workspaceID]map[shellID]struct{} guarded by exitedMu; now a sync.Map keyed by shellID only (UUID is globally unique, no need for the nested map). Helpers IsExited / MarkAsExited / ClearExited. - shellOrder stays a plain map; read+written under sessionsMu. - waitExit remains the sole caller of MarkAsExited; Start / Stop / Restart call ClearExited. - New TestShellIsExitedHelpers covers the helper semantics. go test -race -count=2 ./... clean. Same caveat as the process package: at this app's concurrency level, neither sync.RWMutex nor sync.Map measurably beats the previous pair — the change is mostly stylistic (one fewer lock, no nested maps, more idiomatic Go). Conversation: 019f3673-d2d5-78f0-a7a9-5e3e91b65933 --- internal/shell/manager.go | 128 +++++++++++++++------------------ internal/shell/manager_test.go | 47 ++++++++++++ 2 files changed, 104 insertions(+), 71 deletions(-) diff --git a/internal/shell/manager.go b/internal/shell/manager.go index 4aa9f9f..3edab51 100644 --- a/internal/shell/manager.go +++ b/internal/shell/manager.go @@ -47,14 +47,21 @@ type Manager interface { } // LocalManager implements Manager using local OS processes. +// +// Concurrency model: +// - sessionsMu is an RWMutex. Read-only paths (Status, Subscribe lookup, +// Stdin lookup, ExitStatus lookup, Resize) take RLock; mutating paths +// (Start, Stop, Restart, List) take Lock. +// - exited is a sync.Map keyed by shellID (a UUID, globally unique). +// waitExit is the only goroutine that calls sess.Cmd.Wait() and the only +// one that writes here. Reads from anywhere are lock-free. type LocalManager struct { command string args []string - mu sync.Mutex + sessionsMu sync.RWMutex sessions map[string]map[string]*Session // workspaceID -> shellID -> Session - exitedMu sync.Mutex - exited map[string]map[string]struct{} // workspaceID -> shellID -> exited shellOrder map[string][]string // workspaceID -> shellIDs in creation order + exited sync.Map // shellID -> struct{} } // NewManager creates a LocalManager with the given command and arguments. @@ -64,7 +71,6 @@ func NewManager(command string, args []string) *LocalManager { command: normalizeCommand(command), args: args, sessions: make(map[string]map[string]*Session), - exited: make(map[string]map[string]struct{}), shellOrder: make(map[string][]string), } } @@ -76,10 +82,26 @@ func normalizeCommand(command string) string { return command } +// IsExited reports whether the shell has exited. Lock-free; safe from any goroutine. +func (m *LocalManager) IsExited(shellID string) bool { + _, ok := m.exited.Load(shellID) + return ok +} + +// MarkAsExited records the shell as exited. Called only from waitExit. +func (m *LocalManager) MarkAsExited(shellID string) { + m.exited.Store(shellID, struct{}{}) +} + +// ClearExited removes the exited marker. Called from Start and Stop/Restart. +func (m *LocalManager) ClearExited(shellID string) { + m.exited.Delete(shellID) +} + // Start launches a new shell for the given workspace and returns its shellID. func (m *LocalManager) Start(workspaceID string, workspaceRoot string) (string, error) { - m.mu.Lock() - defer m.mu.Unlock() + m.sessionsMu.Lock() + defer m.sessionsMu.Unlock() shellID := uuid.NewString() @@ -110,13 +132,10 @@ func (m *LocalManager) Start(workspaceID string, workspaceRoot string) (string, } m.sessions[workspaceID][shellID] = sess - if m.exited[workspaceID] == nil { - m.exited[workspaceID] = make(map[string]struct{}) - } - delete(m.exited[workspaceID], shellID) - m.shellOrder[workspaceID] = append(m.shellOrder[workspaceID], shellID) + m.ClearExited(shellID) + outputDone := make(chan struct{}) go m.captureOutput(sess, ptyF, outputDone) go m.waitExit(sess, outputDone) @@ -160,12 +179,7 @@ func (m *LocalManager) waitExit(sess *Session, outputDone <-chan struct{}) { _ = sess.Cmd.Wait() <-outputDone - m.exitedMu.Lock() - if m.exited[sess.WorkspaceID] == nil { - m.exited[sess.WorkspaceID] = make(map[string]struct{}) - } - m.exited[sess.WorkspaceID][sess.ShellID] = struct{}{} - m.exitedMu.Unlock() + m.MarkAsExited(sess.ShellID) sess.mu.Lock() defer sess.mu.Unlock() @@ -185,8 +199,8 @@ func (m *LocalManager) waitExit(sess *Session, outputDone <-chan struct{}) { // Stop kills the shell for the given workspace and shellID. // Returns CodeNotFound if no session exists. func (m *LocalManager) Stop(workspaceID string, shellID string) error { - m.mu.Lock() - defer m.mu.Unlock() + m.sessionsMu.Lock() + defer m.sessionsMu.Unlock() ws, ok := m.sessions[workspaceID] if !ok { @@ -206,44 +220,26 @@ func (m *LocalManager) Stop(workspaceID string, shellID string) error { delete(m.sessions, workspaceID) } - m.exitedMu.Lock() - if m.exited[workspaceID] != nil { - delete(m.exited[workspaceID], shellID) - if len(m.exited[workspaceID]) == 0 { - delete(m.exited, workspaceID) - } - } - m.exitedMu.Unlock() + m.ClearExited(shellID) return nil } // Restart stops the identified shell (if running) then starts a new one. func (m *LocalManager) Restart(workspaceID string, shellID string, workspaceRoot string) (string, error) { - m.mu.Lock() - ws, ok := m.sessions[workspaceID] - if ok { + m.sessionsMu.Lock() + if ws, ok := m.sessions[workspaceID]; ok { if sess, ok := ws[shellID]; ok { - m.exitedMu.Lock() - _, exited := m.exited[workspaceID][shellID] - m.exitedMu.Unlock() - if !exited && sess.Cmd.Process != nil { + if sess.Cmd.Process != nil { _ = sess.Cmd.Process.Kill() } delete(ws, shellID) if len(ws) == 0 { delete(m.sessions, workspaceID) } - m.exitedMu.Lock() - if m.exited[workspaceID] != nil { - delete(m.exited[workspaceID], shellID) - if len(m.exited[workspaceID]) == 0 { - delete(m.exited, workspaceID) - } - } - m.exitedMu.Unlock() + m.ClearExited(shellID) } } - m.mu.Unlock() + m.sessionsMu.Unlock() return m.Start(workspaceID, workspaceRoot) } @@ -251,8 +247,8 @@ func (m *LocalManager) Restart(workspaceID string, shellID string, workspaceRoot // Resize resizes the PTY for the given shell. // Returns CodeNotFound if no session exists, or CodeBadRequest for invalid dimensions. func (m *LocalManager) Resize(workspaceID string, shellID string, cols, rows int) error { - m.mu.Lock() - defer m.mu.Unlock() + m.sessionsMu.RLock() + defer m.sessionsMu.RUnlock() ws, ok := m.sessions[workspaceID] if !ok { @@ -274,22 +270,19 @@ func (m *LocalManager) Resize(workspaceID string, shellID string, cols, rows int // Status returns the current shell status for the workspace. func (m *LocalManager) Status(workspaceID string, shellID string) Status { - m.mu.Lock() - defer m.mu.Unlock() - + m.sessionsMu.RLock() ws, ok := m.sessions[workspaceID] if !ok { + m.sessionsMu.RUnlock() return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false} } sess, ok := ws[shellID] + m.sessionsMu.RUnlock() if !ok { return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false} } - m.exitedMu.Lock() - _, exited := m.exited[workspaceID][shellID] - m.exitedMu.Unlock() - if exited { + if m.IsExited(shellID) { return Status{WorkspaceID: workspaceID, ShellID: shellID, Running: false} } @@ -304,14 +297,14 @@ func (m *LocalManager) Status(workspaceID string, shellID string) Status { // Subscribe creates a new output subscription for the identified shell. // Returns CodeNotFound if the shell has no session. func (m *LocalManager) Subscribe(workspaceID string, shellID string) (Subscription, error) { - m.mu.Lock() - defer m.mu.Unlock() - + m.sessionsMu.RLock() ws, ok := m.sessions[workspaceID] if !ok { + m.sessionsMu.RUnlock() return nil, util.New(util.CodeNotFound, "shell not found") } sess, ok := ws[shellID] + m.sessionsMu.RUnlock() if !ok { return nil, util.New(util.CodeNotFound, "shell not found") } @@ -321,10 +314,7 @@ func (m *LocalManager) Subscribe(workspaceID string, shellID string) (Subscripti sess.Subscribers[sub] = struct{}{} sess.mu.Unlock() - m.exitedMu.Lock() - _, exited := m.exited[workspaceID][shellID] - m.exitedMu.Unlock() - if exited { + if m.IsExited(shellID) { sub.closeChan() } @@ -343,8 +333,8 @@ func (m *LocalManager) ExitStatus(workspaceID string, shellID string) (ExitInfo, // List returns all shells for the workspace in creation order. func (m *LocalManager) List(workspaceID string) []ShellInfo { - m.mu.Lock() - defer m.mu.Unlock() + m.sessionsMu.Lock() + defer m.sessionsMu.Unlock() var infos []ShellInfo for _, shellID := range m.shellOrder[workspaceID] { @@ -363,11 +353,7 @@ func (m *LocalManager) List(workspaceID string) []ShellInfo { CreatedAt: sess.CreatedAt, } - m.exitedMu.Lock() - _, exited := m.exited[workspaceID][shellID] - m.exitedMu.Unlock() - - if !exited && sess.Cmd.Process != nil { + if !m.IsExited(shellID) && sess.Cmd.Process != nil { info.Running = true info.PID = sess.Cmd.Process.Pid } @@ -384,8 +370,8 @@ func (m *LocalManager) List(workspaceID string) []ShellInfo { // stdinOf returns the session's stdin writer or CodeNotFound. func (m *LocalManager) stdinOf(workspaceID string, shellID string) (io.WriteCloser, error) { - m.mu.Lock() - defer m.mu.Unlock() + m.sessionsMu.RLock() + defer m.sessionsMu.RUnlock() ws, ok := m.sessions[workspaceID] if !ok { @@ -400,14 +386,14 @@ func (m *LocalManager) stdinOf(workspaceID string, shellID string) (io.WriteClos // exitOf returns the session's recorded exit info or CodeNotFound. func (m *LocalManager) exitOf(workspaceID string, shellID string) (ExitInfo, error) { - m.mu.Lock() - defer m.mu.Unlock() - + m.sessionsMu.RLock() ws, ok := m.sessions[workspaceID] if !ok { + m.sessionsMu.RUnlock() return ExitInfo{}, util.New(util.CodeNotFound, "shell not found") } sess, ok := ws[shellID] + m.sessionsMu.RUnlock() if !ok { return ExitInfo{}, util.New(util.CodeNotFound, "shell not found") } diff --git a/internal/shell/manager_test.go b/internal/shell/manager_test.go index b9640dc..b660b6e 100644 --- a/internal/shell/manager_test.go +++ b/internal/shell/manager_test.go @@ -201,3 +201,50 @@ func TestShellMultiInstance(t *testing.T) { t.Fatalf("List returned shell %q, want %q", list[0].ShellID, shellID2) } } + +func TestShellIsExitedHelpers(t *testing.T) { + mgr := NewManager("bash", []string{"-i"}) + root := t.TempDir() + + shellIDA, err := mgr.Start("test-ws", root) + if err != nil { + t.Fatalf("Start shell A failed: %v", err) + } + shellIDB, err := mgr.Start("test-ws", root) + if err != nil { + t.Fatalf("Start shell B failed: %v", err) + } + + t.Cleanup(func() { + _ = mgr.Stop("test-ws", shellIDA) + _ = mgr.Stop("test-ws", shellIDB) + }) + + if mgr.IsExited(shellIDA) { + t.Errorf("IsExited(A) = true, want false") + } + if mgr.IsExited(shellIDB) { + t.Errorf("IsExited(B) = true, want false") + } + + if err := mgr.Stop("test-ws", shellIDA); err != nil { + t.Fatalf("Stop shell A failed: %v", err) + } + + if mgr.IsExited(shellIDA) { + t.Errorf("IsExited(A) = true after Stop, want false") + } + if mgr.IsExited(shellIDB) { + t.Errorf("IsExited(B) = true after stopping A, want false") + } + + mgr.MarkAsExited(shellIDB) + if !mgr.IsExited(shellIDB) { + t.Errorf("IsExited(B) = false after MarkAsExited, want true") + } + + mgr.ClearExited(shellIDB) + if mgr.IsExited(shellIDB) { + t.Errorf("IsExited(B) = true after ClearExited, want false") + } +}