coordinator: multi-channel priority queue + stream affinity + multi-subscription consume

osprey_coordinator/src/consumer/pubsub.rs

@@ -17,7 +17,9 @@ use crate::metrics::MetricsClientBuilder;
 use crate::{
     consumer::message_decoder,
     coordinator_metrics::OspreyCoordinatorMetrics,
-    priority_queue::{AckOrNack, AckableAction, PriorityQueueSender},
+    priority_queue::{
+        AckOrNack, AckableAction, Channel as PriorityQueueChannel, PriorityQueueSender,
+    },
     proto,
     pub_sub_streaming_pull::DetachedMessage,
     pub_sub_streaming_pull::{FlowControl, SpawnTaskPerMessageHandler, StreamingPullManager},
@@ -151,15 +153,14 @@ pub async fn start_pubsub_subscriber(
     snowflake_client: Arc<SnowflakeClient>,
     priority_queue_sender: PriorityQueueSender,
     metrics: Arc<OspreyCoordinatorMetrics>,
+    subscription_id: String,
+    channel: PriorityQueueChannel,
 ) -> Result<()> {
     let subscriber_client = create_pubsub_subscription_client().await;
     let subscription_name = {
         let project_id =
             std::env::var("PUBSUB_SUBSCRIPTION_PROJECT_ID").unwrap_or("osprey-dev".to_string());
 
-        let subscription_id = std::env::var("PUBSUB_SUBSCRIPTION_ID")
-            .unwrap_or("osprey-coordinator-actions".to_string());
-
         PubSubSubscription::new(project_id, subscription_id)
     };
 
@@ -242,7 +243,7 @@ pub async fn start_pubsub_subscriber(
                 let send_start_time = Instant::now();
                 match timeout(
                     max_time_to_send_to_async_queue,
-                    priority_queue_sender.send_async(ackable_action),
+                    priority_queue_sender.send(ackable_action, channel),
                 )
                 .await
                 {

osprey_coordinator/src/coordinator_metrics.rs

@@ -11,6 +11,8 @@ define_metrics!(OspreyCoordinatorMetrics, [
     // How many messages are currently buffered in the priority queue
     priority_queue_size_sync => StaticGauge("priority_queue_size",["type" => "sync"]),
     priority_queue_size_async => StaticGauge("priority_queue_size",["type" => "async"]),
+    priority_queue_size_notif_steady => StaticGauge("priority_queue_size",["type" => "notif_steady"]),
+    priority_queue_size_notif_batch => StaticGauge("priority_queue_size",["type" => "notif_batch"]),
 
     // How many receivers are open for the priority queue
     // can be used as a proxy for number of connections open from the osprey worker

osprey_coordinator/src/main.rs

@@ -41,6 +41,69 @@ use tokio::join;
 use crate::osprey_bidirectional_stream::OspreyCoordinatorServer;
 use crate::proto::osprey_coordinator_service_server::OspreyCoordinatorServiceServer;
 
+/// Parse `OSPREY_COORDINATOR_SUBSCRIPTIONS` = "sub_a:async,sub_b:notif_steady".
+/// Falls back to `[(PUBSUB_SUBSCRIPTION_ID, Async)]` when unset (legacy behavior).
+fn parse_subscriptions(
+    raw: Option<String>,
+    legacy_sub: Option<String>,
+) -> Result<Vec<(String, priority_queue::Channel)>> {
+    use priority_queue::Channel;
+    if let Some(raw) = raw.filter(|s| !s.trim().is_empty()) {
+        raw.split(',')
+            .map(|pair| {
+                let (sub, chan) = pair.split_once(':').ok_or_else(|| {
+                    anyhow::anyhow!("bad subscription spec '{pair}', want 'sub:channel'")
+                })?;
+                let channel = match chan.trim() {
+                    "async" => Channel::Async,
+                    "notif_steady" => Channel::NotifSteady,
+                    "notif_batch" => Channel::NotifBatch,
+                    other => anyhow::bail!("unknown channel '{other}'"),
+                };
+                Ok((sub.trim().to_string(), channel))
+            })
+            .collect()
+    } else {
+        let sub = legacy_sub.ok_or_else(|| {
+            anyhow::anyhow!(
+                "neither OSPREY_COORDINATOR_SUBSCRIPTIONS nor PUBSUB_SUBSCRIPTION_ID set"
+            )
+        })?;
+        Ok(vec![(sub, Channel::Async)])
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use priority_queue::Channel;
+
+    #[test]
+    fn parses_multi_subscription_spec() {
+        let got = parse_subscriptions(Some("s1:async,s2:notif_steady,s3:notif_batch".into()), None)
+            .unwrap();
+        assert_eq!(
+            got,
+            vec![
+                ("s1".into(), Channel::Async),
+                ("s2".into(), Channel::NotifSteady),
+                ("s3".into(), Channel::NotifBatch),
+            ]
+        );
+    }
+
+    #[test]
+    fn falls_back_to_legacy_single_sub() {
+        let got = parse_subscriptions(None, Some("legacy-sub".into())).unwrap();
+        assert_eq!(got, vec![("legacy-sub".into(), Channel::Async)]);
+    }
+
+    #[test]
+    fn errors_when_no_subscription_configured() {
+        assert!(parse_subscriptions(None, None).is_err());
+    }
+}
+
 #[derive(Debug, Parser)]
 struct CliOptions {
     #[arg(
@@ -117,12 +180,25 @@ async fn main() -> Result<()> {
                 as std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send>>
         }
         Some("pubsub") => {
-            tracing::info!("starting PubSub subscriber");
-            Box::pin(start_pubsub_subscriber(
-                snowflake_client.clone(),
-                priority_queue_sender.clone(),
-                metrics.clone(),
-            ))
+            let subs = parse_subscriptions(
+                std::env::var("OSPREY_COORDINATOR_SUBSCRIPTIONS").ok(),
+                std::env::var("PUBSUB_SUBSCRIPTION_ID").ok(),
+            )?;
+            tracing::info!(?subs, "starting PubSub subscribers");
+            let mut futs = Vec::new();
+            for (sub_id, channel) in subs {
+                futs.push(start_pubsub_subscriber(
+                    snowflake_client.clone(),
+                    priority_queue_sender.clone(),
+                    metrics.clone(),
+                    sub_id,
+                    channel,
+                ));
+            }
+            Box::pin(async move {
+                futures::future::try_join_all(futs).await?;
+                Ok(())
+            })
                 as std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send>>
         }
         Some(invalid) => {