Additional attributes for worker AI agent (#8387)

RohitRanjanMS · web-flow · commit 2360468596d0 · 2022-05-16T14:42:17.000-07:00
Additional attributes for worker AI agent
diff --git a/src/WebJobs.Script.Grpc/Channel/GrpcWorkerChannel.cs b/src/WebJobs.Script.Grpc/Channel/GrpcWorkerChannel.cs
@@ -5,13 +5,15 @@
 using System.Collections;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Reactive.Linq;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Dataflow;
+using Google.Protobuf.Collections;
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
@@ -487,6 +489,7 @@ internal async Task SendInvocationRequest(ScriptInvocationContext context)
                         return;
                     }
                     var invocationRequest = await context.ToRpcInvocationRequest(_workerChannelLogger, _workerCapabilities, _isSharedMemoryDataTransferEnabled, _sharedMemoryManager);
+                    AddAdditionalTraceContext(invocationRequest.TraceContext.Attributes, context);
                     _executingInvocations.TryAdd(invocationRequest.InvocationId, context);
 
                     SendStreamingMessage(new StreamingMessage
@@ -1002,5 +1005,27 @@ private void AddSample<T>(List<T> samples, T sample)
                 samples.Add(sample);
             }
         }
+
+        private void AddAdditionalTraceContext(MapField<string, string> attributes, ScriptInvocationContext context)
+        {
+            // This is only applicable for AI agents running along side worker
+            if (_environment.IsApplicationInsightsAgentEnabled())
+            {
+                attributes[ScriptConstants.LogPropertyProcessIdKey] = Convert.ToString(_rpcWorkerProcess.Id);
+                if (context.FunctionMetadata.Properties.ContainsKey(ScriptConstants.LogPropertyHostInstanceIdKey))
+                {
+                    attributes[ScriptConstants.LogPropertyHostInstanceIdKey] = Convert.ToString(context.FunctionMetadata.Properties[ScriptConstants.LogPropertyHostInstanceIdKey]);
+                }
+                if (context.FunctionMetadata.Properties.ContainsKey(LogConstants.CategoryNameKey))
+                {
+                    attributes[LogConstants.CategoryNameKey] = Convert.ToString(context.FunctionMetadata.Properties[LogConstants.CategoryNameKey]);
+                }
+                string sessionid = Activity.Current?.GetBaggageItem(ScriptConstants.LiveLogsSessionAIKey);
+                if (!string.IsNullOrEmpty(sessionid))
+                {
+                    attributes[ScriptConstants.LiveLogsSessionAIKey] = sessionid;
+                }
+            }
+        }
     }
 }
diff --git a/src/WebJobs.Script/Environment/EnvironmentExtensions.cs b/src/WebJobs.Script/Environment/EnvironmentExtensions.cs
@@ -16,7 +16,8 @@ namespace Microsoft.Azure.WebJobs.Script
 {
     internal static class EnvironmentExtensions
     {
-        // For testing
+        private static bool? isApplicationInsightsAgentEnabled;
+
         internal static string BaseDirectory { get; set; }
 
         public static string GetEnvironmentVariableOrDefault(this IEnvironment environment, string name, string defaultValue)
@@ -508,5 +509,21 @@ public static HashSet<string> GetLanguageWorkerListToStartInPlaceholder(this IEn
             }
             return placeholderRuntimeSet;
         }
+
+        public static bool IsApplicationInsightsAgentEnabled(this IEnvironment environment)
+        {
+            // cache the value of the environment variable
+            if (isApplicationInsightsAgentEnabled.HasValue)
+            {
+                return isApplicationInsightsAgentEnabled.Value;
+            }
+            else if (!environment.IsPlaceholderModeEnabled())
+            {
+                bool.TryParse(environment.GetEnvironmentVariable(AppInsightsAgent), out bool isEnabled);
+                isApplicationInsightsAgentEnabled = isEnabled;
+                return isApplicationInsightsAgentEnabled.Value;
+            }
+            return false;
+        }
     }
 }
