| 1 | //! 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 | |
| 24 | use worker::wasm_bindgen::JsCast; |
| 25 | use worker::{D1Database, D1DatabaseSession, Env, Request, Response, Result}; |
| 26 | |
| 27 | /// The header that carries a session's constraint or bookmark, both ways. |
| 28 | pub const BOOKMARK: &str = "x-d1-bookmark"; |
| 29 | |
| 30 | /// Longest bookmark accepted; D1's are about 70 characters. |
| 31 | const 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. |
| 37 | pub 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. |
| 53 | pub 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`. |
| 63 | pub 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 | |
| 77 | impl 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)] |
| 101 | mod 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 | } |