use crate::event::{ChangeEvent, ConversionError};
use bson::Document;
use futures::Stream;
use mongodb::{
change_stream::event::{ChangeStreamEvent, ResumeToken},
error::{Error as MongoError, ErrorKind as MongoErrorKind},
options::ChangeStreamOptions,
Collection,
};
use std::{
future::Future,
pin::Pin,
sync::{Arc, Mutex},
task::{Context, Poll},
time::Duration,
};
use thiserror::Error;
use tokio::sync::mpsc;
use tracing::{debug, error, info, warn};
#[derive(Debug, Error)]
pub enum StreamError {
#[error("Connection error: {message}")]
Connection {
message: String,
#[source]
source: Option<Box<dyn std::error::Error + Send + Sync>>,
code: Option<i32>,
labels: Vec<String>,
},
#[error("Event conversion failed: {0}")]
Conversion(#[from] ConversionError),
#[error("Resume token persistence failed: {0}")]
ResumeTokenPersistence(String),
#[error("Stream invalidated: {reason}")]
Invalidated { reason: String },
#[error("Max reconnection attempts ({0}) exceeded")]
MaxReconnectAttemptsExceeded(u32),
#[error("Invalid resume token (code {code}): oplog may be truncated")]
InvalidResumeToken { code: i32 },
#[error("Configuration error: {0}")]
Configuration(String),
}
impl From<MongoError> for StreamError {
fn from(err: MongoError) -> Self {
Self::from_mongo_error(err)
}
}
impl StreamError {
pub fn from_mongo_error(err: MongoError) -> Self {
let code = match err.kind.as_ref() {
MongoErrorKind::Command(cmd_err) => Some(cmd_err.code),
_ => None,
};
if code == Some(286) {
return Self::InvalidResumeToken { code: 286 };
}
let labels: Vec<String> = err.labels().iter().cloned().collect();
Self::Connection {
message: err.to_string(),
source: Some(Box::new(err)),
code,
labels,
}
}
#[must_use]
pub fn is_retryable(&self) -> bool {
match self {
Self::Connection { code, labels, .. } => {
if labels.iter().any(|l| {
l == "RetryableWriteError"
|| l == "TransientTransactionError"
|| l == "NetworkError"
}) {
return true;
}
if let Some(c) = code {
matches!(
c,
6 | 7 | 89 | 91 | 10107 | 11600 | 11602 | 13435 | 13436 | 43 )
} else {
false
}
}
Self::Invalidated { .. } => false,
Self::InvalidResumeToken { .. } => false,
Self::MaxReconnectAttemptsExceeded(_) => false,
Self::Conversion(_) => false,
Self::ResumeTokenPersistence(_) => false,
Self::Configuration(_) => false,
}
}
#[must_use]
pub fn category(&self) -> &'static str {
match self {
Self::Connection { .. } => "connection",
Self::Conversion(_) => "conversion",
Self::ResumeTokenPersistence(_) => "persistence",
Self::Invalidated { .. } => "invalidated",
Self::MaxReconnectAttemptsExceeded(_) => "max_retries",
Self::InvalidResumeToken { .. } => "invalid_token",
Self::Configuration(_) => "configuration",
}
}
}
#[derive(Debug, Clone)]
pub struct AckableEvent {
pub event: ChangeEvent,
resume_token: Document,
ack_sender: mpsc::UnboundedSender<Document>,
last_resume_token: Arc<Mutex<Option<Document>>>,
}
impl AckableEvent {
pub fn ack(self) {
if let Ok(mut last_token) = self.last_resume_token.lock() {
*last_token = Some(self.resume_token.clone());
}
if self.ack_sender.send(self.resume_token).is_err() {
warn!("Failed to send ack - persistence task may have stopped");
}
}
#[must_use]
pub fn event_ref(&self) -> &ChangeEvent {
&self.event
}
#[must_use]
pub fn into_event(self) -> ChangeEvent {
self.event
}
}
#[derive(Debug, Clone)]
pub struct ChangeStreamConfig {
pub pipeline: Vec<Document>,
pub full_document_on_update: bool,
pub full_document_before_change: bool,
pub initial_backoff_ms: u64,
pub max_backoff_ms: u64,
pub max_reconnect_attempts: u32,
pub batch_size: Option<u32>,
pub resume_token: Option<Document>,
pub backoff_jitter: f64,
}
impl Default for ChangeStreamConfig {
fn default() -> Self {
Self {
pipeline: Vec::new(),
full_document_on_update: false,
full_document_before_change: false,
initial_backoff_ms: 100,
max_backoff_ms: 30_000, max_reconnect_attempts: 5,
batch_size: None,
resume_token: None,
backoff_jitter: 0.1, }
}
}
impl ChangeStreamConfig {
#[must_use]
pub fn builder() -> ChangeStreamConfigBuilder {
ChangeStreamConfigBuilder::default()
}
pub fn validate(&self) -> Result<(), StreamError> {
if self.initial_backoff_ms == 0 {
return Err(StreamError::Configuration(
"initial_backoff_ms must be greater than 0".to_string(),
));
}
if self.initial_backoff_ms > self.max_backoff_ms {
return Err(StreamError::Configuration(format!(
"initial_backoff_ms ({}) must be <= max_backoff_ms ({})",
self.initial_backoff_ms, self.max_backoff_ms
)));
}
if !(0.0..=1.0).contains(&self.backoff_jitter) {
return Err(StreamError::Configuration(format!(
"backoff_jitter ({}) must be between 0.0 and 1.0",
self.backoff_jitter
)));
}
Ok(())
}
fn calculate_backoff(&self, attempt: u32) -> Duration {
let base_ms = self
.initial_backoff_ms
.saturating_mul(1_u64 << attempt.saturating_sub(1))
.min(self.max_backoff_ms);
if self.backoff_jitter > 0.0 {
let jitter_range = (base_ms as f64) * self.backoff_jitter;
let jitter = (rand::random::<f64>() * jitter_range) - (jitter_range / 2.0);
let final_ms = ((base_ms as f64) + jitter).max(0.0) as u64;
Duration::from_millis(final_ms)
} else {
Duration::from_millis(base_ms)
}
}
fn to_mongo_options(&self) -> ChangeStreamOptions {
let mut options = ChangeStreamOptions::default();
if self.full_document_on_update {
options.full_document = Some(mongodb::options::FullDocumentType::UpdateLookup);
}
if self.full_document_before_change {
options.full_document_before_change =
Some(mongodb::options::FullDocumentBeforeChangeType::WhenAvailable);
}
options.batch_size = self.batch_size;
if let Some(ref token_doc) = self.resume_token {
if let Ok(bytes) = bson::to_vec(token_doc) {
if let Ok(resume_token) = bson::from_slice::<ResumeToken>(&bytes) {
options.resume_after = Some(resume_token);
}
}
}
options
}
}
#[derive(Debug, Default)]
pub struct ChangeStreamConfigBuilder {
pipeline: Vec<Document>,
full_document_on_update: bool,
full_document_before_change: bool,
initial_backoff_ms: Option<u64>,
max_backoff_ms: Option<u64>,
max_reconnect_attempts: Option<u32>,
batch_size: Option<u32>,
resume_token: Option<Document>,
backoff_jitter: Option<f64>,
}
impl ChangeStreamConfigBuilder {
#[must_use]
pub fn pipeline(mut self, pipeline: Vec<Document>) -> Self {
self.pipeline = pipeline;
self
}
#[must_use]
pub fn full_document_update_lookup(mut self) -> Self {
self.full_document_on_update = true;
self
}
#[must_use]
pub fn full_document_before_change(mut self) -> Self {
self.full_document_before_change = true;
self
}
#[must_use]
pub fn initial_backoff_ms(mut self, ms: u64) -> Self {
self.initial_backoff_ms = Some(ms);
self
}
#[must_use]
pub fn max_backoff_ms(mut self, ms: u64) -> Self {
self.max_backoff_ms = Some(ms);
self
}
#[must_use]
pub fn max_reconnect_attempts(mut self, attempts: u32) -> Self {
self.max_reconnect_attempts = Some(attempts);
self
}
#[must_use]
pub fn batch_size(mut self, size: u32) -> Self {
self.batch_size = Some(size);
self
}
#[must_use]
pub fn resume_token(mut self, token: Document) -> Self {
self.resume_token = Some(token);
self
}
#[must_use]
pub fn backoff_jitter(mut self, jitter: f64) -> Self {
self.backoff_jitter = Some(jitter);
self
}
pub fn build(self) -> Result<ChangeStreamConfig, StreamError> {
let config = ChangeStreamConfig {
pipeline: self.pipeline,
full_document_on_update: self.full_document_on_update,
full_document_before_change: self.full_document_before_change,
initial_backoff_ms: self.initial_backoff_ms.unwrap_or(100),
max_backoff_ms: self.max_backoff_ms.unwrap_or(30_000),
max_reconnect_attempts: self.max_reconnect_attempts.unwrap_or(5),
batch_size: self.batch_size,
resume_token: self.resume_token,
backoff_jitter: self.backoff_jitter.unwrap_or(0.1),
};
config.validate()?;
Ok(config)
}
}
pub type ResumeTokenCallback =
Box<dyn Fn(Document) -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>;
type ReconnectFuture = Pin<
Box<
dyn Future<
Output = Result<
mongodb::change_stream::ChangeStream<ChangeStreamEvent<Document>>,
StreamError,
>,
> + Send,
>,
>;
enum StreamState {
Active,
Reconnecting(ReconnectFuture),
Closed,
}
pub struct ChangeStreamListener {
collection: Collection<Document>,
config: ChangeStreamConfig,
stream: Option<mongodb::change_stream::ChangeStream<ChangeStreamEvent<Document>>>,
resume_token_callback: ResumeTokenCallback,
state: StreamState,
reconnect_attempts: u32,
last_resume_token: Arc<Mutex<Option<Document>>>,
ack_sender: mpsc::UnboundedSender<Document>,
_persistence_task: tokio::task::JoinHandle<()>,
}
impl ChangeStreamListener {
pub async fn new<F>(
collection: Collection<Document>,
config: ChangeStreamConfig,
resume_token_callback: F,
) -> Result<Self, StreamError>
where
F: Fn(Document) -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>>
+ Send
+ Sync
+ 'static,
{
info!(
"Initializing change stream for {}.{}",
collection.namespace().db,
collection.namespace().coll
);
let options = config.to_mongo_options();
let stream = if config.pipeline.is_empty() {
collection.watch().with_options(options).await?
} else {
collection
.watch()
.pipeline(config.pipeline.clone())
.with_options(options)
.await?
};
let (ack_sender, mut ack_receiver) = mpsc::unbounded_channel::<Document>();
let resume_token_callback_clone = Box::new(resume_token_callback);
let persistence_task = tokio::spawn(async move {
while let Some(token) = ack_receiver.recv().await {
if let Err(e) = resume_token_callback_clone(token).await {
warn!("Failed to persist resume token: {}", e);
}
}
debug!("Token persistence task shutting down");
});
Ok(Self {
collection,
config,
stream: Some(stream),
resume_token_callback: Box::new(|_| Box::pin(async { Ok(()) })),
state: StreamState::Active,
reconnect_attempts: 0,
last_resume_token: Arc::new(Mutex::new(None)),
ack_sender,
_persistence_task: persistence_task,
})
}
async fn reconnect_async(
collection: Collection<Document>,
config: ChangeStreamConfig,
last_resume_token: Option<Document>,
mut reconnect_attempts: u32,
) -> Result<mongodb::change_stream::ChangeStream<ChangeStreamEvent<Document>>, StreamError>
{
reconnect_attempts += 1;
if config.max_reconnect_attempts > 0 && reconnect_attempts > config.max_reconnect_attempts {
error!(
attempts = reconnect_attempts,
"Max reconnection attempts exceeded"
);
return Err(StreamError::MaxReconnectAttemptsExceeded(
config.max_reconnect_attempts,
));
}
let backoff = config.calculate_backoff(reconnect_attempts);
warn!(
attempt = reconnect_attempts,
backoff_ms = backoff.as_millis(),
"Reconnecting to change stream"
);
tokio::time::sleep(backoff).await;
let mut options = config.to_mongo_options();
if let Some(ref token_doc) = last_resume_token {
debug!("Resuming from token: {:?}", token_doc);
if let Ok(bytes) = bson::to_vec(token_doc) {
if let Ok(resume_token) = bson::from_slice::<ResumeToken>(&bytes) {
options.resume_after = Some(resume_token);
}
}
}
let stream = if config.pipeline.is_empty() {
collection.watch().with_options(options).await?
} else {
collection
.watch()
.pipeline(config.pipeline.clone())
.with_options(options)
.await?
};
info!(
attempt = reconnect_attempts,
"Successfully reconnected to change stream"
);
Ok(stream)
}
#[allow(dead_code)]
async fn reconnect(&mut self) -> Result<(), StreamError> {
self.reconnect_attempts += 1;
if self.config.max_reconnect_attempts > 0
&& self.reconnect_attempts > self.config.max_reconnect_attempts
{
error!(
attempts = self.reconnect_attempts,
"Max reconnection attempts exceeded"
);
return Err(StreamError::MaxReconnectAttemptsExceeded(
self.config.max_reconnect_attempts,
));
}
let backoff_ms = self
.config
.initial_backoff_ms
.saturating_mul(2_u64.saturating_pow(self.reconnect_attempts - 1))
.min(self.config.max_backoff_ms);
warn!(
attempt = self.reconnect_attempts,
backoff_ms = backoff_ms,
"Reconnecting to change stream"
);
tokio::time::sleep(Duration::from_millis(backoff_ms)).await;
let mut options = self.config.to_mongo_options();
if let Ok(last_token) = self.last_resume_token.lock() {
if let Some(ref token_doc) = *last_token {
debug!("Resuming from token: {:?}", token_doc);
if let Ok(bytes) = bson::to_vec(token_doc) {
if let Ok(resume_token) = bson::from_slice::<ResumeToken>(&bytes) {
options.resume_after = Some(resume_token);
}
}
}
}
let stream = if self.config.pipeline.is_empty() {
self.collection.watch().with_options(options).await?
} else {
self.collection
.watch()
.pipeline(self.config.pipeline.clone())
.with_options(options)
.await?
};
info!(
attempt = self.reconnect_attempts,
"Successfully reconnected to change stream"
);
self.stream = Some(stream);
self.state = StreamState::Active;
self.reconnect_attempts = 0;
Ok(())
}
#[allow(dead_code)]
async fn process_event(
&mut self,
event: ChangeStreamEvent<Document>,
) -> Result<ChangeEvent, StreamError> {
let resume_token = bson::to_document(&event.id)
.map_err(|e| StreamError::ResumeTokenPersistence(e.to_string()))?;
let change_event = ChangeEvent::try_from(event)?;
if change_event.is_invalidate() {
let reason = format!(
"Collection {}.{} was dropped or renamed",
self.collection.namespace().db,
self.collection.namespace().coll
);
error!("{}", reason);
self.state = StreamState::Closed;
return Err(StreamError::Invalidated { reason });
}
(self.resume_token_callback)(resume_token.clone())
.await
.map_err(StreamError::ResumeTokenPersistence)?;
if let Ok(mut last_token) = self.last_resume_token.lock() {
*last_token = Some(resume_token);
}
Ok(change_event)
}
pub async fn close(&mut self) {
info!("Closing change stream");
self.state = StreamState::Closed;
self.stream = None;
}
}
impl Stream for ChangeStreamListener {
type Item = Result<AckableEvent, StreamError>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
if matches!(this.state, StreamState::Closed) {
return Poll::Ready(None);
}
if let StreamState::Reconnecting(ref mut reconnect_fut) = this.state {
match reconnect_fut.as_mut().poll(cx) {
Poll::Ready(Ok(new_stream)) => {
info!("Reconnection successful, resuming stream");
this.stream = Some(new_stream);
this.state = StreamState::Active;
this.reconnect_attempts = 0;
cx.waker().wake_by_ref();
return Poll::Pending;
}
Poll::Ready(Err(e)) => {
error!("Reconnection failed: {}", e);
this.state = StreamState::Closed;
return Poll::Ready(Some(Err(e)));
}
Poll::Pending => {
return Poll::Pending;
}
}
}
if let Some(ref mut stream) = this.stream {
use futures::StreamExt;
let mut stream_pin = Pin::new(stream);
match stream_pin.poll_next_unpin(cx) {
Poll::Ready(Some(Ok(event))) => {
let resume_token = match bson::to_document(&event.id) {
Ok(token) => token,
Err(e) => {
return Poll::Ready(Some(Err(StreamError::ResumeTokenPersistence(
e.to_string(),
))))
}
};
let change_event = match ChangeEvent::try_from(event) {
Ok(evt) => evt,
Err(e) => return Poll::Ready(Some(Err(StreamError::Conversion(e)))),
};
if change_event.is_invalidate() {
let reason = format!(
"Collection {}.{} was dropped or renamed",
this.collection.namespace().db,
this.collection.namespace().coll
);
error!("{}", reason);
this.state = StreamState::Closed;
return Poll::Ready(Some(Err(StreamError::Invalidated { reason })));
}
let ackable = AckableEvent {
event: change_event,
resume_token,
ack_sender: this.ack_sender.clone(),
last_resume_token: this.last_resume_token.clone(),
};
Poll::Ready(Some(Ok(ackable)))
}
Poll::Ready(Some(Err(e))) => {
let stream_err = StreamError::from_mongo_error(e);
if stream_err.is_retryable() {
warn!("Retryable error occurred: {}", stream_err);
let collection = this.collection.clone();
let config = this.config.clone();
let last_resume_token =
this.last_resume_token.lock().ok().and_then(|t| t.clone());
let reconnect_attempts = this.reconnect_attempts;
let reconnect_fut = Box::pin(async move {
Self::reconnect_async(
collection,
config,
last_resume_token,
reconnect_attempts,
)
.await
});
this.state = StreamState::Reconnecting(reconnect_fut);
cx.waker().wake_by_ref();
Poll::Pending
} else {
error!("Fatal error occurred: {}", stream_err);
this.state = StreamState::Closed;
Poll::Ready(Some(Err(stream_err)))
}
}
Poll::Ready(None) => {
warn!("Change stream ended unexpectedly");
this.state = StreamState::Closed;
Poll::Ready(None)
}
Poll::Pending => Poll::Pending,
}
} else {
this.state = StreamState::Closed;
Poll::Ready(None)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use tokio::sync::Mutex;
#[test]
fn test_default_config() {
let config = ChangeStreamConfig::default();
assert_eq!(config.pipeline.len(), 0);
assert!(!config.full_document_on_update);
assert!(!config.full_document_before_change);
assert_eq!(config.initial_backoff_ms, 100);
assert_eq!(config.max_backoff_ms, 30_000);
assert_eq!(config.max_reconnect_attempts, 5);
}
#[test]
fn test_config_builder() {
let config = ChangeStreamConfig::builder()
.full_document_update_lookup()
.full_document_before_change()
.max_reconnect_attempts(10)
.initial_backoff_ms(200)
.max_backoff_ms(60_000)
.batch_size(100)
.build()
.unwrap();
assert!(config.full_document_on_update);
assert!(config.full_document_before_change);
assert_eq!(config.max_reconnect_attempts, 10);
assert_eq!(config.initial_backoff_ms, 200);
assert_eq!(config.max_backoff_ms, 60_000);
assert_eq!(config.batch_size, Some(100));
}
#[test]
fn test_config_with_pipeline() {
use bson::doc;
let pipeline = vec![doc! { "$match": { "operationType": "insert" } }];
let config = ChangeStreamConfig::builder()
.pipeline(pipeline.clone())
.build()
.unwrap();
assert_eq!(config.pipeline.len(), 1);
assert_eq!(config.pipeline[0], pipeline[0]);
}
#[test]
fn test_stream_error_is_retryable() {
let err = StreamError::Connection {
message: "connection refused".to_string(),
source: None,
code: Some(6), labels: vec![],
};
assert!(err.is_retryable());
let err = StreamError::Invalidated {
reason: "collection dropped".to_string(),
};
assert!(!err.is_retryable());
let err = StreamError::MaxReconnectAttemptsExceeded(5);
assert!(!err.is_retryable());
}
#[test]
fn test_stream_error_category() {
let err = StreamError::Invalidated {
reason: "test".to_string(),
};
assert_eq!(err.category(), "invalidated");
let err = StreamError::MaxReconnectAttemptsExceeded(5);
assert_eq!(err.category(), "max_retries");
let err = StreamError::Configuration("test".to_string());
assert_eq!(err.category(), "configuration");
}
#[tokio::test]
async fn test_resume_token_callback() {
use bson::doc;
let saved_token = Arc::new(Mutex::new(None));
let saved_token_clone = saved_token.clone();
let callback = move |token: Document| {
let saved = saved_token_clone.clone();
Box::pin(async move {
*saved.lock().await = Some(token);
Ok(())
}) as Pin<Box<dyn Future<Output = Result<(), String>> + Send>>
};
let test_token = doc! { "_data": "test123" };
callback(test_token.clone()).await.unwrap();
let saved = saved_token.lock().await;
assert_eq!(*saved, Some(test_token));
}
#[test]
fn test_config_resume_token() {
use bson::doc;
let token = doc! { "_data": "token123" };
let config = ChangeStreamConfig::builder()
.resume_token(token.clone())
.build()
.unwrap();
assert_eq!(config.resume_token, Some(token));
}
}