Fixing race condition during token generation (#5306)

* Fixing race condition during token generation

* Correcting changes

* Fixing function name and error handling

* Minor fixes

* Restructuring the UpdateStateOnServer function

Author: surajitshil-03

src/Agent.Worker/ExecutionContext.cs

@@ -268,13 +268,15 @@ public IExecutionContext CreateChild(Guid recordId, string displayName, string r
             return child;
         }
 
+
         public void Start(string currentOperation = null)
         {
             _record.CurrentOperation = currentOperation ?? _record.CurrentOperation;
             _record.StartTime = DateTime.UtcNow;
             _record.State = TimelineRecordState.InProgress;
 
-            _jobServerQueue.QueueTimelineRecordUpdate(_mainTimelineId, _record);
+            //update the state immediately on server
+            _jobServerQueue.UpdateStateOnServer(_mainTimelineId, _record);
 
             if (_logsStreamingOptions.HasFlag(LogsStreamingOptions.StreamToFiles))
             {
@@ -819,7 +821,8 @@ private void InitializeTimelineRecord(Guid timelineId, Guid timelineRecordId, Gu
             _record.WorkerName = configuration.GetSettings().AgentName;
             _record.Variables.Add(TaskWellKnownItems.AgentVersionTimelineVariable, BuildConstants.AgentPackage.Version);
 
-            _jobServerQueue.QueueTimelineRecordUpdate(_mainTimelineId, _record);
+            //update the state immediately on server
+            _jobServerQueue.UpdateStateOnServer(_mainTimelineId, _record);
         }
 
         private void JobServerQueueThrottling_EventReceived(object sender, ThrottlingEventArgs data)

src/Microsoft.VisualStudio.Services.Agent/JobServerQueue.cs

@@ -22,10 +22,12 @@ public interface IJobServerQueue : IAgentService, IThrottlingReporter
         bool ForceDrainTimelineQueue { get; set; }
         event EventHandler<ThrottlingEventArgs> JobServerQueueThrottling;
         Task ShutdownAsync();
+        Task SendTimelineRecordUpdateAsync(Guid timelineId, TimelineRecord timelineRecord);
         void Start(Pipelines.AgentJobRequestMessage jobRequest);
         void QueueWebConsoleLine(Guid stepRecordId, string line, long lineNumber);
         void QueueFileUpload(Guid timelineId, Guid timelineRecordId, string type, string name, string path, bool deleteSource);
         void QueueTimelineRecordUpdate(Guid timelineId, TimelineRecord timelineRecord);
+        void UpdateStateOnServer(Guid timelineId, TimelineRecord timelineRecord);
         void UpdateWebConsoleLineRate(Int32 rateInMillis);
     }
 
@@ -232,6 +234,42 @@ public void QueueTimelineRecordUpdate(Guid timelineId, TimelineRecord timelineRe
             Trace.Verbose(StringUtil.Format("Enqueue timeline {0} update queue: {1}", timelineId, timelineRecord.Id));
             _timelineUpdateQueue[timelineId].Enqueue(timelineRecord.Clone());
         }
+        public void UpdateStateOnServer(Guid timelineId, TimelineRecord timelineRecord)
+        {
+            ArgUtil.NotEmpty(timelineId, nameof(timelineId));
+            ArgUtil.NotNull(timelineRecord, nameof(timelineRecord));
+            ArgUtil.NotEmpty(timelineRecord.Id, nameof(timelineRecord.Id));
+
+            //sending immediate server update for the job timeline records to server
+            if (string.Equals(timelineRecord.RecordType, "Job", StringComparison.OrdinalIgnoreCase))
+            {
+                try
+                {
+                    // Attempting to send immediate update for job records
+                    SendTimelineRecordUpdateAsync(timelineId, timelineRecord).GetAwaiter().GetResult();
+                }
+                catch (Exception ex)
+                {
+                    Trace.Warning($"Failed to send immediate timeline record update: {ex.Message}. Falling back to queue mechanism.");
+                    QueueTimelineRecordUpdate(timelineId, timelineRecord);
+                }
+            }
+            else
+            {
+                // All other record types use queue mechanism
+                QueueTimelineRecordUpdate(timelineId, timelineRecord);
+                Trace.Verbose($"Timeline record {timelineRecord.Id} queued for update (RecordType: {timelineRecord.RecordType})");
+            }
+        }
+        public async Task SendTimelineRecordUpdateAsync(Guid timelineId, TimelineRecord timelineRecord)
+        {
+            ArgUtil.NotEmpty(timelineId, nameof(timelineId));
+            ArgUtil.NotNull(timelineRecord, nameof(timelineRecord));
+            ArgUtil.NotEmpty(timelineRecord.Id, nameof(timelineRecord.Id));
+            var jobtimelinerecord = new List<TimelineRecord> { timelineRecord.Clone() };
+            await _jobServer.UpdateTimelineRecordsAsync(_scopeIdentifier, _hubName, _planId, timelineId, jobtimelinerecord, CancellationToken.None);
+            Trace.Info($"Job timeline record {timelineRecord.Id} sent successfully to server");
+        }
 
         public void ReportThrottling(TimeSpan delay, DateTime expiration)
         {