88use core:: fmt;
99use std:: collections:: hash_map:: Entry ;
1010use std:: collections:: HashMap ;
11- use std:: mem:: replace ;
11+ use std:: mem:: swap ;
1212use std:: ops:: { Add , AddAssign , DerefMut , Sub } ;
1313use std:: sync:: atomic:: { AtomicBool , AtomicI64 , AtomicU64 , Ordering } ;
1414use std:: sync:: { Arc , Mutex , RwLock } ;
6464 sorted_attribs : Mutex < HashMap < Vec < KeyValue > , Arc < A > > > ,
6565 /// Configuration for an Aggregator
6666 config : A :: InitConfig ,
67+ /// Swap with `sorted_attribs` on every `collect_and_reset`.
68+ for_collect_after_reset : Mutex < HashMap < Vec < KeyValue > , Arc < A > > > ,
6769}
6870
6971impl < A > ValueMap < A >
7678 tracker : A :: create ( & config) ,
7779 is_set : AtomicBool :: new ( false ) ,
7880 } ,
79- all_attribs : RwLock :: new ( HashMap :: new ( ) ) ,
80- sorted_attribs : Mutex :: new ( HashMap :: new ( ) ) ,
81+ all_attribs : RwLock :: new ( Default :: default ( ) ) ,
82+ sorted_attribs : Mutex :: new ( Default :: default ( ) ) ,
8183 config,
84+ for_collect_after_reset : Mutex :: new ( Default :: default ( ) ) ,
8285 }
8386 }
8487
@@ -168,11 +171,14 @@ where
168171 where
169172 MapFn : FnMut ( Vec < KeyValue > , A ) -> Res ,
170173 {
174+ let mut to_collect = self
175+ . for_collect_after_reset
176+ . lock ( )
177+ . unwrap_or_else ( |err| err. into_inner ( ) ) ;
171178 // reset sorted trackers so new attributes set will be written into new hashmap
172- let trackers = match self . sorted_attribs . lock ( ) {
179+ match self . sorted_attribs . lock ( ) {
173180 Ok ( mut trackers) => {
174- let new = HashMap :: with_capacity ( trackers. len ( ) ) ;
175- replace ( trackers. deref_mut ( ) , new)
181+ swap ( trackers. deref_mut ( ) , to_collect. deref_mut ( ) ) ;
176182 }
177183 Err ( _) => return ,
178184 } ;
@@ -182,7 +188,7 @@ where
182188 Err ( _) => return ,
183189 } ;
184190
185- prepare_data ( dest, trackers . len ( ) ) ;
191+ prepare_data ( dest, to_collect . len ( ) ) ;
186192
187193 if self . no_attribs . is_set . swap ( false , Ordering :: AcqRel ) {
188194 dest. push ( map_fn (
@@ -191,7 +197,7 @@ where
191197 ) ) ;
192198 }
193199
194- for ( attrs, tracker) in trackers . into_iter ( ) {
200+ for ( attrs, tracker) in to_collect . drain ( ) {
195201 let tracker = Arc :: into_inner ( tracker) . expect ( "the only instance" ) ;
196202 dest. push ( map_fn ( attrs, tracker) ) ;
197203 }
0 commit comments