use std::collections::VecDeque;
#[cfg(not(test))]
use std::sync::atomic::Ordering;
use std::sync::atomic::{AtomicBool, AtomicU64};
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::{Semaphore, oneshot};
use tracing::{Instrument as _, instrument::WithSubscriber as _};
use super::{
Transaction, TransactionCommitOutcome, commit_transaction_cohort,
transaction_is_file_cohort_eligible, transactions_can_share_cohort,
};
use crate::LixError;
use crate::functions::FunctionContext;
use crate::observe_invalidation::ObserveInvalidation;
use crate::storage_adapter::Storage;
use crate::telemetry::{
ActiveTelemetrySpan, SpanContext, TelemetryAttribute, TelemetryContext, TelemetrySink,
Status, TRANSACTION_NOTIFY, TRANSACTION_STORAGE, TRANSACTION_WAIT,
current_telemetry_context, next_commit_cohort_id,
};
const COMMIT_QUEUE_CAPACITY: usize = 256;
const COMMIT_COHORT_CAPACITY: usize = 256;
struct CommitRequest<StorageImpl>
where
StorageImpl: Storage + 'static,
{
transaction: Transaction<StorageImpl>,
runtime_functions: FunctionContext,
result: oneshot::Sender<Result<TransactionCommitOutcome, LixError>>,
file_cohort_eligible: bool,
telemetry_context: Option<TelemetryContext>,
wait_span: Option<ActiveTelemetrySpan>,
tracing_parent: tracing::Span,
tracing_dispatch: tracing::Dispatch,
_capacity: tokio::sync::OwnedSemaphorePermit,
}
#[derive(Clone)]
pub(crate) struct CommitCoordinator<StorageImpl>
where
StorageImpl: Storage + 'static,
{
inner: Arc<CommitCoordinatorInner<StorageImpl>>,
}
struct CommitCoordinatorInner<StorageImpl>
where
StorageImpl: Storage + 'static,
{
collaboration_write_gate: Arc<tokio::sync::Mutex<()>>,
observe_invalidation: Arc<ObserveInvalidation>,
capacity: Arc<Semaphore>,
telemetry: Option<Arc<dyn TelemetrySink>>,
checkpoint_gc_running: AtomicBool,
checkpoint_gc_not_before_sequence: AtomicU64,
state: Mutex<CommitCoordinatorState<StorageImpl>>,
#[cfg(test)]
stats: CommitCoordinatorStats,
}
struct CommitCoordinatorState<StorageImpl>
where
StorageImpl: Storage + 'static,
{
running: bool,
queue: VecDeque<CommitRequest<StorageImpl>>,
}
impl<StorageImpl> Default for CommitCoordinatorState<StorageImpl>
where
StorageImpl: Storage + 'static,
{
fn default() -> Self {
Self {
running: false,
queue: VecDeque::new(),
}
}
}
#[cfg(test)]
#[derive(Default)]
struct CommitCoordinatorStats {
cohort_count: AtomicUsize,
commit_count: AtomicUsize,
max_cohort_size: AtomicUsize,
checkpoint_gc_post_commit_hooks: AtomicUsize,
}
impl<StorageImpl> CommitCoordinator<StorageImpl>
where
StorageImpl: Storage + Clone + Send + Sync + 'static,
{
pub(crate) fn new(
collaboration_write_gate: Arc<tokio::sync::Mutex<()>>,
observe_invalidation: Arc<ObserveInvalidation>,
telemetry: Option<Arc<dyn TelemetrySink>>,
) -> Self {
Self {
inner: Arc::new(CommitCoordinatorInner {
collaboration_write_gate,
observe_invalidation,
capacity: Arc::new(Semaphore::new(COMMIT_QUEUE_CAPACITY)),
telemetry,
checkpoint_gc_running: AtomicBool::new(false),
checkpoint_gc_not_before_sequence: AtomicU64::new(0),
state: Mutex::new(CommitCoordinatorState::default()),
#[cfg(test)]
stats: CommitCoordinatorStats::default(),
}),
}
}
pub(crate) fn try_begin_checkpoint_gc(&self, checkpoint_sequence: u64) -> bool {
if checkpoint_sequence
< self
.inner
.checkpoint_gc_not_before_sequence
.load(Ordering::Acquire)
{
return false;
}
self.inner
.checkpoint_gc_running
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
}
#[cfg(test)]
pub(crate) fn record_checkpoint_gc_post_commit_hook(&self) {
self.inner
.stats
.checkpoint_gc_post_commit_hooks
.fetch_add(1, Ordering::Relaxed);
}
#[cfg(test)]
pub(crate) fn checkpoint_gc_post_commit_hooks(&self) -> usize {
self.inner
.stats
.checkpoint_gc_post_commit_hooks
.load(Ordering::Relaxed)
}
pub(crate) fn defer_checkpoint_gc_until(&self, checkpoint_sequence: u64) {
self.inner
.checkpoint_gc_not_before_sequence
.fetch_max(checkpoint_sequence, Ordering::AcqRel);
}
pub(crate) fn finish_checkpoint_gc(&self) {
self.inner
.checkpoint_gc_running
.store(false, Ordering::Release);
}
pub(crate) async fn commit(
&self,
transaction: Transaction<StorageImpl>,
runtime_functions: FunctionContext,
) -> Result<TransactionCommitOutcome, LixError> {
let wait_span = self.inner.telemetry.as_ref().and_then(|sink| {
ActiveTelemetrySpan::start_if_enabled(
sink,
&TRANSACTION_WAIT,
vec![TelemetryAttribute::string(
"lix.wait.reason",
"commit_coordinator",
)],
)
});
let capacity_future = Arc::clone(&self.inner.capacity).acquire_owned();
let capacity_result = match wait_span.as_ref() {
Some(span) => span.instrument(capacity_future).await,
None => capacity_future.await,
};
let capacity = match capacity_result {
Ok(capacity) => capacity,
Err(_) => {
if let Some(span) = wait_span {
span.finish(Status::error("commit coordinator closed"), Vec::new());
}
return Err(coordinator_closed());
}
};
let file_cohort_eligible = transaction_is_file_cohort_eligible(&transaction);
let telemetry_context = wait_span
.as_ref()
.map(ActiveTelemetrySpan::telemetry_context)
.or_else(current_telemetry_context);
let (result, receive) = oneshot::channel();
let leads = self.enqueue(CommitRequest {
transaction,
runtime_functions,
result,
file_cohort_eligible,
telemetry_context,
wait_span,
tracing_parent: tracing::Span::current(),
tracing_dispatch: tracing::dispatcher::get_default(Clone::clone),
_capacity: capacity,
});
if leads {
#[cfg(not(target_family = "wasm"))]
if let Err(error) = self.spawn_driver() {
self.fail_queued(error);
}
#[cfg(target_family = "wasm")]
self.drive().await;
}
let receive_result = receive.await;
let outcome = match receive_result {
Ok(outcome) => outcome,
Err(_) => {
return Err(coordinator_closed());
}
};
outcome
}
#[cfg(not(target_family = "wasm"))]
fn spawn_driver(&self) -> Result<(), LixError> {
let coordinator = self.clone();
crate::background_task::spawn("lix-commit-coordinator", move || async move {
Box::pin(coordinator.drive()).await;
})
}
fn enqueue(&self, request: CommitRequest<StorageImpl>) -> bool {
{
let mut state = self
.inner
.state
.lock()
.unwrap_or_else(|error| error.into_inner());
state.queue.push_back(request);
if state.running {
false
} else {
state.running = true;
true
}
}
}
#[cfg(not(target_family = "wasm"))]
fn fail_queued(&self, error: LixError) {
let requests = {
let mut state = self
.inner
.state
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
state.running = false;
state.queue.drain(..).collect::<Vec<_>>()
};
for request in requests {
if let Some(span) = request.wait_span {
span.finish(
Status::error(error.code.clone()),
vec![TelemetryAttribute::string("error.type", error.code.clone())],
);
}
let _ = request.result.send(Err(error.clone()));
}
}
async fn drive(&self) {
let mut driver = CommitDriverGuard::new(&self.inner);
#[cfg(target_family = "wasm")]
tokio::task::yield_now().await;
loop {
let mut cohort = {
let mut state = self
.inner
.state
.lock()
.unwrap_or_else(|error| error.into_inner());
let take = state.queue.len().min(COMMIT_COHORT_CAPACITY);
if take == 0 {
state.running = false;
driver.disarm();
return;
}
let compatible = state
.queue
.front()
.map(|leader| {
state
.queue
.iter()
.take(take)
.take_while(|candidate| {
transactions_can_share_cohort(
&leader.transaction,
&candidate.transaction,
leader.file_cohort_eligible,
candidate.file_cohort_eligible,
)
})
.count()
})
.unwrap_or(0);
let take = compatible.max(1);
state.queue.drain(..take).collect::<Vec<_>>()
};
#[cfg(test)]
{
self.inner
.stats
.cohort_count
.fetch_add(1, Ordering::Relaxed);
self.inner
.stats
.commit_count
.fetch_add(cohort.len(), Ordering::Relaxed);
self.inner
.stats
.max_cohort_size
.fetch_max(cohort.len(), Ordering::Relaxed);
}
let _gate = self
.inner
.collaboration_write_gate
.lock()
.instrument(tracing::debug_span!(
target: "lix_transaction",
"lix.transaction.commit_cohort",
cohort_size = cohort.len(),
))
.await;
for request in &mut cohort {
if let Some(span) = request.wait_span.take() {
span.finish(Status::Unset, Vec::new());
}
}
let transaction_count = cohort.len();
let telemetry_context = cohort_telemetry_context(&cohort);
let tracing_parent = cohort.first().map_or_else(tracing::Span::none, |request| {
request.tracing_parent.clone()
});
let tracing_dispatch = cohort
.first()
.map(|request| request.tracing_dispatch.clone());
let mut senders = Vec::with_capacity(transaction_count);
let mut inputs = Vec::with_capacity(cohort.len());
let mut checkpoint_gc_sequences = Vec::with_capacity(cohort.len());
for request in cohort {
senders.push((request.result, request._capacity));
checkpoint_gc_sequences.push(request.transaction.checkpoint_gc_sequence());
inputs.push((request.transaction, request.runtime_functions));
}
let commit_and_notify = async {
let outcomes = Box::pin(commit_transaction_cohort(inputs)).await;
if let Some(outcome) = outcomes.iter().find_map(|result| result.as_ref().ok()) {
let notify = ActiveTelemetrySpan::start_current(
&TRANSACTION_NOTIFY,
vec![TelemetryAttribute::i64(
"lix.transaction.count",
i64::try_from(transaction_count).unwrap_or(i64::MAX),
)],
);
let _entered = notify.as_ref().map(ActiveTelemetrySpan::enter);
self.inner
.observe_invalidation
.bump_if_storage_changed(&outcome.storage_stats);
drop(_entered);
if let Some(notify) = notify {
notify.finish(Status::Unset, Vec::new());
}
}
outcomes
}
.instrument(tracing_parent);
let commit_and_notify = match tracing_dispatch {
Some(dispatch) => commit_and_notify.with_subscriber(dispatch),
None => commit_and_notify.with_current_subscriber(),
};
let mut outcomes = match telemetry_context.as_ref() {
Some(context) => Box::pin(context.instrument(commit_and_notify)).await,
None => Box::pin(commit_and_notify).await,
};
for (outcome, checkpoint_gc_sequence) in
outcomes.iter_mut().zip(checkpoint_gc_sequences)
{
if let Ok(outcome) = outcome {
*outcome = TransactionCommitOutcome {
checkpoint_gc_sequence,
..TransactionCommitOutcome::default()
};
}
}
debug_assert_eq!(outcomes.len(), senders.len());
for ((sender, _capacity), outcome) in senders.into_iter().zip(outcomes) {
let _ = sender.send(outcome);
}
}
}
#[cfg(test)]
fn stats(&self) -> (usize, usize, usize) {
(
self.inner.stats.cohort_count.load(Ordering::Relaxed),
self.inner.stats.commit_count.load(Ordering::Relaxed),
self.inner.stats.max_cohort_size.load(Ordering::Relaxed),
)
}
}
struct CommitDriverGuard<'a, StorageImpl>
where
StorageImpl: Storage + 'static,
{
inner: &'a CommitCoordinatorInner<StorageImpl>,
armed: bool,
}
impl<'a, StorageImpl> CommitDriverGuard<'a, StorageImpl>
where
StorageImpl: Storage + 'static,
{
fn new(inner: &'a CommitCoordinatorInner<StorageImpl>) -> Self {
Self { inner, armed: true }
}
fn disarm(&mut self) {
self.armed = false;
}
}
impl<StorageImpl> Drop for CommitDriverGuard<'_, StorageImpl>
where
StorageImpl: Storage + 'static,
{
fn drop(&mut self) {
if !self.armed {
return;
}
let mut state = self
.inner
.state
.lock()
.unwrap_or_else(|error| error.into_inner());
state.running = false;
state.queue.clear();
}
}
fn cohort_telemetry_context<StorageImpl>(
cohort: &[CommitRequest<StorageImpl>],
) -> Option<TelemetryContext>
where
StorageImpl: Storage + 'static,
{
attach_cohort_parent_contexts(cohort.iter().filter_map(|request| request.telemetry_context.clone()))
}
fn attach_cohort_parent_contexts(
contexts: impl IntoIterator<Item = TelemetryContext>,
) -> Option<TelemetryContext> {
let contexts = contexts.into_iter().collect::<Vec<_>>();
let links = contexts
.iter()
.skip(1)
.filter_map(TelemetryContext::as_link)
.collect::<Vec<SpanContext>>();
contexts.into_iter().next().map(|context| {
context
.with_commit_cohort_id(next_commit_cohort_id())
.with_links(links)
})
}
fn coordinator_closed() -> LixError {
LixError::new(
LixError::CODE_INTERNAL_ERROR,
"transaction commit coordinator closed unexpectedly",
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage_adapter::Memory;
use crate::telemetry::{CallbackTelemetrySink, TelemetryContext, new_span_context};
#[test]
fn cohort_context_uses_leader_as_parent_and_links_other_transactions() {
let completed = Mutex::new(Vec::new());
let captured = Arc::new(completed);
let sink: Arc<dyn TelemetrySink> = Arc::new(CallbackTelemetrySink::new({
let captured = Arc::clone(&captured);
move |span| captured.lock().expect("spans").push(span)
}));
let parent_a = new_span_context(None);
let parent_b = new_span_context(None);
let context = attach_cohort_parent_contexts([
TelemetryContext::for_test(Arc::clone(&sink), parent_a.clone()),
TelemetryContext::for_test(Arc::clone(&sink), parent_b.clone()),
])
.expect("cohort context");
futures_lite::future::block_on(TelemetryContext::instrument(
&context,
async {
let span = ActiveTelemetrySpan::start_current(
&TRANSACTION_STORAGE,
vec![TelemetryAttribute::i64("lix.transaction.count", 2)],
)
.expect("storage enabled");
span.finish(Status::Unset, Vec::new());
},
));
let spans = captured.lock().expect("spans");
assert_eq!(spans.len(), 1);
assert_eq!(spans[0].start.name, "lix.transaction.storage");
assert_eq!(
spans[0]
.start
.parent_span_context
.as_ref()
.map(SpanContext::span_id),
Some(parent_a.span_id())
);
assert_eq!(spans[0].start.links, vec![parent_b]);
assert!(
spans[0]
.start
.attributes
.iter()
.any(|attribute| attribute.key == "lix.commit_cohort_id")
);
}
#[test]
fn coordinator_capacity_accepts_realtime_collaboration_wave() {
assert!(COMMIT_COHORT_CAPACITY >= 100);
assert!(COMMIT_QUEUE_CAPACITY >= COMMIT_COHORT_CAPACITY);
let coordinator = CommitCoordinator::<Memory>::new(
Arc::new(tokio::sync::Mutex::new(())),
Arc::new(ObserveInvalidation::new()),
None,
);
assert_eq!(coordinator.stats(), (0, 0, 0));
}
}