Skip to content

Fix compression chaining with vector interceptor #151

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dkropachev wants to merge 1 commit into
base: vector-1
Choose a base branch
from
Open

Fix compression chaining with vector interceptor #151

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 47 additions & 39 deletions src/main/java/com/scylladb/alternator/GzipRequestInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@ public class GzipRequestInterceptor implements ExecutionInterceptor {

private static final ExecutionAttribute<byte[]> ORIGINAL_BODY_BYTES =
new ExecutionAttribute<>("GzipRequestInterceptor.originalBodyBytes");
private static final ExecutionAttribute<byte[]> COMPRESSED_BODY_BYTES =
new ExecutionAttribute<>("GzipRequestInterceptor.compressedBodyBytes");
private static final ExecutionAttribute<Boolean> SHOULD_COMPRESS =
new ExecutionAttribute<>("GzipRequestInterceptor.shouldCompress");

Expand All @@ -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<RequestBody> 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
Expand All @@ -91,26 +79,51 @@ public Optional<RequestBody> 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<AsyncRequestBody> 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);
Expand All @@ -120,17 +133,12 @@ public Optional<AsyncRequestBody> 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 {
Expand Down
113 changes: 106 additions & 7 deletions src/test/java/com/scylladb/alternator/GzipRequestInterceptorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -111,14 +126,8 @@ private byte[] generateTestData(int size) {

private byte[] readRequestBody(Optional<RequestBody> 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<AsyncRequestBody> body) {
Expand All @@ -136,6 +145,56 @@ private byte[] readAsyncRequestBody(Optional<AsyncRequestBody> 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);
Expand Down Expand Up @@ -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";
}
}
}
68 changes: 68 additions & 0 deletions src/test/java/com/scylladb/alternator/ResponseCompressionInterceptorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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")
Expand Down
Loading
Loading