Use tags to assign responses to result list

Author: zack13532

src/main/java/ai/tecton/client/transport/TectonHttpClient.java

@@ -17,7 +17,6 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.stream.Collectors;
 import okhttp3.Call;
 import okhttp3.Callback;
 import okhttp3.ConnectionPool;
@@ -33,6 +32,15 @@
 
 public class TectonHttpClient {
 
+  // Class to be used as a tag for each request to keep track of the index of the request.
+  static class RequestIndex {
+    final int index;
+
+    RequestIndex(int index) {
+      this.index = index;
+    }
+  }
+
   private HttpUrl url;
   private final String apiKey;
   private final OkHttpClient client;
@@ -96,7 +104,7 @@ public boolean isClosed() {
   public HttpResponse performRequest(String endpoint, HttpMethod method, String requestBody) {
     HttpRequest httpRequest =
         new HttpRequest(url.url().toString(), endpoint, method, apiKey, requestBody);
-    Request request = buildRequestWithDefaultHeaders(httpRequest);
+    Request request = buildRequestWithDefaultHeaders(httpRequest, 0);
     Call call = client.newCall(request);
     try (Response response = call.execute()) {
       return new HttpResponse(response);
@@ -113,13 +121,12 @@ public List<HttpResponse> performParallelRequests(
 
     // Map request body to OkHttp Request
     // ordering of requests is maintained
-    List<Request> requestList =
-        requestBodyList.stream()
-            .map(
-                requestBody ->
-                    new HttpRequest(url.url().toString(), endpoint, method, apiKey, requestBody))
-            .map(this::buildRequestWithDefaultHeaders)
-            .collect(Collectors.toList());
+    List<Request> requestList = new ArrayList<>();
+    for (int i = 0; i < requestBodyList.size(); i++) {
+      HttpRequest httpRequest =
+          new HttpRequest(url.url().toString(), endpoint, method, apiKey, requestBodyList.get(i));
+      requestList.add(buildRequestWithDefaultHeaders(httpRequest, i));
+    }
 
     // Initialize a countdown latch for numberOfCalls.
     CountDownLatch countDownLatch = new CountDownLatch(requestBodyList.size());
@@ -141,7 +148,8 @@ public void onResponse(Call call, Response response) {
             try (ResponseBody responseBody = response.body()) {
               // Add response to corresponding index
               parallelCallHandler.set(
-                  requestList.indexOf(call.request()), new HttpResponse(response, responseBody));
+                  call.request().tag(RequestIndex.class).index,
+                  new HttpResponse(response, responseBody));
             } catch (Exception e) {
               throw new TectonServiceException(e.getMessage());
             } finally {
@@ -171,7 +179,7 @@ public void onResponse(Call call, Response response) {
     }
   }
 
-  public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest) {
+  public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest, int index) {
     // Construct url
     Request.Builder requestBuilder = new Request.Builder().url(httpRequest.getUrl());
 
@@ -183,16 +191,10 @@ public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest) {
     requestBuilder.header(HttpHeader.AUTHORIZATION.getName(), apiKeyHeader);
 
     // Add request body
-    switch (httpRequest.getMethod()) {
-      case POST:
-      default:
-        {
-          okhttp3.MediaType mediaType =
-              okhttp3.MediaType.parse(MediaType.APPLICATION_JSON.getName());
-          RequestBody requestBody = RequestBody.create(httpRequest.getJsonBody(), mediaType);
-          requestBuilder.post(requestBody);
-        }
-    }
+    okhttp3.MediaType mediaType = okhttp3.MediaType.parse(MediaType.APPLICATION_JSON.getName());
+    RequestBody requestBody = RequestBody.create(httpRequest.getJsonBody(), mediaType);
+    requestBuilder.post(requestBody);
+    requestBuilder.tag(RequestIndex.class, new RequestIndex(index));
     return requestBuilder.build();
   }
 

src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java

@@ -75,7 +75,7 @@ public void testDefaultTectonRequest() {
     HttpRequest httpRequest =
         new HttpRequest(url, endpoint, TectonHttpClient.HttpMethod.POST, apiKey, body);
 
-    Request request = httpClient.buildRequestWithDefaultHeaders(httpRequest);
+    Request request = httpClient.buildRequestWithDefaultHeaders(httpRequest, 0);
     Assert.assertNotNull(request);
 
     Assert.assertEquals(POST, request.method());