Sign in

flagon-io/g1t

public

g1t/crates/kit/src/d1.rs

125 lines5,394 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.

Fast pages, required checks on the branch, self-hosted runners, honest incidents1//! Reading a service's D1 database near the caller, with D1's Sessions API,
2//! and saying how long an RPC took.
3//!
4//! The caller chooses, per request, with the `x-d1-bookmark` header:
5//!
6//! | Header | Session | Reads go to |
7//! | --- | --- | --- |
8//! | absent | none (the plain binding) | the primary, as always |
9//! | `first-primary` | started on the primary | the primary, then any copy at least as new |
10//! | `first-unconstrained` | started anywhere | the nearest copy |
11//! | a bookmark | started at that bookmark | any copy at least as new as the bookmark |
12//!
13//! Writes always go to the primary. A session is sequentially consistent:
14//! what it wrote, it reads back. Its latest bookmark comes back in the
15//! response's `x-d1-bookmark` header, so the caller can start the next
16//! request where this one left off. Without the header nothing changes,
17//! so service-to-service calls ([`crate::call`]), queues and crons read the
18//! primary as before. docs/PERFORMANCE.md explains who sends what.
19//!
20//! Every response that goes through [`Served::finish`] also carries
21//! `server-timing: svc;dur=<ms>;desc="<how D1 was read>"`, which the site
22//! adds up per service for its own `Server-Timing` header.
23
24use worker::wasm_bindgen::JsCast;
25use worker::{D1Database, D1DatabaseSession, Env, Request, Response, Result};
26
27/// The header that carries a session's constraint or bookmark, both ways.
28pub const BOOKMARK: &str = "x-d1-bookmark";
29
30/// Longest bookmark accepted; D1's are about 70 characters.
31const MAX_BOOKMARK: usize = 256;
32
33/// What a request's `x-d1-bookmark` header asks for: `None` for no session
34/// (the primary, as without the header), otherwise what `withSession` is
35/// given. Anything that is not a constraint or a well-formed bookmark
36/// starts on the primary: never staler than asked.
37pub fn constraint(header: Option<&str>) -> Option<String> {
38 let value = header?.trim();
39 if value.is_empty() {
40 return None;
41 }
42 if value == "first-primary" || value == "first-unconstrained" {
43 return Some(value.to_owned());
44 }
45 let well_formed = value.len() <= MAX_BOOKMARK
46 && value
47 .chars()
48 .all(|c| c.is_ascii_alphanumeric() || c == '-');
49 Some(if well_formed { value.to_owned() } else { "first-primary".to_owned() })
50}
51
52/// How one RPC read its database, for its response.
53pub struct Served {
54 session: Option<D1DatabaseSession>,
55 started: u64,
56}
57
58/// The database `binding` for an RPC `request`: a D1 session when the
59/// caller asked for one, the plain binding otherwise. The session is
60/// handed back as a [`D1Database`] so a service's code is the same either
61/// way; it answers `prepare` and `batch` (all a request path uses), not
62/// `exec`, `dump` or `withSession`.
63pub fn open(env: &Env, binding: &str, request: &Request) -> Result<(D1Database, Served)> {
64 let started = crate::now_ms();
65 let db = env.d1(binding)?;
66 let asked = constraint(request.headers().get(BOOKMARK)?.as_deref());
67 let Some(asked) = asked else {
68 return Ok((db, Served { session: None, started }));
69 };
70 let session = db.with_session(Some(&asked))?;
71 // The same JavaScript object, seen as a database: D1Database's methods
72 // are structural, so `prepare` and `batch` call the session's own.
73 let as_database = D1Database::unchecked_from_js(AsRef::<worker::wasm_bindgen::JsValue>::as_ref(&session).clone());
74 Ok((as_database, Served { session: Some(session), started }))
75}
76
77impl Served {
78 /// For a request that touches no database: only its timing.
79 pub fn timing_only() -> Self {
80 Served { session: None, started: crate::now_ms() }
81 }
82
83 /// Adds the session's bookmark and the time taken to `response`.
84 /// Errors pass through untouched.
85 pub fn finish(&self, response: Result<Response>) -> Result<Response> {
86 let mut response = response?;
87 let took = crate::now_ms().saturating_sub(self.started);
88 let how = if self.session.is_some() { "session" } else { "primary" };
89 let headers = response.headers_mut();
90 headers.append("server-timing", &format!("svc;dur={took};desc=\"{how}\""))?;
91 if let Some(session) = &self.session
92 && let Ok(Some(bookmark)) = session.get_bookmark()
93 {
94 headers.set(BOOKMARK, &bookmark)?;
95 }
96 Ok(response)
97 }
98}
99
100#[cfg(test)]
101mod tests {
102 use super::constraint;
103
104 #[test]
105 fn no_header_means_the_primary_without_a_session() {
106 assert_eq!(constraint(None), None);
107 assert_eq!(constraint(Some("")), None);
108 assert_eq!(constraint(Some(" ")), None);
109 }
110
111 #[test]
112 fn constraints_and_bookmarks_pass_through() {
113 assert_eq!(constraint(Some("first-primary")).as_deref(), Some("first-primary"));
114 assert_eq!(constraint(Some("first-unconstrained")).as_deref(), Some("first-unconstrained"));
115 let bookmark = "0000002c-00000004-00004f95-c7f4a9b2e8d1f0c3b6a5d4e3f2a1b0c9";
116 assert_eq!(constraint(Some(bookmark)).as_deref(), Some(bookmark));
117 }
118
119 #[test]
120 fn anything_else_starts_on_the_primary() {
121 assert_eq!(constraint(Some("x; DROP")).as_deref(), Some("first-primary"));
122 assert_eq!(constraint(Some(&"a".repeat(300))).as_deref(), Some("first-primary"));
123 assert_eq!(constraint(Some("first primary")).as_deref(), Some("first-primary"));
124 }
125}