From b8ff08b1b0d2dbd8dd546dc9cbdea2f11304d5c8 Mon Sep 17 00:00:00 2001
From: Florian Klink <flokli@flokli.de>
Date: Thu, 18 May 2023 21:43:33 +0300
Subject: refactor(tvix/store/directorysvc): move from Vec<u8> to B3Digest

This introduces a new struct, B3Digest, which internally holds a
Vec<u8>, but only allows construction with 32 bytes.

It also implements display, which will print the base64 representation.
This should reduce some boilerplate when parsing Vec<u8>.

Change-Id: Ia91aa40cb691916773abc8f93e6ed79a5fd34863
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8592
Reviewed-by: tazjin <tazjin@tvl.su>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
---
 tvix/store/src/directoryservice/utils.rs | 36 ++++++++++++++------------------
 1 file changed, 16 insertions(+), 20 deletions(-)

(limited to 'tvix/store/src/directoryservice/utils.rs')

diff --git a/tvix/store/src/directoryservice/utils.rs b/tvix/store/src/directoryservice/utils.rs
index 7d41e8d3ba01..3661808734f3 100644
--- a/tvix/store/src/directoryservice/utils.rs
+++ b/tvix/store/src/directoryservice/utils.rs
@@ -1,6 +1,7 @@
 use super::DirectoryPutter;
 use super::DirectoryService;
 use crate::proto;
+use crate::B3Digest;
 use crate::Error;
 use std::collections::{HashSet, VecDeque};
 use tracing::{debug_span, instrument, warn};
@@ -13,17 +14,17 @@ pub struct DirectoryTraverser<DS: DirectoryService> {
     /// The list of all directories that still need to be traversed. The next
     /// element is picked from the front, new elements are enqueued at the
     /// back.
-    worklist_directory_digests: VecDeque<[u8; 32]>,
+    worklist_directory_digests: VecDeque<B3Digest>,
     /// The list of directory digests already sent to the consumer.
     /// We omit sending the same directories multiple times.
-    sent_directory_digests: HashSet<[u8; 32]>,
+    sent_directory_digests: HashSet<B3Digest>,
 }
 
 impl<DS: DirectoryService> DirectoryTraverser<DS> {
-    pub fn with(directory_service: DS, root_directory_digest: &[u8; 32]) -> Self {
+    pub fn with(directory_service: DS, root_directory_digest: &B3Digest) -> Self {
         Self {
             directory_service,
-            worklist_directory_digests: VecDeque::from([*root_directory_digest]),
+            worklist_directory_digests: VecDeque::from([root_directory_digest.clone()]),
             sent_directory_digests: HashSet::new(),
         }
     }
@@ -33,12 +34,8 @@ impl<DS: DirectoryService> DirectoryTraverser<DS> {
     // This panics if the digest looks invalid, it's supposed to be checked first.
     fn enqueue_child_directories(&mut self, directory: &proto::Directory) {
         for child_directory_node in &directory.directories {
-            let child_digest: [u8; 32] = child_directory_node
-                .digest
-                .as_slice()
-                .try_into()
-                .map_err(|_e| Error::StorageError("invalid digest length".to_string()))
-                .unwrap();
+            // TODO: propagate error
+            let child_digest = B3Digest::from_vec(child_directory_node.digest.clone()).unwrap();
 
             if self.worklist_directory_digests.contains(&child_digest)
                 || self.sent_directory_digests.contains(&child_digest)
@@ -59,8 +56,7 @@ impl<DS: DirectoryService> Iterator for DirectoryTraverser<DS> {
         match self.worklist_directory_digests.pop_front() {
             None => None,
             Some(current_directory_digest) => {
-                let current_directory_b64 = data_encoding::BASE64.encode(&current_directory_digest);
-                let span = debug_span!("directory.digest", current_directory_b64);
+                let span = debug_span!("directory.digest", "{}", current_directory_digest);
                 let _ = span.enter();
 
                 // look up the directory itself.
@@ -73,24 +69,24 @@ impl<DS: DirectoryService> Iterator for DirectoryTraverser<DS> {
                             warn!("directory failed validation: {}", e.to_string());
                             return Some(Err(Error::StorageError(format!(
                                 "invalid directory: {}",
-                                current_directory_b64
+                                current_directory_digest
                             ))));
                         }
                         current_directory
                     }
                     // if it's not there, we have an inconsistent store!
                     Ok(None) => {
-                        warn!("directory {} does not exist", current_directory_b64);
+                        warn!("directory {} does not exist", current_directory_digest);
                         return Some(Err(Error::StorageError(format!(
                             "directory {} does not exist",
-                            current_directory_b64
+                            current_directory_digest
                         ))));
                     }
                     Err(e) => {
                         warn!("failed to look up directory");
                         return Some(Err(Error::StorageError(format!(
                             "unable to look up directory {}: {}",
-                            current_directory_b64, e
+                            current_directory_digest, e
                         ))));
                     }
                 };
@@ -110,7 +106,7 @@ impl<DS: DirectoryService> Iterator for DirectoryTraverser<DS> {
 /// TODO: verify connectivity? Factor out these checks into generic helpers?
 pub struct SimplePutter<DS: DirectoryService> {
     directory_service: DS,
-    last_directory_digest: Option<[u8; 32]>,
+    last_directory_digest: Option<B3Digest>,
 }
 
 impl<DS: DirectoryService> SimplePutter<DS> {
@@ -133,9 +129,9 @@ impl<DS: DirectoryService> DirectoryPutter for SimplePutter<DS> {
     }
 
     /// We need to be mutable here, as that's the signature of the trait.
-    fn close(&mut self) -> Result<[u8; 32], Error> {
-        match self.last_directory_digest {
-            Some(last_digest) => Ok(last_digest),
+    fn close(&mut self) -> Result<B3Digest, Error> {
+        match &self.last_directory_digest {
+            Some(last_digest) => Ok(last_digest.clone()),
             None => Err(Error::InvalidRequest(
                 "no directories sent, can't show root digest".to_string(),
             )),
-- 
cgit 1.4.1