OXYGENHUB
# 941abdd6bb4c oxygenhub / crates / hub / src / db / flow.rs
11.9 KiB · Rust · last changed Write transactions begin IMMEDIATE, so a concurrent writer is waited for (busy timeout) instead of failing at once with database is locked; a test of 16 concurrent log appends fails without it 1 month ago source blame history raw
1
//! OxygenFlow runs, jobs and logs (hub spec §13.2).
2
3
use super::{now_ms, Pool};
4
use sqlx::FromRow;
5
6
#[derive(Clone, Debug, FromRow)]
7
pub struct Run {
8
    pub id: i64,
9
    pub repo_id: i64,
10
    pub bookmark: String,
11
    pub state: String,
12
    pub heads: String,
13
    pub pipeline_hash: String,
14
    pub status: String,
15
    pub error: String,
16
    pub actor: String,
17
    pub created_ms: i64,
18
    pub finished_ms: Option<i64>,
19
}
20
21
impl Run {
22
    pub fn head_list(&self) -> Vec<&str> {
23
        self.heads.split('+').filter(|h| !h.is_empty()).collect()
24
    }
25
26
    pub fn live(&self) -> bool {
27
        matches!(self.status.as_str(), "queued" | "running")
28
    }
29
}
30
31
#[derive(Clone, Debug, FromRow)]
32
pub struct Job {
33
    pub id: i64,
34
    pub run_id: i64,
35
    pub name: String,
36
    pub depth: i64,
37
    pub depends: String,
38
    pub spec: String,
39
    pub sandbox: String,
40
    pub runtime: String,
41
    pub status: String,
42
    pub runner: String,
43
    pub token_hash: String,
44
    pub started_ms: Option<i64>,
45
    pub finished_ms: Option<i64>,
46
    pub exit_code: Option<i64>,
47
    pub error: String,
48
    pub receipt: String,
49
    pub receipt_hash: String,
50
    pub log_count: i64,
51
    /// Comma-separated ids of other repositories the job token may read.
52
    pub also_read: String,
53
}
54
55
impl Job {
56
    pub fn depends_on(&self) -> Vec<String> {
57
        serde_json::from_str(&self.depends).unwrap_or_default()
58
    }
59
60
    pub fn done(&self) -> bool {
61
        matches!(self.status.as_str(), "passed" | "failed" | "skipped" | "error")
62
    }
63
}
64
65
#[derive(Clone, Debug, FromRow)]
66
pub struct LogLine {
67
    pub seq: i64,
68
    pub t: i64,
69
    pub stream: String,
70
    pub text: String,
71
}
72
73
/// A job as the trigger creates it.
74
pub struct NewJob {
75
    pub name: String,
76
    pub depth: i64,
77
    pub depends: Vec<String>,
78
    pub spec: String,
79
    pub sandbox: String,
80
    pub runtime: String,
81
}
82
83
pub struct NewRun<'a> {
84
    pub repo_id: i64,
85
    pub bookmark: &'a str,
86
    pub state: &'a str,
87
    pub heads: &'a str,
88
    pub pipeline_hash: &'a str,
89
    pub actor: &'a str,
90
    /// Problems when the pipeline is invalid: the run is recorded in error.
91
    pub error: Option<&'a str>,
