@@ -28,7 +28,7 @@ use std::fmt;
28
28
use std:: sync:: { Mutex , PoisonError } ;
29
29
30
30
use futures:: future;
31
- use rand:: Rng ;
31
+ use rand:: { self , rngs :: SmallRng , Error as RandError , Rng , SeedableRng } ;
32
32
33
33
use crate :: options:: { ConflictRangeType , MutationType , TransactionOption } ;
34
34
use crate :: tuple:: { PackError , Subspace } ;
@@ -41,6 +41,7 @@ pub enum HcaError {
41
41
PackError ( PackError ) ,
42
42
InvalidDirectoryLayerMetadata ,
43
43
PoisonError ,
44
+ RandError ( RandError ) ,
44
45
}
45
46
46
47
impl fmt:: Debug for HcaError {
@@ -52,6 +53,7 @@ impl fmt::Debug for HcaError {
52
53
write ! ( f, "invalid directory layer metadata" )
53
54
}
54
55
HcaError :: PoisonError => write ! ( f, "mutex poisoned" ) ,
56
+ HcaError :: RandError ( err) => err. fmt ( f) ,
55
57
}
56
58
}
57
59
}
@@ -71,6 +73,11 @@ impl<T> From<PoisonError<T>> for HcaError {
71
73
Self :: PoisonError
72
74
}
73
75
}
76
+ impl From < RandError > for HcaError {
77
+ fn from ( err : RandError ) -> Self {
78
+ Self :: RandError ( err)
79
+ }
80
+ }
74
81
75
82
impl TransactError for HcaError {
76
83
fn try_into_fdb_error ( self ) -> Result < FdbError , Self > {
@@ -114,7 +121,7 @@ impl HighContentionAllocator {
114
121
reverse : true ,
115
122
..RangeOption :: default ( )
116
123
} ;
117
- let mut rng = rand:: thread_rng ( ) ;
124
+ let mut rng = SmallRng :: from_rng ( & mut rand:: thread_rng ( ) ) ? ;
118
125
119
126
loop {
120
127
let kvs = trx. get_range ( & counters_range, 1 , true ) . await ?;
@@ -130,17 +137,18 @@ impl HighContentionAllocator {
130
137
let window = loop {
131
138
let counters_start = self . counters . subspace ( & start) ;
132
139
133
- let mutex_guard = self . allocation_mutex . lock ( ) ?;
134
- if window_advanced {
135
- trx. clear_range ( self . counters . bytes ( ) , counters_start. bytes ( ) ) ;
136
- trx. set_option ( TransactionOption :: NextWriteNoWriteConflictRange ) ?;
137
- trx. clear_range ( self . recent . bytes ( ) , self . recent . subspace ( & start) . bytes ( ) ) ;
138
- }
139
-
140
- // Increment the allocation count for the current window
141
- trx. atomic_op ( counters_start. bytes ( ) , ONE_BYTES , MutationType :: Add ) ;
142
- let count_future = trx. get ( counters_start. bytes ( ) , true ) ;
143
- drop ( mutex_guard) ;
140
+ let count_future = {
141
+ let _mutex_guard = self . allocation_mutex . lock ( ) ?;
142
+ if window_advanced {
143
+ trx. clear_range ( self . counters . bytes ( ) , counters_start. bytes ( ) ) ;
144
+ trx. set_option ( TransactionOption :: NextWriteNoWriteConflictRange ) ?;
145
+ trx. clear_range ( self . recent . bytes ( ) , self . recent . subspace ( & start) . bytes ( ) ) ;
146
+ } ;
147
+
148
+ // Increment the allocation count for the current window
149
+ trx. atomic_op ( counters_start. bytes ( ) , ONE_BYTES , MutationType :: Add ) ;
150
+ trx. get ( counters_start. bytes ( ) , true )
151
+ } ;
144
152
145
153
let count_value = count_future. await ?;
146
154
let count = if let Some ( count_value) = count_value {
@@ -171,12 +179,14 @@ impl HighContentionAllocator {
171
179
let candidate: i64 = rng. gen_range ( start, start + window) ;
172
180
let recent_candidate = self . recent . subspace ( & candidate) ;
173
181
174
- let mutex_guard = self . allocation_mutex . lock ( ) ?;
175
- let latest_counter = trx. get_range ( & counters_range, 1 , true ) ;
176
- let candidate_value = trx. get ( recent_candidate. bytes ( ) , false ) ;
177
- trx. set_option ( TransactionOption :: NextWriteNoWriteConflictRange ) ?;
178
- trx. set ( recent_candidate. bytes ( ) , & [ ] ) ;
179
- drop ( mutex_guard) ;
182
+ let ( latest_counter, candidate_value) = {
183
+ let _mutex_guard = self . allocation_mutex . lock ( ) ?;
184
+ let latest_counter = trx. get_range ( & counters_range, 1 , true ) ;
185
+ let candidate_value = trx. get ( recent_candidate. bytes ( ) , false ) ;
186
+ trx. set_option ( TransactionOption :: NextWriteNoWriteConflictRange ) ?;
187
+ trx. set ( recent_candidate. bytes ( ) , & [ ] ) ;
188
+ ( latest_counter, candidate_value)
189
+ } ;
180
190
181
191
let ( latest_counter, candidate_value) =
182
192
future:: try_join ( latest_counter, candidate_value) . await ?;
0 commit comments