Skip to content
Sign in

flagon-io/g1t

public
0

g1t/services/work/src/prefetch.rs

376 lines15,128 bytesCodeBlame

Pick any line to see why it is the way it is: the commit, the pull request and issue it came from, and what the agent was thinking.

Git storage hardened, pages in tens of milliseconds, honest security alerts, and costs reconciled daily1//! Reading a page's worth of rows in one round trip.
2//!
3//! A pull request's page used to cost about twenty queries one after
4//! another: the repository, the pull request, its issue, then where it
5//! stands (its progress, latest review, settings, statuses, verdicts,
6//! queue entry, the confidence signals), its comments, checks and
7//! messages. Each is a round trip to the database, and from most places a
8//! round trip is 20 to 40 ms. Here they go as one D1 batch, keyed by the
9//! repository's id and the pull request's number, and the batch starts
10//! while the repos service is still deciding whether the viewer may see
11//! the repository ([`Work::repo_then`]).
12//!
13//! The helpers that answer those questions (`settings`, `statuses`,
14//! `review_pending`, `approvals_gap` and the rest) look here first, so the
15//! code that decides a pull request's lifecycle is the same code whether
16//! its rows were batched or read one at a time. Only reads are kept; a
17//! helper that writes and reads back (mergeability) reads the database.
18
19use std::cell::RefCell;
20use std::collections::HashMap;
21use std::future::Future;
22use std::rc::Rc;
23
24use g1t_contracts::repos::{Repo, RepoPath};
25use g1t_contracts::time::rfc3339;
26use g1t_contracts::{Outcome, Viewer};
27use g1t_kit::now_ms;
28use serde::de::DeserializeOwned;
29use worker::wasm_bindgen::JsValue;
30use worker::{D1PreparedStatement, D1Result, Result};
31
32use crate::reviews::AGENT_ID;
33use crate::rows::PULL_COLUMNS;
34use crate::{ISSUE_COLUMNS, Work};
35
36/// The pull request's id, from the batch's `?1` (repository) and `?2` (number).
37const PULL_ID: &str = "(SELECT id FROM pulls WHERE repo_id = ?1 AND number = ?2)";
38
39/// How long a review that never reported is waited for (reviews.rs).
40pub(crate) const REVIEW_PENDING_MS: u64 = 30 * 60 * 1000;
41
42/// Each statement of a pull request's batch, in order.
43#[derive(Clone, Copy)]
44pub(crate) enum Slot {
45 Pull,
46 Issue,
47 Comments,
48 Verdicts,
49 Runs,
50 RunHistory,
51 Review,
52 ReviewComments,
53 ReviewPending,
54 Statuses,
55 Settings,
56 Hold,
57 Messages,
58 Others,
59 Queued,
60 LatestRun,
61 Halted,
62 Denials,
63 Unanswered,
64 Planned,
Merge rulesets: branch and tag rules, agent-first, enforced on push and merge65 Rulesets,
66 Adopted,
Git storage hardened, pages in tens of milliseconds, honest security alerts, and costs reconciled daily67}
68
69/// How many runs `latest_checks` and `earlier_checks` show together.
70pub(crate) const RECENT_RUNS: u32 = 11;
71
72/// One pull request's rows, as read for this request.
73pub(crate) struct Prefetched {
74 pub(crate) pull_id: String,
75 pub(crate) repo_id: String,
Merge rulesets: branch and tag rules, agent-first, enforced on push and merge76 /// The workspace its rulesets were read for, lowercase.
77 pub(crate) namespace: String,
Git storage hardened, pages in tens of milliseconds, honest security alerts, and costs reconciled daily78 /// Its head commit when read, which its statuses are for.
79 pub(crate) head: Option<String>,
80 results: Vec<D1Result>,
81}
82
83impl Prefetched {
84 /// A slot's rows, as the helper that reads them would have.
85 pub(crate) fn rows<T: DeserializeOwned>(&self, slot: Slot) -> Result<Vec<T>> {
86 match self.results.get(slot as usize) {
87 Some(result) => result.results::<T>(),
88 None => Ok(Vec::new()),
89 }
90 }
91
92 /// A slot's first row.
93 pub(crate) fn first<T: DeserializeOwned>(&self, slot: Slot) -> Result<Option<T>> {
94 Ok(self.rows::<T>(slot)?.into_iter().next())
95 }
96}
97
98#[derive(serde::Deserialize)]
99struct IdRow {
100 id: String,
101 head_commit: Option<String>,
102}
103
104thread_local! {
105 /// Repository ids by path, as the repos service last answered: a guess
106 /// that lets a page's batch start before the answer comes. Never
107 /// trusted: nothing read with a guess is used unless the repos service
108 /// then says the viewer may see that very repository.
109 static REPO_IDS: RefCell<HashMap<String, String>> = RefCell::new(HashMap::new());
110}
111
112/// Guesses kept at most; the map is emptied past this.
113const MAX_GUESSES: usize = 1024;
114
115fn path_key(path: &RepoPath) -> String {
116 format!("{}/{}", path.namespace.to_lowercase(), path.name.to_lowercase())
117}
118
119fn guess(path: &RepoPath) -> Option<String> {
120 REPO_IDS.with(|ids| ids.borrow().get(&path_key(path)).cloned())
121}
122
123fn learn(path: &RepoPath, id: Option<&str>) {
124 REPO_IDS.with(|ids| {
125 let mut ids = ids.borrow_mut();
126 match id {
127 Some(id) => {
128 if ids.len() >= MAX_GUESSES {
129 ids.clear();
130 }
131 ids.insert(path_key(path), id.to_owned());
132 }
133 None => {
134 ids.remove(&path_key(path));
135 }
136 }
137 });
138}
139
140impl Work {
141 /// The repository at `path` if `viewer` may see it, and what `load`
142 /// reads for it. When this isolate has seen the repository before,
143 /// `load` starts with its id at once, beside the access check, instead
144 /// of after it: one round trip instead of two. What it read is
145 /// discarded unless the check passes for that same repository.
146 pub(crate) async fn repo_then<T, F, Fut>(&self, path: &RepoPath, viewer: &Viewer, load: F) -> Result<Outcome<(Repo, T)>>
147 where
148 F: Fn(String) -> Fut,
149 Fut: Future<Output = Result<T>>,
150 {
151 let guessed = guess(path);
152 let (found, early) = match &guessed {
153 Some(id) => {
154 let (found, early) = futures_util::future::join(self.repo(path, viewer), load(id.clone())).await;
155 (found?, Some(early))
156 }
157 None => (self.repo(path, viewer).await?, None),
158 };
159 let repo = match found {
160 Outcome::Ok(repo) => repo,
161 Outcome::Fail(failure) => {
162 learn(path, None);
163 return Ok(Outcome::Fail(failure));
164 }
165 };
166 if let (Some(id), Some(Ok(value))) = (&guessed, early)
167 && *id == repo.id
168 {
169 return Ok(Outcome::Ok((repo, value)));
170 }
171 learn(path, Some(&repo.id));
172 let value = load(repo.id.clone()).await?;
173 Ok(Outcome::Ok((repo, value)))
174 }
175
176 fn statement(&self, sql: &str, binds: &[JsValue]) -> Result<D1PreparedStatement> {
177 self.db.prepare(sql).bind(binds)
178 }
179
180 /// Everything a pull request's page and its lifecycle read, in one
181 /// batch. `None` when there is no such pull request.
Merge rulesets: branch and tag rules, agent-first, enforced on push and merge182 pub(crate) async fn prefetch_pull(&self, repo_id: String, namespace: String, number: u32) -> Result<Option<Prefetched>> {
Git storage hardened, pages in tens of milliseconds, honest security alerts, and costs reconciled daily183 let key = || -> [JsValue; 2] { [repo_id.as_str().into(), number.into()] };
184 let with = |extra: JsValue| -> [JsValue; 3] { [repo_id.as_str().into(), number.into(), extra] };
185 let repo_only = || -> [JsValue; 1] { [repo_id.as_str().into()] };
186 let statements = vec![
187 // Slot::Pull: the row, with everything the lifecycle tracks on it.
188 self.statement(&format!("SELECT {PULL_COLUMNS} FROM pulls WHERE repo_id = ?1 AND number = ?2"), &key())?,
189 // Slot::Issue
190 self.statement(
191 &format!(
192 "SELECT {ISSUE_COLUMNS} FROM issues
193 WHERE repo_id = ?1 AND number = (SELECT issue_number FROM pulls WHERE repo_id = ?1 AND number = ?2)"
194 ),
195 &key(),
196 )?,
197 // Slot::Comments (lib.rs `comments`)
198 self.statement("SELECT * FROM comments WHERE repo_id = ?1 AND number = ?2 ORDER BY id LIMIT 500", &key())?,
199 // Slot::Verdicts: every verdict, oldest first, for the approval
200 // rule, a person's request for changes and a person's approval.
201 self.statement(
202 "SELECT author_id, author_name, verdict, created_at FROM comments
203 WHERE repo_id = ?1 AND number = ?2 AND verdict IS NOT NULL ORDER BY id",
204 &key(),
205 )?,
206 // Slot::Runs: the latest and the ten before it (checks.rs).
207 self.statement(
208 &format!("SELECT * FROM check_runs WHERE pull_id = {PULL_ID} ORDER BY id DESC LIMIT ?3"),
209 &with(RECENT_RUNS.into()),
210 )?,
211 // Slot::RunHistory (confidence.rs `signals`)
212 self.statement(
213 &format!("SELECT head_commit, status FROM check_runs WHERE pull_id = {PULL_ID} ORDER BY id LIMIT 50"),
214 &key(),
215 )?,
216 // Slot::Review (lifecycle.rs `assess_now`)
217 self.statement(
218 &format!(
219 "SELECT finished_at, verdict FROM review_runs
220 WHERE pull_id = {PULL_ID} AND finished_at IS NOT NULL ORDER BY id DESC LIMIT 1"
221 ),
222 &key(),
223 )?,
224 // Slot::ReviewComments: lines that latest review commented on
225 // (confidence.rs `review_comments`).
226 self.statement(
227 &format!(
228 "SELECT count(*) AS n FROM comments
229 WHERE repo_id = ?1 AND number = ?2 AND author_id = ?3 AND path IS NOT NULL
230 AND created_at = (SELECT finished_at FROM review_runs
231 WHERE pull_id = {PULL_ID} AND finished_at IS NOT NULL ORDER BY id DESC LIMIT 1)"
232 ),
233 &with(AGENT_ID.into()),
234 )?,
235 // Slot::ReviewPending (reviews.rs `review_pending`)
236 self.statement(
237 &format!(
238 "SELECT id AS value FROM review_runs
239 WHERE pull_id = {PULL_ID} AND finished_at IS NULL AND created_at > ?3 LIMIT 1"
240 ),
241 &with(rfc3339(now_ms().saturating_sub(REVIEW_PENDING_MS)).into()),
242 )?,
243 // Slot::Statuses: on its head (statuses.rs `statuses`).
244 self.statement(
Merge checks: statuses and check runs on every commit245 "SELECT context, state, description, target_url, updated_at, source, check_run_id FROM commit_statuses
Git storage hardened, pages in tens of milliseconds, honest security alerts, and costs reconciled daily246 WHERE repo_id = ?1 AND sha = (SELECT head_commit FROM pulls WHERE repo_id = ?1 AND number = ?2)
247 ORDER BY context",
248 &key(),
249 )?,
250 // Slot::Settings and Slot::Hold (settings.rs `settings`)
251 self.statement("SELECT * FROM repo_settings WHERE repo_id = ?1", &repo_only())?,
252 self.statement("SELECT hold_low AS n FROM confidence_rules WHERE repo_id = ?1", &repo_only())?,
253 // Slot::Messages (messages.rs `messages`)
254 self.statement(
255 &format!("SELECT * FROM agent_messages WHERE pull_id = {PULL_ID} ORDER BY created_at, id"),
256 &key(),
257 )?,
258 // Slot::Others (reviews.rs `overlaps`)
259 self.statement(
260 "SELECT number, title, issue_number, files FROM pulls
261 WHERE repo_id = ?1 AND number != ?2 AND status IN ('draft', 'open')
262 ORDER BY number LIMIT 200",
263 &key(),
264 )?,
265 // Slot::Queued (queue.rs `queued_entry`)
266 self.statement(
267 &format!(
268 "SELECT * FROM queue_entries
269 WHERE pull_id = {PULL_ID} AND state IN ('waiting', 'testing', 'passed') LIMIT 1"
270 ),
271 &key(),
272 )?,
273 // Slot::LatestRun, Halted, Denials, Unanswered, Planned
274 // (confidence.rs `signals`)
275 self.statement(
276 &format!(
277 "SELECT r.id, r.cost_usd, r.budget_usd, r.time_cap_minutes,
278 (julianday(COALESCE(r.finished_at, r.updated_at)) - julianday(COALESCE(r.started_at, r.created_at))) * 1440 AS minutes,
279 c.self_level, c.uncertain_about
280 FROM agent_runs r LEFT JOIN run_confidence c ON c.run_id = r.id
281 WHERE r.pull_id = {PULL_ID} AND r.kind IN ('implement', 'revise')
282 ORDER BY r.created_at DESC LIMIT 1"
283 ),
284 &key(),
285 )?,
286 self.statement(
287 &format!(
288 "SELECT halted FROM agent_runs WHERE pull_id = {PULL_ID} AND halted IS NOT NULL
289 ORDER BY created_at DESC LIMIT 1"
290 ),
291 &key(),
292 )?,
293 self.statement(
294 &format!(
295 "SELECT count(*) AS n FROM session_entries
296 WHERE pull_id = {PULL_ID} AND kind = 'note' AND text LIKE 'Denied:%'"
297 ),
298 &key(),
299 )?,
300 self.statement(
301 "SELECT count(*) AS n FROM agent_messages
302 WHERE repo_id = ?1 AND from_number = ?2 AND kind IN ('question', 'handoff')
303 AND answered_at IS NULL",
304 &key(),
305 )?,
306 self.statement(
307 "SELECT json_extract(planned.value, '$.files') AS files
308 FROM plans, json_each(plans.issues) AS planned
309 WHERE plans.repo_id = ?1 AND plans.status = 'applied'
310 AND json_extract(planned.value, '$.number') = (SELECT issue_number FROM pulls WHERE repo_id = ?1 AND number = ?2)
311 LIMIT 1",
312 &key(),
313 )?,
Merge rulesets: branch and tag rules, agent-first, enforced on push and merge314 // Slot::Rulesets and Slot::Adopted (rulesets.rs `rulesets_for`)
315 self.statement(crate::rulesets::prefetch_sql().0, &[repo_id.as_str().into(), namespace.to_lowercase().into()])?,
316 self.statement(crate::rulesets::prefetch_sql().1, &repo_only())?,
Git storage hardened, pages in tens of milliseconds, honest security alerts, and costs reconciled daily317 ];
318 let count = statements.len() as u32;
319 let results = self.timing.db(count, self.db.batch(statements)).await?;
320 let Some(row) = results
321 .first()
322 .map(|result| result.results::<IdRow>())
323 .transpose()?
324 .and_then(|rows| rows.into_iter().next())
325 else {
326 return Ok(None);
327 };
328 Ok(Some(Prefetched {
329 pull_id: row.id,
330 repo_id,
Merge rulesets: branch and tag rules, agent-first, enforced on push and merge331 namespace: namespace.to_lowercase(),
Git storage hardened, pages in tens of milliseconds, honest security alerts, and costs reconciled daily332 head: row.head_commit,
333 results,
334 }))
335 }
336
337 /// The rows read for pull request `pull_id` in this request, if any.
338 pub(crate) fn prefetched_pull(&self, pull_id: &str) -> Option<Rc<Prefetched>> {
339 self.prefetched
340 .borrow()
341 .as_ref()
342 .filter(|found| found.pull_id == pull_id)
343 .cloned()
344 }
345
346 /// The rows read in this request for a pull request of `repo_id`.
347 pub(crate) fn prefetched_repo(&self, repo_id: &str) -> Option<Rc<Prefetched>> {
348 self.prefetched
349 .borrow()
350 .as_ref()
351 .filter(|found| found.repo_id == repo_id)
352 .cloned()
353 }
354
355 /// Keeps `found` for the rest of this request.
356 pub(crate) fn keep_prefetched(&self, found: Option<Prefetched>) {
357 *self.prefetched.borrow_mut() = found.map(Rc::new);
358 }
359}
360
361#[cfg(test)]
362mod tests {
363 use super::*;
364
365 fn at(namespace: &str, name: &str) -> RepoPath {
366 RepoPath { namespace: namespace.to_owned(), name: name.to_owned() }
367 }
368
369 #[test]
370 fn guesses_ignore_case_and_are_forgotten_on_refusal() {
371 learn(&at("Flagon-IO", "G1T"), Some("repo_1"));
372 assert_eq!(guess(&at("flagon-io", "g1t")).as_deref(), Some("repo_1"));
373 learn(&at("flagon-io", "g1t"), None);
374 assert_eq!(guess(&at("flagon-io", "g1t")), None);
375 }
376}

This file's history is long; its oldest lines are credited to the oldest commit read.