iota_single_node_benchmark/
mock_storage.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::HashMap,
7    sync::{Arc, RwLock},
8};
9
10use iota_storage::package_object_cache::PackageObjectCache;
11use iota_types::{
12    base_types::{EpochId, ObjectID, SequenceNumber, VersionNumber},
13    error::{IotaError, IotaResult},
14    inner_temporary_store::InnerTemporaryStore,
15    object::{Object, Owner},
16    storage::{
17        BackingPackageStore, ChildObjectResolver, GetSharedLocks, ObjectStore, PackageObject,
18        get_module_by_id,
19    },
20    transaction::{InputObjectKind, InputObjects, ObjectReadResult, TransactionKey},
21};
22use move_binary_format::CompiledModule;
23use move_bytecode_utils::module_cache::GetModule;
24use move_core_types::language_storage::ModuleId;
25use once_cell::unsync::OnceCell;
26use prometheus::core::{Atomic, AtomicU64};
27
28// TODO: We won't need a special purpose InMemoryObjectStore once the
29// InMemoryCache is ready.
30#[derive(Clone)]
31pub(crate) struct InMemoryObjectStore {
32    objects: Arc<RwLock<HashMap<ObjectID, Object>>>,
33    package_cache: Arc<PackageObjectCache>,
34    num_object_reads: Arc<AtomicU64>,
35}
36
37impl InMemoryObjectStore {
38    pub(crate) fn new(objects: HashMap<ObjectID, Object>) -> Self {
39        Self {
40            objects: Arc::new(RwLock::new(objects)),
41            package_cache: PackageObjectCache::new(),
42            num_object_reads: Arc::new(AtomicU64::new(0)),
43        }
44    }
45
46    pub(crate) fn get_num_object_reads(&self) -> u64 {
47        self.num_object_reads.get()
48    }
49
50    // TODO: This function is out-of-sync with read_objects_for_execution from
51    // transaction_input_loader.rs. For instance, it does not support the use of
52    // deleted shared objects. We will need a trait to unify the these
53    // functions. (similarly the one in simulacrum)
54    pub(crate) fn read_objects_for_execution(
55        &self,
56        shared_locks: &dyn GetSharedLocks,
57        tx_key: &TransactionKey,
58        input_object_kinds: &[InputObjectKind],
59    ) -> IotaResult<InputObjects> {
60        let shared_locks_cell: OnceCell<Option<HashMap<_, _>>> = OnceCell::new();
61        let mut input_objects = Vec::new();
62        for kind in input_object_kinds {
63            let obj: Option<Object> = match kind {
64                InputObjectKind::MovePackage(id) => self.get_package_object(id)?.map(|o| o.into()),
65                InputObjectKind::ImmOrOwnedMoveObject(objref) => {
66                    self.get_object_by_key(&objref.0, objref.1)?
67                }
68
69                InputObjectKind::SharedMoveObject { id, .. } => {
70                    let shared_locks = shared_locks_cell
71                        .get_or_init(|| {
72                            shared_locks
73                                .get_shared_locks(tx_key)
74                                .expect("get_shared_locks should not fail")
75                                .map(|l| l.into_iter().collect())
76                        })
77                        .as_ref()
78                        .ok_or_else(|| IotaError::GenericAuthority {
79                            error: "Shared object locks should have been set.".to_string(),
80                        })?;
81                    let version = shared_locks.get(id).unwrap_or_else(|| {
82                        panic!("Shared object locks should have been set. key: {tx_key:?}, obj id: {id:?}")
83                    });
84
85                    self.get_object_by_key(id, *version)?
86                }
87            };
88
89            input_objects.push(ObjectReadResult::new(
90                *kind,
91                obj.ok_or_else(|| kind.object_not_found_error())?.into(),
92            ));
93        }
94
95        Ok(input_objects.into())
96    }
97
98    pub(crate) fn commit_objects(&self, inner_temp_store: InnerTemporaryStore) {
99        let mut objects = self.objects.write().unwrap();
100        for (object_id, _) in inner_temp_store.mutable_inputs {
101            if !inner_temp_store.written.contains_key(&object_id) {
102                objects.remove(&object_id);
103            }
104        }
105        for (object_id, object) in inner_temp_store.written {
106            objects.insert(object_id, object);
107        }
108    }
109}
110
111impl ObjectStore for InMemoryObjectStore {
112    fn get_object(
113        &self,
114        object_id: &ObjectID,
115    ) -> Result<Option<Object>, iota_types::storage::error::Error> {
116        self.num_object_reads.inc_by(1);
117        Ok(self.objects.read().unwrap().get(object_id).cloned())
118    }
119
120    fn get_object_by_key(
121        &self,
122        object_id: &ObjectID,
123        version: VersionNumber,
124    ) -> Result<Option<Object>, iota_types::storage::error::Error> {
125        Ok(self.get_object(object_id).unwrap().and_then(|o| {
126            if o.version() == version {
127                Some(o.clone())
128            } else {
129                None
130            }
131        }))
132    }
133}
134
135impl BackingPackageStore for InMemoryObjectStore {
136    fn get_package_object(&self, package_id: &ObjectID) -> IotaResult<Option<PackageObject>> {
137        self.package_cache.get_package_object(package_id, self)
138    }
139}
140
141impl ChildObjectResolver for InMemoryObjectStore {
142    fn read_child_object(
143        &self,
144        parent: &ObjectID,
145        child: &ObjectID,
146        child_version_upper_bound: SequenceNumber,
147    ) -> IotaResult<Option<Object>> {
148        Ok(self.get_object(child).unwrap().and_then(|o| {
149            if o.version() <= child_version_upper_bound
150                && o.owner == Owner::ObjectOwner((*parent).into())
151            {
152                Some(o.clone())
153            } else {
154                None
155            }
156        }))
157    }
158
159    fn get_object_received_at_version(
160        &self,
161        _owner: &ObjectID,
162        _receiving_object_id: &ObjectID,
163        _receive_object_at_version: SequenceNumber,
164        _epoch_id: EpochId,
165    ) -> IotaResult<Option<Object>> {
166        unimplemented!()
167    }
168}
169
170impl GetModule for InMemoryObjectStore {
171    type Error = IotaError;
172    type Item = CompiledModule;
173
174    fn get_module_by_id(&self, id: &ModuleId) -> Result<Option<Self::Item>, Self::Error> {
175        get_module_by_id(self, id)
176    }
177}
178
179impl GetSharedLocks for InMemoryObjectStore {
180    fn get_shared_locks(
181        &self,
182        _key: &TransactionKey,
183    ) -> IotaResult<Option<Vec<(ObjectID, SequenceNumber)>>> {
184        unreachable!()
185    }
186}
iota_single_node_benchmark/mock_storage.rs

iota_single_node_benchmark/
mock_storage.rs
