add implementation for common cases in local partition reader (#6692)

Author: Enjection

ydb/core/backup/impl/local_partition_reader.cpp

@@ -11,7 +11,13 @@
 using namespace NActors;
 using namespace NKikimr::NReplication::NService;
 
+namespace {
+
 constexpr static char OFFLOAD_ACTOR_CLIENT_ID[] = "__OFFLOAD_ACTOR__";
+constexpr static ui64 READ_TIMEOUT_MS = 1000;
+constexpr static ui64 READ_LIMIT_BYTES = 1_MB;
+
+} // anonymous namespace
 
 namespace NKikimr::NBackup::NImpl {
 
@@ -50,6 +56,10 @@ class TLocalPartitionReader
         return request;
     }
 
+    void HandleInit() {
+        Send(PQTablet, CreateGetOffsetRequest().Release());
+    }
+
     void HandleInit(TEvWorker::TEvHandshake::TPtr& ev) {
         Worker = ev->Sender;
         LOG_D("Handshake"
@@ -62,8 +72,8 @@ class TLocalPartitionReader
         LOG_D("Handle " << ev->Get()->ToString());
         auto& record = ev->Get()->Record;
         if (record.GetErrorCode() == NPersQueue::NErrorCode::INITIALIZING) {
-            // TODO reschedule
-            Y_ABORT("Unimplemented!");
+            Schedule(TDuration::Seconds(1), new NActors::TEvents::TEvWakeup);
+            return;
         }
         Y_VERIFY_S(record.GetErrorCode() == NPersQueue::NErrorCode::OK, "Unimplemented!");
         Y_VERIFY_S(record.HasPartitionResponse() && record.GetPartitionResponse().HasCmdGetClientOffsetResult(), "Unimplemented!");
@@ -85,8 +95,8 @@ class TLocalPartitionReader
         auto& read = *req.MutableCmdRead();
         read.SetOffset(Offset);
         read.SetClientId(OFFLOAD_ACTOR_CLIENT_ID);
-        read.SetTimeoutMs(0);
-        read.SetBytes(1_MB);
+        read.SetTimeoutMs(READ_TIMEOUT_MS);
+        read.SetBytes(READ_LIMIT_BYTES);
 
         return request;
     }
@@ -114,6 +124,12 @@ class TLocalPartitionReader
 
         const auto& readResult = record.GetPartitionResponse().GetCmdReadResult();
 
+        if (!readResult.ResultSize()) {
+            Y_ABORT_UNLESS(PQTablet);
+            Send(PQTablet, CreateReadRequest().Release());
+            return;
+        }
+
         auto gotOffset = Offset;
         TVector<TEvWorker::TEvData::TRecord> records(::Reserve(readResult.ResultSize()));
 
@@ -147,6 +163,7 @@ class TLocalPartitionReader
         switch (ev->GetTypeRewrite()) {
             hFunc(TEvWorker::TEvHandshake, HandleInit);
             hFunc(TEvPersQueue::TEvResponse, HandleInit);
+            sFunc(TEvents::TEvWakeup, HandleInit);
             sFunc(TEvents::TEvPoison, PassAway);
         default:
             Y_VERIFY_S(false, "Unhandled event type: " << ev->GetTypeRewrite()

ydb/core/backup/impl/local_partition_reader_ut.cpp

@@ -15,27 +15,90 @@ using namespace NReplication::NService;
 using namespace NPersQueue;
 
 Y_UNIT_TEST_SUITE(LocalPartitionReader) {
+    constexpr static ui64 INITIAL_OFFSET = 3;
+    constexpr static ui64 PARTITION = 7;
+    constexpr static const char* PARTITION_STR = "7";
+    constexpr static const char* OFFLOAD_ACTOR_NAME = "__OFFLOAD_ACTOR__";
+
+    void GrabInitialPQRequest(TTestActorRuntime& runtime) {
+        TAutoPtr<IEventHandle> handle;
+        auto getOffset = runtime.GrabEdgeEventRethrow<TEvPersQueue::TEvRequest>(handle);
+        UNIT_ASSERT(getOffset->Record.HasPartitionRequest());
+        UNIT_ASSERT_VALUES_EQUAL(getOffset->Record.GetPartitionRequest().GetPartition(), PARTITION);
+        UNIT_ASSERT(getOffset->Record.GetPartitionRequest().HasCmdGetClientOffset());
+        auto getOffsetCmd = getOffset->Record.GetPartitionRequest().GetCmdGetClientOffset();
+        UNIT_ASSERT_VALUES_EQUAL(getOffsetCmd.GetClientId(), OFFLOAD_ACTOR_NAME);
+    }
+
+    void GrabReadPQRequest(TTestActorRuntime& runtime, ui32 expectedOffset) {
+        TAutoPtr<IEventHandle> handle;
+        auto read = runtime.GrabEdgeEventRethrow<TEvPersQueue::TEvRequest>(handle);
+        UNIT_ASSERT(read->Record.HasPartitionRequest());
+        UNIT_ASSERT_VALUES_EQUAL(read->Record.GetPartitionRequest().GetPartition(), PARTITION);
+        UNIT_ASSERT(read->Record.GetPartitionRequest().HasCmdRead());
+        auto readCmd = read->Record.GetPartitionRequest().GetCmdRead();
+        UNIT_ASSERT_VALUES_EQUAL(readCmd.GetClientId(), OFFLOAD_ACTOR_NAME);
+        UNIT_ASSERT_VALUES_EQUAL(readCmd.GetOffset(), expectedOffset);
+    }
+
+    void GrabDataEvent(TTestActorRuntime& runtime, ui32 dataPatternCookie) {
+        TAutoPtr<IEventHandle> handle;
+        auto data = runtime.GrabEdgeEventRethrow<TEvWorker::TEvData>(handle);
+        UNIT_ASSERT_VALUES_EQUAL(data->Source, PARTITION_STR);
+        UNIT_ASSERT_VALUES_EQUAL(data->Records.size(), 2);
+        UNIT_ASSERT_VALUES_EQUAL(data->Records[0].Offset, INITIAL_OFFSET + dataPatternCookie * 2);
+        UNIT_ASSERT_VALUES_EQUAL(data->Records[0].Data, Sprintf("1-%d", dataPatternCookie));
+        UNIT_ASSERT_VALUES_EQUAL(data->Records[1].Offset, INITIAL_OFFSET + dataPatternCookie * 2 + 1);
+        UNIT_ASSERT_VALUES_EQUAL(data->Records[1].Data, Sprintf("2-%d", dataPatternCookie));
+    }
+
+    TEvPersQueue::TEvResponse* GenerateData(ui32 dataPatternCookie) {
+        auto* readResponse = new TEvPersQueue::TEvResponse;
+        readResponse->Record.SetErrorCode(NPersQueue::NErrorCode::OK);
+        auto& cmdReadResult = *readResponse->Record.MutablePartitionResponse()->MutableCmdReadResult();
+        auto& readResult1 = *cmdReadResult.AddResult();
+        NKikimrPQClient::TDataChunk msg1;
+        msg1.SetData(Sprintf("1-%d", dataPatternCookie));
+        TString msg1Str;
+        UNIT_ASSERT(msg1.SerializeToString(&msg1Str));
+        readResult1.SetOffset(INITIAL_OFFSET + dataPatternCookie * 2);
+        readResult1.SetData(msg1Str);
+        auto& readResult2 = *cmdReadResult.AddResult();
+        NKikimrPQClient::TDataChunk msg2;
+        msg2.SetData(Sprintf("2-%d", dataPatternCookie));
+        TString msg2Str;
+        UNIT_ASSERT(msg2.SerializeToString(&msg2Str));
+        readResult2.SetOffset(INITIAL_OFFSET + dataPatternCookie * 2 + 1);
+        readResult2.SetData(msg2Str);
+
+        return readResponse;
+    }
+
+    TEvPersQueue::TEvResponse* GenerateEmptyData() {
+        auto* readResponse = new TEvPersQueue::TEvResponse;
+        readResponse->Record.SetErrorCode(NPersQueue::NErrorCode::OK);
+        readResponse->Record.MutablePartitionResponse()->MutableCmdReadResult();
+
+        return readResponse;
+    }
+
     Y_UNIT_TEST(Simple) {
         TTestActorRuntime runtime;
         runtime.Initialize(NKikimr::TAppPrepare().Unwrap());
 
         TActorId pqtablet = runtime.AllocateEdgeActor();
-        TActorId reader = runtime.Register(CreateLocalPartitionReader(pqtablet, 7));
+        TActorId reader = runtime.Register(CreateLocalPartitionReader(pqtablet, PARTITION));
+        runtime.EnableScheduleForActor(reader);
         TActorId worker = runtime.AllocateEdgeActor();
         TAutoPtr<IEventHandle> handle;
 
         runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvHandshake));
 
-        auto getOffset = runtime.GrabEdgeEventRethrow<TEvPersQueue::TEvRequest>(handle);
-        UNIT_ASSERT(getOffset->Record.HasPartitionRequest());
-        UNIT_ASSERT_VALUES_EQUAL(getOffset->Record.GetPartitionRequest().GetPartition(), 7);
-        UNIT_ASSERT(getOffset->Record.GetPartitionRequest().HasCmdGetClientOffset());
-        auto getOffsetCmd = getOffset->Record.GetPartitionRequest().GetCmdGetClientOffset();
-        UNIT_ASSERT_VALUES_EQUAL(getOffsetCmd.GetClientId(), "__OFFLOAD_ACTOR__"); // hardcoded for now
+        GrabInitialPQRequest(runtime);
 
         auto* getOffsetResponse = new TEvPersQueue::TEvResponse;
         getOffsetResponse->Record.SetErrorCode(NPersQueue::NErrorCode::OK);
-        getOffsetResponse->Record.MutablePartitionResponse()->MutableCmdGetClientOffsetResult()->SetOffset(3);
+        getOffsetResponse->Record.MutablePartitionResponse()->MutableCmdGetClientOffsetResult()->SetOffset(INITIAL_OFFSET);
         runtime.Send(new IEventHandle(reader, pqtablet, getOffsetResponse));
 
         auto handshake = runtime.GrabEdgeEventRethrow<TEvWorker::TEvHandshake>(handle);
@@ -44,40 +107,113 @@ Y_UNIT_TEST_SUITE(LocalPartitionReader) {
         for (auto i = 0; i < 3; ++i) {
             runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvPoll));
 
-            auto read = runtime.GrabEdgeEventRethrow<TEvPersQueue::TEvRequest>(handle);
-            UNIT_ASSERT(read->Record.HasPartitionRequest());
-            UNIT_ASSERT_VALUES_EQUAL(read->Record.GetPartitionRequest().GetPartition(), 7);
-            UNIT_ASSERT(read->Record.GetPartitionRequest().HasCmdRead());
-            auto readCmd = read->Record.GetPartitionRequest().GetCmdRead();
-            UNIT_ASSERT_VALUES_EQUAL(readCmd.GetClientId(), "__OFFLOAD_ACTOR__"); // hardcoded for now
-            UNIT_ASSERT_VALUES_EQUAL(readCmd.GetOffset(), 3 + i * 2);
-
-            auto* readResponse = new TEvPersQueue::TEvResponse;
-            readResponse->Record.SetErrorCode(NPersQueue::NErrorCode::OK);
-            auto& cmdReadResult = *readResponse->Record.MutablePartitionResponse()->MutableCmdReadResult();
-            auto& readResult1 = *cmdReadResult.AddResult();
-            NKikimrPQClient::TDataChunk msg1;
-            msg1.SetData(Sprintf("1-%d", i));
-            TString msg1Str;
-            UNIT_ASSERT(msg1.SerializeToString(&msg1Str));
-            readResult1.SetOffset(3 + i * 2);
-            readResult1.SetData(msg1Str);
-            auto& readResult2 = *cmdReadResult.AddResult();
-            NKikimrPQClient::TDataChunk msg2;
-            msg2.SetData(Sprintf("2-%d", i));
-            TString msg2Str;
-            UNIT_ASSERT(msg2.SerializeToString(&msg2Str));
-            readResult2.SetOffset(3 + i * 2 + 1);
-            readResult2.SetData(msg2Str);
-            runtime.Send(new IEventHandle(reader, pqtablet, readResponse));
-
-            auto data = runtime.GrabEdgeEventRethrow<TEvWorker::TEvData>(handle);
-            UNIT_ASSERT_VALUES_EQUAL(data->Source, "7");
-            UNIT_ASSERT_VALUES_EQUAL(data->Records.size(), 2);
-            UNIT_ASSERT_VALUES_EQUAL(data->Records[0].Offset, 3 + i * 2);
-            UNIT_ASSERT_VALUES_EQUAL(data->Records[0].Data, Sprintf("1-%d", i));
-            UNIT_ASSERT_VALUES_EQUAL(data->Records[1].Offset, 3 + i * 2 + 1);
-            UNIT_ASSERT_VALUES_EQUAL(data->Records[1].Data, Sprintf("2-%d", i));
+            GrabReadPQRequest(runtime, INITIAL_OFFSET + i * 2);
+
+            runtime.Send(new IEventHandle(reader, pqtablet, GenerateData(i)));
+
+            GrabDataEvent(runtime, i);
+
+            // TODO check commit
+        }
+    }
+
+    Y_UNIT_TEST(Booting) {
+        TTestActorRuntime runtime;
+        runtime.Initialize(NKikimr::TAppPrepare().Unwrap());
+
+        TActorId pqtablet = runtime.AllocateEdgeActor();
+        TActorId reader = runtime.Register(CreateLocalPartitionReader(pqtablet, PARTITION));
+        runtime.EnableScheduleForActor(reader);
+        TActorId worker = runtime.AllocateEdgeActor();
+        TAutoPtr<IEventHandle> handle;
+
+        runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvHandshake));
+
+        for (int i = 0; i < 5; ++i) {
+            GrabInitialPQRequest(runtime);
+
+            auto* getOffsetResponse = new TEvPersQueue::TEvResponse;
+            getOffsetResponse->Record.SetErrorCode(NPersQueue::NErrorCode::INITIALIZING);
+            runtime.Send(new IEventHandle(reader, pqtablet, getOffsetResponse));
+        }
+
+        GrabInitialPQRequest(runtime);
+
+        auto* getOffsetResponse = new TEvPersQueue::TEvResponse;
+        getOffsetResponse->Record.SetErrorCode(NPersQueue::NErrorCode::OK);
+        getOffsetResponse->Record.MutablePartitionResponse()->MutableCmdGetClientOffsetResult()->SetOffset(INITIAL_OFFSET);
+        runtime.Send(new IEventHandle(reader, pqtablet, getOffsetResponse));
+
+        auto handshake = runtime.GrabEdgeEventRethrow<TEvWorker::TEvHandshake>(handle);
+        Y_UNUSED(handshake);
+    }
+
+    Y_UNIT_TEST(FeedSlowly) {
+        TTestActorRuntime runtime;
+        runtime.Initialize(NKikimr::TAppPrepare().Unwrap());
+
+        TActorId pqtablet = runtime.AllocateEdgeActor();
+        TActorId reader = runtime.Register(CreateLocalPartitionReader(pqtablet, PARTITION));
+        runtime.EnableScheduleForActor(reader);
+        TActorId worker = runtime.AllocateEdgeActor();
+        TAutoPtr<IEventHandle> handle;
+
+        runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvHandshake));
+
+        GrabInitialPQRequest(runtime);
+
+        auto* getOffsetResponse = new TEvPersQueue::TEvResponse;
+        getOffsetResponse->Record.SetErrorCode(NPersQueue::NErrorCode::OK);
+        getOffsetResponse->Record.MutablePartitionResponse()->MutableCmdGetClientOffsetResult()->SetOffset(INITIAL_OFFSET);
+        runtime.Send(new IEventHandle(reader, pqtablet, getOffsetResponse));
+
+        auto handshake = runtime.GrabEdgeEventRethrow<TEvWorker::TEvHandshake>(handle);
+        Y_UNUSED(handshake);
+
+        runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvPoll));
+
+        for (auto i = 0; i < 3; ++i) {
+            GrabReadPQRequest(runtime, INITIAL_OFFSET);
+
+            runtime.SimulateSleep(TDuration::Seconds(1));
+
+            runtime.Send(new IEventHandle(reader, pqtablet, GenerateEmptyData()));
+        }
+
+        for (auto i = 0; i < 3; ++i) {
+            if (i != 0) {
+                runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvPoll));
+            }
+
+            GrabReadPQRequest(runtime, INITIAL_OFFSET + i * 2);
+
+            runtime.Send(new IEventHandle(reader, pqtablet, GenerateData(i)));
+
+            GrabDataEvent(runtime, i);
+
+            // TODO check commit
+        }
+
+        runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvPoll));
+
+        for (auto i = 0; i < 2; ++i) {
+            GrabReadPQRequest(runtime, 9);
+
+            runtime.SimulateSleep(TDuration::Seconds(1));
+
+            runtime.Send(new IEventHandle(reader, pqtablet, GenerateEmptyData()));
+        }
+
+        for (auto i = 3; i < 5; ++i) {
+            if (i != 3) {
+                runtime.Send(new IEventHandle(reader, worker, new TEvWorker::TEvPoll));
+            }
+
+            GrabReadPQRequest(runtime, INITIAL_OFFSET + i * 2);
+
+            runtime.Send(new IEventHandle(reader, pqtablet, GenerateData(i)));
+
+            GrabDataEvent(runtime, i);
 
             // TODO check commit
         }