[Messenger][RateLimiter] fix additional message handled when using a rate limiter

Jean-Beru · Jean-Beru · commit 0ccdc527fd55 · 2024-11-05T10:27:02.000+01:00
diff --git a/src/Symfony/Component/Messenger/Tests/WorkerTest.php b/src/Symfony/Component/Messenger/Tests/WorkerTest.php
@@ -14,6 +14,7 @@
 use PHPUnit\Framework\TestCase;
 use Psr\EventDispatcher\EventDispatcherInterface;
 use Psr\Log\LoggerInterface;
+use Symfony\Bridge\PhpUnit\ClockMock;
 use Symfony\Component\Clock\MockClock;
 use Symfony\Component\EventDispatcher\EventDispatcher;
 use Symfony\Component\HttpKernel\DependencyInjection\ServicesResetter;
@@ -47,8 +48,8 @@
 use Symfony\Component\Messenger\Transport\Receiver\ReceiverInterface;
 use Symfony\Component\Messenger\Worker;
 use Symfony\Component\RateLimiter\RateLimiterFactory;
+use Symfony\Component\RateLimiter\Reservation;
 use Symfony\Component\RateLimiter\Storage\InMemoryStorage;
-use Symfony\Contracts\Service\ResetInterface;
 
 /**
  * @group time-sensitive
@@ -438,21 +439,21 @@ public function testWorkerRateLimitMessages()
         $envelope = [
             new Envelope(new DummyMessage('message1')),
             new Envelope(new DummyMessage('message2')),
+            new Envelope(new DummyMessage('message3')),
+            new Envelope(new DummyMessage('message4')),
         ];
         $receiver = new DummyReceiver([$envelope]);
 
         $bus = $this->createMock(MessageBusInterface::class);
         $bus->method('dispatch')->willReturnArgument(0);
 
         $eventDispatcher = new EventDispatcher();
-        $eventDispatcher->addSubscriber(new StopWorkerOnMessageLimitListener(2));
+        $eventDispatcher->addSubscriber(new StopWorkerOnMessageLimitListener(4));
 
         $rateLimitCount = 0;
-        $listener = function (WorkerRateLimitedEvent $event) use (&$rateLimitCount) {
+        $eventDispatcher->addListener(WorkerRateLimitedEvent::class, static function () use (&$rateLimitCount) {
             ++$rateLimitCount;
-            $event->getLimiter()->reset(); // Reset limiter to continue test
-        };
-        $eventDispatcher->addListener(WorkerRateLimitedEvent::class, $listener);
+        });
 
         $rateLimitFactory = new RateLimiterFactory([
             'id' => 'bus',
@@ -461,11 +462,14 @@ public function testWorkerRateLimitMessages()
             'interval' => '1 minute',
         ], new InMemoryStorage());
 
+        ClockMock::register(Reservation::class);
+        ClockMock::register(InMemoryStorage::class);
+
         $worker = new Worker(['bus' => $receiver], $bus, $eventDispatcher, null, ['bus' => $rateLimitFactory], new MockClock());
         $worker->run();
 
-        $this->assertCount(2, $receiver->getAcknowledgedEnvelopes());
-        $this->assertEquals(1, $rateLimitCount);
+        $this->assertSame(4, $receiver->getAcknowledgeCount());
+        $this->assertSame(3, $rateLimitCount);
     }
 
     public function testWorkerShouldLogOnStop()
diff --git a/src/Symfony/Component/Messenger/Worker.php b/src/Symfony/Component/Messenger/Worker.php
@@ -242,6 +242,7 @@ private function rateLimit(string $transportName): void
 
         $this->eventDispatcher?->dispatch(new WorkerRateLimitedEvent($rateLimiter, $transportName));
         $rateLimiter->reserve()->wait();
+        $rateLimiter->consume();
     }
 
     private function flush(bool $force): bool

Original file line number	Diff line number	Diff line change
`@@ -242,6 +242,7 @@ private function rateLimit(string $transportName): void`
`242`	`242`
`243`	`243`	`$this->eventDispatcher?->dispatch(new WorkerRateLimitedEvent($rateLimiter, $transportName));`
`244`	`244`	`$rateLimiter->reserve()->wait();`
	`245`	`+ $rateLimiter->consume();`
`245`	`246`	`}`
`246`	`247`
`247`	`248`	`private function flush(bool $force): bool`