Skip to main content

platform_runtime/
schedules.rs

1use crate::EnqueueFunctionRequest;
2use crate::functions::{RuntimeClient, map_runtime_error};
3use chrono::Utc;
4use lenso_contracts::CronSchedule;
5use platform_core::{
6    ActorContext, AppError, AppResult, CorrelationId, DbPool, ErrorCode, TraceContext,
7};
8use serde_json::Value;
9use uuid::Uuid;
10
11#[derive(Debug, Clone)]
12pub struct ScheduledFunctionDefinition {
13    pub schedule_key: String,
14    pub module_name: String,
15    pub schedule_name: String,
16    pub function_name: String,
17    pub cron: String,
18    pub schedule: CronSchedule,
19    pub input_json: Value,
20    pub max_attempts: i32,
21}
22
23#[derive(Debug, Clone)]
24pub struct RuntimeScheduler {
25    pool: DbPool,
26    worker_id: String,
27}
28
29impl RuntimeScheduler {
30    #[must_use]
31    pub fn new(pool: DbPool, worker_id: impl Into<String>) -> Self {
32        Self {
33            pool,
34            worker_id: worker_id.into(),
35        }
36    }
37
38    pub async fn enqueue_due(
39        &self,
40        schedules: &[ScheduledFunctionDefinition],
41    ) -> AppResult<Vec<String>> {
42        let mut run_ids = Vec::new();
43        let client = RuntimeClient::new(self.pool.clone());
44
45        for schedule in schedules {
46            if let Some(run_id) = self.enqueue_due_schedule(&client, schedule).await? {
47                run_ids.push(run_id);
48            }
49        }
50
51        Ok(run_ids)
52    }
53
54    async fn enqueue_due_schedule(
55        &self,
56        client: &RuntimeClient,
57        schedule: &ScheduledFunctionDefinition,
58    ) -> AppResult<Option<String>> {
59        let next_run_at = schedule.schedule.next_after(Utc::now()).ok_or_else(|| {
60            AppError::new(
61                ErrorCode::Validation,
62                "Scheduled runtime function has no run within the lookahead window",
63            )
64        })?;
65
66        let mut tx = self.pool.begin().await.map_err(map_runtime_error)?;
67        let existing_cron: Option<String> = sqlx::query_scalar(
68            r#"
69            select cron_expression
70            from runtime.scheduled_functions
71            where schedule_key = $1
72            for update
73            "#,
74        )
75        .bind(&schedule.schedule_key)
76        .fetch_optional(&mut *tx)
77        .await
78        .map_err(map_runtime_error)?;
79
80        let reset_next_run = existing_cron.as_deref() != Some(schedule.cron.as_str());
81        if existing_cron.is_none() {
82            sqlx::query(
83                r#"
84                insert into runtime.scheduled_functions (
85                    schedule_key,
86                    module_name,
87                    schedule_name,
88                    function_name,
89                    cron_expression,
90                    input_json,
91                    max_attempts,
92                    next_run_at
93                )
94                values ($1, $2, $3, $4, $5, $6, $7, $8)
95                "#,
96            )
97            .bind(&schedule.schedule_key)
98            .bind(&schedule.module_name)
99            .bind(&schedule.schedule_name)
100            .bind(&schedule.function_name)
101            .bind(&schedule.cron)
102            .bind(&schedule.input_json)
103            .bind(schedule.max_attempts)
104            .bind(next_run_at)
105            .execute(&mut *tx)
106            .await
107            .map_err(map_runtime_error)?;
108        } else {
109            sqlx::query(
110                r#"
111                update runtime.scheduled_functions
112                set module_name = $2,
113                    schedule_name = $3,
114                    function_name = $4,
115                    cron_expression = $5,
116                    input_json = $6,
117                    max_attempts = $7,
118                    next_run_at = case when $8 then $9 else next_run_at end,
119                    updated_at = now()
120                where schedule_key = $1
121                "#,
122            )
123            .bind(&schedule.schedule_key)
124            .bind(&schedule.module_name)
125            .bind(&schedule.schedule_name)
126            .bind(&schedule.function_name)
127            .bind(&schedule.cron)
128            .bind(&schedule.input_json)
129            .bind(schedule.max_attempts)
130            .bind(reset_next_run)
131            .bind(next_run_at)
132            .execute(&mut *tx)
133            .await
134            .map_err(map_runtime_error)?;
135        }
136
137        let due: Option<String> = sqlx::query_scalar(
138            r#"
139            select schedule_key
140            from runtime.scheduled_functions
141            where schedule_key = $1
142                and next_run_at <= now()
143            for update skip locked
144            "#,
145        )
146        .bind(&schedule.schedule_key)
147        .fetch_optional(&mut *tx)
148        .await
149        .map_err(map_runtime_error)?;
150
151        if due.is_none() {
152            tx.commit().await.map_err(map_runtime_error)?;
153            return Ok(None);
154        }
155
156        let run = client
157            .enqueue_function_in_tx(
158                &mut tx,
159                EnqueueFunctionRequest {
160                    function_name: schedule.function_name.clone(),
161                    input_json: schedule.input_json.clone(),
162                    correlation_id: CorrelationId::new(format!("corr_schedule_{}", Uuid::now_v7())),
163                    actor: ActorContext::Service {
164                        service_id: self.worker_id.clone(),
165                        scopes: vec!["runtime.functions.enqueue".to_owned()],
166                    },
167                    trace: TraceContext::default(),
168                    causation_id: Some(format!("runtime_schedule:{}", schedule.schedule_key)),
169                    max_attempts: Some(schedule.max_attempts),
170                },
171            )
172            .await?;
173
174        sqlx::query(
175            r#"
176            update runtime.scheduled_functions
177            set next_run_at = $2,
178                last_enqueued_at = now(),
179                updated_at = now()
180            where schedule_key = $1
181            "#,
182        )
183        .bind(&schedule.schedule_key)
184        .bind(next_run_at)
185        .execute(&mut *tx)
186        .await
187        .map_err(map_runtime_error)?;
188
189        tx.commit().await.map_err(map_runtime_error)?;
190        client.record_function_enqueued(&run).await;
191        Ok(Some(run.id))
192    }
193}