diff --git a/orange-sdk/src/event.rs b/orange-sdk/src/event.rs index c00f053..6fa2610 100644 --- a/orange-sdk/src/event.rs +++ b/orange-sdk/src/event.rs @@ -17,9 +17,8 @@ use ldk_node::{CustomTlvRecord, UserChannelId}; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; -use std::task::{Poll, Waker}; use std::time::SystemTime; -use tokio::sync::{Mutex, watch}; +use tokio::sync::{Mutex, Notify, watch}; /// The event queue will be persisted under this key. pub(crate) const EVENT_QUEUE_PERSISTENCE_PRIMARY_NAMESPACE: &str = ""; @@ -210,7 +209,11 @@ impl_writeable_tlv_based_enum!(Event, pub struct EventQueue { queue: Arc>>, pending_mpp_events: Arc>>>, - waker: Arc>>, + /// Signals a waiting [`EventQueue::next_event_async`] that the queue changed. A push or an + /// acknowledgement calls [`Notify::notify_one`] after releasing the queue lock; the waiter + /// registers before it checks the queue, so no change can slip between its check and its + /// wait. + notify: Arc, kv_store: Arc, rebalance_enabled: RebalanceEnabledCache, tx_metadata: TxMetadataStore, @@ -223,9 +226,9 @@ impl EventQueue { ) -> Self { let queue = Arc::new(Mutex::new(VecDeque::new())); let pending_mpp_events = Arc::new(Mutex::new(HashMap::new())); - let waker = Arc::new(Mutex::new(None)); + let notify = Arc::new(Notify::new()); let rebalance_enabled = RebalanceEnabledCache::new(Arc::clone(&kv_store)); - Self { queue, pending_mpp_events, waker, kv_store, rebalance_enabled, tx_metadata, logger } + Self { queue, pending_mpp_events, notify, kv_store, rebalance_enabled, tx_metadata, logger } } pub(crate) async fn get_rebalance_enabled(&self) -> bool { @@ -333,9 +336,7 @@ impl EventQueue { self.persist_queue(&locked_queue).await?; } - if let Some(waker) = self.waker.lock().await.take() { - waker.wake(); - } + self.notify.notify_one(); Ok(()) } @@ -344,8 +345,25 @@ impl EventQueue { locked_queue.front().cloned() } + /// Waits until the queue has an event and returns it (without removing it). + /// + /// The wait is registered before the queue is checked, so a push that lands between the + /// check and the wait is not lost: [`EventQueue::push_event`] and + /// [`EventQueue::event_handled`] call [`Notify::notify_one`] after releasing the queue lock, + /// and a `notify_one` with no registered waiter leaves a permit that completes the next wait + /// at once. A wait that is dropped after being notified (the caller raced it against a + /// shutdown, say) passes the notification on to the next waiter, so cancelling this future + /// cannot swallow an event either. Waking spuriously is harmless: the loop checks the queue + /// again. pub(crate) async fn next_event_async(&self) -> Event { - EventFuture { event_queue: Arc::clone(&self.queue), waker: Arc::clone(&self.waker) }.await + loop { + let mut notified = std::pin::pin!(self.notify.notified()); + notified.as_mut().enable(); + if let Some(event) = self.queue.lock().await.front().cloned() { + return event; + } + notified.await; + } } pub(crate) async fn event_handled(&self) -> Result<(), ldk_node::lightning::io::Error> { @@ -355,9 +373,7 @@ impl EventQueue { self.persist_queue(&locked_queue).await?; } - if let Some(waker) = self.waker.lock().await.take() { - waker.wake(); - } + self.notify.notify_one(); Ok(()) } @@ -402,7 +418,12 @@ mod tests { use crate::logging::LoggerType; use crate::store::{PaymentId, PaymentType, TxMetadata, TxMetadataStore, TxType}; use ldk_node::io::sqlite_store::SqliteStore; + use ldk_node::lightning::io; + use ldk_node::lightning::util::persist::{PageToken, PaginatedListResponse}; + use std::future::Future; use std::path::PathBuf; + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::{Duration, UNIX_EPOCH}; fn temp_sqlite_store() -> (PathBuf, Arc) { @@ -484,6 +505,167 @@ mod tests { }) ); } + + fn failed_payment(tag: u8) -> Event { + Event::PaymentFailed { + payment_id: PaymentId::SelfCustodial([tag; 32]), + payment_hash: Some(PaymentHash([tag; 32])), + reason: None, + } + } + + /// A store whose writes wait for the test to open a gate, so a push can be held inside + /// `persist_queue` with the queue lock taken, where a slow remote store (VSS over a mobile + /// link) holds it for hundreds of milliseconds. + struct GatedStore { + inner: Arc, + gate: watch::Receiver, + } + + impl DynStore for GatedStore { + fn read_async( + &self, p: &str, s: &str, k: &str, + ) -> Pin, io::Error>> + Send + 'static>> { + self.inner.read_async(p, s, k) + } + + fn write_async( + &self, p: &str, s: &str, k: &str, buf: Vec, + ) -> Pin> + Send + 'static>> { + let inner = Arc::clone(&self.inner); + let mut gate = self.gate.clone(); + let (p, s, k) = (p.to_owned(), s.to_owned(), k.to_owned()); + Box::pin(async move { + gate.wait_for(|open| *open) + .await + .map_err(|_| io::Error::new(io::ErrorKind::Other, "gate dropped"))?; + inner.write_async(&p, &s, &k, buf).await + }) + } + + fn remove_async( + &self, p: &str, s: &str, k: &str, lazy: bool, + ) -> Pin> + Send + 'static>> { + self.inner.remove_async(p, s, k, lazy) + } + + fn list_async( + &self, p: &str, s: &str, + ) -> Pin, io::Error>> + Send + 'static>> { + self.inner.list_async(p, s) + } + + fn list_paginated_async( + &self, p: &str, s: &str, page_token: Option, + ) -> Pin> + Send + 'static>> + { + self.inner.list_paginated_async(p, s, page_token) + } + } + + /// Stands in for the thread `Wallet::wait_next_event` parks in `block_on`: a wake-up polls + /// the consumer again at once, so the poll lands wherever the wake-up came from. + struct PollOnWake { + consumer: std::sync::Mutex + Send>>>>, + delivered: std::sync::Mutex>, + polls: AtomicUsize, + reentrant_wakes: AtomicUsize, + } + + impl PollOnWake { + fn new(consumer: Pin + Send>>) -> Self { + Self { + consumer: std::sync::Mutex::new(Some(consumer)), + delivered: std::sync::Mutex::new(None), + polls: AtomicUsize::new(0), + reentrant_wakes: AtomicUsize::new(0), + } + } + + fn poll_consumer(self: &Arc) { + let Ok(mut slot) = self.consumer.try_lock() else { + self.reentrant_wakes.fetch_add(1, Ordering::SeqCst); + return; + }; + let Some(consumer) = slot.as_mut() else { return }; + self.polls.fetch_add(1, Ordering::SeqCst); + let waker = std::task::Waker::from(Arc::clone(self)); + let mut cx = std::task::Context::from_waker(&waker); + if let std::task::Poll::Ready(event) = consumer.as_mut().poll(&mut cx) { + *self.delivered.lock().unwrap() = Some(event); + *slot = None; + } + } + + fn delivered(&self) -> Option { + self.delivered.lock().unwrap().clone() + } + } + + impl std::task::Wake for PollOnWake { + fn wake(self: Arc) { + self.poll_consumer(); + } + + fn wake_by_ref(self: &Arc) { + self.poll_consumer(); + } + } + + /// The consumer is woken while one pusher still holds its post-persist section and a second + /// pusher holds the queue inside its own persist. Its poll must leave it registered for the + /// next change. Before the `Notify` rewrite, `EventFuture::poll` read both `try_lock` + /// failures as "nothing to do" and returned `Pending` without a waker, so the second push, + /// and every push after it, woke nobody: the wallet kept queueing events that no + /// `wait_next_event` would ever return. + #[tokio::test] + async fn a_poll_that_lands_on_two_busy_pushers_is_still_woken() { + let (_path, sqlite) = temp_sqlite_store(); + let (open_gate, gate) = watch::channel(false); + let store: Arc = Arc::new(GatedStore { inner: Arc::clone(&sqlite), gate }); + let logger = Arc::new(Logger::new(&LoggerType::LogFacade).expect("logger")); + let tx_metadata = TxMetadataStore::new(sqlite).await; + let queue = Arc::new(EventQueue::new(store, tx_metadata, logger)); + + // The consumer waits on an empty queue. + let consumer = Arc::new(PollOnWake::new({ + let queue = Arc::clone(&queue); + Box::pin(async move { queue.next_event_async().await }) + })); + consumer.poll_consumer(); + assert_eq!(consumer.delivered(), None); + + // Pusher A takes the queue lock and parks in its persist; pusher B queues up behind it. + let a = tokio::spawn({ + let queue = Arc::clone(&queue); + async move { queue.add_event(failed_payment(1)).await } + }); + for _ in 0..8 { + tokio::task::yield_now().await; + } + let b = tokio::spawn({ + let queue = Arc::clone(&queue); + async move { queue.add_event(failed_payment(2)).await } + }); + for _ in 0..8 { + tokio::task::yield_now().await; + } + assert!(!a.is_finished() && !b.is_finished()); + + // A's persist completes: the queue lock passes to B, and A wakes the consumer, whose poll + // lands while B holds the queue. Then B's persist completes. + open_gate.send(true).expect("gate"); + a.await.expect("join").expect("push A"); + b.await.expect("join").expect("push B"); + + assert_eq!(consumer.reentrant_wakes.load(Ordering::SeqCst), 0); + assert_eq!( + consumer.delivered(), + Some(failed_payment(1)), + "polled {} time(s), never handed the event", + consumer.polls.load(Ordering::SeqCst) + ); + } } struct EventQueueSerWrapper<'a>(&'a VecDeque); @@ -498,28 +680,6 @@ impl Writeable for EventQueueSerWrapper<'_> { } } -struct EventFuture { - event_queue: Arc>>, - waker: Arc>>, -} - -impl Future for EventFuture { - type Output = Event; - - fn poll( - self: core::pin::Pin<&mut Self>, cx: &mut core::task::Context<'_>, - ) -> Poll { - if let Some(event) = self.event_queue.try_lock().ok().and_then(|q| q.front().cloned()) { - Poll::Ready(event) - } else { - if let Ok(mut waker) = self.waker.try_lock() { - *waker = Some(cx.waker().clone()); - } - Poll::Pending - } - } -} - #[derive(Clone)] pub(crate) struct LdkEventHandler { pub(crate) event_queue: Arc,