From 6a988a1598ac81a06454f9b18d36e4e3ac33cac6 Mon Sep 17 00:00:00 2001 From: Yureka Date: Wed, 24 Jul 2024 20:54:37 +0200 Subject: [PATCH] fix(tvix/castore/GRPCDirectorySvc): fix a bug in the get_recursive fn 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 --- tvix/castore/src/directoryservice/grpc.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tvix/castore/src/directoryservice/grpc.rs b/tvix/castore/src/directoryservice/grpc.rs index 6d81be0e7..4dc3931ed 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