use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::MutexGuard;
use std::sync::PoisonError;
use std::sync::Weak;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::time::Duration;
use futures::StreamExt;
use leptoskit_db::Db;
use leptoskit_db::Dialect;
use leptoskit_db::ReadEntry;
use leptoskit_db::TableName;
use leptoskit_db::live::Change;
use leptoskit_db::live::ChangeLog;
use leptoskit_db::live::LogEntry;
use serde_json::Value;
use tokio::sync::Notify;
use tokio::sync::mpsc;
use tokio::sync::watch;
use super::Identity;
use super::run_query;
use crate::live::Patch;
use crate::live::Update;
use crate::live::WireError;
use crate::live::patch;
const RERUNS: usize = 16;
const MAX_RESUME_COMMITS: usize = 10_000;
const RESUME_READ: Duration = Duration::from_secs(5);
#[derive(Debug)]
pub(super) struct Batch {
pub to: u64,
pub full: bool,
pub revalidate: bool,
pub changes: Vec<(u32, Changed)>,
}
#[derive(Clone, Debug)]
pub(super) enum Changed {
Whole(Arc<Value>),
Patch(Arc<Patch>, Arc<Value>),
Failed(WireError),
}
impl Changed {
pub(super) fn into_update(self) -> (Update, Option<Arc<Value>>) {
match self {
Self::Whole(value) => (Update::Ok((*value).clone()), Some(value)),
Self::Patch(patch, value) => (Update::Patch((*patch).clone()), Some(value)),
Self::Failed(error) => (Update::Err(error), None),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(super) enum ShareKey {
Shared {
query: String,
args: String,
who: Who,
revision: Option<String>,
},
Alone { connection: u32, id: u32 },
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(super) enum Who {
Caller(String, Option<leptoskit_db::Actor>),
Anonymous,
}
impl ShareKey {
pub(super) fn new(
query: &str,
args: &Value,
identity: &Identity,
connection: u32,
id: u32,
) -> Self {
let who = match &identity.caller {
Some(caller) => Who::Caller(caller.clone(), identity.context.actor()),
None if identity.context.is_empty() => Who::Anonymous,
None => return Self::Alone { connection, id },
};
match serde_json::to_string(args) {
Ok(args) => {
Self::Shared {
query: query.to_owned(),
args,
who,
revision: None,
}
}
Err(_) => Self::Alone { connection, id },
}
}
pub(super) fn with_revision(mut self, key: Option<String>) -> Self {
if let Self::Shared { revision, .. } = &mut self {
*revision = key.filter(|key| !key.is_empty());
}
self
}
}
pub(super) struct Registration {
pub key: ShareKey,
pub query: String,
pub args: Value,
pub identity: Identity,
pub db: Db,
pub result: Arc<Value>,
pub reads: Vec<ReadEntry>,
}
pub(super) struct Hub {
core: Arc<Core>,
#[cfg(feature = "relay")]
relay: std::sync::OnceLock<crate::server::Relay>,
started: tokio::sync::OnceCell<()>,
stop: Arc<Notify>,
kick: Arc<Notify>,
}
struct Core {
db: Db,
state: Mutex<State>,
applied: watch::Sender<u64>,
reruns: AtomicU64,
commits: AtomicU64,
replica_reruns: AtomicU64,
relay_fills: Arc<AtomicU64>,
use_replicas: std::sync::atomic::AtomicBool,
replica_wait: AtomicU64,
floor: std::sync::atomic::AtomicI64,
floor_stale: std::sync::atomic::AtomicBool,
}
pub(super) const REPLICA_WAIT: Duration = Duration::from_millis(20);
#[derive(Default)]
struct State {
applied: u64,
next_group: u64,
groups: HashMap<u64, Group>,
by_key: HashMap<ShareKey, u64>,
by_table: HashMap<TableName, HashSet<u64>>,
everywhere: HashSet<u64>,
members: HashMap<(u32, u32), u64>,
fresh: Vec<u64>,
forced: HashSet<u64>,
connections: HashMap<u32, mpsc::UnboundedSender<Batch>>,
authentication: HashMap<u32, Vec<ReadEntry>>,
}
struct Group {
key: ShareKey,
query: String,
args: Value,
identity: Identity,
db: Db,
result: Arc<Value>,
reads: Vec<ReadEntry>,
tables: HashSet<TableName>,
members: HashMap<(u32, u32), Option<Arc<Value>>>,
}
struct Job {
group: u64,
query: String,
args: Value,
identity: Identity,
db: Db,
}
type Outcome = Result<(Value, Vec<ReadEntry>), WireError>;
impl Hub {
pub(super) fn new(db: Db) -> Self {
Self {
core: Arc::new(Core {
db,
state: Mutex::new(State::default()),
applied: watch::channel(0).0,
reruns: AtomicU64::new(0),
commits: AtomicU64::new(0),
replica_reruns: AtomicU64::new(0),
relay_fills: Arc::new(AtomicU64::new(0)),
use_replicas: std::sync::atomic::AtomicBool::new(true),
replica_wait: AtomicU64::new(
u64::try_from(REPLICA_WAIT.as_micros()).unwrap_or(u64::MAX),
),
floor: std::sync::atomic::AtomicI64::new(i64::MIN),
floor_stale: std::sync::atomic::AtomicBool::new(true),
}),
#[cfg(feature = "relay")]
relay: std::sync::OnceLock::new(),
started: tokio::sync::OnceCell::new(),
stop: Arc::new(Notify::new()),
kick: Arc::new(Notify::new()),
}
}
#[cfg(feature = "relay")]
pub(super) fn use_relay(&self, relay: crate::server::Relay) {
let _ = self.relay.set(relay);
}
pub(super) fn use_replicas(&self, enabled: bool, wait: Duration) {
self.core.use_replicas.store(enabled, Ordering::Relaxed);
self.core.replica_wait.store(
u64::try_from(wait.as_micros()).unwrap_or(u64::MAX),
Ordering::Relaxed,
);
}
pub(super) fn is_started(&self) -> bool {
self.started.initialized()
}
pub(super) async fn start(&self) -> Result<(), WireError> {
self.started
.get_or_try_init(|| {
async {
let unreadable = |error: &dyn std::fmt::Display| {
WireError {
status: 503,
message: format!("the change log can't be read: {error}"),
}
};
let source = self.source().await.map_err(|error| unreadable(&error))?;
let version = u64::try_from(source.version()).unwrap_or(0);
self.core.state().applied = version;
self.core.applied.send_replace(version);
tokio::spawn(read(
Arc::downgrade(&self.core),
source,
Arc::clone(&self.stop),
Arc::clone(&self.kick),
));
Ok(())
}
})
.await
.map(drop)
}
async fn source(&self) -> Result<Source, Box<dyn std::error::Error + Send + Sync>> {
#[cfg(feature = "relay")]
if let Some(relay) = self.relay.get() {
let version = self.core.db.live_version().await?;
let feed = crate::server::relay::Feed::start(
relay.clone(),
self.core.db.clone(),
version,
Arc::clone(&self.core.relay_fills),
)
.await?;
return Ok(Source::Relay(Box::new(feed)));
}
Ok(Source::Log(self.core.db.change_log(None).await?))
}
pub(super) fn applied(&self) -> u64 {
*self.core.applied.borrow()
}
pub(super) fn watch(&self) -> watch::Receiver<u64> {
self.core.applied.subscribe()
}
pub(super) fn counts(&self) -> (u64, u64, u64, u64) {
(
self.core.reruns.load(Ordering::Relaxed),
self.core.commits.load(Ordering::Relaxed),
self.core.replica_reruns.load(Ordering::Relaxed),
self.core.relay_fills.load(Ordering::Relaxed),
)
}
pub(super) fn connect(&self, connection: u32, events: mpsc::UnboundedSender<Batch>) {
self.core.state().connections.insert(connection, events);
}
pub(super) fn watch_identity(
&self,
connection: u32,
reads: Vec<ReadEntry>,
since: u64,
) -> bool {
let mut state = self.core.state();
if since != state.applied {
return false;
}
state.authentication.insert(connection, reads);
true
}
pub(super) fn disconnect(&self, connection: u32, ids: impl IntoIterator<Item = u32>) {
let mut state = self.core.state();
state.connections.remove(&connection);
state.authentication.remove(&connection);
for id in ids {
state.leave(connection, id);
}
}
pub(super) fn leave(&self, connection: u32, id: u32) {
self.core.state().leave(connection, id);
}
pub(super) fn peek(&self, key: &ShareKey) -> Option<(Arc<Value>, Vec<ReadEntry>)> {
let state = self.core.state();
let group = state.groups.get(state.by_key.get(key)?)?;
Some((Arc::clone(&group.result), group.reads.clone()))
}
pub(super) fn join(
&self,
connection: u32,
id: u32,
key: &ShareKey,
known: bool,
) -> Option<Arc<Value>> {
let mut state = self.core.state();
let group_id = *state.by_key.get(key)?;
state.leave(connection, id);
let group = state.groups.get_mut(&group_id)?;
let result = Arc::clone(&group.result);
group
.members
.insert((connection, id), known.then(|| Arc::clone(&result)));
state.members.insert((connection, id), group_id);
Some(result)
}
pub(super) fn register(
&self,
connection: u32,
id: u32,
since: Option<u64>,
known: bool,
registration: Registration,
) -> bool {
let mut state = self.core.state();
if since.is_some_and(|since| state.applied != since) {
return false;
}
self.core.floor_stale.store(true, Ordering::Relaxed);
state.leave(connection, id);
let base = known.then(|| Arc::clone(®istration.result));
let group_id = if let Some(&group_id) = state.by_key.get(®istration.key)
&& let Some(group) = state.groups.get_mut(&group_id)
{
group.members.insert((connection, id), base);
state.members.insert((connection, id), group_id);
group_id
} else {
state.next_group += 1;
let group_id = state.next_group;
let tables = tables(®istration.reads);
let mut members = HashMap::new();
members.insert((connection, id), base);
state.index(
group_id,
&HashSet::new(),
&tables,
registration.reads.is_empty(),
);
state.by_key.insert(registration.key.clone(), group_id);
state.members.insert((connection, id), group_id);
state.groups.insert(
group_id,
Group {
key: registration.key,
query: registration.query,
args: registration.args,
identity: registration.identity,
db: registration.db,
result: registration.result,
reads: registration.reads,
tables,
members,
},
);
state.fresh.push(group_id);
group_id
};
if since.is_none() {
state.forced.insert(group_id);
drop(state);
self.kick.notify_one();
}
true
}
pub(super) fn known(&self, connection: u32, id: u32, value: &Arc<Value>) {
let mut state = self.core.state();
let Some(&group_id) = state.members.get(&(connection, id)) else {
return;
};
let Some(group) = state.groups.get_mut(&group_id) else {
return;
};
if Arc::ptr_eq(&group.result, value)
&& let Some(base @ None) = group.members.get_mut(&(connection, id))
{
*base = Some(Arc::clone(value));
}
}
pub(super) async fn changes_between(&self, from: u64, to: u64) -> Option<Vec<Change>> {
if from == to {
return Some(Vec::new());
}
if from > to {
return None;
}
let deadline = tokio::time::Instant::now() + RESUME_READ;
let mut after = i64::try_from(from).ok()?;
let to = i64::try_from(to).ok()?;
let mut changes = Vec::new();
let mut commits = 0;
while after < to {
let page = tokio::time::timeout_at(deadline, self.core.db.change_log_after(after))
.await
.ok()?
.ok()?;
if page.is_empty() {
return None;
}
for entry in page {
match entry {
LogEntry::Commit(commit) if commit.version <= to => {
commits += 1;
if commits > MAX_RESUME_COMMITS {
return None;
}
after = commit.version;
changes.extend(commit.changes);
}
LogEntry::Commit(_) => return Some(changes),
LogEntry::Pruned { .. } => return None,
}
}
}
Some(changes)
}
pub(super) fn dialect(&self) -> Dialect {
self.core.db.dialect()
}
}
impl Drop for Hub {
fn drop(&mut self) {
self.stop.notify_one();
}
}
impl Core {
fn state(&self) -> MutexGuard<'_, State> {
self.state.lock().unwrap_or_else(PoisonError::into_inner)
}
async fn round(&self, entries: Vec<LogEntry>) {
let dialect = self.db.dialect();
let mut to = self.state().applied;
let mut full = false;
let mut changes: Vec<Change> = Vec::new();
let mut fence = Some(i64::MIN);
for entry in entries {
match entry {
LogEntry::Commit(commit) => {
to = to.max(u64::try_from(commit.version).unwrap_or(0));
fence = fence.zip(commit.lsn).map(|(a, b)| a.max(b));
changes.extend(commit.changes);
}
LogEntry::Pruned { through } => {
to = to.max(u64::try_from(through).unwrap_or(0));
full = true;
fence = None;
}
}
}
let fence = fence.filter(|fence| *fence != i64::MIN);
let mut replica: Option<Option<Db>> = None;
let mut first = true;
let mut results: HashMap<u64, Outcome> = HashMap::new();
loop {
let jobs = {
let mut state = self.state();
let jobs = state.jobs(&changes, full, first, dialect, &results);
first = false;
if jobs.is_empty() {
state.apply(results, to, full, &changes, dialect);
self.applied.send_replace(to);
return;
}
jobs
};
self.reruns.fetch_add(jobs.len() as u64, Ordering::Relaxed);
let first_pass = replica.is_none();
if first_pass {
replica = Some(self.replica(fence).await);
}
let pass = if first_pass {
replica.as_ref().and_then(Option::as_ref)
} else {
None
};
if pass.is_none() {
self.floor_stale.store(true, Ordering::Relaxed);
}
let jobs: Vec<Job> = match pass {
Some(replica) => {
self.replica_reruns
.fetch_add(jobs.len() as u64, Ordering::Relaxed);
jobs.into_iter()
.map(|job| {
Job {
db: job.db.reading_like(replica),
..job
}
})
.collect()
}
None => jobs,
};
let ran: Vec<(u64, Outcome)> = futures::stream::iter(jobs)
.map(|job| {
async move {
(
job.group,
run_query(&job.query, job.args, job.db, job.identity).await,
)
}
})
.buffer_unordered(RERUNS)
.collect()
.await;
results.extend(ran);
if let Some(replica) = pass
&& self.db.replica_count() > 1
&& let Some(position) = replica.replica_position().await
{
self.floor.fetch_max(position, Ordering::Relaxed);
}
}
}
}
impl Core {
async fn replica(&self, fence: Option<i64>) -> Option<Db> {
let fence = fence?;
if self.db.replica_count() == 0 || !self.use_replicas.load(Ordering::Relaxed) {
return None;
}
if self.floor_stale.swap(false, Ordering::Relaxed) {
if let Ok(Some(position)) = self.db.wal_position().await {
self.floor.fetch_max(position, Ordering::Relaxed);
} else {
self.floor_stale.store(true, Ordering::Relaxed);
return None;
}
}
let wait = Duration::from_micros(self.replica_wait.load(Ordering::Relaxed));
let needed = fence.max(self.floor.load(Ordering::Relaxed));
self.db.on_replica_at(needed, wait).await
}
}
impl State {
fn jobs(
&mut self,
changes: &[Change],
full: bool,
first: bool,
dialect: Dialect,
done: &HashMap<u64, Outcome>,
) -> Vec<Job> {
let forced: Vec<u64> = self.forced.drain().collect();
let candidates: Vec<u64> = if first {
self.fresh.clear();
if full {
self.groups.keys().copied().collect()
} else {
let mut candidates: HashSet<u64> = self.everywhere.clone();
for change in changes {
if let Some(groups) = self.by_table.get(&change.table) {
candidates.extend(groups);
}
}
candidates.into_iter().collect()
}
} else {
std::mem::take(&mut self.fresh)
};
let candidates: HashSet<u64> = candidates
.into_iter()
.chain(forced.iter().copied())
.collect();
candidates
.into_iter()
.filter(|group| forced.contains(group) || !done.contains_key(group))
.filter_map(|id| {
let group = self.groups.get(&id)?;
let dirty = full
|| forced.contains(&id)
|| group.reads.is_empty()
|| group.reads.iter().any(|read| {
changes
.iter()
.any(|change| read.could_change(change, dialect))
});
dirty.then(|| {
Job {
group: id,
query: group.query.clone(),
args: group.args.clone(),
identity: group.identity.clone(),
db: group.db.clone(),
}
})
})
.collect()
}
fn apply(
&mut self,
results: HashMap<u64, Outcome>,
to: u64,
full: bool,
changes: &[Change],
dialect: Dialect,
) {
let mut batches: HashMap<u32, Vec<(u32, Changed)>> = HashMap::new();
let mut reindex = Vec::new();
for (id, outcome) in results {
let Some(group) = self.groups.get_mut(&id) else {
continue;
};
match outcome {
Ok((value, reads)) => {
let new = Arc::new(value);
let shared =
(*new != *group.result).then(|| shared_change(&group.result, &new));
for (&(connection, member), base) in &mut group.members {
let change = if full {
Some(Changed::Whole(Arc::clone(&new)))
} else {
match base {
Some(base) if Arc::ptr_eq(base, &group.result) => shared.clone(),
Some(base) if **base == *new => None,
Some(base) => Some(shared_change(base, &new)),
None => Some(Changed::Whole(Arc::clone(&new))),
}
};
if let Some(change) = change {
batches
.entry(connection)
.or_default()
.push((member, change));
}
*base = Some(Arc::clone(&new));
}
group.result = new;
let tables = tables(&reads);
if tables != group.tables || reads.is_empty() != group.reads.is_empty() {
reindex.push((
id,
std::mem::replace(&mut group.tables, tables.clone()),
tables,
reads.is_empty(),
));
}
group.reads = reads;
}
Err(error) => {
for &(connection, member) in group.members.keys() {
batches
.entry(connection)
.or_default()
.push((member, Changed::Failed(error.clone())));
}
}
}
}
for (id, old, new, everywhere) in reindex {
self.index(id, &old, &new, everywhere);
}
let revalidate: HashSet<u32> = self
.authentication
.iter()
.filter(|(_, reads)| {
!reads.is_empty()
&& (full
|| reads.iter().any(|read| {
changes
.iter()
.any(|change| read.could_change(change, dialect))
}))
})
.map(|(&connection, _)| connection)
.collect();
for connection in &revalidate {
batches.entry(*connection).or_default();
}
for (connection, changes) in batches {
if let Some(events) = self.connections.get(&connection) {
let _ = events.send(Batch {
to,
full,
changes,
revalidate: revalidate.contains(&connection),
});
}
}
self.applied = to;
}
fn index(
&mut self,
id: u64,
old: &HashSet<TableName>,
new: &HashSet<TableName>,
everywhere: bool,
) {
for table in old.difference(new) {
if let Some(groups) = self.by_table.get_mut(table) {
groups.remove(&id);
if groups.is_empty() {
self.by_table.remove(table);
}
}
}
for table in new {
self.by_table.entry(table.clone()).or_default().insert(id);
}
if everywhere {
self.everywhere.insert(id);
} else {
self.everywhere.remove(&id);
}
}
fn leave(&mut self, connection: u32, id: u32) {
let Some(group_id) = self.members.remove(&(connection, id)) else {
return;
};
let Some(group) = self.groups.get_mut(&group_id) else {
return;
};
group.members.remove(&(connection, id));
if !group.members.is_empty() {
return;
}
let Some(group) = self.groups.remove(&group_id) else {
return;
};
self.by_key.remove(&group.key);
self.index(group_id, &group.tables, &HashSet::new(), false);
}
}
fn tables(reads: &[ReadEntry]) -> HashSet<TableName> {
reads.iter().map(|read| read.table.clone()).collect()
}
fn shared_change(old: &Value, new: &Arc<Value>) -> Changed {
match patch(old, new) {
Update::Patch(patch) => Changed::Patch(Arc::new(patch), Arc::clone(new)),
_ => Changed::Whole(Arc::clone(new)),
}
}
enum Source {
Log(ChangeLog),
#[cfg(feature = "relay")]
Relay(Box<crate::server::relay::Feed>),
}
impl Source {
fn version(&self) -> i64 {
match self {
Self::Log(log) => log.version(),
#[cfg(feature = "relay")]
Self::Relay(feed) => feed.version(),
}
}
async fn next(&mut self) -> leptoskit_db::Result<LogEntry> {
match self {
Self::Log(log) => log.next().await,
#[cfg(feature = "relay")]
Self::Relay(feed) => feed.next().await,
}
}
fn next_ready(&mut self) -> Option<leptoskit_db::Result<LogEntry>> {
match self {
Self::Log(log) => log.next_ready(),
#[cfg(feature = "relay")]
Self::Relay(feed) => feed.next_ready().map(Ok),
}
}
}
async fn read(core: Weak<Core>, mut log: Source, stop: Arc<Notify>, kick: Arc<Notify>) {
loop {
let next = tokio::select! {
next = log.next() => Some(next),
() = kick.notified() => None,
() = stop.notified() => return,
};
let Some(core) = core.upgrade() else {
return;
};
let mut entries = match next {
None => Vec::new(),
Some(Ok(entry)) => vec![entry],
Some(Err(_)) => {
drop(core);
tokio::time::sleep(Duration::from_millis(500)).await;
continue;
}
};
while let Some(Ok(entry)) = log.next_ready() {
entries.push(entry);
}
core.commits
.fetch_add(entries.len() as u64, Ordering::Relaxed);
core.round(entries).await;
}
}