Merge pull request #34 from justenwalker/stop-hooks

Add preStop and postStop hooks

README.md

@@ -46,6 +46,8 @@ The format of the JSON file configuration is as follows:
   "consul": "consul:8500",
   "onStart": "/opt/containerbuddy/onStart-script.sh",
   "stopTimeout": 5,
+  "preStop": "/opt/containerbuddy/preStop-script.sh",
+  "postStop": "/opt/containerbuddy/postStop-script.sh",
   "services": [
     {
       "name": "app",
@@ -90,6 +92,8 @@ Other fields:
 
 - `consul` is the hostname and port of the Consul discovery service.
 - `onStart` is the executable (and its arguments) that will be called immediately prior to starting the shimmed application. This field is optional. If the `onStart` handler returns a non-zero exit code, Containerbuddy will exit.
+- `preStop` is the executable (and its arguments) that will be called immediately **before** the shimmed application exits. This field is optional. Containerbuddy will wait until this program exits before terminating the shimmed application.
+- `postStop` is the executable (and its arguments) that will be called immediately **after** the shimmed application exits. This field is optional. If the `postStop` handler returns a non-zero exit code, Containerbuddy will exit with this code rather than the application's exit code.
 - `stopTimeout` Optional amount of time in seconds to wait before killing the application. (defaults to `5`). Providing `-1` will kill the application immediately.
 
 *Note that if you're using `curl` to check HTTP endpoints for health checks, that it doesn't return a non-zero exit code on 404s or similar failure modes by default. Use the `--fail` flag for curl if you need to catch those cases.*

src/containerbuddy/config.go

@@ -20,8 +20,9 @@ var (
 type Config struct {
 	Consul       string `json:"consul,omitempty"`
 	OnStart      string `json:"onStart"`
+	PreStop      string `json:"preStop"`
+	PostStop     string `json:"postStop"`
 	StopTimeout  int    `json:"stopTimeout"`
-	onStartArgs  []string
 	Command      *exec.Cmd
 	QuitChannels []chan bool
 	Services     []*ServiceConfig `json:"services"`
@@ -37,7 +38,6 @@ type ServiceConfig struct {
 	TTL              int      `json:"ttl"`
 	Interfaces       []string `json:"interfaces"`
 	discoveryService DiscoveryService
-	healthArgs       []string
 	ipAddress        string
 }
 
@@ -46,7 +46,6 @@ type BackendConfig struct {
 	Poll             int    `json:"poll"` // time in seconds
 	OnChangeExec     string `json:"onChange"`
 	discoveryService DiscoveryService
-	onChangeArgs     []string
 	lastState        interface{}
 }
 
@@ -123,18 +122,14 @@ func loadConfig() (*Config, error) {
 		config.StopTimeout = defaultStopTimeout
 	}
 
-	config.onStartArgs = strings.Split(config.OnStart, " ")
-
 	for _, backend := range config.Backends {
 		backend.discoveryService = discovery
-		backend.onChangeArgs = strings.Split(backend.OnChangeExec, " ")
 	}
 
 	hostname, _ := os.Hostname()
 	for _, service := range config.Services {
 		service.Id = fmt.Sprintf("%s-%s", service.Name, hostname)
 		service.discoveryService = discovery
-		service.healthArgs = strings.Split(service.HealthCheckExec, " ")
 		if service.ipAddress, err = getIp(service.Interfaces); err != nil {
 			return nil, err
 		}

src/containerbuddy/main.go

@@ -5,6 +5,7 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"strings"
 	"syscall"
 	"time"
 )
@@ -17,23 +18,19 @@ func main() {
 	}
 
 	// Run the onStart handler, if any, and exit if it returns an error
-	if config.OnStart != "" {
-		code, err := run(config.onStartArgs)
-		if err != nil {
-			log.Println(err)
-			os.Exit(code)
-		}
+	if onStartCode, err := run(config.OnStart); err != nil {
+		os.Exit(onStartCode)
 	}
 
 	// Set up signal handler for placing instance into maintenance mode
 	handleSignals(config)
 
 	var quit []chan bool
 	for _, backend := range config.Backends {
-		quit = append(quit, poll(backend, checkForChanges, backend.onChangeArgs))
+		quit = append(quit, poll(backend, checkForChanges, backend.OnChangeExec))
 	}
 	for _, service := range config.Services {
-		quit = append(quit, poll(service, checkHealth, service.healthArgs))
+		quit = append(quit, poll(service, checkHealth, service.HealthCheckExec))
 	}
 	config.QuitChannels = quit
 
@@ -54,6 +51,10 @@ func main() {
 		if err != nil {
 			log.Println(err)
 		}
+		// Run the PostStop handler, if any, and exit if it returns an error
+		if postStopCode, err := run(config.PostStop); err != nil {
+			os.Exit(postStopCode)
+		}
 		os.Exit(code)
 	}
 
@@ -62,19 +63,19 @@ func main() {
 	select {}
 }
 
-type pollingFunc func(Pollable, []string)
+type pollingFunc func(Pollable, string)
 
 // Every `pollTime` seconds, run the `pollingFunc` function.
 // Expect a bool on the quit channel to stop gracefully.
-func poll(config Pollable, fn pollingFunc, args []string) chan bool {
+func poll(config Pollable, fn pollingFunc, command string) chan bool {
 	ticker := time.NewTicker(time.Duration(config.PollTime()) * time.Second)
 	quit := make(chan bool)
 	go func() {
 		for {
 			select {
 			case <-ticker.C:
 				if !inMaintenanceMode() {
-					fn(config, args)
+					fn(config, command)
 				}
 			case <-quit:
 				return
@@ -87,19 +88,19 @@ func poll(config Pollable, fn pollingFunc, args []string) chan bool {
 // Implements `pollingFunc`; args are the executable we use to check the
 // application health and its arguments. If the error code on that exectable is
 // 0, we write a TTL health check to the health check store.
-func checkHealth(pollable Pollable, args []string) {
+func checkHealth(pollable Pollable, command string) {
 	service := pollable.(*ServiceConfig) // if we pass a bad type here we crash intentionally
-	if code, _ := run(args); code == 0 {
+	if code, _ := run(command); code == 0 {
 		service.SendHeartbeat()
 	}
 }
 
 // Implements `pollingFunc`; args are the executable we run if the values in
 // the central store have changed since the last run.
-func checkForChanges(pollable Pollable, args []string) {
+func checkForChanges(pollable Pollable, command string) {
 	backend := pollable.(*BackendConfig) // if we pass a bad type here we crash intentionally
 	if backend.CheckForUpstreamChanges() {
-		run(args)
+		run(command)
 	}
 }
 
@@ -129,6 +130,14 @@ func executeAndWait(cmd *exec.Cmd) (int, error) {
 
 // Runs an arbitrary string of arguments as an executable and its arguments.
 // Returns the exit code and error message (if any).
-func run(args []string) (int, error) {
-	return executeAndWait(getCmd(args))
+func run(command string) (int, error) {
+	if command != "" {
+		args := strings.Split(command, " ")
+		code, err := executeAndWait(getCmd(args))
+		if err != nil {
+			log.Println(err)
+		}
+		return code, err
+	}
+	return 0, nil
 }

src/containerbuddy/main_test.go

@@ -9,17 +9,17 @@ import (
 // a closed channel immediately as expected and gracefully.
 func TestPoll(t *testing.T) {
 	service := &ServiceConfig{Poll: 1}
-	quit := poll(service, func(service Pollable, args []string) {
+	quit := poll(service, func(service Pollable, command string) {
 		time.Sleep(5 * time.Second)
 		t.Errorf("We should never reach this code because the channel should close.")
 		return
-	}, []string{"exec", "arg1"})
+	}, "exec arg1")
 	close(quit)
 }
 
 func TestRunSuccess(t *testing.T) {
-	args := []string{"/root/examples/test/test.sh", "doStuff", "--debug"}
-	if exitCode, _ := run(args); exitCode != 0 {
+	command := "/root/examples/test/test.sh doStuff --debug"
+	if exitCode, _ := run(command); exitCode != 0 {
 		t.Errorf("Expected exit code 0 but got %d", exitCode)
 	}
 }
@@ -30,8 +30,14 @@ func TestRunFailed(t *testing.T) {
 			t.Errorf("Expected panic but did not.")
 		}
 	}()
-	args := []string{"/root/examples/test/test.sh", "failStuff", "--debug"}
-	if exitCode, _ := run(args); exitCode != 255 {
+	command := "/root/examples/test/test.sh failStuff --debug"
+	if exitCode, _ := run(command); exitCode != 255 {
 		t.Errorf("Expected exit code 255 but got %d", exitCode)
 	}
 }
+
+func TestRunNothing(t *testing.T) {
+	if code, err := run(""); code != 0 || err != nil {
+		t.Errorf("Expected exit (0,nil) but got (%d,%s)", code, err)
+	}
+}

src/containerbuddy/signals.go

@@ -42,6 +42,9 @@ func terminate(config *Config) {
 		service.Deregister()
 	})
 
+	// Run and wait for preStop command to exit
+	run(config.PreStop)
+
 	cmd := config.Command
 	if cmd == nil || cmd.Process == nil {
 		// Not managing the process, so don't do anything