KIKIMR-19051: refactor RetryContext

In order to reduce number of typenames:
* added Sync and Async namespaces
* added RetryOperation return type deduction
* TRetryContextBase does not use Client

7 files changed, 103 insertions, 96 deletions

diff --git a/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry.h b/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry.h
index 09c3259171e..937884b6735 100644
--- a/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry.h
+++ b/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry.h
@@ -31,17 +31,14 @@ enum class NextStep {
     Finish,
 };
 
-template <typename TClient>
 class TRetryContextBase : TNonCopyable {
 protected:
-    TClient Client;
     TRetryOperationSettings Settings;
     ui32 RetryNumber;
 
 protected:
-    TRetryContextBase(const TClient& client, const TRetryOperationSettings& settings)
-        : Client(client)
-        , Settings(settings)
+    TRetryContextBase(const TRetryOperationSettings& settings)
+        : Settings(settings)
         , RetryNumber(0)
     {}
 
diff --git a/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_async.h b/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_async.h
index bf928f42587..e07e6cdd7a2 100644
--- a/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_async.h
+++ b/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_async.h
@@ -2,39 +2,46 @@
 
 #include <ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry.h>
 
-namespace NYdb::NRetry {
+#include <util/generic/function.h>
 
-template <typename TClient, typename TStatusType>
-class TRetryContextAsync : public TThrRefBase, public TRetryContextBase<TClient> {
+namespace NYdb::NRetry::Async {
+
+template <typename TClient, typename TAsyncStatusType>
+class TRetryContext : public TThrRefBase, public TRetryContextBase {
 public:
-    using TPtr = TIntrusivePtr<NYdb::NRetry::TRetryContextAsync<TClient, TStatusType>>;
-    using TAsyncStatusType = typename NThreading::TFuture<TStatusType>;
+    using TStatusType = typename TAsyncStatusType::value_type;
+    using TPtr = TIntrusivePtr<Async::TRetryContext<TClient, TAsyncStatusType>>;
 
 protected:
+    TClient Client;
     NThreading::TPromise<TStatusType> Promise;
 
 public:
-    TAsyncStatusType GetFuture() {
+    TAsyncStatusType Execute() {
+        this->Retry();
         return this->Promise.GetFuture();
     }
 
-    virtual void Execute() = 0;
-
 protected:
-    explicit TRetryContextAsync(const TClient& client, const TRetryOperationSettings& settings)
-        : TRetryContextBase<TClient>(client, settings)
+    explicit TRetryContext(const TClient& client, const TRetryOperationSettings& settings)
+        : TRetryContextBase(settings)
+        , Client(client)
         , Promise(NThreading::NewPromise<TStatusType>())
     {}
 
-    static void DoExecute(TPtr self) {
-        self->Execute();
+    virtual void Retry() = 0;
+
+    virtual TAsyncStatusType RunOperation() = 0;
+
+    static void DoRetry(TPtr self) {
+        self->Retry();
     }
 
     static void DoBackoff(TPtr self, bool fast) {
         auto backoffSettings = fast ? self->Settings.FastBackoffSettings_
                                     : self->Settings.SlowBackoffSettings_;
         AsyncBackoff(self->Client.Impl_, backoffSettings, self->RetryNumber,
-            [self]() {DoExecute(self);});
+            [self]() {DoRetry(self);});
     }
 
     static void HandleExceptionAsync(TPtr self, std::exception_ptr e) {
@@ -50,7 +57,7 @@ protected:
         }
         switch (nextStep) {
             case NextStep::RetryImmediately:
-                return DoExecute(self);
+                return DoRetry(self);
             case NextStep::RetryFastBackoff:
                 return DoBackoff(self, true);
             case NextStep::RetrySlowBackoff:
@@ -71,42 +78,39 @@ protected:
             }
         );
     }
-
-    virtual TAsyncStatusType RunOperation() = 0;
 };
 
-template <typename TClient, typename TOperation, typename TStatusType>
-class TRetryWithoutSessionAsync : public TRetryContextAsync<TClient, TStatusType> {
-    using TRetryContextAsync = TRetryContextAsync<TClient, TStatusType>;
-    using TPtr = typename TRetryContextAsync::TPtr;
-    using TAsyncStatusType = typename TRetryContextAsync::TAsyncStatusType;
+template <typename TClient, typename TOperation, typename TAsyncStatusType = TFunctionResult<TOperation>>
+class TRetryWithoutSession : public TRetryContext<TClient, TAsyncStatusType> {
+    using TRetryContext = TRetryContext<TClient, TAsyncStatusType>;
+    using TPtr = typename TRetryContext::TPtr;
 
 private:
     TOperation Operation;
 
 public:
-    explicit TRetryWithoutSessionAsync(const TClient& client, TOperation&& operation,
+    explicit TRetryWithoutSession(const TClient& client, TOperation&& operation,
                                        const TRetryOperationSettings& settings)
-        : TRetryContextAsync(client, settings)
+        : TRetryContext(client, settings)
         , Operation(operation)
     {}
 
-    void Execute() override {
+    void Retry() override {
         TPtr self(this);
-        TRetryContextAsync::DoRunOperation(self);
+        TRetryContext::DoRunOperation(self);
     }
 
-private:
+protected:
     TAsyncStatusType RunOperation() override {
         return Operation(this->Client);
     }
 };
 
-template <typename TClient, typename TOperation, typename TStatusType = typename std::result_of<TOperation>::type>
-class TRetryWithSessionAsync : public TRetryContextAsync<TClient, TStatusType> {
-    using TRetryContextAsync = TRetryContextAsync<TClient, TStatusType>;
+template <typename TClient, typename TOperation, typename TAsyncStatusType = TFunctionResult<TOperation>>
+class TRetryWithSession : public TRetryContext<TClient, TAsyncStatusType> {
+    using TRetryContextAsync = TRetryContext<TClient, TAsyncStatusType>;
     using TPtr = typename TRetryContextAsync::TPtr;
-    using TAsyncStatusType = typename TRetryContextAsync::TAsyncStatusType;
+    using TStatusType = typename TRetryContextAsync::TStatusType;
     using TSession = typename TClient::TSession;
     using TCreateSessionSettings = typename TClient::TCreateSessionSettings;
     using TAsyncCreateSessionResult = typename TClient::TAsyncCreateSessionResult;
@@ -116,13 +120,13 @@ private:
     TMaybe<TSession> Session;
 
 public:
-    explicit TRetryWithSessionAsync(const TClient& client, TOperation&& operation,
+    explicit TRetryWithSession(const TClient& client, TOperation&& operation,
                                     const TRetryOperationSettings& settings)
         : TRetryContextAsync(client, settings)
         , Operation(operation)
     {}
 
-    void Execute() override {
+    void Retry() override {
         TPtr self(this);
         if (!Session) {
             auto settings = TCreateSessionSettings().ClientTimeout(this->Settings.GetSessionClientTimeout_);
@@ -134,7 +138,7 @@ public:
                             return TRetryContextAsync::HandleStatusAsync(self, TStatusType(TStatus(result)));
                         }
 
-                        auto* myself = dynamic_cast<TRetryWithSessionAsync*>(self.Get());
+                        auto* myself = dynamic_cast<TRetryWithSession*>(self.Get());
                         myself->Session = result.GetSession();
                         myself->DoRunOperation(self);
                     } catch (...) {
@@ -157,4 +161,4 @@ private:
     }
 };
 
-} // namespace NYdb::NRetry
+} // namespace NYdb::NRetry::Async
diff --git a/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_sync.h b/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_sync.h
index 41aad0d31d1..f62e4a468a6 100644
--- a/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_sync.h
+++ b/ydb/public/sdk/cpp/client/impl/ydb_internal/retry/retry_sync.h
@@ -6,17 +6,16 @@
 
 #include <util/generic/maybe.h>
 
-namespace NYdb::NRetry {
+namespace NYdb::NRetry::Sync {
 
-template <typename TClient>
-class TRetryContextSync : public TRetryContextBase<TClient> {
-public:
-    TRetryContextSync(TClient& client, const TRetryOperationSettings& settings)
-        : TRetryContextBase<TClient>(client, settings)
-    {}
+template <typename TClient, typename TStatusType>
+class TRetryContext : public TRetryContextBase {
+protected:
+    TClient Client;
 
-    TStatus Execute() {
-        TStatus status = RunOperation(); // first attempt
+public:
+    TStatusType Execute() {
+        TStatusType status = Retry(); // first attempt
         for (this->RetryNumber = 0; this->RetryNumber <= this->Settings.MaxRetries_;) {
             auto nextStep = this->GetNextStep(status);
             switch (nextStep) {
@@ -35,40 +34,51 @@ public:
             this->RetryNumber++;
             this->LogRetry(status);
             this->Client.Impl_->CollectRetryStatSync(status.GetStatus());
-            status = RunOperation();
+            status = Retry();
         }
         return status;
     }
 
 protected:
+    TRetryContext(TClient& client, const TRetryOperationSettings& settings)
+        : TRetryContextBase(settings)
+        , Client(client)
+    {}
+
+    virtual TStatusType Retry() = 0;
+
+    virtual TStatusType RunOperation() = 0;
+
     void DoBackoff(bool fast) {
         const auto &settings = fast ? this->Settings.FastBackoffSettings_
                                     : this->Settings.SlowBackoffSettings_;
         Backoff(settings, this->RetryNumber);
     }
-
-    virtual TStatus RunOperation() = 0;
 };
 
-template<typename TClient, typename TOperation>
-class TRetryWithoutSessionSync : public TRetryContextSync<TClient> {
+template<typename TClient, typename TOperation, typename TStatusType = TFunctionResult<TOperation>>
+class TRetryWithoutSession : public TRetryContext<TClient, TStatusType> {
 private:
     const TOperation& Operation;
 
 public:
-    TRetryWithoutSessionSync(TClient& client, const TOperation& operation, const TRetryOperationSettings& settings)
-        : TRetryContextSync<TClient>(client, settings)
+    TRetryWithoutSession(TClient& client, const TOperation& operation, const TRetryOperationSettings& settings)
+        : TRetryContext<TClient, TStatusType>(client, settings)
         , Operation(operation)
     {}
 
 protected:
-    TStatus RunOperation() override {
+    TStatusType Retry() override {
+        return RunOperation();
+    }
+
+    TStatusType RunOperation() override {
         return Operation(this->Client);
     }
 };
 
-template<typename TClient, typename TOperation>
-class TRetryWithSessionSync : public TRetryContextSync<TClient> {
+template<typename TClient, typename TOperation, typename TStatusType = TFunctionResult<TOperation>>
+class TRetryWithSession : public TRetryContext<TClient, TStatusType> {
     using TSession = typename TClient::TSession;
     using TCreateSessionSettings = typename TClient::TCreateSessionSettings;
 
@@ -77,14 +87,14 @@ private:
     TMaybe<TSession> Session;
 
 public:
-    TRetryWithSessionSync(TClient& client, const TOperation& operation, const TRetryOperationSettings& settings)
-        : TRetryContextSync<TClient>(client, settings)
+    TRetryWithSession(TClient& client, const TOperation& operation, const TRetryOperationSettings& settings)
+        : TRetryContext<TClient, TStatusType>(client, settings)
         , Operation(operation)
     {}
 
 protected:
-    TStatus RunOperation() override {
-        TMaybe<NYdb::TStatus> status;
+    TStatusType Retry() override {
+        TMaybe<TStatusType> status;
 
         if (!Session) {
             auto settings = TCreateSessionSettings().ClientTimeout(this->Settings.GetSessionClientTimeout_);
@@ -92,19 +102,23 @@ protected:
             if (sessionResult.IsSuccess()) {
                 Session = sessionResult.GetSession();
             }
-            status = sessionResult;
+            status = TStatusType(TStatus(sessionResult));
         }
 
         if (Session) {
-            status = Operation(Session.GetRef());
+            status = RunOperation();
         }
 
         return *status;
     }
 
+    TStatusType RunOperation() override {
+        return Operation(this->Session.GetRef());
+    }
+
     void Reset() override {
         Session.Clear();
     }
 };
 
-} // namespace NYdb::NRetry
+} // namespace NYdb::NRetry::Sync
diff --git a/ydb/public/sdk/cpp/client/ydb_query/client.cpp b/ydb/public/sdk/cpp/client/ydb_query/client.cpp
index a223f874ec5..b0987066e7e 100644
--- a/ydb/public/sdk/cpp/client/ydb_query/client.cpp
+++ b/ydb/public/sdk/cpp/client/ydb_query/client.cpp
@@ -18,8 +18,7 @@
 
 namespace NYdb::NQuery {
 
-using TRetryContext = NRetry::TRetryContextAsync<TQueryClient, TExecuteQueryResult>;
-using TRetryContextWithSession = NRetry::TRetryWithSessionAsync<TQueryClient, TQueryClient::TQueryFunc, TExecuteQueryResult>;
+using TRetryContextAsync = NRetry::Async::TRetryContext<TQueryClient, TAsyncExecuteQueryResult>;
 
 TCreateSessionSettings::TCreateSessionSettings() {
     ClientTimeout_ = TDuration::Seconds(5);
@@ -380,7 +379,6 @@ TAsyncExecuteQueryResult TQueryClient::ExecuteQuery(const TString& query, const
     return Impl_->ExecuteQuery(query, txControl, params, settings);
 }
 
-
 TAsyncExecuteQueryIterator TQueryClient::StreamExecuteQuery(const TString& query, const TTxControl& txControl,
     const TExecuteQuerySettings& settings)
 {
@@ -424,9 +422,8 @@ i64 TQueryClient::GetCurrentPoolSize() const {
 
 TAsyncExecuteQueryResult TQueryClient::RetryQuery(TQueryFunc&& queryFunc, TRetryOperationSettings settings)
 {
-    TRetryContext::TPtr ctx(new TRetryContextWithSession(*this, std::move(queryFunc), settings));
-    ctx->Execute();
-    return ctx->GetFuture();
+    TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(queryFunc), settings));
+    return ctx->Execute();
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/ydb/public/sdk/cpp/client/ydb_query/client.h b/ydb/public/sdk/cpp/client/ydb_query/client.h
index dc981bdc8f9..31069a26c03 100644
--- a/ydb/public/sdk/cpp/client/ydb_query/client.h
+++ b/ydb/public/sdk/cpp/client/ydb_query/client.h
@@ -14,9 +14,9 @@
 namespace NYdb {
     class TProtoAccessor;
 
-    namespace NRetry {
-        template <typename TClient, typename TStatusType>
-        class TRetryContextAsync;
+    namespace NRetry::Async {
+        template <typename TClient, typename TAsyncStatusType>
+        class TRetryContext;
     }
 }
 
@@ -57,7 +57,7 @@ struct TClientSettings : public TCommonClientSettingsBase<TClientSettings> {
 class TSession;
 class TQueryClient {
     friend class TSession;
-    friend class NRetry::TRetryContextAsync<TQueryClient, TExecuteQueryResult>;
+    friend class NRetry::Async::TRetryContext<TQueryClient, TAsyncExecuteQueryResult>;
 
 public:
     using TQueryFunc = std::function<TAsyncExecuteQueryResult(TSession session)>;
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.cpp b/ydb/public/sdk/cpp/client/ydb_table/table.cpp
index 07d926857e0..2ef0878da5e 100644
--- a/ydb/public/sdk/cpp/client/ydb_table/table.cpp
+++ b/ydb/public/sdk/cpp/client/ydb_table/table.cpp
@@ -36,12 +36,7 @@ namespace NTable {
 using namespace NThreading;
 using namespace NSessionPool;
 
-using TRetryContext = NRetry::TRetryContextAsync<TTableClient, TStatus>;
-using TRetryWithSession = NRetry::TRetryWithSessionAsync<TTableClient, TTableClient::TOperationFunc, TStatus>;
-using TRetryWithoutSession = NRetry::TRetryWithoutSessionAsync<TTableClient, TTableClient::TOperationWithoutSessionFunc, TStatus>;
-
-using TRetryWithSessionSync = NRetry::TRetryWithSessionSync<TTableClient, TTableClient::TOperationSyncFunc>;
-using TRetryWithoutSessionSync = NRetry::TRetryWithoutSessionSync<TTableClient, TTableClient::TOperationWithoutSessionSyncFunc>;
+using TRetryContextAsync = NRetry::Async::TRetryContext<TTableClient, TAsyncStatus>;
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -1384,24 +1379,22 @@ TTypeBuilder TTableClient::GetTypeBuilder() {
 ////////////////////////////////////////////////////////////////////////////////
 
 TAsyncStatus TTableClient::RetryOperation(TOperationFunc&& operation, const TRetryOperationSettings& settings) {
-    TRetryContext::TPtr ctx(new TRetryWithSession(*this, std::move(operation), settings));
-    ctx->Execute();
-    return ctx->GetFuture();
+    TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithSession(*this, std::move(operation), settings));
+    return ctx->Execute();
 }
 
 TAsyncStatus TTableClient::RetryOperation(TOperationWithoutSessionFunc&& operation, const TRetryOperationSettings& settings) {
-    TRetryContext::TPtr ctx(new TRetryWithoutSession(*this, std::move(operation), settings));
-    ctx->Execute();
-    return ctx->GetFuture();
+    TRetryContextAsync::TPtr ctx(new NRetry::Async::TRetryWithoutSession(*this, std::move(operation), settings));
+    return ctx->Execute();
 }
 
 TStatus TTableClient::RetryOperationSync(const TOperationWithoutSessionSyncFunc& operation, const TRetryOperationSettings& settings) {
-    TRetryWithoutSessionSync ctx(*this, operation, settings);
+    NRetry::Sync::TRetryWithoutSession ctx(*this, operation, settings);
     return ctx.Execute();
 }
 
 TStatus TTableClient::RetryOperationSync(const TOperationSyncFunc& operation, const TRetryOperationSettings& settings) {
-    TRetryWithSessionSync ctx(*this, operation, settings);
+    NRetry::Sync::TRetryWithSession ctx(*this, operation, settings);
     return ctx.Execute();
 }
 
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.h b/ydb/public/sdk/cpp/client/ydb_table/table.h
index 435c5716bb0..06300e9ce75 100644
--- a/ydb/public/sdk/cpp/client/ydb_table/table.h
+++ b/ydb/public/sdk/cpp/client/ydb_table/table.h
@@ -35,13 +35,15 @@ class ValueSinceUnixEpochModeSettings;
 
 namespace NYdb {
 
-namespace NRetry {
+namespace NRetry::Async {
 template <typename TClient, typename TStatusType>
-class TRetryContextAsync;
+class TRetryContext;
+} // namespace NRetry::Async
 
-template <typename>
-class TRetryContextSync;
-} // namespace NRetry
+namespace NRetry::Sync {
+template <typename TClient, typename TStatusType>
+class TRetryContext;
+} // namespace NRetry::Sync
 
 namespace NScheme {
 struct TPermissions;
@@ -978,8 +980,8 @@ class TTableClient {
     friend class TSession;
     friend class TTransaction;
     friend class TSessionPool;
-    friend class NRetry::TRetryContextAsync<TTableClient, TStatus>;
-    friend class NRetry::TRetryContextSync<TTableClient>;
+    friend class NRetry::Sync::TRetryContext<TTableClient, TStatus>;
+    friend class NRetry::Async::TRetryContext<TTableClient, TAsyncStatus>;
 
 public:
     using TOperationFunc = std::function<TAsyncStatus(TSession session)>;