Moved commit "Memory leak in Topic SDK" from ydb repo

Author: Alek5andr-Kotov

src/client/topic/impl/read_session_impl.h

@@ -140,7 +140,7 @@ class TDeferredActions {
     }
 
     void DeferReadFromProcessor(const typename IProcessor<UseMigrationProtocol>::TPtr& processor, TServerMessage<UseMigrationProtocol>* dst, typename IProcessor<UseMigrationProtocol>::TReadCallback callback);
-    void DeferStartExecutorTask(const typename IAExecutor<UseMigrationProtocol>::TPtr& executor, typename IAExecutor<UseMigrationProtocol>::TFunction task);
+    void DeferStartExecutorTask(const typename IAExecutor<UseMigrationProtocol>::TPtr& executor, typename IAExecutor<UseMigrationProtocol>::TFunction&& task);
     void DeferAbortSession(TCallbackContextPtr<UseMigrationProtocol> cbContext, TASessionClosedEvent<UseMigrationProtocol>&& closeEvent);
     void DeferAbortSession(TCallbackContextPtr<UseMigrationProtocol> cbContext, EStatus statusCode, NYql::TIssues&& issues);
     void DeferAbortSession(TCallbackContextPtr<UseMigrationProtocol> cbContext, EStatus statusCode, const std::string& message);
@@ -206,6 +206,8 @@ class TDataDecompressionInfo : public std::enable_shared_from_this<TDataDecompre
     void PlanDecompressionTasks(double averageCompressionRatio,
                                 TIntrusivePtr<TPartitionStreamImpl<UseMigrationProtocol>> partitionStream);
 
+    void OnDestroyReadSession();
+
     bool IsReady() const {
         return SourceDataNotProcessed == 0;
     }
@@ -303,6 +305,8 @@ class TDataDecompressionInfo : public std::enable_shared_from_this<TDataDecompre
             return EstimatedDecompressedSize;
         }
 
+        void ClearParent();
+
     private:
         TDataDecompressionInfo::TPtr Parent;
         TIntrusivePtr<TPartitionStreamImpl<UseMigrationProtocol>> PartitionStream;
@@ -1096,6 +1100,11 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext<TSingleClusterRe
 
     void OnCreateNewDecompressionTask();
     void OnDecompressionInfoDestroy(i64 compressedSize, i64 decompressedSize, i64 messagesCount, i64 serverBytesSize);
+    void OnDecompressionInfoDestroyImpl(i64 compressedSize,
+                                        i64 decompressedSize,
+                                        i64 messagesCount,
+                                        i64 serverBytesSize,
+                                        TDeferredActions<UseMigrationProtocol>& deferred);
 
     void OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount, i64 serverBytesSize = 0);
 
@@ -1285,6 +1294,8 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext<TSingleClusterRe
         {
         }
 
+        void OnDestroyReadSession();
+
         TDataDecompressionInfoPtr<UseMigrationProtocol> BatchInfo;
         TIntrusivePtr<TPartitionStreamImpl<UseMigrationProtocol>> PartitionStream;
     };

src/client/topic/impl/read_session_impl.ipp

@@ -218,6 +218,15 @@ void TRawPartitionStreamEventQueue<UseMigrationProtocol>::DeleteNotReadyTail(TDe
     swap(ready, NotReady);
 }
 
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+// TDecompressionQueueItem
+
+template <bool UseMigrationProtocol>
+void TSingleClusterReadSessionImpl<UseMigrationProtocol>::TDecompressionQueueItem::OnDestroyReadSession()
+{
+    BatchInfo->OnDestroyReadSession();
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 // TSingleClusterReadSessionImpl
 
@@ -226,6 +235,10 @@ TSingleClusterReadSessionImpl<UseMigrationProtocol>::~TSingleClusterReadSessionI
     for (auto&& [_, partitionStream] : PartitionStreams) {
         partitionStream->ClearQueue();
     }
+
+    for (auto& e : DecompressionQueue) {
+        e.OnDestroyReadSession();
+    }
 }
 
 template<bool UseMigrationProtocol>
@@ -1565,6 +1578,7 @@ void TSingleClusterReadSessionImpl<UseMigrationProtocol>::OnDecompressionInfoDes
 
 template<bool UseMigrationProtocol>
 void TSingleClusterReadSessionImpl<UseMigrationProtocol>::OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount, i64 serverBytesSize) {
+
     TDeferredActions<UseMigrationProtocol> deferred;
 
     Y_ABORT_UNLESS(DecompressionTasksInflight > 0);
@@ -2512,6 +2526,14 @@ void TDataDecompressionInfo<UseMigrationProtocol>::PlanDecompressionTasks(double
     }
 }
 
