Add base implementation for auto partitioning settings

Daniil Zulin · Daniil Zulin · commit 913ba2922e39 · 2025-05-06T18:28:58.000+07:00
diff --git a/topic/src/main/java/tech/ydb/topic/impl/TopicClientImpl.java b/topic/src/main/java/tech/ydb/topic/impl/TopicClientImpl.java
@@ -35,9 +35,12 @@
 import tech.ydb.topic.read.SyncReader;
 import tech.ydb.topic.read.impl.AsyncReaderImpl;
 import tech.ydb.topic.read.impl.SyncReaderImpl;
+import tech.ydb.topic.settings.AlterAutoPartitioningWriteStrategySettings;
 import tech.ydb.topic.settings.AlterConsumerSettings;
 import tech.ydb.topic.settings.AlterPartitioningSettings;
 import tech.ydb.topic.settings.AlterTopicSettings;
+import tech.ydb.topic.settings.AutoPartitioningStrategy;
+import tech.ydb.topic.settings.AutoPartitioningWriteStrategySettings;
 import tech.ydb.topic.settings.CommitOffsetSettings;
 import tech.ydb.topic.settings.CreateTopicSettings;
 import tech.ydb.topic.settings.DescribeConsumerSettings;
@@ -104,7 +107,19 @@ public CompletableFuture<Status> createTopic(String path, CreateTopicSettings se
         if (partitioningSettings != null) {
             requestBuilder.setPartitioningSettings(YdbTopic.PartitioningSettings.newBuilder()
                     .setMinActivePartitions(partitioningSettings.getMinActivePartitions())
-                    .setPartitionCountLimit(partitioningSettings.getPartitionCountLimit()));
+                    .setPartitionCountLimit(partitioningSettings.getPartitionCountLimit())
+                    .setAutoPartitioningSettings(YdbTopic.AutoPartitioningSettings.newBuilder()
+                            .setStrategy(partitioningSettings.getAutoPartitioningStrategy().getProtoReference())));
+
+            AutoPartitioningWriteStrategySettings writeStrategySettings = partitioningSettings.getWriteStrategySettings();
+            if (writeStrategySettings != null) {
+                requestBuilder.getPartitioningSettingsBuilder().getAutoPartitioningSettingsBuilder()
+                        .setPartitionWriteSpeed(YdbTopic.AutoPartitioningWriteSpeedStrategy.newBuilder()
+                                .setStabilizationWindow(ProtobufUtils.durationToProto(writeStrategySettings.getStabilizationWindow()))
+                                .setDownUtilizationPercent(writeStrategySettings.getDownUtilizationPercent())
+                                .setUpUtilizationPercent(writeStrategySettings.getUpUtilizationPercent())
+                        );
+            }
         }
 
         Duration retentionPeriod = settings.getRetentionPeriod();
@@ -145,6 +160,28 @@ public CompletableFuture<Status> alterTopic(String path, AlterTopicSettings sett
             if (partitionCountLimit != null) {
                 builder.setSetPartitionCountLimit(partitionCountLimit);
             }
+            AutoPartitioningStrategy autoPartitioningStrategy = partitioningSettings.getAutoPartitioningStrategy();
+            if (autoPartitioningStrategy != null) {
+                builder.getAlterAutoPartitioningSettingsBuilder().setSetStrategy(autoPartitioningStrategy.getProtoReference());
+            }
+            AlterAutoPartitioningWriteStrategySettings writeStrategySettings = partitioningSettings.getWriteStrategySettings();
+            if (writeStrategySettings != null) {
+                Duration stabilizationWindow = writeStrategySettings.getStabilizationWindow();
+                if (stabilizationWindow != null) {
+                    builder.getAlterAutoPartitioningSettingsBuilder().getSetPartitionWriteSpeedBuilder()
+                            .setSetStabilizationWindow(ProtobufUtils.durationToProto(stabilizationWindow));
+                }
+                Integer upUtilizationPercent = writeStrategySettings.getUpUtilizationPercent();
+                if (upUtilizationPercent != null) {
+                    builder.getAlterAutoPartitioningSettingsBuilder().getSetPartitionWriteSpeedBuilder()
+                            .setSetUpUtilizationPercent(upUtilizationPercent);
+                }
+                Integer downUtilizationPercent = writeStrategySettings.getDownUtilizationPercent();
+                if (downUtilizationPercent != null) {
+                    builder.getAlterAutoPartitioningSettingsBuilder().getSetPartitionWriteSpeedBuilder()
+                            .setSetDownUtilizationPercent(downUtilizationPercent);
+                }
+            }
             requestBuilder.setAlterPartitioningSettings(builder);
         }
 
