@@ -25,7 +25,9 @@ pub trait Prioritized {
2525#[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
2626pub enum Priority {
2727 High ,
28+
2829 Normal ,
30+
2931 Low ,
3032}
3133
@@ -54,6 +56,7 @@ pub struct StealingQueue<TTask:Prioritized<Kind = Priority>> {
5456/// Contains all necessary components for a single worker thread to operate.
5557///
5658/// This includes the thread-local `Worker` deques, which are not safe to share,
59+
5760/// making this `Context` object the sole owner of a worker's private queues.
5861pub struct Context < TTask > {
5962 /// A unique identifier for the worker, used to avoid self-stealing.
@@ -79,56 +82,72 @@ impl<TTask:Prioritized<Kind = Priority>> StealingQueue<TTask> {
7982 /// 2. A `Vec` of `Context` objects, one for each worker thread to own.
8083 pub fn Create ( Count : usize ) -> ( Self , Vec < Context < TTask > > ) {
8184 let mut High : Vec < Worker < TTask > > = Vec :: with_capacity ( Count ) ;
85+
8286 let mut Normal : Vec < Worker < TTask > > = Vec :: with_capacity ( Count ) ;
87+
8388 let mut Low : Vec < Worker < TTask > > = Vec :: with_capacity ( Count ) ;
8489
8590 // For each priority level, create a thread-local worker queue and its
8691 // corresponding shared stealer.
8792 let StealerHigh : Vec < Stealer < TTask > > = ( 0 ..Count )
8893 . map ( |_| {
8994 let Worker = Worker :: new_fifo ( ) ;
95+
9096 let Stealer = Worker . stealer ( ) ;
97+
9198 High . push ( Worker ) ;
99+
92100 Stealer
93101 } )
94102 . collect ( ) ;
95103
96104 let StealerNormal : Vec < Stealer < TTask > > = ( 0 ..Count )
97105 . map ( |_| {
98106 let Worker = Worker :: new_fifo ( ) ;
107+
99108 let Stealer = Worker . stealer ( ) ;
109+
100110 Normal . push ( Worker ) ;
111+
101112 Stealer
102113 } )
103114 . collect ( ) ;
104115
105116 let StealerLow : Vec < Stealer < TTask > > = ( 0 ..Count )
106117 . map ( |_| {
107118 let Worker = Worker :: new_fifo ( ) ;
119+
108120 let Stealer = Worker . stealer ( ) ;
121+
109122 Low . push ( Worker ) ;
123+
110124 Stealer
111125 } )
112126 . collect ( ) ;
113127
114128 // Bundle all shared components into an Arc for safe sharing.
115129 let Share = Arc :: new ( Share {
116130 Injector : ( Injector :: new ( ) , Injector :: new ( ) , Injector :: new ( ) ) ,
131+
117132 Stealer : ( StealerHigh , StealerNormal , StealerLow ) ,
118133 } ) ;
119134
120135 // Create a unique context for each worker, giving it ownership of its
121136 // local queues and a reference to the shared components.
122137 let mut Contexts = Vec :: with_capacity ( Count ) ;
138+
123139 for Identifier in 0 ..Count {
124140 Contexts . push ( Context {
125141 Identifier ,
142+
126143 Local : ( High . remove ( 0 ) , Normal . remove ( 0 ) , Low . remove ( 0 ) ) ,
144+
127145 Share : Share . clone ( ) ,
128146 } ) ;
129147 }
130148
131149 let Queue = Self { Share } ;
150+
132151 ( Queue , Contexts )
133152 }
134153
@@ -138,7 +157,9 @@ impl<TTask:Prioritized<Kind = Priority>> StealingQueue<TTask> {
138157 pub fn Submit ( & self , Task : TTask ) {
139158 match Task . Rank ( ) {
140159 Priority :: High => self . Share . Injector . 0 . push ( Task ) ,
160+
141161 Priority :: Normal => self . Share . Injector . 1 . push ( Task ) ,
162+
142163 Priority :: Low => self . Share . Injector . 2 . push ( Task ) ,
143164 }
144165 }
@@ -169,8 +190,11 @@ impl<TTask> Context<TTask> {
169190 /// peer worker to ensure fair distribution and avoid contention hotspots.
170191 pub fn Steal < ' a > (
171192 & self ,
193+
172194 Injector : & ' a Injector < TTask > ,
195+
173196 Stealers : & ' a [ Stealer < TTask > ] ,
197+
174198 Local : & ' a Worker < TTask > ,
175199 ) -> Option < TTask > {
176200 // First, try to steal a batch from the global injector.
@@ -182,13 +206,15 @@ impl<TTask> Context<TTask> {
182206
183207 // If the global queue is empty, try stealing from peers.
184208 let Count = Stealers . len ( ) ;
209+
185210 if Count <= 1 {
186211 // Cannot steal if there are no other workers.
187212 return None ;
188213 }
189214
190215 // Allocation-free random iteration: pick a random starting point.
191216 let mut Rng = rand:: rng ( ) ;
217+
192218 let Start = Rng . random_range ( 0 ..Count ) ;
193219
194220 // Iterate through all peers starting from the random offset.
0 commit comments