Add scheduler stats to axiom/datadog integration. (#38288)

Add to function event and rename scheduler_lag event to scheduler_stats (keeping the old event type for backwards compatibility).

Also updated docs.

GitOrigin-RevId: 6310ad16c8dabc2c7925c34909abc452fe9720cf

Author: nipunn1313

crates/application/src/function_log.rs

@@ -35,6 +35,7 @@ use common::{
         FunctionEventSource,
         LogEvent,
         LogSender,
+        SchedulerInfo,
         StructuredLogEvent,
     },
     runtime::{
@@ -163,6 +164,7 @@ pub struct FunctionExecution {
     // Number of retries prior to a successful execution. Only applicable for mutations.
     pub mutation_retry_count: Option<usize>,
 
+    // If this execution resulted in an OCC error, this will be Some.
     pub occ_info: Option<OccInfo>,
 }
 
@@ -265,6 +267,7 @@ impl FunctionExecution {
                     }),
                     None => None,
                 },
+                scheduler_info: self.caller.parent_scheduled_job().map(|_| SchedulerInfo {}),
                 usage_stats: log_streaming::AggregatedFunctionUsageStats {
                     database_read_bytes: self.usage_stats.database_read_bytes,
                     database_write_bytes: self.usage_stats.database_write_bytes,
@@ -1111,11 +1114,16 @@ impl<RT: Runtime> FunctionExecutionLog<RT> {
 
     /// Indicates that as of now (`timestamp`), the next scheduled job is at
     /// `next_job_ts` (None if there are no pending jobs)
-    pub fn log_scheduled_job_lag(&self, next_job_ts: Option<SystemTime>, timestamp: SystemTime) {
-        if let Err(mut e) = self
-            .inner
-            .lock()
-            .log_scheduled_job_lag(next_job_ts, timestamp)
+    pub fn log_scheduled_job_stats(
+        &self,
+        next_job_ts: Option<SystemTime>,
+        timestamp: SystemTime,
+        num_running_jobs: u64,
+    ) {
+        if let Err(mut e) =
+            self.inner
+                .lock()
+                .log_scheduled_job_stats(next_job_ts, timestamp, num_running_jobs)
         {
             report_error_sync(&mut e);
         }
@@ -1670,10 +1678,11 @@ impl<RT: Runtime> Inner<RT> {
         Ok(())
     }
 
-    fn log_scheduled_job_lag(
+    fn log_scheduled_job_stats(
         &mut self,
         next_job_ts: Option<SystemTime>,
         now: SystemTime,
+        num_running_jobs: u64,
     ) -> anyhow::Result<()> {
         let name = scheduled_job_next_ts_metric();
         // -Infinity means there is no scheduled job
@@ -1686,6 +1695,15 @@ impl<RT: Runtime> Inner<RT> {
                 },
             }]);
         }
+        if value > 0.0 || num_running_jobs > 0 {
+            self.log_manager.send_logs(vec![LogEvent {
+                timestamp: UnixTimestamp::from_system_time(now).context("now < UNIX_EPOCH?")?,
+                event: StructuredLogEvent::SchedulerStats {
+                    lag_seconds: Duration::from_secs_f32(value.max(0.0)),
+                    num_running_jobs,
+                },
+            }]);
+        }
         match self.metrics.add_gauge(name, now, value) {
             Ok(()) => (),
             Err(UdfMetricsError::SamplePrecedesCutoff { ts: _, cutoff }) => {

crates/application/src/scheduled_jobs/mod.rs

@@ -157,6 +157,8 @@ pub struct ScheduledJobExecutor<RT: Runtime> {
     next_job_ready_time: Option<Timestamp>,
     job_finished_tx: mpsc::Sender<ResolvedDocumentId>,
     job_finished_rx: mpsc::Receiver<ResolvedDocumentId>,
+    /// The last time we logged stats, used to rate limit logging
+    last_stats_log: SystemTime,
 }
 
 #[derive(Clone)]
@@ -196,7 +198,7 @@ impl<RT: Runtime> ScheduledJobExecutor<RT> {
             mpsc::channel(*SCHEDULED_JOB_EXECUTION_PARALLELISM);
         let mut executor = Self {
             context: ScheduledJobContext {
-                rt,
+                rt: rt.clone(),
                 database,
                 runner,
                 function_log,
@@ -206,6 +208,7 @@ impl<RT: Runtime> ScheduledJobExecutor<RT> {
             next_job_ready_time: None,
             job_finished_tx,
             job_finished_rx,
+            last_stats_log: rt.system_time(),
         };
         let mut backoff = Backoff::new(*SCHEDULED_JOB_INITIAL_BACKOFF, *SCHEDULED_JOB_MAX_BACKOFF);
         tracing::info!("Starting scheduled job executor");
@@ -244,11 +247,13 @@ impl<RT: Runtime> ScheduledJobExecutor<RT> {
             self.query_and_start_jobs(&mut tx).await?
         };
 
-        metrics::log_num_running_jobs(self.running_job_ids.len());
         let now = self.context.rt.system_time();
         let next_job_ready_time = self.next_job_ready_time.map(SystemTime::from);
-        self.context
-            .log_scheduled_job_execution_lag(next_job_ready_time, now);
+        // Only log stats if at least 30 seconds have elapsed since the last log
+        if now.duration_since(self.last_stats_log).unwrap_or_default() >= Duration::from_secs(30) {
+            self.log_scheduled_job_stats(next_job_ready_time, now);
+            self.last_stats_log = now;
+        }
         let next_job_future = if let Some(next_job_ts) = next_job_ready_time {
             let wait_time = next_job_ts.duration_since(now).unwrap_or_else(|_| {
                 // If we're behind, re-run this loop every 5 seconds to log the gauge above and
@@ -287,6 +292,22 @@ impl<RT: Runtime> ScheduledJobExecutor<RT> {
         Ok(())
     }
 
+    fn log_scheduled_job_stats(&self, next_job_ready_time: Option<SystemTime>, now: SystemTime) {
+        metrics::log_num_running_jobs(self.running_job_ids.len());
+        if let Some(next_job_ts) = next_job_ready_time {
+            metrics::log_scheduled_job_execution_lag(
+                now.duration_since(next_job_ts).unwrap_or(Duration::ZERO),
+            );
+        } else {
+            metrics::log_scheduled_job_execution_lag(Duration::ZERO);
+        }
+        self.context.function_log.log_scheduled_job_stats(
+            next_job_ready_time,
+            now,
+            self.running_job_ids.len() as u64,
+        );
+    }
+
     /// Reads through scheduled jobs in timestamp ascending order and starts any
     /// that are allowed by our concurrency limit and the jobs' scheduled
     /// time.
@@ -393,22 +414,6 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
         }
     }
 
-    fn log_scheduled_job_execution_lag(
-        &self,
-        next_job_ready_time: Option<SystemTime>,
-        now: SystemTime,
-    ) {
-        if let Some(next_job_ts) = next_job_ready_time {
-            metrics::log_scheduled_job_execution_lag(
-                now.duration_since(next_job_ts).unwrap_or(Duration::ZERO),
-            );
-        } else {
-            metrics::log_scheduled_job_execution_lag(Duration::ZERO);
-        }
-        self.function_log
-            .log_scheduled_job_lag(next_job_ready_time, now);
-    }
-
     // This handles re-running the scheduled function on transient errors. It
     // guarantees that the job was successfully run or the job state changed.
     pub async fn execute_job(&self, job: ScheduledJob, job_id: ResolvedDocumentId) {

crates/common/src/log_streaming.rs

@@ -84,6 +84,11 @@ pub struct OccInfo {
     pub retry_count: u64,
 }
 
+// Nothing yet. Can add information like parent scheduled job, scheduler lag,
+// etc.
+#[derive(Serialize, Debug, Clone)]
+pub struct SchedulerInfo {}
+
 #[derive(Debug, Clone)]
 pub enum StructuredLogEvent {
     /// Topic for verification logs. These are issued on sink startup and are
@@ -103,6 +108,7 @@ pub enum StructuredLogEvent {
         execution_time: Duration,
         usage_stats: AggregatedFunctionUsageStats,
         occ_info: Option<OccInfo>,
+        scheduler_info: Option<SchedulerInfo>,
     },
     /// Topic for exceptions. These happen when a UDF raises an exception from
     /// JS
@@ -118,6 +124,12 @@ pub enum StructuredLogEvent {
         action: String,
         metadata: serde_json::Map<String, JsonValue>,
     },
+    /// Topic for global stats from the scheduler. For function-specific stats,
+    /// look in FunctionExecution
+    SchedulerStats {
+        lag_seconds: Duration,
+        num_running_jobs: u64,
+    },
     ScheduledJobLag {
         lag_seconds: Duration,
     },
@@ -254,7 +266,8 @@ impl LogEvent {
                     error,
                     execution_time,
                     usage_stats,
-                    ..
+                    occ_info: _,
+                    scheduler_info: _,
                 } => {
                     let (reason, status) = match error {
                         Some(err) => (Some(err.to_string()), "failure"),
@@ -307,6 +320,14 @@ impl LogEvent {
                         "actionMetadata": metadata
                     })
                 },
+                StructuredLogEvent::SchedulerStats {
+                    lag_seconds,
+                    num_running_jobs,
+                } => serialize_map!({
+                    "_timestamp": ms,
+                    "_topic":  "_scheduler_stats",
+                    "lag_seconds": lag_seconds.as_secs(), "num_running_jobs": num_running_jobs
+                }),
                 StructuredLogEvent::ScheduledJobLag { lag_seconds } => {
                     serialize_map!({
                         "_timestamp": ms,
@@ -351,6 +372,7 @@ impl LogEvent {
                     execution_time,
                     usage_stats,
                     occ_info,
+                    scheduler_info,
                 } => {
                     let function_source = source.to_json_map();
                     let (status, error_message) = match error {
@@ -376,6 +398,7 @@ impl LogEvent {
                         "status": status,
                         "error_message": error_message,
                         "occ_info": occ_info,
+                        "scheduler_info": scheduler_info,
                         "usage": Usage {
                             database_read_bytes: usage_stats.database_read_bytes,
                             database_write_bytes: usage_stats.database_write_bytes,
@@ -422,6 +445,17 @@ impl LogEvent {
                         "audit_log_metadata": serde_json::to_string(metadata).map_err(serde::ser::Error::custom)?
                     })
                 },
+                StructuredLogEvent::SchedulerStats {
+                    lag_seconds,
+                    num_running_jobs,
+                } => {
+                    serialize_map!({
+                        "topic": "scheduler_stats",
+                        "timestamp": ms,
+                        "lag_seconds": lag_seconds.as_secs(),
+                        "num_running_jobs": num_running_jobs
+                    })
+                },
                 StructuredLogEvent::ScheduledJobLag { lag_seconds } => {
                     serialize_map!({
                         "timestamp": ms,

npm-packages/docs/docs/production/integrations/log-streams/log-streams.mdx

@@ -164,6 +164,9 @@ Schema:
     `table_name`
   - `retry_count`: the number of previously failed attempts before the current
     function execution
+- `scheduler_info`: object, if set, indicates that the function execution was
+  originally invoked by a scheduled job either directly, or via a subfunction
+  execution.
 - `usage`:
   - `database_read_bytes`: number
   - `database_write_bytes`: number, this and `database_read_bytes` make up the
@@ -220,17 +223,18 @@ The following fields are added under `function` for all `console` and
   [request ID](/functions/debugging.mdx#finding-relevant-logs-by-request-id) of
   the function.
 
-### `scheduled_job_lag` events
+### `scheduler_stats` events
 
-These events are periodically sent by the scheduler reporting the delay in
-executing your scheduled functions.
+These events are periodically sent by the scheduler reporting statistics from
+the scheduled function executor.
 
 Schema:
 
-- `topic`: `"scheduled_job_lag"`
+- `topic`: `"scheduler_stats"`
 - `timestamp`: Unix epoch timestamp in milliseconds
 - `lag_seconds`: The difference between `timestamp` and the scheduled run time
   of the oldest overdue scheduled job, in seconds.
+- `num_running_jobs`: number, the number of scheduled jobs currently running
 
 ### `audit_log` events