From 37aac8a97eaea61d07a3d05534b77d39d79d488d Mon Sep 17 00:00:00 2001
From: Peter Waller <peter@scraperwiki.com>
Date: Thu, 22 May 2014 17:40:09 +0100
Subject: [PATCH] Refactor to use Barrier

---
 barrier.go | 37 +++++++++++++++++++++++++++++++++++++
 start.go   | 39 ++++++++++++++++-----------------------
 2 files changed, 53 insertions(+), 23 deletions(-)
 create mode 100644 barrier.go

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/start.go b/start.go
index e836af1..d63fd76 100644
--- a/start.go
+++ b/start.go
@@ -72,36 +72,26 @@ func parseConcurrency(value string) (map[string]int, error) {
 type Forego struct {
 	outletFactory *OutletFactory
 
-	shutdown    sync.Once     // Closes teardown exactly once
-	teardown    chan struct{} // barrier: closed when shutting down
-	teardownNow chan struct{} // barrier: second CTRL-C. More urgent.
+	teardown, teardownNow Barrier // signal shutting down
 
 	wg sync.WaitGroup
 }
 
-func (f *Forego) SignalShutdown() {
-	f.shutdown.Do(func() {
-		f.outletFactory.SystemOutput("shutting down")
-		close(f.teardown)
-	})
-}
-
 func (f *Forego) monitorInterrupt() {
 	handler := make(chan os.Signal, 1)
 	signal.Notify(handler, os.Interrupt)
 
 	first := true
-	var once sync.Once
 
 	for sig := range handler {
 		switch sig {
 		case os.Interrupt:
 			fmt.Println("      | ctrl-c detected")
 
+			f.teardown.Fall()
 			if !first {
-				once.Do(func() { close(f.teardownNow) })
+				f.teardownNow.Fall()
 			}
-			f.SignalShutdown()
 			first = false
 		}
 	}
@@ -152,7 +142,7 @@ func (f *Forego) startProcess(idx, procNum int, proc ProcfileEntry, env Env, of
 
 		// Prevent goroutine from exiting before process has finished.
 		defer func() { <-finished }()
-		defer f.SignalShutdown()
+		defer f.teardown.Fall()
 
 		select {
 		case <-finished:
@@ -161,7 +151,7 @@ func (f *Forego) startProcess(idx, procNum int, proc ProcfileEntry, env Env, of
 				return
 			}
 
-		case <-f.teardown:
+		case <-f.teardown.Barrier():
 			// Forego tearing down
 
 			if !osHaveSigTerm {
@@ -175,10 +165,7 @@ func (f *Forego) startProcess(idx, procNum int, proc ProcfileEntry, env Env, of
 
 			// Give the process a chance to exit, otherwise kill it.
 			select {
-			case <-time.After(shutdownGraceTime):
-				of.SystemOutput(fmt.Sprintf("Killing %s", procName))
-				ps.SendSigKill()
-			case <-f.teardownNow:
+			case <-f.teardownNow.Barrier():
 				of.SystemOutput(fmt.Sprintf("Killing %s", procName))
 				ps.SendSigKill()
 			case <-finished:
@@ -208,13 +195,19 @@ func runStart(cmd *Command, args []string) {
 
 	f := &Forego{
 		outletFactory: of,
-
-		teardown:    make(chan struct{}),
-		teardownNow: make(chan struct{}),
 	}
 
 	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 {
 		singleton = args[0]
@@ -235,7 +228,7 @@ func runStart(cmd *Command, args []string) {
 		}
 	}
 
-	<-f.teardown
+	<-f.teardown.Barrier()
 
 	f.wg.Wait()
 }