intermediate changes

ref:f9eefd8eed3595f41a195fc7080bbf602b0075af

3 files changed, 148 insertions, 34 deletions

diff --git a/library/cpp/retry/retry.cpp b/library/cpp/retry/retry.cpp
index 92466cdeca..f41c97e179 100644
--- a/library/cpp/retry/retry.cpp
+++ b/library/cpp/retry/retry.cpp
@@ -2,19 +2,44 @@
 
 #include <util/stream/output.h>
 
-void DoWithRetry(std::function<void()> func, TRetryOptions retryOptions) {
-    DoWithRetry(func, retryOptions, true);
-}
+namespace {
+class TRetryOptionsWithRetCodePolicy : public IRetryPolicy<bool> {
+public:
+    explicit TRetryOptionsWithRetCodePolicy(const TRetryOptions& opts)
+        : Opts(opts)
+    {
+    }
 
-bool DoWithRetryOnRetCode(std::function<bool()> func, TRetryOptions retryOptions) {
-    for (ui32 attempt = 0; attempt <= retryOptions.RetryCount; ++attempt) {
-        if (func()) {
-            return true;
+    class TRetryState : public IRetryState {
+    public:
+        explicit TRetryState(const TRetryOptions& opts)
+            : Opts(opts)
+        {
+        }
+
+        TMaybe<TDuration> GetNextRetryDelay(bool ret) override {
+            if (ret || Attempt == Opts.RetryCount) {
+                return Nothing();
+            }
+            return Opts.GetTimeToSleep(Attempt++);
         }
-        auto sleep = retryOptions.SleepFunction;
-        sleep(retryOptions.GetTimeToSleep(attempt));
+
+    private:
+        const TRetryOptions Opts;
+        size_t Attempt = 0;
+    };
+
+    IRetryState::TPtr CreateRetryState() const override {
+        return std::make_unique<TRetryState>(Opts);
     }
-    return false;
+
+private:
+    const TRetryOptions Opts;
+};
+} // namespace
+
+bool DoWithRetryOnRetCode(std::function<bool()> func, TRetryOptions retryOptions) {
+    return DoWithRetryOnRetCode<bool>(std::move(func), std::make_shared<TRetryOptionsWithRetCodePolicy>(retryOptions), retryOptions.SleepFunction);
 }
 
 TRetryOptions MakeRetryOptions(const NRetry::TRetryOptionsPB& retryOptions) {
diff --git a/library/cpp/retry/retry.h b/library/cpp/retry/retry.h
index c47ff5070f..cb1f6a4ba4 100644
--- a/library/cpp/retry/retry.h
+++ b/library/cpp/retry/retry.h
@@ -1,5 +1,5 @@
 #pragma once
-
+#include "retry_policy.h"
 #include "utils.h"
 
 #include <library/cpp/retry/protos/retry_options.pb.h>
@@ -81,20 +81,88 @@ struct TRetryOptions {
     }
 };
 
+TRetryOptions MakeRetryOptions(const NRetry::TRetryOptionsPB& retryOptions);
+
+
+namespace NRetryDetails {
+
+template <class TException>
+class TRetryOptionsPolicy : public IRetryPolicy<const TException&> {
+public:
+    explicit TRetryOptionsPolicy(const TRetryOptions& opts)
+        : Opts(opts)
+    {
+    }
+
+    using IRetryState = typename IRetryPolicy<const TException&>::IRetryState;
+
+    class TRetryState : public IRetryState {
+    public:
+        explicit TRetryState(const TRetryOptions& opts)
+            : Opts(opts)
+        {
+        }
+
+        TMaybe<TDuration> GetNextRetryDelay(const TException&) override {
+            if (Attempt == Opts.RetryCount) {
+                return Nothing();
+            }
+            return Opts.GetTimeToSleep(Attempt++);
+        }
+
+    private:
+        const TRetryOptions Opts;
+        size_t Attempt = 0;
+    };
+
+    typename IRetryState::TPtr CreateRetryState() const override {
+        return std::make_unique<TRetryState>(Opts);
+    }
+
+private:
+    const TRetryOptions Opts;
+};
+
+} // namespace NRetryDetails
+
+template <class TException>
+typename IRetryPolicy<const TException&>::TPtr MakeRetryPolicy(const TRetryOptions& opts) {
+    return std::make_shared<NRetryDetails::TRetryOptionsPolicy<TException>>(opts);
+}
+
+template <class TException>
+typename IRetryPolicy<const TException&>::TPtr MakeRetryPolicy(const NRetry::TRetryOptionsPB& opts) {
+    return MakeRetryPolicy<TException>(MakeRetryOptions(opts));
+}
+
 template <typename TResult, typename TException = yexception>
-TMaybe<TResult> DoWithRetry(std::function<TResult()> func, std::function<void(const TException&)> onFail, TRetryOptions retryOptions, bool throwLast) {
-    for (ui32 attempt = 0; attempt <= retryOptions.RetryCount; ++attempt) {
+TMaybe<TResult> DoWithRetry(std::function<TResult()> func, const typename IRetryPolicy<const TException&>::TPtr& retryPolicy, bool throwLast = true, std::function<void(const TException&)> onFail = {}, std::function<void(TDuration)> sleepFunction = {}) {
+    typename IRetryPolicy<const TException&>::IRetryState::TPtr retryState;
+    while (true) {
         try {
             return func();
-        } catch (TException& ex) {
-            onFail(ex);
-            if (attempt == retryOptions.RetryCount) {
+        } catch (const TException& ex) {
+            if (onFail) {
+                onFail(ex);
+            }
+
+            if (!retryState) {
+                retryState = retryPolicy->CreateRetryState();
+            }
+
+            if (const TMaybe<TDuration> delay = retryState->GetNextRetryDelay(ex)) {
+                if (*delay) {
+                    if (sleepFunction) {
+                        sleepFunction(*delay);
+                    } else {
+                        Sleep(*delay);
+                    }
+                }
+            } else {
                 if (throwLast) {
                     throw;
                 }
-            } else {
-                auto sleep = retryOptions.SleepFunction;
-                sleep(retryOptions.GetTimeToSleep(attempt));
+                break;
             }
         }
     }
@@ -102,32 +170,53 @@ TMaybe<TResult> DoWithRetry(std::function<TResult()> func, std::function<void(co
 }
 
 template <typename TResult, typename TException = yexception>
-TMaybe<TResult> DoWithRetry(std::function<TResult()> func, TRetryOptions retryOptions, bool throwLast) {
-    return DoWithRetry<TResult, TException>(func, [](const TException&){}, retryOptions, throwLast);
+TMaybe<TResult> DoWithRetry(std::function<TResult()> func, std::function<void(const TException&)> onFail, TRetryOptions retryOptions, bool throwLast = true) {
+    return DoWithRetry<TResult, TException>(std::move(func), MakeRetryPolicy<TException>(retryOptions), throwLast, std::move(onFail), retryOptions.SleepFunction);
+}
+
+template <typename TResult, typename TException = yexception>
+TMaybe<TResult> DoWithRetry(std::function<TResult()> func, TRetryOptions retryOptions, bool throwLast = true) {
+    return DoWithRetry<TResult, TException>(std::move(func), MakeRetryPolicy<TException>(retryOptions), throwLast, {}, retryOptions.SleepFunction);
 }
 
 template <typename TException = yexception>
-bool DoWithRetry(std::function<void()> func, std::function<void(const TException&)> onFail, TRetryOptions retryOptions, bool throwLast) {
+bool DoWithRetry(std::function<void()> func, const typename IRetryPolicy<const TException&>::TPtr& retryPolicy, bool throwLast = true, std::function<void(const TException&)> onFail = {}, std::function<void(TDuration)> sleepFunction = {}) {
     auto f = [&]() {
         func();
         return nullptr;
     };
-    return DoWithRetry<void*, TException>(f, onFail, retryOptions, throwLast).Defined();
+    return DoWithRetry<void*, TException>(f, retryPolicy, throwLast, std::move(onFail), std::move(sleepFunction)).Defined();
 }
 
 template <typename TException = yexception>
-bool DoWithRetry(std::function<void()> func, TRetryOptions retryOptions, bool throwLast) {
-    auto f = [&]() {
-        func();
-        return nullptr;
-    };
-    return DoWithRetry<void*, TException>(f, [](const TException&){}, retryOptions, throwLast).Defined();
+bool DoWithRetry(std::function<void()> func, std::function<void(const TException&)> onFail, TRetryOptions retryOptions, bool throwLast) {
+    return DoWithRetry<TException>(std::move(func), MakeRetryPolicy<TException>(retryOptions), throwLast, onFail, retryOptions.SleepFunction);
+}
+
+template <typename TException = yexception>
+bool DoWithRetry(std::function<void()> func, TRetryOptions retryOptions, bool throwLast = true) {
+    return DoWithRetry<TException>(std::move(func), MakeRetryPolicy<TException>(retryOptions), throwLast, {}, retryOptions.SleepFunction);
 }
 
-void DoWithRetry(std::function<void()> func, TRetryOptions retryOptions);
+template <class TRetCode>
+TRetCode DoWithRetryOnRetCode(std::function<TRetCode()> func, const typename IRetryPolicy<TRetCode>::TPtr& retryPolicy, std::function<void(TDuration)> sleepFunction = {}) {
+    auto retryState = retryPolicy->CreateRetryState();
+    while (true) {
+        TRetCode code = func();
+        if (const TMaybe<TDuration> delay = retryState->GetNextRetryDelay(code)) {
+            if (*delay) {
+                if (sleepFunction) {
+                    sleepFunction(*delay);
+                } else {
+                    Sleep(*delay);
+                }
+            }
+        } else {
+            return code;
+        }
+    }
+}
 
 bool DoWithRetryOnRetCode(std::function<bool()> func, TRetryOptions retryOptions);
 
 Y_DECLARE_PODTYPE(TRetryOptions);
-
-TRetryOptions MakeRetryOptions(const NRetry::TRetryOptionsPB& retryOptions);
diff --git a/library/cpp/retry/retry_policy.h b/library/cpp/retry/retry_policy.h
index a2f94840b8..c7a3156e10 100644
--- a/library/cpp/retry/retry_policy.h
+++ b/library/cpp/retry/retry_policy.h
@@ -40,7 +40,7 @@ struct IRetryPolicy {
 
         //! Calculate delay before next retry if next retry is allowed.
         //! Returns empty maybe if retry is not allowed anymore.
-        virtual TMaybe<TDuration> GetNextRetryDelay(typename TTypeTraits<TArgs>::TFuncParam... args) = 0;
+        [[nodiscard]] virtual TMaybe<TDuration> GetNextRetryDelay(typename TTypeTraits<TArgs>::TFuncParam... args) = 0;
     };
 
     virtual ~IRetryPolicy() = default;
@@ -48,7 +48,7 @@ struct IRetryPolicy {
     //! Function that is called after first error
     //! to find out a futher retry behaviour.
     //! Retry state is expected to be created for the whole single retry session.
-    virtual typename IRetryState::TPtr CreateRetryState() const = 0;
+    [[nodiscard]] virtual typename IRetryState::TPtr CreateRetryState() const = 0;
 
     //!
     //! Default implementations.