fix: produce only on available partitions (#4559)

* fix: produce only on available partitions

* chore: move available partition check to a cache

* chore: terminate flag to indicate end loop

Author: fraidev

crates/fluvio/src/producer/mod.rs

@@ -7,17 +7,24 @@
 //! to specific topics and partitions, respectively.
 //!
 use std::collections::HashMap;
+use std::sync::atomic::AtomicU32;
 use std::sync::Arc;
 
+use event_listener::Event;
+use tokio::select;
 use tracing::instrument;
 use async_lock::RwLock;
 use anyhow::Result;
 
+use fluvio_future::task::spawn;
+use fluvio_future::timer::sleep;
 use fluvio_protocol::record::ReplicaKey;
 use fluvio_protocol::record::Record;
 use fluvio_compression::Compression;
 #[cfg(feature = "compress")]
 use fluvio_sc_schema::topic::CompressionAlgorithm;
+use fluvio_sc_schema::topic::TopicSpec;
+use fluvio_sc_schema::partition::PartitionSpec;
 use fluvio_types::PartitionId;
 use fluvio_types::event::StickyEvent;
 
@@ -36,6 +43,7 @@ pub use fluvio_protocol::record::{RecordKey, RecordData};
 
 use crate::spu::SpuPool;
 use crate::spu::SpuSocketPool;
+use crate::sync::StoreContext;
 use crate::FluvioError;
 use crate::metrics::ClientMetrics;
 use crate::producer::accumulator::{RecordAccumulator, PushRecord};
@@ -234,6 +242,7 @@ where
     topic: String,
     spu_pool: Arc<S>,
     record_accumulator: Arc<RecordAccumulator>,
+    partition_tracker: Arc<PartitionAvailabilityTracker>,
     producer_pool: Arc<RwLock<ProducerPool>>,
     metrics: Arc<ClientMetrics>,
 }
@@ -249,15 +258,16 @@ where
     }
 
     async fn push_record(self: Arc<Self>, record: Record) -> Result<PushRecord> {
-        let topics = self.spu_pool.topics();
+        let partition_count = self.partition_tracker.partition_count();
+        let available_partitions = self.partition_tracker.available_partitions();
+        let available_partitions_lock = available_partitions.read().await;
 
-        let topic_spec = topics
-            .lookup_by_key(&self.topic)
-            .await?
-            .ok_or_else(|| FluvioError::TopicNotFound(self.topic.to_string()))?
-            .spec;
-        let partition_count = topic_spec.partitions();
-        let partition_config = PartitionerConfig { partition_count };
+        let partition_config = PartitionerConfig {
+            partition_count,
+            available_partitions: available_partitions_lock.clone(),
+        };
+
+        drop(available_partitions_lock);
 
         let key = record.key.as_ref().map(|k| k.as_ref());
         let value = record.value.as_ref();
@@ -402,6 +412,103 @@ cfg_if::cfg_if! {
     }
 }
 
+/// Tracks the availability of partitions for a given topic
+struct PartitionAvailabilityTracker {
+    topic_name: String,
+    available_partitions: Arc<RwLock<Vec<PartitionId>>>,
+    partition_count: AtomicU32,
+    partitions: StoreContext<PartitionSpec>,
+    topics: StoreContext<TopicSpec>,
+    terminate: Arc<Event>,
+}
+
+impl Drop for PartitionAvailabilityTracker {
+    fn drop(&mut self) {
+        self.terminate.notify(usize::MAX);
+    }
+}
+
+impl PartitionAvailabilityTracker {
+    // Spawn a task to update available partitions
+    fn start(
+        initial_partition_count: u32,
+        topic_name: String,
+        partitions: StoreContext<PartitionSpec>,
+        topics: StoreContext<TopicSpec>,
+    ) -> Arc<Self> {
+        let tracker = Arc::new(Self {
+            topic_name,
+            available_partitions: Arc::new(RwLock::new(vec![])),
+            partition_count: AtomicU32::new(initial_partition_count),
+            partitions: partitions.clone(),
+            topics: topics.clone(),
+            terminate: Arc::new(Event::new()),
+        });
+        let shared_tracker = tracker.clone();
+
+        spawn(async move {
+            loop {
+                select! {
+                    _ = async {
+                        shared_tracker.update_available_partitions().await;
+                        sleep(std::time::Duration::from_secs(1)).await;
+                    } => {}
+                    _ = shared_tracker.terminate.listen() => {
+                        break;
+                    },
+                }
+            }
+        });
+
+        tracker
+    }
+
+    async fn update_available_partitions(&self) {
+        let mut available_partitions = vec![];
+        if let Ok(Some(topic)) = self.topics.lookup_by_key(&self.topic_name).await {
+            let partition_count = topic.spec().partitions();
+            // Update the partition count if it has changed
+            if partition_count
+                != self
+                    .partition_count
+                    .load(std::sync::atomic::Ordering::Relaxed)
+            {
+                self.partition_count
+                    .store(partition_count, std::sync::atomic::Ordering::Relaxed);
+            }
+
+            for partition_id in 0..partition_count {
+                if let Ok(Some(partition)) = self
+                    .partitions
+                    .lookup_by_key(&ReplicaKey::new(&self.topic_name, partition_id))
+                    .await
+                {
+                    if partition.status.is_online() {
+                        available_partitions.push(partition_id);
+                    }
+                }
+            }
+            // Update the shared available partitions if they have changed
+            {
+                let read_lock = self.available_partitions.read().await;
+                if available_partitions == *read_lock {
+                    return; // No change needed
+                }
+            }
+            *self.available_partitions.write().await = available_partitions;
+        }
+    }
+
+    fn available_partitions(&self) -> Arc<RwLock<Vec<PartitionId>>> {
+        self.available_partitions.clone()
+    }
+
+    fn partition_count(&self) -> u32 {
+        self.partition_count
+            .load(std::sync::atomic::Ordering::Relaxed)
+    }
+}
+
 impl<S> TopicProducer<S>
 where
     S: SpuPool + Send + Sync + 'static,
