From 414deefc1996df1bc4dbbd151e10c02dd7c80bbb Mon Sep 17 00:00:00 2001
From: Frediano Ziglio <frediano.ziglio@cloud.com>
Date: Wed, 27 Nov 2024 18:45:44 +0000
Subject: [PATCH] Fix removing elements from Imperative priority queue

The queue is implemented using heap tree saved as an array.
This is pretty standard.
To do a removal you replace the element to remove with the last element and
adjust the heap order.
This can result in having to move the moved element up or down base its the value.
Code only moved down, so also move it up if needed.

Signed-off-by: Frediano Ziglio <frediano.ziglio@cloud.com>
---
 .../xapi-stdext/lib/xapi-stdext-threads/ipq.ml | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/ipq.ml b/ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/ipq.ml
index cba404293e..9fd99c5989 100644
--- a/ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/ipq.ml
+++ b/ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/ipq.ml
@@ -67,9 +67,17 @@ let remove h s =
   if s < 0 || s >= h.size then
     invalid_arg (Printf.sprintf "%s: index %d out of bounds" __FUNCTION__ s) ;
   let n = h.size - 1 in
-  h.size <- n ;
   let d = h.data in
   let x = d.(n) in
+  (* moving [x] up in the heap *)
+  let rec moveup i =
+    let fi = (i - 1) / 2 in
+    if i > 0 && Mtime.is_later d.(fi).time ~than:x.time then (
+      d.(i) <- d.(fi) ;
+      moveup fi
+    ) else
+      d.(i) <- x
+  in
   (* moving [x] down in the heap *)
   let rec movedown i =
     let j = (2 * i) + 1 in
@@ -86,7 +94,13 @@ let remove h s =
     else
       d.(i) <- x
   in
-  movedown s
+  if s = n then
+    ()
+  else if Mtime.is_later d.(s).time ~than:x.time then
+    moveup s
+  else
+    movedown s ;
+  h.size <- n
 
 let find h ev =
   let rec iter n =