sc_client_db/
offchain.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5
6// This program is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// This program is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with this program. If not, see <https://www.gnu.org/licenses/>.
18
19//! RocksDB-based offchain workers local storage.
20
21use std::{collections::HashMap, sync::Arc};
22
23use crate::{columns, Database, DbHash, Transaction};
24use log::error;
25use parking_lot::Mutex;
26
27/// Offchain local storage
28#[derive(Clone)]
29pub struct LocalStorage {
30	db: Arc<dyn Database<DbHash>>,
31	locks: Arc<Mutex<HashMap<Vec<u8>, Arc<Mutex<()>>>>>,
32}
33
34impl std::fmt::Debug for LocalStorage {
35	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
36		fmt.debug_struct("LocalStorage").finish()
37	}
38}
39
40impl LocalStorage {
41	/// Create new offchain storage for tests (backed by memorydb)
42	#[cfg(any(feature = "test-helpers", test))]
43	pub fn new_test() -> Self {
44		let db = kvdb_memorydb::create(crate::utils::NUM_COLUMNS);
45		let db = sp_database::as_database(db);
46		Self::new(db as _)
47	}
48
49	/// Create offchain local storage with given `KeyValueDB` backend.
50	pub fn new(db: Arc<dyn Database<DbHash>>) -> Self {
51		Self { db, locks: Default::default() }
52	}
53}
54
55impl sp_core::offchain::OffchainStorage for LocalStorage {
56	fn set(&mut self, prefix: &[u8], key: &[u8], value: &[u8]) {
57		let mut tx = Transaction::new();
58		tx.set(columns::OFFCHAIN, &concatenate_prefix_and_key(prefix, key), value);
59
60		if let Err(err) = self.db.commit(tx) {
61			error!("Error setting on local storage: {}", err)
62		}
63	}
64
65	fn remove(&mut self, prefix: &[u8], key: &[u8]) {
66		let mut tx = Transaction::new();
67		tx.remove(columns::OFFCHAIN, &concatenate_prefix_and_key(prefix, key));
68
69		if let Err(err) = self.db.commit(tx) {
70			error!("Error removing on local storage: {}", err)
71		}
72	}
73
74	fn get(&self, prefix: &[u8], key: &[u8]) -> Option<Vec<u8>> {
75		self.db.get(columns::OFFCHAIN, &concatenate_prefix_and_key(prefix, key))
76	}
77
78	fn compare_and_set(
79		&mut self,
80		prefix: &[u8],
81		item_key: &[u8],
82		old_value: Option<&[u8]>,
83		new_value: &[u8],
84	) -> bool {
85		let key = concatenate_prefix_and_key(prefix, item_key);
86		let key_lock = {
87			let mut locks = self.locks.lock();
88			locks.entry(key.clone()).or_default().clone()
89		};
90
91		let is_set;
92		{
93			let _key_guard = key_lock.lock();
94			let val = self.db.get(columns::OFFCHAIN, &key);
95			is_set = val.as_deref() == old_value;
96
97			if is_set {
98				self.set(prefix, item_key, new_value)
99			}
100		}
101
102		// clean the lock map if we're the only entry
103		let mut locks = self.locks.lock();
104		{
105			drop(key_lock);
106			let key_lock = locks.get_mut(&key);
107			if key_lock.and_then(Arc::get_mut).is_some() {
108				locks.remove(&key);
109			}
110		}
111		is_set
112	}
113}
114
115/// Concatenate the prefix and key to create an offchain key in the db.
116pub(crate) fn concatenate_prefix_and_key(prefix: &[u8], key: &[u8]) -> Vec<u8> {
117	prefix.iter().chain(key.iter()).cloned().collect()
118}
119
120#[cfg(test)]
121mod tests {
122	use super::*;
123	use sp_core::offchain::OffchainStorage;
124
125	#[test]
126	fn should_compare_and_set_and_clear_the_locks_map() {
127		let mut storage = LocalStorage::new_test();
128		let prefix = b"prefix";
129		let key = b"key";
130		let value = b"value";
131
132		storage.set(prefix, key, value);
133		assert_eq!(storage.get(prefix, key), Some(value.to_vec()));
134
135		assert_eq!(storage.compare_and_set(prefix, key, Some(value), b"asd"), true);
136		assert_eq!(storage.get(prefix, key), Some(b"asd".to_vec()));
137		assert!(storage.locks.lock().is_empty(), "Locks map should be empty!");
138	}
139
140	#[test]
141	fn should_compare_and_set_on_empty_field() {
142		let mut storage = LocalStorage::new_test();
143		let prefix = b"prefix";
144		let key = b"key";
145
146		assert_eq!(storage.compare_and_set(prefix, key, None, b"asd"), true);
147		assert_eq!(storage.get(prefix, key), Some(b"asd".to_vec()));
148		assert!(storage.locks.lock().is_empty(), "Locks map should be empty!");
149	}
150}
sc_client_db/offchain.rs

sc_client_db/
offchain.rs
