Merge pull request #297 from sharwell/default-init

Remove reliance on s_isTracking to detect recursion

Author: tonerdo

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

@@ -20,33 +20,36 @@ public static class ModuleTrackerTemplate
         public static string HitsFilePath;
         public static int[] HitsArray;
 
-        // Special case when instrumenting CoreLib, the static below prevents infinite loop in CoreLib
-        // while allowing the tracker template to call any of its types and functions.
-        private static bool s_isTracking;
-
         static ModuleTrackerTemplate()
         {
-            s_isTracking = true;
-            AppDomain.CurrentDomain.ProcessExit += new EventHandler(UnloadModule);
-            AppDomain.CurrentDomain.DomainUnload += new EventHandler(UnloadModule);
-            s_isTracking = false;
-
             // 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)
         {
-            if (s_isTracking)
+            // Make sure to avoid recording if this is a call to RecordHit within the AppDomain setup code in an
+            // instrumented build of System.Private.CoreLib.
+            if (HitsArray is null)
                 return;
 
             Interlocked.Increment(ref HitsArray[hitLocationIndex]);
         }
 
         public static void UnloadModule(object sender, EventArgs e)
         {
-            s_isTracking = true;
+            // Claim the current hits array and reset it to prevent double-counting scenarios.
+            var hitsArray = Interlocked.Exchange(ref HitsArray, new int[HitsArray.Length]);
 
             // The same module can be unloaded multiple times in the same process via different app domains.
             // Use a global mutex to ensure no concurrent access.
@@ -61,8 +64,8 @@ public static void UnloadModule(object sender, EventArgs e)
                     using (var fs = new FileStream(HitsFilePath, FileMode.CreateNew))
                     using (var bw = new BinaryWriter(fs))
                     {
-                        bw.Write(HitsArray.Length);
-                        foreach (int hitCount in HitsArray)
+                        bw.Write(hitsArray.Length);
+                        foreach (int hitCount in hitsArray)
                         {
                             bw.Write(hitCount);
                         }
@@ -82,25 +85,21 @@ public static void UnloadModule(object sender, EventArgs e)
                     using (var bw = new BinaryWriter(fs))
                     {
                         int hitsLength = br.ReadInt32();
-                        if (hitsLength != HitsArray.Length)
+                        if (hitsLength != hitsArray.Length)
                         {
                             throw new InvalidOperationException(
-                                $"{HitsFilePath} has {hitsLength} entries but on memory {nameof(HitsArray)} has {HitsArray.Length}");
+                                $"{HitsFilePath} has {hitsLength} entries but on memory {nameof(HitsArray)} has {hitsArray.Length}");
                         }
 
                         for (int i = 0; i < hitsLength; ++i)
                         {
                             int oldHitCount = br.ReadInt32();
                             bw.Seek(-sizeof(int), SeekOrigin.Current);
-                            bw.Write(HitsArray[i] + oldHitCount);
+                            bw.Write(hitsArray[i] + oldHitCount);
                         }
                     }
                 }
 
-                // Prevent any double counting scenario, i.e.: UnloadModule called twice (not sure if this can happen in practice ...)
-                // Only an issue if DomainUnload and ProcessExit can both happens, perhaps can be removed...
-                Array.Clear(HitsArray, 0, HitsArray.Length);
-
                 // On purpose this is not under a try-finally: it is better to have an exception if there was any error writing the hits file
                 // this case is relevant when instrumenting corelib since multiple processes can be running against the same instrumented dll.
                 mutex.ReleaseMutex();