diff --git a/ts_dataplane/src/lib.rs b/ts_dataplane/src/lib.rs index 6e972dad..023c237c 100644 --- a/ts_dataplane/src/lib.rs +++ b/ts_dataplane/src/lib.rs @@ -353,23 +353,16 @@ impl DataPlane { } fn ensure_wg(&mut self) { - if let Some(next) = self.wireguard.next_event() - && let Some(prev) = self - .wg_next - .replace(self.events.add(next, Subsystem::Wireguard)) - { - prev.cancel(); - } + self.wg_next = self + .wireguard + .next_event() + .map(|tr| self.events.add(tr, Subsystem::Wireguard)); } fn ensure_peer_gc(&mut self, now: Instant) { - if self.active_peers.is_empty() - && let Some(evt) = self.peer_gc_next.take() - { - evt.cancel(); - } - - if !self.active_peers.is_empty() && self.peer_gc_next.is_none() { + if self.active_peers.is_empty() { + self.peer_gc_next = None; + } else if self.peer_gc_next.is_none() { self.peer_gc_next = Some(self.events.add( TimeRange::new_around(now + Duration::from_secs(10), Duration::from_millis(2500)), Subsystem::PeerGc, diff --git a/ts_time/proptest-regressions/lib.txt b/ts_time/proptest-regressions/lib.txt index d42a2112..5c077c59 100644 --- a/ts_time/proptest-regressions/lib.txt +++ b/ts_time/proptest-regressions/lib.txt @@ -8,3 +8,4 @@ cc 8cea7093c7127919ae8582096c089541566b0801672a136fee39b8e0ee41071b # shrinks to cc 7c7de3f790d4e8aa56234c84f14c689816db54f355a1163f7e50ed1e0a4cecdc # shrinks to actions = [Add((1ms, 1ms)), Reschedule((0, (1ms, 1ms)))] cc 1c877f42a2d9c7e5985ca555bdb0c0a710189d9005f98e2af9ef5c8d4373eae5 # shrinks to actions = [Add((7.888s, 55.678s)), Dispatch, Add((1ms, 1ms)), Dispatch] cc 686a3feb4e14c52f7588160a984c55708dbc887fee1ac5f649a2dc860a87a2d3 # shrinks to times = [(15.845s, 15.845s), (15.845s, 15.845s)] +cc f2da4a01a38904cd67af3bed916c63495b00148052b45befba1bd1f40ff03316 # shrinks to actions = [Dispatch] diff --git a/ts_time/src/lib.rs b/ts_time/src/lib.rs index c3faabd6..aa759aa4 100644 --- a/ts_time/src/lib.rs +++ b/ts_time/src/lib.rs @@ -117,8 +117,8 @@ impl Scheduler { /// Schedule an event to occur at a future point in time. /// - /// Returns a [`Handle`] which may be used to cancel or reschedule the event. The caller need - /// not retain the Handle if cancellation and rescheduling are not required. + /// Returns a [`Handle`] which cancels the event when dropped, unless you [`Handle::forget`] it. + #[must_use] pub fn add(&mut self, when: TimeRange, what: E) -> Handle { let event = Arc::new(FutureEvent { when, what }); let weak_event = Arc::downgrade(&event); @@ -131,6 +131,18 @@ impl Scheduler { } } + /// Schedule an event to occur at a future point in time. + /// + /// Unlike [`Scheduler::add`], the scheduled event cannot be canceled. + /// + /// This is equivalent to `scheduler.add(...).forget()`. + pub fn add_uncancelable(&mut self, when: TimeRange, what: E) { + let event = Arc::new(FutureEvent { when, what }); + let mut events = self.events.lock().unwrap(); + let idx = Scheduler::partition_point(&events, when.start); + events.insert(idx, event); + } + /// Cancel all pending events, leaving the scheduler idle. pub fn clear(&mut self) { self.events.lock().unwrap().clear(); @@ -232,12 +244,21 @@ pub struct Handle { } impl Handle { - /// Attempts to cancel the event. + /// Abandon this handle without canceling the event. /// - /// If the event hasn't yet occurred when cancel is called, it is canceled and will not be - /// returned by [`Scheduler::dispatch`]. Cancelling an event that has already been dispatched - /// is a no-op. - pub fn cancel(self) { + /// If you're always `forget`ing an event immediately after scheduling it, consider + /// [`Scheduler::add_uncancelable`] to make that intent more obvious. + pub fn forget(mut self) { + // Arc's Weak has a handy Default value that can never be upgraded. This makes the Drop + // impl exit early without touching scheduler state (which it can't anyway, since the + // `events` ref is gone). + std::mem::take(&mut self.events); + std::mem::take(&mut self.event); + } +} + +impl Drop for Handle { + fn drop(&mut self) { let Some(events) = self.events.upgrade() else { return; }; @@ -250,33 +271,6 @@ impl Handle { }; events.remove(idx); } - - /// Attempts to reschedule the event to a new time range. - /// - /// Returns an updated Handle if rescheduling succeeds, or None if the event has already - /// been dispatched. - pub fn reschedule(self, when: TimeRange) -> Option> { - let events = self.events.upgrade()?; - let mut events = events.lock().unwrap(); - let mut event = self.event.upgrade()?; - drop(self.event); - let idx = Scheduler::find(&events, &event)?; - drop(events.remove(idx)); - // Invariant: At most 3 refs to the event exist (see doc on SchedulerInner struct). - // We dropped the Handle's Weak and events's Arc above, leaving `event` as the sole Arc - // for this event. Thus, get_mut always succeeds. - Arc::get_mut(&mut event).unwrap().when = when; - let weak = Arc::downgrade(&event); - - let idx = Scheduler::partition_point(&events, when.start); - events.insert(idx, event); - drop(events); - - Some(Handle { - events: self.events, - event: weak, - }) - } } #[cfg(test)] @@ -321,9 +315,36 @@ mod tests { fn test_basic() { let datum = Instant::now(); let mut sched = Scheduler::default(); - sched.add(TimeRange::new(datum, datum), Event::Foo); + + let _handle = sched.add(TimeRange::new(datum, datum), Event::Foo); check_next(&mut sched, TimeRange::new(datum, datum), vec![Event::Foo]); check_empty(&mut sched); + + let handle = sched.add(TimeRange::new(datum, datum), Event::Foo); + handle.forget(); + check_next(&mut sched, TimeRange::new(datum, datum), vec![Event::Foo]); + check_empty(&mut sched); + + sched.add_uncancelable(TimeRange::new(datum, datum), Event::Foo); + check_next(&mut sched, TimeRange::new(datum, datum), vec![Event::Foo]); + check_empty(&mut sched); + } + + #[test] + fn test_cancel() { + let datum = Instant::now(); + let mut sched = Scheduler::default(); + + let handle = sched.add(TimeRange::new(datum, datum), Event::Foo); + drop(handle); + check_empty(&mut sched); + + // Dropping the handle after the event fires is fine, it just does nothing. + let handle = sched.add(TimeRange::new(datum, datum), Event::Foo); + check_next(&mut sched, TimeRange::new(datum, datum), vec![Event::Foo]); + check_empty(&mut sched); + drop(handle); + check_empty(&mut sched); } #[test] @@ -344,7 +365,7 @@ mod tests { let start = datum + Duration::from_secs(*start); let end = datum + Duration::from_secs(*end); let range = TimeRange::new(start, end); - sched.add(range, Event::Bar(i)); + sched.add_uncancelable(range, Event::Bar(i)); } // First wakeup at 4, all events except (5,5). check_next( @@ -402,8 +423,8 @@ mod proptests { /// number of scheduled events so far in the run, so will try to cancel a uniformly sampled /// prior event (which may have already been canceled). Cancel(f64), - /// Reschedule a previously added event. The f64 is rescaled as with Cancel. - Reschedule((f64, (Duration, Duration))), + /// Forget a previously added event. The f64 is rescaled as with Cancel. + Forget(f64), } /// Convert a random 0-1 float value into an index in the range 0..max. @@ -418,7 +439,7 @@ mod proptests { Just(Action::Dispatch), arb_timerange().prop_map(Action::Add), (0f64..1f64).prop_map(Action::Cancel), - ((0f64..1f64), arb_timerange()).prop_map(Action::Reschedule), + (0f64..1f64).prop_map(Action::Forget), ] } @@ -432,11 +453,6 @@ mod proptests { let range = Mutex::new(TimeRange::new(*DATUM + start, *DATUM + end)); Self { id, range } } - - fn update(&mut self, start: Duration, end: Duration) { - let mut range = self.range.lock().unwrap(); - *range = TimeRange::new(*DATUM + start, *DATUM + end); - } } impl Debug for Event { @@ -459,7 +475,7 @@ mod proptests { let mut sched = Scheduler::default(); for (start, end) in × { let tr = TimeRange::new(*DATUM+*start, *DATUM+*end); - sched.add(tr, (start, end, tr)); + sched.add_uncancelable(tr, (start, end, tr)); sched.assert_consistent(); } @@ -495,6 +511,7 @@ mod proptests { let mut now = *DATUM; let mut total_scheduled = 0; let mut total_canceled = 0; + let mut total_forgotten = 0; let mut total_dispatched = 0; println!("\nSTART, now=0s"); for action in actions { @@ -543,28 +560,25 @@ mod proptests { let idx = sample(idx, events.len()); if let Some(handle) = handles[idx].take() { println!("Cancel({})", idx); - handle.cancel(); + drop(handle); events[idx] = None; total_canceled += 1; } else { - println!("Cancel({}) (already canceled)", idx); + println!("Cancel({}) (already canceled or forgotten)", idx); }; } - Action::Reschedule((idx, (start, end))) => { + Action::Forget(idx) => { if events.is_empty() { - println!("Reschedule() (no events yet)"); + println!("Forget() (no events yet)"); continue; } let idx = sample(idx, events.len()); - if let Some(event) = &mut events[idx] { - event.update(start, end); - let tr = { - *event.range.lock().unwrap() - }; - println!("Reschedule({}) event={:?}", idx, event); - handles[idx] = handles[idx].take().and_then(|handle| handle.reschedule(tr)); + if let Some(handle) = handles[idx].take() { + println!("Forget({})", idx); + handle.forget(); + total_forgotten += 1; } else { - println!("Reschedule({}) (no such event)", idx); + println!("Forget({}) (already canceled or forgotten)", idx); } } } @@ -573,6 +587,7 @@ mod proptests { assert_eq!(total_scheduled, events.len()); assert!(total_dispatched <= total_scheduled); assert!(total_canceled <= total_scheduled); + assert!(total_forgotten <= total_scheduled); assert!(total_pending <= total_scheduled); // Cancellations can cause double-counting, when cancelling an already dispatched event. // So, best we can do is bracket the values. diff --git a/ts_tunnel/src/endpoint.rs b/ts_tunnel/src/endpoint.rs index 2f8e1849..cdc3606a 100644 --- a/ts_tunnel/src/endpoint.rs +++ b/ts_tunnel/src/endpoint.rs @@ -250,11 +250,10 @@ impl Peer { } fn shutdown(&mut self) { + // TODO: should this just be a Drop impl instead of an explicit shutdown that leaves a moribund value? self.session.deactivate(); self.handshake.abandon(); - if let Some(handle) = self.keepalive.take() { - handle.cancel(); - } + self.keepalive = None; } #[tracing::instrument(skip_all, fields(now, peer_id = ?self.config.id))] diff --git a/ts_tunnel/src/handshake.rs b/ts_tunnel/src/handshake.rs index 18aec09d..3ffbf1d0 100644 --- a/ts_tunnel/src/handshake.rs +++ b/ts_tunnel/src/handshake.rs @@ -68,6 +68,9 @@ impl ReceivedHandshake { struct SentHandshake { responder_to_initiator_handle: SessionHandle, noise: ikpsk2::SentHandshake, + // timeout is a handle that cancels a retransmit event when the SentHandshake is abandoned, + // but is otherwise unused. + #[allow(dead_code)] timeout: Handle, // The mac1 of the transmitted handshake, required to process CookieReply messages. mac1: Mac, @@ -299,8 +302,6 @@ impl Handshake { now, ); - sent_handshake.timeout.cancel(); - Some(session) } diff --git a/ts_tunnel/src/session.rs b/ts_tunnel/src/session.rs index 395c65b6..cea6f97f 100644 --- a/ts_tunnel/src/session.rs +++ b/ts_tunnel/src/session.rs @@ -407,7 +407,7 @@ impl Session { ); endpoint .scheduler - .add(cleanup, Event::ExpireSession(peer_id)); + .add_uncancelable(cleanup, Event::ExpireSession(peer_id)); } /// Discard all state for this session.