sui_adapter_latest/data_store/
cached_package_store.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::data_store::PackageStore;
5use indexmap::IndexMap;
6use move_core_types::identifier::IdentStr;
7use std::{
8    cell::RefCell,
9    collections::{BTreeMap, BTreeSet},
10    rc::Rc,
11};
12use sui_types::{
13    base_types::ObjectID,
14    error::{ExecutionError, SuiResult},
15    move_package::MovePackage,
16    storage::BackingPackageStore,
17};
18
19/// A package store that caches packages in memory and indexes type origins. This is useful for
20/// speeding up package loading and type resolution.
21///
22/// There is a `max_package_cache_size` that determines how many packages we will cache. If the cache is
23/// full, we will drop the cache.
24///
25/// The `max_type_cache_size` determines how many types we will cache. If the cache is full, we
26/// will drop the cache.
27///
28/// FUTURE: Most/all of this will be replaced by the VM runtime cache in the new VM. For now
29/// though, we need to use this.
30pub struct CachedPackageStore<'state> {
31    /// Underlying store the fetch packages from
32    pub package_store: Box<dyn BackingPackageStore + 'state>,
33    /// A cache of packages that we've loaded so far. This is used to speed up package loading.
34    /// Elements in this are safe to be evicted based on cache decisions.
35    package_cache: RefCell<BTreeMap<ObjectID, Option<Rc<MovePackage>>>>,
36    /// A cache of type origins that we've loaded so far. This is used to speed up type resolution.
37    /// Elements in this are safe to be evicted based on cache decisions.
38    type_origin_cache: RefCell<CachedTypeOriginMap>,
39    /// Elements in this are packages that are being published or have been published in the current
40    /// transaction.
41    /// Elements in this are _not_ safe to evict, unless evicted through `pop_package`.
42    new_packages: RefCell<IndexMap<ObjectID, Rc<MovePackage>>>,
43    /// Maximum size (in number of packages) that the `package_cache` can grow to before it is
44    /// cleared.
45    max_package_cache_size: usize,
46    /// Maximum size (in number of keys) that the `type_origin_cache` can grow to before it is
47    /// cleared.,
48    max_type_cache_size: usize,
49}
50
51type TypeOriginMap = BTreeMap<ObjectID, BTreeMap<(String, String), ObjectID>>;
52
53#[derive(Debug)]
54struct CachedTypeOriginMap {
55    /// Tracker of all packages that we've loaded so far. This is used to determine if the
56    /// `TypeOriginMap` needs to be updated when loading a package, or if that package has already
57    /// contributed to the `TypeOriginMap`.
58    pub cached_type_origins: BTreeSet<ObjectID>,
59    /// A mapping of the (original package ID)::<module_name>::<type_name> to the defining ID for
60    /// that type.
61    pub type_origin_map: TypeOriginMap,
62}
63
64impl CachedTypeOriginMap {
65    pub fn new() -> Self {
66        Self {
67            cached_type_origins: BTreeSet::new(),
68            type_origin_map: TypeOriginMap::new(),
69        }
70    }
71}
72
73impl<'state> CachedPackageStore<'state> {
74    pub const DEFAULT_MAX_PACKAGE_CACHE_SIZE: usize = 200;
75    pub const DEFAULT_MAX_TYPE_ORIGIN_CACHE_SIZE: usize = 1000;
76
77    pub fn new(package_store: Box<dyn BackingPackageStore + 'state>) -> Self {
78        Self {
79            package_store,
80            package_cache: RefCell::new(BTreeMap::new()),
81            type_origin_cache: RefCell::new(CachedTypeOriginMap::new()),
82            new_packages: RefCell::new(IndexMap::new()),
83            max_package_cache_size: Self::DEFAULT_MAX_PACKAGE_CACHE_SIZE,
84            max_type_cache_size: Self::DEFAULT_MAX_TYPE_ORIGIN_CACHE_SIZE,
85        }
86    }
87
88    /// Push a new package into the new packages. This is used to track packages that are being
89    /// published or have been published.
90    pub fn push_package(
91        &self,
92        id: ObjectID,
93        package: Rc<MovePackage>,
94    ) -> Result<(), ExecutionError> {
95        // Check that the package ID is not already present anywhere.
96        debug_assert!(self.fetch_package(&id).unwrap().is_none());
97
98        // Insert the package into the new packages
99        // If the package already exists, we will overwrite it and signal an error.
100        if self.new_packages.borrow_mut().insert(id, package).is_some() {
101            invariant_violation!(
102                "Package with ID {} already exists in the new packages. This should never happen.",
103                id
104            );
105        }
106
107        Ok(())
108    }
109
110    /// Rollback a package that was pushed into the new packages. We keep the invariant that:
111    /// * You can only pop the most recent package that was pushed.
112    /// * The element being popped _must_ exist in the new packages.
113    ///
114    /// Otherwise this returns an invariant violation.
115    pub fn pop_package(&self, id: ObjectID) -> Result<Rc<MovePackage>, ExecutionError> {
116        if self
117            .new_packages
118            .borrow()
119            .last()
120            .is_none_or(|(pkg_id, _)| *pkg_id != id)
121        {
122            make_invariant_violation!(
123                "Tried to pop package {} from new packages, but new packages was empty or \
124                it is not the most recent package inserted. This should never happen.",
125                id
126            );
127        }
128
129        let Some((pkg_id, pkg)) = self.new_packages.borrow_mut().pop() else {
130            unreachable!(
131                "We just checked that new packages is not empty, so this should never happen."
132            );
133        };
134        assert_eq!(
135            pkg_id, id,
136            "Popped package ID {} does not match requested ID {}. This should never happen as was checked above.",
137            pkg_id, id
138        );
139
140        Ok(pkg)
141    }
142
143    pub fn to_new_packages(&self) -> Vec<MovePackage> {
144        self.new_packages
145            .borrow()
146            .iter()
147            .map(|(_, pkg)| pkg.as_ref().clone())
148            .collect()
149    }
150
151    /// Get a package by its package ID (i.e., not original ID). This will first look in the new
152    /// packages, then in the cache, and then finally try and fetch the package from the underlying
153    /// package store.
154    ///
155    /// Once the package is fetched it will be added to the type origin cache if it is not already
156    /// present in the type origin cache.
157    pub fn get_package(&self, object_id: &ObjectID) -> SuiResult<Option<Rc<MovePackage>>> {
158        let Some(pkg) = self.fetch_package(object_id)? else {
159            return Ok(None);
160        };
161
162        let package_id = pkg.id();
163
164        // If the number of type origins that we have cached exceeds the max size, drop the cache.
165        if self.type_origin_cache.borrow().cached_type_origins.len() >= self.max_type_cache_size {
166            *self.type_origin_cache.borrow_mut() = CachedTypeOriginMap::new();
167        }
168
169        if !self
170            .type_origin_cache
171            .borrow()
172            .cached_type_origins
173            .contains(&package_id)
174        {
175            let cached_type_origin_map = &mut self.type_origin_cache.borrow_mut();
176            cached_type_origin_map
177                .cached_type_origins
178                .insert(package_id);
179            let original_package_id = pkg.original_package_id();
180            let package_types = cached_type_origin_map
181                .type_origin_map
182                .entry(original_package_id)
183                .or_default();
184            for ((module_name, type_name), defining_id) in pkg.type_origin_map().into_iter() {
185                if let Some(other) = package_types.insert(
186                    (module_name.to_string(), type_name.to_string()),
187                    defining_id,
188                ) {
189                    assert_eq!(
190                        other, defining_id,
191                        "type origin map should never have conflicting entries"
192                    );
193                }
194            }
195        }
196        Ok(Some(pkg))
197    }
198
199    /// Get a package by its package ID (i.e., not original ID). This will first look in the new
200    /// packages, then in the cache, and then finally try and fetch the package from the underlying
201    /// package store. NB: this does not do any indexing of the package.
202    fn fetch_package(&self, id: &ObjectID) -> SuiResult<Option<Rc<MovePackage>>> {
203        // Look for package in new packages
204        if let Some(pkg) = self.new_packages.borrow().get(id).cloned() {
205            return Ok(Some(pkg));
206        }
207
208        // Look for package in cache
209        if let Some(pkg) = self.package_cache.borrow().get(id).cloned() {
210            return Ok(pkg);
211        }
212
213        if self.package_cache.borrow().len() >= self.max_package_cache_size {
214            self.package_cache.borrow_mut().clear();
215        }
216
217        let pkg = self
218            .package_store
219            .get_package_object(id)?
220            .map(|pkg_obj| Rc::new(pkg_obj.move_package().clone()));
221        self.package_cache.borrow_mut().insert(*id, pkg.clone());
222        Ok(pkg)
223    }
224}
225
226impl PackageStore for CachedPackageStore<'_> {
227    fn get_package(&self, id: &ObjectID) -> SuiResult<Option<Rc<MovePackage>>> {
228        self.get_package(id)
229    }
230
231    fn resolve_type_to_defining_id(
232        &self,
233        module_address: ObjectID,
234        module_name: &IdentStr,
235        type_name: &IdentStr,
236    ) -> SuiResult<Option<ObjectID>> {
237        let Some(pkg) = self.get_package(&module_address)? else {
238            return Ok(None);
239        };
240
241        Ok(self
242            .type_origin_cache
243            .borrow()
244            .type_origin_map
245            .get(&pkg.original_package_id())
246            .and_then(|module_map| {
247                module_map
248                    .get(&(module_name.to_string(), type_name.to_string()))
249                    .copied()
250            }))
251    }
252}
sui_adapter_latest/data_store/cached_package_store.rs

sui_adapter_latest/data_store/
cached_package_store.rs
