[improve][broker] PIP-486: rebucket rollover and the segments-vs-buckets auto-scale policy

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java

@@ -1468,6 +1468,15 @@ The max allowed delay for delayed delivery (in milliseconds). If the broker rece
     )
     private int scalableTopicEntryBucketBudget = 4;
 
+    @FieldContext(
+            dynamic = true,
+            category = CATEGORY_POLICIES,
+            doc = "Hard ceiling on a single segment's entry-bucket count (PIP-486). Bounds both the "
+                    + "manual rebucket operation and the controller's auto rebucket-up; a segment's "
+                    + "bucket count caps how many consumers can share it."
+    )
+    private int scalableTopicEntryBucketMaxPerSegment = 1024;
+
     @FieldContext(
             dynamic = true,
             category = CATEGORY_POLICIES,
@@ -1486,6 +1495,24 @@ The max allowed delay for delayed delivery (in milliseconds). If the broker rece
     )
     private int scalableTopicSplitCooldownSeconds = 60;
 
+    @FieldContext(
+            dynamic = true,
+            category = CATEGORY_POLICIES,
+            doc = "PIP-486 segments-vs-buckets lever: on consumer-driven scale-up, split only if the "
+                    + "busiest segment's inbound msg/s is at or above this floor; below it the "
+                    + "controller grows the segment's entry-buckets instead (a low-throughput topic "
+                    + "should not materialize physical segments just for consumer count)."
+    )
+    private double scalableTopicSplitVsRebucketMinMsgRateInThreshold = 1_000;
+
+    @FieldContext(
+            dynamic = true,
+            category = CATEGORY_POLICIES,
+            doc = "Minimum time (seconds) between automatic entry-bucket rollovers (rebuckets) on a "
+                    + "topic. Coalesces consumer-join bursts, like the split cooldown."
+    )
+    private int scalableTopicRebucketCooldownSeconds = 60;
+
     @FieldContext(
             dynamic = true,
             category = CATEGORY_POLICIES,

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

@@ -935,6 +935,50 @@ public void splitSegment(
                 });
     }
 
+    @POST
+    @Path("/{tenant}/{namespace}/{topic}/rebucket/{segmentId}")
+    @Operation(summary = "Rebucket a segment: roll it over to a same-range successor with a new "
+            + "entry-bucket count.")
+    @ApiResponses(value = {
+            @ApiResponse(responseCode = "204", description = "Segment rebucketed successfully"),
+            @ApiResponse(responseCode = "404", description = "Scalable topic or segment doesn't exist"),
+            @ApiResponse(responseCode = "412",
+                    description = "Segment is not active or the bucket count is invalid"),
+            @ApiResponse(responseCode = "500", description = "Internal server error")})
+    public void rebucketSegment(
+            @Suspended final AsyncResponse asyncResponse,
+            @Parameter(description = "Specify the tenant", required = true)
+            @PathParam("tenant") String tenant,
+            @Parameter(description = "Specify the namespace", required = true)
+            @PathParam("namespace") String namespace,
+            @Parameter(description = "Specify topic name", required = true)
+            @PathParam("topic") @Encoded String encodedTopic,
+            @Parameter(description = "Segment ID to rebucket", required = true)
+            @PathParam("segmentId") long segmentId,
+            @Parameter(description = "Entry-bucket count for the successor segment", required = true)
+            @QueryParam("bucketCount") int bucketCount) {
+        validateNamespaceName(tenant, namespace);
+        TopicName tn = TopicName.get(TopicDomain.topic.value(), namespaceName, encodedTopic);
+
+        validateSuperUserAccessAsync()
+                .thenCompose(__ -> onControllerLeader(tn,
+                        svc -> svc.rebucketSegment(tn, segmentId, bucketCount)))
+                .thenAccept(__ -> {
+                    log.info().attr("clientAppId", clientAppId())
+                            .attr("segmentId", segmentId).attr("bucketCount", bucketCount)
+                            .attr("topic", tn)
+                            .log("Rebucketed segment of scalable topic");
+                    asyncResponse.resume(Response.noContent().build());
+                })
+                .exceptionally(ex -> {
+                    log.error().attr("clientAppId", clientAppId())
+                            .attr("segmentId", segmentId).attr("topic", tn)
+                            .exception(ex).log("Failed to rebucket segment");
+                    resumeAsyncResponseExceptionally(asyncResponse, ex);
+                    return null;
+                });
+    }
+
     @POST
     @Path("/{tenant}/{namespace}/{topic}/merge/{segmentId1}/{segmentId2}")
     @Operation(summary = "Merge two adjacent segments into one.")

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

