background.go

  1package shell
  2
  3import (
  4	"bytes"
  5	"context"
  6	"fmt"
  7	"sync"
  8	"sync/atomic"
  9
 10	"github.com/charmbracelet/crush/internal/csync"
 11)
 12
 13// BackgroundShell represents a shell running in the background.
 14type BackgroundShell struct {
 15	ID          string
 16	Command     string
 17	Description string
 18	Shell       *Shell
 19	WorkingDir  string
 20	ctx         context.Context
 21	cancel      context.CancelFunc
 22	stdout      *bytes.Buffer
 23	stderr      *bytes.Buffer
 24	done        chan struct{}
 25	exitErr     error
 26}
 27
 28// BackgroundShellManager manages background shell instances.
 29type BackgroundShellManager struct {
 30	shells *csync.Map[string, *BackgroundShell]
 31}
 32
 33var (
 34	backgroundManager     *BackgroundShellManager
 35	backgroundManagerOnce sync.Once
 36	idCounter             atomic.Uint64
 37)
 38
 39// GetBackgroundShellManager returns the singleton background shell manager.
 40func GetBackgroundShellManager() *BackgroundShellManager {
 41	backgroundManagerOnce.Do(func() {
 42		backgroundManager = &BackgroundShellManager{
 43			shells: csync.NewMap[string, *BackgroundShell](),
 44		}
 45	})
 46	return backgroundManager
 47}
 48
 49// Start creates and starts a new background shell with the given command.
 50func (m *BackgroundShellManager) Start(ctx context.Context, workingDir string, blockFuncs []BlockFunc, command string, description string) (*BackgroundShell, error) {
 51	id := fmt.Sprintf("%03X", idCounter.Add(1))
 52
 53	shell := NewShell(&Options{
 54		WorkingDir: workingDir,
 55		BlockFuncs: blockFuncs,
 56	})
 57
 58	shellCtx, cancel := context.WithCancel(ctx)
 59
 60	bgShell := &BackgroundShell{
 61		ID:          id,
 62		Command:     command,
 63		Description: description,
 64		WorkingDir:  workingDir,
 65		Shell:       shell,
 66		ctx:         shellCtx,
 67		cancel:      cancel,
 68		stdout:      &bytes.Buffer{},
 69		stderr:      &bytes.Buffer{},
 70		done:        make(chan struct{}),
 71	}
 72
 73	m.shells.Set(id, bgShell)
 74
 75	go func() {
 76		defer close(bgShell.done)
 77
 78		err := shell.ExecStream(shellCtx, command, bgShell.stdout, bgShell.stderr)
 79
 80		bgShell.exitErr = err
 81	}()
 82
 83	return bgShell, nil
 84}
 85
 86// Get retrieves a background shell by ID.
 87func (m *BackgroundShellManager) Get(id string) (*BackgroundShell, bool) {
 88	return m.shells.Get(id)
 89}
 90
 91// Remove removes a background shell from the manager without terminating it.
 92// This is useful when a shell has already completed and you just want to clean up tracking.
 93func (m *BackgroundShellManager) Remove(id string) error {
 94	_, ok := m.shells.Take(id)
 95	if !ok {
 96		return fmt.Errorf("background shell not found: %s", id)
 97	}
 98	return nil
 99}
100
101// Kill terminates a background shell by ID.
102func (m *BackgroundShellManager) Kill(id string) error {
103	shell, ok := m.shells.Take(id)
104	if !ok {
105		return fmt.Errorf("background shell not found: %s", id)
106	}
107
108	shell.cancel()
109	<-shell.done
110	return nil
111}
112
113// BackgroundShellInfo contains information about a background shell.
114type BackgroundShellInfo struct {
115	ID          string
116	Command     string
117	Description string
118}
119
120// List returns all background shell IDs.
121func (m *BackgroundShellManager) List() []string {
122	ids := make([]string, 0, m.shells.Len())
123	for id := range m.shells.Seq2() {
124		ids = append(ids, id)
125	}
126	return ids
127}
128
129// KillAll terminates all background shells.
130func (m *BackgroundShellManager) KillAll() {
131	shells := make([]*BackgroundShell, 0, m.shells.Len())
132	for shell := range m.shells.Seq() {
133		shells = append(shells, shell)
134	}
135	m.shells.Reset(map[string]*BackgroundShell{})
136
137	for _, shell := range shells {
138		shell.cancel()
139		<-shell.done
140	}
141}
142
143// GetOutput returns the current output of a background shell.
144func (bs *BackgroundShell) GetOutput() (stdout string, stderr string, done bool, err error) {
145	select {
146	case <-bs.done:
147		return bs.stdout.String(), bs.stderr.String(), true, bs.exitErr
148	default:
149		return bs.stdout.String(), bs.stderr.String(), false, nil
150	}
151}
152
153// IsDone checks if the background shell has finished execution.
154func (bs *BackgroundShell) IsDone() bool {
155	select {
156	case <-bs.done:
157		return true
158	default:
159		return false
160	}
161}
162
163// Wait blocks until the background shell completes.
164func (bs *BackgroundShell) Wait() {
165	<-bs.done
166}