feat(postgres): Change the Postgres Adapter to be Partition Aware

pg_migrations/0001_create_inflight_activations.sql

@@ -18,3 +18,5 @@ CREATE TABLE IF NOT EXISTS inflight_taskactivations (
     taskname TEXT NOT NULL,
     on_attempts_exceeded INTEGER NOT NULL DEFAULT 1
 );
+
+CREATE INDEX idx_activation_partition ON inflight_taskactivations (partition);

src/fetch/tests.rs

@@ -50,6 +50,10 @@ impl MockStore {
 
 #[async_trait]
 impl InflightActivationStore for MockStore {
+    fn assign_partitions(&self, _partitions: Vec<i32>) -> Result<(), Error> {
+        Ok(())
+    }
+
     async fn vacuum_db(&self) -> Result<(), Error> {
         unimplemented!()
     }

src/kafka/consumer.rs

@@ -1,3 +1,4 @@
+use crate::store::inflight_activation::InflightActivationStore;
 use anyhow::{Error, anyhow};
 use futures::{
     Stream, StreamExt,
@@ -44,6 +45,7 @@ use tracing::{debug, error, info, instrument, warn};
 pub async fn start_consumer(
     topics: &[&str],
     kafka_client_config: &ClientConfig,
+    activation_store: Arc<dyn InflightActivationStore>,
     spawn_actors: impl FnMut(
         Arc<StreamConsumer<KafkaContext>>,
         &BTreeSet<(String, i32)>,
@@ -68,6 +70,7 @@ pub async fn start_consumer(
     handle_events(
         consumer,
         event_receiver,
+        activation_store,
         client_shutdown_sender,
         spawn_actors,
     )
@@ -340,6 +343,7 @@ enum ConsumerState {
 pub async fn handle_events(
     consumer: Arc<StreamConsumer<KafkaContext>>,
     events: UnboundedReceiver<(Event, SyncSender<()>)>,
+    activation_store: Arc<dyn InflightActivationStore>,
     shutdown_client: oneshot::Sender<()>,
     mut spawn_actors: impl FnMut(
         Arc<StreamConsumer<KafkaContext>>,
@@ -372,6 +376,12 @@ pub async fn handle_events(
                 state = match (state, event) {
                     (ConsumerState::Ready, Event::Assign(tpl)) => {
                         metrics::gauge!("arroyo.consumer.current_partitions").set(tpl.len() as f64);
+                        // Note: This assumes we only process one topic per consumer.
+                        let mut partitions = Vec::<i32>::new();
+                        for (_, partition) in tpl.iter() {
+                            partitions.push(*partition);
+                        }
+                        activation_store.assign_partitions(partitions).unwrap();
                         ConsumerState::Consuming(spawn_actors(consumer.clone(), &tpl), tpl)
                     }
                     (ConsumerState::Ready, Event::Revoke(_)) => {
@@ -386,11 +396,13 @@ pub async fn handle_events(
                             tpl == revoked,
                             "Revoked TPL should be equal to the subset of TPL we're consuming from"
                         );
+                        activation_store.assign_partitions(vec![]).unwrap();
                         handles.shutdown(CALLBACK_DURATION).await;
                         metrics::gauge!("arroyo.consumer.current_partitions").set(0);
                         ConsumerState::Ready
                     }
                     (ConsumerState::Consuming(handles, _), Event::Shutdown) => {
+                        activation_store.assign_partitions(vec![]).unwrap();
                         handles.shutdown(CALLBACK_DURATION).await;
                         debug!("Signaling shutdown to client...");
                         shutdown_client.take();

src/main.rs

@@ -164,6 +164,7 @@ async fn main() -> Result<(), Error> {
             start_consumer(
                 &[&consumer_config.kafka_topic],
                 &consumer_config.kafka_consumer_config(),
+                consumer_store.clone(),
                 processing_strategy!({
                     err:
                         OsStreamWriter::new(

src/store/inflight_activation.rs

@@ -375,21 +375,12 @@ pub struct DepthCounts {
 
 #[async_trait]
 pub trait InflightActivationStore: Send + Sync {
-    /// Trigger incremental vacuum to reclaim free pages in the database
-    async fn vacuum_db(&self) -> Result<(), Error>;
-
-    /// Perform a full vacuum on the database
-    async fn full_vacuum_db(&self) -> Result<(), Error>;
-
-    /// Get the size of the database in bytes
-    async fn db_size(&self) -> Result<u64, Error>;
-
-    /// Get an activation by id
-    async fn get_by_id(&self, id: &str) -> Result<Option<InflightActivation>, Error>;
-
+    /// CONSUMER OPERATIONS
     /// Store a batch of activations
     async fn store(&self, batch: Vec<InflightActivation>) -> Result<QueryResult, Error>;
 
+    fn assign_partitions(&self, partitions: Vec<i32>) -> Result<(), Error>;
+
     /// Get `limit` pending activations, optionally filtered by namespaces and bucket subrange.
     /// If `mark_processing` is true, sets status to `Processing` and `processing_deadline`; otherwise `Claimed` and `claim_expires_at`.
     /// If no limit is provided, all matching activations will be returned.
@@ -458,6 +449,14 @@ pub trait InflightActivationStore: Send + Sync {
     /// Record successful push.
     async fn mark_activation_processing(&self, id: &str) -> Result<(), Error>;
 
+    /// Update the status of a specific activation
+    async fn set_status(
+        &self,
+        id: &str,
+        status: InflightActivationStatus,
+    ) -> Result<Option<InflightActivation>, Error>;
+
+    /// COUNT OPERATIONS
     /// Get the age of the oldest pending activation in seconds
     async fn pending_activation_max_lag(&self, now: &DateTime<Utc>) -> f64;
 
@@ -473,6 +472,10 @@ pub trait InflightActivationStore: Send + Sync {
     /// Count all activations
     async fn count(&self) -> Result<usize, Error>;
 
+    /// ACTIVATION OPERATIONS
+    /// Get an activation by id
+    async fn get_by_id(&self, id: &str) -> Result<Option<InflightActivation>, Error>;
+
     /// Queue depths for pending, delay, and processing (writer backpressure and upkeep gauges).
     /// Default implementation uses separate calls, but stores may override with a single query.
     async fn count_depths(&self) -> Result<DepthCounts, Error> {
@@ -491,13 +494,6 @@ pub trait InflightActivationStore: Send + Sync {
         })
     }
 
-    /// Update the status of a specific activation
-    async fn set_status(
-        &self,
-        id: &str,
-        status: InflightActivationStatus,
-    ) -> Result<Option<InflightActivation>, Error>;
-
     /// Set the processing deadline for a specific activation
     async fn set_processing_deadline(
         &self,
@@ -508,12 +504,20 @@ pub trait InflightActivationStore: Send + Sync {
     /// Delete an activation by id
     async fn delete_activation(&self, id: &str) -> Result<(), Error>;
 
+    /// DATABASE OPERATIONS
+    /// Trigger incremental vacuum to reclaim free pages in the database
+    async fn vacuum_db(&self) -> Result<(), Error>;
+
+    /// Perform a full vacuum on the database
+    async fn full_vacuum_db(&self) -> Result<(), Error>;
+
+    /// Get the size of the database in bytes
+    async fn db_size(&self) -> Result<u64, Error>;
+
+    /// UPKEEP OPERATIONS
     /// Get all activations with status Retry
     async fn get_retry_activations(&self) -> Result<Vec<InflightActivation>, Error>;
 
-    /// Clear all activations from the store
-    async fn clear(&self) -> Result<(), Error>;
-
     /// Revert expired push claims back to pending status.
     async fn handle_claim_expiration(&self) -> Result<u64, Error>;
 
@@ -541,6 +545,10 @@ pub trait InflightActivationStore: Send + Sync {
     /// Remove killswitched tasks
     async fn remove_killswitched(&self, killswitched_tasks: Vec<String>) -> Result<u64, Error>;
 
+    /// TEST OPERATIONS
+    /// Clear all activations from the store
+    async fn clear(&self) -> Result<(), Error>;
+
     /// Remove the database, used only in tests
     async fn remove_db(&self) -> Result<(), Error> {
         Ok(())
@@ -810,6 +818,11 @@ impl InflightActivationStore for SqliteActivationStore {
         Ok(Some(row.into()))
     }
 
+    fn assign_partitions(&self, partitions: Vec<i32>) -> Result<(), Error> {
+        warn!("assign_partitions: {:?}", partitions);
+        Ok(())
+    }
+
     #[instrument(skip_all)]
     async fn store(&self, batch: Vec<InflightActivation>) -> Result<QueryResult, Error> {
         if batch.is_empty() {

src/store/inflight_activation_tests.rs

@@ -815,6 +815,105 @@ async fn test_set_activation_status(#[case] adapter: &str) {
     store.remove_db().await.unwrap();
 }
 
+#[tokio::test]
+#[rstest]
+#[case::postgres("postgres")]
+async fn test_set_activation_status_with_partitions(#[case] adapter: &str) {
+    let store = create_test_store(adapter).await;
+
+    let mut batch = make_activations(2);
+    batch[1].partition = 1;
+    assert!(store.store(batch).await.is_ok());
+    assert_counts(
+        StatusCount {
+            pending: 1,
+            ..StatusCount::default()
+        },
+        store.as_ref(),
+    )
+    .await;
+
+    assert!(
+        store
+            .set_status("id_0", InflightActivationStatus::Failure)
+            .await
+            .is_ok()
+    );
+    assert_counts(
+        StatusCount {
+            failure: 1,
+            ..StatusCount::default()
+        },
+        store.as_ref(),
+    )
+    .await;
+
+    assert!(
+        store
+            .set_status("id_0", InflightActivationStatus::Pending)
+            .await
+            .is_ok()
+    );
+    assert_counts(
+        StatusCount {
+            pending: 1,
+            ..StatusCount::default()
+        },
+        store.as_ref(),
+    )
+    .await;
+    assert!(
+        store
+            .set_status("id_0", InflightActivationStatus::Failure)
+            .await
+            .is_ok()
+    );
+    assert!(
+        store
+            .set_status("id_1", InflightActivationStatus::Failure)
+            .await
+            .is_ok()
+    );
+    // The broker can update the status of an activation in a different partition, but
+    // it still should not be counted in its upkeep.
+    assert_counts(
+        StatusCount {
+            pending: 0,
+            failure: 1,
+            ..StatusCount::default()
+        },
+        store.as_ref(),
+    )
+    .await;
+    assert!(
+        store
+            .claim_activations(None, None, Some(1), None, true)
+            .await
+            .unwrap()
+            .is_empty()
+    );
+
+    let result = store
+        .set_status("not_there", InflightActivationStatus::Complete)
+        .await;
+    assert!(result.is_ok(), "no query error");
+
+    let activation = result.unwrap();
+    assert!(activation.is_none(), "no activation found");
+
+    let result = store
+        .set_status("id_0", InflightActivationStatus::Complete)
+        .await;
+    assert!(result.is_ok(), "no query error");
+
+    let result_opt = result.unwrap();
+    assert!(result_opt.is_some(), "activation should be returned");
+    let inflight = result_opt.unwrap();
+    assert_eq!(inflight.id, "id_0");
+    assert_eq!(inflight.status, InflightActivationStatus::Complete);
+    store.remove_db().await.unwrap();
+}
+
 #[tokio::test]
 #[rstest]
 #[case::sqlite("sqlite")]