92
}
93
94
pub async fn create_run(pool: &Pool, r: NewRun<'_>, jobs: &[NewJob]) -> sqlx::Result<i64> {
95
    let mut tx = super::begin_write(pool).await?;
96
    let now = now_ms();
97
    let status = if r.error.is_some() { "error" } else { "queued" };
98
    let id: i64 = sqlx::query_scalar(
99
        "INSERT INTO flow_runs (repo_id, bookmark, state, heads, pipeline_hash, status, error, actor, created_ms, finished_ms) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING id",
100
    )
101
    .bind(r.repo_id)
102
    .bind(r.bookmark)
103
    .bind(r.state)
104
    .bind(r.heads)
105
    .bind(r.pipeline_hash)
106
    .bind(status)
107
    .bind(r.error.unwrap_or(""))
108
    .bind(r.actor)
109
    .bind(now)
110
    .bind(r.error.map(|_| now))
111
    .fetch_one(&mut *tx)
112
    .await?;
113
    for j in jobs {
114
        sqlx::query(
115
            "INSERT INTO flow_jobs (run_id, name, depth, depends, spec, sandbox, runtime, status) VALUES (?, ?, ?, ?, ?, ?, ?, 'queued')",
116
        )
117
        .bind(id)
118
        .bind(&j.name)
119
        .bind(j.depth)
120
        .bind(serde_json::to_string(&j.depends).unwrap_or_else(|_| "[]".into()))
121
        .bind(&j.spec)
122
        .bind(&j.sandbox)
123
        .bind(&j.runtime)
124
        .execute(&mut *tx)
125
        .await?;
126
    }
127
    tx.commit().await?;
128
    Ok(id)
129
}
130
131
pub async fn runs(pool: &Pool, repo_id: i64, limit: i64, before: Option<i64>) -> sqlx::Result<Vec<Run>> {
132
    sqlx::query_as("SELECT * FROM flow_runs WHERE repo_id = ? AND id < ? ORDER BY id DESC LIMIT ?")
133
        .bind(repo_id)
134
        .bind(before.unwrap_or(i64::MAX))
135
        .bind(limit)
136
        .fetch_all(pool)
137
        .await
138
}
139
140
pub async fn run(pool: &Pool, repo_id: i64, id: i64) -> sqlx::Result<Option<Run>> {
141
    sqlx::query_as("SELECT * FROM flow_runs WHERE repo_id = ? AND id = ?").bind(repo_id).bind(id).fetch_optional(pool).await
142
}
143
144
pub async fn run_by_id(pool: &Pool, id: i64) -> sqlx::Result<Option<Run>> {
145
    sqlx::query_as("SELECT * FROM flow_runs WHERE id = ?").bind(id).fetch_optional(pool).await
146
}
147
148
pub async fn jobs(pool: &Pool, run_id: i64) -> sqlx::Result<Vec<Job>> {
149
    sqlx::query_as("SELECT * FROM flow_jobs WHERE run_id = ? ORDER BY depth, name").bind(run_id).fetch_all(pool).await
150
}
151
152
pub async fn job(pool: &Pool, id: i64) -> sqlx::Result<Option<Job>> {
153
    sqlx::query_as("SELECT * FROM flow_jobs WHERE id = ?").bind(id).fetch_optional(pool).await
154
}
155
156
/// The running job a job token belongs to, while it may still be used.
157
pub async fn job_by_token(pool: &Pool, token_hash: &str, max_age_ms: i64) -> sqlx::Result<Option<Job>> {
158
    sqlx::query_as("SELECT * FROM flow_jobs WHERE token_hash = ? AND status = 'running' AND started_ms >= ?")
159
        .bind(token_hash)
160
        .bind(now_ms() - max_age_ms)
161
        .fetch_optional(pool)
162
        .await
163
}
164
165
/// Take the oldest queued job whose dependencies all passed and whose tier
166
/// the runner offers, and mark it running for `runner` with a token.
167
/// Let a job's token also read these repositories (its pinned subrepos).
168
pub async fn set_also_read(pool: &Pool, job_id: i64, repo_ids: &[i64]) -> sqlx::Result<()> {
169
    let list: Vec<String> = repo_ids.iter().map(|i| i.to_string()).collect();
170
    sqlx::query("UPDATE flow_jobs SET also_read = ? WHERE id = ?").bind(list.join(",")).bind(job_id).execute(pool).await?;
171
    Ok(())
172
}
173
174
impl Job {
175
    pub fn also_read_ids(&self) -> Vec<i64> {
176
        self.also_read.split(',').filter_map(|x| x.trim().parse().ok()).collect()
177
    }
178
}
179
180
/// What a runner can take on: free memory (MiB) and vCPUs, `None` for no limit.
181
#[derive(Clone, Copy, Debug, Default)]
182
pub struct Room {
183
    pub memory_mib: Option<u32>,
184
    pub cpus: Option<u32>,
185
}
186
187
impl Room {
188
    /// Whether a job's sandbox (its stored JSON) fits.
189
    pub fn fits(&self, sandbox_json: &str) -> bool {
190
        let v: serde_json::Value = serde_json::from_str(sandbox_json).unwrap_or_default();
191
        let need = |k: &str| v[k].as_u64().unwrap_or(0);
192
        self.memory_mib.is_none_or(|m| need("memory_mib") <= m as u64) && self.cpus.is_none_or(|c| need("cpus") <= c as u64)
193
    }
194
}
195
196
pub async fn claim(pool: &Pool, runtimes: &[String], runner: &str, token_hash: &str) -> sqlx::Result<Option<Job>> {
197
    claim_within(pool, runtimes, runner, token_hash, Room::default()).await
198
}
199
200
/// The oldest runnable job that fits in `room`.
201
pub async fn claim_within(pool: &Pool, runtimes: &[String], runner: &str, token_hash: &str, room: Room) -> sqlx::Result<Option<Job>> {
202
    let mut tx = super::begin_write(pool).await?;
203
    let queued: Vec<Job> = sqlx::query_as(
204
        "SELECT j.* FROM flow_jobs j JOIN flow_runs r ON r.id = j.run_id WHERE j.status = 'queued' AND r.status IN ('queued', 'running') ORDER BY j.run_id, j.depth, j.id LIMIT 200",
205
    )
206
    .fetch_all(&mut *tx)
207
    .await?;
208
    for j in queued {
209
        if !runtimes.iter().any(|r| r == &j.runtime) || !room.fits(&j.sandbox) {
210
            continue;
211
        }
212
        let deps = j.depends_on();
213
        if !deps.is_empty() {
214
            let mut passed = 0i64;
215
            for d in &deps {
216
                let st: Option<String> = sqlx::query_scalar("SELECT status FROM flow_jobs WHERE run_id = ? AND name = ?")
217
                    .bind(j.run_id)
218
                    .bind(d)
219
                    .fetch_optional(&mut *tx)
220
                    .await?;
221
                if st.as_deref() == Some("passed") {
222
                    passed += 1;
223
                }
224
            }
225
            if passed != deps.len() as i64 {
226
                continue;
227
            }
228
        }
229
        let now = now_ms();
230
        let done = sqlx::query(
231
            "UPDATE flow_jobs SET status = 'running', runner = ?, token_hash = ?, started_ms = ? WHERE id = ? AND status = 'queued'",
232
        )
233
        .bind(runner)
234
        .bind(token_hash)
235
        .bind(now)
236
        .bind(j.id)
237
        .execute(&mut *tx)
238
        .await?;
239
        if done.rows_affected() == 1 {
240
            sqlx::query("UPDATE flow_runs SET status = 'running' WHERE id = ? AND status = 'queued'")
241
                .bind(j.run_id)
242
                .execute(&mut *tx)
243
                .await?;
244
            tx.commit().await?;
245
            return job(pool, j.id).await;
246
        }
247
    }
248
    tx.commit().await?;
249
    Ok(None)
250
}
251
252
/// Append records; returns how many were kept (the rest exceed the cap).
253
pub async fn append_log(pool: &Pool, job_id: i64, records: &[flow_core::api::Record], cap: i64) -> sqlx::Result<usize> {
254
    let mut tx = super::begin_write(pool).await?;
255
    let count: i64 = sqlx::query_scalar("SELECT log_count FROM flow_jobs WHERE id = ?").bind(job_id).fetch_one(&mut *tx).await?;
256
    let room = (cap - count).max(0) as usize;
257
    let keep = &records[..records.len().min(room)];
258
    for (i, r) in keep.iter().enumerate() {
259
        sqlx::query("INSERT INTO flow_logs (job_id, seq, t, stream, text) VALUES (?, ?, ?, ?, ?)")
260
            .bind(job_id)
261
            .bind(count + i as i64 + 1)
262
            .bind(r.t as i64)
263
            .bind(&r.stream)
264
            .bind(&r.text)
265
            .execute(&mut *tx)
266
            .await?;
267
    }
268
    sqlx::query("UPDATE flow_jobs SET log_count = log_count + ? WHERE id = ?")
269
        .bind(keep.len() as i64)
270
        .bind(job_id)
271
        .execute(&mut *tx)
272
        .await?;
273
    tx.commit().await?;
274
    Ok(keep.len())
275
}
276
277
pub async fn log(pool: &Pool, job_id: i64, after: i64, limit: i64) -> sqlx::Result<Vec<LogLine>> {
278
    sqlx::query_as("SELECT seq, t, stream, text FROM flow_logs WHERE job_id = ? AND seq > ? ORDER BY seq LIMIT ?")
279
        .bind(job_id)
280
        .bind(after)
281
        .bind(limit)
282
        .fetch_all(pool)
283
        .await
284
}
285
286
pub struct Outcome<'a> {
287
    pub status: &'a str,
