Synchronization of the read and write order for compactified topics

ydb/core/persqueue/events/internal.h

@@ -201,6 +201,9 @@ struct TEvPQ {
         EvTransactionCompleted,
         EvListAllTopicsResponse,
         EvRunCompaction,
+        EvAcquireExclusiveLock,
+        EvExclusiveLockAcquired,
+        EvReleaseExclusiveLock,
         EvEnd
     };
 
@@ -1222,6 +1225,15 @@ struct TEvPQ {
         ui64 MaxBlobSize = 0;
         ui64 CumulativeSize = 0;
     };
+
+    struct TEvAcquireExclusiveLock : public TEventLocal<TEvAcquireExclusiveLock, EvAcquireExclusiveLock> {
+    };
+
+    struct TEvExclusiveLockAcquired : public TEventLocal<TEvExclusiveLockAcquired, EvExclusiveLockAcquired> {
+    };
+
+    struct TEvReleaseExclusiveLock : public TEventLocal<TEvReleaseExclusiveLock, EvReleaseExclusiveLock> {
+    };
 };
 
 } //NKikimr

ydb/core/persqueue/partition_init.cpp

@@ -943,7 +943,7 @@ void TPartition::Initialize(const TActorContext& ctx) {
     WriteCycleStartTime = ctx.Now();
 
     ReadQuotaTrackerActor = Register(new TReadQuoter(
-        ctx,
+        AppData(ctx)->PQConfig,
         TopicConverter,
         Config,
         Partition,

ydb/core/persqueue/read_quoter.cpp

@@ -21,11 +21,14 @@ void TPartitionQuoterBase::HandleQuotaRequest(TEvPQ::TEvRequestQuota::TPtr& ev,
     QuotaRequestedTimes.emplace(ev->Get()->Cookie, ctx.Now());
     TRequestContext context{ev->Release(), ev->Sender};
     HandleQuotaRequestImpl(context);
-    if (RequestsInflight >= MaxInflightRequests || !WaitingInflightRequests.empty()) {
+    if ((ExclusiveLockState != EExclusiveLockState::EReleased) || (RequestsInflight >= MaxInflightRequests) || !WaitingInflightRequests.empty()) {
         if (WaitingInflightRequests.empty())
             InflightIsFullStartTime = ctx.Now();
         WaitingInflightRequests.push_back(std::move(context));
     } else {
+        // ExclusiveLockState == EExclusiveLockState::EReleased
+        // RequestsInflight < MaxInflightRequests
+        // WaitingInflightRequests.empty()
         StartQuoting(std::move(context));
     }
 }
@@ -98,14 +101,19 @@ void TPartitionQuoterBase::HandleConsumed(TEvPQ::TEvConsumed::TPtr& ev, const TA
                         " partition " << Partition <<
                         " readCookie " << ev->Get()->RequestCookie);
     }
+
+    if (!RequestsInflight && WaitingInflightRequests.empty() && (ExclusiveLockState == EExclusiveLockState::EAcquiring)) {
+        ExclusiveLockState = EExclusiveLockState::EAcquired;
+        ReplyExclusiveLockAcquired(ExclusiveLockRequester);
+    }
 }
 
 void TPartitionQuoterBase::ProcessInflightQueue() {
     auto now = ActorContext().Now();
     while (!WaitingInflightRequests.empty() && RequestsInflight < MaxInflightRequests) {
         StartQuoting(std::move(WaitingInflightRequests.front()));
         WaitingInflightRequests.pop_front();
-        if (WaitingInflightRequests.size() == 0) {
+        if (WaitingInflightRequests.empty()) {
             InflightLimitSlidingWindow.Update((now - InflightIsFullStartTime).MicroSeconds(), now);
             UpdateCounters(ActorContext());
         }
@@ -140,6 +148,47 @@ void TPartitionQuoterBase::ScheduleWakeUp(const TActorContext& ctx) {
     ctx.Schedule(WAKE_UP_TIMEOUT, new TEvents::TEvWakeup());
 }
 
+void TPartitionQuoterBase::HandleAcquireExclusiveLock(TEvPQ::TEvAcquireExclusiveLock::TPtr& ev, const TActorContext& ctx)
+{
+    switch (ExclusiveLockState) {
+    case EExclusiveLockState::EReleased:
+        ExclusiveLockState = EExclusiveLockState::EAcquiring;
+        ExclusiveLockRequester = ev->Sender;
+
+        [[fallthrough]];
+
+    case EExclusiveLockState::EAcquiring:
+        if (RequestsInflight) {
+            return;
+        }
+
+        ExclusiveLockState = EExclusiveLockState::EAcquired;
+
+        [[fallthrough]];
+
+    case EExclusiveLockState::EAcquired:
+        ReplyExclusiveLockAcquired(ev->Sender);
+        return;
+    }
+
+    Y_UNUSED(ctx);
+}
+
+void TPartitionQuoterBase::HandleReleaseExclusiveLock(TEvPQ::TEvReleaseExclusiveLock::TPtr& ev, const TActorContext& ctx)
+{
+    ExclusiveLockState = EExclusiveLockState::EReleased;
+
+    ProcessInflightQueue();
+
+    Y_UNUSED(ev);
+    Y_UNUSED(ctx);
+}
+
+void TPartitionQuoterBase::ReplyExclusiveLockAcquired(const TActorId& receiver)
+{
+    Send(receiver, new TEvPQ::TEvExclusiveLockAcquired());
+}
+
 void TReadQuoter::HandleQuotaRequestImpl(TRequestContext& context) {
     auto* readRequest = context.Request->Request->CastAsLocal<TEvPQ::TEvRead>();
     GetOrCreateConsumerQuota(readRequest->ClientId, ActorContext());

ydb/core/persqueue/read_quoter.h

@@ -143,6 +143,8 @@ const TDuration WAKE_UP_TIMEOUT = TDuration::Seconds(1);
     void HandleConfigUpdate(TEvPQ::TEvChangePartitionConfig::TPtr& ev, const TActorContext& ctx);
     virtual void HandlePoisonPill(TEvents::TEvPoisonPill::TPtr& ev, const TActorContext& ctx) = 0;
     virtual void HandleUpdateAccountQuotaCounters(NAccountQuoterEvents::TEvCounters::TPtr& ev, const TActorContext& ctx) = 0;
+    void HandleAcquireExclusiveLock(TEvPQ::TEvAcquireExclusiveLock::TPtr& ev, const TActorContext& ctx);
+    void HandleReleaseExclusiveLock(TEvPQ::TEvReleaseExclusiveLock::TPtr& ev, const TActorContext& ctx);
 
 protected:
     virtual void HandleQuotaRequestImpl(TRequestContext& context) = 0;
@@ -177,7 +179,7 @@ const TDuration WAKE_UP_TIMEOUT = TDuration::Seconds(1);
 
     void ScheduleWakeUp(const TActorContext& ctx);
 
-STFUNC(StateWork)
+    STFUNC(StateWork)
     {
         switch (ev->GetTypeRewrite()) {
             HFunc(TEvPQ::TEvRequestQuota, HandleQuotaRequest);
@@ -186,13 +188,17 @@ STFUNC(StateWork)
             HFunc(TEvPQ::TEvConsumed, HandleConsumed);
             HFunc(TEvPQ::TEvChangePartitionConfig, HandleConfigUpdate);
             HFunc(NAccountQuoterEvents::TEvCounters, HandleUpdateAccountQuotaCounters);
+            HFunc(TEvPQ::TEvAcquireExclusiveLock, HandleAcquireExclusiveLock);
+            HFunc(TEvPQ::TEvReleaseExclusiveLock, HandleReleaseExclusiveLock);
             HFunc(TEvents::TEvPoisonPill, HandlePoisonPill);
         default:
             ProcessEventImpl(ev);
             break;
         };
     }
 
+    void ReplyExclusiveLockAcquired(const TActorId& receiver);
+
 protected:
     std::deque<TRequestContext> WaitingInflightRequests;
     TMicrosecondsSlidingWindow InflightLimitSlidingWindow;
@@ -205,6 +211,12 @@ STFUNC(StateWork)
     TTabletCountersBase Counters;
 
 private:
+    enum class EExclusiveLockState {
+        EReleased,
+        EAcquiring,
+        EAcquired,
+    };
+
     TActorId TabletActor;
     std::deque<TRequestContext> WaitingTotalPartitionQuotaRequests;
     THashMap<ui64, TRequestContext> PendingAccountQuotaRequests;
@@ -213,6 +225,8 @@ STFUNC(StateWork)
     ui64 MaxInflightRequests;
     bool TotalPartitionQuotaEnabled;
     TVector<TEvPQ::TEvRequestQuota::TPtr> PendingQuotaRequests;
+    EExclusiveLockState ExclusiveLockState = EExclusiveLockState::EReleased;
+    TActorId ExclusiveLockRequester;
 };
 
 
@@ -222,7 +236,7 @@ const static ui64 DEFAULT_READ_SPEED_AND_BURST = 1'000'000'000;
 
 public:
     TReadQuoter(
-        const TActorContext& ctx,
+        const NKikimrPQ::TPQConfig& pqConfig,
         const NPersQueue::TTopicConverterPtr& topicConverter,
         const NKikimrPQ::TPQTabletConfig& config,
         const TPartitionId& partition,
@@ -233,7 +247,7 @@ const static ui64 DEFAULT_READ_SPEED_AND_BURST = 1'000'000'000;
     )
         : TPartitionQuoterBase(
                 topicConverter, config, partition, tabletActor, true, tabletId, counters,
-                AppData(ctx)->PQConfig.GetMaxInflightReadRequestsPerPartition()
+                pqConfig.GetMaxInflightReadRequestsPerPartition()
         )
         , Parent(parent)
     {

ydb/core/persqueue/ut/pqtablet_ut.cpp

@@ -1,6 +1,7 @@
 #include <ydb/core/keyvalue/keyvalue_events.h>
 #include <ydb/core/persqueue/events/internal.h>
 #include <ydb/core/persqueue/partition.h>
+#include <ydb/core/persqueue/read_quoter.h>
 #include <ydb/core/persqueue/ut/common/pq_ut_common.h>
 #include <ydb/core/protos/counters_keyvalue.pb.h>
 #include <ydb/core/protos/pqconfig.pb.h>
@@ -291,6 +292,17 @@ class TPQTabletFixture : public NUnitTest::TBaseFixture {
     template<class EventType>
     void AddOneTimeEventObserver(bool& seenEvent, std::function<TTestActorRuntimeBase::EEventAction(TAutoPtr<IEventHandle>&)> callback = [](){return TTestActorRuntimeBase::EEventAction::PROCESS;});
 
+    void ExpectNoExclusiveLockAcquired();
+    void ExpectNoReadQuotaAcquired();
+    void SendAcquireExclusiveLock();
+    void SendAcquireReadQuota();
+    void SendReadQuotaConsumed();
+    void SendReleaseExclusiveLock();
+    void WaitExclusiveLockAcquired();
+    void WaitReadQuotaAcquired();
+
+    void EnsureReadQuoterExists();
+
     //
     // TODO(abcdef): для тестирования повторных вызовов нужны примитивы Send+Wait
     //
@@ -303,6 +315,17 @@ class TPQTabletFixture : public NUnitTest::TBaseFixture {
     TTestActorRuntimeBase::TEventObserver PrevEventObserver;
 
     TActorId Pipe;
+
+    struct TReadQuoter {
+        NKikimrPQ::TPQConfig PQConfig;
+        NPersQueue::TTopicConverterPtr TopicConverter;
+        NKikimrPQ::TPQTabletConfig PQTabletConfig;
+        TPartitionId PartitionId;
+        TTabletCountersBase Counters;
+        TActorId Quoter;
+    };
+
+    TMaybe<TReadQuoter> ReadQuoter;
 };
 
 void TPQTabletFixture::SetUp(NUnitTest::TTestContext&)
@@ -2920,6 +2943,102 @@ Y_UNIT_TEST_F(PQTablet_App_SendReadSet_Invalid_Step, TPQTabletFixture)
     WaitForAppSendRsResponse({.Status = false,});
 }
 
+void TPQTabletFixture::ExpectNoExclusiveLockAcquired()
+{
+    EnsureReadQuoterExists();
+}
+
+void TPQTabletFixture::ExpectNoReadQuotaAcquired()
+{
+    EnsureReadQuoterExists();
+}
+
+void TPQTabletFixture::SendAcquireExclusiveLock()
+{
+    EnsureReadQuoterExists();
+
+    Ctx->Runtime->Send(ReadQuoter->Quoter,
+                       Ctx->Edge,
+                       new TEvPQ::TEvAcquireExclusiveLock());
+}
+
+void TPQTabletFixture::SendAcquireReadQuota(ui64 cookie)
+{
+    EnsureReadQuoterExists();
+
+    auto request = MakeHolder<TEvPQ::TEvRead>();
+
+    Ctx->Runtime->Send(ReadQuoter->Quoter,
+                       Ctx->Edge,
+                       new TEvPQ::TEvRequestQuota(cookie, std::move(request)));
+}
+
+void TPQTabletFixture::SendReadQuotaConsumed()
+{
+    EnsureReadQuoterExists();
+
+    Ctx->Runtime->Send(ReadQuoter->Quoter,
+                       Ctx->Edge,
+                       new TEvPQ::TEvConsumed());
+}
+
+void TPQTabletFixture::SendReleaseExclusiveLock()
+{
+    EnsureReadQuoterExists();
+
+    Ctx->Runtime->Send(ReadQuoter->Quoter,
+                       Ctx->Edge,
+                       new TEvPQ::TEvReleaseExclusiveLock());
+}
+
+void TPQTabletFixture::WaitExclusiveLockAcquired()
+{
+    EnsureReadQuoterExists();
+}
+
+void TPQTabletFixture::WaitReadQuotaAcquired()
+{
+    EnsureReadQuoterExists();
+}
+
+void TPQTabletFixture::EnsureReadQuoterExists()
+{
+    if (ReadQuoter) {
+        return;
+    }
+
+    Cerr << "Ctx->Edge=" << Ctx->Edge << Endl;
+
+    ReadQuoter.ConstructInPlace();
+    ReadQuoter->Quoter = Ctx->Runtime->Register(new NPQ::TReadQuoter(ReadQuoter->PQConfig,
+                                                                     ReadQuoter->TopicConverter,
+                                                                     ReadQuoter->PQTabletConfig,
+                                                                     ReadQuoter->PartitionId,
+                                                                     TActorId{}, // TabletActor
+                                                                     Ctx->Edge,
+                                                                     1234567890, // TabletId
+                                                                     ReadQuoter->Counters));
+}
+
+Y_UNIT_TEST_F(ReadQuoter_ExclusiveLock, TPQTabletFixture)
+{
+    SendAcquireReadQuota();
+    WaitReadQuotaAcquired();
+
+    SendAcquireExclusiveLock();
+    ExpectNoExclusiveLockAcquired();
+
+    SendReadQuotaConsumed();
+    WaitExclusiveLockAcquired();
+
+    SendAcquireReadQuota();
+    ExpectNoReadQuotaAcquired();
+
+    SendReleaseExclusiveLock();
+    WaitReadQuotaAcquired();
+
+    Y_FAIL("not implemented");
+}
 
 }