diff --git a/seastar/src/clocks.rs b/seastar/src/clocks.rs index 1758b14..51c2b55 100644 --- a/seastar/src/clocks.rs +++ b/seastar/src/clocks.rs @@ -1,4 +1,5 @@ use crate::cxx_async_futures::VoidFuture; +use crate::SchedulingGroup; use core::cmp::Ordering; use cxx::UniquePtr; use std::fmt; @@ -37,6 +38,13 @@ mod ffi { fn manual_sleep(nanos: i64) -> VoidFuture; } + #[namespace = "seastar_ffi::scheduling"] + unsafe extern "C++" { + include!("seastar/src/scheduling.hh"); + + type scheduling_group = crate::scheduling_group; + } + #[namespace = "seastar_ffi::timer::steady_clock"] unsafe extern "C++" { include!("seastar/src/timer.hh"); @@ -52,6 +60,14 @@ mod ffi { dropper: unsafe fn(*mut u8), ); + unsafe fn sct_set_callback_under_group( + timer: Pin<&mut steady_clock_timer>, + callback: *mut u8, + caller: unsafe fn(*mut u8), + dropper: unsafe fn(*mut u8), + sg: &scheduling_group, + ); + fn sct_arm_at(timer: Pin<&mut steady_clock_timer>, at: i64); fn sct_arm_at_periodic(timer: Pin<&mut steady_clock_timer>, at: i64, period: i64); @@ -80,6 +96,14 @@ mod ffi { dropper: unsafe fn(*mut u8), ); + unsafe fn lct_set_callback_under_group( + timer: Pin<&mut lowres_clock_timer>, + callback: *mut u8, + caller: unsafe fn(*mut u8), + dropper: unsafe fn(*mut u8), + sg: &scheduling_group, + ); + fn lct_arm_at(timer: Pin<&mut lowres_clock_timer>, at: i64); fn lct_arm_at_periodic(timer: Pin<&mut lowres_clock_timer>, at: i64, period: i64); @@ -108,6 +132,14 @@ mod ffi { dropper: unsafe fn(*mut u8), ); + unsafe fn mct_set_callback_under_group( + timer: Pin<&mut manual_clock_timer>, + callback: *mut u8, + caller: unsafe fn(*mut u8), + dropper: unsafe fn(*mut u8), + sg: &scheduling_group, + ); + fn mct_arm_at(timer: Pin<&mut manual_clock_timer>, at: i64); fn mct_arm_at_periodic(timer: Pin<&mut manual_clock_timer>, at: i64, period: i64); @@ -574,6 +606,14 @@ mod clock_implementation { dropper: fn(*mut u8), ); + fn set_callback_under_group( + cpp_timer: &mut CppTimer, + callback: *mut u8, + caller: fn(*mut u8), + dropper: fn(*mut u8), + sg: &SchedulingGroup, + ); + fn arm_at(cpp_timer: &mut CppTimer, at: i64); fn arm_at_periodic(cpp_timer: &mut CppTimer, at: i64, period: i64); @@ -620,6 +660,27 @@ macro_rules! timer_impl { }; } + fn set_callback_under_group( + cpp_timer: &mut CppTimer, + callback: *mut u8, + caller: fn(*mut u8), + dropper: fn(*mut u8), + sg: &SchedulingGroup, + ) { + match cpp_timer { + CppTimer::$cpp_timer_variant(timer) => unsafe { + [<$ffi_pref _set_callback_under_group>]( + timer.pin_mut(), + callback, + caller, + dropper, + &sg.inner, + ); + }, + _ => panic!(), + }; + } + fn arm_at(cpp_timer: &mut CppTimer, at: i64) { match cpp_timer { CppTimer::$cpp_timer_variant(timer) => { diff --git a/seastar/src/scheduling.rs b/seastar/src/scheduling.rs index ebecc2b..0d8b1dd 100644 --- a/seastar/src/scheduling.rs +++ b/seastar/src/scheduling.rs @@ -35,6 +35,7 @@ mod ffi { } } +pub(crate) use ffi::scheduling_group; use ffi::*; /// Identifies function calls that are accounted as a group. diff --git a/seastar/src/timer.cc b/seastar/src/timer.cc index 30bc194..bef8f68 100644 --- a/seastar/src/timer.cc +++ b/seastar/src/timer.cc @@ -75,6 +75,16 @@ void sct_set_callback( timer.set_callback(callback_object(callback, caller, dropper)); } +void sct_set_callback_under_group( + steady_clock_timer& timer, + uint8_t* callback, + rust::Fn caller, + rust::Fn dropper, + const scheduling_group& sg +) { + timer.set_callback(sg, callback_object(callback, caller, dropper)); +} + void sct_arm_at(steady_clock_timer& timer, int64_t at) { timer.arm(to_sc_time_point(at)); } @@ -124,6 +134,16 @@ void lct_set_callback( timer.set_callback(callback_object(callback, caller, dropper)); } +void lct_set_callback_under_group( + lowres_clock_timer& timer, + uint8_t* callback, + rust::Fn caller, + rust::Fn dropper, + const scheduling_group& sg +) { + timer.set_callback(sg, callback_object(callback, caller, dropper)); +} + void lct_arm_at(lowres_clock_timer& timer, int64_t at) { timer.arm(to_lc_time_point(at)); } @@ -173,6 +193,16 @@ void mct_set_callback( timer.set_callback(callback_object(callback, caller, dropper)); } +void mct_set_callback_under_group( + manual_clock_timer& timer, + uint8_t* callback, + rust::Fn caller, + rust::Fn dropper, + const scheduling_group& sg +) { + timer.set_callback(sg, callback_object(callback, caller, dropper)); +} + void mct_arm_at(manual_clock_timer& timer, int64_t at) { timer.arm(to_mc_time_point(at)); } diff --git a/seastar/src/timer.hh b/seastar/src/timer.hh index d37182f..0c754af 100644 --- a/seastar/src/timer.hh +++ b/seastar/src/timer.hh @@ -7,6 +7,8 @@ namespace seastar_ffi { namespace timer { +using scheduling_group = seastar::scheduling_group; + namespace steady_clock { using steady_clock_timer = seastar::timer; @@ -20,6 +22,14 @@ void sct_set_callback( rust::Fn dropper ); +void sct_set_callback_under_group( + steady_clock_timer& timer, + uint8_t* callback, + rust::Fn caller, + rust::Fn dropper, + const scheduling_group& sg +); + void sct_arm_at(steady_clock_timer& timer, int64_t at); void sct_arm_at_periodic(steady_clock_timer& timer, int64_t at, int64_t period); @@ -49,6 +59,14 @@ void lct_set_callback( rust::Fn dropper ); +void lct_set_callback_under_group( + lowres_clock_timer& timer, + uint8_t* callback, + rust::Fn caller, + rust::Fn dropper, + const scheduling_group& sg +); + void lct_arm_at(lowres_clock_timer& timer, int64_t at); void lct_arm_at_periodic(lowres_clock_timer& timer, int64_t ay, int64_t period); @@ -78,6 +96,14 @@ void mct_set_callback( rust::Fn dropper ); +void mct_set_callback_under_group( + manual_clock_timer& timer, + uint8_t* callback, + rust::Fn caller, + rust::Fn dropper, + const scheduling_group& sg +); + void mct_arm_at(manual_clock_timer& timer, int64_t at); void mct_arm_at_periodic(manual_clock_timer& timer, int64_t at, int64_t period); diff --git a/seastar/src/timer.rs b/seastar/src/timer.rs index c641ccb..c4a857a 100644 --- a/seastar/src/timer.rs +++ b/seastar/src/timer.rs @@ -1,6 +1,7 @@ use crate::assert_runtime_is_running; use crate::ffi_utils::{get_dropper, get_fn_mut_void_caller}; use crate::CppTimer; +use crate::SchedulingGroup; use crate::{Clock, Duration, Instant}; use std::marker::PhantomData; @@ -94,6 +95,23 @@ impl Timer { ClockType::set_callback(&mut self.inner, boxed_callback, caller, dropper); } + /// Sets the callback function to be called under specified group when + /// the timer expires. + /// + /// # Arguments + /// * `callback` - The callback to be executed when the timer expires. + /// * `sg` - The scheduling group under which the callback will be executed. + pub fn set_callback_under_group( + &mut self, + callback: Func, + sg: &SchedulingGroup, + ) { + let caller = get_fn_mut_void_caller(&callback); + let dropper = get_dropper(&callback); + let boxed_callback = Box::into_raw(Box::new(callback)) as *mut u8; + ClockType::set_callback_under_group(&mut self.inner, boxed_callback, caller, dropper, &sg); + } + /// Sets the timer expiration time. /// /// It is illegal to arm a timer that has already been armed (and not disarmed by