Merge pull request #43 from raisedapp/develop

Develop To Master

Author: felixclase

src/main/Hangfire.Storage.SQLite/Entities/DistributedLock.cs

@@ -23,6 +23,12 @@ public class DistributedLock
         [Unique]
         public string Resource { get; set; }
 
+        /// <summary>
+        /// The owner key for this resource.
+        /// Prevents race conditions and changes to locks that are owned by other entities.
+        /// </summary>
+        public string ResourceKey { get; set; }
+
         /// <summary>
         /// The timestamp for when the lock expires.
         /// This is used if the lock is not maintained or 

src/main/Hangfire.Storage.SQLite/HangfireSQLiteConnection.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Text;
 using System.Threading;
 using Hangfire.Common;
 using Hangfire.Server;
@@ -330,10 +329,15 @@ public override void Heartbeat(string serverId)
 
             var server = DbContext.HangfireServerRepository.FirstOrDefault(_ => _.Id == serverId);
             if (server == null)
-                return;
+	            throw new BackgroundServerGoneException();
 
             server.LastHeartbeat = DateTime.UtcNow;
-            DbContext.Database.Update(server);
+            var affected = DbContext.Database.Update(server);
+            
+            if (affected == 0)
+            {
+	            throw new BackgroundServerGoneException();
+            }
         }
 
         public override void RemoveServer(string serverId)

src/main/Hangfire.Storage.SQLite/SQLiteDistributedLock.cs

@@ -19,6 +19,7 @@ private static readonly ThreadLocal<Dictionary<string, int>> AcquiredLocks
                     = new ThreadLocal<Dictionary<string, int>>(() => new Dictionary<string, int>(StringComparer.OrdinalIgnoreCase));
 
         private readonly string _resource;
+        private readonly string _resourceKey;
 
         private readonly HangfireDbContext _dbContext;
 
