chore: use builder pattern for raft clusters

cmd/raft-tester/main.go

@@ -73,8 +73,9 @@ func main() {
 	kctx := kong.Parse(&cli)
 	ctx, _ := signal.NotifyContext(context.Background(), os.Interrupt)
 
-	cluster := raft.New(&cli.RaftConfig)
-	shard := raft.AddShard(ctx, cluster, 1, &IntStateMachine{})
+	builder := raft.NewBuilder(&cli.RaftConfig)
+	shard := raft.AddShard(ctx, builder, 1, &IntStateMachine{})
+	cluster := builder.Build(ctx)
 
 	wg, ctx := errgroup.WithContext(ctx)
 	messages := make(chan int)

common/plugin/serve.go

@@ -22,6 +22,7 @@ import (
 
 	_ "github.com/block/ftl/internal/automaxprocs" // Set GOMAXPROCS to match Linux container CPU quota.
 	ftlhttp "github.com/block/ftl/internal/http"
+	"github.com/block/ftl/internal/local"
 	"github.com/block/ftl/internal/log"
 	"github.com/block/ftl/internal/rpc"
 )
@@ -170,15 +171,11 @@ func Start[Impl any, Iface any, Config any](
 }
 
 func AllocatePort() (*net.TCPAddr, error) {
-	l, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)})
-
+	addresses, err := local.FreeTCPAddresses(1)
 	if err != nil {
 		return nil, fmt.Errorf("failed to allocate port: %w", err)
 	}
-	if err := l.Close(); err != nil {
-		return nil, fmt.Errorf("could not close connection during port check: %w", err)
-	}
-	return l.Addr().(*net.TCPAddr), nil //nolint:forcetypeassert
+	return addresses[0], nil
 }
 
 func cleanup(logger *log.Logger, pidFile string) error {

internal/local/addresses.go

@@ -0,0 +1,20 @@
+package local
+
+import (
+	"fmt"
+	"net"
+)
+
+// FreeTCPAddresses returns a list of local tcp addresses that are free to listen to at the time of the call.
+func FreeTCPAddresses(count int) ([]*net.TCPAddr, error) {
+	addresses := make([]*net.TCPAddr, count)
+	for i := range count {
+		l, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)})
+		if err != nil {
+			return nil, fmt.Errorf("failed to allocate port: %w", err)
+		}
+		defer l.Close()
+		addresses[i] = l.Addr().(*net.TCPAddr) //nolint:forcetypeassert
+	}
+	return addresses, nil
+}

internal/raft/cluster.go

@@ -27,13 +27,57 @@ type RaftConfig struct {
 	CompactionOverhead uint64        `help:"Compaction overhead" default:"100"`
 }
 
+// Builder for a Raft Cluster.
+type Builder struct {
+	config *RaftConfig
+	shards map[uint64]statemachine.CreateStateMachineFunc
+
+	handles []*ShardHandle[Event, any, any]
+}
+
+func NewBuilder(cfg *RaftConfig) *Builder {
+	return &Builder{
+		config: cfg,
+		shards: map[uint64]statemachine.CreateStateMachineFunc{},
+	}
+}
+
+// AddShard adds a shard to the cluster Builder.
+func AddShard[Q any, R any, E Event, EPtr Unmarshallable[E]](
+	ctx context.Context,
+	to *Builder,
+	shardID uint64,
+	sm StateMachine[Q, R, E, EPtr],
+) *ShardHandle[E, Q, R] {
+	to.shards[shardID] = newStateMachineShim[Q, R, E, EPtr](sm)
+
+	handle := &ShardHandle[E, Q, R]{
+		shardID: shardID,
+	}
+	to.handles = append(to.handles, (*ShardHandle[Event, any, any])(handle))
+	return handle
+}
+
 // Cluster of dragonboat nodes.
 type Cluster struct {
 	config *RaftConfig
 	nh     *dragonboat.NodeHost
 	shards map[uint64]statemachine.CreateStateMachineFunc
 }
 
+func (b *Builder) Build(ctx context.Context) *Cluster {
+	cluster := &Cluster{
+		config: b.config,
+		shards: b.shards,
+	}
+
+	for _, handle := range b.handles {
+		handle.cluster = cluster
+	}
+
+	return cluster
+}
+
 // ShardHandle is a handle to a shard in the cluster.
 // It is the interface to update and query the state of a shard.
 //
@@ -88,32 +132,6 @@ func (s *ShardHandle[E, Q, R]) Query(ctx context.Context, query Q) (R, error) {
 	return response, nil
 }
 
