fix(connectors): close the source instance a failed start leaves behind

core/connectors/runtime/src/main.rs

@@ -18,6 +18,7 @@
 use crate::configs::connectors::{
     ConnectorKey, ConnectorsConfig, ConnectorsConfigProvider, create_connectors_config_provider,
 };
+use crate::metrics::ConnectorType;
 use ::configs::ConfigProvider;
 use clap::Parser;
 use configs::connectors::ConfigFormat;
@@ -186,7 +187,7 @@ async fn main() -> Result<(), RuntimeError> {
     let mut source_wrappers = vec![];
     let mut source_containers_by_key: HashMap<String, Arc<Container<SourceApi>>> = HashMap::new();
     for (_path, source) in sources {
-        let container = Arc::new(source.container);
+        let container = source.container;
         let handle_callback = container.iggy_source_handle_v2;
         let batch_result_callback = container.iggy_source_batch_result;
         for plugin in &source.plugins {
@@ -456,8 +457,30 @@ struct SinkConnectorWrapper {
     plugins: Vec<SinkConnectorPlugin>,
 }
 
+/// Closes a plugin instance whose setup did not finish, reporting a refusal
+/// rather than returning it: every caller is already on a failure path with an
+/// error of its own to surface.
+///
+/// The two sides had this body inline, one word apart. That word is the label
+/// [`ConnectorType`] already defines, so it is taken as the enum rather than a
+/// string nothing constrains.
+pub(crate) fn close_plugin_instance(
+    close: &dyn Fn(u32) -> i32,
+    kind: ConnectorType,
+    plugin_id: u32,
+    key: &str,
+) {
+    let close_result = close(plugin_id);
+    if close_result != 0 {
+        let kind = kind.as_label();
+        warn!(
+            "iggy_{kind}_close returned {close_result} while cleaning up failed {kind} connector with ID: {plugin_id} ({key})"
+        );
+    }
+}
+
 struct SourceConnector {
-    container: Container<SourceApi>,
+    container: Arc<Container<SourceApi>>,
     plugins: Vec<SourceConnectorPlugin>,
 }
 

core/connectors/runtime/src/manager/source.rs

@@ -82,28 +82,23 @@ impl SourceManager {
         metrics: Option<&Arc<Metrics>>,
     ) {
         if let Some(source) = self.sources.get(key) {
-            let mut source = source.lock().await;
-            let old_status = source.info.status;
-            source.info.status = status;
-            if matches!(status, ConnectorStatus::Running | ConnectorStatus::Stopped) {
-                source.info.last_error = None;
-            }
-            if let Some(metrics) = metrics {
-                if old_status != ConnectorStatus::Running && status == ConnectorStatus::Running {
-                    metrics.increment_sources_running();
-                } else if old_status == ConnectorStatus::Running
-                    && status != ConnectorStatus::Running
-                {
-                    metrics.decrement_sources_running();
-                }
-            }
+            source.lock().await.apply_status(status, metrics);
         }
     }
 
-    pub async fn set_error(&self, key: &str, error_message: &str) {
+    pub async fn set_error(&self, key: &str, error_message: &str, metrics: Option<&Arc<Metrics>>) {
         if let Some(source) = self.sources.get(key) {
             let mut source = source.lock().await;
-            source.info.status = ConnectorStatus::Error;
+            // Through the shared transition, so leaving `Running` moves the
+            // gauge. Skipping it left an errored instance counted as running,
+            // and the loop's later `Stopped` could not correct that either,
+            // because by then the old status was `Error` and neither branch
+            // fires.
+            //
+            // The message is assigned after the transition, and that ordering is
+            // what preserves it. `Error` being outside the set that clears
+            // `last_error` is belt and braces here, not the mechanism.
+            source.apply_status(ConnectorStatus::Error, metrics);
             source.info.last_error = Some(ConnectorError::new(error_message));
         }
     }
@@ -236,35 +231,64 @@ impl SourceManager {
             state,
         )?;
         info!("Source connector with ID: {plugin_id} for plugin: {key} initialized successfully.");
+        // Armed from here until the id is recorded below. `SourceInstanceGuard`
+        // carries why that window strands the instance.
+        let instance_guard =
+            source::SourceInstanceGuard::for_container(container.clone(), plugin_id, key);
 
         let (producer, encoder, transforms) =
-            source::setup_source_producer(key, config, iggy_client).await?;
+            match source::setup_source_producer(key, config, iggy_client).await {
+                Ok(parts) => parts,
+                Err(error) => {
+                    // Awaited rather than left to `drop`, so this error reaches
+                    // the caller after teardown. `drop` stays the net for a
+                    // cancellation and for a `?` added here later.
+                    instance_guard.close().await;
+                    return Err(error);
+                }
+            };
 
         let handle_callback = container.iggy_source_handle_v2;
         let batch_result_callback = container.iggy_source_batch_result;
-        let handler_tasks = source::spawn_source_handler(
-            plugin_id,
-            key,
-            config.verbose,
-            config.benchmark,
-            producer,
-            encoder,
-            transforms,
-            state_storage,
-            handle_callback,
-            batch_result_callback,
-            context.clone(),
-        );
 
+        // The lock is taken before the spawn so nothing can await between
+        // registering the tasks and recording the id that reaches them. A
+        // cancellation in that gap left the `SOURCE_SENDERS` entry and both
+        // spawned tasks behind with no id naming them, and the forwarding loop
+        // then ran for the life of the process. The guard closes the plugin
+        // instance on that path but cannot reach either of those.
+        //
+        // The forwarding loop's own first act is to take this lock, so it waits
+        // for this block to end rather than racing it.
         {
             let mut details = details.lock().await;
+            // Nothing between these three statements may await. The spawn used
+            // to be passed in as a closure so the compiler refused one; inlined
+            // here that is a rule rather than a check, so keep it: an await
+            // between the spawn and the id strands the `SOURCE_SENDERS` entry
+            // and both tasks with nothing naming them.
+            details.handler_tasks = source::spawn_source_handler(
+                plugin_id,
+                key,
+                config.verbose,
+                config.benchmark,
+                producer,
+                encoder,
+                transforms,
+                state_storage,
+                handle_callback,
+                batch_result_callback,
+                context.clone(),
+            );
             details.info.id = plugin_id;
-            details.info.status = ConnectorStatus::Running;
-            details.info.last_error = None;
             details.config = config.clone();
-            details.handler_tasks = handler_tasks;
-            metrics.increment_sources_running();
+            // In the same hold as the id record, not after it. Released first,
+            // this transition raced the forwarding loop's own report and could
+            // overwrite an `Error` the loop had already set.
+            details.apply_status(ConnectorStatus::Running, Some(metrics));
         }
+        // `details.info.id` now names this instance, so a later stop reaches it.
+        instance_guard.disarm();
 
         Ok(())
     }
@@ -328,6 +352,36 @@ pub struct SourceDetails {
     pub restart_guard: Arc<Mutex<()>>,
 }
 
+impl SourceDetails {
+    /// Applies a status transition and the gauge move that belongs with it.
+    ///
+    /// On `&mut self` rather than behind a key, so it can run inside a lock the
+    /// caller already holds. `update_status` takes the lock and delegates;
+    /// `start_connector` calls it in the same hold as the id record.
+    ///
+    /// That matters: applied after releasing that lock, the initial `Running`
+    /// could land after the forwarding loop had already reported `Running` and
+    /// then failed its first batch, overwriting `Error`, clearing `last_error`,
+    /// and counting the instance a second time.
+    fn apply_status(&mut self, status: ConnectorStatus, metrics: Option<&Arc<Metrics>>) {
+        let old_status = self.info.status;
+        self.info.status = status;
+        if matches!(status, ConnectorStatus::Running | ConnectorStatus::Stopped) {
+            self.info.last_error = None;
+        }
+        let Some(metrics) = metrics else {
+            return;
+        };
+        // Only a real crossing of `Running` moves the gauge, so repeated
+        // reports of a status the connector already holds cost nothing.
+        if old_status != ConnectorStatus::Running && status == ConnectorStatus::Running {
+            metrics.increment_sources_running();
+        } else if old_status == ConnectorStatus::Running && status != ConnectorStatus::Running {
+            metrics.decrement_sources_running();
+        }
+    }
+}
+
 impl fmt::Debug for SourceDetails {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("SourceDetails")
@@ -463,6 +517,100 @@ mod tests {
         assert_eq!(metrics.get_sources_running(), 1);
     }
 
+    #[tokio::test]
+    async fn should_not_double_count_when_an_error_falls_between_two_running_reports() {
+        // The interleaving that double counts: the forwarding loop reports
+        // `Running`, fails its first batch, and a second `Running` report lands
+        // afterwards. That second report crosses into `Running` again, so it
+        // increments a gauge the error never gave back, and the instance is
+        // counted twice.
+        //
+        // What the consecutive-`Running` test cannot see: there the second
+        // report finds the status already `Running`, so no crossing happens and
+        // no arithmetic is exercised. The error in the middle is the whole point.
+        let metrics = Arc::new(Metrics::init());
+        let mut details = create_test_source_details("pg", 1);
+        details.info.status = ConnectorStatus::Stopped;
+        let manager = SourceManager::new(vec![details]);
+
+        manager
+            .update_status("pg", ConnectorStatus::Running, Some(&metrics))
+            .await;
+        manager
+            .set_error("pg", "first batch failed", Some(&metrics))
+            .await;
+        assert_eq!(
+            metrics.get_sources_running(),
+            0,
+            "an instance that has failed is not running, and the gauge has to say so \
+             or nothing later can correct it"
+        );
+
+        manager
+            .update_status("pg", ConnectorStatus::Running, Some(&metrics))
+            .await;
+
+        assert_eq!(
+            metrics.get_sources_running(),
+            1,
+            "one instance, however many times its status crossed Running"
+        );
+    }
+
+    #[tokio::test]
+    async fn should_keep_the_error_message_when_the_status_becomes_error() {
+        // `set_error` routes through a transition that clears `last_error` for
+        // some statuses, so its observable contract is worth pinning: the
+        // status ends `Error` and the message survives.
+        //
+        // Worth knowing what this does NOT pin. The message is assigned after
+        // the transition, so widening the clear to include `Error` leaves this
+        // green; checked, and the mutant survives. The ordering is the
+        // mechanism, and no unit test can see an ordering inside one lock hold.
+        let metrics = Arc::new(Metrics::init());
+        let manager = SourceManager::new(vec![create_test_source_details("pg", 1)]);
+
+        manager
+            .set_error("pg", "producer setup failed", Some(&metrics))
+            .await;
+
+        let source = manager.get("pg").await.expect("source must exist");
+        let source = source.lock().await;
+        assert_eq!(source.info.status, ConnectorStatus::Error);
+        assert_eq!(
+            source
+                .info
+                .last_error
+                .as_ref()
+                .map(|error| error.message.as_str()),
+            Some("producer setup failed"),
+            "the transition must not clear the message set right after it"
+        );
+    }
+
+    #[tokio::test]
+    async fn should_increment_metrics_once_when_running_is_reported_twice() {
+        // Both a start and the forwarding loop report `Running` for the same
+        // instance, so the gauge has to count instances rather than reports.
+        let metrics = Arc::new(Metrics::init());
+        let mut details = create_test_source_details("pg", 1);
+        details.info.status = ConnectorStatus::Stopped;
+        let manager = SourceManager::new(vec![details]);
+
+        manager
+            .update_status("pg", ConnectorStatus::Running, Some(&metrics))
+            .await;
+        manager
+            .update_status("pg", ConnectorStatus::Running, Some(&metrics))
+            .await;
+
+        assert_eq!(
+            metrics.get_sources_running(),
+            1,
+            "a second report of a status the connector already has must not move the gauge"
+        );
+    }
+
     #[tokio::test]
     async fn should_decrement_metrics_when_leaving_running() {
         let metrics = Arc::new(Metrics::init());
@@ -479,7 +627,7 @@ mod tests {
     #[tokio::test]
     async fn should_clear_error_when_status_becomes_running() {
         let manager = SourceManager::new(vec![create_test_source_details("pg", 1)]);
-        manager.set_error("pg", "some error").await;
+        manager.set_error("pg", "some error", None).await;
 
         manager
             .update_status("pg", ConnectorStatus::Running, None)
@@ -494,7 +642,7 @@ mod tests {
     async fn should_set_error_status_and_message() {
         let manager = SourceManager::new(vec![create_test_source_details("pg", 1)]);
 
-        manager.set_error("pg", "connection failed").await;
+        manager.set_error("pg", "connection failed", None).await;
 
         let source = manager.get("pg").await.unwrap();
         let details = source.lock().await;
@@ -557,7 +705,7 @@ mod tests {
     #[tokio::test]
     async fn should_clear_error_when_status_becomes_stopped() {
         let manager = SourceManager::new(vec![create_test_source_details("pg", 1)]);
-        manager.set_error("pg", "some error").await;
+        manager.set_error("pg", "some error", None).await;
 
         manager
             .update_status("pg", ConnectorStatus::Stopped, None)
@@ -609,6 +757,6 @@ mod tests {
     async fn set_error_should_be_noop_for_unknown_key() {
         let manager = SourceManager::new(vec![]);
 
-        manager.set_error("nonexistent", "some error").await;
+        manager.set_error("nonexistent", "some error", None).await;
     }
 }

core/connectors/runtime/src/metrics.rs

@@ -39,7 +39,7 @@ pub enum ConnectorType {
 }
 
 impl ConnectorType {
-    fn as_label(&self) -> &'static str {
+    pub(crate) fn as_label(&self) -> &'static str {
         match self {
             ConnectorType::Source => "source",
             ConnectorType::Sink => "sink",

core/connectors/runtime/src/sink.rs

@@ -19,10 +19,11 @@ use crate::benchmark;
 use crate::configs::connectors::SinkConfig;
 use crate::context::RuntimeContext;
 use crate::log::LOG_CALLBACK;
-use crate::metrics::{Metrics, SinkLabels};
+use crate::metrics::{ConnectorType, Metrics, SinkLabels};
 use crate::{
     FailedPlugin, PLUGIN_ID, RuntimeError, SinkApi, SinkConnector, SinkConnectorConsumer,
-    SinkConnectorPlugin, SinkConnectorWrapper, resolve_plugin_path, transform,
+    SinkConnectorPlugin, SinkConnectorWrapper, close_plugin_instance, resolve_plugin_path,
+    transform,
 };
 use dlopen2::wrapper::Container;
 use futures::StreamExt;
@@ -182,12 +183,8 @@ pub async fn init(
                 let connector = sink_connectors
                     .get_mut(&path)
                     .expect("sink connector was inserted above");
-                let close_result = (connector.container.iggy_sink_close)(plugin_id);
-                if close_result != 0 {
-                    warn!(
-                        "iggy_sink_close returned {close_result} while cleaning up failed sink connector with ID: {plugin_id} ({key})"
-                    );
-                }
+                let close = connector.container.iggy_sink_close;
+                close_plugin_instance(&|id| close(id), ConnectorType::Sink, plugin_id, &key);
                 if let Some(plugin) = connector
                     .plugins
                     .iter_mut()