@@ -44,6 +45,7 @@ public SQLiteDistributedLock(string resource, TimeSpan timeout, HangfireDbContex
             _resource = resource ?? throw new ArgumentNullException(nameof(resource));
             _dbContext = database ?? throw new ArgumentNullException(nameof(database));
             _storageOptions = storageOptions ?? throw new ArgumentNullException(nameof(storageOptions));
+            _resourceKey = Guid.NewGuid().ToString();
 
             if (string.IsNullOrEmpty(resource))
             {
@@ -92,7 +94,7 @@ public void Dispose()
             }
 
             // Timer callback may be invoked after the Dispose method call,
-            // so we are using lock to avoid unsynchronized calls.
+            // but since we use the resource key, we will not disturb other owners.
             AcquiredLocks.Value.Remove(_resource);
 
             if (_heartbeatTimer != null)
@@ -110,32 +112,33 @@ private void Acquire(TimeSpan timeout)
         {
             try
             {
-                // If result is null, then it means we acquired the lock
                 var isLockAcquired = false;
                 var now = DateTime.UtcNow;
                 var lockTimeoutTime = now.Add(timeout);
 
-                while (!isLockAcquired && (lockTimeoutTime >= now))
+                while (lockTimeoutTime >= now)
                 {
-                    DistributedLock result;
+                    Cleanup();
 
                     lock (EventWaitHandleName)
                     {
-                        result = _dbContext.DistributedLockRepository.FirstOrDefault(_ => _.Resource == _resource);
-                        var distributedLock = result ?? new DistributedLock();
+                        var result = _dbContext.DistributedLockRepository.FirstOrDefault(_ => _.Resource == _resource);
 
-                        if (string.IsNullOrWhiteSpace(distributedLock.Id))
-                            distributedLock.Id = Guid.NewGuid().ToString();
-
-                        distributedLock.Resource = _resource;
-                        distributedLock.ExpireAt = DateTime.UtcNow.Add(_storageOptions.DistributedLockLifetime);
-
-                        var rowsAffected = _dbContext.Database.Update(distributedLock);
-                        if (rowsAffected == 0)
+                        if (result == null)
                         {
                             try
                             {
+                                var distributedLock = new DistributedLock();
+                                distributedLock.Id = Guid.NewGuid().ToString();
+                                distributedLock.Resource = _resource;
+                                distributedLock.ResourceKey = _resourceKey;
+                                distributedLock.ExpireAt = DateTime.UtcNow.Add(_storageOptions.DistributedLockLifetime);
+
                                 _dbContext.Database.Insert(distributedLock);
+
+                                // we were able to acquire the lock - break the loop
+                                isLockAcquired = true;
+                                break;
                             }
                             catch (SQLiteException e) when (e.Result == SQLite3.Result.Constraint)
                             {
@@ -145,24 +148,17 @@ private void Acquire(TimeSpan timeout)
                         }
                     }
 
-                    // If result is null, then it means we acquired the lock
-                    if (result == null)
-                    {
-                        isLockAcquired = true;
-                    }
-                    else
-                    {
-                        var waitTime = (int)timeout.TotalMilliseconds / 10;
-                        lock (EventWaitHandleName)
-                            Monitor.Wait(EventWaitHandleName, waitTime);
+                    // we couldn't acquire the lock - wait a bit and try again
+                    var waitTime = (int)timeout.TotalMilliseconds / 10;
+                    lock (EventWaitHandleName)
+                        Monitor.Wait(EventWaitHandleName, waitTime);
 
-                        now = DateTime.UtcNow;
-                    }
+                    now = DateTime.UtcNow;
                 }
 
                 if (!isLockAcquired)
                 {
-                    throw new DistributedLockTimeoutException($"Could not place a lock on the resource \'{_resource}\': The lock request timed out.");
+                    throw new DistributedLockTimeoutException(_resource);
                 }
             }
             catch (DistributedLockTimeoutException ex)
@@ -183,8 +179,8 @@ private void Release()
         {
             try
             {
-                // Remove resource lock
-                _dbContext.DistributedLockRepository.Delete(_ => _.Resource == _resource);
+                // Remove resource lock (if it's still ours)
+                _dbContext.DistributedLockRepository.Delete(_ => _.Resource == _resource && _.ResourceKey == _resourceKey);
                 lock (EventWaitHandleName)
                     Monitor.Pulse(EventWaitHandleName);
             }
@@ -198,7 +194,7 @@ private void Cleanup()
         {
             try
             {
-                // Delete expired locks
+                // Delete expired locks (of any owner)
                 _dbContext.DistributedLockRepository.
                     Delete(x => x.Resource == _resource && x.ExpireAt < DateTime.UtcNow);
             }
@@ -218,13 +214,20 @@ private void StartHeartBeat()
             _heartbeatTimer = new Timer(state =>
             {
                 // Timer callback may be invoked after the Dispose method call,
-                // so we are using lock to avoid unsynchronized calls.
+                // but since we use the resource key, we will not disturb other owners.
                 try
                 {
-                    var distributedLock = _dbContext.DistributedLockRepository.FirstOrDefault(x => x.Resource == _resource);
-                    distributedLock.ExpireAt = DateTime.UtcNow.Add(_storageOptions.DistributedLockLifetime);
+                    var distributedLock = _dbContext.DistributedLockRepository.FirstOrDefault(x => x.Resource == _resource && x.ResourceKey == _resourceKey);
+                    if (distributedLock != null)
+                    {
+                        distributedLock.ExpireAt = DateTime.UtcNow.Add(_storageOptions.DistributedLockLifetime);
 
-                    _dbContext.Database.Update(distributedLock);
+                        _dbContext.Database.Update(distributedLock);
+                    }
+                    else
+                    {
+                        Logger.ErrorFormat("Unable to update heartbeat on the resource '{0}'. The resource is not locked or is locked by another owner.", _resource);
+                    }
                 }
                 catch (Exception ex)
                 {

src/test/Hangfire.Storage.SQLite.Test/SQLiteConnectionFacts.cs

@@ -8,7 +8,6 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
-using System.Text;
 using System.Threading;
 using Xunit;
 
@@ -627,6 +626,13 @@ public void RemoveServer_RemovesAServerRecord()
             });
         }
 
+        [Fact, CleanDatabase]
+        public void Heartbeat_ThrowsBackgroundServerGoneException_WhenGivenServerDoesNotExist()
+        {
+	        UseConnection((database, connection) => Assert.Throws<BackgroundServerGoneException>(
+		        () => connection.Heartbeat(Guid.NewGuid().ToString())));
+        }
+        
         [Fact, CleanDatabase]
         public void Heartbeat_ThrowsAnException_WhenServerIdIsNull()
         {

src/test/Hangfire.Storage.SQLite.Test/SQLiteDistributedLockFacts.cs

@@ -201,6 +201,30 @@ public void Ctor_SetLockExpireAtWorks_WhenResourceIsNotLocked()
             });
         }
 
+        // see https://github.com/raisedapp/Hangfire.Storage.SQLite/issues/38
+        [Fact, CleanDatabase]
+        public void Ctor_SetLockExpireAtWorks_WhenResourceIsLockedAndExpiring()
+        {
+            UseConnection(database =>
+            {
+                // add a lock (taken by another process who is now killed) which will expire in 3 seconds from now
+                database.Database.Insert(new DistributedLock
+                {
+                    Id = Guid.NewGuid().ToString(),
+                    Resource = "resource1",
+                    ExpireAt = DateTime.UtcNow.Add(TimeSpan.FromSeconds(3))
+                });
+
+                // try to get the lock in the next 10 seconds
+                // ideally, after ~3 seconds, the constructor should succeed
+                using (new SQLiteDistributedLock("resource1", TimeSpan.FromSeconds(10), database, new SQLiteStorageOptions() { DistributedLockLifetime = TimeSpan.FromSeconds(3) }))
+                {
+                    DistributedLock lockEntry = database.DistributedLockRepository.FirstOrDefault(_ => _.Resource == "resource1");
+                    Assert.NotNull(lockEntry);
+                }
+            });
+        }
+
         private static void UseConnection(Action<HangfireDbContext> action)
         {
             var connection = ConnectionUtils.CreateConnection();