diff --git a/build.gradle b/build.gradle index 23ef82d2ab..68614beb80 100644 --- a/build.gradle +++ b/build.gradle @@ -84,6 +84,7 @@ ext.libraries = [ grpcProtobuf: "io.grpc:grpc-protobuf:${grpcVersion}", grpcServices: "io.grpc:grpc-services:${grpcVersion}", grpcStub: "io.grpc:grpc-stub:${grpcVersion}", + grpcTesting: "io.grpc:grpc-testing:${grpcVersion}", hadoopCommon: "org.apache.hadoop:hadoop-common:${hadoopVersion}", httpAsyncClient: 'org.apache.httpcomponents:httpasyncclient:4.1.5', httpClient5: 'org.apache.httpcomponents.client5:httpclient5:5.3', @@ -486,6 +487,12 @@ subprojects { doLast { parseJacocoXml("$buildDir/reports/jacoco/test/jacocoTestReport.xml") } +// This is required to remove generated code from the report but it fails build in all-modules +// afterEvaluate { +// classDirectories.setFrom(files(classDirectories.files.collect { fileTree(dir: it, exclude: [ +// '**/com/linkedin/venice/protocols/**', +// ])})) +// } } afterEvaluate { @@ -501,6 +508,12 @@ subprojects { value = 'COVEREDRATIO' minimum = threshold } + + afterEvaluate { + classDirectories.setFrom(files(classDirectories.files.collect { fileTree(dir: it, exclude: [ + '**/com/linkedin/venice/protocols/**', + ])})) + } } } } diff --git a/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java index 30029701f2..07c48d9998 100644 --- a/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java +++ b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java @@ -267,17 +267,24 @@ public VeniceGrpcStreamObserver(CompletableFuture respo @Override public void onNext(VeniceServerResponse value) { - if (value.getErrorCode() != VeniceReadResponseStatus.OK) { - handleResponseError(value); + int statusCode = value.getErrorCode(); + // Successful response + if (statusCode == VeniceReadResponseStatus.OK.getCode()) { + complete( + new TransportClientResponse( + value.getSchemaId(), + CompressionStrategy.valueOf(value.getCompressionStrategy()), + value.getData().toByteArray()), + null); return; } - - complete( - new TransportClientResponse( - value.getSchemaId(), - CompressionStrategy.valueOf(value.getCompressionStrategy()), - value.getData().toByteArray()), - null); + // Key not found is a valid response + if (statusCode == VeniceReadResponseStatus.KEY_NOT_FOUND.getCode()) { + complete(null, null); + return; + } + // Handle the cases where the status code doesn't match healthy response codes + handleResponseError(value); } @Override @@ -308,30 +315,29 @@ void handleResponseError(VeniceServerResponse response) { int statusCode = response.getErrorCode(); String errorMessage = response.getErrorMessage(); Exception exception; - - switch (statusCode) { - case VeniceReadResponseStatus.BAD_REQUEST: - exception = new VeniceClientHttpException(errorMessage, statusCode); - break; - case VeniceReadResponseStatus.TOO_MANY_REQUESTS: - exception = new VeniceClientRateExceededException(errorMessage); - break; - case VeniceReadResponseStatus.KEY_NOT_FOUND: - exception = null; - break; - default: - exception = new VeniceClientException( - String - .format("An unexpected error occurred with status code: %d, message: %s", statusCode, errorMessage)); - break; - } - - if (exception != null) { - LOGGER.error("Got error in response due to", exception); + try { + switch (VeniceReadResponseStatus.fromCode(statusCode)) { + case BAD_REQUEST: + exception = new VeniceClientHttpException(errorMessage, statusCode); + break; + case TOO_MANY_REQUESTS: + exception = new VeniceClientRateExceededException(errorMessage); + break; + default: + exception = new VeniceClientException( + String.format( + "An unexpected error occurred with status code: %d, message: %s", + statusCode, + errorMessage)); + break; + } + } catch (IllegalArgumentException e) { + // Handle the case where the status code doesn't match any known values + exception = new VeniceClientException( + String.format("Unknown status code: %d, message: %s", statusCode, errorMessage), + e); } - - // In the event of record not found, we treat that as a successful response and complete the future with a null - // value and the exception is set to null as well. + LOGGER.error("Received error response with status code: {}, message: {}", statusCode, errorMessage); complete(null, exception); } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java index 44020f9eed..a13a01c683 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java @@ -1,5 +1,6 @@ package com.linkedin.venice.grpc; +import com.google.protobuf.ByteString; import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.security.SSLConfig; import com.linkedin.venice.security.SSLFactory; @@ -7,6 +8,7 @@ import io.grpc.Grpc; import io.grpc.ServerCall; import io.grpc.Status; +import io.netty.buffer.ByteBuf; import io.netty.handler.codec.http.HttpResponseStatus; import java.io.IOException; import java.io.InputStream; @@ -92,4 +94,23 @@ private static KeyStore loadStore(String path, char[] password, String type) } return keyStore; } + + /** + * Converts a Netty ByteBuf to a ByteString, checking if it has a backing array to avoid manual copying. + * + * @param body The ByteBuf to be converted to ByteString. + * @return The resulting ByteString. + */ + public static ByteString toByteString(ByteBuf body) { + if (body.hasArray()) { + // Directly use the backing array to avoid copying + return ByteString.copyFrom(body.array(), body.arrayOffset() + body.readerIndex(), body.readableBytes()); + } + // Fallback to nioBuffer() to handle the conversion efficiently + return ByteString.copyFrom(body.nioBuffer()); + } + + public static ByteString toByteString(byte[] bytes) { + return ByteString.copyFrom(bytes); + } } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/listener/response/HttpShortcutResponse.java b/internal/venice-common/src/main/java/com/linkedin/venice/listener/response/HttpShortcutResponse.java index 3fb6892b28..56d1ee60cc 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/listener/response/HttpShortcutResponse.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/listener/response/HttpShortcutResponse.java @@ -10,8 +10,6 @@ public class HttpShortcutResponse { private final String message; private final HttpResponseStatus status; - private boolean misroutedStoreVersion = false; - public HttpShortcutResponse(String message, HttpResponseStatus status) { this.message = message; this.status = status; @@ -28,12 +26,4 @@ public String getMessage() { public HttpResponseStatus getStatus() { return status; } - - public boolean isMisroutedStoreVersion() { - return misroutedStoreVersion; - } - - public void setMisroutedStoreVersion(boolean misroutedStoreVersion) { - this.misroutedStoreVersion = misroutedStoreVersion; - } } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/ServerAdminAction.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/ServerAdminAction.java index 91ad4004b3..c5c8ec899c 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/meta/ServerAdminAction.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/ServerAdminAction.java @@ -1,8 +1,19 @@ package com.linkedin.venice.meta; +import java.util.HashMap; +import java.util.Map; + + public enum ServerAdminAction { DUMP_INGESTION_STATE(0), DUMP_SERVER_CONFIGS(1); + private static final Map ADMIN_ACTION_MAP = new HashMap<>(2); + + static { + for (ServerAdminAction action: values()) { + ADMIN_ACTION_MAP.put(action.getValue(), action); + } + } private final int value; ServerAdminAction(int value) { @@ -12,4 +23,12 @@ public enum ServerAdminAction { public int getValue() { return this.value; } + + public static ServerAdminAction fromValue(int value) { + ServerAdminAction action = ADMIN_ACTION_MAP.get(value); + if (action == null) { + throw new IllegalArgumentException("Unknown server admin action value: " + value); + } + return action; + } } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/response/VeniceReadResponseStatus.java b/internal/venice-common/src/main/java/com/linkedin/venice/response/VeniceReadResponseStatus.java index 1ed20a5489..80e87bb49c 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/response/VeniceReadResponseStatus.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/response/VeniceReadResponseStatus.java @@ -1,19 +1,49 @@ package com.linkedin.venice.response; +import io.netty.handler.codec.http.HttpResponseStatus; +import java.util.HashMap; +import java.util.Map; + + /** - * Enumeration of response status codes for Venice read requests. - *

- * **Positive values** correspond to standard HTTP status codes and can be used directly in HTTP responses. - * **Negative values** represent custom Venice-specific error codes. - *

- * For example, a status code of `200` indicates a successful read, while a status code of `-100` might indicate a specific Venice-related error. + * Defines response status codes for Venice read requests. This wrapper around {@link HttpResponseStatus} allows + * for the inclusion of custom status codes that extend beyond the standard HTTP status codes. */ -public class VeniceReadResponseStatus { - public static final int KEY_NOT_FOUND = -420; - - public static final int OK = 200; - public static final int BAD_REQUEST = 400; - public static final int INTERNAL_ERROR = 500; - public static final int TOO_MANY_REQUESTS = 429; - public static final int SERVICE_UNAVAILABLE = 503; +public enum VeniceReadResponseStatus { + KEY_NOT_FOUND(HttpResponseStatus.NOT_FOUND), OK(HttpResponseStatus.OK), BAD_REQUEST(HttpResponseStatus.BAD_REQUEST), + FORBIDDEN(HttpResponseStatus.FORBIDDEN), METHOD_NOT_ALLOWED(HttpResponseStatus.METHOD_NOT_ALLOWED), + REQUEST_TIMEOUT(HttpResponseStatus.REQUEST_TIMEOUT), TOO_MANY_REQUESTS(HttpResponseStatus.TOO_MANY_REQUESTS), + INTERNAL_SERVER_ERROR(HttpResponseStatus.INTERNAL_SERVER_ERROR), + SERVICE_UNAVAILABLE(HttpResponseStatus.SERVICE_UNAVAILABLE), + MISROUTED_STORE_VERSION(new HttpResponseStatus(570, "Misrouted request")); + + private static final Map STATUS_MAP = new HashMap<>(16); + + static { + for (VeniceReadResponseStatus status: values()) { + STATUS_MAP.put(status.getCode(), status); + } + } + + private final HttpResponseStatus httpResponseStatus; + + VeniceReadResponseStatus(HttpResponseStatus httpResponseStatus) { + this.httpResponseStatus = httpResponseStatus; + } + + public HttpResponseStatus getHttpResponseStatus() { + return httpResponseStatus; + } + + public int getCode() { + return httpResponseStatus.code(); + } + + public static VeniceReadResponseStatus fromCode(int code) { + VeniceReadResponseStatus status = STATUS_MAP.get(code); + if (status == null) { + throw new IllegalArgumentException("Unknown status venice read response status code: " + code); + } + return status; + } } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/utils/NettyUtils.java b/internal/venice-common/src/main/java/com/linkedin/venice/utils/NettyUtils.java index 8ea7c30379..8dd0d0a40e 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/utils/NettyUtils.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/utils/NettyUtils.java @@ -10,6 +10,7 @@ import io.netty.channel.ChannelHandlerContext; import io.netty.handler.codec.http.DefaultFullHttpResponse; import io.netty.handler.codec.http.FullHttpResponse; +import io.netty.handler.codec.http.HttpRequest; import io.netty.handler.codec.http.HttpResponseStatus; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -38,4 +39,8 @@ public static void setupResponseAndFlush( response.headers().set(CONTENT_LENGTH, response.content().readableBytes()); ctx.writeAndFlush(response); } + + public static boolean containRetryHeader(HttpRequest request) { + return request.headers().contains(HttpConstants.VENICE_RETRY); + } } diff --git a/internal/venice-common/src/main/proto/VeniceReadService.proto b/internal/venice-common/src/main/proto/VeniceReadService.proto index 8116718807..c12e2117d1 100644 --- a/internal/venice-common/src/main/proto/VeniceReadService.proto +++ b/internal/venice-common/src/main/proto/VeniceReadService.proto @@ -1,11 +1,36 @@ syntax = 'proto3'; package com.linkedin.venice.protocols; +import "google/protobuf/empty.proto"; option java_multiple_files = true; service VeniceReadService { - rpc get (VeniceClientRequest) returns (VeniceServerResponse) {} - rpc batchGet(VeniceClientRequest) returns (VeniceServerResponse) {} + + rpc get (VeniceClientRequest) returns (VeniceServerResponse) { + option deprecated = true; + } + + rpc batchGet(VeniceClientRequest) returns (VeniceServerResponse) { + option deprecated = true; + } + + rpc singleGet(SingleGetRequest) returns (SingleGetResponse) {} + + rpc multiGet(MultiGetRequest) returns (MultiKeyResponse) {} + + rpc compute(ComputeRequest) returns (MultiKeyResponse) {} + + rpc isServerHealthy(HealthCheckRequest) returns (HealthCheckResponse) {} + + rpc getCompressionDictionary(CompressionDictionaryRequest) returns (CompressionDictionaryResponse) {} + + rpc handleAdminRequest(AdminRequest) returns (AdminResponse) {} + + rpc getMetadata(MetadataRequest) returns (MetadataResponse) {} + + rpc getCurrentVersionInfo(CurrentVersionInfoRequest) returns (CurrentVersionInfoResponse) {} + + rpc getIngestionContext(IngestionContextRequest) returns (IngestionContextResponse) {} } message VeniceClientRequest { @@ -29,4 +54,131 @@ message VeniceServerResponse { uint32 errorCode = 6; string errorMessage = 7; +} + + +/* + * Note: The following message formats will be evolved in the future. The current format is used + * for the initial implementation and help refactor the existing code. + */ +message SingleGetRequest { + string resourceName = 1; + uint32 partition = 2; + string key = 3; + bool isRetryRequest = 4; +} + +message SingleGetResponse { + int32 statusCode = 1; + bytes value = 2; + sint32 schemaId = 3; + uint32 compressionStrategy = 4; + optional string errorMessage = 5; + string contentType = 6; + uint32 contentLength = 7; + uint32 rcu = 8; +} + +message MultiGetRequest { + string resourceName = 1; + uint32 partition = 2; + bytes keyBytes = 3; // used for batch get + uint32 keyCount = 4; + bool isRetryRequest = 5; +} + +message MultiKeyResponse { + int32 statusCode = 1; + bytes value = 2; + sint32 schemaId = 3; + uint32 compressionStrategy = 4; + optional string errorMessage = 5; + string contentType = 6; + uint32 contentLength = 7; + uint32 rcu = 8; +} + + +message ComputeRequest { + string resourceName = 1; + bytes payload = 2; + sint32 schemaId = 3; + bool isRetryRequest = 4; + bool isStreamingRequest = 5; + uint32 apiVersion = 6; + sint32 computeValueSchemaId = 7; +} + +message HealthCheckRequest { +} + +message HealthCheckResponse { + int32 statusCode = 1; + string message = 2; +} + +message CompressionDictionaryRequest { + string storeName = 1; + uint32 storeVersion = 2; +} + +message CompressionDictionaryResponse { + int32 statusCode = 1; + bytes value = 2; + string contentType = 3; + uint32 contentLength = 4; + string errorMessage = 5; +} + +message AdminRequest { + string resourceName = 1; + optional uint32 partition = 2; + string serverAdminAction = 3; +} + +message AdminResponse { + int32 statusCode = 1; + bytes value = 2; + sint32 schemaId = 3; + string contentType = 4; + uint32 contentLength = 5; + string errorMessage = 6; +} + +message CurrentVersionInfoRequest { + string storeName = 1; +} + +message CurrentVersionInfoResponse { + int32 statusCode = 1; + sint32 currentVersion = 2; + string errorMessage = 3; + string contentType = 4; +} + +message MetadataRequest { + string storeName = 1; +} + +message MetadataResponse { + int32 statusCode = 1; + bytes value = 2; + sint32 schemaId = 3; + string contentType = 4; + uint32 contentLength = 5; + string errorMessage = 6; +} + +message IngestionContextRequest { + string versionTopicName = 1; + string topicName = 2; + uint32 partition = 3; +} + +message IngestionContextResponse { + int32 statusCode = 1; + bytes value = 2; + string contentType = 3; + uint32 contentLength = 4; + string errorMessage = 5; } \ No newline at end of file diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java index e877b4e480..2a4f1615a5 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java @@ -1,12 +1,19 @@ package com.linkedin.venice.grpc; -import static org.mockito.Mockito.*; -import static org.testng.Assert.*; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; +import com.google.protobuf.ByteString; import com.linkedin.venice.security.SSLFactory; import com.linkedin.venice.utils.SslUtils; import io.grpc.Status; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.buffer.Unpooled; import io.netty.handler.codec.http.HttpResponseStatus; +import java.nio.ByteBuffer; import javax.net.ssl.KeyManager; import javax.net.ssl.TrustManager; import org.testng.annotations.BeforeTest; @@ -71,4 +78,98 @@ public void testHttpResponseStatusToGrpcStatus() { grpcStatus.getDescription(), "Mismatch in error description for the mapped grpc status"); } + + @Test + public void testToByteStringWithArrayBacking() { + byte[] byteArray = new byte[] { 1, 2, 3, 4, 5 }; + ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray); + assertTrue(byteBuf.hasArray()); + + ByteString byteString = GrpcUtils.toByteString(byteBuf); + assertEquals(byteString.toByteArray(), byteArray); + } + + @Test + public void testToByteStringWithArrayBackingAndReaderIndex() { + byte[] byteArray = new byte[] { 1, 2, 3, 4, 5 }; + ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray); + byteBuf.readerIndex(2); // Set reader index to 2 + + assertTrue(byteBuf.hasArray()); + ByteString byteString = GrpcUtils.toByteString(byteBuf); + + byte[] expectedArray = new byte[] { 3, 4, 5 }; // Expected array after reader index + assertEquals(byteString.toByteArray(), expectedArray); + } + + @Test + public void testToByteStringWithEmptyByteBuf() { + ByteBuf byteBuf = Unpooled.EMPTY_BUFFER; + + ByteString byteString = GrpcUtils.toByteString(byteBuf); + + assertTrue(byteString.isEmpty()); + } + + @Test + public void testToByteStringWithNonArrayBacking() { + // Create a direct ByteBuf with an initial capacity of 16 bytes + ByteBuf byteBuf = Unpooled.directBuffer(16); + // Write some data into the ByteBuf + byteBuf.writeBytes(new byte[] { 1, 2, 3, 4, 5 }); + assertFalse(byteBuf.hasArray()); + + ByteString byteString = GrpcUtils.toByteString(byteBuf); + + byte[] expectedArray = new byte[] { 1, 2, 3, 4, 5 }; + assertEquals(byteString.toByteArray(), expectedArray); + } + + @Test + public void testToByteStringWithDifferentOffsetsAndLengths() { + byte[] byteArray = new byte[] { 10, 20, 30, 40, 50 }; + ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray); + byteBuf.readerIndex(1); // Start from index 1 + byteBuf.writerIndex(4); // End at index 4 + + ByteString byteString = GrpcUtils.toByteString(byteBuf); + + byte[] expectedArray = new byte[] { 20, 30, 40 }; // Expected byte array + assertEquals(byteString.toByteArray(), expectedArray); + } + + @Test + public void testToByteStringWithNioBuffer() { + byte[] byteArray = new byte[] { 100, (byte) 200, (byte) 255 }; + ByteBuf byteBuf = Unpooled.wrappedBuffer(byteArray); + + ByteBuffer nioBuffer = byteBuf.nioBuffer(); // Get NIO buffer + ByteString byteString = GrpcUtils.toByteString(Unpooled.wrappedBuffer(nioBuffer)); + + assertEquals(byteString.toByteArray(), byteArray); + } + + @Test + public void testToByteStringWithCompositeByteBuf() { + // Create some ByteBufs + ByteBuf buffer1 = Unpooled.wrappedBuffer(new byte[] { 10, 20 }); + ByteBuf buffer2 = Unpooled.wrappedBuffer(new byte[] { 30, 40, 50 }); + + // Create a CompositeByteBuf that combines the above ByteBufs + CompositeByteBuf compositeByteBuf = Unpooled.compositeBuffer(); + compositeByteBuf.addComponents(true, buffer1, buffer2); + + // The expected ByteString combining all bytes from buffer1 and buffer2 + byte[] expectedBytes = new byte[] { 10, 20, 30, 40, 50 }; + + // Verify that the CompositeByteBuf has the combined array + byte[] combinedArray = new byte[5]; + compositeByteBuf.getBytes(0, combinedArray); + assertEquals(combinedArray, expectedBytes); + + // Convert CompositeByteBuf to ByteString using the method + ByteString byteString = GrpcUtils.toByteString(compositeByteBuf); + + assertEquals(byteString.toByteArray(), expectedBytes); + } } diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/listener/response/HttpShortcutResponseTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/listener/response/HttpShortcutResponseTest.java new file mode 100644 index 0000000000..3630ec5f90 --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/listener/response/HttpShortcutResponseTest.java @@ -0,0 +1,27 @@ +package com.linkedin.venice.listener.response; + +import static org.testng.Assert.assertEquals; + +import io.netty.handler.codec.http.HttpResponseStatus; +import org.testng.annotations.Test; + + +public class HttpShortcutResponseTest { + @Test + public void testHttpShortcutResponse() { + // Case 1: Test constructor with both message and status + HttpShortcutResponse responseWithMessage = new HttpShortcutResponse("Success", HttpResponseStatus.OK); + assertEquals(responseWithMessage.getMessage(), "Success", "Message should be 'Success'"); + assertEquals(responseWithMessage.getStatus(), HttpResponseStatus.OK, "Status should be OK"); + + // Case 2: Test constructor with empty message and status + HttpShortcutResponse responseWithEmptyMessage = new HttpShortcutResponse("", HttpResponseStatus.BAD_REQUEST); + assertEquals(responseWithEmptyMessage.getMessage(), "", "Message should be empty"); + assertEquals(responseWithEmptyMessage.getStatus(), HttpResponseStatus.BAD_REQUEST, "Status should be BAD_REQUEST"); + + // Case 3: Test constructor with only status (default message) + HttpShortcutResponse responseWithoutMessage = new HttpShortcutResponse(HttpResponseStatus.NOT_FOUND); + assertEquals(responseWithoutMessage.getMessage(), "", "Message should be empty by default"); + assertEquals(responseWithoutMessage.getStatus(), HttpResponseStatus.NOT_FOUND, "Status should be NOT_FOUND"); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/meta/ServerAdminActionTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/meta/ServerAdminActionTest.java new file mode 100644 index 0000000000..74628cab12 --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/meta/ServerAdminActionTest.java @@ -0,0 +1,20 @@ +package com.linkedin.venice.meta; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.expectThrows; + +import org.testng.annotations.Test; + + +public class ServerAdminActionTest { + @Test + public void testFromValue() { + // Test all valid values + assertEquals(ServerAdminAction.fromValue(0), ServerAdminAction.DUMP_INGESTION_STATE); + assertEquals(ServerAdminAction.fromValue(1), ServerAdminAction.DUMP_SERVER_CONFIGS); + + // Test invalid values + expectThrows(IllegalArgumentException.class, () -> ServerAdminAction.fromValue(-1)); + expectThrows(IllegalArgumentException.class, () -> ServerAdminAction.fromValue(2)); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/response/VeniceReadResponseStatusTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/response/VeniceReadResponseStatusTest.java new file mode 100644 index 0000000000..52210b64b6 --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/response/VeniceReadResponseStatusTest.java @@ -0,0 +1,84 @@ +package com.linkedin.venice.response; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.expectThrows; + +import io.netty.handler.codec.http.HttpResponseStatus; +import org.testng.annotations.Test; + + +public class VeniceReadResponseStatusTest { + @Test + public void testEnumValues() { + assertEquals(VeniceReadResponseStatus.KEY_NOT_FOUND.getHttpResponseStatus(), HttpResponseStatus.NOT_FOUND); + assertEquals(VeniceReadResponseStatus.OK.getHttpResponseStatus(), HttpResponseStatus.OK); + assertEquals(VeniceReadResponseStatus.BAD_REQUEST.getHttpResponseStatus(), HttpResponseStatus.BAD_REQUEST); + assertEquals(VeniceReadResponseStatus.FORBIDDEN.getHttpResponseStatus(), HttpResponseStatus.FORBIDDEN); + assertEquals( + VeniceReadResponseStatus.METHOD_NOT_ALLOWED.getHttpResponseStatus(), + HttpResponseStatus.METHOD_NOT_ALLOWED); + assertEquals(VeniceReadResponseStatus.REQUEST_TIMEOUT.getHttpResponseStatus(), HttpResponseStatus.REQUEST_TIMEOUT); + assertEquals( + VeniceReadResponseStatus.TOO_MANY_REQUESTS.getHttpResponseStatus(), + HttpResponseStatus.TOO_MANY_REQUESTS); + assertEquals( + VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getHttpResponseStatus(), + HttpResponseStatus.INTERNAL_SERVER_ERROR); + assertEquals( + VeniceReadResponseStatus.SERVICE_UNAVAILABLE.getHttpResponseStatus(), + HttpResponseStatus.SERVICE_UNAVAILABLE); + assertEquals(VeniceReadResponseStatus.MISROUTED_STORE_VERSION.getHttpResponseStatus().code(), 570); + } + + @Test + public void testGetCode() { + assertEquals(VeniceReadResponseStatus.KEY_NOT_FOUND.getCode(), HttpResponseStatus.NOT_FOUND.code()); + assertEquals(VeniceReadResponseStatus.OK.getCode(), HttpResponseStatus.OK.code()); + assertEquals(VeniceReadResponseStatus.BAD_REQUEST.getCode(), HttpResponseStatus.BAD_REQUEST.code()); + assertEquals(VeniceReadResponseStatus.FORBIDDEN.getCode(), HttpResponseStatus.FORBIDDEN.code()); + assertEquals(VeniceReadResponseStatus.METHOD_NOT_ALLOWED.getCode(), HttpResponseStatus.METHOD_NOT_ALLOWED.code()); + assertEquals(VeniceReadResponseStatus.REQUEST_TIMEOUT.getCode(), HttpResponseStatus.REQUEST_TIMEOUT.code()); + assertEquals(VeniceReadResponseStatus.TOO_MANY_REQUESTS.getCode(), HttpResponseStatus.TOO_MANY_REQUESTS.code()); + assertEquals( + VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode(), + HttpResponseStatus.INTERNAL_SERVER_ERROR.code()); + assertEquals(VeniceReadResponseStatus.SERVICE_UNAVAILABLE.getCode(), HttpResponseStatus.SERVICE_UNAVAILABLE.code()); + assertEquals(VeniceReadResponseStatus.MISROUTED_STORE_VERSION.getCode(), 570); + } + + @Test + public void testFromCode() { + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.NOT_FOUND.code()), + VeniceReadResponseStatus.KEY_NOT_FOUND); + assertEquals(VeniceReadResponseStatus.fromCode(HttpResponseStatus.OK.code()), VeniceReadResponseStatus.OK); + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.BAD_REQUEST.code()), + VeniceReadResponseStatus.BAD_REQUEST); + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.FORBIDDEN.code()), + VeniceReadResponseStatus.FORBIDDEN); + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.METHOD_NOT_ALLOWED.code()), + VeniceReadResponseStatus.METHOD_NOT_ALLOWED); + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.REQUEST_TIMEOUT.code()), + VeniceReadResponseStatus.REQUEST_TIMEOUT); + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.TOO_MANY_REQUESTS.code()), + VeniceReadResponseStatus.TOO_MANY_REQUESTS); + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.INTERNAL_SERVER_ERROR.code()), + VeniceReadResponseStatus.INTERNAL_SERVER_ERROR); + assertEquals( + VeniceReadResponseStatus.fromCode(HttpResponseStatus.SERVICE_UNAVAILABLE.code()), + VeniceReadResponseStatus.SERVICE_UNAVAILABLE); + assertEquals(VeniceReadResponseStatus.fromCode(570), VeniceReadResponseStatus.MISROUTED_STORE_VERSION); + } + + @Test + public void testFromCodeForInvalidCode() { + Exception exception = expectThrows(IllegalArgumentException.class, () -> VeniceReadResponseStatus.fromCode(999)); + assertEquals(exception.getMessage(), "Unknown status venice read response status code: 999"); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/utils/NettyUtilsTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/utils/NettyUtilsTest.java new file mode 100644 index 0000000000..9433e0b2fe --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/utils/NettyUtilsTest.java @@ -0,0 +1,31 @@ +package com.linkedin.venice.utils; + +import static com.linkedin.venice.HttpConstants.VENICE_RETRY; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +import io.netty.handler.codec.http.DefaultHttpRequest; +import io.netty.handler.codec.http.HttpMethod; +import io.netty.handler.codec.http.HttpRequest; +import io.netty.handler.codec.http.HttpVersion; +import org.testng.annotations.Test; + + +public class NettyUtilsTest { + @Test + public void testContainRetryHeader() { + // 1. Header is present + HttpRequest httpRequest = new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/test"); + httpRequest.headers().set(VENICE_RETRY, "1"); + assertTrue(NettyUtils.containRetryHeader(httpRequest), "Request should contain the retry header"); + + // 2. Header is present with value "true" + httpRequest.headers().clear(); + httpRequest.headers().set(VENICE_RETRY, "true"); + assertTrue(NettyUtils.containRetryHeader(httpRequest), "Request should contain the retry header"); + + // 3. Header is not present + httpRequest.headers().clear(); + assertFalse(NettyUtils.containRetryHeader(httpRequest), "Request should not contain the retry header"); + } +} diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/fastclient/FastClientIndividualFeatureConfigurationTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/fastclient/FastClientIndividualFeatureConfigurationTest.java index 8426dd84e0..38f2a144b1 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/fastclient/FastClientIndividualFeatureConfigurationTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/fastclient/FastClientIndividualFeatureConfigurationTest.java @@ -379,7 +379,7 @@ public void testStreamingBatchGetServerStats() throws Exception { } genericFastClient.batchGet(keys).get(); /** - * Both multi-get and streaming multi-get is emitted under the multi-get request type on server side. See {@link com.linkedin.venice.listener.ServerStatsContext} + * Both multi-get and streaming multi-get is emitted under the multi-get request type on server side. See {@link RequestStatsRecorder} */ String multiGetRequestKeyCountMetric = ".total--" + RequestType.MULTI_GET.getMetricPrefix() + "request_key_count.Rate"; diff --git a/services/venice-server/build.gradle b/services/venice-server/build.gradle index 8f1d314e52..3356bd516b 100644 --- a/services/venice-server/build.gradle +++ b/services/venice-server/build.gradle @@ -63,6 +63,8 @@ dependencies { testImplementation libraries.fastUtil testImplementation project(':internal:alpini:router:alpini-router-api') testImplementation project(':services:venice-router') + + testImplementation libraries.grpcTesting } jar { diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcIoRequestProcessor.java b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcIoRequestProcessor.java new file mode 100644 index 0000000000..c3904ac366 --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcIoRequestProcessor.java @@ -0,0 +1,111 @@ +package com.linkedin.venice.grpc; + +import static com.linkedin.venice.listener.QuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED; +import static com.linkedin.venice.listener.ReadQuotaEnforcementHandler.INVALID_REQUEST_RESOURCE_MSG; +import static com.linkedin.venice.listener.ReadQuotaEnforcementHandler.SERVER_OVER_CAPACITY_MSG; + +import com.linkedin.venice.listener.QuotaEnforcementHandler; +import com.linkedin.venice.listener.QuotaEnforcementHandler.QuotaEnforcementResult; +import com.linkedin.venice.listener.StorageReadRequestHandler; +import com.linkedin.venice.listener.request.RouterRequest; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * This class is responsible for processing IO requests i.e. single-get, multi-get, and compute operations, + * in a gRPC server. It manages the request lifecycle by enforcing quotas, delegating valid requests to the storage + * read handler, and ensuring appropriate error responses for requests that exceed quota limits or encounter other issues. + * + * The core responsibilities of this class include: + *

+ * + * The {@link #processRequest(GrpcRequestContext)} method is the central entry point for processing gRPC IO requests. + * It checks for quota violations, processes valid requests, and sends responses back to the client, either success or error, + * depending on the result of the quota enforcement. + */ +public class GrpcIoRequestProcessor { + private static final Logger LOGGER = LogManager.getLogger(GrpcIoRequestProcessor.class); + private final QuotaEnforcementHandler quotaEnforcementHandler; + private final StorageReadRequestHandler storageReadRequestHandler; + private final GrpcReplyProcessor replyProcessor; + + public GrpcIoRequestProcessor(GrpcServiceDependencies services) { + this.quotaEnforcementHandler = services.getQuotaEnforcementHandler(); + this.storageReadRequestHandler = services.getStorageReadRequestHandler(); + this.replyProcessor = services.getGrpcReplyProcessor(); + } + + /** + * Processes a gRPC request by enforcing quota limits, delegating the request for storage read handling, or + * setting an appropriate error response depending on the result of quota enforcement. + * + * The request is first passed to the {@link QuotaEnforcementHandler}, which determines if the request is allowed + * based on the current quota limits. If the quota enforcement result is {@link QuotaEnforcementResult#ALLOWED}, + * the request is handed off to the {@link StorageReadRequestHandler} for further processing. + * The response handling is done asynchronously via the {@link GrpcStorageResponseHandlerCallback}. + * + * If the quota enforcement result is not {@link QuotaEnforcementResult#ALLOWED}, an error response is set based on the specific result. + * The error responses can be: + * + * + * After determining the appropriate response, the method calls {@link GrpcReplyProcessor#sendResponse(GrpcRequestContext)} to + * finalize and send the response to the client. + * + * This method is executed in the gRPC executor thread, so it should not perform any blocking operations. + * + * @param requestContext The {@link GrpcRequestContext} containing the request details, response context, and + * associated metrics/stats to be reported. + */ + public void processRequest(GrpcRequestContext requestContext) { + RouterRequest request = requestContext.getRouterRequest(); + + QuotaEnforcementResult result = quotaEnforcementHandler.enforceQuota(request); + // If the request is allowed, hand it off to the storage read request handler + if (result == ALLOWED) { + storageReadRequestHandler.queueIoRequestForAsyncProcessing( + request, + GrpcStorageResponseHandlerCallback.create(requestContext, replyProcessor)); + return; + } + + // Otherwise, set an error response based on the quota enforcement result + switch (result) { + case BAD_REQUEST: + requestContext.setErrorMessage(INVALID_REQUEST_RESOURCE_MSG + request.getResourceName()); + requestContext.setReadResponseStatus(VeniceReadResponseStatus.BAD_REQUEST); + break; + case REJECTED: + requestContext.setReadResponseStatus(VeniceReadResponseStatus.TOO_MANY_REQUESTS); + requestContext.setErrorMessage("Quota exceeded for resource: " + request.getResourceName()); + break; + case OVER_CAPACITY: + requestContext.setReadResponseStatus(VeniceReadResponseStatus.SERVICE_UNAVAILABLE); + requestContext.setErrorMessage(SERVER_OVER_CAPACITY_MSG); + break; + default: + requestContext.setReadResponseStatus(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR); + requestContext.setErrorMessage("Unknown quota enforcement result: " + result); + LOGGER.error("Unknown quota enforcement result: {}", result); + } + + replyProcessor.sendResponse(requestContext); + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcReplyProcessor.java b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcReplyProcessor.java new file mode 100644 index 0000000000..e0e3520500 --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcReplyProcessor.java @@ -0,0 +1,184 @@ +package com.linkedin.venice.grpc; + +import com.google.protobuf.ByteString; +import com.linkedin.davinci.listener.response.ReadResponse; +import com.linkedin.venice.HttpConstants; +import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.listener.RequestStatsRecorder; +import com.linkedin.venice.listener.response.AbstractReadResponse; +import com.linkedin.venice.protocols.MultiKeyResponse; +import com.linkedin.venice.protocols.SingleGetResponse; +import com.linkedin.venice.protocols.VeniceServerResponse; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import io.grpc.stub.StreamObserver; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * This class is responsible for sending responses to the client and recording request statistics. + * Though it has methods which do not have any side effects, we have not made them static to allow + * for easier testing of the callers. + */ +public class GrpcReplyProcessor { + private static final Logger LOGGER = LogManager.getLogger(GrpcIoRequestProcessor.class); + + /** + * Callers must ensure that all fields in the request context are properly set before invoking this method. + * Callers must also use the appropriate {@link GrpcRequestContext#readResponseStatus} to comply with the API contract. + * + * @param requestContext The context of the request for which a response is being sent + * @param The type of the response observer + */ + void sendResponse(GrpcRequestContext requestContext) { + GrpcRequestContext.GrpcRequestType grpcRequestType = requestContext.getGrpcRequestType(); + switch (grpcRequestType) { + case SINGLE_GET: + sendSingleGetResponse((GrpcRequestContext) requestContext); + break; + case MULTI_GET: + case COMPUTE: + sendMultiKeyResponse((GrpcRequestContext) requestContext); + break; + case LEGACY: + sendVeniceServerResponse((GrpcRequestContext) requestContext); + break; + default: + VeniceException veniceException = new VeniceException("Unknown response type: " + grpcRequestType); + LOGGER.error("Unknown response type: {}", grpcRequestType, veniceException); + throw veniceException; + } + } + + /** + * Sends a single get response to the client and records the request statistics via {@link #reportRequestStats}. + * Since {@link io.grpc.stub.StreamObserver} is not thread-safe, synchronization is required before invoking + * {@link io.grpc.stub.StreamObserver#onNext} and {@link io.grpc.stub.StreamObserver#onCompleted}. + * + * @param requestContext The context of the gRPC request, which includes the response and stats recorder to be updated. + */ + void sendSingleGetResponse(GrpcRequestContext requestContext) { + ReadResponse readResponse = requestContext.getReadResponse(); + SingleGetResponse.Builder builder = SingleGetResponse.newBuilder(); + VeniceReadResponseStatus responseStatus = requestContext.getReadResponseStatus(); + + if (readResponse == null) { + builder.setStatusCode(requestContext.getReadResponseStatus().getCode()); + builder.setErrorMessage(requestContext.getErrorMessage()); + } else if (readResponse.isFound()) { + builder.setRcu(readResponse.getRCU()) + .setStatusCode(responseStatus.getCode()) + .setSchemaId(readResponse.getResponseSchemaIdHeader()) + .setCompressionStrategy(readResponse.getCompressionStrategy().getValue()) + .setContentLength(readResponse.getResponseBody().readableBytes()) + .setContentType(HttpConstants.AVRO_BINARY) + .setValue(GrpcUtils.toByteString(readResponse.getResponseBody())); + } else { + builder.setStatusCode(responseStatus.getCode()) + .setRcu(readResponse.getRCU()) + .setErrorMessage("Key not found") + .setContentLength(0); + } + + StreamObserver responseObserver = requestContext.getResponseObserver(); + synchronized (responseObserver) { + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + reportRequestStats(requestContext); + } + + /** + * Sends a multi key response (multiGet and compute requests) to the client and records the request statistics via {@link #reportRequestStats}. + * Since {@link StreamObserver} is not thread-safe, synchronization is required before invoking + * {@link StreamObserver#onNext} and {@link StreamObserver#onCompleted}. + * + * @param requestContext The context of the gRPC request, which includes the response and stats recorder to be updated. + */ + void sendMultiKeyResponse(GrpcRequestContext requestContext) { + ReadResponse readResponse = requestContext.getReadResponse(); + MultiKeyResponse.Builder builder = MultiKeyResponse.newBuilder(); + VeniceReadResponseStatus responseStatus = requestContext.getReadResponseStatus(); + + if (readResponse == null) { + builder.setStatusCode(responseStatus.getCode()); + builder.setErrorMessage(requestContext.getErrorMessage()); + } else if (readResponse.isFound()) { + builder.setStatusCode(responseStatus.getCode()) + .setRcu(readResponse.getRCU()) + .setCompressionStrategy(readResponse.getCompressionStrategy().getValue()) + .setContentLength(readResponse.getResponseBody().readableBytes()) + .setContentType(HttpConstants.AVRO_BINARY) + .setValue(GrpcUtils.toByteString(readResponse.getResponseBody())); + } else { + builder.setStatusCode(responseStatus.getCode()) + .setRcu(readResponse.getRCU()) + .setErrorMessage("Key not found") + .setContentLength(0); + } + + StreamObserver responseObserver = requestContext.getResponseObserver(); + synchronized (responseObserver) { + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + reportRequestStats(requestContext); + } + + /** + * Sends response (for the legacy API) to the client and records the request statistics via {@link #reportRequestStats}. + * Since {@link StreamObserver} is not thread-safe, synchronization is required before invoking + * {@link StreamObserver#onNext} and {@link StreamObserver#onCompleted}. + * + * @param requestContext The context of the gRPC request, which includes the response and stats recorder to be updated. + */ + void sendVeniceServerResponse(GrpcRequestContext requestContext) { + ReadResponse readResponse = requestContext.getReadResponse(); + VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); + VeniceReadResponseStatus readResponseStatus = requestContext.getReadResponseStatus(); + + if (readResponse == null) { + builder.setErrorCode(readResponseStatus.getCode()); + builder.setErrorMessage(requestContext.getErrorMessage()); + } else if (readResponse.isFound()) { + builder.setErrorCode(readResponseStatus.getCode()) + .setResponseRCU(readResponse.getRCU()) + .setCompressionStrategy(readResponse.getCompressionStrategy().getValue()) + .setIsStreamingResponse(readResponse.isStreamingResponse()) + .setSchemaId(readResponse.getResponseSchemaIdHeader()) + .setData(GrpcUtils.toByteString(readResponse.getResponseBody())); + } else { + builder.setErrorCode(readResponseStatus.getCode()).setErrorMessage("Key not found").setData(ByteString.EMPTY); + } + + StreamObserver responseObserver = requestContext.getResponseObserver(); + synchronized (responseObserver) { + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + reportRequestStats(requestContext); + } + + /** + * Records the request statistics based on the provided {@link GrpcRequestContext}. + * This method updates the {@link RequestStatsRecorder} with statistics from the {@link GrpcRequestContext} and {@link ReadResponse}. + * @param requestContext The context of the gRPC request, which contains the response and stats recorder to be updated. + */ + void reportRequestStats(GrpcRequestContext requestContext) { + ReadResponse readResponse = requestContext.getReadResponse(); + RequestStatsRecorder requestStatsRecorder = requestContext.getRequestStatsRecorder(); + AbstractReadResponse abstractReadResponse = (AbstractReadResponse) readResponse; + if (readResponse == null) { + requestStatsRecorder.setReadResponseStats(null).setResponseSize(0); + } else if (readResponse.isFound()) { + requestStatsRecorder.setReadResponseStats(abstractReadResponse.getReadResponseStatsRecorder()) + .setResponseSize(abstractReadResponse.getResponseBody().readableBytes()); + } else { + requestStatsRecorder.setReadResponseStats(abstractReadResponse.getReadResponseStatsRecorder()).setResponseSize(0); + } + + RequestStatsRecorder.recordRequestCompletionStats(requestContext.getRequestStatsRecorder(), true, -1); + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcRequestContext.java b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcRequestContext.java new file mode 100644 index 0000000000..0b5feb4308 --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcRequestContext.java @@ -0,0 +1,112 @@ +package com.linkedin.venice.grpc; + +import static java.util.Objects.requireNonNull; + +import com.linkedin.davinci.listener.response.ReadResponse; +import com.linkedin.venice.listener.RequestStatsRecorder; +import com.linkedin.venice.listener.request.RouterRequest; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import io.grpc.stub.StreamObserver; + + +/** + * This class is used to store the context of a gRPC request. + * @param the type of the response observer + */ +public class GrpcRequestContext { + private final StreamObserver responseObserver; + private final GrpcRequestType grpcRequestType; + private final RequestStatsRecorder requestStatsRecorder; + private RouterRequest routerRequest; + private ReadResponse readResponse = null; + private VeniceReadResponseStatus readResponseStatus = null; + private String errorMessage; + + enum GrpcRequestType { + LEGACY, SINGLE_GET, MULTI_GET, COMPUTE + } + + private GrpcRequestContext( + RequestStatsRecorder requestStatsRecorder, + StreamObserver responseObserver, + GrpcRequestType grpcRequestType) { + this.requestStatsRecorder = + requireNonNull(requestStatsRecorder, "RequestStatsRecorder cannot be null in GrpcRequestContext"); + this.responseObserver = requireNonNull(responseObserver, "ResponseObserver cannot be null in GrpcRequestContext"); + this.grpcRequestType = requireNonNull(grpcRequestType, "GrpcRequestType cannot be null in GrpcRequestContext"); + } + + public static GrpcRequestContext create( + GrpcServiceDependencies services, + StreamObserver responseObserver, + GrpcRequestType grpcRequestType) { + return new GrpcRequestContext<>( + new RequestStatsRecorder(services.getSingleGetStats(), services.getMultiGetStats(), services.getComputeStats()), + responseObserver, + grpcRequestType); + } + + public RequestStatsRecorder getRequestStatsRecorder() { + return requestStatsRecorder; + } + + public StreamObserver getResponseObserver() { + return responseObserver; + } + + public RouterRequest getRouterRequest() { + return routerRequest; + } + + public void setRouterRequest(RouterRequest routerRequest) { + this.routerRequest = routerRequest; + } + + public ReadResponse getReadResponse() { + return readResponse; + } + + public void setReadResponse(ReadResponse readResponse) { + this.readResponse = readResponse; + } + + public VeniceReadResponseStatus getReadResponseStatus() { + // If the readResponseStatus is set, return it. + if (readResponseStatus != null) { + return readResponseStatus; + } + + // If the readResponse is set, return the appropriate status based on the response. + if (readResponse != null && readResponse.isFound()) { + return VeniceReadResponseStatus.OK; + } + + // If the readResponse is set and the key is not found, return the appropriate status. + if (readResponse != null && !readResponse.isFound()) { + return VeniceReadResponseStatus.KEY_NOT_FOUND; + } + + // If the readResponse is not set, return an internal server error. + return VeniceReadResponseStatus.INTERNAL_SERVER_ERROR; + } + + public void setReadResponseStatus(VeniceReadResponseStatus readResponseStatus) { + this.readResponseStatus = readResponseStatus; + if (requestStatsRecorder != null) { + // Setting response status is an important step to ensure that metrics are recorded correctly + requestStatsRecorder.setResponseStatus(readResponseStatus); + } + } + + public String getErrorMessage() { + return errorMessage; + } + + public void setErrorMessage(String errorMessage) { + this.errorMessage = errorMessage; + } + + public GrpcRequestType getGrpcRequestType() { + return grpcRequestType; + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcServiceDependencies.java b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcServiceDependencies.java new file mode 100644 index 0000000000..3cba03c8e0 --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcServiceDependencies.java @@ -0,0 +1,121 @@ +package com.linkedin.venice.grpc; + +import com.linkedin.davinci.storage.DiskHealthCheckService; +import com.linkedin.venice.listener.NoOpReadQuotaEnforcementHandler; +import com.linkedin.venice.listener.QuotaEnforcementHandler; +import com.linkedin.venice.listener.StorageReadRequestHandler; +import com.linkedin.venice.stats.AggServerHttpRequestStats; +import java.util.Objects; + + +public class GrpcServiceDependencies { + private final DiskHealthCheckService diskHealthCheckService; + private final StorageReadRequestHandler storageReadRequestHandler; + private final QuotaEnforcementHandler quotaEnforcementHandler; + private final AggServerHttpRequestStats singleGetStats; + private final AggServerHttpRequestStats multiGetStats; + private final AggServerHttpRequestStats computeStats; + private final GrpcReplyProcessor grpcReplyProcessor; + + private GrpcServiceDependencies(Builder builder) { + this.diskHealthCheckService = builder.diskHealthCheckService; + this.storageReadRequestHandler = builder.storageReadRequestHandler; + this.quotaEnforcementHandler = builder.quotaEnforcementHandler; + this.singleGetStats = builder.singleGetStats; + this.multiGetStats = builder.multiGetStats; + this.computeStats = builder.computeStats; + this.grpcReplyProcessor = builder.grpcReplyProcessor; + } + + public DiskHealthCheckService getDiskHealthCheckService() { + return diskHealthCheckService; + } + + public StorageReadRequestHandler getStorageReadRequestHandler() { + return storageReadRequestHandler; + } + + public QuotaEnforcementHandler getQuotaEnforcementHandler() { + return quotaEnforcementHandler; + } + + public AggServerHttpRequestStats getSingleGetStats() { + return singleGetStats; + } + + public AggServerHttpRequestStats getMultiGetStats() { + return multiGetStats; + } + + public AggServerHttpRequestStats getComputeStats() { + return computeStats; + } + + public GrpcReplyProcessor getGrpcReplyProcessor() { + return grpcReplyProcessor; + } + + public static class Builder { + private DiskHealthCheckService diskHealthCheckService; + private StorageReadRequestHandler storageReadRequestHandler; + private QuotaEnforcementHandler quotaEnforcementHandler; + private AggServerHttpRequestStats singleGetStats; + private AggServerHttpRequestStats multiGetStats; + private AggServerHttpRequestStats computeStats; + private GrpcReplyProcessor grpcReplyProcessor; + + public Builder setDiskHealthCheckService(DiskHealthCheckService diskHealthCheckService) { + this.diskHealthCheckService = diskHealthCheckService; + return this; + } + + public Builder setStorageReadRequestHandler(StorageReadRequestHandler storageReadRequestHandler) { + this.storageReadRequestHandler = storageReadRequestHandler; + return this; + } + + public Builder setQuotaEnforcementHandler(QuotaEnforcementHandler quotaEnforcementHandler) { + this.quotaEnforcementHandler = quotaEnforcementHandler; + return this; + } + + public Builder setSingleGetStats(AggServerHttpRequestStats singleGetStats) { + this.singleGetStats = singleGetStats; + return this; + } + + public Builder setMultiGetStats(AggServerHttpRequestStats multiGetStats) { + this.multiGetStats = multiGetStats; + return this; + } + + public Builder setComputeStats(AggServerHttpRequestStats computeStats) { + this.computeStats = computeStats; + return this; + } + + public Builder setGrpcReplyProcessor(GrpcReplyProcessor grpcReplyProcessor) { + this.grpcReplyProcessor = grpcReplyProcessor; + return this; + } + + public GrpcServiceDependencies build() { + // Validate that all required fields are set + if (quotaEnforcementHandler == null) { + quotaEnforcementHandler = NoOpReadQuotaEnforcementHandler.getInstance(); + } + if (grpcReplyProcessor == null) { + grpcReplyProcessor = new GrpcReplyProcessor(); + } + + singleGetStats = Objects.requireNonNull(singleGetStats, "singleGetStats cannot be null"); + multiGetStats = Objects.requireNonNull(multiGetStats, "multiGetStats cannot be null"); + computeStats = Objects.requireNonNull(computeStats, "computeStats cannot be null"); + storageReadRequestHandler = + Objects.requireNonNull(storageReadRequestHandler, "storageReadRequestHandler cannot be null"); + diskHealthCheckService = Objects.requireNonNull(diskHealthCheckService, "diskHealthCheckService cannot be null"); + + return new GrpcServiceDependencies(this); + } + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcStorageResponseHandlerCallback.java b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcStorageResponseHandlerCallback.java new file mode 100644 index 0000000000..fc0a48629d --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/grpc/GrpcStorageResponseHandlerCallback.java @@ -0,0 +1,53 @@ +package com.linkedin.venice.grpc; + +import com.linkedin.davinci.listener.response.ReadResponse; +import com.linkedin.venice.listener.StorageResponseHandlerCallback; +import com.linkedin.venice.response.VeniceReadResponseStatus; + + +/** + * This class is a callback that is used to handle the response from the storage layer. + * + * Except for the cases where the response was returned before the storage layer was invoked, this class is used to + * handle the response from the storage layer and pass it to the next handler in the pipeline. + */ +public class GrpcStorageResponseHandlerCallback implements StorageResponseHandlerCallback { + private final GrpcRequestContext requestContext; + private final GrpcReplyProcessor grpcReplyProcessor; + + private GrpcStorageResponseHandlerCallback(GrpcRequestContext requestContext, GrpcReplyProcessor grpcReplyProcessor) { + this.requestContext = requestContext; + this.grpcReplyProcessor = grpcReplyProcessor; + } + + // Factory method for creating an instance of this class. + public static GrpcStorageResponseHandlerCallback create( + GrpcRequestContext requestContext, + GrpcReplyProcessor grpcReplyProcessor) { + return new GrpcStorageResponseHandlerCallback(requestContext, grpcReplyProcessor); + } + + @Override + public void onReadResponse(ReadResponse readResponse) { + if (readResponse == null) { + onError(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR, "StorageHandler returned a unexpected null response"); + return; + } + + if (readResponse.isFound()) { + requestContext.setReadResponseStatus(VeniceReadResponseStatus.OK); + } else { + requestContext.setReadResponseStatus(VeniceReadResponseStatus.KEY_NOT_FOUND); + } + requestContext.setReadResponse(readResponse); + grpcReplyProcessor.sendResponse(requestContext); + } + + @Override + public void onError(VeniceReadResponseStatus readResponseStatus, String message) { + requestContext.setReadResponseStatus(readResponseStatus); + requestContext.setErrorMessage(message); + requestContext.setReadResponse(null); + grpcReplyProcessor.sendResponse(requestContext); + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcReadServiceImpl.java b/services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcReadServiceImpl.java new file mode 100644 index 0000000000..4e74f24e4e --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcReadServiceImpl.java @@ -0,0 +1,409 @@ +package com.linkedin.venice.grpc; + +import static com.linkedin.venice.grpc.GrpcRequestContext.GrpcRequestType.COMPUTE; +import static com.linkedin.venice.grpc.GrpcRequestContext.GrpcRequestType.LEGACY; +import static com.linkedin.venice.grpc.GrpcRequestContext.GrpcRequestType.MULTI_GET; +import static com.linkedin.venice.grpc.GrpcRequestContext.GrpcRequestType.SINGLE_GET; +import static com.linkedin.venice.listener.StorageReadRequestHandler.VENICE_STORAGE_NODE_HARDWARE_IS_NOT_HEALTHY_MSG; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.protobuf.ByteString; +import com.linkedin.davinci.listener.response.ServerCurrentVersionResponse; +import com.linkedin.davinci.listener.response.TopicPartitionIngestionContextResponse; +import com.linkedin.davinci.storage.DiskHealthCheckService; +import com.linkedin.venice.HttpConstants; +import com.linkedin.venice.listener.StorageReadRequestHandler; +import com.linkedin.venice.listener.request.ComputeRouterRequestWrapper; +import com.linkedin.venice.listener.request.CurrentVersionRequest; +import com.linkedin.venice.listener.request.DictionaryFetchRequest; +import com.linkedin.venice.listener.request.GetRouterRequest; +import com.linkedin.venice.listener.request.MetadataFetchRequest; +import com.linkedin.venice.listener.request.MultiGetRouterRequestWrapper; +import com.linkedin.venice.listener.request.RouterRequest; +import com.linkedin.venice.listener.request.TopicPartitionIngestionContextRequest; +import com.linkedin.venice.listener.response.BinaryResponse; +import com.linkedin.venice.meta.Version; +import com.linkedin.venice.protocols.AdminRequest; +import com.linkedin.venice.protocols.AdminResponse; +import com.linkedin.venice.protocols.CompressionDictionaryRequest; +import com.linkedin.venice.protocols.CompressionDictionaryResponse; +import com.linkedin.venice.protocols.ComputeRequest; +import com.linkedin.venice.protocols.CurrentVersionInfoRequest; +import com.linkedin.venice.protocols.CurrentVersionInfoResponse; +import com.linkedin.venice.protocols.HealthCheckRequest; +import com.linkedin.venice.protocols.HealthCheckResponse; +import com.linkedin.venice.protocols.IngestionContextRequest; +import com.linkedin.venice.protocols.IngestionContextResponse; +import com.linkedin.venice.protocols.MetadataRequest; +import com.linkedin.venice.protocols.MetadataResponse; +import com.linkedin.venice.protocols.MultiGetRequest; +import com.linkedin.venice.protocols.MultiKeyResponse; +import com.linkedin.venice.protocols.SingleGetRequest; +import com.linkedin.venice.protocols.SingleGetResponse; +import com.linkedin.venice.protocols.VeniceClientRequest; +import com.linkedin.venice.protocols.VeniceReadServiceGrpc; +import com.linkedin.venice.protocols.VeniceServerResponse; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import com.linkedin.venice.utils.ObjectMapperFactory; +import io.grpc.stub.StreamObserver; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * This class implements the Venice gRPC read service. It manages the processing of read requests + * in a multi-threaded environment using Netty-based gRPC. To optimize performance and prevent deadlocks, + * it handles IO and non-IO operations in different thread pools. + * + * ### Thread Pools in Netty-based gRPC: + * 1. **Netty Boss Group**: Handles incoming connections. + * 2. **Netty Worker Group**: Handles IO-related work (Netty IO threads). + * 3. **gRPC Executor**: Executes gRPC service methods. The default gRPC executor is a caching thread pool, + * which creates a new thread for each incoming request. However, this is inefficient for handling + * a large number of requests, as it leads to thread creation overhead. + * + * Note: Netty threads and abstractions are not directly exposed to us. + * + * ### Custom gRPC Executor: + * To avoid creating a new thread per request, we provide a fixed-size thread pool executor to the gRPC server builder. + * The fixed thread pool ensures a controlled request processing concurrency. + * However, to prevent deadlocks, we must avoid blocking the gRPC executor threads with long-running or blocking operations. + * + * ### Offloading IO Requests: + * Since IO requests (like reading from storage) are blocking, they are offloaded to a separate thread pool. + * The {@link StorageReadRequestHandler#executor} is responsible for processing these requests. Offloading IO requests + * improves performance by ensuring that the gRPC executor is not tied up with slow operations, allowing it to process + * new incoming requests efficiently. + * + * ### Non-IO Requests: + * Non-IO requests, which are not expected to block, are processed directly in the gRPC executor threads. + * For read requests, quota enforcement is performed in the gRPC executor thread. If the request passes quota enforcement, + * it is then handed off to the IO thread pool for further processing. + * + * ### Rough Request Flow: + * ``` + * +-------------------+ 1. Incoming conn +--------------------------+ + * | Netty Boss Group | <---------------------------- | Client | + * +-------------------+ +--------------------------+ + * | + * v + * +-------------------+ 2. Handles Req IO work +--------------------------+ + * | Netty Worker Group | <---------------------------- | Storage / IO Operations | + * +-------------------+ +--------------------------+ + * | + * v + * +-------------------+ 3. Execute gRPC Method +--------------------------+ + * | gRPC Executor | <---------------------------- | gRPC Service Method | + * +-------------------+ +--------------------------+ + * | + * v + * +-------------------+ 4. Offload Blocking IO +--------------------------+ + * | IO Thread Pool | <---------------------------- | StorageReadRequestHandler| + * +-------------------+ +--------------------------+ + * ``` + */ +public class VeniceGrpcReadServiceImpl extends VeniceReadServiceGrpc.VeniceReadServiceImplBase { + private static final Logger LOGGER = LogManager.getLogger(VeniceGrpcReadServiceImpl.class); + private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getInstance(); + + private final DiskHealthCheckService diskHealthCheckService; + private final StorageReadRequestHandler storageReadRequestHandler; + private final GrpcServiceDependencies dependencies; + private final GrpcIoRequestProcessor requestProcessor; + + public VeniceGrpcReadServiceImpl(GrpcServiceDependencies dependencies) { + this(dependencies, new GrpcIoRequestProcessor(dependencies)); + } + + VeniceGrpcReadServiceImpl(GrpcServiceDependencies dependencies, GrpcIoRequestProcessor requestProcessor) { + this.dependencies = dependencies; + this.requestProcessor = requestProcessor; + this.diskHealthCheckService = dependencies.getDiskHealthCheckService(); + this.storageReadRequestHandler = dependencies.getStorageReadRequestHandler(); + } + + /** + * @deprecated This method is deprecated and will be removed in the future. Use the {@link #singleGet(SingleGetRequest, StreamObserver)} method instead. + */ + @Deprecated + @Override + public void get(VeniceClientRequest singleGetRequest, StreamObserver streamObserver) { + GrpcRequestContext clientRequestCtx = + GrpcRequestContext.create(dependencies, streamObserver, LEGACY); + try { + RouterRequest routerRequest = GetRouterRequest.parseSingleGetGrpcRequest(singleGetRequest); + clientRequestCtx.getRequestStatsRecorder() + .setRequestInfo(routerRequest) + .setRequestSize(singleGetRequest.getSerializedSize()); + clientRequestCtx.setRouterRequest(routerRequest); + requestProcessor.processRequest(clientRequestCtx); + } catch (Exception e) { + LOGGER.debug("Error while processing single get request", e); + VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); + builder.setErrorCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage(e.getMessage()); + } + streamObserver.onNext(builder.build()); + streamObserver.onCompleted(); + } + } + + /** + * @deprecated This method is deprecated and will be removed in the future. Use the {@link #multiGet(MultiGetRequest, StreamObserver)} method instead. + */ + @Deprecated + @Override + public void batchGet(VeniceClientRequest batchGetRequest, StreamObserver streamObserver) { + GrpcRequestContext requestContext = + GrpcRequestContext.create(dependencies, streamObserver, LEGACY); + try { + RouterRequest routerRequest = MultiGetRouterRequestWrapper.parseMultiGetGrpcRequest(batchGetRequest); + requestContext.getRequestStatsRecorder() + .setRequestInfo(routerRequest) + .setRequestSize(batchGetRequest.getSerializedSize()); + requestContext.setRouterRequest(routerRequest); + requestProcessor.processRequest(requestContext); + } catch (Exception e) { + LOGGER.debug("Error while processing batch-get request", e); + VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); + builder.setErrorCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage(e.getMessage()); + } + streamObserver.onNext(builder.build()); + streamObserver.onCompleted(); + } + } + + @Override + public void singleGet(SingleGetRequest singleGetRequest, StreamObserver streamObserver) { + GrpcRequestContext requestContext = + GrpcRequestContext.create(dependencies, streamObserver, SINGLE_GET); + try { + RouterRequest routerRequest = GetRouterRequest.parseSingleGetGrpcRequest(singleGetRequest); + requestContext.getRequestStatsRecorder() + .setRequestInfo(routerRequest) + .setRequestSize(singleGetRequest.getSerializedSize()); + requestContext.setRouterRequest(routerRequest); + requestProcessor.processRequest(requestContext); + } catch (Exception e) { + LOGGER.debug("Error while processing single get request", e); + SingleGetResponse.Builder builder = SingleGetResponse.newBuilder(); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage(e.getMessage()); + } + streamObserver.onNext(builder.build()); + streamObserver.onCompleted(); + } + } + + @Override + public void multiGet(MultiGetRequest multiGetRequest, StreamObserver streamObserver) { + GrpcRequestContext requestContext = + GrpcRequestContext.create(dependencies, streamObserver, MULTI_GET); + try { + RouterRequest routerRequest = MultiGetRouterRequestWrapper.parseMultiGetGrpcRequest(multiGetRequest); + requestContext.getRequestStatsRecorder() + .setRequestInfo(routerRequest) + .setRequestSize(multiGetRequest.getSerializedSize()); + requestContext.setRouterRequest(routerRequest); + requestProcessor.processRequest(requestContext); + } catch (Exception e) { + LOGGER.debug("Error while processing multi get request", e); + MultiKeyResponse.Builder builder = MultiKeyResponse.newBuilder(); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage(e.getMessage()); + } + streamObserver.onNext(builder.build()); + streamObserver.onCompleted(); + } + } + + @Override + public void compute(ComputeRequest computeRequest, StreamObserver responseObserver) { + GrpcRequestContext requestContext = + GrpcRequestContext.create(dependencies, responseObserver, COMPUTE); + try { + RouterRequest routerRequest = ComputeRouterRequestWrapper.parseComputeGrpcRequest(computeRequest); + requestContext.getRequestStatsRecorder() + .setRequestInfo(routerRequest) + .setRequestSize(computeRequest.getSerializedSize()); + requestContext.setRouterRequest(routerRequest); + requestProcessor.processRequest(requestContext); + } catch (Exception e) { + LOGGER.debug("Error while processing compute request", e); + MultiKeyResponse.Builder builder = MultiKeyResponse.newBuilder(); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage(e.getMessage()); + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + } + + @Override + public void isServerHealthy( + HealthCheckRequest healthCheckRequest, + StreamObserver responseObserver) { + HealthCheckResponse.Builder builder = HealthCheckResponse.newBuilder(); + if (diskHealthCheckService.isDiskHealthy()) { + builder.setStatusCode(VeniceReadResponseStatus.OK.getCode()); + } else { + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()) + .setMessage(VENICE_STORAGE_NODE_HARDWARE_IS_NOT_HEALTHY_MSG); + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + @Override + public void getCompressionDictionary( + CompressionDictionaryRequest dictionaryRequest, + StreamObserver responseObserver) { + CompressionDictionaryResponse.Builder builder = CompressionDictionaryResponse.newBuilder(); + try { + String topicName = + Version.composeKafkaTopic(dictionaryRequest.getStoreName(), dictionaryRequest.getStoreVersion()); + BinaryResponse binaryResponse = storageReadRequestHandler + .handleDictionaryFetchRequest(DictionaryFetchRequest.create(dictionaryRequest.getStoreName(), topicName)); + builder.setStatusCode(binaryResponse.getStatus().getCode()) + .setContentType(HttpConstants.BINARY) + .setContentLength(binaryResponse.getBody().readableBytes()) + .setValue(GrpcUtils.toByteString(binaryResponse.getBody())); + } catch (Exception e) { + LOGGER.error("Error while processing dictionary request", e); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage("Error while processing dictionary request: " + e.getMessage()); + } + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + @Override + public void handleAdminRequest(AdminRequest request, StreamObserver responseObserver) { + AdminResponse.Builder builder = AdminResponse.newBuilder(); + try { + com.linkedin.davinci.listener.response.AdminResponse adminResponse = storageReadRequestHandler + .handleServerAdminRequest(com.linkedin.venice.listener.request.AdminRequest.parseAdminGrpcRequest(request)); + if (!adminResponse.isError()) { + builder.setSchemaId(com.linkedin.davinci.listener.response.AdminResponse.getResponseSchemaIdHeader()) + .setStatusCode(VeniceReadResponseStatus.OK.getCode()) + .setValue(GrpcUtils.toByteString(adminResponse.getResponseBody())) + .setContentType(HttpConstants.AVRO_BINARY) + .setContentLength(adminResponse.getResponseBody().readableBytes()); + } else { + String errorMessage = adminResponse.getMessage() != null ? adminResponse.getMessage() : "Unknown error"; + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()) + .setErrorMessage(errorMessage) + .setContentType(HttpConstants.TEXT_PLAIN); + } + } catch (Exception e) { + LOGGER.error("Error while processing admin request", e); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage("Error while processing admin request: " + e.getMessage()); + } + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + @Override + public void getMetadata(MetadataRequest request, StreamObserver responseObserver) { + MetadataResponse.Builder builder = MetadataResponse.newBuilder(); + try { + MetadataFetchRequest metadataFetchRequest = MetadataFetchRequest.parseGetGrpcRequest(request); + com.linkedin.davinci.listener.response.MetadataResponse metadataResponse = + storageReadRequestHandler.handleMetadataFetchRequest(metadataFetchRequest); + if (!metadataResponse.isError()) { + builder.setStatusCode(VeniceReadResponseStatus.OK.getCode()) + .setValue(GrpcUtils.toByteString(metadataResponse.getResponseBody())) + .setContentType(HttpConstants.AVRO_BINARY) + .setContentLength(metadataResponse.getResponseBody().readableBytes()) + .setSchemaId(metadataResponse.getResponseSchemaIdHeader()); + } else { + String errorMessage = metadataResponse.getMessage() != null ? metadataResponse.getMessage() : "Unknown error"; + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()) + .setErrorMessage(errorMessage) + .setContentType(HttpConstants.TEXT_PLAIN); + } + } catch (Exception e) { + LOGGER.error("Error while processing metadata request", e); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage("Error while processing metadata request: " + e.getMessage()); + } + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + @Override + public void getCurrentVersionInfo( + CurrentVersionInfoRequest request, + StreamObserver responseObserver) { + CurrentVersionInfoResponse.Builder builder = CurrentVersionInfoResponse.newBuilder(); + try { + CurrentVersionRequest currentVersionRequest = CurrentVersionRequest.parseGetGrpcRequest(request); + ServerCurrentVersionResponse currentVersionResponse = + storageReadRequestHandler.handleCurrentVersionRequest(currentVersionRequest); + if (!currentVersionResponse.isError()) { + builder.setStatusCode(VeniceReadResponseStatus.OK.getCode()) + .setCurrentVersion(currentVersionResponse.getCurrentVersion()); + } else { + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()) + .setErrorMessage( + currentVersionResponse.getMessage() != null ? currentVersionResponse.getMessage() : "Unknown error"); + } + } catch (Exception e) { + LOGGER.error("Error while processing current version info request", e); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage("Error while processing current version info request: " + e.getMessage()); + } + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + @Override + public void getIngestionContext( + IngestionContextRequest request, + StreamObserver responseObserver) { + IngestionContextResponse.Builder builder = IngestionContextResponse.newBuilder(); + try { + TopicPartitionIngestionContextRequest ingestionContextRequest = + TopicPartitionIngestionContextRequest.parseGetGrpcRequest(request); + TopicPartitionIngestionContextResponse response = + storageReadRequestHandler.handleTopicPartitionIngestionContextRequest(ingestionContextRequest); + if (!response.isError()) { + ByteString body = GrpcUtils.toByteString(response.getTopicPartitionIngestionContext()); + builder.setValue(body).setContentLength(body.size()).setStatusCode(VeniceReadResponseStatus.OK.getCode()); + } else { + String errorMessage = response.getMessage() != null ? response.getMessage() : "Unknown error"; + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()) + .setErrorMessage(errorMessage) + .setContentType(HttpConstants.TEXT_PLAIN); + } + } catch (Exception e) { + LOGGER.error("Error while processing ingestion context request", e); + builder.setStatusCode(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + if (e.getMessage() != null) { + builder.setErrorMessage("Error while processing ingestion context request: " + e.getMessage()); + } + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + @Override + public String toString() { + return this.getClass().getSimpleName(); + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/ErrorCatchingHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/ErrorCatchingHandler.java index 29ce71cf28..64e23464bb 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/ErrorCatchingHandler.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/ErrorCatchingHandler.java @@ -1,10 +1,11 @@ package com.linkedin.venice.listener; +import static com.linkedin.venice.response.VeniceReadResponseStatus.INTERNAL_SERVER_ERROR; + import com.linkedin.venice.listener.response.HttpShortcutResponse; import io.netty.channel.ChannelHandler; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelInboundHandlerAdapter; -import io.netty.handler.codec.http.HttpResponseStatus; /*** @@ -16,7 +17,7 @@ public class ErrorCatchingHandler extends ChannelInboundHandlerAdapter { @Override public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { - ctx.writeAndFlush(new HttpShortcutResponse(cause.getMessage(), HttpResponseStatus.INTERNAL_SERVER_ERROR)); + ctx.writeAndFlush(new HttpShortcutResponse(cause.getMessage(), INTERNAL_SERVER_ERROR.getHttpResponseStatus())); ctx.close(); } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/HttpChannelInitializer.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/HttpChannelInitializer.java index f118c58faf..e970a6333f 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/HttpChannelInitializer.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/HttpChannelInitializer.java @@ -9,13 +9,6 @@ import com.linkedin.venice.authorization.IdentityParser; import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.helix.HelixCustomizedViewOfflinePushRepository; -import com.linkedin.venice.listener.grpc.handlers.GrpcOutboundResponseHandler; -import com.linkedin.venice.listener.grpc.handlers.GrpcOutboundStatsHandler; -import com.linkedin.venice.listener.grpc.handlers.GrpcReadQuotaEnforcementHandler; -import com.linkedin.venice.listener.grpc.handlers.GrpcRouterRequestHandler; -import com.linkedin.venice.listener.grpc.handlers.GrpcStatsHandler; -import com.linkedin.venice.listener.grpc.handlers.GrpcStorageReadRequestHandler; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcRequestProcessor; import com.linkedin.venice.meta.ReadOnlyStoreRepository; import com.linkedin.venice.read.RequestType; import com.linkedin.venice.security.SSLFactory; @@ -151,9 +144,6 @@ public HttpChannelInitializer( this.identityParser = ReflectUtils.callConstructor(identityParserClass, new Class[0], new Object[0]); } - /* - Test only - */ protected ReadQuotaEnforcementHandler getQuotaEnforcer() { return quotaEnforcer; } @@ -176,8 +166,14 @@ public void initChannel(SocketChannel ch) { ServerConnectionStatsHandler serverConnectionStatsHandler = new ServerConnectionStatsHandler(serverConnectionStats, serverConfig.getRouterPrincipalName()); pipeline.addLast(serverConnectionStatsHandler); - StatsHandler statsHandler = new StatsHandler(singleGetStats, multiGetStats, computeStats); - pipeline.addLast(statsHandler); + /** + * In the Netty pipeline, we create one {@link RequestStatsRecorder} per channel. Since only one request is processed at a time + * per channel (with each HTTP/2 stream having its own child channel, see {@link io.netty.handler.codec.http2.Http2MultiplexHandler}), + * the same instance of {@link RequestStatsRecorder} can be reused for all requests processed within that channel. + * The {@link RequestStatsRecorder} is reset before processing each new request. + */ + RequestStatsRecorder requestStatsRecorder = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); + pipeline.addLast(new StatsHandler(requestStatsRecorder)); if (whetherNeedServerCodec) { pipeline.addLast(new HttpServerCodec()); } else { @@ -201,7 +197,7 @@ public void initChannel(SocketChannel ch) { } pipeline.addLast(new HttpObjectAggregator(serverConfig.getMaxRequestSize())) - .addLast(new OutboundHttpWrapperHandler(statsHandler)) + .addLast(new OutboundHttpWrapperHandler(requestStatsRecorder)) .addLast(new IdleStateHandler(0, 0, serverConfig.getNettyIdleTimeInSeconds())); if (sslFactory.isPresent()) { pipeline.addLast(verifySsl); @@ -215,8 +211,8 @@ public void initChannel(SocketChannel ch) { pipeline.addLast(storeAclHandler.get()); } } - pipeline - .addLast(new RouterRequestHttpHandler(statsHandler, serverConfig.getStoreToEarlyTerminationThresholdMSMap())); + pipeline.addLast( + new RouterRequestHttpHandler(requestStatsRecorder, serverConfig.getStoreToEarlyTerminationThresholdMSMap())); if (quotaEnforcer != null) { pipeline.addLast(quotaEnforcer); } @@ -232,32 +228,16 @@ public void initChannel(SocketChannel ch) { } } - public VeniceServerGrpcRequestProcessor initGrpcRequestProcessor() { - VeniceServerGrpcRequestProcessor grpcServerRequestProcessor = new VeniceServerGrpcRequestProcessor(); - - StatsHandler statsHandler = new StatsHandler(singleGetStats, multiGetStats, computeStats); - GrpcStatsHandler grpcStatsHandler = new GrpcStatsHandler(statsHandler); - grpcServerRequestProcessor.addHandler(grpcStatsHandler); - - GrpcRouterRequestHandler grpcRouterRequestHandler = new GrpcRouterRequestHandler(); - grpcServerRequestProcessor.addHandler(grpcRouterRequestHandler); - - if (quotaEnforcer != null) { - GrpcReadQuotaEnforcementHandler grpcReadQuotaEnforcementHandler = - new GrpcReadQuotaEnforcementHandler(quotaEnforcer); - grpcServerRequestProcessor.addHandler(grpcReadQuotaEnforcementHandler); - } - - GrpcStorageReadRequestHandler storageReadRequestHandler = new GrpcStorageReadRequestHandler(requestHandler); - grpcServerRequestProcessor.addHandler(storageReadRequestHandler); - - GrpcOutboundResponseHandler grpcOutboundResponseHandler = new GrpcOutboundResponseHandler(); - grpcServerRequestProcessor.addHandler(grpcOutboundResponseHandler); + public AggServerHttpRequestStats getSingleGetStats() { + return singleGetStats; + } - GrpcOutboundStatsHandler grpcOutboundStatsHandler = new GrpcOutboundStatsHandler(); - grpcServerRequestProcessor.addHandler(grpcOutboundStatsHandler); + public AggServerHttpRequestStats getMultiGetStats() { + return multiGetStats; + } - return grpcServerRequestProcessor; + public AggServerHttpRequestStats getComputeStats() { + return computeStats; } /** diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/HttpStorageResponseHandlerCallback.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/HttpStorageResponseHandlerCallback.java new file mode 100644 index 0000000000..d9f085270c --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/HttpStorageResponseHandlerCallback.java @@ -0,0 +1,34 @@ +package com.linkedin.venice.listener; + +import com.linkedin.davinci.listener.response.ReadResponse; +import com.linkedin.venice.listener.response.HttpShortcutResponse; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import io.netty.channel.ChannelHandlerContext; + + +/** + * This is used in REST/HTTP Netty handlers to handle the response from the {@link StorageReadRequestHandler#queueIoRequestForAsyncProcessing} method. + */ +public class HttpStorageResponseHandlerCallback implements StorageResponseHandlerCallback { + private final ChannelHandlerContext context; + + private HttpStorageResponseHandlerCallback(ChannelHandlerContext context) { + this.context = context; + } + + // Factory method for creating an instance of this class. + public static HttpStorageResponseHandlerCallback create(ChannelHandlerContext context) { + return new HttpStorageResponseHandlerCallback(context); + } + + @Override + public void onReadResponse(ReadResponse readResponse) { + context.writeAndFlush(readResponse); + } + + @Override + public void onError(VeniceReadResponseStatus readResponseStatus, String message) { + HttpShortcutResponse response = new HttpShortcutResponse(message, readResponseStatus.getHttpResponseStatus()); + context.writeAndFlush(response); + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/ListenerService.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/ListenerService.java index 65bfc44f3d..26d9defa80 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/ListenerService.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/ListenerService.java @@ -9,11 +9,11 @@ import com.linkedin.venice.acl.DynamicAccessController; import com.linkedin.venice.acl.StaticAccessController; import com.linkedin.venice.cleaner.ResourceReadUsageTracker; +import com.linkedin.venice.grpc.GrpcServiceDependencies; +import com.linkedin.venice.grpc.VeniceGrpcReadServiceImpl; import com.linkedin.venice.grpc.VeniceGrpcServer; import com.linkedin.venice.grpc.VeniceGrpcServerConfig; import com.linkedin.venice.helix.HelixCustomizedViewOfflinePushRepository; -import com.linkedin.venice.listener.grpc.VeniceReadServiceImpl; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcRequestProcessor; import com.linkedin.venice.meta.ReadOnlySchemaRepository; import com.linkedin.venice.meta.ReadOnlyStoreRepository; import com.linkedin.venice.security.SSLFactory; @@ -105,7 +105,7 @@ public ListenerService( new ThreadPoolStats(metricsRepository, this.sslHandshakeExecutor, "ssl_handshake_thread_pool"); } - StorageReadRequestHandler requestHandler = createRequestHandler( + StorageReadRequestHandler storageReadRequestHandler = createRequestHandler( executor, computeExecutor, storageEngineRepository, @@ -116,7 +116,6 @@ public ListenerService( diskHealthService, compressorFactory, resourceReadUsageTracker); - HttpChannelInitializer channelInitializer = new HttpChannelInitializer( storeMetadataRepository, customizedViewRepository, @@ -126,7 +125,7 @@ public ListenerService( serverConfig, routerAccessController, storeAccessController, - requestHandler); + storageReadRequestHandler); Class serverSocketChannelClass = NioServerSocketChannel.class; boolean epollEnabled = serverConfig.isRestServiceEpollEnabled(); @@ -165,11 +164,29 @@ public ListenerService( if (isGrpcEnabled && grpcServer == null) { List interceptors = channelInitializer.initGrpcInterceptors(); - VeniceServerGrpcRequestProcessor requestProcessor = channelInitializer.initGrpcRequestProcessor(); - grpcExecutor = createThreadPool(serverConfig.getGrpcWorkerThreadCount(), "GrpcWorkerThread", nettyBacklogSize); + grpcExecutor = createThreadPool( + serverConfig.getGrpcWorkerThreadCount(), + "GrpcWorkerThread", + serverConfig.getDatabaseLookupQueueCapacity()); + + /** + * Ideally, we would not pass the NettyHandlers to the gRPC service, but since the common code resides within + * these handlers, we are passing it for now. This should be refactored in the future so that the common code is + * shared between both services, and the handlers are only used for their respective services. + * For example, common functionality should be extracted from {@link StorageReadRequestHandler} and shared between + * {@link HttpChannelInitializer} and {@link VeniceGrpcReadServiceImpl}. + */ + GrpcServiceDependencies dependencies = + new GrpcServiceDependencies.Builder().setDiskHealthCheckService(diskHealthService) + .setQuotaEnforcementHandler(channelInitializer.getQuotaEnforcer()) + .setStorageReadRequestHandler(storageReadRequestHandler) + .setSingleGetStats(channelInitializer.getSingleGetStats()) + .setMultiGetStats(channelInitializer.getMultiGetStats()) + .setComputeStats(channelInitializer.getComputeStats()) + .build(); VeniceGrpcServerConfig.Builder grpcServerBuilder = new VeniceGrpcServerConfig.Builder().setPort(grpcPort) - .setService(new VeniceReadServiceImpl(requestProcessor)) + .setService(new VeniceGrpcReadServiceImpl(dependencies)) .setExecutor(grpcExecutor) .setInterceptors(interceptors); diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/NoOpReadQuotaEnforcementHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/NoOpReadQuotaEnforcementHandler.java new file mode 100644 index 0000000000..a42fd3eddc --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/NoOpReadQuotaEnforcementHandler.java @@ -0,0 +1,23 @@ +package com.linkedin.venice.listener; + +import com.linkedin.venice.listener.request.RouterRequest; + + +/** + * A no-op implementation of {@link QuotaEnforcementHandler} that allows all requests. + */ +public class NoOpReadQuotaEnforcementHandler implements QuotaEnforcementHandler { + private static final NoOpReadQuotaEnforcementHandler INSTANCE = new NoOpReadQuotaEnforcementHandler(); + + private NoOpReadQuotaEnforcementHandler() { + } + + public static NoOpReadQuotaEnforcementHandler getInstance() { + return INSTANCE; + } + + @Override + public QuotaEnforcementResult enforceQuota(RouterRequest request) { + return QuotaEnforcementResult.ALLOWED; + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/OutboundHttpWrapperHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/OutboundHttpWrapperHandler.java index 84b413ee2c..e4d294b5aa 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/OutboundHttpWrapperHandler.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/OutboundHttpWrapperHandler.java @@ -1,10 +1,11 @@ package com.linkedin.venice.listener; +import static com.linkedin.venice.response.VeniceReadResponseStatus.INTERNAL_SERVER_ERROR; +import static com.linkedin.venice.response.VeniceReadResponseStatus.KEY_NOT_FOUND; +import static com.linkedin.venice.response.VeniceReadResponseStatus.MISROUTED_STORE_VERSION; +import static com.linkedin.venice.response.VeniceReadResponseStatus.OK; import static io.netty.handler.codec.http.HttpHeaders.Names.CONTENT_LENGTH; import static io.netty.handler.codec.http.HttpHeaders.Names.CONTENT_TYPE; -import static io.netty.handler.codec.http.HttpResponseStatus.INTERNAL_SERVER_ERROR; -import static io.netty.handler.codec.http.HttpResponseStatus.NOT_FOUND; -import static io.netty.handler.codec.http.HttpResponseStatus.OK; import static io.netty.handler.codec.http.HttpVersion.HTTP_1_1; import com.fasterxml.jackson.databind.ObjectMapper; @@ -17,6 +18,7 @@ import com.linkedin.venice.listener.response.AbstractReadResponse; import com.linkedin.venice.listener.response.BinaryResponse; import com.linkedin.venice.listener.response.HttpShortcutResponse; +import com.linkedin.venice.response.VeniceReadResponseStatus; import com.linkedin.venice.utils.ExceptionUtils; import com.linkedin.venice.utils.ObjectMapperFactory; import io.netty.buffer.ByteBuf; @@ -26,7 +28,6 @@ import io.netty.channel.ChannelPromise; import io.netty.handler.codec.http.DefaultFullHttpResponse; import io.netty.handler.codec.http.FullHttpResponse; -import io.netty.handler.codec.http.HttpResponseStatus; import java.nio.charset.StandardCharsets; @@ -35,19 +36,19 @@ */ public class OutboundHttpWrapperHandler extends ChannelOutboundHandlerAdapter { - private final StatsHandler statsHandler; + private final RequestStatsRecorder requestStatsRecorder; private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getInstance(); - public OutboundHttpWrapperHandler(StatsHandler handler) { + public OutboundHttpWrapperHandler(RequestStatsRecorder requestStatsRecorder) { super(); - statsHandler = handler; + this.requestStatsRecorder = requestStatsRecorder; } @Override public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) { ByteBuf body; String contentType = HttpConstants.AVRO_BINARY; - HttpResponseStatus responseStatus = OK; + VeniceReadResponseStatus responseStatus = OK; int schemaIdHeader = -1; int responseRcu = 1; CompressionStrategy compressionStrategy = CompressionStrategy.NO_OP; @@ -55,25 +56,23 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) try { if (msg instanceof AbstractReadResponse) { AbstractReadResponse obj = (AbstractReadResponse) msg; - ServerStatsContext statsContext = statsHandler.getServerStatsContext(); - setStats(statsContext, obj); - + requestStatsRecorder.setReadResponseStats(obj.getReadResponseStatsRecorder()); compressionStrategy = obj.getCompressionStrategy(); if (obj.isFound()) { body = obj.getResponseBody(); schemaIdHeader = obj.getResponseSchemaIdHeader(); - statsContext.setResponseSize(body.readableBytes()); + requestStatsRecorder.setResponseSize(body.readableBytes()); } else { body = Unpooled.EMPTY_BUFFER; - responseStatus = NOT_FOUND; - statsContext.setResponseSize(0); + responseStatus = KEY_NOT_FOUND; + requestStatsRecorder.setResponseSize(0); } isStreamingResponse = obj.isStreamingResponse(); responseRcu = obj.getRCU(); } else if (msg instanceof HttpShortcutResponse) { // For Early terminated requests HttpShortcutResponse shortcutResponse = (HttpShortcutResponse) msg; - responseStatus = shortcutResponse.getStatus(); + responseStatus = VeniceReadResponseStatus.fromCode(shortcutResponse.getStatus().code()); String message = shortcutResponse.getMessage(); if (message == null) { body = Unpooled.EMPTY_BUFFER; @@ -81,10 +80,13 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) body = Unpooled.wrappedBuffer(message.getBytes(StandardCharsets.UTF_8)); } contentType = HttpConstants.TEXT_PLAIN; - if (shortcutResponse.getStatus().equals(VeniceRequestEarlyTerminationException.getHttpResponseStatus())) { - statsHandler.setRequestTerminatedEarly(); + if (shortcutResponse.getStatus() + .equals(VeniceRequestEarlyTerminationException.getResponseStatusCode().getHttpResponseStatus())) { + requestStatsRecorder.setRequestTerminatedEarly(); + } + if (shortcutResponse.getStatus() == MISROUTED_STORE_VERSION.getHttpResponseStatus()) { + requestStatsRecorder.setMisroutedStoreVersion(true); } - statsHandler.setMisroutedStoreVersionRequest(shortcutResponse.isMisroutedStoreVersion()); } else if (msg instanceof BinaryResponse) { // For dictionary Fetch requests body = ((BinaryResponse) msg).getBody(); @@ -165,10 +167,11 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) .getBytes(StandardCharsets.UTF_8)); contentType = HttpConstants.TEXT_PLAIN; } finally { - statsHandler.setResponseStatus(responseStatus); + // Setting response status is an important step to ensure that metrics are recorded correctly + requestStatsRecorder.setResponseStatus(responseStatus); } - FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, responseStatus, body); + FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, responseStatus.getHttpResponseStatus(), body); response.headers().set(CONTENT_TYPE, contentType); response.headers().set(CONTENT_LENGTH, body.readableBytes()); response.headers().set(HttpConstants.VENICE_COMPRESSION_STRATEGY, compressionStrategy.getValue()); @@ -185,8 +188,4 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) */ ctx.writeAndFlush(response); } - - public void setStats(ServerStatsContext statsContext, AbstractReadResponse obj) { - statsContext.setReadResponseStats(obj.getStatsRecorder()); - } } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/QuotaEnforcementHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/QuotaEnforcementHandler.java new file mode 100644 index 0000000000..d022edcee0 --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/QuotaEnforcementHandler.java @@ -0,0 +1,18 @@ +package com.linkedin.venice.listener; + +import com.linkedin.venice.listener.request.RouterRequest; + + +/** + * Interface for enforcing quota on requests. + */ +public interface QuotaEnforcementHandler { + QuotaEnforcementResult enforceQuota(RouterRequest request); + + enum QuotaEnforcementResult { + ALLOWED, // request is allowed + REJECTED, // too many requests (store level quota enforcement) + OVER_CAPACITY, // server over capacity (server level quota enforcement) + BAD_REQUEST, // bad request + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandler.java index d22e8657a2..1f76e5746f 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandler.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandler.java @@ -1,5 +1,8 @@ package com.linkedin.venice.listener; +import static com.linkedin.venice.response.VeniceReadResponseStatus.BAD_REQUEST; +import static com.linkedin.venice.response.VeniceReadResponseStatus.SERVICE_UNAVAILABLE; +import static com.linkedin.venice.response.VeniceReadResponseStatus.TOO_MANY_REQUESTS; import static com.linkedin.venice.throttle.EventThrottler.REJECT_STRATEGY; import com.linkedin.davinci.config.VeniceServerConfig; @@ -28,7 +31,6 @@ import io.netty.channel.ChannelHandler; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.SimpleChannelInboundHandler; -import io.netty.handler.codec.http.HttpResponseStatus; import io.netty.util.ReferenceCountUtil; import io.tehuti.metrics.MetricsRepository; import java.time.Clock; @@ -43,7 +45,7 @@ @ChannelHandler.Sharable public class ReadQuotaEnforcementHandler extends SimpleChannelInboundHandler - implements RoutingDataRepository.RoutingDataChangedListener, StoreDataChangedListener { + implements RoutingDataRepository.RoutingDataChangedListener, StoreDataChangedListener, QuotaEnforcementHandler { private static final Logger LOGGER = LogManager.getLogger(ReadQuotaEnforcementHandler.class); public static final String SERVER_OVER_CAPACITY_MSG = "Server over capacity"; public static final String INVALID_REQUEST_RESOURCE_MSG = "Invalid request resource: "; @@ -193,31 +195,24 @@ public boolean isInitialized() { return false; } - public enum QuotaEnforcementResult { - ALLOWED, // request is allowed - REJECTED, // too many requests (store level quota enforcement) - OVER_CAPACITY, // server over capacity (server level quota enforcement) - BAD_REQUEST, // bad request - } - /** * Enforce quota for a given request. This is common to both HTTP and GRPC handlers. Respective handlers will * take actions such as retaining the request and passing it to the next handler, or sending an error response. * @param request RouterRequest * @return QuotaEnforcementResult */ - public QuotaEnforcementResult enforceQuota(RouterRequest request) { + public QuotaEnforcementHandler.QuotaEnforcementResult enforceQuota(RouterRequest request) { String storeName = request.getStoreName(); Store store = storeRepository.getStore(storeName); if (store == null) { - return QuotaEnforcementResult.BAD_REQUEST; + return QuotaEnforcementHandler.QuotaEnforcementResult.BAD_REQUEST; } /* * If we haven't completed initialization or store does not have SN read quota enabled, allow all requests */ if (!isInitialized() || !store.isStorageNodeReadQuotaEnabled()) { - return QuotaEnforcementResult.ALLOWED; + return QuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED; } int readCapacityUnits = getRcu(request); @@ -229,7 +224,7 @@ public QuotaEnforcementResult enforceQuota(RouterRequest request) { if (veniceRateLimiter != null) { if (!request.isRetryRequest() && !veniceRateLimiter.tryAcquirePermit(readCapacityUnits)) { stats.recordRejected(request.getStoreName(), readCapacityUnits); - return QuotaEnforcementResult.REJECTED; + return QuotaEnforcementHandler.QuotaEnforcementResult.REJECTED; } } else { // If this happens it is probably due to a short-lived race condition where the resource is being accessed before @@ -242,32 +237,33 @@ public QuotaEnforcementResult enforceQuota(RouterRequest request) { * retried requests need to be throttled at node capacity level */ if (!storageNodeRateLimiter.tryAcquirePermit(readCapacityUnits)) { - return QuotaEnforcementResult.OVER_CAPACITY; + return QuotaEnforcementHandler.QuotaEnforcementResult.OVER_CAPACITY; } stats.recordAllowed(storeName, readCapacityUnits); - return QuotaEnforcementResult.ALLOWED; + return QuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED; } @Override public void channelRead0(ChannelHandlerContext ctx, RouterRequest request) { - QuotaEnforcementResult result = enforceQuota(request); + QuotaEnforcementHandler.QuotaEnforcementResult result = enforceQuota(request); - if (result == QuotaEnforcementResult.BAD_REQUEST) { + if (result == QuotaEnforcementHandler.QuotaEnforcementResult.BAD_REQUEST) { ctx.writeAndFlush( new HttpShortcutResponse( INVALID_REQUEST_RESOURCE_MSG + request.getResourceName(), - HttpResponseStatus.BAD_REQUEST)); + BAD_REQUEST.getHttpResponseStatus())); return; } - if (result == QuotaEnforcementResult.REJECTED) { - ctx.writeAndFlush(new HttpShortcutResponse(HttpResponseStatus.TOO_MANY_REQUESTS)); + if (result == QuotaEnforcementHandler.QuotaEnforcementResult.REJECTED) { + ctx.writeAndFlush(new HttpShortcutResponse(TOO_MANY_REQUESTS.getHttpResponseStatus())); return; } - if (result == QuotaEnforcementResult.OVER_CAPACITY) { - ctx.writeAndFlush(new HttpShortcutResponse(SERVER_OVER_CAPACITY_MSG, HttpResponseStatus.SERVICE_UNAVAILABLE)); + if (result == QuotaEnforcementHandler.QuotaEnforcementResult.OVER_CAPACITY) { + ctx.writeAndFlush( + new HttpShortcutResponse(SERVER_OVER_CAPACITY_MSG, SERVICE_UNAVAILABLE.getHttpResponseStatus())); return; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/RequestStatsRecorder.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/RequestStatsRecorder.java new file mode 100644 index 0000000000..739b22ef2a --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/RequestStatsRecorder.java @@ -0,0 +1,371 @@ +package com.linkedin.venice.listener; + +import static com.linkedin.venice.listener.response.stats.ResponseStatsUtil.consumeDoubleIfAbove; +import static com.linkedin.venice.listener.response.stats.ResponseStatsUtil.consumeIntIfAbove; +import static com.linkedin.venice.response.VeniceReadResponseStatus.MISROUTED_STORE_VERSION; + +import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.listener.request.RouterRequest; +import com.linkedin.venice.listener.response.stats.ReadResponseStatsRecorder; +import com.linkedin.venice.read.RequestType; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import com.linkedin.venice.stats.AggServerHttpRequestStats; +import com.linkedin.venice.stats.ServerHttpRequestStats; +import com.linkedin.venice.utils.LatencyUtils; +import io.netty.channel.ChannelHandlerContext; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * We need to be able to record server side statistics for gRPC requests. The current StatsHandler in the Netty + * Pipeline maintains instance variables per channel, and guarantees that each request will be handled by the same + * thread, thus we cannot augment StatsHandler in the same way as other handlers for gRPC requests. We create a + * new StatsContext for each gRPC request, so we can maintain per request stats which will be aggregated on request + * completion using references to the proper Metrics Repository in AggServerHttpRequestStats. This class is almost a + * direct copy of StatsHandler, without Netty Channel Read/Write logic. + */ +public class RequestStatsRecorder { + private static final Logger LOGGER = LogManager.getLogger(RequestStatsRecorder.class); + + private ReadResponseStatsRecorder responseStatsRecorder = null; + private long startTimeInNS = System.nanoTime(); + private VeniceReadResponseStatus responseStatus = null; + private String storeName = null; + private boolean isMetadataRequest = false; + private int requestKeyCount = -1; + private int requestSizeInBytes = -1; + private boolean isRequestTerminatedEarly = false; + private final AggServerHttpRequestStats singleGetStats; + private final AggServerHttpRequestStats multiGetStats; + private final AggServerHttpRequestStats computeStats; + private AggServerHttpRequestStats currentStats; + + // a flag that indicates if this is a new HttpRequest. Netty is TCP-based, so a HttpRequest is chunked into packages. + // Set the startTimeInNS in ChannelRead if it is the first package within a HttpRequest. + private boolean newRequest = true; + + /** + * To indicate whether the stat callback has been triggered or not for a given request. + * This is mostly to bypass the issue that stat callback could be triggered multiple times for one single request. + */ + private boolean statCallbackExecuted = false; + + /** + * For Netty pipeline: + * Normally, one multi-get request will be split into two parts, and it means + * {@link StatsHandler#channelRead(ChannelHandlerContext, Object)} will be invoked twice. + * + * 'firstPartLatency' will measure the time took by: + * {@link StatsHandler} + * {@link io.netty.handler.codec.http.HttpServerCodec} + * {@link io.netty.handler.codec.http.HttpObjectAggregator} + * + * 'partsInvokeDelayLatency' will measure the delay between the invocation of part1 + * and the invocation of part2; + * + * 'secondPartLatency' will measure the time took by: + * {@link StatsHandler} + * {@link io.netty.handler.codec.http.HttpServerCodec} + * {@link io.netty.handler.codec.http.HttpObjectAggregator} + * {@link VerifySslHandler} + * {@link ServerAclHandler} + * {@link RouterRequestHttpHandler} + * {@link StorageReadRequestHandler} + * + * For gRPC pipeline: Request is not split into parts, so part latency is not applicable. + */ + private double firstPartLatency = -1; + private double secondPartLatency = -1; + private double partsInvokeDelayLatency = -1; + private int requestPartCount = 1; + private boolean isMisroutedStoreVersion = false; + + /** + * TODO: We need to figure out how to record the flush latency for gRPC requests. + */ + private double flushLatency = -1; + private int responseSize = -1; + + public boolean isNewRequest() { + return newRequest; + } + + public RequestStatsRecorder setSecondPartLatency(double secondPartLatency) { + this.secondPartLatency = secondPartLatency; + return this; + } + + public RequestStatsRecorder setPartsInvokeDelayLatency(double partsInvokeDelayLatency) { + this.partsInvokeDelayLatency = partsInvokeDelayLatency; + return this; + } + + public RequestStatsRecorder incrementRequestPartCount() { + this.requestPartCount++; + return this; + } + + public RequestStatsRecorder( + AggServerHttpRequestStats singleGetStats, + AggServerHttpRequestStats multiGetStats, + AggServerHttpRequestStats computeStats) { + this.singleGetStats = singleGetStats; + this.multiGetStats = multiGetStats; + this.computeStats = computeStats; + // default to current stats + this.currentStats = singleGetStats; + } + + public void resetContext() { + this.responseStatsRecorder = null; + storeName = null; + startTimeInNS = System.nanoTime(); + partsInvokeDelayLatency = -1; + secondPartLatency = -1; + requestPartCount = 1; + isMetadataRequest = false; + responseStatus = null; + statCallbackExecuted = false; + requestKeyCount = -1; + requestSizeInBytes = -1; + isRequestTerminatedEarly = false; + isMisroutedStoreVersion = false; + flushLatency = -1; + responseSize = -1; + newRequest = false; // set to false after the first package of a HttpRequest + } + + public RequestStatsRecorder setReadResponseStats(ReadResponseStatsRecorder responseStatsRecorder) { + this.responseStatsRecorder = responseStatsRecorder; + return this; + } + + public RequestStatsRecorder setFirstPartLatency(double firstPartLatency) { + this.firstPartLatency = firstPartLatency; + return this; + } + + public RequestStatsRecorder setNewRequest() { + this.newRequest = true; + return this; + } + + public boolean isMetadataRequest() { + return isMetadataRequest; + } + + public boolean isStatCallBackExecuted() { + return statCallbackExecuted; + } + + public void setStatCallBackExecuted(boolean statCallbackExecuted) { + this.statCallbackExecuted = statCallbackExecuted; + } + + public RequestStatsRecorder setResponseStatus(VeniceReadResponseStatus status) { + this.responseStatus = status; + return this; + } + + public String getStoreName() { + return storeName; + } + + public RequestStatsRecorder setStoreName(String name) { + this.storeName = name; + return this; + } + + public RequestStatsRecorder setMetadataRequest(boolean metadataRequest) { + this.isMetadataRequest = metadataRequest; + return this; + } + + public RequestStatsRecorder setRequestTerminatedEarly() { + this.isRequestTerminatedEarly = true; + return this; + } + + public VeniceReadResponseStatus getVeniceReadResponseStatus() { + return responseStatus; + } + + public RequestStatsRecorder setRequestType(RequestType requestType) { + switch (requestType) { + case MULTI_GET: + case MULTI_GET_STREAMING: + currentStats = multiGetStats; + break; + case COMPUTE: + case COMPUTE_STREAMING: + currentStats = computeStats; + break; + default: + currentStats = singleGetStats; + } + return this; + } + + public RequestStatsRecorder setRequestKeyCount(int keyCount) { + this.requestKeyCount = keyCount; + return this; + } + + public AggServerHttpRequestStats getCurrentStats() { + return currentStats; + } + + public RequestStatsRecorder setRequestInfo(RouterRequest request) { + setStoreName(request.getStoreName()); + setRequestType(request.getRequestType()); + setRequestKeyCount(request.getKeyCount()); + return this; + } + + public RequestStatsRecorder setRequestSize(int requestSizeInBytes) { + this.requestSizeInBytes = requestSizeInBytes; + return this; + } + + public long getRequestStartTimeInNS() { + return this.startTimeInNS; + } + + public RequestStatsRecorder setFlushLatency(double latency) { + this.flushLatency = latency; + return this; + } + + public void setResponseSize(int size) { + this.responseSize = size; + } + + public RequestStatsRecorder recordBasicMetrics(ServerHttpRequestStats serverHttpRequestStats) { + if (serverHttpRequestStats == null) { + return this; + } + + if (this.responseStatsRecorder != null) { + this.responseStatsRecorder.recordMetrics(serverHttpRequestStats); + } + + consumeIntIfAbove(serverHttpRequestStats::recordRequestKeyCount, this.requestKeyCount, 0); + consumeIntIfAbove(serverHttpRequestStats::recordRequestSizeInBytes, this.requestSizeInBytes, 0); + consumeDoubleIfAbove(serverHttpRequestStats::recordRequestFirstPartLatency, this.firstPartLatency, 0); + consumeDoubleIfAbove(serverHttpRequestStats::recordRequestPartsInvokeDelayLatency, this.partsInvokeDelayLatency, 0); + consumeDoubleIfAbove(serverHttpRequestStats::recordRequestSecondPartLatency, this.secondPartLatency, 0); + consumeIntIfAbove(serverHttpRequestStats::recordRequestPartCount, this.requestPartCount, 0); + + if (this.isRequestTerminatedEarly) { + serverHttpRequestStats.recordEarlyTerminatedEarlyRequest(); + } + if (flushLatency >= 0) { + serverHttpRequestStats.recordFlushLatency(flushLatency); + } + if (responseSize >= 0) { + serverHttpRequestStats.recordResponseSize(responseSize); + } + + return this; + } + + /** + * Records a successful request in the statistics. + * This method is called when the {@link responseStatus} is either OK or KEY_NOT_FOUND. + * + * Note: Synchronization is not required here, as operations in Tehuti are already synchronized. + * However, reconsider potential race conditions if new logic is introduced. + */ + public RequestStatsRecorder successRequest(ServerHttpRequestStats stats, double elapsedTime) { + if (stats != null) { + stats.recordSuccessRequest(); + stats.recordSuccessRequestLatency(elapsedTime); + } + return this; + } + + /** + * Records an error request in the statistics. + * This method is called when the {@link responseStatus} is neither OK, KEY_NOT_FOUND, nor TOO_MANY_REQUESTS. + */ + public RequestStatsRecorder errorRequest(ServerHttpRequestStats stats, double elapsedTime) { + if (stats == null) { + currentStats.recordErrorRequest(); + currentStats.recordErrorRequestLatency(elapsedTime); + if (isMisroutedStoreVersion) { + currentStats.recordMisroutedStoreVersionRequest(); + } + } else { + stats.recordErrorRequest(); + stats.recordErrorRequestLatency(elapsedTime); + if (isMisroutedStoreVersion) { + stats.recordMisroutedStoreVersionRequest(); + } + } + + return this; + } + + public int getRequestKeyCount() { + return requestKeyCount; + } + + public void setMisroutedStoreVersion(boolean misroutedStoreVersion) { + isMisroutedStoreVersion = misroutedStoreVersion; + } + + /** + * Records the completion of a request in the statistics. + * This method should be called at the end of the request processing. + */ + public static void recordRequestCompletionStats( + RequestStatsRecorder requestStatsRecorder, + boolean isSuccess, + long flushLatencyNs) { + VeniceReadResponseStatus readResponseStatus = requestStatsRecorder.getVeniceReadResponseStatus(); + if (readResponseStatus == null) { + LOGGER.error( + "Failed to record request stats: response status of request cannot be null", + new VeniceException("response status of request cannot be null")); + return; + } + + String storeName = requestStatsRecorder.getStoreName(); + if (storeName == null) { + LOGGER.error( + "Failed to record request stats: store name of request cannot be null", + new VeniceException("store name of request cannot be null")); + return; + } + + if (readResponseStatus == VeniceRequestEarlyTerminationException.getResponseStatusCode()) { + requestStatsRecorder.setRequestTerminatedEarly(); + } + if (readResponseStatus == MISROUTED_STORE_VERSION) { + requestStatsRecorder.setMisroutedStoreVersion(true); + } + + // Calculate the latency + double requestLatency = LatencyUtils.getElapsedTimeFromNSToMS(requestStatsRecorder.getRequestStartTimeInNS()); + + // Optionally include flush latency if available + if (flushLatencyNs > 0) { + requestStatsRecorder.setFlushLatency(LatencyUtils.getElapsedTimeFromNSToMS(flushLatencyNs)); + } + + ServerHttpRequestStats serverHttpRequestStats = requestStatsRecorder.getCurrentStats().getStoreStats(storeName); + requestStatsRecorder.recordBasicMetrics(serverHttpRequestStats); + + // Record success or error based on the response status + if (isSuccess && (readResponseStatus == VeniceReadResponseStatus.OK + || readResponseStatus == VeniceReadResponseStatus.KEY_NOT_FOUND)) { + requestStatsRecorder.successRequest(serverHttpRequestStats, requestLatency); + } else if (readResponseStatus != VeniceReadResponseStatus.TOO_MANY_REQUESTS) { + requestStatsRecorder.errorRequest(serverHttpRequestStats, requestLatency); + } + + // Mark the callback as executed + requestStatsRecorder.setStatCallBackExecuted(true); + } +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/RouterRequestHttpHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/RouterRequestHttpHandler.java index 3b4918037a..10192e9c10 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/RouterRequestHttpHandler.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/RouterRequestHttpHandler.java @@ -1,5 +1,8 @@ package com.linkedin.venice.listener; +import static com.linkedin.venice.response.VeniceReadResponseStatus.BAD_REQUEST; +import static com.linkedin.venice.response.VeniceReadResponseStatus.INTERNAL_SERVER_ERROR; + import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.listener.request.AdminRequest; import com.linkedin.venice.listener.request.ComputeRouterRequestWrapper; @@ -20,7 +23,6 @@ import io.netty.handler.codec.http.FullHttpRequest; import io.netty.handler.codec.http.HttpMethod; import io.netty.handler.codec.http.HttpRequest; -import io.netty.handler.codec.http.HttpResponseStatus; import io.netty.handler.timeout.IdleState; import io.netty.handler.timeout.IdleStateEvent; import java.net.URI; @@ -42,18 +44,20 @@ */ public class RouterRequestHttpHandler extends SimpleChannelInboundHandler { private static final Logger LOGGER = LogManager.getLogger(RouterRequestHttpHandler.class); - private final StatsHandler statsHandler; + private final RequestStatsRecorder requestStatsRecorder; private final Map storeToEarlyTerminationThresholdMSMap; - public RouterRequestHttpHandler(StatsHandler handler, Map storeToEarlyTerminationThresholdMSMap) { + public RouterRequestHttpHandler( + RequestStatsRecorder requestStatsRecorder, + Map storeToEarlyTerminationThresholdMSMap) { super(); - this.statsHandler = handler; + this.requestStatsRecorder = requestStatsRecorder; this.storeToEarlyTerminationThresholdMSMap = storeToEarlyTerminationThresholdMSMap; } @Override public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { - ctx.writeAndFlush(new HttpShortcutResponse(cause.getMessage(), HttpResponseStatus.INTERNAL_SERVER_ERROR)); + ctx.writeAndFlush(new HttpShortcutResponse(cause.getMessage(), INTERNAL_SERVER_ERROR.getHttpResponseStatus())); ctx.close(); } @@ -67,7 +71,7 @@ private void setupRequestTimeout(RouterRequest routerRequest) { Integer timeoutThresholdInMS = storeToEarlyTerminationThresholdMSMap.get(storeName); if (timeoutThresholdInMS != null) { routerRequest.setRequestTimeoutInNS( - statsHandler.getRequestStartTimeInNS() + TimeUnit.MILLISECONDS.toNanos(timeoutThresholdInMS)); + requestStatsRecorder.getRequestStartTimeInNS() + TimeUnit.MILLISECONDS.toNanos(timeoutThresholdInMS)); } } @@ -77,7 +81,7 @@ protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) thro URI uri = URI.create(req.uri()); String[] requestParts = RequestHelper.getRequestParts(uri); QueryAction action = getQueryActionFromRequest(req, requestParts); - statsHandler.setRequestSize(req.content().readableBytes()); + requestStatsRecorder.setRequestSize(req.content().readableBytes()); switch (action) { case STORAGE: // GET /storage/store/partition/key HttpMethod requestMethod = req.method(); @@ -85,14 +89,14 @@ protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) thro // TODO: evaluate whether we can replace single-get by multi-get GetRouterRequest getRouterRequest = GetRouterRequest.parseGetHttpRequest(req, requestParts); setupRequestTimeout(getRouterRequest); - statsHandler.setRequestInfo(getRouterRequest); + requestStatsRecorder.setRequestInfo(getRouterRequest); ctx.fireChannelRead(getRouterRequest); } else if (requestMethod.equals(HttpMethod.POST)) { // Multi-get MultiGetRouterRequestWrapper multiGetRouterReq = MultiGetRouterRequestWrapper.parseMultiGetHttpRequest(req, requestParts); setupRequestTimeout(multiGetRouterReq); - statsHandler.setRequestInfo(multiGetRouterReq); + requestStatsRecorder.setRequestInfo(multiGetRouterReq); ctx.fireChannelRead(multiGetRouterReq); } else { throw new VeniceException("Unknown request method: " + requestMethod + " for " + QueryAction.STORAGE); @@ -101,53 +105,53 @@ protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest req) thro case COMPUTE: // compute request if (req.method().equals(HttpMethod.POST)) { ComputeRouterRequestWrapper computeRouterReq = - ComputeRouterRequestWrapper.parseComputeRequest(req, requestParts); + ComputeRouterRequestWrapper.parseComputeHttpRequest(req, requestParts); setupRequestTimeout(computeRouterReq); - statsHandler.setRequestInfo(computeRouterReq); + requestStatsRecorder.setRequestInfo(computeRouterReq); ctx.fireChannelRead(computeRouterReq); } else { throw new VeniceException("Only support POST method for " + QueryAction.COMPUTE); } break; case HEALTH: - statsHandler.setMetadataRequest(true); + requestStatsRecorder.setMetadataRequest(true); HealthCheckRequest healthCheckRequest = new HealthCheckRequest(); ctx.fireChannelRead(healthCheckRequest); break; case DICTIONARY: DictionaryFetchRequest dictionaryFetchRequest = DictionaryFetchRequest.parseGetHttpRequest(uri, requestParts); - statsHandler.setStoreName(dictionaryFetchRequest.getStoreName()); + requestStatsRecorder.setStoreName(dictionaryFetchRequest.getStoreName()); ctx.fireChannelRead(dictionaryFetchRequest); break; case ADMIN: AdminRequest adminRequest = AdminRequest.parseAdminHttpRequest(req, uri); - statsHandler.setStoreName(adminRequest.getStoreName()); + requestStatsRecorder.setStoreName(adminRequest.getStoreName()); ctx.fireChannelRead(adminRequest); break; case METADATA: - statsHandler.setMetadataRequest(true); + requestStatsRecorder.setMetadataRequest(true); MetadataFetchRequest metadataFetchRequest = MetadataFetchRequest.parseGetHttpRequest(uri.getPath(), requestParts); - statsHandler.setStoreName(metadataFetchRequest.getStoreName()); + requestStatsRecorder.setStoreName(metadataFetchRequest.getStoreName()); ctx.fireChannelRead(metadataFetchRequest); break; case CURRENT_VERSION: - statsHandler.setMetadataRequest(true); + requestStatsRecorder.setMetadataRequest(true); CurrentVersionRequest currentVersionRequest = CurrentVersionRequest.parseGetHttpRequest(uri, requestParts); - statsHandler.setStoreName(currentVersionRequest.getStoreName()); + requestStatsRecorder.setStoreName(currentVersionRequest.getStoreName()); ctx.fireChannelRead(currentVersionRequest); break; case TOPIC_PARTITION_INGESTION_CONTEXT: TopicPartitionIngestionContextRequest topicPartitionIngestionContextRequest = TopicPartitionIngestionContextRequest.parseGetHttpRequest(uri.getPath(), requestParts); - statsHandler.setStoreName( + requestStatsRecorder.setStoreName( Version.parseStoreFromVersionTopic(topicPartitionIngestionContextRequest.getVersionTopic())); ctx.fireChannelRead(topicPartitionIngestionContextRequest); default: throw new VeniceException("Unrecognized query action"); } } catch (VeniceException e) { - ctx.writeAndFlush(new HttpShortcutResponse(e.getMessage(), HttpResponseStatus.BAD_REQUEST)); + ctx.writeAndFlush(new HttpShortcutResponse(e.getMessage(), BAD_REQUEST.getHttpResponseStatus())); } } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerAclHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerAclHandler.java index 755822c58a..fb86d85ae2 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerAclHandler.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerAclHandler.java @@ -1,6 +1,7 @@ package com.linkedin.venice.listener; import static com.linkedin.venice.listener.ServerHandlerUtils.extractClientCert; +import static com.linkedin.venice.response.VeniceReadResponseStatus.FORBIDDEN; import static io.grpc.Metadata.ASCII_STRING_MARSHALLER; import static io.grpc.Metadata.Key; @@ -20,7 +21,6 @@ import io.netty.channel.ChannelHandlerContext; import io.netty.channel.SimpleChannelInboundHandler; import io.netty.handler.codec.http.HttpRequest; -import io.netty.handler.codec.http.HttpResponseStatus; import io.netty.util.AttributeKey; import io.netty.util.ReferenceCountUtil; import java.security.cert.X509Certificate; @@ -79,7 +79,7 @@ public void channelRead0(ChannelHandlerContext ctx, HttpRequest req) throws SSLP String errLine = String.format("%s requested %s %s", client, method, req.uri()); LOGGER.debug("Unauthorized access rejected: {}", errLine); } - NettyUtils.setupResponseAndFlush(HttpResponseStatus.FORBIDDEN, new byte[0], false, ctx); + NettyUtils.setupResponseAndFlush(FORBIDDEN.getHttpResponseStatus(), new byte[0], false, ctx); } } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerStatsContext.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerStatsContext.java deleted file mode 100644 index b990d4ef91..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerStatsContext.java +++ /dev/null @@ -1,279 +0,0 @@ -package com.linkedin.venice.listener; - -import static com.linkedin.venice.listener.response.stats.ResponseStatsUtil.consumeDoubleIfAbove; -import static com.linkedin.venice.listener.response.stats.ResponseStatsUtil.consumeIntIfAbove; - -import com.linkedin.venice.exceptions.VeniceException; -import com.linkedin.venice.listener.request.RouterRequest; -import com.linkedin.venice.listener.response.stats.ReadResponseStatsRecorder; -import com.linkedin.venice.read.RequestType; -import com.linkedin.venice.stats.AggServerHttpRequestStats; -import com.linkedin.venice.stats.ServerHttpRequestStats; -import io.netty.channel.ChannelHandlerContext; -import io.netty.handler.codec.http.HttpResponseStatus; - - -/** - * We need to be able to record server side statistics for gRPC requests. The current StatsHandler in the Netty - * Pipeline maintains instance variables per channel, and guarantees that each request will be handled by the same - * thread, thus we cannot augment StatsHandler in the same way as other handlers for gRPC requests. We create a - * new StatsContext for each gRPC request, so we can maintain per request stats which will be aggregated on request - * completion using references to the proper Metrics Repository in AggServerHttpRequestStats. This class is almost a - * direct copy of StatsHandler, without Netty Channel Read/Write logic. - */ -public class ServerStatsContext { - private ReadResponseStatsRecorder responseStatsRecorder; - private long startTimeInNS; - private HttpResponseStatus responseStatus; - private String storeName = null; - private boolean isMetadataRequest; - private int requestKeyCount = -1; - private int requestSizeInBytes = -1; - private boolean isRequestTerminatedEarly = false; - private final AggServerHttpRequestStats singleGetStats; - private final AggServerHttpRequestStats multiGetStats; - private final AggServerHttpRequestStats computeStats; - private AggServerHttpRequestStats currentStats; - - // a flag that indicates if this is a new HttpRequest. Netty is TCP-based, so a HttpRequest is chunked into packages. - // Set the startTimeInNS in ChannelRead if it is the first package within a HttpRequest. - private boolean newRequest = true; - /** - * To indicate whether the stat callback has been triggered or not for a given request. - * This is mostly to bypass the issue that stat callback could be triggered multiple times for one single request. - */ - private boolean statCallbackExecuted = false; - - /** - * Normally, one multi-get request will be split into two parts, and it means - * {@link StatsHandler#channelRead(ChannelHandlerContext, Object)} will be invoked twice. - * - * 'firstPartLatency' will measure the time took by: - * {@link StatsHandler} - * {@link io.netty.handler.codec.http.HttpServerCodec} - * {@link io.netty.handler.codec.http.HttpObjectAggregator} - * - * 'partsInvokeDelayLatency' will measure the delay between the invocation of part1 - * and the invocation of part2; - * - * 'secondPartLatency' will measure the time took by: - * {@link StatsHandler} - * {@link io.netty.handler.codec.http.HttpServerCodec} - * {@link io.netty.handler.codec.http.HttpObjectAggregator} - * {@link VerifySslHandler} - * {@link ServerAclHandler} - * {@link RouterRequestHttpHandler} - * {@link StorageReadRequestHandler} - * - */ - private double firstPartLatency = -1; - private double secondPartLatency = -1; - private double partsInvokeDelayLatency = -1; - private int requestPartCount = -1; - private boolean isMisroutedStoreVersion = false; - private double flushLatency = -1; - private int responseSize = -1; - - public boolean isNewRequest() { - return newRequest; - } - - public void setSecondPartLatency(double secondPartLatency) { - this.secondPartLatency = secondPartLatency; - } - - public void setPartsInvokeDelayLatency(double partsInvokeDelayLatency) { - this.partsInvokeDelayLatency = partsInvokeDelayLatency; - } - - public void incrementRequestPartCount() { - this.requestPartCount++; - } - - public ServerStatsContext( - AggServerHttpRequestStats singleGetStats, - AggServerHttpRequestStats multiGetStats, - AggServerHttpRequestStats computeStats) { - this.singleGetStats = singleGetStats; - this.multiGetStats = multiGetStats; - this.computeStats = computeStats; - // default to current stats - this.currentStats = singleGetStats; - } - - public void resetContext() { - this.responseStatsRecorder = null; - storeName = null; - startTimeInNS = System.nanoTime(); - partsInvokeDelayLatency = -1; - secondPartLatency = -1; - requestPartCount = 1; - isMetadataRequest = false; - responseStatus = null; - statCallbackExecuted = false; - requestKeyCount = -1; - requestSizeInBytes = -1; - isRequestTerminatedEarly = false; - isMisroutedStoreVersion = false; - flushLatency = -1; - responseSize = -1; - - newRequest = false; - } - - public void setReadResponseStats(ReadResponseStatsRecorder responseStatsRecorder) { - this.responseStatsRecorder = responseStatsRecorder; - } - - public void setFirstPartLatency(double firstPartLatency) { - this.firstPartLatency = firstPartLatency; - } - - public void setNewRequest() { - this.newRequest = true; - } - - public boolean isMetadataRequest() { - return isMetadataRequest; - } - - public boolean isStatCallBackExecuted() { - return statCallbackExecuted; - } - - public void setStatCallBackExecuted(boolean statCallbackExecuted) { - this.statCallbackExecuted = statCallbackExecuted; - } - - public void setResponseStatus(HttpResponseStatus status) { - this.responseStatus = status; - } - - public String getStoreName() { - return storeName; - } - - public void setStoreName(String name) { - this.storeName = name; - } - - public void setMetadataRequest(boolean metadataRequest) { - this.isMetadataRequest = metadataRequest; - } - - public void setRequestTerminatedEarly() { - this.isRequestTerminatedEarly = true; - } - - public HttpResponseStatus getResponseStatus() { - return responseStatus; - } - - public void setRequestType(RequestType requestType) { - switch (requestType) { - case MULTI_GET: - case MULTI_GET_STREAMING: - currentStats = multiGetStats; - break; - case COMPUTE: - case COMPUTE_STREAMING: - currentStats = computeStats; - break; - default: - currentStats = singleGetStats; - } - } - - public void setRequestKeyCount(int keyCount) { - this.requestKeyCount = keyCount; - } - - public AggServerHttpRequestStats getCurrentStats() { - return currentStats; - } - - public void setRequestInfo(RouterRequest request) { - setStoreName(request.getStoreName()); - setRequestType(request.getRequestType()); - setRequestKeyCount(request.getKeyCount()); - } - - public void setRequestSize(int requestSizeInBytes) { - this.requestSizeInBytes = requestSizeInBytes; - } - - public long getRequestStartTimeInNS() { - return this.startTimeInNS; - } - - public void setFlushLatency(double latency) { - this.flushLatency = latency; - } - - public void setResponseSize(int size) { - this.responseSize = size; - } - - public void recordBasicMetrics(ServerHttpRequestStats serverHttpRequestStats) { - if (serverHttpRequestStats != null) { - if (this.responseStatsRecorder != null) { - this.responseStatsRecorder.recordMetrics(serverHttpRequestStats); - } - - consumeIntIfAbove(serverHttpRequestStats::recordRequestKeyCount, this.requestKeyCount, 0); - consumeIntIfAbove(serverHttpRequestStats::recordRequestSizeInBytes, this.requestSizeInBytes, 0); - consumeDoubleIfAbove(serverHttpRequestStats::recordRequestFirstPartLatency, this.firstPartLatency, 0); - consumeDoubleIfAbove( - serverHttpRequestStats::recordRequestPartsInvokeDelayLatency, - this.partsInvokeDelayLatency, - 0); - consumeDoubleIfAbove(serverHttpRequestStats::recordRequestSecondPartLatency, this.secondPartLatency, 0); - consumeIntIfAbove(serverHttpRequestStats::recordRequestPartCount, this.requestPartCount, 0); - - if (this.isRequestTerminatedEarly) { - serverHttpRequestStats.recordEarlyTerminatedEarlyRequest(); - } - if (flushLatency >= 0) { - serverHttpRequestStats.recordFlushLatency(flushLatency); - } - if (responseSize >= 0) { - serverHttpRequestStats.recordResponseSize(responseSize); - } - } - } - - // This method does not have to be synchronized since operations in Tehuti are already synchronized. - // Please re-consider the race condition if new logic is added. - public void successRequest(ServerHttpRequestStats stats, double elapsedTime) { - if (stats != null) { - stats.recordSuccessRequest(); - stats.recordSuccessRequestLatency(elapsedTime); - } else { - throw new VeniceException("store name could not be null if request succeeded"); - } - } - - public void errorRequest(ServerHttpRequestStats stats, double elapsedTime) { - if (stats == null) { - currentStats.recordErrorRequest(); - currentStats.recordErrorRequestLatency(elapsedTime); - if (isMisroutedStoreVersion) { - currentStats.recordMisroutedStoreVersionRequest(); - } - } else { - stats.recordErrorRequest(); - stats.recordErrorRequestLatency(elapsedTime); - if (isMisroutedStoreVersion) { - stats.recordMisroutedStoreVersionRequest(); - } - } - } - - public int getRequestKeyCount() { - return requestKeyCount; - } - - public void setMisroutedStoreVersion(boolean misroutedStoreVersion) { - isMisroutedStoreVersion = misroutedStoreVersion; - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/StatsHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/StatsHandler.java index f21bc3b32b..fc06f87e52 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/StatsHandler.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/StatsHandler.java @@ -1,101 +1,47 @@ package com.linkedin.venice.listener; -import static io.netty.handler.codec.http.HttpResponseStatus.NOT_FOUND; -import static io.netty.handler.codec.http.HttpResponseStatus.OK; -import static io.netty.handler.codec.http.HttpResponseStatus.TOO_MANY_REQUESTS; - import com.linkedin.venice.exceptions.VeniceException; -import com.linkedin.venice.listener.request.RouterRequest; -import com.linkedin.venice.stats.AggServerHttpRequestStats; -import com.linkedin.venice.stats.ServerHttpRequestStats; import com.linkedin.venice.utils.LatencyUtils; import io.netty.channel.ChannelDuplexHandler; import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelPromise; -import io.netty.handler.codec.http.HttpResponseStatus; +import java.util.Objects; +/** + * This class is responsible for handling the stats for the HTTP request handled by the Netty server. + */ public class StatsHandler extends ChannelDuplexHandler { - private final ServerStatsContext serverStatsContext; - private final AggServerHttpRequestStats singleGetStats; - private final AggServerHttpRequestStats multiGetStats; - private final AggServerHttpRequestStats computeStats; - - public StatsHandler( - AggServerHttpRequestStats singleGetStats, - AggServerHttpRequestStats multiGetStats, - AggServerHttpRequestStats computeStats) { - this.singleGetStats = singleGetStats; - this.multiGetStats = multiGetStats; - this.computeStats = computeStats; - - this.serverStatsContext = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); - } + private final RequestStatsRecorder requestStatsRecorder; - public ServerStatsContext getNewStatsContext() { - return new ServerStatsContext(singleGetStats, multiGetStats, computeStats); - } - - public void setResponseStatus(HttpResponseStatus status) { - serverStatsContext.setResponseStatus(status); - } - - public void setStoreName(String name) { - serverStatsContext.setStoreName(name); - } - - public void setMetadataRequest(boolean metadataRequest) { - serverStatsContext.setMetadataRequest(metadataRequest); - } - - public void setRequestTerminatedEarly() { - serverStatsContext.setRequestTerminatedEarly(); - } - - public void setRequestInfo(RouterRequest request) { - serverStatsContext.setRequestInfo(request); - } - - public void setRequestSize(int requestSizeInBytes) { - serverStatsContext.setRequestSize(requestSizeInBytes); - } - - public long getRequestStartTimeInNS() { - return serverStatsContext.getRequestStartTimeInNS(); - } - - public ServerStatsContext getServerStatsContext() { - return serverStatsContext; - } - - public void setMisroutedStoreVersionRequest(boolean misroutedStoreVersionRequest) { - serverStatsContext.setMisroutedStoreVersion(misroutedStoreVersionRequest); + public StatsHandler(RequestStatsRecorder requestStatsRecorder) { + this.requestStatsRecorder = Objects.requireNonNull(requestStatsRecorder, "RequestStatsContext cannot be null"); } @Override public void channelRead(ChannelHandlerContext ctx, Object msg) { - if (serverStatsContext.isNewRequest()) { + if (requestStatsRecorder.isNewRequest()) { // Reset for every request - serverStatsContext.resetContext(); + requestStatsRecorder.resetContext(); /** * For a single 'channelRead' invocation, Netty will guarantee all the following 'channelRead' functions * registered by the pipeline to be executed in the same thread. */ ctx.fireChannelRead(msg); - double firstPartLatency = LatencyUtils.getElapsedTimeFromNSToMS(serverStatsContext.getRequestStartTimeInNS()); - serverStatsContext.setFirstPartLatency(firstPartLatency); + double firstPartLatency = LatencyUtils.getElapsedTimeFromNSToMS(requestStatsRecorder.getRequestStartTimeInNS()); + requestStatsRecorder.setFirstPartLatency(firstPartLatency); } else { // Only works for multi-get request. long startTimeOfPart2InNS = System.nanoTime(); - long startTimeInNS = serverStatsContext.getRequestStartTimeInNS(); + long startTimeInNS = requestStatsRecorder.getRequestStartTimeInNS(); - serverStatsContext.setPartsInvokeDelayLatency(LatencyUtils.convertNSToMS(startTimeOfPart2InNS - startTimeInNS)); + requestStatsRecorder.setPartsInvokeDelayLatency(LatencyUtils.convertNSToMS(startTimeOfPart2InNS - startTimeInNS)); ctx.fireChannelRead(msg); - serverStatsContext.setSecondPartLatency(LatencyUtils.getElapsedTimeFromNSToMS(startTimeOfPart2InNS)); - serverStatsContext.incrementRequestPartCount(); + requestStatsRecorder.setSecondPartLatency(LatencyUtils.getElapsedTimeFromNSToMS(startTimeOfPart2InNS)); + requestStatsRecorder.incrementRequestPartCount(); } } @@ -106,14 +52,8 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) future.addListener((result) -> { // reset the StatsHandler for the new request. This is necessary since instances are channel-based // and channels are ready for the future requests as soon as the current has been handled. - serverStatsContext.setNewRequest(); - - if (serverStatsContext.getResponseStatus() == null) { - throw new VeniceException("request status could not be null"); - } - - // we don't record if it is a metatadata request - if (serverStatsContext.isMetadataRequest()) { + requestStatsRecorder.setNewRequest(); + if (requestStatsRecorder.isMetadataRequest()) { return; } @@ -121,24 +61,9 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) * TODO: Need to do more investigation to figure out why this callback could be triggered * multiple times for a single request */ - if (!serverStatsContext.isStatCallBackExecuted()) { - serverStatsContext.setFlushLatency(LatencyUtils.getElapsedTimeFromNSToMS(beforeFlushTimestampNs)); - ServerHttpRequestStats serverHttpRequestStats = serverStatsContext.getStoreName() == null - ? null - : serverStatsContext.getCurrentStats().getStoreStats(serverStatsContext.getStoreName()); - serverStatsContext.recordBasicMetrics(serverHttpRequestStats); - double elapsedTime = LatencyUtils.getElapsedTimeFromNSToMS(serverStatsContext.getRequestStartTimeInNS()); - // if ResponseStatus is either OK or NOT_FOUND and the channel write is succeed, - // records a successRequest in stats. Otherwise, records a errorRequest in stats - // For TOO_MANY_REQUESTS do not record either success or error. Recording as success would give out - // wrong interpretation of latency, recording error would give out impression that server failed to serve - if (result.isSuccess() && (serverStatsContext.getResponseStatus().equals(OK) - || serverStatsContext.getResponseStatus().equals(NOT_FOUND))) { - serverStatsContext.successRequest(serverHttpRequestStats, elapsedTime); - } else if (!serverStatsContext.getResponseStatus().equals(TOO_MANY_REQUESTS)) { - serverStatsContext.errorRequest(serverHttpRequestStats, elapsedTime); - } - serverStatsContext.setStatCallBackExecuted(true); + if (!requestStatsRecorder.isStatCallBackExecuted()) { + RequestStatsRecorder + .recordRequestCompletionStats(requestStatsRecorder, result.isSuccess(), beforeFlushTimestampNs); } }); } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/StorageReadRequestHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/StorageReadRequestHandler.java index 90ba8a49a1..7853a92e13 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/StorageReadRequestHandler.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/StorageReadRequestHandler.java @@ -56,6 +56,7 @@ import com.linkedin.venice.meta.Version; import com.linkedin.venice.read.protocol.request.router.MultiGetRouterRequestKeyV1; import com.linkedin.venice.read.protocol.response.MultiGetResponseRecordV1; +import com.linkedin.venice.response.VeniceReadResponseStatus; import com.linkedin.venice.schema.SchemaData; import com.linkedin.venice.schema.SchemaEntry; import com.linkedin.venice.serialization.AvroStoreDeserializerCache; @@ -74,7 +75,6 @@ import io.netty.channel.ChannelHandler; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelInboundHandlerAdapter; -import io.netty.handler.codec.http.HttpResponseStatus; import java.nio.ByteBuffer; import java.util.HashMap; import java.util.LinkedHashMap; @@ -106,6 +106,8 @@ public class StorageReadRequestHandler extends ChannelInboundHandlerAdapter { private static final Logger LOGGER = LogManager.getLogger(StorageReadRequestHandler.class); private static final RedundantExceptionFilter REDUNDANT_LOGGING_FILTER = RedundantExceptionFilter.getRedundantExceptionFilter(); + public static final String VENICE_STORAGE_NODE_HARDWARE_IS_NOT_HEALTHY_MSG = + "Venice storage node hardware is not healthy!"; private final DiskHealthCheckService diskHealthCheckService; private final ThreadPoolExecutor executor; private final ThreadPoolExecutor computeExecutor; @@ -269,95 +271,39 @@ public StorageReadRequestHandler( @Override public void channelRead(ChannelHandlerContext context, Object message) throws Exception { if (message instanceof RouterRequest) { - RouterRequest request = (RouterRequest) message; - this.resourceReadUsageTracker.accept(request.getResourceName()); - // Check before putting the request to the intermediate queue - if (request.shouldRequestBeTerminatedEarly()) { - // Try to make the response short - context.writeAndFlush( - new HttpShortcutResponse( - VeniceRequestEarlyTerminationException.getMessage(request.getStoreName()), - VeniceRequestEarlyTerminationException.getHttpResponseStatus())); - return; - } - - CompletableFuture responseFuture; - switch (request.getRequestType()) { - case SINGLE_GET: - responseFuture = handleSingleGetRequest((GetRouterRequest) request); - break; - case MULTI_GET: - responseFuture = this.multiGetHandler.apply((MultiGetRouterRequestWrapper) request); - break; - case COMPUTE: - responseFuture = this.computeHandler.apply((ComputeRouterRequestWrapper) request); - break; - default: - throw new VeniceException("Unknown request type: " + request.getRequestType()); - } - - responseFuture.whenComplete((response, throwable) -> { - if (throwable == null) { - response.setRCU(ReadQuotaEnforcementHandler.getRcu(request)); - if (request.isStreamingRequest()) { - response.setStreamingResponse(); - } - context.writeAndFlush(response); - return; - } - if (throwable instanceof CompletionException && throwable.getCause() != null) { - throwable = throwable.getCause(); - } - if (throwable instanceof VeniceNoStoreException) { - VeniceNoStoreException e = (VeniceNoStoreException) throwable; - String msg = "No storage exists for store: " + e.getStoreName(); - if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) { - LOGGER.error(msg, e); - } - HttpResponseStatus status = getHttpResponseStatus(e); - context.writeAndFlush(new HttpShortcutResponse("No storage exists for: " + e.getStoreName(), status)); - } else if (throwable instanceof VeniceRequestEarlyTerminationException) { - VeniceRequestEarlyTerminationException e = (VeniceRequestEarlyTerminationException) throwable; - String msg = "Request timed out for store: " + e.getStoreName(); - if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) { - LOGGER.error(msg, e); - } - context.writeAndFlush(new HttpShortcutResponse(e.getMessage(), HttpResponseStatus.REQUEST_TIMEOUT)); - } else if (throwable instanceof OperationNotAllowedException) { - OperationNotAllowedException e = (OperationNotAllowedException) throwable; - String msg = "METHOD_NOT_ALLOWED: " + e.getMessage(); - if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) { - LOGGER.error(msg, e); - } - context.writeAndFlush(new HttpShortcutResponse(e.getMessage(), HttpResponseStatus.METHOD_NOT_ALLOWED)); - } else { - LOGGER.error("Exception thrown for {}", request.getResourceName(), throwable); - HttpShortcutResponse shortcutResponse = - new HttpShortcutResponse(throwable.getMessage(), HttpResponseStatus.INTERNAL_SERVER_ERROR); - shortcutResponse.setMisroutedStoreVersion(checkMisroutedStoreVersionRequest(request)); - context.writeAndFlush(shortcutResponse); - } - }); + // IO requests are processed in a separate thread pool + queueIoRequestForAsyncProcessing((RouterRequest) message, HttpStorageResponseHandlerCallback.create(context)); + return; + } - } else if (message instanceof HealthCheckRequest) { + if (message instanceof HealthCheckRequest) { if (diskHealthCheckService.isDiskHealthy()) { - context.writeAndFlush(new HttpShortcutResponse("OK", HttpResponseStatus.OK)); + context.writeAndFlush(new HttpShortcutResponse("OK", VeniceReadResponseStatus.OK.getHttpResponseStatus())); } else { context.writeAndFlush( new HttpShortcutResponse( - "Venice storage node hardware is not healthy!", - HttpResponseStatus.INTERNAL_SERVER_ERROR)); + VENICE_STORAGE_NODE_HARDWARE_IS_NOT_HEALTHY_MSG, + VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getHttpResponseStatus())); LOGGER.error( "Disk is not healthy according to the disk health check service: {}", diskHealthCheckService.getErrorMessage()); } - } else if (message instanceof DictionaryFetchRequest) { + return; + } + + if (message instanceof DictionaryFetchRequest) { BinaryResponse response = handleDictionaryFetchRequest((DictionaryFetchRequest) message); context.writeAndFlush(response); - } else if (message instanceof AdminRequest) { + return; + } + + if (message instanceof AdminRequest) { AdminResponse response = handleServerAdminRequest((AdminRequest) message); context.writeAndFlush(response); - } else if (message instanceof MetadataFetchRequest) { + return; + } + + if (message instanceof MetadataFetchRequest) { try { MetadataResponse response = handleMetadataFetchRequest((MetadataFetchRequest) message); context.writeAndFlush(response); @@ -365,35 +311,128 @@ public void channelRead(ChannelHandlerContext context, Object message) throws Ex LOGGER.warn( "Metadata requested by a storage node read quota not enabled store: {}", ((MetadataFetchRequest) message).getStoreName()); - context.writeAndFlush(new HttpShortcutResponse(e.getMessage(), HttpResponseStatus.FORBIDDEN)); + context.writeAndFlush( + new HttpShortcutResponse(e.getMessage(), VeniceReadResponseStatus.FORBIDDEN.getHttpResponseStatus())); } - } else if (message instanceof CurrentVersionRequest) { + return; + } + + if (message instanceof CurrentVersionRequest) { ServerCurrentVersionResponse response = handleCurrentVersionRequest((CurrentVersionRequest) message); context.writeAndFlush(response); - } else if (message instanceof TopicPartitionIngestionContextRequest) { + return; + } + + if (message instanceof TopicPartitionIngestionContextRequest) { TopicPartitionIngestionContextResponse response = handleTopicPartitionIngestionContextRequest((TopicPartitionIngestionContextRequest) message); context.writeAndFlush(response); - } else { - context.writeAndFlush( - new HttpShortcutResponse( - "Unrecognized object in StorageExecutionHandler", - HttpResponseStatus.INTERNAL_SERVER_ERROR)); + return; } + + context.writeAndFlush( + new HttpShortcutResponse( + "Unrecognized request type", + VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getHttpResponseStatus())); } - private HttpResponseStatus getHttpResponseStatus(VeniceNoStoreException e) { + /** + * Handles requests that require a storage engine lookup. + */ + public void queueIoRequestForAsyncProcessing(RouterRequest request, StorageResponseHandlerCallback responseCallback) { + this.resourceReadUsageTracker.accept(request.getResourceName()); + + // Check if timeout has occurred before processing the request; if so, return early with an error response + if (request.shouldRequestBeTerminatedEarly()) { + responseCallback.onError( + VeniceRequestEarlyTerminationException.getResponseStatusCode(), + VeniceRequestEarlyTerminationException.getMessage(request.getStoreName())); + return; + } + + CompletableFuture responseFuture; + switch (request.getRequestType()) { + case SINGLE_GET: + responseFuture = handleSingleGetRequest((GetRouterRequest) request); + break; + case MULTI_GET: + responseFuture = this.multiGetHandler.apply((MultiGetRouterRequestWrapper) request); + break; + case COMPUTE: + responseFuture = this.computeHandler.apply((ComputeRouterRequestWrapper) request); + break; + default: + throw new VeniceException("Unknown request type: " + request.getRequestType()); + } + + responseFuture.whenComplete((response, throwable) -> { + // If the throwable is null, it indicates that the storage engine lookup was successful + // (regardless of whether the value was found or not), and we have a response ready to send. + if (throwable == null) { + response.setRCU(ReadQuotaEnforcementHandler.getRcu(request)); + if (request.isStreamingRequest()) { + response.setStreamingResponse(); + } + responseCallback.onReadResponse(response); + return; + } + + if (throwable instanceof CompletionException && throwable.getCause() != null) { + throwable = throwable.getCause(); + } + + if (throwable instanceof VeniceNoStoreException) { + VeniceNoStoreException e = (VeniceNoStoreException) throwable; + String msg = "No storage exists for store: " + e.getStoreName(); + if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) { + LOGGER.error(msg, e); + } + VeniceReadResponseStatus status = getVeniceReadResponseStatus(e); + responseCallback.onError(status, "No storage exists for: " + e.getStoreName()); + return; + } + + if (throwable instanceof VeniceRequestEarlyTerminationException) { + VeniceRequestEarlyTerminationException e = (VeniceRequestEarlyTerminationException) throwable; + String msg = "Request timed out for store: " + e.getStoreName(); + if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) { + LOGGER.error(msg, e); + } + responseCallback.onError(VeniceReadResponseStatus.REQUEST_TIMEOUT, e.getMessage()); + return; + } + + if (throwable instanceof OperationNotAllowedException) { + OperationNotAllowedException e = (OperationNotAllowedException) throwable; + String msg = "METHOD_NOT_ALLOWED: " + e.getMessage(); + if (!REDUNDANT_LOGGING_FILTER.isRedundantException(msg)) { + LOGGER.error(msg, e); + } + responseCallback.onError(VeniceReadResponseStatus.METHOD_NOT_ALLOWED, e.getMessage()); + return; + } + + LOGGER.error("Exception thrown for {}", request.getResourceName(), throwable); + if (checkMisroutedStoreVersionRequest(request)) { + responseCallback.onError(VeniceReadResponseStatus.MISROUTED_STORE_VERSION, throwable.getMessage()); + return; + } + responseCallback.onError(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR, throwable.getMessage()); + }); + } + + private VeniceReadResponseStatus getVeniceReadResponseStatus(VeniceNoStoreException e) { String topic = e.getStoreName(); String storeName = Version.parseStoreFromKafkaTopicName(topic); int version = Version.parseVersionFromKafkaTopicName(topic); Store store = metadataRepository.getStore(storeName); if (store == null || store.getCurrentVersion() != version) { - return HttpResponseStatus.BAD_REQUEST; + return VeniceReadResponseStatus.BAD_REQUEST; } // return SERVICE_UNAVAILABLE to kick off error retry in router when store version resource exists - return HttpResponseStatus.SERVICE_UNAVAILABLE; + return VeniceReadResponseStatus.SERVICE_UNAVAILABLE; } /** @@ -768,16 +807,19 @@ record = new ComputeResponseRecordV1(); incrementOperatorCounters(response.getStats(), requestContext.operations, hits); } - private BinaryResponse handleDictionaryFetchRequest(DictionaryFetchRequest request) { + /** + * TODO: Refactor these methods into a a common service class. + */ + public BinaryResponse handleDictionaryFetchRequest(DictionaryFetchRequest request) { ByteBuffer dictionary = ingestionMetadataRetriever.getStoreVersionCompressionDictionary(request.getResourceName()); return new BinaryResponse(dictionary); } - private MetadataResponse handleMetadataFetchRequest(MetadataFetchRequest request) { + public MetadataResponse handleMetadataFetchRequest(MetadataFetchRequest request) { return readMetadataRetriever.getMetadata(request.getStoreName()); } - private ServerCurrentVersionResponse handleCurrentVersionRequest(CurrentVersionRequest request) { + public ServerCurrentVersionResponse handleCurrentVersionRequest(CurrentVersionRequest request) { return readMetadataRetriever.getCurrentVersionResponse(request.getStoreName()); } @@ -822,7 +864,7 @@ private static void incrementOperatorCounters( } } - private AdminResponse handleServerAdminRequest(AdminRequest adminRequest) { + public AdminResponse handleServerAdminRequest(AdminRequest adminRequest) { switch (adminRequest.getServerAdminAction()) { case DUMP_INGESTION_STATE: String topicName = adminRequest.getStoreVersion(); @@ -844,7 +886,7 @@ private AdminResponse handleServerAdminRequest(AdminRequest adminRequest) { } } - private TopicPartitionIngestionContextResponse handleTopicPartitionIngestionContextRequest( + public TopicPartitionIngestionContextResponse handleTopicPartitionIngestionContextRequest( TopicPartitionIngestionContextRequest topicPartitionIngestionContextRequest) { Integer partition = topicPartitionIngestionContextRequest.getPartition(); String versionTopic = topicPartitionIngestionContextRequest.getVersionTopic(); diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/StorageResponseHandlerCallback.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/StorageResponseHandlerCallback.java new file mode 100644 index 0000000000..44b4712428 --- /dev/null +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/StorageResponseHandlerCallback.java @@ -0,0 +1,19 @@ +package com.linkedin.venice.listener; + +import com.linkedin.davinci.listener.response.ReadResponse; +import com.linkedin.venice.response.VeniceReadResponseStatus; + + +/** + * Callback interface to handle the response from the {@link StorageReadRequestHandler#processRequest} method. + */ +public interface StorageResponseHandlerCallback { + void onReadResponse(ReadResponse readResponse); + + /** + * Use this API to send error responses to the client. + * @param readResponseStatus the status of the response + * @param message the message to send to the client + */ + void onError(VeniceReadResponseStatus readResponseStatus, String message); +} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/VeniceRequestEarlyTerminationException.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/VeniceRequestEarlyTerminationException.java index 86bcc81a48..ca7eca0586 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/VeniceRequestEarlyTerminationException.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/VeniceRequestEarlyTerminationException.java @@ -1,7 +1,7 @@ package com.linkedin.venice.listener; import com.linkedin.venice.exceptions.VeniceException; -import io.netty.handler.codec.http.HttpResponseStatus; +import com.linkedin.venice.response.VeniceReadResponseStatus; public class VeniceRequestEarlyTerminationException extends VeniceException { @@ -14,11 +14,11 @@ public VeniceRequestEarlyTerminationException(String storeName) { @Override public int getHttpStatusCode() { - return HttpResponseStatus.REQUEST_TIMEOUT.code(); + return VeniceReadResponseStatus.REQUEST_TIMEOUT.getHttpResponseStatus().code(); } - public static HttpResponseStatus getHttpResponseStatus() { - return HttpResponseStatus.REQUEST_TIMEOUT; + public static VeniceReadResponseStatus getResponseStatusCode() { + return VeniceReadResponseStatus.REQUEST_TIMEOUT; } public static String getMessage(String storeName) { diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/GrpcRequestContext.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/GrpcRequestContext.java deleted file mode 100644 index be28f31c45..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/GrpcRequestContext.java +++ /dev/null @@ -1,99 +0,0 @@ -package com.linkedin.venice.listener.grpc; - -import com.linkedin.davinci.listener.response.ReadResponse; -import com.linkedin.venice.listener.ServerStatsContext; -import com.linkedin.venice.listener.request.RouterRequest; -import com.linkedin.venice.protocols.VeniceClientRequest; -import com.linkedin.venice.protocols.VeniceServerResponse; -import io.grpc.stub.StreamObserver; - - -/** - * We need to keep track of each request as it goes through the pipeline so that we can record the necessary metrics - * and separate different parts of the logic for the response. If a specific handler raises an error, we set - * the hasError flag to true and stop executing the rest of the pipeline excluding the stats collection. - */ -public class GrpcRequestContext { - private VeniceClientRequest veniceClientRequest; - private VeniceServerResponse.Builder veniceServerResponseBuilder; - private StreamObserver responseObserver; - - private boolean isCompleted = false; - private boolean hasError = false; - private RouterRequest routerRequest; - private ReadResponse readResponse; - private ServerStatsContext serverStatsContext; - - public GrpcRequestContext( - VeniceClientRequest veniceClientRequest, - VeniceServerResponse.Builder veniceServerResponseBuilder, - StreamObserver responseObserver) { - this.veniceClientRequest = veniceClientRequest; - this.veniceServerResponseBuilder = veniceServerResponseBuilder; - this.responseObserver = responseObserver; - } - - public void setGrpcStatsContext(ServerStatsContext serverStatsContext) { - this.serverStatsContext = serverStatsContext; - } - - public ServerStatsContext getGrpcStatsContext() { - return serverStatsContext; - } - - public VeniceClientRequest getVeniceClientRequest() { - return veniceClientRequest; - } - - public void setVeniceClientRequest(VeniceClientRequest veniceClientRequest) { - this.veniceClientRequest = veniceClientRequest; - } - - public VeniceServerResponse.Builder getVeniceServerResponseBuilder() { - return veniceServerResponseBuilder; - } - - public void setVeniceServerResponseBuilder(VeniceServerResponse.Builder veniceServerResponseBuilder) { - this.veniceServerResponseBuilder = veniceServerResponseBuilder; - } - - public StreamObserver getResponseObserver() { - return responseObserver; - } - - public void setResponseObserver(StreamObserver responseObserver) { - this.responseObserver = responseObserver; - } - - public RouterRequest getRouterRequest() { - return routerRequest; - } - - public void setRouterRequest(RouterRequest routerRequest) { - this.routerRequest = routerRequest; - } - - public ReadResponse getReadResponse() { - return readResponse; - } - - public void setReadResponse(ReadResponse readResponse) { - this.readResponse = readResponse; - } - - public void setCompleted() { - isCompleted = true; - } - - public boolean isCompleted() { - return isCompleted; - } - - public boolean hasError() { - return hasError; - } - - public void setError() { - hasError = true; - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/VeniceReadServiceImpl.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/VeniceReadServiceImpl.java deleted file mode 100644 index ee52d2f917..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/VeniceReadServiceImpl.java +++ /dev/null @@ -1,43 +0,0 @@ -package com.linkedin.venice.listener.grpc; - -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcRequestProcessor; -import com.linkedin.venice.protocols.VeniceClientRequest; -import com.linkedin.venice.protocols.VeniceReadServiceGrpc; -import com.linkedin.venice.protocols.VeniceServerResponse; -import com.linkedin.venice.response.VeniceReadResponseStatus; -import io.grpc.stub.StreamObserver; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - - -public class VeniceReadServiceImpl extends VeniceReadServiceGrpc.VeniceReadServiceImplBase { - private static final Logger LOGGER = LogManager.getLogger(VeniceReadServiceImpl.class); - - private final VeniceServerGrpcRequestProcessor requestProcessor; - - public VeniceReadServiceImpl(VeniceServerGrpcRequestProcessor requestProcessor) { - this.requestProcessor = requestProcessor; - } - - @Override - public void get(VeniceClientRequest request, StreamObserver responseObserver) { - handleRequest(request, responseObserver); - } - - @Override - public void batchGet(VeniceClientRequest request, StreamObserver responseObserver) { - handleRequest(request, responseObserver); - } - - private void handleRequest(VeniceClientRequest request, StreamObserver responseObserver) { - VeniceServerResponse.Builder responseBuilder = - VeniceServerResponse.newBuilder().setErrorCode(VeniceReadResponseStatus.OK); - GrpcRequestContext ctx = new GrpcRequestContext(request, responseBuilder, responseObserver); - requestProcessor.process(ctx); - } - - @Override - public String toString() { - return this.getClass().getSimpleName(); - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcOutboundResponseHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcOutboundResponseHandler.java deleted file mode 100644 index 561575e1ef..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcOutboundResponseHandler.java +++ /dev/null @@ -1,60 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import static io.netty.handler.codec.http.HttpResponseStatus.NOT_FOUND; -import static io.netty.handler.codec.http.HttpResponseStatus.OK; - -import com.google.protobuf.ByteString; -import com.linkedin.davinci.listener.response.ReadResponse; -import com.linkedin.venice.compression.CompressionStrategy; -import com.linkedin.venice.listener.ServerStatsContext; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.protocols.VeniceServerResponse; -import com.linkedin.venice.response.VeniceReadResponseStatus; -import io.netty.buffer.ByteBuf; -import io.netty.handler.codec.http.HttpResponseStatus; - - -public class GrpcOutboundResponseHandler extends VeniceServerGrpcHandler { - @Override - public void processRequest(GrpcRequestContext ctx) { - ByteBuf body; - CompressionStrategy compressionStrategy = CompressionStrategy.NO_OP; - - ReadResponse obj = ctx.getReadResponse(); - ServerStatsContext statsContext = ctx.getGrpcStatsContext(); - VeniceServerResponse.Builder veniceServerResponseBuilder = ctx.getVeniceServerResponseBuilder(); - if (ctx.hasError()) { - statsContext.setResponseStatus(HttpResponseStatus.BAD_REQUEST); - veniceServerResponseBuilder.setData(ByteString.EMPTY).setCompressionStrategy(compressionStrategy.getValue()); - invokeNextHandler(ctx); - return; - } - - compressionStrategy = obj.getCompressionStrategy(); - - veniceServerResponseBuilder.setCompressionStrategy(compressionStrategy.getValue()); - veniceServerResponseBuilder.setResponseRCU(obj.getRCU()); - veniceServerResponseBuilder.setIsStreamingResponse(obj.isStreamingResponse()); - - if (obj.isFound()) { - body = obj.getResponseBody(); - - byte[] array = new byte[body.readableBytes()]; - body.getBytes(body.readerIndex(), array); - veniceServerResponseBuilder.setData(ByteString.copyFrom(array)) - .setCompressionStrategy(compressionStrategy.getValue()); - - veniceServerResponseBuilder.setSchemaId(obj.getResponseSchemaIdHeader()); - statsContext.setResponseStatus(OK); - invokeNextHandler(ctx); - return; - } - - ctx.setError(); - statsContext.setResponseStatus(NOT_FOUND); - veniceServerResponseBuilder.setData(ByteString.EMPTY); - veniceServerResponseBuilder.setErrorCode(VeniceReadResponseStatus.KEY_NOT_FOUND); - veniceServerResponseBuilder.setErrorMessage("Key not found"); - invokeNextHandler(ctx); - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcOutboundStatsHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcOutboundStatsHandler.java deleted file mode 100644 index 433e45b91e..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcOutboundStatsHandler.java +++ /dev/null @@ -1,48 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import com.linkedin.venice.exceptions.VeniceException; -import com.linkedin.venice.listener.ServerStatsContext; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.stats.ServerHttpRequestStats; -import com.linkedin.venice.utils.LatencyUtils; -import io.netty.handler.codec.http.HttpResponseStatus; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - - -public class GrpcOutboundStatsHandler extends VeniceServerGrpcHandler { - private static final Logger LOGGER = LogManager.getLogger(GrpcOutboundStatsHandler.class); - - @Override - public void processRequest(GrpcRequestContext ctx) { - writeResponse(ctx); - - // do not exit early if there is an error, we want to record stats for error cases as well - ServerStatsContext statsContext = ctx.getGrpcStatsContext(); - HttpResponseStatus responseStatus = statsContext.getResponseStatus(); - if (statsContext.getResponseStatus() == null) { - LOGGER.error("Received error in outbound gRPC Stats Handler: response status could not be null"); - throw new VeniceException("response status could not be null"); - } - - if (statsContext.getStoreName() == null) { - LOGGER.error("Received error in outbound gRPC Stats Handler: store name could not be null"); - throw new VeniceException("store name could not be null"); - } - - ServerHttpRequestStats serverHttpRequestStats = statsContext.getStoreName() == null - ? null - : statsContext.getCurrentStats().getStoreStats(statsContext.getStoreName()); - - statsContext.recordBasicMetrics(serverHttpRequestStats); - - double elapsedTime = LatencyUtils.getElapsedTimeFromNSToMS(statsContext.getRequestStartTimeInNS()); - - if (!ctx.hasError() && !responseStatus.equals(HttpResponseStatus.OK) - || responseStatus.equals(HttpResponseStatus.NOT_FOUND)) { - statsContext.successRequest(serverHttpRequestStats, elapsedTime); - } else { - statsContext.errorRequest(serverHttpRequestStats, elapsedTime); - } - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcReadQuotaEnforcementHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcReadQuotaEnforcementHandler.java deleted file mode 100644 index 236d85e934..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcReadQuotaEnforcementHandler.java +++ /dev/null @@ -1,48 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import static com.linkedin.venice.listener.ReadQuotaEnforcementHandler.INVALID_REQUEST_RESOURCE_MSG; -import static com.linkedin.venice.listener.ReadQuotaEnforcementHandler.SERVER_OVER_CAPACITY_MSG; - -import com.linkedin.venice.listener.ReadQuotaEnforcementHandler; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.request.RouterRequest; -import com.linkedin.venice.response.VeniceReadResponseStatus; -import java.util.Objects; - - -public class GrpcReadQuotaEnforcementHandler extends VeniceServerGrpcHandler { - private final ReadQuotaEnforcementHandler readQuotaHandler; - - public GrpcReadQuotaEnforcementHandler(ReadQuotaEnforcementHandler readQuotaEnforcementHandler) { - readQuotaHandler = - Objects.requireNonNull(readQuotaEnforcementHandler, "ReadQuotaEnforcementHandler cannot be null"); - } - - @Override - public void processRequest(GrpcRequestContext context) { - RouterRequest request = context.getRouterRequest(); - ReadQuotaEnforcementHandler.QuotaEnforcementResult result = readQuotaHandler.enforceQuota(request); - if (result == ReadQuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED) { - invokeNextHandler(context); - return; - } - - context.setError(); - if (result == ReadQuotaEnforcementHandler.QuotaEnforcementResult.BAD_REQUEST) { - context.getVeniceServerResponseBuilder() - .setErrorCode(VeniceReadResponseStatus.BAD_REQUEST) - .setErrorMessage(INVALID_REQUEST_RESOURCE_MSG + request.getResourceName()); - } else if (result == ReadQuotaEnforcementHandler.QuotaEnforcementResult.REJECTED) { - context.getVeniceServerResponseBuilder() - .setErrorCode(VeniceReadResponseStatus.TOO_MANY_REQUESTS) - .setErrorMessage(""); - } else if (result == ReadQuotaEnforcementHandler.QuotaEnforcementResult.OVER_CAPACITY) { - context.getVeniceServerResponseBuilder() - .setErrorCode(VeniceReadResponseStatus.SERVICE_UNAVAILABLE) - .setErrorMessage(SERVER_OVER_CAPACITY_MSG); - } - - // If we reach here, the request is allowed; retain the request and pass it to the next handler - invokeNextHandler(context); - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcRouterRequestHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcRouterRequestHandler.java deleted file mode 100644 index 816213028f..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcRouterRequestHandler.java +++ /dev/null @@ -1,34 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import com.linkedin.venice.listener.ServerStatsContext; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.request.GetRouterRequest; -import com.linkedin.venice.listener.request.MultiGetRouterRequestWrapper; -import com.linkedin.venice.listener.request.RouterRequest; -import com.linkedin.venice.protocols.VeniceClientRequest; - - -public class GrpcRouterRequestHandler extends VeniceServerGrpcHandler { - public GrpcRouterRequestHandler() { - } - - @Override - public void processRequest(GrpcRequestContext ctx) { - if (ctx.hasError()) { - invokeNextHandler(ctx); - return; - } - - VeniceClientRequest clientRequest = ctx.getVeniceClientRequest(); - ServerStatsContext statsContext = ctx.getGrpcStatsContext(); - - RouterRequest routerRequest = clientRequest.getIsBatchRequest() - ? MultiGetRouterRequestWrapper.parseMultiGetGrpcRequest(clientRequest) - : GetRouterRequest.grpcGetRouterRequest(clientRequest); - - statsContext.setRequestInfo(routerRequest); - - ctx.setRouterRequest(routerRequest); - invokeNextHandler(ctx); - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcStatsHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcStatsHandler.java deleted file mode 100644 index 1f9d729dfa..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcStatsHandler.java +++ /dev/null @@ -1,27 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import com.linkedin.venice.listener.ServerStatsContext; -import com.linkedin.venice.listener.StatsHandler; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; - - -public class GrpcStatsHandler extends VeniceServerGrpcHandler { - private final StatsHandler statsHandler; - - public GrpcStatsHandler(StatsHandler statsHandler) { - this.statsHandler = statsHandler; - } - - @Override - public void processRequest(GrpcRequestContext ctx) { - if (ctx.getGrpcStatsContext() == null) { - // new request - ServerStatsContext statsContext = statsHandler.getNewStatsContext(); - statsContext.resetContext(); - - ctx.setGrpcStatsContext(statsContext); - } - - invokeNextHandler(ctx); - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcStorageReadRequestHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcStorageReadRequestHandler.java deleted file mode 100644 index 833ecddcac..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/GrpcStorageReadRequestHandler.java +++ /dev/null @@ -1,81 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import com.linkedin.davinci.listener.response.ReadResponse; -import com.linkedin.venice.exceptions.VeniceNoStoreException; -import com.linkedin.venice.listener.ReadQuotaEnforcementHandler; -import com.linkedin.venice.listener.StorageReadRequestHandler; -import com.linkedin.venice.listener.VeniceRequestEarlyTerminationException; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.request.GetRouterRequest; -import com.linkedin.venice.listener.request.MultiGetRouterRequestWrapper; -import com.linkedin.venice.listener.request.RouterRequest; -import com.linkedin.venice.response.VeniceReadResponseStatus; -import io.netty.channel.ChannelHandlerContext; - - -/** - * This class is an incomplete copypasta of the logic in {@link StorageReadRequestHandler#channelRead(ChannelHandlerContext, Object)}. - * - * Besides the maintenance issue of the repeated code, and the incomplete functionality support, another potentially big - * issue is that the threading model seems to be significantly different. This class does all the work in-line, in a - * blocking fashion. All of these disparities are likely to cause significant issues in terms of trying to ramp the gRPC - * path. - * - * TODO: Refactor with better abstractions so that gRPC and legacy endpoints have better code reuse and behavior parity. - */ -public class GrpcStorageReadRequestHandler extends VeniceServerGrpcHandler { - private final StorageReadRequestHandler storage; - - public GrpcStorageReadRequestHandler(StorageReadRequestHandler storage) { - this.storage = storage; - } - - @Override - public void processRequest(GrpcRequestContext ctx) { - RouterRequest request = ctx.getRouterRequest(); - ReadResponse response = null; - - try { - if (request.shouldRequestBeTerminatedEarly()) { - throw new VeniceRequestEarlyTerminationException(request.getStoreName()); - } - - switch (request.getRequestType()) { - case SINGLE_GET: - // TODO: get rid of blocking here - response = storage.handleSingleGetRequest((GetRouterRequest) request).get(); - break; - case MULTI_GET: - // TODO: get rid of blocking here - response = storage.handleMultiGetRequest((MultiGetRouterRequestWrapper) request).get(); - break; - default: - ctx.setError(); - ctx.getVeniceServerResponseBuilder() - .setErrorCode(VeniceReadResponseStatus.BAD_REQUEST) - .setErrorMessage("Unknown request type: " + request.getRequestType()); - } - } catch (VeniceNoStoreException e) { - ctx.setError(); - ctx.getVeniceServerResponseBuilder() - .setErrorCode(VeniceReadResponseStatus.BAD_REQUEST) - .setErrorMessage("No storage exists for: " + e.getStoreName()); - } catch (Exception e) { - ctx.setError(); - ctx.getVeniceServerResponseBuilder() - .setErrorCode(VeniceReadResponseStatus.INTERNAL_ERROR) - .setErrorMessage(String.format("Internal Error: %s", e.getMessage())); - } - - if (!ctx.hasError() && response != null) { - response.setRCU(ReadQuotaEnforcementHandler.getRcu(request)); - if (request.isStreamingRequest()) { - response.setStreamingResponse(); - } - - ctx.setReadResponse(response); - } - - invokeNextHandler(ctx); - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/VeniceServerGrpcHandler.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/VeniceServerGrpcHandler.java deleted file mode 100644 index 9021c58616..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/VeniceServerGrpcHandler.java +++ /dev/null @@ -1,41 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.protocols.VeniceServerResponse; -import io.grpc.stub.StreamObserver; - - -public abstract class VeniceServerGrpcHandler { - protected VeniceServerGrpcHandler next; - - public VeniceServerGrpcHandler() { - next = null; - } - - public VeniceServerGrpcHandler getNext() { - return next; - } - - public void addNextHandler(VeniceServerGrpcHandler next) { - this.next = next; - } - - public abstract void processRequest(GrpcRequestContext ctx); - - public void invokeNextHandler(GrpcRequestContext ctx) { - if (next != null) { - next.processRequest(ctx); - } else { - writeResponse(ctx); - } - } - - public void writeResponse(GrpcRequestContext ctx) { - StreamObserver responseObserver = ctx.getResponseObserver(); - VeniceServerResponse response = ctx.getVeniceServerResponseBuilder().build(); - - responseObserver.onNext(response); - responseObserver.onCompleted(); - } - -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/VeniceServerGrpcRequestProcessor.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/VeniceServerGrpcRequestProcessor.java deleted file mode 100644 index 268136eb30..0000000000 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/grpc/handlers/VeniceServerGrpcRequestProcessor.java +++ /dev/null @@ -1,29 +0,0 @@ -package com.linkedin.venice.listener.grpc.handlers; - -import com.linkedin.venice.listener.grpc.GrpcRequestContext; - - -public class VeniceServerGrpcRequestProcessor { - private VeniceServerGrpcHandler head = null; - - public VeniceServerGrpcRequestProcessor() { - } - - public void addHandler(VeniceServerGrpcHandler handler) { - if (head == null) { - head = handler; - return; - } - - VeniceServerGrpcHandler current = head; - while (current.getNext() != null) { - current = current.getNext(); - } - - current.addNextHandler(handler); - } - - public void process(GrpcRequestContext ctx) { - head.processRequest(ctx); - } -} diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/AdminRequest.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/AdminRequest.java index 4627e15413..516e81efe7 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/AdminRequest.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/AdminRequest.java @@ -40,6 +40,20 @@ public static AdminRequest parseAdminHttpRequest(HttpRequest request, URI fullUr } } + public static AdminRequest parseAdminGrpcRequest(com.linkedin.venice.protocols.AdminRequest request) { + try { + String topicName = request.getResourceName(); + ServerAdminAction serverAdminAction = ServerAdminAction.valueOf(request.getServerAdminAction().toUpperCase()); + Integer partition = request.hasPartition() ? request.getPartition() : null; + if (!Version.isVersionTopic(topicName)) { + throw new VeniceException("Invalid store version for an ADMIN action: " + serverAdminAction); + } + return new AdminRequest(topicName, serverAdminAction, partition); + } catch (IllegalArgumentException e) { + throw new VeniceException("Invalid server admin action: " + request.getServerAdminAction()); + } + } + public String getStoreName() { return this.storeName; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/ComputeRouterRequestWrapper.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/ComputeRouterRequestWrapper.java index dfdb0b3279..f33ca4f32d 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/ComputeRouterRequestWrapper.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/ComputeRouterRequestWrapper.java @@ -10,8 +10,9 @@ import com.linkedin.venice.read.RequestType; import com.linkedin.venice.serializer.FastSerializerDeserializerFactory; import com.linkedin.venice.serializer.RecordDeserializer; +import com.linkedin.venice.streaming.StreamingUtils; +import com.linkedin.venice.utils.NettyUtils; import io.netty.handler.codec.http.FullHttpRequest; -import io.netty.handler.codec.http.HttpRequest; import java.util.List; import org.apache.avro.io.BinaryDecoder; import org.apache.avro.io.OptimizedBinaryDecoderFactory; @@ -31,45 +32,99 @@ private ComputeRouterRequestWrapper( String resourceName, ComputeRequest computeRequest, List keys, - HttpRequest request, - String schemaId) { - super(resourceName, keys, request); + String schemaId, + boolean isRetryRequest, + boolean isStreamingRequest) { + super(resourceName, keys, isRetryRequest, isStreamingRequest); this.computeRequest = computeRequest; if (schemaId != null) { this.valueSchemaId = Integer.parseInt(schemaId); } } - public static ComputeRouterRequestWrapper parseComputeRequest(FullHttpRequest httpRequest, String[] requestParts) { + public static ComputeRouterRequestWrapper parseComputeHttpRequest( + FullHttpRequest httpRequest, + String[] requestParts) { if (requestParts.length != 3) { // [0]""/[1]"compute"/[2]{$resourceName} throw new VeniceException("Invalid request: " + httpRequest.uri()); } String resourceName = requestParts[2]; + int apiVersion = parseApiVersion(httpRequest.headers().get(HttpConstants.VENICE_API_VERSION)); + byte[] requestContent = getRequestContent(httpRequest); - // Validate API version - String apiVersionStr = httpRequest.headers().get(HttpConstants.VENICE_API_VERSION); + return buildComputeRouterRequestWrapper( + resourceName, + apiVersion, + requestContent, + httpRequest.headers().get(HttpConstants.VENICE_COMPUTE_VALUE_SCHEMA_ID), + NettyUtils.containRetryHeader(httpRequest), + StreamingUtils.isStreamingEnabled(httpRequest)); + } + + // Helper method to handle common logic + private static ComputeRouterRequestWrapper buildComputeRouterRequestWrapper( + String resourceName, + int apiVersion, + byte[] requestContent, + String schemaId, + boolean isRetryRequest, + boolean isStreamingRequest) { + validateApiVersion(apiVersion); + + BinaryDecoder decoder = OptimizedBinaryDecoderFactory.defaultFactory() + .createOptimizedBinaryDecoder(requestContent, 0, requestContent.length); + + ComputeRequest computeRequest = ComputeUtils.deserializeComputeRequest(decoder, null); + List keys = DESERIALIZER.deserializeObjects(decoder); + + return new ComputeRouterRequestWrapper( + resourceName, + computeRequest, + keys, + schemaId, + isRetryRequest, + isStreamingRequest); + } + + // Extracted version validation logic + private static int parseApiVersion(String apiVersionStr) { if (apiVersionStr == null) { throw new VeniceException("Header: " + HttpConstants.VENICE_API_VERSION + " is missing"); } - int apiVersion = Integer.parseInt(apiVersionStr); + + return Integer.parseInt(apiVersionStr); + } + + // Extracted API version validation method + private static void validateApiVersion(int apiVersion) { if (apiVersion <= 0 || apiVersion > LATEST_SCHEMA_VERSION_FOR_COMPUTE_REQUEST) { throw new VeniceException( - "Compute API version " + apiVersion + " is invalid. " + "Latest version is " + "Compute API version " + apiVersion + " is invalid. Latest version is " + LATEST_SCHEMA_VERSION_FOR_COMPUTE_REQUEST); } + } - // TODO: xplore the possibility of streaming in the request bytes, and processing it in pipelined fashion + // Helper to extract content from HTTP request + private static byte[] getRequestContent(FullHttpRequest httpRequest) { byte[] requestContent = new byte[httpRequest.content().readableBytes()]; httpRequest.content().readBytes(requestContent); + return requestContent; + } - BinaryDecoder decoder = OptimizedBinaryDecoderFactory.defaultFactory() - .createOptimizedBinaryDecoder(requestContent, 0, requestContent.length); - ComputeRequest computeRequest = ComputeUtils.deserializeComputeRequest(decoder, null); + public static ComputeRouterRequestWrapper parseComputeGrpcRequest( + com.linkedin.venice.protocols.ComputeRequest grpcRequest) { + String resourceName = grpcRequest.getResourceName(); + int apiVersion = grpcRequest.getApiVersion(); + byte[] requestContent = grpcRequest.getPayload().toByteArray(); - List keys = DESERIALIZER.deserializeObjects(decoder); - String schemaId = httpRequest.headers().get(HttpConstants.VENICE_COMPUTE_VALUE_SCHEMA_ID); - return new ComputeRouterRequestWrapper(resourceName, computeRequest, keys, httpRequest, schemaId); + return buildComputeRouterRequestWrapper( + resourceName, + apiVersion, + requestContent, + Integer.toString(grpcRequest.getComputeValueSchemaId()), + grpcRequest.getIsRetryRequest(), + grpcRequest.getIsStreamingRequest()); } public ComputeRequest getComputeRequest() { diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/CurrentVersionRequest.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/CurrentVersionRequest.java index d3a4dfbb56..b2c6e39b44 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/CurrentVersionRequest.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/CurrentVersionRequest.java @@ -1,6 +1,7 @@ package com.linkedin.venice.listener.request; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.protocols.CurrentVersionInfoRequest; import java.net.URI; @@ -21,6 +22,10 @@ public static CurrentVersionRequest parseGetHttpRequest(URI uri, String[] reques } } + public static CurrentVersionRequest parseGetGrpcRequest(CurrentVersionInfoRequest request) { + return new CurrentVersionRequest(request.getStoreName()); + } + public String getStoreName() { return storeName; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/DictionaryFetchRequest.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/DictionaryFetchRequest.java index 824bba70d4..b3d3640a50 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/DictionaryFetchRequest.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/DictionaryFetchRequest.java @@ -30,6 +30,10 @@ public static DictionaryFetchRequest parseGetHttpRequest(URI uri, String[] reque } } + public static DictionaryFetchRequest create(String storeName, String resourceName) { + return new DictionaryFetchRequest(storeName, resourceName); + } + public String getResourceName() { return resourceName; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/GetRouterRequest.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/GetRouterRequest.java index 57188d710f..a34efb6bd4 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/GetRouterRequest.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/GetRouterRequest.java @@ -3,9 +3,12 @@ import com.linkedin.venice.HttpConstants; import com.linkedin.venice.RequestConstants; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.protocols.SingleGetRequest; import com.linkedin.venice.protocols.VeniceClientRequest; import com.linkedin.venice.read.RequestType; +import com.linkedin.venice.streaming.StreamingUtils; import com.linkedin.venice.utils.EncodingUtils; +import com.linkedin.venice.utils.NettyUtils; import io.netty.handler.codec.http.HttpHeaders; import io.netty.handler.codec.http.HttpRequest; import io.netty.handler.codec.http.QueryStringDecoder; @@ -19,16 +22,13 @@ public class GetRouterRequest extends RouterRequest { private final int partition; private final byte[] keyBytes; - private GetRouterRequest(String resourceName, int partition, byte[] keyBytes, HttpRequest request) { - super(resourceName, request); - - this.partition = partition; - this.keyBytes = keyBytes; - } - - private GetRouterRequest(String resourceName, int partition, byte[] keyBytes) { - super(resourceName, false, false); - + private GetRouterRequest( + String resourceName, + int partition, + byte[] keyBytes, + boolean isRetryRequest, + boolean isStreamingRequest) { + super(resourceName, isRetryRequest, isStreamingRequest); this.partition = partition; this.keyBytes = keyBytes; } @@ -57,18 +57,36 @@ public static GetRouterRequest parseGetHttpRequest(HttpRequest request, String[] String topicName = requestParts[2]; int partition = Integer.parseInt(requestParts[3]); byte[] keyBytes = getKeyBytesFromUrlKeyString(requestParts[4]); - return new GetRouterRequest(topicName, partition, keyBytes, request); + return new GetRouterRequest( + topicName, + partition, + keyBytes, + NettyUtils.containRetryHeader(request), + StreamingUtils.isStreamingEnabled(request)); } else { throw new VeniceException("Not a valid request for a STORAGE action: " + request.uri()); } } - public static GetRouterRequest grpcGetRouterRequest(VeniceClientRequest request) { + /** + * @deprecated This method has been deprecated and will be removed once the corresponding legacy gRPC code is removed. + */ + @Deprecated + public static GetRouterRequest parseSingleGetGrpcRequest(VeniceClientRequest request) { String resourceName = request.getResourceName(); int partition = request.getPartition(); byte[] keyBytes = getKeyBytesFromUrlKeyString(request.getKeyString()); + boolean isRetryRequest = request.getIsRetryRequest(); + boolean isStreamingRequest = request.getIsStreamingRequest(); + return new GetRouterRequest(resourceName, partition, keyBytes, isRetryRequest, isStreamingRequest); + } - return new GetRouterRequest(resourceName, partition, keyBytes); + public static GetRouterRequest parseSingleGetGrpcRequest(SingleGetRequest singleGetRequest) { + String resourceName = singleGetRequest.getResourceName(); + int partition = singleGetRequest.getPartition(); + byte[] keyBytes = getKeyBytesFromUrlKeyString(singleGetRequest.getKey()); + boolean isRetryRequest = singleGetRequest.getIsRetryRequest(); + return new GetRouterRequest(resourceName, partition, keyBytes, isRetryRequest, false); } public static byte[] getKeyBytesFromUrlKeyString(String keyString) { diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MetadataFetchRequest.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MetadataFetchRequest.java index fb744e9bd0..a3d0dd8f1e 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MetadataFetchRequest.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MetadataFetchRequest.java @@ -1,6 +1,7 @@ package com.linkedin.venice.listener.request; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.protocols.MetadataRequest; /** @@ -24,6 +25,10 @@ public static MetadataFetchRequest parseGetHttpRequest(String uri, String[] requ } } + public static MetadataFetchRequest parseGetGrpcRequest(MetadataRequest request) { + return new MetadataFetchRequest(request.getStoreName()); + } + public String getStoreName() { return storeName; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiGetRouterRequestWrapper.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiGetRouterRequestWrapper.java index 2943865356..deb876fe7e 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiGetRouterRequestWrapper.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiGetRouterRequestWrapper.java @@ -2,14 +2,16 @@ import com.linkedin.venice.HttpConstants; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.protocols.MultiGetRequest; import com.linkedin.venice.protocols.VeniceClientRequest; import com.linkedin.venice.read.RequestType; import com.linkedin.venice.read.protocol.request.router.MultiGetRouterRequestKeyV1; import com.linkedin.venice.schema.avro.ReadAvroProtocolDefinition; import com.linkedin.venice.serializer.FastSerializerDeserializerFactory; import com.linkedin.venice.serializer.RecordDeserializer; +import com.linkedin.venice.streaming.StreamingUtils; +import com.linkedin.venice.utils.NettyUtils; import io.netty.handler.codec.http.FullHttpRequest; -import io.netty.handler.codec.http.HttpRequest; import java.util.List; import org.apache.avro.io.OptimizedBinaryDecoderFactory; @@ -21,13 +23,6 @@ public class MultiGetRouterRequestWrapper extends MultiKeyRouterRequestWrapper DESERIALIZER = FastSerializerDeserializerFactory .getFastAvroSpecificDeserializer(MultiGetRouterRequestKeyV1.SCHEMA$, MultiGetRouterRequestKeyV1.class); - private MultiGetRouterRequestWrapper( - String resourceName, - List keys, - HttpRequest request) { - super(resourceName, keys, request); - } - private MultiGetRouterRequestWrapper( String resourceName, List keys, @@ -57,16 +52,29 @@ public static MultiGetRouterRequestWrapper parseMultiGetHttpRequest( byte[] content = new byte[httpRequest.content().readableBytes()]; httpRequest.content().readBytes(content); keys = parseKeys(content); + boolean isRetryRequest = NettyUtils.containRetryHeader(httpRequest); + boolean isStreamingRequest = StreamingUtils.isStreamingEnabled(httpRequest); - return new MultiGetRouterRequestWrapper(requestParts[2], keys, httpRequest); + return new MultiGetRouterRequestWrapper(requestParts[2], keys, isRetryRequest, isStreamingRequest); } + /** + * @deprecated This method has been deprecated and will be removed once the corresponding legacy gRPC code is removed. + */ + @Deprecated public static MultiGetRouterRequestWrapper parseMultiGetGrpcRequest(VeniceClientRequest grpcRequest) { String resourceName = grpcRequest.getResourceName(); List keys = parseKeys(grpcRequest.getKeyBytes().toByteArray()); + boolean isRetryRequest = grpcRequest.getIsRetryRequest(); + boolean isStreamingRequest = grpcRequest.getIsStreamingRequest(); + return new MultiGetRouterRequestWrapper(resourceName, keys, isRetryRequest, isStreamingRequest); + } - // isRetryRequest set to false for now, retry functionality is a later milestone - return new MultiGetRouterRequestWrapper(resourceName, keys, false, grpcRequest.getIsStreamingRequest()); + public static MultiGetRouterRequestWrapper parseMultiGetGrpcRequest(MultiGetRequest grpcRequest) { + String resourceName = grpcRequest.getResourceName(); + List keys = parseKeys(grpcRequest.getKeyBytes().toByteArray()); + boolean isRetryRequest = grpcRequest.getIsRetryRequest(); + return new MultiGetRouterRequestWrapper(resourceName, keys, isRetryRequest, false); } private static List parseKeys(byte[] content) { diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiKeyRouterRequestWrapper.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiKeyRouterRequestWrapper.java index c4ee1b1a85..303aa0ceab 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiKeyRouterRequestWrapper.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/MultiKeyRouterRequestWrapper.java @@ -1,6 +1,5 @@ package com.linkedin.venice.listener.request; -import io.netty.handler.codec.http.HttpRequest; import java.util.List; @@ -10,11 +9,6 @@ public abstract class MultiKeyRouterRequestWrapper extends RouterRequest { private final List keys; - protected MultiKeyRouterRequestWrapper(String resourceName, List keys, HttpRequest request) { - super(resourceName, request); - this.keys = keys; - } - protected MultiKeyRouterRequestWrapper( String resourceName, List keys, diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/RouterRequest.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/RouterRequest.java index c53c30f08c..7dd4f5c192 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/RouterRequest.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/RouterRequest.java @@ -1,10 +1,7 @@ package com.linkedin.venice.listener.request; -import com.linkedin.venice.HttpConstants; import com.linkedin.venice.meta.Version; import com.linkedin.venice.read.RequestType; -import com.linkedin.venice.streaming.StreamingUtils; -import io.netty.handler.codec.http.HttpRequest; /** @@ -22,13 +19,6 @@ public abstract class RouterRequest { private final String storeName; private final boolean isStreamingRequest; - public RouterRequest(String resourceName, HttpRequest request) { - this.isRetryRequest = containRetryHeader(request); - this.isStreamingRequest = StreamingUtils.isStreamingEnabled(request); - this.resourceName = resourceName; - this.storeName = Version.parseStoreFromKafkaTopicName(resourceName); - } - public RouterRequest(String resourceName, boolean isRetryRequest, boolean isStreamingRequest) { this.resourceName = resourceName; this.storeName = Version.parseStoreFromKafkaTopicName(resourceName); @@ -60,10 +50,6 @@ public boolean isStreamingRequest() { return isStreamingRequest; } - private static boolean containRetryHeader(HttpRequest request) { - return request.headers().contains(HttpConstants.VENICE_RETRY); - } - public boolean shouldRequestBeTerminatedEarly() { return requestTimeoutInNS != NO_REQUEST_TIMEOUT && System.nanoTime() > requestTimeoutInNS; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/TopicPartitionIngestionContextRequest.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/TopicPartitionIngestionContextRequest.java index ab58e1a7a4..64697aeda0 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/request/TopicPartitionIngestionContextRequest.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/request/TopicPartitionIngestionContextRequest.java @@ -1,6 +1,7 @@ package com.linkedin.venice.listener.request; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.protocols.IngestionContextRequest; public class TopicPartitionIngestionContextRequest { @@ -26,6 +27,13 @@ public static TopicPartitionIngestionContextRequest parseGetHttpRequest(String u } } + public static TopicPartitionIngestionContextRequest parseGetGrpcRequest(IngestionContextRequest request) { + return new TopicPartitionIngestionContextRequest( + request.getVersionTopicName(), + request.getTopicName(), + request.getPartition()); + } + public String getVersionTopic() { return versionTopic; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/AbstractReadResponse.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/AbstractReadResponse.java index 5088114248..115029f96d 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/AbstractReadResponse.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/AbstractReadResponse.java @@ -54,5 +54,5 @@ public boolean isFound() { public abstract int getResponseSchemaIdHeader(); - public abstract ReadResponseStatsRecorder getStatsRecorder(); + public abstract ReadResponseStatsRecorder getReadResponseStatsRecorder(); } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/BinaryResponse.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/BinaryResponse.java index 066b0c42bc..7020303207 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/BinaryResponse.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/BinaryResponse.java @@ -1,22 +1,22 @@ package com.linkedin.venice.listener.response; +import com.linkedin.venice.response.VeniceReadResponseStatus; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; -import io.netty.handler.codec.http.HttpResponseStatus; import java.nio.ByteBuffer; public class BinaryResponse { private final ByteBuf body; - private final HttpResponseStatus status; + private final VeniceReadResponseStatus status; public BinaryResponse(ByteBuf body) { if (body == null) { this.body = Unpooled.EMPTY_BUFFER; - this.status = HttpResponseStatus.NOT_FOUND; + this.status = VeniceReadResponseStatus.KEY_NOT_FOUND; } else { this.body = body; - this.status = HttpResponseStatus.OK; + this.status = VeniceReadResponseStatus.OK; } } @@ -28,7 +28,7 @@ public ByteBuf getBody() { return body; } - public HttpResponseStatus getStatus() { + public VeniceReadResponseStatus getStatus() { return status; } } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/MultiKeyResponseWrapper.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/MultiKeyResponseWrapper.java index 50bde34bb7..f24bd93667 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/MultiKeyResponseWrapper.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/MultiKeyResponseWrapper.java @@ -59,7 +59,7 @@ public ReadResponseStats getStats() { } @Override - public ReadResponseStatsRecorder getStatsRecorder() { + public ReadResponseStatsRecorder getReadResponseStatsRecorder() { return this.responseStats; } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/ParallelMultiKeyResponseWrapper.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/ParallelMultiKeyResponseWrapper.java index 96a3e283e2..0a0388c2ab 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/ParallelMultiKeyResponseWrapper.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/ParallelMultiKeyResponseWrapper.java @@ -58,7 +58,7 @@ public ReadResponseStats getStats() { } @Override - public ReadResponseStatsRecorder getStatsRecorder() { + public ReadResponseStatsRecorder getReadResponseStatsRecorder() { return new CompositeReadResponseStatsRecorder(this.chunks); } @@ -93,7 +93,7 @@ private static final class CompositeReadResponseStatsRecorder implements ReadRes private final ReadResponseStatsRecorder[] statsRecorders; CompositeReadResponseStatsRecorder(MultiKeyResponseWrapper[] responseChunks) { - this.mergedStats = responseChunks[0].getStatsRecorder(); + this.mergedStats = responseChunks[0].getReadResponseStatsRecorder(); /** * This array can be one element shorter than {@param responseChunks} because the first chunk's storage exec sub @@ -102,7 +102,7 @@ private static final class CompositeReadResponseStatsRecorder implements ReadRes this.statsRecorders = new ReadResponseStatsRecorder[responseChunks.length - 1]; ReadResponseStatsRecorder statsRecorder; for (int i = 1; i < responseChunks.length; i++) { - statsRecorder = responseChunks[i].getStatsRecorder(); + statsRecorder = responseChunks[i].getReadResponseStatsRecorder(); // We merge the stats of all chunks from the 2nd one to the last one into the stats of the 1st chunk. this.mergedStats.merge(statsRecorder); diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/SingleGetResponseWrapper.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/SingleGetResponseWrapper.java index 14e3b65e55..bc72391572 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/response/SingleGetResponseWrapper.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/response/SingleGetResponseWrapper.java @@ -29,7 +29,7 @@ public ReadResponseStats getStats() { } @Override - public ReadResponseStatsRecorder getStatsRecorder() { + public ReadResponseStatsRecorder getReadResponseStatsRecorder() { return this.responseStats; } diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcIoRequestProcessorTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcIoRequestProcessorTest.java new file mode 100644 index 0000000000..8763432a76 --- /dev/null +++ b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcIoRequestProcessorTest.java @@ -0,0 +1,94 @@ +package com.linkedin.venice.grpc; + +import static com.linkedin.venice.listener.QuotaEnforcementHandler.QuotaEnforcementResult.REJECTED; +import static com.linkedin.venice.listener.ReadQuotaEnforcementHandler.INVALID_REQUEST_RESOURCE_MSG; +import static com.linkedin.venice.listener.ReadQuotaEnforcementHandler.SERVER_OVER_CAPACITY_MSG; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.linkedin.venice.listener.QuotaEnforcementHandler; +import com.linkedin.venice.listener.QuotaEnforcementHandler.QuotaEnforcementResult; +import com.linkedin.venice.listener.StorageReadRequestHandler; +import com.linkedin.venice.listener.request.RouterRequest; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class GrpcIoRequestProcessorTest { + private QuotaEnforcementHandler quotaEnforcementHandler; + private StorageReadRequestHandler storageReadRequestHandler; + private GrpcReplyProcessor grpcReplyProcessor; + private GrpcServiceDependencies grpcServiceDependencies; + private GrpcIoRequestProcessor processor; + private GrpcRequestContext requestContext; + private RouterRequest request; + + @BeforeMethod + public void setUp() { + quotaEnforcementHandler = mock(QuotaEnforcementHandler.class); + storageReadRequestHandler = mock(StorageReadRequestHandler.class); + grpcReplyProcessor = mock(GrpcReplyProcessor.class); + grpcServiceDependencies = mock(GrpcServiceDependencies.class); + when(grpcServiceDependencies.getQuotaEnforcementHandler()).thenReturn(quotaEnforcementHandler); + when(grpcServiceDependencies.getStorageReadRequestHandler()).thenReturn(storageReadRequestHandler); + when(grpcServiceDependencies.getGrpcReplyProcessor()).thenReturn(grpcReplyProcessor); + + request = mock(RouterRequest.class); + when(request.getResourceName()).thenReturn("testResource_v1"); + requestContext = mock(GrpcRequestContext.class); + when(requestContext.getRouterRequest()).thenReturn(request); + when(requestContext.getGrpcRequestType()).thenReturn(GrpcRequestContext.GrpcRequestType.SINGLE_GET); + + processor = new GrpcIoRequestProcessor(grpcServiceDependencies); + } + + @Test + public void testProcessRequestAllowed() { + // Case when quota enforcement result is ALLOWED + when(quotaEnforcementHandler.enforceQuota(request)).thenReturn(QuotaEnforcementResult.ALLOWED); + + processor.processRequest(requestContext); + + // Verify that the request is handed off to the storage read request handler + verify(storageReadRequestHandler) + .queueIoRequestForAsyncProcessing(eq(request), any(GrpcStorageResponseHandlerCallback.class)); + verify(requestContext, never()).setErrorMessage(anyString()); + verify(requestContext, never()).setReadResponseStatus(any()); + + // should not call grpcReplyProcessor.sendResponse() because the request is handed off to the storage read request + // handler which is mocked + verify(grpcReplyProcessor, never()).sendResponse(requestContext); + } + + @Test + public void testProcessRequestQuotaEnforcementErrors() { + // BAD_REQUEST case + when(quotaEnforcementHandler.enforceQuota(request)).thenReturn(QuotaEnforcementResult.BAD_REQUEST); + when(request.getResourceName()).thenReturn("testResource"); + processor.processRequest(requestContext); + verify(requestContext).setErrorMessage(INVALID_REQUEST_RESOURCE_MSG + "testResource"); + verify(requestContext).setReadResponseStatus(VeniceReadResponseStatus.BAD_REQUEST); + verify(grpcReplyProcessor).sendResponse(requestContext); + + // REJECTED case + when(quotaEnforcementHandler.enforceQuota(request)).thenReturn(REJECTED); + processor.processRequest(requestContext); + verify(requestContext).setErrorMessage("Quota exceeded for resource: testResource"); + verify(requestContext).setReadResponseStatus(VeniceReadResponseStatus.TOO_MANY_REQUESTS); + verify(grpcReplyProcessor, times(2)).sendResponse(requestContext); + + // OVER_CAPACITY case + when(quotaEnforcementHandler.enforceQuota(request)).thenReturn(QuotaEnforcementResult.OVER_CAPACITY); + processor.processRequest(requestContext); + verify(requestContext).setErrorMessage(SERVER_OVER_CAPACITY_MSG); + verify(requestContext).setReadResponseStatus(VeniceReadResponseStatus.SERVICE_UNAVAILABLE); + verify(grpcReplyProcessor, times(3)).sendResponse(requestContext); + } +} diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcReplyProcessorTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcReplyProcessorTest.java new file mode 100644 index 0000000000..53de0ef56e --- /dev/null +++ b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcReplyProcessorTest.java @@ -0,0 +1,271 @@ +package com.linkedin.venice.grpc; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +import com.google.protobuf.ByteString; +import com.linkedin.davinci.listener.response.ReadResponse; +import com.linkedin.venice.compression.CompressionStrategy; +import com.linkedin.venice.protocols.MultiKeyResponse; +import com.linkedin.venice.protocols.SingleGetResponse; +import com.linkedin.venice.protocols.VeniceServerResponse; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import io.grpc.stub.StreamObserver; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import org.mockito.ArgumentCaptor; +import org.mockito.InOrder; +import org.testng.annotations.Test; + + +public class GrpcReplyProcessorTest { + @Test + public void testSendResponse() { + GrpcReplyProcessor replyProcessor = new GrpcReplyProcessor(); + + // Spy on the GrpcReplyProcessor to verify the method calls + GrpcReplyProcessor spyReplyProcessor = spy(replyProcessor); + + // Case 1: Test for SINGLE_GET case + GrpcRequestContext singleGetRequestContext = mock(GrpcRequestContext.class); + when(singleGetRequestContext.getGrpcRequestType()).thenReturn(GrpcRequestContext.GrpcRequestType.SINGLE_GET); + doNothing().when(spyReplyProcessor).sendSingleGetResponse(singleGetRequestContext); + spyReplyProcessor.sendResponse(singleGetRequestContext); + verify(spyReplyProcessor).sendSingleGetResponse(singleGetRequestContext); + verify(spyReplyProcessor, never()).sendMultiKeyResponse(any()); + verify(spyReplyProcessor, never()).sendVeniceServerResponse(any()); + + // Case 2: Test for MULTI_GET case + spyReplyProcessor = spy(replyProcessor); + GrpcRequestContext multiGetRequestContext = mock(GrpcRequestContext.class); + when(multiGetRequestContext.getGrpcRequestType()).thenReturn(GrpcRequestContext.GrpcRequestType.MULTI_GET); + doNothing().when(spyReplyProcessor).sendMultiKeyResponse(multiGetRequestContext); + spyReplyProcessor.sendResponse(multiGetRequestContext); + verify(spyReplyProcessor).sendMultiKeyResponse(multiGetRequestContext); + verify(spyReplyProcessor, never()).sendSingleGetResponse(any()); + verify(spyReplyProcessor, never()).sendVeniceServerResponse(any()); + + when(multiGetRequestContext.getGrpcRequestType()).thenReturn(GrpcRequestContext.GrpcRequestType.COMPUTE); + doNothing().when(spyReplyProcessor).sendMultiKeyResponse(multiGetRequestContext); + spyReplyProcessor.sendResponse(multiGetRequestContext); + verify(spyReplyProcessor, times(2)).sendMultiKeyResponse(multiGetRequestContext); + verify(spyReplyProcessor, never()).sendSingleGetResponse(any()); + verify(spyReplyProcessor, never()).sendVeniceServerResponse(any()); + + // Case 3: Test for LEGACY case + spyReplyProcessor = spy(replyProcessor); + GrpcRequestContext legacyRequestContext = mock(GrpcRequestContext.class); + when(legacyRequestContext.getGrpcRequestType()).thenReturn(GrpcRequestContext.GrpcRequestType.LEGACY); + doNothing().when(spyReplyProcessor).sendVeniceServerResponse(legacyRequestContext); + spyReplyProcessor.sendResponse(legacyRequestContext); + verify(spyReplyProcessor).sendVeniceServerResponse(legacyRequestContext); + verify(spyReplyProcessor, never()).sendSingleGetResponse(any()); + verify(spyReplyProcessor, never()).sendMultiKeyResponse(any()); + } + + @Test + public void testSendSingleGetResponse() { + GrpcReplyProcessor replyProcessor = new GrpcReplyProcessor(); + GrpcReplyProcessor spyReplyProcessor = spy(replyProcessor); + doNothing().when(spyReplyProcessor).reportRequestStats(any()); + + // 1. Test Case: readResponse is null + GrpcRequestContext requestContext = mock(GrpcRequestContext.class); + StreamObserver responseObserver = mock(StreamObserver.class); + when(requestContext.getResponseObserver()).thenReturn(responseObserver); + VeniceReadResponseStatus status = VeniceReadResponseStatus.TOO_MANY_REQUESTS; + when(requestContext.getReadResponseStatus()).thenReturn(status); + when(requestContext.getReadResponse()).thenReturn(null); + when(requestContext.getErrorMessage()).thenReturn("Some error"); + + spyReplyProcessor.sendSingleGetResponse(requestContext); + + ArgumentCaptor captor = ArgumentCaptor.forClass(SingleGetResponse.class); + verify(responseObserver).onNext(captor.capture()); + SingleGetResponse capturedResponse = captor.getValue(); + assertEquals(capturedResponse.getStatusCode(), status.getCode()); + assertEquals(capturedResponse.getErrorMessage(), "Some error"); + verify(responseObserver).onCompleted(); + + InOrder inOrder = inOrder(responseObserver); + inOrder.verify(responseObserver).onNext(any()); + inOrder.verify(responseObserver).onCompleted(); + + // 2. Test Case: readResponse is found + status = VeniceReadResponseStatus.OK; + when(requestContext.getReadResponseStatus()).thenReturn(status); + ReadResponse readResponse = mock(ReadResponse.class); + when(readResponse.isFound()).thenReturn(true); + when(requestContext.getReadResponse()).thenReturn(readResponse); + when(readResponse.getRCU()).thenReturn(1); + when(readResponse.getResponseSchemaIdHeader()).thenReturn(1); + when(readResponse.getCompressionStrategy()).thenReturn(CompressionStrategy.GZIP); + + ByteBuf responseBody = Unpooled.EMPTY_BUFFER; + when(readResponse.getResponseBody()).thenReturn(responseBody); + + spyReplyProcessor.sendSingleGetResponse(requestContext); + + verify(responseObserver, times(2)).onNext(captor.capture()); // Capturing the second call + capturedResponse = captor.getValue(); + assertEquals(capturedResponse.getStatusCode(), status.getCode()); + assertEquals(capturedResponse.getRcu(), 1); + assertEquals(capturedResponse.getSchemaId(), 1); + assertEquals(capturedResponse.getCompressionStrategy(), CompressionStrategy.GZIP.getValue()); + assertEquals(capturedResponse.getContentLength(), 0); + assertEquals(capturedResponse.getValue(), ByteString.EMPTY); + verify(responseObserver, times(2)).onCompleted(); + + inOrder.verify(responseObserver).onNext(any()); + inOrder.verify(responseObserver).onCompleted(); + + // 3. Test Case: readResponse is not found + status = VeniceReadResponseStatus.KEY_NOT_FOUND; + when(requestContext.getReadResponseStatus()).thenReturn(status); + when(readResponse.isFound()).thenReturn(false); + when(readResponse.getRCU()).thenReturn(5); + when(readResponse.getResponseSchemaIdHeader()).thenReturn(-1); + spyReplyProcessor.sendSingleGetResponse(requestContext); + + verify(responseObserver, times(3)).onNext(captor.capture()); // Capturing the third call + capturedResponse = captor.getValue(); + assertEquals(capturedResponse.getStatusCode(), status.getCode()); + assertEquals(capturedResponse.getRcu(), 5); + assertEquals(capturedResponse.getErrorMessage(), "Key not found"); + assertEquals(capturedResponse.getContentLength(), 0); + verify(responseObserver, times(3)).onCompleted(); + inOrder.verify(responseObserver).onNext(any()); + inOrder.verify(responseObserver).onCompleted(); + + verify(spyReplyProcessor, times(3)).reportRequestStats(requestContext); + } + + @Test + public void testSendVeniceServerResponse() { + GrpcReplyProcessor replyProcessor = new GrpcReplyProcessor(); + GrpcReplyProcessor spyReplyProcessor = spy(replyProcessor); + doNothing().when(spyReplyProcessor).reportRequestStats(any()); + + GrpcRequestContext requestContext = mock(GrpcRequestContext.class); + StreamObserver responseObserver = mock(StreamObserver.class); + + // 1. Test Case: readResponse is null + when(requestContext.getResponseObserver()).thenReturn(responseObserver); + VeniceReadResponseStatus status = VeniceReadResponseStatus.BAD_REQUEST; + when(requestContext.getReadResponseStatus()).thenReturn(status); + when(requestContext.getReadResponse()).thenReturn(null); + when(requestContext.getErrorMessage()).thenReturn("Null read response"); + + spyReplyProcessor.sendVeniceServerResponse(requestContext); + + ArgumentCaptor captor = ArgumentCaptor.forClass(VeniceServerResponse.class); + verify(responseObserver).onNext(captor.capture()); + VeniceServerResponse capturedResponse = captor.getValue(); + assertEquals(capturedResponse.getErrorCode(), status.getCode()); + assertEquals(capturedResponse.getErrorMessage(), "Null read response"); + verify(responseObserver).onCompleted(); + + // 2. Test Case: readResponse is found + ReadResponse readResponse = mock(ReadResponse.class); + when(readResponse.getResponseBody()).thenReturn(Unpooled.EMPTY_BUFFER); + when(requestContext.getReadResponse()).thenReturn(readResponse); + status = VeniceReadResponseStatus.OK; + when(requestContext.getReadResponseStatus()).thenReturn(status); + when(requestContext.getErrorMessage()).thenReturn(null); + when(readResponse.isFound()).thenReturn(true); + when(requestContext.getReadResponse()).thenReturn(readResponse); + when(readResponse.getRCU()).thenReturn(120); + when(readResponse.getCompressionStrategy()).thenReturn(CompressionStrategy.GZIP); + when(readResponse.isStreamingResponse()).thenReturn(true); + when(readResponse.getResponseSchemaIdHeader()).thenReturn(2); + + spyReplyProcessor.sendVeniceServerResponse(requestContext); + + verify(responseObserver, times(2)).onNext(captor.capture()); // Capturing the second call + capturedResponse = captor.getValue(); + assertEquals(capturedResponse.getErrorCode(), status.getCode()); + assertEquals(capturedResponse.getResponseRCU(), 120); + assertEquals(capturedResponse.getCompressionStrategy(), CompressionStrategy.GZIP.getValue()); + assertEquals(capturedResponse.getSchemaId(), 2); + assertTrue(capturedResponse.getIsStreamingResponse()); + assertEquals(capturedResponse.getData(), ByteString.EMPTY); + verify(responseObserver, times(2)).onCompleted(); + + // 3. Test Case: readResponse is not found + when(readResponse.isFound()).thenReturn(false); + status = VeniceReadResponseStatus.KEY_NOT_FOUND; + when(requestContext.getReadResponseStatus()).thenReturn(status); + + spyReplyProcessor.sendVeniceServerResponse(requestContext); + + verify(responseObserver, times(3)).onNext(captor.capture()); // Capturing the third call + capturedResponse = captor.getValue(); + assertEquals(capturedResponse.getErrorCode(), status.getCode()); + assertEquals(capturedResponse.getErrorMessage(), "Key not found"); + assertEquals(capturedResponse.getData(), ByteString.EMPTY); + verify(responseObserver, times(3)).onCompleted(); + + // 4. Verify synchronization on responseObserver + InOrder inOrder = inOrder(responseObserver); + inOrder.verify(responseObserver).onNext(any()); + inOrder.verify(responseObserver).onCompleted(); + + // 5. Verify reportRequestStats is called at the end + verify(spyReplyProcessor, times(3)).reportRequestStats(requestContext); + } + + // @Test + // public void testReportRequestStats() { + // GrpcReplyProcessor processor = new GrpcReplyProcessor(); + // + // // Mock the GrpcRequestContext and its dependencies + // GrpcRequestContext requestContext = mock(GrpcRequestContext.class); + // ReadResponse readResponse = mock(ReadResponse.class); + // RequestStatsRecorder requestStatsRecorder = mock(RequestStatsRecorder.class); + // AbstractReadResponse abstractReadResponse = mock(AbstractReadResponse.class); + // ByteBuf responseBody = mock(ByteBuf.class); + // + // // Mock behavior for requestContext + // when(requestContext.getReadResponse()).thenReturn(readResponse); + // when(requestContext.getRequestStatsRecorder()).thenReturn(requestStatsRecorder); + // + // // 1. Test Case: readResponse is null + // when(requestContext.getReadResponse()).thenReturn(null); + // + // processor.reportRequestStats(requestContext); + // + // // Verify that recorder sets read response stats to null and response size to 0 + // verify(requestStatsRecorder).setReadResponseStats(null); + // verify(requestStatsRecorder).setResponseSize(0); + // + // // 2. Test Case: readResponse.isFound() is true + // when(readResponse.isFound()).thenReturn(true); + // when(requestContext.getReadResponse()).thenReturn(abstractReadResponse); + // when(abstractReadResponse.getResponseBody()).thenReturn(responseBody); + // when(responseBody.readableBytes()).thenReturn(512); + // when(abstractReadResponse.getReadResponseStatsRecorder()).thenReturn(mock(ReadResponseStatsRecorder.class)); + // + // processor.reportRequestStats(requestContext); + // + // // Verify that recorder sets the stats and the correct response size + // verify(requestStatsRecorder, times(2)).setReadResponseStats(abstractReadResponse.getReadResponseStatsRecorder()); + // verify(requestStatsRecorder, times(2)).setResponseSize(512); + // + // // 3. Test Case: readResponse.isFound() is false + // when(readResponse.isFound()).thenReturn(false); + // processor.reportRequestStats(requestContext); + // + // // Verify that recorder sets the stats and the response size to 0 + // verify(requestStatsRecorder, times(3)).setReadResponseStats(abstractReadResponse.getReadResponseStatsRecorder()); + // verify(requestStatsRecorder).setResponseSize(0); + // } +} diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcRequestContextTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcRequestContextTest.java deleted file mode 100644 index 5f21307859..0000000000 --- a/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcRequestContextTest.java +++ /dev/null @@ -1,107 +0,0 @@ -package com.linkedin.venice.grpc; - -import static org.mockito.Mockito.mock; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertTrue; - -import com.linkedin.davinci.listener.response.ReadResponse; -import com.linkedin.venice.listener.ServerStatsContext; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.request.RouterRequest; -import com.linkedin.venice.protocols.VeniceClientRequest; -import com.linkedin.venice.protocols.VeniceServerResponse; -import io.grpc.stub.StreamObserver; -import org.testng.annotations.Test; - - -public class GrpcRequestContextTest { - @Test - public void testSetAndGetGrpcStatsContext() { - ServerStatsContext context = new ServerStatsContext(null, null, null); - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - handlerContext.setGrpcStatsContext(context); - - assertEquals(context, handlerContext.getGrpcStatsContext()); - } - - @Test - public void testSetAndGetVeniceClientRequest() { - VeniceClientRequest request = VeniceClientRequest.newBuilder().build(); - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - handlerContext.setVeniceClientRequest(request); - - assertEquals(request, handlerContext.getVeniceClientRequest()); - } - - @Test - public void testSetAndGetVeniceServerResponseBuilder() { - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - handlerContext.setVeniceServerResponseBuilder(builder); - - assertEquals(builder, handlerContext.getVeniceServerResponseBuilder()); - } - - @Test - public void testSetAndGetResponseObserver() { - StreamObserver observer = new StreamObserver() { - @Override - public void onNext(VeniceServerResponse value) { - } - - @Override - public void onError(Throwable t) { - } - - @Override - public void onCompleted() { - } - }; - - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - handlerContext.setResponseObserver(observer); - - assertEquals(observer, handlerContext.getResponseObserver()); - } - - @Test - public void testSetAndGetRouterRequest() { - RouterRequest request = mock(RouterRequest.class); - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - handlerContext.setRouterRequest(request); - - assertEquals(request, handlerContext.getRouterRequest()); - } - - @Test - public void testSetAndGetReadResponse() { - ReadResponse response = mock(ReadResponse.class); - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - handlerContext.setReadResponse(response); - - assertEquals(response, handlerContext.getReadResponse()); - } - - @Test - public void testIsCompleted() { - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - - assertFalse(handlerContext.isCompleted()); - - handlerContext.setCompleted(); - - assertTrue(handlerContext.isCompleted()); - } - - @Test - public void testHasError() { - GrpcRequestContext handlerContext = new GrpcRequestContext(null, null, null); - - assertFalse(handlerContext.hasError()); - - handlerContext.setError(); - - assertTrue(handlerContext.hasError()); - } -} diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcServiceDependenciesTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcServiceDependenciesTest.java new file mode 100644 index 0000000000..f96c5421e7 --- /dev/null +++ b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcServiceDependenciesTest.java @@ -0,0 +1,80 @@ +package com.linkedin.venice.grpc; + +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + +import com.linkedin.davinci.storage.DiskHealthCheckService; +import com.linkedin.venice.listener.NoOpReadQuotaEnforcementHandler; +import com.linkedin.venice.listener.StorageReadRequestHandler; +import com.linkedin.venice.stats.AggServerHttpRequestStats; +import org.mockito.Mockito; +import org.testng.annotations.Test; + + +public class GrpcServiceDependenciesTest { + @Test + public void testBuilderValidation() { + DiskHealthCheckService diskHealthCheckService = Mockito.mock(DiskHealthCheckService.class); + StorageReadRequestHandler storageReadRequestHandler = Mockito.mock(StorageReadRequestHandler.class); + NoOpReadQuotaEnforcementHandler quotaEnforcementHandler = Mockito.mock(NoOpReadQuotaEnforcementHandler.class); + AggServerHttpRequestStats singleGetStats = Mockito.mock(AggServerHttpRequestStats.class); + AggServerHttpRequestStats multiGetStats = Mockito.mock(AggServerHttpRequestStats.class); + AggServerHttpRequestStats computeStats = Mockito.mock(AggServerHttpRequestStats.class); + GrpcReplyProcessor grpcReplyProcessor = Mockito.mock(GrpcReplyProcessor.class); + + // Test with all fields set + GrpcServiceDependencies dependencies = + new GrpcServiceDependencies.Builder().setDiskHealthCheckService(diskHealthCheckService) + .setStorageReadRequestHandler(storageReadRequestHandler) + .setQuotaEnforcementHandler(quotaEnforcementHandler) + .setSingleGetStats(singleGetStats) + .setMultiGetStats(multiGetStats) + .setComputeStats(computeStats) + .setGrpcReplyProcessor(grpcReplyProcessor) + .build(); + + assertNotNull(dependencies); + assertSame(dependencies.getDiskHealthCheckService(), diskHealthCheckService); + assertSame(dependencies.getStorageReadRequestHandler(), storageReadRequestHandler); + assertSame(dependencies.getQuotaEnforcementHandler(), quotaEnforcementHandler); + assertSame(dependencies.getSingleGetStats(), singleGetStats); + assertSame(dependencies.getMultiGetStats(), multiGetStats); + assertSame(dependencies.getComputeStats(), computeStats); + assertSame(dependencies.getGrpcReplyProcessor(), grpcReplyProcessor); + } + + @Test + public void testBuilderValidationWithMissingFields() { + assertThrows( + NullPointerException.class, + () -> new GrpcServiceDependencies.Builder().setDiskHealthCheckService(null).build()); + } + + @Test + public void testBuilderValidationWithDefaultValues() { + DiskHealthCheckService diskHealthCheckService = Mockito.mock(DiskHealthCheckService.class); + StorageReadRequestHandler storageReadRequestHandler = Mockito.mock(StorageReadRequestHandler.class); + AggServerHttpRequestStats singleGetStats = Mockito.mock(AggServerHttpRequestStats.class); + AggServerHttpRequestStats multiGetStats = Mockito.mock(AggServerHttpRequestStats.class); + AggServerHttpRequestStats computeStats = Mockito.mock(AggServerHttpRequestStats.class); + + GrpcServiceDependencies dependencies = + new GrpcServiceDependencies.Builder().setDiskHealthCheckService(diskHealthCheckService) + .setStorageReadRequestHandler(storageReadRequestHandler) + .setSingleGetStats(singleGetStats) + .setMultiGetStats(multiGetStats) + .setComputeStats(computeStats) + .build(); + + assertNotNull(dependencies); + assertTrue(dependencies.getQuotaEnforcementHandler() instanceof NoOpReadQuotaEnforcementHandler); + assertNotNull(dependencies.getGrpcReplyProcessor()); + assertSame(dependencies.getDiskHealthCheckService(), diskHealthCheckService); + assertSame(dependencies.getStorageReadRequestHandler(), storageReadRequestHandler); + assertSame(dependencies.getSingleGetStats(), singleGetStats); + assertSame(dependencies.getMultiGetStats(), multiGetStats); + assertSame(dependencies.getComputeStats(), computeStats); + } +} diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcStorageResponseHandlerCallbackTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcStorageResponseHandlerCallbackTest.java new file mode 100644 index 0000000000..e1651a1dd7 --- /dev/null +++ b/services/venice-server/src/test/java/com/linkedin/venice/grpc/GrpcStorageResponseHandlerCallbackTest.java @@ -0,0 +1,86 @@ +package com.linkedin.venice.grpc; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.reset; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.linkedin.davinci.listener.response.ReadResponse; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class GrpcStorageResponseHandlerCallbackTest { + private GrpcRequestContext requestContext; + private GrpcReplyProcessor grpcReplyProcessor; + private GrpcStorageResponseHandlerCallback callback; + + @BeforeMethod + public void setUp() { + requestContext = mock(GrpcRequestContext.class); + grpcReplyProcessor = mock(GrpcReplyProcessor.class); + callback = GrpcStorageResponseHandlerCallback.create(requestContext, grpcReplyProcessor); + } + + @Test + public void testOnReadResponseVariousScenarios() { + // Case 1: Response found + ReadResponse foundResponse = mock(ReadResponse.class); + when(foundResponse.isFound()).thenReturn(true); + ByteBuf foundResponseBody = Unpooled.EMPTY_BUFFER; + when(foundResponse.getResponseBody()).thenReturn(foundResponseBody); + + callback.onReadResponse(foundResponse); + + verify(requestContext).setReadResponseStatus(VeniceReadResponseStatus.OK); + verify(requestContext).setReadResponse(foundResponse); + verify(grpcReplyProcessor).sendResponse(requestContext); + + reset(requestContext, grpcReplyProcessor); // Resetting mocks before next scenario + + // Case 2: Response not found (key not found) + ReadResponse notFoundResponse = mock(ReadResponse.class); + when(notFoundResponse.isFound()).thenReturn(false); + + callback.onReadResponse(notFoundResponse); + + verify(requestContext).setReadResponseStatus(VeniceReadResponseStatus.KEY_NOT_FOUND); + verify(requestContext).setReadResponse(notFoundResponse); + verify(grpcReplyProcessor).sendResponse(requestContext); + + reset(requestContext, grpcReplyProcessor); // Resetting mocks before next scenario + + // Case 3: Null ReadResponse + callback.onReadResponse(null); + + verify(requestContext).setReadResponseStatus(VeniceReadResponseStatus.INTERNAL_SERVER_ERROR); + verify(requestContext).setReadResponse(null); + verify(grpcReplyProcessor).sendResponse(requestContext); + } + + @Test + public void testOnErrorVariousScenarios() { + // Case 1: Standard error case + VeniceReadResponseStatus errorStatus = VeniceReadResponseStatus.INTERNAL_SERVER_ERROR; + String errorMessage = "An error occurred"; + + callback.onError(errorStatus, errorMessage); + + verify(requestContext).setReadResponseStatus(errorStatus); + verify(requestContext).setErrorMessage(errorMessage); + verify(requestContext).setReadResponse(null); + verify(grpcReplyProcessor).sendResponse(requestContext); + + reset(requestContext, grpcReplyProcessor); + + // Case 2: Null error message + callback.onError(errorStatus, null); + verify(requestContext).setReadResponseStatus(errorStatus); + verify(requestContext).setErrorMessage(null); + verify(requestContext).setReadResponse(null); + verify(grpcReplyProcessor).sendResponse(requestContext); + } +} diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/ServerStatsContextTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/RequestStatsRecorderTest.java similarity index 84% rename from services/venice-server/src/test/java/com/linkedin/venice/grpc/ServerStatsContextTest.java rename to services/venice-server/src/test/java/com/linkedin/venice/grpc/RequestStatsRecorderTest.java index 0c1d0922e2..557a202641 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/grpc/ServerStatsContextTest.java +++ b/services/venice-server/src/test/java/com/linkedin/venice/grpc/RequestStatsRecorderTest.java @@ -5,7 +5,7 @@ import static org.mockito.Mockito.verify; import static org.testng.Assert.assertEquals; -import com.linkedin.venice.listener.ServerStatsContext; +import com.linkedin.venice.listener.RequestStatsRecorder; import com.linkedin.venice.listener.request.RouterRequest; import com.linkedin.venice.listener.response.stats.ComputeResponseStats; import com.linkedin.venice.read.RequestType; @@ -16,7 +16,7 @@ import org.testng.annotations.Test; -public class ServerStatsContextTest { +public class RequestStatsRecorderTest { private AggServerHttpRequestStats singleGetStats; private AggServerHttpRequestStats multiGetStats; @@ -32,14 +32,14 @@ public void setUp() { @Test public void testSetStoreName() { - ServerStatsContext context = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); + RequestStatsRecorder context = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); context.setStoreName("testStore"); assertEquals("testStore", context.getStoreName()); } @Test public void testSetRequestType() { - ServerStatsContext context = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); + RequestStatsRecorder context = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); context.setRequestType(RequestType.SINGLE_GET); assertEquals(singleGetStats, context.getCurrentStats()); @@ -53,7 +53,7 @@ public void testSetRequestType() { @Test public void testSuccessRequest() { - ServerStatsContext context = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); + RequestStatsRecorder context = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); context.setStoreName("testStore"); ServerHttpRequestStats stats = mock(ServerHttpRequestStats.class); context.successRequest(stats, 10.5); @@ -64,7 +64,7 @@ public void testSuccessRequest() { @Test public void testErrorRequest() { - ServerStatsContext context = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); + RequestStatsRecorder context = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); ServerHttpRequestStats stats = mock(ServerHttpRequestStats.class); context.setRequestType(RequestType.SINGLE_GET); @@ -88,7 +88,7 @@ public void testRequestInfo() { doReturn(RequestType.MULTI_GET).when(routerRequest).getRequestType(); doReturn(123).when(routerRequest).getKeyCount(); - ServerStatsContext context = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); + RequestStatsRecorder context = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); context.setRequestInfo(routerRequest); assertEquals("testStore", context.getStoreName()); @@ -98,7 +98,7 @@ public void testRequestInfo() { @Test public void setRequestType() { - ServerStatsContext context = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); + RequestStatsRecorder context = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); context.setRequestType(RequestType.SINGLE_GET); assertEquals(singleGetStats, context.getCurrentStats()); @@ -112,7 +112,7 @@ public void setRequestType() { @Test public void testRecordBasicMetrics() { - ServerStatsContext context = new ServerStatsContext(singleGetStats, multiGetStats, computeStats); + RequestStatsRecorder context = new RequestStatsRecorder(singleGetStats, multiGetStats, computeStats); ServerHttpRequestStats stats = mock(ServerHttpRequestStats.class); context.setStoreName("testStore"); context.setRequestType(RequestType.COMPUTE); diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcReadServiceImplTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcReadServiceImplTest.java new file mode 100644 index 0000000000..97255e5e0d --- /dev/null +++ b/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcReadServiceImplTest.java @@ -0,0 +1,277 @@ +package com.linkedin.venice.grpc; + +import static com.linkedin.venice.listener.StorageReadRequestHandler.VENICE_STORAGE_NODE_HARDWARE_IS_NOT_HEALTHY_MSG; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + +import com.linkedin.davinci.listener.response.ServerCurrentVersionResponse; +import com.linkedin.davinci.listener.response.TopicPartitionIngestionContextResponse; +import com.linkedin.davinci.storage.DiskHealthCheckService; +import com.linkedin.venice.HttpConstants; +import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.listener.QuotaEnforcementHandler; +import com.linkedin.venice.listener.StorageReadRequestHandler; +import com.linkedin.venice.listener.response.BinaryResponse; +import com.linkedin.venice.meta.ServerAdminAction; +import com.linkedin.venice.protocols.AdminRequest; +import com.linkedin.venice.protocols.AdminResponse; +import com.linkedin.venice.protocols.CompressionDictionaryRequest; +import com.linkedin.venice.protocols.CompressionDictionaryResponse; +import com.linkedin.venice.protocols.CurrentVersionInfoRequest; +import com.linkedin.venice.protocols.CurrentVersionInfoResponse; +import com.linkedin.venice.protocols.HealthCheckRequest; +import com.linkedin.venice.protocols.HealthCheckResponse; +import com.linkedin.venice.protocols.IngestionContextRequest; +import com.linkedin.venice.protocols.IngestionContextResponse; +import com.linkedin.venice.protocols.MetadataRequest; +import com.linkedin.venice.protocols.MetadataResponse; +import com.linkedin.venice.protocols.VeniceReadServiceGrpc; +import com.linkedin.venice.protocols.VeniceReadServiceGrpc.VeniceReadServiceBlockingStub; +import com.linkedin.venice.response.VeniceReadResponseStatus; +import com.linkedin.venice.stats.AggServerHttpRequestStats; +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.netty.buffer.Unpooled; +import java.io.IOException; +import java.nio.ByteBuffer; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class VeniceGrpcReadServiceImplTest { + private Server grpcServer; + private ManagedChannel grpcChannel; + private GrpcServiceDependencies grpcServiceDependencies; + private DiskHealthCheckService diskHealthCheckService; + private StorageReadRequestHandler storageReadRequestHandler; + private QuotaEnforcementHandler quotaEnforcementHandler; + private AggServerHttpRequestStats singleGetStats; + private AggServerHttpRequestStats multiGetStats; + private AggServerHttpRequestStats computeStats; + + @BeforeMethod + public void setUp() throws IOException { + diskHealthCheckService = mock(DiskHealthCheckService.class); + storageReadRequestHandler = mock(StorageReadRequestHandler.class); + quotaEnforcementHandler = mock(QuotaEnforcementHandler.class); + singleGetStats = mock(AggServerHttpRequestStats.class); + multiGetStats = mock(AggServerHttpRequestStats.class); + computeStats = mock(AggServerHttpRequestStats.class); + + grpcServiceDependencies = new GrpcServiceDependencies.Builder().setDiskHealthCheckService(diskHealthCheckService) + .setStorageReadRequestHandler(storageReadRequestHandler) + .setQuotaEnforcementHandler(quotaEnforcementHandler) + .setSingleGetStats(singleGetStats) + .setMultiGetStats(multiGetStats) + .setComputeStats(computeStats) + .build(); + + // Create a unique server name for the in-process server + String serverName = InProcessServerBuilder.generateName(); + + // Start the gRPC server in-process + grpcServer = InProcessServerBuilder.forName(serverName) + .directExecutor() + .addService(new VeniceGrpcReadServiceImpl(grpcServiceDependencies)) + .build() + .start(); + + // Create a channel to communicate with the server + grpcChannel = InProcessChannelBuilder.forName(serverName).directExecutor().build(); + } + + @AfterMethod + public void cleanUp() { + // Shut down the channel and server after each test + if (grpcChannel != null) { + grpcChannel.shutdownNow(); + } + if (grpcServer != null) { + grpcServer.shutdownNow(); + } + } + + @Test + public void testIsServerHealthy() { + VeniceReadServiceBlockingStub blockingStub = VeniceReadServiceGrpc.newBlockingStub(grpcChannel); + + when(diskHealthCheckService.isDiskHealthy()).thenReturn(true); + HealthCheckRequest request = HealthCheckRequest.newBuilder().build(); + assertEquals(blockingStub.isServerHealthy(request).getStatusCode(), VeniceReadResponseStatus.OK.getCode()); + + when(diskHealthCheckService.isDiskHealthy()).thenReturn(false); + HealthCheckResponse response = blockingStub.isServerHealthy(request); + assertEquals(response.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + assertNotNull(response.getMessage()); + assertTrue(response.getMessage().contains(VENICE_STORAGE_NODE_HARDWARE_IS_NOT_HEALTHY_MSG)); + } + + @Test + public void testGetCompressionDictionary() { + // Case 1: Non-empty buffer response means dictionary found + VeniceReadServiceBlockingStub blockingStub = VeniceReadServiceGrpc.newBlockingStub(grpcChannel); + CompressionDictionaryRequest request = + CompressionDictionaryRequest.newBuilder().setStoreName("testStore").setStoreVersion(1).build(); + BinaryResponse binaryResponse = new BinaryResponse(ByteBuffer.wrap(new byte[] { 4, 5, 6 })); + when(storageReadRequestHandler.handleDictionaryFetchRequest(any())).thenReturn(binaryResponse); + + CompressionDictionaryResponse actualResponse = blockingStub.getCompressionDictionary(request); + assertEquals(actualResponse.getStatusCode(), VeniceReadResponseStatus.OK.getCode()); + assertEquals(actualResponse.getValue().asReadOnlyByteBuffer(), ByteBuffer.wrap(new byte[] { 4, 5, 6 })); + assertEquals(actualResponse.getContentType(), HttpConstants.BINARY); + assertEquals(actualResponse.getContentLength(), 3); + + // Case 2: Empty buffer response means dictionary not found + BinaryResponse emptyBufferResponse = new BinaryResponse((ByteBuffer) null); + when(storageReadRequestHandler.handleDictionaryFetchRequest(any())).thenReturn(emptyBufferResponse); + CompressionDictionaryResponse emptyBufferActualResponse = blockingStub.getCompressionDictionary(request); + assertEquals(emptyBufferActualResponse.getStatusCode(), VeniceReadResponseStatus.KEY_NOT_FOUND.getCode()); + + // Case 3: Exception was thrown when handling the request hence return INTERNAL_SERVER_ERROR + when(storageReadRequestHandler.handleDictionaryFetchRequest(any())) + .thenThrow(new VeniceException("Test exception")); + CompressionDictionaryResponse exceptionActualResponse = blockingStub.getCompressionDictionary(request); + assertEquals(exceptionActualResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + } + + @Test + public void testHandleAdminRequest() { + // Case 1: Admin request is handled successfully with a response + AdminRequest adminRequest = AdminRequest.newBuilder() + .setResourceName("testStore_v1") + .setServerAdminAction(ServerAdminAction.DUMP_SERVER_CONFIGS.name()) + .build(); + VeniceReadServiceBlockingStub blockingStub = VeniceReadServiceGrpc.newBlockingStub(grpcChannel); + com.linkedin.davinci.listener.response.AdminResponse adminResponse = + mock(com.linkedin.davinci.listener.response.AdminResponse.class); + when(adminResponse.isError()).thenReturn(false); + when(storageReadRequestHandler.handleServerAdminRequest(any())).thenReturn(adminResponse); + when(adminResponse.getResponseBody()).thenReturn(Unpooled.wrappedBuffer(new byte[] { 1, 2, 3 })); + + AdminResponse grpcAdminResponse = blockingStub.handleAdminRequest(adminRequest); + assertEquals(grpcAdminResponse.getStatusCode(), VeniceReadResponseStatus.OK.getCode()); + assertEquals(grpcAdminResponse.getValue().asReadOnlyByteBuffer(), ByteBuffer.wrap(new byte[] { 1, 2, 3 })); + assertEquals(grpcAdminResponse.getContentType(), HttpConstants.AVRO_BINARY); + assertEquals(grpcAdminResponse.getContentLength(), 3); + + // Case 2: Admin request is handled successfully with error response + when(adminResponse.isError()).thenReturn(true); + when(adminResponse.getMessage()).thenReturn("Test error message"); + when(adminResponse.getResponseBody()).thenReturn(Unpooled.EMPTY_BUFFER); + AdminResponse grpcAdminErrorResponse = blockingStub.handleAdminRequest(adminRequest); + assertEquals(grpcAdminErrorResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + assertEquals(grpcAdminErrorResponse.getValue().asReadOnlyByteBuffer(), ByteBuffer.wrap(new byte[0])); + assertEquals(grpcAdminErrorResponse.getContentType(), HttpConstants.TEXT_PLAIN); + assertEquals(grpcAdminErrorResponse.getContentLength(), 0); + + // Case 3: Exception was thrown when handling the request hence return INTERNAL_SERVER_ERROR + when(storageReadRequestHandler.handleServerAdminRequest(any())).thenThrow(new VeniceException("Test exception")); + AdminResponse exceptionActualResponse = blockingStub.handleAdminRequest(adminRequest); + assertEquals(exceptionActualResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + } + + @Test + public void testGetMetadata() { + // Case 1: Metadata request is handled successfully with a response + MetadataRequest request = MetadataRequest.newBuilder().setStoreName("testStore").build(); + VeniceReadServiceBlockingStub blockingStub = VeniceReadServiceGrpc.newBlockingStub(grpcChannel); + + com.linkedin.davinci.listener.response.MetadataResponse metadataResponseMock = + mock(com.linkedin.davinci.listener.response.MetadataResponse.class); + when(metadataResponseMock.isError()).thenReturn(false); + when(metadataResponseMock.getResponseBody()).thenReturn(Unpooled.wrappedBuffer(new byte[] { 1, 2, 3 })); + when(metadataResponseMock.getResponseSchemaIdHeader()).thenReturn(-11); + when(storageReadRequestHandler.handleMetadataFetchRequest(any())).thenReturn(metadataResponseMock); + + MetadataResponse metadataResponse = blockingStub.getMetadata(request); + assertEquals(metadataResponse.getStatusCode(), VeniceReadResponseStatus.OK.getCode()); + assertEquals(metadataResponse.getValue().asReadOnlyByteBuffer(), ByteBuffer.wrap(new byte[] { 1, 2, 3 })); + assertEquals(metadataResponse.getContentType(), HttpConstants.AVRO_BINARY); + assertEquals(metadataResponse.getContentLength(), 3); + assertEquals(metadataResponse.getSchemaId(), -11); + + // Case 2: Metadata request is handled successfully with error response + when(metadataResponseMock.isError()).thenReturn(true); + when(metadataResponseMock.getMessage()).thenReturn("Test error message"); + when(metadataResponseMock.getResponseBody()).thenReturn(Unpooled.EMPTY_BUFFER); + MetadataResponse grpcMetadataErrorResponse = blockingStub.getMetadata(request); + assertEquals(grpcMetadataErrorResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + assertEquals(grpcMetadataErrorResponse.getValue().asReadOnlyByteBuffer(), ByteBuffer.wrap(new byte[0])); + assertEquals(grpcMetadataErrorResponse.getContentType(), HttpConstants.TEXT_PLAIN); + assertEquals(grpcMetadataErrorResponse.getContentLength(), 0); + + // Case 3: Exception was thrown when handling the request hence return INTERNAL_SERVER_ERROR + when(storageReadRequestHandler.handleMetadataFetchRequest(any())).thenThrow(new VeniceException("Test exception")); + MetadataResponse exceptionActualResponse = blockingStub.getMetadata(request); + assertEquals(exceptionActualResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + } + + @Test + public void testGetCurrentVersionInfo() { + // Case 1: Current version info request is handled successfully with a response + CurrentVersionInfoRequest request = CurrentVersionInfoRequest.newBuilder().setStoreName("testStore").build(); + VeniceReadServiceBlockingStub blockingStub = VeniceReadServiceGrpc.newBlockingStub(grpcChannel); + ServerCurrentVersionResponse serverCurrentVersionResponseMock = new ServerCurrentVersionResponse(); + serverCurrentVersionResponseMock.setCurrentVersion(2); + serverCurrentVersionResponseMock.setError(false); + when(storageReadRequestHandler.handleCurrentVersionRequest(any())).thenReturn(serverCurrentVersionResponseMock); + CurrentVersionInfoResponse response = blockingStub.getCurrentVersionInfo(request); + assertEquals(response.getStatusCode(), VeniceReadResponseStatus.OK.getCode()); + assertEquals(response.getCurrentVersion(), 2); + + // Case 2: Current version info request is handled successfully with error response + serverCurrentVersionResponseMock.setError(true); + serverCurrentVersionResponseMock.setCurrentVersion(-1); + serverCurrentVersionResponseMock.setMessage("Test error message"); + CurrentVersionInfoResponse errorResponse = blockingStub.getCurrentVersionInfo(request); + assertEquals(errorResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + assertEquals(errorResponse.getErrorMessage(), "Test error message"); + + // Case 3: Exception was thrown when handling the request hence return INTERNAL_SERVER_ERROR + when(storageReadRequestHandler.handleCurrentVersionRequest(any())).thenThrow(new VeniceException("Test exception")); + CurrentVersionInfoResponse exceptionActualResponse = blockingStub.getCurrentVersionInfo(request); + assertEquals(exceptionActualResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + } + + @Test + public void testGetIngestionContext() { + // Case 1: Ingestion context request is handled successfully with a response + IngestionContextRequest request = + IngestionContextRequest.newBuilder().setTopicName("testStore_v1").setPartition(34).build(); + VeniceReadServiceBlockingStub blockingStub = VeniceReadServiceGrpc.newBlockingStub(grpcChannel); + String jsonStr = "{\n" + "\"kafkaUrl\" : {\n" + " TP(topic: \"testStore_v1\", partition: 34) " + ": {\n" + + " \"latestOffset\" : 0,\n" + " \"offsetLag\" : 1,\n" + " \"msgRate\" : 2.0,\n" + + " \"byteRate\" : 4.0,\n" + " \"consumerIdx\" : 6,\n" + + " \"elapsedTimeSinceLastPollInMs\" : 7\n" + " }\n" + " }\n" + "}"; + byte[] ingestionCtxBytes = jsonStr.getBytes(); + TopicPartitionIngestionContextResponse responseMock = new TopicPartitionIngestionContextResponse(); + responseMock.setTopicPartitionIngestionContext(ingestionCtxBytes); + responseMock.setError(false); + when(storageReadRequestHandler.handleTopicPartitionIngestionContextRequest(any())).thenReturn(responseMock); + IngestionContextResponse response = blockingStub.getIngestionContext(request); + assertEquals(response.getStatusCode(), VeniceReadResponseStatus.OK.getCode()); + assertEquals(response.getValue().asReadOnlyByteBuffer(), ByteBuffer.wrap(ingestionCtxBytes)); + assertEquals(response.getContentLength(), ingestionCtxBytes.length); + + // Case 2: Ingestion context request is handled successfully with error response + responseMock.setError(true); + responseMock.setMessage("Test error message"); + responseMock.setTopicPartitionIngestionContext(null); + IngestionContextResponse errorResponse = blockingStub.getIngestionContext(request); + assertEquals(errorResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + assertEquals(errorResponse.getErrorMessage(), "Test error message"); + + // Case 3: Exception was thrown when handling the request hence return INTERNAL_SERVER_ERROR + when(storageReadRequestHandler.handleTopicPartitionIngestionContextRequest(any())) + .thenThrow(new VeniceException("Test exception")); + IngestionContextResponse exceptionActualResponse = blockingStub.getIngestionContext(request); + assertEquals(exceptionActualResponse.getStatusCode(), VeniceReadResponseStatus.INTERNAL_SERVER_ERROR.getCode()); + } +} diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java index 47a6ef204b..3ed2e98124 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java +++ b/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java @@ -5,10 +5,11 @@ import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; +import com.linkedin.davinci.storage.DiskHealthCheckService; import com.linkedin.venice.exceptions.VeniceException; -import com.linkedin.venice.listener.grpc.VeniceReadServiceImpl; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcRequestProcessor; +import com.linkedin.venice.listener.StorageReadRequestHandler; import com.linkedin.venice.security.SSLFactory; +import com.linkedin.venice.stats.AggServerHttpRequestStats; import com.linkedin.venice.utils.SslUtils; import com.linkedin.venice.utils.TestUtils; import io.grpc.InsecureServerCredentials; @@ -21,14 +22,21 @@ public class VeniceGrpcServerTest { private VeniceGrpcServer grpcServer; private VeniceGrpcServerConfig.Builder serverConfig; - private VeniceServerGrpcRequestProcessor grpcRequestProcessor; + private GrpcIoRequestProcessor grpcRequestProcessor; @BeforeMethod void setUp() { - grpcRequestProcessor = mock(VeniceServerGrpcRequestProcessor.class); + grpcRequestProcessor = mock(GrpcIoRequestProcessor.class); + GrpcServiceDependencies grpcServiceDependencies = + new GrpcServiceDependencies.Builder().setSingleGetStats(mock(AggServerHttpRequestStats.class)) + .setMultiGetStats(mock(AggServerHttpRequestStats.class)) + .setComputeStats(mock(AggServerHttpRequestStats.class)) + .setStorageReadRequestHandler(mock(StorageReadRequestHandler.class)) + .setDiskHealthCheckService(mock(DiskHealthCheckService.class)) + .build(); serverConfig = new VeniceGrpcServerConfig.Builder().setPort(TestUtils.getFreePort()) .setNumThreads(10) - .setService(new VeniceReadServiceImpl(grpcRequestProcessor)); + .setService(new VeniceGrpcReadServiceImpl(grpcServiceDependencies, grpcRequestProcessor)); } @Test diff --git a/services/venice-server/src/test/java/com/linkedin/venice/listener/HttpChannelInitializerTest.java b/services/venice-server/src/test/java/com/linkedin/venice/listener/HttpChannelInitializerTest.java index d838ca59ba..f41115b527 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/listener/HttpChannelInitializerTest.java +++ b/services/venice-server/src/test/java/com/linkedin/venice/listener/HttpChannelInitializerTest.java @@ -9,7 +9,6 @@ import com.linkedin.venice.acl.StaticAccessController; import com.linkedin.venice.authorization.DefaultIdentityParser; import com.linkedin.venice.helix.HelixCustomizedViewOfflinePushRepository; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcRequestProcessor; import com.linkedin.venice.meta.ReadOnlyStoreRepository; import com.linkedin.venice.security.SSLConfig; import com.linkedin.venice.security.SSLFactory; @@ -111,25 +110,6 @@ public void testInitChannelWithSSLExecutor() { initializer.initChannel(ch); } - @Test - public void initGrpcRequestProcessor() { - SSLConfig sslConfig = new SSLConfig(); - doReturn(sslConfig).when(sslFactory).getSSLConfig(); - HttpChannelInitializer initializer = new HttpChannelInitializer( - storeMetadataRepository, - customizedViewRepository, - metricsRepository, - sslFactoryOptional, - sslHandshakeExecutor, - serverConfig, - accessController, - storeAccessController, - requestHandler); - - VeniceServerGrpcRequestProcessor processor = initializer.initGrpcRequestProcessor(); - Assert.assertNotNull(processor); - } - @Test public void initGprcHandlersTestInterceptors() { SSLConfig sslConfig = new SSLConfig(); diff --git a/services/venice-server/src/test/java/com/linkedin/venice/listener/OutboundHttpWrapperHandlerTest.java b/services/venice-server/src/test/java/com/linkedin/venice/listener/OutboundHttpWrapperHandlerTest.java index d67aaeb86f..ed83780fc6 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/listener/OutboundHttpWrapperHandlerTest.java +++ b/services/venice-server/src/test/java/com/linkedin/venice/listener/OutboundHttpWrapperHandlerTest.java @@ -5,24 +5,16 @@ import static io.netty.handler.codec.http.HttpVersion.HTTP_1_1; import static org.mockito.Mockito.any; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.protobuf.ByteString; import com.linkedin.davinci.listener.response.MetadataResponse; -import com.linkedin.davinci.listener.response.ReadResponse; import com.linkedin.davinci.listener.response.ServerCurrentVersionResponse; import com.linkedin.davinci.listener.response.TopicPartitionIngestionContextResponse; import com.linkedin.venice.HttpConstants; import com.linkedin.venice.compression.CompressionStrategy; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.grpc.handlers.GrpcOutboundResponseHandler; -import com.linkedin.venice.protocols.VeniceServerResponse; import com.linkedin.venice.utils.ObjectMapperFactory; -import io.grpc.stub.StreamObserver; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.channel.ChannelHandlerContext; @@ -32,17 +24,24 @@ import java.nio.charset.StandardCharsets; import java.util.Collections; import org.testng.Assert; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class OutboundHttpWrapperHandlerTest { private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getInstance(); + private RequestStatsRecorder requestStatsRecorder; + + @BeforeMethod + public void setUp() { + requestStatsRecorder = mock(RequestStatsRecorder.class); + } + @Test public void testWriteMetadataResponse() { MetadataResponse msg = new MetadataResponse(); msg.setVersions(Collections.emptyList()); - StatsHandler statsHandler = mock(StatsHandler.class); ChannelHandlerContext mockCtx = mock(ChannelHandlerContext.class); FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, HttpResponseStatus.OK, msg.getResponseBody()); @@ -52,7 +51,7 @@ public void testWriteMetadataResponse() { response.headers().set(HttpConstants.VENICE_SCHEMA_ID, msg.getResponseSchemaIdHeader()); response.headers().set(HttpConstants.VENICE_REQUEST_RCU, 1); - OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(statsHandler); + OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(requestStatsRecorder); when(mockCtx.writeAndFlush(any())).then(i -> { FullHttpResponse actualResponse = (DefaultFullHttpResponse) i.getArguments()[0]; @@ -69,7 +68,6 @@ public void testWriteMetadataResponse() { public void testWriteCurrentVersionResponse() throws JsonProcessingException { ServerCurrentVersionResponse msg = new ServerCurrentVersionResponse(); msg.setCurrentVersion(2); - StatsHandler statsHandler = mock(StatsHandler.class); ChannelHandlerContext mockCtx = mock(ChannelHandlerContext.class); ByteBuf body = Unpooled.wrappedBuffer(OBJECT_MAPPER.writeValueAsBytes(msg)); FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, HttpResponseStatus.OK, body); @@ -79,7 +77,7 @@ public void testWriteCurrentVersionResponse() throws JsonProcessingException { response.headers().set(HttpConstants.VENICE_SCHEMA_ID, -1); response.headers().set(HttpConstants.VENICE_REQUEST_RCU, 1); - OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(statsHandler); + OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(requestStatsRecorder); when(mockCtx.writeAndFlush(any())).then(i -> { FullHttpResponse actualResponse = (DefaultFullHttpResponse) i.getArguments()[0]; @@ -97,7 +95,6 @@ public void testWriteMetadataErrorResponse() { msg.setError(true); msg.setMessage("test-error"); ByteBuf body = Unpooled.wrappedBuffer(msg.getMessage().getBytes(StandardCharsets.UTF_8)); - StatsHandler statsHandler = mock(StatsHandler.class); ChannelHandlerContext mockCtx = mock(ChannelHandlerContext.class); FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, HttpResponseStatus.INTERNAL_SERVER_ERROR, body); @@ -107,7 +104,7 @@ public void testWriteMetadataErrorResponse() { response.headers().set(HttpConstants.VENICE_SCHEMA_ID, -1); response.headers().set(HttpConstants.VENICE_REQUEST_RCU, 1); - OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(statsHandler); + OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(requestStatsRecorder); when(mockCtx.writeAndFlush(any())).then(i -> { FullHttpResponse actualResponse = (DefaultFullHttpResponse) i.getArguments()[0]; @@ -125,7 +122,6 @@ public void testWriteMetadataUnknownErrorResponse() { MetadataResponse msg = new MetadataResponse(); msg.setError(true); ByteBuf body = Unpooled.wrappedBuffer("Unknown error".getBytes(StandardCharsets.UTF_8)); - StatsHandler statsHandler = mock(StatsHandler.class); ChannelHandlerContext mockCtx = mock(ChannelHandlerContext.class); FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, HttpResponseStatus.INTERNAL_SERVER_ERROR, body); @@ -135,7 +131,7 @@ public void testWriteMetadataUnknownErrorResponse() { response.headers().set(HttpConstants.VENICE_SCHEMA_ID, -1); response.headers().set(HttpConstants.VENICE_REQUEST_RCU, 1); - OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(statsHandler); + OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(requestStatsRecorder); when(mockCtx.writeAndFlush(any())).then(i -> { FullHttpResponse actualResponse = (DefaultFullHttpResponse) i.getArguments()[0]; @@ -151,12 +147,11 @@ public void testWriteMetadataUnknownErrorResponse() { @Test public void testWriteDefaultFullHttpResponse() { FullHttpResponse msg = new DefaultFullHttpResponse(HTTP_1_1, HttpResponseStatus.OK); - StatsHandler statsHandler = mock(StatsHandler.class); ChannelHandlerContext mockCtx = mock(ChannelHandlerContext.class); FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, HttpResponseStatus.OK); - OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(statsHandler); + OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(requestStatsRecorder); when(mockCtx.writeAndFlush(any())).then(i -> { FullHttpResponse actualResponse = (DefaultFullHttpResponse) i.getArguments()[0]; @@ -169,30 +164,6 @@ public void testWriteDefaultFullHttpResponse() { outboundHttpWrapperHandler.write(mockCtx, msg, null); } - @Test - public void testGrpcWrite() { - ByteBuf mockBody = mock(ByteBuf.class); - - ReadResponse readResponse = mock(ReadResponse.class); - when(readResponse.getResponseBody()).thenReturn(mockBody); - when(readResponse.getCompressionStrategy()).thenReturn(CompressionStrategy.NO_OP); - when(readResponse.isStreamingResponse()).thenReturn(false); - - GrpcRequestContext context = new GrpcRequestContext(null, VeniceServerResponse.newBuilder(), getStreamObserver()); - context.setReadResponse(readResponse); - - VeniceServerResponse.Builder responseBuilder = VeniceServerResponse.newBuilder(); - ServerStatsContext statsContext = mock(ServerStatsContext.class); - context.setGrpcStatsContext(statsContext); - GrpcOutboundResponseHandler grpcHandler = spy(new GrpcOutboundResponseHandler()); - - grpcHandler.processRequest(context); - - Assert.assertEquals(ByteString.empty(), responseBuilder.getData()); - - verify(grpcHandler).processRequest(context); - } - @Test public void testWriteTopicPartitionIngestionContextResponse() throws JsonProcessingException { TopicPartitionIngestionContextResponse msg = new TopicPartitionIngestionContextResponse(); @@ -203,11 +174,10 @@ public void testWriteTopicPartitionIngestionContextResponse() throws JsonProcess + " \"byteRate\" : 4.0,\n" + " \"consumerIdx\" : 6,\n" + " \"elapsedTimeSinceLastPollInMs\" : 7\n" + " }\n" + " }\n" + "}"; msg.setTopicPartitionIngestionContext(jsonStr.getBytes()); - StatsHandler statsHandler = mock(StatsHandler.class); ChannelHandlerContext mockCtx = mock(ChannelHandlerContext.class); ByteBuf body = Unpooled.wrappedBuffer(OBJECT_MAPPER.writeValueAsBytes(msg)); FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, HttpResponseStatus.OK, body); - OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(statsHandler); + OutboundHttpWrapperHandler outboundHttpWrapperHandler = new OutboundHttpWrapperHandler(requestStatsRecorder); when(mockCtx.writeAndFlush(any())).then(i -> { FullHttpResponse actualResponse = (DefaultFullHttpResponse) i.getArguments()[0]; @@ -216,23 +186,4 @@ public void testWriteTopicPartitionIngestionContextResponse() throws JsonProcess }); outboundHttpWrapperHandler.write(mockCtx, msg, null); } - - private StreamObserver getStreamObserver() { - return new StreamObserver() { - @Override - public void onNext(VeniceServerResponse value) { - - } - - @Override - public void onError(Throwable t) { - - } - - @Override - public void onCompleted() { - - } - }; - } } diff --git a/services/venice-server/src/test/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandlerTest.java b/services/venice-server/src/test/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandlerTest.java index af1c3945ab..f965932539 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandlerTest.java +++ b/services/venice-server/src/test/java/com/linkedin/venice/listener/ReadQuotaEnforcementHandlerTest.java @@ -1,6 +1,5 @@ package com.linkedin.venice.listener; -import static com.linkedin.venice.response.VeniceReadResponseStatus.BAD_REQUEST; import static com.linkedin.venice.throttle.VeniceRateLimiter.RateLimiterType; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; @@ -9,7 +8,6 @@ import static org.mockito.Mockito.any; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -26,10 +24,6 @@ import com.linkedin.davinci.config.VeniceServerConfig; import com.linkedin.venice.helix.HelixCustomizedViewOfflinePushRepository; -import com.linkedin.venice.listener.ReadQuotaEnforcementHandler.QuotaEnforcementResult; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.grpc.handlers.GrpcReadQuotaEnforcementHandler; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcHandler; import com.linkedin.venice.listener.request.RouterRequest; import com.linkedin.venice.listener.response.HttpShortcutResponse; import com.linkedin.venice.meta.Instance; @@ -38,9 +32,7 @@ import com.linkedin.venice.meta.ReadOnlyStoreRepository; import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.Version; -import com.linkedin.venice.protocols.VeniceServerResponse; import com.linkedin.venice.read.RequestType; -import com.linkedin.venice.response.VeniceReadResponseStatus; import com.linkedin.venice.routerapi.ReplicaState; import com.linkedin.venice.stats.AggServerQuotaUsageStats; import com.linkedin.venice.throttle.GuavaRateLimiter; @@ -75,12 +67,10 @@ public class ReadQuotaEnforcementHandlerTest { private ReadOnlyStoreRepository storeRepository; private HelixCustomizedViewOfflinePushRepository customizedViewRepository; private ReadQuotaEnforcementHandler quotaEnforcementHandler; - private GrpcReadQuotaEnforcementHandler grpcQuotaEnforcementHandler; private AggServerQuotaUsageStats stats; private MetricsRepository metricsRepository; private RouterRequest routerRequest; - private VeniceServerGrpcHandler mockNextHandler; @BeforeMethod public void setUp() { @@ -105,10 +95,6 @@ public void setUp() { stats, metricsRepository, clock); - grpcQuotaEnforcementHandler = new GrpcReadQuotaEnforcementHandler(quotaEnforcementHandler); - mockNextHandler = mock(VeniceServerGrpcHandler.class); - grpcQuotaEnforcementHandler.addNextHandler(mockNextHandler); - doNothing().when(mockNextHandler).processRequest(any()); routerRequest = mock(RouterRequest.class); } @@ -122,7 +108,9 @@ public void testEnforceQuotaWhenResourceIsNotFound() { String storeName = "test_store"; when(routerRequest.getStoreName()).thenReturn(storeName); when(storeRepository.getStore(storeName)).thenReturn(null); - assertEquals(quotaEnforcementHandler.enforceQuota(routerRequest), QuotaEnforcementResult.BAD_REQUEST); + assertEquals( + quotaEnforcementHandler.enforceQuota(routerRequest), + QuotaEnforcementHandler.QuotaEnforcementResult.BAD_REQUEST); verify(stats, never()).recordAllowed(eq(storeName), eq(1L)); verify(stats, never()).recordRejected(eq(storeName), eq(1L)); @@ -137,15 +125,6 @@ public void testEnforceQuotaWhenResourceIsNotFound() { HttpShortcutResponse shortcutResponse = (HttpShortcutResponse) response; assertEquals(shortcutResponse.getStatus(), HttpResponseStatus.BAD_REQUEST); assertNotNull(shortcutResponse.getMessage()); - - // for gRPC handler - GrpcRequestContext grpcCtx = mock(GrpcRequestContext.class); - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - when(grpcCtx.getRouterRequest()).thenReturn(routerRequest); - when(grpcCtx.getVeniceServerResponseBuilder()).thenReturn(builder); - grpcQuotaEnforcementHandler.processRequest(grpcCtx); - assertEquals(builder.getErrorCode(), BAD_REQUEST); - assertNotNull(builder.getErrorMessage()); } @Test @@ -157,7 +136,9 @@ public void testEnforceQuotaDuringInitialization() { Store store = mock(Store.class); when(routerRequest.getStoreName()).thenReturn(storeName); when(storeRepository.getStore(storeName)).thenReturn(store); - assertEquals(quotaEnforcementHandler.enforceQuota(routerRequest), QuotaEnforcementResult.ALLOWED); + assertEquals( + quotaEnforcementHandler.enforceQuota(routerRequest), + QuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED); verify(stats, never()).recordAllowed(eq(storeName), eq(1L)); verify(stats, never()).recordRejected(eq(storeName), eq(1L)); @@ -170,15 +151,6 @@ public void testEnforceQuotaDuringInitialization() { Object receivedRequest = responseCaptor.getValue(); assertTrue(receivedRequest instanceof RouterRequest); assertEquals(receivedRequest, routerRequest); - - // for gRPC handler - GrpcRequestContext grpcCtx = mock(GrpcRequestContext.class); - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - when(grpcCtx.getRouterRequest()).thenReturn(routerRequest); - when(grpcCtx.getVeniceServerResponseBuilder()).thenReturn(builder); - grpcQuotaEnforcementHandler.processRequest(grpcCtx); - assertEquals(builder.getErrorCode(), 0); - verify(mockNextHandler).processRequest(grpcCtx); } @Test @@ -189,7 +161,9 @@ public void testEnforceQuotaWhenStorageNodeReadQuotaIsDisabled() { when(routerRequest.getStoreName()).thenReturn(storeName); when(storeRepository.getStore(storeName)).thenReturn(store); when(store.isStorageNodeReadQuotaEnabled()).thenReturn(false); - assertEquals(quotaEnforcementHandler.enforceQuota(routerRequest), QuotaEnforcementResult.ALLOWED); + assertEquals( + quotaEnforcementHandler.enforceQuota(routerRequest), + QuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED); verify(stats, never()).recordAllowed(eq(storeName), eq(1L)); verify(stats, never()).recordRejected(eq(storeName), eq(1L)); @@ -202,15 +176,6 @@ public void testEnforceQuotaWhenStorageNodeReadQuotaIsDisabled() { Object receivedRequest = responseCaptor.getValue(); assertTrue(receivedRequest instanceof RouterRequest); assertEquals(receivedRequest, routerRequest); - - // for gRPC handler - GrpcRequestContext grpcCtx = mock(GrpcRequestContext.class); - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - when(grpcCtx.getRouterRequest()).thenReturn(routerRequest); - when(grpcCtx.getVeniceServerResponseBuilder()).thenReturn(builder); - grpcQuotaEnforcementHandler.processRequest(grpcCtx); - assertEquals(builder.getErrorCode(), 0); - verify(mockNextHandler).processRequest(grpcCtx); } @Test @@ -229,7 +194,9 @@ public void testQuotaEnforcementWhenStoreVersionQuotaIsExceeded() { VeniceRateLimiter veniceRateLimiter = mock(VeniceRateLimiter.class); quotaEnforcementHandler.setStoreVersionRateLimiter(resourceName, veniceRateLimiter); when(veniceRateLimiter.tryAcquirePermit(1)).thenReturn(false); - assertEquals(quotaEnforcementHandler.enforceQuota(routerRequest), QuotaEnforcementResult.REJECTED); + assertEquals( + quotaEnforcementHandler.enforceQuota(routerRequest), + QuotaEnforcementHandler.QuotaEnforcementResult.REJECTED); verify(stats).recordRejected(eq(storeName), eq(1L)); verify(stats, never()).recordAllowed(eq(storeName), eq(1L)); @@ -243,18 +210,11 @@ public void testQuotaEnforcementWhenStoreVersionQuotaIsExceeded() { assertTrue(response instanceof HttpShortcutResponse); assertEquals(((HttpShortcutResponse) response).getStatus(), HttpResponseStatus.TOO_MANY_REQUESTS); - // for gRPC handler - GrpcRequestContext grpcCtx = mock(GrpcRequestContext.class); - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - when(grpcCtx.getRouterRequest()).thenReturn(routerRequest); - when(grpcCtx.getVeniceServerResponseBuilder()).thenReturn(builder); - grpcQuotaEnforcementHandler.processRequest(grpcCtx); - assertEquals(builder.getErrorCode(), VeniceReadResponseStatus.TOO_MANY_REQUESTS); - assertNotNull(builder.getErrorMessage()); - // Case 2: If request is a retry request, it should be allowed when(routerRequest.isRetryRequest()).thenReturn(true); - assertEquals(quotaEnforcementHandler.enforceQuota(routerRequest), QuotaEnforcementResult.ALLOWED); + assertEquals( + quotaEnforcementHandler.enforceQuota(routerRequest), + QuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED); verify(stats).recordAllowed(eq(storeName), eq(1L)); } @@ -274,7 +234,9 @@ public void testQuotaEnforcementWhenServerIsOverCapacity() { VeniceRateLimiter storageNodeRateLimiter = mock(VeniceRateLimiter.class); quotaEnforcementHandler.setStorageNodeRateLimiter(storageNodeRateLimiter); when(storageNodeRateLimiter.tryAcquirePermit(1)).thenReturn(false); - assertEquals(quotaEnforcementHandler.enforceQuota(routerRequest), QuotaEnforcementResult.OVER_CAPACITY); + assertEquals( + quotaEnforcementHandler.enforceQuota(routerRequest), + QuotaEnforcementHandler.QuotaEnforcementResult.OVER_CAPACITY); verify(stats, never()).recordAllowed(eq(storeName), eq(1L)); verify(stats, never()).recordRejected(eq(storeName), eq(1L)); @@ -287,15 +249,6 @@ public void testQuotaEnforcementWhenServerIsOverCapacity() { Object response = responseCaptor.getValue(); assertTrue(response instanceof HttpShortcutResponse); assertSame(((HttpShortcutResponse) response).getStatus(), HttpResponseStatus.SERVICE_UNAVAILABLE); - - // for gRPC handler - GrpcRequestContext grpcCtx = mock(GrpcRequestContext.class); - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - when(grpcCtx.getRouterRequest()).thenReturn(routerRequest); - when(grpcCtx.getVeniceServerResponseBuilder()).thenReturn(builder); - grpcQuotaEnforcementHandler.processRequest(grpcCtx); - assertEquals(builder.getErrorCode(), VeniceReadResponseStatus.SERVICE_UNAVAILABLE); - assertNotNull(builder.getErrorMessage()); } // both storeVersion quota and server capacity are sufficient to allow the request @@ -317,7 +270,9 @@ public void testQuotaEnforcement() { VeniceRateLimiter storageNodeRateLimiter = mock(VeniceRateLimiter.class); when(storageNodeRateLimiter.tryAcquirePermit(1)).thenReturn(true); quotaEnforcementHandler.setStorageNodeRateLimiter(storageNodeRateLimiter); - assertEquals(quotaEnforcementHandler.enforceQuota(routerRequest), QuotaEnforcementResult.ALLOWED); + assertEquals( + quotaEnforcementHandler.enforceQuota(routerRequest), + QuotaEnforcementHandler.QuotaEnforcementResult.ALLOWED); verify(stats).recordAllowed(eq(storeName), eq(1L)); verify(stats, never()).recordRejected(eq(storeName), eq(1L)); @@ -329,15 +284,6 @@ public void testQuotaEnforcement() { verify(ctx).fireChannelRead(responseCaptor.capture()); Object receivedRequest = responseCaptor.getValue(); assertTrue(receivedRequest instanceof RouterRequest); - - // for gRPC handler - GrpcRequestContext grpcCtx = mock(GrpcRequestContext.class); - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - when(grpcCtx.getRouterRequest()).thenReturn(routerRequest); - when(grpcCtx.getVeniceServerResponseBuilder()).thenReturn(builder); - grpcQuotaEnforcementHandler.processRequest(grpcCtx); - assertEquals(builder.getErrorCode(), 0); - verify(mockNextHandler).processRequest(grpcCtx); } @Test diff --git a/services/venice-server/src/test/java/com/linkedin/venice/listener/RouterRequestHttpHandlerTest.java b/services/venice-server/src/test/java/com/linkedin/venice/listener/RouterRequestHttpHandlerTest.java index 070dc760ae..71f1f8a5c7 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/listener/RouterRequestHttpHandlerTest.java +++ b/services/venice-server/src/test/java/com/linkedin/venice/listener/RouterRequestHttpHandlerTest.java @@ -1,22 +1,14 @@ package com.linkedin.venice.listener; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import com.linkedin.venice.HttpConstants; import com.linkedin.venice.exceptions.VeniceException; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.grpc.handlers.GrpcRouterRequestHandler; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcHandler; import com.linkedin.venice.listener.request.GetRouterRequest; import com.linkedin.venice.listener.request.HealthCheckRequest; -import com.linkedin.venice.listener.request.MultiGetRouterRequestWrapper; import com.linkedin.venice.listener.response.HttpShortcutResponse; import com.linkedin.venice.meta.QueryAction; -import com.linkedin.venice.protocols.VeniceClientRequest; -import com.linkedin.venice.protocols.VeniceServerResponse; import com.linkedin.venice.request.RequestHelper; import io.netty.channel.ChannelHandlerContext; import io.netty.handler.codec.http.DefaultFullHttpRequest; @@ -32,10 +24,18 @@ import java.util.Collections; import org.mockito.ArgumentCaptor; import org.testng.Assert; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; public class RouterRequestHttpHandlerTest { + private RequestStatsRecorder requestStatsRecorder; + + @BeforeMethod + public void setUp() { + requestStatsRecorder = mock(RequestStatsRecorder.class); + } + @Test public void parsesRequests() throws Exception { testRequestParsing("/storage/store_v1/1/key1", "store_v1", 1, "key1".getBytes(StandardCharsets.UTF_8)); @@ -48,12 +48,11 @@ public void parsesRequests() throws Exception { @Test public void respondsToHealthCheck() throws Exception { - RouterRequestHttpHandler testHander = - new RouterRequestHttpHandler(mock(StatsHandler.class), Collections.emptyMap()); + RouterRequestHttpHandler testHandler = new RouterRequestHttpHandler(requestStatsRecorder, Collections.emptyMap()); ChannelHandlerContext mockContext = mock(ChannelHandlerContext.class); ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(HealthCheckRequest.class); HttpRequest healthMsg = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/health"); - testHander.channelRead(mockContext, healthMsg); + testHandler.channelRead(mockContext, healthMsg); verify(mockContext).fireChannelRead(argumentCaptor.capture()); HealthCheckRequest requestObject = argumentCaptor.getValue(); Assert.assertNotNull(requestObject); @@ -63,12 +62,11 @@ public void testRequestParsing(String path, String expectedStore, int expectedPa throws Exception { // Test handler - RouterRequestHttpHandler testHander = - new RouterRequestHttpHandler(mock(StatsHandler.class), Collections.emptyMap()); + RouterRequestHttpHandler testHandler = new RouterRequestHttpHandler(requestStatsRecorder, Collections.emptyMap()); ChannelHandlerContext mockContext = mock(ChannelHandlerContext.class); ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(GetRouterRequest.class); HttpRequest msg = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, path); - testHander.channelRead(mockContext, msg); + testHandler.channelRead(mockContext, msg); verify(mockContext).fireChannelRead(argumentCaptor.capture()); GetRouterRequest requestObject = argumentCaptor.getValue(); Assert.assertEquals( @@ -99,12 +97,11 @@ public void testRequestParsing(String path, String expectedStore, int expectedPa } public void testBadRequest(String path, HttpMethod method) throws Exception { - RouterRequestHttpHandler testHander = - new RouterRequestHttpHandler(mock(StatsHandler.class), Collections.emptyMap()); + RouterRequestHttpHandler testHandler = new RouterRequestHttpHandler(requestStatsRecorder, Collections.emptyMap()); ChannelHandlerContext mockContext = mock(ChannelHandlerContext.class); ArgumentCaptor argumentCaptor = ArgumentCaptor.forClass(HttpShortcutResponse.class); HttpRequest msg = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, method, path); - testHander.channelRead(mockContext, msg); + testHandler.channelRead(mockContext, msg); verify(mockContext).writeAndFlush(argumentCaptor.capture()); HttpShortcutResponse httpShortcutResponse = argumentCaptor.getValue(); Assert.assertEquals(httpShortcutResponse.getStatus(), HttpResponseStatus.BAD_REQUEST); @@ -168,24 +165,4 @@ public void verifyGoodApiVersionOk() { headers.add(HttpConstants.VENICE_API_VERSION, "1"); GetRouterRequest.verifyApiVersion(headers, "1"); } - - @Test - public void testGrpcRead() { - VeniceClientRequest request = - VeniceClientRequest.newBuilder().setResourceName("teststore_v1").setIsBatchRequest(true).build(); - GrpcRequestContext ctx = new GrpcRequestContext(request, VeniceServerResponse.newBuilder(), null); - - ServerStatsContext statsContextMock = mock(ServerStatsContext.class); - ctx.setGrpcStatsContext(statsContextMock); - - GrpcRouterRequestHandler grpcRouterRequestHandler = new GrpcRouterRequestHandler(); - VeniceServerGrpcHandler mockNextHandler = mock(VeniceServerGrpcHandler.class); - grpcRouterRequestHandler.addNextHandler(mockNextHandler); - doNothing().when(mockNextHandler).processRequest(any()); - grpcRouterRequestHandler.processRequest(ctx); - - Assert.assertTrue(ctx.getRouterRequest() instanceof MultiGetRouterRequestWrapper); - Assert.assertEquals(ctx.getRouterRequest().getStoreName(), "teststore"); - Assert.assertEquals(ctx.getRouterRequest().getKeyCount(), 0); - } } diff --git a/services/venice-server/src/test/java/com/linkedin/venice/listener/StorageReadRequestHandlerTest.java b/services/venice-server/src/test/java/com/linkedin/venice/listener/StorageReadRequestHandlerTest.java index c0171a0d2d..7f08852fe4 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/listener/StorageReadRequestHandlerTest.java +++ b/services/venice-server/src/test/java/com/linkedin/venice/listener/StorageReadRequestHandlerTest.java @@ -8,14 +8,12 @@ import static org.mockito.ArgumentMatchers.intThat; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyInt; -import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.reset; -import static org.mockito.Mockito.spy; import static org.mockito.Mockito.timeout; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -50,9 +48,6 @@ import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.guid.JavaUtilGuidV4Generator; import com.linkedin.venice.kafka.protocol.GUID; -import com.linkedin.venice.listener.grpc.GrpcRequestContext; -import com.linkedin.venice.listener.grpc.handlers.GrpcStorageReadRequestHandler; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcHandler; import com.linkedin.venice.listener.request.AdminRequest; import com.linkedin.venice.listener.request.ComputeRouterRequestWrapper; import com.linkedin.venice.listener.request.GetRouterRequest; @@ -79,8 +74,6 @@ import com.linkedin.venice.metadata.response.VersionProperties; import com.linkedin.venice.offsets.OffsetRecord; import com.linkedin.venice.partitioner.VenicePartitioner; -import com.linkedin.venice.protocols.VeniceClientRequest; -import com.linkedin.venice.protocols.VeniceServerResponse; import com.linkedin.venice.read.RequestType; import com.linkedin.venice.read.protocol.request.router.MultiGetRouterRequestKeyV1; import com.linkedin.venice.read.protocol.response.MultiGetResponseRecordV1; @@ -458,7 +451,7 @@ public void merge(ReadResponseStatsRecorder other) { } ServerHttpRequestStats stats = mock(ServerHttpRequestStats.class); - multiGetResponseWrapper.getStatsRecorder().recordMetrics(stats); + multiGetResponseWrapper.getReadResponseStatsRecorder().recordMetrics(stats); if (responseProvider == validResponseProvider) { verify(stats).recordSuccessRequestKeyCount(recordCount); } @@ -647,7 +640,7 @@ public void testUnrecognizedRequestInStorageExecutionHandler() throws Exception verify(context, times(1)).writeAndFlush(argumentCaptor.capture()); HttpShortcutResponse shortcutResponse = (HttpShortcutResponse) argumentCaptor.getValue(); assertEquals(shortcutResponse.getStatus(), HttpResponseStatus.INTERNAL_SERVER_ERROR); - assertEquals(shortcutResponse.getMessage(), "Unrecognized object in StorageExecutionHandler"); + assertEquals(shortcutResponse.getMessage(), "Unrecognized request type"); } @Test(dataProvider = "True-and-False", dataProviderClass = DataProviderUtils.class) @@ -734,7 +727,7 @@ public void testHandleComputeRequest(boolean readComputationEnabled) throws Exce double expectedReadComputeEfficiency = (double) valueBytes.length / (double) expectedReadComputeOutputSize; ServerHttpRequestStats stats = mock(ServerHttpRequestStats.class); - computeResponse.getStatsRecorder().recordMetrics(stats); + computeResponse.getReadResponseStatsRecorder().recordMetrics(stats); verify(stats).recordSuccessRequestKeyCount(keySet.size()); verify(stats).recordDotProductCount(1); verify(stats).recordHadamardProduct(1); @@ -801,23 +794,6 @@ public void testMissingStorageEngine() throws Exception { assertEquals(responseObject.getValueRecord().getDataInBytes(), valueString.getBytes()); } - @Test - public void testGrpcReadReturnsInternalErrorWhenRouterRequestIsNull() { - VeniceClientRequest clientRequest = - VeniceClientRequest.newBuilder().setIsBatchRequest(true).setResourceName("testStore_v1").build(); - VeniceServerResponse.Builder builder = VeniceServerResponse.newBuilder(); - GrpcRequestContext ctx = new GrpcRequestContext(clientRequest, builder, null); - StorageReadRequestHandler requestHandler = createStorageReadRequestHandler(); - GrpcStorageReadRequestHandler grpcReadRequestHandler = spy(new GrpcStorageReadRequestHandler(requestHandler)); - VeniceServerGrpcHandler mockNextHandler = mock(VeniceServerGrpcHandler.class); - grpcReadRequestHandler.addNextHandler(mockNextHandler); - doNothing().when(mockNextHandler).processRequest(any()); - grpcReadRequestHandler.processRequest(ctx); // will cause np exception - - assertEquals(builder.getErrorCode(), VeniceReadResponseStatus.INTERNAL_ERROR); - assertTrue(builder.getErrorMessage().contains("Internal Error")); - } - @Test public void testMisRoutedStoreVersion() throws Exception { String storeName = "testStore"; @@ -835,7 +811,8 @@ public void testMisRoutedStoreVersion() throws Exception { ArgumentCaptor shortcutResponseArgumentCaptor = ArgumentCaptor.forClass(HttpShortcutResponse.class); verify(context).writeAndFlush(shortcutResponseArgumentCaptor.capture()); - Assert.assertTrue(shortcutResponseArgumentCaptor.getValue().isMisroutedStoreVersion()); + HttpResponseStatus response = shortcutResponseArgumentCaptor.getValue().getStatus(); + assertEquals(response, VeniceReadResponseStatus.MISROUTED_STORE_VERSION.getHttpResponseStatus()); } @Test diff --git a/services/venice-server/src/test/java/com/linkedin/venice/listener/response/BinaryResponseTest.java b/services/venice-server/src/test/java/com/linkedin/venice/listener/response/BinaryResponseTest.java new file mode 100644 index 0000000000..7983224282 --- /dev/null +++ b/services/venice-server/src/test/java/com/linkedin/venice/listener/response/BinaryResponseTest.java @@ -0,0 +1,36 @@ +package com.linkedin.venice.listener.response; + +import com.linkedin.venice.response.VeniceReadResponseStatus; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import java.nio.ByteBuffer; +import org.testng.Assert; +import org.testng.annotations.Test; + + +public class BinaryResponseTest { + @Test + public void testBinaryResponseBodyAndStatus() { + // Case 1: Null ByteBuf should default to EMPTY_BUFFER and status KEY_NOT_FOUND + BinaryResponse responseWithNullByteBuf = new BinaryResponse((ByteBuf) null); + Assert.assertEquals(responseWithNullByteBuf.getBody(), Unpooled.EMPTY_BUFFER); + Assert.assertEquals(responseWithNullByteBuf.getStatus(), VeniceReadResponseStatus.KEY_NOT_FOUND); + + // Case 2: Non-null ByteBuf should preserve body and set status to OK + ByteBuf nonNullByteBuf = Unpooled.wrappedBuffer(new byte[] { 1, 2, 3 }); + BinaryResponse responseWithNonNullByteBuf = new BinaryResponse(nonNullByteBuf); + Assert.assertEquals(responseWithNonNullByteBuf.getBody(), nonNullByteBuf); + Assert.assertEquals(responseWithNonNullByteBuf.getStatus(), VeniceReadResponseStatus.OK); + + // Case 3: Null ByteBuffer should default to EMPTY_BUFFER and status KEY_NOT_FOUND + BinaryResponse responseWithNullByteBuffer = new BinaryResponse((ByteBuffer) null); + Assert.assertEquals(responseWithNullByteBuffer.getBody(), Unpooled.EMPTY_BUFFER); + Assert.assertEquals(responseWithNullByteBuffer.getStatus(), VeniceReadResponseStatus.KEY_NOT_FOUND); + + // Case 4: Non-null ByteBuffer should wrap into ByteBuf and set status to OK + ByteBuffer nonNullByteBuffer = ByteBuffer.wrap(new byte[] { 4, 5, 6 }); + BinaryResponse responseWithNonNullByteBuffer = new BinaryResponse(nonNullByteBuffer); + Assert.assertEquals(responseWithNonNullByteBuffer.getBody(), Unpooled.wrappedBuffer(nonNullByteBuffer)); + Assert.assertEquals(responseWithNonNullByteBuffer.getStatus(), VeniceReadResponseStatus.OK); + } +}