- Use a common scratchpad for serializing the different parts of the query.

- Use a threadlocal for the scratchpad. This avoids costly resizing, or initialiing too large buffer
  for every query. Using a thread local is fine now that we limit the number of search threads to a
  reasonable number = #cores * 2.

4 files changed, 45 insertions, 48 deletions

diff --git a/container-search/src/main/java/com/yahoo/prelude/query/Item.java b/container-search/src/main/java/com/yahoo/prelude/query/Item.java
index 7f421832d5f..8c154072a42 100644
--- a/container-search/src/main/java/com/yahoo/prelude/query/Item.java
+++ b/container-search/src/main/java/com/yahoo/prelude/query/Item.java
@@ -1,7 +1,6 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.prelude.query;
 
-
 import com.yahoo.collections.CopyOnWriteHashMap;
 import com.yahoo.compress.IntegerCompressor;
 import com.yahoo.language.Language;
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java
index 4bea68dc403..7e54afcc070 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java
@@ -18,14 +18,10 @@ import java.util.function.Consumer;
  */
 public class MapConverter {
 
-    public static void convertMapTensors(Map<String, Object> map, Consumer<TensorProperty.Builder> inserter) {
-        GrowableByteBuffer buffer = null;
+    public static void convertMapTensors(GrowableByteBuffer buffer, Map<String, Object> map, Consumer<TensorProperty.Builder> inserter) {
         for (var entry : map.entrySet()) {
             var value = entry.getValue();
             if (value instanceof Tensor tensor) {
-                if (buffer == null) {
-                    buffer = new GrowableByteBuffer(4096);
-                }
                 buffer.clear();
                 TypedBinaryFormat.encode(tensor, buffer);
                 inserter.accept(TensorProperty.newBuilder().setName(entry.getKey()).setValue(ByteString.copyFrom(buffer.getByteBuffer().flip())));
@@ -51,10 +47,10 @@ public class MapConverter {
         }
     }
 
-    public static void convertMultiMap(Map<String, List<Object>> map,
+    public static void convertMultiMap(GrowableByteBuffer buffer,
+                                       Map<String, List<Object>> map,
                                        Consumer<StringProperty.Builder> stringInserter,
                                        Consumer<TensorProperty.Builder> tensorInserter) {
-        GrowableByteBuffer buffer = null;
         for (var entry : map.entrySet()) {
             if (entry.getValue() != null) {
                 var key = entry.getKey();
@@ -62,9 +58,6 @@ public class MapConverter {
                 for (var value : entry.getValue()) {
                     if (value != null) {
                         if (value instanceof Tensor tensor) {
-                            if (buffer == null) {
-                                buffer = new GrowableByteBuffer(4096);
-                            }
                             buffer.clear();
                             TypedBinaryFormat.encode(tensor, buffer);
                             tensorInserter.accept(TensorProperty.newBuilder().setName(key).setValue(ByteString.copyFrom(buffer.getByteBuffer().flip())));
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
index de4f4f45eed..90e16e9dd44 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
@@ -38,7 +38,14 @@ import java.util.function.Consumer;
 
 public class ProtobufSerialization {
 
-    private static final int INITIAL_SERIALIZATION_BUFFER_SIZE = 10 * 1024;
+    /*
+     * This is a thread local buffer that is used as scratchpad during serialization.
+     * - avoids the unnecessary cost of allocating and initializing a buffer that is too large.
+     * - avoids resizing for large queries.
+     * - Reduces garbage creation.
+     * There is a limited number of threads that will use this so the upper bound should be fine.
+     */
+    private static final ThreadLocal<GrowableByteBuffer> threadLocalBuffer = ThreadLocal.withInitial(() -> new GrowableByteBuffer(4096));
 
     static byte[] serializeSearchRequest(Query query, int hits, String serverId, double requestTimeout) {
         return convertFromQuery(query, hits, serverId, requestTimeout).toByteArray();
@@ -58,7 +65,8 @@ public class ProtobufSerialization {
         if (documentDb != null) {
             builder.setDocumentType(documentDb);
         }
-        builder.setQueryTreeBlob(serializeQueryTree(query.getModel().getQueryTree()));
+        GrowableByteBuffer scratchPad = threadLocalBuffer.get();
+        builder.setQueryTreeBlob(serializeQueryTree(query.getModel().getQueryTree(), scratchPad));
 
         if (query.getGroupingSessionCache() || query.getRanking().getQueryCache()) {
             // TODO verify that the session key is included whenever rank properties would have been
@@ -69,7 +77,8 @@ public class ProtobufSerialization {
         }
         if (GroupingExecutor.hasGroupingList(query)) {
             List<Grouping> groupingList = GroupingExecutor.getGroupingList(query);
-            BufferSerializer gbuf = new BufferSerializer(new GrowableByteBuffer());
+            scratchPad.clear();
+            BufferSerializer gbuf = new BufferSerializer(scratchPad);
             gbuf.putInt(null, groupingList.size());
             for (Grouping g : groupingList) {
                 g.serialize(gbuf);
@@ -84,7 +93,7 @@ public class ProtobufSerialization {
         builder.setTraceLevel(getTraceLevelForBackend(query));
         builder.setProfileDepth(query.getTrace().getProfileDepth());
 
-        mergeToSearchRequestFromRanking(query.getRanking(), builder);
+        mergeToSearchRequestFromRanking(query.getRanking(), scratchPad, builder);
 
         return builder.build();
     }
@@ -100,7 +109,7 @@ public class ProtobufSerialization {
         return traceLevel;
     }
 
-    private static void mergeToSearchRequestFromRanking(Ranking ranking, SearchProtocol.SearchRequest.Builder builder) {
+    private static void mergeToSearchRequestFromRanking(Ranking ranking, GrowableByteBuffer scratchPad, SearchProtocol.SearchRequest.Builder builder) {
         builder.setRankProfile(ranking.getProfile());
 
         if (ranking.getQueryCache()) {
@@ -115,8 +124,8 @@ public class ProtobufSerialization {
 
         var featureMap = ranking.getFeatures().asMap();
         MapConverter.convertMapPrimitives(featureMap, builder::addFeatureOverrides);
-        MapConverter.convertMapTensors(featureMap, builder::addTensorFeatureOverrides);
-        mergeRankProperties(ranking, builder::addRankProperties, builder::addTensorRankProperties);
+        MapConverter.convertMapTensors(scratchPad, featureMap, builder::addTensorFeatureOverrides);
+        mergeRankProperties(ranking, scratchPad, builder::addRankProperties, builder::addTensorRankProperties);
     }
 
     private static void mergeToSearchRequestFromSorting(Sorting sorting, SearchProtocol.SearchRequest.Builder builder) {
@@ -160,8 +169,9 @@ public class ProtobufSerialization {
         if (ranking.getLocation() != null) {
             builder.setGeoLocation(ranking.getLocation().backendString());
         }
+        GrowableByteBuffer scratchPad = threadLocalBuffer.get();
         if (includeQueryData) {
-            mergeQueryDataToDocsumRequest(query, builder);
+            mergeQueryDataToDocsumRequest(query, scratchPad, builder);
         }
         if (query.getTrace().getLevel() >= 3) {
             query.trace((includeQueryData ? "ProtoBuf: Resending " : "Not resending ") + "query during document summary fetching", 3);
@@ -178,18 +188,18 @@ public class ProtobufSerialization {
         return builder.build().toByteArray();
     }
 
-    private static void mergeQueryDataToDocsumRequest(Query query, SearchProtocol.DocsumRequest.Builder builder) {
+    private static void mergeQueryDataToDocsumRequest(Query query, GrowableByteBuffer scratchPad, SearchProtocol.DocsumRequest.Builder builder) {
         var ranking = query.getRanking();
         var featureMap = ranking.getFeatures().asMap();
 
-        builder.setQueryTreeBlob(serializeQueryTree(query.getModel().getQueryTree()));
+        builder.setQueryTreeBlob(serializeQueryTree(query.getModel().getQueryTree(), scratchPad));
 
         MapConverter.convertMapPrimitives(featureMap, builder::addFeatureOverrides);
-        MapConverter.convertMapTensors(featureMap, builder::addTensorFeatureOverrides);
+        MapConverter.convertMapTensors(scratchPad, featureMap, builder::addTensorFeatureOverrides);
         if (query.getPresentation().getHighlight() != null) {
             MapConverter.convertStringMultiMap(query.getPresentation().getHighlight().getHighlightTerms(), builder::addHighlightTerms);
         }
-        mergeRankProperties(ranking, builder::addRankProperties, builder::addTensorRankProperties);
+        mergeRankProperties(ranking, scratchPad, builder::addRankProperties, builder::addTensorRankProperties);
     }
     static byte[] serializeResult(Result searchResult) {
         return convertFromResult(searchResult).toByteArray();
@@ -297,24 +307,25 @@ public class ProtobufSerialization {
         return builder.build();
     }
 
-    private static ByteString serializeQueryTree(QueryTree queryTree) {
-        int bufferSize = INITIAL_SERIALIZATION_BUFFER_SIZE;
+    private static ByteString serializeQueryTree(QueryTree queryTree, GrowableByteBuffer scratchPad) {
         while (true) {
             try {
-                ByteBuffer treeBuffer = ByteBuffer.allocate(bufferSize);
+                scratchPad.clear();
+                ByteBuffer treeBuffer = scratchPad.getByteBuffer();
                 queryTree.encode(treeBuffer);
-                treeBuffer.flip();
-                return ByteString.copyFrom(treeBuffer);
+                return ByteString.copyFrom(treeBuffer.flip());
             } catch (java.nio.BufferOverflowException e) {
-                bufferSize *= 2;
+                scratchPad.clear();
+                scratchPad.grow(scratchPad.capacity()*2);
             }
         }
     }
 
     private static void mergeRankProperties(Ranking ranking,
+                                            GrowableByteBuffer scratchPad,
                                             Consumer<StringProperty.Builder> stringProperties,
                                             Consumer<TensorProperty.Builder> tensorProperties) {
-        MapConverter.convertMultiMap(ranking.getProperties().asMap(), propB -> {
+        MapConverter.convertMultiMap(scratchPad, ranking.getProperties().asMap(), propB -> {
             if (!GetDocSumsPacket.sessionIdKey.equals(propB.getName())) {
                 stringProperties.accept(propB);
             }
diff --git a/container-search/src/main/java/com/yahoo/search/query/QueryTree.java b/container-search/src/main/java/com/yahoo/search/query/QueryTree.java
index 6326097d9bd..1cc2b98c65b 100644
--- a/container-search/src/main/java/com/yahoo/search/query/QueryTree.java
+++ b/container-search/src/main/java/com/yahoo/search/query/QueryTree.java
@@ -140,22 +140,18 @@ public class QueryTree extends CompositeItem {
         else if (b == null || b instanceof NullItem) {
             return a;
         }
-        else if (a instanceof NotItem && b instanceof NotItem) {
-            NotItem notItemA = (NotItem)a;
-            NotItem notItemB = (NotItem)b;
+        else if (a instanceof NotItem notItemA && b instanceof NotItem notItemB) {
             NotItem combined = new NotItem();
             combined.addPositiveItem(and(notItemA.getPositiveItem(), notItemB.getPositiveItem()));
             notItemA.negativeItems().forEach(item -> combined.addNegativeItem(item));
             notItemB.negativeItems().forEach(item -> combined.addNegativeItem(item));
             return combined;
         }
-        else if (a instanceof NotItem){
-            NotItem notItem = (NotItem)a;
+        else if (a instanceof NotItem notItem){
             notItem.addPositiveItem(b);
             return a;
         }
-        else if (b instanceof NotItem){
-            NotItem notItem = (NotItem)b;
+        else if (b instanceof NotItem notItem){
             notItem.addPositiveItem(a);
             return notItem;
         }
@@ -179,17 +175,16 @@ public class QueryTree extends CompositeItem {
     }
 
     private static void getPositiveTerms(Item item, List<IndexedItem> terms) {
-        if (item instanceof NotItem) {
-            getPositiveTerms(((NotItem) item).getPositiveItem(), terms);
-        } else if (item instanceof PhraseItem) {
-            PhraseItem pItem = (PhraseItem)item;
-            terms.add(pItem);
-        } else if (item instanceof CompositeItem) {
-            for (Iterator<Item> i = ((CompositeItem) item).getItemIterator(); i.hasNext();) {
+        if (item instanceof NotItem notItem) {
+            getPositiveTerms(notItem.getPositiveItem(), terms);
+        } else if (item instanceof PhraseItem phraseItem) {
+            terms.add(phraseItem);
+        } else if (item instanceof CompositeItem compositeItem) {
+            for (Iterator<Item> i = compositeItem.getItemIterator(); i.hasNext();) {
                 getPositiveTerms(i.next(), terms);
             }
-        } else if (item instanceof TermItem) {
-            terms.add((TermItem)item);
+        } else if (item instanceof TermItem termItem) {
+            terms.add(termItem);
         }
     }
 
@@ -203,8 +198,7 @@ public class QueryTree extends CompositeItem {
 
     private int countItemsRecursively(Item item) {
         int children = 0;
-        if (item instanceof CompositeItem) {
-            CompositeItem composite = (CompositeItem)item;
+        if (item instanceof CompositeItem composite) {
             for (ListIterator<Item> i = composite.getItemIterator(); i.hasNext(); ) {
                 children += countItemsRecursively(i.next());
             }