oversight: improve handling of tree termination by trying to wait for children goroutines to stop before returning control

Author: ucirello

child_process.go

@@ -133,10 +133,15 @@ func Transient() Restart { return func(err error) bool { return err != nil } }
 // are signaled to stop.
 type Shutdown func() (context.Context, context.CancelFunc)
 
+type shutdownContextValue string
+
+var detachableContext = shutdownContextValue("detachable")
+
 // Infinity will wait until the process naturally dies.
 func Infinity() Shutdown {
 	return func() (context.Context, context.CancelFunc) {
-		return context.WithCancel(context.Background())
+		ctx := context.WithValue(context.Background(), detachableContext, false)
+		return context.WithCancel(ctx)
 	}
 }
 
@@ -148,6 +153,7 @@ const DefaultChildProcessTimeout = 5 * time.Second
 // detaching from the winding process.
 func Timeout(d time.Duration) Shutdown {
 	return func() (context.Context, context.CancelFunc) {
-		return context.WithTimeout(context.Background(), d)
+		ctx := context.WithValue(context.Background(), detachableContext, true)
+		return context.WithTimeout(ctx, d)
 	}
 }

tree.go

@@ -66,13 +66,14 @@ type Tree struct {
 	stopped        chan struct{}
 
 	// semaphore must be held when adding/deleting dynamic processes
-	semaphore      sync.Mutex
-	strategy       Strategy
-	maxR           int
-	maxT           time.Duration
-	children       map[string]childProcess // map of children name to child process
-	childrenOrder  []string
-	processChanged chan struct{} // indicates that some change to process slice has been made
+	semaphore         sync.Mutex
+	strategy          Strategy
+	maxR              int
+	maxT              time.Duration
+	children          map[string]childProcess // map of children name to child process
+	childrenOrder     []string
+	childrenWaitGroup sync.WaitGroup
+	processChanged    chan struct{} // indicates that some change to process slice has been made
 
 	logger Logger
 
@@ -251,6 +252,7 @@ func (t *Tree) Start(rootCtx context.Context) error {
 	if err := t.err(); err != nil {
 		return err
 	}
+	defer t.childrenWaitGroup.Wait()
 	ctx, cancel := context.WithCancel(rootCtx)
 	defer cancel()
 	t.gracefulCancel = cancel
@@ -348,7 +350,14 @@ func (t *Tree) handleTreeChanges(ctx context.Context, cancel context.CancelFunc)
 
 func (t *Tree) startChildProcess(ctx context.Context, processID int, p *ChildProcessSpecification, startSemaphore <-chan struct{}) {
 	childCtx, childWg, procState := t.plugStop(ctx, processID, p)
+	detachable := childCtx.Value(detachableContext) == true
+	if !detachable {
+		t.childrenWaitGroup.Add(1)
+	}
 	go func() {
+		if !detachable {
+			defer t.childrenWaitGroup.Done()
+		}
 		defer childWg.Done()
 		<-startSemaphore
 		t.logger.Println(p.Name, "child started")
@@ -369,13 +378,16 @@ func (t *Tree) startChildProcess(ctx context.Context, processID int, p *ChildPro
 type oversightValue string
 
 func (t *Tree) plugStop(ctx context.Context, processID int, p *ChildProcessSpecification) (context.Context, *sync.WaitGroup, *state) {
-	childCtx, childCancel := context.WithCancel(context.WithValue(ctx, oversightValue("name"), p.Name))
+	stopCtx, stopCancel := p.Shutdown()
+	baseCtx := ctx
+	baseCtx = context.WithValue(baseCtx, oversightValue("name"), p.Name)
+	baseCtx = context.WithValue(baseCtx, detachableContext, stopCtx.Value(detachableContext))
+	childCtx, childCancel := context.WithCancel(baseCtx)
 	var childWg sync.WaitGroup
 	childWg.Add(1)
 	childProc := t.children[p.Name]
 	childProc.state.setRunning(func() {
 		t.logger.Println(p.Name, "stopping")
-		stopCtx, stopCancel := p.Shutdown()
 		defer stopCancel()
 		wgComplete := make(chan struct{})
 		childCancel()

tree_test.go

@@ -1063,10 +1063,43 @@ func TestTree_shutdownOrder(t *testing.T) {
 }
 
 func TestPanicDoubleStart(t *testing.T) {
+	t.Parallel()
 	var tree oversight.Tree
 	oversight.NeverHalt()(&tree)
 	ctx, cancel := context.WithCancel(context.Background())
 	tree.Add(oversight.ChildProcessSpecification{Name: "child", Start: func(context.Context) error { cancel(); return nil }})
 	tree.Start(ctx)
 	tree.Start(ctx)
 }
+
+func TestWaitAfterStart(t *testing.T) {
+	t.Parallel()
+	tree := oversight.New(oversight.NeverHalt())
+	var (
+		mu    sync.Mutex
+		count int
+	)
+	ctx, cancel := context.WithCancel(context.Background())
+	for i := 0; i < 10; i++ {
+		mu.Lock()
+		count++
+		mu.Unlock()
+		tree.Add(oversight.ChildProcessSpecification{
+			Start: func(ctx context.Context) error {
+				<-ctx.Done()
+				mu.Lock()
+				count--
+				mu.Unlock()
+				return nil
+			},
+			Restart: oversight.Temporary(),
+		})
+	}
+	time.AfterFunc(1*time.Second, cancel)
+	if err := tree.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+	if count != 0 {
+		t.Fatal("run returned before all children returned")
+	}
+}