diff --git a/src/WebJobs.Script/Environment/EnvironmentSettingNames.cs b/src/WebJobs.Script/Environment/EnvironmentSettingNames.cs
@@ -27,6 +27,7 @@ public static class EnvironmentSettingNames
         public const string AppInsightsInstrumentationKey = "APPINSIGHTS_INSTRUMENTATIONKEY";
         public const string AppInsightsConnectionString = "APPLICATIONINSIGHTS_CONNECTION_STRING";
         public const string AppInsightsQuickPulseAuthApiKey = "APPINSIGHTS_QUICKPULSEAUTHAPIKEY";
+        public const string AppInsightsAgent = "APPLICATIONINSIGHTS_ENABLE_AGENT";
         public const string FunctionsExtensionVersion = "FUNCTIONS_EXTENSION_VERSION";
         public const string FunctionWorkerRuntime = "FUNCTIONS_WORKER_RUNTIME";
         public const string ContainerName = "CONTAINER_NAME";
diff --git a/src/WebJobs.Script/ScriptConstants.cs b/src/WebJobs.Script/ScriptConstants.cs
@@ -33,6 +33,7 @@ public static class ScriptConstants
         public const string LogPropertyIsSystemLogKey = "MS_IsSystemLog";
         public const string LogPropertyFunctionInvocationIdKey = "MS_FunctionInvocationId";
         public const string LogPropertyHostInstanceIdKey = "HostInstanceId";
+        public const string LogPropertyProcessIdKey = "ProcessId";
         public const string LogPropertyActivityIdKey = "MS_ActivityId";
 
         public const string TraceSourceSecretManagement = "SecretManagement";
diff --git a/src/WebJobs.Script/Workers/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs b/src/WebJobs.Script/Workers/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs
@@ -10,6 +10,7 @@
 using System.Threading.Tasks;
 using System.Threading.Tasks.Dataflow;
 using Microsoft.AspNetCore.Hosting;
+using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.Eventing;
@@ -314,6 +315,7 @@ public async Task InitializeAsync(IEnumerable<FunctionMetadata> functions, Cance
                     }
                 }
             }
+            AddLogUserCategory(functions);
         }
 
         // Gets metadata from worker
@@ -353,6 +355,7 @@ public async Task FinishInitialization(IEnumerable<FunctionMetadata> functions,
                 }
                 SetFunctionDispatcherStateToInitializedAndLog();
             }
+            AddLogUserCategory(functions);
         }
 
         public async Task<IDictionary<string, WorkerStatus>> GetWorkerStatusesAsync()
@@ -655,5 +658,18 @@ public async Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)
             }
             return false;
         }
+
+        private void AddLogUserCategory(IEnumerable<FunctionMetadata> functions)
+        {
+            // Add category, this is only needed for workers running AI agent
+            if (_environment.IsApplicationInsightsAgentEnabled())
+            {
+                foreach (FunctionMetadata metadata in functions)
+                {
+                    metadata.Properties[LogConstants.CategoryNameKey] = LogCategories.CreateFunctionUserCategory(metadata.Name);
+                    metadata.Properties[ScriptConstants.LogPropertyHostInstanceIdKey] = _scriptOptions.InstanceId;
+                }
+            }
+        }
     }
 }
diff --git a/test/WebJobs.Script.Tests/Workers/Rpc/GrpcWorkerChannelTests.cs b/test/WebJobs.Script.Tests/Workers/Rpc/GrpcWorkerChannelTests.cs
@@ -3,11 +3,13 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.Eventing;
@@ -36,6 +38,7 @@ public class GrpcWorkerChannelTests : IDisposable
         private readonly string _workerId = "testWorkerId";
         private readonly string _scriptRootPath = "c:\testdir";
         private readonly IScriptEventManager _eventManager = new ScriptEventManager();
+        private readonly Mock<IScriptEventManager> _eventManagerMock = new Mock<IScriptEventManager>();
         private readonly TestMetricsLogger _metricsLogger = new TestMetricsLogger();
         private readonly Mock<IWorkerConsoleLogSource> _mockConsoleLogger = new Mock<IWorkerConsoleLogSource>();
         private readonly Mock<FunctionRpc.FunctionRpcBase> _mockFunctionRpcService = new Mock<FunctionRpc.FunctionRpcBase>();
