Remove the need to register events when instrumenting System.Private.CoreLib

Author: sharwell

src/coverlet.core/Instrumentation/Instrumenter.cs

@@ -29,6 +29,7 @@ internal class Instrumenter
         private FieldDefinition _customTrackerHitsFilePath;
         private ILProcessor _customTrackerClassConstructorIl;
         private TypeDefinition _customTrackerTypeDef;
+        private MethodReference _customTrackerRegisterUnloadEventsMethod;
         private MethodReference _customTrackerRecordHitMethod;
 
         public Instrumenter(string module, string identifier, string[] excludeFilters, string[] includeFilters, string[] excludedFiles, string[] excludedAttributes)
@@ -76,6 +77,7 @@ private void InstrumentModule()
 
                 using (var module = ModuleDefinition.ReadModule(stream, parameters))
                 {
+                    var containsAppContext = module.GetType(nameof(System), nameof(AppContext)) != null;
                     var types = module.GetTypes();
                     AddCustomModuleTrackerToModule(module);
 
@@ -95,13 +97,49 @@ private void InstrumentModule()
                     }
 
                     // Fixup the custom tracker class constructor, according to all instrumented types
+                    if (_customTrackerRegisterUnloadEventsMethod == null)
+                    {
+                        _customTrackerRegisterUnloadEventsMethod = new MethodReference(
+                            nameof(ModuleTrackerTemplate.RegisterUnloadEvents), module.TypeSystem.Void, _customTrackerTypeDef);
+                    }
+
                     Instruction lastInstr = _customTrackerClassConstructorIl.Body.Instructions.Last();
+
+                    if (!containsAppContext)
+                    {
+                        // For "normal" cases, where the instrumented assembly is not the core library, we add a call to
+                        // RegisterUnloadEvents to the static constructor of the generated custom tracker. Due to static
+                        // initialization constraints, the core library is handled separately below.
+                        _customTrackerClassConstructorIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Call, _customTrackerRegisterUnloadEventsMethod));
+                    }
+
                     _customTrackerClassConstructorIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Ldc_I4, _result.HitCandidates.Count));
                     _customTrackerClassConstructorIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Newarr, module.TypeSystem.Int32));
                     _customTrackerClassConstructorIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Stsfld, _customTrackerHitsArray));
                     _customTrackerClassConstructorIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Ldstr, _result.HitsFilePath));
                     _customTrackerClassConstructorIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Stsfld, _customTrackerHitsFilePath));
 
+                    if (containsAppContext)
+                    {
+                        // Handle the core library by instrumenting System.AppContext.OnProcessExit to directly call
+                        // the UnloadModule method of the custom tracker type. This avoids loops between the static
+                        // initialization of the custom tracker and the static initialization of the hosting AppDomain
+                        // (which for the core library case will be instrumented code).
+                        var eventArgsType = new TypeReference(nameof(System), nameof(EventArgs), module, module.TypeSystem.CoreLibrary);
+                        var customTrackerUnloadModule = new MethodReference(nameof(ModuleTrackerTemplate.UnloadModule), module.TypeSystem.Void, _customTrackerTypeDef);
+                        customTrackerUnloadModule.Parameters.Add(new ParameterDefinition(module.TypeSystem.Object));
+                        customTrackerUnloadModule.Parameters.Add(new ParameterDefinition(eventArgsType));
+
+                        var appContextType = new TypeReference(nameof(System), nameof(AppContext), module, module.TypeSystem.CoreLibrary);
+                        var onProcessExitMethod = new MethodReference("OnProcessExit", module.TypeSystem.Void, appContextType).Resolve();
+                        var onProcessExitIl = onProcessExitMethod.Body.GetILProcessor();
+
+                        lastInstr = onProcessExitIl.Body.Instructions.Last();
+                        onProcessExitIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Ldnull));
+                        onProcessExitIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Ldnull));
+                        onProcessExitIl.InsertBefore(lastInstr, Instruction.Create(OpCodes.Call, customTrackerUnloadModule));
+                    }
+
                     module.Write(stream);
                 }
             }
@@ -135,12 +173,9 @@ private void AddCustomModuleTrackerToModule(ModuleDefinition module)
                 {
                     MethodDefinition methodOnCustomType = new MethodDefinition(methodDef.Name, methodDef.Attributes, methodDef.ReturnType);
 
-                    if (methodDef.Name == "RecordHit")
+                    foreach (var parameter in methodDef.Parameters)
                     {
-                        foreach (var parameter in methodDef.Parameters)
-                        {
-                            methodOnCustomType.Parameters.Add(new ParameterDefinition(module.ImportReference(parameter.ParameterType)));
-                        }
+                        methodOnCustomType.Parameters.Add(new ParameterDefinition(module.ImportReference(parameter.ParameterType)));
                     }
 
                     foreach (var variable in methodDef.Body.Variables)
@@ -166,8 +201,11 @@ private void AddCustomModuleTrackerToModule(ModuleDefinition module)
                             else
                             {
                                 // Move to the custom type
-                                instr.Operand = new MethodReference(
-                                    methodReference.Name, methodReference.ReturnType, _customTrackerTypeDef);
+                                var updatedMethodReference = new MethodReference(methodReference.Name, methodReference.ReturnType, _customTrackerTypeDef);
+                                foreach (var parameter in methodReference.Parameters)
+                                    updatedMethodReference.Parameters.Add(new ParameterDefinition(parameter.Name, parameter.Attributes, module.ImportReference(parameter.ParameterType)));
+
+                                instr.Operand = updatedMethodReference;
                             }
                         }
                         else if (instr.Operand is FieldReference fieldReference)

src/coverlet.template/ModuleTrackerTemplate.cs

@@ -22,14 +22,20 @@ public static class ModuleTrackerTemplate
 
         static ModuleTrackerTemplate()
         {
-            AppDomain.CurrentDomain.ProcessExit += new EventHandler(UnloadModule);
-            AppDomain.CurrentDomain.DomainUnload += new EventHandler(UnloadModule);
-
             // At the end of the instrumentation of a module, the instrumenter needs to add code here
             // to initialize the static fields according to the values derived from the instrumentation of
             // the module.
         }
 
+        // A call to this method will be injected in the static constructor above for most cases. However, if the
+        // current assembly is System.Private.CoreLib (or more specifically, defines System.AppDomain), a call directly
+        // to UnloadModule will be injected in System.AppContext.OnProcessExit.
+        public static void RegisterUnloadEvents()
+        {
+            AppDomain.CurrentDomain.ProcessExit += new EventHandler(UnloadModule);
+            AppDomain.CurrentDomain.DomainUnload += new EventHandler(UnloadModule);
+        }
+
         public static void RecordHit(int hitLocationIndex)
         {
             // Make sure to avoid recording if this is a call to RecordHit within the AppDomain setup code in an