feat(core): Event loop middlewares for Extensions (#13816)

1 files changed, 29 insertions, 0 deletions

diff --git a/core/extensions.rs b/core/extensions.rs
index 587a2d061..031cb073a 100644
--- a/core/extensions.rs
+++ b/core/extensions.rs
@@ -1,12 +1,14 @@
 use crate::OpFn;
 use crate::OpState;
 use anyhow::Error;
+use std::task::Context;
 
 pub type SourcePair = (&'static str, Box<SourceLoadFn>);
 pub type SourceLoadFn = dyn Fn() -> Result<String, Error>;
 pub type OpPair = (&'static str, Box<OpFn>);
 pub type OpMiddlewareFn = dyn Fn(&'static str, Box<OpFn>) -> Box<OpFn>;
 pub type OpStateFn = dyn Fn(&mut OpState) -> Result<(), Error>;
+pub type OpEventLoopFn = dyn Fn(&mut OpState, &mut Context) -> bool;
 
 #[derive(Default)]
 pub struct Extension {
@@ -14,6 +16,7 @@ pub struct Extension {
   ops: Option<Vec<OpPair>>,
   opstate_fn: Option<Box<OpStateFn>>,
   middleware_fn: Option<Box<OpMiddlewareFn>>,
+  event_loop_middleware: Option<Box<OpEventLoopFn>>,
   initialized: bool,
 }
 
@@ -56,6 +59,22 @@ impl Extension {
   pub fn init_middleware(&mut self) -> Option<Box<OpMiddlewareFn>> {
     self.middleware_fn.take()
   }
+
+  pub fn init_event_loop_middleware(&mut self) -> Option<Box<OpEventLoopFn>> {
+    self.event_loop_middleware.take()
+  }
+
+  pub fn run_event_loop_middleware(
+    &self,
+    op_state: &mut OpState,
+    cx: &mut Context,
+  ) -> bool {
+    self
+      .event_loop_middleware
+      .as_ref()
+      .map(|f| f(op_state, cx))
+      .unwrap_or(false)
+  }
 }
 
 // Provides a convenient builder pattern to declare Extensions
@@ -65,6 +84,7 @@ pub struct ExtensionBuilder {
   ops: Vec<OpPair>,
   state: Option<Box<OpStateFn>>,
   middleware: Option<Box<OpMiddlewareFn>>,
+  event_loop_middleware: Option<Box<OpEventLoopFn>>,
 }
 
 impl ExtensionBuilder {
@@ -94,6 +114,14 @@ impl ExtensionBuilder {
     self
   }
 
+  pub fn event_loop_middleware<F>(&mut self, middleware_fn: F) -> &mut Self
+  where
+    F: Fn(&mut OpState, &mut Context) -> bool + 'static,
+  {
+    self.event_loop_middleware = Some(Box::new(middleware_fn));
+    self
+  }
+
   pub fn build(&mut self) -> Extension {
     let js_files = Some(std::mem::take(&mut self.js));
     let ops = Some(std::mem::take(&mut self.ops));
@@ -102,6 +130,7 @@ impl ExtensionBuilder {
       ops,
       opstate_fn: self.state.take(),
       middleware_fn: self.middleware.take(),
+      event_loop_middleware: self.event_loop_middleware.take(),
       initialized: false,
     }
   }