background.go

  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// 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)
 36}
 37
 38// BackgroundShellManager manages background shell instances.
 39type BackgroundShellManager struct {
 40	shells *csync.Map[string, *BackgroundShell]
 41}
 42
 43var (
 44	backgroundManager     *BackgroundShellManager
 45	backgroundManagerOnce sync.Once
 46	idCounter             atomic.Uint64
 47)
 48
 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
 57}
 58
 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	}
 65
 66	id := fmt.Sprintf("%03X", idCounter.Add(1))
 67
 68	shell := NewShell(&Options{
 69		WorkingDir: workingDir,
 70		BlockFuncs: blockFuncs,
 71	})
 72
 73	shellCtx, cancel := context.WithCancel(ctx)
 74
 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	}
 87
 88	m.shells.Set(id, bgShell)
 89
 90	go func() {
 91		defer close(bgShell.done)
 92
 93		err := shell.ExecStream(shellCtx, command, bgShell.stdout, bgShell.stderr)
 94
 95		bgShell.exitErr = err
 96		atomic.StoreInt64(&bgShell.completedAt, time.Now().Unix())
 97	}()
 98
 99	return bgShell, nil
100}
101
102// Get retrieves a background shell by ID.
103func (m *BackgroundShellManager) Get(id string) (*BackgroundShell, bool) {
104	return m.shells.Get(id)
105}
106
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
115}
116
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	}
123
124	shell.cancel()
125	<-shell.done
126	return nil
127}
128
129// BackgroundShellInfo contains information about a background shell.
130type BackgroundShellInfo struct {
131	ID          string
132	Command     string
133	Description string
134}
135
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
143}
144
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)
149
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	}
157
158	for _, id := range toRemove {
159		m.Remove(id)
160	}
161
162	return len(toRemove)
163}
164
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	}()
181
182	select {
183	case <-done:
184		return
185	case <-time.After(time.Second * 5):
186		return
187	}
188}
189
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	}
198}
199
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	}
208}
209
210// Wait blocks until the background shell completes.
211func (bs *BackgroundShell) Wait() {
212	<-bs.done
213}