From 7ffe402b87f96b438fe378e64b0f5b2e8f67262a Mon Sep 17 00:00:00 2001
From: Abhishek Agarwal <1477457+abhishekagarwal87@users.noreply.github.com>
Date: Mon, 7 Jun 2021 15:27:27 +0530
Subject: [PATCH] Fix priority parsing in HiLoQueryLaningStrategy (#11302)

---
 .../scheduling/HiLoQueryLaningStrategy.java    |  8 ++++++--
 .../HiLoQueryLaningStrategyTest.java           | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java
index 0a974228f66..3af5d10ac9a 100644
--- a/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java
+++ b/server/src/main/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategy.java
@@ -67,8 +67,12 @@ public class HiLoQueryLaningStrategy implements QueryLaningStrategy
   public <T> Optional<String> computeLane(QueryPlus<T> query, Set<SegmentServerSelector> segments)
   {
     final Query<T> theQuery = query.getQuery();
-    // QueryContexts.getPriority gives a default, since we are setting priority
-    final Integer priority = theQuery.getContextValue(QueryContexts.PRIORITY_KEY);
+    // QueryContexts.getPriority gives a default, but it can parse the value to integer. Before calling QueryContexts.getPriority
+    // we make sure that priority has been set.
+    Integer priority = null;
+    if (null != theQuery.getContextValue(QueryContexts.PRIORITY_KEY)) {
+      priority = QueryContexts.getPriority(theQuery);
+    }
     final String lane = theQuery.getContextValue(QueryContexts.LANE_KEY);
     if (lane == null && priority != null && priority < 0) {
       return Optional.of(LOW);
diff --git a/server/src/test/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategyTest.java b/server/src/test/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategyTest.java
index 3bc7c3606b1..2dc8594ce50 100644
--- a/server/src/test/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategyTest.java
+++ b/server/src/test/java/org/apache/druid/server/scheduling/HiLoQueryLaningStrategyTest.java
@@ -155,6 +155,13 @@ public class HiLoQueryLaningStrategyTest
     Assert.assertFalse(strategy.computeLane(QueryPlus.wrap(query), ImmutableSet.of()).isPresent());
   }
 
+  @Test
+  public void testLaningInteractivePriority_String()
+  {
+    TimeseriesQuery query = queryBuilder.context(ImmutableMap.of(QueryContexts.PRIORITY_KEY, "100")).build();
+    Assert.assertFalse(strategy.computeLane(QueryPlus.wrap(query), ImmutableSet.of()).isPresent());
+  }
+
   @Test
   public void testLaningLowPriority()
   {
@@ -166,6 +173,17 @@ public class HiLoQueryLaningStrategyTest
     );
   }
 
+  @Test
+  public void testLaningLowPriority_String()
+  {
+    TimeseriesQuery query = queryBuilder.context(ImmutableMap.of(QueryContexts.PRIORITY_KEY, "-1")).build();
+    Assert.assertTrue(strategy.computeLane(QueryPlus.wrap(query), ImmutableSet.of()).isPresent());
+    Assert.assertEquals(
+        HiLoQueryLaningStrategy.LOW,
+        strategy.computeLane(QueryPlus.wrap(query), ImmutableSet.of()).get()
+    );
+  }
+
   @Test
   public void testLaningPreservesManualSetLane()
   {