Added RetryConfig to configure retries behavior

alex268 · alex268 · commit 173ad53d6531 · 2025-03-10T13:46:14.000Z
diff --git a/topic/src/main/java/tech/ydb/topic/impl/GrpcStreamRetrier.java b/topic/src/main/java/tech/ydb/topic/impl/GrpcStreamRetrier.java
@@ -8,10 +8,10 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.BiConsumer;
 
 import org.slf4j.Logger;
 
+import tech.ydb.common.retry.RetryConfig;
 import tech.ydb.core.Status;
 
 /**
@@ -27,18 +27,18 @@ public abstract class GrpcStreamRetrier {
     private static final char[] ID_ALPHABET = "abcdefghijklmnopqrstuvwxyzABSDEFGHIJKLMNOPQRSTUVWXYZ1234567890"
             .toCharArray();
 
+    private final RetryConfig retryConfig;
     protected final String id;
     protected final AtomicBoolean isReconnecting = new AtomicBoolean(false);
     protected final AtomicBoolean isStopped = new AtomicBoolean(false);
     protected final AtomicInteger reconnectCounter = new AtomicInteger(0);
 
     private final ScheduledExecutorService scheduler;
-    private final BiConsumer<Status, Throwable> errorsHandler;
 
-    protected GrpcStreamRetrier(ScheduledExecutorService scheduler, BiConsumer<Status, Throwable> errorsHandler) {
+    protected GrpcStreamRetrier(RetryConfig retryConfig, ScheduledExecutorService scheduler) {
+        this.retryConfig = retryConfig;
         this.scheduler = scheduler;
         this.id = generateRandomId(ID_LENGTH);
-        this.errorsHandler = errorsHandler;
     }
 
     protected abstract Logger getLogger();
@@ -131,10 +131,6 @@ protected void onSessionClosed(Status status, Throwable th) {
             }
         }
 
-        if (errorsHandler != null) {
-            errorsHandler.accept(status, th);
-        }
-
         if (!isStopped.get()) {
             tryScheduleReconnect();
         } else  {
diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReaderImpl.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReaderImpl.java
@@ -55,7 +55,7 @@ public abstract class ReaderImpl extends GrpcStreamRetrier {
     private final String consumerName;
 
     public ReaderImpl(TopicRpc topicRpc, ReaderSettings settings) {
-        super(topicRpc.getScheduler(), settings.getErrorsHandler());
+        super(settings.getRetryConfig(), topicRpc.getScheduler());
         this.topicRpc = topicRpc;
         this.settings = settings;
         this.session = new ReadSessionImpl();
diff --git a/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java b/topic/src/main/java/tech/ydb/topic/settings/ReaderSettings.java
@@ -9,7 +9,10 @@
 
 import com.google.common.collect.ImmutableList;
 
+import tech.ydb.common.retry.RetryConfig;
+import tech.ydb.common.retry.RetryPolicy;
 import tech.ydb.core.Status;
+import tech.ydb.core.StatusCode;
 
 /**
  * @author Nikolay Perfilov
@@ -20,17 +23,17 @@ public class ReaderSettings {
     private final String consumerName;
     private final String readerName;
     private final List<TopicReadSettings> topics;
+    private final RetryConfig retryConfig;
     private final long maxMemoryUsageBytes;
     private final Executor decompressionExecutor;
-    private final BiConsumer<Status, Throwable> errorsHandler;
 
     private ReaderSettings(Builder builder) {
         this.consumerName = builder.consumerName;
         this.readerName = builder.readerName;
         this.topics = ImmutableList.copyOf(builder.topics);
+        this.retryConfig = builder.retryConfig;
         this.maxMemoryUsageBytes = builder.maxMemoryUsageBytes;
         this.decompressionExecutor = builder.decompressionExecutor;
-        this.errorsHandler = builder.errorsHandler;
     }
 
     public String getConsumerName() {
@@ -42,12 +45,17 @@ public String getReaderName() {
         return readerName;
     }
 
+    public RetryConfig getRetryConfig() {
+        return retryConfig;
+    }
+
     public List<TopicReadSettings> getTopics() {
         return topics;
     }
 
+    @Deprecated
     public BiConsumer<Status, Throwable> getErrorsHandler() {
-        return errorsHandler;
+        return null;
     }
 
     public long getMaxMemoryUsageBytes() {
@@ -70,9 +78,9 @@ public static class Builder {
         private boolean readWithoutConsumer = false;
         private String readerName = null;
         private List<TopicReadSettings> topics = new ArrayList<>();
+        private RetryConfig retryConfig = RetryConfig.idempotentRetryForever();
         private long maxMemoryUsageBytes = MAX_MEMORY_USAGE_BYTES_DEFAULT;
         private Executor decompressionExecutor = null;
-        private BiConsumer<Status, Throwable> errorsHandler = null;
 
         public Builder setConsumerName(String consumerName) {
             this.consumerName = consumerName;
@@ -91,6 +99,7 @@ public Builder withoutConsumer() {
 
         /**
          * Set reader name for debug purposes
+         * @param readerName name of reader
          * @return settings builder
          */
         public Builder setReaderName(String readerName) {
@@ -108,13 +117,42 @@ public Builder setTopics(List<TopicReadSettings> topics) {
             return this;
         }
 
+        /**
+         * Set {@link RetryConfig} to define behavior of the stream internal retries
+         * @param config retry mode
+         * @return settings builder
+         */
+        public Builder setRetryConfig(RetryConfig config) {
+            this.retryConfig = config;
+            return this;
+        }
+
         public Builder setMaxMemoryUsageBytes(long maxMemoryUsageBytes) {
             this.maxMemoryUsageBytes = maxMemoryUsageBytes;
             return this;
         }
 
+        /**
+         * @param handler
+         * @return builder
+         * @deprecated use {@link Builder#setRetryConfig(tech.ydb.common.retry.RetryConfig)} instead
+         */
+        @Deprecated
         public Builder setErrorsHandler(BiConsumer<Status, Throwable> handler) {
-            this.errorsHandler = handler;
+            final RetryConfig currentConfig = retryConfig;
+            retryConfig = new RetryConfig() {
+                @Override
+                public RetryPolicy isStatusRetryable(StatusCode code) {
+                    handler.accept(Status.of(code), null);
+                    return currentConfig.isStatusRetryable(code);
+                }
+
+                @Override
+                public RetryPolicy isThrowableRetryable(Throwable th) {
+                    handler.accept(null, th);
+                    return currentConfig.isThrowableRetryable(th);
+                }
+            };
             return this;
         }
 
diff --git a/topic/src/main/java/tech/ydb/topic/settings/WriterSettings.java b/topic/src/main/java/tech/ydb/topic/settings/WriterSettings.java
@@ -2,7 +2,10 @@
 
 import java.util.function.BiConsumer;
 
+import tech.ydb.common.retry.RetryConfig;
+import tech.ydb.common.retry.RetryPolicy;
 import tech.ydb.core.Status;
+import tech.ydb.core.StatusCode;
 import tech.ydb.topic.description.Codec;
 
 /**
@@ -17,19 +20,19 @@ public class WriterSettings {
     private final String messageGroupId;
     private final Long partitionId;
     private final Codec codec;
+    private final RetryConfig retryConfig;
     private final long maxSendBufferMemorySize;
     private final int maxSendBufferMessagesCount;
-    private final BiConsumer<Status, Throwable> errorsHandler;
 
     private WriterSettings(Builder builder) {
         this.topicPath = builder.topicPath;
         this.producerId = builder.producerId;
         this.messageGroupId = builder.messageGroupId;
         this.partitionId = builder.partitionId;
         this.codec = builder.codec;
+        this.retryConfig = builder.retryConfig;
         this.maxSendBufferMemorySize = builder.maxSendBufferMemorySize;
         this.maxSendBufferMessagesCount = builder.maxSendBufferMessagesCount;
-        this.errorsHandler = builder.errorsHandler;
     }
 
     public static Builder newBuilder() {
@@ -48,8 +51,9 @@ public String getMessageGroupId() {
         return messageGroupId;
     }
 
+    @Deprecated
     public BiConsumer<Status, Throwable> getErrorsHandler() {
-        return errorsHandler;
+        return null;
     }
 
     public Long getPartitionId() {
@@ -60,6 +64,10 @@ public Codec getCodec() {
         return codec;
     }
 
+    public RetryConfig getRetryConfig() {
+        return retryConfig;
+    }
+
     public long getMaxSendBufferMemorySize() {
         return maxSendBufferMemorySize;
     }
@@ -77,9 +85,9 @@ public static class Builder {
         private String messageGroupId = null;
         private Long partitionId = null;
         private Codec codec = Codec.GZIP;
+        private RetryConfig retryConfig = RetryConfig.idempotentRetryForever();
         private long maxSendBufferMemorySize = MAX_MEMORY_USAGE_BYTES_DEFAULT;
         private int maxSendBufferMessagesCount = MAX_IN_FLIGHT_COUNT_DEFAULT;
-        private BiConsumer<Status, Throwable> errorsHandler = null;
 
         /**
          * Set path to a topic to write to
@@ -135,6 +143,16 @@ public Builder setCodec(Codec codec) {
             return this;
         }
 
+        /**
+         * Set {@link RetryConfig} to define behavior of the stream internal retries
+         * @param config retry mode
+         * @return settings builder
+         */
+        public Builder setRetryConfig(RetryConfig config) {
+            this.retryConfig = config;
+            return this;
+        }
+
         /**
          * Set memory usage limit for send buffer.
          * Writer will not accept new messages if memory usage exceeds this limit.
@@ -158,8 +176,27 @@ public Builder setMaxSendBufferMessagesCount(int maxMessagesCount) {
             return this;
         }
 
+        /**
+         * @param handler
+         * @return builder
+         * @deprecated use {@link Builder#setRetryConfig(tech.ydb.common.retry.RetryConfig)} instead
+         */
+        @Deprecated
         public Builder setErrorsHandler(BiConsumer<Status, Throwable> handler) {
-            this.errorsHandler = handler;
+            final RetryConfig currentConfig = retryConfig;
+            retryConfig = new RetryConfig() {
+                @Override
+                public RetryPolicy isStatusRetryable(StatusCode code) {
+                    handler.accept(Status.of(code), null);
+                    return currentConfig.isStatusRetryable(code);
+                }
+
+                @Override
+                public RetryPolicy isThrowableRetryable(Throwable th) {
+                    handler.accept(null, th);
+                    return currentConfig.isThrowableRetryable(th);
+                }
+            };
             return this;
         }
 
diff --git a/topic/src/main/java/tech/ydb/topic/write/impl/WriterImpl.java b/topic/src/main/java/tech/ydb/topic/write/impl/WriterImpl.java
@@ -66,7 +66,7 @@ public abstract class WriterImpl extends GrpcStreamRetrier {
     private CompletableFuture<WriteAck> lastAcceptedMessageFuture;
 
     public WriterImpl(TopicRpc topicRpc, WriterSettings settings, Executor compressionExecutor) {
-        super(topicRpc.getScheduler(), settings.getErrorsHandler());
+        super(settings.getRetryConfig(), topicRpc.getScheduler());
         this.topicRpc = topicRpc;
         this.settings = settings;
         this.session = new WriteSessionImpl();
