kvbm_physical/transfer/notifications/
notification.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Transfer completion notification handle.
5
6use anyhow::Result;
7use futures::future::{Either, Ready, ready};
8use std::{
9    pin::Pin,
10    sync::Arc,
11    task::{Context, Poll},
12};
13use velo::{Event, EventAwaiter, EventManager};
14
15pub enum TransferAwaiter {
16    Local(EventAwaiter),
17    // Sync(SyncResult),
18}
19
20impl std::future::Future for TransferAwaiter {
21    type Output = Result<()>;
22
23    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
24        match self.get_mut() {
25            Self::Local(waiter) => Pin::new(waiter).poll(cx),
26            // Self::Sync(sync) => Pin::new(sync).poll(cx),
27        }
28    }
29}
30
31/// Notification handle for an in-progress transfer.
32///
33/// This object can be awaited to block until the transfer completes.
34/// The transfer is tracked by a background handler that polls for completion
35/// or processes notification events.
36///
37/// Uses `futures::Either` to avoid event system overhead for synchronous completions.
38/// Pending transfers use `LocalEventWaiter` which avoids heap allocation and repeated
39/// DashMap lookups when awaiting.
40pub struct TransferCompleteNotification {
41    awaiter: Either<Ready<Result<()>>, TransferAwaiter>,
42}
43
44impl TransferCompleteNotification {
45    /// Create a notification that is already completed (for synchronous transfers).
46    ///
47    /// This is useful for transfers that complete immediately without needing
48    /// background polling, such as memcpy operations.
49    ///
50    /// This is extremely efficient - no allocations, locks, or event system overhead.
51    pub fn completed() -> Self {
52        Self {
53            awaiter: Either::Left(ready(Ok(()))),
54        }
55    }
56
57    /// Create a notification from a `LocalEventWaiter`.
58    ///
59    /// This is the primary way to construct a notification when you already
60    /// have an event waiter from the event system.
61    pub fn from_awaiter(awaiter: EventAwaiter) -> Self {
62        Self {
63            awaiter: Either::Right(TransferAwaiter::Local(awaiter)),
64        }
65    }
66
67    // /// Create a notification from a synchronous active message result.
68    // pub fn from_sync_result(sync: SyncResult) -> Self {
69    //     Self {
70    //         awaiter: Either::Right(TransferAwaiter::Sync(sync)),
71    //     }
72    // }
73
74    /// Check if the notification can yield the current task.
75    ///
76    /// The internal ::Left arm is guaranteed to be ready, while the ::Right arm is not.
77    pub fn could_yield(&self) -> bool {
78        matches!(self.awaiter, Either::Right(_))
79    }
80
81    /// Aggregate multiple notifications into one that completes when all are done.
82    ///
83    /// This is useful when a transfer is split across multiple workers and you want
84    /// to wait for all of them to complete.
85    ///
86    /// # Arguments
87    /// * `notifications` - The notifications to aggregate
88    /// * `events` - The event system to create the aggregate event
89    /// * `runtime` - The tokio runtime handle to spawn the aggregation task
90    ///
91    /// # Behavior
92    /// - If the list is empty, returns an already-completed notification
93    /// - If there's only one, returns it directly
94    /// - Otherwise, creates a new event and spawns a task to await all notifications
95    pub fn aggregate(
96        notifications: Vec<Self>,
97        events: &Arc<EventManager>,
98        runtime: &tokio::runtime::Handle,
99    ) -> Result<Self> {
100        if notifications.is_empty() {
101            return Ok(Self::completed());
102        }
103        if notifications.len() == 1 {
104            return Ok(notifications.into_iter().next().unwrap());
105        }
106
107        // Check if all notifications are already complete (no yielding needed)
108        if notifications.iter().all(|n| !n.could_yield()) {
109            return Ok(Self::completed());
110        }
111
112        // Create a new event for the aggregate completion
113        let event = events.new_event()?;
114        let awaiter = events.awaiter(event.handle())?;
115
116        // Spawn task that awaits all notifications and triggers/poisons the event
117        runtime.spawn(await_all_notifications(notifications, event));
118
119        Ok(Self::from_awaiter(awaiter))
120    }
121}
122
123/// Awaits all transfer notifications and signals completion via the event.
124///
125/// This function awaits ALL notifications regardless of individual failures,
126/// then triggers the event on success or poisons it with error details on failure.
127async fn await_all_notifications(
128    notifications: Vec<TransferCompleteNotification>,
129    local_event: Event,
130) {
131    // Await all notifications, collecting results
132    let results: Vec<Result<()>> =
133        futures::future::join_all(notifications.into_iter().map(|n| n.into_future())).await;
134
135    // Check for any failures
136    let errors: Vec<_> = results.into_iter().filter_map(|r| r.err()).collect();
137
138    if errors.is_empty() {
139        // Ignore trigger error - if event system is shutdown, nothing to do
140        let _ = local_event.trigger();
141    } else {
142        let error_msg = errors
143            .iter()
144            .map(|e| e.to_string())
145            .collect::<Vec<_>>()
146            .join("; ");
147        // Ignore poison error - if event system is shutdown, nothing to do
148        let _ = local_event.poison(error_msg);
149    }
150}
151
152impl std::future::IntoFuture for TransferCompleteNotification {
153    type Output = Result<()>;
154    type IntoFuture = Either<Ready<Result<()>>, TransferAwaiter>;
155
156    fn into_future(self) -> Self::IntoFuture {
157        self.awaiter
158    }
159}
kvbm_physical/transfer/notifications/notification.rs

kvbm_physical/transfer/notifications/
notification.rs
