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-rust.git
The following commit(s) were added to refs/heads/main by this push:
new ee6a5ac3 [rust] Replace raw i64/i32 with TableId/PartitionId/BucketId
type aliases in write path (#635)
ee6a5ac3 is described below
commit ee6a5ac3b64bddb76196018b15b47e2051ea30ed
Author: Pengcheng Huang <[email protected]>
AuthorDate: Tue Jun 23 18:53:20 2026 +0800
[rust] Replace raw i64/i32 with TableId/PartitionId/BucketId type aliases
in write path (#635)
The write path (sender, accumulator, bucket_assigner) and RPC messages
(produce_log, put_kv) used raw i64/i32 for table/partition/bucket IDs
instead of the semantic type aliases defined in lib.rs. This makes the
code inconsistent with newer modules that already use the aliases.
Changes:
- cluster.rs: get_table_id() returns Option<TableId>, get_table_id_by_path()
returns &HashMap<TablePath, TableId>, get_bucket_locations() takes TableId
- produce_log.rs, put_kv.rs: new() takes TableId
- sender.rs:
build_write_request/send_and_handle_response/handle_write_response
take TableId; BucketResponse trait returns BucketId
- bucket_assigner.rs: BucketAssigner trait uses BucketId throughout
- accumulator.rs: test helper uses BucketId
- remote_log.rs: test helper uses TableId/BucketId
Co-authored-by: warmbupt <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
---
crates/fluss/src/client/table/remote_log.rs | 4 ++--
crates/fluss/src/client/write/accumulator.rs | 2 +-
crates/fluss/src/client/write/bucket_assigner.rs | 19 ++++++++++---------
crates/fluss/src/client/write/sender.rs | 14 +++++++-------
crates/fluss/src/cluster/cluster.rs | 6 +++---
crates/fluss/src/rpc/message/produce_log.rs | 4 ++--
crates/fluss/src/rpc/message/put_kv.rs | 4 ++--
7 files changed, 27 insertions(+), 26 deletions(-)
diff --git a/crates/fluss/src/client/table/remote_log.rs
b/crates/fluss/src/client/table/remote_log.rs
index c48bdcca..c843d1a2 100644
--- a/crates/fluss/src/client/table/remote_log.rs
+++ b/crates/fluss/src/client/table/remote_log.rs
@@ -988,10 +988,10 @@ mod tests {
use super::*;
use crate::metadata::TablePath;
use crate::test_utils::test_scanner_metrics;
+ use crate::{BucketId, TableId};
use std::sync::atomic::{AtomicUsize, Ordering};
- /// Helper function to create a TableBucket for testing
- fn create_table_bucket(table_id: i64, bucket_id: i32) -> TableBucket {
+ fn create_table_bucket(table_id: TableId, bucket_id: BucketId) ->
TableBucket {
TableBucket::new(table_id, bucket_id)
}
diff --git a/crates/fluss/src/client/write/accumulator.rs
b/crates/fluss/src/client/write/accumulator.rs
index b41759f2..ed5d31a3 100644
--- a/crates/fluss/src/client/write/accumulator.rs
+++ b/crates/fluss/src/client/write/accumulator.rs
@@ -1172,7 +1172,7 @@ mod tests {
accumulator: &RecordAccumulator,
cluster: &Arc<crate::cluster::Cluster>,
table_path: &TablePath,
- bucket_id: i32,
+ bucket_id: BucketId,
) -> Result<ReadyWriteBatch> {
let table_info = Arc::new(build_table_info(table_path.clone(), 1, 2));
let physical_table_path =
Arc::new(PhysicalTablePath::of(Arc::new(table_path.clone())));
diff --git a/crates/fluss/src/client/write/bucket_assigner.rs
b/crates/fluss/src/client/write/bucket_assigner.rs
index 8ad38e3d..0868a902 100644
--- a/crates/fluss/src/client/write/bucket_assigner.rs
+++ b/crates/fluss/src/client/write/bucket_assigner.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use crate::BucketId;
use crate::bucketing::BucketingFunction;
use crate::cluster::Cluster;
use crate::error::Error::IllegalArgument;
@@ -28,9 +29,9 @@ use std::sync::atomic::{AtomicI32, Ordering};
pub trait BucketAssigner: Sync + Send {
fn abort_if_batch_full(&self) -> bool;
- fn on_new_batch(&self, cluster: &Cluster, prev_bucket_id: i32);
+ fn on_new_batch(&self, cluster: &Cluster, prev_bucket_id: BucketId);
- fn assign_bucket(&self, bucket_key: Option<&Bytes>, cluster: &Cluster) ->
Result<i32>;
+ fn assign_bucket(&self, bucket_key: Option<&Bytes>, cluster: &Cluster) ->
Result<BucketId>;
}
#[derive(Debug)]
@@ -47,7 +48,7 @@ impl StickyBucketAssigner {
}
}
- fn next_bucket(&self, cluster: &Cluster, prev_bucket_id: i32) -> i32 {
+ fn next_bucket(&self, cluster: &Cluster, prev_bucket_id: BucketId) ->
BucketId {
let old_bucket = self.current_bucket_id.load(Ordering::Relaxed);
let mut new_bucket = old_bucket;
if old_bucket < 0 || old_bucket == prev_bucket_id {
@@ -92,11 +93,11 @@ impl BucketAssigner for StickyBucketAssigner {
true
}
- fn on_new_batch(&self, cluster: &Cluster, prev_bucket_id: i32) {
+ fn on_new_batch(&self, cluster: &Cluster, prev_bucket_id: BucketId) {
self.next_bucket(cluster, prev_bucket_id);
}
- fn assign_bucket(&self, _bucket_key: Option<&Bytes>, cluster: &Cluster) ->
Result<i32> {
+ fn assign_bucket(&self, _bucket_key: Option<&Bytes>, cluster: &Cluster) ->
Result<BucketId> {
let bucket_id = self.current_bucket_id.load(Ordering::Relaxed);
if bucket_id < 0 {
Ok(self.next_bucket(cluster, bucket_id))
@@ -130,9 +131,9 @@ impl BucketAssigner for RoundRobinBucketAssigner {
false
}
- fn on_new_batch(&self, _cluster: &Cluster, _prev_bucket_id: i32) {}
+ fn on_new_batch(&self, _cluster: &Cluster, _prev_bucket_id: BucketId) {}
- fn assign_bucket(&self, _bucket_key: Option<&Bytes>, cluster: &Cluster) ->
Result<i32> {
+ fn assign_bucket(&self, _bucket_key: Option<&Bytes>, cluster: &Cluster) ->
Result<BucketId> {
let next_value = self.counter.fetch_add(1, Ordering::Relaxed);
let available_buckets =
cluster.get_available_buckets_for_table_path(&self.table_path);
if available_buckets.is_empty() {
@@ -175,11 +176,11 @@ impl BucketAssigner for HashBucketAssigner {
false
}
- fn on_new_batch(&self, _: &Cluster, _: i32) {
+ fn on_new_batch(&self, _: &Cluster, _: BucketId) {
// do nothing
}
- fn assign_bucket(&self, bucket_key: Option<&Bytes>, _: &Cluster) ->
Result<i32> {
+ fn assign_bucket(&self, bucket_key: Option<&Bytes>, _: &Cluster) ->
Result<BucketId> {
let key = bucket_key.ok_or_else(|| IllegalArgument {
message: "no bucket key provided".to_string(),
})?;
diff --git a/crates/fluss/src/client/write/sender.rs
b/crates/fluss/src/client/write/sender.rs
index baf5a9b3..2a4f34e6 100644
--- a/crates/fluss/src/client/write/sender.rs
+++ b/crates/fluss/src/client/write/sender.rs
@@ -30,7 +30,7 @@ use crate::proto::{
use crate::record::{NO_BATCH_SEQUENCE, NO_WRITER_ID};
use crate::rpc::ServerConnection;
use crate::rpc::message::{InitWriterRequest, ProduceLogRequest, PutKvRequest};
-use crate::{PartitionId, TableId};
+use crate::{BucketId, PartitionId, TableId};
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use log::{debug, warn};
@@ -405,7 +405,7 @@ impl Sender {
}
fn build_write_request(
- table_id: i64,
+ table_id: TableId,
acks: i16,
timeout_ms: i32,
request_batches: &mut [ReadyWriteBatch],
@@ -457,7 +457,7 @@ impl Sender {
&self,
connection: &ServerConnection,
write_request: WriteRequest,
- table_id: i64,
+ table_id: TableId,
table_buckets: &[TableBucket],
records_by_bucket: &mut HashMap<TableBucket, ReadyWriteBatch>,
) -> Result<()> {
@@ -502,7 +502,7 @@ impl Sender {
async fn handle_write_response<R: WriteResponse>(
&self,
- table_id: i64,
+ table_id: TableId,
request_buckets: &[TableBucket],
records_by_bucket: &mut HashMap<TableBucket, ReadyWriteBatch>,
response: R,
@@ -1007,7 +1007,7 @@ enum WriteRequest {
}
trait BucketResponse {
- fn bucket_id(&self) -> i32;
+ fn bucket_id(&self) -> BucketId;
fn error_code(&self) -> Option<i32>;
fn error_message(&self) -> Option<&String>;
@@ -1015,7 +1015,7 @@ trait BucketResponse {
}
impl BucketResponse for PbProduceLogRespForBucket {
- fn bucket_id(&self) -> i32 {
+ fn bucket_id(&self) -> BucketId {
self.bucket_id
}
fn error_code(&self) -> Option<i32> {
@@ -1031,7 +1031,7 @@ impl BucketResponse for PbProduceLogRespForBucket {
}
impl BucketResponse for PbPutKvRespForBucket {
- fn bucket_id(&self) -> i32 {
+ fn bucket_id(&self) -> BucketId {
self.bucket_id
}
fn error_code(&self) -> Option<i32> {
diff --git a/crates/fluss/src/cluster/cluster.rs
b/crates/fluss/src/cluster/cluster.rs
index 1b6a3685..4370db67 100644
--- a/crates/fluss/src/cluster/cluster.rs
+++ b/crates/fluss/src/cluster/cluster.rs
@@ -338,7 +338,7 @@ impl Cluster {
self.partition_name_by_id.get(&partition_id)
}
- pub fn get_table_id(&self, table_path: &TablePath) -> Option<i64> {
+ pub fn get_table_id(&self, table_path: &TablePath) -> Option<TableId> {
self.table_id_by_path.get(table_path).copied()
}
@@ -352,7 +352,7 @@ impl Cluster {
&self.table_info_by_path
}
- pub fn get_table_id_by_path(&self) -> &HashMap<TablePath, i64> {
+ pub fn get_table_id_by_path(&self) -> &HashMap<TablePath, TableId> {
&self.table_id_by_path
}
@@ -411,7 +411,7 @@ impl Cluster {
fn get_bucket_locations(
servers: &mut HashMap<i32, ServerNode>,
bucket_metadata: &[PbBucketMetadata],
- table_id: i64,
+ table_id: TableId,
partition_id: Option<PartitionId>,
physical_table_path: &Arc<PhysicalTablePath>,
) -> Vec<BucketLocation> {
diff --git a/crates/fluss/src/rpc/message/produce_log.rs
b/crates/fluss/src/rpc/message/produce_log.rs
index dc3acb72..de9dce11 100644
--- a/crates/fluss/src/rpc/message/produce_log.rs
+++ b/crates/fluss/src/rpc/message/produce_log.rs
@@ -23,7 +23,7 @@ use crate::client::ReadyWriteBatch;
use crate::rpc::api_key::ApiKey;
use crate::rpc::frame::WriteError;
use crate::rpc::message::{ReadType, RequestBody, WriteType};
-use crate::{impl_read_type, impl_write_type, proto};
+use crate::{TableId, impl_read_type, impl_write_type, proto};
use bytes::{Buf, BufMut};
use prost::Message;
@@ -33,7 +33,7 @@ pub struct ProduceLogRequest {
impl ProduceLogRequest {
pub fn new(
- table_id: i64,
+ table_id: TableId,
ack: i16,
max_request_timeout_ms: i32,
ready_batches: &mut [ReadyWriteBatch],
diff --git a/crates/fluss/src/rpc/message/put_kv.rs
b/crates/fluss/src/rpc/message/put_kv.rs
index d9f2ad3f..8ffdaeef 100644
--- a/crates/fluss/src/rpc/message/put_kv.rs
+++ b/crates/fluss/src/rpc/message/put_kv.rs
@@ -21,7 +21,7 @@ use crate::rpc::api_key::ApiKey;
use crate::rpc::frame::ReadError;
use crate::rpc::frame::WriteError;
use crate::rpc::message::{ReadType, RequestBody, WriteType};
-use crate::{impl_read_type, impl_write_type, proto};
+use crate::{TableId, impl_read_type, impl_write_type, proto};
use bytes::{Buf, BufMut};
use prost::Message;
@@ -33,7 +33,7 @@ pub struct PutKvRequest {
#[allow(dead_code)]
impl PutKvRequest {
pub fn new(
- table_id: i64,
+ table_id: TableId,
ack: i16,
max_request_timeout_ms: i32,
target_columns: Vec<i32>,
