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:
+57
-71
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user