Moved commit "[C++ SDK] Remove session from session pool if stream session was closed by server side." from ydb repo

Author: dcherednik

src/client/impl/ydb_internal/kqp_session_common/kqp_session_common.cpp

@@ -38,6 +38,7 @@ TKqpSessionCommon::TKqpSessionCommon(
     , State_(S_STANDALONE)
     , TimeToTouch_(TInstant::Now())
     , TimeInPast_(TInstant::Now())
+    , CloseHandler_(nullptr)
     , NeedUpdateActiveCounter_(false)
 {}
 
@@ -114,7 +115,7 @@ void TKqpSessionCommon::ScheduleTimeToTouch(TDuration interval,
     if (updateTimeInPast) {
             TimeInPast_ = now;
     }
-    TimeToTouch_ = now + interval;
+    TimeToTouch_.store(now + interval, std::memory_order_relaxed);
 }
 
 void TKqpSessionCommon::ScheduleTimeToTouchFast(TDuration interval,
@@ -124,11 +125,11 @@ void TKqpSessionCommon::ScheduleTimeToTouchFast(TDuration interval,
     if (updateTimeInPast) {
         TimeInPast_ = now;
     }
-    TimeToTouch_ = now + interval;
+    TimeToTouch_.store(now + interval, std::memory_order_relaxed);
 }
 
 TInstant TKqpSessionCommon::GetTimeToTouchFast() const {
-    return TimeToTouch_;
+    return TimeToTouch_.load(std::memory_order_relaxed);
 }
 
 TInstant TKqpSessionCommon::GetTimeInPastFast() const {
@@ -144,6 +145,24 @@ TDuration TKqpSessionCommon::GetTimeInterval() const {
     return TimeInterval_;
 }
 
+void TKqpSessionCommon::UpdateServerCloseHandler(IServerCloseHandler* handler) {
+    CloseHandler_.store(handler);
+}
+
+void TKqpSessionCommon::CloseFromServer(std::weak_ptr<ISessionClient> client) noexcept {
+    auto strong = client.lock();
+    if (!strong) {
+        // Session closed on the server after stopping client - do nothing
+        // moreover pool maybe destoyed now
+        return;
+    }
+
+    IServerCloseHandler* h = CloseHandler_.load();
+    if (h) {
+        h->OnCloseSession(this, strong);
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 std::function<void(TKqpSessionCommon*)> TKqpSessionCommon::GetSmartDeleter(

src/client/impl/ydb_internal/kqp_session_common/kqp_session_common.h

@@ -13,6 +13,15 @@ namespace NYdb::inline V3 {
 ////////////////////////////////////////////////////////////////////////////////
 ui64 GetNodeIdFromSession(const std::string& sessionId);
 
+class TKqpSessionCommon;
+
+class IServerCloseHandler {
+public:
+    virtual ~IServerCloseHandler() = default;
+    // called when session should be closed by server signal
+    virtual void OnCloseSession(const TKqpSessionCommon*, std::shared_ptr<ISessionClient>) = 0;
+};
+
 class TKqpSessionCommon : public TEndpointObj {
 public:
     TKqpSessionCommon(const std::string& sessionId, const std::string& endpoint,
@@ -54,6 +63,12 @@ class TKqpSessionCommon : public TEndpointObj {
     static std::function<void(TKqpSessionCommon*)>
         GetSmartDeleter(std::shared_ptr<ISessionClient> client);
 
+    // Shoult be called under session pool lock
+    void UpdateServerCloseHandler(IServerCloseHandler*);
+
+    // Called asynchronously from grpc thread.
+    void CloseFromServer(std::weak_ptr<ISessionClient> client) noexcept;
+
 protected:
     TAdaptiveLock Lock_;
 
@@ -63,10 +78,14 @@ class TKqpSessionCommon : public TEndpointObj {
     const bool IsOwnedBySessionPool_;
 
     EState State_;
-    TInstant TimeToTouch_;
+    // This time is used during async close session handling which does not lock the session
+    // so we need to be able to read this value atomicaly
+    std::atomic<TInstant> TimeToTouch_;
     TInstant TimeInPast_;
     // Is used to implement progressive timeout for settler keep alive call
     TDuration TimeInterval_;
+
+    std::atomic<IServerCloseHandler*> CloseHandler_;
     // Indicate session was in active state, but state was changed
     // (need to decrement active session counter)
     // TODO: suboptimal because need lock for atomic change from interceptor

src/client/impl/ydb_internal/session_pool/session_pool.cpp

@@ -141,6 +141,7 @@ void TSessionPool::GetSession(std::unique_ptr<IGetSessionCtx> ctx)
         }
         if (!Sessions_.empty()) {
             auto it = std::prev(Sessions_.end());
+            it->second->UpdateServerCloseHandler(nullptr);
             sessionImpl = std::move(it->second);
             Sessions_.erase(it);
         }
@@ -206,6 +207,7 @@ bool TSessionPool::ReturnSession(TKqpSessionCommon* impl, bool active) {
             if (!active)
                 IncrementActiveCounterUnsafe();
         } else {
+            impl->UpdateServerCloseHandler(this);
             Sessions_.emplace(std::make_pair(
                 impl->GetTimeToTouchFast(),
                 impl));
@@ -242,6 +244,7 @@ void TSessionPool::Drain(std::function<bool(std::unique_ptr<TKqpSessionCommon>&&
     std::lock_guard guard(Mtx_);
     Closed_ = close;
     for (auto it = Sessions_.begin(); it != Sessions_.end();) {
+        it->second->UpdateServerCloseHandler(nullptr);
         const bool cont = cb(std::move(it->second));
         it = Sessions_.erase(it);
         if (!cont)
@@ -283,9 +286,11 @@ TPeriodicCb TSessionPool::CreatePeriodicTask(std::weak_ptr<ISessionClient> weakC
                             break;
 
                         if (deletePredicate(it->second.get(), sessions.size())) {
+                            it->second->UpdateServerCloseHandler(nullptr);
                             sessionsToDelete.emplace_back(std::move(it->second));
                             sessions.erase(it++);
                         } else if (cmd) {
+                            it->second->UpdateServerCloseHandler(nullptr);
                             sessionsToTouch.emplace_back(std::move(it->second));
                             sessions.erase(it++);
                         } else {
@@ -338,6 +343,32 @@ i64 TSessionPool::GetCurrentPoolSize() const {
     return Sessions_.size();
 }
 
+void TSessionPool::OnCloseSession(const TKqpSessionCommon* s, std::shared_ptr<ISessionClient> client) {
+    std::unique_ptr<TKqpSessionCommon> session;
+    {
+        std::lock_guard guard(Mtx_);
+        const auto timeToTouch = s->GetTimeToTouchFast();
+        const auto id = s->GetId();
+        auto it = Sessions_.find(timeToTouch);
+        // Sessions_ is multimap of sessions sorted by scheduled time to run periodic task
+        // Scan sessions with same scheduled time to find needed one. In most cases only one session here
+        while (it != Sessions_.end() && it->first == timeToTouch) {
+            if (id != it->second->GetId()) {
+                it++;
+                continue;
+            }
+            session = std::move(it->second);
+            Sessions_.erase(it);
+            break;
+        }
+    }
+
+    if (session) {
+        Y_ABORT_UNLESS(session->GetState() == TKqpSessionCommon::S_IDLE);
+        CloseAndDeleteSession(std::move(session), client);
+    }
+}
+
 void TSessionPool::SetStatCollector(NSdkStats::TStatCollector::TSessionPoolStatCollector statCollector) {
     ActiveSessionsCounter_.Set(statCollector.ActiveSessions);
     InPoolSessionsCounter_.Set(statCollector.InPoolSessions);

src/client/impl/ydb_internal/session_pool/session_pool.h

@@ -82,7 +82,7 @@ NThreading::TFuture<TResponse> InjectSessionStatusInterception(
     return promise.GetFuture();
 }
 
-class TSessionPool {
+class TSessionPool : public IServerCloseHandler {
 private:
     class TWaitersQueue {
     public:
@@ -125,6 +125,8 @@ class TSessionPool {
     void Drain(std::function<bool(std::unique_ptr<TKqpSessionCommon>&&)> cb, bool close);
     void SetStatCollector(NSdkStats::TStatCollector::TSessionPoolStatCollector collector);
 
+    void OnCloseSession(const TKqpSessionCommon*, std::shared_ptr<ISessionClient> client) override;
+
 private:
     void UpdateStats();
     static void ReplySessionToUser(TKqpSessionCommon* session, std::unique_ptr<IGetSessionCtx> ctx);

src/client/query/client.cpp

@@ -344,7 +344,7 @@ class TQueryClient::TImpl: public TClientImplCommon<TQueryClient::TImpl>, public
         const auto sessionId = resp->session_id();
         request.set_session_id(sessionId);
 
-        auto args = std::make_shared<TSession::TImpl::TAttachSessionArgs>(promise, sessionId, endpoint, client);
+        auto args = std::make_shared<TSession::TImpl::TAttachSessionArgs>(promise, sessionId, endpoint, client, client);
 
         // Do not pass client timeout here. Session must be alive
         TRpcRequestSettings rpcSettings;

src/client/query/impl/client_session.cpp

@@ -5,20 +5,106 @@
 #undef INCLUDE_YDB_INTERNAL_H
 
 #include <src/library/issue/yql_issue_message.h>
+#include <thread>
 
 namespace NYdb::inline V3::NQuery {
 
-TSession::TImpl::TImpl(TStreamProcessorPtr ptr, const std::string& sessionId, const std::string& endpoint)
+// Custom lock primitive to protect session from destroying
+// during async read execution.
+// The problem is TSession::TImpl holds grpc stream processor by IntrusivePtr
+// and this processor alredy refcounted by internal code.
+// That mean during TSession::TImpl dtor no gurantee to grpc procerrot will be destroyed.
+// StreamProcessor_->Cancel() doesn't help it just start async cancelation but we have no way
+// to wait cancelation has done.
+// So we need some way to protect access to row session impl pointer
+// from async reader (processor callback). We can't use shared/weak ptr here because TSessionImpl
+// stores as uniq ptr inside session pool and as shared ptr in the TSession
+// when user got session (see GetSmartDeleter related code).
+
+// Why just not std::mutex? - Requirement do not destroy a mutex while it is locked
+// makes it difficult to use here. Moreover we need to allow recursive lock.
+
+// Why recursive lock? - In happy path we destroy session from CloseFromServer call,
+// so the session dtor called from thread which already got the lock.
+
+// TODO: Proably we can add sync version of Cancel method in to grpc stream procesor to make sure
+// no more callback will be called.
+
+class TSafeTSessionImplHolder {
+    TSession::TImpl* Ptr;
+    std::atomic_uint32_t Semaphore;
+    std::atomic<std::thread::id> OwnerThread;
+public:
+    TSafeTSessionImplHolder(TSession::TImpl* p)
+        : Ptr(p)
+        , Semaphore(0)
+    {}
+
+    TSession::TImpl* TrySharedOwning() noexcept {
+        auto old = Semaphore.fetch_add(1); 
+        if (old == 0) {
+            OwnerThread.store(std::this_thread::get_id());
+            return Ptr;
+        } else {
+            return nullptr;
+        }
+    }
+
+    void Release() noexcept {
+        OwnerThread.store(std::thread::id());
+        Semaphore.store(0);
+    }
+
+    void WaitAndLock() noexcept {
+        if (OwnerThread.load() == std::this_thread::get_id()) {
+            return;
+        }
+
+        uint32_t cur = 0;
+        uint32_t newVal = 1;
+        while (!Semaphore.compare_exchange_weak(cur, newVal,
+            std::memory_order_release, std::memory_order_relaxed)) {
+                std::this_thread::yield();
+                cur = 0;
+        }
+    }
+};
+
+void TSession::TImpl::StartAsyncRead(TStreamProcessorPtr ptr, std::weak_ptr<ISessionClient> client,
+    std::shared_ptr<TSafeTSessionImplHolder> holder)
+{
+    auto resp = std::make_shared<Ydb::Query::SessionState>();
+    ptr->Read(resp.get(), [resp, ptr, client, holder](NYdbGrpc::TGrpcStatus grpcStatus) mutable {
+        switch (grpcStatus.GRpcStatusCode) {
+            case grpc::StatusCode::OK:
+                StartAsyncRead(ptr, client, holder);
+                break;
+            case grpc::StatusCode::OUT_OF_RANGE: {
+                auto impl = holder->TrySharedOwning();
+                if (impl) {
+                    impl->CloseFromServer(client);
+                    holder->Release();
+                }
+                break;
+            }
+        }
+    });
+}
+
+TSession::TImpl::TImpl(TStreamProcessorPtr ptr, const std::string& sessionId, const std::string& endpoint, std::weak_ptr<ISessionClient> client)
     : TKqpSessionCommon(sessionId, endpoint, true)
     , StreamProcessor_(ptr)
+    , SessionHolder(std::make_shared<TSafeTSessionImplHolder>(this))
 {
     MarkActive();
     SetNeedUpdateActiveCounter(true);
+    StartAsyncRead(StreamProcessor_, client, SessionHolder);
 }
 
 TSession::TImpl::~TImpl()
 {
     StreamProcessor_->Cancel();
+    SessionHolder->WaitAndLock();
 }
 
 void TSession::TImpl::MakeImplAsync(TStreamProcessorPtr ptr,
@@ -53,7 +139,7 @@ void TSession::TImpl::NewSmartShared(TStreamProcessorPtr ptr,
             std::move(st),
             TSession(
                 args->Client,
-                new TSession::TImpl(ptr, args->SessionId, args->Endpoint)
+                new TSession::TImpl(ptr, args->SessionId, args->Endpoint, args->SessionClient)
             )
         )
     );

src/client/query/impl/client_session.h

@@ -7,36 +7,44 @@
 
 namespace NYdb::inline V3::NQuery {
 
+class TSafeTSessionImplHolder;
+
 class TSession::TImpl : public TKqpSessionCommon {
 public:
     struct TAttachSessionArgs {
         TAttachSessionArgs(NThreading::TPromise<TCreateSessionResult> promise,
             std::string sessionId,
             std::string endpoint,
-            std::shared_ptr<TQueryClient::TImpl> client)
+            std::shared_ptr<TQueryClient::TImpl> client,
+            std::weak_ptr<ISessionClient> sessionClient)
             : Promise(promise)
             , SessionId(sessionId)
             , Endpoint(endpoint)
             , Client(client)
+            , SessionClient(sessionClient)
         { }
         NThreading::TPromise<TCreateSessionResult> Promise;
         std::string SessionId;
         std::string Endpoint;
         std::shared_ptr<TQueryClient::TImpl> Client;
+        std::weak_ptr<ISessionClient> SessionClient;
     };
 
     using TResponse = Ydb::Query::SessionState;
     using TStreamProcessorPtr = NYdbGrpc::IStreamRequestReadProcessor<TResponse>::TPtr;
-    TImpl(TStreamProcessorPtr ptr, const std::string& id, const std::string& endpoint);
+    TImpl(TStreamProcessorPtr ptr, const std::string& id, const std::string& endpoint, std::weak_ptr<ISessionClient> client);
     ~TImpl();
 
     static void MakeImplAsync(TStreamProcessorPtr processor, std::shared_ptr<TAttachSessionArgs> args);
 
 private:
     static void NewSmartShared(TStreamProcessorPtr ptr, std::shared_ptr<TAttachSessionArgs> args, NYdb::TStatus status);
 
+    static void StartAsyncRead(TStreamProcessorPtr ptr, std::weak_ptr<ISessionClient> client, std::shared_ptr<TSafeTSessionImplHolder> session);
+
 private:
     TStreamProcessorPtr StreamProcessor_;
+    std::shared_ptr<TSafeTSessionImplHolder> SessionHolder;
 };
 
 }