orch/engine/
relay.rs

1//! Client for orch-relay, which receives the GitHub App's webhooks on Cloudflare. orch connects
2//! out, so it needs no endpoint. The contract is PROTOCOL.md in the orch-relay repository
3
4use crate::store::TaskStore;
5use futures::{SinkExt, StreamExt};
6use std::sync::Arc;
7use std::time::{Duration, Instant};
8use tokio::sync::Notify;
9use tokio_tungstenite::tungstenite::{Message, client::IntoClientRequest};
10
11const CURSOR_KEY: &str = "relay_cursor";
12const PING_EVERY: Duration = Duration::from_secs(300);
13const MAX_BACKOFF_SECS: u64 = 300;
14// A shorter session keeps the backoff growing, so a relay that closes each socket at once is not
15// hit every second, every request counts toward its daily limit
16const STABLE_SESSION: Duration = Duration::from_secs(60);
17
18#[derive(serde::Deserialize)]
19struct Incoming {
20    #[serde(rename = "type")]
21    kind: String,
22    #[serde(default)]
23    seq: i64,
24    #[serde(default)]
25    event: String,
26    #[serde(default)]
27    action: String,
28    #[serde(default)]
29    repo: String,
30    #[serde(default)]
31    payload: Option<serde_json::Value>,
32}
33
34/// Keeps one stream open for the life of the daemon, reconnecting with backoff
35pub async fn run(url: String, token: String, store: Arc<TaskStore>, notify: Arc<Notify>) {
36    let mut backoff = 1;
37
38    loop {
39        let started = Instant::now();
40        match session(&url, &token, &store, &notify).await {
41            Ok(()) => tracing::warn!("relay stream closed, reconnecting"),
42            Err(e) => tracing::warn!(err = %e, "relay stream failed, reconnecting"),
43        }
44
45        let wait = retry_after(backoff, started.elapsed());
46        tokio::time::sleep(Duration::from_secs(wait)).await;
47        backoff = (wait * 2).min(MAX_BACKOFF_SECS);
48    }
49}
50
51/// Seconds to wait before the next connect
52fn retry_after(backoff: u64, lasted: Duration) -> u64 {
53    if lasted >= STABLE_SESSION { 1 } else { backoff }
54}
55
56async fn session(url: &str, token: &str, store: &TaskStore, notify: &Notify) -> anyhow::Result<()> {
57    let ws_url = format!("{}/ws", url.trim_end_matches('/'))
58        .replacen("https://", "wss://", 1)
59        .replacen("http://", "ws://", 1);
60    let mut request = ws_url.into_client_request()?;
61    request
62        .headers_mut()
63        .insert("authorization", format!("Bearer {token}").parse()?);
64    let (mut ws, _) = tokio_tungstenite::connect_async(request).await?;
65
66    let after = cursor(store).await;
67    let hello = serde_json::json!({ "v": 1, "type": "hello", "after": after });
68    ws.send(Message::Text(hello.to_string())).await?;
69    tracing::info!(after, "relay stream connected");
70
71    let mut ping = tokio::time::interval(PING_EVERY);
72    ping.tick().await;
73
74    loop {
75        tokio::select! {
76            _ = ping.tick() => ws.send(Message::Text("ping".to_string())).await?,
77            msg = ws.next() => match msg {
78                None => return Ok(()),
79                Some(msg) => match msg? {
80                    Message::Text(text) => handle(&text, store, notify).await,
81                    Message::Close(_) => return Ok(()),
82                    _ => {}
83                },
84            },
85        }
86    }
87}
88
89async fn cursor(store: &TaskStore) -> i64 {
90    store
91        .kv_get(CURSOR_KEY)
92        .await
93        .ok()
94        .flatten()
95        .and_then(|v| v.parse().ok())
96        .unwrap_or(0)
97}
98
99/// An event is a hint to re-read GitHub, never a command, so orch acts on what the sync reads
100pub(crate) async fn handle(text: &str, store: &TaskStore, notify: &Notify) {
101    if text == "pong" {
102        return;
103    }
104
105    let msg: Incoming = match serde_json::from_str(text) {
106        Ok(msg) => msg,
107        Err(e) => {
108            tracing::warn!(err = %e, "unreadable relay message");
109            return;
110        }
111    };
112
113    match msg.kind.as_str() {
114        "gap" => {
115            tracing::warn!("relay deleted events this client never read, syncing every project");
116            notify.notify_one();
117        }
118        "event" => {
119            // An orch PR closes its issue, so `issues` `closed` also covers a merged PR
120            if msg.event == "issues" && msg.action == "closed" {
121                let number = msg
122                    .payload
123                    .as_ref()
124                    .and_then(|p| p.pointer("/issue/number"))
125                    .and_then(serde_json::Value::as_i64);
126
127                if let Some(number) = number {
128                    let key = format!("dep_blocker_read_at:{}#{number}", msg.repo);
129
130                    if let Err(e) = store.kv_delete(&key).await {
131                        tracing::warn!(err = %e, key, "failed to release a closed blocker");
132                    }
133                }
134            }
135
136            if let Err(e) = store.kv_set(CURSOR_KEY, &msg.seq.to_string()).await {
137                tracing::warn!(err = %e, seq = msg.seq, "failed to store the relay cursor");
138            }
139
140            notify.notify_one();
141        }
142        _ => {}
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    async fn woke(notify: &Notify) -> bool {
151        tokio::time::timeout(Duration::from_millis(200), notify.notified())
152            .await
153            .is_ok()
154    }
155
156    #[tokio::test]
157    async fn an_event_stores_the_cursor_and_wakes_the_sync() {
158        let store = TaskStore::open_memory().await.unwrap();
159        let notify = Notify::new();
160        handle(
161            r#"{"type":"event","seq":7,"event":"pull_request","action":"opened","repo":"o/r","payload":{}}"#,
162            &store,
163            &notify,
164        )
165        .await;
166
167        assert_eq!(cursor(&store).await, 7);
168        assert!(woke(&notify).await);
169    }
170
171    #[tokio::test]
172    async fn a_closed_issue_releases_its_blocker_cache_entry() {
173        let store = TaskStore::open_memory().await.unwrap();
174        store
175            .kv_set("dep_blocker_read_at:o/r#12", "1")
176            .await
177            .unwrap();
178        store
179            .kv_set("dep_blocker_read_at:o/r#13", "1")
180            .await
181            .unwrap();
182        handle(
183            r#"{"type":"event","seq":1,"event":"issues","action":"closed","repo":"o/r","payload":{"issue":{"number":12}}}"#,
184            &store,
185            &Notify::new(),
186        )
187        .await;
188
189        assert_eq!(
190            store.kv_get("dep_blocker_read_at:o/r#12").await.unwrap(),
191            None
192        );
193        assert!(
194            store
195                .kv_get("dep_blocker_read_at:o/r#13")
196                .await
197                .unwrap()
198                .is_some()
199        );
200    }
201
202    #[tokio::test]
203    async fn a_gap_wakes_the_sync_and_keeps_the_cursor() {
204        let store = TaskStore::open_memory().await.unwrap();
205        store.kv_set(CURSOR_KEY, "3").await.unwrap();
206        let notify = Notify::new();
207        handle(r#"{"type":"gap","oldest":50}"#, &store, &notify).await;
208
209        assert_eq!(cursor(&store).await, 3);
210        assert!(woke(&notify).await);
211    }
212
213    #[tokio::test]
214    async fn pong_and_unknown_messages_change_nothing() {
215        let store = TaskStore::open_memory().await.unwrap();
216        let notify = Notify::new();
217        handle("pong", &store, &notify).await;
218        handle("not json", &store, &notify).await;
219
220        assert_eq!(cursor(&store).await, 0);
221        assert!(!woke(&notify).await);
222    }
223
224    #[expect(
225        clippy::result_large_err,
226        reason = "the handshake callback returns tungstenite's large error type"
227    )]
228    #[tokio::test]
229    async fn a_session_says_hello_with_the_cursor_and_applies_events() {
230        use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
231
232        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
233        let addr = listener.local_addr().unwrap();
234        let server = tokio::spawn(async move {
235            let (tcp, _) = listener.accept().await.unwrap();
236            let mut auth = String::new();
237            let mut ws =
238                tokio_tungstenite::accept_hdr_async(tcp, |req: &Request, res: Response| {
239                    auth = req
240                        .headers()
241                        .get("authorization")
242                        .and_then(|v| v.to_str().ok())
243                        .unwrap_or_default()
244                        .to_string();
245                    Ok(res)
246                })
247                .await
248                .unwrap();
249            let hello = ws.next().await.unwrap().unwrap().into_text().unwrap();
250            let event = r#"{"type":"event","seq":9,"event":"issues","action":"opened","repo":"o/r","payload":{}}"#;
251            ws.send(Message::Text(event.to_string())).await.unwrap();
252            ws.close(None).await.unwrap();
253            (auth, hello)
254        });
255
256        let store = TaskStore::open_memory().await.unwrap();
257        store.kv_set(CURSOR_KEY, "4").await.unwrap();
258        let notify = Notify::new();
259        session(&format!("http://{addr}"), "tok", &store, &notify)
260            .await
261            .unwrap();
262
263        let (auth, hello) = server.await.unwrap();
264        assert_eq!(auth, "Bearer tok");
265        let hello: serde_json::Value = serde_json::from_str(&hello).unwrap();
266        assert_eq!(
267            hello,
268            serde_json::json!({ "v": 1, "type": "hello", "after": 4 })
269        );
270        assert_eq!(cursor(&store).await, 9);
271        assert!(woke(&notify).await);
272    }
273
274    #[test]
275    fn a_socket_closed_at_once_keeps_the_backoff_growing() {
276        let quick = Duration::from_secs(1);
277        let mut backoff = 1;
278        let mut waits = Vec::new();
279
280        for _ in 0..12 {
281            let wait = retry_after(backoff, quick);
282            waits.push(wait);
283            backoff = (wait * 2).min(MAX_BACKOFF_SECS);
284        }
285
286        assert_eq!(waits[..9], [1, 2, 4, 8, 16, 32, 64, 128, 256]);
287        assert_eq!(waits[9..], [300, 300, 300]);
288        assert_eq!(retry_after(300, STABLE_SESSION), 1);
289    }
290}
orch/engine/relay.rs

orch/engine/
relay.rs
