diff --git a/core/src/subgraph/context/instance/mod.rs b/core/src/subgraph/context/instance/mod.rs index 494ae227e97..9372791ccda 100644 --- a/core/src/subgraph/context/instance/mod.rs +++ b/core/src/subgraph/context/instance/mod.rs @@ -1,7 +1,7 @@ mod hosts; +mod module_cache; use anyhow::ensure; -use graph::futures01::sync::mpsc::Sender; use graph::{ blockchain::{Blockchain, TriggerData as _}, data_source::{ @@ -11,7 +11,7 @@ use graph::{ prelude::*, }; use hosts::{OffchainHosts, OnchainHosts}; -use std::collections::HashMap; +use module_cache::ModuleCache; pub(crate) struct SubgraphInstance> { subgraph_id: DeploymentHash, @@ -36,7 +36,7 @@ pub(crate) struct SubgraphInstance> { offchain_hosts: OffchainHosts, /// Maps the hash of a module to a channel to the thread in which the module is instantiated. - module_cache: HashMap<[u8; 32], Sender>, + module_cache: ModuleCache, /// This manages the sequence of causality regions for the subgraph. causality_region_seq: CausalityRegionSeq, @@ -85,7 +85,7 @@ where onchain_hosts: OnchainHosts::new(), subgraph_hosts: OnchainHosts::new(), offchain_hosts: OffchainHosts::new(), - module_cache: HashMap::new(), + module_cache: ModuleCache::new(), templates, host_metrics, causality_region_seq, @@ -104,7 +104,7 @@ where }; let mapping_request_sender = { - let module_hash = alloy::primitives::keccak256(module_bytes.as_ref()).0; + let module_hash = self.module_cache.hash(&module_bytes); if let Some(sender) = self.module_cache.get(&module_hash) { sender.clone() } else { diff --git a/core/src/subgraph/context/instance/module_cache.rs b/core/src/subgraph/context/instance/module_cache.rs new file mode 100644 index 00000000000..1440f557bee --- /dev/null +++ b/core/src/subgraph/context/instance/module_cache.rs @@ -0,0 +1,189 @@ +use graph::futures01::sync::mpsc::Sender; +use std::collections::HashMap; +use std::sync::Arc; + +/// Maps the hash of a module to a channel to the thread in which the module +/// is instantiated. +/// +/// Hashing a module reads all of its bytes, and a subgraph can have millions +/// of data sources that share a handful of modules. The hash is therefore +/// memoized by the identity of the `Arc` that holds the bytes: every data +/// source created from the same manifest mapping or template shares that +/// mapping's `Arc`. The memo holds a clone of the `Arc`, so the allocation +/// lives as long as the entry and an address found in the memo always +/// designates that same `Arc`. Bytes in another `Arc` are hashed again: equal +/// contents give the same hash, hence the same module. +pub(super) struct ModuleCache { + senders: HashMap<[u8; 32], Sender>, + hashes: HashMap>, [u8; 32])>, +} + +impl ModuleCache { + pub fn new() -> Self { + ModuleCache { + senders: HashMap::new(), + hashes: HashMap::new(), + } + } + + /// The keccak256 hash of `module_bytes`, computed once per `Arc`. + pub fn hash(&mut self, module_bytes: &Arc>) -> [u8; 32] { + let key = Arc::as_ptr(module_bytes) as usize; + if let Some((_, hash)) = self.hashes.get(&key) { + return *hash; + } + let hash = alloy::primitives::keccak256(module_bytes.as_slice()).0; + self.hashes.insert(key, (module_bytes.clone(), hash)); + hash + } + + pub fn get(&self, module_hash: &[u8; 32]) -> Option<&Sender> { + self.senders.get(module_hash) + } + + pub fn insert(&mut self, module_hash: [u8; 32], sender: Sender) { + self.senders.insert(module_hash, sender); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use graph::futures01::sync::mpsc::{Receiver, channel}; + + /// A cache used like `new_host` uses it, whose spawns are recorded: the + /// receivers, by spawn order. + struct Spawns { + cache: ModuleCache, + receivers: Vec>, + } + + impl Spawns { + fn new() -> Self { + Spawns { + cache: ModuleCache::new(), + receivers: Vec::new(), + } + } + + /// Memoize `mark` as the hash of `bytes`. + fn seed(&mut self, bytes: &Arc>, mark: [u8; 32]) { + let key = Arc::as_ptr(bytes) as usize; + self.cache.hashes.insert(key, (bytes.clone(), mark)); + } + + fn sender(&mut self, bytes: &Arc>) -> Sender { + let module_hash = self.cache.hash(bytes); + if let Some(sender) = self.cache.get(&module_hash) { + return sender.clone(); + } + let (tx, rx) = channel(1); + self.receivers.push(rx); + self.cache.insert(module_hash, tx.clone()); + tx + } + } + + /// Whether `sender` sends to the receiver of spawn `idx`: drop that + /// receiver and see if the sender is disconnected. + fn sends_to(mut spawns: Spawns, sender: &mut Sender, idx: usize) -> bool { + drop(spawns.receivers.remove(idx)); + sender.try_send(0).is_err_and(|e| e.is_disconnected()) + } + + fn module(n: u8) -> Vec { + (0..66_000u32).map(|i| (i as u8) ^ n).collect() + } + + #[test] + fn two_templates_two_hashes() { + let a = Arc::new(module(1)); + let b = Arc::new(module(2)); + let mut s = Spawns::new(); + let ha = s.cache.hash(&a); + let hb = s.cache.hash(&b); + assert_ne!(ha, hb); + assert_eq!(ha, alloy::primitives::keccak256(a.as_slice()).0); + assert_eq!(hb, alloy::primitives::keccak256(b.as_slice()).0); + + s.sender(&a); + let mut sb = s.sender(&b); + assert_eq!(s.receivers.len(), 2, "one module per template"); + assert!(sends_to(s, &mut sb, 1)); + } + + #[test] + fn sources_of_one_template_share_hash_and_sender() { + let template = Arc::new(module(1)); + let other = Arc::new(module(2)); + let mut s = Spawns::new(); + s.sender(&other); + s.sender(&template); + let mut second = s.sender(&template.clone()); + assert_eq!(s.cache.hashes.len(), 2); + assert_eq!(s.receivers.len(), 2); + assert!(sends_to(s, &mut second, 1), "the template's sender"); + } + + #[test] + fn many_sources_hash_once_per_template() { + const N: usize = 100_000; + // A recomputed hash overwrites its entry, which `hashes.len()` can't + // show. Marks are not real hashes on purpose: only a hit returns them. + let (mark_a, mark_b) = ([0xaa; 32], [0xbb; 32]); + let a = Arc::new(module(1)); + let b = Arc::new(module(2)); + let mut s = Spawns::new(); + s.seed(&a, mark_a); + s.seed(&b, mark_b); + for i in 0..N { + let (bytes, mark) = if i % 2 == 0 { + (a.clone(), mark_a) + } else { + (b.clone(), mark_b) + }; + assert_eq!(s.cache.hash(&bytes), mark); + s.sender(&bytes); + } + assert_eq!(s.cache.hashes.len(), 2); + assert_eq!(s.receivers.len(), 2); + + let mut s = Spawns::new(); + for _ in 0..N { + s.sender(&a.clone()); + } + assert_eq!(s.cache.hashes.len(), 1); + assert_eq!(s.receivers.len(), 1); + } + + #[test] + fn equal_unshared_bytes_give_same_hash_and_module() { + let shared = Arc::new(module(1)); + let copy = Arc::new(module(1)); + assert!(!Arc::ptr_eq(&shared, ©)); + let mut s = Spawns::new(); + s.sender(&Arc::new(module(2))); + s.sender(&shared); + let hash = s.cache.hash(&shared); + assert_eq!(s.cache.hash(©), hash); + assert_eq!(s.cache.hashes.len(), 3, "the copy is hashed"); + let mut from_copy = s.sender(©); + assert_eq!(s.receivers.len(), 2, "no new module for the copy"); + assert!(sends_to(s, &mut from_copy, 1), "the shared bytes' sender"); + } + + #[test] + fn dropped_bytes_are_not_confused_with_new_ones() { + // Allocate and drop many `Arc`s: the memo holds each one, so no + // later allocation can take the address of a memoized one. + let mut s = Spawns::new(); + for n in 0..64u8 { + let bytes = Arc::new(vec![n; 64]); + assert_eq!( + s.cache.hash(&bytes), + alloy::primitives::keccak256(bytes.as_slice()).0 + ); + } + assert_eq!(s.cache.hashes.len(), 64); + } +}