@@ -273,11 +310,20 @@ private TopicDescription mapDescribeTopic(YdbTopic.DescribeTopicResult result) {
                 .setMeteringMode(fromProto(result.getMeteringMode()));
 
         YdbTopic.PartitioningSettings partitioningSettings = result.getPartitioningSettings();
-        description.setPartitioningSettings(PartitioningSettings.newBuilder()
+        PartitioningSettings.Builder partitioningDescription = PartitioningSettings.newBuilder()
                 .setMinActivePartitions(partitioningSettings.getMinActivePartitions())
                 .setPartitionCountLimit(partitioningSettings.getPartitionCountLimit())
+                .setAutoPartitioningStrategy(AutoPartitioningStrategy.fromProto(partitioningSettings.getAutoPartitioningSettings().getStrategy()));
+
+        YdbTopic.AutoPartitioningWriteSpeedStrategy partitionWriteSpeed = partitioningSettings.getAutoPartitioningSettings().getPartitionWriteSpeed();
+        partitioningDescription.setWriteStrategySettings(AutoPartitioningWriteStrategySettings.newBuilder()
+                .setStabilizationWindow(ProtobufUtils.protoToDuration(partitionWriteSpeed.getStabilizationWindow()))
+                .setUpUtilizationPercent(partitionWriteSpeed.getUpUtilizationPercent())
+                .setDownUtilizationPercent(partitionWriteSpeed.getDownUtilizationPercent())
                 .build());
 
+        description.setPartitioningSettings(partitioningDescription.build());
+
         List<PartitionInfo> partitions = new ArrayList<>();
         for (YdbTopic.DescribeTopicResult.PartitionInfo partition : result.getPartitionsList()) {
             PartitionInfo.Builder partitionBuilder = PartitionInfo.newBuilder()
diff --git a/topic/src/main/java/tech/ydb/topic/settings/AlterAutoPartitioningWriteStrategySettings.java b/topic/src/main/java/tech/ydb/topic/settings/AlterAutoPartitioningWriteStrategySettings.java
@@ -0,0 +1,77 @@
+package tech.ydb.topic.settings;
+
+import javax.annotation.Nullable;
+
+import java.time.Duration;
+
+// TODO add proper javadocs
+public class AlterAutoPartitioningWriteStrategySettings {
+    @Nullable
+    private final Duration stabilizationWindow;
+    @Nullable
+    private final Integer upUtilizationPercent;
+    @Nullable
+    private final Integer downUtilizationPercent;
+
+    public AlterAutoPartitioningWriteStrategySettings(Builder builder) {
+        this.stabilizationWindow = builder.stabilizationWindow;
+        this.upUtilizationPercent = builder.upUtilizationPercent;
+        this.downUtilizationPercent = builder.downUtilizationPercent;
+    }
+
+    @Nullable
+    public Duration getStabilizationWindow() {
+        return stabilizationWindow;
+    }
+
+    @Nullable
+    public Integer getUpUtilizationPercent() {
+        return upUtilizationPercent;
+    }
+
+    @Nullable
+    public Integer getDownUtilizationPercent() {
+        return downUtilizationPercent;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    public static class Builder {
+        private Duration stabilizationWindow = null;
+        private Integer upUtilizationPercent = null;
+        private Integer downUtilizationPercent = null;
+
+        /**
+         * @param stabilizationWindow
+         * @return strategy builder
+         */
+        public Builder setStabilizationWindow(Duration stabilizationWindow) {
+            this.stabilizationWindow = stabilizationWindow;
+            return this;
+        }
+
+        /**
+         * @param upUtilizationPercent
+         * @return strategy builder
+         */
+        public Builder setUpUtilizationPercent(int upUtilizationPercent) {
+            this.upUtilizationPercent = upUtilizationPercent;
+            return this;
+        }
+
+        /**
+         * @param downUtilizationPercent
+         * @return strategy builder
+         */
+        public Builder setDownUtilizationPercent(int downUtilizationPercent) {
+            this.downUtilizationPercent = downUtilizationPercent;
+            return this;
+        }
+
+        public AlterAutoPartitioningWriteStrategySettings build() {
+            return new AlterAutoPartitioningWriteStrategySettings(this);
+        }
+    }
+}
diff --git a/topic/src/main/java/tech/ydb/topic/settings/AlterPartitioningSettings.java b/topic/src/main/java/tech/ydb/topic/settings/AlterPartitioningSettings.java
@@ -10,10 +10,16 @@ public class AlterPartitioningSettings {
     private final Long minActivePartitions;
     @Nullable
     private final Long partitionCountLimit;
+    @Nullable
+    private final AutoPartitioningStrategy autoPartitioningStrategy;
+    @Nullable
+    private final AlterAutoPartitioningWriteStrategySettings writeStrategySettings;
 
     private AlterPartitioningSettings(Builder builder) {
         this.minActivePartitions = builder.minActivePartitions;
         this.partitionCountLimit = builder.partitionCountLimit;
+        this.autoPartitioningStrategy = builder.autoPartitioningStrategy;
+        this.writeStrategySettings = builder.writeStrategySettings;
     }
 
     public static Builder newBuilder() {
@@ -30,12 +36,24 @@ public Long getPartitionCountLimit() {
         return partitionCountLimit;
     }
 
+    @Nullable
+    public AutoPartitioningStrategy getAutoPartitioningStrategy() {
+        return autoPartitioningStrategy;
+    }
+
+    @Nullable
+    public AlterAutoPartitioningWriteStrategySettings getWriteStrategySettings() {
+        return writeStrategySettings;
+    }
+
     /**
      * BUILDER
      */
     public static class Builder {
         private Long minActivePartitions = null;
         private Long partitionCountLimit = null;
+        private AutoPartitioningStrategy autoPartitioningStrategy = null;
+        private AlterAutoPartitioningWriteStrategySettings writeStrategySettings = null;
 
         /**
          * @param minActivePartitions  minimum partition count auto merge would stop working at.
@@ -58,6 +76,28 @@ public Builder setPartitionCountLimit(long partitionCountLimit) {
             return this;
         }
 
+        /**
+         * @param autoPartitioningStrategy  Strategy for auto partitioning.
+         *                                  Auto partitioning is disabled by default.
+         * @return settings builder
+         * @see AutoPartitioningStrategy#DISABLED
+         */
+        public Builder setAutoPartitioningStrategy(AutoPartitioningStrategy autoPartitioningStrategy) {
+            this.autoPartitioningStrategy = autoPartitioningStrategy;
+            return this;
+        }
+
+        /**
+         * @param writeStrategySettings     Settings for auto partitioning write strategy.
+         *                                  Does not have any effect if auto partitioning is disabled.
+         *                                  See {@link AlterAutoPartitioningWriteStrategySettings} for defaults
+         * @return settings builder
+         */
+        public Builder setWriteStrategySettings(AlterAutoPartitioningWriteStrategySettings writeStrategySettings) {
+            this.writeStrategySettings = writeStrategySettings;
+            return this;
+        }
+
         public AlterPartitioningSettings build() {
             return new AlterPartitioningSettings(this);
         }
diff --git a/topic/src/main/java/tech/ydb/topic/settings/AutoPartitioningStrategy.java b/topic/src/main/java/tech/ydb/topic/settings/AutoPartitioningStrategy.java
@@ -0,0 +1,48 @@
+package tech.ydb.topic.settings;
+
+import tech.ydb.proto.topic.YdbTopic;
+
+import javax.annotation.Nullable;
+
+public enum AutoPartitioningStrategy {
+    /**
+     * The auto partitioning is disabled.
+     * You cannot disable the auto partitioning after it has been enabled.
+     * @see AutoPartitioningStrategy#PAUSED
+     */
+    DISABLED(YdbTopic.AutoPartitioningStrategy.AUTO_PARTITIONING_STRATEGY_DISABLED),
+    /**
+     * The auto partitioning algorithm will increase the partition count depending on the load characteristics.
+     * The auto partitioning algorithm will never decrease the number of partitions.
+     * @see AlterAutoPartitioningWriteStrategySettings
+     */
+    SCALE_UP(YdbTopic.AutoPartitioningStrategy.AUTO_PARTITIONING_STRATEGY_SCALE_UP),
+    /**
+     * The auto partitioning algorithm will both increase and decrease partitions count depending on the load characteristics.
+     * @see AlterAutoPartitioningWriteStrategySettings
+     */
+    SCALE_UP_AND_DOWN(YdbTopic.AutoPartitioningStrategy.AUTO_PARTITIONING_STRATEGY_SCALE_UP_AND_DOWN),
+    /**
+     * The auto partitioning is paused.
+     */
+    PAUSED(YdbTopic.AutoPartitioningStrategy.AUTO_PARTITIONING_STRATEGY_PAUSED);
+
+    private final YdbTopic.AutoPartitioningStrategy protoReference;
+
+    AutoPartitioningStrategy(YdbTopic.AutoPartitioningStrategy protoReference) {
+        this.protoReference = protoReference;
+    }
+
+    public YdbTopic.AutoPartitioningStrategy getProtoReference() {
+        return protoReference;
+    }
+
+    public static @Nullable AutoPartitioningStrategy fromProto(YdbTopic.AutoPartitioningStrategy protoReference) {
+        for (AutoPartitioningStrategy value : values()) {
+            if (value.getProtoReference() == protoReference) {
+                return value;
+            }
+        }
+        return null;
+    }
+}
diff --git a/topic/src/main/java/tech/ydb/topic/settings/AutoPartitioningWriteStrategySettings.java b/topic/src/main/java/tech/ydb/topic/settings/AutoPartitioningWriteStrategySettings.java
@@ -0,0 +1,68 @@
+package tech.ydb.topic.settings;
+
+import java.time.Duration;
+
+public class AutoPartitioningWriteStrategySettings {
+    private final Duration stabilizationWindow;
+    private final int upUtilizationPercent;
+    private final int downUtilizationPercent;
+
+    public AutoPartitioningWriteStrategySettings(Builder builder) {
+        this.stabilizationWindow = builder.stabilizationWindow;
+        this.upUtilizationPercent = builder.upUtilizationPercent;
+        this.downUtilizationPercent = builder.downUtilizationPercent;
+    }
+
+    public Duration getStabilizationWindow() {
+        return stabilizationWindow;
+    }
+
+    public int getUpUtilizationPercent() {
+        return upUtilizationPercent;
+    }
+
+    public int getDownUtilizationPercent() {
+        return downUtilizationPercent;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    public static class Builder {
+        private Duration stabilizationWindow = Duration.ofMinutes(5);
+        private int upUtilizationPercent = 90;
+        private int downUtilizationPercent = 30;
+
+        /**
+         * @param stabilizationWindow
+         * @return strategy builder
+         */
+        public Builder setStabilizationWindow(Duration stabilizationWindow) {
+            this.stabilizationWindow = stabilizationWindow;
+            return this;
+        }
+
+        /**
+         * @param upUtilizationPercent
+         * @return strategy builder
+         */
+        public Builder setUpUtilizationPercent(int upUtilizationPercent) {
+            this.upUtilizationPercent = upUtilizationPercent;
+            return this;
+        }
+
+        /**
+         * @param downUtilizationPercent
+         * @return strategy builder
+         */
+        public Builder setDownUtilizationPercent(int downUtilizationPercent) {
+            this.downUtilizationPercent = downUtilizationPercent;
+            return this;
+        }
+
+        public AutoPartitioningWriteStrategySettings build() {
+            return new AutoPartitioningWriteStrategySettings(this);
+        }
+    }
+}
diff --git a/topic/src/main/java/tech/ydb/topic/settings/PartitioningSettings.java b/topic/src/main/java/tech/ydb/topic/settings/PartitioningSettings.java
diff --git a/topic/src/test/java/tech/ydb/topic/impl/YdbTopicsIntegrationTest.java b/topic/src/test/java/tech/ydb/topic/impl/YdbTopicsIntegrationTest.java
