simplify startup by using the new wg.Go feature (#65)

* simplify code by using the new wg.Go feature

* add more test coverage

Author: robbyt

examples/http/main_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"io"
 	"log/slog"
 	"net/http"
@@ -40,8 +41,15 @@ func TestRunServer(t *testing.T) {
 		errCh <- sv.Run()
 	}()
 
-	// Give the server a moment to start
-	time.Sleep(100 * time.Millisecond)
+	// Wait for the server to be ready by checking if it responds to requests
+	assert.Eventually(t, func() bool {
+		resp, err := http.Get("http://localhost:8080/status")
+		if err != nil {
+			return false
+		}
+		defer func() { assert.NoError(t, resp.Body.Close()) }()
+		return resp.StatusCode == http.StatusOK
+	}, 2*time.Second, 50*time.Millisecond, "Server should become ready")
 
 	// Make a request to the server
 	resp, err := http.Get("http://localhost:8080/status")
@@ -56,12 +64,14 @@ func TestRunServer(t *testing.T) {
 	// Stop the supervisor
 	sv.Shutdown()
 
-	// Check that Run() didn't return an error
+	// Wait for Run() to complete and check the result
 	select {
 	case err := <-errCh:
-		require.NoError(t, err, "Run() should not return an error")
-	case <-time.After(100 * time.Millisecond):
-		// This is expected - the server is still running
+		if err != nil && !errors.Is(err, context.Canceled) {
+			require.NoError(t, err, "Run() should not return an error")
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Run() should have completed within timeout")
 	}
 }
 

supervisor/reload.go

@@ -25,7 +25,6 @@ func (p *PIDZero) ReloadAll() {
 // and calls the reload method on all reloadable services. This will also prevent
 // multiple reloads from happening concurrently.
 func (p *PIDZero) startReloadManager() {
-	defer p.wg.Done()
 	p.logger.Debug("Starting reload manager...")
 
 	// iterate all the runnables, and find the ones that are can send reload notifications

supervisor/shutdown.go

@@ -5,7 +5,6 @@ import "sync"
 // startShutdownManager starts goroutines to listen for shutdown notifications
 // from any runnables that implement ShutdownSender. It blocks until the context is done.
 func (p *PIDZero) startShutdownManager() {
-	defer p.wg.Done()
 	p.logger.Debug("Starting shutdown manager...")
 
 	var shutdownWg sync.WaitGroup

supervisor/shutdown_test.go

@@ -34,9 +34,8 @@ func TestPIDZero_StartShutdownManager_TriggersShutdown(t *testing.T) {
 	pidZero, err := New(WithContext(supervisorCtx), WithRunnables(mockService))
 	require.NoError(t, err)
 
-	// Start the shutdown manager in a goroutine
-	pidZero.wg.Add(1)
-	go pidZero.startShutdownManager()
+	// Start the shutdown manager using wg.Go
+	pidZero.wg.Go(pidZero.startShutdownManager)
 
 	time.Sleep(200 * time.Millisecond)
 
@@ -81,9 +80,8 @@ func TestPIDZero_StartShutdownManager_ContextCancel(t *testing.T) {
 	pidZero, err := New(WithContext(ctx), WithRunnables(mockService))
 	require.NoError(t, err)
 
-	// Start the shutdown manager in a goroutine
-	pidZero.wg.Add(1)
-	go pidZero.startShutdownManager()
+	// Start the shutdown manager using wg.Go
+	pidZero.wg.Go(pidZero.startShutdownManager)
 
 	// Give the manager a moment to start its internal listener goroutine
 	time.Sleep(100 * time.Millisecond) // Increased sleep duration
@@ -128,9 +126,8 @@ func TestPIDZero_StartShutdownManager_NoSenders(t *testing.T) {
 	pidZero, err := New(WithContext(ctx), WithRunnables(nonSenderRunnable))
 	require.NoError(t, err)
 
-	// Start the shutdown manager in a goroutine
-	pidZero.wg.Add(1)
-	go pidZero.startShutdownManager()
+	// Start the shutdown manager using wg.Go
+	pidZero.wg.Go(pidZero.startShutdownManager)
 
 	// Give the manager a moment to start
 	time.Sleep(50 * time.Millisecond)

supervisor/state.go

@@ -163,7 +163,6 @@ func (p *PIDZero) broadcastState() {
 // Stateable runnable. It blocks until the context is done, coordinating state updates
 // from all state-emitting services.
 func (p *PIDZero) startStateMonitor() {
-	defer p.wg.Done()
 	p.logger.Debug("Starting state monitor...")
 
 	// Create a WaitGroup to track state monitoring goroutines

supervisor/supervisor.go

@@ -190,34 +190,36 @@ func (p *PIDZero) Run() error {
 	// Start a single reload manager if any runnable is reloadable
 	for _, r := range p.runnables {
 		if _, ok := r.(Reloadable); ok {
-			p.wg.Add(1)
-			go p.startReloadManager()
+			p.wg.Go(p.startReloadManager)
 			break
 		}
 	}
 
 	// Start a single state monitor if any runnable reports state
 	for _, r := range p.runnables {
 		if _, ok := r.(Stateable); ok {
-			p.wg.Add(1)
-			go p.startStateMonitor()
+			p.wg.Go(p.startStateMonitor)
 			break
 		}
 	}
 
 	// Start a single shutdown manager if any runnable can trigger shutdown
 	for _, r := range p.runnables {
 		if _, ok := r.(ShutdownSender); ok {
-			p.wg.Add(1)
-			go p.startShutdownManager()
+			p.wg.Go(p.startShutdownManager)
 			break
 		}
 	}
 
 	// Start each service in sequence
 	for _, r := range p.runnables {
-		p.wg.Add(1)
-		go p.startRunnable(r) // start this runnable in a separate goroutine
+		p.wg.Go(func() {
+			err := p.startRunnable(r)
+			if err != nil {
+				p.logger.Error("Runnable exited with error", "runnable", r, "error", err)
+				p.errorChan <- err
+			}
+		})
 
 		// if this Runnable implements the Stateable block here until IsRunning()
 		if stateable, ok := r.(Stateable); ok {
@@ -263,9 +265,7 @@ func (p *PIDZero) blockUntilRunnableReady(r Stateable) error {
 		)
 		select {
 		case err := <-p.errorChan:
-			// error received from `startRunnable`- put it back in the channel for reap() to process it later
-			p.errorChan <- err
-			return fmt.Errorf("runnable failed to start: %w", err)
+			return err
 		case <-startupCtx.Done():
 			return fmt.Errorf("timeout waiting for runnable to start: %w", startupCtx.Err())
 		case <-p.ctx.Done():
@@ -362,9 +362,7 @@ func (p *PIDZero) listenForSignals() {
 }
 
 // startRunnable starts a service and sends any errors to the error channel
-func (p *PIDZero) startRunnable(r Runnable) {
-	defer p.wg.Done()
-
+func (p *PIDZero) startRunnable(r Runnable) error {
 	// Log the initial state if available
 	if stateable, ok := r.(Stateable); ok {
 		initialState := stateable.GetState()
@@ -378,31 +376,16 @@ func (p *PIDZero) startRunnable(r Runnable) {
 
 	// Run the runnable with the child context
 	err := r.Run(runCtx)
-	logger := p.logger.With("runnable", r)
-	if err == nil {
-		logger.Debug("Runnable completed without error")
-		return
-	}
-	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-		// Filter out expected cancellation errors
-		logger.Debug("Runnable stopped gracefully", "reason", err)
-		return
-	}
-
-	// Unexpected error - log and send to the errorChan
-	logger = logger.With("error", err)
-	if stateable, ok := r.(Stateable); ok {
-		logger.Error("Service failed", "state", stateable.GetState())
-	} else {
-		logger.Error("Service failed")
-	}
+	if err != nil {
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			// Filter out expected cancellation errors
+			p.logger.Debug("Runnable stopped gracefully", "reason", err, "runnable", r)
+			return nil
+		}
 
-	select {
-	case p.errorChan <- fmt.Errorf("failed to start runnable: %w", err):
-		// Error sent successfully
-	default:
-		logger.Warn("Unable to send error to errorChan (full or closed?)")
+		return err
 	}
+	return nil
 }
 
 // reap listens indefinitely for errors or OS signals and handles them appropriately.