From 917a88540107c86ec0df686e72bfb48db1fb6dcc Mon Sep 17 00:00:00 2001 From: Mayur Kale Date: Wed, 30 Jul 2025 17:29:05 +0000 Subject: [PATCH 1/3] chore: refactor code to reduce object creations --- .../adapter/DriverConnectionHandler.java | 128 ++++++++---------- .../spanner/adapter/PreparePayloadResult.java | 11 +- 2 files changed, 58 insertions(+), 81 deletions(-) diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java index 93e92498..c438017c 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java @@ -21,6 +21,7 @@ import static com.google.cloud.spanner.adapter.util.StringUtils.startsWith; import com.datastax.oss.driver.internal.core.protocol.ByteBufPrimitiveCodec; +import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting; import com.datastax.oss.protocol.internal.Compressor; import com.datastax.oss.protocol.internal.Frame; import com.datastax.oss.protocol.internal.FrameCodec; @@ -65,12 +66,15 @@ final class DriverConnectionHandler implements Runnable { private final Socket socket; private final AdapterClientWrapper adapterClientWrapper; private final Optional maxCommitDelayMillis; - private final GrpcCallContext defaultContext; - private final GrpcCallContext defaultContextWithLAR; - private static final Map> ROUTE_TO_LEADER_HEADER_MAP = - ImmutableMap.of(ROUTE_TO_LEADER_HEADER_KEY, Collections.singletonList("true")); private static final int defaultStreamId = -1; + // These contexts are thread-safe and can be reused across all instances. + private static final GrpcCallContext DEFAULT_CONTEXT = GrpcCallContext.createDefault(); + private static final Map> ROUTE_TO_LEADER_HEADER_MAP = + ImmutableMap.of(ROUTE_TO_LEADER_HEADER_KEY, Collections.singletonList("true")); + private static final GrpcCallContext DEFAULT_CONTEXT_WITH_LAR = + GrpcCallContext.createDefault().withExtraHeaders(ROUTE_TO_LEADER_HEADER_MAP); + /** * Constructor for DriverConnectionHandler. * @@ -82,9 +86,6 @@ public DriverConnectionHandler( Socket socket, AdapterClientWrapper adapterClientWrapper, Optional maxCommitDelay) { this.socket = socket; this.adapterClientWrapper = adapterClientWrapper; - this.defaultContext = GrpcCallContext.createDefault(); - this.defaultContextWithLAR = - GrpcCallContext.createDefault().withExtraHeaders(ROUTE_TO_LEADER_HEADER_MAP); if (maxCommitDelay.isPresent()) { this.maxCommitDelayMillis = Optional.of(String.valueOf(maxCommitDelay.get().toMillis())); } else { @@ -92,6 +93,7 @@ public DriverConnectionHandler( } } + @VisibleForTesting public DriverConnectionHandler(Socket socket, AdapterClientWrapper adapterClientWrapper) { this(socket, adapterClientWrapper, Optional.empty()); } @@ -123,7 +125,6 @@ private void processRequestsLoop(InputStream inputStream, OutputStream outputStr throws IOException { // Keep processing until End-Of-Stream is reached on the input while (true) { - byte[] responseToWrite; int streamId = defaultStreamId; // Initialize with a default value. try { // 1. Read and construct the payload from the input stream @@ -137,27 +138,24 @@ private void processRequestsLoop(InputStream inputStream, OutputStream outputStr // 3. Prepare the payload. PreparePayloadResult prepareResult = preparePayload(payload); streamId = prepareResult.getStreamId(); - Optional response = prepareResult.getAttachmentErrorResponse(); // 4. If attachment preparation didn't yield an immediate response, send the gRPC request. - if (!response.isPresent()) { - responseToWrite = - adapterClientWrapper.sendGrpcRequest( - payload, prepareResult.getAttachments(), prepareResult.getContext(), streamId); + if (prepareResult.getAttachmentErrorResponse() != null) { + outputStream.write(prepareResult.getAttachmentErrorResponse()); + } else { + outputStream.write(adapterClientWrapper.sendGrpcRequest( + payload, prepareResult.getAttachments(), prepareResult.getContext(), streamId)); // Now response holds the gRPC result, which might still be empty. - } else { - responseToWrite = response.get(); } } catch (RuntimeException e) { // 5. Handle any error during payload construction or attachment processing. // Create a server error response to send back to the client. LOG.error("Error processing request: ", e); - responseToWrite = + outputStream.write( serverErrorResponse( - streamId, "Server error during request processing: " + e.getMessage()); + streamId, "Server error during request processing: " + e.getMessage())); } - outputStream.write(responseToWrite); outputStream.flush(); } } @@ -174,23 +172,12 @@ private static int readNBytesJava8(InputStream in, byte[] b, int off, int len) } int totalBytesRead = 0; - int bytesReadInCurrentLoop; - - // Loop until the desired number of bytes are read or EOF is reached while (totalBytesRead < len) { - // Calculate how many bytes are still needed - int remaining = len - totalBytesRead; - // Calculate the current offset in the buffer - int currentOffset = off + totalBytesRead; - - // Attempt to read the remaining bytes - bytesReadInCurrentLoop = in.read(b, currentOffset, remaining); - + int bytesReadInCurrentLoop = in.read(b, off + totalBytesRead, len - totalBytesRead); if (bytesReadInCurrentLoop == -1) { - // End Of Stream (EOF) reached before 'len' bytes were read. + // EOF reached before 'len' bytes were read. break; } - totalBytesRead += bytesReadInCurrentLoop; } @@ -202,28 +189,23 @@ private byte[] constructPayload(InputStream socketInputStream) byte[] header = new byte[HEADER_LENGTH]; int bytesRead = readNBytesJava8(socketInputStream, header, 0, HEADER_LENGTH); if (bytesRead == 0) { - // EOF - return new byte[0]; + return new byte[0]; // EOF } else if (bytesRead < HEADER_LENGTH) { throw new IllegalArgumentException("Payload is not well formed."); } // Extract the body length from the header. int bodyLength = load32BigEndian(header, 5); - if (bodyLength < 0) { throw new IllegalArgumentException("Payload is not well formed."); } - byte[] body = new byte[bodyLength]; - if (readNBytesJava8(socketInputStream, body, 0, bodyLength) < bodyLength) { - throw new IllegalArgumentException("Payload is not well formed."); - } - - // Combine the header and body into the payload. byte[] payload = new byte[HEADER_LENGTH + bodyLength]; System.arraycopy(header, 0, payload, 0, HEADER_LENGTH); - System.arraycopy(body, 0, payload, HEADER_LENGTH, bodyLength); + if (bodyLength > 0 + && readNBytesJava8(socketInputStream, payload, HEADER_LENGTH, bodyLength) < bodyLength) { + throw new IllegalArgumentException("Payload is not well formed."); + } return payload; } @@ -250,79 +232,75 @@ private PreparePayloadResult preparePayload(byte[] payload) { Frame frame = serverFrameCodec.decode(payloadBuf); payloadBuf.release(); - Map attachments = new HashMap<>(); if (frame.message instanceof Execute) { - return prepareExecuteMessage((Execute) frame.message, frame.streamId, attachments); + return prepareExecuteMessage((Execute) frame.message, frame.streamId); } else if (frame.message instanceof Batch) { - return prepareBatchMessage((Batch) frame.message, frame.streamId, attachments); + return prepareBatchMessage((Batch) frame.message, frame.streamId); } else if (frame.message instanceof Query) { - return prepareQueryMessage((Query) frame.message, frame.streamId, attachments); + return prepareQueryMessage((Query) frame.message, frame.streamId); } else { - return new PreparePayloadResult(defaultContext, frame.streamId); + return new PreparePayloadResult(DEFAULT_CONTEXT, frame.streamId); } } - private PreparePayloadResult prepareExecuteMessage( - Execute message, int streamId, Map attachments) { + private PreparePayloadResult prepareExecuteMessage(Execute message, int streamId) { ApiCallContext context; + Map attachments = new HashMap<>(); if (message.queryId != null && message.queryId.length > 0 && message.queryId[0] == WRITE_ACTION_QUERY_ID_PREFIX) { - context = defaultContextWithLAR; - if (maxCommitDelayMillis.isPresent()) { - attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, maxCommitDelayMillis.get()); - } + context = DEFAULT_CONTEXT_WITH_LAR; + maxCommitDelayMillis.ifPresent( + delay -> attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, delay)); } else { - context = defaultContext; + context = DEFAULT_CONTEXT; } - Optional errorResponse = - prepareAttachmentForQueryId(streamId, attachments, message.queryId); + byte[] errorResponse = prepareAttachmentForQueryId(streamId, attachments, message.queryId); return new PreparePayloadResult(context, streamId, attachments, errorResponse); } - private PreparePayloadResult prepareBatchMessage( - Batch message, int streamId, Map attachments) { - Optional attachmentErrorResponse = Optional.empty(); + private PreparePayloadResult prepareBatchMessage(Batch message, int streamId) { + Map attachments = new HashMap<>(); for (Object obj : message.queriesOrIds) { if (obj instanceof byte[]) { - Optional errorResponse = - prepareAttachmentForQueryId(streamId, attachments, (byte[]) obj); - if (errorResponse.isPresent()) { - attachmentErrorResponse = errorResponse; - break; + byte[] errorResponse = prepareAttachmentForQueryId(streamId, attachments, (byte[]) obj); + if (errorResponse != null) { + // Return immediately on the first unprepared statement. + return new PreparePayloadResult( + DEFAULT_CONTEXT_WITH_LAR, streamId, attachments, errorResponse); } } } - if (maxCommitDelayMillis.isPresent()) { - attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, maxCommitDelayMillis.get()); - } - return new PreparePayloadResult( - defaultContextWithLAR, streamId, attachments, attachmentErrorResponse); + maxCommitDelayMillis.ifPresent( + delay -> attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, delay)); + // No error, return with populated attachments. + return new PreparePayloadResult(DEFAULT_CONTEXT_WITH_LAR, streamId, attachments, null); } - private PreparePayloadResult prepareQueryMessage( - Query message, int streamId, Map attachments) { + private PreparePayloadResult prepareQueryMessage(Query message, int streamId) { ApiCallContext context; + Map attachments = Collections.emptyMap(); if (startsWith(message.query, "SELECT")) { - context = defaultContext; + context = DEFAULT_CONTEXT; } else { - context = defaultContextWithLAR; + context = DEFAULT_CONTEXT_WITH_LAR; if (maxCommitDelayMillis.isPresent()) { + attachments = new HashMap<>(); attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, maxCommitDelayMillis.get()); } } return new PreparePayloadResult(context, streamId, attachments); } - private Optional prepareAttachmentForQueryId( + private byte[] prepareAttachmentForQueryId( int streamId, Map attachments, byte[] queryId) { String key = constructKey(queryId); Optional val = adapterClientWrapper.getAttachmentsCache().get(key); if (!val.isPresent()) { - return Optional.of(unpreparedResponse(streamId, queryId)); + return unpreparedResponse(streamId, queryId); } attachments.put(key, val.get()); - return Optional.empty(); + return null; } private static String constructKey(byte[] queryId) { diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java index 9c38ce3f..78e19c26 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java @@ -19,7 +19,6 @@ import com.google.api.gax.rpc.ApiCallContext; import java.util.Collections; import java.util.Map; -import java.util.Optional; /** * An object used to encapsulate the result of preparing the Adapter payload prior to sending the @@ -29,14 +28,14 @@ public class PreparePayloadResult { private ApiCallContext context; private int streamId; private Map attachments; - private Optional attachmentErrorResponse; + private byte[] attachmentErrorResponse; private static final Map EMPTY_ATTACHMENTS = Collections.emptyMap(); public PreparePayloadResult( ApiCallContext context, int streamId, Map attachments, - Optional attachmentErrorResponse) { + byte[] attachmentErrorResponse) { this.context = context; this.streamId = streamId; this.attachments = attachments; @@ -45,18 +44,18 @@ public PreparePayloadResult( public PreparePayloadResult( ApiCallContext context, int streamId, Map attachments) { - this(context, streamId, attachments, Optional.empty()); + this(context, streamId, attachments, null); } public PreparePayloadResult(ApiCallContext context, int streamId) { - this(context, streamId, EMPTY_ATTACHMENTS, Optional.empty()); + this(context, streamId, EMPTY_ATTACHMENTS, null); } public Map getAttachments() { return attachments; } - public Optional getAttachmentErrorResponse() { + public byte[] getAttachmentErrorResponse() { return attachmentErrorResponse; } From d3d4cd8b14380e321e4ecc58e47aef2f9dc11906 Mon Sep 17 00:00:00 2001 From: Mayur Kale Date: Wed, 30 Jul 2025 17:33:40 +0000 Subject: [PATCH 2/3] chore: format code --- .../cloud/spanner/adapter/DriverConnectionHandler.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java index c438017c..bbd72df8 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java @@ -70,7 +70,7 @@ final class DriverConnectionHandler implements Runnable { // These contexts are thread-safe and can be reused across all instances. private static final GrpcCallContext DEFAULT_CONTEXT = GrpcCallContext.createDefault(); - private static final Map> ROUTE_TO_LEADER_HEADER_MAP = + private static final Map> ROUTE_TO_LEADER_HEADER_MAP = ImmutableMap.of(ROUTE_TO_LEADER_HEADER_KEY, Collections.singletonList("true")); private static final GrpcCallContext DEFAULT_CONTEXT_WITH_LAR = GrpcCallContext.createDefault().withExtraHeaders(ROUTE_TO_LEADER_HEADER_MAP); @@ -142,8 +142,9 @@ private void processRequestsLoop(InputStream inputStream, OutputStream outputStr // 4. If attachment preparation didn't yield an immediate response, send the gRPC request. if (prepareResult.getAttachmentErrorResponse() != null) { outputStream.write(prepareResult.getAttachmentErrorResponse()); - } else { - outputStream.write(adapterClientWrapper.sendGrpcRequest( + } else { + outputStream.write( + adapterClientWrapper.sendGrpcRequest( payload, prepareResult.getAttachments(), prepareResult.getContext(), streamId)); // Now response holds the gRPC result, which might still be empty. } From 9d2c146517d2ae552881a2943663269a0ba9101c Mon Sep 17 00:00:00 2001 From: Mayur Kale Date: Wed, 30 Jul 2025 19:05:55 +0000 Subject: [PATCH 3/3] chore: fix code reviews --- .../adapter/DriverConnectionHandler.java | 27 ++++++++++--------- .../spanner/adapter/PreparePayloadResult.java | 11 ++++---- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java index bbd72df8..0aca3f9b 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/DriverConnectionHandler.java @@ -140,8 +140,8 @@ private void processRequestsLoop(InputStream inputStream, OutputStream outputStr streamId = prepareResult.getStreamId(); // 4. If attachment preparation didn't yield an immediate response, send the gRPC request. - if (prepareResult.getAttachmentErrorResponse() != null) { - outputStream.write(prepareResult.getAttachmentErrorResponse()); + if (prepareResult.getAttachmentErrorResponse().isPresent()) { + outputStream.write(prepareResult.getAttachmentErrorResponse().get()); } else { outputStream.write( adapterClientWrapper.sendGrpcRequest( @@ -256,26 +256,29 @@ private PreparePayloadResult prepareExecuteMessage(Execute message, int streamId } else { context = DEFAULT_CONTEXT; } - byte[] errorResponse = prepareAttachmentForQueryId(streamId, attachments, message.queryId); + Optional errorResponse = + prepareAttachmentForQueryId(streamId, attachments, message.queryId); return new PreparePayloadResult(context, streamId, attachments, errorResponse); } private PreparePayloadResult prepareBatchMessage(Batch message, int streamId) { + Optional attachmentErrorResponse = Optional.empty(); Map attachments = new HashMap<>(); for (Object obj : message.queriesOrIds) { if (obj instanceof byte[]) { - byte[] errorResponse = prepareAttachmentForQueryId(streamId, attachments, (byte[]) obj); - if (errorResponse != null) { - // Return immediately on the first unprepared statement. - return new PreparePayloadResult( - DEFAULT_CONTEXT_WITH_LAR, streamId, attachments, errorResponse); + Optional errorResponse = + prepareAttachmentForQueryId(streamId, attachments, (byte[]) obj); + if (errorResponse.isPresent()) { + attachmentErrorResponse = errorResponse; + break; } } } maxCommitDelayMillis.ifPresent( delay -> attachments.put(MAX_COMMIT_DELAY_ATTACHMENT_KEY, delay)); // No error, return with populated attachments. - return new PreparePayloadResult(DEFAULT_CONTEXT_WITH_LAR, streamId, attachments, null); + return new PreparePayloadResult( + DEFAULT_CONTEXT_WITH_LAR, streamId, attachments, attachmentErrorResponse); } private PreparePayloadResult prepareQueryMessage(Query message, int streamId) { @@ -293,15 +296,15 @@ private PreparePayloadResult prepareQueryMessage(Query message, int streamId) { return new PreparePayloadResult(context, streamId, attachments); } - private byte[] prepareAttachmentForQueryId( + private Optional prepareAttachmentForQueryId( int streamId, Map attachments, byte[] queryId) { String key = constructKey(queryId); Optional val = adapterClientWrapper.getAttachmentsCache().get(key); if (!val.isPresent()) { - return unpreparedResponse(streamId, queryId); + return Optional.of(unpreparedResponse(streamId, queryId)); } attachments.put(key, val.get()); - return null; + return Optional.empty(); } private static String constructKey(byte[] queryId) { diff --git a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java index 78e19c26..9c38ce3f 100644 --- a/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java +++ b/google-cloud-spanner-cassandra/src/main/java/com/google/cloud/spanner/adapter/PreparePayloadResult.java @@ -19,6 +19,7 @@ import com.google.api.gax.rpc.ApiCallContext; import java.util.Collections; import java.util.Map; +import java.util.Optional; /** * An object used to encapsulate the result of preparing the Adapter payload prior to sending the @@ -28,14 +29,14 @@ public class PreparePayloadResult { private ApiCallContext context; private int streamId; private Map attachments; - private byte[] attachmentErrorResponse; + private Optional attachmentErrorResponse; private static final Map EMPTY_ATTACHMENTS = Collections.emptyMap(); public PreparePayloadResult( ApiCallContext context, int streamId, Map attachments, - byte[] attachmentErrorResponse) { + Optional attachmentErrorResponse) { this.context = context; this.streamId = streamId; this.attachments = attachments; @@ -44,18 +45,18 @@ public PreparePayloadResult( public PreparePayloadResult( ApiCallContext context, int streamId, Map attachments) { - this(context, streamId, attachments, null); + this(context, streamId, attachments, Optional.empty()); } public PreparePayloadResult(ApiCallContext context, int streamId) { - this(context, streamId, EMPTY_ATTACHMENTS, null); + this(context, streamId, EMPTY_ATTACHMENTS, Optional.empty()); } public Map getAttachments() { return attachments; } - public byte[] getAttachmentErrorResponse() { + public Optional getAttachmentErrorResponse() { return attachmentErrorResponse; }