Prejit support on FunctionsNetHost

eng/ci/templates/steps/install-dotnet.yml

@@ -17,3 +17,10 @@ steps:
   inputs:
     packageType: sdk
     useGlobalJson: true
+
+- task: UseDotNet@2 
+  displayName: Install .NET 9
+  inputs:
+    packageType: sdk
+    version: '9.0.100-preview.6.24328.19'
+    includePreviewVersions: true

host/src/FunctionsNetHost.Shared/Constants.cs

@@ -0,0 +1,13 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+namespace FunctionsNetHost.Shared
+{
+    public static class Constants
+    {
+        public const string LogCategory = "FunctionsNetHost";
+        public const string DefaultLogPrefix = "LanguageWorkerConsoleLog";
+        public const string NetHostWaitHandleName = "AzureFunctionsNetHostSpecializationWaitHandle";
+        public const string LogTimeStampFormat = "yyyy-MM-dd HH:mm:ss.fff";
+    }
+}

host/src/FunctionsNetHost.Shared/EnvironmentVariables.cs

@@ -0,0 +1,23 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+namespace FunctionsNetHost.Shared
+{
+    public static class EnvironmentVariables
+    {
+        /// <summary>
+        /// The environment variable which is used to specify the specialized (function app payload) entry assembly.
+        /// </summary>
+        public const string SpecializedEntryAssembly = "AZURE_FUNCTIONS_FUNCTIONSNETHOST_SPECIALIZED_ENTRY_ASSEMBLY";
+
+        /// <summary>
+        /// The environment variable which is used to specify the path to the jittrace file which will be used for prejitting.
+        /// </summary>
+        public const string PreJitFilePath = "AZURE_FUNCTIONS_FUNCTIONSNETHOST_PREJIT_FILE_PATH";
+
+        /// <summary>
+        /// The .NET startup hooks environment variable.
+        /// </summary>
+        public const string DotnetStartupHooks = "DOTNET_STARTUP_HOOKS";
+    }
+}

host/src/FunctionsNetHost.Shared/FunctionsNetHost.Shared.csproj

@@ -0,0 +1,5 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <TargetFramework>netstandard2.0</TargetFramework>
+  </PropertyGroup>
+</Project>

host/src/FunctionsNetHost.sln

@@ -5,6 +5,8 @@ VisualStudioVersion = 17.5.33627.172
 MinimumVisualStudioVersion = 10.0.40219.1
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "FunctionsNetHost", "FunctionsNetHost\FunctionsNetHost.csproj", "{6C05D0AC-F6AC-45FB-8A73-A3F44DF131BC}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "FunctionsNetHost.Shared", "FunctionsNetHost.Shared\FunctionsNetHost.Shared.csproj", "{91867A34-8B79-4E01-81BC-592F1CD1CCCE}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -15,6 +17,10 @@ Global
 		{6C05D0AC-F6AC-45FB-8A73-A3F44DF131BC}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{6C05D0AC-F6AC-45FB-8A73-A3F44DF131BC}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{6C05D0AC-F6AC-45FB-8A73-A3F44DF131BC}.Release|Any CPU.Build.0 = Release|Any CPU
+		{91867A34-8B79-4E01-81BC-592F1CD1CCCE}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{91867A34-8B79-4E01-81BC-592F1CD1CCCE}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{91867A34-8B79-4E01-81BC-592F1CD1CCCE}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{91867A34-8B79-4E01-81BC-592F1CD1CCCE}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE

host/src/FunctionsNetHost/Configuration/Configuration.cs

