1use anyhow::Result;
2use futures::executor::block_on as sync_block_on;
3use hashtree_core::store::Store;
4use hashtree_core::{to_hex, types::Hash, Cid, HashTree, HashTreeConfig, HashTreeError, LinkType};
5use serde::de::{self, IgnoredAny, MapAccess, SeqAccess, Visitor};
6use serde::{Deserialize, Serialize};
7use std::collections::HashSet;
8use std::path::{Path, PathBuf};
9use std::time::{SystemTime, UNIX_EPOCH};
10
11use super::{BlobMetadata, HashtreeStore, PRIORITY_FOLLOWED, PRIORITY_OWN};
12
13const MAX_PINNED_TREE_NODES: usize = 10_000_000;
14const MAX_UNBOUNDED_PINNED_TREE_BYTES: u64 = 1 << 50;
15
16#[derive(Debug, Clone, Copy)]
18pub struct TreeIndexLimits {
19 pub max_nodes: usize,
20 pub max_bytes: u64,
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct PinTreeResult {
26 pub indexed_hashes: usize,
27 pub total_size: u64,
28 pub already_pinned: bool,
29}
30
31#[derive(Debug, thiserror::Error)]
32pub enum PinTreeError {
33 #[error("root blob {hash} is missing")]
34 MissingRoot { hash: String },
35 #[error("descendant blob {hash} is missing")]
36 MissingDescendant { hash: String },
37 #[error("invalid DAG node {hash}: {message}")]
38 InvalidDag { hash: String, message: String },
39 #[error("DAG exceeds the {max_nodes} node limit")]
40 NodeLimitExceeded { max_nodes: usize },
41 #[error("DAG exceeds the {max_bytes} byte limit")]
42 ByteLimitExceeded { max_bytes: u64 },
43 #[error("storage error: {0}")]
44 Storage(String),
45}
46
47struct TreeIndexPlan {
48 tracked_hashes: HashSet<Hash>,
49 total_size: u64,
50}
51
52#[derive(Debug, Clone, Serialize)]
54pub struct TreeMeta {
55 pub owner: String,
57 pub name: Option<String>,
59 pub synced_at: u64,
61 pub total_size: u64,
63 pub priority: u8,
65}
66
67impl<'de> Deserialize<'de> for TreeMeta {
68 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
69 where
70 D: serde::Deserializer<'de>,
71 {
72 const FIELDS: &[&str] = &[
73 "owner",
74 "name",
75 "synced_at",
76 "last_accessed_at",
77 "total_size",
78 "priority",
79 ];
80
81 struct TreeMetaVisitor;
82
83 impl<'de> Visitor<'de> for TreeMetaVisitor {
84 type Value = TreeMeta;
85
86 fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
87 formatter.write_str("TreeMeta as current or legacy metadata")
88 }
89
90 fn visit_seq<A>(self, mut seq: A) -> std::result::Result<Self::Value, A::Error>
91 where
92 A: SeqAccess<'de>,
93 {
94 let has_accidental_access_field = matches!(seq.size_hint(), Some(6));
95 let owner = seq
96 .next_element()?
97 .ok_or_else(|| de::Error::invalid_length(0, &self))?;
98 let name = seq
99 .next_element()?
100 .ok_or_else(|| de::Error::invalid_length(1, &self))?;
101 let synced_at = seq
102 .next_element()?
103 .ok_or_else(|| de::Error::invalid_length(2, &self))?;
104
105 if has_accidental_access_field {
106 let _: IgnoredAny = seq
107 .next_element()?
108 .ok_or_else(|| de::Error::invalid_length(3, &self))?;
109 }
110
111 let total_size = seq
112 .next_element()?
113 .ok_or_else(|| de::Error::invalid_length(3, &self))?;
114 let priority = seq
115 .next_element()?
116 .ok_or_else(|| de::Error::invalid_length(4, &self))?;
117
118 Ok(TreeMeta {
119 owner,
120 name,
121 synced_at,
122 total_size,
123 priority,
124 })
125 }
126
127 fn visit_map<A>(self, mut map: A) -> std::result::Result<Self::Value, A::Error>
128 where
129 A: MapAccess<'de>,
130 {
131 let mut owner = None;
132 let mut name = None;
133 let mut synced_at = None;
134 let mut total_size = None;
135 let mut priority = None;
136
137 while let Some(key) = map.next_key::<String>()? {
138 match key.as_str() {
139 "owner" => owner = Some(map.next_value()?),
140 "name" => name = Some(map.next_value()?),
141 "synced_at" => synced_at = Some(map.next_value()?),
142 "last_accessed_at" => {
143 let _: IgnoredAny = map.next_value()?;
144 }
145 "total_size" => total_size = Some(map.next_value()?),
146 "priority" => priority = Some(map.next_value()?),
147 _ => {
148 let _: IgnoredAny = map.next_value()?;
149 }
150 }
151 }
152
153 Ok(TreeMeta {
154 owner: owner.ok_or_else(|| de::Error::missing_field("owner"))?,
155 name: name.unwrap_or(None),
156 synced_at: synced_at.ok_or_else(|| de::Error::missing_field("synced_at"))?,
157 total_size: total_size.ok_or_else(|| de::Error::missing_field("total_size"))?,
158 priority: priority.ok_or_else(|| de::Error::missing_field("priority"))?,
159 })
160 }
161 }
162
163 deserializer.deserialize_struct("TreeMeta", FIELDS, TreeMetaVisitor)
164 }
165}
166
167#[derive(Debug)]
168pub struct StorageStats {
169 pub total_dags: usize,
170 pub pinned_dags: usize,
171 pub total_bytes: u64,
172}
173
174#[derive(Debug, Clone)]
176pub struct StorageByPriority {
177 pub own: u64,
179 pub followed: u64,
181 pub other: u64,
183}
184
185#[derive(Debug, Clone)]
186pub struct PinnedItem {
187 pub cid: String,
188 pub name: String,
189 pub is_directory: bool,
190 pub size_bytes: u64,
191}
192
193#[derive(Debug, Clone)]
194pub struct OwnedBlobStats {
195 pub owner: [u8; 32],
196 pub count: usize,
197 pub total_bytes: u64,
198}
199
200fn pinned_item_name(hash: &Hash, meta: Option<&TreeMeta>) -> String {
201 let Some(meta) = meta else {
202 return to_hex(hash);
203 };
204
205 match (meta.owner.as_str(), meta.name.as_deref()) {
206 ("pinned", Some(name)) => name.to_string(),
207 ("", Some(name)) => name.to_string(),
208 (owner, Some(name)) if !owner.is_empty() => format!("{owner}/{name}"),
209 (owner, None) if !owner.is_empty() && owner != "pinned" => owner.to_string(),
210 _ => to_hex(hash),
211 }
212}
213
214fn unix_timestamp_now() -> u64 {
215 SystemTime::now()
216 .duration_since(UNIX_EPOCH)
217 .unwrap_or_default()
218 .as_secs()
219}
220
221impl HashtreeStore {
222 fn socialgraph_root_files(&self) -> [PathBuf; 4] {
223 let socialgraph = self.base_path().join("socialgraph");
224 [
225 socialgraph.join("events-root.msgpack"),
226 socialgraph.join("events-root-ambient.msgpack"),
227 socialgraph.join("profile-search-root.msgpack"),
228 socialgraph.join("profiles-by-pubkey-root.msgpack"),
229 ]
230 }
231
232 fn read_stored_cid(path: &Path) -> Result<Option<Hash>> {
233 #[derive(Deserialize)]
234 struct StoredCid {
235 hash: [u8; 32],
236 #[allow(dead_code)]
237 key: Option<[u8; 32]>,
238 }
239
240 let Ok(bytes) = std::fs::read(path) else {
241 return Ok(None);
242 };
243 let stored: StoredCid = rmp_serde::from_slice(&bytes)
244 .map_err(|e| anyhow::anyhow!("Failed to decode root file {}: {}", path.display(), e))?;
245 Ok(Some(stored.hash))
246 }
247
248 async fn collect_tree_hashes<S: Store>(
249 &self,
250 tree: &HashTree<S>,
251 root: &Hash,
252 ) -> Result<HashSet<Hash>> {
253 let mut hashes = HashSet::new();
254 let mut stack = vec![*root];
255
256 while let Some(hash) = stack.pop() {
257 if !hashes.insert(hash) {
258 continue;
259 }
260
261 let is_tree = tree
262 .is_tree(&hash)
263 .await
264 .map_err(|e| anyhow::anyhow!("Failed to check tree: {}", e))?;
265
266 if !is_tree {
267 continue;
268 }
269
270 if let Some(node) = tree
271 .get_tree_node(&hash)
272 .await
273 .map_err(|e| anyhow::anyhow!("Failed to get tree node: {}", e))?
274 {
275 for link in &node.links {
276 stack.push(link.hash);
277 }
278 }
279 }
280
281 Ok(hashes)
282 }
283
284 fn protected_hashes(&self) -> Result<HashSet<Hash>> {
285 let mut protected = HashSet::new();
286
287 let rtxn = self.env.read_txn()?;
288 for (key_bytes, _) in self.blob_trees.iter(&rtxn)?.flatten() {
289 if key_bytes.len() >= 32 {
290 let hash: Hash = key_bytes[..32].try_into().unwrap();
291 protected.insert(hash);
292 }
293 }
294 drop(rtxn);
295
296 let tree = HashTree::new(HashTreeConfig::new(self.store_arc()).public());
297 for path in self.socialgraph_root_files() {
298 let Some(root_hash) = Self::read_stored_cid(&path)? else {
299 continue;
300 };
301 protected.extend(sync_block_on(self.collect_tree_hashes(&tree, &root_hash))?);
302 }
303
304 Ok(protected)
305 }
306
307 fn evict_disposable_orphans_to_target(&self, target_bytes: u64) -> Result<u64> {
308 let stats = self
309 .router
310 .writable_stats()
311 .map_err(|e| anyhow::anyhow!("Failed to get writable stats: {}", e))?;
312 let mut current_size = stats.total_bytes;
313 if current_size <= target_bytes {
314 return Ok(0);
315 }
316
317 let rtxn = self.env.read_txn()?;
318 let pinned: HashSet<Hash> = self
319 .pins
320 .iter(&rtxn)?
321 .filter_map(|item| item.ok())
322 .filter_map(|(hash_bytes, _)| {
323 if hash_bytes.len() == 32 {
324 let mut hash = [0u8; 32];
325 hash.copy_from_slice(hash_bytes);
326 Some(hash)
327 } else {
328 None
329 }
330 })
331 .collect();
332 drop(rtxn);
333
334 let protected_hashes = self.protected_hashes()?;
335 let all_hashes = self
336 .router
337 .list_writable()
338 .map_err(|e| anyhow::anyhow!("Failed to list writable hashes: {}", e))?;
339
340 let mut freed = 0u64;
341 for hash in all_hashes {
342 if current_size <= target_bytes {
343 break;
344 }
345
346 if pinned.contains(&hash) || protected_hashes.contains(&hash) {
347 continue;
348 }
349
350 if self.blob_has_owners(&hash)? {
351 continue;
352 }
353
354 let Some(size) = self
355 .router
356 .blob_size_sync(&hash)
357 .map_err(|e| anyhow::anyhow!("Failed to get blob size: {}", e))?
358 else {
359 continue;
360 };
361
362 if self
363 .router
364 .delete_local_only(&hash)
365 .map_err(|e| anyhow::anyhow!("Failed to delete orphaned blob: {}", e))?
366 {
367 freed = freed.saturating_add(size);
368 current_size = current_size.saturating_sub(size);
369 tracing::debug!(
370 "Deleted disposable orphaned blob {} ({} bytes)",
371 &to_hex(&hash)[..8],
372 size
373 );
374 }
375 }
376
377 Ok(freed)
378 }
379
380 pub fn make_room_for_cached_blob(&self, incoming_bytes: u64) -> Result<u64> {
381 if self.max_size_bytes == 0 {
382 return Ok(0);
383 }
384
385 let stats = self
386 .router
387 .writable_stats()
388 .map_err(|e| anyhow::anyhow!("Failed to get writable stats: {}", e))?;
389 if stats.total_bytes.saturating_add(incoming_bytes) <= self.max_size_bytes {
390 return Ok(0);
391 }
392
393 let target = if incoming_bytes >= self.max_size_bytes {
394 0
395 } else {
396 (self.max_size_bytes.saturating_mul(9) / 10)
397 .min(self.max_size_bytes.saturating_sub(incoming_bytes))
398 };
399 self.evict_disposable_orphans_to_target(target)
400 }
401
402 pub fn enforce_cached_blob_budget_after_insert(&self, inserted_bytes: u64) -> Result<u64> {
403 if self.max_size_bytes == 0 || inserted_bytes == 0 {
404 return Ok(0);
405 }
406
407 let stats = self
408 .router
409 .writable_stats()
410 .map_err(|e| anyhow::anyhow!("Failed to get writable stats: {}", e))?;
411 if stats.total_bytes <= self.max_size_bytes {
412 return Ok(0);
413 }
414
415 let target = if inserted_bytes >= self.max_size_bytes {
416 inserted_bytes
417 } else {
418 (self.max_size_bytes.saturating_mul(9) / 10)
419 .saturating_add(inserted_bytes)
420 .min(self.max_size_bytes)
421 };
422 self.evict_disposable_orphans_to_target(target)
423 }
424
425 pub fn make_room_for_durable_blob(&self, incoming_bytes: u64) -> Result<u64> {
426 if self.max_size_bytes == 0 || incoming_bytes == 0 {
427 return Ok(0);
428 }
429
430 if incoming_bytes > self.max_size_bytes {
431 anyhow::bail!(
432 "storage limit exceeded: incoming blob is {} bytes but limit is {} bytes",
433 incoming_bytes,
434 self.max_size_bytes
435 );
436 }
437
438 let stats = self
439 .router
440 .writable_stats()
441 .map_err(|e| anyhow::anyhow!("Failed to get writable stats: {}", e))?;
442 if stats.total_bytes.saturating_add(incoming_bytes) <= self.max_size_bytes {
443 return Ok(0);
444 }
445
446 let target = (self.max_size_bytes.saturating_mul(9) / 10)
447 .min(self.max_size_bytes.saturating_sub(incoming_bytes));
448 let freed = self.evict_with_policy_to_target(stats.total_bytes, target)?;
449
450 let next_stats = self
451 .router
452 .writable_stats()
453 .map_err(|e| anyhow::anyhow!("Failed to get writable stats after eviction: {}", e))?;
454 if next_stats.total_bytes.saturating_add(incoming_bytes) > self.max_size_bytes {
455 anyhow::bail!(
456 "storage limit exceeded: {} bytes used, {} byte incoming blob, {} byte limit",
457 next_stats.total_bytes,
458 incoming_bytes,
459 self.max_size_bytes
460 );
461 }
462
463 Ok(freed)
464 }
465
466 pub fn enforce_durable_blob_budget_after_insert(&self, inserted_bytes: u64) -> Result<u64> {
467 if self.max_size_bytes == 0 || inserted_bytes == 0 {
468 return Ok(0);
469 }
470
471 if inserted_bytes > self.max_size_bytes {
472 anyhow::bail!(
473 "storage limit exceeded: inserted blobs are {} bytes but limit is {} bytes",
474 inserted_bytes,
475 self.max_size_bytes
476 );
477 }
478
479 let stats = self
480 .router
481 .writable_stats()
482 .map_err(|e| anyhow::anyhow!("Failed to get writable stats: {}", e))?;
483 if stats.total_bytes <= self.max_size_bytes {
484 return Ok(0);
485 }
486
487 let target = (self.max_size_bytes.saturating_mul(9) / 10)
488 .saturating_add(inserted_bytes)
489 .min(self.max_size_bytes);
490 let freed = self.evict_with_policy_to_target(stats.total_bytes, target)?;
491
492 let next_stats = self
493 .router
494 .writable_stats()
495 .map_err(|e| anyhow::anyhow!("Failed to get writable stats after eviction: {}", e))?;
496 if next_stats.total_bytes > self.max_size_bytes {
497 anyhow::bail!(
498 "storage limit exceeded: {} bytes used after inserting {} bytes, {} byte limit",
499 next_stats.total_bytes,
500 inserted_bytes,
501 self.max_size_bytes
502 );
503 }
504
505 Ok(freed)
506 }
507
508 pub fn relieve_cached_blob_write_pressure(&self, incoming_bytes: u64) -> Result<u64> {
509 let stats = self
510 .router
511 .writable_stats()
512 .map_err(|e| anyhow::anyhow!("Failed to get writable stats: {}", e))?;
513 if stats.total_bytes == 0 {
514 return Ok(0);
515 }
516
517 let headroom = incoming_bytes.max(stats.total_bytes / 10).max(1);
518 let target = stats.total_bytes.saturating_sub(headroom);
519 self.evict_disposable_orphans_to_target(target)
520 }
521
522 pub fn pin(&self, hash: &[u8; 32]) -> Result<()> {
524 let mut wtxn = self.env.write_txn()?;
525 self.pins.put(&mut wtxn, hash.as_slice(), &())?;
526 wtxn.commit()?;
527 Ok(())
528 }
529
530 pub fn unpin(&self, hash: &[u8; 32]) -> Result<()> {
532 let mut wtxn = self.env.write_txn()?;
533 self.pins.delete(&mut wtxn, hash.as_slice())?;
534 wtxn.commit()?;
535 Ok(())
536 }
537
538 pub fn is_pinned(&self, hash: &[u8; 32]) -> Result<bool> {
540 let rtxn = self.env.read_txn()?;
541 Ok(self.pins.get(&rtxn, hash.as_slice())?.is_some())
542 }
543
544 pub fn list_pins_raw(&self) -> Result<Vec<[u8; 32]>> {
546 let rtxn = self.env.read_txn()?;
547 let mut pins = Vec::new();
548
549 for item in self.pins.iter(&rtxn)? {
550 let (hash_bytes, _) = item?;
551 if hash_bytes.len() == 32 {
552 let mut hash = [0u8; 32];
553 hash.copy_from_slice(hash_bytes);
554 pins.push(hash);
555 }
556 }
557
558 Ok(pins)
559 }
560
561 pub fn list_pins_with_names(&self) -> Result<Vec<PinnedItem>> {
563 let rtxn = self.env.read_txn()?;
564 let store = self.store_arc();
565 let tree = HashTree::new(HashTreeConfig::new(store).public());
566 let mut pins = Vec::new();
567
568 for item in self.pins.iter(&rtxn)? {
569 let (hash_bytes, _) = item?;
570 if hash_bytes.len() != 32 {
571 continue;
572 }
573 let mut hash = [0u8; 32];
574 hash.copy_from_slice(hash_bytes);
575
576 let is_directory =
578 sync_block_on(async { tree.is_directory(&hash).await.unwrap_or(false) });
579
580 let meta = self
581 .tree_meta
582 .get(&rtxn, hash.as_slice())?
583 .map(|bytes| {
584 rmp_serde::from_slice::<TreeMeta>(bytes)
585 .map_err(|e| anyhow::anyhow!("Failed to deserialize TreeMeta: {}", e))
586 })
587 .transpose()?;
588 let size_bytes = if let Some(meta) = meta.as_ref() {
589 meta.total_size
590 } else {
591 self.router
592 .blob_size_sync(&hash)
593 .map_err(|e| anyhow::anyhow!("Failed to get pinned blob size: {}", e))?
594 .unwrap_or(0)
595 };
596
597 pins.push(PinnedItem {
598 cid: to_hex(&hash),
599 name: pinned_item_name(&hash, meta.as_ref()),
600 is_directory,
601 size_bytes,
602 });
603 }
604
605 Ok(pins)
606 }
607
608 pub fn owned_blob_stats(&self) -> Result<Vec<OwnedBlobStats>> {
609 let rtxn = self.env.read_txn()?;
610 let mut owners = Vec::new();
611
612 for item in self.pubkey_blobs.iter(&rtxn)? {
613 let (owner_bytes, blobs_bytes) = item?;
614 if owner_bytes.len() != 32 {
615 continue;
616 }
617
618 let blobs: Vec<BlobMetadata> = serde_json::from_slice(blobs_bytes)
619 .map_err(|e| anyhow::anyhow!("Failed to deserialize blob metadata: {}", e))?;
620 let mut owner = [0u8; 32];
621 owner.copy_from_slice(owner_bytes);
622 let total_bytes = blobs
623 .iter()
624 .fold(0u64, |total, blob| total.saturating_add(blob.size));
625 owners.push(OwnedBlobStats {
626 owner,
627 count: blobs.len(),
628 total_bytes,
629 });
630 }
631
632 owners.sort_by_key(|stats| stats.owner);
633 Ok(owners)
634 }
635
636 pub fn tree_index_limits(&self) -> TreeIndexLimits {
642 TreeIndexLimits {
643 max_nodes: MAX_PINNED_TREE_NODES,
644 max_bytes: if self.max_size_bytes == 0 {
645 MAX_UNBOUNDED_PINNED_TREE_BYTES
646 } else {
647 self.max_size_bytes
648 },
649 }
650 }
651
652 pub fn pin_and_index_tree(
656 &self,
657 root: &Cid,
658 owner: &str,
659 name: Option<&str>,
660 priority: u8,
661 limits: TreeIndexLimits,
662 ) -> std::result::Result<PinTreeResult, PinTreeError> {
663 let store = self.store_arc();
664 let tree = HashTree::new(HashTreeConfig::new(store).public());
665 let plan = sync_block_on(self.collect_tree_index(&tree, root, limits))?;
666 let already_pinned = self
667 .write_tree_index(
668 &root.hash,
669 &plan.tracked_hashes,
670 plan.total_size,
671 owner,
672 name,
673 priority,
674 None,
675 true,
676 )
677 .map_err(|error| PinTreeError::Storage(error.to_string()))?;
678
679 Ok(PinTreeResult {
680 indexed_hashes: plan.tracked_hashes.len(),
681 total_size: plan.total_size,
682 already_pinned,
683 })
684 }
685
686 pub fn index_tree(
691 &self,
692 root_hash: &Hash,
693 owner: &str,
694 name: Option<&str>,
695 priority: u8,
696 ref_key: Option<&str>,
697 ) -> Result<()> {
698 let root_hex = to_hex(root_hash);
699
700 if let Some(key) = ref_key {
702 let rtxn = self.env.read_txn()?;
703 if let Some(old_hash_bytes) = self.tree_refs.get(&rtxn, key)? {
704 if old_hash_bytes != root_hash.as_slice() {
705 let old_hash: Hash = old_hash_bytes
706 .try_into()
707 .map_err(|_| anyhow::anyhow!("Invalid hash in tree_refs"))?;
708 drop(rtxn);
709 let _ = self.unpin(&old_hash);
710 let _ = self.unindex_tree(&old_hash);
712 tracing::debug!("Replaced old tree for ref {}", key);
713 }
714 }
715 }
716
717 let store = self.store_arc();
718 let tree = HashTree::new(HashTreeConfig::new(store).public());
719
720 let plan = sync_block_on(self.collect_tree_index(
721 &tree,
722 &Cid::public(*root_hash),
723 TreeIndexLimits {
724 max_nodes: MAX_PINNED_TREE_NODES,
725 max_bytes: MAX_UNBOUNDED_PINNED_TREE_BYTES,
726 },
727 ))?;
728 self.write_tree_index(
729 root_hash,
730 &plan.tracked_hashes,
731 plan.total_size,
732 owner,
733 name,
734 priority,
735 ref_key,
736 false,
737 )?;
738
739 tracing::debug!(
740 "Indexed tree {} ({} blobs, {} bytes, priority {})",
741 &root_hex[..8],
742 plan.tracked_hashes.len(),
743 plan.total_size,
744 priority
745 );
746
747 Ok(())
748 }
749
750 #[allow(clippy::too_many_arguments)]
751 fn write_tree_index(
752 &self,
753 root_hash: &Hash,
754 tracked_hashes: &HashSet<Hash>,
755 total_size: u64,
756 owner: &str,
757 name: Option<&str>,
758 priority: u8,
759 ref_key: Option<&str>,
760 pin: bool,
761 ) -> Result<bool> {
762 let mut wtxn = self.env.write_txn()?;
763 let already_pinned = self.pins.get(&wtxn, root_hash.as_slice())?.is_some();
764
765 for tracked_hash in tracked_hashes {
766 let mut key = [0u8; 64];
767 key[..32].copy_from_slice(tracked_hash);
768 key[32..].copy_from_slice(root_hash);
769 self.blob_trees.put(&mut wtxn, &key[..], &())?;
770 }
771
772 let meta = TreeMeta {
773 owner: owner.to_string(),
774 name: name.map(str::to_string),
775 synced_at: unix_timestamp_now(),
776 total_size,
777 priority,
778 };
779 let meta_bytes = rmp_serde::to_vec(&meta)
780 .map_err(|error| anyhow::anyhow!("Failed to serialize TreeMeta: {error}"))?;
781 self.tree_meta
782 .put(&mut wtxn, root_hash.as_slice(), &meta_bytes)?;
783
784 if let Some(key) = ref_key {
785 self.tree_refs.put(&mut wtxn, key, root_hash.as_slice())?;
786 }
787 if pin {
788 self.pins.put(&mut wtxn, root_hash.as_slice(), &())?;
789 }
790
791 wtxn.commit()?;
792 Ok(already_pinned)
793 }
794
795 async fn collect_tree_index<S: Store>(
796 &self,
797 tree: &HashTree<S>,
798 root: &Cid,
799 limits: TreeIndexLimits,
800 ) -> std::result::Result<TreeIndexPlan, PinTreeError> {
801 let mut hashes = HashSet::new();
802 let mut visited = HashSet::new();
803 let mut total_size = 0u64;
804 let mut stored_size = 0u64;
805 let mut stack = vec![(root.clone(), true, true, false)];
807
808 while let Some((cid, count_bytes, follow_tree, require_tree)) = stack.pop() {
809 let visit_key = (cid.hash, cid.key, follow_tree);
810 if !visited.insert(visit_key) {
811 continue;
812 }
813 if visited.len() > limits.max_nodes {
814 return Err(PinTreeError::NodeLimitExceeded {
815 max_nodes: limits.max_nodes,
816 });
817 }
818
819 let size = self
820 .router
821 .blob_size_sync(&cid.hash)
822 .map_err(|error| PinTreeError::Storage(error.to_string()))?
823 .ok_or_else(|| {
824 if cid.hash == root.hash {
825 PinTreeError::MissingRoot {
826 hash: to_hex(&cid.hash),
827 }
828 } else {
829 PinTreeError::MissingDescendant {
830 hash: to_hex(&cid.hash),
831 }
832 }
833 })?;
834 if hashes.insert(cid.hash) {
835 stored_size = stored_size
836 .checked_add(size)
837 .filter(|size| *size <= limits.max_bytes)
838 .ok_or(PinTreeError::ByteLimitExceeded {
839 max_bytes: limits.max_bytes,
840 })?;
841 }
842
843 if !follow_tree {
844 continue;
845 }
846
847 let node = tree.get_node(&cid).await.map_err(|error| match error {
848 HashTreeError::Store(message) => PinTreeError::Storage(message),
849 error => PinTreeError::InvalidDag {
850 hash: to_hex(&cid.hash),
851 message: error.to_string(),
852 },
853 })?;
854 let Some(node) = node else {
855 if require_tree {
856 return Err(PinTreeError::InvalidDag {
857 hash: to_hex(&cid.hash),
858 message: "directory link does not contain a tree node".to_string(),
859 });
860 }
861 if count_bytes {
862 total_size = total_size
863 .checked_add(size)
864 .filter(|size| *size <= limits.max_bytes)
865 .ok_or(PinTreeError::ByteLimitExceeded {
866 max_bytes: limits.max_bytes,
867 })?;
868 }
869 continue;
870 };
871
872 if visited
873 .len()
874 .saturating_add(stack.len())
875 .saturating_add(node.links.len())
876 > limits.max_nodes
877 {
878 return Err(PinTreeError::NodeLimitExceeded {
879 max_nodes: limits.max_nodes,
880 });
881 }
882
883 for link in &node.links {
884 match link.link_type {
885 LinkType::Blob => {
886 if count_bytes {
887 total_size = total_size
888 .checked_add(link.size)
889 .filter(|size| *size <= limits.max_bytes)
890 .ok_or(PinTreeError::ByteLimitExceeded {
891 max_bytes: limits.max_bytes,
892 })?;
893 }
894 stack.push((link.to_cid(), false, false, false));
895 }
896 LinkType::File => {
897 if count_bytes {
898 total_size = total_size
899 .checked_add(link.size)
900 .filter(|size| *size <= limits.max_bytes)
901 .ok_or(PinTreeError::ByteLimitExceeded {
902 max_bytes: limits.max_bytes,
903 })?;
904 }
905 stack.push((link.to_cid(), false, true, false));
906 }
907 LinkType::Dir | LinkType::Fanout => {
908 stack.push((link.to_cid(), count_bytes, true, true));
909 }
910 }
911 }
912 }
913
914 Ok(TreeIndexPlan {
915 tracked_hashes: hashes,
916 total_size,
917 })
918 }
919
920 pub fn unindex_tree(&self, root_hash: &Hash) -> Result<u64> {
923 let root_hex = to_hex(root_hash);
924
925 let store = self.store_arc();
926 let tree = HashTree::new(HashTreeConfig::new(store).public());
927
928 let tracked_hashes = sync_block_on(self.collect_tree_hashes(&tree, root_hash))?;
930
931 let mut wtxn = self.env.write_txn()?;
932 let mut freed = 0u64;
933
934 for tracked_hash in &tracked_hashes {
936 let mut key = [0u8; 64];
938 key[..32].copy_from_slice(tracked_hash);
939 key[32..].copy_from_slice(root_hash);
940 self.blob_trees.delete(&mut wtxn, &key[..])?;
941
942 let mut has_other_tree = false;
944 for item in self.blob_trees.prefix_iter(&wtxn, &tracked_hash[..])? {
945 if item.is_ok() {
946 has_other_tree = true;
947 break;
948 }
949 }
950
951 if !has_other_tree {
953 if let Some(size) = self
954 .router
955 .blob_size_sync(tracked_hash)
956 .map_err(|e| anyhow::anyhow!("Failed to get blob size: {}", e))?
957 {
958 freed += size;
959 self.router
961 .delete_local_only(tracked_hash)
962 .map_err(|e| anyhow::anyhow!("Failed to delete blob: {}", e))?;
963 }
964 }
965 }
966
967 self.tree_meta.delete(&mut wtxn, root_hash.as_slice())?;
969
970 wtxn.commit()?;
971
972 tracing::debug!("Unindexed tree {} ({} bytes freed)", &root_hex[..8], freed);
973
974 Ok(freed)
975 }
976
977 pub fn get_tree_meta(&self, root_hash: &Hash) -> Result<Option<TreeMeta>> {
979 let rtxn = self.env.read_txn()?;
980 if let Some(bytes) = self.tree_meta.get(&rtxn, root_hash.as_slice())? {
981 let meta: TreeMeta = rmp_serde::from_slice(bytes)
982 .map_err(|e| anyhow::anyhow!("Failed to deserialize TreeMeta: {}", e))?;
983 Ok(Some(meta))
984 } else {
985 Ok(None)
986 }
987 }
988
989 pub fn get_tree_ref(&self, key: &str) -> Result<Option<Hash>> {
990 let rtxn = self.env.read_txn()?;
991 let Some(bytes) = self.tree_refs.get(&rtxn, key)? else {
992 return Ok(None);
993 };
994
995 let hash: Hash = bytes
996 .try_into()
997 .map_err(|_| anyhow::anyhow!("Invalid hash in tree_refs"))?;
998 Ok(Some(hash))
999 }
1000
1001 pub fn list_indexed_trees(&self) -> Result<Vec<(Hash, TreeMeta)>> {
1003 let rtxn = self.env.read_txn()?;
1004 let mut trees = Vec::new();
1005
1006 for item in self.tree_meta.iter(&rtxn)? {
1007 let (hash_bytes, meta_bytes) = item?;
1008 let hash: Hash = hash_bytes
1009 .try_into()
1010 .map_err(|_| anyhow::anyhow!("Invalid hash in tree_meta"))?;
1011 let meta: TreeMeta = rmp_serde::from_slice(meta_bytes)
1012 .map_err(|e| anyhow::anyhow!("Failed to deserialize TreeMeta: {}", e))?;
1013 trees.push((hash, meta));
1014 }
1015
1016 Ok(trees)
1017 }
1018
1019 pub fn tracked_size(&self) -> Result<u64> {
1021 let rtxn = self.env.read_txn()?;
1022 let mut total = 0u64;
1023
1024 for item in self.tree_meta.iter(&rtxn)? {
1025 let (_, bytes) = item?;
1026 let meta: TreeMeta = rmp_serde::from_slice(bytes)
1027 .map_err(|e| anyhow::anyhow!("Failed to deserialize TreeMeta: {}", e))?;
1028 total += meta.total_size;
1029 }
1030
1031 Ok(total)
1032 }
1033
1034 fn get_evictable_trees(&self) -> Result<Vec<(Hash, TreeMeta)>> {
1040 let mut trees = self.list_indexed_trees()?;
1041
1042 trees.sort_by(|a, b| match a.1.priority.cmp(&b.1.priority) {
1044 std::cmp::Ordering::Equal => a.1.synced_at.cmp(&b.1.synced_at),
1045 other => other,
1046 });
1047
1048 Ok(trees)
1049 }
1050
1051 pub fn evict_if_needed(&self) -> Result<u64> {
1058 let stats = self
1060 .router
1061 .writable_stats()
1062 .map_err(|e| anyhow::anyhow!("Failed to get writable stats: {}", e))?;
1063 let current = stats.total_bytes;
1064
1065 if current <= self.max_size_bytes {
1066 return Ok(0);
1067 }
1068
1069 let target = self.max_size_bytes * 90 / 100;
1071 self.evict_with_policy_to_target(current, target)
1072 }
1073
1074 fn evict_with_policy_to_target(&self, current: u64, target: u64) -> Result<u64> {
1075 let mut freed = 0u64;
1076 let mut current_size = current;
1077
1078 if self.evict_orphans {
1080 let orphan_freed = self.evict_disposable_orphans_to_target(target)?;
1081 freed += orphan_freed;
1082 current_size = current_size.saturating_sub(orphan_freed);
1083
1084 if orphan_freed > 0 {
1085 tracing::info!("Evicted orphaned blobs: {} bytes freed", orphan_freed);
1086 }
1087 } else {
1088 tracing::debug!("Skipping orphan blob eviction; storage.evict_orphans=false");
1089 }
1090
1091 if current_size <= target {
1093 if freed > 0 {
1094 tracing::info!("Eviction complete: {} bytes freed", freed);
1095 }
1096 return Ok(freed);
1097 }
1098
1099 let evictable = self.get_evictable_trees()?;
1102
1103 for (root_hash, meta) in evictable {
1104 if current_size <= target {
1105 break;
1106 }
1107
1108 let root_hex = to_hex(&root_hash);
1109
1110 if self.is_pinned(&root_hash)? {
1112 continue;
1113 }
1114
1115 let tree_freed = self.unindex_tree(&root_hash)?;
1116 freed += tree_freed;
1117 current_size = current_size.saturating_sub(tree_freed);
1118
1119 tracing::info!(
1120 "Evicted tree {} (owner={}, priority={}, {} bytes)",
1121 &root_hex[..8],
1122 &meta.owner[..8.min(meta.owner.len())],
1123 meta.priority,
1124 tree_freed
1125 );
1126 }
1127
1128 if freed > 0 {
1129 tracing::info!("Eviction complete: {} bytes freed", freed);
1130 }
1131
1132 Ok(freed)
1133 }
1134
1135 pub fn max_size_bytes(&self) -> u64 {
1137 self.max_size_bytes
1138 }
1139
1140 pub fn storage_by_priority(&self) -> Result<StorageByPriority> {
1142 let rtxn = self.env.read_txn()?;
1143 let mut own = 0u64;
1144 let mut followed = 0u64;
1145 let mut other = 0u64;
1146
1147 for item in self.tree_meta.iter(&rtxn)? {
1148 let (_, bytes) = item?;
1149 let meta: TreeMeta = rmp_serde::from_slice(bytes)
1150 .map_err(|e| anyhow::anyhow!("Failed to deserialize TreeMeta: {}", e))?;
1151
1152 if meta.priority == PRIORITY_OWN {
1153 own += meta.total_size;
1154 } else if meta.priority >= PRIORITY_FOLLOWED {
1155 followed += meta.total_size;
1156 } else {
1157 other += meta.total_size;
1158 }
1159 }
1160
1161 Ok(StorageByPriority {
1162 own,
1163 followed,
1164 other,
1165 })
1166 }
1167
1168 pub fn get_storage_stats(&self) -> Result<StorageStats> {
1170 let rtxn = self.env.read_txn()?;
1171 let total_pins = self.pins.len(&rtxn)? as usize;
1172
1173 let stats = self
1174 .router
1175 .stats()
1176 .map_err(|e| anyhow::anyhow!("Failed to get stats: {}", e))?;
1177
1178 Ok(StorageStats {
1179 total_dags: stats.count,
1180 pinned_dags: total_pins,
1181 total_bytes: stats.total_bytes,
1182 })
1183 }
1184}
1185
1186#[cfg(test)]
1187mod tests {
1188 use super::*;
1189 use hashtree_core::Cid;
1190 use hashtree_index::{BTree, BTreeOptions};
1191 use tempfile::TempDir;
1192
1193 use crate::storage::PRIORITY_OTHER;
1194
1195 fn write_root_file(path: &Path, cid: &Cid) {
1196 #[derive(Serialize)]
1197 struct StoredCid {
1198 hash: [u8; 32],
1199 key: Option<[u8; 32]>,
1200 }
1201
1202 std::fs::create_dir_all(path.parent().expect("root file parent")).expect("create dir");
1203 let bytes = rmp_serde::to_vec_named(&StoredCid {
1204 hash: cid.hash,
1205 key: cid.key,
1206 })
1207 .expect("encode cid");
1208 std::fs::write(path, bytes).expect("write root file");
1209 }
1210
1211 fn build_test_tree(store: &HashtreeStore) -> Cid {
1212 let index = BTree::new(store.store_arc(), BTreeOptions { order: Some(8) });
1213 sync_block_on(index.build(vec![
1214 ("alpha".to_string(), "one".to_string()),
1215 ("beta".to_string(), "two".to_string()),
1216 ("gamma".to_string(), "three".to_string()),
1217 ]))
1218 .expect("build btree")
1219 .expect("non-empty root")
1220 }
1221
1222 #[test]
1223 fn orphan_cleanup_keeps_indexed_tree_hashes() {
1224 let temp_dir = TempDir::new().expect("temp dir");
1225 let store = HashtreeStore::with_options(temp_dir.path(), None, 1024).expect("store");
1226 let cid = build_test_tree(&store);
1227
1228 store
1229 .index_tree(
1230 &cid.hash,
1231 "owner",
1232 Some("tree"),
1233 PRIORITY_OTHER,
1234 Some("owner/tree"),
1235 )
1236 .expect("index tree");
1237 let freed = store
1238 .evict_disposable_orphans_to_target(0)
1239 .expect("orphan cleanup");
1240
1241 assert!(freed < 1024);
1242 assert!(store.blob_exists(&cid.hash).expect("root exists"));
1243 }
1244
1245 #[test]
1246 fn list_pins_with_names_uses_indexed_tree_metadata() {
1247 let temp_dir = TempDir::new().expect("temp dir");
1248 let store = HashtreeStore::with_options(temp_dir.path(), None, 1024 * 1024).expect("store");
1249 let cid = build_test_tree(&store);
1250
1251 store.pin(&cid.hash).expect("pin tree");
1252 store
1253 .index_tree(
1254 &cid.hash,
1255 "npub1example",
1256 Some("playlist"),
1257 PRIORITY_OTHER,
1258 Some("npub1example/playlist"),
1259 )
1260 .expect("index tree");
1261
1262 let pins = store.list_pins_with_names().expect("list pins");
1263
1264 assert_eq!(pins.len(), 1);
1265 assert_eq!(pins[0].name, "npub1example/playlist");
1266 assert!(pins[0].size_bytes > 0);
1267 }
1268
1269 #[test]
1270 fn index_tree_records_multilevel_file_size_from_links() {
1271 let temp_dir = TempDir::new().expect("temp dir");
1272 let store = HashtreeStore::with_options(temp_dir.path(), None, 1024 * 1024).expect("store");
1273 let tree = HashTree::new(
1274 HashTreeConfig::new(store.store_arc())
1275 .public()
1276 .with_chunk_size(4)
1277 .with_max_links(2),
1278 );
1279 let data = (0u8..31).collect::<Vec<_>>();
1280 let (cid, size) = sync_block_on(tree.put(&data)).expect("put file");
1281
1282 store
1283 .index_tree(
1284 &cid.hash,
1285 "npub1example",
1286 Some("large-file"),
1287 PRIORITY_OTHER,
1288 Some("npub1example/large-file"),
1289 )
1290 .expect("index tree");
1291
1292 let meta = store
1293 .get_tree_meta(&cid.hash)
1294 .expect("tree meta")
1295 .expect("indexed meta");
1296 assert_eq!(size, data.len() as u64);
1297 assert_eq!(meta.total_size, data.len() as u64);
1298 }
1299
1300 #[test]
1301 fn get_tree_ref_returns_stored_root() {
1302 let temp_dir = TempDir::new().expect("temp dir");
1303 let store = HashtreeStore::with_options(temp_dir.path(), None, 1024 * 1024).expect("store");
1304 let cid = build_test_tree(&store);
1305
1306 store
1307 .index_tree(
1308 &cid.hash,
1309 "npub1example",
1310 Some("playlist"),
1311 PRIORITY_OTHER,
1312 Some("npub1example/playlist"),
1313 )
1314 .expect("index tree");
1315
1316 assert_eq!(
1317 store
1318 .get_tree_ref("npub1example/playlist")
1319 .expect("tree ref lookup"),
1320 Some(cid.hash)
1321 );
1322 }
1323
1324 #[test]
1325 fn tree_meta_deserializes_metadata_without_tree_access_field() {
1326 #[derive(Serialize)]
1327 struct LegacyTreeMeta {
1328 owner: String,
1329 name: Option<String>,
1330 synced_at: u64,
1331 total_size: u64,
1332 priority: u8,
1333 }
1334
1335 let bytes = rmp_serde::to_vec(&LegacyTreeMeta {
1336 owner: "owner".to_string(),
1337 name: Some("tree".to_string()),
1338 synced_at: 123,
1339 total_size: 456,
1340 priority: PRIORITY_OTHER,
1341 })
1342 .expect("serialize legacy metadata");
1343 let meta: TreeMeta = rmp_serde::from_slice(&bytes).expect("deserialize tree metadata");
1344
1345 assert_eq!(meta.owner, "owner");
1346 assert_eq!(meta.name.as_deref(), Some("tree"));
1347 assert_eq!(meta.synced_at, 123);
1348 assert_eq!(meta.total_size, 456);
1349 assert_eq!(meta.priority, PRIORITY_OTHER);
1350 }
1351
1352 #[test]
1353 fn tree_meta_deserializes_accidental_access_field_but_drops_it_on_write() {
1354 #[derive(Serialize)]
1355 struct AccidentalTreeMeta {
1356 owner: String,
1357 name: Option<String>,
1358 synced_at: u64,
1359 last_accessed_at: u64,
1360 total_size: u64,
1361 priority: u8,
1362 }
1363
1364 let bytes = rmp_serde::to_vec(&AccidentalTreeMeta {
1365 owner: "owner".to_string(),
1366 name: Some("tree".to_string()),
1367 synced_at: 123,
1368 last_accessed_at: 999,
1369 total_size: 456,
1370 priority: PRIORITY_OTHER,
1371 })
1372 .expect("serialize accidental metadata");
1373 let meta: TreeMeta = rmp_serde::from_slice(&bytes).expect("deserialize tree metadata");
1374 let encoded = rmp_serde::to_vec(&meta).expect("serialize current metadata");
1375 let reparsed: (String, Option<String>, u64, u64, u8) =
1376 rmp_serde::from_slice(&encoded).expect("parse current metadata shape");
1377
1378 assert_eq!(meta.owner, "owner");
1379 assert_eq!(meta.name.as_deref(), Some("tree"));
1380 assert_eq!(meta.synced_at, 123);
1381 assert_eq!(meta.total_size, 456);
1382 assert_eq!(meta.priority, PRIORITY_OTHER);
1383 assert_eq!(reparsed.0, "owner");
1384 assert_eq!(reparsed.3, 456);
1385 assert_eq!(reparsed.4, PRIORITY_OTHER);
1386 }
1387
1388 #[test]
1389 fn eviction_prefers_oldest_tree_within_priority() {
1390 let temp_dir = TempDir::new().expect("temp dir");
1391 let store = HashtreeStore::with_options(temp_dir.path(), None, 500).expect("store");
1392
1393 let hash1 = hashtree_core::sha256(&[1u8; 200]);
1394 let hash2 = hashtree_core::sha256(&[2u8; 200]);
1395 let hash3 = hashtree_core::sha256(&[3u8; 200]);
1396 store.put_blob(&[1u8; 200]).expect("put blob 1");
1397 store.put_blob(&[2u8; 200]).expect("put blob 2");
1398 store.put_blob(&[3u8; 200]).expect("put blob 3");
1399 store
1400 .index_tree(&hash1, "owner1", Some("tree1"), PRIORITY_OTHER, None)
1401 .expect("index tree 1");
1402 store
1403 .index_tree(&hash2, "owner2", Some("tree2"), PRIORITY_OTHER, None)
1404 .expect("index tree 2");
1405 store
1406 .index_tree(&hash3, "owner3", Some("tree3"), PRIORITY_OTHER, None)
1407 .expect("index tree 3");
1408
1409 let freed = store.evict_if_needed().expect("evict");
1410
1411 assert!(freed > 0);
1412 assert!(
1413 store.get_tree_meta(&hash3).expect("tree meta").is_some(),
1414 "newest tree should survive before older peers at the same priority"
1415 );
1416 }
1417
1418 #[test]
1419 fn orphan_cleanup_keeps_socialgraph_root_hashes() {
1420 let temp_dir = TempDir::new().expect("temp dir");
1421 let store = HashtreeStore::with_options(temp_dir.path(), None, 1024).expect("store");
1422 let cid = build_test_tree(&store);
1423 write_root_file(
1424 &temp_dir.path().join("socialgraph/events-root.msgpack"),
1425 &cid,
1426 );
1427
1428 let freed = store
1429 .evict_disposable_orphans_to_target(0)
1430 .expect("orphan cleanup");
1431
1432 assert!(freed < 1024);
1433 assert!(store.blob_exists(&cid.hash).expect("root exists"));
1434 }
1435}