From f64949582813c2d3813276576357ece9e4507db7 Mon Sep 17 00:00:00 2001
From: Sergey Belyakov <serg-belyakov@ydb.tech>
Date: Thu, 28 Dec 2023 16:15:20 +0000
Subject: [PATCH 1/4] Introduce InitialAllocation option for storage load actor

---
 ydb/core/load_test/gen.h           |   2 +
 ydb/core/load_test/group_write.cpp | 348 +++++++++++++++++++++++------
 ydb/core/protos/load_test.proto    |  12 +
 3 files changed, 293 insertions(+), 69 deletions(-)

diff --git a/ydb/core/load_test/gen.h b/ydb/core/load_test/gen.h
index b5ad5af2f4ba..b8e07e18f5db 100644
--- a/ydb/core/load_test/gen.h
+++ b/ydb/core/load_test/gen.h
@@ -18,6 +18,8 @@ namespace NKikimr {
         double AccumWeight = 0;
 
     public:
+        TGenerator() = default;
+
         template<typename T>
         TGenerator(const google::protobuf::RepeatedPtrField<T>& setting) {
             for (const auto& item : setting) {
diff --git a/ydb/core/load_test/group_write.cpp b/ydb/core/load_test/group_write.cpp
index 34ea48e295b2..532459b2ded7 100644
--- a/ydb/core/load_test/group_write.cpp
+++ b/ydb/core/load_test/group_write.cpp
@@ -86,11 +86,126 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         NKikimrBlobStorage::EPutHandleClass PutHandleClass;
     };
 
-    friend struct THardRateDelayManager;
+    class TInitialAllocation {
+    public:
+        using TProtoSettings = NKikimr::TEvLoadTestRequest::TStorageLoad::TInitialBlobAllocation;
+
+        TInitialAllocation() = default;
+
+        TInitialAllocation(const TProtoSettings& proto)
+            : SizeGenerator(proto.GetBlobSizes())
+            , SizeToWrite(proto.GetTotalSize())
+            , BlobsToWrite(proto.GetBlobsNumber())
+            , WritesInFlight(proto.GetWritesInFlight())
+        {
+            if (proto.HasPutHandleClass()) {
+                PutHandleClass = proto.GetPutHandleClass();
+            }
+        }
+
+    public:
+        bool IsEmpty() {
+            return ConfirmedBlobs.empty();
+        }
+
+        bool Size() {
+            return ConfirmedBlobs.size();
+        }
+
+        const TLogoBlobID& operator[](ui32 idx) const {
+            return ConfirmedBlobs[idx];
+        }
+
+        operator bool() const {
+            return SizeToWrite > 0 || BlobsToWrite > 0;
+        }
+
+        bool EnoughBlobsWritten(bool countPending = false) {
+            if (SizeToWrite > 0) {
+                return ConfirmedDataSize + PendingWriteSize * countPending >= SizeToWrite;
+            } else if (BlobsToWrite > 0) {
+                return ConfirmedBlobs.size() + PendingWrites * countPending >= BlobsToWrite;
+            }
+            return true;
+        }
+
+        void ConfirmBlob(const TLogoBlobID& id, bool success) {
+            PendingWriteSize -= id.BlobSize();
+            PendingWrites -= 1;
+            if (success) {
+                ConfirmedBlobs.push_back(id);
+                ConfirmedDataSize += id.BlobSize();
+            }
+        }
+
+        std::unique_ptr<TEvBlobStorage::TEvPut> MakePutMessage(ui64 tabletId, ui32 gen, ui32 step, ui32 channel) {
+            ui32 blobSize = SizeGenerator.Generate();
+            const TLogoBlobID id(tabletId, gen, step, channel, blobSize, BlobCookie++);
+            const TSharedData buffer = GenerateBuffer<TSharedData>(id);
+            auto ev = std::make_unique<TEvBlobStorage::TEvPut>(id, buffer, TInstant::Max(), PutHandleClass);
+            PendingWriteSize += blobSize;
+            PendingWrites += 1;
+            return std::move(ev);
+        }
+
+        std::unique_ptr<TEvBlobStorage::TEvCollectGarbage> ManageKeepFlags(ui64 tabletId, ui32 gen, ui32 step,
+                ui32 channel, bool keep) {
+            auto blobsWritten = std::make_unique<TVector<TLogoBlobID>>(ConfirmedBlobs);
+
+            if (keep) {
+                return std::make_unique<TEvBlobStorage::TEvCollectGarbage>(tabletId, gen, step, channel,
+                        false, gen, step, blobsWritten.release(), nullptr, TInstant::Max(), false);
+            } else {
+                ConfirmedDataSize = 0;
+                ConfirmedBlobs.clear();
+                return std::make_unique<TEvBlobStorage::TEvCollectGarbage>(tabletId, gen, step, channel,
+                        true, gen, step, nullptr, blobsWritten.release(), TInstant::Max(), false);
+            }
+        }
+
+        TLogoBlobID GetRandomBlobId() {
+            Y_VERIFY(!IsEmpty());
+            auto iter = ConfirmedBlobs.begin();
+            std::advance(iter, RandomNumber(ConfirmedBlobs.size()));
+            return *iter;
+        }
+
+        TString ToString() {
+            return TStringBuilder() << "TInitialAllocation# {"
+                << " PutHandleClass# " << NKikimrBlobStorage::EPutHandleClass_Name(PutHandleClass)
+                << " SizeToWrite# " << SizeToWrite
+                << " BlobsToWrite# " << BlobsToWrite
+                << " WritesInFlight# " << WritesInFlight
+                << " ConfirmedSize# " << ConfirmedDataSize
+                << " ConfirmedBlobs.size()# " << ConfirmedBlobs.size()
+                << " PendingWrites# " << PendingWrites
+                << " PendingWriteSize# " << PendingWriteSize << " }";
+        }
+
+    ui32 GetWritesInFlight() {
+        return WritesInFlight;
+    }
+
+    private:
+        TSizeGenerator SizeGenerator;
+        NKikimrBlobStorage::EPutHandleClass PutHandleClass = NKikimrBlobStorage::EPutHandleClass::UserData;
+
+        uint64_t SizeToWrite;
+        uint64_t BlobsToWrite;
+        ui32 WritesInFlight = 1;
+        uint64_t ConfirmedDataSize = 0;
+        TVector<TLogoBlobID> ConfirmedBlobs;
+
+        uint64_t PendingWriteSize = 0;
+        uint64_t PendingWrites = 0;
+
+        ui64 BlobCookie = 0;
+    };
 
     struct TRequestDelayManager {
         virtual ~TRequestDelayManager() = default;
 
+        virtual void Start(TMonotonic now) = 0;
         virtual TDuration CalculateDelayForNextRequest(TMonotonic now) = 0;
         virtual void CountResponse() = 0;
 
@@ -106,6 +221,10 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             : IntervalGenerator(intervalGenerator) {
         }
 
+        void Start(TMonotonic/* now*/) override {
+            return;
+        }
+
         TDuration CalculateDelayForNextRequest(TMonotonic/* now*/) override {
             return IntervalGenerator.Generate();
         }
@@ -120,15 +239,11 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
     };
 
     struct THardRateDelayManager : public TRequestDelayManager {
-        THardRateDelayManager(double requestsPerSecondAtStart, double requestsPerSecondOnFinish, TMonotonic now, const TMaybe<TDuration>& duration)
+        THardRateDelayManager(double requestsPerSecondAtStart, double requestsPerSecondOnFinish, const TMaybe<TDuration>& duration)
             : EpochDuration(TDuration::MilliSeconds(100))
             , RequestRateAtStart((EpochDuration / TDuration::Seconds(1)) * requestsPerSecondAtStart)
             , RequestRateOnFinish((EpochDuration / TDuration::Seconds(1)) * requestsPerSecondOnFinish)
-            , CurrentEpochEnd(now)
-            , LoadStart(now)
-            , LoadDuration(duration)
-            , PlannedForCurrentEpoch(std::max(1., CalculateRequestRate(now))) {
-            CalculateDelayForNextRequest(now);
+            , LoadDuration(duration) {
         }
 
         const TDuration EpochDuration;
@@ -142,12 +257,19 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         TDuration CurrentDelay = TDuration::Seconds(0);
         TMonotonic Now = TMonotonic::Max();
 
-        const TMonotonic LoadStart;
+        TMonotonic LoadStart;
         const TMaybe<TDuration> LoadDuration;
 
         double PlannedForCurrentEpoch;
         ui32 ResponsesAwaiting = 0;
 
+        void Start(TMonotonic now) override {
+            LoadStart = now;
+            CurrentEpochEnd = now;
+            PlannedForCurrentEpoch = std::max(1., CalculateRequestRate(now));
+            CalculateDelayForNextRequest(now);
+        }
+
         double CalculateRequestRate(TMonotonic now) {
             double ratio = LoadDuration ? (now - LoadStart) / *LoadDuration : 0;
             return ratio * (RequestRateOnFinish - RequestRateAtStart) + RequestRateAtStart;
@@ -159,8 +281,6 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 CurrentEpochEnd += EpochDuration;
                 RequestsPerEpoch = CalculateRequestRate(now); 
                 PlannedForCurrentEpoch += RequestsPerEpoch;
-                PlannedForCurrentEpoch += ResponsesAwaiting;
-                ResponsesAwaiting = 0;
             }
 
             if (PlannedForCurrentEpoch < 1) {
@@ -203,10 +323,12 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
     };
 
+    friend class TTabletWriter;
     class TTabletWriter {
     public:
         struct TRequestDispatchingSettings {
-            TSizeGenerator SizeGen;
+            bool LoadEnabled = true;
+            std::optional<TSizeGenerator> SizeGen;
             std::shared_ptr<TRequestDelayManager> DelayManager;
             const ui32 MaxRequestsInFlight;
             const ui64 MaxBytesInFlight;
@@ -222,10 +344,9 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
 
         const TDuration ExposePeriod = TDuration::Seconds(10);
 
+        TLogWriterLoadTestActor& Self;
         TIntrusivePtr<::NMonitoring::TDynamicCounters> TagCounters;
         TIntrusivePtr<::NMonitoring::TDynamicCounters> Counters;
-        TWakeupQueue& WakeupQueue;
-        TQueryDispatcher& QueryDispatcher;
         const ui64 TabletId;
         const ui32 Channel;
         ui32 Generation;
@@ -286,18 +407,23 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         ui64 ScriptedRound;
         TVector<TReqInfo> ScriptedRequests;
 
+        // Initial allocation
+        TInitialAllocation InitialAllocation;
+
+        bool MainCycleStarted = false;
+
     public:
-        TTabletWriter(ui64 tag, TIntrusivePtr<::NMonitoring::TDynamicCounters> counters,
-                TWakeupQueue& wakeupQueue, TQueryDispatcher& queryDispatcher, ui64 tabletId, ui32 channel,
+        TTabletWriter(TIntrusivePtr<::NMonitoring::TDynamicCounters> counters,
+                TLogWriterLoadTestActor& self, ui64 tabletId, ui32 channel,
                 TMaybe<ui32> generation, ui32 groupId,
                 NKikimrBlobStorage::EPutHandleClass putHandleClass, const TRequestDispatchingSettings& writeSettings,
                 NKikimrBlobStorage::EGetHandleClass getHandleClass, const TRequestDispatchingSettings& readSettings,
                 TIntervalGenerator garbageCollectIntervalGen,
-                TDuration scriptedRoundDuration, TVector<TReqInfo>&& scriptedRequests)
-            : TagCounters(counters->GetSubgroup("tag", Sprintf("%" PRIu64, tag)))
+                TDuration scriptedRoundDuration, TVector<TReqInfo>&& scriptedRequests,
+                const TInitialAllocation& initialAllocation)
+            : Self(self)
+            , TagCounters(counters->GetSubgroup("tag", Sprintf("%" PRIu64, Self.Tag)))
             , Counters(TagCounters->GetSubgroup("channel", Sprintf("%" PRIu32, channel)))
-            , WakeupQueue(wakeupQueue)
-            , QueryDispatcher(queryDispatcher)
             , TabletId(tabletId)
             , Channel(channel)
             , Generation(generation ? *generation : 0)
@@ -331,6 +457,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             , ScriptedCounter(0)
             , ScriptedRound(0)
             , ScriptedRequests(std::move(scriptedRequests))
+            , InitialAllocation(initialAllocation)
         {
             *Counters->GetCounter("tabletId") = tabletId;
             const auto& percCounters = Counters->GetSubgroup("sensor", "microseconds");
@@ -376,7 +503,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 Generation = res->BlockedGeneration + 1;
                 IssueTEvBlock(ctx);
             };
-            SendToBSProxy(ctx, GroupId, ev.release(), QueryDispatcher.ObtainCookie(std::move(callback)));
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
         }
 
         void IssueTEvBlock(const TActorContext& ctx) {
@@ -398,7 +525,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 ++Generation;
                 IssueLastBlob(ctx);
             };
-            SendToBSProxy(ctx, GroupId, ev.release(), QueryDispatcher.ObtainCookie(std::move(callback)));
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
         }
 
         void IssueLastBlob(const TActorContext& ctx) {
@@ -418,7 +545,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 IssueTEvCollectGarbage(ctx);
             };
 
-            SendToBSProxy(ctx, GroupId, ev.release(), QueryDispatcher.ObtainCookie(std::move(callback)));
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
         }
 
         void IssueTEvCollectGarbage(const TActorContext& ctx) {
@@ -433,20 +560,73 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                     return;
                 }
                 LOG_INFO_S(ctx, NKikimrServices::BS_LOAD_TEST, PrintMe() << " recieved " << res->ToString());
-                StartWorking(ctx);
+                MakeInitialAllocation(ctx);
             };
 
-            SendToBSProxy(ctx, GroupId, ev.Release(), QueryDispatcher.ObtainCookie(std::move(callback)));
+            SendToBSProxy(ctx, GroupId, ev.Release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
         }
 
         void StartWorking(const TActorContext& ctx) {
+            MainCycleStarted = true;
             StartTimestamp = TActivationContext::Monotonic();
+            if (Self.TestDuration) {
+                ctx.Schedule(*Self.TestDuration, new TEvents::TEvPoisonPill());
+            }
             InitializeTrackers(StartTimestamp);
+            WriteSettings.DelayManager->Start(StartTimestamp);
             IssueWriteIfPossible(ctx);
+            ReadSettings.DelayManager->Start(StartTimestamp);
+            IssueReadIfPossible(ctx);
             ScheduleGarbageCollect(ctx);
             ExposeCounters(ctx);
         }
 
+        void MakeInitialAllocation(const TActorContext& ctx) {
+            if (InitialAllocation.EnoughBlobsWritten(false)) {
+                Self.InitialAllocationCompleted(ctx);
+                return;
+            }
+            LOG_DEBUG_S(ctx, NKikimrServices::BS_LOAD_TEST, PrintMe() << " going to make initial allocation,"
+                    << InitialAllocation.ToString());
+            for (ui32 i = 0; i < InitialAllocation.GetWritesInFlight() && !InitialAllocation.EnoughBlobsWritten(true); ++i) {
+                IssueInitialPut(ctx);
+            }
+        }
+
+        void IssueInitialPut(const TActorContext& ctx) {
+            auto ev = InitialAllocation.MakePutMessage(TabletId, Generation, GarbageCollectStep, Channel);
+
+            auto callback = [this](IEventBase *event, const TActorContext& ctx) {
+                auto *res = dynamic_cast<TEvBlobStorage::TEvPutResult*>(event);
+                Y_ABORT_UNLESS(res);
+
+                InitialAllocation.ConfirmBlob(res->Id, CheckStatus(ctx, res, {}));
+                if (!InitialAllocation.EnoughBlobsWritten(true)) {
+                    IssueInitialPut(ctx);
+                } else if (InitialAllocation.EnoughBlobsWritten(false)) {
+                    SetKeepFlagsOnInitialAllocation(ctx);
+                }
+            };
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
+        }
+
+        void SetKeepFlagsOnInitialAllocation(const TActorContext& ctx) {
+            auto ev = InitialAllocation.ManageKeepFlags(TabletId, Generation, GarbageCollectStep, Channel, true);
+
+            LOG_DEBUG_S(ctx, NKikimrServices::BS_LOAD_TEST, PrintMe() << " going to set keep flags on initally allocated blobs, ev#"
+                    << ev->Print(false));
+            auto callback = [this](IEventBase *event, const TActorContext& ctx) {
+                auto *res = dynamic_cast<TEvBlobStorage::TEvCollectGarbageResult*>(event);
+                Y_ABORT_UNLESS(res);
+                
+                if (!MainCycleStarted) {
+                    StartWorking(ctx);
+                }
+            };
+
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
+        }
+
         void StopWorking(const TActorContext& ctx) {
             auto ev = TEvBlobStorage::TEvCollectGarbage::CreateHardBarrier(TabletId, Generation, GarbageCollectStep,
                     Channel, Generation, Max<ui32>(), TInstant::Max());
@@ -464,7 +644,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                     ctx.Send(ctx.SelfID, new TEvStopTest());
                 }
             };
-            SendToBSProxy(ctx, GroupId, ev.Release(), QueryDispatcher.ObtainCookie(std::move(callback)));
+            SendToBSProxy(ctx, GroupId, ev.Release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
         }
 
         void InitializeTrackers(TMonotonic now) {
@@ -521,7 +701,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             ReadBytesInFlightQT.CalculateQuantiles();
 
             using namespace std::placeholders;
-            WakeupQueue.Put(TActivationContext::Monotonic() + ExposePeriod,
+            Self.WakeupQueue.Put(TActivationContext::Monotonic() + ExposePeriod,
                     std::bind(&TTabletWriter::ExposeCounters, this, _1), ctx);
         }
 
@@ -569,6 +749,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 DUMP_PARAM_FINAL(ReadSettings.MaxRequestsInFlight)
                 DUMP_PARAM_FINAL(ReadSettings.MaxBytesInFlight)
                 DUMP_PARAM(ConfirmedBlobIds.size())
+                DUMP_PARAM(InitialAllocation.ToString())
 
                 static constexpr size_t count = 5;
                 std::array<size_t, count> nums{{9000, 9900, 9990, 9999, 10000}};
@@ -613,20 +794,21 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         void UpdateNextWakeups(const TActorContext& ctx, const TMonotonic& now) {
             if (now < NextWriteTimestamp && !NextWriteInQueue) {
                 using namespace std::placeholders;
-                WakeupQueue.Put(NextWriteTimestamp, std::bind(&TTabletWriter::IssueWriteIfPossible, this, _1), ctx);
+                Self.WakeupQueue.Put(NextWriteTimestamp, std::bind(&TTabletWriter::IssueWriteIfPossible, this, _1), ctx);
                 NextWriteInQueue = true;
             }
 
             if (now < NextReadTimestamp && !NextReadInQueue) {
                 using namespace std::placeholders;
-                WakeupQueue.Put(NextReadTimestamp, std::bind(&TTabletWriter::IssueReadIfPossible, this, _1), ctx);
+                Self.WakeupQueue.Put(NextReadTimestamp, std::bind(&TTabletWriter::IssueReadIfPossible, this, _1), ctx);
                 NextReadInQueue = true;
             }
         }
 
         void IssueWriteIfPossible(const TActorContext& ctx) {
             const TMonotonic now = TActivationContext::Monotonic();
-            while ((WritesInFlight < WriteSettings.MaxRequestsInFlight || !WriteSettings.MaxRequestsInFlight) &&
+            while (WriteSettings.LoadEnabled &&
+                    (WritesInFlight < WriteSettings.MaxRequestsInFlight || !WriteSettings.MaxRequestsInFlight) &&
                     (WriteBytesInFlight < WriteSettings.MaxBytesInFlight || !WriteSettings.MaxBytesInFlight) &&
                     (TotalBytesWritten + WriteBytesInFlight < WriteSettings.MaxTotalBytes || !WriteSettings.MaxTotalBytes) &&
                     now >= NextWriteTimestamp &&
@@ -648,7 +830,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 size = req.Size;
                 putHandleClass = req.PutHandleClass;
             } else {
-                size = WriteSettings.SizeGen.Generate();
+                size = WriteSettings.SizeGen->Generate();
                 putHandleClass = PutHandleClass;
             }
             const TLogoBlobID id(TabletId, Generation, WriteStep, Channel, size, Cookie);
@@ -696,14 +878,14 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 ResponseQT->Increment(response.MicroSeconds());
                 IssueWriteIfPossible(ctx);
 
-                if (ConfirmedBlobIds.size() == 1) {
+                if (ConfirmedBlobIds.size() == 1 && InitialAllocation.IsEmpty()) {
                     if (NextReadTimestamp == TMonotonic()) {
                         NextReadTimestamp = TActivationContext::Monotonic();
                     }
                     IssueReadIfPossible(ctx);
                 }
             };
-            SendToBSProxy(ctx, GroupId, ev.release(), QueryDispatcher.ObtainCookie(std::move(writeCallback)));
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(writeCallback)));
             const auto nowCycles = GetCycleCountFast();
             WritesInFlightTimestamps.emplace_back(writeQueryId, nowCycles);
             SentTimestamp.emplace(writeQueryId, nowCycles);
@@ -757,7 +939,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             TDuration duration = GarbageCollectIntervalGen.Generate();
             if (duration != TDuration()) {
                 using namespace std::placeholders;
-                WakeupQueue.Put(TActivationContext::Monotonic() + duration,
+                Self.WakeupQueue.Put(TActivationContext::Monotonic() + duration,
                         std::bind(&TTabletWriter::IssueGarbageCollectRequest, this, _1), ctx);
             }
         }
@@ -769,7 +951,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 auto *res = dynamic_cast<TEvBlobStorage::TEvCollectGarbageResult *>(event);
                 Y_ABORT_UNLESS(res);
             };
-            SendToBSProxy(ctx, GroupId, ev.release(), QueryDispatcher.ObtainCookie(std::move(callback)));
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(callback)));
 
             // just as we have sent this request, we have to trim all confirmed blobs which are going to be deleted
             const auto it = std::lower_bound(ConfirmedBlobIds.begin(), ConfirmedBlobIds.end(),
@@ -786,10 +968,10 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         void IssueReadIfPossible(const TActorContext& ctx) {
             const TMonotonic now = TActivationContext::Monotonic();
 
-            while (ReadsInFlight < ReadSettings.MaxRequestsInFlight &&
+            while (ReadSettings.LoadEnabled && ReadsInFlight < ReadSettings.MaxRequestsInFlight &&
                     (ReadBytesInFlight < ReadSettings.MaxBytesInFlight || !ReadSettings.MaxBytesInFlight) &&
                     now >= NextReadTimestamp &&
-                    ConfirmedBlobIds &&
+                    ConfirmedBlobIds.size() + InitialAllocation.Size() > 0 &&
                     (!ScriptedRequests || ScriptedRequests[ScriptedCounter].EvType == TEvBlobStorage::EvGet)) {
                 IssueReadRequest(ctx);
             }
@@ -801,16 +983,26 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
         void IssueReadRequest(const TActorContext& ctx) {
-            auto iter = ConfirmedBlobIds.begin();
-            std::advance(iter, RandomNumber(ConfirmedBlobIds.size()));
-            const TLogoBlobID &id = *iter;
+            TLogoBlobID id;
+            ui32 confirmedBlobs = ConfirmedBlobIds.size();
+            ui32 initialBlobs = InitialAllocation.Size();
+            Y_VERIFY(confirmedBlobs + initialBlobs > 0);
+            ui32 blobIdx = RandomNumber(confirmedBlobs + initialBlobs);
+    
+            if (blobIdx < confirmedBlobs) {
+                auto iter = ConfirmedBlobIds.begin();
+                std::advance(iter, blobIdx);
+                id = *iter;
+            } else {
+                id = InitialAllocation[blobIdx - confirmedBlobs];
+            }
 
-            ui32 size;
+            ui32 size = Max<ui32>();
             if (ScriptedRequests) {
                 const auto& req = ScriptedRequests[ScriptedCounter];
                 size = req.Size ? req.Size : id.BlobSize();
-            } else {
-                size = ReadSettings.SizeGen.Generate();
+            } else if (ReadSettings.SizeGen) {
+                size = ReadSettings.SizeGen->Generate();
             }
             size = Min(size, id.BlobSize());
 
@@ -842,7 +1034,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 IssueReadIfPossible(ctx);
             };
 
-            SendToBSProxy(ctx, GroupId, ev.release(), QueryDispatcher.ObtainCookie(std::move(readCallback)));
+            SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(readCallback)));
             ReadSentTimestamp.emplace(readQueryId, GetCycleCountFast());
 
             ++ReadsInFlight;
@@ -856,11 +1048,6 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             }
             NextReadInQueue = false;
         }