+template <bool UseMigrationProtocol>
+void TDataDecompressionInfo<UseMigrationProtocol>::OnDestroyReadSession()
+{
+    for (auto& task : Tasks) {
+        task.ClearParent();
+    }
+}
+
 template<bool UseMigrationProtocol>
 void TDataDecompressionEvent<UseMigrationProtocol>::TakeData(TIntrusivePtr<TPartitionStreamImpl<UseMigrationProtocol>> partitionStream,
                                                              std::vector<typename TADataReceivedEvent<UseMigrationProtocol>::TMessage>& messages,
@@ -2661,19 +2683,23 @@ TDataDecompressionInfo<UseMigrationProtocol>::TDecompressionTask::TDecompression
 
 template<bool UseMigrationProtocol>
 void TDataDecompressionInfo<UseMigrationProtocol>::TDecompressionTask::operator()() {
+    auto parent = Parent;
+    if (!parent) {
+        return;
+    }
     i64 minOffset = Max<i64>();
     i64 maxOffset = 0;
-    const i64 partition_id = [this](){
+    const i64 partition_id = [parent](){
         if constexpr (UseMigrationProtocol) {
-            return Parent->ServerMessage.partition();
+            return parent->ServerMessage.partition();
         } else {
-            return Parent->ServerMessage.partition_session_id();
+            return parent->ServerMessage.partition_session_id();
         }
     }();
     i64 dataProcessed = 0;
     size_t messagesProcessed = 0;
     for (const TMessageRange& messages : Messages) {
-        auto& batch = *Parent->ServerMessage.mutable_batches(messages.Batch);
+        auto& batch = *parent->ServerMessage.mutable_batches(messages.Batch);
         for (size_t i = messages.MessageRange.first; i < messages.MessageRange.second; ++i) {
             auto& data = *batch.mutable_message_data(i);
 
@@ -2684,7 +2710,7 @@ void TDataDecompressionInfo<UseMigrationProtocol>::TDecompressionTask::operator(
 
             try {
                 if constexpr (UseMigrationProtocol) {
-                    if (Parent->DoDecompress
+                    if (parent->DoDecompress
                         && data.codec() != Ydb::PersQueue::V1::CODEC_RAW
                         && data.codec() != Ydb::PersQueue::V1::CODEC_UNSPECIFIED
                     ) {
@@ -2694,7 +2720,7 @@ void TDataDecompressionInfo<UseMigrationProtocol>::TDecompressionTask::operator(
                         data.set_codec(Ydb::PersQueue::V1::CODEC_RAW);
                     }
                 } else {
-                    if (Parent->DoDecompress
+                    if (parent->DoDecompress
                         && static_cast<Ydb::Topic::Codec>(batch.codec()) != Ydb::Topic::CODEC_RAW
                         && static_cast<Ydb::Topic::Codec>(batch.codec()) != Ydb::Topic::CODEC_UNSPECIFIED
                     ) {
@@ -2706,32 +2732,38 @@ void TDataDecompressionInfo<UseMigrationProtocol>::TDecompressionTask::operator(
 
                 DecompressedSize += data.data().size();
             } catch (...) {
-                Parent->PutDecompressionError(std::current_exception(), messages.Batch, i);
+                parent->PutDecompressionError(std::current_exception(), messages.Batch, i);
                 data.clear_data(); // Free memory, because we don't count it.
 
-                if (auto session = Parent->CbContext->LockShared()) {
+                if (auto session = parent->CbContext->LockShared()) {
                     session->GetLog() << TLOG_INFO << "Error decompressing data: " << CurrentExceptionMessage();
                 }
             }
         }
     }
-    if (auto session = Parent->CbContext->LockShared()) {
+    if (auto session = parent->CbContext->LockShared()) {
         LOG_LAZY(session->GetLog(), TLOG_DEBUG, TStringBuilder() << "Decompression task done. Partition/PartitionSessionId: "
                                                                  << partition_id << " (" << minOffset << "-"
                                                                  << maxOffset << ")");
     }
     Y_ASSERT(dataProcessed == SourceDataSize);
 
-    Parent->OnDataDecompressed(SourceDataSize, EstimatedDecompressedSize, DecompressedSize, messagesProcessed);
+    parent->OnDataDecompressed(SourceDataSize, EstimatedDecompressedSize, DecompressedSize, messagesProcessed);
 
-    Parent->SourceDataNotProcessed -= dataProcessed;
+    parent->SourceDataNotProcessed -= dataProcessed;
     Ready->Ready = true;
 
-    if (auto session = Parent->CbContext->LockShared()) {
+    if (auto session = parent->CbContext->LockShared()) {
         session->GetEventsQueue()->SignalReadyEvents(PartitionStream);
     }
 }
 
+template<bool UseMigrationProtocol>
+void TDataDecompressionInfo<UseMigrationProtocol>::TDecompressionTask::ClearParent()
+{
+    Parent = nullptr;
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 // TUserRetrievedEventsInfoAccumulator
 
@@ -2769,7 +2801,7 @@ void TDeferredActions<UseMigrationProtocol>::DeferReadFromProcessor(const typena
 }
 
 template<bool UseMigrationProtocol>
-void TDeferredActions<UseMigrationProtocol>::DeferStartExecutorTask(const typename IAExecutor<UseMigrationProtocol>::TPtr& executor, typename IAExecutor<UseMigrationProtocol>::TFunction task) {
+void TDeferredActions<UseMigrationProtocol>::DeferStartExecutorTask(const typename IAExecutor<UseMigrationProtocol>::TPtr& executor, typename IAExecutor<UseMigrationProtocol>::TFunction&& task) {
     ExecutorsTasks.emplace_back(executor, std::move(task));
 }