From c6f2cc2fa7dd32a94e724440b2893dd4b3af7561 Mon Sep 17 00:00:00 2001
From: Zack <zachary.siegel@tecton.ai>
Date: Sun, 1 Dec 2024 23:39:02 -0500
Subject: [PATCH] benchmark

---
 .vscode/settings.json                         |  3 +
 build.gradle                                  |  1 +
 .../ai/tecton/client/TectonClientTest.java    | 70 +++++++++++++++++++
 3 files changed, 74 insertions(+)
 create mode 100644 .vscode/settings.json

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 00000000..04cd6188
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,3 @@
+{
+    "java.configuration.updateBuildConfiguration": "automatic"
+}
diff --git a/build.gradle b/build.gradle
index 83e60b90..ff71ec4e 100644
--- a/build.gradle
+++ b/build.gradle
@@ -61,6 +61,7 @@ compileJava.dependsOn generateVersionFile
 
 // Run tests during build and output status
 test {
+    jvmArgs '-XX:ActiveProcessorCount=2'
     useJUnit()
     testLogging {
         showStandardStreams = true
diff --git a/src/test/java/ai/tecton/client/TectonClientTest.java b/src/test/java/ai/tecton/client/TectonClientTest.java
index 9190092a..4aeb1550 100644
--- a/src/test/java/ai/tecton/client/TectonClientTest.java
+++ b/src/test/java/ai/tecton/client/TectonClientTest.java
@@ -326,6 +326,76 @@ public void testParallelGetFeaturesCall() {
     Assert.assertFalse(batchResponse.getBatchSloInformation().isPresent());
   }
 
+  @Test
+  public void benchmarkParallelGetFeaturesCall() {
+    int ITERATIONS = 50;
+
+    // Enqueue mock responses
+    for (int i = 0; i < ITERATIONS; i++) {
+      sampleResponses.forEach(
+          sampleResponse ->
+              mockWebServer.enqueue(
+                  new MockResponse().setResponseCode(200).setBody(sampleResponse)));
+    }
+
+    // Create request with a large number of request data in the list
+    GetFeaturesBatchRequest.Builder batchRequestBuilder =
+        new GetFeaturesBatchRequest.Builder()
+            .featureServiceName("fraud_detection_feature_service")
+            .workspaceName("prod")
+            .metadataOptions(RequestConstants.ALL_METADATA_OPTIONS);
+
+    for (int i = 0; i < ITERATIONS; i++) {
+      batchRequestBuilder
+          .addRequestData(
+              new GetFeaturesRequestData()
+                  .addJoinKey("user_id", "1_" + i)
+                  .addRequestContext("amount", 55.0))
+          .addRequestData(
+              new GetFeaturesRequestData()
+                  .addJoinKey("user_id", "2_" + i)
+                  .addRequestContext("amount", 155.0))
+          .addRequestData(
+              new GetFeaturesRequestData()
+                  .addJoinKey("user_id", "3" + i)
+                  .addRequestContext("amount", 1000.0));
+    }
+    GetFeaturesBatchRequest batchRequest = batchRequestBuilder.build();
+
+    // Measure start time
+    long startTime = System.nanoTime();
+
+    // Send request and receive response
+    GetFeaturesBatchResponse batchResponse = tectonClient.getFeaturesBatch(batchRequest);
+    List<GetFeaturesResponse> responseList = batchResponse.getBatchResponseList();
+
+    // Measure end time
+    long endTime = System.nanoTime();
+
+    // Calculate elapsed time in milliseconds
+    long elapsedTime = (endTime - startTime) / 1_000_000;
+    System.out.println("Elapsed time: " + elapsedTime + " ms");
+
+    // Verify the responses
+    Assert.assertEquals(3 * ITERATIONS, responseList.size());
+
+    IntStream.range(0, sampleResponses.size())
+        .forEach(
+            i -> {
+              GetFeaturesResponse getFeaturesResponse = responseList.get(i);
+              Assert.assertEquals(14, getFeaturesResponse.getFeatureValues().size());
+              for (FeatureValue value : getFeaturesResponse.getFeatureValues()) {
+                Assert.assertTrue(value.getFeatureStatus().isPresent());
+                Assert.assertTrue(
+                    value.getFeatureStatus().get() == FeatureStatus.PRESENT
+                        || value.getFeatureStatus().get() == FeatureStatus.CACHED_PRESENT);
+              }
+              Assert.assertTrue(getFeaturesResponse.getSloInformation().isPresent());
+            });
+
+    Assert.assertFalse(batchResponse.getBatchSloInformation().isPresent());
+  }
+
   @Test
   public void testGetFeaturesBatchCall() throws IOException {
     sampleBatchResponses.forEach(