@@ -9,14 +9,24 @@ package main
99
1010import (
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+
2030var (
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+
3165func 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