-
-        template <class ResultContainer = TString>
-        static ResultContainer GenerateBuffer(const TLogoBlobID& id) {
-            return GenDataForLZ4<ResultContainer>(id.BlobSize());
-        }
     };
 
     TString ConfingString;
@@ -888,6 +1075,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
     TMonotonic LastWakeupTime = TMonotonic::Max();
 
     static constexpr ui64 DefaultTabletId = 5000;
+    ui32 WorkersInInitialState = 0;
 
 public:
     static constexpr NKikimrServices::TActivity::EType ActorActivityType() {
@@ -910,33 +1098,32 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             if (!profile.TabletsSize()) {
                 ythrow TLoadActorException() << "TPerTabletProfile.Tablets must have at least one item";
             }
-            if (!profile.WriteSizesSize()) {
-                ythrow TLoadActorException() << "TPerTabletProfile.Sizes must have at least one item";
-            }
-            if (!profile.WriteIntervalsSize() && !profile.HasWriteHardRateDispatcher()) {
-                ythrow TLoadActorException() << "Either TPerTabletProfile.WriteHardRateDispatcher or "
-                        "at least one item in TPerTabletProfile.WriteIntervals must be specified";
-            }
+            bool enableWrites = profile.WriteSizesSize() && profile.GetPutHandleClass() &&
+                    (profile.WriteIntervalsSize() || profile.HasWriteHardRateDispatcher());
 
-            if (!profile.HasPutHandleClass()) {
-                ythrow TLoadActorException() << "missing mandatory TPerTabletProfile.PutHandleClass";
+            TInitialAllocation initialAllocation;
+            if (profile.HasInitialAllocation()) {
+                initialAllocation = TInitialAllocation(profile.GetInitialAllocation());
             }
 
-            NKikimrBlobStorage::EPutHandleClass putHandleClass = profile.GetPutHandleClass();
+            NKikimrBlobStorage::EPutHandleClass putHandleClass = NKikimrBlobStorage::EPutHandleClass::UserData;
+            if (profile.HasPutHandleClass()) {
+                putHandleClass = profile.GetPutHandleClass();
+            }
 
             // object may be shared across multiple writers
             std::shared_ptr<TRequestDelayManager> writeDelayManager;
             if (profile.HasWriteHardRateDispatcher()) {
-                auto now = TActivationContext::Monotonic();
                 const auto& dispatcherSettings = profile.GetWriteHardRateDispatcher();
                 double atStart = dispatcherSettings.GetRequestsPerSecondAtStart();
                 double onFinish = dispatcherSettings.GetRequestsPerSecondOnFinish();
-                writeDelayManager = std::make_shared<THardRateDelayManager>(atStart, onFinish, now, TestDuration);
+                writeDelayManager = std::make_shared<THardRateDelayManager>(atStart, onFinish, TestDuration);
             } else {
                 writeDelayManager = std::make_shared<TRandomIntervalDelayManager>(TIntervalGenerator(profile.GetWriteIntervals()));
             }
 
             TTabletWriter::TRequestDispatchingSettings writeSettings{
+                .LoadEnabled = enableWrites,
                 .SizeGen = TSizeGenerator(profile.GetWriteSizes()),
                 .DelayManager = std::move(writeDelayManager),
                 .MaxRequestsInFlight = profile.GetMaxInFlightWriteRequests(),
@@ -944,6 +1131,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 .MaxTotalBytes = profile.GetMaxTotalBytesWritten(),
             };
 
+            bool enableReads = profile.ReadIntervalsSize() || profile.HasReadHardRateDispatcher();
             NKikimrBlobStorage::EGetHandleClass getHandleClass = NKikimrBlobStorage::EGetHandleClass::FastRead;
             if (profile.HasGetHandleClass()) {
                 getHandleClass = profile.GetGetHandleClass();
@@ -951,17 +1139,22 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
 
             std::shared_ptr<TRequestDelayManager> readDelayManager;
             if (profile.HasReadHardRateDispatcher()) {
-                auto now = TActivationContext::Monotonic();
                 const auto& dispatcherSettings = profile.GetReadHardRateDispatcher();
                 double atStart = dispatcherSettings.GetRequestsPerSecondAtStart();
                 double onFinish = dispatcherSettings.GetRequestsPerSecondOnFinish();
-                readDelayManager = std::make_shared<THardRateDelayManager>(atStart, onFinish, now, TestDuration);
+                readDelayManager = std::make_shared<THardRateDelayManager>(atStart, onFinish, TestDuration);
             } else {
                 readDelayManager = std::make_shared<TRandomIntervalDelayManager>(TIntervalGenerator(profile.GetReadIntervals()));
             }
 
+            std::optional<TSizeGenerator> readSizeGen;
+            if (profile.ReadSizesSize() > 0) {
+                readSizeGen.emplace(profile.GetReadSizes());
+            }
+
             TTabletWriter::TRequestDispatchingSettings readSettings{
-                .SizeGen = TSizeGenerator(profile.GetReadSizes()),
+                .LoadEnabled = enableReads,
+                .SizeGen = readSizeGen,
                 .DelayManager = std::move(readDelayManager),
                 .MaxRequestsInFlight = profile.GetMaxInFlightReadRequests(),
                 .MaxBytesInFlight = profile.GetMaxInFlightReadBytes(),
@@ -1009,12 +1202,27 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                     Y_FAIL();
                 }
 
-                TabletWriters.emplace_back(std::make_unique<TTabletWriter>(Tag, counters, WakeupQueue, QueryDispatcher, tabletId,
+                TabletWriters.emplace_back(std::make_unique<TTabletWriter>(counters, *this, tabletId,
                     tablet.GetChannel(), tablet.HasGeneration() ?  TMaybe<ui32>(tablet.GetGeneration()) : TMaybe<ui32>(),
                     tablet.GetGroupId(), putHandleClass, writeSettings,
                     getHandleClass, readSettings,
                     garbageCollectIntervalGen,
-                    scriptedRoundDuration, std::move(scriptedRequests)));
+                    scriptedRoundDuration, std::move(scriptedRequests),
+                    initialAllocation));
+
+                WorkersInInitialState++;
+            }
+        }
+    }
+
+    void InitialAllocationCompleted(const TActorContext& ctx) {
+        WorkersInInitialState--;
+        if (!WorkersInInitialState) {
+            if (TestDuration) {
+                ctx.Schedule(*TestDuration, new TEvents::TEvPoisonPill());
+            }
+            for (auto& writer : TabletWriters) {
+                writer->StartWorking(ctx);
             }
         }
     }
