feat: enable graceful shutdown of persistent subscription consumer

docs/server/features/persistent-subscriptions.md

@@ -33,6 +33,30 @@ Just as with message brokers, processing events in a group of consumers running
 
 Clients must acknowledge (or not acknowledge) messages as they are handled. If messages aren't acknowledged before they time out on the server, the server will retry them. If a message has been retried more than the `maxRetryCount` setting for the persistent subscription, then the message will be parked and processing will continue.
 
+## Graceful shutdown
+
+When a consumer needs to shut down, simply closing the connection works but is abrupt: any events the consumer was still working on are immediately retried by another consumer in the group, even if the original consumer was about to ack them. This can lead to duplicate processing.
+
+To shut down cleanly, a consumer can send an explicit `Stop` message on its subscription stream. This:
+
+- Removes the consumer from the consumer pool, so the server stops sending it new events.
+- Leaves the connection open so the consumer can still acknowledge or not-acknowledge any events it already received.
+- Does **not** redistribute the consumer's in-flight events to other consumers — they stay assigned to the stopping consumer until it acks/nacks them or the message timeout fires.
+
+The recommended sequence is:
+
+1. **Send `Stop`.** The server marks the consumer as stopped and routes new events to the remaining active consumers in the group.
+2. **Drain in-flight events.** Continue receiving event-appeared messages already in flight on the stream, processing them normally, and sending acks or nacks back to the server.
+3. **Unsubscribe.** Once the consumer has nothing left in flight, close the subscription stream to fully unsubscribe.
+
+::: tip
+If the consumer disconnects without sending `Stop`, or disconnects after `Stop` while events are still in flight, those unconfirmed events follow the normal disconnect behavior: they are retried to other consumers as soon as the server detects the disconnect. `Stop` is therefore an optimization for clean shutdown, not a correctness requirement — at-least-once delivery is preserved either way.
+:::
+
+::: note
+If a stopping consumer never acks an in-flight event, that event will eventually hit the `messageTimeoutMilliseconds` and be retried to another consumer in the group, just like any other timed-out message. A consumer that intends to stop quickly should ack or nack what it can rather than relying on the timeout.
+:::
+
 ## Parked messages
 
 Messages that have been retried too many times will often be parked in the persistent subscription's parked message stream. This stream is named `$persistentsubscription-{streamname}::{groupname}-parked`. You can easily see the number of parked events in the persistent subscription statistics or browse the parked messages in the admin UI.

proto.lock

@@ -3059,6 +3059,11 @@
                 "id": 3,
                 "name": "nack",
                 "type": "Nack"
+              },
+              {
+                "id": 4,
+                "name": "stop",
+                "type": "Stop"
               }
             ],
             "messages": [
@@ -3150,6 +3155,9 @@
                     "type": "string"
                   }
                 ]
+              },
+              {
+                "name": "Stop"
               }
             ]
           },

src/KurrentDB.Core.Tests/Services/PersistentSubscription/PersistentSubscriptionTests.cs

@@ -2425,6 +2425,337 @@ public void disconnecting_a_client_with_no_persistent_subscription() {
 	}
 }
 
