fix(tvix/castore/GRPCDirectorySvc): fix a bug in the get_recursive fn r/8412

When retrieving a closure with get_recursive, the following could happen in the GRPC client:
- The first reference to the deduplicated directory is added to expected_directory_digests
- The deduplicated directory is obtained removed from expected_directory_digests
- The second reference to the deduplicated directory is added to expected_directory_digests
- The deduplicated directory has already been sent, but is still in the
  expected_directory_digests. It looks to the GRPC client like the
  closure is incomplete and the stream ended prematurely.

Change-Id: Ic62bca12e7f8fb85af5fa4dacd199f0f3b8eea8c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12033
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

1 files changed, 9 insertions, 3 deletions

diff --git a/tvix/castore/src/directoryservice/grpc.rs b/tvix/castore/src/directoryservice/grpc.rs
index 6d81be0e70db..4dc3931ed410 100644
--- a/tvix/castore/src/directoryservice/grpc.rs
+++ b/tvix/castore/src/directoryservice/grpc.rs
@@ -176,11 +176,17 @@ where
                         return
                     }
                     Ok(None) => {
-                        // If we were still expecting something, that's an error.
-                        if !expected_directory_digests.is_empty() {
+                        // The stream has ended
+                        let diff_len = expected_directory_digests
+                            // Account for directories which have been referenced more than once,
+                            // but only received once since they were deduplicated
+                            .difference(&received_directory_digests)
+                            .count();
+                        // If this is not empty, then the closure is incomplete
+                        if diff_len != 0 {
                             Err(crate::Error::StorageError(format!(
                                 "still expected {} directories, but got premature end of stream",
-                                expected_directory_digests.len(),
+                                diff_len
                             )))?
                         } else {
                             return