-// New creates a new cluster.
-func New(cfg *RaftConfig) *Cluster {
-	return &Cluster{
-		config: cfg,
-		shards: make(map[uint64]statemachine.CreateStateMachineFunc),
-	}
-}
-
-// AddShard adds a shard to the cluster.
-// This can be only called before the cluster is started.
-func AddShard[Q any, R any, E Event, EPtr Unmarshallable[E]](
-	ctx context.Context,
-	to *Cluster,
-	shardID uint64,
-	sm StateMachine[Q, R, E, EPtr],
-) *ShardHandle[E, Q, R] {
-	if to.nh != nil {
-		panic("cluster already started")
-	}
-	to.shards[shardID] = newStateMachineShim[Q, R, E, EPtr](sm)
-	return &ShardHandle[E, Q, R]{
-		shardID: shardID,
-		cluster: to,
-	}
-}
-
 // Start the cluster. Blocks until the cluster instance is ready.
 func (c *Cluster) Start(ctx context.Context) error {
 	return c.start(ctx, false)

internal/raft/cluster_test.go

@@ -4,10 +4,12 @@ import (
 	"context"
 	"encoding/binary"
 	"io"
+	"net"
 	"testing"
 	"time"
 
 	"github.com/alecthomas/assert/v2"
+	"github.com/block/ftl/internal/local"
 	"github.com/block/ftl/internal/raft"
 	"golang.org/x/sync/errgroup"
 )
@@ -43,15 +45,18 @@ func TestCluster(t *testing.T) {
 	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(60*time.Second))
 	defer cancel()
 
-	members := []string{"localhost:51001", "localhost:51002"}
+	members, err := local.FreeTCPAddresses(2)
+	assert.NoError(t, err)
 
-	cluster1 := testCluster(t, members, 1, members[0])
-	shard1_1 := raft.AddShard(ctx, cluster1, 1, &IntStateMachine{})
-	shard1_2 := raft.AddShard(ctx, cluster1, 2, &IntStateMachine{})
+	builder1 := testBuilder(t, members, 1, members[0].String())
+	shard1_1 := raft.AddShard(ctx, builder1, 1, &IntStateMachine{})
+	shard1_2 := raft.AddShard(ctx, builder1, 2, &IntStateMachine{})
+	cluster1 := builder1.Build(ctx)
 
-	cluster2 := testCluster(t, members, 2, members[1])
-	shard2_1 := raft.AddShard(ctx, cluster2, 1, &IntStateMachine{})
-	shard2_2 := raft.AddShard(ctx, cluster2, 2, &IntStateMachine{})
+	builder2 := testBuilder(t, members, 2, members[1].String())
+	shard2_1 := raft.AddShard(ctx, builder2, 1, &IntStateMachine{})
+	shard2_2 := raft.AddShard(ctx, builder2, 2, &IntStateMachine{})
+	cluster2 := builder2.Build(ctx)
 
 	wg, wctx := errgroup.WithContext(ctx)
 	wg.Go(func() error { return cluster1.Start(wctx) })
@@ -74,13 +79,16 @@ func TestJoiningExistingCluster(t *testing.T) {
 	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(60*time.Second))
 	defer cancel()
 
-	members := []string{"localhost:51001", "localhost:51002"}
+	members, err := local.FreeTCPAddresses(4)
+	assert.NoError(t, err)
 
-	cluster1 := testCluster(t, members, 1, members[0])
-	shard1 := raft.AddShard(ctx, cluster1, 1, &IntStateMachine{})
+	builder1 := testBuilder(t, members[:2], 1, members[0].String())
+	shard1 := raft.AddShard(ctx, builder1, 1, &IntStateMachine{})
+	cluster1 := builder1.Build(ctx)
 
-	cluster2 := testCluster(t, members, 2, members[1])
-	shard2 := raft.AddShard(ctx, cluster2, 1, &IntStateMachine{})
+	builder2 := testBuilder(t, members[:2], 2, members[1].String())
+	shard2 := raft.AddShard(ctx, builder2, 1, &IntStateMachine{})
+	cluster2 := builder2.Build(ctx)
 
 	wg, wctx := errgroup.WithContext(ctx)
 	wg.Go(func() error { return cluster1.Start(wctx) })
@@ -90,10 +98,11 @@ func TestJoiningExistingCluster(t *testing.T) {
 	defer cluster2.Stop()
 
 	// join to the existing cluster as a new member
-	cluster3 := testCluster(t, nil, 3, "localhost:51003")
-	shard3 := raft.AddShard(ctx, cluster3, 1, &IntStateMachine{})
+	builder3 := testBuilder(t, nil, 3, members[2].String())
+	shard3 := raft.AddShard(ctx, builder3, 1, &IntStateMachine{})
+	cluster3 := builder3.Build(ctx)
 
-	assert.NoError(t, cluster1.AddMember(ctx, 1, 3, "localhost:51003"))
+	assert.NoError(t, cluster1.AddMember(ctx, 1, 3, members[2].String()))
 
 	assert.NoError(t, cluster3.Join(ctx))
 	defer cluster3.Stop()
@@ -103,10 +112,11 @@ func TestJoiningExistingCluster(t *testing.T) {
 	assertShardValue(ctx, t, 1, shard1, shard2, shard3)
 
 	// join through the new member
-	cluster4 := testCluster(t, nil, 4, "localhost:51004")
-	shard4 := raft.AddShard(ctx, cluster4, 1, &IntStateMachine{})
+	builder4 := testBuilder(t, nil, 4, members[3].String())
+	shard4 := raft.AddShard(ctx, builder4, 1, &IntStateMachine{})
+	cluster4 := builder4.Build(ctx)
 
-	assert.NoError(t, cluster3.AddMember(ctx, 1, 4, "localhost:51004"))
+	assert.NoError(t, cluster3.AddMember(ctx, 1, 4, members[3].String()))
 	assert.NoError(t, cluster4.Join(ctx))
 	defer cluster4.Stop()
 
@@ -115,8 +125,13 @@ func TestJoiningExistingCluster(t *testing.T) {
 	assertShardValue(ctx, t, 2, shard1, shard2, shard3, shard4)
 }
 
-func testCluster(t *testing.T, members []string, id uint64, address string) *raft.Cluster {
-	return raft.New(&raft.RaftConfig{
+func testBuilder(t *testing.T, addresses []*net.TCPAddr, id uint64, address string) *raft.Builder {
+	members := make([]string, len(addresses))
+	for i, member := range addresses {
+		members[i] = member.String()
+	}
+
+	return raft.NewBuilder(&raft.RaftConfig{
 		ReplicaID:          id,
 		RaftAddress:        address,
 		DataDir:            t.TempDir(),

internal/raft/eventview.go

@@ -84,13 +84,14 @@ func (s *eventStreamStateMachine[V, VPrt, E, EPtr]) Recover(reader io.Reader) er
 	return nil
 }
 
-func NewRaftEventStream[
+// AddEventView to the Builder
+func AddEventView[
 	V encoding.BinaryMarshaler,
 	VPtr Unmarshallable[V],
 	E RaftStreamEvent[V, VPtr],
 	EPtr Unmarshallable[E],
-](ctx context.Context, cluster *Cluster, shardID uint64) eventstream.EventView[V, E] {
+](ctx context.Context, builder *Builder, shardID uint64) eventstream.EventView[V, E] {
 	sm := &eventStreamStateMachine[V, VPtr, E, EPtr]{}
-	shard := AddShard[UnitQuery, V, E, EPtr](ctx, cluster, shardID, sm)
+	shard := AddShard[UnitQuery, V, E, EPtr](ctx, builder, shardID, sm)
 	return &RaftEventView[V, VPtr, E]{shard: shard}
 }

internal/raft/eventview_test.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/alecthomas/assert/v2"
+	"github.com/block/ftl/internal/local"
 	"github.com/block/ftl/internal/raft"
 	"golang.org/x/sync/errgroup"
 )
@@ -41,16 +42,20 @@ func (v *IntSumView) UnmarshalBinary(data []byte) error {
 	return nil
 }
 
-func TestEventStream(t *testing.T) {
+func TestEventView(t *testing.T) {
 	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(60*time.Second))
 	defer cancel()
 
-	members := []string{"localhost:51001", "localhost:51002"}
+	members, err := local.FreeTCPAddresses(2)
+	assert.NoError(t, err)
+
+	builder1 := testBuilder(t, members, 1, members[0].String())
+	view1 := raft.AddEventView[IntSumView, *IntSumView, IntStreamEvent](ctx, builder1, 1)
+	cluster1 := builder1.Build(ctx)
 
-	cluster1 := testCluster(t, members, 1, members[0])
-	stream1 := raft.NewRaftEventStream[IntSumView, *IntSumView, IntStreamEvent](ctx, cluster1, 1)
-	cluster2 := testCluster(t, members, 2, members[1])
-	stream2 := raft.NewRaftEventStream[IntSumView, *IntSumView, IntStreamEvent](ctx, cluster2, 1)
+	builder2 := testBuilder(t, members, 2, members[1].String())
+	view2 := raft.AddEventView[IntSumView, *IntSumView, IntStreamEvent](ctx, builder2, 1)
+	cluster2 := builder2.Build(ctx)
 
 	eg, wctx := errgroup.WithContext(ctx)
 	eg.Go(func() error { return cluster1.Start(wctx) })
@@ -59,13 +64,13 @@ func TestEventStream(t *testing.T) {
 	defer cluster1.Stop()
 	defer cluster2.Stop()
 
-	assert.NoError(t, stream1.Publish(ctx, IntStreamEvent{Value: 1}))
+	assert.NoError(t, view1.Publish(ctx, IntStreamEvent{Value: 1}))
 
-	view, err := stream1.View(ctx)
+	view, err := view1.View(ctx)
 	assert.NoError(t, err)
 	assert.Equal(t, IntSumView{Sum: 1}, view)
 
-	view, err = stream2.View(ctx)
+	view, err = view2.View(ctx)
 	assert.NoError(t, err)
 	assert.Equal(t, IntSumView{Sum: 1}, view)
 }