Memory leak in Topic SDK (#6249)

Alek5andr-Kotov · web-flow · commit be2302d1f095 · 2024-07-10T11:35:02.000+03:00
diff --git a/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.h b/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.h
@@ -143,7 +143,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 TString& message);
@@ -209,6 +209,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;
     }
@@ -306,6 +308,8 @@ class TDataDecompressionInfo : public std::enable_shared_from_this<TDataDecompre
             return EstimatedDecompressedSize;
         }
 
+        void ClearParent();
+
     private:
         TDataDecompressionInfo::TPtr Parent;
         TIntrusivePtr<TPartitionStreamImpl<UseMigrationProtocol>> PartitionStream;
@@ -1098,6 +1102,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);
 
@@ -1287,6 +1296,8 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext<TSingleClusterRe
         {
         }
 
+        void OnDestroyReadSession();
+
         TDataDecompressionInfoPtr<UseMigrationProtocol> BatchInfo;
         TIntrusivePtr<TPartitionStreamImpl<UseMigrationProtocol>> PartitionStream;
     };
diff --git a/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.ipp b/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.ipp
@@ -216,6 +216,15 @@ void TRawPartitionStreamEventQueue<UseMigrationProtocol>::DeleteNotReadyTail(TDe
     swap(ready, NotReady);
 }
 
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+// TDecompressionQueueItem
+
+template <bool UseMigrationProtocol>
+void TSingleClusterReadSessionImpl<UseMigrationProtocol>::TDecompressionQueueItem::OnDestroyReadSession()
+{
+    BatchInfo->OnDestroyReadSession();
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 // TSingleClusterReadSessionImpl
 
@@ -224,6 +233,10 @@ TSingleClusterReadSessionImpl<UseMigrationProtocol>::~TSingleClusterReadSessionI
     for (auto&& [_, partitionStream] : PartitionStreams) {
         partitionStream->ClearQueue();
     }
+
+    for (auto& e : DecompressionQueue) {
+        e.OnDestroyReadSession();
+    }
 }
 
 
@@ -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);
@@ -2524,6 +2538,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,
                                                              TVector<typename TADataReceivedEvent<UseMigrationProtocol>::TMessage>& messages,
@@ -2673,19 +2695,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);
 
@@ -2696,7 +2722,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
                     ) {
@@ -2706,7 +2732,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
                     ) {
@@ -2718,32 +2744,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
 
@@ -2781,7 +2813,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));
 }
 
diff --git a/ydb/services/persqueue_v1/persqueue_ut.cpp b/ydb/services/persqueue_v1/persqueue_ut.cpp
@@ -3376,6 +3376,54 @@ TPersQueueV1TestServer server{{.CheckACL=true, .NodeCount=1}};
         DumpCounters("End");
     }
 
+    Y_UNIT_TEST(NoDecompressionMemoryLeaks) {
+
+        NPersQueue::TTestServer server;
+        server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY, NKikimrServices::PQ_READ_PROXY});
+        server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1);
+
+        auto driver = server.AnnoyingClient->GetDriver();
+        auto decompressor = CreateThreadPoolExecutorWrapper(2);
+
+        NYdb::NPersQueue::TReadSessionSettings settings;
+        settings.ConsumerName("shared/user").AppendTopics(SHORT_TOPIC_NAME).ReadOriginal({"dc1"});
+        settings.DecompressionExecutor(decompressor);
+        settings.MaxMemoryUsageBytes(5_MB);
+        settings.Decompress(true);
+
+        auto reader = CreateReader(*driver, settings);
+
+        //
+        // there should be 1 TCreatePartitionStreamEvent events in the queue
+        //
+        {
+            auto msg = reader->GetEvent(true, 1);
+            UNIT_ASSERT(msg);
+
+            Cerr << ">>>> message: " << NYdb::NPersQueue::DebugString(*msg) << Endl;
+
+            auto ev = std::get_if<NYdb::NPersQueue::TReadSessionEvent::TCreatePartitionStreamEvent>(&*msg);
+            UNIT_ASSERT(ev);
+
+            ev->Confirm();
+        }
+
+        for (ui32 i = 0; i < 10; ++i) {
+            auto writer = CreateSimpleWriter(*driver, SHORT_TOPIC_NAME, TStringBuilder() << "source" << i);
+
+            std::string message(1_MB - 1_KB, 'x');
+
+            bool res = writer->Write(message, 1);
+            UNIT_ASSERT(res);
+
+            res = writer->Close(TDuration::Seconds(10));
+            UNIT_ASSERT(res);
+        }
+
+        decompressor->StartFuncs({0, 1, 2});
+        Sleep(TDuration::Seconds(1));
+    }
+
     enum WhenTheTopicIsDeletedMode {
         AFTER_WRITES,
         AFTER_START_TASKS,
