Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add new chunk count column for tracking #342

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

feat: Add new chunk count column for tracking #342

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
282084a
add new chunk count column for tracking
snowmead Jan 27, 2025
c9c22c7
refactor column definitions to have consistent methods retrieving the…
snowmead Jan 27, 2025
e2d9f2d
Merge branch 'main' into feat/chunk-count-column
snowmead Jan 30, 2025
c7d315b
Merge branch 'main' into feat/chunk-count-column
snowmead Feb 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ scale-info = { version = "2.11.0", default-features = false, features = [
serde = { version = "1.0.210", default-features = false }
serde_json = { version = "1.0.121", default-features = false }
smallvec = "1.11.0"
strum = { version = "0.26.3", features = ["derive"] }
thiserror = "1.0.48"
tokio = "1.36.0"
toml = "0.8.19"
Expand Down
1 change: 1 addition & 0 deletions client/file-manager/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ kvdb-rocksdb = { workspace = true }
kvdb-memorydb = { workspace = true }
log = { workspace = true }
serde_json = { workspace = true }
strum = { workspace = true }
thiserror = { workspace = true }
trie-db = { workspace = true }

Expand Down
118 changes: 64 additions & 54 deletions client/file-manager/src/in_memory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,16 +31,6 @@ impl<T: TrieLayout> FileDataTrie<T> for InMemoryFileDataTrie<T> {
&self.root
}

fn stored_chunks_count(&self) -> Result<u64, FileStorageError> {
let trie = TrieDBBuilder::<T>::new(&self.memdb, &self.root).build();
let trie_iter = trie
.iter()
.map_err(|_| FileStorageError::FailedToConstructTrieIter)?;
let stored_chunks = trie_iter.count() as u64;

Ok(stored_chunks)
}

fn generate_proof(&self, chunk_ids: &Vec<ChunkId>) -> Result<FileProof, FileStorageError> {
let recorder: Recorder<T::Hash> = Recorder::default();

Expand Down Expand Up @@ -152,6 +142,7 @@ where
pub file_data: HashMap<HasherOutT<T>, InMemoryFileDataTrie<T>>,
pub bucket_prefix_map: HashSet<[u8; 64]>,
pub exclude_list: HashMap<ExcludeType, HashSet<HasherOutT<T>>>,
pub chunk_counts: HashMap<HasherOutT<T>, u64>,
}

impl<T: TrieLayout> InMemoryFileStorage<T>
Expand All @@ -172,6 +163,7 @@ where
file_data: HashMap::new(),
bucket_prefix_map: HashSet::new(),
exclude_list,
chunk_counts: HashMap::new(),
}
}
}
Expand Down Expand Up @@ -204,7 +196,7 @@ where
.as_str(),
);

let stored_chunks = file_data.stored_chunks_count()?;
let stored_chunks = self.stored_chunks_count(file_key)?;
if metadata.chunks_count() != stored_chunks {
return Err(FileStorageError::IncompleteFile);
}
Expand All @@ -224,9 +216,17 @@ where
.to_file_key_proof(metadata.clone()))
}

fn stored_chunks_count(&self, key: &HasherOutT<T>) -> Result<u64, FileStorageError> {
self.chunk_counts
.get(key)
.copied()
.ok_or(FileStorageError::FileDoesNotExist)
}

fn delete_file(&mut self, key: &HasherOutT<T>) -> Result<(), FileStorageError> {
self.metadata.remove(key);
self.file_data.remove(key);
self.chunk_counts.remove(key);

Ok(())
}
Expand All @@ -252,22 +252,11 @@ where
.as_str(),
);

let stored_chunks = file_data.stored_chunks_count()?;
if metadata.chunks_count() != stored_chunks {
return Ok(false);
}

if metadata.fingerprint
!= file_data
.get_root()
.as_ref()
.try_into()
.expect("Hasher output mismatch!")
{
if metadata.fingerprint != file_data.get_root().as_ref().into() {
return Ok(false);
}

Ok(true)
Ok(metadata.chunks_count() == self.stored_chunks_count(key)?)
}

