From d7f6ea35fcc65a0fbb6a878d2ee1e512b03ed60b Mon Sep 17 00:00:00 2001
From: Chris Staite <christopher.staite@menlosecurity.com>
Date: Fri, 20 Oct 2023 08:34:46 +0000
Subject: [PATCH] Close on complete in GrpcScheduler

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
---
 cas/scheduler/grpc_scheduler.rs | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/cas/scheduler/grpc_scheduler.rs b/cas/scheduler/grpc_scheduler.rs
index 63ababfbe..27f1fbb34 100644
--- 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)) {