jiengup commented on code in PR #3886: URL: https://github.com/apache/iggy/pull/3886#discussion_r4134126577
########## core/connectors/sources/iggy_source/src/lib.rs: ########## @@ -0,0 +1,1197 @@ +// 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::{ + Client, CompressionAlgorithm, Consumer, Identifier, IggyClient, IggyError, IggyMessage, + MessageClient, PollingStrategy, StreamClient, TopicClient, TopicCreateOptions, +}; +use iggy_connector_sdk::{ + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, + retry::{parse_duration, retry_backoff}, + source::SourceBatchResult, + source_connector, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use std::{ + collections::HashMap, + fmt, + str::FromStr, + sync::atomic::{AtomicU64, Ordering}, + time::Duration, +}; +use tokio::{sync::Mutex, time::sleep}; +use tracing::{debug, error, info, warn}; + +source_connector!(IggySource); + +const CONNECTOR_NAME: &str = "Iggy source"; +const DEFAULT_POLL_INTERVAL: &str = "2s"; +const DEFAULT_RETRY_INTERVAL: &str = "1s"; +const DEFAULT_MAX_RETRY_INTERVAL: &str = "60s"; +const DEFAULT_BATCH_SIZE: u32 = 100; +const MAX_BATCH_SIZE: u32 = 10_000; +const AUTO_CREATED_PARTITIONS_COUNT: u32 = 1; + +/// Configuration for the Iggy source connector, replicating a topic from an +/// upstream Iggy cluster. `connection_string` points at the upstream cluster. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct IggySourceConfig { + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub connection_string: SecretString, + pub upstream_stream: String, + pub upstream_topic: String, + pub poll_interval: Option<String>, + pub batch_size: Option<u32>, + pub initial_offset: Option<String>, + pub include_user_headers: Option<bool>, + #[serde(default)] + pub malformed_message_policy: MalformedMessagePolicy, + pub retry_interval: Option<String>, + pub max_retry_interval: Option<String>, + pub verbose_logging: Option<bool>, +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MalformedMessagePolicy { + #[default] + Block, + DropHeaders, +} + +impl fmt::Display for MalformedMessagePolicy { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Block => formatter.write_str("block"), + Self::DropHeaders => formatter.write_str("drop_headers"), + } + } +} + +/// Starting point for a partition that has no saved offset in the state yet. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InitialOffset { + Earliest, + Latest, + Offset(u64), +} + +impl FromStr for InitialOffset { + type Err = (); + + fn from_str(s: &str) -> Result<Self, Self::Err> { + match s.trim().to_lowercase().as_str() { + "earliest" => Ok(Self::Earliest), + "latest" => Ok(Self::Latest), + other => other.parse::<u64>().map(Self::Offset).map_err(|_| ()), + } + } +} + +/// Committed state. `offsets` maps each upstream partition to the offset of +/// the last message confirmed by the runtime after both the downstream send +/// and checkpoint save succeeded. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct State { + offsets: HashMap<u32, u64>, + messages_synced: u64, + errors_count: u64, + topic_created_at: Option<u64>, +} + +impl State { + fn reset_offsets_if_topic_changed(&mut self, created_at: u64) -> bool { + if self.topic_created_at == Some(created_at) { + return false; + } + + // Old offsets can be past the new topic's frontier, where polls return + // empty batches and silently skip its history. + let reset_offsets = !self.offsets.is_empty(); + self.offsets.clear(); + self.topic_created_at = Some(created_at); + reset_offsets + } +} + +#[derive(Debug, Default)] +struct PollCycleCounts { + successful_upstream_polls: u64, + upstream_poll_errors: u64, + conversion_errors: u64, +} + +impl PollCycleCounts { + fn total_errors(&self) -> u64 { + self.upstream_poll_errors + self.conversion_errors + } + + fn requires_backoff(&self) -> bool { + self.upstream_poll_errors > 0 && self.successful_upstream_polls == 0 + } +} + +#[derive(Debug)] +pub struct IggySource { + id: u32, + config: IggySourceConfig, + client: Option<IggyClient>, + state: Mutex<State>, + pending_state: Mutex<Option<State>>, + partitions: Vec<u32>, + stream_id: Option<Identifier>, + topic_id: Option<Identifier>, + poll_interval: Duration, + retry_interval: Duration, + max_retry_interval: Duration, + consecutive_failed_poll_cycles: AtomicU64, + initial_offset: InitialOffset, + batch_size: u32, + include_user_headers: bool, + malformed_message_policy: MalformedMessagePolicy, + verbose: bool, +} + +impl IggySource { + pub fn new(id: u32, config: IggySourceConfig, state: Option<ConnectorState>) -> Self { + let verbose = config.verbose_logging.unwrap_or(false); + let restored_state = state + .and_then(|s| s.deserialize::<State>(CONNECTOR_NAME, id)) + .inspect(|s| { + info!( + "Restored state for {CONNECTOR_NAME} connector ID: {id}. \ + Offsets: {:?}, messages synced: {}, errors: {}", + s.offsets, s.messages_synced, s.errors_count + ); + }); + + let poll_interval = parse_duration(config.poll_interval.as_deref(), DEFAULT_POLL_INTERVAL); + let retry_interval = + parse_duration(config.retry_interval.as_deref(), DEFAULT_RETRY_INTERVAL); + let max_retry_interval = parse_duration( + config.max_retry_interval.as_deref(), + DEFAULT_MAX_RETRY_INTERVAL, + ); + + let initial_offset = config + .initial_offset + .as_deref() + .map(InitialOffset::from_str) + .and_then(Result::ok) + .unwrap_or_else(|| { + warn!( + "Invalid initial offset {:?} for {CONNECTOR_NAME} connector ID: {id}, \ + defaulting to earliest", + config.initial_offset + ); + InitialOffset::Earliest + }); + + let batch_size = config.batch_size.unwrap_or(DEFAULT_BATCH_SIZE); + let include_user_headers = config.include_user_headers.unwrap_or(true); + let malformed_message_policy = config.malformed_message_policy; + + IggySource { + id, + config, + client: None, + state: Mutex::new(restored_state.unwrap_or(State { + offsets: HashMap::new(), + messages_synced: 0, + errors_count: 0, + topic_created_at: None, + })), + pending_state: Mutex::new(None), + partitions: Vec::new(), + stream_id: None, + topic_id: None, + poll_interval, + retry_interval, + max_retry_interval, + consecutive_failed_poll_cycles: AtomicU64::new(0), + initial_offset, + batch_size, + include_user_headers, + malformed_message_policy, + verbose, + } + } + + fn serialize_state(&self, state: &State) -> Option<ConnectorState> { + ConnectorState::serialize(state, CONNECTOR_NAME, self.id) + } + + async fn ensure_stream_and_topic( + &self, + client: &IggyClient, + stream_id: &Identifier, + topic_id: &Identifier, + ) -> Result<(), Error> { + match client.get_stream(stream_id).await { + Ok(Some(_)) => {} + Ok(None) => { + warn!( + "Upstream stream '{}' does not exist, creating it for {CONNECTOR_NAME} \ + connector ID: {}", + self.config.upstream_stream, self.id + ); + match client.create_stream(&self.config.upstream_stream).await { + Ok(_) | Err(IggyError::StreamNameAlreadyExists(_)) => {} + Err(e) => { + return Err(Error::InitError(format!( + "Failed to create upstream stream '{}': {e}", + self.config.upstream_stream + ))); + } + } + } + Err(e) => { + return Err(Error::InitError(format!( + "Failed to check upstream stream '{}': {e}", + self.config.upstream_stream + ))); + } + } + + match client.get_topic(stream_id, topic_id).await { + Ok(Some(_)) => Ok(()), + Ok(None) => { + warn!( + "Upstream topic '{}' does not exist, creating it with \ + {AUTO_CREATED_PARTITIONS_COUNT} partition(s) for {CONNECTOR_NAME} \ + connector ID: {}", + self.config.upstream_topic, self.id + ); + match client + .create_topic( + stream_id, + &self.config.upstream_topic, + &TopicCreateOptions { + partitions_count: Some(AUTO_CREATED_PARTITIONS_COUNT), + compression_algorithm: Some(CompressionAlgorithm::None), + ..TopicCreateOptions::default() + }, + ) + .await + { + Ok(_) | Err(IggyError::TopicNameAlreadyExists(_, _)) => {} + Err(e) => { + return Err(Error::InitError(format!( + "Failed to create upstream topic '{}': {e}", + self.config.upstream_topic + ))); + } + } + Ok(()) + } + Err(e) => Err(Error::InitError(format!( + "Failed to check upstream topic '{}': {e}", + self.config.upstream_topic + ))), + } + } +} + +#[async_trait] +impl Source for IggySource { + async fn open(&mut self) -> Result<(), Error> { + validate_batch_size(self.batch_size)?; + + let redacted = redact_connection_string(self.config.connection_string.expose_secret()); + info!( + "Opening {CONNECTOR_NAME} connector ID: {}, upstream: {}/{} at {}", + self.id, self.config.upstream_stream, self.config.upstream_topic, redacted + ); + + let client = + IggyClient::from_connection_string(self.config.connection_string.expose_secret()) + .map_err(|e| { + Error::InitError(format!("Failed to parse upstream connection string: {e}")) + })?; + + client.connect().await.map_err(|e| { + Error::InitError(format!("Failed to connect to upstream Iggy cluster: {e}")) + })?; + + let stream_id = Identifier::named(&self.config.upstream_stream).map_err(|_| { + Error::InvalidConfigValue(format!( + "Invalid upstream stream name '{}'", + self.config.upstream_stream + )) + })?; + let topic_id = Identifier::named(&self.config.upstream_topic).map_err(|_| { + Error::InvalidConfigValue(format!( + "Invalid upstream topic name '{}'", + self.config.upstream_topic + )) + })?; + + self.stream_id = Some(stream_id.clone()); + self.topic_id = Some(topic_id.clone()); + + self.ensure_stream_and_topic(&client, &stream_id, &topic_id) + .await?; + + let topic = client + .get_topic(&stream_id, &topic_id) + .await + .map_err(|e| Error::InitError(format!("Failed to fetch upstream topic details: {e}")))? + .ok_or_else(|| { + Error::InitError(format!( + "Upstream topic '{}/{}' not found after creation", + self.config.upstream_stream, self.config.upstream_topic + )) + })?; + self.partitions = topic + .partitions + .iter() + .map(|partition| partition.id) + .collect(); + + self.client = Some(client); + info!( + "Opened {CONNECTOR_NAME} connector ID: {}, partitions: {}, initial offset: {:?}, \ + poll interval: {:?}, batch size: {}, malformed message policy: {}", + self.id, + self.partitions.len(), + self.initial_offset, + self.poll_interval, + self.batch_size, + self.malformed_message_policy + ); + Ok(()) + } + + async fn poll(&self) -> Result<ProducedMessages, Error> { + sleep(self.poll_interval).await; + + let client = self + .client + .as_ref() + .ok_or_else(|| Error::InitError("Upstream client not connected".to_string()))?; + let stream_id = self + .stream_id + .as_ref() + .ok_or_else(|| Error::InitError("Upstream stream not initialized".to_string()))?; + let topic_id = self + .topic_id + .as_ref() + .ok_or_else(|| Error::InitError("Upstream topic not initialized".to_string()))?; + + let failed_cycles = self.consecutive_failed_poll_cycles.load(Ordering::Relaxed); + if failed_cycles > 0 { + let delay = retry_backoff( + self.retry_interval, + failed_cycles as u32, + self.max_retry_interval, + ); + debug!( + "Backing off for {delay:?} after {failed_cycles} consecutive poll cycles with no \ + successful partitions for \ + {CONNECTOR_NAME} connector ID: {}", + self.id + ); + sleep(delay).await; + } + + let mut candidate_state = self.state.lock().await.clone(); + let topic = client + .get_topic(stream_id, topic_id) + .await + .map_err(|error| { + Error::Connection(format!("Failed to fetch upstream topic details: {error}")) + })?; + let current_partitions = topic.as_ref().map(|topic| { + if candidate_state.reset_offsets_if_topic_changed(topic.created_at.as_micros()) { + warn!( + "Upstream topic '{}/{}' changed for {CONNECTOR_NAME} connector ID: {}, \ + resetting saved offsets to initial offset", + self.config.upstream_stream, self.config.upstream_topic, self.id + ); + } + topic + .partitions + .iter() + .map(|partition| partition.id) + .collect::<Vec<_>>() + }); + let partitions = current_partitions.as_deref().unwrap_or(&self.partitions); + let consumer = Consumer::default(); + + let mut messages = Vec::with_capacity(self.batch_size as usize); + let mut cycle_counts = PollCycleCounts::default(); + + for &partition_id in partitions { + let strategy = next_strategy( + self.initial_offset, + candidate_state.offsets.get(&partition_id).copied(), + ); + let polled = client + .poll_messages( + stream_id, + topic_id, + Some(partition_id), + &consumer, + &strategy, + self.batch_size, + false, + ) + .await; + + match polled { + Ok(polled) => { + cycle_counts.successful_upstream_polls += 1; + if polled.messages.is_empty() { + continue; + } + let PartitionBuildOutcome { + messages: partition_messages, + last_produced_offset, + errors: conversion_errors, + } = build_produced_messages( + &polled.messages, + self.include_user_headers, + self.malformed_message_policy, + ); + + cycle_counts.conversion_errors += conversion_errors.len() as u64; + for conversion_error in conversion_errors { + error!( + "Failed to convert upstream message at offset {} on partition \ + {partition_id} for {CONNECTOR_NAME} connector ID: {}. \ + Malformed message policy: {}. {}", + conversion_error.offset, + self.id, + self.malformed_message_policy, + conversion_error.error + ); + } + + messages.extend(partition_messages); Review Comment: I kept `batch_size` as the per-partition request limit and added a fixed 10,000-message limit for each poll cycle. The final request is reduced to the remaining capacity, and the next cycle resumes at the first partition that was not polled. I updated the README to clarify that this bounds message count rather than payload bytes, and added boundary tests for the limit and partition cursor -- 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]
