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::{client::IntoClientRequest, Message};
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 {
54        1
55    } else {
56        backoff
57    }
58}
59
60async fn session(url: &str, token: &str, store: &TaskStore, notify: &Notify) -> anyhow::Result<()> {
61    let ws_url = format!("{}/ws", url.trim_end_matches('/'))
62        .replacen("https://", "wss://", 1)
63        .replacen("http://", "ws://", 1);
64    let mut request = ws_url.into_client_request()?;
65    request
66        .headers_mut()
67        .insert("authorization", format!("Bearer {token}").parse()?);
68    let (mut ws, _) = tokio_tungstenite::connect_async(request).await?;
69
70    let after = cursor(store).await;
71    let hello = serde_json::json!({ "v": 1, "type": "hello", "after": after });
72    ws.send(Message::Text(hello.to_string())).await?;
73    tracing::info!(after, "relay stream connected");
74
75    let mut ping = tokio::time::interval(PING_EVERY);
76    ping.tick().await;
77
78    loop {
79        tokio::select! {
80            _ = ping.tick() => ws.send(Message::Text("ping".to_string())).await?,
81            msg = ws.next() => match msg {
82                None => return Ok(()),
83                Some(msg) => match msg? {
84                    Message::Text(text) => handle(&text, store, notify).await,
85                    Message::Close(_) => return Ok(()),
86                    _ => {}
87                },
88            },
89        }
90    }
91}
92
93async fn cursor(store: &TaskStore) -> i64 {
94    store
95        .kv_get(CURSOR_KEY)
96        .await
97        .ok()
98        .flatten()
99        .and_then(|v| v.parse().ok())
100        .unwrap_or(0)
101}
102
103/// An event is a hint to re-read GitHub, never a command, so orch acts on what the sync reads
104pub(crate) async fn handle(text: &str, store: &TaskStore, notify: &Notify) {
105    if text == "pong" {
106        return;
107    }
108
109    let msg: Incoming = match serde_json::from_str(text) {
110        Ok(msg) => msg,
111        Err(e) => {
112            tracing::warn!(err = %e, "unreadable relay message");
113            return;
114        }
115    };
116
117    match msg.kind.as_str() {
118        "gap" => {
119            tracing::warn!("relay deleted events this client never read, syncing every project");
120            notify.notify_one();
121        }
122        "event" => {
123            // An orch PR closes its issue, so `issues` `closed` also covers a merged PR
124            if msg.event == "issues" && msg.action == "closed" {
125                let number = msg
126                    .payload
127                    .as_ref()
128                    .and_then(|p| p.pointer("/issue/number"))
129                    .and_then(|v| v.as_i64());
130
131                if let Some(number) = number {
132                    let key = format!("dep_blocker_read_at:{}#{number}", msg.repo);
133
134                    if let Err(e) = store.kv_delete(&key).await {
135                        tracing::warn!(err = %e, key, "failed to release a closed blocker");
136                    }
137                }
138            }
139
140            if let Err(e) = store.kv_set(CURSOR_KEY, &msg.seq.to_string()).await {
141                tracing::warn!(err = %e, seq = msg.seq, "failed to store the relay cursor");
142            }
143
144            notify.notify_one();
145        }
146        _ => {}
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    async fn woke(notify: &Notify) -> bool {
155        tokio::time::timeout(Duration::from_millis(200), notify.notified())
156            .await
157            .is_ok()
158    }
159
160    #[tokio::test]
161    async fn an_event_stores_the_cursor_and_wakes_the_sync() {
162        let store = TaskStore::open_memory().await.unwrap();
163        let notify = Notify::new();
164        handle(
165            r#"{"type":"event","seq":7,"event":"pull_request","action":"opened","repo":"o/r","payload":{}}"#,
166            &store,
167            &notify,
168        )
169        .await;
170
171        assert_eq!(cursor(&store).await, 7);
172        assert!(woke(&notify).await);
173    }
174
175    #[tokio::test]
176    async fn a_closed_issue_releases_its_blocker_cache_entry() {
177        let store = TaskStore::open_memory().await.unwrap();
178        store
179            .kv_set("dep_blocker_read_at:o/r#12", "1")
180            .await
181            .unwrap();
182        store
183            .kv_set("dep_blocker_read_at:o/r#13", "1")
184            .await
185            .unwrap();
186        handle(
187            r#"{"type":"event","seq":1,"event":"issues","action":"closed","repo":"o/r","payload":{"issue":{"number":12}}}"#,
188            &store,
189            &Notify::new(),
190        )
191        .await;
192
193        assert_eq!(
194            store.kv_get("dep_blocker_read_at:o/r#12").await.unwrap(),
195            None
196        );
197        assert!(store
198            .kv_get("dep_blocker_read_at:o/r#13")
199            .await
200            .unwrap()
201            .is_some());
202    }
203
204    #[tokio::test]
205    async fn a_gap_wakes_the_sync_and_keeps_the_cursor() {
206        let store = TaskStore::open_memory().await.unwrap();
207        store.kv_set(CURSOR_KEY, "3").await.unwrap();
208        let notify = Notify::new();
209        handle(r#"{"type":"gap","oldest":50}"#, &store, &notify).await;
210
211        assert_eq!(cursor(&store).await, 3);
212        assert!(woke(&notify).await);
213    }
214
215    #[tokio::test]
216    async fn pong_and_unknown_messages_change_nothing() {
217        let store = TaskStore::open_memory().await.unwrap();
218        let notify = Notify::new();
219        handle("pong", &store, &notify).await;
220        handle("not json", &store, &notify).await;
221
222        assert_eq!(cursor(&store).await, 0);
223        assert!(!woke(&notify).await);
224    }
225
226    // The handshake callback's error type is tungstenite's, too large for the lint
227    #[allow(clippy::result_large_err)]
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
