1package shell
2
3import (
4 "bytes"
5 "context"
6 "fmt"
7 "slices"
8 "sync"
9 "sync/atomic"
10 "time"
11
12 "github.com/charmbracelet/crush/internal/csync"
13)
14
15const (
16 // MaxBackgroundJobs is the maximum number of concurrent background jobs allowed
17 MaxBackgroundJobs = 50
18 // CompletedJobRetentionMinutes is how long to keep completed jobs before auto-cleanup (8 hours)
19 CompletedJobRetentionMinutes = 8 * 60
20)
21
22// syncBuffer is a thread-safe wrapper around bytes.Buffer.
23type syncBuffer struct {
24 buf bytes.Buffer
25 mu sync.RWMutex
26}
27
28func (sb *syncBuffer) Write(p []byte) (n int, err error) {
29 sb.mu.Lock()
30 defer sb.mu.Unlock()
31 return sb.buf.Write(p)
32}
33
34func (sb *syncBuffer) WriteString(s string) (n int, err error) {
35 sb.mu.Lock()
36 defer sb.mu.Unlock()
37 return sb.buf.WriteString(s)
38}
39
40func (sb *syncBuffer) String() string {
41 sb.mu.RLock()
42 defer sb.mu.RUnlock()
43 return sb.buf.String()
44}
45
46// BackgroundShell represents a shell running in the background.
47type BackgroundShell struct {
48 ID string
49 Command string
50 Description string
51 Shell *Shell
52 WorkingDir string
53 ctx context.Context
54 cancel context.CancelFunc
55 stdout *syncBuffer
56 stderr *syncBuffer
57 done chan struct{}
58 exitErr error
59 completedAt int64 // Unix timestamp when job completed (0 if still running)
60}
61
62// BackgroundShellManager manages background shell instances.
63type BackgroundShellManager struct {
64 shells *csync.Map[string, *BackgroundShell]
65}
66
67var (
68 backgroundManager *BackgroundShellManager
69 backgroundManagerOnce sync.Once
70 idCounter atomic.Uint64
71)
72
73// newBackgroundShellManager creates a new BackgroundShellManager instance.
74func newBackgroundShellManager() *BackgroundShellManager {
75 return &BackgroundShellManager{
76 shells: csync.NewMap[string, *BackgroundShell](),
77 }
78}
79
80// GetBackgroundShellManager returns the singleton background shell manager.
81func GetBackgroundShellManager() *BackgroundShellManager {
82 backgroundManagerOnce.Do(func() {
83 backgroundManager = newBackgroundShellManager()
84 })
85 return backgroundManager
86}
87
88// Start creates and starts a new background shell with the given command.
89func (m *BackgroundShellManager) Start(ctx context.Context, workingDir string, blockFuncs []BlockFunc, command string, description string) (*BackgroundShell, error) {
90 // Check job limit
91 if m.shells.Len() >= MaxBackgroundJobs {
92 return nil, fmt.Errorf("maximum number of background jobs (%d) reached. Please terminate or wait for some jobs to complete", MaxBackgroundJobs)
93 }
94
95 id := fmt.Sprintf("%03X", idCounter.Add(1))
96
97 shell := NewShell(&Options{
98 WorkingDir: workingDir,
99 BlockFuncs: blockFuncs,
100 })
101
102 shellCtx, cancel := context.WithCancel(ctx)
103
104 bgShell := &BackgroundShell{
105 ID: id,
106 Command: command,
107 Description: description,
108 WorkingDir: workingDir,
109 Shell: shell,
110 ctx: shellCtx,
111 cancel: cancel,
112 stdout: &syncBuffer{},
113 stderr: &syncBuffer{},
114 done: make(chan struct{}),
115 }
116
117 m.shells.Set(id, bgShell)
118
119 go func() {
120 defer close(bgShell.done)
121
122 err := shell.ExecStream(shellCtx, command, bgShell.stdout, bgShell.stderr)
123
124 bgShell.exitErr = err
125 atomic.StoreInt64(&bgShell.completedAt, time.Now().Unix())
126 }()
127
128 return bgShell, nil
129}
130
131// Get retrieves a background shell by ID.
132func (m *BackgroundShellManager) Get(id string) (*BackgroundShell, bool) {
133 return m.shells.Get(id)
134}
135
136// Remove removes a background shell from the manager without terminating it.
137// This is useful when a shell has already completed and you just want to clean up tracking.
138func (m *BackgroundShellManager) Remove(id string) error {
139 _, ok := m.shells.Take(id)
140 if !ok {
141 return fmt.Errorf("background shell not found: %s", id)
142 }
143 return nil
144}
145
146// Kill terminates a background shell by ID.
147func (m *BackgroundShellManager) Kill(id string) error {
148 shell, ok := m.shells.Take(id)
149 if !ok {
150 return fmt.Errorf("background shell not found: %s", id)
151 }
152
153 shell.cancel()
154 <-shell.done
155 return nil
156}
157
158// BackgroundShellInfo contains information about a background shell.
159type BackgroundShellInfo struct {
160 ID string
161 Command string
162 Description string
163}
164
165// List returns all background shell IDs.
166func (m *BackgroundShellManager) List() []string {
167 ids := make([]string, 0, m.shells.Len())
168 for id := range m.shells.Seq2() {
169 ids = append(ids, id)
170 }
171 return ids
172}
173
174// Cleanup removes completed jobs that have been finished for more than the retention period
175func (m *BackgroundShellManager) Cleanup() int {
176 now := time.Now().Unix()
177 retentionSeconds := int64(CompletedJobRetentionMinutes * 60)
178
179 var toRemove []string
180 for shell := range m.shells.Seq() {
181 completedAt := atomic.LoadInt64(&shell.completedAt)
182 if completedAt > 0 && now-completedAt > retentionSeconds {
183 toRemove = append(toRemove, shell.ID)
184 }
185 }
186
187 for _, id := range toRemove {
188 m.Remove(id)
189 }
190
191 return len(toRemove)
192}
193
194// KillAll terminates all background shells.
195func (m *BackgroundShellManager) KillAll() {
196 shells := slices.Collect(m.shells.Seq())
197 m.shells.Reset(map[string]*BackgroundShell{})
198 done := make(chan struct{}, 1)
199 go func() {
200 var wg sync.WaitGroup
201 for _, shell := range shells {
202 wg.Go(func() {
203 shell.cancel()
204 <-shell.done
205 })
206 }
207 wg.Wait()
208 done <- struct{}{}
209 }()
210
211 select {
212 case <-done:
213 return
214 case <-time.After(time.Second * 5):
215 return
216 }
217}
218
219// GetOutput returns the current output of a background shell.
220func (bs *BackgroundShell) GetOutput() (stdout string, stderr string, done bool, err error) {
221 select {
222 case <-bs.done:
223 return bs.stdout.String(), bs.stderr.String(), true, bs.exitErr
224 default:
225 return bs.stdout.String(), bs.stderr.String(), false, nil
226 }
227}
228
229// IsDone checks if the background shell has finished execution.
230func (bs *BackgroundShell) IsDone() bool {
231 select {
232 case <-bs.done:
233 return true
234 default:
235 return false
236 }
237}
238
239// Wait blocks until the background shell completes.
240func (bs *BackgroundShell) Wait() {
241 <-bs.done
242}