@@ -412,12 +519,13 @@ where
         config: Arc<TopicProducerConfig>,
         metrics: Arc<ClientMetrics>,
     ) -> Result<Self> {
-        let topics = spu_pool.topics();
-        let topic_spec: fluvio_sc_schema::topic::TopicSpec = topics
+        let topic_store = spu_pool.topics();
+        let topic_spec = topic_store
             .lookup_by_key(&topic)
             .await?
             .ok_or_else(|| FluvioError::TopicNotFound(topic.to_string()))?
             .spec;
+
         let partition_count = topic_spec.partitions();
 
         cfg_if::cfg_if! {
@@ -435,6 +543,9 @@ where
             partition_count,
             compression,
         );
+
+        let partitions = spu_pool.partitions().clone();
+
         let producer_pool = ProducerPool::new(
             config.clone(),
             topic.clone(),
@@ -444,13 +555,21 @@ where
             config.callback.clone(),
         );
 
+        let partition_tracker = PartitionAvailabilityTracker::start(
+            partition_count,
+            topic.clone(),
+            partitions,
+            spu_pool.topics().clone(),
+        );
+
         Ok(Self {
             inner: Arc::new(InnerTopicProducer {
                 config,
                 topic,
                 spu_pool,
                 producer_pool: Arc::new(RwLock::new(producer_pool)),
                 record_accumulator: Arc::new(record_accumulator),
+                partition_tracker,
                 metrics: metrics.clone(),
             }),
             #[cfg(feature = "smartengine")]
@@ -641,11 +760,16 @@ mod tests {
     use std::sync::Arc;
 
     use async_trait::async_trait;
-    use fluvio_protocol::record::RecordKey;
-    use fluvio_sc_schema::{partition::PartitionSpec, store::MetadataStoreObject, topic::TopicSpec};
+    use fluvio_future::timer::sleep;
+    use fluvio_protocol::record::{RecordKey, ReplicaKey};
+    use fluvio_sc_schema::{
+        partition::{PartitionSpec},
+        store::MetadataStoreObject,
+        topic::TopicSpec,
+    };
     use fluvio_socket::{ClientConfig, SocketError, StreamSocket, VersionedSerialSocket};
     use fluvio_stream_dispatcher::metadata::local::LocalMetadataItem;
-    use fluvio_types::SpuId;
+    use fluvio_types::{PartitionId, SpuId};
 
     use crate::{
         metrics::ClientMetrics,
@@ -656,6 +780,7 @@ mod tests {
 
     struct SpuPoolMock {
         topics: StoreContext<TopicSpec>,
+        partitions: StoreContext<PartitionSpec>,
     }
 
     #[async_trait]
@@ -691,7 +816,7 @@ mod tests {
         }
 
         fn partitions(&self) -> &StoreContext<PartitionSpec> {
-            todo!()
+            &self.partitions
         }
     }
 
@@ -708,10 +833,21 @@ mod tests {
                 (partitions_count, 2, false).into(), // 2 partitions, 2 replicas, not ignore rack
             ),
         ];
-
+        let partition_2 = vec![
+            MetadataStoreObject::<PartitionSpec, LocalMetadataItem>::with_spec(
+                ReplicaKey::new(topic.clone(), 0 as PartitionId),
+                vec![0, 1].into(),
+            ),
+            MetadataStoreObject::<PartitionSpec, LocalMetadataItem>::with_spec(
+                ReplicaKey::new(topic.clone(), 1 as PartitionId),
+                vec![0, 1].into(),
+            ),
+        ];
         let topics = StoreContext::<TopicSpec>::new();
-        let spu_pool = Arc::new(SpuPoolMock { topics });
+        let partitions = StoreContext::<PartitionSpec>::new();
+        let spu_pool = Arc::new(SpuPoolMock { topics, partitions });
         spu_pool.topics().store().sync_all(topic_2_partitions).await;
+        spu_pool.partitions().store().sync_all(partition_2).await;
         let producer = TopicProducer::new(topic.clone(), spu_pool.clone(), config, metrics)
             .await
             .expect("producer");
@@ -750,6 +886,8 @@ mod tests {
 
         spu_pool.topics.store().sync_all(topic_3_partitions).await;
 
+        sleep(std::time::Duration::from_secs(2)).await;
+
         let _ = producer
             .send(RecordKey::NULL, "789".to_string())
             .await

crates/fluvio/src/producer/partitioning.rs

@@ -24,9 +24,11 @@ pub trait Partitioner {
 
 pub struct PartitionerConfig {
     pub partition_count: PartitionCount,
+    pub available_partitions: Vec<PartitionCount>,
 }
 
 impl PartitionerConfig {
+    /// Creates a new `PartitionerConfig` with the given partition count and available partitions.
     pub fn partition_count(&self) -> PartitionCount {
         self.partition_count
     }
@@ -60,8 +62,12 @@ impl Partitioner for SiphashRoundRobinPartitioner {
             None => {
                 // Atomic increment. This will wrap on overflow, which is fine
                 // because we are only interested in the modulus anyway
-                let partition = self.index.fetch_add(1, Ordering::Relaxed);
-                partition % config.partition_count
+                let index = self.index.fetch_add(1, Ordering::Relaxed);
+                if config.available_partitions.is_empty() {
+                    return index % config.partition_count;
+                }
+                let partition = index as usize % config.available_partitions.len();
+                config.available_partitions[partition]
             }
         }
     }
@@ -110,7 +116,10 @@ mod tests {
     /// Ensure that feeding keyless records one-at-a-time does not assign the same partition
     #[test]
     fn test_round_robin_individual() {
-        let config = PartitionerConfig { partition_count: 3 };
+        let config = PartitionerConfig {
+            partition_count: 3,
+            available_partitions: vec![0, 1, 2],
+        };
         let partitioner = SiphashRoundRobinPartitioner::new();
 
         let key1_partition = partitioner.partition(&config, None, &[]);
@@ -133,7 +142,10 @@ mod tests {
 
         let (tx, rx) = std::sync::mpsc::channel();
         let partitioner = Arc::new(SiphashRoundRobinPartitioner::new());
-        let config = Arc::new(PartitionerConfig { partition_count: 4 });
+        let config = Arc::new(PartitionerConfig {
+            partition_count: 4,
+            available_partitions: vec![0, 1, 2, 3],
+        });
 
         // We have 5 threads calculating partitions 400 times each for NULL key (aka round-robin).
         // This is 20,000 records total, among 4 partitions. If it is evenly distributed like we
@@ -162,4 +174,26 @@ mod tests {
             assert_eq!(count, 500);
         }
     }
+
+    #[test]
+    fn test_available_partitions() {
+        let config = PartitionerConfig {
+            partition_count: 3,
+            available_partitions: vec![0, 2], // Only partitions 0 and 2 are available
+        };
+        let partitioner = SiphashRoundRobinPartitioner::new();
+
+        let key1_partition = partitioner.partition(&config, None, &[]);
+        assert_eq!(key1_partition, 0);
+        let key2_partition = partitioner.partition(&config, None, &[]);
+        assert_eq!(key2_partition, 2);
+        let key3_partition = partitioner.partition(&config, None, &[]);
+        assert_eq!(key3_partition, 0);
+        let key4_partition = partitioner.partition(&config, None, &[]);
+        assert_eq!(key4_partition, 2);
+        let key5_partition = partitioner.partition(&config, None, &[]);
+        assert_eq!(key5_partition, 0);
+        let key6_partition = partitioner.partition(&config, None, &[]);
+        assert_eq!(key6_partition, 2);
+    }
 }