Concurrency support (#5)

This commit provides a new option named SetConcurrency that allows the user of the package
to specify the number of queued commands to process concurrently. This effectively adds
batching with some slight differences.

Mainly, commands in a batch are all managed within a single database transaction, so are
dependent on each other's success. If any one command in the batch fails then all commands
in the batch will return to the queue. This makes idempotency quite important.

I wanted to add this kind of batching as I figured you could end up in situations where all
instances of the FSM get blocked waiting on a chain of dependent commands, so allowing nodes
to process large amounts of commands in a batch can help mitigate that problem.

Signed-off-by: David Bond <[email protected]>

Author: davidsbond

README.md

@@ -17,12 +17,13 @@ any other tables applications built on top of this package may use.
 
 Commands use an ordinal identifier to ensure that they're processed in order of insertion. This is an implementation
 detail not exposed to users of the package. Once processed, their entry in the `pgfsm.command` table is removed,
-providing once-only processing of individual commands.
+providing once-only processing of individual commands. When using a concurrency of greater than one, this package
+provides at-least-once processing of commands in batches. Batches of commands are dependent on eachother as they are
+handled within the same transaction. If a single command within the batch fails, all commands in the batch are returned
+to the database.
 
-When reading commands, a transaction is used to ensure that should command processing fail, the contents of the command
-are returned to the `pgfsm.command` table with the same identifier. This package uses a registration system for commands
-utilising parameterised types that allows commands to be directly decoded into their concrete types without excessive
-usage of reflection.
+This package uses a registration system for commands utilising parameterised types that allows commands to be directly
+decoded into their concrete types without excessive usage of reflection.
 
 When handling a command, you have the option to return a command as a result. This allows commands to act as a graph of
 sorts. Where the successful processing of one command creates zero or more child commands. This can be used to define

database.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	_ "embed"
 	"errors"
-	"log/slog"
 	"strings"
 
 	"github.com/jackc/pgx/v5"
@@ -46,51 +45,55 @@ func insert(ctx context.Context, tx pgx.Tx, encoder Encoding, cmd Command) error
 	return err
 }
 