@@ -52,6 +55,7 @@ public class GrpcWorkerChannelTests : IDisposable
         private readonly IFunctionDataCache _functionDataCache;
         private readonly IOptions<WorkerConcurrencyOptions> _workerConcurrencyOptions;
         private GrpcWorkerChannel _workerChannel;
+        private GrpcWorkerChannel _workerChannelwithMockEventManager;
 
         public GrpcWorkerChannelTests()
         {
@@ -104,6 +108,23 @@ public GrpcWorkerChannelTests()
                _sharedMemoryManager,
                _functionDataCache,
                _workerConcurrencyOptions);
+
+            _eventManagerMock.Setup(proxy => proxy.Publish(It.IsAny<OutboundGrpcEvent>())).Verifiable();
+            _workerChannelwithMockEventManager = new GrpcWorkerChannel(
+               _workerId,
+               _eventManagerMock.Object,
+               _testWorkerConfig,
+               _mockrpcWorkerProcess.Object,
+               _logger,
+               _metricsLogger,
+               0,
+               _testEnvironment,
+               _hostOptionsMonitor,
+               _sharedMemoryManager,
+               _functionDataCache,
+               _workerConcurrencyOptions);
+
+            _testEnvironment.SetEnvironmentVariable("APPLICATIONINSIGHTS_ENABLE_AGENT", "true");
         }
 
         public void Dispose()
@@ -730,6 +751,57 @@ public async Task GetLatencies_DoesNot_StartTimer_WhenDynamicConcurrencyDisabled
             Assert.True(latencyHistory.Count() == 0);
         }
 
+        [Fact]
+        public async Task SendInvocationRequest_ValidateTraceContext()
+        {
+            ScriptInvocationContext scriptInvocationContext = GetTestScriptInvocationContext(Guid.NewGuid(), null);
+            await _workerChannelwithMockEventManager.SendInvocationRequest(scriptInvocationContext);
+            if (_testEnvironment.IsApplicationInsightsAgentEnabled())
+            {
+                _eventManagerMock.Verify(proxy => proxy.Publish(It.Is<OutboundGrpcEvent>(
+                    grpcEvent => grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(ScriptConstants.LogPropertyProcessIdKey)
+                    && grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(ScriptConstants.LogPropertyHostInstanceIdKey)
+                    && grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(LogConstants.CategoryNameKey)
+                    && grpcEvent.Message.InvocationRequest.TraceContext.Attributes[LogConstants.CategoryNameKey].Equals("testcat1")
+                    && grpcEvent.Message.InvocationRequest.TraceContext.Attributes.Count == 3)));
+            }
+            else
+            {
+                _eventManagerMock.Verify(proxy => proxy.Publish(It.Is<OutboundGrpcEvent>(
+                    grpcEvent => !grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(ScriptConstants.LogPropertyProcessIdKey)
+                    && !grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(ScriptConstants.LogPropertyHostInstanceIdKey)
+                    && !grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(LogConstants.CategoryNameKey))));
+            }
+        }
+
+        [Fact]
+        public async Task SendInvocationRequest_ValidateTraceContext_SessionId()
+        {
+            string sessionId = "sessionId1234";
+            Activity activity = new Activity("testActivity");
+            activity.AddBaggage(ScriptConstants.LiveLogsSessionAIKey, sessionId);
+            activity.Start();
+            ScriptInvocationContext scriptInvocationContext = GetTestScriptInvocationContext(Guid.NewGuid(), null);
+            await _workerChannelwithMockEventManager.SendInvocationRequest(scriptInvocationContext);
+            activity.Stop();
+            _eventManagerMock.Verify(p => p.Publish(It.Is<OutboundGrpcEvent>(grpcEvent => ValidateInvocationRequest(grpcEvent, sessionId))));
+        }
+
+        private bool ValidateInvocationRequest(OutboundGrpcEvent grpcEvent, string sessionId)
+        {
+            if (_testEnvironment.IsApplicationInsightsAgentEnabled())
+            {
+                return grpcEvent.Message.InvocationRequest.TraceContext.Attributes[ScriptConstants.LiveLogsSessionAIKey].Equals(sessionId)
+                && grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(LogConstants.CategoryNameKey)
+                && grpcEvent.Message.InvocationRequest.TraceContext.Attributes[LogConstants.CategoryNameKey].Equals("testcat1")
+                && grpcEvent.Message.InvocationRequest.TraceContext.Attributes.Count == 4;
+            }
+            else
+            {
+                return !grpcEvent.Message.InvocationRequest.TraceContext.Attributes.ContainsKey(LogConstants.CategoryNameKey);
+            }
+        }
+
         private IEnumerable<FunctionMetadata> GetTestFunctionsList(string runtime)
         {
             var metadata1 = new FunctionMetadata()
@@ -739,15 +811,17 @@ private IEnumerable<FunctionMetadata> GetTestFunctionsList(string runtime)
             };
 
             metadata1.SetFunctionId("TestFunctionId1");
