|
| 1 | +// This file is part of Substrate. |
| 2 | + |
| 3 | +// Copyright (C) Parity Technologies (UK) Ltd. |
| 4 | +// SPDX-License-Identifier: Apache-2.0 |
| 5 | + |
| 6 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 7 | +// you may not use this file except in compliance with the License. |
| 8 | +// You may obtain a copy of the License at |
| 9 | +// |
| 10 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | +// |
| 12 | +// Unless required by applicable law or agreed to in writing, software |
| 13 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | +// See the License for the specific language governing permissions and |
| 16 | +// limitations under the License. |
| 17 | + |
| 18 | +//! Provides functions to interact with the dispatch context. |
| 19 | +//! |
| 20 | +//! A Dispatch context is created by calling [`run_in_context`] and then the given closure will be |
| 21 | +//! executed in this dispatch context. Everyting run in this `closure` will have access to the same |
| 22 | +//! dispatch context. This also applies to nested calls of [`run_in_context`]. The dispatch context |
| 23 | +//! can be used to store and retrieve information locally in this context. The dispatch context can |
| 24 | +//! be accessed by using [`with_context`]. This function will execute the given closure and give it |
| 25 | +//! access to the value stored in the dispatch context. |
| 26 | +//! |
| 27 | +//! # FRAME integration |
| 28 | +//! |
| 29 | +//! The FRAME macros implement [`UnfilteredDispatchable`](crate::traits::UnfilteredDispatchable) for |
| 30 | +//! each pallet `Call` enum. Part of this implementation is the call to [`run_in_context`], so that |
| 31 | +//! each call to |
| 32 | +//! [`UnfilteredDispatchable::dispatch_bypass_filter`](crate::traits::UnfilteredDispatchable::dispatch_bypass_filter) |
| 33 | +//! or [`Dispatchable::dispatch`](crate::dispatch::Dispatchable::dispatch) will run in a dispatch |
| 34 | +//! context. |
| 35 | +//! |
| 36 | +//! # Example |
| 37 | +//! |
| 38 | +//! ``` |
| 39 | +//! use frame_support::dispatch_context::{with_context, run_in_context}; |
| 40 | +//! |
| 41 | +//! // Not executed in a dispatch context, so it should return `None`. |
| 42 | +//! assert!(with_context::<(), _>(|_| println!("Hello")).is_none()); |
| 43 | +//! |
| 44 | +//! // Run it in a dispatch context and `with_context` returns `Some(_)`. |
| 45 | +//! run_in_context(|| { |
| 46 | +//! assert!(with_context::<(), _>(|_| println!("Hello")).is_some()); |
| 47 | +//! }); |
| 48 | +//! |
| 49 | +//! #[derive(Default)] |
| 50 | +//! struct CustomContext(i32); |
| 51 | +//! |
| 52 | +//! run_in_context(|| { |
| 53 | +//! with_context::<CustomContext, _>(|v| { |
| 54 | +//! // Intitialize the value to the default value. |
| 55 | +//! assert_eq!(0, v.or_default().0); |
| 56 | +//! v.or_default().0 = 10; |
| 57 | +//! }); |
| 58 | +//! |
| 59 | +//! with_context::<CustomContext, _>(|v| { |
| 60 | +//! // We are still in the same context and can still access the set value. |
| 61 | +//! assert_eq!(10, v.or_default().0); |
| 62 | +//! }); |
| 63 | +//! |
| 64 | +//! run_in_context(|| { |
| 65 | +//! with_context::<CustomContext, _>(|v| { |
| 66 | +//! // A nested call of `run_in_context` stays in the same dispatch context |
| 67 | +//! assert_eq!(10, v.or_default().0); |
| 68 | +//! }) |
| 69 | +//! }) |
| 70 | +//! }); |
| 71 | +//! |
| 72 | +//! run_in_context(|| { |
| 73 | +//! with_context::<CustomContext, _>(|v| { |
| 74 | +//! // We left the other context and created a new one, so we should be back |
| 75 | +//! // to our default value. |
| 76 | +//! assert_eq!(0, v.or_default().0); |
| 77 | +//! }); |
| 78 | +//! }); |
| 79 | +//! ``` |
| 80 | +//! |
| 81 | +//! In your pallet you will only have to use [`with_context`], because as described above |
| 82 | +//! [`run_in_context`] will be handled by FRAME for you. |
| 83 | +
|
| 84 | +use sp_std::{ |
| 85 | + any::{Any, TypeId}, |
| 86 | + boxed::Box, |
| 87 | + collections::btree_map::{BTreeMap, Entry}, |
| 88 | +}; |
| 89 | + |
| 90 | +environmental::environmental!(DISPATCH_CONTEXT: BTreeMap<TypeId, Box<dyn Any>>); |
| 91 | + |
| 92 | +/// Abstraction over some optional value `T` that is stored in the dispatch context. |
| 93 | +pub struct Value<'a, T> { |
| 94 | + value: Option<&'a mut T>, |
| 95 | + new_value: Option<T>, |
| 96 | +} |
| 97 | + |
| 98 | +impl<T> Value<'_, T> { |
| 99 | + /// Get the value as reference. |
| 100 | + pub fn get(&self) -> Option<&T> { |
| 101 | + self.new_value.as_ref().or_else(|| self.value.as_ref().map(|v| *v as &T)) |
| 102 | + } |
| 103 | + |
| 104 | + /// Get the value as mutable reference. |
| 105 | + pub fn get_mut(&mut self) -> Option<&mut T> { |
| 106 | + self.new_value.as_mut().or_else(|| self.value.as_mut().map(|v| *v as &mut T)) |
| 107 | + } |
| 108 | + |
| 109 | + /// Set to the given value. |
| 110 | + /// |
| 111 | + /// [`Self::get`] and [`Self::get_mut`] will return `new_value` afterwards. |
| 112 | + pub fn set(&mut self, new_value: T) { |
| 113 | + self.value = None; |
| 114 | + self.new_value = Some(new_value); |
| 115 | + } |
| 116 | + |
| 117 | + /// Returns a mutable reference to the value. |
| 118 | + /// |
| 119 | + /// If the internal value isn't initialized, this will set it to [`Default::default()`] before |
| 120 | + /// returning the mutable reference. |
| 121 | + pub fn or_default(&mut self) -> &mut T |
| 122 | + where |
| 123 | + T: Default, |
| 124 | + { |
| 125 | + if let Some(v) = &mut self.value { |
| 126 | + return v |
| 127 | + } |
| 128 | + |
| 129 | + self.new_value.get_or_insert_with(|| Default::default()) |
| 130 | + } |
| 131 | + |
| 132 | + /// Clear the internal value. |
| 133 | + /// |
| 134 | + /// [`Self::get`] and [`Self::get_mut`] will return `None` afterwards. |
| 135 | + pub fn clear(&mut self) { |
| 136 | + self.new_value = None; |
| 137 | + self.value = None; |
| 138 | + } |
| 139 | +} |
| 140 | + |
| 141 | +/// Runs the given `callback` in the dispatch context and gives access to some user defined value. |
| 142 | +/// |
| 143 | +/// Passes the a mutable reference of [`Value`] to the callback. The value will be of type `T` and |
| 144 | +/// is identified using the [`TypeId`] of `T`. This means that `T` should be some unique type to |
| 145 | +/// make the value unique. If no value is set yet [`Value::get()`] and [`Value::get_mut()`] will |
| 146 | +/// return `None`. It is totally valid to have some `T` that is shared between different callers to |
| 147 | +/// have access to the same value. |
| 148 | +/// |
| 149 | +/// Returns `None` if the current context is not a dispatch context. To create a context it is |
| 150 | +/// required to call [`run_in_context`] with the closure to execute in this context. So, for example |
| 151 | +/// in tests it could be that there isn't any dispatch context or when calling a dispatchable like a |
| 152 | +/// normal Rust function from some FRAME hook. |
| 153 | +pub fn with_context<T: 'static, R>(callback: impl FnOnce(&mut Value<T>) -> R) -> Option<R> { |
| 154 | + DISPATCH_CONTEXT::with(|c| match c.entry(TypeId::of::<T>()) { |
| 155 | + Entry::Occupied(mut o) => { |
| 156 | + let value = o.get_mut().downcast_mut::<T>(); |
| 157 | + |
| 158 | + if value.is_none() { |
| 159 | + log::error!( |
| 160 | + "Failed to downcast value for type {} in dispatch context!", |
| 161 | + sp_std::any::type_name::<T>(), |
| 162 | + ); |
| 163 | + } |
| 164 | + |
| 165 | + let mut value = Value { value, new_value: None }; |
| 166 | + let res = callback(&mut value); |
| 167 | + |
| 168 | + if value.value.is_none() && value.new_value.is_none() { |
| 169 | + o.remove(); |
| 170 | + } else if let Some(new_value) = value.new_value { |
| 171 | + o.insert(Box::new(new_value) as Box<_>); |
| 172 | + } |
| 173 | + |
| 174 | + res |
| 175 | + }, |
| 176 | + Entry::Vacant(v) => { |
| 177 | + let mut value = Value { value: None, new_value: None }; |
| 178 | + |
| 179 | + let res = callback(&mut value); |
| 180 | + |
| 181 | + if let Some(new_value) = value.new_value { |
| 182 | + v.insert(Box::new(new_value) as Box<_>); |
| 183 | + } |
| 184 | + |
| 185 | + res |
| 186 | + }, |
| 187 | + }) |
| 188 | +} |
| 189 | + |
| 190 | +/// Run the given closure `run` in a dispatch context. |
| 191 | +/// |
| 192 | +/// Nested calls to this function will execute `run` in the same dispatch context as the initial |
| 193 | +/// call to this function. In other words, all nested calls of this function will be done in the |
| 194 | +/// same dispatch context. |
| 195 | +pub fn run_in_context<R>(run: impl FnOnce() -> R) -> R { |
| 196 | + DISPATCH_CONTEXT::using_once(&mut Default::default(), run) |
| 197 | +} |
| 198 | + |
| 199 | +#[cfg(test)] |
| 200 | +mod tests { |
| 201 | + use super::*; |
| 202 | + |
| 203 | + #[test] |
| 204 | + fn dispatch_context_works() { |
| 205 | + // No context, so we don't execute |
| 206 | + assert!(with_context::<(), _>(|_| ()).is_none()); |
| 207 | + |
| 208 | + let ret = run_in_context(|| with_context::<(), _>(|_| 1).unwrap()); |
| 209 | + assert_eq!(1, ret); |
| 210 | + |
| 211 | + #[derive(Default)] |
| 212 | + struct Context(i32); |
| 213 | + |
| 214 | + let res = run_in_context(|| { |
| 215 | + with_context::<Context, _>(|v| { |
| 216 | + assert_eq!(0, v.or_default().0); |
| 217 | + |
| 218 | + v.or_default().0 = 100; |
| 219 | + }); |
| 220 | + |
| 221 | + run_in_context(|| { |
| 222 | + run_in_context(|| { |
| 223 | + run_in_context(|| with_context::<Context, _>(|v| v.or_default().0).unwrap()) |
| 224 | + }) |
| 225 | + }) |
| 226 | + }); |
| 227 | + |
| 228 | + // Ensure that the initial value set in the context is also accessible after nesting the |
| 229 | + // `run_in_context` calls. |
| 230 | + assert_eq!(100, res); |
| 231 | + } |
| 232 | +} |
0 commit comments