Move http2 compression handling from connection to stream level

Author: zyclonite

vertx-core/src/main/java/io/vertx/core/http/impl/http2/codec/VertxCompressorHttp2ConnectionEncoder.java

@@ -35,23 +35,26 @@
 
 public class VertxCompressorHttp2ConnectionEncoder implements Http2FrameWriter, Http2ConnectionEncoder, Http2SettingsReceivedConsumer {
 
-  private Http2ConnectionEncoder delegate;
+  private final Http2ConnectionEncoder delegate;
   private final Http2ConnectionEncoder plainEncoder;
+  private final Http2Connection.PropertyKey encoderKey;
 
   public VertxCompressorHttp2ConnectionEncoder(Http2ConnectionEncoder plainEncoder, CompressionOptions[] compressionOptions) {
     this.delegate = new CompressorHttp2ConnectionEncoder(plainEncoder, compressionOptions);
     this.plainEncoder = plainEncoder;
+    this.encoderKey = plainEncoder.connection().newKey();
   }
 
-  private void beforeWritingHeaders(ChannelHandlerContext ctx, int streamId, Http2Headers responseHeaders) {
+  private Http2ConnectionEncoder selectEncoder(ChannelHandlerContext ctx, int streamId, Http2Headers responseHeaders) {
     String contentEncodingToApply = determineContentEncodingToApply(ctx, streamId, responseHeaders);
     if (contentEncodingToApply == null || contentEncodingToApply.equalsIgnoreCase(IDENTITY.toString())) {
       if (responseHeaders.contains(CONTENT_ENCODING, IDENTITY)) {
         responseHeaders.remove(CONTENT_ENCODING);
       }
-      delegate = plainEncoder;
+      return plainEncoder;
     } else {
       responseHeaders.set(CONTENT_ENCODING, contentEncodingToApply);
+      return delegate;
     }
   }
 
@@ -69,16 +72,42 @@ private <T, R> R ifType(Object obj, Class<T> type, Function<T, R> then) {
     return obj != null && type.isAssignableFrom(obj.getClass()) ? then.apply(type.cast(obj)) : null;
   }
 
+  private void storeEncoder(int streamId, Http2ConnectionEncoder encoder) {
+    io.netty.handler.codec.http2.Http2Stream stream = delegate.connection().stream(streamId);
+    if (stream != null) {
+      stream.setProperty(encoderKey, encoder);
+    }
+  }
+
+  private Http2ConnectionEncoder getStoredEncoder(int streamId) {
+    io.netty.handler.codec.http2.Http2Stream stream = delegate.connection().stream(streamId);
+    if (stream != null) {
+      Http2ConnectionEncoder encoder = stream.getProperty(encoderKey);
+      if (encoder != null) {
+        return encoder;
+      }
+    }
+    return delegate;
+  }
+
   @Override
   public ChannelFuture writeHeaders(ChannelHandlerContext ctx, int streamId, Http2Headers headers, int padding, boolean endStream, ChannelPromise promise) {
-    beforeWritingHeaders(ctx, streamId, headers);
-    return delegate.writeHeaders(ctx, streamId, headers, padding, endStream, promise);
+    Http2ConnectionEncoder encoder = selectEncoder(ctx, streamId, headers);
+    storeEncoder(streamId, encoder);
+    return encoder.writeHeaders(ctx, streamId, headers, padding, endStream, promise);
   }
 
   @Override
   public ChannelFuture writeHeaders(ChannelHandlerContext ctx, int streamId, Http2Headers headers, int streamDependency, short weight, boolean exclusive, int padding, boolean endStream, ChannelPromise promise) {
-    beforeWritingHeaders(ctx, streamId, headers);
-    return delegate.writeHeaders(ctx, streamId, headers, streamDependency, weight, exclusive, padding, endStream, promise);
+    Http2ConnectionEncoder encoder = selectEncoder(ctx, streamId, headers);
+    storeEncoder(streamId, encoder);
+    return encoder.writeHeaders(ctx, streamId, headers, streamDependency, weight, exclusive, padding, endStream, promise);
+  }
+
+  @Override
+  public ChannelFuture writeData(ChannelHandlerContext ctx, int streamId, ByteBuf data, int padding, boolean endStream, ChannelPromise promise) {
+    Http2ConnectionEncoder encoder = getStoredEncoder(streamId);
+    return encoder.writeData(ctx, streamId, data, padding, endStream, promise);
   }
 
   @Override
@@ -166,11 +195,6 @@ public void close() {
     delegate.close();
   }
 
-  @Override
-  public ChannelFuture writeData(ChannelHandlerContext ctx, int streamId, ByteBuf data, int padding, boolean endStream, ChannelPromise promise) {
-    return delegate.writeData(ctx, streamId, data, padding, endStream, promise);
-  }
-
   @Override
   public void consumeReceivedSettings(Http2Settings settings) {
     if (delegate instanceof Http2SettingsReceivedConsumer) {

vertx-core/src/test/java/io/vertx/tests/http/Http2ServerTest.java

@@ -89,6 +89,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -2257,6 +2258,88 @@ public int onDataRead(ChannelHandlerContext ctx, int streamId, ByteBuf data, int
     await();
   }
 
+  @Test
+  public void testResponseCompressionEnabledMixedStreams() throws Exception {
+    waitFor(6);
+    String expected = TestUtils.randomAlphaString(1000);
+    server.close();
+    server = vertx.createHttpServer(new HttpServerOptions(serverOptions).setCompressionSupported(true));
+    server.requestHandler(req -> {
+      req.response().end(expected);
+    });
+    startServer();
+    TestClient client = new TestClient();
+    client.connect(DEFAULT_HTTPS_PORT, DEFAULT_HTTPS_HOST, request -> {
+      // Stream 1: with compression
+      int id1 = request.nextStreamId();
+      // Stream 2: without compression
+      int id2 = request.nextStreamId();
+      // Stream 3: with compression
+      int id3 = request.nextStreamId();
+      Map<Integer, ByteArrayOutputStream> streamData = new ConcurrentHashMap<>();
+      streamData.put(id1, new ByteArrayOutputStream());
+      streamData.put(id2, new ByteArrayOutputStream());
+      streamData.put(id3, new ByteArrayOutputStream());
+      request.decoder.frameListener(new Http2EventAdapter() {
+        @Override
+        public void onHeadersRead(ChannelHandlerContext ctx, int streamId, Http2Headers headers, int streamDependency, short weight, boolean exclusive, int padding, boolean endStream) throws Http2Exception {
+          vertx.runOnContext(v -> {
+            if (streamId == id1 || streamId == id3) {
+              // Stream with accept-encoding: gzip should be compressed
+              Assert.assertEquals("gzip", headers.get(HttpHeaderNames.CONTENT_ENCODING).toString());
+            } else {
+              // Stream without accept-encoding should not be compressed
+              Assert.assertFalse(headers.contains(HttpHeaderNames.CONTENT_ENCODING));
+            }
+            complete();
+          });
+        }
+        @Override
+        public int onDataRead(ChannelHandlerContext ctx, int streamId, ByteBuf data, int padding, boolean endOfStream) throws Http2Exception {
+          byte[] bytes = new byte[data.readableBytes()];
+          data.readBytes(bytes);
+          ByteArrayOutputStream buf = streamData.get(streamId);
+          buf.write(bytes, 0, bytes.length);
+          if (endOfStream) {
+            byte[] allBytes = buf.toByteArray();
+            vertx.runOnContext(v -> {
+              if (streamId == id1 || streamId == id3) {
+                // Compressed stream - decode gzip
+                String decoded;
+                try {
+                  GZIPInputStream in = new GZIPInputStream(new ByteArrayInputStream(allBytes));
+                  ByteArrayOutputStream baos = new ByteArrayOutputStream();
+                  while (true) {
+                    int i = in.read();
+                    if (i == -1) {
+                      break;
+                    }
+                    baos.write(i);
+                  }
+                  decoded = baos.toString();
+                } catch (IOException e) {
+                  fail(e);
+                  return;
+                }
+                Assert.assertEquals(expected, decoded);
+              } else {
+                // Uncompressed stream - plain text
+                Assert.assertEquals(expected, new String(allBytes, StandardCharsets.UTF_8));
+              }
+              complete();
+            });
+          }
+          return super.onDataRead(ctx, streamId, data, padding, endOfStream);
+        }
+      });
+      request.encoder.writeHeaders(request.context, id1, GET("/").add("accept-encoding", "gzip"), 0, true, request.context.newPromise());
+      request.encoder.writeHeaders(request.context, id2, GET("/"), 0, true, request.context.newPromise());
+      request.encoder.writeHeaders(request.context, id3, GET("/").add("accept-encoding", "gzip"), 0, true, request.context.newPromise());
+      request.context.flush();
+    });
+    await();
+  }
+
   @Test
   public void testRequestCompressionEnabled() throws Exception {
     String expected = TestUtils.randomAlphaString(1000);