From 2176e58bd4e3e77357376fde0b3a6a436757d22f Mon Sep 17 00:00:00 2001
From: Pankaj Saini <pasaini@microsoft.com>
Date: Wed, 10 Apr 2024 14:01:17 -0700
Subject: [PATCH] Adding logs around semaphore waitasync and release.

---
 ...zureStorageOrchestrationServiceSettings.cs |  5 +++
 .../DurableTask.AzureStorage.csproj           |  2 +-
 .../Storage/AzureStorageClient.cs             | 36 +++++++++++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs
index 609a8b35d..c23144eb1 100644
--- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs
+++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs
@@ -123,6 +123,11 @@ public class AzureStorageOrchestrationServiceSettings
         /// </summary>
         public int MaxStorageOperationConcurrency { get; set; } = Environment.ProcessorCount * 25;
 
+        /// <summary>
+        /// Max wait time to wait for semaphore to process after which log should be added.
+        /// </summary>
+        public int MaxWaitForSemaphoreLoggingInSeconds { get; set; } = 20;
+
         /// <summary>
         /// Gets the maximum number of orchestrator actions to checkpoint at a time.
         /// </summary>
diff --git a/src/DurableTask.AzureStorage/DurableTask.AzureStorage.csproj b/src/DurableTask.AzureStorage/DurableTask.AzureStorage.csproj
index 7e281163d..49efcaad8 100644
--- a/src/DurableTask.AzureStorage/DurableTask.AzureStorage.csproj
+++ b/src/DurableTask.AzureStorage/DurableTask.AzureStorage.csproj
@@ -24,7 +24,7 @@
     <PatchVersion>1</PatchVersion>
 
     <VersionPrefix>$(MajorVersion).$(MinorVersion).$(PatchVersion)</VersionPrefix>
-    <FileVersion>$(VersionPrefix).0</FileVersion>
+    <FileVersion>$(VersionPrefix)-semaphorelogs.4</FileVersion>
     <!-- FileVersionRevision is expected to be set by the CI. This is useful for distinguishing between multiple builds of the same version. -->
     <FileVersion Condition="'$(FileVersionRevision)' != ''">$(VersionPrefix).$(FileVersionRevision)</FileVersion>
     <!-- The assembly version is only the major/minor pair, making it easier to do in-place upgrades -->
diff --git a/src/DurableTask.AzureStorage/Storage/AzureStorageClient.cs b/src/DurableTask.AzureStorage/Storage/AzureStorageClient.cs
index 004f09c71..92cc41b38 100644
--- a/src/DurableTask.AzureStorage/Storage/AzureStorageClient.cs
+++ b/src/DurableTask.AzureStorage/Storage/AzureStorageClient.cs
@@ -14,6 +14,7 @@
 namespace DurableTask.AzureStorage.Storage
 {
     using System;
+    using System.Diagnostics;
     using System.Threading;
     using System.Threading.Tasks;
     using DurableTask.AzureStorage.Monitoring;
@@ -127,9 +128,16 @@ public Task MakeTableStorageRequest(Func<OperationContext, CancellationToken, Ta
 
         private async Task<T> MakeStorageRequest<T>(Func<OperationContext, CancellationToken, Task<T>> storageRequest, string accountName, string operationName, string? clientRequestId = null, bool force = false)
         {
+            Guid guid = Guid.NewGuid();
+
             if (!force)
             {
+                var cTkWait = new CancellationTokenSource();
+                _ = WaitAndLog(guid, "Wait", cTkWait.Token);
+
                 await requestThrottleSemaphore.WaitAsync();
+
+                cTkWait.Cancel();
             }
 
             try
@@ -144,7 +152,12 @@ private async Task<T> MakeStorageRequest<T>(Func<OperationContext, CancellationT
             {
                 if (!force)
                 {
+                    var cTkRelease = new CancellationTokenSource();
+                    _ = WaitAndLog(guid, "Release", cTkRelease.Token);
+
                     requestThrottleSemaphore.Release();
+
+                    cTkRelease.Cancel();
                 }
             }
         }
@@ -160,5 +173,28 @@ private Task MakeStorageRequest(Func<OperationContext, CancellationToken, Task>
 
         private static string GetAccountName(StorageCredentials credentials, AzureStorageOrchestrationServiceSettings settings, StorageUri serviceUri, string service) =>
             credentials.AccountName ?? settings.StorageAccountDetails?.AccountName ?? serviceUri.GetAccountName(service) ?? "(unknown)";
+        
+        private async Task WaitAndLog(Guid guid, string waitOrRelease, CancellationToken cTk)
+        {
+            if (this.Settings.MaxWaitForSemaphoreLoggingInSeconds >= 0)
+            {
+                await Task.Delay(this.Settings.MaxWaitForSemaphoreLoggingInSeconds * 1000, cTk);
+
+                if (!cTk.IsCancellationRequested)
+                {
+                    StackTrace stackTrace = new StackTrace();
+
+                    this.Settings?.Logger.PartitionManagerInfo(
+                        this.Settings?.StorageAccountDetails?.AccountName,
+                        this.Settings?.TaskHubName,
+                        this.Settings?.WorkerId,
+                        string.Empty,
+                        $"Before {waitOrRelease} Id:{guid.ToString()}, processId: {Thread.CurrentThread.ManagedThreadId}" +
+                        $"| Semaphore currentCount = {requestThrottleSemaphore.CurrentCount}" +
+                        $"| Semaphore initialCount = {this.Settings?.MaxStorageOperationConcurrency}" +
+                        $"| StackTrace : {stackTrace.ToString()}");
+                }
+            }
+        }
     }
 }