@@ -1023,9 +1231,6 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         Become(&TLogWriterLoadTestActor::StateFunc);
         EarlyStop = false;
         TestStartTime = TActivationContext::Monotonic();
-        if (TestDuration) {
-            ctx.Schedule(*TestDuration, new TEvents::TEvPoisonPill());
-        }
         for (auto& writer : TabletWriters) {
             writer->Bootstrap(ctx);
         }
@@ -1177,6 +1382,11 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         Y_ABORT("TEvUndelivered# 0x%08" PRIx32 " ActorId# %s", ev->Get()->SourceType, ev->Sender.ToString().data());
     }
 
+    template <class ResultContainer = TString>
+    static ResultContainer GenerateBuffer(const TLogoBlobID& id) {
+        return GenDataForLZ4<ResultContainer>(id.BlobSize());
+    }
+
     STRICT_STFUNC(StateFunc,
         CFunc(EvStopTest, HandleStopTest);
         CFunc(EvUpdateQuantile, HandleUpdateQuantile);
@@ -1187,7 +1397,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         HFunc(TEvBlobStorage::TEvPutResult, HandleDispatcher);
         HFunc(TEvBlobStorage::TEvGetResult, HandleDispatcher);
         HFunc(TEvBlobStorage::TEvCollectGarbageResult, HandleDispatcher);
-        HFunc(NMon::TEvHttpInfo, Handle)
+        HFunc(NMon::TEvHttpInfo, Handle);
         HFunc(TEvents::TEvUndelivered, Handle);
     )
 };
