Add Batch function

This commit adds a `Batch` function that allows reads and writes to the FSM to contain multiple `Command`
implementations. This works nicely as a response to a `Handler` as it provides the transactionality where
the success of once command produces 2 or more.

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

Author: davidsbond

go.sum

@@ -74,8 +74,6 @@ go.opentelemetry.io/otel/trace v1.34.0 h1:+ouXS2V8Rd4hp4580a8q23bg0azF2nI8cqLYnC
 go.opentelemetry.io/otel/trace v1.34.0/go.mod h1:Svm7lSjQD7kG7KJ/MUHPVXSDGz2OX4h0M2jHBhmSfRE=
 go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
-golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ=
-golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w=
 golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU=

pgfsm.go

@@ -20,7 +20,6 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
-	"time"
 )
 
 type (
@@ -95,26 +94,38 @@ func (fsm *FSM) Handle(handler Handler) {
 // Write a Command to the FSM. This Command will be encoded using the Encoding implementation and stored within the
 // database, where it can then be read and the relevant Handler invoked.
 func (fsm *FSM) Write(ctx context.Context, cmd Command) error {
-	fsm.options.logger.
-		With(slog.String("command_kind", cmd.Kind())).
-		InfoContext(ctx, "writing command")
-
 	return transaction(ctx, fsm.db, func(ctx context.Context, tx *sql.Tx) error {
-		return insert(ctx, tx, fsm.options.encoder, cmd)
+		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)
+		}
+
+		return nil
 	})
 }
 
 // Read commands from the FSM. For each command read, the relevant Handler implementation will be invoked. This method
 // blocks until the provided context is cancelled, or a Handler implementation returns an error.
 func (fsm *FSM) Read(ctx context.Context) error {
-	ticker := time.NewTicker(time.Second / 10)
-	defer ticker.Stop()
-
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-ticker.C:
+		default:
 			if err := fsm.next(ctx); err != nil {
 				return err
 			}
@@ -154,10 +165,26 @@ func (fsm *FSM) next(ctx context.Context) error {
 		}
 
 		if cmd != nil {
-			log.With(slog.String("received_command_kind", cmd.Kind())).
-				InfoContext(ctx, "received additional command")
 
-			if err = insert(ctx, tx, fsm.options.encoder, cmd); err != nil {
+			switch command := cmd.(type) {
+			case batchCommand:
+				for _, cmd = range command {
+					log.With(slog.String("received_command_kind", cmd.Kind())).
+						InfoContext(ctx, "received additional command")
+
+					if err = insert(ctx, tx, fsm.options.encoder, cmd); err != nil {
+						return err
+					}
+				}
+
+			default:
+				log.With(slog.String("received_command_kind", cmd.Kind())).
+					InfoContext(ctx, "received additional command")
+
+				err = insert(ctx, tx, fsm.options.encoder, cmd)
+			}
+
+			if err != nil {
 				return err
 			}
 		}
@@ -249,3 +276,18 @@ func UseEncoding(e Encoding) Option {
 func UseLogger(l *slog.Logger) Option {
 	return func(o *options) { o.logger = l.WithGroup("pgfsm") }
 }
+
+type (
+	batchCommand []Command
+)
+
+func (cmd batchCommand) Kind() string {
+	return ""
+}
+
+// Batch returns a single Command implementation that wraps multiple other Command implementations. This can be used
+// to return multiple commands at once when returning from a CommandHandler function. Or to send multiple commands to
+// the FSM at once using FSM.Write.
+func Batch(commands ...Command) Command {
+	return batchCommand(commands)
+}

pgfsm_test.go

@@ -69,23 +69,33 @@ func TestFSM_ReadWrite(t *testing.T) {
 
 	fsm.Handle(pgfsm.CommandHandler[TestCommandA](func(ctx context.Context, a TestCommandA) (pgfsm.Command, error) {
 		handledA = true
+
 		return TestCommandB{Foo: a.Foo + 1}, nil
 	}))
 
 	fsm.Handle(pgfsm.CommandHandler[TestCommandB](func(ctx context.Context, b TestCommandB) (pgfsm.Command, error) {
 		handledB = true
-		return TestCommandC{Foo: b.Foo + 1}, nil
+
+		return pgfsm.Batch(
+			TestCommandC{Foo: b.Foo + 1},
+			TestCommandC{Foo: b.Foo + 1},
+		), nil
 	}))
 
 	fsm.Handle(pgfsm.CommandHandler[TestCommandC](func(ctx context.Context, c TestCommandC) (pgfsm.Command, error) {
 		handledC = true
 		assert.EqualValues(t, 3, c.Foo)
+
 		return nil, nil
 	}))
 
 	group, ctx := errgroup.WithContext(ctx)
 	group.Go(func() error {
-		return fsm.Write(ctx, TestCommandA{Foo: 1})
+		return fsm.Write(ctx, pgfsm.Batch(
+			TestCommandA{Foo: 1},
+			TestCommandA{Foo: 1},
+			TestCommandA{Foo: 1},
+		))
 	})
 
 	group.Go(func() error {