-func next(ctx context.Context, tx pgx.Tx) (int64, string, []byte, error) {
-	const q = `
-		SELECT id, kind, data FROM pgfsm.command 
-		ORDER BY id ASC
-		FOR UPDATE SKIP LOCKED
-		LIMIT 1
-	`
-
-	var (
-		id   int64
+type (
+	record struct {
+		id   uint64
 		kind string
 		data []byte
-	)
+	}
+)
+
+func next(ctx context.Context, tx pgx.Tx, limit uint) ([]record, error) {
+	const q = `
+		DELETE FROM pgfsm.command WHERE id IN (
+		    SELECT id FROM pgfsm.command 
+			ORDER BY id ASC
+			FOR UPDATE SKIP LOCKED
+			LIMIT $1
+		) RETURNING id, kind, data
+	`
 
-	if err := tx.QueryRow(ctx, q).Scan(&id, &kind, &data); err != nil {
-		return 0, "", []byte{}, err
+	rows, err := tx.Query(ctx, q, limit)
+	if err != nil {
+		return nil, err
 	}
 
-	return id, kind, data, nil
-}
+	records := make([]record, 0, limit)
+	defer rows.Close()
 
-func remove(ctx context.Context, tx pgx.Tx, id int64) error {
-	const q = `DELETE FROM pgfsm.command WHERE id = $1`
+	for rows.Next() {
+		var r record
+		if err = rows.Scan(&r.id, &r.kind, &r.data); err != nil {
+			return nil, err
+		}
 
-	_, err := tx.Exec(ctx, q, id)
-	return err
+		records = append(records, r)
+	}
+
+	return records, rows.Err()
 }
 
 //go:embed migrate.sql
 var migration string
 
-func migrateUp(ctx context.Context, db *pgxpool.Pool, logger *slog.Logger) error {
-	logger.DebugContext(ctx, "performing migrations")
-
+func migrateUp(ctx context.Context, db *pgxpool.Pool) error {
 	statements := strings.Split(migration, ";")
 
 	for _, statement := range statements {
 		if strings.TrimSpace(statement) == "" {
 			continue
 		}
 
-		logger.
-			With(slog.String("statement", statement)).
-			DebugContext(ctx, "executing statement")
-
 		if _, err := db.Exec(ctx, statement); err != nil {
 			return err
 		}

pgfsm.go

@@ -18,11 +18,12 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"log/slog"
+	"sync"
 	"time"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"golang.org/x/sync/errgroup"
 )
 
 type (
@@ -79,7 +80,7 @@ func New(ctx context.Context, db *pgxpool.Pool, options ...Option) (*FSM, error)
 		o(&opts)
 	}
 
-	if err := migrateUp(ctx, db, opts.logger); err != nil {
+	if err := migrateUp(ctx, db); err != nil {
 		return nil, err
 	}
 
@@ -113,19 +114,11 @@ func (fsm *FSM) Write(ctx context.Context, cmd Command) error {
 		switch command := cmd.(type) {
 		case batchCommand:
 			for _, cmd = range command {
-				fsm.options.logger.
-					With(slog.String("command_kind", cmd.Kind())).
-					InfoContext(ctx, "writing command")
-
 				if err := insert(ctx, tx, fsm.options.encoder, cmd); err != nil {
 					return err
 				}
 			}
 		default:
-			fsm.options.logger.
-				With(slog.String("command_kind", cmd.Kind())).
-				InfoContext(ctx, "writing command")
-
 			return insert(ctx, tx, fsm.options.encoder, cmd)
 		}
 
@@ -174,76 +167,92 @@ var (
 )
 
 func (fsm *FSM) next(ctx context.Context, h Handler) error {
-	fsm.options.logger.DebugContext(ctx, "checking for new commands")
-
 	return transaction(ctx, fsm.db, func(ctx context.Context, tx pgx.Tx) error {
-		id, kind, data, err := next(ctx, tx)
+		records, err := next(ctx, tx, fsm.options.concurrency)
 		switch {
-		case errors.Is(err, pgx.ErrNoRows):
+		case len(records) == 0:
 			return errNoCommands
 		case err != nil:
 			return err
 		}
 
-		log := fsm.options.logger.With(
-			slog.String("command_kind", kind),
-			slog.Int64("command_id", id),
-		)
-
-		factory, ok := fsm.commandFactories[kind]
-		switch {
-		case !ok && fsm.options.skipUnknownCommands:
-			log.WarnContext(ctx, "skipping unknown command")
-			return remove(ctx, tx, id)
-		case !ok && !fsm.options.skipUnknownCommands:
-			return UnknownCommandError{Kind: kind}
-		}
+		return fsm.handleCommands(ctx, tx, records, h)
+	})
+}
 
-		cmd := factory()
-		if err = fsm.options.encoder.Decode(data, cmd); err != nil {
-			return fmt.Errorf("failed to decode command %d: %w", id, err)
-		}
+func (fsm *FSM) handleCommands(ctx context.Context, tx pgx.Tx, records []record, h Handler) error {
+	var mux sync.Mutex
+	commands := make([]Command, 0)
+	group, gCtx := errgroup.WithContext(ctx)
 
-		log.InfoContext(ctx, "handling command")
-		returned, err := h(ctx, cmd)
-		if err != nil {
-			log.ErrorContext(ctx, "error handling command")
-			return err
-		}
+	for _, r := range records {
+		group.Go(func() error {
+			cmd, err := fsm.handleCommand(gCtx, r, h)
+			switch {
+			case err != nil:
+				return err
+			case cmd == nil:
+				return nil
+			}
 
-		if returned != nil {
-			switch command := returned.(type) {
+			mux.Lock()
+			defer mux.Unlock()
+			switch command := cmd.(type) {
 			case batchCommand:
 				for _, batched := range command {
-					log.With(slog.String("received_command_kind", batched.Kind())).
-						InfoContext(ctx, "received additional command")
-
-					if err = insert(ctx, tx, fsm.options.encoder, batched); err != nil {
-						return err
-					}
+					commands = append(commands, batched)
 				}
 
 			default:
-				log.With(slog.String("received_command_kind", returned.Kind())).
-					InfoContext(ctx, "received additional command")
-
-				if err = insert(ctx, tx, fsm.options.encoder, command); err != nil {
-					return err
-				}
+				commands = append(commands, command)
 			}
+
+			return nil
+		})
+	}
+
+	if err := group.Wait(); err != nil {
+		return err
+	}
+
+	for _, cmd := range commands {
+		if err := insert(ctx, tx, fsm.options.encoder, cmd); err != nil {
+			return err
 		}
+	}
 
-		return remove(ctx, tx, id)
-	})
+	return nil
+}
+
+func (fsm *FSM) handleCommand(ctx context.Context, r record, h Handler) (Command, error) {
+	factory, ok := fsm.commandFactories[r.kind]
+	switch {
+	case !ok && fsm.options.skipUnknownCommands:
+		return nil, nil
+	case !ok && !fsm.options.skipUnknownCommands:
+		return nil, UnknownCommandError{Kind: r.kind}
+	}
+
+	cmd := factory()
+	if err := fsm.options.encoder.Decode(r.data, cmd); err != nil {
+		return nil, fmt.Errorf("failed to decode command %d: %w", r.id, err)
+	}
+
+	returned, err := h(ctx, cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	return returned, nil
 }
 
 type (
 	options struct {
 		skipUnknownCommands bool
 		encoder             Encoding
-		logger              *slog.Logger
 		minPollInterval     time.Duration
 		maxPollInterval     time.Duration
+		concurrency         uint
 	}
 
 	// The Option type is a function that can modify the behaviour of the FSM.
@@ -254,9 +263,9 @@ func defaultOptions() options {
 	return options{
 		skipUnknownCommands: false,
 		encoder:             &JSON{},
-		logger:              slog.New(slog.DiscardHandler),
 		minPollInterval:     time.Millisecond,
 		maxPollInterval:     time.Second * 5,
+		concurrency:         1,
 	}
 }
 
@@ -278,6 +287,14 @@ func UseEncoding(e Encoding) Option {
 	return func(o *options) { o.encoder = e }
 }
 
+// SetConcurrency is an Option implementation that changes the number of commands that can be processed at the same time.
+// It defaults to 1. Concurrent commands are processed within the same transaction. This means that within a batch of
+// commands, should any fail then all commands within the batch will be returned to the database and any subsequent
+// commands rolled back.
+func SetConcurrency(concurrency uint) Option {
+	return func(o *options) { o.concurrency = concurrency }
+}
+
 // PollInterval is an Option implementation that configures the minimum and maximum frequency at which Command implementations
 // will be read from the database. Each time the FSM checks for commands and finds none, it will half the frequency at
 // which it checks up to the maximum value. This is done to prevent excessive load on the database at times where there
@@ -289,12 +306,6 @@ func PollInterval(min, max time.Duration) Option {
 	}
 }
 
-// UseLogger is an Option implementation that modifies the logger used by the FSM. By default, the FSM uses
-// slog.DiscardHandler and will not write any logs.
-func UseLogger(l *slog.Logger) Option {
-	return func(o *options) { o.logger = l.WithGroup("pgfsm") }
-}
-
 type (
 	batchCommand []Command
 )

pgfsm_test.go

@@ -3,9 +3,8 @@ package pgfsm_test
 import (
 	"context"
 	"errors"
-	"log/slog"
 	"net/url"
-	"os"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -48,20 +47,18 @@ func TestFSM_ReadWrite(t *testing.T) {
 
 	db := testDB(t)
 
-	logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))
-
 	fsm, err := pgfsm.New(t.Context(), db,
 		pgfsm.SkipUnknownCommands(),
 		pgfsm.UseEncoding(&pgfsm.GOB{}),
-		pgfsm.UseLogger(logger),
 		pgfsm.PollInterval(time.Millisecond, time.Minute),
+		pgfsm.SetConcurrency(10),
 	)
 	require.NoError(t, err)
 
 	var (
-		handledA bool
-		handledB bool
-		handledC bool
+		handledA atomic.Bool
+		handledB atomic.Bool
+		handledC atomic.Bool
 	)
 
 	ctx, cancel := context.WithTimeout(t.Context(), time.Second*5)
@@ -84,16 +81,16 @@ func TestFSM_ReadWrite(t *testing.T) {
 		return fsm.Read(ctx, func(ctx context.Context, cmd any) (pgfsm.Command, error) {
 			switch msg := cmd.(type) {
 			case *TestCommandA:
-				handledA = true
+				handledA.Store(true)
 				return TestCommandB{Foo: msg.Foo + 1}, nil
 			case *TestCommandB:
-				handledB = true
+				handledB.Store(true)
 				return pgfsm.Batch(
 					TestCommandC{Foo: msg.Foo + 1},
 					TestCommandC{Foo: msg.Foo + 1},
 				), nil
 			case *TestCommandC:
-				handledC = true
+				handledC.Store(true)
 				return nil, nil
 			default:
 				assert.Fail(t, "should be skipping unknown commands")
@@ -107,9 +104,9 @@ func TestFSM_ReadWrite(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	assert.True(t, handledA)
-	assert.True(t, handledB)
-	assert.True(t, handledC)
+	assert.True(t, handledA.Load())
+	assert.True(t, handledB.Load())
+	assert.True(t, handledC.Load())
 }
 
 func testDB(t *testing.T) *pgxpool.Pool {