diff --git a/ydb/core/protos/load_test.proto b/ydb/core/protos/load_test.proto
index 3bc3bb84b72b..a487971b6e7c 100644
--- a/ydb/core/protos/load_test.proto
+++ b/ydb/core/protos/load_test.proto
@@ -32,6 +32,16 @@ message TEvLoadTestRequest {
         optional double Weight = 3;
     }
     message TStorageLoad {
+        message TInitialBlobAllocation {
+            oneof DataSize {
+                uint64 TotalSize = 1;
+                uint64 BlobsNumber = 2;
+            }
+            repeated TSizeInfo BlobSizes = 3;
+            optional uint32 WritesInFlight = 4 [default = 1];
+            optional NKikimrBlobStorage.EPutHandleClass PutHandleClass = 5;
+        }
+
         message TRequestInfo {
             optional float SendTime = 1;
             optional uint64 Type = 2;
@@ -72,6 +82,8 @@ message TEvLoadTestRequest {
 
             optional THardRateDispatcherInfo WriteHardRateDispatcher = 15;
             optional THardRateDispatcherInfo ReadHardRateDispatcher = 16;
+
+            optional TInitialBlobAllocation InitialAllocation = 17;
         };
         optional uint64 Tag = 1;
         optional uint32 DurationSeconds = 2;

From fbc6fd1e20ce585526f1c5528f2d4a0cd773145b Mon Sep 17 00:00:00 2001
From: Sergey Belyakov <serg-belyakov@ydb.tech>
Date: Fri, 29 Dec 2023 12:20:53 +0000
Subject: [PATCH 2/4] Improve write throttling, add delay after initial write

---
 ydb/core/load_test/group_write.cpp | 194 ++++++++++++++++++-----------
 ydb/core/protos/load_test.proto    |   6 +-
 2 files changed, 122 insertions(+), 78 deletions(-)

diff --git a/ydb/core/load_test/group_write.cpp b/ydb/core/load_test/group_write.cpp
index 532459b2ded7..e35208ed2c8a 100644
--- a/ydb/core/load_test/group_write.cpp
+++ b/ydb/core/load_test/group_write.cpp
@@ -86,6 +86,44 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         NKikimrBlobStorage::EPutHandleClass PutHandleClass;
     };
 
+    struct TInFlightTracker {
+    public:
+        TInFlightTracker(ui32 maxRequestsInFlight = 0, ui64 maxBytesInFlight = 0)
+            : MaxRequestsInFlight(maxRequestsInFlight)
+            , MaxBytesInFlight(maxBytesInFlight)
+            , RequestsInFlight(0)
+            , BytesInFlight(0)
+        {}
+
+        operator bool() const {
+            return (!MaxRequestsInFlight || RequestsInFlight < MaxRequestsInFlight) &&
+                    (!MaxBytesInFlight || BytesInFlight < MaxBytesInFlight);
+        }
+
+        void Request(ui64 size) {
+            BytesInFlight += size;
+            ++RequestsInFlight;
+        }
+
+        void Response(ui64 size) {
+            Y_DEBUG_ABORT_UNLESS(BytesInFlight >= size && RequestsInFlight >= 0);
+            BytesInFlight -= size;
+            --RequestsInFlight;
+        }
+    
+        TString ToString()  const{
+            return TStringBuilder() << "{"
+                << " Requests# " << RequestsInFlight << "/" << MaxRequestsInFlight
+                << " Bytes# " << BytesInFlight << "/" << MaxBytesInFlight
+                << " }";
+        }
+
+        ui32 MaxRequestsInFlight;
+        ui64 MaxBytesInFlight;
+        ui32 RequestsInFlight;
+        ui64 BytesInFlight;
+    };
+
     class TInitialAllocation {
     public:
         using TProtoSettings = NKikimr::TEvLoadTestRequest::TStorageLoad::TInitialBlobAllocation;
@@ -96,7 +134,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             : SizeGenerator(proto.GetBlobSizes())
             , SizeToWrite(proto.GetTotalSize())
             , BlobsToWrite(proto.GetBlobsNumber())
-            , WritesInFlight(proto.GetWritesInFlight())
+            , InFlightTracker(proto.GetMaxWritesInFlight(), proto.GetMaxWriteBytesInFlight())
         {
             if (proto.HasPutHandleClass()) {
                 PutHandleClass = proto.GetPutHandleClass();
@@ -104,10 +142,6 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
     public:
-        bool IsEmpty() {
-            return ConfirmedBlobs.empty();
-        }
-
         bool Size() {
             return ConfirmedBlobs.size();
         }
@@ -120,18 +154,24 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             return SizeToWrite > 0 || BlobsToWrite > 0;
         }
 
+        bool CanSendRequest() {
+            if (!InFlightTracker) {
+                return false;
+            }
+            return !EnoughBlobsWritten(true);
+        }
+
         bool EnoughBlobsWritten(bool countPending = false) {
             if (SizeToWrite > 0) {
-                return ConfirmedDataSize + PendingWriteSize * countPending >= SizeToWrite;
+                return ConfirmedDataSize + InFlightTracker.BytesInFlight * countPending >= SizeToWrite;
             } else if (BlobsToWrite > 0) {
-                return ConfirmedBlobs.size() + PendingWrites * countPending >= BlobsToWrite;
+                return ConfirmedBlobs.size() + InFlightTracker.RequestsInFlight * countPending >= BlobsToWrite;
             }
             return true;
         }
 
         void ConfirmBlob(const TLogoBlobID& id, bool success) {
-            PendingWriteSize -= id.BlobSize();
-            PendingWrites -= 1;
+            InFlightTracker.Response(id.BlobSize());
             if (success) {
                 ConfirmedBlobs.push_back(id);
                 ConfirmedDataSize += id.BlobSize();
@@ -139,12 +179,12 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
         std::unique_ptr<TEvBlobStorage::TEvPut> MakePutMessage(ui64 tabletId, ui32 gen, ui32 step, ui32 channel) {
+            Y_DEBUG_ABORT_UNLESS(CanSendRequest());
             ui32 blobSize = SizeGenerator.Generate();
             const TLogoBlobID id(tabletId, gen, step, channel, blobSize, BlobCookie++);
             const TSharedData buffer = GenerateBuffer<TSharedData>(id);
             auto ev = std::make_unique<TEvBlobStorage::TEvPut>(id, buffer, TInstant::Max(), PutHandleClass);
-            PendingWriteSize += blobSize;
-            PendingWrites += 1;
+            InFlightTracker.Request(blobSize);
             return std::move(ev);
         }
 
@@ -164,7 +204,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
         TLogoBlobID GetRandomBlobId() {
-            Y_VERIFY(!IsEmpty());
+            Y_ABORT_UNLESS(!ConfirmedBlobs.empty());
             auto iter = ConfirmedBlobs.begin();
             std::advance(iter, RandomNumber(ConfirmedBlobs.size()));
             return *iter;
@@ -175,29 +215,21 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 << " PutHandleClass# " << NKikimrBlobStorage::EPutHandleClass_Name(PutHandleClass)
                 << " SizeToWrite# " << SizeToWrite
                 << " BlobsToWrite# " << BlobsToWrite
-                << " WritesInFlight# " << WritesInFlight
+                << " PendingWrites# " << InFlightTracker.ToString()
                 << " ConfirmedSize# " << ConfirmedDataSize
-                << " ConfirmedBlobs.size()# " << ConfirmedBlobs.size()
-                << " PendingWrites# " << PendingWrites
-                << " PendingWriteSize# " << PendingWriteSize << " }";
+                << " ConfirmedBlobs.size()# " << ConfirmedBlobs.size() << " }";
         }
 
-    ui32 GetWritesInFlight() {
-        return WritesInFlight;
-    }
-
     private:
         TSizeGenerator SizeGenerator;
         NKikimrBlobStorage::EPutHandleClass PutHandleClass = NKikimrBlobStorage::EPutHandleClass::UserData;
 
         uint64_t SizeToWrite;
         uint64_t BlobsToWrite;
-        ui32 WritesInFlight = 1;
         uint64_t ConfirmedDataSize = 0;
         TVector<TLogoBlobID> ConfirmedBlobs;
 
-        uint64_t PendingWriteSize = 0;
-        uint64_t PendingWrites = 0;
+        TInFlightTracker InFlightTracker;
 
         ui64 BlobCookie = 0;
     };
@@ -330,8 +362,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             bool LoadEnabled = true;
             std::optional<TSizeGenerator> SizeGen;
             std::shared_ptr<TRequestDelayManager> DelayManager;
-            const ui32 MaxRequestsInFlight;
-            const ui64 MaxBytesInFlight;
+            TInFlightTracker InFlightTracker;
             const ui64 MaxTotalBytes;
         };
 
@@ -360,8 +391,6 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         const NKikimrBlobStorage::EPutHandleClass PutHandleClass;
         TRequestDispatchingSettings WriteSettings;
         TMonotonic NextWriteTimestamp;
-        ui32 WritesInFlight = 0;
-        ui64 WriteBytesInFlight = 0;
         ui64 TotalBytesWritten = 0;
         THashMap<ui64, ui64> SentTimestamp;
         ui64 WriteQueryId = 0;
@@ -378,8 +407,6 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         const NKikimrBlobStorage::EGetHandleClass GetHandleClass;
         TRequestDispatchingSettings ReadSettings;
         TMonotonic NextReadTimestamp;
-        ui32 ReadsInFlight = 0;
-        ui64 ReadBytesInFlight = 0;
         ui64 TotalBytesRead = 0;
         THashMap<ui64, ui64> ReadSentTimestamp;
         ui64 ReadQueryId = 0;
@@ -582,13 +609,13 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
         void MakeInitialAllocation(const TActorContext& ctx) {
-            if (InitialAllocation.EnoughBlobsWritten(false)) {
+            if (!InitialAllocation) {
                 Self.InitialAllocationCompleted(ctx);
                 return;
             }
             LOG_DEBUG_S(ctx, NKikimrServices::BS_LOAD_TEST, PrintMe() << " going to make initial allocation,"
                     << InitialAllocation.ToString());
-            for (ui32 i = 0; i < InitialAllocation.GetWritesInFlight() && !InitialAllocation.EnoughBlobsWritten(true); ++i) {
+            while (InitialAllocation.CanSendRequest()) {
                 IssueInitialPut(ctx);
             }
         }
@@ -601,9 +628,10 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 Y_ABORT_UNLESS(res);
 
                 InitialAllocation.ConfirmBlob(res->Id, CheckStatus(ctx, res, {}));
-                if (!InitialAllocation.EnoughBlobsWritten(true)) {
+                while (InitialAllocation.CanSendRequest()) {
                     IssueInitialPut(ctx);
-                } else if (InitialAllocation.EnoughBlobsWritten(false)) {
+                }
+                if (InitialAllocation.EnoughBlobsWritten()) {
                     SetKeepFlagsOnInitialAllocation(ctx);
                 }
             };
@@ -672,10 +700,10 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             if (ReadMegabytesPerSecondST.CalculateSpeed(&speed)) {
                 ReadMegabytesPerSecondQT.Add(now, speed);
             }
-            WritesInFlightQT.Add(now, WritesInFlight);
-            WriteBytesInFlightQT.Add(now, WriteBytesInFlight);
-            ReadsInFlightQT.Add(now, ReadsInFlight);
-            ReadBytesInFlightQT.Add(now, ReadBytesInFlight);
+            WritesInFlightQT.Add(now, WriteSettings.InFlightTracker.RequestsInFlight);
+            WriteBytesInFlightQT.Add(now, WriteSettings.InFlightTracker.BytesInFlight);
+            ReadsInFlightQT.Add(now, ReadSettings.InFlightTracker.RequestsInFlight);
+            ReadBytesInFlightQT.Add(now, ReadSettings.InFlightTracker.BytesInFlight);
             if (now > LastLatencyTrackerUpdate + TDuration::Seconds(1)) {
                 LastLatencyTrackerUpdate = now;
                 ResponseQT->Update();
@@ -736,18 +764,12 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                     DUMP_PARAM(EarliestTimestamp)
                 }
                 DUMP_PARAM(NextWriteTimestamp)
-                DUMP_PARAM(WritesInFlight)
-                DUMP_PARAM(WriteBytesInFlight)
-                DUMP_PARAM_FINAL(WriteSettings.MaxRequestsInFlight)
-                DUMP_PARAM_FINAL(WriteSettings.MaxBytesInFlight)
+                DUMP_PARAM(WriteSettings.InFlightTracker.ToString())
                 DUMP_PARAM_FINAL(TotalBytesWritten)
                 DUMP_PARAM_FINAL(WriteSettings.MaxTotalBytes)
                 DUMP_PARAM_FINAL(TotalBytesRead)
                 DUMP_PARAM(NextReadTimestamp)
-                DUMP_PARAM(ReadsInFlight)
-                DUMP_PARAM(ReadBytesInFlight)
-                DUMP_PARAM_FINAL(ReadSettings.MaxRequestsInFlight)
-                DUMP_PARAM_FINAL(ReadSettings.MaxBytesInFlight)
+                DUMP_PARAM(ReadSettings.InFlightTracker.ToString())
                 DUMP_PARAM(ConfirmedBlobIds.size())
                 DUMP_PARAM(InitialAllocation.ToString())
 
@@ -807,10 +829,8 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
 
         void IssueWriteIfPossible(const TActorContext& ctx) {
             const TMonotonic now = TActivationContext::Monotonic();
-            while (WriteSettings.LoadEnabled &&
-                    (WritesInFlight < WriteSettings.MaxRequestsInFlight || !WriteSettings.MaxRequestsInFlight) &&
-                    (WriteBytesInFlight < WriteSettings.MaxBytesInFlight || !WriteSettings.MaxBytesInFlight) &&
-                    (TotalBytesWritten + WriteBytesInFlight < WriteSettings.MaxTotalBytes || !WriteSettings.MaxTotalBytes) &&
+            while (WriteSettings.LoadEnabled && WriteSettings.InFlightTracker &&
+                    (TotalBytesWritten + WriteSettings.InFlightTracker.BytesInFlight < WriteSettings.MaxTotalBytes || !WriteSettings.MaxTotalBytes) &&
                     now >= NextWriteTimestamp &&
                     (!ScriptedRequests || ScriptedRequests[ScriptedCounter].EvType == TEvBlobStorage::EvPut)) {
                 IssueWriteRequest(ctx);
@@ -858,9 +878,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                     ConfirmedBlobIds.insert(std::lower_bound(ConfirmedBlobIds.begin(), ConfirmedBlobIds.end(), id), id);
                 }
 
-                Y_ABORT_UNLESS(WritesInFlight >= 1 && WriteBytesInFlight >= size);
-                --WritesInFlight;
-                WriteBytesInFlight -= size;
+                WriteSettings.InFlightTracker.Response(size);
 
                 TotalBytesWritten += size;
 
@@ -878,7 +896,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 ResponseQT->Increment(response.MicroSeconds());
                 IssueWriteIfPossible(ctx);
 
-                if (ConfirmedBlobIds.size() == 1 && InitialAllocation.IsEmpty()) {
+                if (ConfirmedBlobIds.size() == 1 && !InitialAllocation) {
                     if (NextReadTimestamp == TMonotonic()) {
                         NextReadTimestamp = TActivationContext::Monotonic();
                     }
@@ -896,8 +914,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
 
             ++Cookie;
 
-            ++WritesInFlight;
-            WriteBytesInFlight += size;
+            WriteSettings.InFlightTracker.Request(size);
 
             if (ScriptedRequests) {
                 UpdateNextTimestemps(true);
@@ -968,8 +985,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         void IssueReadIfPossible(const TActorContext& ctx) {
             const TMonotonic now = TActivationContext::Monotonic();
 
-            while (ReadSettings.LoadEnabled && ReadsInFlight < ReadSettings.MaxRequestsInFlight &&
-                    (ReadBytesInFlight < ReadSettings.MaxBytesInFlight || !ReadSettings.MaxBytesInFlight) &&
+            while (ReadSettings.LoadEnabled && ReadSettings.InFlightTracker &&
                     now >= NextReadTimestamp &&
                     ConfirmedBlobIds.size() + InitialAllocation.Size() > 0 &&
                     (!ScriptedRequests || ScriptedRequests[ScriptedCounter].EvType == TEvBlobStorage::EvGet)) {
@@ -1020,9 +1036,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                     return;
                 }
 
-                Y_ABORT_UNLESS(ReadsInFlight >= 1 && ReadBytesInFlight >= size);
-                --ReadsInFlight;
-                ReadBytesInFlight -= size;
+                ReadSettings.InFlightTracker.Response(size);
                 TotalBytesRead += size;
 
                 auto it = ReadSentTimestamp.find(readQueryId);
@@ -1037,8 +1051,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             SendToBSProxy(ctx, GroupId, ev.release(), Self.QueryDispatcher.ObtainCookie(std::move(readCallback)));
             ReadSentTimestamp.emplace(readQueryId, GetCycleCountFast());
 
-            ++ReadsInFlight;
-            ReadBytesInFlight += size;
+            ReadSettings.InFlightTracker.Request(size);
 
             // calculate time of next read request
             if (ScriptedRequests) {
@@ -1050,6 +1063,12 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
     };
 
+    enum EWakeupType : ui32 {
+        MAIN_CYCLE = 0,
+        DELAY_AFTER_INITIAL_WRITE,
+    };
+
+
     TString ConfingString;
     const ui64 Tag;
     const TActorId Parent;
@@ -1077,6 +1096,8 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
     static constexpr ui64 DefaultTabletId = 5000;
     ui32 WorkersInInitialState = 0;
 
+    ui32 DelayAfterInitialWrite = 0;
+
 public:
     static constexpr NKikimrServices::TActivity::EType ActorActivityType() {
         return NKikimrServices::TActivity::BS_LOAD_ACTOR;
@@ -1103,7 +1124,11 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
 
             TInitialAllocation initialAllocation;
             if (profile.HasInitialAllocation()) {
+                auto initialAllocationProto = profile.GetInitialAllocation();
                 initialAllocation = TInitialAllocation(profile.GetInitialAllocation());
+                if (initialAllocationProto.HasDelayAfterCompletionSec()) {
+                    DelayAfterInitialWrite = std::max(DelayAfterInitialWrite, initialAllocationProto.GetDelayAfterCompletionSec());
+                }
             }
 
             NKikimrBlobStorage::EPutHandleClass putHandleClass = NKikimrBlobStorage::EPutHandleClass::UserData;
@@ -1126,8 +1151,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 .LoadEnabled = enableWrites,
                 .SizeGen = TSizeGenerator(profile.GetWriteSizes()),
                 .DelayManager = std::move(writeDelayManager),
-                .MaxRequestsInFlight = profile.GetMaxInFlightWriteRequests(),
-                .MaxBytesInFlight = profile.GetMaxInFlightWriteBytes(),
+                .InFlightTracker = TInFlightTracker(profile.GetMaxInFlightWriteRequests(), profile.GetMaxInFlightWriteBytes()),
                 .MaxTotalBytes = profile.GetMaxTotalBytesWritten(),
             };
 
@@ -1156,8 +1180,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 .LoadEnabled = enableReads,
                 .SizeGen = readSizeGen,
                 .DelayManager = std::move(readDelayManager),
-                .MaxRequestsInFlight = profile.GetMaxInFlightReadRequests(),
-                .MaxBytesInFlight = profile.GetMaxInFlightReadBytes(),
+                .InFlightTracker = TInFlightTracker(profile.GetMaxInFlightReadRequests(), profile.GetMaxInFlightReadBytes()),
                 .MaxTotalBytes = ::Max<ui64>(),
             };
 
@@ -1215,14 +1238,22 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
     }
 
+    void StartWorkers(const TActorContext& ctx) {
+        if (TestDuration) {
+            ctx.Schedule(*TestDuration, new TEvents::TEvPoisonPill());
+        }
+        for (auto& writer : TabletWriters) {
+            writer->StartWorking(ctx);
+        }
+    }
+
     void InitialAllocationCompleted(const TActorContext& ctx) {
         WorkersInInitialState--;
-        if (!WorkersInInitialState) {
-            if (TestDuration) {
-                ctx.Schedule(*TestDuration, new TEvents::TEvPoisonPill());
-            }
-            for (auto& writer : TabletWriters) {
-                writer->StartWorking(ctx);
+        if (!WorkersInInitialState)  {
+            if (DelayAfterInitialWrite) {
+                ctx.Schedule(TDuration::Seconds(DelayAfterInitialWrite), new TEvents::TEvWakeup);
+            } else {
+                StartWorkers(ctx);
             }
         }
     }
@@ -1282,9 +1313,20 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         ctx.Schedule(TDuration::MilliSeconds(5), new TEvUpdateQuantile);
     }
 
-    void HandleWakeup(const TActorContext& ctx) {
-        --WakeupsScheduled;
-        UpdateWakeupQueue(ctx);
+    void HandleWakeup(TEvents::TEvWakeup::TPtr& ev, const TActorContext& ctx) {
+        switch (ev->Get()->Tag) {
+        case MAIN_CYCLE:
+            --WakeupsScheduled;
+            UpdateWakeupQueue(ctx);
+            break;
+
+        case DELAY_AFTER_INITIAL_WRITE:
+            StartWorkers(ctx);
+            break;
+
+        default:
+            Y_FAIL_S("Unexpected wakeup tag# " << ev->Get()->Tag);
+        }
     }
 
     void UpdateWakeupQueue(const TActorContext& ctx) {
@@ -1390,7 +1432,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
     STRICT_STFUNC(StateFunc,
         CFunc(EvStopTest, HandleStopTest);
         CFunc(EvUpdateQuantile, HandleUpdateQuantile);
-        CFunc(TEvents::TSystem::Wakeup, HandleWakeup);
+        HFunc(TEvents::TEvWakeup, HandleWakeup);
         CFunc(TEvents::TSystem::PoisonPill, HandlePoison);
         HFunc(TEvBlobStorage::TEvDiscoverResult, HandleDispatcher);
         HFunc(TEvBlobStorage::TEvBlockResult, HandleDispatcher);
diff --git a/ydb/core/protos/load_test.proto b/ydb/core/protos/load_test.proto
index a487971b6e7c..1f70c224b3b2 100644
--- a/ydb/core/protos/load_test.proto
+++ b/ydb/core/protos/load_test.proto
@@ -38,8 +38,10 @@ message TEvLoadTestRequest {
                 uint64 BlobsNumber = 2;
             }
             repeated TSizeInfo BlobSizes = 3;
-            optional uint32 WritesInFlight = 4 [default = 1];
-            optional NKikimrBlobStorage.EPutHandleClass PutHandleClass = 5;
+            optional uint32 MaxWritesInFlight = 4;
+            optional uint32 MaxWriteBytesInFlight = 5;
+            optional NKikimrBlobStorage.EPutHandleClass PutHandleClass = 6;
+            optional uint32 DelayAfterCompletionSec = 7;
         }
 
         message TRequestInfo {

From 2081b7bf27363413a9429dc66bb9b916f1432900 Mon Sep 17 00:00:00 2001
From: Sergey Belyakov <serg-belyakov@ydb.tech>
Date: Fri, 29 Dec 2023 16:32:27 +0000
Subject: [PATCH 3/4] Fix time in UI

---
 ydb/core/load_test/group_write.cpp | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/ydb/core/load_test/group_write.cpp b/ydb/core/load_test/group_write.cpp
index e35208ed2c8a..a1f931883700 100644
--- a/ydb/core/load_test/group_write.cpp
+++ b/ydb/core/load_test/group_write.cpp
@@ -648,7 +648,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 Y_ABORT_UNLESS(res);
                 
                 if (!MainCycleStarted) {
-                    StartWorking(ctx);
+                    Self.InitialAllocationCompleted(ctx);
                 }
             };
 
@@ -1245,13 +1245,16 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         for (auto& writer : TabletWriters) {
             writer->StartWorking(ctx);
         }
+        TestStartTime = TActivationContext::Monotonic();
+        UpdateWakeupQueue(ctx);
     }
 
     void InitialAllocationCompleted(const TActorContext& ctx) {
         WorkersInInitialState--;
-        if (!WorkersInInitialState)  {
+        if (!WorkersInInitialState) {
             if (DelayAfterInitialWrite) {
-                ctx.Schedule(TDuration::Seconds(DelayAfterInitialWrite), new TEvents::TEvWakeup);
+                ctx.Schedule(TDuration::Seconds(DelayAfterInitialWrite),
+                        new TEvents::TEvWakeup(DELAY_AFTER_INITIAL_WRITE));
             } else {
                 StartWorkers(ctx);
             }
@@ -1261,11 +1264,9 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
     void Bootstrap(const TActorContext& ctx) {
         Become(&TLogWriterLoadTestActor::StateFunc);
         EarlyStop = false;
-        TestStartTime = TActivationContext::Monotonic();
         for (auto& writer : TabletWriters) {
             writer->Bootstrap(ctx);
         }
-        UpdateWakeupQueue(ctx);
         HandleUpdateQuantile(ctx);
     }
 
@@ -1392,7 +1393,12 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                             str << "Passed/Total, sec";
                         }
                         TABLED() {
-                            str << (TActivationContext::Monotonic() - TestStartTime).Seconds() << " / ";
+                            if (TestStartTime != TMonotonic()) {
+                                str << (TActivationContext::Monotonic() - TestStartTime).Seconds();
+                            } else {
+                                str << "0";
+                            }
+                            str << " / ";
                             if (TestDuration.Defined()) {
                                 str << TestDuration->Seconds();
                             } else {

From 817ecac990641ad1d71a9235f1f0fcf6ea09e839 Mon Sep 17 00:00:00 2001
From: Sergey Belyakov <serg-belyakov@ydb.tech>
Date: Wed, 10 Jan 2024 13:30:22 +0000
Subject: [PATCH 4/4] Adress comments

---
 ydb/core/load_test/group_write.cpp | 35 +++++++++++++++---------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/ydb/core/load_test/group_write.cpp b/ydb/core/load_test/group_write.cpp
index a1f931883700..a57e1584bb9c 100644
--- a/ydb/core/load_test/group_write.cpp
+++ b/ydb/core/load_test/group_write.cpp
@@ -95,9 +95,9 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             , BytesInFlight(0)
         {}
 
-        operator bool() const {
-            return (!MaxRequestsInFlight || RequestsInFlight < MaxRequestsInFlight) &&
-                    (!MaxBytesInFlight || BytesInFlight < MaxBytesInFlight);
+        bool LimitReached() const {
+            return (MaxRequestsInFlight && RequestsInFlight >= MaxRequestsInFlight) ||
+                    (MaxBytesInFlight && BytesInFlight >= MaxBytesInFlight);
         }
 
         void Request(ui64 size) {
@@ -106,7 +106,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
         void Response(ui64 size) {
-            Y_DEBUG_ABORT_UNLESS(BytesInFlight >= size && RequestsInFlight >= 0);
+            Y_DEBUG_ABORT_UNLESS(BytesInFlight >= size && RequestsInFlight > 0);
             BytesInFlight -= size;
             --RequestsInFlight;
         }
@@ -142,7 +142,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
     public:
-        bool Size() {
+        bool ConfirmedSize() {
             return ConfirmedBlobs.size();
         }
 
@@ -150,12 +150,12 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
             return ConfirmedBlobs[idx];
         }
 
-        operator bool() const {
-            return SizeToWrite > 0 || BlobsToWrite > 0;
+        bool IsEmpty() const {
+            return SizeToWrite == 0 && BlobsToWrite == 0;
         }
 
         bool CanSendRequest() {
-            if (!InFlightTracker) {
+            if (InFlightTracker.LimitReached()) {
                 return false;
             }
             return !EnoughBlobsWritten(true);
@@ -205,9 +205,8 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
 
         TLogoBlobID GetRandomBlobId() {
             Y_ABORT_UNLESS(!ConfirmedBlobs.empty());
-            auto iter = ConfirmedBlobs.begin();
-            std::advance(iter, RandomNumber(ConfirmedBlobs.size()));
-            return *iter;
+            auto idx = RandomNumber(ConfirmedBlobs.size());
+            return ConfirmedBlobs[idx];
         }
 
         TString ToString() {
@@ -609,7 +608,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         }
 
         void MakeInitialAllocation(const TActorContext& ctx) {
-            if (!InitialAllocation) {
+            if (InitialAllocation.IsEmpty()) {
                 Self.InitialAllocationCompleted(ctx);
                 return;
             }
@@ -829,7 +828,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
 
         void IssueWriteIfPossible(const TActorContext& ctx) {
             const TMonotonic now = TActivationContext::Monotonic();
-            while (WriteSettings.LoadEnabled && WriteSettings.InFlightTracker &&
+            while (WriteSettings.LoadEnabled && !WriteSettings.InFlightTracker.LimitReached() &&
                     (TotalBytesWritten + WriteSettings.InFlightTracker.BytesInFlight < WriteSettings.MaxTotalBytes || !WriteSettings.MaxTotalBytes) &&
                     now >= NextWriteTimestamp &&
                     (!ScriptedRequests || ScriptedRequests[ScriptedCounter].EvType == TEvBlobStorage::EvPut)) {
@@ -896,7 +895,7 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
                 ResponseQT->Increment(response.MicroSeconds());
                 IssueWriteIfPossible(ctx);
 
-                if (ConfirmedBlobIds.size() == 1 && !InitialAllocation) {
+                if (ConfirmedBlobIds.size() == 1 && !InitialAllocation.IsEmpty()) {
                     if (NextReadTimestamp == TMonotonic()) {
                         NextReadTimestamp = TActivationContext::Monotonic();
                     }
@@ -985,9 +984,9 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         void IssueReadIfPossible(const TActorContext& ctx) {
             const TMonotonic now = TActivationContext::Monotonic();
 
-            while (ReadSettings.LoadEnabled && ReadSettings.InFlightTracker &&
+            while (ReadSettings.LoadEnabled && !ReadSettings.InFlightTracker.LimitReached() &&
                     now >= NextReadTimestamp &&
-                    ConfirmedBlobIds.size() + InitialAllocation.Size() > 0 &&
+                    ConfirmedBlobIds.size() + InitialAllocation.ConfirmedSize() > 0 &&
                     (!ScriptedRequests || ScriptedRequests[ScriptedCounter].EvType == TEvBlobStorage::EvGet)) {
                 IssueReadRequest(ctx);
             }
@@ -1001,8 +1000,8 @@ class TLogWriterLoadTestActor : public TActorBootstrapped<TLogWriterLoadTestActo
         void IssueReadRequest(const TActorContext& ctx) {
             TLogoBlobID id;
             ui32 confirmedBlobs = ConfirmedBlobIds.size();
-            ui32 initialBlobs = InitialAllocation.Size();
-            Y_VERIFY(confirmedBlobs + initialBlobs > 0);
+            ui32 initialBlobs = InitialAllocation.ConfirmedSize();
+            Y_ABORT_UNLESS(confirmedBlobs + initialBlobs > 0);
             ui32 blobIdx = RandomNumber(confirmedBlobs + initialBlobs);
     
             if (blobIdx < confirmedBlobs) {