fix: simplify WaitForCleanup

internal/api/cleanup.go

@@ -3,6 +3,8 @@ package api
 import (
 	"context"
 	"sync"
+
+	"github.com/supabase/auth/internal/utilities"
 )
 
 var (
@@ -12,19 +14,5 @@ var (
 // WaitForCleanup waits until all API servers are shut down cleanly or until
 // the provided context signals done, whichever comes first.
 func WaitForCleanup(ctx context.Context) {
-	cleanupDone := make(chan struct{})
-
-	go func() {
-		defer close(cleanupDone)
-
-		cleanupWaitGroup.Wait()
-	}()
-
-	select {
-	case <-ctx.Done():
-		return
-
-	case <-cleanupDone:
-		return
-	}
+	utilities.WaitForCleanup(ctx, &cleanupWaitGroup)
 }

internal/observability/cleanup.go

@@ -3,6 +3,8 @@ package observability
 import (
 	"context"
 	"sync"
+
+	"github.com/supabase/auth/internal/utilities"
 )
 
 var (
@@ -12,19 +14,5 @@ var (
 // WaitForCleanup waits until all observability long-running goroutines shut
 // down cleanly or until the provided context signals done.
 func WaitForCleanup(ctx context.Context) {
-	cleanupDone := make(chan struct{})
-
-	go func() {
-		defer close(cleanupDone)
-
-		cleanupWaitGroup.Wait()
-	}()
-
-	select {
-	case <-ctx.Done():
-		return
-
-	case <-cleanupDone:
-		return
-	}
+	utilities.WaitForCleanup(ctx, &cleanupWaitGroup)
 }

internal/utilities/context.go

@@ -1,6 +1,9 @@
 package utilities
 
-import "context"
+import (
+	"context"
+	"sync"
+)
 
 type contextKey string
 
@@ -26,3 +29,23 @@ func GetRequestID(ctx context.Context) string {
 
 	return obj.(string)
 }
+
+// WaitForCleanup waits until all long-running goroutines shut
+// down cleanly or until the provided context signals done.
+func WaitForCleanup(ctx context.Context, wg *sync.WaitGroup) {
+	cleanupDone := make(chan struct{})
+
+	go func() {
+		defer close(cleanupDone)
+
+		wg.Wait()
+	}()
+
+	select {
+	case <-ctx.Done():
+		return
+
+	case <-cleanupDone:
+		return
+	}
+}