Add deferred reads to the Rust channel pipeline

Summary:
Add a reusable `DeferredRead` primitive that parks one type-erased inbound message together with its exact pipeline continuation. Completion and cancellation are marshalled to the originating EventBase, closed pipelines suppress late delivery, and an EventBase keepalive prevents late wakeups from targeting destroyed scheduler state.

Expose `CallbackContext::defer_read` and `spawn_deferred_read` to Rust handlers without revealing the native message layout or changing the existing synchronous fast path.

Reviewed By: robertroeser

Differential Revision: D117026425

fbshipit-source-id: b3ff5e2c3f462eea68d2da0dadc38ecdbbada9c3

Author: pranavtbhat

thrift/lib/rust/channel_pipeline/src/context.rs

@@ -25,6 +25,7 @@ use std::rc::Rc;
 
 use crate::adapter::BytesPtr;
 use crate::adapter::RustMessageAdapter;
+use crate::erased::BorrowedMessageAdapter;
 use crate::event_base::EventBaseTask;
 use crate::event_base::FirstPoll;
 use crate::ffi::ffi::FfiCallbackContext;
@@ -135,6 +136,74 @@ impl Drop for ContextHandle {
     }
 }
 
+/// Move-only ownership of an inbound message suspended at its pipeline
+/// position.
+///
+/// The original type-erased C++ message remains intact inside this token. It
+/// may be inspected or mutated on the originating EventBase and then resumed
+/// exactly once. Dropping the token cancels the read. Resume and cancellation
+/// are safe from any thread; native code performs delivery and destruction on
+/// the EventBase that owns the pipeline.
+pub struct DeferredRead {
+    storage: MaybeUninit<usize>,
+    _not_sync: PhantomData<Cell<()>>,
+}
+
+// SAFETY: the native token has unique ownership of both the message and its
+// pipeline guard. Cross-thread resume and destruction consume the token and
+// enqueue it onto the originating EventBase before touching or destroying the
+// EventBase-owned values.
+unsafe impl Send for DeferredRead {}
+
+impl DeferredRead {
+    /// Borrow a typed view of the intact message on its originating EventBase.
+    ///
+    /// Returns `None` off the owning EventBase. The returned borrow prevents
+    /// this token from being resumed or dropped while the message is in use.
+    pub fn borrow<M: BorrowedMessageAdapter>(&mut self) -> Option<M::View<'_>> {
+        // SAFETY: this token uniquely owns one live native deferred-read token.
+        // Native code returns its message only when called on the owning
+        // EventBase, and the resulting borrow is tied to `&mut self`.
+        let message =
+            unsafe { crate::ffi::ffi::deferred_read_message(self.storage.as_mut_ptr().cast()) };
+        if message.is_null() {
+            return None;
+        }
+        // SAFETY: native code returned the address of the live message owned by
+        // this token. `&mut self` guarantees exclusive access for the lifetime
+        // of the pinned borrow and the token cannot move or be consumed then.
+        let message = unsafe { &mut *message };
+        assert!(
+            M::holds(message),
+            "DeferredRead::borrow: box does not hold the requested type"
+        );
+        // SAFETY: the unconditional `M::holds` check establishes the adapter's
+        // exact C++ type, and the view remains tied to this exclusive borrow.
+        Some(unsafe { M::borrow(Pin::new_unchecked(message)) })
+    }
+
+    /// Resume the original inbound message from its captured pipeline
+    /// position. Delivery is suppressed if the pipeline has closed.
+    pub fn resume(self) {
+        let mut deferred = std::mem::ManuallyDrop::new(self);
+        // SAFETY: `deferred` owns one live token and ManuallyDrop prevents its
+        // destructor from consuming that token a second time.
+        unsafe {
+            crate::ffi::ffi::resume_deferred_read(deferred.storage.as_mut_ptr().cast());
+        }
+    }
+}
+
+impl Drop for DeferredRead {
+    fn drop(&mut self) {
+        // SAFETY: `DeferredRead` uniquely owns one token initialized by
+        // `CallbackContext::defer_read`; native code consumes it exactly once.
+        unsafe {
+            crate::ffi::ffi::destroy_deferred_read(self.storage.as_mut_ptr().cast());
+        }
+    }
+}
+
 /// Borrowed, callback-scoped view of the live C++ pipeline context.
 ///
 /// Each data or lifecycle callback receives an exclusive mutable reference to a
