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)
@@ -95,13 +96,42 @@ 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 (!module.Equals(module.TypeSystem.Void.Module))
+                    {
+                        _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 (module.Equals(module.TypeSystem.Void.Module))
+                    {
+                        var eventArgsType = new TypeReference(nameof(System), nameof(EventArgs), module, module.TypeSystem.CoreLibrary).Resolve();
+                        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).Resolve();
+                        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 +165,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)

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