Reject operations in column shards by probability (ydb-platform#26783)

Author: neyrox

ydb/core/protos/feature_flags.proto

@@ -233,5 +233,6 @@ message TFeatureFlags {
     optional bool EnableTinyDisks = 207 [default = false];
     optional bool EnableMetricsLevel = 208 [default = false];
     optional bool EnableSecureScriptExecutions = 209 [default = false];
+    optional bool EnableOlapRejectProbability = 213 [default = false];
     optional bool DisableMissingDefaultColumnsInBulkUpsert = 215 [default = false];
 }

ydb/core/protos/memory_controller_config.proto

@@ -28,4 +28,6 @@ message TMemoryControllerConfig {
 
     optional float CompactionLimitPercent = 130 [default = 10];
     optional uint64 CompactionLimitBytes = 131;
+
+    optional uint64 MaxTxInFly = 132 [default = 10000];
 }

ydb/core/tablet_flat/flat_executor.cpp

@@ -31,6 +31,7 @@
 #include <ydb/core/base/hive.h>
 #include <ydb/core/base/tablet_pipecache.h>
 #include <ydb/core/control/lib/immediate_control_board_impl.h>
+#include <ydb/core/protos/memory_controller_config.pb.h>
 #include <ydb/core/scheme/scheme_type_registry.h>
 #include <ydb/core/tablet/tablet_counters_aggregator.h>
 #include <ydb/library/actors/core/hfunc.h>
@@ -45,10 +46,10 @@
 #include <util/generic/xrange.h>
 #include <util/generic/ymath.h>
 
+
 namespace NKikimr {
 namespace NTabletFlatExecutor {
 
-static constexpr ui64 MaxTxInFly = 10000;
 
 LWTRACE_USING(TABLET_FLAT_PROVIDER)
 
@@ -132,6 +133,7 @@ TExecutor::TExecutor(
     , Time(TAppData::TimeProvider)
     , Owner(owner)
     , OwnerActorId(ownerActorId)
+    , MaxTxInFly(AppData()->MemoryControllerConfig.GetMaxTxInFly())
     , Emitter(new TIdEmitter)
     , CounterEventsInFlight(new TEvTabletCounters::TInFlightCookie)
     , Stats(new TExecutorStatsImpl())
@@ -140,7 +142,6 @@ TExecutor::TExecutor(
 {}
 
 TExecutor::~TExecutor() {
-
 }
 
 bool TExecutor::OnUnhandledException(const std::exception& e) {

ydb/core/tablet_flat/flat_executor.h

@@ -385,6 +385,8 @@ class TExecutor
     // Counts the number of times LeaseDuration was increased
     size_t LeaseDurationIncreases = 0;
 
+    const ui64 MaxTxInFly;
+
     struct TLeaseCommit : public TIntrusiveListItem<TLeaseCommit> {
         using TByEndMap = std::multimap<TMonotonic, TLeaseCommit*>;
 

ydb/core/tx/columnshard/columnshard__overload.cpp

@@ -24,6 +24,18 @@ TColumnShard::EOverloadStatus TColumnShard::CheckOverloadedImmediate(const TInte
             "limit", txLimit);
         return EOverloadStatus::ShardTxInFly;
     }
+
+    if (AppData()->FeatureFlags.GetEnableOlapRejectProbability()) {
+        const float rejectProbabilty = Executor()->GetRejectProbability();
+        if (rejectProbabilty > 0) {
+            const float rnd = TAppData::RandomProvider->GenRandReal2();
+            if (rnd < rejectProbabilty) {
+                AFL_WARN(NKikimrServices::TX_COLUMNSHARD_WRITE)("event", "shard_overload")("reason", "reject_probality")("RP", rejectProbabilty);
+                return EOverloadStatus::RejectProbability;
+            }
+        }
+    }
+
     return EOverloadStatus::None;
 }
 

ydb/core/tx/columnshard/columnshard__write.cpp

@@ -46,6 +46,9 @@ void TColumnShard::OverloadWriteFail(const EOverloadStatus overloadReason, const
         case EOverloadStatus::ShardWritesSizeInFly:
             Counters.OnWriteOverloadShardWritesSize(writeSize);
             break;
+        case EOverloadStatus::RejectProbability:
+            Counters.OnWriteOverloadRejectProbability(writeSize);
+            break;
         case EOverloadStatus::None:
             Y_ABORT("invalid function usage");
     }
@@ -345,7 +348,6 @@ void TColumnShard::Handle(NEvents::TDataEvents::TEvWrite::TPtr& ev, const TActor
     const auto source = ev->Sender;
     const auto cookie = ev->Cookie;
 
-
     std::optional<TDuration> writeTimeout;
     if (record.HasTimeoutSeconds()) {
         writeTimeout = TDuration::Seconds(record.GetTimeoutSeconds());

ydb/core/tx/columnshard/counters/columnshard.cpp

@@ -43,6 +43,8 @@ TCSCounters::TCSCounters()
     OverloadShardWritesCount = TBase::GetDeriviative("Overload/Shard/Writes/Count");
     OverloadShardWritesSizeBytes = TBase::GetDeriviative("Overload/Shard/WritesSize/Bytes");
     OverloadShardWritesSizeCount = TBase::GetDeriviative("Overload/Shard/WritesSize/Count");
+    OverloadRejectProbabilityBytes = TBase::GetDeriviative("Overload/RejectProbability/Bytes");
+    OverloadRejectProbabilityCount = TBase::GetDeriviative("Overload/RejectProbability/Count");
 
     InternalCompactionGranuleBytes = TBase::GetValueAutoAggregationsClient("InternalCompaction/Bytes");
     InternalCompactionGranulePortionsCount = TBase::GetValueAutoAggregationsClient("InternalCompaction/PortionsCount");

ydb/core/tx/columnshard/counters/columnshard.h

@@ -18,7 +18,8 @@ enum class EOverloadStatus {
     OverloadMetadata /* "overload_metadata" */,
     Disk /* "disk" */,
     None /* "none" */,
-    OverloadCompaction /* "overload_compaction" */
+    OverloadCompaction /* "overload_compaction" */,
+    RejectProbability,
 };
 
 enum class EWriteFailReason {
@@ -106,6 +107,8 @@ class TCSCounters: public TCommonCountersOwner {
     NMonitoring::TDynamicCounters::TCounterPtr OverloadShardWritesCount;
     NMonitoring::TDynamicCounters::TCounterPtr OverloadShardWritesSizeBytes;
     NMonitoring::TDynamicCounters::TCounterPtr OverloadShardWritesSizeCount;
+    NMonitoring::TDynamicCounters::TCounterPtr OverloadRejectProbabilityBytes;
+    NMonitoring::TDynamicCounters::TCounterPtr OverloadRejectProbabilityCount;
 
     std::shared_ptr<TValueAggregationClient> InternalCompactionGranuleBytes;
     std::shared_ptr<TValueAggregationClient> InternalCompactionGranulePortionsCount;
@@ -222,6 +225,11 @@ class TCSCounters: public TCommonCountersOwner {
         OverloadShardWritesSizeCount->Add(1);
     }
 
+    void OnWriteOverloadRejectProbability(const ui64 size) const {
+        OverloadRejectProbabilityBytes->Add(size);
+        OverloadRejectProbabilityCount->Add(1);
+    }
+
     void SkipIndexationInputDueToSplitCompaction(const ui64 size) const {
         SkipIndexationInputDueToSplitCompactionBytes->Add(size);
         SkipIndexationInputDueToSplitCompactionCount->Add(1);

ydb/core/tx/columnshard/counters/counters_manager.h

@@ -79,6 +79,11 @@ class TCountersManager {
         CSCounters.OnWriteOverloadShardWritesSize(size);
     }
 
+    void OnWriteOverloadRejectProbability(const ui64 size) const {
+        TabletCounters->IncCounter(COUNTER_WRITE_OVERLOAD);
+        CSCounters.OnWriteOverloadRejectProbability(size);
+    }
+
     void FillTableStats(TInternalPathId pathId, ::NKikimrTableStats::TTableStats& tableStats) {
         ColumnTablesCounters->GetPathIdCounter(pathId)->FillStats(tableStats);
         BackgroundControllerCounters->FillStats(pathId, tableStats);

ydb/tests/olap/test_overloads.py

@@ -4,6 +4,7 @@
 import logging
 import yatest.common
 import ydb
+import random
 
 from ydb.tests.library.harness.kikimr_config import KikimrConfigGenerator
 from ydb.tests.library.harness.kikimr_runner import KiKiMR
@@ -103,6 +104,20 @@ def _setup_ydb(cls, writing_in_flight_requests_count_limit, writing_in_flight_re
         cls.ydb_client = YdbClient(endpoint=f"grpc://{node.host}:{node.port}", database=f"/{config.domain_name}")
         cls.ydb_client.wait_connection(timeout=60)
 
+    @classmethod
+    def _setup_ydb_rp(cls):
+        ydb_path = yatest.common.build_path(os.environ.get("YDB_DRIVER_BINARY"))
+        logger.info(yatest.common.execute([ydb_path, "-V"], wait=True).stdout.decode("utf-8"))
+        config = KikimrConfigGenerator(
+            extra_feature_flags={"enable_olap_reject_probability": True},
+            memory_controller_config={"max_tx_in_fly": 0},
+        )
+        cls.cluster = KiKiMR(config)
+        cls.cluster.start()
+        node = cls.cluster.nodes[1]
+        cls.ydb_client = YdbClient(endpoint=f"grpc://{node.host}:{node.port}", database=f"/{config.domain_name}")
+        cls.ydb_client.wait_connection(timeout=60)
+
     def get_row_count(self) -> int:
         return self.ydb_client.query(f"select count(*) as Rows from `{self.table_name}`")[0].rows[0]["Rows"]
 
@@ -156,6 +171,7 @@ def test_overloads_bulk_upsert(self, writing_in_flight_requests_count_limit, wri
     @pytest.mark.parametrize('writing_in_flight_requests_count_limit, writing_in_flight_request_bytes_limit', [(1, 10000), (2, 10000), (1000, 1), (1000, 2), (1, 1), (2, 2)])
     def test_overloads_workload(self, writing_in_flight_requests_count_limit, writing_in_flight_request_bytes_limit):
         self._setup_ydb(writing_in_flight_requests_count_limit, writing_in_flight_request_bytes_limit)
+
         wait_time: int = 60
         self.table_name: str = f"log_{writing_in_flight_requests_count_limit}_{writing_in_flight_request_bytes_limit}"
 
@@ -191,3 +207,55 @@ def test_overloads_workload(self, writing_in_flight_requests_count_limit, writin
 
         logging.info(f"Count rows after insert {self.get_row_count()}")
         assert self.get_row_count() != 0
+
+    def test_overloads_reject_probability(self):
+        self._setup_ydb_rp()
+
+        table_path = f"{self.ydb_client.database}/table_for_test_overloads_reject_probability"
+        self.ydb_client.query(
+            f"""
+            CREATE TABLE `{table_path}` (
+                id Uint64 NOT NULL,
+                v1 Int64,
+                v2 Int64,
+                PRIMARY KEY(id),
+            )
+            WITH (
+                STORE = COLUMN,
+                PARTITION_COUNT = 1,
+                AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = 1
+            )
+            """
+        )
+
+        column_types = ydb.BulkUpsertColumns()
+        column_types.add_column("id", ydb.PrimitiveType.Uint64)
+        column_types.add_column("v1", ydb.PrimitiveType.Int64)
+        column_types.add_column("v2", ydb.PrimitiveType.Int64)
+
+        rows_count = 1000
+
+        data = [
+            {
+                "id": i,
+                "v1": 1,
+                "v2": -1,
+            }
+            for i in range(rows_count)
+        ]
+
+        self.ydb_client.bulk_upsert(table_path, column_types, data)
+
+        futures = []
+
+        for _ in range(19):
+            lb = random.randint(0, rows_count)
+            futures.append(self.ydb_client.query_async(f"UPDATE `{table_path}` SET v1 = v1 + 1, v2 = v2 - 1 WHERE id > {lb};"))
+
+        for future in futures:
+            future.result()
+
+        monitor = self.cluster.monitors[0].fetch()
+        _, rejectProbabilityCount = monitor.get_by_name('Deriviative/Overload/RejectProbability/Count')[0]
+
+        assert rejectProbabilityCount > 0