Merge pull request #2007 from Netflix/kavitha/fix-scheduled-dataloaders

Fix: Close the scheduled data loader registry upon completion of request
Netflix · Sep 12, 2024 · 2942c51 · 2942c51
2 parents 3004820 + 5586e4d
commit 2942c51
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 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
@@ -68,10 +68,18 @@ class DgsWebMvcGraphQLInterceptor(
         graphQLContextFuture.complete(request.toExecutionInput().graphQLContext)
 
         return if (dgsSpringConfigurationProperties.webmvc.asyncdispatch.enabled) {
-            chain.next(request)
+            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)
         }
     }
 }