chore(cyclotron): modify janitor metric, expose depth limit (#25060)

plugin-server/package.json

@@ -147,6 +147,6 @@
     },
     "cyclotron": {
         "//This is a short term workaround to ensure that cyclotron changes trigger a rebuild": true,
-        "version": "0.1.3"
+        "version": "0.1.4"
     }
 }

plugin-server/src/cdp/cdp-consumers.ts

@@ -592,8 +592,10 @@ export class CdpProcessedEventsConsumer extends CdpConsumerBase {
             handleBatch: (messages) => this._handleKafkaBatch(messages),
         })
 
+        const shardDepthLimit = this.hub.CYCLOTRON_SHARD_DEPTH_LIMIT ?? 1000000
+
         this.cyclotronManager = this.hub.CYCLOTRON_DATABASE_URL
-            ? new CyclotronManager({ shards: [{ dbUrl: this.hub.CYCLOTRON_DATABASE_URL }] })
+            ? new CyclotronManager({ shards: [{ dbUrl: this.hub.CYCLOTRON_DATABASE_URL }], shardDepthLimit })
             : undefined
 
         await this.cyclotronManager?.connect()

plugin-server/src/config/config.ts

@@ -197,6 +197,8 @@ export function getDefaultConfig(): PluginsServerConfig {
             : isDevEnv()
             ? 'postgres://posthog:posthog@localhost:5432/cyclotron'
             : '',
+
+        CYCLOTRON_SHARD_DEPTH_LIMIT: 1000000,
     }
 }
 

plugin-server/src/types.ts

@@ -287,6 +287,7 @@ export interface PluginsServerConfig extends CdpConfig {
     SESSION_RECORDING_KAFKA_CONSUMPTION_STATISTICS_EVENT_INTERVAL_MS: number
 
     CYCLOTRON_DATABASE_URL: string
+    CYCLOTRON_SHARD_DEPTH_LIMIT: number
 }
 
 export interface Hub extends PluginsServerConfig {

rust/cyclotron-core/src/config.rs

@@ -31,7 +31,7 @@ impl PoolConfig {
     }
 }
 
-pub const DEFAULT_QUEUE_DEPTH_LIMIT: u64 = 10_000;
+pub const DEFAULT_QUEUE_DEPTH_LIMIT: u64 = 1_000_000;
 pub const DEFAULT_SHARD_HEALTH_CHECK_INTERVAL: u64 = 10;
 
 #[derive(Debug, Serialize, Deserialize)]

rust/cyclotron-core/src/janitor.rs

@@ -59,7 +59,7 @@ impl Janitor {
         Ok(poison.len() as u64)
     }
 
-    pub async fn waiting_jobs(&self) -> Result<u64, QueueError> {
+    pub async fn waiting_jobs(&self) -> Result<Vec<(u64, String)>, QueueError> {
         count_total_waiting_jobs(&self.pool).await
     }
 

rust/cyclotron-core/src/manager.rs

@@ -216,7 +216,8 @@ impl Shard {
         }
 
         let pending = count_total_waiting_jobs(&self.pool).await?;
-        let is_full = pending >= self.depth_limit;
+        let total_pending = pending.iter().map(|(count, _)| count).sum::<u64>();
+        let is_full = total_pending >= self.depth_limit;
         if !is_full {
             *last_healthy = Utc::now();
         }

rust/cyclotron-core/src/ops/meta.rs

@@ -6,18 +6,26 @@ use crate::{
     DEAD_LETTER_QUEUE,
 };
 
-pub async fn count_total_waiting_jobs<'c, E>(executor: E) -> Result<u64, QueueError>
+pub async fn count_total_waiting_jobs<'c, E>(executor: E) -> Result<Vec<(u64, String)>, QueueError>
 where
     E: sqlx::Executor<'c, Database = sqlx::Postgres>,
 {
-    let res = sqlx::query!(
-        "SELECT COUNT(*) FROM cyclotron_jobs WHERE state = 'available' AND scheduled <= NOW()",
+    struct Count {
+        count: Option<i64>,
+        queue_name: String,
+    }
+
+    let res = sqlx::query_as!(
+        Count,
+        "SELECT COUNT(*), queue_name FROM cyclotron_jobs WHERE state = 'available' AND scheduled <= NOW() GROUP BY queue_name",
     )
-    .fetch_one(executor)
+    .fetch_all(executor)
     .await?;
 
-    let res = res.count.unwrap_or(0);
-    Ok(res as u64)
+    Ok(res
+        .into_iter()
+        .map(|r| (r.count.unwrap_or(0) as u64, r.queue_name))
+        .collect())
 }
 
 // Returns an InvalidLock error if the query run did not affect any rows.

rust/cyclotron-janitor/src/janitor.rs

@@ -135,7 +135,13 @@ impl Janitor {
             let _time = common_metrics::timing_guard(AVAILABLE_DEPTH_TIME, &self.metrics_labels);
             self.inner.waiting_jobs().await?
         };
-        common_metrics::gauge(AVAILABLE_DEPTH, &self.metrics_labels, available as f64);
+
+        let mut available_labels = self.metrics_labels.clone();
+        for (count, queue_name) in available {
+            available_labels.push(("queue_name".to_string(), queue_name));
+            common_metrics::gauge(AVAILABLE_DEPTH, &available_labels, count as f64);
+            available_labels.pop();
+        }
 
         let dlq_depth = {
             let _time = common_metrics::timing_guard(DLQ_DEPTH_TIME, &self.metrics_labels);

rust/cyclotron-node/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@posthog/cyclotron",
-    "version": "0.1.3",
+    "version": "0.1.4",
     "description": "Node bindings for cyclotron",
     "main": "dist/index.js",
     "types": "dist/index.d.ts",

rust/cyclotron-node/src/manager.ts

@@ -5,14 +5,15 @@ import { convertToInternalPoolConfig, serializeObject } from './helpers'
 import { CyclotronJobInit, CyclotronPoolConfig } from './types'
 
 export class CyclotronManager {
-    constructor(private config: { shards: CyclotronPoolConfig[] }) {
+    constructor(private config: { shards: CyclotronPoolConfig[], shardDepthLimit: number }) {
         this.config = config
     }
 
     async connect(): Promise<void> {
         return await cyclotron.maybeInitManager(
             JSON.stringify({
                 shards: this.config.shards.map((shard) => convertToInternalPoolConfig(shard)),
+                shard_depth_limit: this.config.shardDepthLimit,
             })
         )
     }