Merge pull request #27 from pwaller/process-refactor

Big process refactor

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
+}

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
+}

fixtures/writey/Procfile

@@ -0,0 +1,2 @@
+writey1: writey
+writey2: writey

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")
+	}
+
+}

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()
+	}
+}

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
 }

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)
 }

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()
 }

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)
 }

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)
 }