Fix leak in ProxyWithAuthorizationTest

apache · Nov 3, 2023 · e31e499 · e31e499
1 parent eb475ae
commit e31e499
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyService.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyService.java
@@ -34,6 +34,7 @@
 import io.netty.resolver.dns.DnsAddressResolverGroup;
 import io.netty.resolver.dns.DnsNameResolverBuilder;
 import io.netty.util.concurrent.DefaultThreadFactory;
+import io.netty.util.concurrent.Future;
 import io.prometheus.client.Counter;
 import io.prometheus.client.Gauge;
 import java.io.Closeable;
@@ -151,6 +152,8 @@ public class ProxyService implements Closeable {
     @Getter
     private final ConnectionController connectionController;
 
+    private boolean gracefulShutdown = true;
+
     public ProxyService(ProxyConfiguration proxyConfig,
                         AuthenticationService authenticationService) throws Exception {
         requireNonNull(proxyConfig);
@@ -373,7 +376,7 @@ public void close() throws IOException {
 
         // Don't accept any new connections
         try {
-            acceptorGroup.shutdownGracefully().sync();
+            shutdownEventLoop(acceptorGroup).sync();
         } catch (InterruptedException e) {
             LOG.info("Shutdown of acceptorGroup interrupted");
             Thread.currentThread().interrupt();
@@ -413,14 +416,14 @@ public void close() throws IOException {
             }
         }
         try {
-            workerGroup.shutdownGracefully().sync();
+            shutdownEventLoop(workerGroup).sync();
         } catch (InterruptedException e) {
             LOG.info("Shutdown of workerGroup interrupted");
             Thread.currentThread().interrupt();
         }
         for (EventLoopGroup group : extensionsWorkerGroups) {
             try {
-                group.shutdownGracefully().sync();
+                shutdownEventLoop(group).sync();
             } catch (InterruptedException e) {
                 LOG.info("Shutdown of {} interrupted", group);
                 Thread.currentThread().interrupt();
@@ -532,4 +535,24 @@ public synchronized void addPrometheusRawMetricsProvider(PrometheusRawMetricsPro
     protected LookupProxyHandler newLookupProxyHandler(ProxyConnection proxyConnection) {
         return new LookupProxyHandler(this, proxyConnection);
     }
+
+    // Shutdown the event loop.
+    // If graceful is true, will wait for the current requests to be completed, up to 15 seconds.
+    // Graceful shutdown can be disabled by setting the gracefulShutdown flag to false. This is used in tests
+    // to speed up the shutdown process.
+    private Future<?> shutdownEventLoop(EventLoopGroup eventLoop) {
+        if (gracefulShutdown) {
+            return eventLoop.shutdownGracefully();
+        } else {
+            return eventLoop.shutdownGracefully(0, 0, TimeUnit.SECONDS);
+        }
+    }
+
+    public boolean isGracefulShutdown() {
+        return gracefulShutdown;
+    }
+
+    public void setGracefulShutdown(boolean gracefulShutdown) {
+        this.gracefulShutdown = gracefulShutdown;
+    }
 }
diff --git a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyWithAuthorizationTest.java b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyWithAuthorizationTest.java
@@ -229,6 +229,7 @@ protected void setup() throws Exception {
         AuthenticationService authService =
                 new AuthenticationService(PulsarConfigurationLoader.convertFrom(proxyConfig));
         proxyService = Mockito.spy(new ProxyService(proxyConfig, authService));
+        proxyService.setGracefulShutdown(false);
         webServer = new WebServer(proxyConfig, authService);
     }
 
@@ -455,9 +456,11 @@ public void tlsCiphersAndProtocols(Set<String> tlsCiphers, Set<String> tlsProtoc
         proxyConfig.setTlsProtocols(tlsProtocols);
         proxyConfig.setTlsCiphers(tlsCiphers);
 
+        @Cleanup
         ProxyService proxyService = Mockito.spy(new ProxyService(proxyConfig,
                                                         new AuthenticationService(
                                                                 PulsarConfigurationLoader.convertFrom(proxyConfig))));
+        proxyService.setGracefulShutdown(false);
         try {
             proxyService.start();
         } catch (Exception ex) {