Merge pull request #10 from acaloiaro/implement-shutdown-functionality

Add Shutdown() functionality

memory_backend.go

@@ -22,6 +22,7 @@ const (
 
 // MemBackend is a memory-backed neoq backend
 type MemBackend struct {
+	cancelFuncs  []context.CancelFunc // A collection of cancel functions to be called upon Shutdown()
 	cron         *cron.Cron
 	logger       Logger
 	mu           *sync.Mutex // mutext to protect mutating state on a pgWorker
@@ -42,6 +43,7 @@ func NewMemBackend(opts ...ConfigOption) (n Neoq, err error) {
 		fingerprints: sync.Map{},
 		logger:       slog.New(slog.NewTextHandler(os.Stdout)),
 		jobCount:     0,
+		cancelFuncs:  []context.CancelFunc{},
 	}
 	mb.cron.Start()
 
@@ -61,7 +63,7 @@ func (m *MemBackend) Enqueue(ctx context.Context, job Job) (jobID int64, err err
 	var ok bool
 
 	if qc, ok = m.queues.Load(job.Queue); !ok {
-		return UnqueuedJobID, fmt.Errorf("queue has no listeners: %w", job.Queue)
+		return UnqueuedJobID, fmt.Errorf("queue has no listeners: %s", job.Queue)
 	}
 
 	queueChan = qc.(chan Job)
@@ -117,6 +119,12 @@ func (m *MemBackend) Listen(ctx context.Context, queue string, h Handler) (err e
 	m.handlers.Store(queue, h)
 	m.queues.Store(queue, make(chan Job, queueCapacity))
 
+	ctx, cancel := context.WithCancel(ctx)
+
+	m.mu.Lock()
+	m.cancelFuncs = append(m.cancelFuncs, cancel)
+	m.mu.Unlock()
+
 	err = m.start(ctx, queue)
 	if err != nil {
 		return
@@ -141,6 +149,11 @@ func (m *MemBackend) ListenCron(ctx context.Context, cronSpec string, h Handler)
 
 	queue := stripNonAlphanum(strcase.ToSnake(*cdStr))
 
+	ctx, cancel := context.WithCancel(ctx)
+	m.mu.Lock()
+	m.cancelFuncs = append(m.cancelFuncs, cancel)
+	m.mu.Unlock()
+
 	err = m.Listen(ctx, queue, h)
 	if err != nil {
 		return
@@ -155,6 +168,12 @@ func (m *MemBackend) ListenCron(ctx context.Context, cronSpec string, h Handler)
 
 // Shutdown halts the worker
 func (m *MemBackend) Shutdown(ctx context.Context) (err error) {
+	for _, f := range m.cancelFuncs {
+		f()
+	}
+
+	m.cancelFuncs = nil
+
 	return
 }
 

postgres_backend.go

@@ -48,14 +48,15 @@ const (
 
 // PgBackend is a Postgres-based Neoq backend
 type PgBackend struct {
-	config     *pgConfig
-	cron       *cron.Cron
-	listenConn *pgx.Conn
-	pool       *pgxpool.Pool
-	handlers   map[string]Handler  // a map of queue names to queue handlers
-	mu         *sync.Mutex         // mutext to protect mutating state on a pgWorker
-	futureJobs map[int64]time.Time // map of future job IDs to their due time
-	logger     Logger
+	cancelFuncs []context.CancelFunc // A collection of cancel functions to be called upon Shutdown()
+	config      *pgConfig
+	cron        *cron.Cron
+	listenConn  *pgx.Conn
+	pool        *pgxpool.Pool
+	handlers    map[string]Handler  // a map of queue names to queue handlers
+	mu          *sync.Mutex         // mutex to protect mutating state on a pgWorker
+	futureJobs  map[int64]time.Time // map of future job IDs to their due time
+	logger      Logger
 }
 
 // NewPgBackend creates a new neoq backend backed by Postgres
@@ -82,21 +83,27 @@ type PgBackend struct {
 // - WithTransactionTimeout(timeout int): configure the idle_in_transaction_timeout for the worker's database
 // connection(s)
 func NewPgBackend(ctx context.Context, connectString string, opts ...ConfigOption) (n Neoq, err error) {
-	w := PgBackend{
-		mu:         &sync.Mutex{},
-		config:     &pgConfig{connectString: connectString},
-		handlers:   make(map[string]Handler),
-		futureJobs: make(map[int64]time.Time),
-		logger:     slog.New(slog.NewTextHandler(os.Stdout)),
-		cron:       cron.New(),
+	w := &PgBackend{
+		mu:          &sync.Mutex{},
+		config:      &pgConfig{connectString: connectString},
+		handlers:    make(map[string]Handler),
+		futureJobs:  make(map[int64]time.Time),
+		logger:      slog.New(slog.NewTextHandler(os.Stdout)),
+		cron:        cron.New(),
+		cancelFuncs: []context.CancelFunc{},
 	}
 	w.cron.Start()
 
 	// Set all options
 	for _, opt := range opts {
-		opt(&w)
+		opt(w)
 	}
 
+	ctx, cancel := context.WithCancel(ctx)
+	w.mu.Lock()
+	w.cancelFuncs = append(w.cancelFuncs, cancel)
+	w.mu.Unlock()
+
 	err = w.initializeDB(ctx)
 	if err != nil {
 		return
@@ -274,7 +281,12 @@ func (w PgBackend) initializeDB(ctx context.Context) (err error) {
 }
 
 // Enqueue adds jobs to the specified queue
-func (w PgBackend) Enqueue(ctx context.Context, job Job) (jobID int64, err error) {
+func (w *PgBackend) Enqueue(ctx context.Context, job Job) (jobID int64, err error) {
+	ctx, cancel := context.WithCancel(ctx)
+	w.mu.Lock()
+	w.cancelFuncs = append(w.cancelFuncs, cancel)
+	w.mu.Unlock()
+
 	conn, err := w.pool.Acquire(ctx)
 	if err != nil {
 		return
@@ -332,8 +344,11 @@ func (w PgBackend) Enqueue(ctx context.Context, job Job) (jobID int64, err error
 // Listen sets the queue handler function for the specified queue.
 //
 // Neoq will not process any queues until Listen() is called at least once.
-func (w PgBackend) Listen(ctx context.Context, queue string, h Handler) (err error) {
+func (w *PgBackend) Listen(ctx context.Context, queue string, h Handler) (err error) {
+	ctx, cancel := context.WithCancel(ctx)
+
 	w.mu.Lock()
+	w.cancelFuncs = append(w.cancelFuncs, cancel)
 	w.handlers[queue] = h
 	w.mu.Unlock()
 
@@ -347,7 +362,7 @@ func (w PgBackend) Listen(ctx context.Context, queue string, h Handler) (err err
 // ListenCron listens for jobs on a cron schedule and handles them with the provided handler
 //
 // See: https://pkg.go.dev/github.com/robfig/cron?#hdr-CRON_Expression_Format for details on the cron spec format
-func (w PgBackend) ListenCron(ctx context.Context, cronSpec string, h Handler) (err error) {
+func (w *PgBackend) ListenCron(ctx context.Context, cronSpec string, h Handler) (err error) {
 	cd, err := crondescriptor.NewCronDescriptor(cronSpec)
 	if err != nil {
 		return err
@@ -359,6 +374,12 @@ func (w PgBackend) ListenCron(ctx context.Context, cronSpec string, h Handler) (
 	}
 
 	queue := stripNonAlphanum(strcase.ToSnake(*cdStr))
+	ctx, cancel := context.WithCancel(ctx)
+
+	w.mu.Lock()
+	w.cancelFuncs = append(w.cancelFuncs, cancel)
+	w.mu.Unlock()
+
 	w.cron.AddFunc(cronSpec, func() {
 		w.Enqueue(ctx, Job{Queue: queue})
 	})
@@ -368,11 +389,15 @@ func (w PgBackend) ListenCron(ctx context.Context, cronSpec string, h Handler) (
 	return
 }
 
-func (w PgBackend) Shutdown(ctx context.Context) (err error) {
-	w.pool.Close()
+func (w *PgBackend) Shutdown(ctx context.Context) (err error) {
+	w.pool.Close() // also closes the hijacked listenConn
 	w.cron.Stop()
 
-	err = w.listenConn.Close(ctx)
+	for _, f := range w.cancelFuncs {
+		f()
+	}
+
+	w.cancelFuncs = nil
 
 	return
 }
@@ -737,6 +762,10 @@ func (w PgBackend) listen(ctx context.Context, queue string) (c chan int64) {
 		for {
 			notification, waitErr := w.listenConn.WaitForNotification(ctx)
 			if waitErr != nil {
+				if errors.Is(waitErr, context.Canceled) {
+					return
+				}
+
 				w.logger.Error("failed to wait for notification", waitErr)
 				time.Sleep(1 * time.Second)
 				continue
@@ -786,9 +815,9 @@ func (w PgBackend) getPendingJobID(ctx context.Context, conn *pgxpool.Conn, queu
 func PgTransactionTimeout(txTimeout int) ConfigOption {
 	return func(n Neoq) {
 		var ok bool
-		var npg PgBackend
+		var npg *PgBackend
 
-		if npg, ok = n.(PgBackend); !ok {
+		if npg, ok = n.(*PgBackend); !ok {
 			return
 		}
 

postgres_backend_test.go

@@ -0,0 +1,78 @@
+package neoq
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/pkg/errors"
+	"golang.org/x/exp/slog"
+)
+
+// TestPgBackendBasicJobProcessing tests that the postgres backend is able to process the most basic jobs with the
+// most basic configuration.
+func TestPgBackendBasicJobProcessing(t *testing.T) {
+	queue := "testing"
+	numJobs := 10
+	doneCnt := 0
+	done := make(chan bool)
+	var timeoutTimer = time.After(5 * time.Second)
+
+	var connString = os.Getenv("TEST_DATABASE_URL")
+	if connString == "" {
+		t.Skip("Skipping: TEST_DATABASE_URL not set")
+		return
+	}
+
+	ctx := context.TODO()
+	nq, err := NewPgBackend(ctx, connString)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	handler := NewHandler(func(ctx context.Context) (err error) {
+		done <- true
+		return
+	})
+
+	nq.Listen(ctx, queue, handler)
+
+	go func() {
+		for i := 0; i < numJobs; i++ {
+			jid, err := nq.Enqueue(ctx, Job{
+				Queue: queue,
+				Payload: map[string]interface{}{
+					"message": fmt.Sprintf("hello world: %d", i),
+				},
+			})
+			if err != nil || jid == DuplicateJobID {
+				slog.Error("job was not enqueued. either it was duplicate or this error caused it:", err)
+			}
+		}
+	}()
+
+	for {
+		select {
+		case <-timeoutTimer:
+			err = errors.New("timed out waiting for job(s)")
+		case <-done:
+			doneCnt++
+		}
+
+		if doneCnt >= numJobs {
+			break
+		}
+
+		if err != nil {
+			break
+		}
+	}
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	nq.Shutdown(ctx)
+}