ModuleTrackerTemplate.cs

// Copyright (c) Toni Solarin-Sodara
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Reflection;
using System.Runtime.CompilerServices;
using System.Threading;

namespace Coverlet.Core.Instrumentation
{
  /// <summary>
  /// This static class will be injected on a module being instrumented in order to direct on module hits
  /// to a single location.
  /// </summary>
  /// <remarks>
  /// As this type is going to be customized for each instrumented module it doesn't follow typical practices
  /// regarding visibility of members, etc.
  /// </remarks>
  [CompilerGenerated]
  [ExcludeFromCodeCoverage]
  internal static class ModuleTrackerTemplate
  {
    public static string HitsFilePath;
    public static int[] HitsArray;
    public static bool SingleHit;
    public static bool FlushHitFile;
    private static readonly bool s_enableLog = int.TryParse(Environment.GetEnvironmentVariable("COVERLET_ENABLETRACKERLOG"), out int result) && result == 1;
    private static readonly string s_sessionId = Guid.NewGuid().ToString();

    static ModuleTrackerTemplate()
    {
      // 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 RecordHitInCoreLibrary(int hitLocationIndex)
    {
      // 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 RecordHit(int hitLocationIndex)
    {
      Interlocked.Increment(ref HitsArray[hitLocationIndex]);
    }

    public static void RecordSingleHitInCoreLibrary(int hitLocationIndex)
    {
      // 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;

      ref int location = ref HitsArray[hitLocationIndex];
      if (location == 0)
        location = 1;
    }

    public static void RecordSingleHit(int hitLocationIndex)
    {
      ref int location = ref HitsArray[hitLocationIndex];
      if (location == 0)
        location = 1;
    }

    public static void UnloadModule(object sender, EventArgs e)
    {
      // 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.
      using var mutex = new Mutex(true, Path.GetFileNameWithoutExtension(HitsFilePath) + "_Mutex", out bool createdNew);
      if (!createdNew)
      {
        mutex.WaitOne();
      }

      if (FlushHitFile)
      {
        try
        {
          // Claim the current hits array and reset it to prevent double-counting scenarios.
          int[] hitsArray = Interlocked.Exchange(ref HitsArray, new int[HitsArray.Length]);

          WriteLog($"Unload called for '{Assembly.GetExecutingAssembly().Location}' by '{sender ?? "null"}'");
          WriteLog($"Flushing hit file '{HitsFilePath}'");

          bool failedToCreateNewHitsFile = false;
          try
          {
            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(hitCount);
            }
          }
          catch (Exception ex)
          {
            WriteLog($"Failed to create new hits file '{HitsFilePath}' -> '{ex.Message}'");
            failedToCreateNewHitsFile = true;
          }

          if (failedToCreateNewHitsFile)
          {
            // Update the number of hits by adding value on disk with the ones on memory.
            // This path should be triggered only in the case of multiple AppDomain unloads.
            using var fs = new FileStream(HitsFilePath, FileMode.Open, FileAccess.ReadWrite, FileShare.None);
            using var br = new BinaryReader(fs);
            using var bw = new BinaryWriter(fs);
            int hitsLength = br.ReadInt32();
            WriteLog($"Current hits found '{hitsLength}'");

            if (hitsLength != hitsArray.Length)
            {
              throw new InvalidOperationException($"{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);
              if (SingleHit)
              {
                bw.Write(hitsArray[i] + oldHitCount > 0 ? 1 : 0);
              }
              else
              {
                bw.Write(hitsArray[i] + oldHitCount);
              }
            }
          }

          WriteHits(sender);

          WriteLog($"Hit file '{HitsFilePath}' flushed, size {new FileInfo(HitsFilePath).Length}");
          WriteLog("--------------------------------");
        }
        catch (Exception ex)
        {
          WriteLog(ex.ToString());
          throw;
        }
      }

      // 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();
    }

    private static void WriteHits(object sender)
    {
      if (s_enableLog)
      {
        var currentAssembly = Assembly.GetExecutingAssembly();
        var location = new DirectoryInfo(Path.Combine(Path.GetDirectoryName(currentAssembly.Location), "TrackersHitsLog"));
        location.Create();
        string logFile = Path.Combine(location.FullName, $"{Path.GetFileName(currentAssembly.Location)}_{DateTime.UtcNow.Ticks}_{s_sessionId}.txt");
        using (var fs = new FileStream(HitsFilePath, FileMode.Open, FileAccess.ReadWrite, FileShare.None))
        using (var log = new FileStream(logFile, FileMode.CreateNew, FileAccess.ReadWrite, FileShare.None))
        using (var logWriter = new StreamWriter(log))
        using (var br = new BinaryReader(fs))
        {
          int hitsLength = br.ReadInt32();
          for (int i = 0; i < hitsLength; ++i)
          {
            logWriter.WriteLine($"{i},{br.ReadInt32()}");
          }
        }

        File.AppendAllText(logFile, $"Hits flushed file path {HitsFilePath} location '{Assembly.GetExecutingAssembly().Location}' by '{sender ?? "null"}'");
      }
    }

    private static void WriteLog(string logText)
    {
      if (s_enableLog)
      {
        // We don't set path as global var to keep benign possible errors inside try/catch
        // I'm not sure that location will be ok in every scenario
        string location = Assembly.GetExecutingAssembly().Location;
        File.AppendAllText(Path.Combine(Path.GetDirectoryName(location), Path.GetFileName(location) + "_tracker.txt"), $"[{DateTime.UtcNow} S:{s_sessionId} T:{Thread.CurrentThread.ManagedThreadId}]{logText}{Environment.NewLine}");
      }
    }
  }
}