diff --git a/src/main/java/com/scylladb/alternator/GzipRequestInterceptor.java b/src/main/java/com/scylladb/alternator/GzipRequestInterceptor.java index 749b709..82fe7f5 100644 --- a/src/main/java/com/scylladb/alternator/GzipRequestInterceptor.java +++ b/src/main/java/com/scylladb/alternator/GzipRequestInterceptor.java @@ -31,6 +31,8 @@ public class GzipRequestInterceptor implements ExecutionInterceptor { private static final ExecutionAttribute ORIGINAL_BODY_BYTES = new ExecutionAttribute<>("GzipRequestInterceptor.originalBodyBytes"); + private static final ExecutionAttribute COMPRESSED_BODY_BYTES = + new ExecutionAttribute<>("GzipRequestInterceptor.compressedBodyBytes"); private static final ExecutionAttribute SHOULD_COMPRESS = new ExecutionAttribute<>("GzipRequestInterceptor.shouldCompress"); @@ -49,38 +51,24 @@ public GzipRequestInterceptor(int minCompressionSizeBytes) { public SdkHttpRequest modifyHttpRequest( Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { - byte[] originalContent; - try { - originalContent = readOriginalBody(context); - } catch (IOException | CompletionException e) { - executionAttributes.putAttribute(SHOULD_COMPRESS, false); - return context.httpRequest(); - } - - if (originalContent == null) { - executionAttributes.putAttribute(SHOULD_COMPRESS, false); + prepareBody(context, executionAttributes); + Boolean shouldCompress = executionAttributes.getAttribute(SHOULD_COMPRESS); + byte[] compressedContent = executionAttributes.getAttribute(COMPRESSED_BODY_BYTES); + if (shouldCompress == null || !shouldCompress || compressedContent == null) { return context.httpRequest(); } - // Cache the original content for modifyHttpContent / modifyAsyncHttpContent. - executionAttributes.putAttribute(ORIGINAL_BODY_BYTES, originalContent); - - // Check if we should compress based on size. - boolean shouldCompress = originalContent.length >= minCompressionSizeBytes; - executionAttributes.putAttribute(SHOULD_COMPRESS, shouldCompress); - - if (shouldCompress) { - // Add Content-Encoding header. - return context.httpRequest().toBuilder().putHeader("Content-Encoding", "gzip").build(); - } - - return context.httpRequest(); + return context.httpRequest().toBuilder() + .putHeader("Content-Encoding", "gzip") + .putHeader("Content-Length", String.valueOf(compressedContent.length)) + .build(); } @Override public Optional modifyHttpContent( Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { + prepareBody(context, executionAttributes); Boolean shouldCompress = executionAttributes.getAttribute(SHOULD_COMPRESS); if (shouldCompress == null || !shouldCompress) { // Return original body from cached bytes if available, otherwise return as-is @@ -91,26 +79,51 @@ public Optional modifyHttpContent( return context.requestBody(); } - byte[] originalContent = executionAttributes.getAttribute(ORIGINAL_BODY_BYTES); - if (originalContent == null) { + byte[] compressedContent = executionAttributes.getAttribute(COMPRESSED_BODY_BYTES); + if (compressedContent == null) { return context.requestBody(); } + return Optional.of(RequestBody.fromBytes(compressedContent)); + } + + private void prepareBody( + Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { + if (executionAttributes.getAttribute(SHOULD_COMPRESS) != null) { + return; + } + + byte[] originalContent; try { - // Compress the content - byte[] compressedContent = gzipCompress(originalContent); - return Optional.of(RequestBody.fromBytes(compressedContent)); + originalContent = readOriginalBody(context); + } catch (IOException | CompletionException e) { + executionAttributes.putAttribute(SHOULD_COMPRESS, false); + return; + } - } catch (IOException e) { - // If compression fails, return original content - return Optional.of(RequestBody.fromBytes(originalContent)); + if (originalContent == null) { + executionAttributes.putAttribute(SHOULD_COMPRESS, false); + return; } + + executionAttributes.putAttribute(ORIGINAL_BODY_BYTES, originalContent); + boolean shouldCompress = originalContent.length >= minCompressionSizeBytes; + if (shouldCompress) { + try { + executionAttributes.putAttribute(COMPRESSED_BODY_BYTES, gzipCompress(originalContent)); + } catch (IOException e) { + executionAttributes.putAttribute(SHOULD_COMPRESS, false); + return; + } + } + executionAttributes.putAttribute(SHOULD_COMPRESS, shouldCompress); } @Override public Optional modifyAsyncHttpContent( Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { + prepareBody(context, executionAttributes); Boolean shouldCompress = executionAttributes.getAttribute(SHOULD_COMPRESS); if (shouldCompress == null || !shouldCompress) { byte[] cachedBytes = executionAttributes.getAttribute(ORIGINAL_BODY_BYTES); @@ -120,17 +133,12 @@ public Optional modifyAsyncHttpContent( return context.asyncRequestBody(); } - byte[] originalContent = executionAttributes.getAttribute(ORIGINAL_BODY_BYTES); - if (originalContent == null) { + byte[] compressedContent = executionAttributes.getAttribute(COMPRESSED_BODY_BYTES); + if (compressedContent == null) { return context.asyncRequestBody(); } - try { - byte[] compressedContent = gzipCompress(originalContent); - return Optional.of(AsyncRequestBody.fromBytes(compressedContent)); - } catch (IOException e) { - return Optional.of(AsyncRequestBody.fromBytes(originalContent)); - } + return Optional.of(AsyncRequestBody.fromBytes(compressedContent)); } private byte[] readOriginalBody(Context.ModifyHttpRequest context) throws IOException { diff --git a/src/test/java/com/scylladb/alternator/GzipRequestInterceptorTest.java b/src/test/java/com/scylladb/alternator/GzipRequestInterceptorTest.java index bace9d0..c8e2d57 100644 --- a/src/test/java/com/scylladb/alternator/GzipRequestInterceptorTest.java +++ b/src/test/java/com/scylladb/alternator/GzipRequestInterceptorTest.java @@ -2,24 +2,39 @@ import static org.junit.Assert.*; +import com.scylladb.alternator.vectorsearch.VectorSearchInterceptor; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.net.URI; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Random; import java.util.zip.GZIPInputStream; import org.junit.Test; +import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.core.SdkRequest; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.interceptor.Context; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.sync.RequestBody; +import software.amazon.awssdk.http.AbortableInputStream; +import software.amazon.awssdk.http.ExecutableHttpRequest; +import software.amazon.awssdk.http.HttpExecuteRequest; +import software.amazon.awssdk.http.HttpExecuteResponse; +import software.amazon.awssdk.http.SdkHttpClient; +import software.amazon.awssdk.http.SdkHttpFullResponse; import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.http.SdkHttpRequest; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; import software.amazon.awssdk.services.dynamodb.model.ListTablesRequest; +import software.amazon.awssdk.services.dynamodb.model.PutItemRequest; /** * Unit tests for GzipRequestInterceptor. @@ -111,14 +126,8 @@ private byte[] generateTestData(int size) { private byte[] readRequestBody(Optional body) throws IOException { assertTrue("Request body should be present", body.isPresent()); - ByteArrayOutputStream bos = new ByteArrayOutputStream(); - byte[] buf = new byte[4096]; - int len; java.io.InputStream is = body.get().contentStreamProvider().newStream(); - while ((len = is.read(buf)) != -1) { - bos.write(buf, 0, len); - } - return bos.toByteArray(); + return readAllBytes(is); } private byte[] readAsyncRequestBody(Optional body) { @@ -136,6 +145,56 @@ private byte[] readAsyncRequestBody(Optional body) { return bos.toByteArray(); } + private byte[] readAllBytes(java.io.InputStream is) throws IOException { + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + byte[] buf = new byte[4096]; + int len; + while ((len = is.read(buf)) != -1) { + bos.write(buf, 0, len); + } + return bos.toByteArray(); + } + + @Test + public void testSyncDynamoDbClientSendsCompressedRequestBody() throws Exception { + RecordingHttpClient httpClient = new RecordingHttpClient(); + DynamoDbClient client = + DynamoDbClient.builder() + .endpointOverride(URI.create("http://localhost:8000")) + .region(Region.US_EAST_1) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .httpClient(httpClient) + .overrideConfiguration( + c -> + c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE) + .addExecutionInterceptor(new ResponseCompressionInterceptor()) + .addExecutionInterceptor(new GzipRequestInterceptor(100))) + .build(); + + try { + StringBuilder largeValue = new StringBuilder(); + for (int i = 0; i < 100; i++) { + largeValue.append("This is a test value that should be compressed. "); + } + client.putItem( + PutItemRequest.builder() + .tableName("items") + .item( + Map.of( + "ID", AttributeValue.builder().s("compression-test").build(), + "LargeData", AttributeValue.builder().s(largeValue.toString()).build())) + .build()); + } finally { + client.close(); + } + + assertEquals("gzip", httpClient.capturedRequest.firstMatchingHeader("Content-Encoding").get()); + String requestJson = + new String(gzipDecompress(httpClient.capturedBody), StandardCharsets.UTF_8); + assertTrue(requestJson.contains("\"TableName\":\"items\"")); + assertTrue(requestJson.contains("compression-test")); + } + @Test public void testContentEncodingHeaderSetOnCompressedRequest() { GzipRequestInterceptor interceptor = new GzipRequestInterceptor(DEFAULT_MIN_COMPRESSION_SIZE); @@ -373,4 +432,44 @@ public void testDecompressedBodyMatchesOriginalExactly() throws IOException { assertEquals("Byte mismatch at index " + i, testData[i], decompressed[i]); } } + + private final class RecordingHttpClient implements SdkHttpClient { + private SdkHttpRequest capturedRequest; + private byte[] capturedBody; + + @Override + public ExecutableHttpRequest prepareRequest(HttpExecuteRequest request) { + capturedRequest = request.httpRequest(); + return new ExecutableHttpRequest() { + @Override + public HttpExecuteResponse call() throws IOException { + capturedBody = + request.contentStreamProvider().isPresent() + ? readAllBytes(request.contentStreamProvider().get().newStream()) + : new byte[0]; + byte[] body = "{}".getBytes(StandardCharsets.UTF_8); + return HttpExecuteResponse.builder() + .response( + SdkHttpFullResponse.builder() + .statusCode(200) + .putHeader("Content-Type", "application/x-amz-json-1.0") + .putHeader("Content-Length", String.valueOf(body.length)) + .build()) + .responseBody(AbortableInputStream.create(new ByteArrayInputStream(body))) + .build(); + } + + @Override + public void abort() {} + }; + } + + @Override + public void close() {} + + @Override + public String clientName() { + return "recording"; + } + } } diff --git a/src/test/java/com/scylladb/alternator/ResponseCompressionInterceptorTest.java b/src/test/java/com/scylladb/alternator/ResponseCompressionInterceptorTest.java index 3ae31c7..acaabc0 100644 --- a/src/test/java/com/scylladb/alternator/ResponseCompressionInterceptorTest.java +++ b/src/test/java/com/scylladb/alternator/ResponseCompressionInterceptorTest.java @@ -2,6 +2,7 @@ import static org.junit.Assert.*; +import com.scylladb.alternator.vectorsearch.VectorSearchInterceptor; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -24,6 +25,8 @@ import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.interceptor.Context; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptorChain; +import software.amazon.awssdk.core.interceptor.InterceptorContext; import software.amazon.awssdk.core.sync.RequestBody; import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.http.SdkHttpRequest; @@ -203,6 +206,71 @@ public void testDeflateAsyncResponseIsDecompressed() throws Exception { assertArrayEquals(original, collect(modifiedPublisher.get())); } + @Test + public void testResponseCompressionRunsBeforeVectorSearchForSyncResponses() throws Exception { + byte[] original = + "{\"Item\":{\"embedding\":{\"FLOAT32VECTOR\":[1.0,2.0]}}}".getBytes(StandardCharsets.UTF_8); + byte[] compressed = gzip(original); + SdkHttpResponse response = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Encoding", "gzip") + .putHeader("Content-Length", Integer.toString(compressed.length)) + .build(); + ExecutionInterceptorChain chain = + new ExecutionInterceptorChain( + Arrays.asList(VectorSearchInterceptor.INSTANCE, new ResponseCompressionInterceptor())); + ExecutionAttributes attrs = new ExecutionAttributes(); + InterceptorContext context = + InterceptorContext.builder() + .request(ListTablesRequest.builder().build()) + .httpRequest(createHttpRequest()) + .httpResponse(response) + .responseBody(new ByteArrayInputStream(compressed)) + .build(); + + InterceptorContext result = chain.modifyHttpResponse(context, attrs); + + assertFalse(result.httpResponse().firstMatchingHeader("Content-Encoding").isPresent()); + assertFalse(result.httpResponse().firstMatchingHeader("Content-Length").isPresent()); + String body = new String(readAll(result.responseBody().get()), StandardCharsets.UTF_8); + assertTrue(body.contains("\"L\"")); + assertFalse(body.contains("FLOAT32VECTOR")); + } + + @Test + public void testResponseCompressionRunsBeforeVectorSearchForAsyncResponses() throws Exception { + byte[] original = + "{\"Item\":{\"embedding\":{\"FLOAT32VECTOR\":[1.0,2.0]}}}".getBytes(StandardCharsets.UTF_8); + byte[] compressed = gzip(original); + SdkHttpResponse response = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Encoding", "gzip") + .putHeader("Content-Length", Integer.toString(compressed.length)) + .build(); + ExecutionInterceptorChain chain = + new ExecutionInterceptorChain( + Arrays.asList(VectorSearchInterceptor.INSTANCE, new ResponseCompressionInterceptor())); + ExecutionAttributes attrs = new ExecutionAttributes(); + InterceptorContext context = + InterceptorContext.builder() + .request(ListTablesRequest.builder().build()) + .httpRequest(createHttpRequest()) + .httpResponse(response) + .responsePublisher(singleBufferPublisher(compressed)) + .build(); + + InterceptorContext headerResult = chain.modifyHttpResponse(context, attrs); + InterceptorContext bodyResult = chain.modifyAsyncHttpResponse(headerResult, attrs); + + assertFalse(headerResult.httpResponse().firstMatchingHeader("Content-Encoding").isPresent()); + assertFalse(headerResult.httpResponse().firstMatchingHeader("Content-Length").isPresent()); + String body = new String(collect(bodyResult.responsePublisher().get()), StandardCharsets.UTF_8); + assertTrue(body.contains("\"L\"")); + assertFalse(body.contains("FLOAT32VECTOR")); + } + private static SdkHttpRequest createHttpRequest() { return SdkHttpRequest.builder() .protocol("http") diff --git a/src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java b/src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java index c1e7638..bd166e7 100644 --- a/src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java +++ b/src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java @@ -8,6 +8,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.scylladb.alternator.GzipRequestInterceptor; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -20,6 +21,7 @@ import java.util.concurrent.TimeUnit; import java.util.zip.CRC32; import java.util.zip.DeflaterOutputStream; +import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; import org.junit.Test; import org.reactivestreams.Publisher; @@ -28,7 +30,10 @@ import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.core.interceptor.Context; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptorChain; import software.amazon.awssdk.core.interceptor.InterceptorContext; +import software.amazon.awssdk.core.sync.RequestBody; +import software.amazon.awssdk.http.ContentStreamProvider; import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.http.SdkHttpRequest; import software.amazon.awssdk.http.SdkHttpResponse; @@ -37,6 +42,7 @@ import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.ListTablesRequest; import software.amazon.awssdk.services.dynamodb.model.QueryRequest; public class VectorSearchHttpInterceptorTest { @@ -470,6 +476,65 @@ public void testDynamoDbAsyncClientReadsUncompressedVectorResponseWithCrc32Heade } } + @Test + public void testVectorSearchBeforeGzipCompressesModifiedRequestBody() throws Exception { + ExecutionInterceptorChain chain = + new ExecutionInterceptorChain( + Arrays.asList(VectorSearchInterceptor.INSTANCE, new GzipRequestInterceptor(1))); + ExecutionAttributes attrs = new ExecutionAttributes(); + attrs.putAttribute( + VectorSearchInterceptor.VECTOR_SEARCH, + VectorSearch.builder().queryVector(1.0f, 2.0f).returnScores(true).build()); + + InterceptorContext context = + InterceptorContext.builder() + .request(ListTablesRequest.builder().build()) + .httpRequest(queryHttpRequest()) + .requestBody(RequestBody.fromString("{\"TableName\":\"items\"}")) + .build(); + + InterceptorContext result = chain.modifyHttpRequestAndHttpContent(context, attrs); + + assertEquals("gzip", result.httpRequest().firstMatchingHeader("Content-Encoding").get()); + byte[] compressed = readRequestBody(result.requestBody()); + byte[] uncompressed = gzipDecompress(compressed); + String json = new String(uncompressed, StandardCharsets.UTF_8); + assertTrue(json.contains("\"VectorSearch\"")); + assertTrue(json.contains("\"FLOAT32VECTOR\"")); + assertEquals( + String.valueOf(compressed.length), + result.httpRequest().firstMatchingHeader("Content-Length").get()); + } + + @Test + public void testVectorSearchBeforeGzipReplaysUnchangedSingleUseRequestBody() throws Exception { + ExecutionInterceptorChain chain = + new ExecutionInterceptorChain( + Arrays.asList(VectorSearchInterceptor.INSTANCE, new GzipRequestInterceptor(1))); + ExecutionAttributes attrs = new ExecutionAttributes(); + byte[] original = bytes("{\"TableName\":\"items\",\"Item\":{\"id\":{\"S\":\"item-1\"}}}"); + ByteArrayInputStream singleUseStream = new ByteArrayInputStream(original); + ContentStreamProvider singleUseProvider = () -> singleUseStream; + + InterceptorContext context = + InterceptorContext.builder() + .request(ListTablesRequest.builder().build()) + .httpRequest(queryHttpRequest()) + .requestBody( + RequestBody.fromContentProvider( + singleUseProvider, original.length, "application/x-amz-json-1.0")) + .build(); + + InterceptorContext result = chain.modifyHttpRequestAndHttpContent(context, attrs); + + assertEquals("gzip", result.httpRequest().firstMatchingHeader("Content-Encoding").get()); + byte[] compressed = readRequestBody(result.requestBody()); + assertArrayEquals(original, gzipDecompress(compressed)); + assertEquals( + String.valueOf(compressed.length), + result.httpRequest().firstMatchingHeader("Content-Length").get()); + } + @Test(expected = IllegalStateException.class) public void testDeleteVectorIndexActionRequiresIndexName() { DeleteVectorIndexAction.builder().build(); @@ -558,6 +623,11 @@ private static Context.ModifyHttpResponse responseContext( .build(); } + private static byte[] readRequestBody(Optional body) throws IOException { + assertTrue(body.isPresent()); + return readAllBytes(body.get().contentStreamProvider().newStream()); + } + private static byte[] gzipCompress(byte[] uncompressed) throws IOException { ByteArrayOutputStream out = new ByteArrayOutputStream(); try (GZIPOutputStream gzip = new GZIPOutputStream(out)) { @@ -566,6 +636,10 @@ private static byte[] gzipCompress(byte[] uncompressed) throws IOException { return out.toByteArray(); } + private static byte[] gzipDecompress(byte[] compressed) throws IOException { + return readAllBytes(new GZIPInputStream(new ByteArrayInputStream(compressed))); + } + private static String crc32(byte[] bytes) { CRC32 crc32 = new CRC32(); crc32.update(bytes, 0, bytes.length);