use std::{
collections::{HashMap, HashSet},
mem,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use dashmap::DashMap;
use reifydb_core::{
interface::{catalog::id::SubscriptionId, change::StagedBatch},
value::column::columns::Columns,
};
use reifydb_runtime::{
context::{clock::Clock, rng::Rng},
sync::mutex::Mutex,
};
use reifydb_subscription::{
batch::BatchId,
delivery::{DeliveryResult, SubscriptionDelivery},
};
use reifydb_value::{
reifydb_assertions,
value::{diff_type::DiffType, duration::Duration, frame::frame::Frame, uuid::Uuid7},
};
use tokio::sync::Notify;
use tracing::{debug, instrument, trace};
use crate::subscription::wire_sink::WireSink;
pub type ConnectionId = Uuid7;
type PromoteDrain<S> = (Vec<StagedBatch>, Option<BatchId>, <S as WireSink>::Format, S);
type ThrottleReady<S> = (SubscriptionId, Vec<StagedBatch>, <S as WireSink>::Format, S);
#[derive(Debug)]
pub enum PromoteResult {
Promoted(usize),
Overflowed,
NotWarming,
NotFound,
Disconnected,
}
struct SubscriptionState<S: WireSink> {
connection_id: ConnectionId,
sink: S,
format: S::Format,
#[allow(dead_code)]
query: String,
batch_id: Option<BatchId>,
warming: Option<WarmingBuffer>,
throttle: ThrottleState,
}
struct FlushGate {
throttle_millis: u64,
linger_millis: u64,
first_pending_at: Option<u64>,
last_sent_at: Option<u64>,
}
impl FlushGate {
fn new(throttle: Duration, linger: Duration) -> Self {
Self {
throttle_millis: throttle.to_std().as_millis() as u64,
linger_millis: linger.to_std().as_millis() as u64,
first_pending_at: None,
last_sent_at: None,
}
}
fn shaping(&self) -> bool {
self.throttle_millis != 0 || self.linger_millis != 0
}
fn on_pending(&mut self, now: u64) {
if self.first_pending_at.is_none() {
self.first_pending_at = Some(now);
}
}
fn throttle_ok(&self, now: u64) -> bool {
match self.last_sent_at {
None => true,
Some(prev) => now.saturating_sub(prev) >= self.throttle_millis,
}
}
fn linger_ok(&self, now: u64) -> bool {
match self.first_pending_at {
None => false,
Some(first) => now.saturating_sub(first) >= self.linger_millis,
}
}
fn ready(&self, now: u64) -> bool {
self.linger_ok(now) && self.throttle_ok(now)
}
fn remaining_millis(&self, now: u64) -> u64 {
let linger_left = match self.first_pending_at {
None => u64::MAX,
Some(first) => self.linger_millis.saturating_sub(now.saturating_sub(first)),
};
let throttle_left = match self.last_sent_at {
None => 0,
Some(prev) => self.throttle_millis.saturating_sub(now.saturating_sub(prev)),
};
linger_left.max(throttle_left)
}
fn on_flush(&mut self, now: u64) {
self.first_pending_at = None;
self.last_sent_at = Some(now);
}
}
struct ThrottleState {
gate: FlushGate,
pending: Vec<StagedBatch>,
}
impl ThrottleState {
fn new(throttle: Duration, linger: Duration) -> Self {
Self {
gate: FlushGate::new(throttle, linger),
pending: Vec::new(),
}
}
fn enabled(&self) -> bool {
self.gate.shaping()
}
}
struct WarmingBuffer {
buffered: Vec<StagedBatch>,
cap: usize,
overflowed: bool,
}
impl WarmingBuffer {
fn new(cap: usize) -> Self {
Self {
buffered: Vec::new(),
cap,
overflowed: false,
}
}
fn push(&mut self, batch: StagedBatch) {
if self.overflowed {
return;
}
if self.buffered.len() >= self.cap {
self.overflowed = true;
self.buffered.clear();
return;
}
self.buffered.push(batch);
}
}
struct MemberPending {
frames: Vec<Frame>,
gate: FlushGate,
}
impl MemberPending {
fn new(linger: Duration, throttle: Duration) -> Self {
Self {
frames: Vec::new(),
gate: FlushGate::new(throttle, linger),
}
}
fn push(&mut self, frame: Frame, now: u64) {
self.gate.on_pending(now);
self.frames.push(frame);
}
fn ready(&self, now: u64) -> bool {
!self.frames.is_empty() && self.gate.ready(now)
}
fn remaining_millis(&self, now: u64) -> u64 {
self.gate.remaining_millis(now)
}
}
struct BatchState<S: WireSink> {
connection_id: ConnectionId,
sink: S,
format: S::Format,
member_ids: Vec<SubscriptionId>,
pending: DashMap<SubscriptionId, MemberPending>,
lingers: HashMap<SubscriptionId, Duration>,
throttles: HashMap<SubscriptionId, Duration>,
}
pub struct SubscriptionRegistry<S: WireSink> {
subscriptions: DashMap<SubscriptionId, SubscriptionState<S>>,
connections: DashMap<ConnectionId, Vec<SubscriptionId>>,
batches: DashMap<BatchId, BatchState<S>>,
connection_batches: DashMap<ConnectionId, Vec<BatchId>>,
clock: Clock,
throttle_pending: AtomicUsize,
wakers: Mutex<Vec<Arc<Notify>>>,
}
impl<S: WireSink> SubscriptionRegistry<S> {
pub fn new(clock: Clock) -> Self {
Self {
subscriptions: DashMap::new(),
connections: DashMap::new(),
batches: DashMap::new(),
connection_batches: DashMap::new(),
clock,
throttle_pending: AtomicUsize::new(0),
wakers: Mutex::new(Vec::new()),
}
}
#[allow(clippy::too_many_arguments)]
pub fn subscribe(
&self,
subscription_id: SubscriptionId,
connection_id: ConnectionId,
query: String,
sink: S,
format: S::Format,
warming_cap: Option<usize>,
throttle: Duration,
linger: Duration,
) {
self.subscriptions.insert(
subscription_id,
SubscriptionState {
connection_id,
sink,
format,
query,
batch_id: None,
warming: warming_cap.map(WarmingBuffer::new),
throttle: ThrottleState::new(throttle, linger),
},
);
self.connections.entry(connection_id).or_default().push(subscription_id);
debug!(
"Registered subscription {} for connection {} (warming_cap={:?})",
subscription_id, connection_id, warming_cap
);
}
pub fn promote_to_live(&self, subscription_id: SubscriptionId) -> PromoteResult {
let mut state = match self.subscriptions.get_mut(&subscription_id) {
Some(s) => s,
None => return PromoteResult::NotFound,
};
let (buffered, batch_id, format, sink) = match self.drain_warming(&mut state) {
Ok(drain) => drain,
Err(result) => return result,
};
drop(state);
let count = buffered.len();
match batch_id {
Some(batch_id) => self.flush_buffered_to_batch(batch_id, subscription_id, buffered, count),
None => self.flush_buffered_to_sink(subscription_id, buffered, format, sink, count),
}
}
#[inline]
fn drain_warming(&self, state: &mut SubscriptionState<S>) -> Result<PromoteDrain<S>, PromoteResult> {
let warming = match state.warming.as_ref() {
Some(w) => w,
None => return Err(PromoteResult::NotWarming),
};
if warming.overflowed {
state.warming = None;
return Err(PromoteResult::Overflowed);
}
let buffered = state.warming.take().expect("warming present").buffered;
let batch_id = state.batch_id;
let format = state.format;
let sink = state.sink.clone();
Ok((buffered, batch_id, format, sink))
}
#[inline]
fn flush_buffered_to_batch(
&self,
batch_id: BatchId,
subscription_id: SubscriptionId,
buffered: Vec<StagedBatch>,
count: usize,
) -> PromoteResult {
let Some(batch) = self.batches.get(&batch_id) else {
return PromoteResult::Disconnected;
};
{
let now = self.clock.now().to_millis();
let linger = batch.lingers.get(&subscription_id).copied().unwrap_or(Duration::zero());
let throttle = batch.throttles.get(&subscription_id).copied().unwrap_or(Duration::zero());
let mut entry = batch
.pending
.entry(subscription_id)
.or_insert_with(|| MemberPending::new(linger, throttle));
for (op, columns) in buffered {
entry.push(Frame::from(columns).with_op(op), now);
}
}
for waker in self.wakers.lock().iter() {
waker.notify_one();
}
PromoteResult::Promoted(count)
}
#[inline]
fn flush_buffered_to_sink(
&self,
subscription_id: SubscriptionId,
buffered: Vec<StagedBatch>,
format: S::Format,
sink: S,
count: usize,
) -> PromoteResult {
for (op, columns) in buffered {
match sink.send_change(subscription_id, op, columns, format) {
DeliveryResult::Delivered => {}
DeliveryResult::Disconnected => return PromoteResult::Disconnected,
}
}
PromoteResult::Promoted(count)
}
#[allow(clippy::too_many_arguments)]
pub fn register_batch(
&self,
connection_id: ConnectionId,
members: Vec<(SubscriptionId, Duration)>,
sink: S,
format: S::Format,
clock: &Clock,
rng: &Rng,
) -> BatchId {
let batch_id = BatchId(Uuid7::generate(clock, rng));
let member_ids: Vec<SubscriptionId> = members.iter().map(|(id, _)| *id).collect();
let lingers: HashMap<SubscriptionId, Duration> = members.into_iter().collect();
let mut throttles: HashMap<SubscriptionId, Duration> = HashMap::new();
for member_id in &member_ids {
if let Some(mut state) = self.subscriptions.get_mut(member_id) {
state.batch_id = Some(batch_id);
let throttle_millis = state.throttle.gate.throttle_millis;
if throttle_millis != 0 {
throttles.insert(
*member_id,
Duration::from_milliseconds(throttle_millis as i64).unwrap(),
);
}
}
}
self.batches.insert(
batch_id,
BatchState {
connection_id,
sink,
format,
member_ids: member_ids.clone(),
pending: DashMap::new(),
lingers,
throttles,
},
);
self.connection_batches.entry(connection_id).or_default().push(batch_id);
debug!(
"Registered batch {} with {} members for connection {}",
batch_id,
member_ids.len(),
connection_id
);
batch_id
}
pub fn unsubscribe_batch(&self, batch_id: BatchId) -> Option<Vec<SubscriptionId>> {
let (_, state) = self.batches.remove(&batch_id)?;
let connection_id = state.connection_id;
let members = state.member_ids.clone();
self.remove_batch_members(&members);
self.prune_empty_connection(connection_id);
self.prune_empty_connection_batches(connection_id, batch_id);
debug!("Unsubscribed batch {} ({} members)", batch_id, members.len());
Some(members)
}
#[inline]
fn remove_batch_members(&self, members: &[SubscriptionId]) {
let mut removed_by_conn: HashMap<ConnectionId, HashSet<SubscriptionId>> = HashMap::new();
for member_id in members {
if let Some((_, sub_state)) = self.subscriptions.remove(member_id) {
removed_by_conn.entry(sub_state.connection_id).or_default().insert(*member_id);
}
}
for (conn_id, removed) in removed_by_conn {
if let Some(mut subs) = self.connections.get_mut(&conn_id) {
subs.retain(|id| !removed.contains(id));
}
}
}
#[inline]
fn prune_empty_connection(&self, connection_id: ConnectionId) {
let remove_connection =
self.connections.get(&connection_id).map(|subs| subs.is_empty()).unwrap_or(false);
if remove_connection {
self.connections.remove(&connection_id);
}
}
#[inline]
fn prune_empty_connection_batches(&self, connection_id: ConnectionId, batch_id: BatchId) {
let batches_empty = {
if let Some(mut batches) = self.connection_batches.get_mut(&connection_id) {
batches.retain(|id| *id != batch_id);
batches.is_empty()
} else {
false
}
};
if batches_empty {
self.connection_batches.remove(&connection_id);
}
}
pub fn batch_for(&self, subscription_id: &SubscriptionId) -> Option<BatchId> {
self.subscriptions.get(subscription_id).and_then(|state| state.batch_id)
}
pub fn batch_count(&self) -> usize {
self.batches.len()
}
pub fn push_batch_frames(
&self,
batch_id: BatchId,
subscription_id: SubscriptionId,
frames: Vec<Frame>,
) -> bool {
let Some(batch) = self.batches.get(&batch_id) else {
return false;
};
{
let now = self.clock.now().to_millis();
let linger = batch.lingers.get(&subscription_id).copied().unwrap_or(Duration::zero());
let throttle = batch.throttles.get(&subscription_id).copied().unwrap_or(Duration::zero());
let mut entry = batch
.pending
.entry(subscription_id)
.or_insert_with(|| MemberPending::new(linger, throttle));
for frame in frames {
entry.push(frame, now);
}
}
for waker in self.wakers.lock().iter() {
waker.notify_one();
}
true
}
pub fn emit_batch_member_closed(&self, batch_id: BatchId, subscription_id: SubscriptionId) -> bool {
let Some(batch) = self.batches.get(&batch_id) else {
return false;
};
matches!(batch.sink.send_batch_member_closed(batch_id, subscription_id), DeliveryResult::Delivered)
}
pub fn unsubscribe(&self, subscription_id: SubscriptionId) -> bool {
if let Some((_, state)) = self.subscriptions.remove(&subscription_id) {
if !state.throttle.pending.is_empty() {
self.throttle_pending.fetch_sub(1, Ordering::AcqRel);
}
let connection_id = state.connection_id;
let should_remove_connection = {
if let Some(mut subs) = self.connections.get_mut(&connection_id) {
subs.retain(|id| *id != subscription_id);
subs.is_empty()
} else {
false
}
};
if should_remove_connection {
self.connections.remove(&connection_id);
}
debug!("Unsubscribed subscription {}", subscription_id);
true
} else {
false
}
}
pub fn cleanup_connection(&self, connection_id: ConnectionId) -> Vec<SubscriptionId> {
if let Some((_, batch_ids)) = self.connection_batches.remove(&connection_id) {
for batch_id in &batch_ids {
self.batches.remove(batch_id);
}
}
if let Some((_, subscription_ids)) = self.connections.remove(&connection_id) {
for sub_id in &subscription_ids {
self.subscriptions.remove(sub_id);
}
debug!("Cleaned up subscriptions for disconnected connection {}", connection_id);
subscription_ids
} else {
Vec::new()
}
}
pub fn close_all(&self) {
self.subscriptions.clear();
self.batches.clear();
self.connections.clear();
self.connection_batches.clear();
self.throttle_pending.store(0, Ordering::Release);
}
#[allow(dead_code)]
pub fn subscription_count(&self) -> usize {
self.subscriptions.len()
}
#[allow(dead_code)]
pub fn connection_count(&self) -> usize {
self.connections.len()
}
#[allow(dead_code)]
pub fn log_stats(&self) {
trace!(
"Registry stats: {} subscriptions, {} connections",
self.subscriptions.len(),
self.connections.len()
);
}
#[inline]
fn deliver_to_batch_member(
&self,
batch_id: BatchId,
subscription_id: &SubscriptionId,
op: DiffType,
columns: Columns,
) -> DeliveryResult {
let Some(batch) = self.batches.get(&batch_id) else {
return DeliveryResult::Disconnected;
};
let now = self.clock.now().to_millis();
let linger = batch.lingers.get(subscription_id).copied().unwrap_or(Duration::zero());
let throttle = batch.throttles.get(subscription_id).copied().unwrap_or(Duration::zero());
batch.pending
.entry(*subscription_id)
.or_insert_with(|| MemberPending::new(linger, throttle))
.push(Frame::from(columns).with_op(op), now);
DeliveryResult::Delivered
}
#[inline]
fn send_now_and_mark(
&self,
subscription_id: &SubscriptionId,
op: DiffType,
columns: Columns,
format: S::Format,
sink: S,
now: u64,
) -> DeliveryResult {
match sink.send_change(*subscription_id, op, columns, format) {
DeliveryResult::Delivered => {
if let Some(mut s) = self.subscriptions.get_mut(subscription_id) {
s.throttle.gate.last_sent_at = Some(now);
}
DeliveryResult::Delivered
}
DeliveryResult::Disconnected => DeliveryResult::Disconnected,
}
}
#[inline]
fn queue_throttled(&self, state: &mut SubscriptionState<S>, op: DiffType, columns: Columns) -> DeliveryResult {
let was_empty = state.throttle.pending.is_empty();
state.throttle.pending.push((op, columns));
state.throttle.gate.on_pending(self.clock.now().to_millis());
if was_empty {
self.throttle_pending.fetch_add(1, Ordering::AcqRel);
}
DeliveryResult::Delivered
}
#[inline]
fn flush_ready_throttled(&self, now: u64, next_deadline: &mut Option<u64>) {
if self.throttle_pending.load(Ordering::Acquire) == 0 {
return;
}
let mut throttle_ready: Vec<ThrottleReady<S>> = Vec::new();
for mut entry in self.subscriptions.iter_mut() {
let sub_id = *entry.key();
let state = entry.value_mut();
if state.batch_id.is_some() || state.warming.is_some() {
continue;
}
if !state.throttle.enabled() || state.throttle.pending.is_empty() {
continue;
}
if !state.throttle.gate.ready(now) {
let rem = state.throttle.gate.remaining_millis(now);
*next_deadline = Some(next_deadline.map_or(rem, |d| d.min(rem)));
continue;
}
let drained = mem::take(&mut state.throttle.pending);
reifydb_assertions! {
let counted = self.throttle_pending.load(Ordering::Acquire);
assert!(
counted > 0,
"throttle_pending counter is {counted} while draining a subscription with \
non-empty pending; queue_throttled must increment once per empty->non-empty \
transition, so an underflowing fetch_sub here wraps the counter and the \
flush gate `load() > 0` would then fire forever, scanning every subscription \
on every flush even when nothing is pending"
);
}
self.throttle_pending.fetch_sub(1, Ordering::AcqRel);
state.throttle.gate.on_flush(now);
throttle_ready.push((sub_id, drained, state.format, state.sink.clone()));
}
let mut dead_subs: Vec<SubscriptionId> = Vec::new();
for (sub_id, drained, format, sink) in throttle_ready {
for (op, columns) in drained {
match sink.send_change(sub_id, op, columns, format) {
DeliveryResult::Delivered => {}
DeliveryResult::Disconnected => {
dead_subs.push(sub_id);
break;
}
}
}
}
for sub_id in dead_subs {
self.unsubscribe(sub_id);
}
}
#[inline]
fn flush_due_batches(&self, now: u64, next_deadline: &mut Option<u64>) {
let mut dead_batches: Vec<BatchId> = Vec::new();
for entry in self.batches.iter() {
let batch_id = *entry.key();
let batch = entry.value();
let mut due: Vec<(SubscriptionId, Vec<Frame>)> = Vec::new();
for mut e in batch.pending.iter_mut() {
let key = *e.key();
let member = e.value_mut();
if member.frames.is_empty() {
continue;
}
if member.ready(now) {
let frames = mem::take(&mut member.frames);
member.gate.on_flush(now);
due.push((key, frames));
} else {
let rem = member.remaining_millis(now);
*next_deadline = Some(next_deadline.map_or(rem, |d| d.min(rem)));
}
}
if due.is_empty() {
continue;
}
match batch.sink.send_batch_envelope(batch_id, batch.format, due) {
DeliveryResult::Delivered => {}
DeliveryResult::Disconnected => dead_batches.push(batch_id),
}
}
for batch_id in dead_batches {
if let Some(members) = self.unsubscribe_batch(batch_id) {
debug!("Batch {} push channel closed; cascaded {} members", batch_id, members.len());
}
}
}
}
impl<S: WireSink> SubscriptionDelivery for SubscriptionRegistry<S> {
fn try_deliver(&self, subscription_id: &SubscriptionId, op: DiffType, columns: Columns) -> DeliveryResult {
if let Some(batch_id) = self.batch_for(subscription_id) {
return self.deliver_to_batch_member(batch_id, subscription_id, op, columns);
}
let mut state = match self.subscriptions.get_mut(subscription_id) {
Some(s) => s,
None => return DeliveryResult::Disconnected,
};
if let Some(buffer) = state.warming.as_mut() {
buffer.push((op, columns));
return DeliveryResult::Delivered;
}
if state.throttle.enabled() {
let now = self.clock.now().to_millis();
if state.throttle.gate.linger_millis == 0
&& state.throttle.gate.throttle_ok(now)
&& state.throttle.pending.is_empty()
{
let format = state.format;
let sink = state.sink.clone();
drop(state);
self.send_now_and_mark(subscription_id, op, columns, format, sink, now)
} else {
self.queue_throttled(&mut state, op, columns)
}
} else {
let format = state.format;
let sink = state.sink.clone();
drop(state);
sink.send_change(*subscription_id, op, columns, format)
}
}
fn active_subscriptions(&self) -> Vec<SubscriptionId> {
self.subscriptions.iter().map(|entry| *entry.key()).collect()
}
fn active_subscriptions_into(&self, out: &mut Vec<SubscriptionId>) {
out.extend(self.subscriptions.iter().map(|entry| *entry.key()));
}
fn register_waker(&self, waker: Arc<Notify>) {
self.wakers.lock().push(waker);
}
fn simulated_clock(&self) -> bool {
self.clock.is_mock()
}
#[instrument(name = "server::flush", level = "debug", skip_all)]
fn flush(&self) -> Option<Duration> {
let now = self.clock.now().to_millis();
let mut next_deadline: Option<u64> = None;
self.flush_ready_throttled(now, &mut next_deadline);
if self.batches.is_empty() {
return next_deadline.map(|ms| Duration::from_milliseconds(ms as i64).unwrap());
}
self.flush_due_batches(now, &mut next_deadline);
next_deadline.map(|ms| Duration::from_milliseconds(ms as i64).unwrap())
}
}