[fix][broker] PIP-468: StreamConsumer / CheckpointConsumer DAG-replay with parent-drain ordering (#25642)

Author: merlimat

pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Segments.java

@@ -28,6 +28,7 @@
 import javax.ws.rs.DELETE;
 import javax.ws.rs.DefaultValue;
 import javax.ws.rs.Encoded;
+import javax.ws.rs.GET;
 import javax.ws.rs.POST;
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
@@ -174,6 +175,59 @@ public void terminateSegment(
                 });
     }
 
+    @GET
+    @Path("/{tenant}/{namespace}/{topic}/{descriptor}/subscription/{subscription}/backlog")
+    @ApiOperation(value = "Number of unconsumed entries in the segment topic for the "
+            + "given subscription. Super-user only.")
+    @ApiResponses(value = {
+            @ApiResponse(code = 401, message = "This operation requires super-user access"),
+            @ApiResponse(code = 403, message = "This operation requires super-user access"),
+            @ApiResponse(code = 404, message = "Segment topic or subscription not found"),
+            @ApiResponse(code = 500, message = "Internal server error")})
+    public void getSubscriptionBacklog(
+            @Suspended final AsyncResponse asyncResponse,
+            @ApiParam(value = "Specify the tenant", required = true)
+            @PathParam("tenant") String tenant,
+            @ApiParam(value = "Specify the namespace", required = true)
+            @PathParam("namespace") String namespace,
+            @ApiParam(value = "Specify the parent topic name", required = true)
+            @PathParam("topic") @Encoded String encodedTopic,
+            @ApiParam(value = "Segment descriptor (e.g. 0000-7fff-1)", required = true)
+            @PathParam("descriptor") String descriptor,
+            @ApiParam(value = "Subscription name", required = true)
+            @PathParam("subscription") String subscription,
+            @ApiParam(value = "Whether leader broker redirected this call to this broker.")
+            @QueryParam("authoritative") @DefaultValue("false") boolean authoritative) {
+        validateNamespaceName(tenant, namespace);
+        TopicName segmentTopic = segmentTopicName(tenant, namespace, encodedTopic, descriptor);
+
+        validateSuperUserAccessAsync()
+                .thenCompose(__ -> validateTopicOwnershipAsync(segmentTopic, authoritative))
+                .thenCompose(__ -> pulsar().getBrokerService().getTopicIfExists(segmentTopic.toString()))
+                .thenAccept(optTopic -> {
+                    if (optTopic.isEmpty()) {
+                        // No topic loaded → no subscription cursor → no backlog. Returning
+                        // 0 here would be wrong (caller might mark the segment drained on
+                        // a topic that simply hasn't loaded yet); a 404 forces the caller
+                        // to retry, which matches our drain-poll contract.
+                        throw new RestException(Response.Status.NOT_FOUND,
+                                "Segment topic not loaded: " + segmentTopic);
+                    }
+                    var sub = optTopic.get().getSubscription(subscription);
+                    if (sub == null) {
+                        throw new RestException(Response.Status.NOT_FOUND,
+                                "Subscription not found on segment: " + subscription);
+                    }
+                    asyncResponse.resume(sub.getNumberOfEntriesInBacklog(false));
+                })
+                .exceptionally(ex -> {
+                    log.error().attr("clientAppId", clientAppId()).attr("segment", segmentTopic)
+                            .exception(ex).log("Failed to get segment subscription backlog");
+                    resumeAsyncResponseExceptionally(asyncResponse, ex);
+                    return null;
+                });
+    }
+
     @DELETE
     @Path("/{tenant}/{namespace}/{topic}/{descriptor}")
     @ApiOperation(value = "Delete a segment topic. Super-user only.")

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java

