88 / 95 lines covered (92.6%)
1unit module BDD::Behave::Parallel::Queue;23use BDD::Behave::Parallel::Distribution;4use BDD::Behave::Parallel::EventStream;5use BDD::Behave::Parallel::Manifest;67constant Bucket = BDD::Behave::Parallel::Distribution::Bucket;89# Queue protocol: the parent process maintains a FIFO queue of buckets,10# sorted by cost descending so the largest buckets dispatch first (this11# minimizes the chance that a slow bucket ends up as a worker's last task12# while the queue is otherwise empty). Workers pull one bucket at a time13# over stdin and emit a `bucket-done` event over stdout when each completes,14# at which point the parent dispatches the next bucket. When the queue15# is empty the parent sends SHUTDOWN to each worker.16#17# Stdin command grammar (line-delimited):18# BUCKET\t<id>\t<file>\t<loc1>,<loc2>,...19# SHUTDOWN20#21# Stdout events from the worker (line-delimited JSON, same as the LPT22# worker) plus two queue-specific events:23# {"type":"worker-ready"} -- emitted once at startup24# {"type":"bucket-done","id":"..."} -- emitted after each bucket2526class QueueScheduler is export {27 has @!queue;28 has Int $!total-dispatched = 0;29 has Int $!total-completed = 0;3031 method enqueue(Bucket $b --> Nil) {32 @!queue.push: $b;33 }3435 method enqueue-sorted(@buckets --> Nil) {36 @!queue.append: @buckets.sort({ $^b.cost <=> $^a.cost }).list;37 }3839 method next-bucket(--> Bucket) {40 return Bucket if @!queue.elems == 0;41 $!total-dispatched++;42 @!queue.shift;43 }4445 method mark-complete(--> Nil) {46 $!total-completed++;47 }4849 method has-pending(--> Bool) {50 @!queue.elems > 0;51 }5253 method pending-count(--> Int) {54 @!queue.elems;55 }5657 method dispatched(--> Int) { $!total-dispatched }58 method completed(--> Int) { $!total-completed }5960 method all-complete(--> Bool) {61 @!queue.elems == 0 && $!total-dispatched == $!total-completed;62 }63}6465our sub format-bucket-command(Bucket $b --> Str) is export {66 my $locs = $b.locations.join(',');67 "BUCKET\t{$b.id}\t{$b.file}\t$locs";68}6970our sub parse-bucket-command(Str $line --> Hash) is export {71 return %( :type<unknown> ) unless $line.defined;72 my $trimmed = $line.chomp;73 return %( :type<shutdown> ) if $trimmed eq 'SHUTDOWN';74 return %( :type<unknown> ) unless $trimmed.starts-with('BUCKET');7576 my @parts = $trimmed.split("\t");77 return %( :type<unknown> ) unless @parts.elems >= 4;7879 my @locations = @parts[3].split(',').grep(*.chars).list;80 %(81 :type<bucket>,82 :id(@parts[1]),83 :file(@parts[2]),84 :locations(@locations),85 );86}8788class QueueWorkerHandle is export {89 has Int $.index is required;90 has Proc::Async $.proc is required;91 has Promise $.exit-promise is required;92 has JsonLineParser $.parser = JsonLineParser.new;93 has Bool $.finished is rw = False;94 has Int $.exit-code is rw = 0;95 has Bool $.shutdown-sent is rw = False;96}9798class QueueWorkerPool is export {99 has Int $.worker-count is required;100 has @.worker-argv is required;101 has @.spec-files;102 has %.base-env;103 has IO::Path $.coverage-log-dir;104 has Bool $.coverage-counts = False;105 has &.on-event = sub ($wi, $event) { };106 has &.on-ready = sub ($wi) { };107 has &.on-done = sub ($wi, $id) { };108109 has Lock $!io-lock = Lock.new;110 has QueueWorkerHandle @!workers;111112 method launch(--> Nil) {113 for ^$!worker-count -> $i {114 my @argv = @!worker-argv.Slip;115 @argv.push: '--queue-worker';116 @argv.append: @!spec-files;117118 my %env = |%!base-env;119 %env<BEHAVE_WORKER_INDEX> = $i.Str;120 %env<BEHAVE_WORKER_COUNT> = $!worker-count.Str;121122 if $!coverage-log-dir.defined {123 %env<MVM_COVERAGE_LOG>124 = $!coverage-log-dir.add("worker-$i.raw").absolute;125 %env<MVM_COVERAGE_CONTROL> = $!coverage-counts ?? '2' !! '0';126 }127128 my $proc = Proc::Async.new(:w, |@argv);129130 my $handle = QueueWorkerHandle.new(131 :index($i),132 :$proc,133 :exit-promise(Promise.new),134 );135136 # The race we are avoiding: Proc::Async's start-promise resolves137 # when the *process* exits, but stdout tap callbacks are scheduled138 # on the thread pool and may not have run yet. If we let wait-all139 # return as soon as the process exit fires, late chunks can dispatch140 # event handlers after the parent has already moved on to summary141 # rendering, dropping example counts. The fix is to delay keeping142 # exit-promise until the stdout supply itself signals `done` (which143 # only happens after every chunk has been delivered to our tap).144 my $stdout-done = Promise.new;145 my $stderr-done = Promise.new;146147 $proc.stdout.tap(148 -> $chunk {149 try {150 my @events = $handle.parser.feed($chunk);151 for @events -> $event {152 self!handle-worker-event($i, $event);153 }154 CATCH { default { note "Queue worker $i stdout tap error: {.message}" } }155 }156 },157 done => { $stdout-done.keep },158 quit => { $stdout-done.keep },159 );160161 $proc.stderr.tap(162 -> $chunk { $*ERR.print($chunk) },163 done => { $stderr-done.keep },164 quit => { $stderr-done.keep },165 );166167 my $start-promise = $proc.start(:ENV(%env));168 Promise.allof($start-promise, $stdout-done, $stderr-done).then({169 try {170 my $proc-result = $start-promise.result;171 # Fold a signal death (exitcode 0, signal set) to 128+N so crash172 # checks downstream never read it as success.173 $handle.exit-code = $proc-result.signal174 ?? 128 + $proc-result.signal175 !! $proc-result.exitcode;176 note "Queue worker $i died with signal {$proc-result.signal}"177 if $proc-result.signal;178 my @tail = $handle.parser.flush;179 for @tail -> $event {180 self!handle-worker-event($i, $event);181 }182 CATCH { default { note "Queue worker $i exit handler error: {.message}" } }183 }184 $handle.finished = True;185 $handle.exit-promise.keep($handle.exit-code);186 });187188 @!workers.push: $handle;189 }190 }191192 method !handle-worker-event(Int $i, %event --> Nil) {193 my $type = (%event<type> // '').Str;194 given $type {195 when 'worker-ready' { &!on-ready($i) }196 when 'bucket-done' { &!on-done($i, (%event<id> // '').Str) }197 default { &!on-event($i, %event) }198 }199 }200201 method send-bucket(Int $i, $bucket --> Nil) {202 return unless 0 <= $i < @!workers.elems;203 my $handle = @!workers[$i];204 return if $handle.finished;205 return if $handle.shutdown-sent;206 my $cmd = format-bucket-command($bucket) ~ "\n";207 $!io-lock.protect: {208 try {209 $handle.proc.print($cmd);210 CATCH { default { note "Queue worker $i: failed to send bucket: {.message}" } }211 }212 }213 }214215 method send-shutdown(Int $i --> Nil) {216 return unless 0 <= $i < @!workers.elems;217 my $handle = @!workers[$i];218 return if $handle.finished;219 return if $handle.shutdown-sent;220 $handle.shutdown-sent = True;221 $!io-lock.protect: {222 try {223 # Write SHUTDOWN and then close the stdin pipe so the worker reads224 # EOF if it's already past the SHUTDOWN line. close-stdin must NOT225 # race with the worker's own exit (which would also close the pipe226 # from the other side and may drop in-flight stdout events on some227 # OSes), so we await the worker's read by giving it a beat.228 $handle.proc.print("SHUTDOWN\n");229 CATCH { default { note "Queue worker $i: failed to send shutdown: {.message}" } }230 }231 }232 }233234 method shutdown-all(--> Nil) {235 for ^$!worker-count -> $i {236 self.send-shutdown($i);237 }238 }239240 method wait-all(--> Nil) {241 await Promise.allof(@!workers.map(*.exit-promise));242 }243244 method workers(--> List) { @!workers.List }245}