feat: implement lazy response body loading for improved memory efficiency

chirino · chirino · commit fdaa09dc16dc · 2025-06-02T16:25:30.000-04:00
• Add HttpResponseBody class for lazy response body management with supplier-based loading
• Add deferred loading mechanism - response body only loaded when WASM modules access it
• Fix request body processing logic to properly handle early returns for non-body plugins
• Update response header filtering logic to skip unnecessary processing when no body handling needed

This implementation extends the existing lazy request body loading concept to response bodies,
providing significant memory and performance improvements. The lazy loading mechanism ensures
response bodies are only read from the output stream when WASM plugins actually call
proxy_get_buffer_bytes to access the data. This prevents unnecessary memory allocation and
I/O operations for plugins that process headers but don't need response body content, while
maintaining full backward compatibility with existing plugin implementations.

Signed-off-by: Hiram Chirino &lt;hiram@hiramchirino.com&gt;
diff --git a/proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/internal/HttpResponseBody.java b/proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/internal/HttpResponseBody.java
@@ -0,0 +1,64 @@
+package io.roastedroot.proxywasm.internal;
+
+import java.util.function.Supplier;
+
+/**
+ * Holds the HTTP response body, loading it from a supplier when first accessed.
+ * Unlike request bodies which come from streams, response bodies can be provided
+ * directly as byte arrays or from suppliers.
+ */
+public class HttpResponseBody {
+
+    private byte[] body;
+    private boolean loaded = false;
+    private final Supplier<byte[]> bodySupplier;
+
+    public HttpResponseBody(Supplier<byte[]> bodySupplier) {
+        this.bodySupplier = bodySupplier;
+    }
+
+    /**
+     * Creates an HttpResponseBody with a fixed byte array (no lazy loading needed).
+     */
+    public HttpResponseBody(byte[] body) {
+        this.body = body;
+        this.loaded = true;
+        this.bodySupplier = null;
+    }
+
+    public byte[] get() {
+        if (!loaded) {
+            if (bodySupplier != null) {
+                body = bodySupplier.get();
+            } else {
+                body = new byte[0];
+            }
+            loaded = true;
+        }
+        return body;
+    }
+
+    /**
+     * Returns true if the response body has been loaded.
+     */
+    public boolean isLoaded() {
+        return loaded;
+    }
+
+    /**
+     * Sets the response body directly, marking it as loaded.
+     * This is used when the body is modified by WASM plugins.
+     */
+    public void setBody(byte[] body) {
+        this.body = body;
+        this.loaded = true;
+    }
+
+    /**
+     * Returns the response body if it has been loaded, null otherwise.
+     * This allows checking if the body was accessed without triggering a load.
+     */
+    public byte[] getBodyIfLoaded() {
+        return loaded ? body : null;
+    }
+}
diff --git a/proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/internal/PluginHttpContext.java b/proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/internal/PluginHttpContext.java
@@ -14,12 +14,12 @@ public class PluginHttpContext {
 
     final HashMap<List<String>, byte[]> properties = new HashMap<>();
     private HttpRequestBody httpRequestBodyState;
+    private HttpResponseBody httpResponseBodyState;
 
     // Other body buffers and state fields (not lazy)
     private byte[] grpcReceiveBuffer = new byte[0];
     private byte[] upstreamData = new byte[0];
     private byte[] downStreamData = new byte[0];
-    private byte[] httpResponseBody = new byte[0];
     private SendResponse sendResponse;
     private Action action;
     private CountDownLatch resumeLatch;
@@ -63,6 +63,39 @@ public void setHttpRequestBody(byte[] body) {
         }
     }
 
+    /**
+     * Sets the HTTP response body state.
+     */
+    public void setHttpResponseBodyState(HttpResponseBody responseBody) {
+        this.httpResponseBodyState = responseBody;
+    }
+
+    /**
+     * Gets the HTTP response body state.
+     */
+    public HttpResponseBody getHttpResponseBodyState() {
+        return httpResponseBodyState;
+    }
+
+    /**
+     * Gets the HTTP response body, triggering lazy loading if needed.
+     */
+    public byte[] getHttpResponseBody() {
+        if (httpResponseBodyState != null) {
+            return httpResponseBodyState.get();
+        }
+        return new byte[0];
+    }
+
+    /**
+     * Sets the HTTP response body, updating the state if present.
+     */
+    public void setHttpResponseBody(byte[] httpResponseBody) {
+        if (httpResponseBodyState != null && httpResponseBody != null) {
+            httpResponseBodyState.setBody(httpResponseBody);
+        }
+    }
+
     public Plugin plugin() {
         return plugin;
     }
