Pause and Resume functionality

SnaffCore/Checkpoint/ScanState.cs

@@ -0,0 +1,310 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.IO;
+using System.Text;
+using SnaffCore.Concurrency;
+
+namespace SnaffCore.Checkpoint
+{
+    /// <summary>
+    /// Tracks scan progress so a run can be resumed after an interruption
+    /// (network drop, laptop sleep, crash, Ctrl-C, etc).
+    ///
+    /// Design goals:
+    ///  - No external dependencies (works fine on .NET Framework 4.5.1).
+    ///  - Crash-safe: every completed unit of work is appended to disk
+    ///    immediately and flushed, instead of being buffered and saved
+    ///    periodically. If the machine dies mid-scan, we only lose the
+    ///    (small) amount of work that was in flight at that instant -
+    ///    nothing that was already marked done.
+    ///  - Same command line every time: point -w/--checkpoint at a file.
+    ///    First run creates it. Any later run against the same file
+    ///    skips everything already recorded there and carries on.
+    ///
+    /// File format is a simple pipe-delimited append log:
+    ///   D|<absolute directory path>              -> directory fully enumerated
+    ///   C|<computer name>                        -> share list for computer fully enumerated
+    ///   S|<computer name>|<share unc path>        -> a share that was found scannable on that computer
+    /// </summary>
+    public class ScanState
+    {
+        private static readonly object _initLock = new object();
+        private static ScanState _instance;
+
+        public static ScanState Instance
+        {
+            get
+            {
+                if (_instance == null)
+                {
+                    lock (_initLock)
+                    {
+                        if (_instance == null)
+                        {
+                            _instance = new ScanState();
+                        }
+                    }
+                }
+                return _instance;
+            }
+        }
+
+        private readonly object _writeLock = new object();
+        private ConcurrentDictionary<string, byte> _completedDirs;
+        private ConcurrentDictionary<string, byte> _completedComputers;
+        private ConcurrentDictionary<string, ConcurrentBag<string>> _computerShares;
+        private ConcurrentDictionary<string, int> _pendingFileCounts;
+
+        private string _checkpointPath;
+        private StreamWriter _writer;
+        public bool Enabled { get; private set; }
+
+        private ScanState()
+        {
+            _completedDirs = new ConcurrentDictionary<string, byte>(StringComparer.OrdinalIgnoreCase);
+            _completedComputers = new ConcurrentDictionary<string, byte>(StringComparer.OrdinalIgnoreCase);
+            _computerShares = new ConcurrentDictionary<string, ConcurrentBag<string>>(StringComparer.OrdinalIgnoreCase);
+            _pendingFileCounts = new ConcurrentDictionary<string, int>(StringComparer.OrdinalIgnoreCase);
+        }
+
+        /// <summary>
+        /// Call once at startup. If the checkpoint file already has content,
+        /// it's replayed into memory so this run resumes instead of restarting.
+        /// </summary>
+        public void Load(string checkpointPath)
+        {
+            if (string.IsNullOrWhiteSpace(checkpointPath))
+            {
+                Enabled = false;
+                return;
+            }
+
+            BlockingMq mq = BlockingMq.GetMq();
+            _checkpointPath = checkpointPath;
+            Enabled = true;
+
+            int resumedDirs = 0;
+            int resumedComputers = 0;
+            int resumedShares = 0;
+
+            if (File.Exists(_checkpointPath))
+            {
+                using (FileStream fs = new FileStream(_checkpointPath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite))
+                using (StreamReader reader = new StreamReader(fs, Encoding.UTF8))
+                {
+                    string line;
+                    while ((line = reader.ReadLine()) != null)
+                    {
+                        if (string.IsNullOrWhiteSpace(line))
+                        {
+                            continue;
+                        }
+
+                        string[] parts = line.Split(new[] { '|' }, 3);
+                        if (parts.Length < 2)
+                        {
+                            continue;
+                        }
+
+                        switch (parts[0])
+                        {
+                            case "D":
+                                if (_completedDirs.TryAdd(parts[1], 0))
+                                {
+                                    resumedDirs++;
+                                }
+                                break;
+                            case "C":
+                                if (_completedComputers.TryAdd(parts[1], 0))
+                                {
+                                    resumedComputers++;
+                                }
+                                break;
+                            case "S":
+                                if (parts.Length == 3)
+                                {
+                                    ConcurrentBag<string> bag = _computerShares.GetOrAdd(parts[1], _ => new ConcurrentBag<string>());
+                                    bag.Add(parts[2]);
+                                    resumedShares++;
+                                }
+                                break;
+                        }
+                    }
+                }
+
+                mq.Info(string.Format(
+                    "Resuming from checkpoint '{0}': {1} directories, {2} computers, and {3} shares already accounted for.",
+                    _checkpointPath, resumedDirs, resumedComputers, resumedShares));
+            }
+            else
+            {
+                mq.Info("Checkpoint file '" + _checkpointPath + "' doesn't exist yet, starting a fresh scan and recording progress to it.");
+            }
+
+            // open for append; each write below is flushed immediately so progress survives a hard kill.
+            FileStream appendStream = new FileStream(_checkpointPath, FileMode.Append, FileAccess.Write, FileShare.Read);
+            _writer = new StreamWriter(appendStream, Encoding.UTF8) { AutoFlush = true };
+        }
+
+        public void Flush()
+        {
+            if (!Enabled || _writer == null)
+            {
+                return;
+            }
+            lock (_writeLock)
+            {
+                try
+                {
+                    _writer.Flush();
+                }
+                catch (Exception e)
+                {
+                    BlockingMq.GetMq().Degub("Failed to flush checkpoint file: " + e.Message);
+                }
+            }
+        }
+
+        public void Close()
+        {
+            if (!Enabled || _writer == null)
+            {
+                return;
+            }
+            lock (_writeLock)
+            {
+                try
+                {
+                    _writer.Flush();
+                    _writer.Close();
+                    _writer = null;
+                }
+                catch (Exception e)
+                {
+                    BlockingMq.GetMq().Degub("Failed to close checkpoint file: " + e.Message);
+                }
+            }
+        }
+
+        private void WriteLine(string line)
+        {
+            if (!Enabled || _writer == null)
+            {
+                return;
+            }
+            lock (_writeLock)
+            {
+                try
+                {
+                    _writer.WriteLine(line);
+                }
+                catch (Exception e)
+                {
+                    // don't let checkpoint I/O problems take down the scan itself
+                    BlockingMq.GetMq().Degub("Failed to write to checkpoint file: " + e.Message);
+                }
+            }
+        }
+
+        public bool IsDirComplete(string dir)
+        {
+            return Enabled && _completedDirs.ContainsKey(dir);
+        }
+
+        public void MarkDirComplete(string dir)
+        {
+            if (!Enabled)
+            {
+                return;
+            }
+            if (_completedDirs.TryAdd(dir, 0))
+            {
+                WriteLine("D|" + dir);
+            }
+        }
+
+        /// <summary>
+        /// Call once per directory, right after Directory.GetFiles() succeeds, with the
+        /// number of files found there. A directory is only checkpointed as complete once
+        /// every one of those files has actually finished being scanned (see
+        /// FileScanFinished) - NOT when the scan tasks are merely queued. Tree-walking
+        /// (listing names) races far ahead of content-scanning (reading/grepping bytes),
+        /// so marking completion at queue-time would checkpoint directories whose files
+        /// were never actually scanned, and a resumed run would skip them entirely.
+        /// </summary>
+        public void BeginDirFiles(string dir, int fileCount)
+        {
+            if (!Enabled)
+            {
+                return;
+            }
+            if (fileCount <= 0)
+            {
+                // no files in this dir to wait on, safe to checkpoint immediately.
+                MarkDirComplete(dir);
+                return;
+            }
+            _pendingFileCounts[dir] = fileCount;
+        }
+
+        /// <summary>
+        /// Call exactly once per file that was counted in BeginDirFiles, when that file's
+        /// scan attempt is done (success OR failure - a permanently-failing file shouldn't
+        /// block the directory from ever checkpointing). Once every file counted for a
+        /// directory has reported in, the directory is marked complete.
+        /// </summary>
+        public void FileScanFinished(string dir)
+        {
+            if (!Enabled)
+            {
+                return;
+            }
+            int remaining = _pendingFileCounts.AddOrUpdate(dir, 0, (_, current) => current - 1);
+            if (remaining <= 0)
+            {
+                int dummy;
+                _pendingFileCounts.TryRemove(dir, out dummy);
+                MarkDirComplete(dir);
+            }
+        }
+
+        public bool IsComputerComplete(string computer)
+        {
+            return Enabled && _completedComputers.ContainsKey(computer);
+        }
+
+        public void MarkComputerComplete(string computer)
+        {
+            if (!Enabled)
+            {
+                return;
+            }
+            if (_completedComputers.TryAdd(computer, 0))
+            {
+                WriteLine("C|" + computer);
+            }
+        }
+
+        public void RecordShare(string computer, string sharePath)
+        {
+            if (!Enabled)
+            {
+                return;
+            }
+            ConcurrentBag<string> bag = _computerShares.GetOrAdd(computer, _ => new ConcurrentBag<string>());
+            bag.Add(sharePath);
+            WriteLine("S|" + computer + "|" + sharePath);
+        }
+
+        public IEnumerable<string> GetComputerShares(string computer)
+        {
+            ConcurrentBag<string> bag;
+            if (_computerShares.TryGetValue(computer, out bag))
+            {
+                return bag;
+            }
+            return new string[0];
+        }
+    }
+}

