Close on complete in GrpcScheduler (#328)

The GrpcScheduler fails to close the downstream Receiver when the upstream one completes. This causes it to hang forever even though its completed. Fixes #323
TraceMachina · Oct 22, 2023 · 6c937da · 6c937da
1 parent 1339e9d
commit 6c937da
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/cas/scheduler/grpc_scheduler.rs b/cas/scheduler/grpc_scheduler.rs
@@ -66,7 +66,12 @@ impl GrpcScheduler {
                             log::info!("Client disconnected in GrpcScheduler");
                             return;
                         }
-                        Ok(Some(response)) = result_stream.message() => {
+                        response = result_stream.message() => {
+                            // When the upstream closes the channel, close the
+                            // downstream too.
+                            let Ok(Some(response)) = response else {
+                                return;
+                            };
                             match response.try_into() {
                                 Ok(response) => {
                                     if let Err(err) = tx.send(Arc::new(response)) {