From 181dd8157d71f0d925a24b174d9f8f25bc74e4fa Mon Sep 17 00:00:00 2001 From: Junbo Wang Date: Thu, 6 Aug 2026 15:56:02 +0800 Subject: [PATCH] [rust] Support cooperative KV backpressure (client side of #3463) --- fluss-rust/bindings/cpp/include/fluss.hpp | 7 +- fluss-rust/bindings/cpp/src/ffi_converter.hpp | 2 + fluss-rust/bindings/cpp/src/lib.rs | 2 + .../bindings/cpp/test/test_ffi_converter.cpp | 11 ++ .../bindings/elixir/lib/fluss/config.ex | 7 ++ fluss-rust/bindings/elixir/lib/fluss/error.ex | 6 +- .../elixir/native/fluss_nif/src/atoms.rs | 2 + .../elixir/native/fluss_nif/src/config.rs | 4 + .../bindings/elixir/test/config_test.exs | 8 ++ .../bindings/elixir/test/error_test.exs | 3 +- fluss-rust/bindings/python/fluss/__init__.pyi | 5 + fluss-rust/bindings/python/src/config.rs | 20 +++ fluss-rust/bindings/python/src/error.rs | 3 + .../bindings/python/test/test_config.py | 32 +++++ .../fluss/src/client/write/accumulator.rs | 119 ++++++++++++++++++ .../crates/fluss/src/client/write/sender.rs | 90 +++++++++++++ fluss-rust/crates/fluss/src/config.rs | 12 ++ fluss-rust/crates/fluss/src/rpc/api_key.rs | 8 +- .../crates/fluss/src/rpc/fluss_api_error.rs | 12 ++ .../crates/fluss/src/rpc/server_connection.rs | 6 +- .../docs/user-guide/cpp/api-reference.md | 1 + .../docs/user-guide/python/api-reference.md | 1 + .../docs/user-guide/rust/api-reference.md | 1 + 23 files changed, 352 insertions(+), 10 deletions(-) create mode 100644 fluss-rust/bindings/python/test/test_config.py diff --git a/fluss-rust/bindings/cpp/include/fluss.hpp b/fluss-rust/bindings/cpp/include/fluss.hpp index 935cf70853b..945b447ed01 100644 --- a/fluss-rust/bindings/cpp/include/fluss.hpp +++ b/fluss-rust/bindings/cpp/include/fluss.hpp @@ -188,6 +188,8 @@ struct ErrorCode { static constexpr int INVALID_ALTER_TABLE_EXCEPTION = 56; /// Deletion operations are disabled on this table. static constexpr int DELETION_DISABLED_EXCEPTION = 57; + /// The server rejected a write due to storage backpressure. + static constexpr int STORAGE_BACKPRESSURE_EXCEPTION = 72; /// Returns true if retrying the request may succeed. Mirrors Java's RetriableException hierarchy. static constexpr bool IsRetriable(int32_t code) { @@ -198,7 +200,8 @@ struct ErrorCode { code == UNKNOWN_TABLE_OR_BUCKET_EXCEPTION || code == REQUEST_TIME_OUT || code == STORAGE_EXCEPTION || code == NOT_ENOUGH_REPLICAS_AFTER_APPEND_EXCEPTION || - code == NOT_ENOUGH_REPLICAS_EXCEPTION || code == LEADER_NOT_AVAILABLE_EXCEPTION; + code == NOT_ENOUGH_REPLICAS_EXCEPTION || code == LEADER_NOT_AVAILABLE_EXCEPTION || + code == STORAGE_BACKPRESSURE_EXCEPTION; } }; @@ -1388,6 +1391,8 @@ struct Configuration { size_t writer_buffer_memory_size{64 * 1024 * 1024}; // Maximum time in milliseconds to block waiting for buffer memory uint64_t writer_buffer_wait_timeout_ms{std::numeric_limits::max()}; + // Maximum KV backpressure throttle in milliseconds + uint64_t writer_kv_backpressure_max_throttle_ms{3000}; // Connect timeout in milliseconds for TCP transport connect uint64_t connect_timeout_ms{120000}; // Security protocol: "PLAINTEXT" (default, no auth) or "sasl" (SASL auth) diff --git a/fluss-rust/bindings/cpp/src/ffi_converter.hpp b/fluss-rust/bindings/cpp/src/ffi_converter.hpp index 584749675e3..35f18830128 100644 --- a/fluss-rust/bindings/cpp/src/ffi_converter.hpp +++ b/fluss-rust/bindings/cpp/src/ffi_converter.hpp @@ -234,6 +234,8 @@ inline ffi::FfiConfig to_ffi_config(const Configuration& config) { config.writer_max_inflight_requests_per_bucket; ffi_config.writer_buffer_memory_size = config.writer_buffer_memory_size; ffi_config.writer_buffer_wait_timeout_ms = config.writer_buffer_wait_timeout_ms; + ffi_config.writer_kv_backpressure_max_throttle_ms = + config.writer_kv_backpressure_max_throttle_ms; ffi_config.connect_timeout_ms = config.connect_timeout_ms; ffi_config.security_protocol = rust::String(config.security_protocol); ffi_config.security_sasl_mechanism = rust::String(config.security_sasl_mechanism); diff --git a/fluss-rust/bindings/cpp/src/lib.rs b/fluss-rust/bindings/cpp/src/lib.rs index 675a28061c1..9bc12bbe8a0 100644 --- a/fluss-rust/bindings/cpp/src/lib.rs +++ b/fluss-rust/bindings/cpp/src/lib.rs @@ -68,6 +68,7 @@ mod ffi { writer_max_inflight_requests_per_bucket: usize, writer_buffer_memory_size: usize, writer_buffer_wait_timeout_ms: u64, + writer_kv_backpressure_max_throttle_ms: u64, connect_timeout_ms: u64, security_protocol: String, security_sasl_mechanism: String, @@ -974,6 +975,7 @@ fn new_connection(config: &ffi::FfiConfig) -> ffi::FfiPtrResult { writer_max_inflight_requests_per_bucket: config.writer_max_inflight_requests_per_bucket, writer_buffer_memory_size: config.writer_buffer_memory_size, writer_buffer_wait_timeout_ms: config.writer_buffer_wait_timeout_ms, + writer_kv_backpressure_max_throttle_ms: config.writer_kv_backpressure_max_throttle_ms, connect_timeout_ms: config.connect_timeout_ms, security_protocol: config.security_protocol.to_string(), security_sasl_mechanism: config.security_sasl_mechanism.to_string(), diff --git a/fluss-rust/bindings/cpp/test/test_ffi_converter.cpp b/fluss-rust/bindings/cpp/test/test_ffi_converter.cpp index b0e4e250818..da37feb112e 100644 --- a/fluss-rust/bindings/cpp/test/test_ffi_converter.cpp +++ b/fluss-rust/bindings/cpp/test/test_ffi_converter.cpp @@ -51,6 +51,17 @@ fluss::ffi::FfiTypeNode Node(int32_t type_id, uint32_t child_count = 0, bool nul } // namespace +TEST(FfiConverterTest, KvBackpressureConfiguration) { + fluss::Configuration config; + EXPECT_EQ(config.writer_kv_backpressure_max_throttle_ms, 3000u); + + config.writer_kv_backpressure_max_throttle_ms = 1500; + auto ffi_config = fluss::utils::to_ffi_config(config); + EXPECT_EQ(ffi_config.writer_kv_backpressure_max_throttle_ms, 1500u); + EXPECT_EQ(fluss::ErrorCode::STORAGE_BACKPRESSURE_EXCEPTION, 72); + EXPECT_TRUE(fluss::ErrorCode::IsRetriable(72)); +} + // --- DataType value semantics --- TEST(DataTypeTest, DefaultNullable) { EXPECT_TRUE(DataType::Int().nullable()); } diff --git a/fluss-rust/bindings/elixir/lib/fluss/config.ex b/fluss-rust/bindings/elixir/lib/fluss/config.ex index d7075b9ec7e..729a95b5cfc 100644 --- a/fluss-rust/bindings/elixir/lib/fluss/config.ex +++ b/fluss-rust/bindings/elixir/lib/fluss/config.ex @@ -52,6 +52,7 @@ defmodule Fluss.Config do writer_bucket_no_key_assigner: nil, writer_buffer_memory_size: nil, writer_buffer_wait_timeout_ms: nil, + writer_kv_backpressure_max_throttle_ms: nil, writer_dynamic_batch_size_enabled: nil, writer_dynamic_batch_size_min: nil, writer_enable_idempotence: nil, @@ -80,6 +81,7 @@ defmodule Fluss.Config do writer_bucket_no_key_assigner: :sticky | :round_robin | nil, writer_buffer_memory_size: non_neg_integer() | nil, writer_buffer_wait_timeout_ms: non_neg_integer() | nil, + writer_kv_backpressure_max_throttle_ms: non_neg_integer() | nil, writer_dynamic_batch_size_enabled: boolean() | nil, writer_dynamic_batch_size_min: non_neg_integer() | nil, writer_enable_idempotence: boolean() | nil, @@ -186,6 +188,11 @@ defmodule Fluss.Config do def set_writer_buffer_wait_timeout_ms(%__MODULE__{} = config, ms) when is_non_neg_integer(ms), do: %{config | writer_buffer_wait_timeout_ms: ms} + @spec set_writer_kv_backpressure_max_throttle_ms(t(), non_neg_integer()) :: t() + def set_writer_kv_backpressure_max_throttle_ms(%__MODULE__{} = config, ms) + when is_non_neg_integer(ms), + do: %{config | writer_kv_backpressure_max_throttle_ms: ms} + @spec set_writer_dynamic_batch_size_enabled(t(), boolean()) :: t() def set_writer_dynamic_batch_size_enabled(%__MODULE__{} = config, enabled) when is_boolean(enabled), diff --git a/fluss-rust/bindings/elixir/lib/fluss/error.ex b/fluss-rust/bindings/elixir/lib/fluss/error.ex index fe5d1ca8b48..b2f1ddc3b69 100644 --- a/fluss-rust/bindings/elixir/lib/fluss/error.ex +++ b/fluss-rust/bindings/elixir/lib/fluss/error.ex @@ -22,7 +22,7 @@ defmodule Fluss.Error do Fields: * `:code` — stable atom for pattern matching. - * `:error_code` — raw integer code. Protocol codes `0..57`, `-1` for + * `:error_code` — raw integer code. Protocol codes are non-negative; `-1` for `:unknown_server_error`, `-2` for `:client_error`. * `:message` — human-readable description. @@ -96,6 +96,7 @@ defmodule Fluss.Error do | :ineligible_replica_exception | :invalid_alter_table_exception | :deletion_disabled_exception + | :storage_backpressure_exception | :client_error @type t :: %__MODULE__{code: code(), error_code: integer(), message: String.t()} @@ -113,7 +114,8 @@ defmodule Fluss.Error do :storage_exception, :not_enough_replicas_after_append_exception, :not_enough_replicas_exception, - :leader_not_available_exception + :leader_not_available_exception, + :storage_backpressure_exception ] @impl true diff --git a/fluss-rust/bindings/elixir/native/fluss_nif/src/atoms.rs b/fluss-rust/bindings/elixir/native/fluss_nif/src/atoms.rs index 45d5aa303ad..7a7116422b2 100644 --- a/fluss-rust/bindings/elixir/native/fluss_nif/src/atoms.rs +++ b/fluss-rust/bindings/elixir/native/fluss_nif/src/atoms.rs @@ -100,6 +100,7 @@ rustler::atoms! { ineligible_replica_exception, invalid_alter_table_exception, deletion_disabled_exception, + storage_backpressure_exception, client_error, } @@ -212,6 +213,7 @@ fn api_error_atom(code: i32) -> Atom { FlussError::IneligibleReplicaException => ineligible_replica_exception(), FlussError::InvalidAlterTableException => invalid_alter_table_exception(), FlussError::DeletionDisabledException => deletion_disabled_exception(), + FlussError::StorageBackpressureException => storage_backpressure_exception(), } } diff --git a/fluss-rust/bindings/elixir/native/fluss_nif/src/config.rs b/fluss-rust/bindings/elixir/native/fluss_nif/src/config.rs index 8c1bab51eb5..1f09f6bfc34 100644 --- a/fluss-rust/bindings/elixir/native/fluss_nif/src/config.rs +++ b/fluss-rust/bindings/elixir/native/fluss_nif/src/config.rs @@ -50,6 +50,7 @@ pub struct NifConfig { pub writer_bucket_no_key_assigner: Option, pub writer_buffer_memory_size: Option, pub writer_buffer_wait_timeout_ms: Option, + pub writer_kv_backpressure_max_throttle_ms: Option, pub writer_dynamic_batch_size_enabled: Option, pub writer_dynamic_batch_size_min: Option, pub writer_enable_idempotence: Option, @@ -130,6 +131,9 @@ impl NifConfig { if let Some(timeout_ms) = self.writer_buffer_wait_timeout_ms { config.writer_buffer_wait_timeout_ms = timeout_ms; } + if let Some(timeout_ms) = self.writer_kv_backpressure_max_throttle_ms { + config.writer_kv_backpressure_max_throttle_ms = timeout_ms; + } if let Some(enabled) = self.writer_enable_idempotence { config.writer_enable_idempotence = enabled; } diff --git a/fluss-rust/bindings/elixir/test/config_test.exs b/fluss-rust/bindings/elixir/test/config_test.exs index 767ac24723f..d3e092fc78a 100644 --- a/fluss-rust/bindings/elixir/test/config_test.exs +++ b/fluss-rust/bindings/elixir/test/config_test.exs @@ -182,6 +182,14 @@ defmodule Fluss.ConfigTest do assert config.writer_buffer_wait_timeout_ms == 5_000 end + test "set_writer_kv_backpressure_max_throttle_ms/2 sets the throttle" do + config = + Fluss.Config.new("localhost:9123") + |> Fluss.Config.set_writer_kv_backpressure_max_throttle_ms(1_500) + + assert config.writer_kv_backpressure_max_throttle_ms == 1_500 + end + test "set_writer_enable_idempotence/2 sets the idempotence flag" do config = Fluss.Config.new("localhost:9123") diff --git a/fluss-rust/bindings/elixir/test/error_test.exs b/fluss-rust/bindings/elixir/test/error_test.exs index d6d40175974..0159df9a082 100644 --- a/fluss-rust/bindings/elixir/test/error_test.exs +++ b/fluss-rust/bindings/elixir/test/error_test.exs @@ -31,7 +31,8 @@ defmodule Fluss.ErrorTest do :storage_exception, :not_enough_replicas_after_append_exception, :not_enough_replicas_exception, - :leader_not_available_exception + :leader_not_available_exception, + :storage_backpressure_exception ] @non_retriable_codes [ diff --git a/fluss-rust/bindings/python/fluss/__init__.pyi b/fluss-rust/bindings/python/fluss/__init__.pyi index 8a2c27f1f99..01245d6ecfa 100644 --- a/fluss-rust/bindings/python/fluss/__init__.pyi +++ b/fluss-rust/bindings/python/fluss/__init__.pyi @@ -232,6 +232,10 @@ class Config: @writer_buffer_wait_timeout_ms.setter def writer_buffer_wait_timeout_ms(self, timeout: int) -> None: ... @property + def writer_kv_backpressure_max_throttle_ms(self) -> int: ... + @writer_kv_backpressure_max_throttle_ms.setter + def writer_kv_backpressure_max_throttle_ms(self, timeout: int) -> None: ... + @property def connect_timeout_ms(self) -> int: ... @connect_timeout_ms.setter def connect_timeout_ms(self, timeout: int) -> None: ... @@ -1280,6 +1284,7 @@ class ErrorCode: INELIGIBLE_REPLICA_EXCEPTION: int INVALID_ALTER_TABLE_EXCEPTION: int DELETION_DISABLED_EXCEPTION: int + STORAGE_BACKPRESSURE_EXCEPTION: int @final class OffsetSpec: diff --git a/fluss-rust/bindings/python/src/config.rs b/fluss-rust/bindings/python/src/config.rs index a4688511b83..5fe40aaa3ec 100644 --- a/fluss-rust/bindings/python/src/config.rs +++ b/fluss-rust/bindings/python/src/config.rs @@ -174,6 +174,14 @@ impl Config { )) })?; } + "writer.kv-backpressure.max-throttle-ms" => { + config.writer_kv_backpressure_max_throttle_ms = + value.parse::().map_err(|e| { + FlussError::new_err(format!( + "Invalid value '{value}' for '{key}': {e}" + )) + })?; + } "writer.bucket.no-key-assigner" => { config.writer_bucket_no_key_assigner = value.parse::().map_err(|e| { @@ -419,6 +427,18 @@ impl Config { self.inner.writer_buffer_wait_timeout_ms = timeout; } + /// Get the maximum KV backpressure throttle in milliseconds + #[getter] + fn writer_kv_backpressure_max_throttle_ms(&self) -> u64 { + self.inner.writer_kv_backpressure_max_throttle_ms + } + + /// Set the maximum KV backpressure throttle in milliseconds + #[setter] + fn set_writer_kv_backpressure_max_throttle_ms(&mut self, timeout: u64) { + self.inner.writer_kv_backpressure_max_throttle_ms = timeout; + } + /// Get the connect timeout in milliseconds #[getter] fn connect_timeout_ms(&self) -> u64 { diff --git a/fluss-rust/bindings/python/src/error.rs b/fluss-rust/bindings/python/src/error.rs index a9da03761a5..9bacea6af2d 100644 --- a/fluss-rust/bindings/python/src/error.rs +++ b/fluss-rust/bindings/python/src/error.rs @@ -273,4 +273,7 @@ impl ErrorCode { /// Deletion operations are disabled on this table. #[classattr] const DELETION_DISABLED_EXCEPTION: i32 = 57; + /// The server rejected a write due to storage backpressure. + #[classattr] + const STORAGE_BACKPRESSURE_EXCEPTION: i32 = 72; } diff --git a/fluss-rust/bindings/python/test/test_config.py b/fluss-rust/bindings/python/test/test_config.py new file mode 100644 index 00000000000..2a7641ca8d6 --- /dev/null +++ b/fluss-rust/bindings/python/test/test_config.py @@ -0,0 +1,32 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import fluss + + +def test_kv_backpressure_configuration(): + config = fluss.Config({"writer.kv-backpressure.max-throttle-ms": "1500"}) + assert config.writer_kv_backpressure_max_throttle_ms == 1500 + + config.writer_kv_backpressure_max_throttle_ms = 750 + assert config.writer_kv_backpressure_max_throttle_ms == 750 + + +def test_storage_backpressure_error_is_retriable(): + assert fluss.ErrorCode.STORAGE_BACKPRESSURE_EXCEPTION == 72 + error = fluss.FlussError("backpressure", 72) + assert error.is_retriable diff --git a/fluss-rust/crates/fluss/src/client/write/accumulator.rs b/fluss-rust/crates/fluss/src/client/write/accumulator.rs index 066e67fa3a7..0d026d83c15 100644 --- a/fluss-rust/crates/fluss/src/client/write/accumulator.rs +++ b/fluss-rust/crates/fluss/src/client/write/accumulator.rs @@ -199,11 +199,17 @@ pub struct RecordAccumulator { /// for its next poll cycle. This is the Rust equivalent of Java's /// `Sender.wakeup()` / Kafka's `RecordAccumulator.wakeup()`. sender_wakeup: Notify, + /// Per-bucket backpressure throttle expiry timestamps in milliseconds. + throttle_expiry_ms: DashMap, + max_throttle_ms: i64, } impl RecordAccumulator { pub fn new(config: Config, idempotence_manager: Arc) -> Self { let batch_timeout_ms = config.writer_batch_timeout_ms; + let max_throttle_ms = config + .writer_kv_backpressure_max_throttle_ms + .min(i64::MAX as u64) as i64; let memory_limiter = Arc::new(MemoryLimiter::new( config.writer_buffer_memory_size, Duration::from_millis(config.writer_buffer_wait_timeout_ms), @@ -221,6 +227,8 @@ impl RecordAccumulator { idempotence_manager, memory_limiter, sender_wakeup: Notify::new(), + throttle_expiry_ms: Default::default(), + max_throttle_ms, } } @@ -405,6 +413,9 @@ impl RecordAccumulator { } pub fn ready(&self, cluster: &Arc) -> Result { + let now = current_time_ms(); + self.throttle_expiry_ms.retain(|_, expiry| *expiry > now); + // Snapshot just the Arcs we need, avoiding cloning the entire BucketAndWriteBatches struct let entries: Vec<(Arc, Option, BucketBatches)> = self .write_batches @@ -496,6 +507,13 @@ impl RecordAccumulator { let deque_size = batch_guard.len(); let full = deque_size > 1 || batch.is_closed(); let table_bucket = cluster.get_table_bucket(physical_table_path, bucket_id)?; + if let Some(expiry) = self.throttle_expiry_ms.get(&table_bucket).map(|e| *e) { + let remaining = expiry.saturating_sub(current_time_ms()); + if remaining > 0 { + next_delay = next_delay.min(remaining); + continue; + } + } if let Some(leader) = cluster.leader_for(&table_bucket) { next_delay = self.batch_ready( leader, @@ -566,6 +584,9 @@ impl RecordAccumulator { first: &WriteBatch, table_bucket: &TableBucket, ) -> bool { + if self.is_throttled(table_bucket) { + return true; + } if !self.idempotence_manager.is_enabled() { return false; } @@ -599,6 +620,42 @@ impl RecordAccumulator { } } + /// Returns whether the bucket is currently throttled. + pub(crate) fn is_throttled(&self, table_bucket: &TableBucket) -> bool { + let expiry = match self.throttle_expiry_ms.get(table_bucket) { + Some(entry) => *entry, + None => return false, + }; + if current_time_ms() < expiry { + return true; + } + self.throttle_expiry_ms.remove(table_bucket); + false + } + + /// Updates the bucket throttle using `max_throttle * pressure²`. + /// Pressure `1.0` represents a hard rejection and applies the full window. + pub(crate) fn update_throttle(&self, table_bucket: &TableBucket, pressure: f32) { + if pressure >= 1f32 { + self.throttle_expiry_ms.insert( + table_bucket.clone(), + current_time_ms().saturating_add(self.max_throttle_ms), + ); + return; + } + if pressure > 0f32 { + let delay = (self.max_throttle_ms as f64 * pressure as f64 * pressure as f64) as i64; + if delay > 0 { + self.throttle_expiry_ms.insert( + table_bucket.clone(), + current_time_ms().saturating_add(delay), + ); + return; + } + } + self.throttle_expiry_ms.remove(table_bucket); + } + fn drain_batches_for_one_node( &self, cluster: &Cluster, @@ -1101,6 +1158,68 @@ mod tests { Arc::new(IdempotenceManager::new(false, 5)) } + #[test] + fn test_update_throttle() { + let accumulator = RecordAccumulator::new(Config::default(), disabled_idempotence()); + let tb = TableBucket::new(1, 0); + + let before = current_time_ms(); + accumulator.update_throttle(&tb, 0.5); + let expiry = *accumulator.throttle_expiry_ms.get(&tb).expect("entry"); + assert!((750..=800).contains(&(expiry - before))); + + let before = current_time_ms(); + accumulator.update_throttle(&tb, 1.0); + let expiry = *accumulator.throttle_expiry_ms.get(&tb).expect("entry"); + assert!((3_000..=3_050).contains(&(expiry - before))); + + accumulator.update_throttle(&tb, 0.0); + assert!(!accumulator.is_throttled(&tb)); + + accumulator + .throttle_expiry_ms + .insert(tb.clone(), current_time_ms() - 1); + assert!(!accumulator.is_throttled(&tb)); + assert!(!accumulator.throttle_expiry_ms.contains_key(&tb)); + } + + #[test] + fn test_throttle_blocks_ready_and_drain() -> Result<()> { + let config = Config { + writer_batch_timeout_ms: 10_000, + ..Config::default() + }; + let accumulator = RecordAccumulator::new(config, disabled_idempotence()); + let table_path = TablePath::new("db".to_string(), "tbl".to_string()); + let cluster = Arc::new(build_cluster(&table_path, 1, 1)); + let table_info = Arc::new(build_table_info(table_path.clone(), 1, 1)); + let physical_table_path = Arc::new(PhysicalTablePath::of(Arc::new(table_path))); + let row = GenericRow { + values: vec![Datum::Int32(1)], + }; + let record = WriteRecord::for_append(table_info, physical_table_path, 1, &row); + accumulator.append(&record, 0, &cluster, false)?; + + let tb = TableBucket::new(1, 0); + accumulator.update_throttle(&tb, 1.0); + let ready = accumulator.ready(&cluster)?; + assert!(ready.ready_nodes.is_empty()); + assert!((1..=3_000).contains(&ready.next_ready_check_delay_ms)); + + let server = cluster.get_tablet_server(1).expect("server"); + let nodes = HashSet::from([server.clone()]); + assert!( + accumulator + .drain(cluster.clone(), &nodes, 1024 * 1024)? + .is_empty() + ); + + accumulator.update_throttle(&tb, 0.0); + let mut batches = accumulator.drain(cluster, &nodes, 1024 * 1024)?; + assert_eq!(batches.remove(&1).expect("drained").len(), 1); + Ok(()) + } + fn enabled_idempotence() -> Arc { Arc::new(IdempotenceManager::new(true, 5)) } diff --git a/fluss-rust/crates/fluss/src/client/write/sender.rs b/fluss-rust/crates/fluss/src/client/write/sender.rs index 2a4f34e6465..f264fac0ebf 100644 --- a/fluss-rust/crates/fluss/src/client/write/sender.rs +++ b/fluss-rust/crates/fluss/src/client/write/sender.rs @@ -517,6 +517,9 @@ impl Sender { bucket_resp.partition_id(), bucket_resp.bucket_id(), ); + if let Some(pressure) = bucket_resp.pressure() { + self.accumulator.update_throttle(&tb, pressure); + } let Some(ready_batch) = records_by_bucket.remove(&tb) else { panic!("Missing ready batch for table bucket {tb}"); }; @@ -655,6 +658,11 @@ impl Sender { ) -> Result>> { let physical_table_path = Arc::clone(ready_write_batch.write_batch.physical_table_path()); + if error == FlussError::StorageBackpressureException { + self.accumulator + .update_throttle(&ready_write_batch.table_bucket, 1.0); + } + if error == FlussError::DuplicateSequenceException { warn!( "Duplicate sequence for {} on bucket {}: {message}", @@ -841,6 +849,7 @@ impl Sender { | FlussError::LogStorageException | FlussError::KvStorageException | FlussError::StorageException + | FlussError::StorageBackpressureException | FlussError::RequestTimeOut | FlussError::NotEnoughReplicasAfterAppendException | FlussError::NotEnoughReplicasException @@ -1012,6 +1021,11 @@ trait BucketResponse { fn error_message(&self) -> Option<&String>; fn partition_id(&self) -> Option; + + /// Backpressure signal carried by PutKv responses. + fn pressure(&self) -> Option { + None + } } impl BucketResponse for PbProduceLogRespForBucket { @@ -1044,6 +1058,10 @@ impl BucketResponse for PbPutKvRespForBucket { fn partition_id(&self) -> Option { self.partition_id } + + fn pressure(&self) -> Option { + self.pressure + } } trait WriteResponse { @@ -1155,6 +1173,78 @@ mod tests { Ok(()) } + #[tokio::test] + async fn kv_backpressure_throttles_pressure_and_hard_rejection() -> Result<()> { + let table_path = Arc::new(TablePath::new("db".to_string(), "tbl".to_string())); + let cluster = build_cluster_arc(table_path.as_ref(), 1, 1); + let metadata = Arc::new(Metadata::new_for_test(cluster.clone())); + let idempotence = disabled_idempotence(); + let accumulator = Arc::new(RecordAccumulator::new( + Config::default(), + Arc::clone(&idempotence), + )); + let sender = Sender::new( + metadata, + accumulator.clone(), + 1024 * 1024, + 1000, + 1, + 1, + idempotence, + Arc::new(crate::metrics::WriterMetrics::new()), + ); + + let (batch, _handle) = + build_ready_batch(accumulator.as_ref(), cluster.clone(), table_path.clone())?; + let tb = batch.table_bucket.clone(); + let mut records_by_bucket = HashMap::new(); + records_by_bucket.insert(tb.clone(), batch); + let request_buckets = vec![tb.clone()]; + + let response = PutKvResponse { + buckets_resp: vec![PbPutKvRespForBucket { + partition_id: None, + bucket_id: tb.bucket_id(), + error_code: None, + error_message: None, + log_end_offset: None, + pressure: Some(0.5), + }], + }; + sender + .handle_write_response( + tb.table_id(), + &request_buckets, + &mut records_by_bucket, + response, + ) + .await?; + + assert!(accumulator.is_throttled(&tb)); + accumulator.update_throttle(&tb, 0.0); + + let (batch, _handle) = + build_ready_batch(accumulator.as_ref(), cluster.clone(), table_path)?; + + sender.handle_write_batch_error( + batch, + FlussError::StorageBackpressureException, + "backpressure".to_string(), + )?; + + assert!(accumulator.is_throttled(&tb)); + let server = cluster.get_tablet_server(1).expect("server"); + let nodes = HashSet::from([server.clone()]); + let batches = accumulator.drain(cluster.clone(), &nodes, 1024 * 1024)?; + assert!(batches.is_empty()); + + accumulator.update_throttle(&tb, 0.0); + let mut batches = accumulator.drain(cluster, &nodes, 1024 * 1024)?; + let batch = batches.remove(&1).expect("drained").pop().expect("batch"); + assert_eq!(batch.write_batch.attempts(), 1); + Ok(()) + } + #[test] fn retriable_error_records_retry_metric() { use metrics_util::debugging::{DebugValue, DebuggingRecorder}; diff --git a/fluss-rust/crates/fluss/src/config.rs b/fluss-rust/crates/fluss/src/config.rs index cad8d9cb559..f2d932f7dd2 100644 --- a/fluss-rust/crates/fluss/src/config.rs +++ b/fluss-rust/crates/fluss/src/config.rs @@ -38,6 +38,7 @@ const DEFAULT_SCANNER_LOG_FETCH_MAX_BYTES_FOR_BUCKET: i32 = 1024 * 1024; const DEFAULT_WRITER_MAX_INFLIGHT_REQUESTS_PER_BUCKET: usize = 5; const DEFAULT_WRITER_BUFFER_MEMORY_SIZE: usize = 64 * 1024 * 1024; // 64MB, matching Java const DEFAULT_WRITER_BUFFER_WAIT_TIMEOUT_MS: u64 = u64::MAX; +const DEFAULT_WRITER_KV_BACKPRESSURE_MAX_THROTTLE_MS: u64 = 3000; const MAX_IN_FLIGHT_REQUESTS_PER_BUCKET_FOR_IDEMPOTENCE: usize = 5; const DEFAULT_ACKS: &str = "all"; @@ -162,6 +163,12 @@ pub struct Config { #[arg(long, default_value_t = DEFAULT_WRITER_BUFFER_WAIT_TIMEOUT_MS)] pub writer_buffer_wait_timeout_ms: u64, + /// Maximum KV backpressure throttle in milliseconds. A pressure `p` delays the bucket by + /// `max_throttle * p²`; a hard rejection uses the full window. + /// Default: 3000 (matching Java `client.writer.kv-backpressure.max-throttle`) + #[arg(long, default_value_t = DEFAULT_WRITER_KV_BACKPRESSURE_MAX_THROTTLE_MS)] + pub writer_kv_backpressure_max_throttle_ms: u64, + /// Connect timeout in milliseconds for TCP transport connect. /// Default: 120000 (120 seconds). #[arg(long, default_value_t = DEFAULT_CONNECT_TIMEOUT_MS)] @@ -264,6 +271,10 @@ impl std::fmt::Debug for Config { "writer_buffer_wait_timeout_ms", &self.writer_buffer_wait_timeout_ms, ) + .field( + "writer_kv_backpressure_max_throttle_ms", + &self.writer_kv_backpressure_max_throttle_ms, + ) .field("connect_timeout_ms", &self.connect_timeout_ms) .field("security_protocol", &self.security_protocol) .field("security_sasl_mechanism", &self.security_sasl_mechanism) @@ -306,6 +317,7 @@ impl Default for Config { DEFAULT_WRITER_MAX_INFLIGHT_REQUESTS_PER_BUCKET, writer_buffer_memory_size: DEFAULT_WRITER_BUFFER_MEMORY_SIZE, writer_buffer_wait_timeout_ms: DEFAULT_WRITER_BUFFER_WAIT_TIMEOUT_MS, + writer_kv_backpressure_max_throttle_ms: DEFAULT_WRITER_KV_BACKPRESSURE_MAX_THROTTLE_MS, connect_timeout_ms: DEFAULT_CONNECT_TIMEOUT_MS, security_protocol: String::from(DEFAULT_SECURITY_PROTOCOL), security_sasl_mechanism: String::from(DEFAULT_SASL_MECHANISM), diff --git a/fluss-rust/crates/fluss/src/rpc/api_key.rs b/fluss-rust/crates/fluss/src/rpc/api_key.rs index cab034d216b..239b1d44823 100644 --- a/fluss-rust/crates/fluss/src/rpc/api_key.rs +++ b/fluss-rust/crates/fluss/src/rpc/api_key.rs @@ -129,10 +129,12 @@ impl ApiKey { | ApiKey::GetClusterHealth | ApiKey::ListRemoteLogManifests | ApiKey::ListKvSnapshots => Some(ApiVersionRange::new(ApiVersion(0), ApiVersion(0))), - // PutKv / Lookup / PrefixLookup support v0 (legacy key encoding) - // and v1 (Paimon BinaryRow key encoding for kv_format_version=2 + // PutKv v2 adds the storage backpressure error code. + ApiKey::PutKv => Some(ApiVersionRange::new(ApiVersion(0), ApiVersion(2))), + // Lookup / PrefixLookup support v0 (legacy key encoding) and v1 + // (Paimon BinaryRow key encoding for kv_format_version=2 // non-default bucket keys). The Rust client encodes both. - ApiKey::PutKv | ApiKey::Lookup | ApiKey::PrefixLookup => { + ApiKey::Lookup | ApiKey::PrefixLookup => { Some(ApiVersionRange::new(ApiVersion(0), ApiVersion(1))) } Unknown(_) => None, diff --git a/fluss-rust/crates/fluss/src/rpc/fluss_api_error.rs b/fluss-rust/crates/fluss/src/rpc/fluss_api_error.rs index 418f5443dc4..10d0dd37969 100644 --- a/fluss-rust/crates/fluss/src/rpc/fluss_api_error.rs +++ b/fluss-rust/crates/fluss/src/rpc/fluss_api_error.rs @@ -171,6 +171,8 @@ pub enum FlussError { InvalidAlterTableException = 56, /// Deletion operations are disabled on this table. DeletionDisabledException = 57, + /// The KV storage engine rejected a write due to backpressure. + StorageBackpressureException = 72, } impl FlussError { @@ -195,6 +197,7 @@ impl FlussError { | FlussError::NotEnoughReplicasAfterAppendException | FlussError::NotEnoughReplicasException | FlussError::LeaderNotAvailableException + | FlussError::StorageBackpressureException ) } @@ -298,6 +301,9 @@ impl FlussError { FlussError::DeletionDisabledException => { "Deletion operations are disabled on this table." } + FlussError::StorageBackpressureException => { + "The tablet server has rejected the write because the KV storage engine has reached its write-pressure threshold." + } } } @@ -372,6 +378,7 @@ impl FlussError { 55 => FlussError::IneligibleReplicaException, 56 => FlussError::InvalidAlterTableException, 57 => FlussError::DeletionDisabledException, + 72 => FlussError::StorageBackpressureException, _ => FlussError::UnknownServerError, } } @@ -410,6 +417,10 @@ mod tests { FlussError::for_code(FlussError::AuthorizationException.code()), FlussError::AuthorizationException ); + assert_eq!( + FlussError::for_code(72), + FlussError::StorageBackpressureException + ); assert_eq!(FlussError::for_code(9999), FlussError::UnknownServerError); } @@ -473,6 +484,7 @@ mod tests { FlussError::NotEnoughReplicasAfterAppendException, FlussError::NotEnoughReplicasException, FlussError::LeaderNotAvailableException, + FlussError::StorageBackpressureException, ]; for err in &retriable { assert!(err.is_retriable(), "{err:?} should be retriable"); diff --git a/fluss-rust/crates/fluss/src/rpc/server_connection.rs b/fluss-rust/crates/fluss/src/rpc/server_connection.rs index cc80fd15b27..885dcbbf60c 100644 --- a/fluss-rust/crates/fluss/src/rpc/server_connection.rs +++ b/fluss-rust/crates/fluss/src/rpc/server_connection.rs @@ -1186,11 +1186,11 @@ mod tests { assert_eq!( resolve_api_version_for(None, ApiKey::PutKv).unwrap(), - ApiVersion(1) + ApiVersion(2) ); let server_versions = vec![ - // PutKv: server v0..v3, client v0 only (v1 key encoding not yet implemented) → negotiated v0 + // PutKv: server v0..v3, client v0..v2 → negotiated v2 PbApiVersion { api_key: 1016, min_version: 0, @@ -1220,7 +1220,7 @@ mod tests { // Successful negotiation cases assert_eq!( negotiated.highest_available_version(ApiKey::PutKv).unwrap(), - ApiVersion(1) + ApiVersion(2) ); assert_eq!( negotiated diff --git a/fluss-rust/website/docs/user-guide/cpp/api-reference.md b/fluss-rust/website/docs/user-guide/cpp/api-reference.md index 621eb7ace50..d8b4fb4256e 100644 --- a/fluss-rust/website/docs/user-guide/cpp/api-reference.md +++ b/fluss-rust/website/docs/user-guide/cpp/api-reference.md @@ -25,6 +25,7 @@ Complete API reference for the Fluss C++ client. | `writer_dynamic_batch_size_enabled` | `bool` | `true` | Enable per-table dynamic batch sizing: target grows 10% above 80% fill, shrinks 5% below 50% | | `writer_dynamic_batch_size_min` | `int32_t` | `262144` (256 KB) | Lower bound for the dynamic batch size estimator (ignored when disabled) | | `writer_batch_timeout_ms` | `int64_t` | `100` | Maximum time in ms to wait for a writer batch to fill up before sending | +| `writer_kv_backpressure_max_throttle_ms` | `uint64_t` | `3000` | Maximum per-bucket KV backpressure throttle in milliseconds | | `writer_bucket_no_key_assigner` | `std::string` | `"sticky"` | Bucket assignment strategy for tables without bucket keys: `"sticky"` or `"round_robin"` | | `scanner_remote_log_prefetch_num` | `size_t` | `4` | Number of remote log segments to prefetch | | `remote_file_download_thread_num` | `size_t` | `3` | Number of threads for remote log downloads | diff --git a/fluss-rust/website/docs/user-guide/python/api-reference.md b/fluss-rust/website/docs/user-guide/python/api-reference.md index 54d8af665fc..b6c55374ba2 100644 --- a/fluss-rust/website/docs/user-guide/python/api-reference.md +++ b/fluss-rust/website/docs/user-guide/python/api-reference.md @@ -18,6 +18,7 @@ Complete API reference for the Fluss Python client. | `writer_dynamic_batch_size_enabled` | `writer.dynamic-batch-size.enabled` | Get/set whether the per-table dynamic batch size estimator is enabled (default `true`) | | `writer_dynamic_batch_size_min` | `writer.dynamic-batch-size-min` | Get/set the lower bound for the dynamic batch size estimator (default 256 KB; ignored when disabled) | | `writer_batch_timeout_ms` | `writer.batch-timeout-ms` | Get/set max time in ms to wait for a writer batch to fill up before sending | +| `writer_kv_backpressure_max_throttle_ms` | `writer.kv-backpressure.max-throttle-ms` | Get/set maximum per-bucket KV backpressure throttle in milliseconds (default `3000`) | | `writer_bucket_no_key_assigner` | `writer.bucket.no-key-assigner` | Get/set bucket assignment strategy (`"sticky"` or `"round_robin"`) | | `scanner_remote_log_prefetch_num` | `scanner.remote-log.prefetch-num` | Get/set number of remote log segments to prefetch | | `remote_file_download_thread_num` | `remote-file.download-thread-num` | Get/set number of threads for remote log downloads | diff --git a/fluss-rust/website/docs/user-guide/rust/api-reference.md b/fluss-rust/website/docs/user-guide/rust/api-reference.md index 09618e5a27c..43194c7972f 100644 --- a/fluss-rust/website/docs/user-guide/rust/api-reference.md +++ b/fluss-rust/website/docs/user-guide/rust/api-reference.md @@ -17,6 +17,7 @@ Complete API reference for the Fluss Rust client. | `writer_dynamic_batch_size_enabled` | `bool` | `true` | Enable per-table dynamic batch sizing: target grows 10% above 80% fill, shrinks 5% below 50%, clamped to `[writer_dynamic_batch_size_min, writer_batch_size]` | | `writer_dynamic_batch_size_min` | `i32` | `262144` (256 KB) | Lower bound for the dynamic batch size estimator (ignored when `writer_dynamic_batch_size_enabled` is `false`) | | `writer_batch_timeout_ms` | `i64` | `100` | Maximum time in ms to wait for a writer batch to fill up before sending | +| `writer_kv_backpressure_max_throttle_ms` | `u64` | `3000` | Maximum per-bucket KV backpressure throttle in milliseconds | | `writer_bucket_no_key_assigner` | `NoKeyAssigner` | `sticky` | Bucket assignment strategy for tables without bucket keys: `sticky` or `round_robin` | | `scanner_remote_log_prefetch_num` | `usize` | `4` | Number of remote log segments to prefetch | | `remote_file_download_thread_num` | `usize` | `3` | Number of threads for remote log downloads |