@@ -8,8 +8,6 @@ namespace FunctionsNetHost
     /// </summary>
     internal static class Configuration
     {
-        private const string DefaultLogPrefix = "LanguageWorkerConsoleLog";
-
         static Configuration()
         {
             Reload();
@@ -22,7 +20,7 @@ internal static void Reload()
         {
             IsTraceLogEnabled = string.Equals(EnvironmentUtils.GetValue(EnvironmentVariables.EnableTraceLogs), "1");
             var disableLogPrefix = string.Equals(EnvironmentUtils.GetValue(EnvironmentVariables.DisableLogPrefix), "1");
-            LogPrefix = disableLogPrefix ? string.Empty : DefaultLogPrefix;
+            LogPrefix = disableLogPrefix ? string.Empty : Shared.Constants.DefaultLogPrefix;
         }
 
         /// <summary>

host/src/FunctionsNetHost/Environment/EnvironmentVariables.cs

@@ -19,10 +19,15 @@ internal static class EnvironmentVariables
     /// <summary>
     /// Application pool Id for the placeholder app. Only available in Windows(when running in IIS).
     /// </summary>
-    internal const string AppPoolId  = "APP_POOL_ID";
+    internal const string AppPoolId = "APP_POOL_ID";
 
     /// <summary>
     /// The worker runtime version. Example value: "8.0" (for a .NET8 placeholder)
     /// </summary>
     internal const string FunctionsWorkerRuntimeVersion = "FUNCTIONS_WORKER_RUNTIME_VERSION";
+
+    /// <summary>
+    /// The environment variable that disables prejit. If set to "1," prejit will be disabled.
+    /// </summary>
+    internal const string DisablePrejit = "AZURE_FUNCTIONS_FUNCTIONSNETHOST_DISABLE_PREJIT";
 }

host/src/FunctionsNetHost/FunctionsNetHost.csproj

@@ -2,7 +2,9 @@
 
   <PropertyGroup>
     <OutputType>Exe</OutputType>
-    <TargetFramework>net8.0</TargetFramework>
+    <TargetFramework>net9.0</TargetFramework>
+    <EnablePreviewFeatures>True</EnablePreviewFeatures>
+    <LangVersion>preview</LangVersion>
     <ImplicitUsings>enable</ImplicitUsings>
     <Nullable>enable</Nullable>
     <AllowUnsafeBlocks>True</AllowUnsafeBlocks>
@@ -28,12 +30,16 @@
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
     </PackageReference>
     <PackageReference Include="System.CommandLine" Version="2.0.0-beta4.22272.1" />
-    <PackageReference Include="System.Text.Json" Version="8.0.0-preview.4.23259.5" />
-    <PackageReference Include="Microsoft.NETCore.DotNetAppHost" Version="8.0.0-preview.4.23259.5" />
+    <PackageReference Include="System.Text.Json" Version="9.0.0-preview.6.24327.7" />
+    <PackageReference Include="Microsoft.NETCore.DotNetAppHost" Version="9.0.0-preview.6.24327.7" />
   </ItemGroup>
 
   <ItemGroup>
     <Protobuf Include="..\..\..\protos\azure-functions-language-worker-protobuf\**\*.proto" ProtoRoot="..\..\..\protos\azure-functions-language-worker-protobuf\src\proto" GrpcServices="Client" Access="internal" />
   </ItemGroup>
 
+  <ItemGroup>
+    <ProjectReference Include="..\FunctionsNetHost.Shared\FunctionsNetHost.Shared.csproj" />
+  </ItemGroup>
+
 </Project>

host/src/FunctionsNetHost/Grpc/GrpcClient.cs

@@ -14,11 +14,11 @@ internal sealed class GrpcClient
     {
         private readonly Channel<StreamingMessage> _outgoingMessageChannel;
         private readonly IncomingGrpcMessageHandler _messageHandler;
-        private readonly GrpcWorkerStartupOptions _grpcWorkerStartupOptions;
+        private readonly NetHostRunOptions _netHostRunOptions;
 
-        internal GrpcClient(GrpcWorkerStartupOptions grpcWorkerStartupOptions, AppLoader appLoader)
+        internal GrpcClient(NetHostRunOptions netHostRunOptions, AppLoader appLoader)
         {
-            _grpcWorkerStartupOptions = grpcWorkerStartupOptions;
+            _netHostRunOptions = netHostRunOptions;
             var channelOptions = new UnboundedChannelOptions
             {
                 SingleWriter = false,
@@ -28,12 +28,12 @@ internal GrpcClient(GrpcWorkerStartupOptions grpcWorkerStartupOptions, AppLoader
 
             _outgoingMessageChannel = Channel.CreateUnbounded<StreamingMessage>(channelOptions);
 
-            _messageHandler = new IncomingGrpcMessageHandler(appLoader, _grpcWorkerStartupOptions);
+            _messageHandler = new IncomingGrpcMessageHandler(appLoader, _netHostRunOptions);
         }
 
         internal async Task InitAsync()
         {
-            var endpoint = _grpcWorkerStartupOptions.ServerUri.AbsoluteUri;
+            var endpoint = _netHostRunOptions.WorkerStartupOptions.ServerUri.AbsoluteUri;
             Logger.LogTrace($"Grpc service endpoint:{endpoint}");
 
             var functionRpcClient = CreateFunctionRpcClient(endpoint);
@@ -69,7 +69,7 @@ private async Task SendStartStreamMessageAsync(IClientStreamWriter<StreamingMess
         {
             var startStreamMsg = new StartStream()
             {
-                WorkerId = _grpcWorkerStartupOptions.WorkerId
+                WorkerId = _netHostRunOptions.WorkerStartupOptions.WorkerId
             };
 
             var startStream = new StreamingMessage()
@@ -89,8 +89,8 @@ private FunctionRpcClient CreateFunctionRpcClient(string endpoint)
 
             var grpcChannel = GrpcChannel.ForAddress(grpcUri, new GrpcChannelOptions()
             {
-                MaxReceiveMessageSize = _grpcWorkerStartupOptions.GrpcMaxMessageLength,
-                MaxSendMessageSize = _grpcWorkerStartupOptions.GrpcMaxMessageLength,
+                MaxReceiveMessageSize = _netHostRunOptions.WorkerStartupOptions.GrpcMaxMessageLength,
+                MaxSendMessageSize = _netHostRunOptions.WorkerStartupOptions.GrpcMaxMessageLength,
                 Credentials = ChannelCredentials.Insecure
             });
 

host/src/FunctionsNetHost/Grpc/IncomingGrpcMessageHandler.cs

@@ -11,12 +11,12 @@ internal sealed class IncomingGrpcMessageHandler
     {
         private bool _specializationDone;
         private readonly AppLoader _appLoader;
-        private readonly GrpcWorkerStartupOptions _grpcWorkerStartupOptions;
+        private readonly NetHostRunOptions _netHostRunOptions;
 
-        internal IncomingGrpcMessageHandler(AppLoader appLoader, GrpcWorkerStartupOptions grpcWorkerStartupOptions)
+        internal IncomingGrpcMessageHandler(AppLoader appLoader, NetHostRunOptions netHostRunOptions)
         {
             _appLoader = appLoader;
-            _grpcWorkerStartupOptions = grpcWorkerStartupOptions;
+            _netHostRunOptions = netHostRunOptions;
         }
 
         internal Task ProcessMessageAsync(StreamingMessage message)
@@ -62,11 +62,14 @@ private async Task Process(StreamingMessage msg)
                     }
                 case StreamingMessage.ContentOneofCase.FunctionEnvironmentReloadRequest:
 
-                    Configuration.Reload();
+                    var envReloadRequest = msg.FunctionEnvironmentReloadRequest;
+                    foreach (var kv in envReloadRequest.EnvironmentVariables)
+                    {
+                        EnvironmentUtils.SetValue(kv.Key, kv.Value);
+                    }
+                    Configuration.Reload();
                     Logger.LogTrace("Specialization request received.");
 
-                    var envReloadRequest = msg.FunctionEnvironmentReloadRequest;
-
                     var workerConfig = await WorkerConfigUtils.GetWorkerConfig(envReloadRequest.FunctionAppDirectory);
 
                     if (workerConfig?.Description is null)
@@ -88,19 +91,20 @@ private async Task Process(StreamingMessage msg)
                     var applicationExePath = Path.Combine(envReloadRequest.FunctionAppDirectory, workerConfig.Description.DefaultWorkerPath!);
                     Logger.LogTrace($"application path {applicationExePath}");
 
-                    foreach (var kv in envReloadRequest.EnvironmentVariables)
+                    if (_netHostRunOptions.IsPreJitSupported)
                     {
-                        EnvironmentUtils.SetValue(kv.Key, kv.Value);
+                        EnvironmentUtils.SetValue(Shared.EnvironmentVariables.SpecializedEntryAssembly, applicationExePath);
+                        // Signal so that startup hook load the payload assembly.
+                        SpecializationSyncManager.WaitHandle.Set();
                     }
-
+                    else
+                    {
 #pragma warning disable CS4014
-                    Task.Run(() =>
+                        Task.Run(() => _appLoader.RunApplication(applicationExePath));
 #pragma warning restore CS4014
-                    {
-                        _ = _appLoader.RunApplication(applicationExePath);
-                    });
+                    }
 
-                    Logger.LogTrace($"Will wait for worker loaded signal.");
+                    Logger.LogTrace("Will wait for worker loaded signal.");
                     WorkerLoadStatusSignalManager.Instance.Signal.WaitOne();
 
                     var logMessage = $"FunctionApp assembly loaded successfully. ProcessId:{Environment.ProcessId}";

host/src/FunctionsNetHost/Grpc/NetHostRunOptions.cs

@@ -0,0 +1,63 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using FunctionsNetHost.Grpc;
+
+namespace FunctionsNetHost
+{
+    /// <summary>
+    /// Encapsulates various configuration options required to run the FunctionsNetHost application.
+    /// </summary>
+    internal sealed class NetHostRunOptions
+    {
+        //.NET 8.0 is the minimum version that supports pre-jitting.
+        private const int MinimumNetTfmToSupportPreJit = 8;
+
+        /// <summary>
+        /// Gets a value indicating whether pre-jitting is supported.
+        /// </summary>
+        public bool IsPreJitSupported { get; }
+
+        /// <summary>
+        /// Gets the worker startup options.
+        /// </summary>
+        public GrpcWorkerStartupOptions WorkerStartupOptions { get; }
+
+        /// <summary>
+        /// Gets the runtime version. This usually corresponds to the .NET runtime version.
+        /// Example value: 8.0.
+        /// </summary>
+        public string RuntimeVersion { get; }
+
+        /// <summary>
+        /// Gets the directory where the FunctionsNetHost executable is located.
+        /// </summary>
+        public string ExecutableDirectory { get; }
+
+        public NetHostRunOptions(GrpcWorkerStartupOptions workerStartupOptions, string executableDirectory)
+        {
+            WorkerStartupOptions = workerStartupOptions;
+            ExecutableDirectory = executableDirectory;
+            RuntimeVersion = EnvironmentUtils.GetValue(EnvironmentVariables.FunctionsWorkerRuntimeVersion)!;
+            IsPreJitSupported = IsPrejitSupported(RuntimeVersion);
+        }
+
+        private static bool IsPrejitSupported(string runtimeVersion)
+        {
+            if (string.IsNullOrEmpty(runtimeVersion))
+            {
+                return false;
+            }
+
+            var disablePrejitEnvironmentVaValue = EnvironmentUtils.GetValue(EnvironmentVariables.DisablePrejit);
+            if (string.Equals(disablePrejitEnvironmentVaValue, "1"))
+            {
+                Logger.Log($"PreJitting is disabled due to the environment variable '{EnvironmentVariables.DisablePrejit}' being set to '{disablePrejitEnvironmentVaValue}'.");
+                return false;
+            }
+
+            return decimal.TryParse(runtimeVersion, out var value) && value >= MinimumNetTfmToSupportPreJit;
+        }
+    }
+}
+

host/src/FunctionsNetHost/Logger.cs

@@ -20,7 +20,7 @@ internal static void LogTrace(string message)
 
         internal static void Log(string message)
         {
-            var ts = DateTime.UtcNow.ToString("yyyy-MM-dd HH:mm:ss.fff", CultureInfo.InvariantCulture);
+            var ts = DateTime.UtcNow.ToString(Shared.Constants.LogTimeStampFormat, CultureInfo.InvariantCulture);
             Console.WriteLine($"{Configuration.LogPrefix}[{ts}] [FunctionsNetHost] {message}");
         }
     }

host/src/FunctionsNetHost/PreJit/PreJitManager.cs

@@ -0,0 +1,44 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+namespace FunctionsNetHost.Prejit
+{
+    internal class PreJitManager
+    {
+        private const string PlaceholderAppDirectory = "PlaceholderApp";
+        private const string PlaceholderAppAssemblyName = "FunctionsNetHost.PlaceholderApp.dll";
+        private const string JitTraceDirectory = "JitTrace";
+#if OS_LINUX
+        private const string JitTraceFileName = "linux.coldstart.jittrace";
+#else
+        private const string JitTraceFileName = "coldstart.jittrace";
+#endif
+
+        /// <summary>
+        /// Starts the placeholder app for the current runtime version.
+        /// Startup hook code is part of this placeholder app.
+        /// </summary>
+        internal static void InitializeAndRunPreJitPlaceholderApp(NetHostRunOptions applicationRunOption, AppLoader appLoader)
+        {
+
+            var placeHolderAppDir = Path.Combine(applicationRunOption.ExecutableDirectory, PlaceholderAppDirectory, applicationRunOption.RuntimeVersion);
+            var placeholderAppAssemblyPath = Path.Combine(placeHolderAppDir, PlaceholderAppAssemblyName);
+            if (!File.Exists(placeholderAppAssemblyPath))
+            {
+                throw new FileNotFoundException($"Placeholder app assembly not found at the specified path:{placeholderAppAssemblyPath}");
+            }
+
+            var preJitFilePath = Path.Combine(placeHolderAppDir, JitTraceDirectory, JitTraceFileName);
+            if (!File.Exists(preJitFilePath))
+            {
+                throw new FileNotFoundException($"Pre-jit file not found at the specified path:{preJitFilePath}");
+            }
+
+            EnvironmentUtils.SetValue(Shared.EnvironmentVariables.PreJitFilePath, preJitFilePath);
+            EnvironmentUtils.SetValue(Shared.EnvironmentVariables.DotnetStartupHooks, placeholderAppAssemblyPath);
+
+            Logger.Log($"Going to run placeholder app:{placeholderAppAssemblyPath}");
+            _ = Task.Run(() => appLoader.RunApplication(placeholderAppAssemblyPath));
+        }
+    }
+}