refactor(shell): sync.RWMutex for sessions + sync.Map for exited (mirror process pkg)

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
This commit is contained in:
tao.chen
2026-07-06 16:14:55 +08:00
parent 5e694aec12
commit 3d34de96bf
2 changed files with 104 additions and 71 deletions
+57 -71
View File
@@ -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")
}
+47
View File
@@ -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")
}
}