Repository: honeydipper/honeydipper
Pull request #489 (2022-08-04 23:57:02)

feat(core): graceful shutdown by draining drivers
#### Description
Upon forceful reload, gracefully stop all drivers before shutdown the
daemon process. Adding a stage in config named `Drained`.

This is useful when we migrate the work from one node to another to allow daemon restarting.

#### This PR fixes the following issues
N/A


wyardley: 🎉 nice!
wyardley: btw, refactor is not a triggering type, so if you want this to be in the changelog, etc., you might consider considering it a `feat` instead.

Patch:
diff --git a/.github/workflows/check-semantic-commits.yml b/.github/workflows/check-semantic-commits.yml
index e88a4d1..cf88d6f 100644
--- a/.github/workflows/check-semantic-commits.yml
+++ b/.github/workflows/check-semantic-commits.yml
@@ -20,7 +20,8 @@ jobs:
           validateSingleCommit: true
           scopes: |
             core
-            database
+            drivers
+            pkg
             deps
             error
             release
diff --git a/drivers/cmd/webhook/main.go b/drivers/cmd/webhook/main.go
index d43eac5..9189228 100644
--- a/drivers/cmd/webhook/main.go
+++ b/drivers/cmd/webhook/main.go
@@ -27,7 +27,7 @@ import (
 )
 
 const (
-	// Timeout (in seconds) for RequestHeaderTimeout.
+	// RequestHeaderTimeoutSecs is the timeout (in seconds) for accepting incoming requests.
 	RequestHeaderTimeoutSecs = 20
 )
 