SnaffCore/Config/Options.cs

@@ -32,6 +32,11 @@ public partial class Options
         public string CurrentUser { get; set; } = WindowsIdentity.GetCurrent().Name;
         public string RuleDir { get; set; }
 
+        // Checkpoint/resume options. Point this at a file and Snaffler will
+        // record progress there and skip already-completed work on a re-run
+        // against the same file (e.g. after a dropped connection or a sleep).
+        public string CheckpointFile { get; set; }
+
         public int TimeOut { get; set; } = 5;
 
         // Concurrency Options

SnaffCore/ShareFind/ShareFinder.cs

@@ -30,6 +30,31 @@ public ShareFinder()
 
         internal void GetComputerShares(string computer)
         {
+            // Resume support: if we already fully enumerated this computer's shares on
+            // a previous (interrupted) run, don't hit the network again - just replay
+            // the shares we already found scannable and let TreeWalker's own resume
+            // logic (ScanState.IsDirComplete) figure out what's left to do inside them.
+            if (SnaffCore.Checkpoint.ScanState.Instance.IsComputerComplete(computer))
+            {
+                Mq.Trace("Skipping share enumeration on " + computer + ", already completed in a previous run.");
+                foreach (string knownShare in SnaffCore.Checkpoint.ScanState.Instance.GetComputerShares(computer))
+                {
+                    TreeTaskScheduler.New(() =>
+                    {
+                        try
+                        {
+                            TreeWalker.WalkTree(knownShare);
+                        }
+                        catch (Exception e)
+                        {
+                            Mq.Error("Exception in TreeWalker task for share " + knownShare);
+                            Mq.Error(e.ToString());
+                        }
+                    });
+                }
+                return;
+            }
+
             // find the shares
             HostShareInfo[] hostShareInfos = GetHostShareInfo(computer);
 
@@ -183,6 +208,10 @@ internal void GetComputerShares(string computer)
 
                             if (MyOptions.ScanFoundShares)
                             {
+                                // remember this share so a resumed run can re-queue it without
+                                // re-hitting the network to enumerate shares on this computer.
+                                SnaffCore.Checkpoint.ScanState.Instance.RecordShare(computer, shareResult.SharePath);
+
                                 Mq.Trace("Creating a TreeWalker task for " + shareResult.SharePath);
                                 TreeTaskScheduler.New(() =>
                                 {
@@ -206,6 +235,10 @@ internal void GetComputerShares(string computer)
                     }
                 }
             }
+
+            // Share enumeration for this computer is done; a resumed run can skip
+            // straight to replaying the shares we already recorded above.
+            SnaffCore.Checkpoint.ScanState.Instance.MarkComputerComplete(computer);
         }
 
         internal bool IsShareReadable(string share)

SnaffCore/SnaffCon.cs

@@ -42,6 +42,8 @@ public SnaffCon(Options options)
             MyOptions = options;
             Mq = BlockingMq.GetMq();
 
+            SnaffCore.Checkpoint.ScanState.Instance.Load(MyOptions.CheckpointFile);
+
             int shareThreads = MyOptions.ShareThreads;
             int treeThreads = MyOptions.TreeThreads;
             int fileThreads = MyOptions.FileThreads;
@@ -156,6 +158,7 @@ public void Execute()
             TimeSpan runSpan = finished.Subtract(StartTime);
             Mq.Info("Finished at " + finished.ToLocalTime());
             Mq.Info("Snafflin' took " + runSpan);
+            SnaffCore.Checkpoint.ScanState.Instance.Close();
             Mq.Finish();
         }