@@ -103,14 +136,6 @@ public void maybePause() {
         }
     }
 
-    public byte[] getHttpResponseBody() {
-        return httpResponseBody;
-    }
-
-    public void setHttpResponseBody(byte[] httpResponseBody) {
-        this.httpResponseBody = httpResponseBody;
-    }
-
     public byte[] getGrpcReceiveBuffer() {
         return grpcReceiveBuffer;
     }
@@ -235,17 +260,18 @@ public WasmResult setDownStreamData(byte[] data) {
 
         @Override
         public byte[] getHttpResponseBody() {
-            return httpResponseBody;
+            return PluginHttpContext.this.getHttpResponseBody();
         }
 
         @Override
         public WasmResult setHttpResponseBody(byte[] body) {
-            httpResponseBody = body;
+            PluginHttpContext.this.setHttpResponseBody(body);
             return WasmResult.OK;
         }
 
         public void appendHttpResponseBody(byte[] body) {
-            httpResponseBody = Helpers.append(httpResponseBody, body);
+            byte[] currentBody = PluginHttpContext.this.getHttpResponseBody();
+            PluginHttpContext.this.setHttpResponseBody(Helpers.append(currentBody, body));
         }
 
         // //////////////////////////////////////////////////////////////////////
diff --git a/proxy-wasm-jaxrs/src/main/java/io/roastedroot/proxywasm/jaxrs/ProxyWasmFilter.java b/proxy-wasm-jaxrs/src/main/java/io/roastedroot/proxywasm/jaxrs/ProxyWasmFilter.java
@@ -5,6 +5,7 @@
 import io.roastedroot.proxywasm.StartException;
 import io.roastedroot.proxywasm.internal.Action;
 import io.roastedroot.proxywasm.internal.HttpRequestBody;
+import io.roastedroot.proxywasm.internal.HttpResponseBody;
 import io.roastedroot.proxywasm.internal.Plugin;
 import io.roastedroot.proxywasm.internal.PluginHttpContext;
 import io.roastedroot.proxywasm.internal.Pool;
@@ -170,32 +171,33 @@ private void filter(ContainerRequestContext requestContext, FilterContext filter
             }
 
             // the plugin may not be interested in the request body.
-            if (httpContext.context().hasOnRequestBody()) {
+            if (!httpContext.context().hasOnRequestBody()) {
+                return;
+            }
 
-                HttpRequestBody httpRequestBodyState = httpContext.getHttpRequestBodyState();
+            HttpRequestBody httpRequestBodyState = httpContext.getHttpRequestBodyState();
 
-                while (true) {
-                    // if we streamed body updates, then endOfStream would be initially false
-                    var action = httpContext.context().callOnRequestBody(true);
+            while (true) {
+                // if we streamed body updates, then endOfStream would be initially false
+                var action = httpContext.context().callOnRequestBody(true);
 
-                    // does the plugin want to respond early?
-                    var sendResponse = httpContext.consumeSentHttpResponse();
-                    if (sendResponse != null) {
-                        requestContext.abortWith(toResponse(sendResponse));
-                        return;
-                    }
-
-                    if (action == Action.CONTINUE) {
-                        break;
-                    }
-                    httpContext.maybePause();
+                // does the plugin want to respond early?
+                var sendResponse = httpContext.consumeSentHttpResponse();
+                if (sendResponse != null) {
+                    requestContext.abortWith(toResponse(sendResponse));
+                    return;
                 }
 
-                // Body was accessed and potentially modified, update the request stream
-                if (httpRequestBodyState.isLoaded()) {
-                    byte[] bytes = httpRequestBodyState.getBodyIfLoaded();
-                    requestContext.setEntityStream(new java.io.ByteArrayInputStream(bytes));
+                if (action == Action.CONTINUE) {
+                    break;
                 }
+                httpContext.maybePause();
+            }
+
+            // Body was accessed and potentially modified, update the request stream
+            if (httpRequestBodyState.isLoaded()) {
+                byte[] bytes = httpRequestBodyState.getBodyIfLoaded();
+                requestContext.setEntityStream(new java.io.ByteArrayInputStream(bytes));
             }
 
         } finally {
@@ -269,19 +271,17 @@ private void filter(
                 }
 
                 // aroundWriteTo won't be called if there is no entity to send.
-                if (responseContext.getEntity() == null
-                        && httpContext.context().hasOnResponseBody()) {
+                if (responseContext.getEntity() != null
+                        || !httpContext.context().hasOnResponseBody()) {
+                    return;
+                }
+
+                // Set up empty response body for plugins that need it
+                HttpResponseBody responseBodyState = new HttpResponseBody(new byte[0]);
+                httpContext.setHttpResponseBodyState(responseBodyState);
 
-                    byte[] bytes = new byte[0];
-                    httpContext.setHttpResponseBody(bytes);
+                while (true) {
                     action = httpContext.context().callOnResponseBody(true);
-                    bytes = httpContext.getHttpResponseBody();
-                    if (action == Action.CONTINUE) {
-                        // continue means plugin is done reading the body.
-                        httpContext.setHttpResponseBody(null);
-                    } else {
-                        httpContext.maybePause();
-                    }
 
                     // does the plugin want to respond early?
                     sendResponse = httpContext.consumeSentHttpResponse();
@@ -292,6 +292,11 @@ private void filter(
                         responseContext.setEntity(response.getEntity());
                         return;
                     }
+
+                    if (action == Action.CONTINUE) {
+                        break;
+                    }
+                    httpContext.maybePause();
                 }
 
             } finally {
@@ -332,41 +337,57 @@ public void aroundWriteTo(WriterInterceptorContext ctx)
         try {
 
             var original = ctx.getOutputStream();
-            ByteArrayOutputStream baos = new ByteArrayOutputStream();
-            ctx.setOutputStream(baos);
-            ctx.proceed();
 
-            byte[] bytes = baos.toByteArray();
+            HttpResponseBody sharedResponseBody =
+                    new HttpResponseBody(
+                            () -> {
+                                ByteArrayOutputStream baos = new ByteArrayOutputStream();
+                                ctx.setOutputStream(baos);
+                                try {
+                                    ctx.proceed();
+                                } catch (Exception e) {
+                                    throw new RuntimeException("Failed to read response body", e);
+                                }
+                                return baos.toByteArray();
+                            });
 
             for (var filterContext : List.copyOf(filterContexts)) {
                 var httpContext = filterContext.httpContext;
 
                 httpContext.plugin().lock();
 
-                // the plugin may not be interested in the request body.
+                // the plugin may not be interested in the response body.
                 if (!httpContext.context().hasOnResponseBody()) {
-                    ctx.proceed();
+                    continue;
                 }
 
-                httpContext.setHttpResponseBody(bytes);
-                var action = httpContext.context().callOnResponseBody(true);
-                bytes = httpContext.getHttpResponseBody();
-                if (action == Action.CONTINUE) {
-                    // continue means plugin is done reading the body.
-                    httpContext.setHttpResponseBody(null);
-                } else {
-                    httpContext.maybePause();
-                }
+                // Set up lazy response body - will only be accessed if plugin needs it
+                httpContext.setHttpResponseBodyState(sharedResponseBody);
 
-                // does the plugin want to respond early?
-                var sendResponse = httpContext.consumeSentHttpResponse();
-                if (sendResponse != null) {
-                    throw new WebApplicationException(toResponse(sendResponse));
+                while (true) {
+                    var action = httpContext.context().callOnResponseBody(true);
+
+                    // does the plugin want to respond early?
+                    var sendResponse = httpContext.consumeSentHttpResponse();
+                    if (sendResponse != null) {
+                        throw new WebApplicationException(toResponse(sendResponse));
+                    }
+
+                    if (action == Action.CONTINUE) {
+                        break;
+                    }
+                    httpContext.maybePause();
                 }
             }
 
-            // plugin may have modified the body
-            original.write(bytes);
+            // Write the response body - if it was accessed and modified, use that,
+            // otherwise continue with the original stream.
+            if (sharedResponseBody.isLoaded()) {
+                original.write(sharedResponseBody.get());
+            } else {
+                // Body was never accessed by any plugin, use original
+                ctx.proceed();
+            }
 
         } finally {
             for (var filterContext : List.copyOf(filterContexts)) {
