Merge #5118

5118: Change the reserve and grant function to accept a closure r=ManyTheFish a=Kerollmops

This simplifies the usage of the grant and commits it at the right time, just after having written in it.

Co-authored-by: Kerollmops <clement@meilisearch.com>

crates/milli/src/update/new/channel.rs

@@ -7,7 +7,7 @@ use std::num::NonZeroU16;
 use std::ops::Range;
 use std::time::Duration;
 
-use bbqueue::framed::{FrameGrantR, FrameGrantW, FrameProducer};
+use bbqueue::framed::{FrameGrantR, FrameProducer};
 use bbqueue::BBBuffer;
 use bytemuck::{checked, CheckedBitPattern, NoUninit};
 use flume::{RecvTimeoutError, SendError};
@@ -454,14 +454,10 @@ impl<'b> ExtractorBbqueueSender<'b> {
         }
 
         // Spin loop to have a frame the size we requested.
-        let mut grant = reserve_grant(&mut producer, total_length, &self.sender)?;
-        payload_header.serialize_into(&mut grant);
-
-        // We only send a wake up message when the channel is empty
-        // so that we don't fill the channel with too many WakeUps.
-        if self.sender.is_empty() {
-            self.sender.send(ReceiverAction::WakeUp).unwrap();
-        }
+        reserve_and_write_grant(&mut producer, total_length, &self.sender, |grant| {
+            payload_header.serialize_into(grant);
+            Ok(())
+        })?;
 
         Ok(())
     }
@@ -500,8 +496,7 @@ impl<'b> ExtractorBbqueueSender<'b> {
         }
 
         // Spin loop to have a frame the size we requested.
-        let mut grant = reserve_grant(&mut producer, total_length, &self.sender)?;
-
+        reserve_and_write_grant(&mut producer, total_length, &self.sender, |grant| {
             let header_size = payload_header.header_size();
             let (header_bytes, remaining) = grant.split_at_mut(header_size);
             payload_header.serialize_into(header_bytes);
@@ -513,11 +508,8 @@ impl<'b> ExtractorBbqueueSender<'b> {
                 }
             }
 
-        // We only send a wake up message when the channel is empty
-        // so that we don't fill the channel with too many WakeUps.
-        if self.sender.is_empty() {
-            self.sender.send(ReceiverAction::WakeUp).unwrap();
-        }
+            Ok(())
+        })?;
 
         Ok(())
     }
@@ -575,19 +567,13 @@ impl<'b> ExtractorBbqueueSender<'b> {
         }
 
         // Spin loop to have a frame the size we requested.
-        let mut grant = reserve_grant(&mut producer, total_length, &self.sender)?;
-
+        reserve_and_write_grant(&mut producer, total_length, &self.sender, |grant| {
             let header_size = payload_header.header_size();
             let (header_bytes, remaining) = grant.split_at_mut(header_size);
             payload_header.serialize_into(header_bytes);
             let (key_buffer, value_buffer) = remaining.split_at_mut(key_length.get() as usize);
-        key_value_writer(key_buffer, value_buffer)?;
-
-        // We only send a wake up message when the channel is empty
-        // so that we don't fill the channel with too many WakeUps.
-        if self.sender.is_empty() {
-            self.sender.send(ReceiverAction::WakeUp).unwrap();
-        }
+            key_value_writer(key_buffer, value_buffer)
+        })?;
 
         Ok(())
     }
@@ -629,37 +615,44 @@ impl<'b> ExtractorBbqueueSender<'b> {
         }
 
         // Spin loop to have a frame the size we requested.
-        let mut grant = reserve_grant(&mut producer, total_length, &self.sender)?;
-
+        reserve_and_write_grant(&mut producer, total_length, &self.sender, |grant| {
             let header_size = payload_header.header_size();
             let (header_bytes, remaining) = grant.split_at_mut(header_size);
             payload_header.serialize_into(header_bytes);
-        key_writer(remaining)?;
-
-        // We only send a wake up message when the channel is empty
-        // so that we don't fill the channel with too many WakeUps.
-        if self.sender.is_empty() {
-            self.sender.send(ReceiverAction::WakeUp).unwrap();
-        }
+            key_writer(remaining)
+        })?;
 
         Ok(())
     }
 }
 
-/// Try to reserve a frame grant of `total_length` by spin looping
-/// on the BBQueue buffer and panics if the receiver has been disconnected.
-fn reserve_grant<'b>(
-    producer: &mut FrameProducer<'b>,
+/// Try to reserve a frame grant of `total_length` by spin
+/// looping on the BBQueue buffer, panics if the receiver
+/// has been disconnected or send a WakeUp message if necessary.
+fn reserve_and_write_grant<F>(
+    producer: &mut FrameProducer,
     total_length: usize,
     sender: &flume::Sender<ReceiverAction>,
-) -> crate::Result<FrameGrantW<'b>> {
+    f: F,
+) -> crate::Result<()>
+where
+    F: FnOnce(&mut [u8]) -> crate::Result<()>,
+{
     loop {
         for _ in 0..10_000 {
             match producer.grant(total_length) {
                 Ok(mut grant) => {
                     // We could commit only the used memory.
-                    grant.to_commit(total_length);
-                    return Ok(grant);
+                    f(&mut grant)?;
+                    grant.commit(total_length);
+
+                    // We only send a wake up message when the channel is empty
+                    // so that we don't fill the channel with too many WakeUps.
+                    if sender.is_empty() {
+                        sender.send(ReceiverAction::WakeUp).unwrap();
+                    }
+
+                    return Ok(());
                 }
                 Err(bbqueue::Error::InsufficientSize) => continue,
                 Err(e) => unreachable!("{e:?}"),