git.iko.soy/crush
on 77ecb39 help sign in
crush/internal/shell/background.go
213 lines 5.3 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// BackgroundShell represents a shell running in the background.
23type BackgroundShell struct {
24 ID string
25 Command string
26 Description string
27 Shell *Shell
28 WorkingDir string
29 ctx context.Context
30 cancel context.CancelFunc
31 stdout *bytes.Buffer
32 stderr *bytes.Buffer
33 done chan struct{}
34 exitErr error
35 completedAt int64 // Unix timestamp when job completed (0 if still running)
38// BackgroundShellManager manages background shell instances.
39type BackgroundShellManager struct {
40 shells *csync.Map[string, *BackgroundShell]
43var (
44 backgroundManager *BackgroundShellManager
45 backgroundManagerOnce sync.Once
46 idCounter atomic.Uint64
49// GetBackgroundShellManager returns the singleton background shell manager.
50func GetBackgroundShellManager() *BackgroundShellManager {
51 backgroundManagerOnce.Do(func() {
52 backgroundManager = &BackgroundShellManager{
53 shells: csync.NewMap[string, *BackgroundShell](),
54 }
55 })
56 return backgroundManager
59// Start creates and starts a new background shell with the given command.
60func (m *BackgroundShellManager) Start(ctx context.Context, workingDir string, blockFuncs []BlockFunc, command string, description string) (*BackgroundShell, error) {
61 // Check job limit
62 if m.shells.Len() >= MaxBackgroundJobs {
63 return nil, fmt.Errorf("maximum number of background jobs (%d) reached. Please terminate or wait for some jobs to complete", MaxBackgroundJobs)
64 }
66 id := fmt.Sprintf("%03X", idCounter.Add(1))
68 shell := NewShell(&Options{
69 WorkingDir: workingDir,
70 BlockFuncs: blockFuncs,
71 })
73 shellCtx, cancel := context.WithCancel(ctx)
75 bgShell := &BackgroundShell{
76 ID: id,
77 Command: command,
78 Description: description,
79 WorkingDir: workingDir,
80 Shell: shell,
81 ctx: shellCtx,
82 cancel: cancel,
83 stdout: &bytes.Buffer{},
84 stderr: &bytes.Buffer{},
85 done: make(chan struct{}),
86 }
88 m.shells.Set(id, bgShell)
90 go func() {
91 defer close(bgShell.done)
93 err := shell.ExecStream(shellCtx, command, bgShell.stdout, bgShell.stderr)
95 bgShell.exitErr = err
96 atomic.StoreInt64(&bgShell.completedAt, time.Now().Unix())
97 }()
99 return bgShell, nil
102// Get retrieves a background shell by ID.
103func (m *BackgroundShellManager) Get(id string) (*BackgroundShell, bool) {
104 return m.shells.Get(id)
107// Remove removes a background shell from the manager without terminating it.
108// This is useful when a shell has already completed and you just want to clean up tracking.
109func (m *BackgroundShellManager) Remove(id string) error {
110 _, ok := m.shells.Take(id)
111 if !ok {
112 return fmt.Errorf("background shell not found: %s", id)
113 }
114 return nil
117// Kill terminates a background shell by ID.
118func (m *BackgroundShellManager) Kill(id string) error {
119 shell, ok := m.shells.Take(id)
120 if !ok {
121 return fmt.Errorf("background shell not found: %s", id)
122 }
124 shell.cancel()
125 <-shell.done
126 return nil
129// BackgroundShellInfo contains information about a background shell.
130type BackgroundShellInfo struct {
131 ID string
132 Command string
133 Description string
136// List returns all background shell IDs.
137func (m *BackgroundShellManager) List() []string {
138 ids := make([]string, 0, m.shells.Len())
139 for id := range m.shells.Seq2() {
140 ids = append(ids, id)
141 }
142 return ids
145// Cleanup removes completed jobs that have been finished for more than the retention period
146func (m *BackgroundShellManager) Cleanup() int {
147 now := time.Now().Unix()
148 retentionSeconds := int64(CompletedJobRetentionMinutes * 60)
150 var toRemove []string
151 for shell := range m.shells.Seq() {
152 completedAt := atomic.LoadInt64(&shell.completedAt)
153 if completedAt > 0 && now-completedAt > retentionSeconds {
154 toRemove = append(toRemove, shell.ID)
155 }
156 }
158 for _, id := range toRemove {
159 m.Remove(id)
160 }
162 return len(toRemove)
165// KillAll terminates all background shells.
166func (m *BackgroundShellManager) KillAll() {
167 shells := slices.Collect(m.shells.Seq())
168 m.shells.Reset(map[string]*BackgroundShell{})
169 done := make(chan struct{}, 1)
170 go func() {
171 var wg sync.WaitGroup
172 for _, shell := range shells {
173 wg.Go(func() {
174 shell.cancel()
175 <-shell.done
176 })
177 }
178 wg.Wait()
179 done <- struct{}{}
180 }()
182 select {
183 case <-done:
184 return
185 case <-time.After(time.Second * 5):
186 return
187 }
190// GetOutput returns the current output of a background shell.
191func (bs *BackgroundShell) GetOutput() (stdout string, stderr string, done bool, err error) {
192 select {
193 case <-bs.done:
194 return bs.stdout.String(), bs.stderr.String(), true, bs.exitErr
195 default:
196 return bs.stdout.String(), bs.stderr.String(), false, nil
197 }
200// IsDone checks if the background shell has finished execution.
201func (bs *BackgroundShell) IsDone() bool {
202 select {
203 case <-bs.done:
204 return true
205 default:
206 return false
207 }
210// Wait blocks until the background shell completes.
211func (bs *BackgroundShell) Wait() {
212 <-bs.done