From bc8818190c3114c1d47803ed3c5b71987cc959ad Mon Sep 17 00:00:00 2001
From: Aleksey Plekhanov <plehanov.alex@gmail.com>
Date: Fri, 10 Apr 2020 02:33:11 +0300
Subject: [PATCH] Thread-safety for ResponseContext.REGISTERED_KEYS

---
 .../druid/query/context/ResponseContext.java  | 34 +++++++++----------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/query/context/ResponseContext.java b/processing/src/main/java/org/apache/druid/query/context/ResponseContext.java
index c0d90b31f559..ffcef1620341 100644
--- a/processing/src/main/java/org/apache/druid/query/context/ResponseContext.java
+++ b/processing/src/main/java/org/apache/druid/query/context/ResponseContext.java
@@ -25,9 +25,10 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
 import org.apache.druid.guice.annotations.PublicApi;
+import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.NonnullPair;
 import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.query.SegmentDescriptor;
@@ -41,8 +42,9 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
-import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.function.BiFunction;
 
 /**
@@ -197,9 +199,11 @@ public enum Key implements BaseKey
     );
 
     /**
-     * TreeMap is used to have the natural ordering of its keys
+     * ConcurrentSkipListMap is used to have the natural ordering of its keys.
+     * Thread-safe structure is required since there is no guarantee that {@link #registerKey(BaseKey)}
+     * would be called only from class static blocks.
      */
-    private static final Map<String, BaseKey> REGISTERED_KEYS = new TreeMap<>();
+    private static final ConcurrentMap<String, BaseKey> REGISTERED_KEYS = new ConcurrentSkipListMap<>();
 
     static {
       for (BaseKey key : values()) {
@@ -211,14 +215,11 @@ public enum Key implements BaseKey
      * Primary way of registering context keys.
      * @throws IllegalArgumentException if the key has already been registered.
      */
-    public static synchronized void registerKey(BaseKey key)
+    public static void registerKey(BaseKey key)
     {
-      Preconditions.checkArgument(
-          !REGISTERED_KEYS.containsKey(key.getName()),
-          "Key [%s] has already been registered as a context key",
-          key.getName()
-      );
-      REGISTERED_KEYS.put(key.getName(), key);
+      if (REGISTERED_KEYS.putIfAbsent(key.getName(), key) != null) {
+        throw new IAE("Key [%s] has already been registered as a context key", key.getName());
+      }
     }
 
     /**
@@ -227,12 +228,11 @@ public static synchronized void registerKey(BaseKey key)
      */
     public static BaseKey keyOf(String name)
     {
-      Preconditions.checkState(
-           REGISTERED_KEYS.containsKey(name),
-          "Key [%s] has not yet been registered as a context key",
-          name
-      );
-      return REGISTERED_KEYS.get(name);
+      BaseKey key = REGISTERED_KEYS.get(name);
+      if (key == null) {
+        throw new ISE("Key [%s] has not yet been registered as a context key", name);
+      }
+      return key;
     }
 
     /**