Use a context parameter in (*Controller).Stop

The Stop method of the orch.Controller type accepted a time.Duration. It
would wait for a graceful shutdown for this amount of time. Then, it
would return an error if the graceful shutdown did not occur.

This commit replaces the time.Duration with a context.Context,
permitting more flexibility.

Author: Ben Reed

testctrl/cmd/orchtest/main.go

@@ -15,6 +15,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"log"
 	"os"
@@ -53,7 +54,10 @@ func main() {
 	if err := c.Start(); err != nil {
 		panic(err)
 	}
-	defer c.Stop(*timeout)
+
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), *timeout)
+	defer cancel()
+	defer c.Stop(shutdownCtx)
 
 	go func() {
 		for i := 0; i < *count; i++ {

testctrl/cmd/svc/main.go

@@ -15,6 +15,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"net"
@@ -92,7 +93,9 @@ func main() {
 		glog.Fatalf("unable to start orchestration controller: %v", err)
 	}
 
-	defer controller.Stop(*shutdownTimeout)
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), *shutdownTimeout)
+	defer cancel()
+	defer controller.Stop(shutdownCtx)
 
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
 	if err != nil {

testctrl/svc/orch/controller.go

@@ -15,6 +15,7 @@
 package orch
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"sync"
@@ -158,15 +159,16 @@ func (c *Controller) Start() error {
 	return nil
 }
 
-// Stop attempts to terminate all orchestration threads spawned by a call to Start. It waits for a
-// graceful shutdown until for a specified timeout.
+// Stop attempts to terminate all orchestration threads spawned by a call to
+// Start. It waits for a graceful shutdown until the context is cancelled.
 //
-// If the timeout is reached before shutdown, an improper shutdown will occur. This may result in
-// unpredictable states for running sessions and their resources. To signal these potential issues,
-// an error is returned when this occurs.
+// If the context is cancelled before a graceful shutdown, an error is returned.
+// This improper shutdown may result in unpredictable states. It should be
+// avoided if possible.
 //
-// If Start was not called prior to Stop, there will be no adverse effects and nil will be returned.
-func (c *Controller) Stop(timeout time.Duration) error {
+// If Start was not called prior to Stop, there will be no adverse effects and
+// nil will be returned.
+func (c *Controller) Stop(ctx context.Context) error {
 	defer c.watcher.Stop()
 
 	c.mux.Lock()
@@ -182,7 +184,7 @@ func (c *Controller) Stop(timeout time.Duration) error {
 	select {
 	case <-done:
 		glog.Infof("controller: executors safely exited")
-	case <-time.After(timeout):
+	case <-ctx.Done():
 		glog.Warning("controller: unable to wait for executors to safely exit, timed out")
 		return fmt.Errorf("executors did not safely exit before timeout")
 	}

testctrl/svc/orch/controller_test.go

@@ -1,6 +1,7 @@
 package orch
 
 import (
+	"context"
 	"errors"
 	"testing"
 	"time"
@@ -58,7 +59,7 @@ func TestControllerSchedule(t *testing.T) {
 
 			if tc.start {
 				controller.Start()
-				defer controller.Stop(0)
+				defer controller.Stop(context.TODO())
 			}
 
 			err := controller.Schedule(tc.session)
@@ -86,7 +87,7 @@ func TestControllerStart(t *testing.T) {
 	t.Run("sets running state", func(t *testing.T) {
 		controller, _ := NewController(fake.NewSimpleClientset(), nil, nil)
 		controller.Start()
-		defer controller.Stop(0)
+		defer controller.Stop(context.TODO())
 		if controller.Stopped() {
 			t.Errorf("Stopped unexpectedly returned true after starting controller")
 		}
@@ -196,7 +197,11 @@ func TestControllerStop(t *testing.T) {
 
 			go controller.loop()
 			time.Sleep(timeout)
-			err := controller.Stop(tc.stopTimeout)
+
+			ctx, cancel := context.WithTimeout(context.TODO(), tc.stopTimeout)
+			defer cancel()
+
+			err := controller.Stop(ctx)
 			if tc.shouldError && err == nil {
 				t.Errorf("executors unexpectedly finished before timeout")
 			} else if !tc.shouldError && err != nil {