use std::{collections::HashMap, sync::Arc, time::Duration};
use reifydb_client::{HydrationConfig as ClientHydrationConfig, SubscriptionConfig as ClientSubscriptionConfig};
use reifydb_core::{
common::CommitVersion,
interface::catalog::{id::SubscriptionId, subscription::HydrationConfig},
metric::ExecutionMetrics,
value::column::columns::Columns,
};
use reifydb_engine::subscription::{HydrateError, SubscriptionServiceRef};
use reifydb_remote_proxy::{RemoteSubscription, connect_remote, proxy_remote_to_sink};
use reifydb_subscription::{batch::BatchId, delivery::DeliveryResult};
use reifydb_transaction::multi::lease::VersionLeaseGuard;
use reifydb_value::value::{frame::frame::Frame, identity::IdentityId};
use tokio::{spawn, sync::watch::Receiver as WatchReceiver, task::JoinHandle};
use tracing::{info, warn};
use crate::{
interceptor::RequestMetadata,
state::AppState,
subscription::{
cleanup::cleanup_subscription,
create::{CreateSubscriptionResult, create_subscription},
errors::CreateSubscriptionError,
hydrate::run_hydrate,
registry::{ConnectionId, PromoteResult, SubscriptionRegistry},
wire_sink::{BatchSubscribedMember, WireSink},
},
};
#[derive(Debug, Clone)]
pub struct BatchMemberInfo {
pub index: usize,
pub subscription_id: SubscriptionId,
}
pub struct BatchAck {
pub batch_id: BatchId,
pub members: Vec<BatchMemberInfo>,
pub remote_handles: Vec<JoinHandle<()>>,
}
pub struct SubscribeAck {
pub subscription_id: SubscriptionId,
pub remote_handle: Option<JoinHandle<()>>,
}
#[derive(Debug)]
pub enum SubscribeError {
Create(CreateSubscriptionError),
RemoteConnect(String),
InvalidRemoteId,
LeaseFailed {
code: &'static str,
message: String,
},
HydrationBackpressure,
HydrationFailed {
error: HydrateError,
rql: String,
max_rows: u64,
},
HydrationServiceUnavailable(String),
}
#[derive(Debug)]
pub enum BatchSubscribeError {
Empty,
Create(CreateSubscriptionError),
RemoteConnect(String),
InvalidRemoteId,
LeaseFailed {
code: &'static str,
message: String,
},
HydrationBackpressure,
HydrationFailed {
error: HydrateError,
rql: String,
max_rows: u64,
},
HydrationServiceUnavailable(String),
}
enum ResolvedBatchMember {
Local {
index: usize,
subscription_id: SubscriptionId,
query: String,
},
Remote {
index: usize,
subscription_id: SubscriptionId,
remote_sub: Box<RemoteSubscription>,
},
}
impl ResolvedBatchMember {
fn index(&self) -> usize {
match self {
Self::Local {
index,
..
}
| Self::Remote {
index,
..
} => *index,
}
}
fn subscription_id(&self) -> SubscriptionId {
match self {
Self::Local {
subscription_id,
..
}
| Self::Remote {
subscription_id,
..
} => *subscription_id,
}
}
}
#[allow(clippy::too_many_arguments)]
pub async fn handle_subscribe<S: WireSink>(
state: &AppState,
connection_id: ConnectionId,
identity: IdentityId,
rql: String,
sink: S,
registry: &Arc<SubscriptionRegistry<S>>,
format: S::Format,
shutdown: WatchReceiver<bool>,
metadata: RequestMetadata,
) -> Result<SubscribeAck, SubscribeError> {
match create_subscription(state, identity, &rql, metadata).await {
Ok(CreateSubscriptionResult::Local {
id: subscription_id,
hydration,
throttle,
}) => {
let server_cap = state.subscribe_max_hydration_rows();
let throttle = state.clamp_throttle(throttle);
let max_rows = match hydration.max_rows {
Some(n) if n > server_cap => {
warn!("clamping hydration.max_rows from {} to server cap {}", n, server_cap);
server_cap
}
Some(n) => n,
None => server_cap,
};
let warming_cap = if hydration.enabled {
Some(max_rows as usize)
} else {
None
};
registry.subscribe(
subscription_id,
connection_id,
rql.clone(),
sink.clone(),
format,
warming_cap,
throttle,
);
if !matches!(sink.send_subscribed(subscription_id), DeliveryResult::Delivered) {
abort_warming(state, registry, subscription_id).await;
return Err(SubscribeError::LeaseFailed {
code: "STREAM_CLOSED",
message: "Client stream closed before Subscribed could be delivered"
.to_string(),
});
}
if hydration.enabled {
let lease = match state.engine().acquire_current_snapshot_lease() {
Ok((_, lease)) => lease,
Err(e) => {
let code = if e.0.code == "TXN_012" {
"HYDRATION_VERSION_EVICTED"
} else {
"PIN_VERSION_FAILED"
};
abort_warming(state, registry, subscription_id).await;
return Err(SubscribeError::LeaseFailed {
code,
message: e.to_string(),
});
}
};
if let Err(err) = run_member_hydrate(
state,
registry,
&sink,
subscription_id,
&rql,
identity,
lease,
max_rows,
format,
)
.await
{
return Err(err.into());
}
} else {
let _ = registry.promote_to_live(subscription_id);
}
info!(
"Connection {} subscribed: subscription_id={} hydration_enabled={}",
connection_id, subscription_id, hydration.enabled
);
Ok(SubscribeAck {
subscription_id,
remote_handle: None,
})
}
Ok(CreateSubscriptionResult::Remote {
address,
body,
token: ns_token,
hydration,
throttle,
}) => {
let client_format = S::client_wire_format(format);
let config = ClientSubscriptionConfig {
hydration: ClientHydrationConfig {
enabled: hydration.enabled,
max_rows: hydration.max_rows,
},
throttle,
};
let remote_sub = connect_remote(&address, &body, config, ns_token.as_deref(), client_format)
.await
.map_err(|e| SubscribeError::RemoteConnect(e.to_string()))?;
let remote_id = remote_sub.subscription_id().to_string();
let subscription_id =
SubscriptionId(remote_id.parse::<u64>().map_err(|_| SubscribeError::InvalidRemoteId)?);
if !matches!(sink.send_subscribed(subscription_id), DeliveryResult::Delivered) {
return Err(SubscribeError::LeaseFailed {
code: "STREAM_CLOSED",
message: "Client stream closed before Subscribed could be delivered"
.to_string(),
});
}
let sink_for_proxy = sink.clone();
let sink_for_close = sink.clone();
let handle = spawn(async move {
proxy_remote_to_sink(remote_sub, shutdown, move |payload| {
matches!(
sink_for_proxy.send_remote_change(subscription_id, payload, format),
DeliveryResult::Delivered
)
})
.await;
let _ = sink_for_close.send_closed(subscription_id);
});
info!("Connection {} subscribed to remote: subscription_id={}", connection_id, subscription_id);
Ok(SubscribeAck {
subscription_id,
remote_handle: Some(handle),
})
}
Err(e) => Err(SubscribeError::Create(e)),
}
}
#[allow(clippy::too_many_arguments)]
pub async fn handle_batch_subscribe<S: WireSink>(
state: &AppState,
connection_id: ConnectionId,
identity: IdentityId,
queries: &[String],
sink: S,
registry: &Arc<SubscriptionRegistry<S>>,
format: S::Format,
shutdown: WatchReceiver<bool>,
metadata: RequestMetadata,
) -> Result<BatchAck, BatchSubscribeError> {
if queries.is_empty() {
return Err(BatchSubscribeError::Empty);
}
let mut resolved: Vec<ResolvedBatchMember> = Vec::with_capacity(queries.len());
let mut local_hydrations: Vec<(SubscriptionId, String, HydrationConfig, Option<Duration>)> = Vec::new();
for (index, user_rql) in queries.iter().enumerate() {
match create_subscription(state, identity, user_rql, metadata.clone()).await {
Ok(CreateSubscriptionResult::Local {
id: subscription_id,
hydration,
throttle,
}) => {
local_hydrations.push((subscription_id, user_rql.clone(), hydration, throttle));
resolved.push(ResolvedBatchMember::Local {
index,
subscription_id,
query: user_rql.clone(),
});
}
Ok(CreateSubscriptionResult::Remote {
address,
body,
token: ns_token,
hydration,
throttle,
}) => {
let client_format = S::client_wire_format(format);
let config = ClientSubscriptionConfig {
hydration: ClientHydrationConfig {
enabled: hydration.enabled,
max_rows: hydration.max_rows,
},
throttle,
};
let remote_sub = match connect_remote(
&address,
&body,
config,
ns_token.as_deref(),
client_format,
)
.await
{
Ok(s) => s,
Err(e) => {
rollback_batch_members(state, &resolved).await;
return Err(BatchSubscribeError::RemoteConnect(e.to_string()));
}
};
let remote_id = remote_sub.subscription_id().to_string();
let subscription_id = match remote_id.parse::<u64>() {
Ok(n) => SubscriptionId(n),
Err(_) => {
rollback_batch_members(state, &resolved).await;
return Err(BatchSubscribeError::InvalidRemoteId);
}
};
resolved.push(ResolvedBatchMember::Remote {
index,
subscription_id,
remote_sub: Box::new(remote_sub),
});
}
Err(e) => {
rollback_batch_members(state, &resolved).await;
return Err(BatchSubscribeError::Create(e));
}
}
}
let server_cap = state.subscribe_max_hydration_rows();
let mut effective_max_rows: HashMap<SubscriptionId, u64> = HashMap::new();
for (sub_id, _, hydration, _) in &local_hydrations {
let max_rows = match hydration.max_rows {
Some(n) if n > server_cap => {
warn!("clamping hydration.max_rows from {} to server cap {}", n, server_cap);
server_cap
}
Some(n) => n,
None => server_cap,
};
effective_max_rows.insert(*sub_id, max_rows);
}
for member in &resolved {
if let ResolvedBatchMember::Local {
subscription_id,
query,
..
} = member
{
let warming_cap = local_hydrations
.iter()
.find(|(sid, _, _, _)| sid == subscription_id)
.and_then(|(_, _, h, _)| {
if h.enabled {
Some(*effective_max_rows.get(subscription_id).unwrap_or(&server_cap)
as usize)
} else {
None
}
});
let throttle = state.clamp_throttle(
local_hydrations
.iter()
.find(|(sid, _, _, _)| sid == subscription_id)
.and_then(|(_, _, _, t)| *t),
);
registry.subscribe(
*subscription_id,
connection_id,
query.clone(),
sink.clone(),
format,
warming_cap,
throttle,
);
}
}
let member_ids: Vec<SubscriptionId> = resolved.iter().map(|m| m.subscription_id()).collect();
let batch_id =
registry.register_batch(connection_id, member_ids, sink.clone(), format, state.clock(), state.rng());
let mut remote_members_taken: Vec<(SubscriptionId, RemoteSubscription)> = Vec::new();
let mut members_for_ack: Vec<BatchMemberInfo> = Vec::with_capacity(resolved.len());
for member in resolved {
members_for_ack.push(BatchMemberInfo {
index: member.index(),
subscription_id: member.subscription_id(),
});
if let ResolvedBatchMember::Remote {
subscription_id,
remote_sub,
..
} = member
{
remote_members_taken.push((subscription_id, *remote_sub));
}
}
let subscribed_members: Vec<BatchSubscribedMember> = members_for_ack
.iter()
.map(|m| BatchSubscribedMember {
index: m.index,
subscription_id: m.subscription_id,
})
.collect();
if !matches!(sink.send_batch_subscribed(batch_id, &subscribed_members), DeliveryResult::Delivered) {
registry.unsubscribe_batch(batch_id);
rollback_batch_members(state, &[]).await;
return Err(BatchSubscribeError::LeaseFailed {
code: "STREAM_CLOSED",
message: "Client stream closed before BatchSubscribed could be delivered".to_string(),
});
}
let any_hydration = local_hydrations.iter().any(|(_, _, h, _)| h.enabled);
if any_hydration {
let lease = match state.engine().acquire_current_snapshot_lease() {
Ok((_, lease)) => lease,
Err(e) => {
let code = if e.0.code == "TXN_012" {
"HYDRATION_VERSION_EVICTED"
} else {
"PIN_VERSION_FAILED"
};
registry.unsubscribe_batch(batch_id);
return Err(BatchSubscribeError::LeaseFailed {
code,
message: e.to_string(),
});
}
};
for (sub_id, rql, hydration, _) in &local_hydrations {
if !hydration.enabled {
continue;
}
let max_rows = *effective_max_rows.get(sub_id).unwrap_or(&server_cap);
if let Err(err) = run_member_hydrate(
state,
registry,
&sink,
*sub_id,
rql,
identity,
lease.clone(),
max_rows,
format,
)
.await
{
registry.unsubscribe_batch(batch_id);
return Err(err.into_batch());
}
}
} else {
for (sub_id, _, _, _) in &local_hydrations {
let _ = registry.promote_to_live(*sub_id);
}
}
let mut remote_handles: Vec<JoinHandle<()>> = Vec::with_capacity(remote_members_taken.len());
for (subscription_id, remote_sub) in remote_members_taken {
let registry_clone = Arc::clone(registry);
let proxy_shutdown = shutdown.clone();
let handle = spawn(async move {
run_batch_remote_proxy(registry_clone, batch_id, subscription_id, remote_sub, proxy_shutdown)
.await;
});
remote_handles.push(handle);
}
info!("Connection {} created batch {} with {} members", connection_id, batch_id, members_for_ack.len());
Ok(BatchAck {
batch_id,
members: members_for_ack,
remote_handles,
})
}
pub async fn handle_batch_unsubscribe<S: WireSink>(
state: &AppState,
registry: &Arc<SubscriptionRegistry<S>>,
batch_id: BatchId,
) -> Option<Vec<SubscriptionId>> {
let members = registry.unsubscribe_batch(batch_id)?;
for subscription_id in &members {
if let Err(e) = cleanup_subscription(state, *subscription_id).await {
warn!("Failed to cleanup batch member subscription {} from database: {:?}", subscription_id, e);
}
}
Some(members)
}
async fn run_batch_remote_proxy<S: WireSink>(
registry: Arc<SubscriptionRegistry<S>>,
batch_id: BatchId,
subscription_id: SubscriptionId,
remote_sub: RemoteSubscription,
shutdown: WatchReceiver<bool>,
) {
let registry_push = Arc::clone(®istry);
proxy_remote_to_sink(remote_sub, shutdown, move |payload| {
let frames = payload.into_frames();
registry_push.push_batch_frames(batch_id, subscription_id, frames)
})
.await;
let _ = registry.emit_batch_member_closed(batch_id, subscription_id);
}
#[allow(clippy::too_many_arguments)]
async fn run_member_hydrate<S: WireSink>(
state: &AppState,
registry: &Arc<SubscriptionRegistry<S>>,
sink: &S,
subscription_id: SubscriptionId,
rql: &str,
identity: IdentityId,
lease: VersionLeaseGuard,
max_rows: u64,
format: S::Format,
) -> Result<(), MemberHydrateError> {
let service: SubscriptionServiceRef = match state.engine().services().ioc.resolve() {
Ok(s) => s,
Err(e) => {
abort_warming(state, registry, subscription_id).await;
return Err(MemberHydrateError {
rql: rql.to_string(),
max_rows,
kind: MemberHydrateErrorKind::ServiceUnavailable(e.to_string()),
});
}
};
let engine = state.engine_clone();
let (version, outcome): (CommitVersion, Option<(Columns, ExecutionMetrics)>) =
match run_hydrate(service, engine, subscription_id, identity, lease, max_rows).await {
Ok(t) => t,
Err(err) => {
abort_warming(state, registry, subscription_id).await;
return Err(MemberHydrateError {
rql: rql.to_string(),
max_rows,
kind: MemberHydrateErrorKind::Failed(err),
});
}
};
if let Some((cols, metrics)) = outcome {
let row_count = cols.row_count();
info!(
subscription_id = subscription_id.0,
version = version.0,
total_us = metrics.total.microseconds().unwrap_or(0),
compute_us = metrics.compute.microseconds().unwrap_or(0),
statement_count = metrics.statements.len(),
row_count = row_count,
fingerprint = %metrics.fingerprint.to_hex(),
"hydrate completed"
);
if let Some(batch_id) = registry.batch_for(&subscription_id) {
let _ = sink.send_batch_envelope(
batch_id,
format,
vec![(subscription_id, vec![Frame::from(cols)])],
);
} else {
let _ = sink.send_change(subscription_id, cols, format);
}
}
match registry.promote_to_live(subscription_id) {
PromoteResult::Promoted(_)
| PromoteResult::NotWarming
| PromoteResult::NotFound
| PromoteResult::Disconnected => Ok(()),
PromoteResult::Overflowed => Err(MemberHydrateError {
rql: rql.to_string(),
max_rows,
kind: MemberHydrateErrorKind::Backpressure,
}),
}
}
struct MemberHydrateError {
rql: String,
max_rows: u64,
kind: MemberHydrateErrorKind,
}
enum MemberHydrateErrorKind {
ServiceUnavailable(String),
Failed(HydrateError),
Backpressure,
}
impl MemberHydrateError {
fn into_batch(self) -> BatchSubscribeError {
match self.kind {
MemberHydrateErrorKind::ServiceUnavailable(msg) => {
BatchSubscribeError::HydrationServiceUnavailable(msg)
}
MemberHydrateErrorKind::Failed(err) => BatchSubscribeError::HydrationFailed {
error: err,
rql: self.rql,
max_rows: self.max_rows,
},
MemberHydrateErrorKind::Backpressure => BatchSubscribeError::HydrationBackpressure,
}
}
}
impl From<MemberHydrateError> for SubscribeError {
fn from(e: MemberHydrateError) -> Self {
match e.kind {
MemberHydrateErrorKind::ServiceUnavailable(msg) => {
SubscribeError::HydrationServiceUnavailable(msg)
}
MemberHydrateErrorKind::Failed(err) => SubscribeError::HydrationFailed {
error: err,
rql: e.rql,
max_rows: e.max_rows,
},
MemberHydrateErrorKind::Backpressure => SubscribeError::HydrationBackpressure,
}
}
}
async fn rollback_batch_members(state: &AppState, resolved: &[ResolvedBatchMember]) {
for member in resolved {
if let ResolvedBatchMember::Local {
subscription_id,
..
} = member && let Err(e) = cleanup_subscription(state, *subscription_id).await
{
warn!("Failed to cleanup partial batch member {} during rollback: {:?}", subscription_id, e);
}
}
}
async fn abort_warming<S: WireSink>(
state: &AppState,
registry: &Arc<SubscriptionRegistry<S>>,
subscription_id: SubscriptionId,
) {
registry.unsubscribe(subscription_id);
if let Err(e) = cleanup_subscription(state, subscription_id).await {
warn!("Failed to cleanup subscription {} after hydrate failure: {:?}", subscription_id, e);
}
}