diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java index 42988cabfd96f..839cff2dc8d33 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java @@ -142,6 +142,7 @@ import org.apache.pulsar.common.api.proto.CommandScalableTopicClose; import org.apache.pulsar.common.api.proto.CommandScalableTopicLookup; import org.apache.pulsar.common.api.proto.CommandScalableTopicSubscribe; +import org.apache.pulsar.common.api.proto.CommandScalableTopicUnsubscribe; import org.apache.pulsar.common.api.proto.CommandSeek; import org.apache.pulsar.common.api.proto.CommandSend; import org.apache.pulsar.common.api.proto.CommandSubscribe; @@ -516,15 +517,22 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { if (!scalableConsumerRegistrations.isEmpty()) { var scalableTopicService = service.getScalableTopicService(); if (scalableTopicService != null) { - scalableConsumerRegistrations.values().forEach(ref -> { - try { - scalableTopicService.onConsumerDisconnect( - ref.topicName(), ref.subscription(), ref.consumerName()); - } catch (Exception e) { - log.warn().attr("consumerName", ref.consumerName()).exceptionMessage(e) - .log("Error notifying scalable controller of consumer disconnect"); - } - }); + scalableConsumerRegistrations.values().forEach(ref -> + // Chained on the registration outcome: a registration still in flight + // when the connection dies creates its session only afterwards, and + // the disconnect report must not race ahead of it (it would no-op on + // a not-yet-existing session and never arm the grace timer). + ref.registration().whenComplete((__, ___) -> { + try { + scalableTopicService.onConsumerDisconnect( + ref.topicName(), ref.subscription(), ref.consumerName()); + } catch (Exception e) { + log.warn().attr("consumerName", ref.consumerName()) + .exceptionMessage(e) + .log("Error notifying scalable controller of consumer " + + "disconnect"); + } + })); } scalableConsumerRegistrations.clear(); } @@ -1093,7 +1101,8 @@ protected void handleCommandScalableTopicClose( private record ScalableConsumerRegistrationRef( TopicName topicName, String subscription, - String consumerName) {} + String consumerName, + CompletableFuture registration) {} @Override protected void handleCommandScalableTopicSubscribe( @@ -1148,24 +1157,32 @@ protected void handleCommandScalableTopicSubscribe( ServerError.AuthorizationError, msg); return; } - scalableTopicService.registerConsumer(topicName, subscription, consumerName, - consumerId, consumerType, this) - .whenCompleteAsync((assignment, ex) -> { - if (ex != null) { - Throwable cause = ex.getCause() != null ? ex.getCause() : ex; - log.warn().attr("topic", topicName).attr("subscription", subscription) - .attr("consumerName", consumerName).exception(cause) - .log("ScalableTopicSubscribe failed"); - getCommandSender().sendScalableTopicSubscribeError(requestId, - ServerError.UnknownError, cause.getMessage()); - return; - } - // Record the registration so we can call onConsumerDisconnect on channelInactive. - scalableConsumerRegistrations.put(consumerId, - new ScalableConsumerRegistrationRef(topicName, subscription, consumerName)); - getCommandSender().sendScalableTopicSubscribeResponse(requestId, - ConsumerSession.toProto(assignment)); - }, ctx.executor()); + // Record the registration BEFORE it resolves, carrying its future: an + // unsubscribe (or the channelInactive sweep) arriving mid-registration + // chains behind it instead of silently missing it. The client's subscribe + // can time out while the broker-side registration is still in flight, so + // this ordering must not depend on how long the client was able to wait. + var registration = scalableTopicService.registerConsumer(topicName, + subscription, consumerName, consumerId, consumerType, this); + var ref = new ScalableConsumerRegistrationRef( + topicName, subscription, consumerName, registration); + scalableConsumerRegistrations.put(consumerId, ref); + registration.whenCompleteAsync((assignment, ex) -> { + if (ex != null) { + Throwable cause = ex.getCause() != null ? ex.getCause() : ex; + log.warn().attr("topic", topicName).attr("subscription", subscription) + .attr("consumerName", consumerName).exception(cause) + .log("ScalableTopicSubscribe failed"); + // Nothing was registered: drop the ref so unsubscribes and the + // disconnect sweep have nothing to report for it. + scalableConsumerRegistrations.remove(consumerId, ref); + getCommandSender().sendScalableTopicSubscribeError(requestId, + ServerError.UnknownError, cause.getMessage()); + return; + } + getCommandSender().sendScalableTopicSubscribeResponse(requestId, + ConsumerSession.toProto(assignment)); + }, ctx.executor()); }) .exceptionally(ex -> { logAuthException(remoteAddress, "scalable-topic-subscribe", getPrincipal(), @@ -1177,6 +1194,49 @@ protected void handleCommandScalableTopicSubscribe( }); } + @Override + protected void handleCommandScalableTopicUnsubscribe( + CommandScalableTopicUnsubscribe commandScalableTopicUnsubscribe) { + checkArgument(state == State.Connected); + final long requestId = commandScalableTopicUnsubscribe.getRequestId(); + final long consumerId = commandScalableTopicUnsubscribe.getConsumerId(); + + // The lookup is scoped to this connection's own registrations, so a client can only + // unregister sessions it created here — no further authorization is needed. + ScalableConsumerRegistrationRef ref = scalableConsumerRegistrations.get(consumerId); + var scalableTopicService = service.getScalableTopicService(); + if (ref == null || scalableTopicService == null) { + // Unknown or already swept by a disconnect: idempotent success. + getCommandSender().sendSuccessResponse(requestId); + return; + } + log.debug().attr("topic", ref.topicName()).attr("subscription", ref.subscription()) + .attr("consumerName", ref.consumerName()).attr("requestId", requestId) + .log("Received ScalableTopicUnsubscribe"); + // Ordered behind the (possibly still in-flight) registration; a failed registration + // has nothing to unregister and the idempotent unregister below tolerates that. + ref.registration().handle((__, ___) -> (Void) null) + .thenCompose(__ -> scalableTopicService.unregisterConsumer( + ref.topicName(), ref.subscription(), ref.consumerName(), consumerId)) + .whenCompleteAsync((__, ex) -> { + if (ex != null) { + // Keep the ref: the channelInactive sweep can still report the + // disconnect, so the grace-period fallback stays alive for a + // registration the explicit unregister failed to delete. + Throwable cause = ex.getCause() != null ? ex.getCause() : ex; + log.warn().attr("consumerName", ref.consumerName()).exceptionMessage(cause) + .log("ScalableTopicUnsubscribe failed"); + getCommandSender().sendErrorResponse(requestId, ServerError.UnknownError, + cause.getMessage()); + return; + } + // Removed only on success; a channelInactive racing the unregister just + // re-reports an already-removed session, which the coordinator ignores. + scalableConsumerRegistrations.remove(consumerId, ref); + getCommandSender().sendSuccessResponse(requestId); + }, ctx.executor()); + } + @Override protected void handlePartitionMetadataRequest(CommandPartitionedTopicMetadata partitionMetadataParam) { checkArgument(state == State.Connected); diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicController.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicController.java index aebeca97c6ab8..a8a599fff8318 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicController.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicController.java @@ -811,13 +811,14 @@ public CompletableFuture registerConsumer(String subscriptio * Explicit unregister: the consumer is leaving the subscription for good. Deletes the * persisted session entry and rebalances remaining consumers. */ - public CompletableFuture unregisterConsumer(String subscription, String consumerName) { + public CompletableFuture unregisterConsumer(String subscription, String consumerName, + long consumerId) { checkLeader(); SubscriptionCoordinator coordinator = subscriptions.get(subscription); if (coordinator == null) { return CompletableFuture.completedFuture(null); } - return coordinator.unregisterConsumer(consumerName) + return coordinator.unregisterConsumer(consumerName, consumerId) .thenAccept(__ -> { if (coordinator.getConsumers().isEmpty()) { subscriptions.remove(subscription); @@ -882,7 +883,8 @@ public CompletableFuture deleteSubscription(String subscription) { private CompletableFuture dropAllConsumers(SubscriptionCoordinator coordinator) { CompletableFuture[] futures = coordinator.getConsumers().stream() - .map(session -> coordinator.unregisterConsumer(session.getConsumerName())) + .map(session -> coordinator.unregisterConsumer( + session.getConsumerName(), session.getConsumerId())) .toArray(CompletableFuture[]::new); return CompletableFuture.allOf(futures); } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicService.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicService.java index 3e5081e3985d2..e53364f001638 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicService.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicService.java @@ -316,6 +316,23 @@ public void onConsumerDisconnect(TopicName topic, String subscription, String co } } + /** + * Explicit clean leave: forwards to the locally-held controller, which deletes the + * persisted registration and rebalances the remaining consumers immediately. No-op when + * no controller entry exists here — which happens only for a consumer that never + * registered on this broker, a deleted topic, or a shutting-down service. (A deposed + * leader keeps its entry and fails via {@code checkLeader()}, taking the error path + * instead, which preserves the caller's registration ref and grace fallback.) + */ + public CompletableFuture unregisterConsumer(TopicName topic, String subscription, + String consumerName, long consumerId) { + CompletableFuture future = controllers.get(topic.toString()); + if (future == null) { + return CompletableFuture.completedFuture(null); + } + return future.thenCompose(c -> c.unregisterConsumer(subscription, consumerName, consumerId)); + } + // --- Internal helpers --- private CompletableFuture createUnderlyingSegmentTopic(TopicName parentTopic, SegmentInfo segment) { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java index 22a8ccea0f511..8166e9a2d0b11 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java @@ -214,19 +214,49 @@ public synchronized CompletableFuture> } /** - * Explicit unregister (consumer asked to leave the subscription). Cancels any pending - * grace timer, deletes the persisted registration, and rebalances. + * Explicit unregister (consumer asked to leave the subscription). Deletes the persisted + * registration first, and only on success removes the in-memory session, cancels its + * grace timer, and rebalances. A failed delete therefore changes nothing: the session + * stays registered and connected, so the channelInactive → grace-period fallback still + * works, and a retried unregister actually retries the deletion instead of short- + * circuiting on an already-removed session. + * + *

{@code expectedConsumerId} guards the removal against a same-name rejoin racing the + * in-flight delete: a re-register attaches a new consumer id to the session, so if the + * id no longer matches when the delete completes, the departed consumer's leave must not + * take the rejoined consumer down with it — the session is kept and the persisted + * registration the delete just erased is restored. */ public synchronized CompletableFuture> unregisterConsumer( - String consumerName) { - ConsumerSession removed = sessions.remove(consumerName); - if (removed == null) { + String consumerName, long expectedConsumerId) { + ConsumerSession session = sessions.get(consumerName); + if (session == null || session.getConsumerId() != expectedConsumerId) { return CompletableFuture.completedFuture(snapshotAssignments()); } - removed.cancelGraceTimer(); return resources.unregisterConsumerAsync(topicName, subscriptionName, consumerName) .thenApply(__ -> { synchronized (this) { + ConsumerSession current = sessions.get(consumerName); + if (current != null && current.getConsumerId() != expectedConsumerId) { + // A same-name consumer rejoined while the delete was in flight + // (the reconnect branch attached a new id). Keep it, and restore + // the persisted registration the delete just erased so a + // controller failover still knows this member. + resources.registerConsumerAsync(topicName, subscriptionName, + consumerName) + .exceptionally(ex -> { + log.warn().attr("consumer", consumerName) + .exceptionMessage(ex) + .log("Failed to restore the rejoined consumer's " + + "persisted registration"); + return null; + }); + return snapshotAssignments(); + } + ConsumerSession removed = sessions.remove(consumerName); + if (removed != null) { + removed.cancelGraceTimer(); + } if (sessions.isEmpty()) { segmentAssignments.clear(); return Map.of(); diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ServerCnxTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ServerCnxTest.java index 7de4c198f30df..542b3f9d85908 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ServerCnxTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/ServerCnxTest.java @@ -1591,6 +1591,118 @@ public void testScalableTopicCommandsRequireTopicAuthorization() throws Exceptio channel.finish(); } + /** + * PIP-486 clean leave ordering: an unsubscribe arriving while the registration is still + * in flight must chain behind it — never answer as an idempotent no-op and let the + * registration complete afterwards (the client's subscribe can time out while the + * broker-side registration is still running). + */ + @Test(timeOut = 30000) + public void testScalableTopicUnsubscribeChainsBehindInFlightRegistration() throws Exception { + var scalableTopicService = + mock(org.apache.pulsar.broker.service.scalable.ScalableTopicService.class); + when(brokerService.getScalableTopicService()).thenReturn(scalableTopicService); + var registration = new CompletableFuture< + org.apache.pulsar.broker.service.scalable.ConsumerAssignment>(); + when(scalableTopicService.registerConsumer(any(), anyString(), anyString(), anyLong(), + any(), any())).thenReturn(registration); + when(scalableTopicService.unregisterConsumer(any(), anyString(), anyString(), anyLong())) + .thenReturn(CompletableFuture.completedFuture(null)); + + resetChannel(); + channel.writeInbound(Commands.newConnect("none", "", null)); + assertTrue(getResponse() instanceof CommandConnected); + + channel.writeInbound(Commands.newScalableTopicSubscribe(400L, + "persistent://public/default/scalable-inflight", "sub", "c1", 7L, + ScalableConsumerType.STREAM)); + channel.writeInbound(Commands.newScalableTopicUnsubscribe(401L, 7L)); + channel.runPendingTasks(); + assertTrue(channel.outboundMessages().isEmpty(), + "no response may be sent while the registration is still in flight"); + verify(scalableTopicService, times(0)) + .unregisterConsumer(any(), anyString(), anyString(), anyLong()); + + // The registration completes: the subscribe answers, and only then does the queued + // unsubscribe run its unregister and answer success. + registration.complete(new org.apache.pulsar.broker.service.scalable.ConsumerAssignment( + 1L, Collections.emptyList())); + channel.runPendingTasks(); + // Both answers are released now (their relative order is a CompletableFuture + // callback-ordering detail): the subscribe response, and the unsubscribe success. + Object first = getResponse(); + Object second = getResponse(); + assertTrue(first instanceof CommandScalableTopicSubscribeResponse + || second instanceof CommandScalableTopicSubscribeResponse, + first + " / " + second); + CommandSuccess success = (CommandSuccess) + (first instanceof CommandSuccess ? first : second); + assertEquals(success.getRequestId(), 401L); + verify(scalableTopicService, times(1)) + .unregisterConsumer(any(), anyString(), anyString(), eq(7L)); + + channel.finish(); + } + + /** + * PIP-486 clean leave: unsubscribe for an unknown consumer id is an idempotent success; a + * failed unregister answers an error and keeps the per-connection registration ref (so the + * grace fallback and a retry both still work); a successful retry then removes it, and a + * repeat unsubscribe is again an idempotent success. + */ + @Test(timeOut = 30000) + public void testScalableTopicUnsubscribeIdempotencyAndErrorPath() throws Exception { + var scalableTopicService = + mock(org.apache.pulsar.broker.service.scalable.ScalableTopicService.class); + when(brokerService.getScalableTopicService()).thenReturn(scalableTopicService); + + resetChannel(); + ByteBuf connect = Commands.newConnect("none", "", null); + channel.writeInbound(connect); + assertTrue(getResponse() instanceof CommandConnected); + + // Unknown consumer id: idempotent success. + channel.writeInbound(Commands.newScalableTopicUnsubscribe(300L, 999L)); + Object response = getResponse(); + assertTrue(response instanceof CommandSuccess, String.valueOf(response)); + assertEquals(((CommandSuccess) response).getRequestId(), 300L); + + // Register a consumer so the connection records its registration ref. + when(scalableTopicService.registerConsumer(any(), anyString(), anyString(), anyLong(), + any(), any())).thenReturn(CompletableFuture.completedFuture( + new org.apache.pulsar.broker.service.scalable.ConsumerAssignment( + 1L, Collections.emptyList()))); + channel.writeInbound(Commands.newScalableTopicSubscribe(301L, + "persistent://public/default/scalable-unsub", "sub", "c1", 5L, + ScalableConsumerType.STREAM)); + assertTrue(getResponse() instanceof CommandScalableTopicSubscribeResponse); + + // Failing unregister: error response, ref retained. + when(scalableTopicService.unregisterConsumer(any(), anyString(), anyString(), anyLong())) + .thenReturn(CompletableFuture.failedFuture(new RuntimeException("store down"))); + channel.writeInbound(Commands.newScalableTopicUnsubscribe(302L, 5L)); + response = getResponse(); + assertTrue(response instanceof CommandError, String.valueOf(response)); + assertEquals(((CommandError) response).getRequestId(), 302L); + + // Retry after the failure actually retries the unregister (the ref survived). + when(scalableTopicService.unregisterConsumer(any(), anyString(), anyString(), anyLong())) + .thenReturn(CompletableFuture.completedFuture(null)); + channel.writeInbound(Commands.newScalableTopicUnsubscribe(303L, 5L)); + response = getResponse(); + assertTrue(response instanceof CommandSuccess, String.valueOf(response)); + assertEquals(((CommandSuccess) response).getRequestId(), 303L); + verify(scalableTopicService, times(2)).unregisterConsumer(any(), anyString(), anyString(), anyLong()); + + // The ref is gone now: one more unsubscribe is an idempotent success with no new call. + channel.writeInbound(Commands.newScalableTopicUnsubscribe(304L, 5L)); + response = getResponse(); + assertTrue(response instanceof CommandSuccess, String.valueOf(response)); + verify(scalableTopicService, times(2)).unregisterConsumer(any(), anyString(), anyString(), anyLong()); + + channel.finish(); + } + @Test public void testRefreshOriginalPrincipalWithAuthDataForwardedFromProxy() throws Exception { AuthenticationService authenticationService = mock(AuthenticationService.class); diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/ScalableTopicControllerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/ScalableTopicControllerTest.java index 52d7a25adb9b7..68e203316dfc8 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/ScalableTopicControllerTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/ScalableTopicControllerTest.java @@ -267,7 +267,7 @@ public void testWriteOperationThrowsWhenNotLeader() { assertThrows(IllegalStateException.class, () -> controller.registerConsumer( "sub", "c1", 1L, ScalableConsumerType.STREAM, mock(TransportCnx.class))); - assertThrows(IllegalStateException.class, () -> controller.unregisterConsumer("sub", "c1")); + assertThrows(IllegalStateException.class, () -> controller.unregisterConsumer("sub", "c1", 1L)); } // --- Consumer registration --- @@ -307,7 +307,7 @@ public void testUnregisterConsumerDeletesPersistedEntry() throws Exception { controller.registerConsumer("sub-a", "c2", 2L, ScalableConsumerType.STREAM, mock(TransportCnx.class)).get(); assertEquals(resources.listConsumersAsync(topicName, "sub-a").get().size(), 2); - controller.unregisterConsumer("sub-a", "c1").get(); + controller.unregisterConsumer("sub-a", "c1", 1L).get(); assertEquals(resources.listConsumersAsync(topicName, "sub-a").get(), List.of("c2")); } @@ -316,7 +316,7 @@ public void testUnregisterConsumerDeletesPersistedEntry() throws Exception { public void testUnregisterConsumerUnknownSubscriptionIsNoop() throws Exception { controller.initialize().get(); // No subscription 'ghost' exists; call should complete without error. - controller.unregisterConsumer("ghost", "c1").get(); + controller.unregisterConsumer("ghost", "c1", 1L).get(); } @Test diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java index b821820e556cc..87bddbf21d55c 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinatorTest.java @@ -129,7 +129,7 @@ public void testUnregisterConsumerRebalances() throws Exception { coordinator.registerConsumer("consumer-2", 2L, mock(TransportCnx.class)).get(); Map result = - coordinator.unregisterConsumer("consumer-2").get(); + coordinator.unregisterConsumer("consumer-2", 2L).get(); assertEquals(result.size(), 1); assertEquals(findByName(result, "consumer-1").assignedSegments().size(), 4); @@ -332,7 +332,7 @@ public void testInstallDrainCheckerAfterRestoreEnablesOrdering() throws Exceptio public void testEmptyAfterAllConsumersRemoved() throws Exception { coordinator.registerConsumer("consumer-1", 1L, mock(TransportCnx.class)).get(); Map result = - coordinator.unregisterConsumer("consumer-1").get(); + coordinator.unregisterConsumer("consumer-1", 1L).get(); assertTrue(result.isEmpty()); assertTrue(coordinator.getConsumers().isEmpty()); @@ -516,6 +516,27 @@ private SubscriptionCoordinator bucketedCoordinator() { // --- Helpers --- + /** + * PIP-486 review: the identity guard on explicit unregister. A same-name rejoin attaches + * a new consumer id to the session; a leave carrying the OLD id (the departed consumer's) + * must not remove the rejoined session, while a leave with the current id must. + */ + @Test + public void testUnregisterWithStaleConsumerIdKeepsRejoinedSession() throws Exception { + coordinator.registerConsumer("c1", 1L, mock(TransportCnx.class)).get(); + // Same-name rejoin: the reconnect branch attaches the new id to the existing session. + coordinator.registerConsumer("c1", 2L, mock(TransportCnx.class)).get(); + + // The departed consumer's leave (id 1) arrives after the rejoin: no-op. + coordinator.unregisterConsumer("c1", 1L).get(); + assertEquals(coordinator.getConsumers().size(), 1, + "a stale-id leave must not remove the rejoined session"); + + // A leave with the live id removes it. + coordinator.unregisterConsumer("c1", 2L).get(); + assertEquals(coordinator.getConsumers().size(), 0); + } + private static ConsumerAssignment findByName(Map m, String name) { return m.entrySet().stream() .filter(e -> name.equals(e.getKey().getConsumerName())) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java index 9e52890b7bda7..705830ae89bf1 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/v5/V5EntryBucketDispatchTest.java @@ -444,21 +444,6 @@ public void testConsumerRejoiningAfterLeaveDoesNotWedgeRelease() throws Exceptio // watermark) must not leak into the last flip's drain: with everything acked and nothing // outstanding, a stale delivered-vs-acked pair would make the drain wait forever and the // rejoining consumer could never attach. - // A consumer's clean close is a disconnect to the controller, which holds its session for - // the grace period before rebalancing. Shrink it (read at coordinator creation, i.e. on - // this subscription's first register) so B1's departure hands the segment back promptly. - int defaultGrace = getPulsar().getConfiguration() - .getScalableTopicConsumerSessionGracePeriodSeconds(); - getPulsar().getConfiguration().setScalableTopicConsumerSessionGracePeriodSeconds(1); - try { - runRejoinAfterLeaveScenario(); - } finally { - getPulsar().getConfiguration() - .setScalableTopicConsumerSessionGracePeriodSeconds(defaultGrace); - } - } - - private void runRejoinAfterLeaveScenario() throws Exception { String topic = newScalableTopic(1); admin.scalableTopics().setAutoScalePolicy(topic, AutoScalePolicyOverride.builder().enabled(false).build()); @@ -497,16 +482,16 @@ private void runRejoinAfterLeaveScenario() throws Exception { } a.acknowledgeCumulative(last.id()); - // Phase 2 — B1 joins (shared, individual acks), both drain, then B1 leaves. B1 gets its - // own client: a departure is only visible to the controller as a connection drop, so - // leaving means closing the whole client (the shared client's pooled connection would - // keep B1's registration alive indefinitely). - PulsarClient b1Client = newV5Client(); - StreamConsumer b1 = b1Client.newStreamConsumer(Schema.string()) + // Phase 2 — B1 joins (shared, individual acks), both drain, then B1 leaves. B1 uses the + // shared client on purpose: its close must reach the controller through the explicit + // unsubscribe (the pooled controller connection stays open, so without it the + // registration would linger for the full disconnect grace period and the segment + // would never be handed back within this test's window). + StreamConsumer b1 = track(v5Client.newStreamConsumer(Schema.string()) .topic(topic) .subscriptionName(subscription) .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) - .subscribeAsync().get(30, TimeUnit.SECONDS); + .subscribeAsync().get(30, TimeUnit.SECONDS)); sendPhase(producer, keys, sent, perKey, perKey * 2); Map> b1Got = new ConcurrentHashMap<>(); Thread ta1 = drainOrdered(a, aGot); @@ -515,9 +500,9 @@ private void runRejoinAfterLeaveScenario() throws Exception { tb1.join(); assertFalse(b1Got.isEmpty(), "consumer B1 received nothing — the segment did not fan out"); b1.close(); - b1Client.close(); - // Wait until the controller has handed the whole segment back to A and A completed the - // flip back to Exclusive — the stale-watermark state only matters once that is done. + // The clean leave unregisters immediately (no grace wait): the controller hands the + // whole segment back to A, which flips back to Exclusive. Only once that is done does + // the stale-watermark state matter for B2's rejoin. String segmentTopic = admin.scalableTopics().getStats(topic) .getSegments().values().iterator().next().name(); Awaitility.await().atMost(Duration.ofSeconds(15)).untilAsserted(() -> { @@ -559,6 +544,58 @@ private void runRejoinAfterLeaveScenario() throws Exception { } } + @Test + public void testRapidSubscribeCloseChurnLeavesGroupClean() throws Exception { + // PIP-486 clean leave under churn: consumers joining and closing in quick succession — + // including closes racing the registration and rebalance machinery — must leave no + // ghost group members behind. A single ghost would make the controller keep the + // segment fanned out (Key_Shared) instead of returning it to the survivor Exclusive. + String topic = newScalableTopic(1); + admin.scalableTopics().setAutoScalePolicy(topic, + AutoScalePolicyOverride.builder().enabled(false).build()); + String subscription = "leave-churn"; + + @Cleanup + StreamConsumer survivor = v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribe(); + for (int i = 0; i < 5; i++) { + v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribeAsync().get(30, TimeUnit.SECONDS) + .close(); + } + // Same-name churn: close-then-resubscribe under one explicit consumerName races the + // leave's delete against the rejoin's reconnect attach — the identity guard must keep + // the rejoined member alive, and the final close must still remove it. + for (int i = 0; i < 3; i++) { + v5Client.newStreamConsumer(Schema.string()) + .topic(topic) + .subscriptionName(subscription) + .consumerName("same-name-churner") + .subscriptionInitialPosition(SubscriptionInitialPosition.EARLIEST) + .subscribeAsync().get(30, TimeUnit.SECONDS) + .close(); + } + + // Every churned consumer unregistered cleanly: the survivor converges back to sole + // Exclusive ownership. With the default 60s grace, any missed unregistration would + // hold the group fanned out well past this window. + String segmentTopic = admin.scalableTopics().getStats(topic) + .getSegments().values().iterator().next().name(); + Awaitility.await().atMost(Duration.ofSeconds(20)).untilAsserted(() -> { + var sub = getTopicReference(segmentTopic).orElseThrow().getSubscription(subscription); + assertNotNull(sub, "segment subscription missing"); + assertEquals(sub.getType(), CommandSubscribe.SubType.Exclusive, + "a churned consumer left a ghost registration behind"); + assertEquals(sub.getConsumers().size(), 1); + }); + } + private void sendPhase(Producer producer, List keys, Map> sent, int from, int to) throws Exception { for (int i = from; i < to; i++) { diff --git a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/CheckpointConsumerBuilderV5.java b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/CheckpointConsumerBuilderV5.java index ce66b6fb608cc..e87a41231dde1 100644 --- a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/CheckpointConsumerBuilderV5.java +++ b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/CheckpointConsumerBuilderV5.java @@ -81,7 +81,14 @@ public CompletableFuture> createAsync() { return session.start() .thenCompose(initialAssignment -> ScalableCheckpointConsumer.createManagedAsync( client, v5Schema, topic.toString(), session, initialAssignment, - startPosition, name)); + startPosition, name)) + .whenComplete((consumer, ex) -> { + if (ex != null) { + // See StreamConsumerBuilderV5: a failed subscribe must still send + // its clean unsubscribe once the in-flight attempt settles. + session.close(); + } + }); } // Unmanaged: read every active segment, no broker-side state. diff --git a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableConsumerClient.java b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableConsumerClient.java index 1597a7b083c68..af2c7242b3b21 100644 --- a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableConsumerClient.java +++ b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/ScalableConsumerClient.java @@ -25,15 +25,17 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import org.apache.pulsar.client.api.PulsarClientException; import org.apache.pulsar.client.impl.ClientCnx; import org.apache.pulsar.client.impl.PulsarClientImpl; import org.apache.pulsar.client.impl.ScalableConsumerSession; import org.apache.pulsar.client.impl.v5.SegmentRouter.ActiveSegment; -import org.apache.pulsar.client.util.TimedCompletableFuture; import org.apache.pulsar.common.api.proto.ScalableAssignedSegment; import org.apache.pulsar.common.api.proto.ScalableConsumerAssignment; import org.apache.pulsar.common.api.proto.ScalableConsumerType; @@ -83,6 +85,15 @@ final class ScalableConsumerClient implements ScalableConsumerSession, AutoClose private volatile AssignmentChangeListener listener; private volatile ClientCnx cnx; private volatile boolean closed = false; + /** + * Every subscribe attempt still in flight; close() awaits all of them before + * unsubscribing (a single last-attempt slot can be overwritten when two reconnect + * attempts overlap, letting the older one register after the unsubscribe was sent). + */ + private final Set> inFlightSubscribes = + ConcurrentHashMap.newKeySet(); + /** Collapses concurrent reconnect triggers (request failure + connectionClosed) into one. */ + private final AtomicBoolean reconnectPending = new AtomicBoolean(false); ScalableConsumerClient(PulsarClientImpl v4Client, TopicName topicName, @@ -136,6 +147,14 @@ CompletableFuture> start() { */ private CompletableFuture connectAndSubscribe() { CompletableFuture result = new CompletableFuture<>(); + // Registered before the closed-flag check below: close() chains its clean unsubscribe + // on every attempt in this set, so either an attempt observes closed and never + // subscribes, or close() observes the attempt and waits for its outcome — never an + // unsubscribe that overtakes an in-flight registration (which would no-op and leave a + // ghost registration behind). Settled attempts drop out; their registrations, if any, + // already exist and are covered by the unsubscribe that follows. + inFlightSubscribes.add(result); + result.whenComplete((__, ___) -> inFlightSubscribes.remove(result)); DagWatchClient watch = new DagWatchClient(v4Client, topicName); watch.start() @@ -186,32 +205,26 @@ private CompletableFuture connectAndSubscribe() { cnx.registerScalableConsumerSession(consumerId, this); long requestId = v4Client.newRequestId(); - var responseFuture = new TimedCompletableFuture(); - cnx.getPendingRequests().put(requestId, responseFuture); - - cnx.ctx().writeAndFlush(Commands.newScalableTopicSubscribe( - requestId, - topicName.toString(), - subscription, - consumerName, - consumerId, - consumerType)) - .addListener(writeFuture -> { - if (!writeFuture.isSuccess()) { - cnx.getPendingRequests().remove(requestId); + cnx.sendScalableSessionRequest( + Commands.newScalableTopicSubscribe( + requestId, + topicName.toString(), + subscription, + consumerName, + consumerId, + consumerType), + requestId) + .whenComplete((assignment, ex) -> { + if (ex != null) { + // Write failure, error response, or request timeout: the + // broker holds no push route for us, so drop the local + // session registration too (retries re-register). cnx.removeScalableConsumerSession(consumerId); - result.completeExceptionally( - new PulsarClientException(writeFuture.cause())); + result.completeExceptionally(ex); + } else { + result.complete(assignment); } }); - - responseFuture.whenComplete((assignment, ex) -> { - if (ex != null) { - result.completeExceptionally(ex); - } else { - result.complete(assignment); - } - }); }) .exceptionally(ex -> { result.completeExceptionally(ex); @@ -280,6 +293,12 @@ private void scheduleReconnect() { if (closed) { return; } + // A connection drop triggers this twice (the failed pending request and the + // connectionClosed callback); collapse to a single scheduled attempt so two + // overlapping connectAndSubscribe() calls never run. + if (!reconnectPending.compareAndSet(false, true)) { + return; + } long delayMs = reconnectBackoff.next().toMillis(); log.info().attr("delayMs", delayMs).log("Scheduling reconnect"); v4Client.timer().newTimeout(timeout -> reconnect(), @@ -287,6 +306,7 @@ private void scheduleReconnect() { } private void reconnect() { + reconnectPending.set(false); if (closed) { return; } @@ -374,11 +394,50 @@ public void close() { return; } closed = true; + // Deregister the local session immediately and unconditionally (as before the clean + // leave existed): the subscribe response routes through pendingRequests, not through + // the session registry, so nothing here needs to wait — and waiting on a subscribe + // that is never answered would retain the closed session (and pin the pooled + // connection non-idle) for the connection's lifetime. ClientCnx c = cnx; if (c != null) { c.removeScalableConsumerSession(consumerId); - // No close command for now — broker reaps registrations via grace timer on - // disconnect. A future refactor can add an explicit unsubscribe. + } + // Clean leave: tell the controller to unregister this consumer and rebalance the group + // immediately, instead of holding the registration for the disconnect grace period (the + // controller connection is pooled, so closing this consumer does not close the channel + // and the broker would otherwise never notice the departure while the client lives). + // Chained on every in-flight subscribe attempt so the unsubscribe can never overtake a + // registration (the broker records it before sending the subscribe response). Sent even + // when a subscribe failed — the command is idempotent, and skipping on a razor-edge + // failure would risk leaving a registration behind. Best-effort throughout: on any + // failure the grace period remains the fallback for the eventual real disconnect. + CompletableFuture[] pending = inFlightSubscribes.toArray(CompletableFuture[]::new); + if (pending.length == 0) { + sendUnsubscribe(); + } else { + CompletableFuture.allOf(pending).whenComplete((__, ___) -> sendUnsubscribe()); + } + } + + private void sendUnsubscribe() { + ClientCnx c = cnx; + if (c == null) { + return; + } + c.removeScalableConsumerSession(consumerId); + try { + long requestId = v4Client.newRequestId(); + c.sendScalableSessionRequest( + Commands.newScalableTopicUnsubscribe(requestId, consumerId), requestId) + .exceptionally(ex -> { + log.debug().exceptionMessage(ex) + .log("Clean unsubscribe failed; relying on the grace period"); + return null; + }); + } catch (Exception e) { + log.debug().exceptionMessage(e) + .log("Clean unsubscribe failed; relying on the grace period"); } } diff --git a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/StreamConsumerBuilderV5.java b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/StreamConsumerBuilderV5.java index 13ab896800c33..8e0d4feca0652 100644 --- a/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/StreamConsumerBuilderV5.java +++ b/pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/StreamConsumerBuilderV5.java @@ -97,7 +97,17 @@ public CompletableFuture> subscribeAsync() { return session.start() .thenCompose(initialAssignment -> ScalableStreamConsumer.createAsync( - client, v5Schema, conf, session, topic.toString(), initialAssignment)); + client, v5Schema, conf, session, topic.toString(), initialAssignment)) + .whenComplete((consumer, ex) -> { + if (ex != null) { + // An abandoned subscribe must not leave a controller registration + // behind: close() sends the clean unsubscribe once the in-flight + // attempt settles — covering a registration the broker completed + // after the client's subscribe timed out. Idempotent if the consumer + // path already closed the session. + session.close(); + } + }); } @Override diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java index 719bd728ff500..4bd1396e032fe 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ClientCnx.java @@ -1468,6 +1468,16 @@ CompletableFuture sendRequestWithId(ByteBuf cmd, long requestI return sendRequestAndHandleTimeout(cmd, requestId, RequestType.Command, true); } + /** + * Send a scalable-topic session request (PIP-468/486) with the standard request + * bookkeeping — pending-request registration, request-timeout tracking, and + * write-failure completion — instead of hand-rolling it at the call site. The matching + * response handler completes the returned future by request id. + */ + public CompletableFuture sendScalableSessionRequest(ByteBuf requestMessage, long requestId) { + return sendRequestAndHandleTimeout(requestMessage, requestId, RequestType.Command, true); + } + private void sendRequestAndHandleTimeout(ByteBuf requestMessage, long requestId, RequestType requestType, boolean flush, TimedCompletableFuture future) { diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java index 575664e8b0930..5fff50798780b 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java @@ -1804,6 +1804,19 @@ public static ByteBuf newScalableTopicSubscribe(long requestId, String topic, return serializeWithSize(cmd); } + /** + * Client -> Broker: a scalable consumer is cleanly leaving its subscription; the + * controller unregisters it and rebalances immediately instead of waiting out the + * disconnect grace period. Acknowledged with {@code CommandSuccess}. + */ + public static ByteBuf newScalableTopicUnsubscribe(long requestId, long consumerId) { + BaseCommand cmd = localCmd(Type.SCALABLE_TOPIC_UNSUBSCRIBE); + cmd.setScalableTopicUnsubscribe() + .setRequestId(requestId) + .setConsumerId(consumerId); + return serializeWithSize(cmd); + } + /** * Broker -> Client: response to a scalable-topic subscribe request. On success the * caller must populate the nested {@link ScalableConsumerAssignment} via diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/PulsarDecoder.java b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/PulsarDecoder.java index 9a412c1def16e..7cf04d830d1c6 100644 --- a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/PulsarDecoder.java +++ b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/PulsarDecoder.java @@ -77,6 +77,7 @@ import org.apache.pulsar.common.api.proto.CommandScalableTopicLookup; import org.apache.pulsar.common.api.proto.CommandScalableTopicSubscribe; import org.apache.pulsar.common.api.proto.CommandScalableTopicSubscribeResponse; +import org.apache.pulsar.common.api.proto.CommandScalableTopicUnsubscribe; import org.apache.pulsar.common.api.proto.CommandScalableTopicUpdate; import org.apache.pulsar.common.api.proto.CommandSeek; import org.apache.pulsar.common.api.proto.CommandSend; @@ -514,6 +515,11 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception handleCommandScalableTopicAssignmentUpdate(cmd.getScalableTopicAssignmentUpdate()); break; + case SCALABLE_TOPIC_UNSUBSCRIBE: + checkArgument(cmd.hasScalableTopicUnsubscribe()); + handleCommandScalableTopicUnsubscribe(cmd.getScalableTopicUnsubscribe()); + break; + case WATCH_SCALABLE_TOPICS: checkArgument(cmd.hasWatchScalableTopics()); handleCommandWatchScalableTopics(cmd.getWatchScalableTopics()); @@ -837,6 +843,11 @@ protected void handleCommandScalableTopicAssignmentUpdate( throw new UnsupportedOperationException(); } + protected void handleCommandScalableTopicUnsubscribe( + CommandScalableTopicUnsubscribe commandScalableTopicUnsubscribe) { + throw new UnsupportedOperationException(); + } + protected void handleCommandWatchScalableTopics( org.apache.pulsar.common.api.proto.CommandWatchScalableTopics commandWatchScalableTopics) { throw new UnsupportedOperationException(); diff --git a/pulsar-common/src/main/proto/PulsarApi.proto b/pulsar-common/src/main/proto/PulsarApi.proto index 5d30972d93791..7d2d6845636f3 100644 --- a/pulsar-common/src/main/proto/PulsarApi.proto +++ b/pulsar-common/src/main/proto/PulsarApi.proto @@ -995,6 +995,16 @@ message CommandScalableTopicAssignmentUpdate { required ScalableConsumerAssignment assignment = 2; } +// Client -> Broker: a scalable consumer is cleanly leaving its subscription. The controller +// deletes the registration and rebalances the group immediately, instead of holding the +// session for the disconnect grace period (which remains the fallback for unclean +// departures). Acknowledged with CommandSuccess; idempotent — an unknown consumer_id (e.g. +// already swept by a disconnect) still succeeds. +message CommandScalableTopicUnsubscribe { + required uint64 request_id = 1; + required uint64 consumer_id = 2; +} + // Multi-topic consumer watcher: subscribes to the union of scalable topics in a // namespace that match a (possibly empty) set of property filters. The broker keeps // pushing updates as topics enter or leave the matching set. See @@ -1356,6 +1366,8 @@ message BaseCommand { WATCH_TC_ASSIGNMENTS = 79; WATCH_TC_ASSIGNMENTS_UPDATE = 80; WATCH_TC_ASSIGNMENTS_CLOSE = 81; + + SCALABLE_TOPIC_UNSUBSCRIBE = 82; } @@ -1455,4 +1467,6 @@ message BaseCommand { optional CommandWatchTcAssignments watchTcAssignments = 79; optional CommandWatchTcAssignmentsUpdate watchTcAssignmentsUpdate = 80; optional CommandWatchTcAssignmentsClose watchTcAssignmentsClose = 81; + + optional CommandScalableTopicUnsubscribe scalableTopicUnsubscribe = 82; } diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/protocol/CommandsScalableTopicTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/protocol/CommandsScalableTopicTest.java index a5d028889eea4..af9db0da4ab4d 100644 --- a/pulsar-common/src/test/java/org/apache/pulsar/common/protocol/CommandsScalableTopicTest.java +++ b/pulsar-common/src/test/java/org/apache/pulsar/common/protocol/CommandsScalableTopicTest.java @@ -169,6 +169,15 @@ public void testNewScalableTopicError() { "Scalable topic not found: topic://t/n/x"); } + @Test + public void testNewScalableTopicUnsubscribe() { + BaseCommand cmd = parseFrame(Commands.newScalableTopicUnsubscribe(42L, 7L)); + assertEquals(cmd.getType(), BaseCommand.Type.SCALABLE_TOPIC_UNSUBSCRIBE); + assertTrue(cmd.hasScalableTopicUnsubscribe()); + assertEquals(cmd.getScalableTopicUnsubscribe().getRequestId(), 42L); + assertEquals(cmd.getScalableTopicUnsubscribe().getConsumerId(), 7L); + } + @Test public void testNewScalableTopicSubscribeResponseSuccess() { ScalableConsumerAssignment assignment = new ScalableConsumerAssignment().setLayoutEpoch(3L);