Feedbacks

VincentLanglet · VincentLanglet · commit e96f8730f9ed · 2024-03-13T10:12:03.000+01:00
diff --git a/src/Symfony/Component/Messenger/Message/LockableMessageInterface.php b/src/Symfony/Component/Messenger/Message/LockableMessageInterface.php
@@ -17,4 +17,9 @@ interface LockableMessageInterface
      * Returns null if you want to force the dispatch of the message.
      */
     public function getKey(): ?string;
+
+    /**
+     * Should we release the lock before calling the handler or after.
+     */
+    public function shouldBeReleasedBeforeHandlerCall(): bool;
 }
diff --git a/src/Symfony/Component/Messenger/Message/TTLAwareLockableMessageInterface.php b/src/Symfony/Component/Messenger/Message/TTLAwareLockableMessageInterface.php
@@ -0,0 +1,17 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\Component\Messenger\Message;
+
+interface TTLAwareLockableMessageInterface extends LockableMessageInterface
+{
+    public function getTTL(): ?int;
+}
diff --git a/src/Symfony/Component/Messenger/Middleware/LockMiddleware.php b/src/Symfony/Component/Messenger/Middleware/LockMiddleware.php
@@ -14,6 +14,7 @@
 use Symfony\Component\Lock\LockFactory;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Message\LockableMessageInterface;
+use Symfony\Component\Messenger\Message\TTLAwareLockableMessageInterface;
 use Symfony\Component\Messenger\Stamp\LockStamp;
 use Symfony\Component\Messenger\Stamp\ReceivedStamp;
 
@@ -35,34 +36,50 @@ public function handle(Envelope $envelope, StackInterface $stack): Envelope
 
         $message = $envelope->getMessage();
 
-        // If we're trying to dispatch a lockable message.
-        if ($message instanceof LockableMessageInterface && null === $envelope->last(ReceivedStamp::class)) {
-            $key = $message->getKey();
+        if (null === $envelope->last(ReceivedStamp::class)) {
+            if ($message instanceof LockableMessageInterface) {
+                // If we're trying to dispatch a lockable message.
+                $key = $message->getKey();
 
-            if (null !== $key) {
-                // The acquire call must be done before stamping the message
-                // in order to have the full state of the key in the stamp.
-                $canAcquire = $this->lockFactory->createLock($key, autoRelease: false)->acquire();
+                if (null !== $key) {
+                    // The acquire call must be done before stamping the message
+                    // in order to have the full state of the key in the stamp.
+                    $lock = $message instanceof TTLAwareLockableMessageInterface
+                        ? $this->lockFactory->createLock($key, $message->getTTL(), autoRelease: false)
+                        : $this->lockFactory->createLock($key, autoRelease: false);
+                    $canAcquire = $lock->acquire();
 
-                $envelope = $envelope->with(new LockStamp($key));
-                if (!$canAcquire) {
-                    return $envelope;
+                    $envelope = $envelope->with(new LockStamp($key, $message->shouldBeReleasedBeforeHandlerCall()));
+                    if (!$canAcquire) {
+                        return $envelope;
+                    }
                 }
             }
+        } else {
+            $this->releaseLock($envelope, true);
         }
 
         try {
             $envelope = $stack->next()->handle($envelope, $stack);
         } finally {
             // If we've received a lockable message, we're releasing it.
             if (null !== $envelope->last(ReceivedStamp::class)) {
-                $stamp = $envelope->last(LockStamp::class);
-                if ($stamp instanceof LockStamp) {
-                    $this->lockFactory->createLockFromKey($stamp->getKey(), autoRelease: false)->release();
-                }
+                $this->releaseLock($envelope, false);
             }
         }
 
         return $envelope;
     }
+
+    private function releaseLock(Envelope $envelope, bool $beforeHandlerCall): void
+    {
+        $stamp = $envelope->last(LockStamp::class);
+        if ($stamp instanceof LockStamp && $stamp->shouldBeReleasedBeforHandlerCall() === $beforeHandlerCall) {
+            $message = $envelope->getMessage();
+            $lock = $message instanceof TTLAwareLockableMessageInterface
+                ? $this->lockFactory->createLockFromKey($stamp->getKey(), $message->getTTL(), autoRelease: false)
+                : $this->lockFactory->createLockFromKey($stamp->getKey(), autoRelease: false);
+            $lock->release();
+        }
+    }
 }
diff --git a/src/Symfony/Component/Messenger/Stamp/LockStamp.php b/src/Symfony/Component/Messenger/Stamp/LockStamp.php
@@ -16,14 +16,21 @@
 final class LockStamp implements StampInterface
 {
     private Key $key;
+    private bool $shouldBeReleasedBeforeHandlerCall;
 
-    public function __construct(Key $key)
+    public function __construct(Key $key, bool $shouldBeReleasedBeforeHandlerCall)
     {
         $this->key = $key;
+        $this->shouldBeReleasedBeforeHandlerCall = $shouldBeReleasedBeforeHandlerCall;
     }
 
     public function getKey(): Key
     {
         return $this->key;
     }
+
+    public function shouldBeReleasedBeforHandlerCall(): bool
+    {
+        return $this->shouldBeReleasedBeforeHandlerCall;
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -16,14 +16,21 @@`
`16`	`16`	`final class LockStamp implements StampInterface`
`17`	`17`	`{`
`18`	`18`	`private Key $key;`
	`19`	`+ private bool $shouldBeReleasedBeforeHandlerCall;`
`19`	`20`
`20`		`- public function __construct(Key $key)`
	`21`	`+ public function __construct(Key $key, bool $shouldBeReleasedBeforeHandlerCall)`
`21`	`22`	`{`
`22`	`23`	`$this->key = $key;`
	`24`	`+ $this->shouldBeReleasedBeforeHandlerCall = $shouldBeReleasedBeforeHandlerCall;`
`23`	`25`	`}`
`24`	`26`
`25`	`27`	`public function getKey(): Key`
`26`	`28`	`{`
`27`	`29`	`return $this->key;`
`28`	`30`	`}`
	`31`	`+`
	`32`	`+ public function shouldBeReleasedBeforHandlerCall(): bool`
	`33`	`+ {`
	`34`	`+ return $this->shouldBeReleasedBeforeHandlerCall;`
	`35`	`+ }`
`29`	`36`	`}`