From c34dd31b97b1c138050de5c87b854b0532baffdd Mon Sep 17 00:00:00 2001 From: Manoj Nagarajan Date: Sun, 25 Aug 2024 23:31:22 -0700 Subject: [PATCH] cleanup --- .../fastclient/InternalAvroStoreClient.java | 64 ++++++++----------- .../StatsAvroGenericStoreClient.java | 36 +---------- 2 files changed, 29 insertions(+), 71 deletions(-) diff --git a/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/InternalAvroStoreClient.java b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/InternalAvroStoreClient.java index 0f08d454b8..dd39401ee8 100644 --- a/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/InternalAvroStoreClient.java +++ b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/InternalAvroStoreClient.java @@ -45,18 +45,27 @@ public final CompletableFuture> batchGet(Set 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 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. @@ -64,19 +73,20 @@ Throwable checkBatchGetPartialFailure(MultiKeyRequestContext multiKeyReque 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> batchGet(BatchGetRequestContext requestContext, Set keys) @@ -88,32 +98,10 @@ protected CompletableFuture> batchGet(BatchGetRequestContext req if (throwable != null) { resultFuture.completeExceptionally(throwable); } else { - Optional 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/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/StatsAvroGenericStoreClient.java b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/StatsAvroGenericStoreClient.java index abf62ad70e..4fc6583c6d 100644 --- a/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/StatsAvroGenericStoreClient.java +++ b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/StatsAvroGenericStoreClient.java @@ -128,39 +128,9 @@ private CompletableFuture 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; } } }