Cluster LeaveAsync() with CancellationToken support

src/core/Akka.API.Tests/CoreAPISpec.ApproveCluster.approved.txt

@@ -35,6 +35,7 @@ namespace Akka.Cluster
         public void JoinSeedNodes(System.Collections.Generic.IEnumerable<Akka.Actor.Address> seedNodes) { }
         public void Leave(Akka.Actor.Address address) { }
         public System.Threading.Tasks.Task LeaveAsync() { }
+        public System.Threading.Tasks.Task LeaveAsync(System.Threading.CancellationToken cancellationToken) { }
         public void RegisterOnMemberRemoved(System.Action callback) { }
         public void RegisterOnMemberUp(System.Action callback) { }
         public Akka.Actor.ActorPath RemotePathOf(Akka.Actor.IActorRef actorRef) { }

src/core/Akka.API.Tests/CoreAPISpec.ApproveCore.approved.txt

@@ -5128,6 +5128,7 @@ namespace Akka.Util.Internal
     public class static TaskExtensions
     {
         public static System.Threading.Tasks.Task<TResult> CastTask<TTask, TResult>(this System.Threading.Tasks.Task<TTask> task) { }
+        public static System.Threading.Tasks.Task WithCancellation(this System.Threading.Tasks.Task task, System.Threading.CancellationToken cancellationToken) { }
     }
 }
 namespace Akka.Util.Internal.Collections

src/core/Akka.Cluster.Tests/ClusterSpec.cs

@@ -8,6 +8,7 @@
 using System;
 using System.Collections.Immutable;
 using System.Linq;
+using System.Threading;
 using Akka.Actor;
 using Akka.Configuration;
 using Akka.TestKit;
@@ -185,6 +186,76 @@ public void A_cluster_must_complete_LeaveAsync_task_upon_being_removed()
             Cluster.Get(sys2).LeaveAsync().IsCompleted.Should().BeTrue();
         }
 