+[TestFixture(EventSource.SingleStream)]
+[TestFixture(EventSource.AllStream)]
+[TestFixture(EventSource.FilteredAllStream)]
+public class StopClientTests {
+	private readonly EventSource _eventSource;
+	public StopClientTests(EventSource eventSource) {
+		_eventSource = eventSource;
+	}
+
+	private KurrentDB.Core.Services.PersistentSubscription.PersistentSubscription BuildSubscription(
+		FakePushScheduler pushScheduler,
+		FakeMessageParker parker = null,
+		Action<IPersistentSubscriptionStreamPosition> onCheckpoint = null) {
+		var reader = new FakeCheckpointReader();
+		var sub = new KurrentDB.Core.Services.PersistentSubscription.PersistentSubscription(
+			Helper.CreatePersistentSubscriptionBuilderFor(_eventSource)
+				.WithEventLoader(new FakeStreamReader())
+				.WithCheckpointReader(reader)
+				.WithMessageParker(parker ?? new FakeMessageParker())
+				.WithPushScheduler(pushScheduler)
+				.PreferRoundRobin()
+				.StartFromCurrent()
+				.WithCheckpointWriter(new FakeCheckpointWriter(onCheckpoint ?? (_ => { }))));
+		reader.Load(null);
+		return sub;
+	}
+
+	[Test]
+	public void stopping_unknown_correlation_id_returns_false() {
+		var sub = BuildSubscription(new FakePushScheduler());
+		Assert.IsFalse(sub.StopClient(Guid.NewGuid()));
+	}
+
+	[Test]
+	public void stopping_a_client_returns_true() {
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+		var corrId = Guid.NewGuid();
+		sub.AddClient(corrId, Guid.NewGuid(), "connection-1", new FakeEnvelope(), 10, "foo", "bar");
+
+		Assert.IsTrue(sub.StopClient(corrId));
+	}
+
+	[Test]
+	public void stopping_a_client_keeps_it_in_the_collection() {
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+		var corrId = Guid.NewGuid();
+		sub.AddClient(corrId, Guid.NewGuid(), "connection-1", new FakeEnvelope(), 10, "foo", "bar");
+
+		sub.StopClient(corrId);
+
+		// The client object still belongs to the subscription so it can keep
+		// processing acks/nacks for in-flight events.
+		Assert.IsTrue(sub.HasClients);
+		Assert.AreEqual(1, sub.ClientCount);
+	}
+
+	[Test]
+	public void stopping_a_client_does_not_send_drop_notification() {
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+		var envelope = new FakeEnvelope();
+		var corrId = Guid.NewGuid();
+		sub.AddClient(corrId, Guid.NewGuid(), "connection-1", envelope, 10, "foo", "bar");
+
+		sub.StopClient(corrId);
+
+		Assert.IsFalse(envelope.Replies.OfType<ClientMessage.SubscriptionDropped>().Any());
+	}
+
+	[Test]
+	public void stopping_a_client_twice_is_idempotent() {
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+		var corrId = Guid.NewGuid();
+		sub.AddClient(corrId, Guid.NewGuid(), "connection-1", new FakeEnvelope(), 10, "foo", "bar");
+
+		Assert.IsTrue(sub.StopClient(corrId));
+		Assert.DoesNotThrow(() => sub.StopClient(corrId));
+		Assert.IsTrue(sub.StopClient(corrId));
+	}
+
+	[Test]
+	public void stopped_client_does_not_receive_new_events() {
+		var stoppedEnvelope = new FakeEnvelope();
+		var liveEnvelope = new FakeEnvelope();
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+
+		var stoppedId = Guid.NewGuid();
+		sub.AddClient(stoppedId, Guid.NewGuid(), "connection-stopped", stoppedEnvelope, 10, "foo", "bar");
+		sub.AddClient(Guid.NewGuid(), Guid.NewGuid(), "connection-live", liveEnvelope, 10, "foo", "bar");
+
+		sub.StopClient(stoppedId);
+
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(0, _eventSource));
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(1, _eventSource));
+		pushScheduler.Push(sub);
+
+		// Both events should land on the live consumer; the stopped one stays silent.
+		Assert.AreEqual(0, stoppedEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+		Assert.AreEqual(2, liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+	}
+
+	[Test]
+	public void stopping_does_not_retry_in_flight_events_to_others() {
+		// Contrast with RemoveClientByCorrelationId, which retries unconfirmed
+		// events immediately. Stop leaves them with the original consumer so
+		// it can finish processing them.
+		var stoppedEnvelope = new FakeEnvelope();
+		var liveEnvelope = new FakeEnvelope();
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+
+		var stoppedId = Guid.NewGuid();
+		sub.AddClient(stoppedId, Guid.NewGuid(), "connection-stopped", stoppedEnvelope, 10, "foo", "bar");
+		sub.AddClient(Guid.NewGuid(), Guid.NewGuid(), "connection-live", liveEnvelope, 10, "foo", "bar");
+
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(0, _eventSource));
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(1, _eventSource));
+		pushScheduler.Push(sub);
+
+		Assert.AreEqual(1, stoppedEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+		Assert.AreEqual(1, liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+
+		// Stop should not schedule any redistribution of in-flight events; the
+		// stopped client keeps them until it acks/nacks or the timeout fires.
+		sub.StopClient(stoppedId);
+
+		Assert.AreEqual(1, liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+		Assert.AreEqual(1, stoppedEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+	}
+
+	[Test]
+	public void stopped_client_can_still_ack_in_flight_events() {
+		IPersistentSubscriptionStreamPosition cp = null;
+		var stoppedEnvelope = new FakeEnvelope();
+		var liveEnvelope = new FakeEnvelope();
+		var pushScheduler = new FakePushScheduler();
+		var reader = new FakeCheckpointReader();
+		var sub = new KurrentDB.Core.Services.PersistentSubscription.PersistentSubscription(
+			Helper.CreatePersistentSubscriptionBuilderFor(_eventSource)
+				.WithEventLoader(new FakeStreamReader())
+				.WithCheckpointReader(reader)
+				.WithMessageParker(new FakeMessageParker())
+				.WithPushScheduler(pushScheduler)
+				.PreferRoundRobin()
+				.StartFromCurrent()
+				.MinimumToCheckPoint(1)
+				.MaximumToCheckPoint(2)
+				.WithCheckpointWriter(new FakeCheckpointWriter(i => cp = i)));
+		reader.Load(null);
+
+		var stoppedId = Guid.NewGuid();
+		sub.AddClient(stoppedId, Guid.NewGuid(), "connection-stopped", stoppedEnvelope, 10, "foo", "bar");
+		sub.AddClient(Guid.NewGuid(), Guid.NewGuid(), "connection-live", liveEnvelope, 10, "foo", "bar");
+
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(0, _eventSource));
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(1, _eventSource));
+		pushScheduler.Push(sub);
+
+		Assert.AreEqual(2, sub.OutstandingMessageCount);
+
+		sub.StopClient(stoppedId);
+
+		// Ack arrives over the still-open connection from the stopped client.
+		sub.AcknowledgeMessagesProcessed(stoppedId, new[] { Helper.GetEventIdFor(0) });
+		sub.AcknowledgeMessagesProcessed(Guid.NewGuid(), new[] { Helper.GetEventIdFor(1) });
+
+		Assert.AreEqual(0, sub.OutstandingMessageCount);
+		// Both events were acked, so the checkpoint should advance.
+		Assert.IsNotNull(cp);
+	}
+
+	[Test]
+	public void stopped_client_can_still_nack_in_flight_events() {
+		var parker = new FakeMessageParker();
+		var stoppedEnvelope = new FakeEnvelope();
+		var liveEnvelope = new FakeEnvelope();
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler, parker);
+
+		var stoppedId = Guid.NewGuid();
+		sub.AddClient(stoppedId, Guid.NewGuid(), "connection-stopped", stoppedEnvelope, 10, "foo", "bar");
+		sub.AddClient(Guid.NewGuid(), Guid.NewGuid(), "connection-live", liveEnvelope, 10, "foo", "bar");
+
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(0, _eventSource));
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(1, _eventSource));
+		pushScheduler.Push(sub);
+
+		sub.StopClient(stoppedId);
+
+		// Nack with park from the stopped client should still land in the parker.
+		sub.NotAcknowledgeMessagesProcessed(stoppedId, new[] { Helper.GetEventIdFor(0) }, NakAction.Park, "stopped consumer parking");
+
+		Assert.AreEqual(1, parker.ParkedEvents.Count);
+		Assert.AreEqual(Helper.GetEventIdFor(0), parker.ParkedEvents[0].OriginalEvent.EventId);
+	}
+
+	[Test]
+	public void stopped_clients_in_flight_event_eventually_times_out_to_live_consumer() {
+		var stoppedEnvelope = new FakeEnvelope();
+		var liveEnvelope = new FakeEnvelope();
+		var pushScheduler = new FakePushScheduler();
+		var reader = new FakeCheckpointReader();
+		var sub = new KurrentDB.Core.Services.PersistentSubscription.PersistentSubscription(
+			Helper.CreatePersistentSubscriptionBuilderFor(_eventSource)
+				.WithEventLoader(new FakeStreamReader())
+				.WithCheckpointReader(reader)
+				.WithMessageParker(new FakeMessageParker())
+				.WithPushScheduler(pushScheduler)
+				.WithMessageTimeoutOf(TimeSpan.FromMilliseconds(1))
+				.PreferRoundRobin()
+				.StartFromCurrent()
+				.WithCheckpointWriter(new FakeCheckpointWriter(_ => { })));
+		reader.Load(null);
+
+		var stoppedId = Guid.NewGuid();
+		var liveId = Guid.NewGuid();
+		sub.AddClient(stoppedId, Guid.NewGuid(), "connection-stopped", stoppedEnvelope, 10, "foo", "bar");
+		sub.AddClient(liveId, Guid.NewGuid(), "connection-live", liveEnvelope, 10, "foo", "bar");
+
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(0, _eventSource));
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(1, _eventSource));
+		pushScheduler.Push(sub);
+
+		Assert.AreEqual(1, stoppedEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+		Assert.AreEqual(1, liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+
+		// Ack the live consumer's event so only the stopped client has anything in flight.
+		sub.AcknowledgeMessagesProcessed(liveId, new[] { Helper.GetEventIdFor(1) });
+
+		sub.StopClient(stoppedId);
+
+		// After the message timeout the stopped client's in-flight event is
+		// retried — and routes to the only remaining active consumer.
+		sub.NotifyClockTick(DateTime.UtcNow.AddSeconds(1));
+		pushScheduler.Push(sub);
+
+		Assert.AreEqual(2, liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+		Assert.AreEqual(1, stoppedEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+	}
+
+	[Test]
+	public void disconnect_after_stop_does_not_double_remove_from_strategy() {
+		// RoundRobin throws InvalidOperationException if ClientRemoved is called
+		// twice for the same client. Stop must mark the client so the eventual
+		// disconnect path skips the second strategy removal.
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+
+		var corrId = Guid.NewGuid();
+		var connectionId = Guid.NewGuid();
+		sub.AddClient(corrId, connectionId, "connection-1", new FakeEnvelope(), 10, "foo", "bar");
+
+		sub.StopClient(corrId);
+
+		Assert.DoesNotThrow(() => sub.RemoveClientByCorrelationId(corrId, sendDropNotification: false));
+		Assert.IsFalse(sub.HasClients);
+	}
+
+	[Test]
+	public void disconnect_after_stop_via_connection_id_does_not_throw() {
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+
+		var corrId = Guid.NewGuid();
+		var connectionId = Guid.NewGuid();
+		sub.AddClient(corrId, connectionId, "connection-1", new FakeEnvelope(), 10, "foo", "bar");
+
+		sub.StopClient(corrId);
+
+		Assert.DoesNotThrow(() => sub.RemoveClientByConnectionId(connectionId));
+		Assert.IsFalse(sub.HasClients);
+	}
+
+	[Test]
+	public void disconnect_after_stop_retries_unconfirmed_events_to_others() {
+		// Acks/nacks may never arrive after a stop (e.g. crashing client). Once
+		// the connection drops, any still-unconfirmed events should be retried
+		// to other consumers via the normal disconnect path.
+		var stoppedEnvelope = new FakeEnvelope();
+		var liveEnvelope = new FakeEnvelope();
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+
+		var stoppedId = Guid.NewGuid();
+		sub.AddClient(stoppedId, Guid.NewGuid(), "connection-stopped", stoppedEnvelope, 10, "foo", "bar");
+		sub.AddClient(Guid.NewGuid(), Guid.NewGuid(), "connection-live", liveEnvelope, 10, "foo", "bar");
+
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(0, _eventSource));
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(1, _eventSource));
+		pushScheduler.Push(sub);
+
+		var liveCountBeforeStop = liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count();
+
+		sub.StopClient(stoppedId);
+		sub.RemoveClientByCorrelationId(stoppedId, sendDropNotification: false);
+		pushScheduler.Push(sub);
+
+		Assert.AreEqual(liveCountBeforeStop + 1,
+			liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+		Assert.AreEqual(1,
+			((ClientMessage.PersistentSubscriptionStreamEventAppeared)liveEnvelope.Replies.Last()).RetryCount);
+	}
+
+	[Test]
+	public void stopping_only_consumer_holds_new_events_until_a_consumer_joins() {
+		var stoppedEnvelope = new FakeEnvelope();
+		var pushScheduler = new FakePushScheduler();
+		var sub = BuildSubscription(pushScheduler);
+
+		var stoppedId = Guid.NewGuid();
+		sub.AddClient(stoppedId, Guid.NewGuid(), "connection-stopped", stoppedEnvelope, 10, "foo", "bar");
+		sub.StopClient(stoppedId);
+
+		sub.NotifyLiveSubscriptionMessage(Helper.GetFakeEventFor(0, _eventSource));
+		pushScheduler.Push(sub);
+
+		Assert.AreEqual(0, stoppedEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+
+		// A new active consumer joining should pick up the buffered event.
+		var liveEnvelope = new FakeEnvelope();
+		sub.AddClient(Guid.NewGuid(), Guid.NewGuid(), "connection-live", liveEnvelope, 10, "foo", "bar");
+		pushScheduler.Push(sub);
+
+		Assert.AreEqual(1, liveEnvelope.Replies.OfType<ClientMessage.PersistentSubscriptionStreamEventAppeared>().Count());
+	}
+}
+
 [TestFixture(EventSource.SingleStream)]
 [TestFixture(EventSource.AllStream)]
 [TestFixture(EventSource.FilteredAllStream)]