Skip to content

Commit 8b5c8d9

Browse files
croakywarp-agent
andauthored
go: consolidate procman.go into main.go
Single-file structure is idiomatic for small CLI tools that aren't intended as libraries. At ~320 lines, the code is easily scannable in one file without jumping between files. This also improves consistency with mdembed, another pinned repo with similar single-purpose CLI tool characteristics. Ordering: imports → const/types → main() → functions by receiver. Also consolidate test files. This mirrors the source file consolidation for consistency. At ~270 lines, the tests are easily scannable in one file. TestMain stays at the top as test infrastructure, followed by unit tests, then the integration test. #2 Co-authored-by: Warp <agent@warp.dev>
1 parent 918ab8f commit 8b5c8d9

4 files changed

Lines changed: 368 additions & 381 deletions

File tree

main.go

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,24 @@ package main
99

1010
import (
1111
"bufio"
12+
"bytes"
1213
"errors"
1314
"fmt"
1415
"io"
1516
"os"
17+
"os/exec"
18+
"os/signal"
1619
"regexp"
1720
"strings"
21+
"sync"
22+
"syscall"
23+
"time"
24+
25+
"github.com/creack/pty"
1826
)
1927

28+
const timeout = 5 * time.Second
29+
2030
var (
2131
colors = []int{2, 3, 4, 5, 6, 42, 130, 103, 129, 108}
2232
procfileRe = regexp.MustCompile(`^([\w-]+):\s+(.+)$`)
@@ -28,6 +38,30 @@ type procDef struct {
2838
cmd string
2939
}
3040

41+
// manager handles overall process management
42+
type manager struct {
43+
output *output
44+
procs []*process
45+
procWg sync.WaitGroup
46+
done chan bool
47+
interrupted chan os.Signal
48+
}
49+
50+
// process represents an individual process to be managed
51+
type process struct {
52+
*exec.Cmd
53+
name string
54+
color int
55+
output *output
56+
}
57+
58+
// output manages the output display of processes
59+
type output struct {
60+
maxNameLength int
61+
mutex sync.Mutex
62+
pipes map[*process]*os.File
63+
}
64+
3165
func main() {
3266
procNames, err := parseArgs(os.Args)
3367
if err != nil {
@@ -115,3 +149,172 @@ func parseProcfile(r io.Reader) ([]procDef, error) {
115149
}
116150
return defs, nil
117151
}
152+
153+
// setupProcesses creates and initializes processes based on the given procDefs.
154+
func (mgr *manager) setupProcesses(defs []procDef, procNames []string) error {
155+
defMap := make(map[string]string)
156+
for _, def := range defs {
157+
defMap[def.name] = def.cmd
158+
}
159+
160+
for i, name := range procNames {
161+
cmd, ok := defMap[name]
162+
if !ok {
163+
return fmt.Errorf("No process named %s in Procfile.dev\n", name)
164+
}
165+
166+
proc := &process{
167+
Cmd: exec.Command("/bin/sh", "-c", cmd),
168+
name: name,
169+
color: colors[i%len(colors)],
170+
output: mgr.output,
171+
}
172+
mgr.procs = append(mgr.procs, proc)
173+
}
174+
return nil
175+
}
176+
177+
// setupSignalHandling configures handling of interrupt signals.
178+
func (mgr *manager) setupSignalHandling() {
179+
mgr.done = make(chan bool, len(mgr.procs))
180+
mgr.interrupted = make(chan os.Signal, 1)
181+
signal.Notify(mgr.interrupted, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
182+
}
183+
184+
// runProcess adds the process to the wait group and starts it.
185+
func (mgr *manager) runProcess(proc *process) {
186+
mgr.procWg.Add(1)
187+
go func() {
188+
defer mgr.procWg.Done()
189+
defer func() { mgr.done <- true }()
190+
proc.run()
191+
}()
192+
}
193+
194+
// waitForExit waits for all processes to exit or for an interruption signal.
195+
func (mgr *manager) waitForExit() {
196+
select {
197+
case <-mgr.done:
198+
case <-mgr.interrupted:
199+
}
200+
201+
for _, proc := range mgr.procs {
202+
proc.interrupt()
203+
}
204+
205+
select {
206+
case <-time.After(timeout):
207+
case <-mgr.interrupted:
208+
}
209+
210+
for _, proc := range mgr.procs {
211+
proc.kill()
212+
}
213+
}
214+
215+
// running inspects the process to see if it is currently running.
216+
func (proc *process) running() bool {
217+
return proc.Process != nil && proc.ProcessState == nil
218+
}
219+
220+
// run starts the execution of the process and handles its output.
221+
func (proc *process) run() {
222+
if proc.Process != nil {
223+
fmt.Fprintf(os.Stderr, "Process %s already started\n", proc.name)
224+
return
225+
}
226+
227+
proc.output.pipeOutput(proc)
228+
defer proc.output.closePipe(proc)
229+
230+
if err := proc.Cmd.Wait(); err != nil {
231+
proc.output.writeErr(proc, err)
232+
}
233+
}
234+
235+
// interrupt sends an interrupt signal to a running process.
236+
func (proc *process) interrupt() {
237+
if proc.running() {
238+
proc.signal(syscall.SIGINT)
239+
}
240+
}
241+
242+
// kill forcefully stops a running process.
243+
func (proc *process) kill() {
244+
if proc.running() {
245+
proc.signal(syscall.SIGKILL)
246+
}
247+
}
248+
249+
// signal sends a specified signal to the process group.
250+
func (proc *process) signal(sig syscall.Signal) {
251+
if err := syscall.Kill(-proc.Process.Pid, sig); err != nil {
252+
proc.output.writeErr(proc, err)
253+
}
254+
}
255+
256+
// init initializes the output handler for all processes.
257+
func (out *output) init(procs []*process) {
258+
out.pipes = make(map[*process]*os.File)
259+
for _, proc := range procs {
260+
if len(proc.name) > out.maxNameLength {
261+
out.maxNameLength = len(proc.name)
262+
}
263+
}
264+
}
265+
266+
// pipeOutput handles the output piping for a process.
267+
func (out *output) pipeOutput(proc *process) {
268+
ptyFile, err := pty.Start(proc.Cmd)
269+
if err != nil {
270+
fmt.Fprintf(os.Stderr, "Error opening PTY: %v\n", err)
271+
os.Exit(1)
272+
}
273+
274+
out.mutex.Lock()
275+
out.pipes[proc] = ptyFile
276+
out.mutex.Unlock()
277+
278+
go func() {
279+
scanner := bufio.NewScanner(ptyFile)
280+
for scanner.Scan() {
281+
out.writeLine(proc, scanner.Bytes())
282+
}
283+
if err := scanner.Err(); err != nil && !errors.Is(err, os.ErrClosed) {
284+
out.writeErr(proc, err)
285+
}
286+
}()
287+
}
288+
289+
// closePipe closes the pseudo-terminal associated with the process.
290+
func (out *output) closePipe(proc *process) {
291+
out.mutex.Lock()
292+
ptyFile := out.pipes[proc]
293+
out.mutex.Unlock()
294+
295+
if ptyFile != nil {
296+
ptyFile.Close()
297+
}
298+
}
299+
300+
// writeLine writes a line of output for the specified process, with color formatting.
301+
func (out *output) writeLine(proc *process, p []byte) {
302+
var buf bytes.Buffer
303+
buf.WriteString(fmt.Sprintf("\033[1;38;5;%vm", proc.color))
304+
buf.WriteString(proc.name)
305+
for i := len(proc.name); i <= out.maxNameLength; i++ {
306+
buf.WriteByte(' ')
307+
}
308+
buf.WriteString("\033[0m| ")
309+
buf.Write(p)
310+
buf.WriteByte('\n')
311+
312+
out.mutex.Lock()
313+
defer out.mutex.Unlock()
314+
buf.WriteTo(os.Stdout)
315+
}
316+
317+
// writeErr writes an error message for the specified process.
318+
func (out *output) writeErr(proc *process, err error) {
319+
out.writeLine(proc, []byte(fmt.Sprintf("\033[0;31m%v\033[0m", err)))
320+
}

0 commit comments

Comments
 (0)