frame_support/storage/
transactional.rs1use sp_io::storage::{commit_transaction, rollback_transaction, start_transaction};
30use sp_runtime::{DispatchError, TransactionOutcome, TransactionalError};
31
32pub type Layer = u32;
34pub const TRANSACTION_LEVEL_KEY: &[u8] = b":transaction_level:";
38pub const TRANSACTIONAL_LIMIT: Layer = 255;
40
41fn get_transaction_level() -> Layer {
43 crate::storage::unhashed::get_or_default::<Layer>(TRANSACTION_LEVEL_KEY)
44}
45
46fn set_transaction_level(level: Layer) {
48 crate::storage::unhashed::put::<Layer>(TRANSACTION_LEVEL_KEY, &level);
49}
50
51fn kill_transaction_level() {
53 crate::storage::unhashed::kill(TRANSACTION_LEVEL_KEY);
54}
55
56fn inc_transaction_level() -> Result<StorageLayerGuard, ()> {
60 let existing_levels = get_transaction_level();
61 if existing_levels >= TRANSACTIONAL_LIMIT {
62 return Err(())
63 }
64 set_transaction_level(existing_levels + 1);
66 Ok(StorageLayerGuard)
67}
68
69fn dec_transaction_level() {
70 let existing_levels = get_transaction_level();
71 if existing_levels == 0 {
72 log::warn!(
73 "We are underflowing with calculating transactional levels. Not great, but let's not panic...",
74 );
75 } else if existing_levels == 1 {
76 kill_transaction_level();
78 } else {
79 set_transaction_level(existing_levels - 1);
81 }
82}
83
84struct StorageLayerGuard;
85
86impl Drop for StorageLayerGuard {
87 fn drop(&mut self) {
88 dec_transaction_level()
89 }
90}
91
92pub fn is_transactional() -> bool {
94 get_transaction_level() > 0
95}
96
97pub fn with_transaction<T, E, F>(f: F) -> Result<T, E>
107where
108 E: From<DispatchError>,
109 F: FnOnce() -> TransactionOutcome<Result<T, E>>,
110{
111 let _guard = inc_transaction_level().map_err(|()| TransactionalError::LimitReached.into())?;
115
116 start_transaction();
117
118 match f() {
119 TransactionOutcome::Commit(res) => {
120 commit_transaction();
121 res
122 },
123 TransactionOutcome::Rollback(res) => {
124 rollback_transaction();
125 res
126 },
127 }
128}
129
130pub fn with_transaction_opaque_err<T, E, F>(f: F) -> Result<Result<T, E>, ()>
134where
135 F: FnOnce() -> TransactionOutcome<Result<T, E>>,
136{
137 with_transaction(move || -> TransactionOutcome<Result<Result<T, E>, DispatchError>> {
138 match f() {
139 TransactionOutcome::Commit(res) => TransactionOutcome::Commit(Ok(res)),
140 TransactionOutcome::Rollback(res) => TransactionOutcome::Rollback(Ok(res)),
141 }
142 })
143 .map_err(|_| ())
144}
145
146pub fn with_transaction_unchecked<R, F>(f: F) -> R
152where
153 F: FnOnce() -> TransactionOutcome<R>,
154{
155 let maybe_guard = inc_transaction_level();
159
160 if maybe_guard.is_err() {
161 log::warn!(
162 "The transactional layer limit has been reached, and new transactional layers are being
163 spawned with `with_transaction_unchecked`. This could be caused by someone trying to
164 attack your chain, and you should investigate usage of `with_transaction_unchecked` and
165 potentially migrate to `with_transaction`, which enforces a transactional limit.",
166 );
167 }
168
169 start_transaction();
170
171 match f() {
172 TransactionOutcome::Commit(res) => {
173 commit_transaction();
174 res
175 },
176 TransactionOutcome::Rollback(res) => {
177 rollback_transaction();
178 res
179 },
180 }
181}
182
183pub fn with_storage_layer<T, E, F>(f: F) -> Result<T, E>
189where
190 E: From<DispatchError>,
191 F: FnOnce() -> Result<T, E>,
192{
193 with_transaction(|| {
194 let r = f();
195 if r.is_ok() {
196 TransactionOutcome::Commit(r)
197 } else {
198 TransactionOutcome::Rollback(r)
199 }
200 })
201}
202
203pub fn in_storage_layer<T, E, F>(f: F) -> Result<T, E>
208where
209 E: From<DispatchError>,
210 F: FnOnce() -> Result<T, E>,
211{
212 if is_transactional() {
213 f()
214 } else {
215 with_storage_layer(f)
216 }
217}
218
219#[cfg(test)]
220mod tests {
221 use super::*;
222 use crate::{assert_noop, assert_ok};
223 use sp_io::TestExternalities;
224 use sp_runtime::DispatchResult;
225
226 #[test]
227 fn is_transactional_should_return_false() {
228 TestExternalities::default().execute_with(|| {
229 assert!(!is_transactional());
230 });
231 }
232
233 #[test]
234 fn is_transactional_should_not_error_in_with_transaction() {
235 TestExternalities::default().execute_with(|| {
236 assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
237 assert!(is_transactional());
238 TransactionOutcome::Commit(Ok(()))
239 }));
240
241 assert_noop!(
242 with_transaction(|| -> TransactionOutcome<DispatchResult> {
243 assert!(is_transactional());
244 TransactionOutcome::Rollback(Err("revert".into()))
245 }),
246 "revert"
247 );
248 });
249 }
250
251 fn recursive_transactional(num: u32) -> DispatchResult {
252 if num == 0 {
253 return Ok(())
254 }
255
256 with_transaction(|| -> TransactionOutcome<DispatchResult> {
257 let res = recursive_transactional(num - 1);
258 TransactionOutcome::Commit(res)
259 })
260 }
261
262 #[test]
263 fn transaction_limit_should_work() {
264 TestExternalities::default().execute_with(|| {
265 assert_eq!(get_transaction_level(), 0);
266
267 assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
268 assert_eq!(get_transaction_level(), 1);
269 TransactionOutcome::Commit(Ok(()))
270 }));
271
272 assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
273 assert_eq!(get_transaction_level(), 1);
274 let res = with_transaction(|| -> TransactionOutcome<DispatchResult> {
275 assert_eq!(get_transaction_level(), 2);
276 TransactionOutcome::Commit(Ok(()))
277 });
278 TransactionOutcome::Commit(res)
279 }));
280
281 assert_ok!(recursive_transactional(255));
282 assert_noop!(
283 recursive_transactional(256),
284 sp_runtime::TransactionalError::LimitReached
285 );
286
287 assert_eq!(get_transaction_level(), 0);
288 });
289 }
290
291 #[test]
292 fn in_storage_layer_works() {
293 TestExternalities::default().execute_with(|| {
294 assert_eq!(get_transaction_level(), 0);
295
296 let res = in_storage_layer(|| -> DispatchResult {
297 assert_eq!(get_transaction_level(), 1);
298 in_storage_layer(|| -> DispatchResult {
299 assert_eq!(get_transaction_level(), 1);
301 Ok(())
302 })
303 });
304
305 assert_ok!(res);
306
307 let res = in_storage_layer(|| -> DispatchResult {
308 assert_eq!(get_transaction_level(), 1);
309 in_storage_layer(|| -> DispatchResult {
310 assert_eq!(get_transaction_level(), 1);
312 Err("epic fail".into())
313 })
314 });
315
316 assert_noop!(res, "epic fail");
317 });
318 }
319}