From 51355ea7221dfaeee9ac88840b1be3215fea0939 Mon Sep 17 00:00:00 2001
From: Nick Hill <nickhill@us.ibm.com>
Date: Tue, 16 May 2023 06:27:34 -0700
Subject: [PATCH] fix to pruning of cancelled requests from queue

---
 router/src/queue.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/router/src/queue.rs b/router/src/queue.rs
index 322fc6b4..5a4f9f08 100644
--- a/router/src/queue.rs
+++ b/router/src/queue.rs
@@ -375,7 +375,7 @@ impl<B: BatchType> State<B> {
             queue_index += 1;
             if entry.response_tx.is_disconnected() {
                 // Eject cancelled entry from queue
-                indices_to_drop.push(queue_index);
+                indices_to_drop.push(queue_index - 1);
                 continue
             }
             // This is the index into the queue after cancelled entries
@@ -465,7 +465,8 @@ impl<B: BatchType> State<B> {
 
         // Drop any cancelled requests
         if !indices_to_drop.is_empty() {
-            indices_to_drop.iter().for_each(|i| {
+            // Iterate in reverse so that indices remain correct
+            indices_to_drop.iter().rev().for_each(|i| {
                 self.entries.remove(*i);
             });
             metrics::gauge!("tgi_queue_size", self.entries.len() as f64);