klukai_agent/api/public/
update.rs

1use std::{collections::HashMap, io::Write, sync::Arc, time::Duration};
2
3use antithesis_sdk::assert_sometimes;
4use axum::Extension;
5use bytes::{BufMut, Bytes, BytesMut};
6use compact_str::ToCompactString;
7use http_body_util::StreamBody;
8use klukai_types::{
9    agent::Agent,
10    api::NotifyEvent,
11    tripwire::Tripwire,
12    updates::{Handle, UpdateCreated, UpdateHandle, UpdatesManager},
13};
14use tokio::sync::{
15    RwLock as TokioRwLock,
16    broadcast::{self, error::RecvError},
17    mpsc,
18};
19use tokio_stream::wrappers::ReceiverStream;
20use tracing::{debug, info, warn};
21use uuid::Uuid;
22
23use crate::api::public::pubsub::MatcherUpsertError;
24
25pub type UpdateBroadcastCache = HashMap<Uuid, broadcast::Sender<Bytes>>;
26pub type SharedUpdateBroadcastCache = Arc<TokioRwLock<UpdateBroadcastCache>>;
27
28// this should be a fraction of the MAX_UNSUB_TIME
29const RECEIVERS_CHECK_INTERVAL: Duration = Duration::from_secs(30);
30
31pub async fn api_v1_updates(
32    Extension(agent): Extension<Agent>,
33    Extension(bcast_cache): Extension<SharedUpdateBroadcastCache>,
34    Extension(tripwire): Extension<Tripwire>,
35    axum::extract::Path(table): axum::extract::Path<String>,
36) -> hyper::Response<
37    StreamBody<impl futures::Stream<Item = Result<hyper::body::Frame<Bytes>, std::io::Error>>>,
38> {
39    info!("Received update request for table: {table}");
40
41    assert_sometimes!(true, "Corrosion receives requests for table updates");
42
43    let mut bcast_write = bcast_cache.write().await;
44    let updates = agent.updates_manager();
45
46    let upsert_res = updates.get_or_insert(
47        &table,
48        &agent.schema().read(),
49        agent.pool(),
50        tripwire.clone(),
51    );
52
53    let (handle, maybe_created) = match upsert_res {
54        Ok(res) => res,
55        Err(e) => {
56            // Create error stream
57            let error_msg = format!("Error: {}", MatcherUpsertError::from(e));
58            let error_json = serde_json::to_vec(&error_msg).expect("could not serialize error");
59            let (error_tx, error_rx) =
60                tokio::sync::mpsc::channel::<Result<hyper::body::Frame<Bytes>, std::io::Error>>(1);
61            let _ = error_tx
62                .send(Ok(hyper::body::Frame::data(Bytes::from(error_json))))
63                .await;
64            let error_stream = ReceiverStream::new(error_rx);
65            let error_body = StreamBody::new(error_stream);
66
67            return hyper::Response::builder()
68                .status(hyper::StatusCode::INTERNAL_SERVER_ERROR)
69                .body(error_body)
70                .expect("could not build error response");
71        }
72    };
73
74    let (update_id, sub_rx) =
75        match upsert_update(handle.clone(), maybe_created, updates, &mut bcast_write).await {
76            Ok(id) => id,
77            Err(e) => {
78                // Create error stream
79                let error_msg = format!("Error: {}", e);
80                let error_json = serde_json::to_vec(&error_msg).expect("could not serialize error");
81                let (error_tx, error_rx) = tokio::sync::mpsc::channel::<
82                    Result<hyper::body::Frame<Bytes>, std::io::Error>,
83                >(1);
84                let _ = error_tx
85                    .send(Ok(hyper::body::Frame::data(Bytes::from(error_json))))
86                    .await;
87                let error_stream = ReceiverStream::new(error_rx);
88                let error_body = StreamBody::new(error_stream);
89
90                return hyper::Response::builder()
91                    .status(hyper::StatusCode::INTERNAL_SERVER_ERROR)
92                    .body(error_body)
93                    .expect("could not build error response");
94            }
95        };
96
97    // Create streaming body
98    let (stream_tx, stream_rx) =
99        tokio::sync::mpsc::channel::<Result<hyper::body::Frame<Bytes>, std::io::Error>>(10240);
100    let stream = ReceiverStream::new(stream_rx);
101    let body = StreamBody::new(stream);
102
103    tokio::spawn(forward_update_bytes_to_body_sender(
104        handle, sub_rx, stream_tx, tripwire,
105    ));
106
107    hyper::Response::builder()
108        .status(hyper::StatusCode::OK)
109        .header("corro-query-id", update_id.to_string())
110        .body(body)
111        .expect("could not generate ok http response for update request")
112}
113
114pub async fn upsert_update(
115    handle: UpdateHandle,
116    maybe_created: Option<UpdateCreated>,
117    updates: &UpdatesManager,
118    bcast_write: &mut UpdateBroadcastCache,
119) -> Result<(Uuid, broadcast::Receiver<Bytes>), MatcherUpsertError> {
120    let sub_rx = if let Some(created) = maybe_created {
121        let (sub_tx, sub_rx) = broadcast::channel(10240);
122        bcast_write.insert(handle.id(), sub_tx.clone());
123        tokio::spawn(process_update_channel(
124            updates.clone(),
125            handle.id(),
126            sub_tx,
127            created.evt_rx,
128        ));
129
130        sub_rx
131    } else {
132        let id = handle.id();
133        let sub_tx = bcast_write
134            .get(&id)
135            .cloned()
136            .ok_or(MatcherUpsertError::MissingBroadcaster)?;
137        debug!("found update handle");
138
139        sub_tx.subscribe()
140    };
141
142    Ok((handle.id(), sub_rx))
143}
144
145pub async fn process_update_channel(
146    updates: UpdatesManager,
147    id: Uuid,
148    tx: broadcast::Sender<Bytes>,
149    mut evt_rx: mpsc::Receiver<NotifyEvent>,
150) {
151    let mut buf = BytesMut::new();
152
153    // interval check for receivers
154    // useful for queries that don't change often so we can cleanup...
155    let mut subs_check = tokio::time::interval(RECEIVERS_CHECK_INTERVAL);
156
157    loop {
158        tokio::select! {
159            biased;
160            Some(query_evt) = evt_rx.recv() => {
161                match make_query_event_bytes(&mut buf, &query_evt) {
162                    Ok(b) => {
163                        if tx.send(b).is_err() {
164                            break;
165                        }
166                    },
167                    Err(e) => {
168                        match make_query_event_bytes(&mut buf, &NotifyEvent::Error(e.to_compact_string())) {
169                            Ok(b) => {
170                                let _ = tx.send(b);
171                            }
172                            Err(e) => {
173                                warn!(update_id = %id, "failed to send error in update channel: {e}");
174                            }
175                        }
176                        break;
177                    }
178                };
179            },
180            _ = subs_check.tick() => {
181                if tx.receiver_count() == 0 {
182                    break;
183                };
184            },
185        };
186    }
187
188    warn!(sub_id = %id, "updates channel done");
189
190    // remove and get handle from the agent's "matchers"
191    let handle = match updates.remove(&id) {
192        Some(h) => {
193            info!(update_id = %id, "Removed update handle from process_update_channel");
194            h
195        }
196        None => {
197            warn!(update_id = %id, "update handle was already gone. odd!");
198            return;
199        }
200    };
201
202    // clean up the subscription
203    handle.cleanup().await;
204}
205
206fn make_query_event_bytes(
207    buf: &mut BytesMut,
208    query_evt: &NotifyEvent,
209) -> serde_json::Result<Bytes> {
210    {
211        let mut writer = buf.writer();
212        serde_json::to_writer(&mut writer, query_evt)?;
213
214        // NOTE: I think that's infaillible...
215        writer
216            .write_all(b"\n")
217            .expect("could not write new line to BytesMut Writer");
218    }
219
220    Ok(buf.split().freeze())
221}
222
223async fn forward_update_bytes_to_body_sender(
224    update: UpdateHandle,
225    mut rx: broadcast::Receiver<Bytes>,
226    tx: tokio::sync::mpsc::Sender<Result<hyper::body::Frame<Bytes>, std::io::Error>>,
227    mut tripwire: Tripwire,
228) {
229    let mut buf = BytesMut::new();
230
231    let send_deadline = tokio::time::sleep(Duration::from_millis(10));
232    tokio::pin!(send_deadline);
233
234    loop {
235        tokio::select! {
236            biased;
237            res = rx.recv() => {
238                match res {
239                    Ok(event_buf) => {
240                        buf.extend_from_slice(&event_buf);
241                        if buf.len() >= 64 * 1024
242                            && let Err(_) = tx.send(Ok(hyper::body::Frame::data(buf.split().freeze()))).await {
243                                warn!(update_id = %update.id(), "could not forward update query event to receiver: stream closed");
244                                return;
245                            }
246                    },
247                    Err(RecvError::Lagged(skipped)) => {
248                        warn!(update_id = %update.id(), "update skipped {} events, aborting", skipped);
249                        return;
250                    },
251                    Err(RecvError::Closed) => {
252                        info!(update_id = %update.id(), "events subcription ran out");
253                        return;
254                    },
255                }
256            },
257            _ = &mut send_deadline => {
258                if !buf.is_empty() {
259                    if tx.send(Ok(hyper::body::Frame::data(buf.split().freeze()))).await.is_err() {
260                        warn!(update_id = %update.id(), "could not forward subscription query event to receiver: stream closed");
261                        return;
262                    }
263                } else {
264                    send_deadline.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(10));
265                    continue;
266                }
267            },
268            _ = update.cancelled() => {
269                info!(update_id = %update.id(), "update cancelled, aborting forwarding bytes to subscriber");
270                return;
271            },
272            _ = &mut tripwire => {
273                break;
274            }
275        }
276    }
277
278    // Send any remaining data
279    while let Ok(event_buf) = rx.try_recv() {
280        buf.extend_from_slice(&event_buf);
281        if tx
282            .send(Ok(hyper::body::Frame::data(buf.split().freeze())))
283            .await
284            .is_err()
285        {
286            warn!(update_id = %update.id(), "could not forward subscription query event to receiver: stream closed");
287            return;
288        }
289    }
290}
klukai_agent/api/public/update.rs

klukai_agent/api/public/
update.rs
