From ee66db900e2645e8fada79662600959c1a398231 Mon Sep 17 00:00:00 2001
From: Andy Sloane <asloane@tetrationanalytics.com>
Date: Mon, 23 Oct 2017 12:10:24 -0700
Subject: [PATCH] Fix binary serialization in caching (#4993)

* Fix binary serialization in caching

The previous caching code just concatenated a list of objects into a
byte array -- this is actually not valid because jackson-databind uses
enumerated references to strings internally, and concatenating multiple
binary serialized objects can throw off the references.

This change uses a single JsonGenerator to serialize the object list
rather than concatenating byte arrays.

* remove unused imports
---
 .../main/java/io/druid/client/CacheUtil.java  | 24 +++++++------------
 .../client/CachingClusteredClientTest.java    | 15 ++++++------
 2 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/server/src/main/java/io/druid/client/CacheUtil.java b/server/src/main/java/io/druid/client/CacheUtil.java
index c8c59e751b8..af59bee8bcb 100644
--- a/server/src/main/java/io/druid/client/CacheUtil.java
+++ b/server/src/main/java/io/druid/client/CacheUtil.java
@@ -19,9 +19,9 @@
 
 package io.druid.client;
 
+import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Throwables;
-import com.google.common.collect.Lists;
 import io.druid.client.cache.Cache;
 import io.druid.client.cache.CacheConfig;
 import io.druid.java.util.common.StringUtils;
@@ -31,9 +31,9 @@ import io.druid.query.QueryContexts;
 import io.druid.query.SegmentDescriptor;
 import org.joda.time.Interval;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.util.List;
 
 public class CacheUtil
 {
@@ -60,22 +60,14 @@ public class CacheUtil
   public static void populate(Cache cache, ObjectMapper mapper, Cache.NamedKey key, Iterable<Object> results)
   {
     try {
-      List<byte[]> bytes = Lists.newArrayList();
-      int size = 0;
-      for (Object result : results) {
-        final byte[] array = mapper.writeValueAsBytes(result);
-        size += array.length;
-        bytes.add(array);
+      ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+      try (JsonGenerator gen = mapper.getFactory().createGenerator(bytes)) {
+        for (Object result : results) {
+          gen.writeObject(result);
+        }
       }
 
-      byte[] valueBytes = new byte[size];
-      int offset = 0;
-      for (byte[] array : bytes) {
-        System.arraycopy(array, 0, valueBytes, offset, array.length);
-        offset += array.length;
-      }
-
-      cache.put(key, valueBytes);
+      cache.put(key, bytes.toByteArray());
     }
     catch (IOException e) {
       throw Throwables.propagate(e);
diff --git a/server/src/test/java/io/druid/client/CachingClusteredClientTest.java b/server/src/test/java/io/druid/client/CachingClusteredClientTest.java
index 88cd3697efa..974f87f96ef 100644
--- a/server/src/test/java/io/druid/client/CachingClusteredClientTest.java
+++ b/server/src/test/java/io/druid/client/CachingClusteredClientTest.java
@@ -1307,9 +1307,9 @@ public class CachingClusteredClientTest
         makeSelectResults(dimensions, metrics, DateTimes.of("2011-01-02"), ImmutableMap.of("a", "c", "rows", 5)),
 
         Intervals.of("2011-01-05/2011-01-10"),
-        makeSelectResults(dimensions, metrics, DateTimes.of("2011-01-05"), ImmutableMap.of("a", "d", "rows", 5),
-            DateTimes.of("2011-01-06"), ImmutableMap.of("a", "e", "rows", 6),
-            DateTimes.of("2011-01-07"), ImmutableMap.of("a", "f", "rows", 7),
+        makeSelectResults(dimensions, metrics, DateTimes.of("2011-01-05"),
+            DateTimes.of("2011-01-06"),
+            DateTimes.of("2011-01-07"), ImmutableMap.of("a", "f", "rows", 7), ImmutableMap.of("a", "ff"),
             DateTimes.of("2011-01-08"), ImmutableMap.of("a", "g", "rows", 8),
             DateTimes.of("2011-01-09"), ImmutableMap.of("a", "h", "rows", 9)
         ),
@@ -1335,11 +1335,11 @@ public class CachingClusteredClientTest
     TestHelper.assertExpectedResults(
         makeSelectResults(dimensions, metrics, DateTimes.of("2011-01-01"), ImmutableMap.of("a", "b", "rows", 1),
             DateTimes.of("2011-01-02"), ImmutableMap.of("a", "c", "rows", 5),
-            DateTimes.of("2011-01-05"), ImmutableMap.of("a", "d", "rows", 5),
+            DateTimes.of("2011-01-05"),
             DateTimes.of("2011-01-05T01"), ImmutableMap.of("a", "d", "rows", 5),
-            DateTimes.of("2011-01-06"), ImmutableMap.of("a", "e", "rows", 6),
+            DateTimes.of("2011-01-06"),
             DateTimes.of("2011-01-06T01"), ImmutableMap.of("a", "e", "rows", 6),
-            DateTimes.of("2011-01-07"), ImmutableMap.of("a", "f", "rows", 7),
+            DateTimes.of("2011-01-07"), ImmutableMap.of("a", "f", "rows", 7), ImmutableMap.of("a", "ff"),
             DateTimes.of("2011-01-07T01"), ImmutableMap.of("a", "f", "rows", 7),
             DateTimes.of("2011-01-08"), ImmutableMap.of("a", "g", "rows", 8),
             DateTimes.of("2011-01-08T01"), ImmutableMap.of("a", "g", "rows", 8),
@@ -2654,7 +2654,8 @@ public class CachingClusteredClientTest
 
       retVal.add(new Result<>(
           timestamp,
-          new SelectResultValue(null, dimensions, metrics, values)
+          new SelectResultValue(ImmutableMap.of(timestamp.toString(), 0),
+                                dimensions, metrics, values)
       ));
     }
     return retVal;