@@ -879,6 +879,8 @@ protected void handleCommandScalableTopicSubscribe(
         final String subscription = commandScalableTopicSubscribe.getSubscription();
         final String consumerName = commandScalableTopicSubscribe.getConsumerName();
         final long consumerId = commandScalableTopicSubscribe.getConsumerId();
+        final org.apache.pulsar.common.api.proto.ScalableConsumerType consumerType =
+                commandScalableTopicSubscribe.getConsumerType();
 
         log.debug().attr("topic", topicStr).attr("subscription", subscription)
                 .attr("consumerName", consumerName).attr("requestId", requestId)
@@ -920,7 +922,7 @@ protected void handleCommandScalableTopicSubscribe(
                         return;
                     }
                     scalableTopicService.registerConsumer(topicName, subscription, consumerName,
-                                    consumerId, this)
+                                    consumerId, consumerType, this)
                             .whenCompleteAsync((assignment, ex) -> {
                                 if (ex != null) {
                                     Throwable cause = ex.getCause() != null ? ex.getCause() : ex;

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

@@ -31,6 +31,7 @@
 import org.apache.pulsar.broker.resources.ScalableTopicResources;
 import org.apache.pulsar.broker.service.BrokerService;
 import org.apache.pulsar.broker.service.TransportCnx;
+import org.apache.pulsar.common.api.proto.ScalableConsumerType;
 import org.apache.pulsar.common.naming.TopicName;
 import org.apache.pulsar.common.scalable.HashRange;
 import org.apache.pulsar.common.scalable.SegmentInfo;
@@ -155,7 +156,32 @@ private CompletableFuture<Void> restoreSubscription(String subscription) {
                 });
     }
 
+    /**
+     * Restore-path entry: consumer type isn't persisted in metadata yet, so we don't
+     * know whether the original subscription was STREAM (needs parent-drain ordering)
+     * or CHECKPOINT / QUEUE (mustn't have it — CHECKPOINT never drains parents because
+     * it doesn't create per-segment cursors). Default to <em>no enforcement</em>; on the
+     * first register-after-restore the controller calls
+     * {@link SubscriptionCoordinator#installDrainChecker} if the type is STREAM.
+     */
     private SubscriptionCoordinator createCoordinator(String subscription) {
+        return createCoordinator(subscription, null);
+    }
+
+    private SubscriptionCoordinator createCoordinator(String subscription,
+            ScalableConsumerType consumerType) {
+        // Parent-drain ordering matters only for STREAM consumers (Exclusive per-segment
+        // subscription with broker-tracked cursors → preserving per-key order across a
+        // split requires waiting for the parent to drain before handing out children).
+        // CHECKPOINT consumers track position client-side via Checkpoints and don't even
+        // create per-segment cursors — their parent never reports as drained, so the
+        // ordering machinery would block their children indefinitely. QUEUE consumers
+        // are shared and accept out-of-order delivery by design. Null type (restore
+        // path) starts without a checker; it's installed lazily on first STREAM
+        // register.
+        SegmentDrainChecker checker =
+                consumerType == ScalableConsumerType.STREAM ? this::isSegmentDrained : null;
+
         // Defensive: PulsarService.getConfig() is null in some unit-test mocks. Fall
         // back to the SubscriptionCoordinator's default grace period in that case.
         var config = brokerService.getPulsar().getConfig();
@@ -175,7 +201,42 @@ private SubscriptionCoordinator createCoordinator(String subscription) {
                 currentLayout,
                 resources,
                 brokerService.getPulsar().getExecutor(),
-                gracePeriod);
+                gracePeriod,
+                checker,
+                SubscriptionCoordinator.DEFAULT_DRAIN_INITIAL_DELAY,
+                SubscriptionCoordinator.DEFAULT_DRAIN_MAX_DELAY);
+    }
+
+    /**
+     * Drain check used by every {@link SubscriptionCoordinator} on this topic. Asks the
+     * segment topic's owning broker for the per-subscription backlog via the
+     * {@code /segments/.../subscription/.../backlog} admin endpoint, which redirects to
+     * the topic owner — works whether the controller and the segment colocate or not.
+     *
+     * <p>Returns {@code false} if the segment topic or subscription is not yet loaded
+     * (the admin endpoint replies 404). The next poll will succeed once the consumer's
+     * subscribe lands the topic on its owning broker.
+     */
+    private CompletableFuture<Boolean> isSegmentDrained(SegmentInfo segment, String subscription) {
+        String segmentTopicName = toSegmentPersistentName(segment);
+        try {
+            return brokerService.getPulsar().getAdminClient()
+                    .scalableTopics()
+                    .getSegmentSubscriptionBacklogAsync(segmentTopicName, subscription)
+                    .thenApply(backlog -> backlog != null && backlog <= 0)
+                    .exceptionally(ex -> {
+                        Throwable cause =
+                                org.apache.pulsar.common.util.FutureUtil.unwrapCompletionException(ex);
+                        if (cause instanceof org.apache.pulsar.client.admin.PulsarAdminException.NotFoundException) {
+                            // Topic or subscription not loaded yet — try again on the
+                            // next poll. The consumer's subscribe will materialize it.
+                            return false;
+                        }
+                        throw org.apache.pulsar.common.util.FutureUtil.wrapToCompletionException(cause);
+                    });
+        } catch (PulsarServerException e) {
+            return CompletableFuture.failedFuture(e);
+        }
     }
 
     private CompletableFuture<Void> electLeader() {
@@ -309,14 +370,42 @@ public CompletableFuture<SegmentLayout> mergeSegments(long segmentId1, long segm
      * <p>If a session with the same {@code consumerName} already exists (for example
      * because the consumer is reconnecting within the grace period), the existing
      * assignment is reused and no rebalance occurs.
+     *
+     * <p>The {@code consumerType} is used at coordinator creation time to decide whether
+     * to enforce parent-drain ordering on assignments — see
+     * {@link SubscriptionCoordinator}. The coordinator's setting is fixed at first
+     * registration (a subscription's type doesn't change in practice); subsequent
+     * registers with a different type still work but won't change the ordering policy.
+     */
+    /**
+     * @deprecated Defaults to {@link ScalableConsumerType#STREAM}
+     *     for backward compatibility. New callers should pass the explicit type.
      */
+    @Deprecated
     public CompletableFuture<ConsumerAssignment> registerConsumer(String subscription,
                                                                    String consumerName,
                                                                    long consumerId,
                                                                    TransportCnx cnx) {
+        return registerConsumer(subscription, consumerName, consumerId,
+                ScalableConsumerType.STREAM, cnx);
+    }
+
+    public CompletableFuture<ConsumerAssignment> registerConsumer(String subscription,
+                                                                   String consumerName,
+                                                                   long consumerId,
+                                                                   ScalableConsumerType
+                                                                           consumerType,
+                                                                   TransportCnx cnx) {
         checkLeader();
         SubscriptionCoordinator coordinator = subscriptions.computeIfAbsent(
-                subscription, this::createCoordinator);
+                subscription, sub -> createCoordinator(sub, consumerType));
+        // The coordinator may have been created on the failover-restore path (consumer
+        // type unknown then; we defaulted to "no parent-drain enforcement"). Now that we
+        // know the type, upgrade if it's STREAM. installDrainChecker is a no-op if the
+        // coordinator already has a checker, so safe to call unconditionally.
+        if (consumerType == ScalableConsumerType.STREAM) {
+            coordinator.installDrainChecker(this::isSegmentDrained);
+        }
         return coordinator.registerConsumer(consumerName, consumerId, cnx)
                 .thenApply(assignments -> {
                     // Look up by name since the key may have been an existing session
@@ -525,6 +614,9 @@ public CompletableFuture<org.apache.pulsar.common.policies.data.ScalableTopicSta
 
     public CompletableFuture<Void> close() {
         closed = true;
+        // Stop each coordinator's drain poller before clearing — otherwise the scheduler
+        // task keeps running after the controller goes away.
+        subscriptions.values().forEach(SubscriptionCoordinator::close);
         subscriptions.clear();
         return leaderElection.asyncClose();
     }

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

@@ -27,6 +27,7 @@
 import org.apache.pulsar.broker.resources.ScalableTopicMetadata;
 import org.apache.pulsar.broker.resources.ScalableTopicResources;
 import org.apache.pulsar.broker.service.BrokerService;
+import org.apache.pulsar.common.api.proto.ScalableConsumerType;
 import org.apache.pulsar.common.naming.TopicDomain;
 import org.apache.pulsar.common.naming.TopicName;
 import org.apache.pulsar.common.policies.data.ScalableTopicStats;
@@ -230,12 +231,35 @@ public CompletableFuture<Void> deleteScalableTopic(TopicName topic) {
     /**
      * Register a scalable consumer with the controller leader for {@code topic}.
      * Persists a durable session and returns the consumer's segment assignment.
+     *
+     * <p>The {@code consumerType} drives broker-side semantics that depend on the
+     * consumer mode — most importantly whether the {@link SubscriptionCoordinator}
+     * enforces parent-drain ordering before handing out children of a split. STREAM
+     * consumers want it (per-key ordering); CHECKPOINT and QUEUE consumers don't
+     * (they either track position client-side or have shared, already-out-of-order
+     * delivery semantics).
      */
     public CompletableFuture<ConsumerAssignment> registerConsumer(TopicName topic, String subscription,
                                                                    String consumerName, long consumerId,
+                                                                   ScalableConsumerType
+                                                                           consumerType,
                                                                    org.apache.pulsar.broker.service.TransportCnx cnx) {
         return getOrCreateController(topic)
-                .thenCompose(controller -> controller.registerConsumer(subscription, consumerName, consumerId, cnx));
+                .thenCompose(controller ->
+                        controller.registerConsumer(subscription, consumerName, consumerId,
+                                consumerType, cnx));
+    }
+
+    /**
+     * @deprecated Defaults to {@link ScalableConsumerType#STREAM}
+     *     for backward compatibility. New callers should pass the explicit consumer type.
+     */
+    @Deprecated
+    public CompletableFuture<ConsumerAssignment> registerConsumer(TopicName topic, String subscription,
+                                                                   String consumerName, long consumerId,
+                                                                   org.apache.pulsar.broker.service.TransportCnx cnx) {
+        return registerConsumer(topic, subscription, consumerName, consumerId,
+                ScalableConsumerType.STREAM, cnx);
     }
 
     /**

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

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.broker.service.scalable;
+
+import java.util.concurrent.CompletableFuture;
+import org.apache.pulsar.common.scalable.SegmentInfo;
+
+/**
+ * Resolves whether a (sealed) segment has been fully drained by a particular subscription.
+ *
+ * <p>Used by {@link SubscriptionCoordinator} to decide when newly-active children of a
+ * split / merge can be assigned to consumers: an active child is only assignable once
+ * <em>every</em> parent has been drained for the subscription, so message order with respect
+ * to the split point is preserved.
+ *
+ * <p>Implementations typically read the segment topic's per-subscription backlog (the
+ * cursor on a sealed topic with {@code msgBacklog == 0} is by definition at the end). For
+ * subscriptions started with {@code Latest}, every sealed segment's cursor is created at
+ * the topic's end, so the drain check completes immediately.
+ */
+@FunctionalInterface
+public interface SegmentDrainChecker {
+
+    /**
+     * Returns {@code true} if the segment's cursor for {@code subscription} has reached the
+     * end of the segment's data, {@code false} otherwise. Errors complete the future
+     * exceptionally; callers should treat them as "not drained yet" and retry.
+     *
+     * @param segment the segment to check (only meaningful for sealed segments — active
+     *                segments still receive new messages, so they're never drained)
+     * @param subscription the subscription whose cursor we're asking about
+     */
+    CompletableFuture<Boolean> isDrained(SegmentInfo segment, String subscription);
+}