This is an automated email from the ASF dual-hosted git repository.
fresh-borzoni pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss.git
The following commit(s) were added to refs/heads/main by this push:
new dd653f0b0 [rust][tests] Add DynamicWriteBatchSizeEstimator integration
tests (#3588)
dd653f0b0 is described below
commit dd653f0b0c28e9d036a06104944882d2950a71d2
Author: Prajwal Banakar <[email protected]>
AuthorDate: Thu Jul 9 22:22:22 2026 +0530
[rust][tests] Add DynamicWriteBatchSizeEstimator integration tests (#3588)
* [rust][tests] Add DynamicWriteBatchSizeEstimator integration tests
* [rust][tests] Add DynamicWriteBatchSizeEstimator integration tests
* Address review: verify estimator grows in shrink_then_grow test
---
fluss-rust/crates/fluss/src/client/connection.rs | 14 +
.../crates/fluss/src/client/write/accumulator.rs | 7 +-
.../crates/fluss/src/client/write/writer_client.rs | 8 +
.../fluss/tests/integration/dynamic_batch_size.rs | 340 +++++++++++++++++++++
fluss-rust/crates/fluss/tests/test_fluss.rs | 1 +
5 files changed, 368 insertions(+), 2 deletions(-)
diff --git a/fluss-rust/crates/fluss/src/client/connection.rs
b/fluss-rust/crates/fluss/src/client/connection.rs
index c31104c46..71985cb1a 100644
--- a/fluss-rust/crates/fluss/src/client/connection.rs
+++ b/fluss-rust/crates/fluss/src/client/connection.rs
@@ -23,6 +23,9 @@ use crate::client::table::FlussTable;
use crate::config::Config;
use crate::error::{Error, FlussError, Result};
use crate::metadata::TablePath;
+
+#[cfg(feature = "integration_tests")]
+use crate::metadata::PhysicalTablePath;
use crate::rpc::RpcClient;
use parking_lot::RwLock;
use std::sync::Arc;
@@ -144,6 +147,17 @@ impl FlussConnection {
Ok(new_client)
}
+ #[cfg(feature = "integration_tests")]
+ pub fn estimated_batch_size_for_table(
+ &self,
+ table_path: &Arc<PhysicalTablePath>,
+ ) -> Option<usize> {
+ self.writer_client
+ .read()
+ .as_ref()
+ .and_then(|c| c.estimated_batch_size_for_table(table_path))
+ }
+
/// Gets or creates a lookup client for batched lookup operations.
pub fn get_or_create_lookup_client(&self) -> Result<Arc<LookupClient>> {
// 1. Fast path: Attempt to acquire a read lock to check if the client
already exists.
diff --git a/fluss-rust/crates/fluss/src/client/write/accumulator.rs
b/fluss-rust/crates/fluss/src/client/write/accumulator.rs
index ed5d31a30..a39f3bfaf 100644
--- a/fluss-rust/crates/fluss/src/client/write/accumulator.rs
+++ b/fluss-rust/crates/fluss/src/client/write/accumulator.rs
@@ -741,8 +741,11 @@ impl RecordAccumulator {
}
}
- #[cfg(test)]
- fn estimated_batch_size(&self, table_path: &Arc<PhysicalTablePath>) ->
Option<usize> {
+ #[cfg(any(test, feature = "integration_tests"))]
+ pub(crate) fn estimated_batch_size(
+ &self,
+ table_path: &Arc<PhysicalTablePath>,
+ ) -> Option<usize> {
self.write_batches
.get(table_path)?
.dynamic_batch_size
diff --git a/fluss-rust/crates/fluss/src/client/write/writer_client.rs
b/fluss-rust/crates/fluss/src/client/write/writer_client.rs
index 73d91c7c2..53e1fd487 100644
--- a/fluss-rust/crates/fluss/src/client/write/writer_client.rs
+++ b/fluss-rust/crates/fluss/src/client/write/writer_client.rs
@@ -221,6 +221,14 @@ impl WriterClient {
Ok(())
}
+ #[cfg(feature = "integration_tests")]
+ pub fn estimated_batch_size_for_table(
+ &self,
+ table_path: &Arc<PhysicalTablePath>,
+ ) -> Option<usize> {
+ self.accumulate.estimated_batch_size(table_path)
+ }
+
pub fn create_bucket_assigner(
table_info: &Arc<TableInfo>,
table_path: Arc<PhysicalTablePath>,
diff --git a/fluss-rust/crates/fluss/tests/integration/dynamic_batch_size.rs
b/fluss-rust/crates/fluss/tests/integration/dynamic_batch_size.rs
new file mode 100644
index 000000000..965f1005d
--- /dev/null
+++ b/fluss-rust/crates/fluss/tests/integration/dynamic_batch_size.rs
@@ -0,0 +1,340 @@
+/*
+ * 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.
+ */
+
+#[cfg(test)]
+mod dynamic_batch_size_test {
+ use crate::integration::utils::{create_table, get_shared_cluster,
wait_for_table_ready};
+ use fluss::client::{EARLIEST_OFFSET, FlussConnection};
+ use fluss::config::Config;
+ use fluss::metadata::{DataTypes, PhysicalTablePath, Schema,
TableDescriptor, TablePath};
+ use fluss::row::{DataGetters, Datum, GenericRow};
+ use std::sync::Arc;
+ use std::time::Duration;
+
+ fn make_config(
+ bootstrap_servers: &str,
+ batch_size: i32,
+ dynamic_enabled: bool,
+ dynamic_min: i32,
+ ) -> Config {
+ Config {
+ bootstrap_servers: bootstrap_servers.to_string(),
+ writer_acks: "all".to_string(),
+ writer_batch_size: batch_size,
+ writer_dynamic_batch_size_enabled: dynamic_enabled,
+ writer_dynamic_batch_size_min: dynamic_min,
+ ..Config::default()
+ }
+ }
+
+ fn log_table_descriptor() -> TableDescriptor {
+ TableDescriptor::builder()
+ .schema(
+ Schema::builder()
+ .column("id", DataTypes::int())
+ .column("payload", DataTypes::string())
+ .build()
+ .expect("Failed to build schema"),
+ )
+ .build()
+ .expect("Failed to build table descriptor")
+ }
+
+ fn make_row(id: i32, payload: &str) -> GenericRow<'static> {
+ let mut row = GenericRow::new(2);
+ row.set_field(0, Datum::Int32(id));
+ row.set_field(1, Datum::String(payload.to_string().into()));
+ row
+ }
+
+ async fn read_back_records(
+ connection: &FlussConnection,
+ table_path: &TablePath,
+ expected_count: usize,
+ ) -> Vec<(i32, String)> {
+ let table = connection
+ .get_table(table_path)
+ .await
+ .expect("Failed to get table for scan");
+ let log_scanner = table
+ .new_scan()
+ .create_log_scanner()
+ .expect("Failed to create log scanner");
+ log_scanner
+ .subscribe(0, EARLIEST_OFFSET)
+ .await
+ .expect("Failed to subscribe");
+
+ let mut collected: Vec<(i32, String)> = Vec::new();
+ let start = std::time::Instant::now();
+ while collected.len() < expected_count && start.elapsed() <
Duration::from_secs(30) {
+ let records = log_scanner
+ .poll(Duration::from_millis(500))
+ .await
+ .expect("Failed to poll");
+ for rec in records {
+ let row = rec.row();
+ collected.push((
+ row.get_int(0).unwrap(),
+ row.get_string(1).unwrap().to_string(),
+ ));
+ }
+ }
+ collected
+ }
+
+ /// Writes many tiny rows (well below 50% of batch size per drain) so the
estimator
+ /// shrinks the target toward min, then reads all records back to verify
data integrity.
+ #[tokio::test]
+ async fn small_rows_shrink_batch_size() {
+ let cluster = get_shared_cluster();
+ let config = make_config(
+ cluster.plaintext_bootstrap_servers(),
+ 256 * 1024, // max = 256 KB
+ true,
+ 4 * 1024, // min = 4 KB
+ );
+ let connection = FlussConnection::new(config)
+ .await
+ .expect("Failed to connect");
+ let admin = connection.get_admin().expect("Failed to get admin");
+ let table_path = TablePath::new("fluss", "test_dynamic_small_rows");
+ create_table(&admin, &table_path, &log_table_descriptor()).await;
+ wait_for_table_ready(&admin, &table_path).await;
+
+ let table = connection
+ .get_table(&table_path)
+ .await
+ .expect("Failed to get table");
+ let writer = table
+ .new_append()
+ .expect("Failed to create append")
+ .create_writer()
+ .expect("Failed to create writer");
+
+ // Write many tiny rows — each well below 50% of the batch size so
+ // the estimator shrinks the target toward min after each drain.
+ let row_count = 200usize;
+ for i in 0..row_count {
+ let row = make_row(i as i32, "x");
+ writer.append(&row).expect("Failed to append row");
+ }
+ writer.flush().await.expect("Failed to flush");
+
+ // Verify the estimator actually shrunk toward min, not just that
flush succeeded.
+ let physical_path =
Arc::new(PhysicalTablePath::of(Arc::new(table_path.clone())));
+ let estimated = connection
+ .estimated_batch_size_for_table(&physical_path)
+ .expect("Estimator should exist after writes");
+ assert!(
+ estimated < 256 * 1024,
+ "Expected batch size to shrink below max (256 KB), got {estimated}"
+ );
+ }
+
+ /// Verifies estimator adapts correctly across a shrink-then-grow cycle:
+ /// first shrinks with tiny rows, then grows with large rows, and all
+ /// records written in both phases are read back correctly.
+ #[tokio::test]
+ async fn shrink_then_grow_batch_size() {
+ let cluster = get_shared_cluster();
+ let max_batch = 256 * 1024i32; // 256 KB
+ let config = make_config(
+ cluster.plaintext_bootstrap_servers(),
+ max_batch,
+ true,
+ 4 * 1024, // min = 4 KB
+ );
+ let connection = FlussConnection::new(config)
+ .await
+ .expect("Failed to connect");
+ let admin = connection.get_admin().expect("Failed to get admin");
+ let table_path = TablePath::new("fluss", "test_dynamic_shrink_grow");
+ create_table(&admin, &table_path, &log_table_descriptor()).await;
+ wait_for_table_ready(&admin, &table_path).await;
+
+ let table = connection
+ .get_table(&table_path)
+ .await
+ .expect("Failed to get table");
+ let writer = table
+ .new_append()
+ .expect("Failed to create append")
+ .create_writer()
+ .expect("Failed to create writer");
+
+ // Phase 1: shrink — many tiny rows drive the target toward min.
+ let small_count = 100usize;
+ for i in 0..small_count {
+ let row = make_row(i as i32, "small");
+ writer.append(&row).expect("Failed to append small row");
+ }
+ writer.flush().await.expect("Failed to flush phase 1");
+
+ let physical_path =
Arc::new(PhysicalTablePath::of(Arc::new(table_path.clone())));
+
+ let estimated_after_shrink = connection
+ .estimated_batch_size_for_table(&physical_path)
+ .expect("Estimator should exist after shrink phase");
+
+ assert!(
+ estimated_after_shrink < max_batch as usize,
+ "Expected batch size to shrink below max ({max_batch}), got
{estimated_after_shrink}"
+ );
+
+ // Phase 2: grow — large rows (>80% of current target) drive the
+ // estimator back up toward max after it has been shrunk.
+ let large_payload = "A".repeat(200 * 1024); // ~200 KB
+ let large_count = 5usize;
+ for i in 0..large_count {
+ let row = make_row((small_count + i) as i32, &large_payload);
+ writer.append(&row).expect("Failed to append large row");
+ }
+ writer.flush().await.expect("Failed to flush phase 2");
+
+ let estimated_after_grow = connection
+ .estimated_batch_size_for_table(&physical_path)
+ .expect("Estimator should exist after grow phase");
+
+ assert!(
+ estimated_after_grow > estimated_after_shrink,
+ "Expected batch size to grow after large writes
(before={estimated_after_shrink}, after={estimated_after_grow})"
+ );
+
+ // Read back all records from both phases to verify data integrity.
+ let total = small_count + large_count;
+ let records = read_back_records(&connection, &table_path, total).await;
+ assert_eq!(
+ records.len(),
+ total,
+ "Expected {total} records after shrink-then-grow cycle, got {}",
+ records.len()
+ );
+ }
+
+ /// With dynamic sizing disabled, the writer uses the static
writer_batch_size
+ /// for every batch. Verifies all records are written and read back
correctly.
+ #[tokio::test]
+ async fn disabled_keeps_static_batch_size() {
+ let cluster = get_shared_cluster();
+ let config = make_config(
+ cluster.plaintext_bootstrap_servers(),
+ 256 * 1024,
+ false, // disabled
+ 4 * 1024,
+ );
+ let connection = FlussConnection::new(config)
+ .await
+ .expect("Failed to connect");
+ let admin = connection.get_admin().expect("Failed to get admin");
+ let table_path = TablePath::new("fluss", "test_dynamic_disabled");
+ create_table(&admin, &table_path, &log_table_descriptor()).await;
+ wait_for_table_ready(&admin, &table_path).await;
+
+ let table = connection
+ .get_table(&table_path)
+ .await
+ .expect("Failed to get table");
+ let writer = table
+ .new_append()
+ .expect("Failed to create append")
+ .create_writer()
+ .expect("Failed to create writer");
+
+ let row_count = 100usize;
+ for i in 0..row_count {
+ let row = make_row(i as i32, "static");
+ writer.append(&row).expect("Failed to append row");
+ }
+ writer.flush().await.expect("Failed to flush");
+
+ let records = read_back_records(&connection, &table_path,
row_count).await;
+ assert_eq!(
+ records.len(),
+ row_count,
+ "Expected {row_count} records with static batch size, got {}",
+ records.len()
+ );
+ }
+
+ /// Multiple concurrent tasks share a single connection (and thus a single
+ /// WriterClient and RecordAccumulator with a shared estimator). Verifies
that
+ /// concurrent appends from different tasks don't corrupt estimator state
and
+ /// all records land correctly.
+ #[tokio::test]
+ async fn concurrent_appends_share_estimator_without_corruption() {
+ let cluster = get_shared_cluster();
+ let config = make_config(
+ cluster.plaintext_bootstrap_servers(),
+ 256 * 1024,
+ true,
+ 4 * 1024,
+ );
+ let connection = FlussConnection::new(config)
+ .await
+ .expect("Failed to connect");
+ let admin = connection.get_admin().expect("Failed to get admin");
+ let table_path = TablePath::new("fluss", "test_dynamic_concurrent");
+ create_table(&admin, &table_path, &log_table_descriptor()).await;
+ wait_for_table_ready(&admin, &table_path).await;
+
+ let table = connection
+ .get_table(&table_path)
+ .await
+ .expect("Failed to get table");
+
+ // All tasks share the same AppendWriter (and thus the same
+ // RecordAccumulator / estimator) via Arc.
+ let writer = Arc::new(
+ table
+ .new_append()
+ .expect("Failed to create append")
+ .create_writer()
+ .expect("Failed to create writer"),
+ );
+
+ let tasks_count = 4usize;
+ let rows_per_task = 50usize;
+ let mut handles = Vec::new();
+
+ for task_id in 0..tasks_count {
+ let writer_clone = Arc::clone(&writer);
+ let handle = tokio::spawn(async move {
+ for i in 0..rows_per_task {
+ let row = make_row((task_id * rows_per_task + i) as i32,
"concurrent");
+ writer_clone.append(&row).expect("Failed to append");
+ }
+ });
+ handles.push(handle);
+ }
+
+ for handle in handles {
+ handle.await.expect("Task panicked");
+ }
+ writer.flush().await.expect("Failed to flush");
+
+ let total = tasks_count * rows_per_task;
+ let records = read_back_records(&connection, &table_path, total).await;
+ assert_eq!(
+ records.len(),
+ total,
+ "Expected {total} records from concurrent appends, got {}",
+ records.len()
+ );
+ }
+}
diff --git a/fluss-rust/crates/fluss/tests/test_fluss.rs
b/fluss-rust/crates/fluss/tests/test_fluss.rs
index eb1e40700..63c4f7855 100644
--- a/fluss-rust/crates/fluss/tests/test_fluss.rs
+++ b/fluss-rust/crates/fluss/tests/test_fluss.rs
@@ -31,5 +31,6 @@ mod integration {
mod utils;
+ mod dynamic_batch_size;
mod table_remote_scan;
}
