fix: various bugs in observability / OTel (#1564)

backend/internal/bootstrap/bootstrap.go

@@ -19,29 +19,18 @@ import (
 )
 
 func Bootstrap(ctx context.Context) error {
-	var (
-		shutdownFns       []servicerunner.Service
-		closeDatabasePool func()
-	)
-	defer func() { //nolint:contextcheck
-		// Invoke all shutdown functions on exit
-		shutdownCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
-		defer cancel()
-		shutdownErr := servicerunner.NewServiceRunner(shutdownFns...).Run(shutdownCtx)
-		if shutdownErr != nil {
-			slog.Error("Error during graceful shutdown", "error", shutdownErr)
-		}
-		// Close the database connection pool only after the shutdown functions have run: some of them (e.g. releasing the application lock) still need to query the database.
-		if closeDatabasePool != nil {
-			closeDatabasePool()
-		}
-	}()
+	// List of services to run
+	services := make([]servicerunner.Service, 0, 3)
+	shutdowns := &shutdownManager{
+		fns: make([]servicerunner.Service, 0, 4),
+	}
 
 	// Initialize the observability stack, including the logger, distributed tracing, and metrics
-	shutdownFns, httpClient, err := initObservability(ctx, common.EnvConfig.MetricsEnabled, common.EnvConfig.TracingEnabled)
+	shutdownFns, httpClient, err := initObservability(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to initialize OpenTelemetry: %w", err)
 	}
+	shutdowns.Add(shutdownFns...)
 
 	slog.InfoContext(ctx, "Pocket ID is starting")
 
@@ -51,7 +40,10 @@ func Bootstrap(ctx context.Context) error {
 		return fmt.Errorf("failed to initialize database: %w", err)
 	}
 	if pg != nil {
-		closeDatabasePool = pg.Close
+		defer func() {
+			// Close the database connection pool only after the shutdown functions have run: some of them (e.g. releasing the application lock) still need to query the database.
+			pg.Close()
+		}()
 	}
 
 	// Init storage
@@ -77,6 +69,7 @@ func Bootstrap(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to initialize services: %w", err)
 	}
+	services = append(services, svc.appLockService.RunRenewal)
 
 	// Acquire the lock from the app lock service
 	waitUntil, err := svc.appLockService.Acquire(ctx, false)
@@ -92,14 +85,13 @@ func Bootstrap(ctx context.Context) error {
 	case <-time.After(time.Until(waitUntil)):
 	}
 
-	shutdownFn := func(shutdownCtx context.Context) error {
+	shutdowns.Add(func(shutdownCtx context.Context) error {
 		sErr := svc.appLockService.Release(shutdownCtx)
 		if sErr != nil {
 			return fmt.Errorf("failed to release application lock: %w", sErr)
 		}
 		return nil
-	}
-	shutdownFns = append(shutdownFns, shutdownFn)
+	})
 
 	// Init the actors
 	actorsOpts := NewActorsOpts{
@@ -121,11 +113,15 @@ func Bootstrap(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to initialize actors: %w", err)
 	}
+	services = append(services, actors.Run)
 
-	// Register scheduled jobs
-	err = registerScheduledJobs(ctx, db, svc, scheduler)
-	if err != nil {
-		return fmt.Errorf("failed to register scheduled jobs: %w", err)
+	// Register scheduled jobs, only in non-test mode
+	if common.EnvConfig.AppEnv != "test" {
+		err = registerScheduledJobs(ctx, db, svc, scheduler)
+		if err != nil {
+			return fmt.Errorf("failed to register scheduled jobs: %w", err)
+		}
+		services = append(services, scheduler.Run)
 	}
 
 	// Init the router
@@ -135,24 +131,18 @@ func Bootstrap(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to initialize router: %w", err)
 	}
+	services = append(services, router)
 
 	// Run all background services
 	// This call blocks until the context is canceled
-	services := []servicerunner.Service{
-		svc.appLockService.RunRenewal,
-		actors.Run,
-		router,
-	}
-
-	if common.EnvConfig.AppEnv != "test" {
-		services = append(services, scheduler.Run)
-	}
-
 	err = servicerunner.NewServiceRunner(services...).Run(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to run services: %w", err)
 	}
 
+	// Run all shutdown functions
+	shutdowns.Run(ctx)
+
 	return nil
 }
 
@@ -183,3 +173,30 @@ func InitStorage(ctx context.Context, db *gorm.DB) (fileStorage storage.FileStor
 
 	return fileStorage, nil
 }
+
+type shutdownManager struct {
+	fns []servicerunner.Service
+}
+
+func (s *shutdownManager) Add(fns ...servicerunner.Service) {
+	for _, fn := range fns {
+		if fn == nil {
+			continue
+		}
+
+		s.fns = append(s.fns, fn)
+	}
+}
+
+func (s *shutdownManager) Run(ctx context.Context) {
+	// Cleanup functions are one-shot and must each run to completion independently, so we set WaitAll to true
+	sr := servicerunner.NewServiceRunner(s.fns...)
+	sr.WaitAll = true
+
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+	defer shutdownCancel()
+	err := sr.Run(shutdownCtx)
+	if err != nil {
+		slog.ErrorContext(ctx, "Error shutting down services", slog.Any("error", err))
+	}
+}