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
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
//! Acknowledgment tracking: waiting for offsets and inspecting unacked records.
//!
//! Transport-agnostic: these only touch the offset generator, watch channels,
//! and the failed-records buffer. They don't move bytes; they observe progress
//! the IO tasks report back.
use std::sync::atomic::Ordering;
use tokio::time::Duration;
use tracing::{debug, error, instrument, trace};
use super::ZerobusStream;
use crate::{EncodedBatch, EncodedRecord, OffsetId, ZerobusError, ZerobusResult};
impl ZerobusStream {
/// Internal method to wait for a specific offset to be acknowledged.
/// Used by both `flush()` and `wait_for_offset()`.
async fn wait_for_offset_internal(
&self,
offset_to_wait: OffsetId,
operation_name: &str,
) -> ZerobusResult<()> {
let wait_operation = async {
let mut offset_receiver = self.logical_last_received_offset_id_tx.subscribe();
let mut error_rx = self.server_error_rx.clone();
loop {
let offset = *offset_receiver.borrow_and_update();
let stream_id = match self.stream_id.as_deref() {
Some(stream_id) => stream_id,
None => {
error!("Stream ID is None during {}", operation_name.to_lowercase());
"None"
}
};
if let Some(offset) = offset {
if offset >= offset_to_wait {
debug!(stream_id = %stream_id, "Stream is caught up to the given offset. {} completed.", operation_name);
return Ok(());
} else {
trace!(
stream_id = %stream_id,
"Stream is caught up to offset {}. Waiting for offset {}.",
offset, offset_to_wait
);
}
} else {
trace!(
stream_id = %stream_id,
"Stream is not caught up to any offset yet. Waiting for the first offset."
);
}
if self.is_closed.load(Ordering::Relaxed) {
// Re-check offset before failing, it might have been updated.
let offset = *offset_receiver.borrow_and_update();
if let Some(offset) = offset {
if offset >= offset_to_wait {
return Ok(());
}
}
// The supervisor always sends the real error to server_error_tx
// before setting is_closed=true, so check error_rx first to
// return the actual error instead of a generic one.
if let Some(server_error) = error_rx.borrow().clone() {
return Err(server_error);
}
return Err(ZerobusError::StreamClosedError(tonic::Status::internal(
format!("Stream closed during {}", operation_name.to_lowercase()),
)));
}
// Race between offset updates and server errors.
tokio::select! {
result = offset_receiver.changed() => {
// If offset_receiver channel is closed, break the loop.
if result.is_err() {
break;
}
// Loop continues to check new offset value.
}
_ = error_rx.changed() => {
// Server error occurred, return it immediately if stream is closed.
if let Some(server_error) = error_rx.borrow().clone() {
if self.is_closed.load(Ordering::Relaxed) {
// Re-check offset before failing, it might have been updated.
let offset = *offset_receiver.borrow_and_update();
if let Some(offset) = offset {
if offset >= offset_to_wait {
return Ok(());
}
}
return Err(server_error);
}
}
}
}
}
if let Some(server_error) = error_rx.borrow().clone() {
if self.is_closed.load(Ordering::Relaxed) {
return Err(server_error);
}
}
Err(ZerobusError::StreamClosedError(tonic::Status::internal(
format!("Stream closed during {}", operation_name.to_lowercase()),
)))
};
match tokio::time::timeout(
Duration::from_millis(self.options.flush_timeout_ms),
wait_operation,
)
.await
{
Ok(Ok(())) => Ok(()),
Ok(Err(e)) => Err(e),
Err(_) => {
if let Some(stream_id) = self.stream_id.as_deref() {
error!(stream_id = %stream_id, table_name = %self.table_properties.table_name, "{} timed out", operation_name);
} else {
error!(table_name = %self.table_properties.table_name, "{} timed out", operation_name);
}
Err(ZerobusError::StreamClosedError(
tonic::Status::deadline_exceeded(format!("{} timed out", operation_name)),
))
}
}
}
/// Flushes all currently pending records and waits for their acknowledgments.
///
/// This method captures the current highest offset and waits until all records up to
/// that offset have been acknowledged by the server. Records ingested during the flush
/// operation are not included in this flush.
///
/// # Returns
///
/// `Ok(())` when all pending records at the time of the call have been acknowledged.
///
/// # Errors
///
/// * `StreamClosedError` - If the stream is closed or times out
///
/// # Examples
///
/// ```no_run
/// # use databricks_zerobus_ingest_sdk::*;
/// # async fn example(stream: ZerobusStream) -> Result<(), ZerobusError> {
/// // Ingest many records
/// for i in 0..1000 {
/// let _offset = stream.ingest_record_offset(vec![i as u8]).await?;
/// }
///
/// // Wait for all to be acknowledged
/// stream.flush().await?;
/// println!("All 1000 records have been acknowledged");
/// # Ok(())
/// # }
/// ```
#[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))]
pub async fn flush(&self) -> ZerobusResult<()> {
let offset_to_wait = match self.logical_offset_id_generator.last() {
Some(offset) => offset,
None => return Ok(()), // Nothing to flush.
};
self.wait_for_offset_internal(offset_to_wait, "Flush").await
}
/// Waits for server acknowledgment of a specific logical offset.
///
/// This method blocks until the server has acknowledged the record or batch at the
/// specified offset. Use this with offsets returned from `ingest_record_offset()` or
/// `ingest_records_offset()` to explicitly control when to wait for acknowledgments.
///
/// # Arguments
///
/// * `offset` - The logical offset ID to wait for (returned from `ingest_record_offset()` or `ingest_records_offset()`)
///
/// # Returns
///
/// `Ok(())` when the record/batch at the specified offset has been acknowledged.
///
/// # Errors
///
/// * `StreamClosedError` - If the stream is closed or times out while waiting
///
/// # Examples
///
/// ```no_run
/// # use databricks_zerobus_ingest_sdk::*;
/// # async fn example(stream: ZerobusStream) -> Result<(), ZerobusError> {
/// # let my_record = vec![1, 2, 3];
/// // Ingest multiple records and collect their offsets
/// let mut offsets = Vec::new();
/// for i in 0..100 {
/// let offset = stream.ingest_record_offset(vec![i as u8]).await?;
/// offsets.push(offset);
/// }
///
/// // Wait for specific offsets
/// for offset in offsets {
/// stream.wait_for_offset(offset).await?;
/// }
/// println!("All records acknowledged");
/// # Ok(())
/// # }
/// ```
pub async fn wait_for_offset(&self, offset: OffsetId) -> ZerobusResult<()> {
self.wait_for_offset_internal(offset, "Waiting for acknowledgement")
.await
}
/// Returns all records that were ingested but not acknowledged by the server.
///
/// This method should only be called after a stream has failed or been closed.
/// It's useful for implementing custom retry logic or persisting failed records.
///
/// **Note:** This method flattens all unacknowledged records into a single iterator,
/// losing the original batch grouping.
/// If you want to preserve the batch grouping, use `ZerobusStream::get_unacked_batches()` instead.
/// If you want to re-ingest unacknowledged records while preserving their batch
/// structure, use `ZerobusSdk::recreate_stream()` instead.
///
///
/// # Returns
///
/// An iterator over individual `EncodedRecord` items. All unacknowledged records are
/// flattened into a single sequence, regardless of how they were originally ingested
/// (via `ingest_record()` or `ingest_records()`).
///
/// # Errors
///
/// * `InvalidStateError` - If called on an active (not closed) stream
///
/// # Examples
///
/// ```no_run
/// # use databricks_zerobus_ingest_sdk::*;
/// # async fn example(sdk: ZerobusSdk, mut stream: ZerobusStream) -> Result<(), ZerobusError> {
/// match stream.close().await {
/// Err(e) => {
/// // Stream failed, get unacked records
/// let unacked = stream.get_unacked_records().await?;
/// let total_records = unacked.into_iter().count();
/// println!("Failed to acknowledge {} records", total_records);
///
/// // For re-ingestion with preserved batch structure, use recreate_stream
/// let new_stream = sdk.recreate_stream(&stream).await?;
/// }
/// Ok(_) => println!("All records acknowledged"),
/// }
/// # Ok(())
/// # }
/// ```
pub async fn get_unacked_records(&self) -> ZerobusResult<impl Iterator<Item = EncodedRecord>> {
Ok(self
.get_unacked_batches()
.await?
.into_iter()
.flat_map(|batch| batch.into_iter()))
}
/// Returns all records that were ingested but not acknowledged by the server, grouped by batch.
///
/// This method should only be called after a stream has failed or been closed.
/// It's useful for implementing custom retry logic or persisting failed records.
///
/// **Note:** This method returns the unacknowledged records as a vector of `EncodedBatch` items,
/// where each batch corresponds to how records were ingested:
/// - Each `ingest_record()` call creates a single batch containing one record
/// - Each `ingest_records()` call creates a single batch containing multiple records
///
/// For alternatives, see `ZerobusStream::get_unacked_records()` and `ZerobusSdk::recreate_stream()`.
///
/// # Returns
///
/// A vector of `EncodedBatch` items. Records are grouped by their original ingestion call.
pub async fn get_unacked_batches(&self) -> ZerobusResult<Vec<EncodedBatch>> {
if self.is_closed.load(Ordering::Relaxed) {
// The supervisor only moves landing-zone records into
// `failed_records` on a stream failure. A stream torn down without
// one (flush timed out during `close`, or `signal_shutdown` from a
// poisoned MultiplexedStream) can still hold unacked records in the
// landing zone; drain them here so they are reported too and so
// repeat calls return the same result.
let mut failed = self.failed_records.write().await;
failed.extend(
self.landing_zone
.remove_all()
.into_iter()
.map(|request| request.payload),
);
return Ok(failed.clone());
}
if let Some(stream_id) = self.stream_id.as_deref() {
error!(stream_id = %stream_id, "Cannot get unacked records from an active stream. Stream must be closed first.");
} else {
error!(
"Cannot get unacked records from an active stream. Stream must be closed first."
);
}
Err(ZerobusError::InvalidStateError(
"Cannot get unacked records from an active stream. Stream must be closed first."
.to_string(),
))
}
}