Skip to content

Commit c86fd62

Browse files
committed
[vqueues] Fix a bug with a lost wakeup on inbox event
The vqueues DRR scheduler wasn't waking up the waker whenever an inbox event happens (e.g. an entry is enqueued). It seems like it was there before, and got removed as part of [this commit](4f22aa677816#diff-02cc9fd99c15c55730f0822a024e9d2c6bd7f4f92d3cc2d8c9d030773f1d6f6d). This is currently not causing any issues because we're constructing a fresh scheduler stream on every `LeaderState::run`, so wakers are not that useful there. However, with the scheduler of the next PR, it relies on the wakers to figure out which streams to poll, and this bug clearly manifests there. For simplicity, I'm callin the waker with every `on_inbox_event` call, though admittedly, some events won't result into different scheduler decisions. The original commit had a more invovled `should_wake` propagation, happy to copy it if deemed necessary. Note: This is a bug that codex found as part of my testing
1 parent 1cd3fcb commit c86fd62

1 file changed

Lines changed: 38 additions & 5 deletions

File tree

  • crates/vqueues/src/scheduler

crates/vqueues/src/scheduler/drr.rs

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,9 @@ impl<S: VQueueStore> DRRScheduler<S> {
354354
}
355355
}
356356
}
357+
358+
// The event may have made a queue eligible after the scheduler returned Pending.
359+
self.waker.wake_by_ref();
357360
}
358361

359362
fn release_lock(&mut self, scope: &Option<Scope>, lock_name: &LockName) {
@@ -423,7 +426,9 @@ impl<S: VQueueStore> DRRScheduler<S> {
423426
#[cfg(test)]
424427
mod tests {
425428
use std::num::{NonZeroU16, NonZeroU32, NonZeroUsize};
426-
use std::task::Poll;
429+
use std::sync::Arc;
430+
use std::sync::atomic::{AtomicBool, Ordering};
431+
use std::task::{Context, Poll, Wake, Waker};
427432

428433
use restate_clock::RoughTimestamp;
429434
use restate_clock::time::MillisSinceEpoch;
@@ -459,6 +464,14 @@ mod tests {
459464

460465
const BASE_RUN_AT_MS: u64 = 1_744_000_000_000;
461466

467+
struct TestWaker(Arc<AtomicBool>);
468+
469+
impl Wake for TestWaker {
470+
fn wake(self: Arc<Self>) {
471+
self.0.store(true, Ordering::Relaxed);
472+
}
473+
}
474+
462475
fn test_qid(partition_key: u64) -> VQueueId {
463476
VQueueId::custom(partition_key, "1")
464477
}
@@ -716,18 +729,38 @@ mod tests {
716729
}
717730

718731
#[restate_core::test]
719-
async fn empty_scheduler_returns_pending() {
720-
let rocksdb = storage_test_environment().await;
732+
async fn inbox_event_wakes_pending_scheduler() {
733+
let mut rocksdb = storage_test_environment().await;
721734
let db = rocksdb.partition_db();
722-
let cache = VQueuesMetaCache::create(db.clone(), TEST_VQUEUES_CAPACITY)
735+
let mut cache = VQueuesMetaCache::create(db.clone(), TEST_VQUEUES_CAPACITY)
723736
.await
724737
.unwrap();
725738

726739
let mut scheduler = create_scheduler(db, &cache).await;
740+
let was_woken = Arc::new(AtomicBool::new(false));
741+
let waker = Waker::from(Arc::new(TestWaker(Arc::clone(&was_woken))));
742+
let mut cx = Context::from_waker(&waker);
727743
assert!(matches!(
728-
poll_scheduler(Pin::new(&mut scheduler), cache.view()),
744+
Pin::new(&mut scheduler).poll_schedule_next(cache.view(), &mut cx),
729745
Poll::Pending
730746
));
747+
748+
let qid = test_qid(1);
749+
let mut events = Vec::new();
750+
let mut txn = rocksdb.transaction();
751+
enqueue_entry(&mut txn, &mut cache, &qid, 1, 0, Some(&mut events)).await;
752+
txn.commit().await.expect("commit should succeed");
753+
drop(txn);
754+
755+
for event in events {
756+
scheduler.on_inbox_event(cache.view(), event);
757+
}
758+
759+
assert!(was_woken.load(Ordering::Relaxed));
760+
assert!(matches!(
761+
poll_scheduler(Pin::new(&mut scheduler), cache.view()),
762+
Poll::Ready(Ok(_))
763+
));
731764
}
732765

733766
#[restate_core::test]

0 commit comments

Comments
 (0)