1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
//! Public ingestion methods and their shared internals.
//!
//! Transport-agnostic: these methods write into the landing zone and
//! oneshot/ack plumbing. Record format dependencies (`EncodedBatch`,
//! `EncodedRecord`) flow through, but no gRPC types appear here — the
//! sender task is what eventually moves bytes onto the wire.
use std::future::Future;
use std::sync::atomic::Ordering;
use tracing::{debug, error};
use super::types::IngestRequest;
use super::ZerobusStream;
use crate::{EncodedBatch, EncodedRecord, OffsetId, ZerobusError, ZerobusResult};
impl ZerobusStream {
/// Ingests a single record and returns its logical offset directly.
///
/// This is an alternative to `ingest_record()` that returns the logical offset directly
/// as an integer (after queuing) instead of wrapping it in a Future. Use `wait_for_offset()`
/// to explicitly wait for server acknowledgment of this offset when needed.
///
/// # Arguments
///
/// * `payload` - A record that can be converted to `EncodedRecord` (either JSON string or protobuf bytes)
///
/// # Returns
///
/// The logical offset ID assigned to this record.
///
/// # Errors
///
/// * `InvalidArgument` - If the record type doesn't match stream configuration
/// * `StreamClosedError` - If the stream has been closed
///
/// # Examples
///
/// ```no_run
/// # use databricks_zerobus_ingest_sdk::*;
/// # use prost::Message;
/// # async fn example(stream: ZerobusStream) -> Result<(), ZerobusError> {
/// # let my_record = vec![1, 2, 3]; // Example protobuf-encoded data
/// // Ingest and get offset immediately
/// let offset = stream.ingest_record_offset(my_record).await?;
///
/// // Later, wait for acknowledgment
/// stream.wait_for_offset(offset).await?;
/// println!("Record at offset {} has been acknowledged", offset);
/// # Ok(())
/// # }
/// ```
pub async fn ingest_record_offset(
&self,
payload: impl Into<EncodedRecord>,
) -> ZerobusResult<OffsetId> {
let encoded_batch = EncodedBatch::try_from_record(payload, self.options.record_type)
.ok_or_else(|| {
ZerobusError::InvalidArgument(
"Record type does not match stream configuration".to_string(),
)
})?;
self.ingest_internal_v2(encoded_batch).await
}
/// Ingests a batch of records and returns the logical offset directly.
///
/// This is an alternative to `ingest_records()` that returns the logical offset directly
/// (after queuing) instead of wrapping it in a Future. Use `wait_for_offset()` to explicitly
/// wait for server acknowledgment when needed.
///
/// # Arguments
///
/// * `payload` - An iterator of records (each item should be convertible to `EncodedRecord`)
///
/// # Returns
///
/// `Some(offset_id)` for non-empty batches, or `None` if the batch is empty.
///
/// # Errors
///
/// * `InvalidArgument` - If record types don't match stream configuration
/// * `StreamClosedError` - If the stream has been closed
///
/// # Examples
///
/// ```no_run
/// # use databricks_zerobus_ingest_sdk::*;
/// # use prost::Message;
/// # async fn example(stream: ZerobusStream) -> Result<(), ZerobusError> {
/// let records = vec![vec![1, 2, 3], vec![4, 5, 6]]; // Example protobuf-encoded data
///
/// // Ingest batch and get offset immediately
/// if let Some(offset) = stream.ingest_records_offset(records).await? {
/// // Later, wait for batch acknowledgment
/// stream.wait_for_offset(offset).await?;
/// println!("Batch at offset {} has been acknowledged", offset);
/// }
/// # Ok(())
/// # }
/// ```
pub async fn ingest_records_offset<I, T>(&self, payload: I) -> ZerobusResult<Option<OffsetId>>
where
I: IntoIterator<Item = T>,
T: Into<EncodedRecord>,
{
let encoded_batch = EncodedBatch::try_from_batch(payload, self.options.record_type)
.ok_or_else(|| {
ZerobusError::InvalidArgument(
"Record type does not match stream configuration".to_string(),
)
})?;
if encoded_batch.is_empty() {
Ok(None)
} else {
self.ingest_internal_v2(encoded_batch)
.await
.map(Option::Some)
}
}
/// Internal unified method for ingesting records and batches.
///
/// Returns a future that resolves once the server has acknowledged the batch.
/// Used by `ZerobusSdk::recreate_stream` to replay unacknowledged batches.
pub(crate) async fn ingest_internal(
&self,
encoded_batch: EncodedBatch,
) -> ZerobusResult<impl Future<Output = ZerobusResult<OffsetId>>> {
if self.is_closed.load(Ordering::Relaxed) {
error!(table_name = %self.table_properties.table_name, "Stream closed");
return Err(ZerobusError::StreamClosedError(tonic::Status::internal(
"Stream closed",
)));
}
let _guard = self.sync_mutex.lock().await;
let offset_id = self.logical_offset_id_generator.next();
debug!(
offset_id = offset_id,
record_count = encoded_batch.get_record_count(),
"Ingesting record(s)"
);
if let Some(stream_id) = self.stream_id.as_ref() {
let (tx, rx) = tokio::sync::oneshot::channel();
{
let mut map = self.oneshot_map.lock().await;
map.insert(offset_id, tx);
}
self.landing_zone
.add(Box::new(IngestRequest {
payload: encoded_batch,
offset_id,
}))
.await;
let stream_id = stream_id.to_string();
Ok(async move {
rx.await.map_err(|err| {
error!(stream_id = %stream_id, "Failed to receive ack: {}", err);
ZerobusError::StreamClosedError(tonic::Status::internal(
"Failed to receive ack",
))
})?
})
} else {
error!("Stream ID is None");
Err(ZerobusError::StreamClosedError(tonic::Status::internal(
"Stream ID is None",
)))
}
}
/// Internal unified method for ingesting records and batches.
///
/// Returns the logical offset directly without waiting for acknowledgment.
/// Used by the public `ingest_*_offset` methods.
async fn ingest_internal_v2(&self, encoded_batch: EncodedBatch) -> ZerobusResult<OffsetId> {
let byte_size = encoded_batch.total_byte_size();
let max_payload_bytes = self.options.max_ingest_payload_bytes;
if byte_size > max_payload_bytes {
return Err(ZerobusError::InvalidArgument(format!(
"Ingest payload too large: {byte_size} bytes exceeds the configured limit of {max_payload_bytes} bytes"
)));
}
if self.is_closed.load(Ordering::Relaxed) {
error!(table_name = %self.table_properties.table_name, "Stream closed");
return Err(ZerobusError::StreamClosedError(tonic::Status::internal(
"Stream closed",
)));
}
let _guard = self.sync_mutex.lock().await;
let offset_id = self.logical_offset_id_generator.next();
debug!(
offset_id = offset_id,
record_count = encoded_batch.get_record_count(),
"Ingesting record(s)"
);
self.landing_zone
.add(Box::new(IngestRequest {
payload: encoded_batch,
offset_id,
}))
.await;
Ok(offset_id)
}
#[cfg(feature = "testing")]
pub(crate) fn has_capacity(&self) -> bool {
self.landing_zone.len() < self.options.max_inflight_requests
}
}