Merge 3c07741 into 5c3d5c3

Author: CyberROFL

ydb/core/tx/datashard/cdc_stream_scan.cpp

@@ -3,6 +3,7 @@
 #include "datashard_impl.h"
 
 #include <ydb/core/protos/datashard_config.pb.h>
+#include <ydb/core/protos/tx_datashard.pb.h>
 
 #include <util/generic/maybe.h>
 #include <util/string/builder.h>
@@ -17,6 +18,11 @@ using namespace NActors;
 using namespace NTable;
 using namespace NTabletFlatExecutor;
 
+void TCdcStreamScanManager::TStats::Serialize(NKikimrTxDataShard::TEvCdcStreamScanResponse_TStats& proto) const {
+    proto.SetRowsProcessed(RowsProcessed);
+    proto.SetBytesProcessed(BytesProcessed);
+}
+
 void TCdcStreamScanManager::Reset() {
     Scans.clear();
     TxIdToPathId.clear();
@@ -95,6 +101,7 @@ void TCdcStreamScanManager::Complete(const TPathId& streamPathId) {
         return;
     }
 
+    CompletedScans[streamPathId] = it->second.Stats;
     TxIdToPathId.erase(it->second.TxId);
     Scans.erase(it);
 }
@@ -104,6 +111,15 @@ void TCdcStreamScanManager::Complete(ui64 txId) {
     Complete(TxIdToPathId.at(txId));
 }
 
+bool TCdcStreamScanManager::IsCompleted(const TPathId& streamPathId) const {
+    return CompletedScans.contains(streamPathId);
+}
+
+const TCdcStreamScanManager::TStats& TCdcStreamScanManager::GetCompletedStats(const TPathId& streamPathId) const {
+    Y_ABORT_UNLESS(CompletedScans.contains(streamPathId));
+    return CompletedScans.at(streamPathId);
+}
+
 TCdcStreamScanManager::TScanInfo* TCdcStreamScanManager::Get(const TPathId& streamPathId) {
     return Scans.FindPtr(streamPathId);
 }
@@ -456,8 +472,7 @@ class TCdcStreamScan: public IActorCallback, public IScan {
         PathIdFromPathId(StreamPathId, response->Record.MutableStreamPathId());
         response->Record.SetStatus(status);
         response->Record.SetErrorDescription(error);
-        response->Record.MutableStats()->SetRowsProcessed(Stats.RowsProcessed);
-        response->Record.MutableStats()->SetBytesProcessed(Stats.BytesProcessed);
+        Stats.Serialize(*response->Record.MutableStats());
 
         Send(ReplyTo, std::move(response));
     }
