Move queue runner to DBOSContext + fix races & leaks (#46)

This is a move towards having the `DBOSContext` federate correctly all
its underlying resources, and, specifically, propagate cancellation to
said resources (including workflows.)

Calling `DBOSContext.Shutdown()` will trigger the cancelation of the
context, which will signal all resources they should stop work.

- New runner struct, attached to DBOSContext. Move queue registry in
runner struct.
- Queue runner cancellation managed by DBOSContext.Shutdown
- Add [go leak checker](https://github.com/uber-go/goleak) in tests
- Simplify queue metadata (call list queues on the runner)
- Fix races in the tests
- Handle context timeout in system DB `DequeueWorkflows`
- Better system DB cleanup, including notification loop cancellation
- Have GHA test for races using `-race`

Author: maxdml

.github/workflows/tests.yml

@@ -6,7 +6,6 @@ on:
         - main
     pull_request:
       branches:
-        - main
       types:
         - ready_for_review
         - opened
@@ -62,7 +61,7 @@ jobs:
       run: go install gotest.tools/gotestsum@latest
 
     - name: Run tests
-      run: gotestsum --format github-action
+      run: gotestsum --format github-action -- -race ./...
       working-directory: ./dbos
       env:
         PGPASSWORD: a!b@c$d()e*_,/:;=?@ff[]22

dbos/admin_server.go

@@ -20,13 +20,6 @@ type adminServer struct {
 	logger *slog.Logger
 }
 
-type queueMetadata struct {
-	Name              string       `json:"name"`
-	Concurrency       *int         `json:"concurrency,omitempty"`
-	WorkerConcurrency *int         `json:"workerConcurrency,omitempty"`
-	RateLimit         *RateLimiter `json:"rateLimit,omitempty"`
-}
-
 func newAdminServer(ctx *dbosContext, port int) *adminServer {
 	mux := http.NewServeMux()
 
@@ -80,19 +73,7 @@ func newAdminServer(ctx *dbosContext, port int) *adminServer {
 			return
 		}
 
-		var queueMetadataArray []queueMetadata
-
-		// Iterate through all queues in the registry
-		for _, queue := range workflowQueueRegistry {
-			queueMetadata := queueMetadata{
-				Name:              queue.name,
-				WorkerConcurrency: queue.workerConcurrency,
-				Concurrency:       queue.globalConcurrency,
-				RateLimit:         queue.limiter,
-			}
-
-			queueMetadataArray = append(queueMetadataArray, queueMetadata)
-		}
+		queueMetadataArray := ctx.queueRunner.listQueues()
 
 		w.Header().Set("Content-Type", "application/json")
 		if err := json.NewEncoder(w).Encode(queueMetadataArray); err != nil {
@@ -125,10 +106,11 @@ func (as *adminServer) Start() error {
 	return nil
 }
 
-func (as *adminServer) Shutdown() error {
+func (as *adminServer) Shutdown(ctx context.Context) error {
 	as.logger.Info("Shutting down admin server")
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	// XXX consider moving the grace period to DBOSContext.Shutdown()
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	if err := as.server.Shutdown(ctx); err != nil {

dbos/admin_server_test.go

@@ -144,7 +144,7 @@ func TestAdminServer(t *testing.T) {
 				endpoint:       "http://localhost:3001" + workflowQueuesMetadataPath,
 				expectedStatus: http.StatusOK,
 				validateResp: func(t *testing.T, resp *http.Response) {
-					var queueMetadata []queueMetadata
+					var queueMetadata []WorkflowQueue
 					if err := json.NewDecoder(resp.Body).Decode(&queueMetadata); err != nil {
 						t.Errorf("Failed to decode response as QueueMetadata array: %v", err)
 					}
@@ -160,8 +160,8 @@ func TestAdminServer(t *testing.T) {
 					for _, queue := range queueMetadata {
 						if queue.Name == _DBOS_INTERNAL_QUEUE_NAME { // Internal queue name
 							foundInternalQueue = true
-							if queue.Concurrency != nil {
-								t.Errorf("Expected internal queue to have no concurrency limit, but got %v", *queue.Concurrency)
+							if queue.GlobalConcurrency != nil {
+								t.Errorf("Expected internal queue to have no concurrency limit, but got %v", *queue.GlobalConcurrency)
 							}
 							if queue.WorkerConcurrency != nil {
 								t.Errorf("Expected internal queue to have no worker concurrency limit, but got %v", *queue.WorkerConcurrency)

dbos/dbos.go

@@ -5,17 +5,19 @@ import (
 	"crypto/sha256"
 	"encoding/gob"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 	"os"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/robfig/cron/v3"
 )
 
-var (
+const (
 	_DEFAULT_ADMIN_SERVER_PORT = 3001
 )
 
@@ -78,18 +80,17 @@ type DBOSContext interface {
 }
 
 type dbosContext struct {
-	ctx context.Context
+	ctx           context.Context
+	ctxCancelFunc context.CancelCauseFunc
 
-	launched bool
+	launched atomic.Bool
 
 	systemDB    SystemDatabase
 	adminServer *adminServer
 	config      *Config
 
-	// Queue runner context and cancel function
-	queueRunnerCtx        context.Context
-	queueRunnerCancelFunc context.CancelFunc
-	queueRunnerDone       chan struct{}
+	// Queue runner
+	queueRunner *queueRunner
 
 	// Application metadata
 	applicationVersion string
@@ -171,9 +172,11 @@ func (c *dbosContext) GetApplicationID() string {
 }
 
 func NewDBOSContext(inputConfig Config) (DBOSContext, error) {
+	ctx, cancelFunc := context.WithCancelCause(context.Background())
 	initExecutor := &dbosContext{
 		workflowsWg:      &sync.WaitGroup{},
-		ctx:              context.Background(),
+		ctx:              ctx,
+		ctxCancelFunc:    cancelFunc,
 		workflowRegistry: make(map[string]workflowRegistryEntry),
 		workflowRegMutex: &sync.RWMutex{},
 	}
@@ -207,24 +210,34 @@ func NewDBOSContext(inputConfig Config) (DBOSContext, error) {
 
 	initExecutor.applicationID = os.Getenv("DBOS__APPID")
 
+	initExecutor.logger = initExecutor.logger.With(
+		//"app_version", initExecutor.applicationVersion, // This is really verbose...
+		"executor_id", initExecutor.executorID,
+		//"app_id", initExecutor.applicationID, // This should stay internal
+	)
+
 	// Create the system database
-	systemDB, err := NewSystemDatabase(config.DatabaseURL, config.Logger)
+	systemDB, err := newSystemDatabase(initExecutor, config.DatabaseURL, initExecutor.logger)
 	if err != nil {
 		return nil, newInitializationError(fmt.Sprintf("failed to create system database: %v", err))
 	}
 	initExecutor.systemDB = systemDB
 	initExecutor.logger.Info("System database initialized")
 
+	// Initialize the queue runner and register DBOS internal queue
+	initExecutor.queueRunner = newQueueRunner()
+	NewWorkflowQueue(initExecutor, _DBOS_INTERNAL_QUEUE_NAME)
+
 	return initExecutor, nil
 }
 
 func (c *dbosContext) Launch() error {
-	if c.launched {
+	if c.launched.Load() {
 		return newInitializationError("DBOS is already launched")
 	}
 
 	// Start the system database
-	c.systemDB.Launch(context.Background())
+	c.systemDB.Launch(c)
 
 	// Start the admin server if configured
 	if c.config.AdminServer {
@@ -238,17 +251,9 @@ func (c *dbosContext) Launch() error {
 		c.adminServer = adminServer
 	}
 
-	// Create context with cancel function for queue runner
-	// FIXME: cancellation now has to go through the DBOSContext
-	ctx, cancel := context.WithCancel(c.ctx)
-	c.queueRunnerCtx = ctx
-	c.queueRunnerCancelFunc = cancel
-	c.queueRunnerDone = make(chan struct{})
-
 	// Start the queue runner in a goroutine
 	go func() {
-		defer close(c.queueRunnerDone)
-		queueRunner(c)
+		c.queueRunner.run(c)
 	}()
 	c.logger.Info("Queue runner started")
 
@@ -268,62 +273,61 @@ func (c *dbosContext) Launch() error {
 	}
 
 	c.logger.Info("DBOS initialized", "app_version", c.applicationVersion, "executor_id", c.executorID)
-	c.launched = true
+	c.launched.Store(true)
 	return nil
 }
 
+// We might consider renaming this to "Cancel" to me more idiomatic
 func (c *dbosContext) Shutdown() {
+	c.logger.Info("Shutting down DBOS context")
+
+	// Cancel the context to signal all resources to stop
+	c.ctxCancelFunc(errors.New("DBOS shutdown initiated"))
+
 	// Wait for all workflows to finish
 	c.logger.Info("Waiting for all workflows to finish")
 	c.workflowsWg.Wait()
+	c.logger.Info("All workflows completed")
 
-	if !c.launched {
-		c.logger.Warn("DBOS is not launched, nothing to shutdown")
-		return
-	}
-
-	// Cancel the context to stop the queue runner
-	if c.queueRunnerCancelFunc != nil {
-		c.logger.Info("Stopping queue runner")
-		c.queueRunnerCancelFunc()
-		// Wait for queue runner to finish
-		<-c.queueRunnerDone
-		c.logger.Info("Queue runner stopped")
-	}
-
-	if c.workflowScheduler != nil {
-		c.logger.Info("Stopping workflow scheduler")
-
-		ctx := c.workflowScheduler.Stop()
-		// Wait for all running jobs to complete with 5-second timeout
-		timeoutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-
-		select {
-		case <-ctx.Done():
-			c.logger.Info("All scheduled jobs completed")
-		case <-timeoutCtx.Done():
-			c.logger.Warn("Timeout waiting for jobs to complete", "timeout", "5s")
-		}
-	}
-
+	// Close the pool and the notification listener if started
 	if c.systemDB != nil {
 		c.logger.Info("Shutting down system database")
-		c.systemDB.Shutdown()
+		c.systemDB.Shutdown(c)
 		c.systemDB = nil
 	}
 
-	if c.adminServer != nil {
-		c.logger.Info("Shutting down admin server")
+	if c.launched.Load() {
+		// Wait for queue runner to finish
+		<-c.queueRunner.completionChan
+		c.logger.Info("Queue runner completed")
+
+		if c.workflowScheduler != nil {
+			c.logger.Info("Stopping workflow scheduler")
+			ctx := c.workflowScheduler.Stop()
+			// Wait for all running jobs to complete with 5-second timeout
+			timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+			defer cancel()
+
+			select {
+			case <-ctx.Done():
+				c.logger.Info("All scheduled jobs completed")
+			case <-timeoutCtx.Done():
+				c.logger.Warn("Timeout waiting for jobs to complete. Moving on", "timeout", "5s")
+			}
+		}
 
-		err := c.adminServer.Shutdown()
-		if err != nil {
-			c.logger.Error("Failed to shutdown admin server", "error", err)
-		} else {
-			c.logger.Info("Admin server shutdown complete")
+		if c.adminServer != nil {
+			c.logger.Info("Shutting down admin server")
+			err := c.adminServer.Shutdown(c)
+			if err != nil {
+				c.logger.Error("Failed to shutdown admin server", "error", err)
+			} else {
+				c.logger.Info("Admin server shutdown complete")
+			}
+			c.adminServer = nil
 		}
-		c.adminServer = nil
 	}
+	c.launched.Store(false)
 }
 
 func GetBinaryHash() (string, error) {