-
+            metadata1.Properties.Add(LogConstants.CategoryNameKey, "testcat1");
+            metadata1.Properties.Add(ScriptConstants.LogPropertyHostInstanceIdKey, "testhostId1");
             var metadata2 = new FunctionMetadata()
             {
                 Language = runtime,
                 Name = "js2",
             };
 
             metadata2.SetFunctionId("TestFunctionId2");
-
+            metadata2.Properties.Add(LogConstants.CategoryNameKey, "testcat2");
+            metadata2.Properties.Add(ScriptConstants.LogPropertyHostInstanceIdKey, "testhostId2");
             return new List<FunctionMetadata>()
             {
                 metadata1,
diff --git a/test/WebJobs.Script.Tests/Workers/Rpc/RpcFunctionInvocationDispatcherTests.cs b/test/WebJobs.Script.Tests/Workers/Rpc/RpcFunctionInvocationDispatcherTests.cs
@@ -8,6 +8,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Hosting;
+using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.Eventing;
@@ -635,6 +636,22 @@ public async Task FunctionDispatcher_DoNot_Restart_ErroredChannels_If_WorkerRunt
             Assert.Equal(expectedProcessCount, functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels().Count());
         }
 
+        [Fact]
+        public async Task FunctionDispatcher_AdditionalAttributes()
+        {
+            int expectedProcessCount = 3;
+            RpcFunctionInvocationDispatcher functionDispatcher = GetTestFunctionDispatcher(expectedProcessCount, false, runtime: RpcWorkerConstants.JavaLanguageWorkerName, workerIndexing: true);
+
+            var functions = GetTestFunctionsList(RpcWorkerConstants.JavaLanguageWorkerName);
+            await functionDispatcher.FinishInitialization(functions);
+            foreach (var item in functions)
+            {
+                Assert.Equal(item.Properties.Count, 2);
+                Assert.True(item.Properties.ContainsKey(LogConstants.CategoryNameKey) &&
+                    item.Properties.ContainsKey(ScriptConstants.LogPropertyHostInstanceIdKey));
+            }
+        }
+
         [Theory]
         [InlineData("node", "node", false, true, true)]
         [InlineData("node", "node", true, false, true)]
