use std::collections::BTreeSet;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use async_graphql::Value;
use futures_util::Stream;
use tokio::sync::{broadcast, mpsc};
use crate::microsvc::Session;
use crate::read_model::ReadModelChange;
use super::compile::{self, RootKind, SelectionNode, SqlPlan};
use super::engine::{execute_plan, EngineInner};
use super::protocol::{ProtocolResponseAccumulator, RequestedLiveResume};
#[derive(Clone, Debug)]
pub struct ChangeHub {
tx: broadcast::Sender<ReadModelChange>,
}
impl ChangeHub {
pub fn new() -> Self {
let (tx, _) = broadcast::channel(256);
Self { tx }
}
pub fn subscribe(&self) -> broadcast::Receiver<ReadModelChange> {
self.tx.subscribe()
}
pub fn publish(&self, change: ReadModelChange) {
if change.is_empty() {
}
let _ = self.tx.send(change);
}
pub fn spawn_forward_from(&self, mut rx: broadcast::Receiver<ReadModelChange>) {
let tx = self.tx.clone();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(change) => {
let _ = tx.send(change);
}
Err(broadcast::error::RecvError::Lagged(_)) => {
let _ = tx.send(ReadModelChange {
tables: BTreeSet::new(),
});
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
});
}
}
impl Default for ChangeHub {
fn default() -> Self {
Self::new()
}
}
type LiveItem = Result<Value, async_graphql::Error>;
pub struct LiveQueryStream {
rx: mpsc::Receiver<LiveItem>,
}
impl Stream for LiveQueryStream {
type Item = LiveItem;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.rx.poll_recv(cx)
}
}
pub(crate) async fn live_query_stream(
inner: Arc<EngineInner>,
session: Session,
role: String,
model: String,
selection: SelectionNode,
protocol: Option<ProtocolResponseAccumulator>,
) -> Result<LiveQueryStream, String> {
let plan: SqlPlan =
compile::compile_root(&inner, &session, &role, &model, RootKind::List, &selection)?;
let footprint = footprint_from_tables(&plan.tables_touched);
let mut change_rx = inner.change_hub.subscribe();
let (tx, rx) = mpsc::channel::<LiveItem>(8);
let debounce = Duration::from_millis(100);
let requested_live_resume = protocol
.as_ref()
.map(ProtocolResponseAccumulator::requested_live_resume)
.transpose()
.map_err(|error| error.to_string())?
.unwrap_or(RequestedLiveResume::Absent);
tokio::spawn(async move {
let mut initial = match execute_list(
&inner,
&role,
&plan,
protocol.as_ref(),
requested_live_resume,
)
.await
{
Ok(executed) => executed,
Err(e) => {
let _ = tx.send(Err(async_graphql::Error::new(e))).await;
return;
}
};
if let Err(error) = initial.record_protocol_metadata(protocol.as_ref()) {
let _ = tx.send(Err(async_graphql::Error::new(error))).await;
return;
}
let mut last_hash = Some(initial.hash);
let mut next_live_resume = initial.next_live_resume;
if tx.send(Ok(initial.value)).await.is_err() {
return;
}
loop {
let change = match change_rx.recv().await {
Ok(c) => c,
Err(broadcast::error::RecvError::Lagged(_)) => ReadModelChange {
tables: BTreeSet::new(),
},
Err(broadcast::error::RecvError::Closed) => break,
};
if !footprint_hits(&footprint, &change) {
continue;
}
tokio::time::sleep(debounce).await;
loop {
match change_rx.try_recv() {
Ok(more) => {
let _ = more;
}
Err(broadcast::error::TryRecvError::Empty) => break,
Err(broadcast::error::TryRecvError::Lagged(_)) => break,
Err(broadcast::error::TryRecvError::Closed) => return,
}
}
match execute_list(
&inner,
&role,
&plan,
protocol.as_ref(),
next_live_resume.clone(),
)
.await
{
Ok(mut executed) => {
next_live_resume = executed.next_live_resume.clone();
if last_hash == Some(executed.hash) {
continue; }
if let Err(error) = executed.record_protocol_metadata(protocol.as_ref()) {
if tx
.send(Err(async_graphql::Error::new(error)))
.await
.is_err()
{
return;
}
continue;
}
last_hash = Some(executed.hash);
if tx.send(Ok(executed.value)).await.is_err() {
return;
}
}
Err(e) => {
if tx.send(Err(async_graphql::Error::new(e))).await.is_err() {
return;
}
}
}
}
});
Ok(LiveQueryStream { rx })
}
struct ExecutedLiveQuery {
value: Value,
hash: u64,
snapshot: Option<super::protocol::DistributedQuerySnapshot>,
live: Option<super::protocol::DistributedLiveMetadata>,
next_live_resume: RequestedLiveResume,
}
impl ExecutedLiveQuery {
fn record_protocol_metadata(
&mut self,
protocol: Option<&ProtocolResponseAccumulator>,
) -> Result<(), String> {
let Some(protocol) = protocol else {
return Ok(());
};
let snapshot = self
.snapshot
.take()
.ok_or_else(|| "causal live query omitted its snapshot metadata".to_string())?;
protocol
.record_query_metadata(snapshot, self.live.take())
.map_err(|error| error.to_string())
}
}
async fn execute_list(
inner: &EngineInner,
role: &str,
plan: &SqlPlan,
protocol: Option<&ProtocolResponseAccumulator>,
requested_live_resume: RequestedLiveResume,
) -> Result<ExecutedLiveQuery, String> {
let Some(protocol) = protocol else {
let value = execute_plan(inner, plan).await?;
return Ok(ExecutedLiveQuery {
hash: response_hash(&value),
value,
snapshot: None,
live: None,
next_live_resume: RequestedLiveResume::Absent,
});
};
let role_surface = inner
.role_surfaces
.get(role)
.cloned()
.ok_or_else(|| "authorized GraphQL role surface is unavailable".to_string())?;
let executed = super::query_protocol::execute_query_with_protocol(
inner,
role_surface,
protocol.clone(),
plan,
Some(requested_live_resume),
)
.await?;
let hash = protocol_response_hash(&executed.value, &executed.snapshot, &executed.live);
let next_live_resume = executed
.live
.as_ref()
.filter(|live| live.supported)
.map(|live| RequestedLiveResume::Cursors(live.cursors.clone()))
.unwrap_or(RequestedLiveResume::Absent);
Ok(ExecutedLiveQuery {
value: executed.value,
hash,
snapshot: Some(executed.snapshot),
live: executed.live,
next_live_resume,
})
}
fn footprint_hits(footprint: &BTreeSet<String>, change: &ReadModelChange) -> bool {
if change.tables.is_empty() {
return true;
}
change.tables.iter().any(|t| footprint.contains(t))
}
pub fn footprint_from_tables(tables: &[String]) -> BTreeSet<String> {
tables.iter().cloned().collect()
}
pub fn response_hash(value: &Value) -> u64 {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut h = DefaultHasher::new();
if let Ok(json) = serde_json::to_string(value) {
json.hash(&mut h);
} else {
format!("{value:?}").hash(&mut h);
}
h.finish()
}
fn protocol_response_hash(
value: &Value,
snapshot: &super::protocol::DistributedQuerySnapshot,
live: &Option<super::protocol::DistributedLiveMetadata>,
) -> u64 {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hash = DefaultHasher::new();
match serde_json::to_string(&(value, snapshot, live)) {
Ok(encoded) => encoded.hash(&mut hash),
Err(_) => format!("{value:?}:{snapshot:?}:{live:?}").hash(&mut hash),
}
hash.finish()
}
pub fn spawn_change_forwarder(hub: ChangeHub, rx: broadcast::Receiver<ReadModelChange>) {
hub.spawn_forward_from(rx);
}