@@ -218,6 +287,20 @@ impl<'callback> CallbackContext<'callback> {
         });
     }
 
+    /// Poll a future now and capture the pipeline continuation only if it suspends.
+    ///
+    /// The future is first placed at its final pinned address and polled inline on
+    /// the current EventBase callback. If it is ready, `ready` runs immediately
+    /// with this borrowed context and its [`HandlerResult`] becomes the callback's
+    /// result; no [`ContextHandle`] is created. If it is pending, the task takes a
+    /// new one-shot `ContextHandle`, later polls remain on the same EventBase, and
+    /// `complete` receives that handle with the output. The current callback then
+    /// returns [`HandlerResult::Success`] because ownership of its in-flight work
+    /// has moved into the task.
+    ///
+    /// Panics are contained by [`EventBaseTask`]. A panic does not invoke either
+    /// completion callback and is reported here as [`HandlerResult::Success`] so
+    /// unwinding never crosses the C++ FFI boundary.
     pub(crate) fn spawn_deferred<T, Fut, Ready, Complete>(
         &mut self,
         future: Fut,
@@ -257,6 +340,67 @@ impl<'callback> CallbackContext<'callback> {
         handle
     }
 
+    /// Suspend the current inbound message without unpacking or copying it.
+    ///
+    /// Returns `None` if this callback has no message, the message is empty, or
+    /// it was already forwarded. Dropping the returned token cancels delivery;
+    /// [`DeferredRead::resume`] continues it from this handler's exact position.
+    pub fn defer_read(
+        &mut self,
+        message: crate::erased::RustTypeErasedBox<'_>,
+    ) -> Option<DeferredRead> {
+        let _ = &message;
+        let mut deferred = DeferredRead {
+            storage: MaybeUninit::uninit(),
+            _not_sync: PhantomData,
+        };
+        // SAFETY: storage is one pointer-aligned word and remains
+        // exclusively owned by `deferred`. On false, native code constructed
+        // nothing, so mem::forget prevents running a destructor on garbage.
+        let initialized = unsafe {
+            self.inner
+                .as_mut()
+                .init_deferred_read(deferred.storage.as_mut_ptr().cast())
+        };
+        if initialized {
+            Some(deferred)
+        } else {
+            std::mem::forget(deferred);
+            None
+        }
+    }
+
+    /// Suspend the current inbound message while a future runs on this
+    /// pipeline's EventBase.
+    ///
+    /// Completion receives the intact message token and may borrow, mutate,
+    /// resume, or cancel it. The task owns the token throughout suspension, so
+    /// panic or task destruction safely cancels the read.
+    ///
+    /// If the current callback has no live message, the message is empty, or it
+    /// was already forwarded, this returns HandlerResult::Error without
+    /// polling future or invoking complete.
+    pub fn spawn_deferred_read<T, Fut, Complete>(
+        &mut self,
+        message: crate::erased::RustTypeErasedBox<'_>,
+        future: Fut,
+        complete: Complete,
+    ) -> HandlerResult
+    where
+        T: Send + 'static,
+        Fut: Future<Output = T> + Send + 'static,
+        Complete: FnOnce(DeferredRead, T) + Send + 'static,
+    {
+        let Some(deferred) = self.defer_read(message) else {
+            return HandlerResult::Error;
+        };
+        let event_base = self.inner.as_ref().get_ref().event_base();
+        EventBaseTask::start(event_base, async move {
+            complete(deferred, future.await);
+        });
+        HandlerResult::Success
+    }
+
     /// Forward the inbound buffer downstream and return the result.
     ///
     /// The buffer is moved into the C++ message box and forwarded via

thrift/lib/rust/channel_pipeline/src/ffi.rs

@@ -136,6 +136,12 @@ pub(crate) mod ffi {
         #[cxx_name = "initContextHandle"]
         unsafe fn init_context_handle(self: Pin<&mut FfiCallbackContext>, storage: *mut u8);
 
+        /// SAFETY: `storage` must point to one pointer-sized, pointer-aligned,
+        /// uninitialized word. On success it contains one live token until
+        /// resumed or destroyed.
+        #[cxx_name = "initDeferredRead"]
+        unsafe fn init_deferred_read(self: Pin<&mut FfiCallbackContext>, storage: *mut u8) -> bool;
+
         #[namespace = "folly"]
         type EventBase;
         #[cxx_name = "eventBase"]
@@ -222,6 +228,19 @@ pub(crate) mod ffi {
         /// live guard back to the originating EventBase when necessary.
         #[cxx_name = "destroyContextHandle"]
         unsafe fn destroy_context_handle(storage: *mut u8);
+
+        /// SAFETY: `storage` must contain one live token constructed by
+        /// `init_deferred_read`. Both functions consume it exactly once.
+        #[cxx_name = "resumeDeferredRead"]
+        unsafe fn resume_deferred_read(storage: *mut u8);
+        /// SAFETY: same token contract as `resume_deferred_read`.
+        #[cxx_name = "destroyDeferredRead"]
+        unsafe fn destroy_deferred_read(storage: *mut u8);
+        /// SAFETY: `storage` must contain a live deferred-read token. The
+        /// returned pointer is non-null only on its originating EventBase and
+        /// remains borrowed from that token.
+        #[cxx_name = "deferredReadMessage"]
+        unsafe fn deferred_read_message(storage: *mut u8) -> *mut TypeErasedBox;
     }
 
     #[namespace = "folly"]

thrift/lib/rust/channel_pipeline/src/handler.rs

@@ -360,6 +360,25 @@ impl RustHandler for ContextHandleTestHandler {
                 });
                 return handler.fire_read(ctx, msg);
             }
