cleanup

m-nagarajan · Aug 26, 2024 · c34dd31 · c34dd31
1 parent b85c571
commit c34dd31
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 71 deletions.
diff --git a/...s/venice-client/src/main/java/com/linkedin/venice/fastclient/InternalAvroStoreClient.java b/...s/venice-client/src/main/java/com/linkedin/venice/fastclient/InternalAvroStoreClient.java
@@ -45,38 +45,48 @@ public final CompletableFuture<Map<K, V>> batchGet(Set<K> keys) throws VeniceCli
     return batchGet(new BatchGetRequestContext<>(keys.size(), false), keys);
   }
 
+  /**
+   * Check for partial failures in the multi-key request context and return the exception if any.
+   * @param multiKeyRequestContext
+   * @return
+   */
   Throwable checkBatchGetPartialFailure(MultiKeyRequestContext<K, V> multiKeyRequestContext) {
-    Throwable throwable = null;
+    Throwable partialResponseException = null;
     // check for partial failures for multi-key requests
-    boolean checkOriginalRequestContext = false;
     if (multiKeyRequestContext.retryContext != null
         && multiKeyRequestContext.retryContext.retryRequestContext != null) {
       // retry is triggered
-      if (multiKeyRequestContext.retryContext.retryRequestContext.isCompletedSuccessfullyWithPartialResponse()) {
-        throwable =
+      if (multiKeyRequestContext.isPartialSuccessAllowed) {
+        if (multiKeyRequestContext.retryContext.retryRequestContext.isCompletedSuccessfullyWithPartialResponse()) {
+          partialResponseException =
+              (Throwable) multiKeyRequestContext.retryContext.retryRequestContext.getPartialResponseException().get();
+        }
+      } else if (multiKeyRequestContext.retryContext.retryRequestContext.getPartialResponseException().isPresent()) {
+        partialResponseException =
             (Throwable) multiKeyRequestContext.retryContext.retryRequestContext.getPartialResponseException().get();
       }
-      if (throwable != null) {
+      if (partialResponseException != null) {
         // if there is no exception in the retry request, everything passed, but if there is an exception in the
         // retry request, that failure might have passed in the original request after the retry started. checking
         // the numKeysCompleted for now.
         int totalKeyCount = multiKeyRequestContext.numKeysInRequest;
         int successKeyCount = multiKeyRequestContext.numKeysCompleted.get()
             + multiKeyRequestContext.retryContext.retryRequestContext.numKeysCompleted.get();
         if (successKeyCount >= totalKeyCount) {
-          throwable = null;
+          partialResponseException = null;
         }
       }
     } else {
       // retry not enabled or not triggered: check the original request context
-      checkOriginalRequestContext = true;
-    }
-    if (checkOriginalRequestContext) {
-      if (multiKeyRequestContext.isCompletedSuccessfullyWithPartialResponse()) {
-        throwable = (Throwable) multiKeyRequestContext.getPartialResponseException().get();
+      if (multiKeyRequestContext.isPartialSuccessAllowed) {
+        if (multiKeyRequestContext.isCompletedSuccessfullyWithPartialResponse()) {
+          partialResponseException = (Throwable) multiKeyRequestContext.getPartialResponseException().get();
+        }
+      } else if (multiKeyRequestContext.getPartialResponseException().isPresent()) {
+        partialResponseException = (Throwable) multiKeyRequestContext.getPartialResponseException().get();
       }
     }
-    return throwable;
+    return partialResponseException;
   }
 
   protected CompletableFuture<Map<K, V>> batchGet(BatchGetRequestContext<K, V> requestContext, Set<K> keys)
@@ -88,32 +98,10 @@ protected CompletableFuture<Map<K, V>> batchGet(BatchGetRequestContext<K, V> req
       if (throwable != null) {
         resultFuture.completeExceptionally(throwable);
       } else {
-        Optional<Throwable> partialResponseException = Optional.empty();
-        if (!requestContext.isPartialSuccessAllowed) {
-          if (requestContext.retryContext != null && requestContext.retryContext.retryRequestContext != null) {
-            // retry triggered
-            if (requestContext.retryContext.retryRequestContext.getPartialResponseException().isPresent()) {
-              // if there is no exception in the retry request, everything passed, but if there is an exception in the
-              // retry request, that failure might have passed in the original request after the retry started. checking
-              // the numKeysCompleted for now.
-              int totalKeyCount = requestContext.numKeysInRequest;
-              int successKeyCount = requestContext.numKeysCompleted.get()
-                  + requestContext.retryContext.retryRequestContext.numKeysCompleted.get();
-              if (successKeyCount < totalKeyCount) {
-                partialResponseException =
-                    requestContext.retryContext.retryRequestContext.getPartialResponseException();
-              }
-            }
-          } else {
-            // retry not enabled or not triggered
-            if (requestContext.getPartialResponseException().isPresent()) {
-              partialResponseException = requestContext.getPartialResponseException();
-            }
-          }
-        }
-        if (partialResponseException.isPresent()) {
-          resultFuture.completeExceptionally(
-              new VeniceClientException("Response was not complete", partialResponseException.get()));
+        Throwable partialResponseException = checkBatchGetPartialFailure(requestContext);
+        if (partialResponseException != null) {
+          resultFuture
+              .completeExceptionally(new VeniceClientException("Response was not complete", partialResponseException));
         } else {
           resultFuture.complete(response);
         }

diff --git a/...nice-client/src/main/java/com/linkedin/venice/fastclient/StatsAvroGenericStoreClient.java b/...nice-client/src/main/java/com/linkedin/venice/fastclient/StatsAvroGenericStoreClient.java
@@ -128,39 +128,9 @@ private <R> CompletableFuture<R> recordRequestMetrics(
       } else {
         // check for partial failures for multi-key requests
         if (requestContext instanceof MultiKeyRequestContext) {
-          MultiKeyRequestContext multiKeyRequestContext = (MultiKeyRequestContext) requestContext;
-          boolean checkOriginalRequestContext = false;
-          if (multiKeyRequestContext.retryContext != null
-              && multiKeyRequestContext.retryContext.retryRequestContext != null) {
-            // retry is triggered
-            if (multiKeyRequestContext.retryContext.retryRequestContext.isCompletedSuccessfullyWithPartialResponse()) {
-              exceptionReceived = true;
-              throwable =
-                  (Throwable) multiKeyRequestContext.retryContext.retryRequestContext.getPartialResponseException()
-                      .get();
-            }
-            if (exceptionReceived) {
-              // if there is no exception in the retry request, everything passed, but if there is an exception in the
-              // retry request, that failure might have passed in the original request after the retry started. checking
-              // the numKeysCompleted for now.
-              int totalKeyCount = multiKeyRequestContext.numKeysInRequest;
-              int successKeyCount = multiKeyRequestContext.numKeysCompleted.get()
-                  + multiKeyRequestContext.retryContext.retryRequestContext.numKeysCompleted.get();
-              if (successKeyCount >= totalKeyCount) {
-                exceptionReceived = false;
-              } else {
-                checkOriginalRequestContext = true;
-              }
-            }
-          } else {
-            // retry not enabled or not triggered: check the original request context
-            checkOriginalRequestContext = true;
-          }
-          if (checkOriginalRequestContext) {
-            if (multiKeyRequestContext.isCompletedSuccessfullyWithPartialResponse()) {
-              exceptionReceived = true;
-              throwable = (Throwable) multiKeyRequestContext.getPartialResponseException().get();
-            }
+          throwable = checkBatchGetPartialFailure((MultiKeyRequestContext) requestContext);
+          if (throwable != null) {
+            exceptionReceived = true;
           }
         }
       }