@@ -46,6 +46,10 @@
  * @param mergeWindow       how long a segment must continuously stay below every merge threshold
  *                          before it becomes merge-eligible (measured from the load record's
  *                          metadata-store last-modified time)
+ * @param rebucketCooldown  minimum interval between automatic entry-bucket rollovers
+ * @param splitVsRebucketMinMsgRateIn consumer-driven scale-up splits only at/above this inbound
+ *                          msg/s on the busiest segment; below it, entry-buckets grow instead
+ * @param maxEntryBucketsPerSegment hard ceiling on a single segment's entry-bucket count
  * @param splitMsgRateIn    inbound msg/s above which a segment is split
  * @param splitBytesRateIn  inbound bytes/s above which a segment is split
  * @param splitMsgRateOut   outbound (dispatched) msg/s above which a segment is split
@@ -62,8 +66,11 @@ public record AutoScaleConfig(
         int minSegments,
         int maxDagDepth,
         Duration splitCooldown,
+        Duration rebucketCooldown,
         Duration mergeCooldown,
         Duration mergeWindow,
+        double splitVsRebucketMinMsgRateIn,
+        int maxEntryBucketsPerSegment,
         double splitMsgRateIn,
         double splitBytesRateIn,
         double splitMsgRateOut,
@@ -112,6 +119,10 @@ private static AutoScaleConfig brokerDefaults(ServiceConfiguration conf) {
                 .minSegments(conf.getScalableTopicMinSegments())
                 .maxDagDepth(conf.getScalableTopicMaxDagDepth())
                 .splitCooldown(Duration.ofSeconds(conf.getScalableTopicSplitCooldownSeconds()))
+                .rebucketCooldown(Duration.ofSeconds(conf.getScalableTopicRebucketCooldownSeconds()))
+                .splitVsRebucketMinMsgRateIn(
+                        conf.getScalableTopicSplitVsRebucketMinMsgRateInThreshold())
+                .maxEntryBucketsPerSegment(conf.getScalableTopicEntryBucketMaxPerSegment())
                 .mergeCooldown(Duration.ofSeconds(conf.getScalableTopicMergeCooldownSeconds()))
                 .mergeWindow(Duration.ofSeconds(conf.getScalableTopicMergeWindowSeconds()))
                 .splitMsgRateIn(conf.getScalableTopicSplitMsgRateInThreshold())
@@ -145,6 +156,12 @@ private static AutoScaleConfig applyOverride(AutoScaleConfig base, AutoScalePoli
         if (o.getSplitCooldownSeconds() != null) {
             b.splitCooldown(Duration.ofSeconds(o.getSplitCooldownSeconds()));
         }
+        if (o.getRebucketCooldownSeconds() != null) {
+            b.rebucketCooldown(Duration.ofSeconds(o.getRebucketCooldownSeconds()));
+        }
+        if (o.getSplitVsRebucketMinMsgRateInThreshold() != null) {
+            b.splitVsRebucketMinMsgRateIn(o.getSplitVsRebucketMinMsgRateInThreshold());
+        }
         if (o.getMergeCooldownSeconds() != null) {
             b.mergeCooldown(Duration.ofSeconds(o.getMergeCooldownSeconds()));
         }

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

@@ -24,7 +24,8 @@
  * carries a short {@code reason} string used for logging and metrics.
  */
 public sealed interface AutoScaleDecision
-        permits AutoScaleDecision.Split, AutoScaleDecision.Merge, AutoScaleDecision.NoAction {
+        permits AutoScaleDecision.Split, AutoScaleDecision.Merge, AutoScaleDecision.Rebucket,
+                AutoScaleDecision.NoAction {
 
     /** Split {@code segmentId} at its midpoint. */
     record Split(long segmentId, String reason) implements AutoScaleDecision {
@@ -34,6 +35,16 @@ record Split(long segmentId, String reason) implements AutoScaleDecision {
     record Merge(long segmentId1, long segmentId2, String reason) implements AutoScaleDecision {
     }
 
+    /**
+     * Roll every segment in {@code segmentIds} over to a same-range successor with
+     * {@code newBucketCount} entry-buckets (PIP-486): consumer scale-up served by buckets
+     * instead of a split. One decision carries the whole batch so a multi-segment topic
+     * converges to a uniform bucketing in a single evaluation, not one segment per cooldown.
+     */
+    record Rebucket(java.util.List<Long> segmentIds, int newBucketCount, String reason)
+            implements AutoScaleDecision {
+    }
+
     /** No action this evaluation. */
     record NoAction() implements AutoScaleDecision {
     }

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

@@ -68,29 +68,115 @@ public static AutoScaleDecision decide(
             AutoScaleConfig config,
             long nowMs,
             long lastSplitAtMs,
-            long lastMergeAtMs) {
+            long lastMergeAtMs,
+            long lastRebucketAtMs) {
 
         if (!config.enabled()) {
             return AutoScaleDecision.NONE;
         }
 
         List<SegmentInfo> active = new ArrayList<>(layout.getActiveSegments().values());
 
-        AutoScaleDecision split = trySplit(active, loadBySegment, streamConsumerCount,
-                config, nowMs, lastSplitAtMs);
+        AutoScaleDecision consumerScale = tryConsumerScale(active, loadBySegment,
+                streamConsumerCount, config, nowMs, lastSplitAtMs, lastRebucketAtMs);
+        if (!(consumerScale instanceof AutoScaleDecision.NoAction)) {
+            return consumerScale;
+        }
+
+        AutoScaleDecision split = trySplit(active, loadBySegment, config, nowMs, lastSplitAtMs);
         if (!(split instanceof AutoScaleDecision.NoAction)) {
             return split;
         }
 
         return tryMerge(active, layout, loadBySegment, config, nowMs, lastMergeAtMs);
     }
 
+    // --- Consumer-driven scale-up: segments vs entry-buckets (PIP-486) ---
+
+    /**
+     * Serve surplus consumers (a subscription with more consumers than active segments) by
+     * adding capacity on one of two axes:
+     * <ul>
+     *   <li><b>Split</b> when traffic justifies a physical segment: the busiest segment's
+     *       inbound rate is at or above {@code splitVsRebucketMinMsgRateIn} and the topic is
+     *       under {@code maxSegments} — today's "segments first" behavior.</li>
+     *   <li><b>Rebucket-up</b> otherwise (a low-throughput topic, or the topic is at the
+     *       segment cap): if the existing entry-bucket capacity cannot absorb the surplus,
+     *       roll the smallest-bucketed segment over to the smallest power of two that lets
+     *       every consumer own a bucket, capped at {@code maxEntryBucketsPerSegment}.
+     *       Raising is fast (one rollover sized to the surplus); lowering is deliberately
+     *       not automated here — spiky consumer counts must not flap the bucketing.</li>
+     * </ul>
+     */
+    private static AutoScaleDecision tryConsumerScale(
+            List<SegmentInfo> active,
+            Map<Long, SegmentLoadSample> loadBySegment,
+            Map<String, Integer> streamConsumerCount,
+            AutoScaleConfig config,
+            long nowMs,
+            long lastSplitAtMs,
+            long lastRebucketAtMs) {
+
+        int consumers = streamConsumerCount.values().stream()
+                .mapToInt(Integer::intValue).max().orElse(0);
+        int segments = active.size();
+        if (consumers <= segments) {
+            return AutoScaleDecision.NONE;
+        }
+
+        SegmentInfo busiest = busiestByMsgRateIn(active, loadBySegment);
+        if (busiest == null) {
+            return AutoScaleDecision.NONE;
+        }
+        boolean atSegmentCap = segments >= config.maxSegments();
+        boolean belowSplitFloor = statsOf(busiest.segmentId(), loadBySegment).msgRateIn()
+                < config.splitVsRebucketMinMsgRateIn();
+
+        if (!atSegmentCap && !belowSplitFloor) {
+            // Traffic justifies a physical segment.
+            if (withinCooldown(nowMs, lastSplitAtMs, config.splitCooldown().toMillis())) {
+                return AutoScaleDecision.NONE;
+            }
+            return new AutoScaleDecision.Split(busiest.segmentId(), "consumer-count");
+        }
+
+        // Bucket lane: absorb the surplus with entry-buckets.
+        long capacity = 0;
+        for (SegmentInfo segment : active) {
+            capacity += segment.bucketCount();
+        }
+        if (consumers <= capacity) {
+            // The existing buckets already absorb the surplus (broker-side fan-out).
+            return AutoScaleDecision.NONE;
+        }
+        if (withinCooldown(nowMs, lastRebucketAtMs, config.rebucketCooldown().toMillis())) {
+            return AutoScaleDecision.NONE;
+        }
+        // One shot: bring every segment below the common per-segment target up to it in a
+        // single decision, so the topic converges to a uniform bucketing in one evaluation —
+        // never one segment per cooldown, and no arrival-history-dependent skew.
+        int target = Math.min(nextPowerOfTwo(ceilDiv(consumers, segments)),
+                config.maxEntryBucketsPerSegment());
+        List<Long> below = new ArrayList<>();
+        for (SegmentInfo segment : active) {
+            if (segment.bucketCount() < target) {
+                below.add(segment.segmentId());
+            }
+        }
+        if (below.isEmpty()) {
+            // Bucket capacity is maxed out; the remaining surplus stays idle.
+            return AutoScaleDecision.NONE;
+        }
+        below.sort(Long::compareTo);
+        return new AutoScaleDecision.Rebucket(below, target,
+                atSegmentCap ? "at-max-segments" : "below-split-rate-floor");
+    }
+
     // --- Split pass ---
 
     private static AutoScaleDecision trySplit(
             List<SegmentInfo> active,
             Map<Long, SegmentLoadSample> loadBySegment,
-            Map<String, Integer> streamConsumerCount,
             AutoScaleConfig config,
             long nowMs,
             long lastSplitAtMs) {
@@ -102,20 +188,7 @@ private static AutoScaleDecision trySplit(
             return AutoScaleDecision.NONE;
         }
 
-        // (a) Consumer-driven: per-subscription max. If any managed subscription has more
-        // consumers than there are active segments, add a segment so the 1:1 assignment can
-        // give the extra consumer its own segment. Split the busiest segment by msgRateIn so
-        // the new pair lands where it relieves the most ingest.
-        int requiredConsumers = streamConsumerCount.values().stream()
-                .mapToInt(Integer::intValue).max().orElse(0);
-        if (requiredConsumers > active.size()) {
-            SegmentInfo target = busiestByMsgRateIn(active, loadBySegment);
-            if (target != null) {
-                return new AutoScaleDecision.Split(target.segmentId(), "consumer-count");
-            }
-        }
-
-        // (b) Load-driven: split the segment with the highest overload score among those over
+        // Load-driven: split the segment with the highest overload score among those over
         // at least one split threshold.
         SegmentInfo hottest = null;
         double hottestScore = 1.0; // strictly over threshold means a per-metric ratio > 1.0
@@ -239,6 +312,17 @@ private static double combinedRate(long segmentId, Map<Long, SegmentLoadSample>
         return s.msgRateIn() + s.bytesRateIn() + s.msgRateOut() + s.bytesRateOut();
     }
 
+    /** Ceiling integer division for positive operands. */
+    private static int ceilDiv(int a, int b) {
+        return (a + b - 1) / b;
+    }
+
+    /** The smallest power of two {@code >= v} (for {@code v >= 1}). */
+    private static int nextPowerOfTwo(int v) {
+        int highest = Integer.highestOneBit(v);
+        return highest == v ? v : highest << 1;
+    }
+
     private static SegmentInfo busiestByMsgRateIn(List<SegmentInfo> active,
                                                   Map<Long, SegmentLoadSample> load) {
         SegmentInfo best = null;