lazily parse scheduled job args (#31053)

currently we parse udf args from Bytes -> ConvexArray whenever we read from the `_scheduled_jobs` table.
This can be unnecessary, for example `SchedulerModel::complete` reads the job, parses it, changes the state, and replaces the document. In that flow we don't need to parse and re-serialize udf arguments.

Only parse arguments when necessary: when returning a virtual table document, when running the function, or when reporting a system error to the logs.

GitOrigin-RevId: bdfadf8ea2adaa9e1ec3d42fbcb60397314f9aed

Author: ldanilek

crates/application/src/scheduled_jobs/mod.rs

@@ -536,7 +536,7 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
                 self.function_log.log_mutation_system_error(
                     &error,
                     path,
-                    job.udf_args,
+                    job.udf_args()?,
                     identity,
                     self.rt.monotonic_now(),
                     caller,
@@ -585,7 +585,7 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
                             )
                             .into(),
                             path,
-                            job.udf_args,
+                            job.udf_args()?,
                             identity,
                             self.rt.monotonic_now(),
                             caller,
@@ -603,7 +603,7 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
                             )
                             .into(),
                             path,
-                            job.udf_args,
+                            job.udf_args()?,
                             identity,
                             self.rt.monotonic_now(),
                             caller,
@@ -636,12 +636,13 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
         let namespace = tx.table_mapping().tablet_namespace(job_id.tablet_id)?;
         let path = job.path.clone();
 
+        let udf_args = job.udf_args()?;
         let result = self
             .runner
             .run_mutation_no_udf_log(
                 tx,
                 PublicFunctionPath::Component(path.clone()),
-                job.udf_args.clone(),
+                udf_args.clone(),
                 caller.allowed_visibility(),
                 context.clone(),
             )
@@ -650,13 +651,7 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
             Ok(r) => r,
             Err(e) => {
                 self.function_log.log_mutation_system_error(
-                    &e,
-                    path,
-                    job.udf_args.clone(),
-                    identity,
-                    start,
-                    caller,
-                    context,
+                    &e, path, udf_args, identity, start, caller, context,
                 )?;
                 return Err(e);
             },
