2
2
3
3
use std:: rc:: Rc ;
4
4
use std:: cell:: RefCell ;
5
- use std:: collections:: { VecDeque , HashMap } ;
5
+ use std:: collections:: { VecDeque , HashMap , hash_map :: Entry } ;
6
6
use std:: sync:: mpsc:: { Sender , Receiver } ;
7
7
8
8
use bytes:: arc:: Bytes ;
@@ -28,7 +28,6 @@ pub struct ProcessBuilder {
28
28
peers : usize , // number of peer allocators.
29
29
pushers : Vec < Receiver < MergeQueue > > , // for pushing bytes at other workers.
30
30
pullers : Vec < Sender < MergeQueue > > , // for pulling bytes from other workers.
31
- // signal: Signal,
32
31
}
33
32
34
33
impl ProcessBuilder {
@@ -80,11 +79,11 @@ impl ProcessBuilder {
80
79
peers : self . peers ,
81
80
events : Rc :: new ( RefCell :: new ( VecDeque :: new ( ) ) ) ,
82
81
canaries : Rc :: new ( RefCell :: new ( Vec :: new ( ) ) ) ,
82
+ channel_id_bound : None ,
83
83
staged : Vec :: new ( ) ,
84
84
sends,
85
85
recvs,
86
86
to_local : HashMap :: new ( ) ,
87
- // _signal: self.signal,
88
87
}
89
88
}
90
89
}
@@ -108,7 +107,8 @@ pub struct ProcessAllocator {
108
107
109
108
canaries : Rc < RefCell < Vec < usize > > > ,
110
109
111
- // _signal: Signal,
110
+ channel_id_bound : Option < usize > ,
111
+
112
112
// sending, receiving, and responding to binary buffers.
113
113
staged : Vec < Bytes > ,
114
114
sends : Vec < Rc < RefCell < SendEndpoint < MergeQueue > > > > , // sends[x] -> goes to thread x.
@@ -121,6 +121,12 @@ impl Allocate for ProcessAllocator {
121
121
fn peers ( & self ) -> usize { self . peers }
122
122
fn allocate < T : Data > ( & mut self , identifier : usize ) -> ( Vec < Box < dyn Push < Message < T > > > > , Box < dyn Pull < Message < T > > > ) {
123
123
124
+ // Assume and enforce in-order identifier allocation.
125
+ if let Some ( bound) = self . channel_id_bound {
126
+ assert ! ( bound < identifier) ;
127
+ }
128
+ self . channel_id_bound = Some ( identifier) ;
129
+
124
130
let mut pushes = Vec :: < Box < dyn Push < Message < T > > > > :: new ( ) ;
125
131
126
132
for target_index in 0 .. self . peers ( ) {
@@ -158,11 +164,15 @@ impl Allocate for ProcessAllocator {
158
164
// Check for channels whose `Puller` has been dropped.
159
165
let mut canaries = self . canaries . borrow_mut ( ) ;
160
166
for dropped_channel in canaries. drain ( ..) {
161
- let dropped =
167
+ let _dropped =
162
168
self . to_local
163
169
. remove ( & dropped_channel)
164
170
. expect ( "non-existent channel dropped" ) ;
165
- assert ! ( dropped. borrow( ) . is_empty( ) ) ;
171
+ // Borrowed channels may be non-empty, if the dataflow was forcibly
172
+ // dropped. The contract is that if a dataflow is dropped, all other
173
+ // workers will drop the dataflow too, without blocking indefinitely
174
+ // on events from it.
175
+ // assert!(dropped.borrow().is_empty());
166
176
}
167
177
std:: mem:: drop ( canaries) ;
168
178
@@ -185,15 +195,23 @@ impl Allocate for ProcessAllocator {
185
195
let _ = peel. extract_to ( 40 ) ;
186
196
187
197
// Increment message count for channel.
198
+ // Safe to do this even if the channel has been dropped.
188
199
events. push_back ( ( header. channel , Event :: Pushed ( 1 ) ) ) ;
189
200
190
201
// Ensure that a queue exists.
191
- // We may receive data before allocating, and shouldn't block.
192
- self . to_local
193
- . entry ( header. channel )
194
- . or_insert_with ( || Rc :: new ( RefCell :: new ( VecDeque :: new ( ) ) ) )
195
- . borrow_mut ( )
196
- . push_back ( peel) ;
202
+ match self . to_local . entry ( header. channel ) {
203
+ Entry :: Vacant ( entry) => {
204
+ // We may receive data before allocating, and shouldn't block.
205
+ if self . channel_id_bound . map ( |b| b < header. channel ) . unwrap_or ( true ) {
206
+ entry. insert ( Rc :: new ( RefCell :: new ( VecDeque :: new ( ) ) ) )
207
+ . borrow_mut ( )
208
+ . push_back ( peel) ;
209
+ }
210
+ }
211
+ Entry :: Occupied ( mut entry) => {
212
+ entry. get_mut ( ) . borrow_mut ( ) . push_back ( peel) ;
213
+ }
214
+ }
197
215
}
198
216
else {
199
217
println ! ( "failed to read full header!" ) ;
0 commit comments