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

orch/engine/
relay.rs