@@ -751,7 +746,7 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
                     .runner
                     .run_action_no_udf_log(
                         PublicFunctionPath::Component(path),
-                        job.udf_args,
+                        job.udf_args()?,
                         identity,
                         caller,
                         usage_tracker.clone(),
@@ -800,7 +795,7 @@ impl<RT: Runtime> ScheduledJobContext<RT> {
                 self.function_log.log_action_system_error(
                     &JsError::from_message(message).into(),
                     path,
-                    job.udf_args.clone(),
+                    job.udf_args()?,
                     identity.into(),
                     self.rt.monotonic_now(),
                     caller,

crates/model/src/scheduled_jobs/mod.rs

@@ -219,18 +219,18 @@ impl<'a, RT: Runtime> SchedulerModel<'a, RT> {
 
         let now: Timestamp = self.tx.runtime().generate_timestamp()?;
         let original_scheduled_ts: Timestamp = ts.as_system_time().try_into()?;
-        let scheduled_job = ScheduledJob {
-            path: path.clone(),
-            udf_args: args.clone(),
-            state: ScheduledJobState::Pending,
+        let scheduled_job = ScheduledJob::new(
+            path.clone(),
+            args.clone(),
+            ScheduledJobState::Pending,
             // Don't set next_ts in the past to avoid scheduler incorrectly logging
             // it is falling behind. We should keep `original_scheduled_ts` intact
             // since this is exposed to the developer via the virtual table.
-            next_ts: Some(original_scheduled_ts.max(now)),
-            completed_ts: None,
+            Some(original_scheduled_ts.max(now)),
+            None,
             original_scheduled_ts,
-            attempts: ScheduledJobAttempts::default(),
-        };
+            ScheduledJobAttempts::default(),
+        )?;
         let job = if let Some(parent_scheduled_job) = context.parent_scheduled_job {
             let table_mapping = self.tx.table_mapping();
             let parent_scheduled_job = parent_scheduled_job
@@ -245,15 +245,15 @@ impl<'a, RT: Runtime> SchedulerModel<'a, RT> {
                     | ScheduledJobState::Success => scheduled_job,
                     ScheduledJobState::Canceled => {
                         let scheduled_ts = self.tx.begin_timestamp();
-                        ScheduledJob {
+                        ScheduledJob::new(
                             path,
-                            udf_args: args,
-                            state: ScheduledJobState::Canceled,
-                            next_ts: None,
-                            completed_ts: Some(*scheduled_ts),
-                            original_scheduled_ts: *scheduled_ts,
-                            attempts: ScheduledJobAttempts::default(),
-                        }
+                            args,
+                            ScheduledJobState::Canceled,
+                            None,
+                            Some(*scheduled_ts),
+                            *scheduled_ts,
+                            ScheduledJobAttempts::default(),
+                        )?
                     },
                 }
             } else {

crates/model/src/scheduled_jobs/types.rs

@@ -27,9 +27,12 @@ pub struct ScheduledJob {
     pub path: CanonicalizedComponentFunctionPath,
     #[cfg_attr(
         any(test, feature = "testing"),
-        proptest(strategy = "proptest::arbitrary::any_with::<ConvexArray>((0..4).into())")
+        proptest(
+            strategy = "proptest::arbitrary::any_with::<ConvexArray>((0..4).into()).\
+                        prop_map(args_to_bytes).prop_filter_map(\"invalid json\", |b| b.ok())"
+        )
     )]
-    pub udf_args: ConvexArray,
+    pub udf_args_bytes: ByteBuf,
 
     pub state: ScheduledJobState,
 
@@ -47,11 +50,47 @@ pub struct ScheduledJob {
     pub attempts: ScheduledJobAttempts,
 }
 
+fn args_to_bytes(args: ConvexArray) -> anyhow::Result<ByteBuf> {
+    let args_json = JsonValue::from(args);
+    let args_bytes = serde_json::to_vec(&args_json)?;
+    Ok(ByteBuf::from(args_bytes))
+}
+
+impl ScheduledJob {
+    pub fn new(
+        path: CanonicalizedComponentFunctionPath,
+        udf_args: ConvexArray,
+        state: ScheduledJobState,
+        next_ts: Option<Timestamp>,
+        completed_ts: Option<Timestamp>,
+        original_scheduled_ts: Timestamp,
+        attempts: ScheduledJobAttempts,
+    ) -> anyhow::Result<Self> {
+        Ok(Self {
+            path,
+            udf_args_bytes: args_to_bytes(udf_args)?,
+            state,
+            next_ts,
+            completed_ts,
+            original_scheduled_ts,
+            attempts,
+        })
+    }
+
+    pub fn udf_args(&self) -> anyhow::Result<ConvexArray> {
+        let args_json: JsonValue = serde_json::from_slice(&self.udf_args_bytes)?;
+        let args = args_json.try_into()?;
+        Ok(args)
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 struct SerializedScheduledJob {
     component: Option<String>,
     udf_path: String,
+    // Serialize the udf arguments as binary since we restrict what
+    // field names can be used in a `Document`'s top-level object.
     udf_args: ByteBuf,
     state: SerializedScheduledJobState,
     next_ts: Option<i64>,
@@ -64,14 +103,10 @@ impl TryFrom<ScheduledJob> for SerializedScheduledJob {
     type Error = anyhow::Error;
 
     fn try_from(job: ScheduledJob) -> anyhow::Result<Self> {
-        // Serialize the udf arguments as binary since we restrict what
-        // field names can be used in a `Document`'s top-level object.
-        let udf_args_json = JsonValue::from(job.udf_args);
-        let udf_args_bytes = serde_json::to_vec(&udf_args_json)?;
         Ok(SerializedScheduledJob {
             component: Some(String::from(job.path.component)),
             udf_path: String::from(job.path.udf_path),
-            udf_args: ByteBuf::from(udf_args_bytes),
+            udf_args: job.udf_args_bytes,
             state: job.state.try_into()?,
             next_ts: job.next_ts.map(|ts| ts.into()),
             completed_ts: job.completed_ts.map(|ts| ts.into()),
@@ -91,8 +126,7 @@ impl TryFrom<SerializedScheduledJob> for ScheduledJob {
             .transpose()?
             .unwrap_or_else(ComponentPath::root);
         let udf_path = value.udf_path.parse()?;
-        let udf_args_json: JsonValue = serde_json::from_slice(&value.udf_args)?;
-        let udf_args = udf_args_json.try_into()?;
+        let udf_args_bytes = value.udf_args;
         let state = value.state.try_into()?;
         let next_ts = value.next_ts.map(|ts| ts.try_into()).transpose()?;
         let completed_ts = value.completed_ts.map(|ts| ts.try_into()).transpose()?;
@@ -114,7 +148,7 @@ impl TryFrom<SerializedScheduledJob> for ScheduledJob {
                 component,
                 udf_path,
             },
-            udf_args,
+            udf_args_bytes,
             state,
             next_ts,
             completed_ts,

crates/model/src/scheduled_jobs/virtual_table.rs

@@ -60,10 +60,11 @@ impl VirtualSystemDocMapper for ScheduledJobsDocMapper {
 
         let job: ParsedDocument<ScheduledJob> = doc.clone().try_into()?;
         let job: ScheduledJob = job.into_value();
+        let udf_args = job.udf_args()?;
         let public_job = PublicScheduledJob {
             // TODO(ENG-6920) include component (job.path.component) in virtual table.
             name: job.path.udf_path,
-            args: job.udf_args,
+            args: udf_args,
             state: job.state,
             scheduled_time: timestamp_to_ms(job.original_scheduled_ts)?,
             completed_time: match job.completed_ts {