Add tryLock method

amphp · Dec 10, 2024 · dafeb20 · dafeb20
1 parent c2839b5
commit dafeb20
Show file tree

Hide file tree

Showing 10 changed files with 155 additions and 36 deletions.
diff --git a/composer-require-check.json b/composer-require-check.json
@@ -21,6 +21,7 @@
         "filesystem",
         "openFile",
         "lock",
+        "tryLock",
         "unlock",
         "eio_cancel",
         "eio_chmod",

diff --git a/src/Driver/BlockingFile.php b/src/Driver/BlockingFile.php
@@ -73,6 +73,16 @@ public function lock(LockType $type, ?Cancellation $cancellation = null): void
         $this->lockMode = $type;
     }
 
+    public function tryLock(LockType $type): bool
+    {
+        $locked = Internal\tryLock($this->path, $this->getFileHandle(), $type);
+        if ($locked) {
+            $this->lockMode = $type;
+        }
+
+        return $locked;
+    }
+
     public function unlock(): void
     {
         Internal\unlock($this->path, $this->getFileHandle());

diff --git a/src/Driver/ParallelFile.php b/src/Driver/ParallelFile.php
@@ -152,6 +152,16 @@ public function lock(LockType $type, ?Cancellation $cancellation = null): void
         $this->lockMode = $type;
     }
 
+    public function tryLock(LockType $type): bool
+    {
+        $locked = $this->flock('try-lock', $type);
+        if ($locked) {
+            $this->lockMode = $type;
+        }
+
+        return $locked;
+    }
+
     public function unlock(): void
     {
         $this->flock('unlock');
@@ -163,7 +173,7 @@ public function getLockType(): ?LockType
         return $this->lockMode;
     }
 
