numinnex commented on code in PR #3973: URL: https://github.com/apache/iggy/pull/3973#discussion_r3870460471
########## core/connectors/sinks/rabbitmq_sink/src/lib.rs: ########## @@ -0,0 +1,621 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use async_trait::async_trait; +use iggy::prelude::HeaderKind; +use iggy_connector_sdk::retry::{exponential_backoff, jitter}; +use iggy_connector_sdk::{ + ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata, sink_connector, +}; +use lapin::{ + BasicProperties, Channel, Connection, ConnectionProperties, ExchangeKind, + options::{ConfirmSelectOptions, ExchangeDeclareOptions}, + publisher_confirm::Confirmation, + types::{AMQPValue, ByteArray, FieldTable, ShortString}, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::time::Duration; +use tokio::sync::Mutex; +use tracing::{debug, info, warn}; + +sink_connector!(RabbitMQSink); + +#[derive(Debug)] +struct RabbitMqState { + connection: Connection, + channel: Channel, +} + +#[derive(Debug)] +pub struct RabbitMQSink { + id: u32, + amqp_url: SecretString, + exchange: String, + exchange_type: String, + routing_key: String, + include_metadata: bool, + verbose: bool, + durable_exchange: bool, + delivery_mode: u8, + state: Mutex<Option<RabbitMqState>>, + reconnecting: AtomicBool, + max_retries: u32, + retry_delay: Duration, + max_retry_delay: Duration, + messages_published: AtomicU64, + publish_errors: AtomicU64, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RabbitMQSinkConfig { + #[serde( + default = "default_amqp_url", + serialize_with = "iggy_common::serde_secret::serialize_secret" + )] + amqp_url: SecretString, + #[serde(default)] + exchange: Option<String>, + #[serde(default = "default_exchange_type")] + exchange_type: Option<String>, + #[serde(default)] + routing_key: Option<String>, + #[serde(default = "default_true")] + include_metadata: Option<bool>, + #[serde(default)] + verbose_logging: Option<bool>, + #[serde(default = "default_max_retries")] + max_retries: Option<u32>, + #[serde(default = "default_retry_delay_secs")] + retry_delay_secs: Option<u64>, + #[serde(default = "default_max_retry_delay_secs")] + max_retry_delay_secs: Option<u64>, + #[serde(default = "default_true")] + durable_exchange: Option<bool>, + #[serde(default = "default_delivery_mode")] + delivery_mode: Option<String>, +} + +fn default_exchange_type() -> Option<String> { + Some("topic".into()) +} + +fn default_amqp_url() -> SecretString { + SecretString::from("amqp://guest:guest@localhost:5672") +} + +fn default_delivery_mode() -> Option<String> { + Some("persistent".into()) +} + +fn default_true() -> Option<bool> { + Some(true) +} + +fn default_max_retries() -> Option<u32> { + Some(3) +} +fn default_retry_delay_secs() -> Option<u64> { + Some(1) +} +fn default_max_retry_delay_secs() -> Option<u64> { + Some(5) +} + +impl RabbitMQSink { + pub fn new(id: u32, config: RabbitMQSinkConfig) -> Self { + let delivery_mode = match config.delivery_mode.as_deref() { + Some("non_persistent") => 1, + Some("persistent") => 2, + Some(other) => { + warn!( + "Unknown delivery_mode: {other}, defaulting to persistent for connector ID: {id}" + ); + 2 + } + None => 2, + }; + RabbitMQSink { + id, + amqp_url: config.amqp_url, + exchange: config.exchange.unwrap_or_else(|| "iggy_events".into()), + exchange_type: config.exchange_type.unwrap_or_else(|| "topic".into()), + routing_key: config.routing_key.unwrap_or_else(|| "iggy.messages".into()), + include_metadata: config.include_metadata.unwrap_or(true), + verbose: config.verbose_logging.unwrap_or(false), + durable_exchange: config.durable_exchange.unwrap_or(true), + delivery_mode, + state: Mutex::new(None), + reconnecting: AtomicBool::new(false), + max_retries: config.max_retries.unwrap_or(3), + retry_delay: Duration::from_secs(config.retry_delay_secs.unwrap_or(1)), + max_retry_delay: Duration::from_secs(config.max_retry_delay_secs.unwrap_or(5)), + messages_published: AtomicU64::new(0), + publish_errors: AtomicU64::new(0), + } + } + + fn exchange_kind(&self) -> Result<ExchangeKind, Error> { + match self.exchange_type.as_str() { + "direct" => Ok(ExchangeKind::Direct), + "topic" => Ok(ExchangeKind::Topic), + "fanout" => Ok(ExchangeKind::Fanout), + "headers" => Ok(ExchangeKind::Headers), + other => Err(Error::InvalidConfigValue(format!( + "unknown exchange_type: {other}. Valid: direct, topic, fanout, headers" + ))), + } + } + + async fn publish_batch_with_retry( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], + ) -> Result<u64, Error> { + let mut attempts = 0u32; + let mut confirmed: usize = 0; + + loop { + let channel = { + let guard = self.state.lock().await; + guard + .as_ref() + .map(|s| s.channel.clone()) + .ok_or_else(|| Error::Connection("RabbitMQ not connected".into()))? + }; + + let mut last_error: Option<Error> = None; + for message in &messages[confirmed..] { + let body = message.payload.try_to_bytes()?; + let mut props = BasicProperties::default().with_delivery_mode(self.delivery_mode); + let headers = self.build_headers(topic_metadata, messages_metadata, message); + if !headers.inner().is_empty() { + props = props.with_headers(headers); + } + + let confirm = match channel + .basic_publish( + &self.exchange, + &self.routing_key, + lapin::options::BasicPublishOptions { + mandatory: true, + ..Default::default() + }, + &body, + props, + ) + .await + { + Ok(confirm) => confirm, + Err(e) => { + last_error = Some(Error::CannotStoreData(e.to_string())); + break; + } + }; + match confirm.await { + Ok(Confirmation::Ack(None)) => confirmed += 1, + Ok(Confirmation::Ack(Some(_))) | Ok(Confirmation::Nack(_)) => { + last_error = Some(Error::InvalidRecordValue( + "message returned as unroutable by RabbitMQ".into(), + )); + break; + } + Ok(Confirmation::NotRequested) => { + last_error = Some(Error::CannotStoreData( + "publisher confirms not enabled".into(), + )); + break; + } + Err(e) => { + last_error = Some(Error::CannotStoreData(format!("publish rejected: {e}"))); + break; + } + } + } + + if last_error.is_none() { + return Ok(confirmed as u64); + } + + let error = last_error.unwrap(); + attempts += 1; + + if !is_publish_retryable(&error) || attempts >= self.max_retries { + self.publish_errors + .fetch_add((messages.len() - confirmed) as u64, Ordering::Relaxed); + return Err(Error::CannotStoreData(format!( + "batch publish failed after {attempts} attempts: {error}" + ))); + } + + match self.reconnect().await { + Ok(_) => {} + Err(reconnect_error) => { + self.publish_errors + .fetch_add((messages.len() - confirmed) as u64, Ordering::Relaxed); + return Err(Error::Connection(format!( + "failed to reconnect: {reconnect_error}" + ))); + } + } + + let delay = jitter(exponential_backoff( + self.retry_delay, + attempts.saturating_sub(1), + self.max_retry_delay, + )); + warn!( + "Transient RabbitMQ publish error for connector ID: {} (attempt {attempts}/{}): {error}. Retrying in {:?}.", + self.id, self.max_retries, delay + ); + tokio::time::sleep(delay).await; + } + } + + fn build_headers( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + message: &ConsumedMessage, + ) -> FieldTable { + let mut headers = FieldTable::default(); + if let Some(user_headers) = &message.headers + && !user_headers.is_empty() + { + for (key, value) in user_headers { + let name = ShortString::from(key.to_string_value()); + let amqp_value = match value.kind() { + HeaderKind::String => AMQPValue::LongString(value.to_string_value().into()), + _ => AMQPValue::ByteArray(ByteArray::from(value.as_bytes())), + }; + headers.insert(name, amqp_value); + } + } + if self.include_metadata { + headers.insert( + "iggy_stream".into(), + AMQPValue::LongString(topic_metadata.stream.clone().into()), + ); + headers.insert( + "iggy_topic".into(), + AMQPValue::LongString(topic_metadata.topic.clone().into()), + ); + headers.insert( + "iggy_partition_id".into(), + AMQPValue::LongUInt(messages_metadata.partition_id), + ); + headers.insert( + "iggy_offset".into(), + AMQPValue::LongLongInt(message.offset as i64), + ); + } + headers + } + + async fn reconnect(&self) -> Result<(), Error> { + if self + .reconnecting + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + tokio::time::sleep(self.retry_delay).await; + return Ok(()); + } + + warn!("Reconnecting RabbitMQ sink ID: {}", self.id); + let result = async { + let conn = Connection::connect( Review Comment: No timeout on `Connection::connect` here, nor on `basic_publish` / `confirm.await` in the publish loop. RabbitMQ blocks publishers on a disk or memory alarm while still answering heartbeats, so the confirm can hang indefinitely, and a blackholed SYN here waits out the kernel TCP timeout while holding the `reconnecting` CAS. Since the FFI entry point is `block_on`, that parks a connectors-runtime worker thread. `surrealdb_sink` and `clickhouse_sink` both expose a timeout config for this. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
