LOGBROKER-9686 Add limits to in flight bytes per partition (ydb-platform#33024)

Co-authored-by: Roman Kuzin <[email protected]>
Co-authored-by: Bulat <[email protected]>

Author: 3 people

ydb/core/persqueue/public/inflight_limiter.cpp

@@ -0,0 +1,77 @@
+#include "inflight_limiter.h"
+#include <ydb/library/actors/core/log.h>
+
+namespace NKikimr::NPQ {
+
+TInFlightController::TInFlightController(ui64 MaxAllowedSize)
+    : LayoutUnitSize(MaxAllowedSize / MAX_LAYOUT_COUNT)
+    , TotalSize(0)
+    , MaxAllowedSize(MaxAllowedSize)
+{
+    if (MaxAllowedSize > 0 && LayoutUnitSize == 0) {
+        LayoutUnitSize = 1;
+    }
+}
+
+bool TInFlightController::Add(ui64 Offset, ui64 Size) {
+    if (MaxAllowedSize == 0) {
+        // means that there are no limits were set
+        return true;
+    }
+
+    if (Size == 0) {
+        return TotalSize < MaxAllowedSize;
+    }
+
+    AFL_ENSURE(Layout.empty() || Offset > Layout.back());
+
+    if (TotalSize % LayoutUnitSize != 0) {
+        AFL_ENSURE(!Layout.empty());
+        Layout.back() = Offset;
+    }
+
+    auto unitsBefore = (TotalSize + LayoutUnitSize - 1) / LayoutUnitSize;
+    TotalSize += Size;
+    auto unitsAfter = std::min(MAX_LAYOUT_COUNT, (TotalSize + LayoutUnitSize - 1) / LayoutUnitSize);
+    for (auto currentUnits = unitsBefore; currentUnits < unitsAfter; currentUnits++) {
+        Layout.push_back(Offset);
+    }
+    
+    AFL_ENSURE(!Layout.empty());
+    Layout.back() = Offset;
+
+    return TotalSize < MaxAllowedSize;
+}
+
+bool TInFlightController::Remove(ui64 Offset) {
+    if (MaxAllowedSize == 0) {
+        // means that there are no limits were set
+        return true;
+    }
+
+    for (auto it = Layout.begin(); it != Layout.end(); it = Layout.erase(it)) {
+        if (*it >= Offset) {
+            break;
+        }
+
+        if (Layout.size() == 1) {
+            TotalSize = 0;
+        } else {
+            AFL_ENSURE(TotalSize >= LayoutUnitSize);
+            TotalSize -= LayoutUnitSize;
+        }
+    }
+
+    return TotalSize < MaxAllowedSize;
+}
+
+bool TInFlightController::IsMemoryLimitReached() const {
+    if (MaxAllowedSize == 0) {
+        // means that there are no limits were set
+        return false;
+    }
+
+    return TotalSize >= MaxAllowedSize;
+}
+
+} // namespace NKikimr::NPQ

ydb/core/persqueue/public/inflight_limiter.h

@@ -0,0 +1,38 @@
+#include <util/system/types.h>
+#include <deque>
+
+namespace NKikimr::NPQ {
+
+// This class is used to control in-flight data.
+// Contoller handles layout of data units, max units count is MAX_LAYOUT_COUNT constant.
+// Layout is a deque of offsets
+// Each item (associated with data unit) in the deque is max offset which intersects with this unit
+// For example, if we have 3 units:
+// Unit 1: [0, 100]
+// Unit 2: [100, 200]
+// Unit 3: [200, 300]
+// Offsets are:
+// 1 - Size 60
+// 2 - Size 60
+// 3 - Size 60
+// 4 - Size 80
+// Then the deque will be [2, 4, 4]
+struct TInFlightController {
+    constexpr static ui64 MAX_LAYOUT_COUNT = 1024;
+
+    TInFlightController() = default;
+    TInFlightController(ui64 MaxAllowedSize);
+    
+    ui64 LayoutUnitSize = 0;
+    std::deque<ui64> Layout;
+    ui64 TotalSize = 0;
+    ui64 MaxAllowedSize = 0;
+
+    // Adds an offset with size
+    bool Add(ui64 Offset, ui64 Size);
+    // Removes offsets <= given offset
+    bool Remove(ui64 Offset);
+    bool IsMemoryLimitReached() const;
+};
+
+} // namespace NKikimr::NPQ

ydb/core/persqueue/public/ya.make

@@ -2,6 +2,7 @@ LIBRARY()
 
 SRCS(
     config.cpp
+    inflight_limiter.cpp
     pq_database.cpp
     pq_rl_helpers.cpp
     utils.cpp