[worker] Introduce a new byte aware self proposer scheduler

Author: MohamedBassem

Cargo.lock

@@ -13,12 +13,12 @@ use std::num::NonZeroUsize;
 use bytes::BytesMut;
 use futures::FutureExt;
 use pin_project::pin_project;
-use restate_memory::{MemoryLease, MemoryPool, NonZeroByteCount};
 use tokio::sync::mpsc::error::TryRecvError;
 use tokio::sync::{mpsc, oneshot};
 use tracing::{trace, warn};
 
 use restate_core::{ShutdownError, TaskCenter, TaskHandle, cancellation_token};
+use restate_memory::{MemoryLease, MemoryPool, NonZeroByteCount, PollMemoryPool};
 use restate_types::logs::Record;
 use restate_types::storage::StorageEncode;
 
@@ -310,15 +310,11 @@ pub struct LogSender<T> {
 }
 
 impl<T: StorageEncode> LogSender<T> {
-    /// Returns whether the next record or batch can be admitted.
-    pub fn has_capacity(&self) -> bool {
-        self.mem_pool.available() > 0
-    }
-
-    /// Waits until the next record or batch can be admitted.
-    pub fn wait_for_capacity(&self) -> impl std::future::Future<Output = ()> + 'static {
-        let mem_pool = self.mem_pool.clone();
-        async move { mem_pool.wait_until_available().await }
+    /// Returns a poll-style watcher over the appender's memory-pool capacity,
+    /// for use in manual [`std::future::Future::poll`] implementations. Waker
+    /// registration survives across polls.
+    pub fn capacity_poller(&self) -> PollMemoryPool {
+        PollMemoryPool::new(self.mem_pool.clone())
     }
 
     fn admit<E>(&self, payload: E) -> Result<E, EnqueueError<E>> {

crates/bifrost/src/background_appender.rs

@@ -531,6 +531,44 @@ impl PollMemoryPool {
             }
         }
     }
+
+    /// Waits until the pool has any available budget, registering `cx` for
+    /// wakeup while it is exhausted.
+    ///
+    /// Like [`MemoryPool::wait_until_available`], there is no guarantee that
+    /// the budget is still available by the time the caller acts on it.
+    pub fn poll_available(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> {
+        // Fast path: unlimited pools always have budget.
+        if self.pool.is_unlimited() {
+            return Poll::Ready(());
+        }
+
+        loop {
+            // Ensure we have a notified future *before* checking availability,
+            // so we don't miss a concurrent `return_memory()` notification.
+            let notified = self.notified.get_or_insert_with(|| {
+                Box::pin(
+                    self.pool
+                        .availability_notified_owned()
+                        .expect("bounded pool must provide notified"),
+                )
+            });
+
+            if self.pool.available() > 0 {
+                self.notified = None;
+                return Poll::Ready(());
+            }
+
+            match notified.as_mut().poll(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(()) => {
+                    // We were notified — discard the consumed future and loop
+                    // to re-check availability with a fresh notified.
+                    self.notified = None;
+                }
+            }
+        }
+    }
 }
 
 const _: () = {

crates/memory/src/pool.rs

@@ -60,6 +60,7 @@ bytes = { workspace = true }
 bytestring = { workspace = true }
 codederror = { workspace = true }
 derive_more = { workspace = true, features = ["debug", "display", "from"] }
+enum-map = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 metrics = { workspace = true }