@@ -48,46 +48,48 @@ pub mod unhashed;
4848pub mod weak_bounded_vec;
4949
5050mod transaction_level_tracker {
51- use core:: sync:: atomic:: { AtomicU32 , Ordering } ;
52-
5351 type Layer = u32 ;
54- static NUM_LEVELS : AtomicU32 = AtomicU32 :: new ( 0 ) ;
52+ const TRANSACTION_LEVEL_KEY : & ' static [ u8 ] = b":transaction_level:" ;
5553 const TRANSACTIONAL_LIMIT : Layer = 255 ;
5654
5755 pub fn get_transaction_level ( ) -> Layer {
58- NUM_LEVELS . load ( Ordering :: SeqCst )
56+ crate :: storage:: unhashed:: get_or_default :: < Layer > ( TRANSACTION_LEVEL_KEY )
57+ }
58+
59+ fn set_transaction_level ( level : & Layer ) {
60+ crate :: storage:: unhashed:: put :: < Layer > ( TRANSACTION_LEVEL_KEY , level) ;
61+ }
62+
63+ fn kill_transaction_level ( ) {
64+ crate :: storage:: unhashed:: kill ( TRANSACTION_LEVEL_KEY ) ;
5965 }
6066
6167 /// Increments the transaction level. Returns an error if levels go past the limit.
6268 ///
6369 /// Returns a guard that when dropped decrements the transaction level automatically.
6470 pub fn inc_transaction_level ( ) -> Result < StorageLayerGuard , ( ) > {
65- NUM_LEVELS
66- . fetch_update ( Ordering :: SeqCst , Ordering :: SeqCst , |existing_levels| {
67- if existing_levels >= TRANSACTIONAL_LIMIT {
68- return None
69- }
70- // Cannot overflow because of check above.
71- Some ( existing_levels + 1 )
72- } )
73- . map_err ( |_| ( ) ) ?;
71+ let existing_levels = get_transaction_level ( ) ;
72+ if existing_levels >= TRANSACTIONAL_LIMIT {
73+ return Err ( ( ) )
74+ }
75+ // Cannot overflow because of check above.
76+ set_transaction_level ( & ( existing_levels + 1 ) ) ;
7477 Ok ( StorageLayerGuard )
7578 }
7679
7780 fn dec_transaction_level ( ) {
78- NUM_LEVELS
79- . fetch_update ( Ordering :: SeqCst , Ordering :: SeqCst , |existing_levels| {
80- if existing_levels == 0 {
81- log:: warn!(
82- "We are underflowing with calculating transactional levels. Not great, but let's not panic..." ,
83- ) ;
84- None
85- } else {
86- // Cannot underflow because of checks above.
87- Some ( existing_levels - 1 )
88- }
89- } )
90- . ok ( ) ;
81+ let existing_levels = get_transaction_level ( ) ;
82+ if existing_levels == 0 {
83+ log:: warn!(
84+ "We are underflowing with calculating transactional levels. Not great, but let's not panic..." ,
85+ ) ;
86+ } else if existing_levels == 1 {
87+ // Don't leave any trace of this storage item.
88+ kill_transaction_level ( ) ;
89+ } else {
90+ // Cannot underflow because of checks above.
91+ set_transaction_level ( & ( existing_levels - 1 ) ) ;
92+ }
9193 }
9294
9395 pub fn is_transactional ( ) -> bool {
0 commit comments