flaky test

LangStream · May 20, 2024 · 6ada992 · 6ada992
1 parent 64e00d8
commit 6ada992
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 9 deletions.
diff --git a/langstream-api-gateway/src/main/java/ai/langstream/apigateway/http/GatewayResource.java b/langstream-api-gateway/src/main/java/ai/langstream/apigateway/http/GatewayResource.java
@@ -270,13 +270,7 @@ private CompletableFuture<ResponseEntity> handleServiceWithTopics(
                             topicConnectionsRuntimeRegistryProvider
                                     .getTopicConnectionsRuntimeRegistry(),
                             clusterRuntimeRegistry);
-            completableFuture.thenRunAsync(
-                    () -> {
-                        if (consumeGateway != null) {
-                            consumeGateway.close();
-                        }
-                    },
-                    consumeThreadPool);
+            completableFuture.thenRunAsync(consumeGateway::close, consumeThreadPool);
 
             final Gateway.ServiceOptions serviceOptions = authContext.gateway().getServiceOptions();
             try {
@@ -297,7 +291,7 @@ record -> {
                 final AtomicBoolean stop = new AtomicBoolean(false);
                 consumeGateway.startReadingAsync(
                         consumeThreadPool,
-                        () -> stop.get(),
+                        stop::get,
                         record -> {
                             stop.set(true);
                             completableFuture.complete(ResponseEntity.ok(record));

diff --git a/langstream-api-gateway/src/main/java/ai/langstream/apigateway/websocket/WebSocketConfig.java b/langstream-api-gateway/src/main/java/ai/langstream/apigateway/websocket/WebSocketConfig.java
@@ -97,7 +97,8 @@ public ServletServerContainerFactoryBean createWebSocketContainer() {
 
     @PreDestroy
     public void onDestroy() {
-        consumeThreadPool.shutdown();
+        log.info("Shutting down WebSocket");
+        consumeThreadPool.shutdownNow();
         clusterRuntimeRegistry.close();
     }
 }