From bfb57b22db42b0805c99f55b161eef03cbf98dcb Mon Sep 17 00:00:00 2001
From: Tamo <tamo@meilisearch.com>
Date: Thu, 23 Jan 2025 01:21:50 +0100
Subject: [PATCH] fix the early exit when rewriting a batch

---
 crates/index-scheduler/src/queue/batches.rs | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/crates/index-scheduler/src/queue/batches.rs b/crates/index-scheduler/src/queue/batches.rs
index b77602e1e..5c8a573ab 100644
--- a/crates/index-scheduler/src/queue/batches.rs
+++ b/crates/index-scheduler/src/queue/batches.rs
@@ -228,14 +228,16 @@ impl BatchQueue {
             })?;
         }
 
-        // Update the enqueued_at, we cannot retrieve the previous enqueued at from the previous batch and
+        // Update the enqueued_at: we cannot retrieve the previous enqueued at from the previous batch, and
         // must instead go through the db looking for it. We cannot look at the task contained in this batch either
         // because they may have been removed.
-        // What we know though is that the task date from before the enqueued_at and only two timestamp have been written
+        // What we know, though, is that the task date is from before the enqueued_at, and max two timestamps have been written
         // to the DB per batches.
         if let Some(ref old_batch) = old_batch {
             let started_at = old_batch.started_at.unix_timestamp_nanos();
-            let mut exit = false;
+
+            // We have either one or two enqueued at to remove
+            let mut exit = old_batch.stats.total_nb_tasks.clamp(0, 2);
             let mut iterator = self.enqueued_at.rev_iter_mut(wtxn)?;
             while let Some(entry) = iterator.next() {
                 let (key, mut value) = entry?;
@@ -243,14 +245,13 @@ impl BatchQueue {
                     continue;
                 }
                 if value.remove(old_batch.uid) {
+                    exit = exit.saturating_sub(1);
                     // Safe because the key and value are owned
                     unsafe {
                         iterator.put_current(&key, &value)?;
                     }
-                    if exit {
+                    if exit == 0 {
                         break;
-                    } else {
-                        exit = true;
                     }
                 }
             }