OXYGENHUB
◆ v0.1.0 oxygenhub / crates / hub / src / db / runners.rs
8.6 KiB · Rust · last changed Runners: registration, Admin -> Runners, history, drain/disable/remove 1 month ago source blame history raw
b6b7e4cf Andy McCall ↶ 1 
//! OxygenFlow runners (hub spec §13.2): who may claim jobs, what they are,
2 
//! whether they are taking work, and what they have run.
3
4 
use super::{now_ms, Pool};
5 
use sqlx::FromRow;
6
7 
#[derive(Clone, Debug, FromRow)]
8 
pub struct Runner {
9 
    pub id: i64,
10 
    pub name: String,
11 
    pub token_hash: String,
12 
    pub source: String,
13 
    pub key: String,
14 
    pub key_name: String,
15 
    pub arch: String,
16 
    pub os: String,
17 
    pub version: String,
18 
    pub labels: String,
19 
    pub max_jobs: i64,
20 
    pub memory_mib: Option<i64>,
21 
    pub cpus: Option<i64>,
22 
    pub state: String,
23 
    pub created_ms: i64,
24 
    pub created_by: String,
25 
    pub last_seen_ms: Option<i64>,
26 
}
27
28 
/// A runner that has not polled for this long is shown offline.
29 
pub const ONLINE_MS: i64 = 90_000;
30
31 
impl Runner {
32 
    pub fn online(&self) -> bool {
33 
        self.last_seen_ms.is_some_and(|t| now_ms() - t < ONLINE_MS)
34 
    }
35
36 
    /// `online`, `offline`, `draining` or `disabled`.
37 
    pub fn status(&self) -> &'static str {
38 
        match self.state.as_str() {
39 
            "disabled" => "disabled",
40 
            "draining" => "draining",
41 
            _ if self.online() => "online",
42 
            _ => "offline",
43 
        }
44 
    }
45
46 
    pub fn label_list(&self) -> Vec<&str> {
47 
        self.labels.split(',').filter(|l| !l.is_empty()).collect()
48 
    }
