Fixed pool limit exceeding after server restart in QueryClient (#13901)

Author: Gazizonoki

include/ydb-cpp-sdk/client/query/client.h

@@ -151,6 +151,7 @@ class TSession {
     class TImpl;
 private:
     TSession();
+    TSession(std::shared_ptr<TQueryClient::TImpl> client); // Create broken session
     TSession(std::shared_ptr<TQueryClient::TImpl> client, TSession::TImpl* sessionImpl);
 
     std::shared_ptr<TQueryClient::TImpl> Client_;

src/client/query/client.cpp

@@ -361,7 +361,7 @@ class TQueryClient::TImpl: public TClientImplCommon<TQueryClient::TImpl>, public
                 TSession::TImpl::MakeImplAsync(processor, args);
             } else {
                 TStatus st(std::move(status));
-                args->Promise.SetValue(TCreateSessionResult(std::move(st), TSession()));
+                args->Promise.SetValue(TCreateSessionResult(std::move(st), TSession(args->Client)));
             }
         },
         &Ydb::Query::V1::QueryService::Stub::AsyncAttachSession,
@@ -384,13 +384,13 @@ class TQueryClient::TImpl: public TClientImplCommon<TQueryClient::TImpl>, public
                     NYdb::NIssue::TIssues opIssues;
                     NYdb::NIssue::IssuesFromMessage(resp->issues(), opIssues);
                     TStatus st(static_cast<EStatus>(resp->status()), std::move(opIssues));
-                    promise.SetValue(TCreateSessionResult(std::move(st), TSession()));
+                    promise.SetValue(TCreateSessionResult(std::move(st), TSession(self)));
                 } else {
                     self->DoAttachSession(resp, promise, status.Endpoint, self);
                 }
             } else {
                 TStatus st(std::move(status));
-                promise.SetValue(TCreateSessionResult(std::move(st), TSession()));
+                promise.SetValue(TCreateSessionResult(std::move(st), TSession(self)));
             }
         };
 
@@ -631,6 +631,14 @@ TSession TCreateSessionResult::GetSession() const {
 TSession::TSession()
 {}
 
+TSession::TSession(std::shared_ptr<TQueryClient::TImpl> client)
+    : Client_(client)
+    , SessionImpl_(
+        new TSession::TImpl(nullptr, "", "", client),
+        TKqpSessionCommon::GetSmartDeleter(client)
+    )
+{}
+
 TSession::TSession(std::shared_ptr<TQueryClient::TImpl> client, TSession::TImpl* session)
     : Client_(client)
     , SessionImpl_(session, TKqpSessionCommon::GetSmartDeleter(client))

src/client/query/impl/client_session.cpp

@@ -79,13 +79,12 @@ void TSession::TImpl::StartAsyncRead(TStreamProcessorPtr ptr, std::weak_ptr<ISes
             case grpc::StatusCode::OK:
                 StartAsyncRead(ptr, client, holder);
                 break;
-            case grpc::StatusCode::OUT_OF_RANGE: {
+            default: {
                 auto impl = holder->TrySharedOwning();
                 if (impl) {
                     impl->CloseFromServer(client);
                     holder->Release();
                 }
-                break;
             }
         }
     });
@@ -96,14 +95,21 @@ TSession::TImpl::TImpl(TStreamProcessorPtr ptr, const std::string& sessionId, co
     , StreamProcessor_(ptr)
     , SessionHolder(std::make_shared<TSafeTSessionImplHolder>(this))
 {
-    MarkActive();
-    SetNeedUpdateActiveCounter(true);
-    StartAsyncRead(StreamProcessor_, client, SessionHolder);
+    if (ptr) {
+        MarkActive();
+        SetNeedUpdateActiveCounter(true);
+        StartAsyncRead(StreamProcessor_, client, SessionHolder);
+    } else {
+        MarkBroken();
+        SetNeedUpdateActiveCounter(true);
+    }
 }
 
 TSession::TImpl::~TImpl()
 {
-    StreamProcessor_->Cancel();
+    if (StreamProcessor_) {
+        StreamProcessor_->Cancel();
+    }
     SessionHolder->WaitAndLock();
 }
 
@@ -114,7 +120,7 @@ void TSession::TImpl::MakeImplAsync(TStreamProcessorPtr ptr,
     ptr->Read(resp.get(), [args, resp, ptr](NYdbGrpc::TGrpcStatus grpcStatus) mutable {
         if (grpcStatus.GRpcStatusCode != grpc::StatusCode::OK) {
             TStatus st(TPlainStatus(grpcStatus, args->Endpoint));
-            args->Promise.SetValue(TCreateSessionResult(std::move(st), TSession()));
+            args->Promise.SetValue(TCreateSessionResult(std::move(st), TSession(args->Client)));
 
         } else {
             if (resp->status() == Ydb::StatusIds::SUCCESS) {
@@ -125,7 +131,7 @@ void TSession::TImpl::MakeImplAsync(TStreamProcessorPtr ptr,
                 NYdb::NIssue::TIssues opIssues;
                 NYdb::NIssue::IssuesFromMessage(resp->issues(), opIssues);
                 TStatus st(static_cast<EStatus>(resp->status()), std::move(opIssues));
-                args->Promise.SetValue(TCreateSessionResult(std::move(st), TSession()));
+                args->Promise.SetValue(TCreateSessionResult(std::move(st), TSession(args->Client)));
             }
         }
     });

tests/integration/basic_example_it/basic_example.cpp renamed to tests/integration/basic_example/basic_example.cpp

@@ -38,7 +38,7 @@ static NYdb::TType MakeOptionalType(NYdb::EPrimitiveType type) {
 }
 
 
-TEST(Integration, BasicExample) {
+TEST(BasicExample, BasicExample) {
     auto [driver, path] = GetRunArgs();
 
     NYdb::NTable::TTableClient client(driver);

tests/integration/basic_example_it/basic_example.h renamed to tests/integration/basic_example/basic_example.h

@@ -17,7 +17,6 @@ static std::string JoinPath(const std::string& basePath, const std::string& path
 }
 
 TRunArgs GetRunArgs() {
-    
     std::string database = std::getenv("YDB_DATABASE");
     std::string endpoint = std::getenv("YDB_ENDPOINT");
 

tests/integration/basic_example_it/basic_example_data.cpp renamed to tests/integration/basic_example/basic_example_data.cpp

@@ -5,8 +5,7 @@
 #include <util/string/cast.h>
 
 
-TEST(Integration, BulkUpsert) {
-
+TEST(BulkUpsert, BulkUpsert) {
     uint32_t correctSumApp = 0;
     uint32_t correctSumHost = 0;
     uint32_t correctRowCount = 0;