Fix: Close the scheduled data loader registry upon completion of requ…

…est.
Netflix · Sep 12, 2024 · 393bef2 · 393bef2
1 parent 3004820
commit 393bef2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/...main/kotlin/com/netflix/graphql/dgs/springgraphql/webflux/DgsWebFluxGraphQLInterceptor.kt b/...main/kotlin/com/netflix/graphql/dgs/springgraphql/webflux/DgsWebFluxGraphQLInterceptor.kt
@@ -60,6 +60,10 @@ class DgsWebFluxGraphQLInterceptor(
                         .build()
                 }
                 graphQLContextFuture.complete(request.toExecutionInput().graphQLContext)
-                chain.next(request)
+                chain.next(request).doFinally {
+                    if (dataLoaderRegistry is AutoCloseable) {
+                        dataLoaderRegistry.close()
+                    }
+                }
             }
 }
diff --git a/...c/main/kotlin/com/netflix/graphql/dgs/springgraphql/webmvc/DgsWebMvcGraphQLInterceptor.kt b/...c/main/kotlin/com/netflix/graphql/dgs/springgraphql/webmvc/DgsWebMvcGraphQLInterceptor.kt
@@ -67,11 +67,19 @@ class DgsWebMvcGraphQLInterceptor(
         }
         graphQLContextFuture.complete(request.toExecutionInput().graphQLContext)
 
-        return if (dgsSpringConfigurationProperties.webmvc.asyncdispatch.enabled) {
-            chain.next(request)
+       return if (dgsSpringConfigurationProperties.webmvc.asyncdispatch.enabled) {
+            chain.next(request).doFinally {
+                if (dataLoaderRegistry is AutoCloseable) {
+                    dataLoaderRegistry.close()
+                }
+            }
         } else {
             @Suppress("BlockingMethodInNonBlockingContext")
-            return Mono.just(chain.next(request).block()!!)
+            val response = chain.next(request).block()!!
+            if(dataLoaderRegistry is AutoCloseable) {
+                dataLoaderRegistry.close()
+            }
+            return Mono.just(response)
         }
     }
 }