Configure Buffer pools for Jetty server

This circumvents jetty/jetty.project#12670

Author: wendigo

http-server/src/main/java/io/airlift/http/server/HttpServer.java

@@ -22,6 +22,7 @@
 import io.airlift.http.server.jetty.MonitoredQueuedThreadPoolMBean;
 import io.airlift.log.Logger;
 import io.airlift.node.NodeInfo;
+import io.airlift.units.DataSize;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 import jakarta.servlet.Filter;
@@ -34,6 +35,8 @@
 import org.eclipse.jetty.http2.server.AuthorityCustomizer;
 import org.eclipse.jetty.http2.server.HTTP2CServerConnectionFactory;
 import org.eclipse.jetty.http2.server.HTTP2ServerConnectionFactory;
+import org.eclipse.jetty.io.ArrayByteBufferPool;
+import org.eclipse.jetty.io.ByteBufferPool;
 import org.eclipse.jetty.io.ConnectionStatistics;
 import org.eclipse.jetty.jmx.MBeanContainer;
 import org.eclipse.jetty.server.ConnectionFactory;
@@ -65,6 +68,7 @@
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.OptionalLong;
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
@@ -75,6 +79,7 @@
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
+import static java.lang.Math.max;
 import static java.lang.Math.toIntExact;
 import static java.time.temporal.ChronoUnit.DAYS;
 import static java.util.Collections.list;
@@ -139,7 +144,19 @@ public HttpServer(
             log.info("Virtual threads support is enabled");
             threadPool.setVirtualThreadsExecutor(executor);
         }
-        server = new Server(threadPool);
+
+        int maxBufferSize = toIntExact(max(
+                toSafeBytes(config.getMaxRequestHeaderSize()).orElse(65536),
+                toSafeBytes(config.getMaxResponseHeaderSize()).orElse(65536)));
+
+        ByteBufferPool byteBufferPool = new ArrayByteBufferPool.Quadratic(
+                0,
+                maxBufferSize,
+                Integer.MAX_VALUE,
+                config.getMaxHeapMemory().toBytes(),
+                config.getMaxDirectMemory().toBytes());
+
+        server = new Server(threadPool, null, byteBufferPool);
         this.monitoredQueuedThreadPoolMBean = new MonitoredQueuedThreadPoolMBean(threadPool);
 
         boolean showStackTrace = config.isShowStackTrace();
@@ -273,7 +290,6 @@ public HttpServer(
         }
 
         server.setHandler(statsHandler);
-
         ErrorHandler errorHandler = new ErrorHandler();
         errorHandler.setShowMessageInTitle(showStackTrace);
         errorHandler.setShowStacks(showStackTrace);
@@ -471,4 +487,13 @@ private static ServerConnector createServerConnector(
         connector.open(channel);
         return connector;
     }
+
+    private static OptionalLong toSafeBytes(DataSize dataSize)
+    {
+        if (dataSize == null) {
+            return OptionalLong.empty();
+        }
+
+        return OptionalLong.of(dataSize.toBytes());
+    }
 }

http-server/src/main/java/io/airlift/http/server/HttpServerConfig.java

@@ -15,8 +15,10 @@
  */
 package io.airlift.http.server;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
+import io.airlift.configuration.ConfigHidden;
 import io.airlift.configuration.DefunctConfig;
 import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
@@ -94,6 +96,16 @@ public class HttpServerConfig
 
     private boolean compressionEnabled = true;
 
+    private DataSize maxHeapMemory = clamp(
+            DataSize.ofBytes(Runtime.getRuntime().maxMemory() / 8),
+            DataSize.of(128, MEGABYTE),
+            DataSize.of(512, MEGABYTE));
+
+    private DataSize maxDirectMemory = clamp(
+            DataSize.ofBytes(Runtime.getRuntime().maxMemory() / 8),
+            DataSize.of(128, MEGABYTE),
+            DataSize.of(512, MEGABYTE));
+
     public boolean isHttpEnabled()
     {
         return httpEnabled;
@@ -473,6 +485,36 @@ public HttpServerConfig setCompressionEnabled(boolean compressionEnabled)
         return this;
     }
 
