git.iko.soy/crush
on efe42e8 help sign in
crush/internal/shell/background.go
245 lines 6.1 kB history raw
1package shell
3import (
4 "bytes"
5 "context"
6 "fmt"
7 "slices"
8 "sync"
9 "sync/atomic"
10 "time"
12 "github.com/charmbracelet/crush/internal/csync"
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
22// syncBuffer is a thread-safe wrapper around bytes.Buffer.
23type syncBuffer struct {
24 buf bytes.Buffer
25 mu sync.RWMutex
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)
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)
40func (sb *syncBuffer) String() string {
41 sb.mu.RLock()
42 defer sb.mu.RUnlock()
43 return sb.buf.String()
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 atomic.Int64 // Unix timestamp when job completed (0 if still running)
62// BackgroundShellManager manages background shell instances.
63type BackgroundShellManager struct {
64 shells *csync.Map[string, *BackgroundShell]
67var (
68 backgroundManager *BackgroundShellManager
69 backgroundManagerOnce sync.Once
70 idCounter atomic.Uint64
73// newBackgroundShellManager creates a new BackgroundShellManager instance.
74func newBackgroundShellManager() *BackgroundShellManager {
75 return &BackgroundShellManager{
76 shells: csync.NewMap[string, *BackgroundShell](),
77 }
80// GetBackgroundShellManager returns the singleton background shell manager.
81func GetBackgroundShellManager() *BackgroundShellManager {
82 backgroundManagerOnce.Do(func() {
83 backgroundManager = newBackgroundShellManager()
84 })
85 return backgroundManager
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 }
95 id := fmt.Sprintf("%03X", idCounter.Add(1))
97 shell := NewShell(&Options{
98 WorkingDir: workingDir,
99 BlockFuncs: blockFuncs,
100 })
102 shellCtx, cancel := context.WithCancel(ctx)
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 }
117 m.shells.Set(id, bgShell)
119 go func() {
120 defer close(bgShell.done)
122 err := shell.ExecStream(shellCtx, command, bgShell.stdout, bgShell.stderr)
124 bgShell.exitErr = err
125 bgShell.completedAt.Store(time.Now().Unix())
126 }()
128 return bgShell, nil
131// Get retrieves a background shell by ID.
132func (m *BackgroundShellManager) Get(id string) (*BackgroundShell, bool) {
133 return m.shells.Get(id)
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
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 }
153 shell.cancel()
154 <-shell.done
155 return nil
158// BackgroundShellInfo contains information about a background shell.
159type BackgroundShellInfo struct {
160 ID string
161 Command string
162 Description string
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
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)
179 var toRemove []string
180 for shell := range m.shells.Seq() {
181 completedAt := shell.completedAt.Load()
182 if completedAt > 0 && now-completedAt > retentionSeconds {
183 toRemove = append(toRemove, shell.ID)
184 }
185 }
187 for _, id := range toRemove {
188 m.Remove(id)
189 }
191 return len(toRemove)
194// KillAll terminates all background shells. The provided context bounds how
195// long the function waits for each shell to exit.
196func (m *BackgroundShellManager) KillAll(ctx context.Context) {
197 shells := slices.Collect(m.shells.Seq())
198 m.shells.Reset(map[string]*BackgroundShell{})
200 var wg sync.WaitGroup
201 for _, shell := range shells {
202 wg.Go(func() {
203 shell.cancel()
204 select {
205 case <-shell.done:
206 case <-ctx.Done():
207 }
208 })
209 }
210 wg.Wait()
213// GetOutput returns the current output of a background shell.
214func (bs *BackgroundShell) GetOutput() (stdout string, stderr string, done bool, err error) {
215 select {
216 case <-bs.done:
217 return bs.stdout.String(), bs.stderr.String(), true, bs.exitErr
218 default:
219 return bs.stdout.String(), bs.stderr.String(), false, nil
220 }
223// IsDone checks if the background shell has finished execution.
224func (bs *BackgroundShell) IsDone() bool {
225 select {
226 case <-bs.done:
227 return true
228 default:
229 return false
230 }
233// Wait blocks until the background shell completes.
234func (bs *BackgroundShell) Wait() {
235 <-bs.done
238func (bs *BackgroundShell) WaitContext(ctx context.Context) bool {
239 select {
240 case <-bs.done:
241 return true
242 case <-ctx.Done():
243 return false
244 }