+            27 | 28 => {
+                let (wake, woke) = WorkerWakeFuture::new();
+                let result = ctx.spawn_deferred_read(msg, wake, |deferred, ()| {
+                    deferred.resume();
+                });
+                woke.recv().expect("worker should issue the EventBase wake");
+                if self.scenario == 28 {
+                    ctx.close();
+                }
+                return result;
+            }
+            29 => {
+                let deferred = ctx
+                    .defer_read(msg)
+                    .expect("the callback should own one live inbound message");
+                std::thread::spawn(move || drop(deferred))
+                    .join()
+                    .expect("off-thread DeferredRead cancellation should not panic");
+            }
             24 => {
                 drop(msg.take::<BytesPtr>());
                 let mut handler = CoroExceptionHandle::new(|error| async move { error });

thrift/lib/rust/channel_pipeline/src/integration_test.rs

@@ -20,6 +20,7 @@ use std::sync::atomic::Ordering;
 use channel_pipeline::BytesPtr;
 use channel_pipeline::CallbackContext;
 use channel_pipeline::ContextHandle;
+use channel_pipeline::DeferredRead;
 use channel_pipeline::HandlerResult;
 use channel_pipeline::RustHandler;
 use channel_pipeline::RustHandlerOpaque;
@@ -720,6 +721,18 @@ fn context_handle_is_send_not_sync_clone_or_copy() {
     });
 }
 
