fix: prevent concurrent access to non-thread-safe maps in SketchBufferAggregatorHelper

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchBufferAggregatorHelper.java

@@ -33,6 +33,20 @@
 /**
  * A helper class used by {@link SketchBufferAggregator} and {@link SketchVectorAggregator}
  * for aggregation operations on byte buffers. Getting the object from value selectors is outside this class.
+ *
+ * <p>Thread-safety: the caches are keyed by {@link ByteBuffer} <em>identity</em>, using
+ * {@link IdentityHashMap}. A {@link java.util.concurrent.ConcurrentHashMap} is unsuitable here
+ * because it keys by {@link ByteBuffer#equals(Object)}/{@link ByteBuffer#hashCode()}, which are
+ * derived from the buffer's <em>remaining contents</em>. The union updates performed by this
+ * helper mutate those contents in place (through the {@link WritableMemory} wrapping the same
+ * buffer), so the hash of an already-inserted key changes as the aggregation progresses: lookups
+ * then miss, {@link #get} degrades to returning {@link SketchHolder#EMPTY}, and distinct buffers
+ * holding equal contents alias onto each other. Identity keying keeps the key stable because it
+ * depends only on the object reference.
+ *
+ * <p>Since {@link IdentityHashMap} is not thread-safe, all access is guarded by {@link #lock}.
+ * The critical sections are plain map operations on the query-processing path, so a single
+ * monitor does not become a bottleneck.
  */
 final class SketchBufferAggregatorHelper
 {
@@ -41,6 +55,12 @@ final class SketchBufferAggregatorHelper
   private final IdentityHashMap<ByteBuffer, Int2ObjectMap<Union>> unions = new IdentityHashMap<>();
   private final IdentityHashMap<ByteBuffer, WritableMemory> memCache = new IdentityHashMap<>();
 
+  /**
+   * Guards {@link #unions} and {@link #memCache}, which are {@link IdentityHashMap}s and must not
+   * be mutated or read concurrently.
+   */
+  private final Object lock = new Object();
+
   public SketchBufferAggregatorHelper(final int size, final int maxIntermediateSize)
   {
     this.size = size;
@@ -62,8 +82,11 @@ public void init(ByteBuffer buf, int position)
    */
   public Object get(ByteBuffer buf, int position)
   {
-    Int2ObjectMap<Union> unionMap = unions.get(buf);
-    Union union = unionMap != null ? unionMap.get(position) : null;
+    final Union union;
+    synchronized (lock) {
+      final Int2ObjectMap<Union> unionMap = unions.get(buf);
+      union = unionMap != null ? unionMap.get(position) : null;
+    }
     if (union == null) {
       return SketchHolder.EMPTY;
     }
@@ -81,13 +104,15 @@ public Object get(ByteBuffer buf, int position)
    */
   public void relocate(int oldPosition, int newPosition, ByteBuffer oldBuffer, ByteBuffer newBuffer)
   {
-    createNewUnion(newBuffer, newPosition, true);
-    Int2ObjectMap<Union> unionMap = unions.get(oldBuffer);
-    if (unionMap != null) {
-      unionMap.remove(oldPosition);
-      if (unionMap.isEmpty()) {
-        unions.remove(oldBuffer);
-        memCache.remove(oldBuffer);
+    synchronized (lock) {
+      createNewUnionLocked(newBuffer, newPosition, true);
+      final Int2ObjectMap<Union> unionMap = unions.get(oldBuffer);
+      if (unionMap != null) {
+        unionMap.remove(oldPosition);
+        if (unionMap.isEmpty()) {
+          unions.remove(oldBuffer);
+          memCache.remove(oldBuffer);
+        }
       }
     }
   }
@@ -99,20 +124,32 @@ public void relocate(int oldPosition, int newPosition, ByteBuffer oldBuffer, Byt
    */
   public Union getOrCreateUnion(ByteBuffer buf, int position)
   {
-    Int2ObjectMap<Union> unionMap = unions.get(buf);
-    Union union = unionMap != null ? unionMap.get(position) : null;
-    if (union != null) {
-      return union;
+    synchronized (lock) {
+      final Int2ObjectMap<Union> unionMap = unions.get(buf);
+      final Union union = unionMap != null ? unionMap.get(position) : null;
+      if (union != null) {
+        return union;
+      }
+      return createNewUnionLocked(buf, position, true);
     }
-    return createNewUnion(buf, position, true);
   }
 
   private Union createNewUnion(ByteBuffer buf, int position, boolean isWrapped)
   {
-    WritableMemory mem = getMemory(buf).writableRegion(position, maxIntermediateSize);
-    Union union = isWrapped
-                  ? (Union) SetOperation.wrap(mem)
-                  : (Union) SetOperation.builder().setNominalEntries(size).build(Family.UNION, mem);
+    synchronized (lock) {
+      return createNewUnionLocked(buf, position, isWrapped);
+    }
+  }
+
+  /**
+   * Creates and caches a {@link Union} at the given buffer location. Callers must hold {@link #lock}.
+   */
+  private Union createNewUnionLocked(ByteBuffer buf, int position, boolean isWrapped)
+  {
+    final WritableMemory mem = getMemoryLocked(buf).writableRegion(position, maxIntermediateSize);
+    final Union union = isWrapped
+                        ? (Union) SetOperation.wrap(mem)
+                        : (Union) SetOperation.builder().setNominalEntries(size).build(Family.UNION, mem);
     Int2ObjectMap<Union> unionMap = unions.get(buf);
     if (unionMap == null) {
       unionMap = new Int2ObjectOpenHashMap<>();
@@ -124,11 +161,17 @@ private Union createNewUnion(ByteBuffer buf, int position, boolean isWrapped)
 
   public void clear()
   {
-    unions.clear();
-    memCache.clear();
+    synchronized (lock) {
+      unions.clear();
+      memCache.clear();
+    }
   }
 
-  private WritableMemory getMemory(ByteBuffer buffer)
+  /**
+   * Returns the memory wrapping the given buffer, creating and caching it on first use.
+   * Callers must hold {@link #lock}.
+   */
+  private WritableMemory getMemoryLocked(ByteBuffer buffer)
   {
     WritableMemory mem = memCache.get(buffer);
     if (mem == null) {
@@ -137,4 +180,4 @@ private WritableMemory getMemory(ByteBuffer buffer)
     }
     return mem;
   }
-}
+}

processing/src/main/java/org/apache/druid/segment/data/ObjectStrategy.java

@@ -52,6 +52,25 @@ public interface ObjectStrategy<T> extends Comparator<T>
   @Nullable
   byte[] toBytes(@Nullable T val);
 
+  /**
+   * Converts an Object to bytes, handling the case where the input is already a byte[].
+   * This is needed when complex types are read from realtime segments, where
+   * {@link org.apache.druid.segment.ColumnValueSelector#getObject()} returns the already-serialized byte[] form
+   * rather than the deserialized object form. In that case, we can return the byte[] directly instead of
+   * trying to cast it to {@code T} and re-serialize it, which would cause a {@link ClassCastException}.
+   *
+   * @see <a href="https://github.com/apache/druid/issues/18340">GitHub issue #18340</a>
+   */
+  @Nullable
+  default byte[] objectToBytes(@Nullable Object val)
+  {
+    if (val instanceof byte[]) {
+      return (byte[]) val;
+    } else {
+      return toBytes((T) val);
+    }
+  }
+
   /**
    * Whether {@link #compare} is valid or not.
    */

processing/src/main/java/org/apache/druid/segment/serde/ComplexMetricSerde.java

@@ -94,7 +94,7 @@ public Function<Object, Long> inputSizeFn()
   public byte[] toBytes(@Nullable Object val)
   {
     if (val != null) {
-      byte[] bytes = getObjectStrategy().toBytes(val);
+      byte[] bytes = getObjectStrategy().objectToBytes(val);
       return bytes != null ? bytes : ByteArrays.EMPTY_ARRAY;
     } else {
       return ByteArrays.EMPTY_ARRAY;