@@ -124,7 +124,7 @@ func startWebhook(m *dipper.Message) {
 	go func() {
 		log.Infof("[%s] start listening for webhook requests", driver.Service)
 		log.Infof("[%s] listener stopped: %+v", driver.Service, server.ListenAndServe())
-		if driver.State != "exit" && driver.State != "cold" {
+		if driver.State != "stopped" && driver.State != "cold" {
 			startWebhook(m)
 		}
 	}()
diff --git a/internal/config/config.go b/internal/config/config.go
index 5146856..b715911 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -30,6 +30,8 @@ const (
 	StageDiscovering
 	// StageServing means all config should be processed and serving.
 	StageServing
+	// StageDrained means that the service has finished all in-fly requests and no longer accepts new requests.
+	StageDrained
 )
 
 var (
@@ -39,6 +41,7 @@ var (
 		"Booting",
 		"Discovering",
 		"Serving",
+		"Drained",
 	}
 
 	// ErrConfigRollback happens when daemon decides to rollback during reload.
@@ -96,13 +99,15 @@ func (c *Config) ResetStage() {
 		}
 	}
 	//nolint:gomnd
-	c.StageWG = make([]*sync.WaitGroup, 3)
+	c.StageWG = make([]*sync.WaitGroup, 4)
 	c.StageWG[StageLoading] = &sync.WaitGroup{}
 	c.StageWG[StageLoading].Add(len(c.Services))
 	c.StageWG[StageBooting] = &sync.WaitGroup{}
 	c.StageWG[StageBooting].Add(len(c.Services))
 	c.StageWG[StageDiscovering] = &sync.WaitGroup{}
 	c.StageWG[StageDiscovering].Add(len(c.Services))
+	c.StageWG[StageServing] = &sync.WaitGroup{}
+	c.StageWG[StageServing].Add(len(c.Services))
 }
 
 // Bootstrap loads the configuration during daemon bootstrap.
diff --git a/internal/config/gitssh.go b/internal/config/gitssh.go
index f4a16d6..ba8e76f 100644
--- a/internal/config/gitssh.go
+++ b/internal/config/gitssh.go
@@ -15,7 +15,7 @@ import (
 	"gopkg.in/src-d/go-git.v4/plumbing/transport/ssh"
 )
 
-var currentSSHAuth map[string]transport.AuthMethod = map[string]transport.AuthMethod{}
+var currentSSHAuth = map[string]transport.AuthMethod{}
 
 // GetGitSSHAuth creates an AuthMethod to be used for various git operations.
 func GetGitSSHAuth(keyfile, keypassEnv string) transport.AuthMethod {
diff --git a/internal/driver/schema.go b/internal/driver/schema.go
index e9aadbd..18fca8c 100644
--- a/internal/driver/schema.go
+++ b/internal/driver/schema.go
@@ -21,4 +21,5 @@ const (
 	DriverReloading
 	DriverAlive
 	DriverFailed
+	DriverStopped
 )
diff --git a/internal/service/api.go b/internal/service/api.go
index 1d45f58..9739712 100644
--- a/internal/service/api.go
+++ b/internal/service/api.go
@@ -24,7 +24,7 @@ const (
 	// APIServerGracefulTimeout is the timeout in seconds for waiting for a http.Server to shutdown gracefully.
 	APIServerGracefulTimeout time.Duration = 10
 
-	// Timeout (in seconds) for RequestHeaderTimeout.
+	// RequestHeaderTimeoutSecs is the timeout (in seconds) for accepting incoming requests.
 	RequestHeaderTimeoutSecs = 20
 )
 
@@ -45,7 +45,6 @@ func StartAPI(cfg *config.Config) {
 	API.ServiceReload = reloadAPI
 	API.DiscoverFeatures = APIFeatures
 	API.addResponder("eventbus:api", handleAPIMessage)
-	Services["api"] = API
 	APIRequestStore = api.NewStore(API)
 	API.start()
 }
diff --git a/internal/service/engine.go b/internal/service/engine.go
index d804ce5..ecea41b 100644
--- a/internal/service/engine.go
+++ b/internal/service/engine.go
@@ -51,7 +51,6 @@ func (h *WorkflowHelper) GetConfig() *config.Config {
 // StartEngine Starts the engine service.
 func StartEngine(cfg *config.Config) {
 	engine = NewService(cfg, "engine")
-	Services["engine"] = engine
 	helper := &WorkflowHelper{engine: engine}
 	sessionStore = workflow.NewSessionStore(helper)
 
diff --git a/internal/service/operator.go b/internal/service/operator.go
index 69ff4b4..5a26a40 100644
--- a/internal/service/operator.go
+++ b/internal/service/operator.go
@@ -28,7 +28,6 @@ var (
 func StartOperator(cfg *config.Config) {
 	operator = NewService(cfg, "operator")
 	operator.Route = operatorRoute
-	Services["operator"] = operator
 	operator.start()
 }
 
diff --git a/internal/service/receiver.go b/internal/service/receiver.go
index 0f8be91..80114dd 100644
--- a/internal/service/receiver.go
+++ b/internal/service/receiver.go
@@ -25,7 +25,6 @@ func StartReceiver(cfg *config.Config) {
 	receiver.Route = receiverRoute
 	receiver.DiscoverFeatures = ReceiverFeatures
 	receiver.EmitMetrics = receiverMetrics
-	Services["receiver"] = receiver
 	setupReceiverAPIs()
 	receiver.start()
 }
diff --git a/internal/service/service.go b/internal/service/service.go
index 4eb55da..ecb69bb 100644
--- a/internal/service/service.go
+++ b/internal/service/service.go
@@ -74,12 +74,16 @@ type Service struct {
 	APIs               map[string]func(*api.Response)
 	ResponseFactory    *api.ResponseFactory
 	healthy            bool
+	drainingGroup      *sync.WaitGroup
 }
 
 var (
 	// Services holds a catalog of running services in this daemon process.
 	Services = map[string]*Service{}
 
+	// masterService the service responsible for reloading and daemon lifecycle.
+	masterService *Service
+
 	// ErrServiceError indicates error condition when rendering service.
 	ErrServiceError = errors.New("service error")
 )
@@ -96,6 +100,7 @@ func NewService(cfg *config.Config, name string) *Service {
 	svc.RPCCallerBase.Init(svc, "rpc", "call")
 
 	svc.responders["state:cold"] = []MessageResponder{coldReloadDriverRuntime}
+	svc.responders["state:stopped"] = []MessageResponder{handleDriverStop}
 	svc.responders["rpc:call"] = []MessageResponder{handleRPCCall}
 	svc.responders["rpc:return"] = []MessageResponder{handleRPCReturn}
 	svc.responders["broadcast:reload"] = []MessageResponder{handleReload}
@@ -104,6 +109,11 @@ func NewService(cfg *config.Config, name string) *Service {
 	svc.ResponseFactory = api.NewResponseFactory()
 	svc.APIs = map[string]func(*api.Response){}
 
+	if len(Services) == 0 {
+		masterService = svc
+	}
+	Services[name] = svc
+
 	return svc
 }
 
@@ -639,6 +649,7 @@ func loadFailedDriverRuntime(d *driver.Runtime, count int) {
 		})
 	}
 
+	dipper.Logger.Warningf("[%s] start loading/reloading driver %s", s.name, driverName)
 	retry := func() {
 		dipper.Logger.Warningf("[%s] failed to load/reload driver %s attempt %d", s.name, driverName, count)
 		if count < DriverRetryCount {
@@ -708,24 +719,37 @@ func handleReload(from *driver.Runtime, m *dipper.Message) {
 	if ok && daemonID != dipper.GetIP() {
 		return
 	}
-	var min string
-	for min = range Services {
-		if from.Service > min {
-			break
+
+	m = dipper.DeserializePayload(m)
+	force := false
+	if f, ok := dipper.GetMapData(m.Payload, "force"); ok && dipper.IsTruthy(f) {
+		force = true
+	}
+
+	if !force {
+		if from.Service == masterService.name {
+			dipper.Logger.Warningf("[%s] reload config on broadcast reload message", from.Service)
+			go masterService.config.Refresh()
 		}
+
+		return
 	}
-	if from.Service <= min {
-		m := dipper.DeserializePayload(m)
-		go func() {
-			time.Sleep(time.Second)
-			if force, ok := dipper.GetMapDataStr(m.Payload, "force"); ok && (force == "yes" || force == "true") {
-				daemon.ShutDown()
-				os.Exit(0)
-			} else {
-				dipper.Logger.Infof("[%s] reload config on broadcast reload message", min)
-				Services[min].config.Refresh()
-			}
-		}()
+
+	go func() {
+		<-time.After(time.Second)
+		dipper.Logger.Warningf("[%s] quiting on broadcast force reload message", from.Service)
+		Services[from.Service].Drain()
+		if from.Service == masterService.name {
+			daemon.ShutDown()
+			os.Exit(0)
+		}
+	}()
+}
+
+func handleDriverStop(from *driver.Runtime, m *dipper.Message) {
+	if from.State != driver.DriverStopped {
+		from.State = driver.DriverStopped
+		Services[from.Service].drainingGroup.Done()
 	}
 }
 
@@ -784,3 +808,35 @@ func (s *Service) metricsLoop() {
 		time.Sleep(time.Minute)
 	}
 }
+
+// Drain stops the service from accepting new requests but allow the remaining requests to complete.
+func (s *Service) Drain() {
+	s.healthy = false
+
+	cnt := 0
+	s.driverLock.Lock()
+	for _, d := range s.driverRuntimes {
+		if d.State != driver.DriverFailed && d.State != driver.DriverStopped {
+			cnt++
+		}
+	}
+	s.driverLock.Unlock()
+
+	if cnt > 0 {
+		s.drainingGroup = &sync.WaitGroup{}
+		s.drainingGroup.Add(cnt)
+
+		for _, d := range s.driverRuntimes {
+			if d.State != driver.DriverFailed && d.State != driver.DriverStopped {
+				d.SendMessage(&dipper.Message{
+					Channel: "command",
+					Subject: "stop",
+				})
+			}
+		}
+
+		dipper.WaitGroupWaitTimeout(s.drainingGroup, time.Second)
+	}
+
+	s.config.AdvanceStage(s.name, config.StageDrained)
+}
diff --git a/pkg/dipper/driver.go b/pkg/dipper/driver.go
index be7ba83..b513dd6 100644
--- a/pkg/dipper/driver.go
+++ b/pkg/dipper/driver.go
@@ -148,12 +148,12 @@ func (d *Driver) start(msg *Message) {
 }
 
 func (d *Driver) stop(msg *Message) {
-	d.State = "exit"
+	d.State = "stopped"
 	if d.Stop != nil {
 		d.Stop(msg)
 	}
 	d.Ping(msg)
-	Logger.Fatalf("[%s] quiting on daemon request", d.Service)
+	Logger.Warningf("[%s] quiting on daemon request", d.Service)
 }
 
 // SendMessage : send a prepared message to daemon.
diff --git a/pkg/dipper/sync.go b/pkg/dipper/sync.go
index 51e57b4..c11ec0d 100644
--- a/pkg/dipper/sync.go
+++ b/pkg/dipper/sync.go
@@ -7,7 +7,10 @@
 // Package dipper is a library used for developing drivers for Honeydipper.
 package dipper
 
-import "sync"
+import (
+	"sync"
+	"time"
+)
 
 // WaitGroupDone is a way to safely decrement the counter for a WaitGroup.
 func WaitGroupDone(wg *sync.WaitGroup) (ok bool) {
@@ -42,3 +45,18 @@ func WaitGroupWait(wg *sync.WaitGroup) {
 
 	wg.Wait()
 }
+
+// WaitGroupWaitTimeout is a wrapper for safely waiting for a WaitGroup with a timeout.
+func WaitGroupWaitTimeout(wg *sync.WaitGroup, t time.Duration) {
+	allDone := make(chan interface{})
+	go func() {
+		WaitGroupWait(wg)
+		close(allDone)
+	}()
+
+	select {
+	case <-time.After(t):
+		WaitGroupDoneAll(wg)
+	case <-allDone:
+	}
+}


Test patch:
diff --git a/cmd/honeydipper/integration_test.go b/cmd/honeydipper/integration_test.go
index c5b603a..6ecc8cc 100644
--- a/cmd/honeydipper/integration_test.go
+++ b/cmd/honeydipper/integration_test.go
@@ -18,6 +18,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -40,6 +41,7 @@ func TestIntegrationStart(t *testing.T) {
 	t.Run("checking processes", intTestProcesses)
 	t.Run("checking API calls", intTestMakingAPICall)
 	t.Run("checking crashed driver", intTestDriverCrash)
+	t.Run("checking draining drivers", intTestDrain)
 }
 
 func intTestDaemonStartup(t *testing.T) {
@@ -172,6 +174,18 @@ func intTestMakingAPICall(t *testing.T) {
 	assert.Equalf(t, 200, resp.StatusCode, "api call should succeed with correct creds")
 }
 
+func intTestDrain(t *testing.T) {
+	waiting := (int32)(len(service.Services))
+	for _, s := range service.Services {
+		go func(s *service.Service) {
+			s.Drain()
+			atomic.AddInt32(&waiting, -1)
+		}(s)
+	}
+
+	assert.Eventually(t, func() bool { return waiting == 0 }, time.Second*3, time.Millisecond*5, "All service should reach drained stage eventually.")
+}
+
 func intTestDaemonShutdown(t *testing.T) {
 	graceful := make(chan bool)
 	go func() {
diff --git a/pkg/dipper/sync_test.go b/pkg/dipper/sync_test.go
new file mode 100644
index 0000000..f7065b1
--- /dev/null
+++ b/pkg/dipper/sync_test.go
@@ -0,0 +1,42 @@
+// Copyright 2022 PayPal Inc.
+
+// This Source Code Form is subject to the terms of the MIT License.
+// If a copy of the MIT License was not distributed with this file,
+// you can obtain one at https://mit-license.org/.
+
+//go:build !integration
+// +build !integration
+
+package dipper
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWaitGroupWaitTimeout(t *testing.T) {
+	subj := &sync.WaitGroup{}
+	subj.Add(3)
+	done := false
+	go func() {
+		WaitGroupWaitTimeout(subj, time.Millisecond*10)
+		done = true
+	}()
+	assert.Eventually(t, func() bool { return done }, time.Second, time.Millisecond*5, "WaitGroupWaitTimeout should eventually quit.")
+}
+
+func TestWaitGroupWait(t *testing.T) {
+	subj := &sync.WaitGroup{}
+	subj.Add(3)
+	done := false
+	go func() {
+		WaitGroupDone(subj)
+		WaitGroupWait(subj)
+		done = true
+	}()
+	assert.Never(t, func() bool { return done }, time.Second, time.Millisecond*5, "WaitGroupWait should never quit until the wait group is done.")
+	WaitGroupDoneAll(subj)
+}