From 68045c689057b819547b788b36ce7db644147645 Mon Sep 17 00:00:00 2001 From: danilo-najkov-db Date: Thu, 25 Jun 2026 18:00:28 +0000 Subject: [PATCH 1/2] changes --- rust/NEXT_CHANGELOG.md | 4 + rust/README.md | 147 +++++++++++++- rust/examples/README.md | 23 ++- rust/examples/proto/Cargo.toml | 4 + rust/examples/proto/README.md | 61 +++++- rust/examples/proto/multiplexed.rs | 155 +++++++++++++++ rust/sdk/Cargo.toml | 6 +- rust/sdk/src/builder/mod.rs | 2 +- rust/sdk/src/builder/stream_builder.rs | 253 ++++++++++++++++++++++--- rust/sdk/src/lib.rs | 6 +- rust/sdk/src/multiplexed_stream.rs | 19 +- 11 files changed, 622 insertions(+), 58 deletions(-) create mode 100644 rust/examples/proto/multiplexed.rs diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 16b6031a..fb9803b0 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -16,6 +16,7 @@ `.tls_config(Arc::new(NoTlsConfig))` when connecting to plaintext `http://` endpoints. Gated behind the `testing` feature flag. - Added a configurable payload size limit per `ingest_record_offset` / `ingest_records_offset` call. Attempts to ingest more than the limit of encoded record data in a single call now return `ZerobusError::InvalidArgument` immediately, before any network I/O. The default is set slightly below the 10 MiB server limit to leave headroom for the request envelope (protobuf framing/stream metadata), so payloads accepted client-side are not later rejected by the server's transport layer. The limit is tunable per stream via `StreamBuilder::max_ingest_payload_bytes` (gRPC JSON/proto streams only; Arrow Flight streams do not enforce it and log a warning if it is set before `build_arrow()`). +- Added stream multiplexing through `StreamBuilder::multiplexed(stream_count).build()`. The builder opens multiple gRPC sub-streams behind one `MultiplexedStream` for near-linear throughput scaling when global record ordering is not required. Multiplexed ingest calls return `MessageId` values, which can be acknowledged with `wait_for_message_id` if needed. - Added stream lifecycle logging to make recovery observable. The SDK now logs (at `info`) when recovery starts and how many records are pending, and when a recovered stream re-sends unacknowledged records and how many. Each failed stream-creation attempt is logged (at `warn`) with its attempt number and retryability, and a non-retryable failure logs (at `error`) how many records were left unacknowledged (these are retained for retrieval via `get_unacked_records`/`get_unacked_batches`). These counts now distinguish in-flight batches from the true record count they carry (a single `ingest_records` can be one batch but many records), and a terminal recovery failure now always emits a single `error` even when no records remain pending. ### Bug Fixes @@ -26,6 +27,8 @@ ### Documentation +- Added Rust README and protobuf example coverage for multiplexed stream ingestion. + ### Internal Changes - Added `ZerobusStream::signal_shutdown` (crate-private), a `&self`-callable @@ -43,3 +46,4 @@ - Added `ZerobusSdkBuilder::token_cache_enabled(bool)` to enable or disable OAuth token caching (default enabled). - Added `ZerobusSdkBuilder::token_refresh_buffer(Duration)` to configure how long before a cached token's expiry it is refreshed (default 5 minutes). - Added `HeadersProvider::invalidate` with a default no-op implementation; the SDK calls it when the server rejects the supplied credentials so a provider can drop cached auth state. Existing trait implementations are unaffected. +- Added `StreamBuilder::multiplexed(usize) -> MultiplexedStreamBuilder` and made `MultiplexedStream` / `MessageId` part of the normal public Rust API instead of `testing`-only exports. diff --git a/rust/README.md b/rust/README.md index f0b7c573..47a5d25d 100644 --- a/rust/README.md +++ b/rust/README.md @@ -42,6 +42,7 @@ The Zerobus Rust SDK provides a robust, async-first interface for ingesting larg - **Built-in Recovery** - Automatic retry and reconnection for transient failures - **High Throughput** - Configurable inflight record limits for optimal performance - **Batch Ingestion** - Ingest multiple records at once with all-or-nothing semantics for maximum throughput +- **Stream Multiplexing** - Distribute protobuf/JSON records across multiple sub-streams for near-linear throughput scaling when global ordering is not required - **Flexible Serialization** - Support for both JSON (simple) and Protocol Buffers (type-safe) data formats - **Type Safety** - Protocol Buffers ensure schema validation at compile time - **Schema Generation** - CLI tool to generate protobuf schemas from Unity Catalog tables @@ -118,6 +119,7 @@ zerobus_rust_sdk/ │ │ ├── stream_options.rs # Shared stream configuration │ │ ├── tls_config.rs # TLS configuration strategies │ │ ├── landing_zone.rs # Inflight record buffer +│ │ ├── multiplexed_stream.rs # JSON/proto stream fan-out │ │ ├── offset_generator.rs # Logical offset tracking │ │ ├── proxy.rs # HTTP proxy support │ │ ├── arrow_stream.rs # Arrow Flight stream (feature: arrow-flight) @@ -162,6 +164,7 @@ zerobus_rust_sdk/ │ │ ├── Cargo.toml │ │ ├── single.rs # Protocol Buffers single-record example │ │ ├── batch.rs # Protocol Buffers batch ingestion example +│ │ ├── multiplexed.rs # Multiplexed Protocol Buffers example │ │ └── output/ # Generated schema files (shared) │ └── arrow/ # Arrow Flight example (feature: arrow-flight, Beta) │ ├── README.md @@ -173,6 +176,7 @@ zerobus_rust_sdk/ │ │ ├── mock_grpc.rs # Mock Zerobus gRPC server │ │ ├── mock_arrow_flight.rs # Mock Arrow Flight server │ │ ├── rust_tests.rs # Core SDK test suite +│ │ ├── multiplexed_stream_tests.rs # Multiplexed stream test suite │ │ ├── proxy_tests.rs # HTTP proxy tests │ │ ├── arrow_tests.rs # Arrow Flight test suite │ │ └── utils.rs # Shared test utilities @@ -490,7 +494,42 @@ let mut stream = sdk .await?; ``` -Setters can be called in any order. The builder validates at `build()` time that both authentication and format have been configured. +#### Multiplexed Stream + +Use `.multiplexed(stream_count)` before `.build()` to open multiple +gRPC sub-streams behind one handle. This can provide near-linear throughput +scaling when your workload does not require global record ordering. + +```rust +use databricks_zerobus_ingest_sdk::ProtoMessage; + +let descriptor_proto = load_descriptor_proto( + "output/orders.descriptor", + "orders.proto", + "table_Orders", +); + +let mut stream = sdk + .stream_builder() + .table("catalog.schema.orders") + .oauth(client_id, client_secret) + .compiled_proto(descriptor_proto) + .max_inflight_requests(10_000) + .multiplexed(4) + .build() + .await?; + +let message_id = stream.ingest_record(ProtoMessage(order)).await?; +stream.wait_for_message_id(message_id).await?; +stream.close().await?; +``` + +`MultiplexedStream` returns `MessageId` values instead of `OffsetId` values +because each id records which sub-stream accepted the record. Per-sub-stream +ordering is preserved, but message ids are not globally ordered across the +multiplexed stream. Arrow Flight streams are not multiplexed by this API. + +Setters can be called in any order before `.multiplexed(...)` or `.build()`. The builder validates at `build()` time that both authentication and format have been configured. ### 5. Ingest Data @@ -555,6 +594,34 @@ for i in 0..100_000 { stream.flush().await?; ``` +#### Multiplexed High Throughput Pattern + +When global ordering is not required, a multiplexed stream distributes records +round-robin across multiple protobuf gRPC sub-streams: + +```rust +let descriptor_proto = load_descriptor_proto( + "output/orders.descriptor", + "orders.proto", + "table_Orders", +); + +let mut stream = sdk + .stream_builder() + .table("catalog.schema.orders") + .oauth(client_id, client_secret) + .compiled_proto(descriptor_proto) + .multiplexed(4) + .build() + .await?; + +let message_id = stream.ingest_record(ProtoMessage(record)).await?; +stream.wait_for_message_id(message_id).await?; + +stream.flush().await?; +stream.close().await?; +``` + See [`examples/`](https://github.com/databricks/zerobus-sdk/tree/main/rust/examples) for complete working examples with all wrapper types, serialization formats, and ingestion patterns. ### 6. Handle Acknowledgments @@ -817,7 +884,7 @@ match stream.ingest_record_offset(payload).await { ### Complete Working Examples -The `examples/` directory contains four working examples covering different serialization formats and ingestion patterns: +The `examples/` directory contains working examples covering different serialization formats and ingestion patterns: | Example | Serialization | Ingestion | Run with | |---------|--------------|-----------|----------| @@ -825,6 +892,7 @@ The `examples/` directory contains four working examples covering different seri | `json/batch.rs` | JSON | Batch | `cargo run -p rust-examples-json --example json_batch` | | `proto/single.rs` | Protocol Buffers | Single-record | `cargo run -p rust-examples-proto --example proto_single` | | `proto/batch.rs` | Protocol Buffers | Batch | `cargo run -p rust-examples-proto --example proto_batch` | +| `proto/multiplexed.rs` | Protocol Buffers | Multiplexed | `cargo run -p rust-examples-proto --example proto_multiplexed` | Check [`examples/README.md`](https://github.com/databricks/zerobus-sdk/blob/main/rust/examples/README.md) for setup instructions and detailed comparisons. @@ -876,12 +944,13 @@ cargo test -p tests -- --nocapture - Use `ingest_record_offset()` when processing records individually - Both return offsets directly; use `wait_for_offset()` to explicitly wait for acknowledgments 4. **Tune Inflight Limits** - Adjust `max_inflight_requests` based on memory and throughput needs -5. **Enable Recovery** - Always set `recovery: true` in production environments -6. **Handle Ack Futures** - Use `tokio::spawn` for fire-and-forget or batch-wait for verification -7. **Monitor Errors** - Log and alert on non-retryable errors -8. **Validate Schemas** - Use the schema generation tool to ensure type safety (for Protocol Buffers) -9. **Secure Credentials** - Never hardcode secrets; use environment variables or secret managers -10. **Test Recovery** - Simulate failures to verify your error handling logic +5. **Use Multiplexing Only Without Global Ordering Requirements** - `.multiplexed(n)` distributes records across sub-streams and returns `MessageId` values, not globally ordered offsets +6. **Enable Recovery** - Always set `recovery: true` in production environments +7. **Handle Ack Futures** - Use `tokio::spawn` for fire-and-forget or batch-wait for verification +8. **Monitor Errors** - Log and alert on non-retryable errors +9. **Validate Schemas** - Use the schema generation tool to ensure type safety (for Protocol Buffers) +10. **Secure Credentials** - Never hardcode secrets; use environment variables or secret managers +11. **Test Recovery** - Simulate failures to verify your error handling logic ## API Reference @@ -963,10 +1032,72 @@ Returns unacknowledged records grouped by batch, preserving the original batch s Only call after stream failure. +### `MultiplexedStream` + +Represents multiple protobuf/JSON gRPC streams behind one round-robin ingestion +handle. Use it when you want higher throughput and do not require global record +ordering. + +**Builder:** +```rust +let descriptor_proto = load_descriptor_proto( + "output/orders.descriptor", + "orders.proto", + "table_Orders", +); + +let stream = sdk + .stream_builder() + .table("catalog.schema.table") + .oauth(client_id, client_secret) + .compiled_proto(descriptor_proto) + .multiplexed(4) + .build() + .await?; +``` + +**Methods:** +```rust +pub async fn ingest_record( + &self, + payload: impl Into +) -> ZerobusResult +``` +Queues a single record on the next sub-stream and returns a `MessageId`. + +```rust +pub async fn ingest_records( + &self, + payload: I +) -> ZerobusResult> +where + I: IntoIterator, + T: Into +``` +Queues a batch on one sub-stream and returns one `MessageId` for the batch, or +`None` for an empty batch. + +```rust +pub async fn wait_for_message_id(&self, message_id: MessageId) -> ZerobusResult<()> +``` +Waits for acknowledgment of the record or batch identified by the `MessageId`. + +```rust +pub async fn flush(&self) -> ZerobusResult<()> +pub async fn close(&mut self) -> ZerobusResult<()> +``` +Flushes or closes every sub-stream. + ### `StreamBuilder` Configure stream parameters via fluent setters; all configuration goes through the builder. See [Create a Stream](#4-create-a-stream) for usage and [Configuration Options](#configuration-options) for the full list of available setters and their defaults. +```rust +pub fn multiplexed(self, stream_count: usize) -> MultiplexedStreamBuilder<'_> +``` +Configures the builder to create a `MultiplexedStream` from `stream_count` +protobuf/JSON gRPC sub-streams. `stream_count` must be between 1 and 64. + ### `AckCallback` Trait for receiving acknowledgment notifications. diff --git a/rust/examples/README.md b/rust/examples/README.md index aeb868b1..3dd6bfed 100644 --- a/rust/examples/README.md +++ b/rust/examples/README.md @@ -30,6 +30,7 @@ The SDK supports two serialization formats and two ingestion methods: **Ingestion Methods:** - **Single-record** (`ingest_record_offset`) - Ingest records one at a time (JSON / Protocol Buffers) - **Batch** (`ingest_records_offset`) - Ingest multiple records at once with all-or-nothing semantics (JSON / Protocol Buffers) +- **Multiplexed** (`multiplexed(n)`) - Distribute protobuf records across multiple sub-streams when global ordering is not required - **Arrow batch** (`ingest_batch` / `ingest_ipc_batch`) - Ingest an Arrow `RecordBatch` (one or many rows) over Arrow Flight **Available Examples:** @@ -40,6 +41,7 @@ The SDK supports two serialization formats and two ingestion methods: | [JSON Batch](json/README.md#batch-example) | JSON | Batch | `cargo run -p rust-examples-json --example json_batch` | | [Proto Single](proto/README.md#single-record-example) | Protocol Buffers | Single-record | `cargo run -p rust-examples-proto --example proto_single` | | [Proto Batch](proto/README.md#batch-example) | Protocol Buffers | Batch | `cargo run -p rust-examples-proto --example proto_batch` | +| [Proto Multiplexed](proto/README.md#multiplexed-example) | Protocol Buffers | Multiplexed | `cargo run -p rust-examples-proto --example proto_multiplexed` | | [Arrow](arrow/README.md) | Arrow Flight (Beta) | `RecordBatch` | `cargo run -p example_arrow` | ## Prerequisites @@ -75,7 +77,7 @@ Replace `catalog.schema.orders` with your actual catalog, schema, and table name ### 3. Configure Credentials -Edit the source file (`batch.rs` or `single.rs`) for your chosen example and update these constants: +Edit the source file (`batch.rs`, `single.rs`, or `multiplexed.rs`) for your chosen example and update these constants: ```rust const DATABRICKS_WORKSPACE_URL: &str = "https://your-workspace.cloud.databricks.com"; @@ -137,6 +139,25 @@ let mut stream = sdk .await?; ``` +**Multiplexed Protocol Buffers:** +```rust +let descriptor_proto = load_descriptor_proto( + "output/orders.descriptor", + "orders.proto", + "table_Orders" +); + +let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .compiled_proto(descriptor_proto) + .max_inflight_requests(100) + .multiplexed(4) + .build() + .await?; +``` + **Arrow Flight (Beta):** ```rust use std::sync::Arc; diff --git a/rust/examples/proto/Cargo.toml b/rust/examples/proto/Cargo.toml index b5b452d2..9f092764 100644 --- a/rust/examples/proto/Cargo.toml +++ b/rust/examples/proto/Cargo.toml @@ -12,6 +12,10 @@ path = "batch.rs" name = "proto_single" path = "single.rs" +[[example]] +name = "proto_multiplexed" +path = "multiplexed.rs" + [dependencies] databricks-zerobus-ingest-sdk = { path = "../../sdk" } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/rust/examples/proto/README.md b/rust/examples/proto/README.md index ba6c5b69..1b59efc8 100644 --- a/rust/examples/proto/README.md +++ b/rust/examples/proto/README.md @@ -12,6 +12,9 @@ This directory contains examples demonstrating Protocol Buffers-based data inges - [Batch Example](#batch-example) - [Running the Example](#running-the-example-1) - [Code Highlights](#code-highlights-1) +- [Multiplexed Example](#multiplexed-example) + - [Running the Example](#running-the-example-2) + - [Code Highlights](#code-highlights-2) - [Adapting for Your Custom Table](#adapting-for-your-custom-table) - [Generate Schema Files](#generate-schema-files) - [Update main.rs](#update-mainrs) @@ -28,6 +31,7 @@ Protocol Buffers examples provide type safety and better performance. **No schem **Available examples:** - **`single.rs`** - Ingest records one at a time using `ingest_record_offset()` / `ingest_record()` - **`batch.rs`** - Ingest multiple records at once using `ingest_records_offset()` / `ingest_records()` +- **`multiplexed.rs`** - Ingest across multiple protobuf streams using `.multiplexed(n)` when global ordering is not required ## Three Ways to Pass Data @@ -177,6 +181,61 @@ if let Some(offset) = stream.ingest_records_offset(batch).await? { - **Single acknowledgment**: One offset ID for the whole batch - **Empty batches**: Returns `None` (no-op) +## Multiplexed Example + +Use multiplexing when the workload does not require global record ordering and +you want to distribute ingestion across multiple protobuf gRPC sub-streams. + +### Running the Example + +1. Configure credentials in `multiplexed.rs` (see [Prerequisites](../README.md#prerequisites)) + +2. Run the example: + ```bash + cargo run -p rust-examples-proto --example proto_multiplexed + ``` + +**Expected output:** +``` +=== Message ID-based API (Multiplexed) === +[Auto-encoding] Record sent with message ID: MessageId(stream=0, offset=0) +[Auto-encoding] Record acknowledged with message ID: MessageId(stream=0, offset=0) +[Pre-encoded] Record sent with message ID: MessageId(stream=1, offset=0) +[Pre-encoded] Record acknowledged with message ID: MessageId(stream=1, offset=0) +[Backward-compatible] Record sent with message ID: MessageId(stream=2, offset=0) +[Backward-compatible] Record acknowledged with message ID: MessageId(stream=2, offset=0) +All records flushed +Stream closed successfully +``` + +### Code Highlights + +```rust +let descriptor_proto = load_descriptor_proto( + "output/orders.descriptor", + "orders.proto", + "table_Orders" +); + +let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .compiled_proto(descriptor_proto) + .multiplexed(4) + .build() + .await?; + +let message_id = stream.ingest_record(ProtoMessage(order)).await?; +stream.wait_for_message_id(message_id).await?; +stream.close().await?; +``` + +`MultiplexedStream` returns `MessageId` values instead of `OffsetId` values +because records may be queued on different sub-streams. Per-sub-stream ordering +is preserved, but message ids are not globally ordered across the multiplexed +stream. + ## Adapting for Your Custom Table To use your own table, you need to generate schema files and update the example code. @@ -196,7 +255,7 @@ cargo run -- \ --output-dir "../../examples/proto/output" ``` -Both `single.rs` and `batch.rs` share the same `output/` directory, so the generated schema files only need to be produced once. +`single.rs`, `batch.rs`, and `multiplexed.rs` share the same `output/` directory, so the generated schema files only need to be produced once. This generates: - `output/.proto` - Protocol Buffer schema definition diff --git a/rust/examples/proto/multiplexed.rs b/rust/examples/proto/multiplexed.rs new file mode 100644 index 00000000..3fc44210 --- /dev/null +++ b/rust/examples/proto/multiplexed.rs @@ -0,0 +1,155 @@ +use std::error::Error; +use std::fs; + +use prost::Message; +use prost_reflect::prost_types; + +use databricks_zerobus_ingest_sdk::{MultiplexedStream, ProtoBytes, ProtoMessage, ZerobusSdk}; + +pub mod orders { + include!("output/orders.rs"); +} +use crate::orders::TableOrders; + +// Change constants to match your data. +const TABLE_NAME: &str = ""; +const DATABRICKS_CLIENT_ID: &str = ""; +const DATABRICKS_CLIENT_SECRET: &str = ""; + +// Uncomment the appropriate lines for your cloud. + +// For AWS: +const DATABRICKS_WORKSPACE_URL: &str = "https://.cloud.databricks.com"; +const SERVER_ENDPOINT: &str = "https://.zerobus..cloud.databricks.com"; + +// For Azure: +// const DATABRICKS_WORKSPACE_URL: &str = "https://.azuredatabricks.net"; +// const SERVER_ENDPOINT: &str = "https://.zerobus..azuredatabricks.net"; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let descriptor_proto = + load_descriptor_proto("output/orders.descriptor", "orders.proto", "table_Orders"); + let sdk_handle = ZerobusSdk::builder() + .endpoint(SERVER_ENDPOINT) + .unity_catalog_url(DATABRICKS_WORKSPACE_URL) + .build()?; + + let mut stream = sdk_handle + .stream_builder() + .table(TABLE_NAME) + .oauth(DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET) + .compiled_proto(descriptor_proto) + .max_inflight_requests(100000) + .multiplexed(4) + .build() + .await?; + + ingest_with_message_id_api(&stream).await?; + + stream.close().await?; + println!("Stream closed successfully"); + + Ok(()) +} + +/// Multiplexed API: returns a message ID after queuing on a sub-stream. +async fn ingest_with_message_id_api(stream: &MultiplexedStream) -> Result<(), Box> { + println!("=== Message ID-based API (Multiplexed) ==="); + + // Delta TIMESTAMP is int64 microseconds since epoch UTC. + let now = chrono::Utc::now().timestamp_micros(); + + // 1. Auto-encoding: ProtoMessage - pass message directly, SDK handles encoding. + let order = TableOrders { + id: Some(1), + customer_name: Some("Alice Smith".to_string()), + product_name: Some("Wireless Mouse".to_string()), + quantity: Some(2), + price: Some(25.99), + status: Some("pending".to_string()), + created_at: Some(now), + updated_at: Some(now), + }; + + let message_id = stream.ingest_record(ProtoMessage(order)).await?; + println!( + "[Auto-encoding] Record sent with message ID: {}", + message_id + ); + stream.wait_for_message_id(message_id).await?; + println!( + "[Auto-encoding] Record acknowledged with message ID: {}", + message_id + ); + + // 2. Pre-encoded: ProtoBytes - pass bytes with explicit wrapper. + let order = TableOrders { + id: Some(2), + customer_name: Some("Bob Johnson".to_string()), + product_name: Some("Mechanical Keyboard".to_string()), + quantity: Some(1), + price: Some(89.99), + status: Some("shipped".to_string()), + created_at: Some(now), + updated_at: Some(now), + }; + let message_id = stream + .ingest_record(ProtoBytes(order.encode_to_vec())) + .await?; + println!("[Pre-encoded] Record sent with message ID: {}", message_id); + stream.wait_for_message_id(message_id).await?; + println!( + "[Pre-encoded] Record acknowledged with message ID: {}", + message_id + ); + + // 3. Backward-compatible: raw Vec - no wrapper needed, works the same as ProtoBytes. + let order = TableOrders { + id: Some(3), + customer_name: Some("Carol Williams".to_string()), + product_name: Some("USB-C Hub".to_string()), + quantity: Some(3), + price: Some(45.00), + status: Some("delivered".to_string()), + created_at: Some(now), + updated_at: Some(now), + }; + let message_id = stream.ingest_record(order.encode_to_vec()).await?; + println!( + "[Backward-compatible] Record sent with message ID: {}", + message_id + ); + stream.wait_for_message_id(message_id).await?; + println!( + "[Backward-compatible] Record acknowledged with message ID: {}", + message_id + ); + + stream.flush().await?; + println!("All records flushed"); + + Ok(()) +} + +fn load_descriptor_proto( + path: &str, + file_name: &str, + message_name: &str, +) -> prost_types::DescriptorProto { + let descriptor_bytes = fs::read(path).expect("Failed to read proto descriptor file"); + let file_descriptor_set = + prost_types::FileDescriptorSet::decode(descriptor_bytes.as_ref()).unwrap(); + + let file_descriptor_proto = file_descriptor_set + .file + .into_iter() + .find(|f| f.name.as_deref() == Some(file_name)) + .unwrap(); + + file_descriptor_proto + .message_type + .into_iter() + .find(|m| m.name.as_deref() == Some(message_name)) + .unwrap() +} diff --git a/rust/sdk/Cargo.toml b/rust/sdk/Cargo.toml index a31fa982..b443a504 100644 --- a/rust/sdk/Cargo.toml +++ b/rust/sdk/Cargo.toml @@ -44,7 +44,7 @@ arrow-flight = { workspace = true, optional = true } arrow-array = { workspace = true, optional = true } arrow-schema = { workspace = true, optional = true } arrow-ipc = { workspace = true, optional = true, features = ["lz4", "zstd"] } -futures = { workspace = true, optional = true } +futures.workspace = true [dev-dependencies] tracing-subscriber.workspace = true @@ -63,10 +63,10 @@ prost-build = { version = "0.14", optional = true } [features] default = [] # Arrow Flight is in Beta -arrow-flight = ["dep:arrow-flight", "dep:arrow-array", "dep:arrow-schema", "dep:arrow-ipc", "dep:futures"] +arrow-flight = ["dep:arrow-flight", "dep:arrow-array", "dep:arrow-schema", "dep:arrow-ipc"] # Zero-copy protobuf parser. zeroparser = ["dep:self_cell", "dep:prost-build"] -testing = ["dep:futures"] +testing = [] [[test]] name = "zeroparser_e2e" diff --git a/rust/sdk/src/builder/mod.rs b/rust/sdk/src/builder/mod.rs index c995c8ae..e4701600 100644 --- a/rust/sdk/src/builder/mod.rs +++ b/rust/sdk/src/builder/mod.rs @@ -33,4 +33,4 @@ mod sdk_builder; mod stream_builder; pub use sdk_builder::ZerobusSdkBuilder; -pub use stream_builder::StreamBuilder; +pub use stream_builder::{MultiplexedStreamBuilder, StreamBuilder}; diff --git a/rust/sdk/src/builder/stream_builder.rs b/rust/sdk/src/builder/stream_builder.rs index 4696e685..27e111d0 100644 --- a/rust/sdk/src/builder/stream_builder.rs +++ b/rust/sdk/src/builder/stream_builder.rs @@ -27,7 +27,9 @@ use crate::databricks::zerobus::RecordType; use crate::headers_provider::NoAuthHeadersProvider; use crate::headers_provider::{HeadersProvider, OAuthHeadersProvider}; use crate::stream_configuration::StreamConfigurationOptions; -use crate::{TableProperties, ZerobusError, ZerobusResult, ZerobusSdk, ZerobusStream}; +use crate::{ + MultiplexedStream, TableProperties, ZerobusError, ZerobusResult, ZerobusSdk, ZerobusStream, +}; #[cfg(feature = "arrow-flight")] use crate::arrow_configuration::ArrowStreamConfigurationOptions; @@ -35,6 +37,7 @@ use crate::arrow_configuration::ArrowStreamConfigurationOptions; use crate::arrow_stream::{ArrowSchema, ArrowTableProperties, ZerobusArrowStream}; /// Internal representation of the authentication configuration. +#[derive(Clone)] enum AuthConfig { OAuth { client_id: String, @@ -46,6 +49,7 @@ enum AuthConfig { } /// Which record format was selected. +#[derive(Clone)] enum FormatConfig { Json, CompiledProto(Box), @@ -102,6 +106,31 @@ pub struct StreamBuilder<'a> { arrow_config: ArrowStreamConfigurationOptions, } +/// Fluent builder for creating a [`MultiplexedStream`]. +/// +/// Created by calling [`StreamBuilder::multiplexed`]. It preserves the same +/// table, authentication, format, and stream options configured on the source +/// [`StreamBuilder`], then opens the requested number of protobuf/JSON gRPC +/// sub-streams when [`build`](Self::build) is called. +/// +/// # Examples +/// +/// ```rust,ignore +/// let stream = sdk +/// .stream_builder() +/// .table("catalog.schema.table") +/// .oauth("client-id", "client-secret") +/// .compiled_proto(descriptor_proto) +/// .multiplexed(4) +/// .build() +/// .await?; +/// ``` +#[must_use = "a MultiplexedStreamBuilder does nothing until `.build()` is called"] +pub struct MultiplexedStreamBuilder<'a> { + inner: StreamBuilder<'a>, + stream_count: usize, +} + impl fmt::Debug for StreamBuilder<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let auth_kind = match &self.auth { @@ -126,6 +155,15 @@ impl fmt::Debug for StreamBuilder<'_> { } } +impl fmt::Debug for MultiplexedStreamBuilder<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MultiplexedStreamBuilder") + .field("inner", &self.inner) + .field("stream_count", &self.stream_count) + .finish() + } +} + const fn missing_auth_error() -> &'static str { #[cfg(feature = "testing")] { @@ -201,6 +239,23 @@ impl<'a> StreamBuilder<'a> { self } + /// Build multiple protobuf/JSON gRPC streams and expose them as one + /// round-robin [`MultiplexedStream`]. + /// + /// Use this for near-linear throughput scaling when the caller does not + /// require global record ordering. All sub-streams use the same table, + /// authentication, format, and stream options configured on this builder. + /// `stream_count` must be between 1 and [`MultiplexedStream::MAX_STREAMS`]. + /// + /// Arrow Flight streams are not multiplexed by this API; use + /// `build_arrow()` for Arrow ingestion. + pub fn multiplexed(self, stream_count: usize) -> MultiplexedStreamBuilder<'a> { + MultiplexedStreamBuilder { + inner: self, + stream_count, + } + } + /// Enable or disable automatic stream recovery. pub fn recovery(mut self, enabled: bool) -> Self { self.grpc_config.recovery = enabled; @@ -385,49 +440,55 @@ impl<'a> StreamBuilder<'a> { } } - /// Build and open a gRPC ingestion stream (JSON or compiled protobuf). - /// - /// Returns an error if table name, authentication, or format has not been set, - /// or if an Arrow format was selected (use `build_arrow()` instead). - pub async fn build(mut self) -> ZerobusResult { - self.validate()?; - let headers_provider = self.resolve_headers_provider()?; - - let (record_type, descriptor_proto) = match self.format { - Some(FormatConfig::Json) => (RecordType::Json, None), - Some(FormatConfig::CompiledProto(desc)) => (RecordType::Proto, Some(*desc)), + fn grpc_record_config( + &self, + _arrow_format_error: &'static str, + ) -> ZerobusResult<(RecordType, Option)> { + match self.format.as_ref() { + Some(FormatConfig::Json) => Ok((RecordType::Json, None)), + Some(FormatConfig::CompiledProto(desc)) => { + Ok((RecordType::Proto, Some((**desc).clone()))) + } #[cfg(feature = "arrow-flight")] Some(FormatConfig::Arrow(_)) => { - return Err(ZerobusError::InvalidArgument( - "Arrow format requires .build_arrow() instead of .build()".into(), - )); - } - None => { - return Err(ZerobusError::InvalidArgument( - "record format is required: call .json() or .compiled_proto() before .build()" - .into(), - )); + Err(ZerobusError::InvalidArgument(_arrow_format_error.into())) } - }; + None => Err(ZerobusError::InvalidArgument( + "record format is required: call .json() or .compiled_proto() before .build()" + .into(), + )), + } + } - self.grpc_config.record_type = record_type; + async fn build_grpc_stream(&self) -> ZerobusResult { + self.validate()?; + let headers_provider = self.resolve_headers_provider()?; + let (record_type, descriptor_proto) = + self.grpc_record_config("Arrow format requires .build_arrow() instead of .build()")?; + + let mut grpc_config = self.grpc_config.clone(); + grpc_config.record_type = record_type; let table_properties = TableProperties { - table_name: self.table_name, + table_name: self.table_name.clone(), descriptor_proto, }; let channel = self.sdk.get_or_create_channel_zerobus_client().await?; - let stream = ZerobusStream::new_stream( - channel, - table_properties, - headers_provider, - self.grpc_config, - ) - .await?; + let stream = + ZerobusStream::new_stream(channel, table_properties, headers_provider, grpc_config) + .await?; crate::client_warnings::record_stream_creation(stream.table_properties.table_name.as_str()); Ok(stream) } + /// Build and open a gRPC ingestion stream (JSON or compiled protobuf). + /// + /// Returns an error if table name, authentication, or format has not been set, + /// or if an Arrow format was selected (use `build_arrow()` instead). + pub async fn build(self) -> ZerobusResult { + self.build_grpc_stream().await + } + /// Build and open an Arrow Flight ingestion stream. /// /// Returns an error if table name, authentication, or format has not been set, @@ -480,6 +541,53 @@ impl<'a> StreamBuilder<'a> { } } +impl<'a> MultiplexedStreamBuilder<'a> { + fn validate_stream_count(&self) -> ZerobusResult<()> { + if self.stream_count == 0 { + return Err(ZerobusError::InvalidArgument( + "multiplexed stream count must be at least 1".into(), + )); + } + if self.stream_count > MultiplexedStream::MAX_STREAMS { + return Err(ZerobusError::InvalidArgument(format!( + "multiplexed stream count must be at most {}", + MultiplexedStream::MAX_STREAMS + ))); + } + Ok(()) + } + + /// Validate that the multiplexed stream builder has all required fields. + /// + /// Returns `Ok(())` when the wrapped stream builder is valid, the selected + /// format is protobuf/JSON gRPC, and the stream count is within the + /// supported range. + pub fn validate(&self) -> ZerobusResult<()> { + self.validate_stream_count()?; + self.inner.validate()?; + self.inner.grpc_record_config( + "Arrow format is not supported by .multiplexed(); use .build_arrow() for Arrow ingestion", + )?; + Ok(()) + } + + /// Build and open a multiplexed protobuf/JSON gRPC ingestion stream. + /// + /// Each sub-stream is created with the same table, authentication, format, + /// and stream options configured before [`StreamBuilder::multiplexed`] was + /// called. + pub async fn build(self) -> ZerobusResult { + self.validate()?; + + let mut streams = Vec::with_capacity(self.stream_count); + for _ in 0..self.stream_count { + streams.push(self.inner.build_grpc_stream().await?); + } + + Ok(MultiplexedStream::new(streams)) + } +} + #[cfg(test)] mod tests { use super::*; @@ -598,6 +706,89 @@ mod tests { ); } + #[test] + fn multiplexed_builder_validates_json() { + let sdk = test_sdk(); + let builder = sdk + .stream_builder() + .table("catalog.schema.table") + .oauth("a", "b") + .json() + .multiplexed(4); + + builder.validate().expect("valid multiplexed builder"); + let debug_str = format!("{:?}", builder); + assert!(debug_str.contains("MultiplexedStreamBuilder")); + assert!(debug_str.contains("stream_count")); + } + + #[test] + fn multiplexed_builder_validates_compiled_proto() { + let sdk = test_sdk(); + let builder = sdk + .stream_builder() + .table("catalog.schema.table") + .oauth("a", "b") + .compiled_proto(prost_types::DescriptorProto::default()) + .multiplexed(2); + + builder.validate().expect("valid multiplexed proto builder"); + } + + #[test] + fn multiplexed_builder_rejects_invalid_stream_count() { + let sdk = test_sdk(); + + let zero = sdk + .stream_builder() + .table("catalog.schema.table") + .oauth("a", "b") + .json() + .multiplexed(0) + .validate(); + assert!(matches!( + zero, + Err(ZerobusError::InvalidArgument(msg)) if msg.contains("at least 1") + )); + + let too_many = sdk + .stream_builder() + .table("catalog.schema.table") + .oauth("a", "b") + .json() + .multiplexed(MultiplexedStream::MAX_STREAMS + 1) + .validate(); + assert!(matches!( + too_many, + Err(ZerobusError::InvalidArgument(msg)) if msg.contains("at most") + )); + } + + #[cfg(feature = "arrow-flight")] + #[test] + fn multiplexed_builder_rejects_arrow() { + use arrow_schema::{DataType, Field, Schema as ArrowSchema}; + + let sdk = test_sdk(); + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + false, + )])); + let result = sdk + .stream_builder() + .table("catalog.schema.table") + .oauth("a", "b") + .arrow(schema) + .multiplexed(2) + .validate(); + + assert!(matches!( + result, + Err(ZerobusError::InvalidArgument(msg)) if msg.contains("not supported by .multiplexed()") + )); + } + #[tokio::test] async fn build_without_auth_returns_error() { let sdk = test_sdk(); diff --git a/rust/sdk/src/lib.rs b/rust/sdk/src/lib.rs index beed5b00..28215a3a 100644 --- a/rust/sdk/src/lib.rs +++ b/rust/sdk/src/lib.rs @@ -48,7 +48,6 @@ mod default_token_factory; mod errors; mod headers_provider; mod landing_zone; -#[cfg(feature = "testing")] mod multiplexed_stream; mod offset_generator; mod proxy; @@ -91,14 +90,13 @@ use landing_zone::{Countable, LandingZone}; pub use arrow_configuration::ArrowStreamConfigurationOptions; #[cfg(feature = "arrow-flight")] pub use arrow_stream::{ArrowSchema, DataType, Field, RecordBatch, TimeUnit, ZerobusArrowStream}; -pub use builder::{StreamBuilder, ZerobusSdkBuilder}; +pub use builder::{MultiplexedStreamBuilder, StreamBuilder, ZerobusSdkBuilder}; pub use callbacks::AckCallback; pub use default_token_factory::DefaultTokenFactory; pub use errors::ZerobusError; #[cfg(feature = "testing")] pub use headers_provider::NoAuthHeadersProvider; pub use headers_provider::{HeadersProvider, OAuthHeadersProvider}; -#[cfg(feature = "testing")] pub use multiplexed_stream::{MessageId, MultiplexedStream}; pub use offset_generator::{OffsetId, OffsetIdGenerator}; pub use proxy::{ConnectorFactory, ProxyConnector}; @@ -1989,7 +1987,6 @@ impl ZerobusStream { )) } - #[cfg(feature = "testing")] pub(crate) fn has_capacity(&self) -> bool { self.landing_zone.len() < self.options.max_inflight_requests } @@ -1999,7 +1996,6 @@ impl ZerobusStream { // cancellation token and `is_closed` flag, both of which are already // interior-mutable. The `JoinHandle`s aren't reaped here; that happens in // `close` or `Drop`. - #[cfg(feature = "testing")] pub(crate) fn signal_shutdown(&self) { self.is_closed.store(true, Ordering::Relaxed); self.cancellation_token.cancel(); diff --git a/rust/sdk/src/multiplexed_stream.rs b/rust/sdk/src/multiplexed_stream.rs index 95d169dd..598c117d 100644 --- a/rust/sdk/src/multiplexed_stream.rs +++ b/rust/sdk/src/multiplexed_stream.rs @@ -1,10 +1,10 @@ //! Fan-out wrapper that distributes ingestion across multiple [`ZerobusStream`]s. //! -//! A single `ZerobusStream` is throughput-limited by its in-flight window; //! `MultiplexedStream` routes records round-robin across a fixed set of -//! sub-streams to raise aggregate throughput. When the chosen sub-stream is at -//! capacity the call awaits drain rather than rerouting, so per-sub-stream -//! ordering is preserved. +//! sub-streams for near-linear throughput scaling when the caller does not +//! require global record ordering. When the chosen sub-stream is at capacity +//! the call awaits drain rather than rerouting, so per-sub-stream ordering is +//! preserved. //! //! Each ingest returns an opaque [`MessageId`] that packs the sub-stream index //! and its offset into a single `i64` (6 bits of stream index → up to 64 @@ -83,8 +83,8 @@ impl MessageId { /// Distributes ingestion round-robin across a fixed set of [`ZerobusStream`]s. /// -/// See the [module-level documentation](self) for routing, `MessageId`, and -/// poisoning semantics. +/// See the crate documentation for routing, `MessageId`, and poisoning +/// semantics. pub struct MultiplexedStream { streams: Vec, round_robin_counter: AtomicUsize, @@ -92,6 +92,9 @@ pub struct MultiplexedStream { } impl MultiplexedStream { + /// Maximum number of sub-streams supported by a multiplexed stream. + pub const MAX_STREAMS: usize = 1 << STREAM_BITS; + /// Creates a multiplexed stream over the given sub-streams. /// /// # Panics @@ -103,9 +106,9 @@ impl MultiplexedStream { "MultiplexedStream requires at least one sub-stream" ); assert!( - streams.len() <= (1 << STREAM_BITS), + streams.len() <= Self::MAX_STREAMS, "MultiplexedStream supports at most {} sub-streams", - 1 << STREAM_BITS + Self::MAX_STREAMS ); Self { streams, From 591120e3a43fd2607b9dfe4c4b449fcf62e4ebd5 Mon Sep 17 00:00:00 2001 From: danilo-najkov-db Date: Mon, 29 Jun 2026 09:45:33 +0000 Subject: [PATCH 2/2] changes --- rust/README.md | 18 +++++++++++++----- rust/examples/proto/README.md | 27 +++++++++++++++++++-------- rust/examples/proto/multiplexed.rs | 28 +++++++++------------------- 3 files changed, 41 insertions(+), 32 deletions(-) diff --git a/rust/README.md b/rust/README.md index 47a5d25d..ac726207 100644 --- a/rust/README.md +++ b/rust/README.md @@ -520,14 +520,21 @@ let mut stream = sdk .await?; let message_id = stream.ingest_record(ProtoMessage(order)).await?; -stream.wait_for_message_id(message_id).await?; + +// As a replacement for waiting on every message ID, call flush() +// occasionally to confirm everything queued so far has been acknowledged. +// If you need to confirm one specific record, wait for its message ID instead. +// stream.wait_for_message_id(message_id).await?; + +stream.flush().await?; stream.close().await?; ``` `MultiplexedStream` returns `MessageId` values instead of `OffsetId` values -because each id records which sub-stream accepted the record. Per-sub-stream -ordering is preserved, but message ids are not globally ordered across the -multiplexed stream. Arrow Flight streams are not multiplexed by this API. +because each id records which sub-stream accepted the record and the offset +within that sub-stream. Per-sub-stream ordering is preserved, but message ids +are not globally ordered across the multiplexed stream. Arrow Flight streams +are not multiplexed by this API. Setters can be called in any order before `.multiplexed(...)` or `.build()`. The builder validates at `build()` time that both authentication and format have been configured. @@ -616,9 +623,10 @@ let mut stream = sdk .await?; let message_id = stream.ingest_record(ProtoMessage(record)).await?; -stream.wait_for_message_id(message_id).await?; +// For periodic confirmation, flush records that have already been queued. stream.flush().await?; + stream.close().await?; ``` diff --git a/rust/examples/proto/README.md b/rust/examples/proto/README.md index 1b59efc8..51460a5d 100644 --- a/rust/examples/proto/README.md +++ b/rust/examples/proto/README.md @@ -197,13 +197,10 @@ you want to distribute ingestion across multiple protobuf gRPC sub-streams. **Expected output:** ``` -=== Message ID-based API (Multiplexed) === +=== Multiplexed Protocol Buffers ingestion === [Auto-encoding] Record sent with message ID: MessageId(stream=0, offset=0) -[Auto-encoding] Record acknowledged with message ID: MessageId(stream=0, offset=0) [Pre-encoded] Record sent with message ID: MessageId(stream=1, offset=0) -[Pre-encoded] Record acknowledged with message ID: MessageId(stream=1, offset=0) [Backward-compatible] Record sent with message ID: MessageId(stream=2, offset=0) -[Backward-compatible] Record acknowledged with message ID: MessageId(stream=2, offset=0) All records flushed Stream closed successfully ``` @@ -227,14 +224,28 @@ let mut stream = sdk .await?; let message_id = stream.ingest_record(ProtoMessage(order)).await?; -stream.wait_for_message_id(message_id).await?; + +// As a replacement for waiting on every message ID, call flush() +// occasionally to confirm everything queued so far has been acknowledged. +// If you need to confirm one specific record, wait for its message ID +// instead. +// stream.wait_for_message_id(message_id).await?; + +stream.flush().await?; stream.close().await?; ``` `MultiplexedStream` returns `MessageId` values instead of `OffsetId` values -because records may be queued on different sub-streams. Per-sub-stream ordering -is preserved, but message ids are not globally ordered across the multiplexed -stream. +because records are queued on different sub-streams. A `MessageId` identifies +the sub-stream and the offset within that sub-stream, for example +`MessageId(stream=1, offset=42)`. Message IDs are not globally ordered: two +different sub-streams can both have offset `42`. + +Single records are assigned to sub-streams round-robin. `ingest_records()` +places the whole batch on one sub-stream and returns one `MessageId` for that +batch. Use `flush()` to wait for all records already queued on all sub-streams. +Use `wait_for_message_id()` only when you need to confirm one specific record +or batch. ## Adapting for Your Custom Table diff --git a/rust/examples/proto/multiplexed.rs b/rust/examples/proto/multiplexed.rs index 3fc44210..2d38d9ee 100644 --- a/rust/examples/proto/multiplexed.rs +++ b/rust/examples/proto/multiplexed.rs @@ -45,7 +45,7 @@ async fn main() -> Result<(), Box> { .build() .await?; - ingest_with_message_id_api(&stream).await?; + ingest_multiplexed(&stream).await?; stream.close().await?; println!("Stream closed successfully"); @@ -53,9 +53,8 @@ async fn main() -> Result<(), Box> { Ok(()) } -/// Multiplexed API: returns a message ID after queuing on a sub-stream. -async fn ingest_with_message_id_api(stream: &MultiplexedStream) -> Result<(), Box> { - println!("=== Message ID-based API (Multiplexed) ==="); +async fn ingest_multiplexed(stream: &MultiplexedStream) -> Result<(), Box> { + println!("=== Multiplexed Protocol Buffers ingestion ==="); // Delta TIMESTAMP is int64 microseconds since epoch UTC. let now = chrono::Utc::now().timestamp_micros(); @@ -77,11 +76,12 @@ async fn ingest_with_message_id_api(stream: &MultiplexedStream) -> Result<(), Bo "[Auto-encoding] Record sent with message ID: {}", message_id ); - stream.wait_for_message_id(message_id).await?; - println!( - "[Auto-encoding] Record acknowledged with message ID: {}", - message_id - ); + + // As a replacement for waiting on every message ID, call flush() + // occasionally to confirm everything queued so far has been acknowledged. + // If you need to confirm one specific record, wait for its message ID + // instead. + // stream.wait_for_message_id(message_id).await?; // 2. Pre-encoded: ProtoBytes - pass bytes with explicit wrapper. let order = TableOrders { @@ -98,11 +98,6 @@ async fn ingest_with_message_id_api(stream: &MultiplexedStream) -> Result<(), Bo .ingest_record(ProtoBytes(order.encode_to_vec())) .await?; println!("[Pre-encoded] Record sent with message ID: {}", message_id); - stream.wait_for_message_id(message_id).await?; - println!( - "[Pre-encoded] Record acknowledged with message ID: {}", - message_id - ); // 3. Backward-compatible: raw Vec - no wrapper needed, works the same as ProtoBytes. let order = TableOrders { @@ -120,11 +115,6 @@ async fn ingest_with_message_id_api(stream: &MultiplexedStream) -> Result<(), Bo "[Backward-compatible] Record sent with message ID: {}", message_id ); - stream.wait_for_message_id(message_id).await?; - println!( - "[Backward-compatible] Record acknowledged with message ID: {}", - message_id - ); stream.flush().await?; println!("All records flushed");