bytedance / monoio

Rust async runtime based on io-uring.
Apache License 2.0
4.02k stars 226 forks source link

suggestion: allow to register hook before executor goes park #296

Open Wireless4024 opened 2 months ago

Wireless4024 commented 2 months ago

Is your feature request related to a problem? Please describe. I have a special use case for my buffer pool,

struct BufferPool {
    keep_capacity: usize,
    buffers: UnsafeCell<VecDeque<BytesMut>>,
}

impl BufferPool {
    fn new(keep_capacity: usize) -> Self { todo!() }
    fn take(&self, capacity: usize) -> BytesMut { todo!() }
    fn store(&self, buffer: BytesMut) {}
    // cleanup extra buffer, deallocation can take some time
    fn compact(&self) {}
}

async fn init_listening_task() {
    let pool = Rc::new(BufferPool::new(16));

    // the idea!
    register_pre_park_hook(Rc::clone(&pool));

    loop {
        let buffer = pool.take(4096);
        let pool = Rc::clone(&pool);
        // accept connection
        spawn(async move {
            // do something with buffer
            pool.store(buffer);
        });
    }
}

Describe the solution you'd like and I would like to have some function to register hook before executor park itself, so I can compact my buffer that might overflow its capacity during brust of connections.

implmentation example

impl PreParkHook for BufferPool {
    fn pre_park(&self) { self.compact(); }
}

trait PreParkHook {
    fn pre_park(&self);
}

// example signature of register_pre_park_hook
fn register_pre_park_hook<T: PreParkHook>(value: Rc<T>) {
    let hook = value as Rc<dyn PreParkHook>;
    // put into runtime context
}

// example of code that will run during pre-park
fn run_pre_park(hooks: &mut VecDeque<Rc<dyn PreParkHook>>) {
    let mut i = 0;
    while i < hooks.len() {
        let hook = &hooks[i];
        if Rc::strong_count(hook) == 1 {
            // cleanup hook that has no strong reference
            // drop here instead of using weak to reduce
            // performance hit of drop call of the hook during execution
            hooks.remove(i);
            continue;
        }
        hook.pre_park();
        i += 1;
    }
}

probably run the hook before this call. https://github.com/bytedance/monoio/blob/master/monoio/src/runtime.rs#L187

ihciah commented 2 months ago

It's good! I wanted to add metrics that expose more detailed behavior of the runtime internals because I encountered some related needs, but it seems that your suggestion is better: let users plug in hooks to complete these statistics and even more things by themselves. It is also zero-cost because it has no overhead when users do not need to perform custom behavior. Adding such a Callback will make the generics of Runtime and RuntimeBuilder more complicated. Maybe using a default type can solve this problem.

Wireless4024 commented 2 months ago

can pre_park hook also have bool in return value as skip park? use case is when having code to spawn or wake something inside hook. I think by default runtime will continue to park?

this idea seem hard for user to get it right, always false make it not wake the runtime and alway true turn runtime into spin lock :smile:.

ihciah commented 2 months ago

The runtime needs syscall when park to get op completion or fd readiness. So skip park is not correct, but you can make the wait time to zero to do spin(currently there's no method to do it, previously I impl a draft to allow users impl io component in https://github.com/bytedance/monoio/pull/269). As you mentioned, control park(timeout=None or Some(0)) by the hook's return value is a good solution too, though it may be misused(return an enum should make it harder to misuse). A poc is like:

enum ParkPolicy {
  Wait,
  Spin,
}
pub trait ParkHook {
  fn before_park(&mut self) -> ParkPolicy;
}

or pass runtime's park as a lambda to the hook, and let it control the park timeout:

pub trait ParkHook {
  fn park(&mut self, f: impl FnOnce(Option<Duration>));
}