Skip to content

Clean up in-flight dedup entries when the request owner is dropped #6577

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ndr-ds wants to merge 1 commit into
base: main
Choose a base branch
from
+195 −97
Draft

Clean up in-flight dedup entries when the request owner is dropped #6577

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
224 changes: 145 additions & 79 deletions linera-core/src/client/requests_scheduler/in_flight_tracker.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,14 @@
// Copyright (c) Zefchain Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::{collections::HashMap, fmt::Debug, sync::Arc};
use std::{
collections::HashMap,
fmt::Debug,
sync::{
atomic::{AtomicU64, Ordering},
Arc, Mutex,
},
};

use linera_base::{data_types::Timestamp, time::Duration};
use tokio::sync::broadcast;
Expand All @@ -17,10 +24,17 @@ use crate::node::NodeError;
/// This structure manages a map of request keys to in-flight entries, each containing
/// broadcast senders for notifying waiters when a request completes, as well as timing
/// information and alternative data sources.
///
/// The entries map is guarded by a synchronous [`Mutex`] (never held across an `.await`)
/// so that an owner's [`InFlightGuard`] can remove its entry from `Drop`, which is what
/// makes deduplication cancel-safe — see [`InFlightTracker::insert_new`].
#[derive(Debug, Clone)]
pub(super) struct InFlightTracker<N> {
/// Maps request keys to in-flight entries containing broadcast senders and metadata
entries: Arc<tokio::sync::RwLock<HashMap<RequestKey, InFlightEntry<N>>>>,
entries: Arc<Mutex<HashMap<RequestKey, InFlightEntry<N>>>>,
/// Source of per-entry generation numbers, so an owner only removes the entry it
/// created and not a later owner that reused the same key.
next_generation: Arc<AtomicU64>,
/// Maximum duration before an in-flight request is considered stale and deduplication is skipped
timeout: Duration,
}
Expand All @@ -32,7 +46,8 @@ impl<N: Clone> InFlightTracker<N> {
/// - `timeout`: Maximum duration before an in-flight request is considered too old to deduplicate against
pub(super) fn new(timeout: Duration) -> Self {
Self {
entries: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
entries: Arc::new(Mutex::new(HashMap::new())),
next_generation: Arc::new(AtomicU64::new(0)),
timeout,
}
}
Expand All @@ -49,12 +64,11 @@ impl<N: Clone> InFlightTracker<N> {
/// # Returns
/// - `None`: No matching in-flight request found. Also returned if the found request is stale (exceeds timeout).
/// - `Some(InFlightMatch::Subsuming { key, outcome })`: Subsuming request found
pub(super) async fn try_subscribe(
&self,
key: &RequestKey,
now: Timestamp,
) -> Option<InFlightMatch> {
let in_flight = self.entries.read().await;
pub(super) fn try_subscribe(&self, key: &RequestKey, now: Timestamp) -> Option<InFlightMatch> {
let in_flight = self
.entries
.lock()
.expect("in-flight tracker mutex poisoned");

if let Some(entry) = in_flight.get(key) {
let elapsed = now.duration_since(entry.started_at);
Expand All @@ -81,61 +95,39 @@ impl<N: Clone> InFlightTracker<N> {
None
}

/// Inserts a new in-flight request entry.
/// Inserts a new in-flight request entry and returns the owner guard.
///
/// Creates a new broadcast channel and in-flight entry for the given key,
/// marking the start time as now.
/// The returned [`InFlightGuard`] owns the entry. Completing the request via
/// [`InFlightGuard::complete_and_broadcast`] broadcasts the result to all
/// subscribers and removes the entry. If the guard is instead dropped without
/// completing — for example because the owning task was cancelled, such as a
/// losing branch of `communicate_with_quorum` — its `Drop` removes the entry,
/// dropping the broadcast sender so subscribers wake with `RecvError::Closed`
/// and execute the request themselves rather than waiting forever.
///
/// # Arguments
/// - `key`: The request key to insert
pub(super) async fn insert_new(&self, key: RequestKey, now: Timestamp) {
/// - `now`: The current time, recorded as the entry's start time
pub(super) fn insert_new(&self, key: RequestKey, now: Timestamp) -> InFlightGuard<N> {
let (sender, _receiver) = broadcast::channel(1);
let mut in_flight = self.entries.write().await;

in_flight.insert(
key,
InFlightEntry {
sender,
started_at: now,
alternative_peers: Arc::new(tokio::sync::RwLock::new(Vec::new())),
},
);
}

/// Completes an in-flight request by removing it and broadcasting the result.
///
/// Removes the entry for the given key and broadcasts the result to all waiting
/// subscribers. Logs the number of waiters that received the notification.
///
/// # Arguments
/// - `key`: The request key to complete
/// - `result`: The result to broadcast to waiters
pub(super) async fn complete_and_broadcast(
&self,
key: &RequestKey,
result: Arc<Result<RequestResult, NodeError>>,
) -> usize {
let mut in_flight = self.entries.write().await;

if let Some(entry) = in_flight.remove(key) {
let waiter_count = entry.sender.receiver_count();
tracing::trace!(
key = ?key,
waiters = waiter_count,
"request completed; broadcasting result to waiters",
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
self.entries
.lock()
.expect("in-flight tracker mutex poisoned")
.insert(
key.clone(),
InFlightEntry {
sender,
started_at: now,
generation,
alternative_peers: Arc::new(tokio::sync::RwLock::new(Vec::new())),
},
);
if waiter_count != 0 {
if let Err(err) = entry.sender.send(result) {
tracing::warn!(
key = ?key,
error = ?err,
"failed to broadcast result to waiters"
);
}
}
return waiter_count;
InFlightGuard {
entries: self.entries.clone(),
key,
generation,
}
0
}

/// Registers an alternative peer for an in-flight request.
Expand All @@ -150,14 +142,12 @@ impl<N: Clone> InFlightTracker<N> {
where
N: PartialEq + Eq,
{
if let Some(entry) = self.entries.read().await.get(key) {
// Register this peer as an alternative source if not already present
{
let mut alt_peers = entry.alternative_peers.write().await;
if !alt_peers.contains(&peer) {
alt_peers.push(peer);
}
}
let Some(alternative_peers) = self.alternative_peers(key) else {
return;
};
let mut alt_peers = alternative_peers.write().await;
if !alt_peers.contains(&peer) {
alt_peers.push(peer);
}
}

Expand All @@ -171,10 +161,8 @@ impl<N: Clone> InFlightTracker<N> {
/// # Returns
/// - `Vec<N>`: List of alternative peers (empty if no entry exists)
pub(super) async fn get_alternative_peers(&self, key: &RequestKey) -> Option<Vec<N>> {
let in_flight = self.entries.read().await;

let entry = in_flight.get(key)?;
let peers = entry.alternative_peers.read().await;
let alternative_peers = self.alternative_peers(key)?;
let peers = alternative_peers.read().await;
Some(peers.clone())
}

Expand All @@ -187,10 +175,10 @@ impl<N: Clone> InFlightTracker<N> {
where
N: PartialEq + Eq,
{
if let Some(entry) = self.entries.read().await.get(key) {
let mut alt_peers = entry.alternative_peers.write().await;
alt_peers.retain(|p| p != peer);
}
let Some(alternative_peers) = self.alternative_peers(key) else {
return;
};
alternative_peers.write().await.retain(|p| p != peer);
}

/// Pops and returns the newest alternative peer from the list.
Expand All @@ -205,12 +193,19 @@ impl<N: Clone> InFlightTracker<N> {
/// - `Some(N)`: The newest alternative peer
/// - `None`: No entry exists or alternatives list is empty
pub(super) async fn pop_alternative_peer(&self, key: &RequestKey) -> Option<N> {
if let Some(entry) = self.entries.read().await.get(key) {
let mut alt_peers = entry.alternative_peers.write().await;
alt_peers.pop()
} else {
None
}
self.alternative_peers(key)?.write().await.pop()
}

/// Returns the shared alternative-peers list for an entry, if it exists.
///
/// The entries lock is released before the caller awaits the returned lock, so the
/// synchronous [`Mutex`] is never held across an `.await`.
fn alternative_peers(&self, key: &RequestKey) -> Option<Arc<tokio::sync::RwLock<Vec<N>>>> {
self.entries
.lock()
.expect("in-flight tracker mutex poisoned")
.get(key)
.map(|entry| entry.alternative_peers.clone())
}
}

Expand All @@ -233,13 +228,84 @@ pub(super) enum InFlightMatch {
#[derive(Debug)]
pub(super) struct Subscribed(pub(super) broadcast::Receiver<Arc<Result<RequestResult, NodeError>>>);

/// Owns an in-flight entry for the lifetime of the request that created it.
///
/// Completing via [`InFlightGuard::complete_and_broadcast`] is the success path. If the
/// guard is dropped first (the owning task was cancelled), `Drop` removes the entry so
/// subscribers stop waiting. Both paths are generation-checked, so a guard never touches
/// an entry a later owner created after this one went stale and was replaced.
pub(super) struct InFlightGuard<N> {
entries: Arc<Mutex<HashMap<RequestKey, InFlightEntry<N>>>>,
key: RequestKey,
generation: u64,
}

impl<N> InFlightGuard<N> {
/// Removes the owned entry and broadcasts the result to all waiting subscribers.
///
/// Returns the number of subscribers that were notified. Does nothing (returns `0`)
/// if a later owner has already replaced this entry, in which case this result is
/// stale and must not be broadcast. After this call the guard's `Drop` is a no-op.
pub(super) fn complete_and_broadcast(
&self,
result: Arc<Result<RequestResult, NodeError>>,
) -> usize {
let mut in_flight = self
.entries
.lock()
.expect("in-flight tracker mutex poisoned");
if in_flight
.get(&self.key)
.is_none_or(|entry| entry.generation != self.generation)
{
return 0;
}
let entry = in_flight.remove(&self.key).expect("checked above");
let waiter_count = entry.sender.receiver_count();
tracing::trace!(
key = ?self.key,
waiters = waiter_count,
"request completed; broadcasting result to waiters",
);
if waiter_count != 0 {
if let Err(err) = entry.sender.send(result) {
tracing::warn!(
key = ?self.key,
error = ?err,
"failed to broadcast result to waiters"
);
}
}
waiter_count
}
}

impl<N> Drop for InFlightGuard<N> {
fn drop(&mut self) {
let Ok(mut in_flight) = self.entries.lock() else {
return; // Poisoned mutex: nothing safe to do from `Drop`.
};
// Remove only if our entry is still the current one for this key. Dropping it
// drops the broadcast sender, waking subscribers with `Closed` so they execute
// the request themselves instead of waiting for a result that will never arrive.
if in_flight
.get(&self.key)
.is_some_and(|entry| entry.generation == self.generation)
{
in_flight.remove(&self.key);
}
}
}

/// In-flight request entry that tracks when the request was initiated.
#[derive(Debug)]
pub(super) struct InFlightEntry<N> {
/// Broadcast sender for notifying waiters when the request completes
sender: broadcast::Sender<Arc<Result<RequestResult, NodeError>>>,
/// Time when this request was initiated
started_at: Timestamp,
/// Generation of this entry, used by the owner guard to remove only its own entry.
generation: u64,
/// Alternative peers that can provide this data if the primary request fails
alternative_peers: Arc<tokio::sync::RwLock<Vec<N>>>,
}
Loading
Loading