f further test

Author: TheBlueMatt

lightning/src/util/wakers.rs

@@ -586,4 +586,42 @@ mod tests {
 		notifier_a.notify();
 		Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
 	}
+
+	#[test]
+	fn sleeper_with_pending_callbacks() {
+		// This is similar to the above `test_multi_future_sleep` test, but in addition registers
+		// "normal" callbacks which will cause the futures to assume notification has occurred,
+		// rather than waiting for a woken sleeper.
+		let notifier_a = Notifier::new();
+		let notifier_b = Notifier::new();
+
+		// Set both notifiers as woken without sleeping yet.
+		notifier_a.notify();
+		notifier_b.notify();
+
+		// After sleeping one future (not guaranteed which one, however) will have its notification
+		// bit cleared.
+		Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
+
+		// By registering a callback on the futures for both notifiers, one will complete
+		// immediately, but one will remain tied to the notifier, and will complete once the
+		// notifier is next woken, which will be considered the completion of the notification.
+		let callback_a = Arc::new(AtomicBool::new(false));
+		let callback_b = Arc::new(AtomicBool::new(false));
+		let callback_a_ref = Arc::clone(&callback_a);
+		let callback_b_ref = Arc::clone(&callback_b);
+		notifier_a.get_future().register_callback(Box::new(move || assert!(!callback_a_ref.fetch_or(true, Ordering::SeqCst))));
+		notifier_b.get_future().register_callback(Box::new(move || assert!(!callback_b_ref.fetch_or(true, Ordering::SeqCst))));
+		assert!(callback_a.load(Ordering::SeqCst) ^ callback_b.load(Ordering::SeqCst));
+
+		// If we now notify both notifiers again, the other callback will fire, completing the
+		// notification, and we'll be back to one pending notification.
+		notifier_a.notify();
+		notifier_b.notify();
+
+		assert!(callback_a.load(Ordering::SeqCst) && callback_b.load(Ordering::SeqCst));
+		Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future()).wait();
+		assert!(!Sleeper::from_two_futures(notifier_a.get_future(), notifier_b.get_future())
+			.wait_timeout(Duration::from_millis(10)));
+	}
 }