+#[test]
+fn deferred_read_is_send_not_sync_clone_or_copy() {
+    run_with_timeout(|| {
+        static_assertions::assert_impl_all!(DeferredRead: Send);
+        static_assertions::assert_not_impl_any!(DeferredRead: Sync, Clone, Copy);
+        assert_eq!(
+            std::mem::size_of::<DeferredRead>(),
+            std::mem::size_of::<usize>()
+        );
+    });
+}
+
 #[test]
 fn native_context_handle_token_traits_match_rust_abi() {
     run_with_timeout(|| {
@@ -896,6 +909,35 @@ fn coro_read_worker_wake_repolls_and_resumes_read() {
     assert_context_handle_read_sandwich(&ffi::run_context_handle_sandwich_test(21), false);
 }
 
+#[test]
+fn deferred_read_worker_wake_resumes_original_erased_message() {
+    assert_context_handle_read_sandwich(&ffi::run_context_handle_sandwich_test(27), false);
+}
+
+#[test]
+fn deferred_read_resume_after_close_is_suppressed() {
+    let result = ffi::run_context_handle_sandwich_test(28);
+    assert_eq!(result.before_reads_before_fence, 1);
+    assert_eq!(result.before_reads_after_fence, 1);
+    assert_eq!(result.after_reads_before_fence, 0);
+    assert_eq!(result.after_reads_after_fence, 0);
+    assert_eq!(result.endpoint_calls_before_fence, 0);
+    assert_eq!(result.endpoint_calls_after_fence, 0);
+    assert!(!result.pointer_identity_preserved);
+}
+
+#[test]
+fn deferred_read_off_thread_drop_cancels_safely() {
+    let result = ffi::run_context_handle_sandwich_test(29);
+    assert_eq!(result.before_reads_before_fence, 1);
+    assert_eq!(result.before_reads_after_fence, 1);
+    assert_eq!(result.after_reads_before_fence, 0);
+    assert_eq!(result.after_reads_after_fence, 0);
+    assert_eq!(result.endpoint_calls_before_fence, 0);
+    assert_eq!(result.endpoint_calls_after_fence, 0);
+    assert!(!result.pointer_identity_preserved);
+}
+
 #[test]
 fn coro_write_ready_future_polls_inline_and_resumes_write() {
     assert_context_handle_write_sandwich(&ffi::run_context_handle_sandwich_test(22), true);

thrift/lib/rust/channel_pipeline/src/lib.rs

@@ -53,6 +53,10 @@
 //!   later wakes schedule further polls back onto that same EventBase. The task
 //!   owns a [`ContextHandle`], which retains the pipeline until the task
 //!   completes or is cancelled.
+//! - [`CallbackContext::defer_read`] moves the intact inbound erased message
+//!   and its continuation into a one-word [`DeferredRead`] token. Typed opaque
+//!   views can be borrowed on the owning EventBase, and consuming `resume`
+//!   forwards the original box without copying the message payload.
 //! - [`CoroReadHandle`], [`CoroWriteHandle`], and [`CoroExceptionHandle`] wrap
 //!   an `async` handler body so the message or error round-trips through a
 //!   future and resumes the pipeline on completion.
@@ -143,6 +147,7 @@
 //! | [`box_handler`] | Type-erases a handler for a downstream CXX factory |
 //! | [`CallbackContext`] | Borrowed pipeline context — `!Send`, `!Sync`, non-escapable |
 //! | [`ContextHandle`] | Move-only captured pipeline position; consuming `fire_read`/`fire_write`/`fire_exception` |
+//! | [`DeferredRead`] | Move-only suspended inbound message plus its exact pipeline continuation |
 //! | [`RustTypeErasedBox`] | Borrowed, type-erased message box; recover the value with `take::<T>()` |
 //! | [`BytesPtr`] | Zero-copy `unique_ptr<folly::IOBuf>` adapter |
 //! | [`PipelineError`] | Owned Rust error converted to `folly::exception_wrapper` |
@@ -256,6 +261,7 @@ pub use adapter::BytesPtr;
 pub use adapter::RustMessageAdapter;
 pub use context::CallbackContext;
 pub use context::ContextHandle;
+pub use context::DeferredRead;
 pub use context::PipelineError;
 pub use coro_handler::ContextReadMessage;
 pub use coro_handler::ContextWriteMessage;