From 4cbb04cb94a497e5103aec73e37252ce0515c960 Mon Sep 17 00:00:00 2001 From: Zack Date: Sun, 1 Dec 2024 23:07:40 -0500 Subject: [PATCH 1/2] Use tags to assign responses to result list --- .../client/transport/TectonHttpClient.java | 44 ++++++++++--------- .../transport/TectonHttpClientTest.java | 2 +- 2 files changed, 24 insertions(+), 22 deletions(-) diff --git a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java b/src/main/java/ai/tecton/client/transport/TectonHttpClient.java index d0330c90..cfe0e22d 100644 --- a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java +++ b/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 performParallelRequests( // Map request body to OkHttp Request // ordering of requests is maintained - List requestList = - requestBodyList.stream() - .map( - requestBody -> - new HttpRequest(url.url().toString(), endpoint, method, apiKey, requestBody)) - .map(this::buildRequestWithDefaultHeaders) - .collect(Collectors.toList()); + List 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(); } diff --git a/src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java b/src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java index 034c8fe7..1d95b2f6 100644 --- a/src/test/java/ai/tecton/client/transport/TectonHttpClientTest.java +++ b/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()); From 77900604f9cee95b78f78ac1dfe48407f26e4147 Mon Sep 17 00:00:00 2001 From: Zack Date: Mon, 2 Dec 2024 00:06:43 -0500 Subject: [PATCH 2/2] Remove RequestIndex and just use Object as the tag key --- .../tecton/client/transport/TectonHttpClient.java | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java b/src/main/java/ai/tecton/client/transport/TectonHttpClient.java index cfe0e22d..ed6dbfc0 100644 --- a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java +++ b/src/main/java/ai/tecton/client/transport/TectonHttpClient.java @@ -31,16 +31,6 @@ import org.apache.commons.lang3.Validate; 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; @@ -148,8 +138,7 @@ public void onResponse(Call call, Response response) { try (ResponseBody responseBody = response.body()) { // Add response to corresponding index parallelCallHandler.set( - call.request().tag(RequestIndex.class).index, - new HttpResponse(response, responseBody)); + (Integer) call.request().tag(), new HttpResponse(response, responseBody)); } catch (Exception e) { throw new TectonServiceException(e.getMessage()); } finally { @@ -194,7 +183,7 @@ public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest, int index 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)); + requestBuilder.tag(index); return requestBuilder.build(); }