49 
}
50
51 
pub async fn by_token(pool: &Pool, token_hash: &str) -> sqlx::Result<Option<Runner>> {
52 
    sqlx::query_as("SELECT * FROM flow_runners WHERE token_hash = ?").bind(token_hash).fetch_optional(pool).await
53 
}
54
55 
pub async fn get(pool: &Pool, id: i64) -> sqlx::Result<Option<Runner>> {
56 
    sqlx::query_as("SELECT * FROM flow_runners WHERE id = ?").bind(id).fetch_optional(pool).await
57 
}
58
59 
pub async fn list(pool: &Pool) -> sqlx::Result<Vec<Runner>> {
60 
    sqlx::query_as("SELECT * FROM flow_runners ORDER BY name, id").fetch_all(pool).await
61 
}
62
63 
/// The row of a runner configured with a hub.toml token, made on first sight.
64 
pub async fn configured(pool: &Pool, token_hash: &str, name: &str) -> sqlx::Result<Runner> {
65 
    sqlx::query("INSERT OR IGNORE INTO flow_runners (name, token_hash, source, created_ms) VALUES (?, ?, 'config', ?)")
66 
        .bind(name)
67 
        .bind(token_hash)
68 
        .bind(now_ms())
69 
        .execute(pool)
70 
        .await?;
71 
    Ok(by_token(pool, token_hash).await?.expect("just made"))
72 
}
73
74 
/// What a claim says about the runner, and that it was seen now. A
75 
/// configured runner takes the name it claims with.
76 
pub async fn seen(pool: &Pool, id: i64, name: Option<&str>, info: Option<&flow_core::api::RunnerInfo>) -> sqlx::Result<()> {
77 
    let now = now_ms();
78 
    match info {
79 
        Some(i) => {
80 
            sqlx::query(
81 
                "UPDATE flow_runners SET arch = ?, os = ?, version = ?, labels = ?, max_jobs = ?, memory_mib = ?, cpus = ?, last_seen_ms = ? WHERE id = ?",
82 
            )
83 
            .bind(&i.arch)
84 
            .bind(&i.os)
85 
            .bind(&i.version)
86 
            .bind(i.labels.join(","))
87 
            .bind(i.max_jobs as i64)
88 
            .bind(i.memory_mib.map(i64::from))
89 
            .bind(i.cpus.map(i64::from))
90 
            .bind(now)
91 
            .bind(id)
92 
            .execute(pool)
93 
            .await?;
94 
        }
95 
        None => {
96 
            sqlx::query("UPDATE flow_runners SET last_seen_ms = ? WHERE id = ?").bind(now).bind(id).execute(pool).await?;
97 
        }
98 
    }
99 
    if let Some(n) = name {
100 
        sqlx::query("UPDATE flow_runners SET name = ? WHERE id = ? AND source = 'config'").bind(n).bind(id).execute(pool).await?;
101 
    }
102 
    Ok(())
103 
}
104
105 
pub async fn set_state(pool: &Pool, id: i64, state: &str) -> sqlx::Result<()> {
106 
    sqlx::query("UPDATE flow_runners SET state = ? WHERE id = ?").bind(state).bind(id).execute(pool).await?;
107 
    Ok(())
108 
}
109
110 
/// Forget a runner: its credential stops working. Its jobs keep its name.
111 
pub async fn remove(pool: &Pool, id: i64) -> sqlx::Result<()> {
112 
    sqlx::query("DELETE FROM flow_runners WHERE id = ?").bind(id).execute(pool).await?;
113 
    Ok(())
114 
}
115
116 
/// The key a runner signs receipts with: the one bound, or `key` (named
117 
/// `name`) when none is yet.
118 
pub async fn bind_key(pool: &Pool, id: i64, key: &str, name: &str) -> sqlx::Result<String> {
119 
    sqlx::query("UPDATE flow_runners SET key = ?, key_name = ? WHERE id = ? AND key = ''")
120 
        .bind(key)
121 
        .bind(name)
122 
        .bind(id)
123 
        .execute(pool)
124 
        .await?;
125 
    sqlx::query_scalar("SELECT key FROM flow_runners WHERE id = ?").bind(id).fetch_one(pool).await
126 
}
127
128 
/// Keys that runners of this hub sign with.
129 
pub async fn keys(pool: &Pool) -> sqlx::Result<Vec<String>> {
130 
    sqlx::query_scalar("SELECT key FROM flow_runners WHERE key != ''").fetch_all(pool).await
131 
}
132
133 
/// A one-time registration token, valid for a day.
134 
pub async fn add_registration(pool: &Pool, token_hash: &str, by: &str) -> sqlx::Result<()> {
135 
    let now = now_ms();
136 
    sqlx::query("INSERT INTO flow_runner_registrations (token_hash, created_by, created_ms, expires_ms) VALUES (?, ?, ?, ?)")
137 
        .bind(token_hash)
138 
        .bind(by)
139 
        .bind(now)
140 
        .bind(now + 24 * 3_600_000)
141 
        .execute(pool)
142 
        .await?;
143 
    Ok(())
144 
}
145
146 
/// Registration tokens not used yet and not expired: (created by, created, expires).
147 
pub async fn pending_registrations(pool: &Pool) -> sqlx::Result<Vec<(String, i64, i64)>> {
148 
    sqlx::query_as("SELECT created_by, created_ms, expires_ms FROM flow_runner_registrations WHERE used_ms IS NULL AND expires_ms > ? ORDER BY created_ms DESC")
149 
        .bind(now_ms())
150 
        .fetch_all(pool)
151 
        .await
152 
}
153
154 
/// Use registration token `token_hash` to register a runner: `None` when
155 
/// the token is unknown, used or expired.
156 
#[allow(clippy::too_many_arguments)]
157 
pub async fn register(
158 
    pool: &Pool,
159 
    registration: &str,
160 
    name: &str,
161 
    token_hash: &str,
162 
    key: &str,
163 
    key_name: &str,
164 
    by_default: &str,
165 
) -> sqlx::Result<Option<Runner>> {
166 
    let mut tx = super::begin_write(pool).await?;
167 
    let now = now_ms();
168 
    let who: Option<String> =
169 
        sqlx::query_scalar("SELECT created_by FROM flow_runner_registrations WHERE token_hash = ? AND used_ms IS NULL AND expires_ms > ?")
170 
            .bind(registration)
171 
            .bind(now)
172 
            .fetch_optional(&mut *tx)
173 
            .await?;
174 
    let Some(who) = who else { return Ok(None) };
175 
    let id: i64 = sqlx::query_scalar(
176 
        "INSERT INTO flow_runners (name, token_hash, source, key, key_name, created_ms, created_by) VALUES (?, ?, 'registered', ?, ?, ?, ?) RETURNING id",
177 
    )
178 
    .bind(name)
179 
    .bind(token_hash)
180 
    .bind(key)
181 
    .bind(key_name)
182 
    .bind(now)
183 
    .bind(if who.is_empty() { by_default } else { &who })
184 
    .fetch_one(&mut *tx)
185 
    .await?;
186 
    sqlx::query("UPDATE flow_runner_registrations SET used_ms = ?, runner_id = ? WHERE token_hash = ?")
187 
        .bind(now)
188 
        .bind(id)
189 
        .bind(registration)
190 
        .execute(&mut *tx)
191 
        .await?;
192 
    tx.commit().await?;
193 
    get(pool, id).await
194 
}
195
196 
/// A job a runner ran, for its history.
197 
#[derive(Clone, Debug, FromRow)]
198 
pub struct RunnerJob {
199 
    pub owner: String,
200 
    pub repo: String,
201 
    pub run_number: i64,
202 
    pub job: String,
203 
    pub status: String,
204 
    pub started_ms: Option<i64>,
205 
    pub finished_ms: Option<i64>,
206 
    pub receipt_attest: String,
207 
}
208
209 
/// Jobs runner `id` ran or is running, newest first.
210 
pub async fn history(pool: &Pool, id: i64, limit: i64, before: Option<i64>) -> sqlx::Result<Vec<RunnerJob>> {
211 
    sqlx::query_as(
212 
        "SELECT p.owner AS owner, p.name AS repo, r.number AS run_number, j.name AS job, j.status AS status, j.started_ms AS started_ms, \
213 
         j.finished_ms AS finished_ms, j.receipt_attest AS receipt_attest \
214 
         FROM flow_jobs j JOIN flow_runs r ON r.id = j.run_id JOIN repositories p ON p.id = r.repo_id \
215 
         WHERE j.runner_id = ? AND (? IS NULL OR j.started_ms < ?) ORDER BY j.started_ms DESC LIMIT ?",
216 
    )
217 
    .bind(id)
218 
    .bind(before)
219 
    .bind(before)
220 
    .bind(limit)
221 
    .fetch_all(pool)
222 
    .await
223 
}
224
225 
/// (jobs run, passed, busy milliseconds) of runner `id`.
226 
pub async fn totals(pool: &Pool, id: i64) -> sqlx::Result<(i64, i64, i64)> {
227 
    sqlx::query_as(
228 
        "SELECT COUNT(*), COALESCE(SUM(status = 'passed'), 0), COALESCE(SUM(finished_ms - started_ms), 0) \
229 
         FROM flow_jobs WHERE runner_id = ? AND finished_ms IS NOT NULL",
230 
    )
231 
    .bind(id)
232 
    .fetch_one(pool)
233 
    .await
234 
}
235
236 
/// Jobs running on each runner now, by runner id.
237 
pub async fn running(pool: &Pool) -> sqlx::Result<Vec<(i64, i64)>> {
238 
    sqlx::query_as("SELECT runner_id, COUNT(*) FROM flow_jobs WHERE status = 'running' AND runner_id IS NOT NULL GROUP BY runner_id")
239 
        .fetch_all(pool)
240 
        .await
241 
}
242
243 
pub async fn set_job_runner(pool: &Pool, job_id: i64, runner_id: i64) -> sqlx::Result<()> {
244 
    sqlx::query("UPDATE flow_jobs SET runner_id = ? WHERE id = ?").bind(runner_id).bind(job_id).execute(pool).await?;
245 
    Ok(())
246 
}