From 5373e80e7c8207ae7c96b2d2cab8cfc11397c4dd Mon Sep 17 00:00:00 2001
From: xdark <me@xdark.dev>
Date: Sat, 1 Apr 2023 03:38:07 +0300
Subject: [PATCH] Make an optimization for BatchingClient

---
 .../twitter/RetrieveCardBatchedStage.java     |  4 +--
 .../pipeline/util/BatchingClient.java         | 28 +++++++++----------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/src/java/com/twitter/search/ingester/pipeline/twitter/RetrieveCardBatchedStage.java b/src/java/com/twitter/search/ingester/pipeline/twitter/RetrieveCardBatchedStage.java
index 705c211c5..7c956d6fc 100644
--- a/src/java/com/twitter/search/ingester/pipeline/twitter/RetrieveCardBatchedStage.java
+++ b/src/java/com/twitter/search/ingester/pipeline/twitter/RetrieveCardBatchedStage.java
@@ -134,7 +134,7 @@ public class RetrieveCardBatchedStage extends TwitterBaseStage
         }));
   }
 
-  private Future<Map<Long, Card2>> batchRetrieveURLs(Set<Long> keys) {
+  private Future<Map<Long, Card2>> batchRetrieveURLs(List<Long> keys) {
     retrieveCardsTimer.start();
     totalTweets.increment(keys.size());
 
@@ -145,7 +145,7 @@ public class RetrieveCardBatchedStage extends TwitterBaseStage
 
     GetTweetsRequest request = new GetTweetsRequest()
         .setOptions(options)
-        .setTweet_ids(new ArrayList<>(keys));
+        .setTweet_ids(keys);
 
     return tweetyPieService.get_tweets(request)
         .onFailure(throwable -> {
diff --git a/src/java/com/twitter/search/ingester/pipeline/util/BatchingClient.java b/src/java/com/twitter/search/ingester/pipeline/util/BatchingClient.java
index 222c6f544..abcf82076 100644
--- a/src/java/com/twitter/search/ingester/pipeline/util/BatchingClient.java
+++ b/src/java/com/twitter/search/ingester/pipeline/util/BatchingClient.java
@@ -2,7 +2,7 @@ package com.twitter.search.ingester.pipeline.util;
 
 import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
+import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 
 import com.google.common.collect.Sets;
@@ -20,7 +20,7 @@ public class BatchingClient<RQ, RP> {
     /**
      * Issue a request to the underlying store which supports batches of requests.
      */
-    Future<Map<RQ, RP>> batchGet(Set<RQ> requests);
+    Future<Map<RQ, RP>> batchGet(List<RQ> requests);
   }
 
   /**
@@ -54,7 +54,7 @@ public class BatchingClient<RQ, RP> {
   }
 
   private void maybeBatchCall(RQ request) {
-    Set<RQ> frozenRequests;
+    List<RQ> frozenRequests;
     synchronized (unsentRequests) {
       unsentRequests.add(request);
       if (unsentRequests.size() < batchSize) {
@@ -63,27 +63,25 @@ public class BatchingClient<RQ, RP> {
 
       // Make a copy of requests so we can modify it inside executeBatchCall without additional
       // synchronization.
-      frozenRequests = new HashSet<>(unsentRequests);
+      frozenRequests = Arrays.asList((RQ[]) unsentRequests.toArray());
       unsentRequests.clear();
     }
 
     executeBatchCall(frozenRequests);
   }
 
-  private void executeBatchCall(Set<RQ> requests) {
+  private void executeBatchCall(List<RQ> requests) {
     batchClient.batchGet(requests)
         .onSuccess(responseMap -> {
-          for (Map.Entry<RQ, RP> entry : responseMap.entrySet()) {
-            Promise<RP> promise = promises.remove(entry.getKey());
-            if (promise != null) {
-              promise.become(Future.value(entry.getValue()));
-            }
-          }
-
-          Set<RQ> outstandingRequests = Sets.difference(requests, responseMap.keySet());
-          for (RQ request : outstandingRequests) {
+          for (RQ request : requests) {
             Promise<RP> promise = promises.remove(request);
-            if (promise != null) {
+            if (promise == null) {
+              continue;
+            }
+            RP response = responseMap.get(request);
+            if (response != null) {
+              promise.become(Future.value(response));
+            } else {
               promise.become(Future.exception(new ResponseNotReturnedException(request)));
             }
           }