Skip to main content

wire/
object_availability.rs

1// SPDX-License-Identifier: Apache-2.0
2use objects::store::ObjectStore;
3
4use crate::{ObjectId, ObjectInfo, ObjectType, Result};
5
6#[derive(Debug, Clone, Default, PartialEq, Eq)]
7pub struct ObjectAvailabilityPlan {
8    pub have_objects: Vec<ObjectId>,
9    pub want_objects: Vec<ObjectId>,
10    pub present_objects: Vec<ObjectId>,
11    pub missing_objects: Vec<ObjectId>,
12    pub resumable_objects: Vec<ObjectId>,
13    pub lazy_objects: Vec<ObjectId>,
14    pub partial_fetch_allowed: bool,
15}
16
17pub fn has_object(store: &impl ObjectStore, info: &ObjectInfo) -> Result<bool> {
18    match (&info.id, info.obj_type) {
19        (ObjectId::Hash(hash), ObjectType::Blob) => Ok(store.has_blob(hash)?),
20        (ObjectId::Hash(hash), ObjectType::Tree) => Ok(store.has_tree(hash)?),
21        (ObjectId::StateId(state_id), ObjectType::State) => Ok(store.has_state(state_id)?),
22        (ObjectId::StateAttachment { state, id, kind: _ }, ObjectType::StateAttachment) => {
23            Ok(store.get_state_attachment(state, id)?.is_some())
24        }
25        // Redactions are keyed by the redacted blob's hash. Two senders
26        // can declare different redactions on the same blob (different
27        // reason / signature / timestamp), so we conservatively report
28        // "do not have" and always re-fetch — `accept_wire_redactions`
29        // deduplicates via the content-addressed `put_redaction`
30        // idempotency rule. Cheap to refetch; correct under merge.
31        (ObjectId::Hash(_), ObjectType::Redaction) => Ok(false),
32        // StateVisibility is a per-state sidecar with append/merge
33        // semantics. Like Redaction, conservatively refetch and let the
34        // repository boundary validate + dedupe.
35        (ObjectId::StateId(_), ObjectType::StateVisibility) => Ok(false),
36        _ => Ok(false),
37    }
38}
39
40pub fn plan_object_availability(
41    store: &impl ObjectStore,
42    objects: &[ObjectInfo],
43) -> Result<ObjectAvailabilityPlan> {
44    let mut plan = ObjectAvailabilityPlan::default();
45
46    for info in objects {
47        if has_object(store, info)? {
48            plan.have_objects.push(info.id.clone());
49            plan.present_objects.push(info.id.clone());
50        } else {
51            plan.want_objects.push(info.id.clone());
52            plan.missing_objects.push(info.id.clone());
53        }
54    }
55
56    Ok(plan)
57}
58
59impl ObjectAvailabilityPlan {
60    pub fn with_partial_fetch_allowed(mut self, allowed: bool) -> Self {
61        self.partial_fetch_allowed = allowed;
62        self
63    }
64
65    pub fn is_complete(&self) -> bool {
66        self.want_objects.is_empty()
67            && self.missing_objects.is_empty()
68            && self.resumable_objects.is_empty()
69            && self.lazy_objects.is_empty()
70    }
71
72    pub fn has_partial_fetch_candidates(&self) -> bool {
73        !self.resumable_objects.is_empty() || !self.lazy_objects.is_empty()
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use objects::{
80        object::{Blob, ContentHash, StateId, Tree},
81        store::{ObjectStore, Result as StoreResult},
82    };
83
84    use super::*;
85
86    #[derive(Default)]
87    struct DummyStore {
88        blob: Option<ContentHash>,
89        state: Option<StateId>,
90    }
91
92    impl ObjectStore for DummyStore {
93        fn get_blob(&self, _hash: &ContentHash) -> StoreResult<Option<Blob>> {
94            Ok(None)
95        }
96
97        fn put_blob(&self, _blob: &Blob) -> StoreResult<ContentHash> {
98            unreachable!("not used in test")
99        }
100
101        fn has_blob(&self, hash: &ContentHash) -> StoreResult<bool> {
102            Ok(self.blob == Some(*hash))
103        }
104
105        fn get_tree(&self, _hash: &ContentHash) -> StoreResult<Option<Tree>> {
106            Ok(None)
107        }
108
109        fn put_tree(&self, _tree: &Tree) -> StoreResult<ContentHash> {
110            unreachable!("not used in test")
111        }
112
113        fn has_tree(&self, _hash: &ContentHash) -> StoreResult<bool> {
114            Ok(false)
115        }
116
117        fn get_state(&self, _id: &StateId) -> StoreResult<Option<objects::object::State>> {
118            Ok(None)
119        }
120
121        fn put_state(&self, _state: &objects::object::State) -> StoreResult<()> {
122            unreachable!("not used in test")
123        }
124
125        fn has_state(&self, id: &StateId) -> StoreResult<bool> {
126            Ok(self.state == Some(*id))
127        }
128
129        fn list_states(&self) -> StoreResult<Vec<StateId>> {
130            Ok(vec![])
131        }
132
133        fn get_action(
134            &self,
135            _id: &objects::object::ActionId,
136        ) -> StoreResult<Option<objects::object::Action>> {
137            Ok(None)
138        }
139
140        fn put_action(
141            &self,
142            _action: &mut objects::object::Action,
143        ) -> StoreResult<objects::object::ActionId> {
144            unreachable!("not used in test")
145        }
146
147        fn list_actions(&self) -> StoreResult<Vec<objects::object::ActionId>> {
148            Ok(vec![])
149        }
150
151        fn list_blobs(&self) -> StoreResult<Vec<ContentHash>> {
152            Ok(vec![])
153        }
154
155        fn list_trees(&self) -> StoreResult<Vec<ContentHash>> {
156            Ok(vec![])
157        }
158    }
159
160    #[test]
161    fn test_plan_tracks_present_and_missing_objects() {
162        let blob = Blob::new(b"hello".to_vec());
163        let blob_hash = blob.hash();
164        let store = DummyStore {
165            blob: Some(blob_hash),
166            state: None,
167        };
168        let missing_hash = ContentHash::from_bytes([7; 32]);
169        let objects = vec![
170            ObjectInfo {
171                id: ObjectId::Hash(blob_hash),
172                obj_type: ObjectType::Blob,
173                size: blob.size() as u64,
174                delta_base: None,
175            },
176            ObjectInfo {
177                id: ObjectId::Hash(missing_hash),
178                obj_type: ObjectType::Tree,
179                size: 0,
180                delta_base: None,
181            },
182        ];
183
184        let plan = plan_object_availability(&store, &objects).unwrap();
185
186        assert_eq!(plan.have_objects.len(), 1);
187        assert_eq!(plan.want_objects.len(), 1);
188        assert_eq!(plan.present_objects.len(), 1);
189        assert_eq!(plan.missing_objects.len(), 1);
190        assert!(!plan.is_complete());
191    }
192
193    #[test]
194    fn missing_state_objects_are_requested() {
195        let store = DummyStore::default();
196        let state = StateId::from_bytes([9; 32]);
197        let objects = vec![ObjectInfo {
198            id: ObjectId::StateId(state),
199            obj_type: ObjectType::State,
200            size: 0,
201            delta_base: None,
202        }];
203
204        let plan = plan_object_availability(&store, &objects).unwrap();
205
206        assert!(plan.have_objects.is_empty());
207        assert_eq!(plan.want_objects, vec![ObjectId::StateId(state)]);
208        assert_eq!(plan.missing_objects, vec![ObjectId::StateId(state)]);
209    }
210
211    #[test]
212    fn immutable_state_objects_are_not_requested_when_present() {
213        let state = StateId::from_bytes([9; 32]);
214        let store = DummyStore {
215            state: Some(state),
216            ..DummyStore::default()
217        };
218        let objects = vec![ObjectInfo {
219            id: ObjectId::StateId(state),
220            obj_type: ObjectType::State,
221            size: 0,
222            delta_base: None,
223        }];
224
225        let plan = plan_object_availability(&store, &objects).unwrap();
226
227        assert_eq!(plan.have_objects, vec![ObjectId::StateId(state)]);
228        assert!(plan.want_objects.is_empty());
229    }
230
231    #[test]
232    fn test_partial_fetch_flag_helpers() {
233        let plan = ObjectAvailabilityPlan::default().with_partial_fetch_allowed(true);
234
235        assert!(plan.partial_fetch_allowed);
236        assert!(!plan.has_partial_fetch_candidates());
237        assert!(plan.is_complete());
238    }
239}