+        [Fact]
+        public void A_cluster_must_return_completed_LeaveAsync_task_if_member_already_removed()
+        {
+            // Join cluster
+            _cluster.Join(_selfAddress);
+            LeaderActions(); // Joining -> Up
+
+            // Subscribe to MemberRemoved and wait for confirmation
+            _cluster.Subscribe(TestActor, typeof(ClusterEvent.MemberRemoved));
+            ExpectMsg<ClusterEvent.CurrentClusterState>();
+
+            // Leave the cluster prior to calling LeaveAsync()
+            _cluster.Leave(_selfAddress);
+
+            Within(TimeSpan.FromSeconds(10), () =>
+            {
+                LeaderActions(); // Leaving --> Exiting
+                LeaderActions(); // Exiting --> Removed
+
+                // Member should leave
+                ExpectMsg<ClusterEvent.MemberRemoved>().Member.Address.Should().Be(_selfAddress);
+            });
+
+            // LeaveAsync() task expected to complete immediately
+            _cluster.LeaveAsync().IsCompleted.Should().BeTrue();
+        }
+
+        [Fact]
+        public void A_cluster_must_cancel_LeaveAsync_task_if_CancellationToken_fired_before_node_left()
+        {
+            // Join cluster
+            _cluster.Join(_selfAddress);
+            LeaderActions(); // Joining -> Up
+
+            // Subscribe to MemberRemoved and wait for confirmation
+            _cluster.Subscribe(TestActor, typeof(ClusterEvent.MemberRemoved));
+            ExpectMsg<ClusterEvent.CurrentClusterState>();
+
+            // Requesting leave with cancellation token
+            var cts = new CancellationTokenSource();
+            var task1 = _cluster.LeaveAsync(cts.Token);
+
+            // Requesting another leave without cancellation
+            var task2 = _cluster.LeaveAsync(new CancellationTokenSource().Token);
+
+            // Cancelling the first task
+            cts.Cancel();
+            task1.Should(t => t.IsCanceled, "Task should be cancelled.");
+
+            Within(TimeSpan.FromSeconds(10), () =>
+            {
+                // Second task should continue awaiting for cluster leave
+                task2.IsCompleted.Should().BeFalse();
+
+                // Waiting for leave
+                LeaderActions(); // Leaving --> Exiting
+                LeaderActions(); // Exiting --> Removed
+
+                // Member should leave even a task was cancelled
+                ExpectMsg<ClusterEvent.MemberRemoved>().Member.Address.Should().Be(_selfAddress);
+
+                // Second task should complete (not cancelled)
+                task2.Should(t => t.IsCompleted && !t.IsCanceled, "Task should be completed, but not cancelled.");
+            });
+
+            // Subsequent LeaveAsync() tasks expected to complete immediately (not cancelled)
+            var task3 = _cluster.LeaveAsync();
+            task3.Should(t => t.IsCompleted && !t.IsCanceled, "Task should be completed, but not cancelled.");
+        }
+
         [Fact]
         public void A_cluster_must_be_allowed_to_join_and_leave_with_local_address()
         {

src/core/Akka.Cluster/Cluster.cs

@@ -265,7 +265,12 @@ public void JoinSeedNodes(IEnumerable<Address> seedNodes)
         /// <param name="address">The address of the node leaving the cluster.</param>
         public void Leave(Address address)
         {
-            ClusterCore.Tell(new ClusterUserAction.Leave(FillLocal(address)));
+            if (FillLocal(address) == SelfAddress)
+            {
+                LeaveSelf();
+            }
+            else
+                ClusterCore.Tell(new ClusterUserAction.Leave(FillLocal(address)));
         }
 
         /// <summary>
@@ -274,20 +279,45 @@ public void Leave(Address address)
         /// Once the returned <see cref="Task"/> completes, it means that the member has successfully been removed
         /// from the cluster.
         /// </summary>
-        /// <returns>A <see cref="Task"/> that will return true upon the current node being removed from the cluster.</returns>
+        /// <returns>A <see cref="Task"/> that will return upon the current node being removed from the cluster.</returns>
         public Task LeaveAsync()
         {
-            var tcs = _leaveTask.Value;
+            return LeaveSelf();
+        }
+
+        /// <summary>
+        /// Causes the CURRENT node, i.e. the one calling this function, to leave the cluster.
+        /// 
+        /// Once the returned <see cref="Task"/> completes in completed or cancelled state, it means that the member has successfully been removed
+        /// from the cluster or cancellation token cancelled the task.
+        /// </summary>
+        /// <param name="cancellationToken">The cancellation token to cancel awaiting.</param>
+        /// <returns>A <see cref="Task"/> that will return upon the current node being removed from the cluster, or if await was cancelled.</returns>
+        /// <remarks>
+        /// The cancellation token doesn't cancel leave from the cluster, it only lets to give up on awating (by timeout for example).
+        /// </remarks>
+        public Task LeaveAsync(CancellationToken cancellationToken)
+        {
+            return LeaveSelf().WithCancellation(cancellationToken);
+        }
+
+        private Task _leaveTask;
+
+        private Task LeaveSelf()
+        {
+            var tcs = new TaskCompletionSource<object>();
+            var leaveTask = Interlocked.CompareExchange(ref _leaveTask, tcs.Task, null);
 
-            // short-circuit - check to see if we've already successfully left.
-            if (tcs.Task.IsCompleted)
-                return tcs.Task;
+            // It's assumed here that once the member left the cluster, it won't get back again.
+            // So, the member removal event being memoized in TaskCompletionSource and never reset.
+            if (leaveTask != null)
+                return leaveTask;
 
-            // Register it such that our TCS is automatically completed when we're removed
-            _clusterDaemons.Tell(new InternalClusterAction.AddOnMemberRemovedListener(() => tcs.TrySetResult(true)));
+            // Subscribe to MemberRemoved events
+            _clusterDaemons.Tell(new InternalClusterAction.AddOnMemberRemovedListener(() => tcs.TrySetResult(null)));
 
-            // Issue the leave command
-            Leave(SelfAddress);
+            // Send leave message
+            ClusterCore.Tell(new ClusterUserAction.Leave(SelfAddress));
 
             return tcs.Task;
         }
@@ -403,7 +433,6 @@ public ImmutableHashSet<string> SelfRoles
         /// </summary>
         public DefaultFailureDetectorRegistry<Address> FailureDetector { get { return _failureDetector; } }
 
-        private Lazy<TaskCompletionSource<bool>> _leaveTask = new Lazy<TaskCompletionSource<bool>>(() => new TaskCompletionSource<bool>(), LazyThreadSafetyMode.ExecutionAndPublication);
         /// <summary>
         /// TBD
         /// </summary>

src/core/Akka/Util/Internal/TaskExtensions.cs

@@ -6,13 +6,14 @@
 //-----------------------------------------------------------------------
 
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Akka.Util.Internal
 {
     /// <summary>
     /// INTERNAL API
-    /// 
+    ///
     /// Extensions for working with <see cref="Task"/> types
     /// </summary>
     public static class TaskExtensions
@@ -50,5 +51,26 @@ public static Task<TResult> CastTask<TTask, TResult>(this Task<TTask> task)
                 }, TaskContinuationOptions.ExecuteSynchronously);
             return tcs.Task;
         }
+
+        /// <summary>
+        /// Returns the task which completes with result of original task if cancellation token not canceled it before completion.
+        /// </summary>
+        /// <param name="task">The original task.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>The task which completes with result of original task or with cancelled state.</returns>
+        public static Task WithCancellation(this Task task, CancellationToken cancellationToken)
+        {
+            if (task.IsCompleted || !cancellationToken.CanBeCanceled)
+                return task;
+
+            var tcs = new TaskCompletionSource<object>();
+            var r = cancellationToken.Register(() => { tcs.SetCanceled(); }, false);
+
+            return Task.WhenAny(task, tcs.Task)
+                // Dispose subscription to cancellation token
+                .ContinueWith(t => { r.Dispose(); }, TaskContinuationOptions.ExecuteSynchronously)
+                // Check cancellation, to return task in cancelled state instead of completed
+                .ContinueWith(t => { }, cancellationToken, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
+        }
     }
 }