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

orch/engine/
relay.rs
