fix: cleanup call fiber on exit (#515)

Co-authored-by: Regis Kuckaertz <[email protected]>
scalapb · Jun 19, 2023 · 64b7ded · 64b7ded
1 parent 7eb76a1
commit 64b7ded
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/core/src/main/scalajvm/scalapb/zio_grpc/server/ListenerDriver.scala b/core/src/main/scalajvm/scalapb/zio_grpc/server/ListenerDriver.scala
@@ -18,6 +18,7 @@ object ListenerDriver {
 
   def unaryInputListener[Req](
       runtime: Runtime[Any],
+      scope: Scope.Closeable,
       call: ZServerCall[_],
       completed: Promise[StatusException, Unit],
       request: Promise[Nothing, Req],
@@ -29,13 +30,10 @@ object ListenerDriver {
         completed.await *>
         call.sendHeaders(new Metadata) *>
         request.await flatMap writeResponse
-    ).onExit(ex => call.close(ListenerDriver.exitToStatus(ex), requestContext.responseMetadata.metadata).ignore)
-      .ignore
-      // Why forkDaemon? we need the driver to keep runnning in the background after we return a listener
-      // back to grpc-java. If it was just fork, the call to unsafeRun would not return control, so grpc-java
-      // won't have a listener to call on.  The driver awaits on the calls to the listener to pass to the user's
-      // service.
-      .forkDaemon
+    ).onExit(ex =>
+      call.close(ListenerDriver.exitToStatus(ex), requestContext.responseMetadata.metadata).ignore *> scope.close(ex)
+    ).ignore
+      .forkIn(scope)
       .map(fiber =>
         new Listener[Req] {
           override def onCancel(): Unit =
@@ -91,8 +89,10 @@ object ListenerDriver {
     for {
       completed <- Promise.make[StatusException, Unit]
       request   <- Promise.make[Nothing, Req]
+      scope     <- Scope.make
       listener  <- unaryInputListener(
                      runtime,
+                     scope,
                      zioCall,
                      completed,
                      request,