[fix][client] PIP-486: explicit unsubscribe for scalable consumers' clean close

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);

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/ScalableTopicController.java

@@ -811,13 +811,14 @@ public CompletableFuture<ConsumerAssignment> registerConsumer(String subscriptio
      * Explicit unregister: the consumer is leaving the subscription for good. Deletes the
      * persisted session entry and rebalances remaining consumers.
      */
-    public CompletableFuture<Void> unregisterConsumer(String subscription, String consumerName) {
+    public CompletableFuture<Void> 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<Void> deleteSubscription(String subscription) {
 
     private CompletableFuture<Void> 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);
     }

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<Void> unregisterConsumer(TopicName topic, String subscription,
+                                                      String consumerName, long consumerId) {
+        CompletableFuture<ScalableTopicController> 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<Void> createUnderlyingSegmentTopic(TopicName parentTopic, SegmentInfo segment) {

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/scalable/SubscriptionCoordinator.java

@@ -214,19 +214,49 @@ public synchronized CompletableFuture<Map<ConsumerSession, ConsumerAssignment>>
     }
 
     /**
-     * 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.
+     *
+     * <p>{@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<Map<ConsumerSession, ConsumerAssignment>> 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();