Merge pull request #40 from sportsball-ai/fix-custom-waiter

make with_custom_wait_for_work actually return self
cksac · Nov 13, 2024 · 998d77f · 998d77f
2 parents 1a10661 + df65cda
commit 998d77f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/cached.rs b/src/cached.rs
@@ -134,8 +134,9 @@ where
     /// Replaces the yielding for work behavior with an arbitrary future. Rather than yielding
     /// the runtime repeatedly this will generate and `.await` a future of your choice.
     /// ***This is incompatible with*** [`Self::with_yield_count()`].
-    pub fn with_custom_wait_for_work(mut self, wait_for_work_fn: impl WaitForWorkFn) {
+    pub fn with_custom_wait_for_work(mut self, wait_for_work_fn: impl WaitForWorkFn) -> Self {
         self.wait_for_work_fn = Arc::new(wait_for_work_fn);
+        self
     }
 
     pub fn max_batch_size(&self) -> usize {

diff --git a/src/non_cached.rs b/src/non_cached.rs
@@ -85,8 +85,9 @@ where
     /// Replaces the yielding for work behavior with an arbitrary future. Rather than yielding
     /// the runtime repeatedly this will generate and `.await` a future of your choice.
     /// ***This is incompatible with*** [`Self::with_yield_count()`].
-    pub fn with_custom_wait_for_work(mut self, wait_for_work_fn: impl WaitForWorkFn) {
+    pub fn with_custom_wait_for_work(mut self, wait_for_work_fn: impl WaitForWorkFn) -> Self {
         self.wait_for_work_fn = Arc::new(wait_for_work_fn);
+        self
     }
 
     pub fn max_batch_size(&self) -> usize {