+    @MinDataSize("128MB")
+    @MaxDataSize("32GB")
+    public DataSize getMaxHeapMemory()
+    {
+        return maxHeapMemory;
+    }
+
+    @Config("http-server.max-heap-memory")
+    @ConfigHidden
+    public HttpServerConfig setMaxHeapMemory(DataSize maxHeapMemory)
+    {
+        this.maxHeapMemory = maxHeapMemory;
+        return this;
+    }
+
+    @MinDataSize("128MB")
+    @MaxDataSize("32GB")
+    public DataSize getMaxDirectMemory()
+    {
+        return maxDirectMemory;
+    }
+
+    @Config("http-server.max-direct-memory")
+    @ConfigHidden
+    public HttpServerConfig setMaxDirectMemory(DataSize maxDirectMemory)
+    {
+        this.maxDirectMemory = maxDirectMemory;
+        return this;
+    }
+
     public enum ProcessForwardedMode
     {
         ACCEPT,
@@ -488,4 +530,16 @@ public static ProcessForwardedMode fromString(String value)
             };
         }
     }
+
+    @VisibleForTesting
+    static DataSize clamp(DataSize value, DataSize min, DataSize max)
+    {
+        if (value.compareTo(min) < 0) {
+            return min;
+        }
+        if (value.compareTo(max) > 0) {
+            return max;
+        }
+        return value;
+    }
 }

http-server/src/test/java/io/airlift/http/server/TestHttpServerConfig.java

@@ -27,6 +27,7 @@
 import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
 import static io.airlift.http.server.HttpServerConfig.ProcessForwardedMode.IGNORE;
 import static io.airlift.http.server.HttpServerConfig.ProcessForwardedMode.REJECT;
+import static io.airlift.http.server.HttpServerConfig.clamp;
 import static io.airlift.units.DataSize.Unit.GIGABYTE;
 import static io.airlift.units.DataSize.Unit.KILOBYTE;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
@@ -67,7 +68,9 @@ public void testDefaults()
                 .setHttp2InputBufferSize(DataSize.of(8, KILOBYTE))
                 .setHttp2InitialStreamReceiveWindowSize(DataSize.of(16, MEGABYTE))
                 .setHttp2StreamIdleTimeout(new Duration(15, SECONDS))
-                .setCompressionEnabled(true));
+                .setCompressionEnabled(true)
+                .setMaxHeapMemory(clamp(DataSize.ofBytes(Runtime.getRuntime().maxMemory() / 8), DataSize.of(128, MEGABYTE), DataSize.of(512, MEGABYTE)))
+                .setMaxDirectMemory(clamp(DataSize.ofBytes(Runtime.getRuntime().maxMemory() / 8), DataSize.of(128, MEGABYTE), DataSize.of(512, MEGABYTE))));
     }
 
     @Test
@@ -103,6 +106,8 @@ public void testExplicitPropertyMappings()
                 .put("http-server.http2.input-buffer-size", "4MB")
                 .put("http-server.http2.stream-idle-timeout", "23s")
                 .put("http-server.compression.enabled", "false")
+                .put("http-server.max-heap-memory", "2GB")
+                .put("http-server.max-direct-memory", "2GB")
                 .build();
 
         HttpServerConfig expected = new HttpServerConfig()
@@ -134,7 +139,9 @@ public void testExplicitPropertyMappings()
                 .setHttp2InitialStreamReceiveWindowSize(DataSize.of(4, MEGABYTE))
                 .setHttp2InputBufferSize(DataSize.of(4, MEGABYTE))
                 .setHttp2StreamIdleTimeout(new Duration(23, SECONDS))
-                .setCompressionEnabled(false);
+                .setCompressionEnabled(false)
+                .setMaxHeapMemory(DataSize.of(2, GIGABYTE))
+                .setMaxDirectMemory(DataSize.of(2, GIGABYTE));
 
         assertFullMapping(properties, expected);
     }