@@ -11,16 +11,36 @@ use concurrent_queue::ConcurrentQueue;
1111///
1212/// For ids that may be re-used, see [`IdFactoryWithReuse`].
1313pub struct IdFactory < T > {
14- next_id : AtomicU64 ,
15- max_id : u64 ,
14+ /// A value starting at 0 and incremented each time a new id is allocated. Regardless of the
15+ /// underlying type, a u64 is used to cheaply detect overflows.
16+ counter : AtomicU64 ,
17+ /// We've overflowed if the `counter > max_count`.
18+ max_count : u64 ,
19+ id_offset : u64 , // added to the value received from `counter`
1620 _phantom_data : PhantomData < T > ,
1721}
1822
1923impl < T > IdFactory < T > {
20- pub const fn new ( start : u64 , max : u64 ) -> Self {
24+ /// Create a factory for ids in the range `start..=max`.
25+ pub fn new ( start : T , max : T ) -> Self
26+ where
27+ T : Into < NonZeroU64 > + Ord ,
28+ {
29+ Self :: new_const ( start. into ( ) , max. into ( ) )
30+ }
31+
32+ /// Create a factory for ids in the range `start..=max`.
33+ ///
34+ /// Provides a less convenient API than [`IdFactory::new`], but skips a type conversion that
35+ /// would make the function non-const.
36+ pub const fn new_const ( start : NonZeroU64 , max : NonZeroU64 ) -> Self {
37+ assert ! ( start. get( ) < max. get( ) ) ;
2138 Self {
22- next_id : AtomicU64 :: new ( start) ,
23- max_id : max,
39+ // Always start `counter` at 0, don't use the value of `start` because `start` could be
40+ // close to `u64::MAX`.
41+ counter : AtomicU64 :: new ( 0 ) ,
42+ max_count : max. get ( ) - start. get ( ) ,
43+ id_offset : start. get ( ) ,
2444 _phantom_data : PhantomData ,
2545 }
2646 }
@@ -32,47 +52,80 @@ where
3252{
3353 /// Return a unique new id.
3454///
35- /// Panics (best-effort) if the id type overflows.
55+ /// Panics if the id type overflows.
3656pub fn get ( & self ) -> T {
37- let new_id = self . next_id . fetch_add ( 1 , Ordering :: Relaxed ) ;
57+ let count = self . counter . fetch_add ( 1 , Ordering :: Relaxed ) ;
58+
59+ #[ cfg( debug_assertions) ]
60+ {
61+ if count == u64:: MAX {
62+ // u64 counter is about to overflow -- this should never happen! A `u64` counter
63+ // starting at 0 should take decades to overflow on a single machine.
64+ //
65+ // This is unrecoverable because other threads may have already read the overflowed
66+ // value, so abort the entire process.
67+ std:: process:: abort ( )
68+ }
69+ }
3870
39- if new_id > self . max_id {
71+ // `max_count` might be something like `u32::MAX`. The extra bits of `u64` are useful to
72+ // detect overflows in that case. We assume the u64 counter is large enough to never
73+ // overflow.
74+ if count > self . max_count {
4075 panic ! (
41- "Max id limit hit while attempting to generate a unique {}" ,
76+ "Max id limit (overflow) hit while attempting to generate a unique {}" ,
4277 type_name:: <T >( ) ,
4378 )
4479 }
4580
46- // Safety: u64 will not overflow. This is *very* unlikely to happen (would take
47- // decades).
48- let new_id = unsafe { NonZeroU64 :: new_unchecked ( new_id) } ;
81+ let new_id_u64 = count + self . id_offset ;
82+ // Safety:
83+ // - `count` is assumed not to overflow.
84+ // - `id_offset` is a non-zero value.
85+ // - `id_offset + count < u64::MAX`.
86+ let new_id = unsafe { NonZeroU64 :: new_unchecked ( new_id_u64) } ;
4987
50- // Use the extra bits of the AtomicU64 as cheap overflow detection when the
51- // value is less than 64 bits.
5288 match new_id. try_into ( ) {
5389 Ok ( id) => id,
90+ // With any sane implementation of `TryFrom`, this shouldn't happen, as we've already
91+ // checked the `max_count` bound. (Could happen with the `new_const` constructor)
5492 Err ( _) => panic ! (
55- "Overflow detected while attempting to generate a unique {}" ,
56- type_name:: <T >( ) ,
93+ "Failed to convert NonZeroU64 value of {} into {}" ,
94+ new_id,
95+ type_name:: <T >( )
5796 ) ,
5897 }
5998 }
6099}
61100
62- /// An [`IdFactory`], but extended with a free list to allow for id reuse for
63- /// ids such as [`BackendJobId`][crate::backend::BackendJobId].
101+ /// An [`IdFactory`], but extended with a free list to allow for id reuse for ids such as
102+ /// [`BackendJobId`][crate::backend::BackendJobId].
64103pub struct IdFactoryWithReuse < T > {
65104 factory : IdFactory < T > ,
66105 free_ids : ConcurrentQueue < T > ,
67106}
68107
69- impl < T > IdFactoryWithReuse < T > {
70- pub const fn new ( start : u64 , max : u64 ) -> Self {
108+ impl < T > IdFactoryWithReuse < T >
109+ where
110+ T : Into < NonZeroU64 > + Ord ,
111+ {
112+ /// Create a factory for ids in the range `start..=max`.
113+ pub fn new ( start : T , max : T ) -> Self {
71114 Self {
72115 factory : IdFactory :: new ( start, max) ,
73116 free_ids : ConcurrentQueue :: unbounded ( ) ,
74117 }
75118 }
119+
120+ /// Create a factory for ids in the range `start..=max`. Provides a less convenient API than
121+ /// [`IdFactoryWithReuse::new`], but skips a type conversion that would make the function
122+ /// non-const.
123+ pub const fn new_const ( start : NonZeroU64 , max : NonZeroU64 ) -> Self {
124+ Self {
125+ factory : IdFactory :: new_const ( start, max) ,
126+ free_ids : ConcurrentQueue :: unbounded ( ) ,
127+ }
128+ }
76129}
77130
78131impl < T > IdFactoryWithReuse < T >
@@ -81,18 +134,18 @@ where
81134{
82135 /// Return a new or potentially reused id.
83136///
84- /// Panics (best-effort) if the id type overflows.
137+ /// Panics if the id type overflows.
85138pub fn get ( & self ) -> T {
86139 self . free_ids . pop ( ) . unwrap_or_else ( |_| self . factory . get ( ) )
87140 }
88141
89- /// Add an id to the free list, allowing it to be re-used on a subsequent
90- /// call to [`IdFactoryWithReuse::get`].
142+ /// Add an id to the free list, allowing it to be re-used on a subsequent call to
143+ /// [`IdFactoryWithReuse::get`].
91144///
92145/// # Safety
93146///
94- /// It must be ensured that the id is no longer used. Id must be a valid id
95- /// that was previously returned by ` get`.
147+ /// The id must no longer be used. Must be a valid id that was previously returned by
148+ /// [`IdFactoryWithReuse:: get`] .
96149pub unsafe fn reuse ( & self , id : T ) {
97150 let _ = self . free_ids . push ( id) ;
98151 }
@@ -105,12 +158,21 @@ mod tests {
105158 use super :: * ;
106159
107160 #[ test]
108- #[ should_panic( expected = "Overflow detected " ) ]
109- fn test_overflow ( ) {
110- let factory = IdFactory :: < NonZeroU8 > :: new ( 1 , u16 :: MAX as u64 ) ;
161+ #[ should_panic( expected = "Max id limit (overflow) " ) ]
162+ fn test_overflow_detection ( ) {
163+ let factory = IdFactory :: new ( NonZeroU8 :: MIN , NonZeroU8 :: MAX ) ;
111164 assert_eq ! ( factory. get( ) , NonZeroU8 :: new( 1 ) . unwrap( ) ) ;
112165 assert_eq ! ( factory. get( ) , NonZeroU8 :: new( 2 ) . unwrap( ) ) ;
113- for _i in 2 ..256 {
166+ for _ in 2 ..256 {
167+ factory. get ( ) ;
168+ }
169+ }
170+
171+ #[ test]
172+ #[ should_panic( expected = "Max id limit (overflow)" ) ]
173+ fn test_overflow_detection_near_u64_max ( ) {
174+ let factory = IdFactory :: new ( NonZeroU64 :: try_from ( u64:: MAX - 5 ) . unwrap ( ) , NonZeroU64 :: MAX ) ;
175+ for _ in 0 ..=6 {
114176 factory. get ( ) ;
115177 }
116178 }
0 commit comments