fn insert_file(
Expand All @@ -286,6 +275,9 @@ where
panic!("Key already associated with File Data, but not with File Metadata. Possible inconsistency between them.");
}

// Initialize chunk count to 0
self.chunk_counts.insert(key, 0);

let full_key = [metadata.bucket_id.as_slice(), key.as_ref()].concat();
self.bucket_prefix_map.insert(full_key.try_into().unwrap());

Expand All @@ -303,6 +295,15 @@ where
}
self.metadata.insert(key, metadata.clone());

// Count all chunks in the file trie
let trie = TrieDBBuilder::<T>::new(&file_data.memdb, &file_data.get_root()).build();
let chunk_count = trie
.iter()
.map_err(|_| FileStorageError::FailedToConstructTrieIter)?
.count();

self.chunk_counts.insert(key, chunk_count as u64);

let previous = self.file_data.insert(key, file_data);
if previous.is_some() {
panic!("Key already associated with File Data, but not with File Metadata. Possible inconsistency between them.");
Expand All @@ -314,13 +315,6 @@ where
Ok(())
}

fn stored_chunks_count(&self, key: &HasherOutT<T>) -> Result<u64, FileStorageError> {
let file_data = self.file_data.get(key);
let file_data = file_data.ok_or(FileStorageError::FileDoesNotExist)?;

file_data.stored_chunks_count()
}

fn get_chunk(
&self,
file_key: &HasherOutT<T>,
Expand Down Expand Up @@ -352,11 +346,13 @@ where
.as_str(),
);

// Check if we have all the chunks for the file.
let stored_chunks = file_data
.stored_chunks_count()
.map_err(|_| FileStorageWriteError::FailedToConstructTrieIter)?;
if metadata.chunks_count() != stored_chunks {
// Increment chunk count
let current_count = self.chunk_counts.get(file_key).copied().unwrap_or(0);
let new_count = current_count + 1;
self.chunk_counts.insert(*file_key, new_count);

// Check if we have all the chunks for the file using the count
if metadata.chunks_count() != new_count {
return Ok(FileStorageWriteOutcome::FileIncomplete);
}

Expand Down Expand Up @@ -393,6 +389,7 @@ where
for key in keys_to_delete {
self.metadata.remove(&key);
self.file_data.remove(&key);
self.chunk_counts.remove(&key);
}

Ok(())
Expand Down Expand Up @@ -451,6 +448,19 @@ mod tests {
use sp_runtime::AccountId32;
use sp_trie::LayoutV1;

fn stored_chunks_count(
file_trie: &InMemoryFileDataTrie<LayoutV1<BlakeTwo256>>,
) -> Result<u64, FileStorageError> {
let trie =
TrieDBBuilder::<LayoutV1<BlakeTwo256>>::new(&file_trie.memdb, &file_trie.root).build();
let trie_iter = trie
.iter()
.map_err(|_| FileStorageError::FailedToConstructTrieIter)?;
let stored_chunks = trie_iter.count() as u64;

Ok(stored_chunks)
}

#[test]
fn file_trie_create_empty_works() {
let file_trie = InMemoryFileDataTrie::<LayoutV1<BlakeTwo256>>::new();
Expand Down Expand Up @@ -501,11 +511,11 @@ mod tests {
let mut file_trie = InMemoryFileDataTrie::<LayoutV1<BlakeTwo256>>::new();

file_trie.write_chunk(&chunk_ids[0], &chunks[0]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 1);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 1);
assert!(file_trie.get_chunk(&chunk_ids[0]).is_ok());

file_trie.write_chunk(&chunk_ids[1], &chunks[1]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 2);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 2);
assert!(file_trie.get_chunk(&chunk_ids[1]).is_ok());
}

Expand All @@ -522,15 +532,15 @@ mod tests {
let mut file_trie = InMemoryFileDataTrie::<LayoutV1<BlakeTwo256>>::new();

file_trie.write_chunk(&chunk_ids[0], &chunks[0]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 1);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 1);
assert!(file_trie.get_chunk(&chunk_ids[0]).is_ok());

file_trie.write_chunk(&chunk_ids[1], &chunks[1]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 2);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 2);
assert!(file_trie.get_chunk(&chunk_ids[1]).is_ok());

