From 87056f6ef7fd2ba5fce294bba59d5daa23bde453 Mon Sep 17 00:00:00 2001
From: Thomas Groh <tgroh@google.com>
Date: Wed, 26 Oct 2016 13:33:52 -0700
Subject: [PATCH] Remove the PriorityQueue in KeyedHolds

Adding and Removing holds is significantly more common than extracting
the current minimum hold.
---
 .../beam/runners/direct/WatermarkManager.java   | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
index f8cbc510c909..4930bb1eb463 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
@@ -1003,11 +1003,9 @@ public String toString() {
 
   private static class PerKeyHolds {
     private final Map<Object, KeyedHold> keyedHolds;
-    private final PriorityQueue<KeyedHold> allHolds;
 
     private PerKeyHolds() {
       this.keyedHolds = new HashMap<>();
-      this.allHolds = new PriorityQueue<>();
     }
 
     /**
@@ -1015,7 +1013,13 @@ private PerKeyHolds() {
      * there are no holds within this {@link PerKeyHolds}.
      */
     public Instant getMinHold() {
-      return allHolds.isEmpty() ? THE_END_OF_TIME.get() : allHolds.peek().getTimestamp();
+      Instant min = THE_END_OF_TIME.get();
+      for (KeyedHold hold : keyedHolds.values()) {
+        if (hold.getTimestamp().isBefore(min)) {
+          min = hold.getTimestamp();
+        }
+      }
+      return min;
     }
 
     /**
@@ -1023,20 +1027,15 @@ public Instant getMinHold() {
      * the same key.
      */
     public void updateHold(@Nullable Object key, Instant newHold) {
-      removeHold(key);
       KeyedHold newKeyedHold = KeyedHold.of(key, newHold);
       keyedHolds.put(key, newKeyedHold);
-      allHolds.offer(newKeyedHold);
     }
 
     /**
      * Removes the hold of the provided key.
      */
     public void removeHold(Object key) {
-      KeyedHold oldHold = keyedHolds.get(key);
-      if (oldHold != null) {
-        allHolds.remove(oldHold);
-      }
+      keyedHolds.remove(key);
     }
   }