288
    pub exit_code: Option<i64>,
289
    pub error: &'a str,
290
    pub receipt: &'a str,
291
    pub receipt_hash: &'a str,
292
}
293
294
/// Record a job's end, skip the jobs that can no longer run, and settle the
295
/// run. Returns the jobs whose status changed (for their checks).
296
pub async fn finish(pool: &Pool, job_id: i64, o: Outcome<'_>) -> sqlx::Result<Vec<Job>> {
297
    let mut tx = super::begin_write(pool).await?;
298
    let now = now_ms();
299
    let j: Job = sqlx::query_as("SELECT * FROM flow_jobs WHERE id = ?").bind(job_id).fetch_one(&mut *tx).await?;
300
    sqlx::query("UPDATE flow_jobs SET status = ?, exit_code = ?, error = ?, receipt = ?, receipt_hash = ?, finished_ms = ?, token_hash = '' WHERE id = ?")
301
        .bind(o.status)
302
        .bind(o.exit_code)
303
        .bind(o.error)
304
        .bind(o.receipt)
305
        .bind(o.receipt_hash)
306
        .bind(now)
307
        .bind(job_id)
308
        .execute(&mut *tx)
309
        .await?;
310
    let mut changed = vec![job_id];
311
    // Dependents of anything that did not pass are skipped, transitively.
312
    loop {
313
        let all: Vec<Job> = sqlx::query_as("SELECT * FROM flow_jobs WHERE run_id = ?").bind(j.run_id).fetch_all(&mut *tx).await?;
314
        let blocked: Vec<i64> = all
315
            .iter()
316
            .filter(|x| x.status == "queued")
317
            .filter(|x| {
318
                x.depends_on()
319
                    .iter()
320
                    .any(|d| all.iter().any(|y| &y.name == d && matches!(y.status.as_str(), "failed" | "skipped" | "error")))
321
            })
322
            .map(|x| x.id)
323
            .collect();
324
        if blocked.is_empty() {
325
            break;
326
        }
327
        for id in blocked {
328
            sqlx::query("UPDATE flow_jobs SET status = 'skipped', finished_ms = ? WHERE id = ?")
329
                .bind(now)
330
                .bind(id)
331
                .execute(&mut *tx)
332
                .await?;
333
            changed.push(id);
334
        }
335
    }
336
    let all: Vec<Job> = sqlx::query_as("SELECT * FROM flow_jobs WHERE run_id = ?").bind(j.run_id).fetch_all(&mut *tx).await?;
337
    if all.iter().all(Job::done) {
338
        let status = if all.iter().all(|x| x.status == "passed") { "passed" } else { "failed" };
339
        sqlx::query("UPDATE flow_runs SET status = ?, finished_ms = ? WHERE id = ?")
340
            .bind(status)
341
            .bind(now)
342
            .bind(j.run_id)
343
            .execute(&mut *tx)
344
            .await?;
345
    }
346
    tx.commit().await?;
347
    let mut out = Vec::new();
348
    for id in changed {
349
        if let Some(x) = job(pool, id).await? {
350
            out.push(x);
351
        }
352
    }
353
    Ok(out)
354
}