refactor(runtime): make poll_task_with_extra consistent (#736)

Author: George-Miao

compio-runtime/src/runtime/future.rs

@@ -146,13 +146,15 @@ impl<T: OpCode + 'static> Future for Submit<T, Extra> {
         let this = unsafe { self.get_unchecked_mut() };
         loop {
             match this.state.take().expect("Cannot poll after ready") {
-                State::Submitted { key, .. } => match this.runtime.poll_task_with_extra(cx, key) {
-                    PushEntry::Pending(key) => {
-                        this.state = Some(State::submitted(key));
-                        return Poll::Pending;
+                State::Submitted { key, .. } => {
+                    match this.runtime.poll_task_with_extra(cx.waker(), key) {
+                        PushEntry::Pending(key) => {
+                            this.state = Some(State::submitted(key));
+                            return Poll::Pending;
+                        }
+                        PushEntry::Ready(res) => return Poll::Ready(res),
                     }
-                    PushEntry::Ready(res) => return Poll::Ready(res),
-                },
+                }
                 State::Idle { op } => {
                     let extra = cx.as_extra(|| this.runtime.default_extra());
                     match this.runtime.submit_raw(op, extra) {

compio-runtime/src/runtime/mod.rs

@@ -316,13 +316,13 @@ impl Runtime {
 
     pub(crate) fn poll_task_with_extra<T: OpCode>(
         &self,
-        cx: &mut Context,
+        waker: &Waker,
         key: Key<T>,
     ) -> PushEntry<Key<T>, (BufResult<usize, T>, Extra)> {
         instrument!(compio_log::Level::DEBUG, "poll_task_with_extra", ?key);
         let mut driver = self.driver.borrow_mut();
         driver.pop_with_extra(key).map_pending(|mut k| {
-            driver.update_waker(&mut k, cx.waker());
+            driver.update_waker(&mut k, waker);
             k
         })
     }