file_trie.write_chunk(&chunk_ids[2], &chunks[2]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 3);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 3);
assert!(file_trie.get_chunk(&chunk_ids[2]).is_ok());

let file_proof = file_trie.generate_proof(&chunk_ids).unwrap();
Expand All @@ -554,23 +564,23 @@ mod tests {
let mut file_trie = InMemoryFileDataTrie::<LayoutV1<BlakeTwo256>>::new();

file_trie.write_chunk(&chunk_ids[0], &chunks[0]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 1);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 1);
assert!(file_trie.get_chunk(&chunk_ids[0]).is_ok());

file_trie.write_chunk(&chunk_ids[1], &chunks[1]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 2);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 2);
assert!(file_trie.get_chunk(&chunk_ids[1]).is_ok());

file_trie.write_chunk(&chunk_ids[2], &chunks[2]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 3);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 3);
assert!(file_trie.get_chunk(&chunk_ids[2]).is_ok());

file_trie.delete().unwrap();
assert!(file_trie.get_chunk(&chunk_ids[0]).is_err());
assert!(file_trie.get_chunk(&chunk_ids[1]).is_err());
assert!(file_trie.get_chunk(&chunk_ids[2]).is_err());

assert_eq!(file_trie.stored_chunks_count().unwrap(), 0);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 0);
}

#[test]
Expand All @@ -590,15 +600,15 @@ mod tests {
let mut file_trie = InMemoryFileDataTrie::<LayoutV1<BlakeTwo256>>::new();

file_trie.write_chunk(&chunk_ids[0], &chunks[0]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 1);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 1);
assert!(file_trie.get_chunk(&chunk_ids[0]).is_ok());

file_trie.write_chunk(&chunk_ids[1], &chunks[1]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 2);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 2);
assert!(file_trie.get_chunk(&chunk_ids[1]).is_ok());

file_trie.write_chunk(&chunk_ids[2], &chunks[2]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 3);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 3);
assert!(file_trie.get_chunk(&chunk_ids[2]).is_ok());

let file_metadata = FileMetadata {
Expand Down Expand Up @@ -637,15 +647,15 @@ mod tests {

let mut file_trie = InMemoryFileDataTrie::<LayoutV1<BlakeTwo256>>::new();
file_trie.write_chunk(&chunk_ids[0], &chunks[0]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 1);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 1);
assert!(file_trie.get_chunk(&chunk_ids[0]).is_ok());

file_trie.write_chunk(&chunk_ids[1], &chunks[1]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 2);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 2);
assert!(file_trie.get_chunk(&chunk_ids[1]).is_ok());

file_trie.write_chunk(&chunk_ids[2], &chunks[2]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 3);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 3);
assert!(file_trie.get_chunk(&chunk_ids[2]).is_ok());

let file_metadata = FileMetadata {
Expand Down Expand Up @@ -690,15 +700,15 @@ mod tests {

let mut file_trie = InMemoryFileDataTrie::<LayoutV1<BlakeTwo256>>::new();
file_trie.write_chunk(&chunk_ids[0], &chunks[0]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 1);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 1);
assert!(file_trie.get_chunk(&chunk_ids[0]).is_ok());

file_trie.write_chunk(&chunk_ids[1], &chunks[1]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 2);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 2);
assert!(file_trie.get_chunk(&chunk_ids[1]).is_ok());

file_trie.write_chunk(&chunk_ids[2], &chunks[2]).unwrap();
assert_eq!(file_trie.stored_chunks_count().unwrap(), 3);
assert_eq!(stored_chunks_count(&file_trie).unwrap(), 3);
assert!(file_trie.get_chunk(&chunk_ids[2]).is_ok());

let file_metadata = FileMetadata {
Expand Down
Loading
Loading