[Messenger] Fix processing batches

HypeMC · HypeMC · commit 840a3d642ddb · 2025-12-26T18:22:23.000+01:00
diff --git a/src/Symfony/Component/Messenger/Handler/BatchHandlerTrait.php b/src/Symfony/Component/Messenger/Handler/BatchHandlerTrait.php
@@ -20,6 +20,10 @@ trait BatchHandlerTrait
 
     public function flush(bool $force): void
     {
+        if (!$force && !$this->shouldFlush()) {
+            return;
+        }
+
         if ($jobs = $this->jobs) {
             $this->jobs = [];
             $this->process($jobs);
diff --git a/src/Symfony/Component/Messenger/Tests/Middleware/HandleMessageMiddlewareTest.php b/src/Symfony/Component/Messenger/Tests/Middleware/HandleMessageMiddlewareTest.php
@@ -323,6 +323,54 @@ private function process(array $jobs): void
         $this->assertSame([$message], $handler->processedMessages);
     }
 
+    public function testBatchHandlerFlushFalseDoesNotFlushPartialBatch()
+    {
+        $handler = new class implements BatchHandlerInterface {
+            public array $processedMessages = [];
+
+            use BatchHandlerTrait;
+
+            public function __invoke(DummyMessage $message, ?Acknowledger $ack = null)
+            {
+                return $this->handle($message, $ack);
+            }
+
+            private function getBatchSize(): int
+            {
+                return 3;
+            }
+
+            private function process(array $jobs): void
+            {
+                $this->processedMessages = array_column($jobs, 0);
+
+                foreach ($jobs as [$job, $ack]) {
+                    $ack->ack($job);
+                }
+            }
+        };
+
+        $middleware = new HandleMessageMiddleware(new HandlersLocator([
+            DummyMessage::class => [new HandlerDescriptor($handler)],
+        ]));
+
+        $ack = static function () {};
+
+        $message = new DummyMessage('Hey');
+        $envelope = $middleware->handle(new Envelope($message, [new AckStamp($ack)]), new StackMiddleware());
+
+        $this->assertEmpty($handler->processedMessages);
+        $this->assertCount(1, $envelope->all(NoAutoAckStamp::class));
+
+        $handler->flush(false);
+
+        $this->assertEmpty($handler->processedMessages);
+
+        $handler->flush(true);
+
+        $this->assertSame([$message], $handler->processedMessages);
+    }
+
     public function testHandlerArgumentsStamp()
     {
         $message = new DummyMessage('Hey');
diff --git a/src/Symfony/Component/Messenger/Tests/WorkerTest.php b/src/Symfony/Component/Messenger/Tests/WorkerTest.php
@@ -545,15 +545,15 @@ public function testFlushBatchOnIdle()
             static $i = 0;
             if (1 < ++$i) {
                 $event->getWorker()->stop();
-                $this->assertSame(1, $receiver->getAcknowledgeCount());
-            } else {
-                $this->assertSame(0, $receiver->getAcknowledgeCount());
             }
+            $this->assertSame(0, $receiver->getAcknowledgeCount());
         });
 
         $worker = new Worker([$receiver], $bus, $dispatcher, clock: new MockClock());
         $worker->run();
 
+        $this->assertSame(1, $receiver->getAcknowledgeCount());
+
         $this->assertSame($expectedMessages, $handler->processedMessages);
     }
 
@@ -619,7 +619,7 @@ public function dispatch(object $message, array $stamps = []): Envelope
         $unacks = $unacksProperty->getValue($worker);
         $dummyHandler = new DummyBatchHandler();
         $envelopeWithNoAutoAck = $envelope->with(new NoAutoAckStamp(new HandlerDescriptor($dummyHandler)));
-        $unacks[$dummyHandler] = [$envelopeWithNoAutoAck, 'transport'];
+        $unacks[$dummyHandler] = [$envelopeWithNoAutoAck, 'transport', false];
 
         $worker->run();
 
diff --git a/src/Symfony/Component/Messenger/Worker.php b/src/Symfony/Component/Messenger/Worker.php
@@ -165,7 +165,7 @@ private function handleMessage(Envelope $envelope, string $transportName): void
         if (!$acked && !$noAutoAckStamp) {
             $this->acks[] = [$transportName, $envelope, $e];
         } elseif ($noAutoAckStamp) {
-            $this->unacks[$noAutoAckStamp->getHandlerDescriptor()->getBatchHandler()] = [$envelope->withoutAll(AckStamp::class), $transportName];
+            $this->unacks[$noAutoAckStamp->getHandlerDescriptor()->getBatchHandler()] = [$envelope->withoutAll(AckStamp::class), $transportName, &$acked];
         }
 
         $this->ack();
@@ -254,13 +254,23 @@ private function flush(bool $force): bool
         $this->unacks = new \SplObjectStorage();
 
         foreach ($unacks as $batchHandler) {
-            [$envelope, $transportName] = $unacks[$batchHandler];
+            [$envelope, $transportName, $acked] = $unacks[$batchHandler];
             try {
+                $e = null;
                 $this->bus->dispatch($envelope->with(new FlushBatchHandlersStamp($force)));
                 unset($unacks[$batchHandler], $batchHandler);
             } catch (\Throwable $e) {
                 $envelope = $envelope->withoutAll(NoAutoAckStamp::class);
                 $this->acks[] = [$transportName, $envelope, $e];
+                continue;
+            }
+
+            $noAutoAckStamp = $envelope->last(NoAutoAckStamp::class);
+
+            if (!$acked && !$noAutoAckStamp) {
+                $this->acks[] = [$transportName, $envelope, $e];
+            } elseif ($noAutoAckStamp) {
+                $this->unacks[$noAutoAckStamp->getHandlerDescriptor()->getBatchHandler()] = [$envelope->withoutAll(AckStamp::class), $transportName, &$acked];
             }
         }
 
