From 2a70861ae9fe33334ec09ce5cb835c71637262da Mon Sep 17 00:00:00 2001
From: Richard K <richard@fonky.dev>
Date: Mon, 9 Sep 2024 12:16:48 +0200
Subject: [PATCH] Azure: Handle blob lease acquire errors from response instead
 of exception

---
 .../AzureBlobLeaseDistributedLock.cs           | 18 ++++++++++--------
 .../BlobLeaseClientWrapper.cs                  | 18 ++++++++++++------
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/src/DistributedLock.Azure/AzureBlobLeaseDistributedLock.cs b/src/DistributedLock.Azure/AzureBlobLeaseDistributedLock.cs
index 32a67f8e..18b3cc07 100644
--- a/src/DistributedLock.Azure/AzureBlobLeaseDistributedLock.cs
+++ b/src/DistributedLock.Azure/AzureBlobLeaseDistributedLock.cs
@@ -12,7 +12,7 @@ namespace Medallion.Threading.Azure;
 public sealed partial class AzureBlobLeaseDistributedLock : IInternalDistributedLock<AzureBlobLeaseDistributedLockHandle>
 {
     /// <summary>
-    /// Metadata marker used to indicate that a blob was created for distributed locking and therefore 
+    /// Metadata marker used to indicate that a blob was created for distributed locking and therefore
     /// should be destroyed upon release
     /// </summary>
     private static readonly string CreatedMetadataKey = $"__DistributedLock";
@@ -53,7 +53,7 @@ internal static string GetSafeName(string name, BlobContainerClient blobContaine
 
         return DistributedLockHelpers.ToSafeName(name, maxLength, s => ConvertToValidName(s));
 
-        // check based on 
+        // check based on
         // https://docs.microsoft.com/en-us/azure/storage/common/storage-use-emulator#connect-to-the-emulator-account-using-the-well-known-account-name-and-key
         bool IsStorageEmulator() => blobContainerClient.Uri.IsAbsoluteUri
             && blobContainerClient.Uri.AbsoluteUri.StartsWith("http://127.0.0.1:10000/devstoreaccount1", StringComparison.Ordinal);
@@ -105,13 +105,15 @@ static string ConvertToValidName(string name)
         );
 
     private async ValueTask<AzureBlobLeaseDistributedLockHandle?> TryAcquireAsync(
-        BlobLeaseClientWrapper leaseClient, 
+        BlobLeaseClientWrapper leaseClient,
         CancellationToken cancellationToken,
         bool isRetryAfterCreate)
     {
-        try  { await leaseClient.AcquireAsync(this._options.duration, cancellationToken).ConfigureAwait(false); }
-        catch (RequestFailedException acquireException)
+        var response = await leaseClient.AcquireAsync(this._options.duration, cancellationToken).ConfigureAwait(false);
+        if (response.IsError)
         {
+            var acquireException = new RequestFailedException(response);
+
             if (acquireException.ErrorCode == AzureErrors.LeaseAlreadyPresent) { return null; }
 
             if (acquireException.ErrorCode == AzureErrors.BlobNotFound)
@@ -144,7 +146,7 @@ static string ConvertToValidName(string name)
                 }
             }
 
-            throw;
+            throw acquireException;
         }
 
         var shouldDeleteBlob = isRetryAfterCreate
@@ -160,7 +162,7 @@ internal sealed class InternalHandle : IDistributedSynchronizationHandle, LeaseM
         private readonly bool _ownsBlob;
         private readonly AzureBlobLeaseDistributedLock _lock;
         private readonly LeaseMonitor _leaseMonitor;
-        
+
         public InternalHandle(BlobLeaseClientWrapper leaseClient, bool ownsBlob, AzureBlobLeaseDistributedLock @lock)
         {
             this._leaseClient = leaseClient;
@@ -193,7 +195,7 @@ public async ValueTask DisposeAsync()
             {
                 await this._lock._blobClient.DeleteIfExistsAsync(leaseId: this._leaseClient.LeaseId).ConfigureAwait(false);
             }
-            else 
+            else
             {
                 await this._leaseClient.ReleaseAsync().ConfigureAwait(false);
             }
diff --git a/src/DistributedLock.Azure/BlobLeaseClientWrapper.cs b/src/DistributedLock.Azure/BlobLeaseClientWrapper.cs
index 11b13ec9..8bffc41f 100644
--- a/src/DistributedLock.Azure/BlobLeaseClientWrapper.cs
+++ b/src/DistributedLock.Azure/BlobLeaseClientWrapper.cs
@@ -1,4 +1,5 @@
-using Azure.Storage.Blobs.Specialized;
+using Azure;
+using Azure.Storage.Blobs.Specialized;
 using Medallion.Threading.Internal;
 
 namespace Medallion.Threading.Azure;
@@ -17,14 +18,19 @@ public BlobLeaseClientWrapper(BlobLeaseClient blobLeaseClient)
 
     public string LeaseId => this._blobLeaseClient.LeaseId;
 
-    public ValueTask AcquireAsync(TimeoutValue duration, CancellationToken cancellationToken)
+    public ValueTask<Response> AcquireAsync(TimeoutValue duration, CancellationToken cancellationToken)
     {
+        var requestContext = new RequestContext
+        {
+            CancellationToken = cancellationToken,
+            ErrorOptions = ErrorOptions.NoThrow
+        };
+
         if (SyncViaAsync.IsSynchronous)
         {
-            this._blobLeaseClient.Acquire(duration.TimeSpan, cancellationToken: cancellationToken);
-            return default;
+            return new ValueTask<Response>(this._blobLeaseClient.Acquire(duration.TimeSpan, conditions: null, requestContext));
         }
-        return new ValueTask(this._blobLeaseClient.AcquireAsync(duration.TimeSpan, cancellationToken: cancellationToken));
+        return new ValueTask<Response>(this._blobLeaseClient.AcquireAsync(duration.TimeSpan, conditions: null, requestContext));
     }
 
     public ValueTask RenewAsync(CancellationToken cancellationToken)
@@ -46,4 +52,4 @@ public ValueTask ReleaseAsync()
         }
         return new ValueTask(this._blobLeaseClient.ReleaseAsync());
     }
-}
+}
\ No newline at end of file