diff --git a/barrier.go b/barrier.go new file mode 100644 index 0000000..3709721 --- /dev/null +++ b/barrier.go @@ -0,0 +1,37 @@ +package main + +import ( + "sync" +) + +// Direct import of https://github.com/pwaller/barrier/blob/master/barrier.go + +// The zero of Barrier is a ready-to-use value +type Barrier struct { + channel chan struct{} + fall, initialize sync.Once + FallHook func() +} + +func (b *Barrier) init() { + b.initialize.Do(func() { b.channel = make(chan struct{}) }) +} + +// `b.Fall()` can be called any number of times and causes the channel returned +// by `b.Barrier()` to become closed (permanently available for immediate reading) +func (b *Barrier) Fall() { + b.init() + b.fall.Do(func() { + if b.FallHook != nil { + b.FallHook() + } + close(b.channel) + }) +} + +// When `b.Fall()` is called, the channel returned by Barrier() is closed +// (and becomes always readable) +func (b *Barrier) Barrier() <-chan struct{} { + b.init() + return b.channel +} diff --git a/env.go b/env.go index 36aa500..c162fd1 100644 --- a/env.go +++ b/env.go @@ -1,6 +1,7 @@ package main import ( + "fmt" "github.com/subosito/gotenv" "os" "regexp" @@ -25,3 +26,13 @@ func ReadEnv(filename string) (Env, error) { } return env, nil } + +func (e *Env) asArray() (env []string) { + for _, pair := range os.Environ() { + env = append(env, pair) + } + for name, val := range *e { + env = append(env, fmt.Sprintf("%s=%s", name, val)) + } + return +} diff --git a/fixtures/writey/Procfile b/fixtures/writey/Procfile new file mode 100644 index 0000000..6e91b25 --- /dev/null +++ b/fixtures/writey/Procfile @@ -0,0 +1,2 @@ +writey1: writey +writey2: writey diff --git a/fixtures/writey/main.go b/fixtures/writey/main.go new file mode 100644 index 0000000..ff652b5 --- /dev/null +++ b/fixtures/writey/main.go @@ -0,0 +1,42 @@ +package main + +import ( + "fmt" + "math/rand" + "os" + "os/signal" + "syscall" + "time" +) + +func main() { + rand.Seed(time.Now().UnixNano()) + print("Foo") + time.Sleep(10 * time.Millisecond) + println("Bar") + + print("Baz") + time.Sleep(10 * time.Millisecond) + println("Qux") + + fmt.Fprintln(os.Stdout, "This is on \x1b[32mstdout") + + os.Stdout.Close() + + s := rand.Intn(3) + 1 + + c := make(chan os.Signal) + signal.Notify(c, os.Interrupt, syscall.SIGTERM) + + select { + case <-c: + if rand.Intn(4) == 1 { + println("IGNORING EXIT") + time.Sleep(100 * time.Second) + } + println("Got SIGTERM") + case <-time.After(time.Duration(s) * time.Second): + println("Timed out") + } + +} diff --git a/outlet.go b/outlet.go index ffc42a8..946c485 100644 --- a/outlet.go +++ b/outlet.go @@ -2,7 +2,6 @@ package main import ( "bufio" - "bytes" "fmt" "github.com/daviddengcn/go-colortext" "io" @@ -11,19 +10,11 @@ import ( ) type OutletFactory struct { - Outlets map[string]*Outlet Padding int -} -type Outlet struct { - Name string - Color ct.Color - IsError bool - Factory *OutletFactory + sync.Mutex } -var mx sync.Mutex - var colors = []ct.Color{ ct.Cyan, ct.Yellow, @@ -34,50 +25,45 @@ var colors = []ct.Color{ } func NewOutletFactory() (of *OutletFactory) { - of = new(OutletFactory) - of.Outlets = make(map[string]*Outlet) - return + return new(OutletFactory) } -func (o *Outlet) Write(b []byte) (num int, err error) { - mx.Lock() - defer mx.Unlock() - scanner := bufio.NewScanner(bytes.NewReader(b)) +func (of *OutletFactory) LineReader(wg *sync.WaitGroup, name string, index int, r io.Reader, isError bool) { + defer wg.Done() + + color := colors[index%len(colors)] + + scanner := bufio.NewScanner(r) for scanner.Scan() { - formatter := fmt.Sprintf("%%-%ds | ", o.Factory.Padding) - ct.ChangeColor(o.Color, true, ct.None, false) - fmt.Printf(formatter, o.Name) - if o.IsError { - ct.ChangeColor(ct.Red, true, ct.None, true) - } else { - ct.ResetColor() - } - fmt.Println(scanner.Text()) - ct.ResetColor() + of.WriteLine(name, scanner.Text(), color, ct.None, isError) } - num = len(b) - return -} - -func ProcessOutput(w io.Writer, str string) { - w.Write([]byte(str)) -} - -func (of *OutletFactory) CreateOutlet(name string, index int, isError bool) *Outlet { - of.Outlets[name] = &Outlet{name, colors[index%len(colors)], isError, of} - return of.Outlets[name] } func (of *OutletFactory) SystemOutput(str string) { - ct.ChangeColor(ct.White, true, ct.None, false) - formatter := fmt.Sprintf("%%-%ds | ", of.Padding) - fmt.Printf(formatter, "forego") - ct.ResetColor() - fmt.Println(str) - ct.ResetColor() + of.WriteLine("forego", str, ct.White, ct.None, false) } func (of *OutletFactory) ErrorOutput(str string) { fmt.Printf("ERROR: %s\n", str) os.Exit(1) } + +// Write out a single coloured line +func (of *OutletFactory) WriteLine(left, right string, leftC, rightC ct.Color, isError bool) { + of.Lock() + defer of.Unlock() + + ct.ChangeColor(leftC, true, ct.None, false) + formatter := fmt.Sprintf("%%-%ds | ", of.Padding) + fmt.Printf(formatter, left) + + if isError { + ct.ChangeColor(ct.Red, true, ct.None, true) + } else { + ct.ResetColor() + } + fmt.Println(right) + if isError { + ct.ResetColor() + } +} diff --git a/process.go b/process.go index e107dca..ae1bcdd 100644 --- a/process.go +++ b/process.go @@ -1,62 +1,36 @@ package main import ( - "fmt" - "io" "os" "os/exec" + "syscall" ) type Process struct { Command string Env Env Interactive bool - Stdin io.Reader - Stdout io.Writer - Stderr io.Writer - Root string - cmd *exec.Cmd + *exec.Cmd } -func NewProcess(command string, env Env) (p *Process) { - p = new(Process) - p.Command = command - p.Env = env - p.Interactive = false - p.Stdin = os.Stdin - p.Stdout = os.Stdout - p.Stderr = os.Stderr - p.Root, _ = os.Getwd() - return -} - -func (p *Process) Running() bool { - return (p.cmd.Process != nil) -} - -func (p *Process) Pid() int { - return p.cmd.Process.Pid -} - -func (p *Process) Wait() { - p.cmd.Wait() -} - -func (p *Process) shellArgument() string { - if p.Interactive { - return "-ic" - } else { - return "-c" +func NewProcess(workdir, command string, env Env, interactive bool) (p *Process) { + argv := ShellInvocationCommand(interactive, workdir, command) + return &Process{ + command, env, interactive, exec.Command(argv[0], argv[1:]...), } } -func (p *Process) envAsArray() (env []string) { - for _, pair := range os.Environ() { - env = append(env, pair) +func (p *Process) Start() error { + p.Cmd.Env = p.Env.asArray() + p.PlatformSpecificInit() + return p.Cmd.Start() +} + +func (p *Process) Signal(signal syscall.Signal) error { + group, err := os.FindProcess(-1 * p.Process.Pid) + if err == nil { + err = group.Signal(signal) } - for name, val := range p.Env { - env = append(env, fmt.Sprintf("%s=%s", name, val)) - } - return + return err } diff --git a/run.go b/run.go index a230a44..d7f067d 100644 --- a/run.go +++ b/run.go @@ -24,19 +24,26 @@ func init() { } func runRun(cmd *Command, args []string) { + workDir, err := os.Getwd() + if err != nil { + handleError(err) + } if flagEnv == "" { - root, _ := os.Getwd() - flagEnv = filepath.Join(root, ".env") + flagEnv = filepath.Join(workDir, ".env") } env, err := ReadEnv(flagEnv) handleError(err) - ps := NewProcess(strings.Join(args, " "), env) - ps.Interactive = true + const interactive = true + ps := NewProcess(workDir, strings.Join(args, " "), env, interactive) ps.Stdin = os.Stdin ps.Stdout = os.Stdout ps.Stderr = os.Stderr - ps.Start() - ps.Wait() + + err = ps.Start() + handleError(err) + + err = ps.Wait() + handleError(err) } diff --git a/start.go b/start.go index e20d8f0..d63fd76 100644 --- a/start.go +++ b/start.go @@ -17,11 +17,6 @@ const shutdownGraceTime = 3 * time.Second var flagPort int var flagConcurrency string var flagRestart bool -var shutdownNow bool - -var processes = map[string]*Process{} -var shutdown_mutex = new(sync.Mutex) -var wg sync.WaitGroup var cmdStart = &Command{ Run: runStart, @@ -74,38 +69,109 @@ func parseConcurrency(value string) (map[string]int, error) { return concurrency, nil } -func startProcess(idx, procNum int, proc ProcfileEntry, env Env, of *OutletFactory) { - shutdown_mutex.Lock() - wg.Add(1) - port := flagPort + (idx * 100) - ps := NewProcess(proc.Command, env) - procName := strings.Join([]string{ - proc.Name, - strconv.FormatInt(int64(procNum+1), 10)}, ".") - processes[procName] = ps - ps.Env["PORT"] = strconv.Itoa(port) - ps.Root = filepath.Dir(flagProcfile) - ps.Stdin = nil - ps.Stdout = of.CreateOutlet(procName, idx, false) - ps.Stderr = of.CreateOutlet(procName, idx, true) - ps.Start() - of.SystemOutput(fmt.Sprintf("starting %s on port %d", procName, port)) - go func(proc ProcfileEntry, ps *Process) { - ps.Wait() +type Forego struct { + outletFactory *OutletFactory - if flagRestart && !shutdownNow { - delete(processes, proc.Name) - startProcess(idx, procNum, proc, env, of) - wg.Done() - return + teardown, teardownNow Barrier // signal shutting down + + wg sync.WaitGroup +} + +func (f *Forego) monitorInterrupt() { + handler := make(chan os.Signal, 1) + signal.Notify(handler, os.Interrupt) + + first := true + + for sig := range handler { + switch sig { + case os.Interrupt: + fmt.Println(" | ctrl-c detected") + + f.teardown.Fall() + if !first { + f.teardownNow.Fall() + } + first = false } + } +} - wg.Done() - delete(processes, proc.Name) - ShutdownProcesses(of) +func (f *Forego) startProcess(idx, procNum int, proc ProcfileEntry, env Env, of *OutletFactory) { + port := flagPort + (idx * 100) - }(proc, ps) - shutdown_mutex.Unlock() + const interactive = false + workDir := filepath.Dir(flagProcfile) + ps := NewProcess(workDir, proc.Command, env, interactive) + procName := fmt.Sprint(proc.Name, ".", procNum+1) + ps.Env["PORT"] = strconv.Itoa(port) + + ps.Stdin = nil + + stdout, err := ps.StdoutPipe() + if err != nil { + panic(err) + } + stderr, err := ps.StderrPipe() + if err != nil { + panic(err) + } + + pipeWait := new(sync.WaitGroup) + pipeWait.Add(2) + go of.LineReader(pipeWait, procName, idx, stdout, false) + go of.LineReader(pipeWait, procName, idx, stderr, true) + + of.SystemOutput(fmt.Sprintf("starting %s on port %d", procName, port)) + + finished := make(chan struct{}) // closed on process exit + + ps.Start() + + f.wg.Add(1) + go func() { + defer f.wg.Done() + defer close(finished) + pipeWait.Wait() + ps.Wait() + }() + + f.wg.Add(1) + go func() { + defer f.wg.Done() + + // Prevent goroutine from exiting before process has finished. + defer func() { <-finished }() + defer f.teardown.Fall() + + select { + case <-finished: + if flagRestart { + f.startProcess(idx, procNum, proc, env, of) + return + } + + case <-f.teardown.Barrier(): + // Forego tearing down + + if !osHaveSigTerm { + of.SystemOutput(fmt.Sprintf("Killing %s", procName)) + ps.Process.Kill() + return + } + + of.SystemOutput(fmt.Sprintf("sending SIGTERM to %s", procName)) + ps.SendSigTerm() + + // Give the process a chance to exit, otherwise kill it. + select { + case <-f.teardownNow.Barrier(): + of.SystemOutput(fmt.Sprintf("Killing %s", procName)) + ps.SendSigKill() + case <-finished: + } + } + }() } func runStart(cmd *Command, args []string) { @@ -127,19 +193,20 @@ func runStart(cmd *Command, args []string) { of := NewOutletFactory() of.Padding = pf.LongestProcessName() - handler := make(chan os.Signal, 1) - signal.Notify(handler, os.Interrupt) + f := &Forego{ + outletFactory: of, + } - go func() { - for sig := range handler { - switch sig { - case os.Interrupt: - shutdownNow = true - fmt.Println(" | ctrl-c detected") - go func() { ShutdownProcesses(of) }() - } - } - }() + go f.monitorInterrupt() + + // When teardown fires, start the grace timer + f.teardown.FallHook = func() { + go func() { + time.Sleep(shutdownGraceTime) + of.SystemOutput("Grace time expired") + f.teardownNow.Fall() + }() + } var singleton string = "" if len(args) > 0 { @@ -156,10 +223,12 @@ func runStart(cmd *Command, args []string) { } for i := 0; i < numProcs; i++ { if (singleton == "") || (singleton == proc.Name) { - startProcess(idx, i, proc, env, of) + f.startProcess(idx, i, proc, env, of) } } } - wg.Wait() + <-f.teardown.Barrier() + + f.wg.Wait() } diff --git a/unix.go b/unix.go index 55c8db6..1cf106b 100644 --- a/unix.go +++ b/unix.go @@ -1,50 +1,38 @@ -package main - // +build darwin freebsd linux netbsd openbsd +package main + import ( "fmt" - "os" - "os/exec" "path/filepath" "syscall" - "time" ) -func (p *Process) Start() { - command := []string{"/bin/bash", p.shellArgument(), fmt.Sprintf("source \"%s\" 2>/dev/null; %s", filepath.Join(p.Root, ".profile"), p.Command)} - p.cmd = exec.Command(command[0], command[1:]...) - p.cmd.Dir = p.Root - p.cmd.Env = p.envAsArray() - p.cmd.Stdin = p.Stdin - p.cmd.Stdout = p.Stdout - p.cmd.Stderr = p.Stderr +const osHaveSigTerm = true + +func ShellInvocationCommand(interactive bool, root, command string) []string { + shellArgument := "-c" + if interactive { + shellArgument = "-ic" + } + profile := filepath.Join(root, ".profile") + shellCommand := fmt.Sprintf("source \"%s\" 2>/dev/null; %s", profile, command) + return []string{"/bin/bash", shellArgument, shellCommand} + +} + +func (p *Process) PlatformSpecificInit() { if !p.Interactive { - p.cmd.SysProcAttr = &syscall.SysProcAttr{} - p.cmd.SysProcAttr.Setsid = true + p.SysProcAttr = &syscall.SysProcAttr{} + p.SysProcAttr.Setsid = true } - p.cmd.Start() + return } -func (p *Process) Signal(signal syscall.Signal) { - if p.Running() { - group, _ := os.FindProcess(-1 * p.Pid()) - group.Signal(signal) - } +func (p *Process) SendSigTerm() { + p.Signal(syscall.SIGTERM) } -func ShutdownProcesses(of *OutletFactory) { - shutdown_mutex.Lock() - of.SystemOutput("shutting down") - for name, ps := range processes { - of.SystemOutput(fmt.Sprintf("sending SIGTERM to %s", name)) - ps.Signal(syscall.SIGTERM) - } - go func() { - time.Sleep(shutdownGraceTime) - for name, ps := range processes { - of.SystemOutput(fmt.Sprintf("sending SIGKILL to %s", name)) - ps.Signal(syscall.SIGKILL) - } - }() +func (p *Process) SendSigKill() { + p.Signal(syscall.SIGKILL) } diff --git a/windows.go b/windows.go index ee14e3b..232c4bc 100644 --- a/windows.go +++ b/windows.go @@ -1,36 +1,26 @@ -package main - // +build windows +package main + import ( - "fmt" - "os" - "os/exec" "syscall" ) -func (p *Process) Start() { - command := []string{"cmd", "/C", p.Command} - p.cmd = exec.Command(command[0], command[1:]...) - p.cmd.Dir = p.Root - p.cmd.Env = p.envAsArray() - p.cmd.Stdin = p.Stdin - p.cmd.Stdout = p.Stdout - p.cmd.Stderr = p.Stderr - p.cmd.Start() +const osHaveSigTerm = false + +func ShellInvocationCommand(interactive bool, root, command string) []string { + return []string{"cmd", "/C", command} } -func (p *Process) Signal(signal syscall.Signal) { - group, _ := os.FindProcess(-1 * p.cmd.Process.Pid) - group.Signal(signal) +func (p *Process) PlatformSpecificInit() { + // NOP on windows for now. + return } -func ShutdownProcesses(of *OutletFactory) { - shutdown_mutex.Lock() - of.SystemOutput("shutting down") - for name, ps := range processes { - of.SystemOutput(fmt.Sprintf("terminating %s", name)) - ps.cmd.Process.Signal(os.Kill) - } - os.Exit(1) +func (p *Process) SendSigTerm() { + panic("SendSigTerm() not implemented on this platform") +} + +func (p *Process) SendSigKill() { + p.Signal(syscall.SIGKILL) }