@@ -572,11 +587,10 @@ class TDataShard::TTxCdcStreamScanRun: public TTransactionBase<TDataShard> {
     TEvDataShard::TEvCdcStreamScanRequest::TPtr Request;
     THolder<IEventHandle> Response; // response to sender or forward to scanner
 
-    THolder<IEventHandle> MakeResponse(const TActorContext& ctx,
-            NKikimrTxDataShard::TEvCdcStreamScanResponse::EStatus status, const TString& error = {}) const
-    {
+    template <typename... Args>
+    THolder<IEventHandle> MakeResponse(const TActorContext& ctx, Args&&... args) const {
         return MakeHolder<IEventHandle>(Request->Sender, ctx.SelfID, new TEvDataShard::TEvCdcStreamScanResponse(
-            Request->Get()->Record, Self->TabletID(), status, error
+            Request->Get()->Record, Self->TabletID(), std::forward<Args>(args)...
         ));
     }
 
@@ -643,6 +657,11 @@ class TDataShard::TTxCdcStreamScanRun: public TTransactionBase<TDataShard> {
             } else if (info->ScanId) {
                 return true; // nop, scan actor will report state when it starts
             }
+        } else if (Self->CdcStreamScanManager.IsCompleted(streamPathId)) {
+            Response = MakeResponse(ctx, NKikimrTxDataShard::TEvCdcStreamScanResponse::DONE);
+            Self->CdcStreamScanManager.GetCompletedStats(streamPathId).Serialize(
+                *Response->Get<TEvDataShard::TEvCdcStreamScanResponse>()->Record.MutableStats());
+            return true;
         } else if (Self->CdcStreamScanManager.Size()) {
             Response = MakeResponse(ctx, NKikimrTxDataShard::TEvCdcStreamScanResponse::OVERLOADED);
             return true;

ydb/core/tx/datashard/cdc_stream_scan.h

@@ -9,13 +9,19 @@
 #include <util/generic/hash.h>
 #include <util/generic/maybe.h>
 
+namespace NKikimrTxDataShard {
+    class TEvCdcStreamScanResponse_TStats;
+}
+
 namespace NKikimr::NDataShard {
 
 class TCdcStreamScanManager {
 public:
     struct TStats {
         ui64 RowsProcessed = 0;
         ui64 BytesProcessed = 0;
+
+        void Serialize(NKikimrTxDataShard::TEvCdcStreamScanResponse_TStats& proto) const;
     };
 
 private:
@@ -39,6 +45,8 @@ class TCdcStreamScanManager {
 
     void Complete(const TPathId& streamPathId);
     void Complete(ui64 txId);
+    bool IsCompleted(const TPathId& streamPathId) const;
+    const TStats& GetCompletedStats(const TPathId& streamPathId) const;
 
     TScanInfo* Get(const TPathId& streamPathId);
     const TScanInfo* Get(const TPathId& streamPathId) const;
@@ -57,6 +65,7 @@ class TCdcStreamScanManager {
 
 private:
     THashMap<TPathId, TScanInfo> Scans;
+    THashMap<TPathId, TStats> CompletedScans;
     THashMap<ui64, TPathId> TxIdToPathId;
 };
 

ydb/core/tx/datashard/datashard_ut_change_exchange.cpp

@@ -2761,6 +2761,49 @@ Y_UNIT_TEST_SUITE(Cdc) {
         });
     }
 
+    Y_UNIT_TEST(InitialScanRacyCompleteAndRequest) {
+        TPortManager portManager;
+        TServer::TPtr server = new TServer(TServerSettings(portManager.GetPort(2134), {}, DefaultPQConfig())
+            .SetUseRealThreads(false)
+            .SetDomainName("Root")
+            .SetEnableChangefeedInitialScan(true)
+        );
+
+        auto& runtime = *server->GetRuntime();
+        const auto edgeActor = runtime.AllocateEdgeActor();
+
+        SetupLogging(runtime);
+        InitRoot(server, edgeActor);
+        CreateShardedTable(server, edgeActor, "/Root", "Table", SimpleTable());
+
+        std::unique_ptr<IEventHandle> doneResponse;
+        auto blockDone = runtime.AddObserver<TEvDataShard::TEvCdcStreamScanResponse>(
+            [&](TEvDataShard::TEvCdcStreamScanResponse::TPtr& ev) {
+                if (ev->Get()->Record.GetStatus() == NKikimrTxDataShard::TEvCdcStreamScanResponse::DONE) {
+                    doneResponse.reset(ev.Release());
+                }
+            }
+        );
+
+        WaitTxNotification(server, edgeActor, AsyncAlterAddStream(server, "/Root", "Table",
+            WithInitialScan(Updates(NKikimrSchemeOp::ECdcStreamFormatJson))));
+        WaitFor(runtime, [&]{ return bool(doneResponse); }, "doneResponse");
+        blockDone.Remove();
+
+        bool done = false;
+        auto waitDone = runtime.AddObserver<TEvDataShard::TEvCdcStreamScanResponse>(
+            [&](TEvDataShard::TEvCdcStreamScanResponse::TPtr& ev) {
+                if (ev->Get()->Record.GetStatus() == NKikimrTxDataShard::TEvCdcStreamScanResponse::DONE) {
+                    done = true;
+                }
+            }
+        );
+
+        const auto& record = doneResponse->Get<TEvDataShard::TEvCdcStreamScanResponse>()->Record;
+        RebootTablet(runtime, record.GetTablePathId().GetOwnerId(), edgeActor);
+        WaitFor(runtime, [&]{ return done; }, "done");
+    }
+
     Y_UNIT_TEST(InitialScanUpdatedRows) {
         TPortManager portManager;
         TServer::TPtr server = new TServer(TServerSettings(portManager.GetPort(2134), {}, DefaultPQConfig())