Expose entry priority in poll/peek in IndexedPriorityQueue

Author: losipiuk

core/trino-main/src/main/java/io/trino/execution/resourcegroups/IndexedPriorityQueue.java

@@ -81,6 +81,24 @@ public boolean remove(E element)
 
     @Override
     public E poll()
+    {
+        Entry<E> entry = pollEntry();
+        if (entry == null) {
+            return null;
+        }
+        return entry.getValue();
+    }
+
+    public Prioritized<E> pollPrioritized()
+    {
+        Entry<E> entry = pollEntry();
+        if (entry == null) {
+            return null;
+        }
+        return new Prioritized<>(entry.getValue(), entry.getPriority());
+    }
+
+    private Entry<E> pollEntry()
     {
         Iterator<Entry<E>> iterator = queue.iterator();
         if (!iterator.hasNext()) {
@@ -89,18 +107,35 @@ public E poll()
         Entry<E> entry = iterator.next();
         iterator.remove();
         checkState(index.remove(entry.getValue()) != null, "Failed to remove entry from index");
-        return entry.getValue();
+        return entry;
     }
 
     @Override
     public E peek()
+    {
+        Entry<E> entry = peekEntry();
+        if (entry == null) {
+            return null;
+        }
+        return entry.getValue();
+    }
+
+    public Prioritized<E> peekPrioritized()
+    {
+        Entry<E> entry = peekEntry();
+        if (entry == null) {
+            return null;
+        }
+        return new Prioritized<>(entry.getValue(), entry.getPriority());
+    }
+
+    public Entry<E> peekEntry()
     {
         Iterator<Entry<E>> iterator = queue.iterator();
         if (!iterator.hasNext()) {
             return null;
         }
-        Entry<E> entry = iterator.next();
-        return entry.getValue();
+        return iterator.next();
     }
 
     @Override
@@ -149,4 +184,26 @@ public long getGeneration()
             return generation;
         }
     }
+
+    public static class Prioritized<V>
+    {
+        private final V value;
+        private final long priority;
+
+        public Prioritized(V value, long priority)
+        {
+            this.value = requireNonNull(value, "value is null");
+            this.priority = priority;
+        }
+
+        public V getValue()
+        {
+            return value;
+        }
+
+        public long getPriority()
+        {
+            return priority;
+        }
+    }
 }

core/trino-main/src/test/java/io/trino/execution/resourcegroups/TestUpdateablePriorityQueue.java

@@ -18,6 +18,7 @@
 
 import java.util.List;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
@@ -35,6 +36,39 @@ public void testIndexedPriorityQueue()
         assertEquals(populateAndExtract(new IndexedPriorityQueue<>()), ImmutableList.of(3, 2, 1));
     }
 
+    @Test
+    public void testPrioritizedPeekPollIndexedPriorityQueue()
+    {
+        IndexedPriorityQueue<Object> queue = new IndexedPriorityQueue<>();
+        queue.addOrUpdate("a", 1);
+        queue.addOrUpdate("b", 3);
+        queue.addOrUpdate("c", 2);
+
+        IndexedPriorityQueue.Prioritized<Object> peek1 = queue.peekPrioritized();
+        assertThat(peek1.getValue()).isEqualTo("b");
+        assertThat(peek1.getPriority()).isEqualTo(3);
+        IndexedPriorityQueue.Prioritized<Object> poll1 = queue.pollPrioritized();
+        assertThat(poll1.getValue()).isEqualTo("b");
+        assertThat(poll1.getPriority()).isEqualTo(3);
+
+        IndexedPriorityQueue.Prioritized<Object> peek2 = queue.peekPrioritized();
+        assertThat(peek2.getValue()).isEqualTo("c");
+        assertThat(peek2.getPriority()).isEqualTo(2);
+        IndexedPriorityQueue.Prioritized<Object> poll2 = queue.pollPrioritized();
+        assertThat(poll2.getValue()).isEqualTo("c");
+        assertThat(poll2.getPriority()).isEqualTo(2);
+
+        IndexedPriorityQueue.Prioritized<Object> peek3 = queue.peekPrioritized();
+        assertThat(peek3.getValue()).isEqualTo("a");
+        assertThat(peek3.getPriority()).isEqualTo(1);
+        IndexedPriorityQueue.Prioritized<Object> poll3 = queue.pollPrioritized();
+        assertThat(poll3.getValue()).isEqualTo("a");
+        assertThat(poll3.getPriority()).isEqualTo(1);
+
+        assertThat(queue.peekPrioritized()).isNull();
+        assertThat(queue.pollPrioritized()).isNull();
+    }
+
     @Test
     public void testStochasticPriorityQueue()
     {