From 32a57ef9d8615bb81399476d3b227048c28b1c5e Mon Sep 17 00:00:00 2001
From: Viktor Hofer <viktor.hofer@microsoft.com>
Date: Thu, 25 Apr 2019 18:11:21 +0200
Subject: [PATCH] Restore backup assemblies when sigkill

---
 .../Helpers/InstrumentationHelper.cs          | 32 ++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/coverlet.core/Helpers/InstrumentationHelper.cs b/src/coverlet.core/Helpers/InstrumentationHelper.cs
index 79d89d4c0..075d6c75f 100644
--- a/src/coverlet.core/Helpers/InstrumentationHelper.cs
+++ b/src/coverlet.core/Helpers/InstrumentationHelper.cs
@@ -14,6 +14,13 @@ namespace Coverlet.Core.Helpers
 {
     internal static class InstrumentationHelper
     {
+        private static readonly Dictionary<string, string> _backupList = new Dictionary<string, string>();
+
+        static InstrumentationHelper()
+        {
+            AppDomain.CurrentDomain.ProcessExit += (s,e) => RestoreOriginalModules();
+        }
+
         public static string[] GetCoverableModules(string module, string[] directories, bool includeTestAssembly)
         {
             Debug.Assert(directories != null);
@@ -87,11 +94,13 @@ public static void BackupOriginalModule(string module, string identifier)
             var backupPath = GetBackupPath(module, identifier);
             var backupSymbolPath = Path.ChangeExtension(backupPath, ".pdb");
             File.Copy(module, backupPath, true);
+            _backupList.Add(module, backupPath);
 
             var symbolFile = Path.ChangeExtension(module, ".pdb");
             if (File.Exists(symbolFile))
             {
                 File.Copy(symbolFile, backupSymbolPath, true);
+                _backupList.Add(symbolFile, backupSymbolPath);
             }
         }
 
@@ -108,18 +117,39 @@ public static void RestoreOriginalModule(string module, string identifier)
             {
                 File.Copy(backupPath, module, true);
                 File.Delete(backupPath);
+                _backupList.Remove(module);
             }, retryStrategy, 10);
 
             RetryHelper.Retry(() =>
             {
                 if (File.Exists(backupSymbolPath))
                 {
-                    File.Copy(backupSymbolPath, Path.ChangeExtension(module, ".pdb"), true);
+                    string symbolFile = Path.ChangeExtension(module, ".pdb");
+                    File.Copy(backupSymbolPath, symbolFile, true);
                     File.Delete(backupSymbolPath);
+                    _backupList.Remove(symbolFile);
                 }
             }, retryStrategy, 10);
         }
 
+        public static void RestoreOriginalModules()
+        {
+            // Restore the original module - retry up to 10 times, since the destination file could be locked
+            // See: https://github.com/tonerdo/coverlet/issues/25
+            var retryStrategy = CreateRetryStrategy();
+
+            foreach (string key in _backupList.Keys.ToList())
+            {
+                string backupPath = _backupList[key];
+                RetryHelper.Retry(() =>
+                {
+                    File.Copy(backupPath, key, true);
+                    File.Delete(backupPath);
+                    _backupList.Remove(key);
+                }, retryStrategy, 10);
+            }
+        }
+
         public static void DeleteHitsFile(string path)
         {
             // Retry hitting the hits file - retry up to 10 times, since the file could be locked