diff --git a/internal/membership_listener.go b/internal/membership_listener.go index b67a450..ea33979 100644 --- a/internal/membership_listener.go +++ b/internal/membership_listener.go @@ -81,6 +81,7 @@ type membershipListener struct { membershipClient MembershipClient modules modules wp *pond.WorkerPool + dispatcher *stackDispatcher } func (c *membershipListener) Start(ctx context.Context) { @@ -91,59 +92,68 @@ func (c *membershipListener) Start(ctx context.Context) { if !ok { return } - - c.wp.Submit(func() { - ctx = grpcclient.ExtractOtelCtxFromMessage(ctx, msg) - - ctx, span := tracer.Start(ctx, "NewOrder") - defer span.End() - - logger := logging.FromContext(ctx). - WithField("traceId", span.SpanContext().TraceID()). - WithField("spanId", span.SpanContext().SpanID()) - logger.Infof("Got message from membership: %T", msg.GetMessage()) - - switch msg := msg.Message.(type) { - case *generated.Order_ExistingStack: - logger = logger.WithField("stack", msg.ExistingStack.ClusterName) - ctx = logging.ContextWithLogger(ctx, logger) - - span.SetName("SyncExistingStack") - span.SetAttributes(attribute.String("stack", msg.ExistingStack.ClusterName)) - - c.syncExistingStack(ctx, msg.ExistingStack) - case *generated.Order_DeletedStack: - logger = logger.WithField("stack", msg.DeletedStack.ClusterName) - ctx = logging.ContextWithLogger(ctx, logger) - - span.SetName("DeleteStack") - span.SetAttributes(attribute.String("stack", msg.DeletedStack.ClusterName)) - - c.deleteStack(ctx, msg.DeletedStack) - case *generated.Order_DisabledStack: - logger = logger.WithField("stack", msg.DisabledStack.ClusterName) - ctx = logging.ContextWithLogger(ctx, logger) - - span.SetName("DisableStack") - span.SetAttributes(attribute.String("stack", msg.DisabledStack.ClusterName)) - - c.disableStack(ctx, msg.DisabledStack) - case *generated.Order_EnabledStack: - logger = logger.WithField("stack", msg.EnabledStack.ClusterName) - ctx = logging.ContextWithLogger(ctx, logger) - - span.SetName("EnableStack") - span.SetAttributes(attribute.String("stack", msg.EnabledStack.ClusterName)) - - c.enableStack(ctx, msg.EnabledStack) - } - }) + c.dispatcher.Dispatch(ctx, msg) case <-ctx.Done(): return } } } +// handleOrder is the per-order worker invoked by the dispatcher. It owns +// trace span creation and routes to the per-type sync method. The +// dispatcher guarantees at most one in-flight call to handleOrder per +// stack, so the sync methods can assume no concurrent peer for the same +// stack. They must, however, return promptly on ctx cancellation — the +// dispatcher cancels the in-flight handler when a superseding order +// arrives (e.g. a DeletedStack arriving while ExistingStack is in +// progress). +func (c *membershipListener) handleOrder(ctx context.Context, msg *generated.Order) { + ctx = grpcclient.ExtractOtelCtxFromMessage(ctx, msg) + + ctx, span := tracer.Start(ctx, "NewOrder") + defer span.End() + + logger := logging.FromContext(ctx). + WithField("traceId", span.SpanContext().TraceID()). + WithField("spanId", span.SpanContext().SpanID()) + logger.Infof("Got message from membership: %T", msg.GetMessage()) + + switch msg := msg.Message.(type) { + case *generated.Order_ExistingStack: + logger = logger.WithField("stack", msg.ExistingStack.ClusterName) + ctx = logging.ContextWithLogger(ctx, logger) + + span.SetName("SyncExistingStack") + span.SetAttributes(attribute.String("stack", msg.ExistingStack.ClusterName)) + + c.syncExistingStack(ctx, msg.ExistingStack) + case *generated.Order_DeletedStack: + logger = logger.WithField("stack", msg.DeletedStack.ClusterName) + ctx = logging.ContextWithLogger(ctx, logger) + + span.SetName("DeleteStack") + span.SetAttributes(attribute.String("stack", msg.DeletedStack.ClusterName)) + + c.deleteStack(ctx, msg.DeletedStack) + case *generated.Order_DisabledStack: + logger = logger.WithField("stack", msg.DisabledStack.ClusterName) + ctx = logging.ContextWithLogger(ctx, logger) + + span.SetName("DisableStack") + span.SetAttributes(attribute.String("stack", msg.DisabledStack.ClusterName)) + + c.disableStack(ctx, msg.DisabledStack) + case *generated.Order_EnabledStack: + logger = logger.WithField("stack", msg.EnabledStack.ClusterName) + ctx = logging.ContextWithLogger(ctx, logger) + + span.SetName("EnableStack") + span.SetAttributes(attribute.String("stack", msg.EnabledStack.ClusterName)) + + c.enableStack(ctx, msg.EnabledStack) + } +} + func (c *membershipListener) syncExistingStack(ctx context.Context, membershipStack *generated.Stack) { versions := membershipStack.Versions if versions == "" { @@ -160,12 +170,27 @@ func (c *membershipListener) syncExistingStack(ctx context.Context, membershipSt }, }) if err != nil { + if errors.Is(err, context.Canceled) { + return + } logging.FromContext(ctx).Errorf("Unable to create stack cluster side: %s", err) return } + // The dispatcher may cancel ctx mid-sync when a superseding order + // arrives. Bail out between steps to avoid pointless work; the + // successor order will reconcile any partial state. + if ctx.Err() != nil { + return + } c.syncModules(ctx, metadata, stack, membershipStack) + if ctx.Err() != nil { + return + } c.syncStargate(ctx, metadata, stack, membershipStack) + if ctx.Err() != nil { + return + } c.syncAuthClients(ctx, metadata, stack, membershipStack.StaticClients) logging.FromContext(ctx).Infof("Stack %s updated cluster side", stack.GetName()) @@ -474,14 +499,17 @@ func NewMembershipListener( membershipClient MembershipClient, modules modules, ) *membershipListener { - return &membershipListener{ + wp := pond.New(5, 5) + l := &membershipListener{ client: client, clientInfo: clientInfo, restMapper: mapper, membershipClient: membershipClient, - wp: pond.New(5, 5), + wp: wp, modules: modules, } + l.dispatcher = newStackDispatcher(wp, l.handleOrder) + return l } func must[T any](t *T, err error) T { diff --git a/internal/stack_dispatcher.go b/internal/stack_dispatcher.go new file mode 100644 index 0000000..5d41b4c --- /dev/null +++ b/internal/stack_dispatcher.go @@ -0,0 +1,204 @@ +package internal + +import ( + "context" + "sync" + + "github.com/alitto/pond" + "github.com/formancehq/stack/components/agent/internal/generated" +) + +// orderHandler processes a single order. Implementations must honor ctx +// cancellation: when ctx is Done, return as soon as it is safe to do so. +// A partially-applied order is fine — the next order for the same stack +// will reconcile the cluster to its own desired state. +type orderHandler func(ctx context.Context, order *generated.Order) + +// stackDispatcher routes membership orders through a per-stack mailbox. +// +// # Why this exists +// +// Orders from membership are stack-scoped (ExistingStack, DeletedStack, +// EnabledStack, DisabledStack). Each handler reads current K8s state +// and writes back desired state. When two orders for the same stack +// run concurrently, the slower one's write lands last and silently +// overwrites the faster one. Submitting every order to a shared worker +// pool — the previous behaviour — makes this race trivially +// reproducible whenever membership emits two orders in close succession. +// +// # Guarantees +// +// - Per-stack FIFO: at most one order per stack is in flight at any +// time. Pending orders for the same stack run in arrival order. +// - Cross-stack parallelism: orders for distinct stacks remain +// concurrent, bounded by the underlying worker pool. +// - Cancellation on supersession: an arriving order may cancel the +// in-flight one when continuing it would be wasteful or wrong. +// Today two cases trigger cancellation: +// (a) a fresher ExistingStack supersedes an in-flight ExistingStack +// (the older snapshot is stale; replaying it would overwrite the +// newer state), +// (b) a DeletedStack supersedes any non-Deleted in-flight order +// (no point creating resources we are about to tear down). +// A DeletedStack in flight is never cancelled. Cancellation only +// fires when a successor is already queued, so the cluster is +// guaranteed to be reconciled by the next order. +// +// # Coalescing +// +// Consecutive ExistingStack entries in the pending queue are collapsed +// to the latest one — each carries the full desired state, so older +// snapshots are pure waste. +type stackDispatcher struct { + wp *pond.WorkerPool + handle orderHandler + + mu sync.Mutex + queues map[string]*stackQueue +} + +// stackQueue tracks the per-stack mailbox state. While running is true, +// cancel is non-nil and refers to the in-flight handler's context. +type stackQueue struct { + running bool + cancel context.CancelFunc + currentType orderType + pending []*generated.Order +} + +func newStackDispatcher(wp *pond.WorkerPool, handle orderHandler) *stackDispatcher { + return &stackDispatcher{ + wp: wp, + handle: handle, + queues: make(map[string]*stackQueue), + } +} + +// Dispatch routes an order through the per-stack mailbox. Orders that +// are not stack-scoped (an unrecognised message variant) bypass the +// mailbox and go straight to the pool. +func (d *stackDispatcher) Dispatch(ctx context.Context, order *generated.Order) { + stackName := stackNameFromOrder(order) + if stackName == "" { + d.wp.Submit(func() { d.handle(ctx, order) }) + return + } + + d.mu.Lock() + q, ok := d.queues[stackName] + if !ok { + q = &stackQueue{} + d.queues[stackName] = q + } + + if q.running { + q.pending = appendCoalesced(q.pending, order) + if shouldCancelInFlight(q.currentType, order) { + q.cancel() + } + d.mu.Unlock() + return + } + + // First order for this stack: start a drain goroutine. The cancel + // func is stored before releasing the lock so a concurrent Dispatch + // observing running=true can always call it safely. + runCtx, cancel := context.WithCancel(ctx) + q.running = true + q.cancel = cancel + q.currentType = orderTypeOf(order) + d.mu.Unlock() + + d.wp.Submit(func() { d.drain(ctx, stackName, runCtx, cancel, order) }) +} + +func (d *stackDispatcher) drain(parent context.Context, stackName string, firstCtx context.Context, firstCancel context.CancelFunc, first *generated.Order) { + current := first + currentCtx := firstCtx + currentCancel := firstCancel + + for { + d.handle(currentCtx, current) + currentCancel() + + d.mu.Lock() + q := d.queues[stackName] + if len(q.pending) == 0 { + delete(d.queues, stackName) + d.mu.Unlock() + return + } + current = q.pending[0] + q.pending = q.pending[1:] + currentCtx, currentCancel = context.WithCancel(parent) + q.cancel = currentCancel + q.currentType = orderTypeOf(current) + d.mu.Unlock() + } +} + +type orderType uint8 + +const ( + orderUnknown orderType = iota + orderExistingStack + orderDeletedStack + orderEnabledStack + orderDisabledStack +) + +func orderTypeOf(o *generated.Order) orderType { + switch o.GetMessage().(type) { + case *generated.Order_ExistingStack: + return orderExistingStack + case *generated.Order_DeletedStack: + return orderDeletedStack + case *generated.Order_EnabledStack: + return orderEnabledStack + case *generated.Order_DisabledStack: + return orderDisabledStack + } + return orderUnknown +} + +func stackNameFromOrder(o *generated.Order) string { + switch m := o.GetMessage().(type) { + case *generated.Order_ExistingStack: + return m.ExistingStack.GetClusterName() + case *generated.Order_DeletedStack: + return m.DeletedStack.GetClusterName() + case *generated.Order_EnabledStack: + return m.EnabledStack.GetClusterName() + case *generated.Order_DisabledStack: + return m.DisabledStack.GetClusterName() + } + return "" +} + +func appendCoalesced(pending []*generated.Order, order *generated.Order) []*generated.Order { + if orderTypeOf(order) != orderExistingStack || len(pending) == 0 { + return append(pending, order) + } + if orderTypeOf(pending[len(pending)-1]) != orderExistingStack { + return append(pending, order) + } + pending[len(pending)-1] = order + return pending +} + +// shouldCancelInFlight reports whether an arriving order should cancel +// the in-flight one for the same stack. See stackDispatcher's doc +// comment for the rules. +func shouldCancelInFlight(current orderType, incoming *generated.Order) bool { + if current == orderDeletedStack { + return false + } + incomingType := orderTypeOf(incoming) + if incomingType == orderDeletedStack { + return true + } + if current == orderExistingStack && incomingType == orderExistingStack { + return true + } + return false +} diff --git a/internal/stack_dispatcher_test.go b/internal/stack_dispatcher_test.go new file mode 100644 index 0000000..9c55bbd --- /dev/null +++ b/internal/stack_dispatcher_test.go @@ -0,0 +1,336 @@ +package internal + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/alitto/pond" + "github.com/formancehq/stack/components/agent/internal/generated" + "github.com/stretchr/testify/require" +) + +// fakeHandler instruments the dispatcher under test. Every handle() call +// signals on `started`, blocks waiting for a value on `release` (or until +// its context is cancelled), and records the resulting outcome. +type fakeHandler struct { + started chan *generated.Order + release chan struct{} + + mu sync.Mutex + done []handlerOutcome +} + +type handlerOutcome struct { + order *generated.Order + cancelled bool +} + +func newFakeHandler(buffer int) *fakeHandler { + return &fakeHandler{ + started: make(chan *generated.Order, buffer), + release: make(chan struct{}, buffer), + } +} + +func (h *fakeHandler) handle(ctx context.Context, order *generated.Order) { + h.started <- order + + select { + case <-ctx.Done(): + h.record(order, true) + return + case <-h.release: + } + h.record(order, false) +} + +func (h *fakeHandler) record(order *generated.Order, cancelled bool) { + h.mu.Lock() + defer h.mu.Unlock() + h.done = append(h.done, handlerOutcome{order: order, cancelled: cancelled}) +} + +func (h *fakeHandler) outcomes() []handlerOutcome { + h.mu.Lock() + defer h.mu.Unlock() + out := make([]handlerOutcome, len(h.done)) + copy(out, h.done) + return out +} + +func existingOrder(stackName string) *generated.Order { + return &generated.Order{Message: &generated.Order_ExistingStack{ + ExistingStack: &generated.Stack{ClusterName: stackName}, + }} +} + +func deletedOrder(stackName string) *generated.Order { + return &generated.Order{Message: &generated.Order_DeletedStack{ + DeletedStack: &generated.DeletedStack{ClusterName: stackName}, + }} +} + +func disabledOrder(stackName string) *generated.Order { + return &generated.Order{Message: &generated.Order_DisabledStack{ + DisabledStack: &generated.DisabledStack{ClusterName: stackName}, + }} +} + +const dispatchTimeout = 2 * time.Second + +func expectStart(t *testing.T, h *fakeHandler, want *generated.Order) { + t.Helper() + select { + case got := <-h.started: + require.Same(t, want, got, "wrong order started") + case <-time.After(dispatchTimeout): + t.Fatalf("timed out waiting for handler to start order") + } +} + +func expectNoStart(t *testing.T, h *fakeHandler, within time.Duration) { + t.Helper() + select { + case got := <-h.started: + t.Fatalf("handler unexpectedly started order %v", got) + case <-time.After(within): + } +} + +// TestStackDispatcher_SerializesSameStack confirms two orders for the +// same stack are processed strictly one-after-the-other. +func TestStackDispatcher_SerializesSameStack(t *testing.T) { + t.Parallel() + h := newFakeHandler(8) + wp := pond.New(5, 5) + defer wp.StopAndWait() + d := newStackDispatcher(wp, h.handle) + + o1 := existingOrder("stack-a") + o2 := disabledOrder("stack-a") // distinct type so it is not coalesced + + d.Dispatch(context.Background(), o1) + expectStart(t, h, o1) + + d.Dispatch(context.Background(), o2) + expectNoStart(t, h, 50*time.Millisecond) + + h.release <- struct{}{} // unblock o1 + expectStart(t, h, o2) + h.release <- struct{}{} // unblock o2 + + require.Eventually(t, func() bool { return len(h.outcomes()) == 2 }, time.Second, 5*time.Millisecond) + outcomes := h.outcomes() + require.Same(t, o1, outcomes[0].order) + require.False(t, outcomes[0].cancelled) + require.Same(t, o2, outcomes[1].order) + require.False(t, outcomes[1].cancelled) +} + +// TestStackDispatcher_ParallelDifferentStacks confirms two orders for +// distinct stacks run concurrently. +func TestStackDispatcher_ParallelDifferentStacks(t *testing.T) { + t.Parallel() + h := newFakeHandler(8) + wp := pond.New(5, 5) + defer wp.StopAndWait() + d := newStackDispatcher(wp, h.handle) + + a := existingOrder("stack-a") + b := existingOrder("stack-b") + + d.Dispatch(context.Background(), a) + d.Dispatch(context.Background(), b) + + // Both should start before either is released. + started := map[*generated.Order]bool{} + for range 2 { + select { + case o := <-h.started: + started[o] = true + case <-time.After(dispatchTimeout): + t.Fatalf("timed out waiting for both orders to start") + } + } + require.True(t, started[a]) + require.True(t, started[b]) + + h.release <- struct{}{} + h.release <- struct{}{} + + require.Eventually(t, func() bool { return len(h.outcomes()) == 2 }, time.Second, 5*time.Millisecond) +} + +// TestStackDispatcher_CoalescesConsecutiveExistingStack confirms that +// when several ExistingStack orders pile up behind an in-flight one, +// only the freshest is replayed. +func TestStackDispatcher_CoalescesConsecutiveExistingStack(t *testing.T) { + t.Parallel() + h := newFakeHandler(8) + wp := pond.New(5, 5) + defer wp.StopAndWait() + d := newStackDispatcher(wp, h.handle) + + first := existingOrder("stack-a") + d.Dispatch(context.Background(), first) + expectStart(t, h, first) + + stale := existingOrder("stack-a") + stale2 := existingOrder("stack-a") + freshest := existingOrder("stack-a") + // They cancel `first` (Existing→Existing) but that is fine for this + // test — `first` records cancelled=true and the loop advances. + d.Dispatch(context.Background(), stale) + d.Dispatch(context.Background(), stale2) + d.Dispatch(context.Background(), freshest) + + // Drain `first` (cancelled by stale). + require.Eventually(t, func() bool { + o := h.outcomes() + return len(o) >= 1 && o[0].cancelled + }, time.Second, 5*time.Millisecond) + + // The next handler invocation should be `freshest`, not stale. + expectStart(t, h, freshest) + h.release <- struct{}{} + + require.Eventually(t, func() bool { return len(h.outcomes()) == 2 }, time.Second, 5*time.Millisecond) + outcomes := h.outcomes() + require.Same(t, first, outcomes[0].order) + require.True(t, outcomes[0].cancelled, "first should be cancelled by the fresher order") + require.Same(t, freshest, outcomes[1].order) + require.False(t, outcomes[1].cancelled) +} + +// TestStackDispatcher_DeleteCancelsInflightExisting confirms that an +// arriving DeletedStack interrupts an in-flight ExistingStack so the +// delete runs promptly. +func TestStackDispatcher_DeleteCancelsInflightExisting(t *testing.T) { + t.Parallel() + h := newFakeHandler(8) + wp := pond.New(5, 5) + defer wp.StopAndWait() + d := newStackDispatcher(wp, h.handle) + + existing := existingOrder("stack-a") + del := deletedOrder("stack-a") + + d.Dispatch(context.Background(), existing) + expectStart(t, h, existing) + + d.Dispatch(context.Background(), del) + + // Existing should be cancelled without anyone releasing it. + expectStart(t, h, del) + h.release <- struct{}{} + + require.Eventually(t, func() bool { return len(h.outcomes()) == 2 }, time.Second, 5*time.Millisecond) + outcomes := h.outcomes() + require.Same(t, existing, outcomes[0].order) + require.True(t, outcomes[0].cancelled, "in-flight Existing should be cancelled by Delete") + require.Same(t, del, outcomes[1].order) + require.False(t, outcomes[1].cancelled) +} + +// TestStackDispatcher_DeleteDoesNotCancelDelete asserts an in-flight +// DeletedStack is never interrupted, even by a later order. +func TestStackDispatcher_DeleteDoesNotCancelDelete(t *testing.T) { + t.Parallel() + h := newFakeHandler(8) + wp := pond.New(5, 5) + defer wp.StopAndWait() + d := newStackDispatcher(wp, h.handle) + + del := deletedOrder("stack-a") + follower := deletedOrder("stack-a") + + d.Dispatch(context.Background(), del) + expectStart(t, h, del) + + d.Dispatch(context.Background(), follower) + expectNoStart(t, h, 50*time.Millisecond) + + h.release <- struct{}{} + expectStart(t, h, follower) + h.release <- struct{}{} + + require.Eventually(t, func() bool { return len(h.outcomes()) == 2 }, time.Second, 5*time.Millisecond) + for _, o := range h.outcomes() { + require.False(t, o.cancelled, "delete handlers should never be cancelled") + } +} + +// TestStackDispatcher_HandlerRespectsParentCancel makes sure the parent +// context still propagates: when the caller's ctx is done, in-flight +// handlers observe ctx.Done() and the dispatcher quiesces. +func TestStackDispatcher_HandlerRespectsParentCancel(t *testing.T) { + t.Parallel() + h := newFakeHandler(4) + wp := pond.New(5, 5) + defer wp.StopAndWait() + d := newStackDispatcher(wp, h.handle) + + ctx, cancel := context.WithCancel(context.Background()) + + o := existingOrder("stack-a") + d.Dispatch(ctx, o) + expectStart(t, h, o) + + cancel() + + require.Eventually(t, func() bool { + out := h.outcomes() + return len(out) == 1 && out[0].cancelled + }, time.Second, 5*time.Millisecond) +} + +// TestShouldCancelInFlight exhaustively covers the cancellation matrix. +func TestShouldCancelInFlight(t *testing.T) { + t.Parallel() + tests := []struct { + name string + current orderType + incoming *generated.Order + want bool + }{ + {"existing supersedes existing", orderExistingStack, existingOrder("x"), true}, + {"delete supersedes existing", orderExistingStack, deletedOrder("x"), true}, + {"delete supersedes enabled", orderEnabledStack, deletedOrder("x"), true}, + {"delete supersedes disabled", orderDisabledStack, deletedOrder("x"), true}, + {"existing does not supersede delete", orderDeletedStack, existingOrder("x"), false}, + {"delete does not supersede delete", orderDeletedStack, deletedOrder("x"), false}, + {"disabled does not supersede existing", orderExistingStack, disabledOrder("x"), false}, + {"disabled does not supersede disabled", orderDisabledStack, disabledOrder("x"), false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, shouldCancelInFlight(tc.current, tc.incoming)) + }) + } +} + +// TestStackDispatcher_UnknownOrderBypassesMailbox asserts that orders +// without a stack scope are still processed (today only stack-scoped +// orders reach the listener, but this defends against future variants). +func TestStackDispatcher_UnknownOrderBypassesMailbox(t *testing.T) { + t.Parallel() + h := newFakeHandler(2) + wp := pond.New(5, 5) + defer wp.StopAndWait() + d := newStackDispatcher(wp, h.handle) + + // Connected is a non-stack-scoped Order variant. + o := &generated.Order{Message: &generated.Order_Connected{Connected: &generated.Connected{}}} + d.Dispatch(context.Background(), o) + expectStart(t, h, o) + h.release <- struct{}{} + + require.Eventually(t, func() bool { return len(h.outcomes()) == 1 }, time.Second, 5*time.Millisecond) + + // Sanity: stackNameFromOrder returned empty so no queue should have been created. + d.mu.Lock() + require.Empty(t, d.queues) + d.mu.Unlock() +}