Add IDistributedLockManager for Kubernetes environment (#7327)

Added IDistributedLockManager for Kubernetes environment.

Author: divyagandhisethi

release_notes.md

@@ -2,9 +2,11 @@
 <!-- Please add your release notes in the following format:
 - My change description (#PR)
 -->
+
 - Added a feature flag to opt out of the default behavior where the host sets the environment name to `Development` when running in debug mode. To disable the behavior, set the app setting: `AzureWebJobsFeatureFlags` to `DisableDevModeInDebug`
 - Reorder CORS and EasyAuth middleware to prevent EasyAuth from blocking CORS requests (#7315)
 - Updated PowerShell Worker (PS7) to [3.0.833](https://github.com/Azure/azure-functions-powershell-worker/releases/tag/v3.0.833)
+- Add IDistributedLockManager for Kubernetes environment (#7327)
 - Call EnableDrainModeAsync on ApplicationStopping (#7262)
 
 **Release sprint:** Sprint 100

src/WebJobs.Script/Environment/EnvironmentExtensions.cs

@@ -448,6 +448,11 @@ public static bool SupportsAzureFileShareMount(this IEnvironment environment)
                 StringComparison.OrdinalIgnoreCase);
         }
 
+        public static string GetHttpLeaderEndpoint(this IEnvironment environment)
+        {
+            return environment.GetEnvironmentVariableOrDefault(HttpLeaderEndpoint, string.Empty);
+        }
+
         public static bool DrainOnApplicationStoppingEnabled(this IEnvironment environment)
         {
             return !string.IsNullOrEmpty(environment.GetEnvironmentVariable(KubernetesServiceHost)) ||

src/WebJobs.Script/Environment/EnvironmentSettingNames.cs

@@ -65,6 +65,7 @@ public static class EnvironmentSettingNames
         public const string PodNamespace = "POD_NAMESPACE";
         public const string PodName = "POD_NAME";
         public const string PodEncryptionKey = "POD_ENCRYPTION_KEY";
+        public const string HttpLeaderEndpoint = "HTTP_LEADER_ENDPOINT";
 
         /// <summary>
         /// Environment variable dynamically set by the platform when it is safe to

src/WebJobs.Script/Host/Kubernetes/KubernetesClient.cs

@@ -0,0 +1,94 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using Newtonsoft.Json;
+
+namespace Microsoft.Azure.WebJobs.Script
+{
+    internal class KubernetesClient
+    {
+        private const int LeaseRenewDeadline = 10;
+        private readonly HttpClient _httpClient;
+        private readonly string _httpLeaderEndpoint;
+
+        internal KubernetesClient(IEnvironment environment, HttpClient httpClient = null)
+        {
+            _httpClient = httpClient ?? new HttpClient();
+            _httpLeaderEndpoint = environment.GetHttpLeaderEndpoint();
+        }
+
+        internal async Task<KubernetesLockHandle> GetLock(string lockName, CancellationToken cancellationToken)
+        {
+            if (string.IsNullOrEmpty(lockName))
+            {
+                throw new ArgumentNullException(nameof(lockName));
+            }
+
+            var request = new HttpRequestMessage()
+            {
+                Method = HttpMethod.Get,
+                RequestUri = GetRequestUri($"?name={lockName}")
+            };
+
+            var response = await _httpClient.SendAsync(request, cancellationToken);
+
+            response.EnsureSuccessStatusCode();
+
+            var responseString = await response.Content.ReadAsStringAsync();
+            return JsonConvert.DeserializeObject<KubernetesLockHandle>(responseString);
+        }
+
+        internal async Task<KubernetesLockHandle> TryAcquireLock (string lockId, string ownerId, TimeSpan lockPeriod, CancellationToken cancellationToken)
+        {
+            if (string.IsNullOrEmpty(lockId))
+            {
+                throw new ArgumentNullException(nameof(lockId));
+            }
+
+            if (string.IsNullOrEmpty(ownerId))
+            {
+                throw new ArgumentNullException(nameof(ownerId));
+            }
+
+            var lockHandle = new KubernetesLockHandle();
+            var request = new HttpRequestMessage()
+            {
+                Method = HttpMethod.Post,
+                RequestUri = GetRequestUri($"/acquire?name={lockId}&owner={ownerId}&duration={lockPeriod.TotalSeconds}&renewDeadline={LeaseRenewDeadline}"),
+            };
+
+            var response = await _httpClient.SendAsync(request, cancellationToken);
+            if (response.IsSuccessStatusCode)
+            {
+                lockHandle.LockId = lockId;
+                lockHandle.Owner = ownerId;
+            }
+            return lockHandle;
+        }
+
+        internal async Task<HttpResponseMessage> ReleaseLock(string lockId, string ownerId)
+        {
+            if (string.IsNullOrEmpty(lockId))
+            {
+                throw new ArgumentNullException(nameof(lockId));
+            }
+
+            var request = new HttpRequestMessage()
+            {
+                Method = HttpMethod.Post,
+                RequestUri = GetRequestUri($"/release?name={lockId}&owner={ownerId}")
+            };
+
+            return await _httpClient.SendAsync(request);
+        }
+
+        internal Uri GetRequestUri(string requestStem)
+        {
+            return new Uri($"{_httpLeaderEndpoint}/lock{requestStem}");
+        }
+    }
+}

src/WebJobs.Script/Host/Kubernetes/KubernetesDistributedLockManager.cs

@@ -0,0 +1,62 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Host;
+using Microsoft.Azure.WebJobs.Script.Config;
+
+namespace Microsoft.Azure.WebJobs.Script
+{
+    // This is an implementation of IDistributedLockManager to be used when running
+    // in Kubernetes environments.
+    internal class KubernetesDistributedLockManager : IDistributedLockManager
+    {
+        private readonly KubernetesClient _kubernetesClient;
+        private readonly string _websiteInstanceId;
+
+        public KubernetesDistributedLockManager(IEnvironment environment,
+            ScriptSettingsManager settingsManager)
+            : this(new KubernetesClient(environment), settingsManager)
+        {
+        }
+
+        internal KubernetesDistributedLockManager(KubernetesClient client, ScriptSettingsManager settingsManager)
+        {
+            _kubernetesClient = client;
+            _websiteInstanceId = settingsManager.AzureWebsiteInstanceId;
+        }
+
+        public async Task<string> GetLockOwnerAsync(string account, string lockId, CancellationToken cancellationToken)
+        {
+            var response = await _kubernetesClient.GetLock(lockId, cancellationToken);
+            return response.Owner;
+        }
+
+        public async Task ReleaseLockAsync(IDistributedLock lockHandle, CancellationToken cancellationToken)
+        {
+            var kubernetesLock = (KubernetesLockHandle)lockHandle;
+            var response = await _kubernetesClient.ReleaseLock(kubernetesLock.LockId, kubernetesLock.Owner);
+            response.EnsureSuccessStatusCode();
+        }
+
+        public async Task<bool> RenewAsync(IDistributedLock lockHandle, CancellationToken cancellationToken)
+        {
+            var kubernetesLock = (KubernetesLockHandle)lockHandle;
+            var renewedLockHandle = await _kubernetesClient.TryAcquireLock(kubernetesLock.LockId, kubernetesLock.Owner, TimeSpan.Parse(kubernetesLock.LockPeriod), cancellationToken);
+            return !string.IsNullOrEmpty(renewedLockHandle.LockId);
+        }
+
+        public async Task<IDistributedLock> TryLockAsync(string account, string lockId, string lockOwnerId, string proposedLeaseId, TimeSpan lockPeriod, CancellationToken cancellationToken)
+        {
+            var ownerId = string.IsNullOrEmpty(lockOwnerId) ? _websiteInstanceId : lockOwnerId;
+            var kubernetesLock = await _kubernetesClient.TryAcquireLock(lockId, ownerId, lockPeriod, cancellationToken);
+            if (string.IsNullOrEmpty(kubernetesLock.LockId))
+            {
+                return null;
+            }
+            return kubernetesLock;
+        }
+    }
+}

src/WebJobs.Script/Host/Kubernetes/KubernetesLockHandle.cs

@@ -0,0 +1,16 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using Microsoft.Azure.WebJobs.Host;
+
+namespace Microsoft.Azure.WebJobs.Script
+{
+    internal class KubernetesLockHandle : IDistributedLock
+    {
+        public string LockId { get; set; }
+
+        public string Owner { get; set; }
+
+        public string LockPeriod { get; set; }
+    }
+}

src/WebJobs.Script/ScriptHostBuilderExtensions.cs

@@ -6,6 +6,7 @@
 using System.Runtime.InteropServices;
 using Microsoft.ApplicationInsights.DependencyCollector;
 using Microsoft.ApplicationInsights.Extensibility;
+using Microsoft.Azure.WebJobs.Host;
 using Microsoft.Azure.WebJobs.Host.Executors;
 using Microsoft.Azure.WebJobs.Hosting;
 using Microsoft.Azure.WebJobs.Logging;
@@ -288,6 +289,12 @@ public static IHostBuilder AddScriptHostCore(this IHostBuilder builder, ScriptAp
                 }
 
                 services.AddSingleton<IHostedService, WorkerConsoleLogService>();
+
+                if (SystemEnvironment.Instance.IsKubernetesManagedHosting())
+                {
+                    services.AddSingleton<IDistributedLockManager, KubernetesDistributedLockManager>();
+                }
+
                 services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, PrimaryHostCoordinator>());
                 services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, FunctionInvocationDispatcherShutdownManager>());