@@ -719,6 +736,7 @@ private static RpcFunctionInvocationDispatcher GetTestFunctionDispatcher(int max
             _javaTestChannel = new TestRpcWorkerChannel(Guid.NewGuid().ToString(), "java", eventManager, _testLogger, false);
             var optionsMonitor = TestHelpers.CreateOptionsMonitor(workerConfigOptions);
 
+            testEnv.SetEnvironmentVariable("APPLICATIONINSIGHTS_ENABLE_AGENT", "true");
             return new RpcFunctionInvocationDispatcher(scriptOptions,
                 metricsLogger.Object,
                 testEnv,

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,8 @@ namespace Microsoft.Azure.WebJobs.Script`
`16`	`16`	`{`
`17`	`17`	`internal static class EnvironmentExtensions`
`18`	`18`	`{`
`19`		`- // For testing`
	`19`	`+ private static bool? isApplicationInsightsAgentEnabled;`
	`20`	`+`
`20`	`21`	`internal static string BaseDirectory { get; set; }`
`21`	`22`
`22`	`23`	`public static string GetEnvironmentVariableOrDefault(this IEnvironment environment, string name, string defaultValue)`
`@@ -508,5 +509,21 @@ public static HashSet<string> GetLanguageWorkerListToStartInPlaceholder(this IEn`
`508`	`509`	`}`
`509`	`510`	`return placeholderRuntimeSet;`
`510`	`511`	`}`
	`512`	`+`
	`513`	`+ public static bool IsApplicationInsightsAgentEnabled(this IEnvironment environment)`
	`514`	`+ {`
	`515`	`+ // cache the value of the environment variable`
	`516`	`+ if (isApplicationInsightsAgentEnabled.HasValue)`
	`517`	`+ {`
	`518`	`+ return isApplicationInsightsAgentEnabled.Value;`
	`519`	`+ }`
	`520`	`+ else if (!environment.IsPlaceholderModeEnabled())`
	`521`	`+ {`
	`522`	`+ bool.TryParse(environment.GetEnvironmentVariable(AppInsightsAgent), out bool isEnabled);`
	`523`	`+ isApplicationInsightsAgentEnabled = isEnabled;`
	`524`	`+ return isApplicationInsightsAgentEnabled.Value;`
	`525`	`+ }`
	`526`	`+ return false;`
	`527`	`+ }`
`511`	`528`	`}`
`512`	`529`	`}`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@`
`10`	`10`	`using System.Threading.Tasks;`
`11`	`11`	`using System.Threading.Tasks.Dataflow;`
`12`	`12`	`using Microsoft.AspNetCore.Hosting;`
	`13`	`+using Microsoft.Azure.WebJobs.Logging;`
`13`	`14`	`using Microsoft.Azure.WebJobs.Script.Description;`
`14`	`15`	`using Microsoft.Azure.WebJobs.Script.Diagnostics;`
`15`	`16`	`using Microsoft.Azure.WebJobs.Script.Eventing;`
`@@ -314,6 +315,7 @@ public async Task InitializeAsync(IEnumerable<FunctionMetadata> functions, Cance`
`314`	`315`	`}`
`315`	`316`	`}`
`316`	`317`	`}`
	`318`	`+ AddLogUserCategory(functions);`
`317`	`319`	`}`
`318`	`320`
`319`	`321`	`// Gets metadata from worker`
`@@ -353,6 +355,7 @@ public async Task FinishInitialization(IEnumerable<FunctionMetadata> functions,`
`353`	`355`	`}`
`354`	`356`	`SetFunctionDispatcherStateToInitializedAndLog();`
`355`	`357`	`}`
	`358`	`+ AddLogUserCategory(functions);`
`356`	`359`	`}`
`357`	`360`
`358`	`361`	`public async Task<IDictionary<string, WorkerStatus>> GetWorkerStatusesAsync()`
`@@ -655,5 +658,18 @@ public async Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)`
`655`	`658`	`}`
`656`	`659`	`return false;`
`657`	`660`	`}`
	`661`	`+`
	`662`	`+ private void AddLogUserCategory(IEnumerable<FunctionMetadata> functions)`
	`663`	`+ {`
	`664`	`+ // Add category, this is only needed for workers running AI agent`
	`665`	`+ if (_environment.IsApplicationInsightsAgentEnabled())`
	`666`	`+ {`
	`667`	`+ foreach (FunctionMetadata metadata in functions)`
	`668`	`+ {`
	`669`	`+ metadata.Properties[LogConstants.CategoryNameKey] = LogCategories.CreateFunctionUserCategory(metadata.Name);`
	`670`	`+ metadata.Properties[ScriptConstants.LogPropertyHostInstanceIdKey] = _scriptOptions.InstanceId;`
	`671`	`+ }`
	`672`	`+ }`
	`673`	`+ }`
`658`	`674`	`}`
`659`	`675`	`}`