-    private function flock(string $action, ?LockType $type = null, ?Cancellation $cancellation = null): void
+    private function flock(string $action, ?LockType $type = null, ?Cancellation $cancellation = null): bool
     {
         if ($this->id === null) {
             throw new ClosedException("The file has been closed");
@@ -172,7 +182,7 @@ private function flock(string $action, ?LockType $type = null, ?Cancellation $ca
         $this->busy = true;
 
         try {
-            $this->worker->execute(new Internal\FileTask('flock', [$type, $action], $this->id), $cancellation);
+            return $this->worker->execute(new Internal\FileTask('flock', [$type, $action], $this->id), $cancellation);
         } catch (TaskFailureException $exception) {
             throw new StreamException("Attempting to lock the file failed", 0, $exception);
         } catch (WorkerException $exception) {

diff --git a/src/Driver/StatusCachingFile.php b/src/Driver/StatusCachingFile.php
@@ -59,6 +59,11 @@ public function lock(LockType $type, ?Cancellation $cancellation = null): void
         $this->file->lock($type, $cancellation);
     }
 
+    public function tryLock(LockType $type): bool
+    {
+        return $this->file->tryLock($type);
+    }
+
     public function unlock(): void
     {
         $this->file->unlock();

diff --git a/src/File.php b/src/File.php
@@ -6,7 +6,6 @@
 use Amp\ByteStream\ReadableStream;
 use Amp\ByteStream\WritableStream;
 use Amp\Cancellation;
-use Amp\Sync\Lock;
 
 interface File extends ReadableStream, WritableStream
 {
@@ -58,13 +57,23 @@ public function getMode(): string;
     public function truncate(int $size): void;
 
     /**
-     * Non-blocking method to obtain a shared or exclusive lock on the file.
+     * Non-blocking method to obtain a shared or exclusive lock on the file. This method must only return once
+     * the lock has been obtained. Use {@see tryLock()} to make a single attempt to get the lock.
      *
      * @throws FilesystemException If there is an error when attempting to lock the file.
      * @throws ClosedException If the file has been closed.
      */
     public function lock(LockType $type, ?Cancellation $cancellation = null): void;
 
+    /**
+     * Make a single non-blocking attempt to obtain a shared or exclusive lock on the file. Returns true if the lock
+     * was obtained, otherwise false. Use {@see lock()} to return only once the lock is obtained.
+     *
+     * @throws FilesystemException If there is an error when attempting to lock the file.
+     * @throws ClosedException If the file has been closed.
+     */
+    public function tryLock(LockType $type): bool;
+
     /**
      * @throws FilesystemException If there is an error when attempting to unlock the file.
      * @throws ClosedException If the file has been closed.

diff --git a/src/Internal/FileTask.php b/src/Internal/FileTask.php
@@ -106,17 +106,20 @@ public function run(Channel $channel, Cancellation $cancellation): mixed
                     switch ($action) {
                         case 'lock':
                             $file->lock($type, $cancellation);
-                            break;
+                            return true;
+
+                        case 'try-lock':
+                            return $file->tryLock($type);
 
                         case 'unlock':
                             $file->unlock();
-                            break;
+                            return true;
 
                         default:
                             throw new \Error("Invalid lock action - " . $action);
                     }
 
-                    return null;
+                    return false; // CS fixer fails without this return.
 
                 default:
                     throw new \Error('Invalid operation');

diff --git a/src/Internal/QueuedWritesFile.php b/src/Internal/QueuedWritesFile.php
@@ -133,13 +133,23 @@ abstract protected function getFileHandle();
 
     public function lock(LockType $type, ?Cancellation $cancellation = null): void
     {
-        lock($this->getPath(), $this->getFileHandle(), $type, $cancellation);
+        lock($this->path, $this->getFileHandle(), $type, $cancellation);
         $this->lockMode = $type;
     }
 
+    public function tryLock(LockType $type): bool
+    {
+        $locked = tryLock($this->path, $this->getFileHandle(), $type);
+        if ($locked) {
+            $this->lockMode = $type;
+        }
+
+        return $locked;
+    }
+
     public function unlock(): void
     {
-        unlock($this->getPath(), $this->getFileHandle());
+        unlock($this->path, $this->getFileHandle());
         $this->lockMode = null;
     }
 

diff --git a/src/Internal/functions.php b/src/Internal/functions.php
@@ -16,44 +16,57 @@
  */
 function lock(string $path, $handle, LockType $type, ?Cancellation $cancellation): void
 {
-    static $latencyTimeout = 0.01;
-    static $delayLimit = 1;
+    for ($attempt = 0; true; ++$attempt) {
+        if (tryLock($path, $handle, $type)) {
+            return;
+        }
 
-    $error = null;
-    $errorHandler = static function (int $type, string $message) use (&$error): bool {
-        $error = $message;
-        return true;
-    };
+        // Exponential back-off with a maximum delay of 1 second.
+        delay(\min(1, 0.01 * (2 ** $attempt)), cancellation: $cancellation);
+    }
+}
 
+/**
+ * @internal
+ *
+ * @param resource $handle
+ *
+ * @throws FilesystemException
+ */
+function tryLock(string $path, $handle, LockType $type): bool
+{
     $flags = \LOCK_NB | match ($type) {
         LockType::Shared => \LOCK_SH,
         LockType::Exclusive => \LOCK_EX,
     };
 
-    for ($attempt = 0; true; ++$attempt) {
-        \set_error_handler($errorHandler);
-        try {
-            $lock = \flock($handle, $flags, $wouldBlock);
-        } finally {
-            \restore_error_handler();
-        }
+    $error = null;
+    \set_error_handler(static function (int $type, string $message) use (&$error): bool {
+        $error = $message;
+        return true;
+    });
 
-        if ($lock) {
-            return;
-        }
+    try {
+        $lock = \flock($handle, $flags, $wouldBlock);
+    } finally {
+        \restore_error_handler();
+    }
 
-        if (!$wouldBlock) {
-            throw new FilesystemException(
-                \sprintf(
-                    'Error attempting to lock file at "%s": %s',
-                    $path,
-                    $error ?? 'Unknown error',
-                )
-            );
-        }
+    if ($lock) {
+        return true;
+    }
 
-        delay(\min($delayLimit, $latencyTimeout * (2 ** $attempt)), cancellation: $cancellation);
+    if (!$wouldBlock) {
+        throw new FilesystemException(
+            \sprintf(
+                'Error attempting to lock file at "%s": %s',
+                $path,
+                $error ?? 'Unknown error',
+            )
+        );
     }
+
+    return false;
 }
 
 /**

diff --git a/test/AsyncFileTest.php b/test/AsyncFileTest.php
@@ -9,6 +9,7 @@
 use Amp\File\PendingOperationError;
 use Revolt\EventLoop;
 use function Amp\async;
+use function Amp\delay;
 
 abstract class AsyncFileTest extends FileTest
 {
@@ -128,4 +129,30 @@ public function testSimultaneousLock(): void
         $handle1->close();
         $handle2->close();
     }
+
+    public function testTryLockLoop(): void
+    {
+        $this->setMinimumRuntime(0.1);
+        $this->setTimeout(0.3);
+
+        $path = Fixture::path() . "/lock";
+        $handle1 = $this->driver->openFile($path, "c+");
+        $handle2 = $this->driver->openFile($path, "c+");
+
+        self::assertTrue($handle1->tryLock(LockType::Exclusive));
+        self::assertSame(LockType::Exclusive, $handle1->getLockType());
+
+        EventLoop::delay(0.1, $handle1->unlock(...));
+
+        $future = async(function () use ($handle2): void {
+            while (!$handle2->tryLock(LockType::Exclusive)) {
+                delay(0.1);
+            }
+        });
+
+        $future->await();
+
+        $handle1->close();
+        $handle2->close();
+    }
 }
diff --git a/test/FileTest.php b/test/FileTest.php
@@ -354,6 +354,37 @@ public function testUnlockAfterClose(): void
         $handle->unlock();
     }
 
+    public function testTryLock(): void
+    {
+        $path = Fixture::path() . "/lock";
+        $handle1 = $this->driver->openFile($path, "c+");
+        $handle2 = $this->driver->openFile($path, "c+");
+
+        self::assertTrue($handle1->tryLock(LockType::Exclusive));
+        self::assertSame(LockType::Exclusive, $handle1->getLockType());
+
+        self::assertFalse($handle2->tryLock(LockType::Exclusive));
+        self::assertNull($handle2->getLockType());
+
+        $handle1->unlock();
+        self::assertNull($handle1->getLockType());
+
+        self::assertTrue($handle2->tryLock(LockType::Shared));
+        self::assertSame(LockType::Shared, $handle2->getLockType());
+
+        self::assertTrue($handle1->tryLock(LockType::Shared));
+        self::assertSame(LockType::Shared, $handle1->getLockType());
+
+        self::assertFalse($handle1->tryLock(LockType::Exclusive));
+        self::assertSame(LockType::Shared, $handle1->getLockType());
+
+        $handle2->unlock();
+        self::assertNull($handle2->getLockType());
+
+        self::assertTrue($handle1->tryLock(LockType::Exclusive));
+        self::assertSame(LockType::Exclusive, $handle1->getLockType());
+    }
+
     abstract protected function createDriver(): File\FilesystemDriver;
 
     protected function setUp(): void