(fix) Handling of batch messages when there are two batch handlers wi…

…th unacked messages.
symfony · ehoutsma · Oct 2, 2024 · Oct 2, 2024 · Oct 2, 2024 · e1e290c6f5eb4244365839cc8cd35a78d4b11a6a
commit e1e290c6f5eb4244365839cc8cd35a78d4b11a6a
diff --git a/src/Symfony/Component/Messenger/Tests/WorkerTest.php b/src/Symfony/Component/Messenger/Tests/WorkerTest.php
@@ -581,6 +581,52 @@ public function testFlushBatchOnStop()
         $this->assertSame($expectedMessages, $handler->processedMessages);
     }
 
+    public function testFlushMultipleBatchOnStop()
+    {
+        $expectedMessages = [
+            new DummyMessage('Hey'),
+        ];
+
+        $secondHandlerExpectedMessages = [
+            new SecondHandlerDummyMessage('Ho'),
+        ];
+
+        $receiver = new DummyReceiver([
+            [new Envelope($expectedMessages[0])],
+        ]);
+
+        $secondHandlerReceiver = new SecondMessageDummyReceiver([
+            [new Envelope($secondHandlerExpectedMessages[0])],
+        ]);
+
+        $handler = new DummyBatchHandler();
+        $secondHandler = new SecondDummyBatchHandler();
+
+        $middleware = new HandleMessageMiddleware(new HandlersLocator([
+            DummyMessage::class => [new HandlerDescriptor($handler)],
+            SecondHandlerDummyMessage::class => [new HandlerDescriptor($secondHandler)],
+        ]));
+
+        $bus = new MessageBus([$middleware]);
+
+        $dispatcher = new EventDispatcher();
+        $dispatcher->addListener(WorkerRunningEvent::class, function (WorkerRunningEvent $event) use ($receiver, $secondHandlerReceiver) {
+            static $i = 0;
+            if (1 < ++$i) {
+                $event->getWorker()->stop();
+            }
+
+            $this->assertSame(0, $receiver->getAcknowledgeCount());
+            $this->assertSame(0, $secondHandlerReceiver->getAcknowledgeCount());
+        });
+
+        $worker = new Worker([$receiver, $secondHandlerReceiver], $bus, $dispatcher, clock: new MockClock());
+        $worker->run();
+
+        $this->assertSame($expectedMessages, $handler->processedMessages);
+        $this->assertSame($secondHandlerExpectedMessages, $secondHandler->processedMessages);
+    }
+
     public function testGcCollectCyclesIsCalledOnMessageHandle()
     {
         $apiMessage = new DummyMessage('API');
@@ -622,7 +668,33 @@ public function __invoke(DummyMessage $message, ?Acknowledger $ack = null)
 
     private function shouldFlush(): bool
     {
-        return 2 <= \count($this->jobs);
+        return 5 <= \count($this->jobs);
+    }
+
+    private function process(array $jobs): void
+    {
+        $this->processedMessages = array_column($jobs, 0);
+
+        foreach ($jobs as [$job, $ack]) {
+            $ack->ack($job);
+        }
+    }
+}
+
+Class SecondDummyBatchHandler implements BatchHandlerInterface
+{
+    use BatchHandlerTrait;
+
+    public array $processedMessages;
+
+    public function __invoke(SecondHandlerDummyMessage $message, ?Acknowledger $ack = null)
+    {
+        return $this->handle($message, $ack);
+    }
+
+    private function shouldFlush(): bool
+    {
+        return 5 <= \count($this->jobs);
     }
 
     private function process(array $jobs): void
@@ -634,3 +706,11 @@ private function process(array $jobs): void
         }
     }
 }
+
+class SecondHandlerDummyMessage extends DummyMessage
+{
+}
+
+class SecondMessageDummyReceiver extends DummyReceiver
+{
+}
diff --git a/src/Symfony/Component/Messenger/Worker.php b/src/Symfony/Component/Messenger/Worker.php
@@ -254,15 +254,17 @@ private function flush(bool $force): bool
 
         $this->unacks = new \SplObjectStorage();
 
-        foreach ($unacks as $batchHandler) {
+        while ($unacks->valid()) {
+            $batchHandler = $unacks->current();
             [$envelope, $transportName] = $unacks[$batchHandler];
             try {
                 $this->bus->dispatch($envelope->with(new FlushBatchHandlersStamp($force)));
                 $envelope = $envelope->withoutAll(NoAutoAckStamp::class);
-                unset($unacks[$batchHandler], $batchHandler);
             } catch (\Throwable $e) {
                 $this->acks[] = [$transportName, $envelope, $e];
             }
+            $unacks->next();
+            $unacks->detach($batchHandler);
         }
 
         return $this->ack();