SERVER-48416 Write index build progress to internal table on clean shutdown

3 files changed, 46 insertions, 23 deletions

diff --git a/src/mongo/db/sorter/sorter.cpp b/src/mongo/db/sorter/sorter.cpp
index bc427c0a71f..a82915c8302 100644
--- a/src/mongo/db/sorter/sorter.cpp
+++ b/src/mongo/db/sorter/sorter.cpp
@@ -537,7 +537,7 @@ public:
     }
 
     ~NoLimitSorter() {
-        if (!_done) {
+        if (!_done && !this->_shouldKeepFilesOnDestruction) {
             // If done() was never called to return a MergeIterator, then this Sorter still owns
             // file deletion.
             DESTRUCTOR_GUARD(boost::filesystem::remove(this->_fileName));
@@ -669,6 +669,10 @@ public:
     }
 
 private:
+    void spill() {
+        invariant(false, "LimitOneSorter does not spill to disk");
+    }
+
     const Comparator _comp;
     Data _best;
     bool _haveData;  // false at start, set to true on first call to add()
@@ -711,7 +715,7 @@ public:
     }
 
     ~TopKSorter() {
-        if (!_done) {
+        if (!_done && !this->_shouldKeepFilesOnDestruction) {
             // If done() was never called to return a MergeIterator, then this Sorter still owns
             // file deletion.
             DESTRUCTOR_GUARD(boost::filesystem::remove(this->_fileName));
@@ -937,6 +941,24 @@ private:
 
 }  // namespace sorter
 
+template <typename Key, typename Value>
+std::vector<SorterRangeInfo> Sorter<Key, Value>::_getRangeInfos() const {
+    std::vector<SorterRangeInfo> ranges;
+    ranges.reserve(_iters.size());
+
+    std::transform(_iters.begin(), _iters.end(), std::back_inserter(ranges), [](const auto it) {
+        return it->getRangeInfo();
+    });
+
+    return ranges;
+}
+
+template <typename Key, typename Value>
+void Sorter<Key, Value>::persistDataForShutdown() {
+    spill();
+    _shouldKeepFilesOnDestruction = true;
+}
+
 //
 // SortedFileWriter
 //
diff --git a/src/mongo/db/sorter/sorter.h b/src/mongo/db/sorter/sorter.h
index 4ee8fb5dbc8..534c6cc4c02 100644
--- a/src/mongo/db/sorter/sorter.h
+++ b/src/mongo/db/sorter/sorter.h
@@ -153,8 +153,8 @@ public:
 };
 
 struct SorterRangeInfo {
-    std::streampos startOffset;
-    std::streampos endOffset;
+    std::streamoff startOffset;
+    std::streamoff endOffset;
     uint32_t checksum;
 };
 
@@ -224,6 +224,12 @@ public:
                       typename Value::SorterDeserializeSettings>
         Settings;
 
+    struct State {
+        std::string tempDir;
+        std::string fileName;
+        std::vector<SorterRangeInfo> ranges;
+    };
+
     template <typename Comparator>
     static Sorter* make(const SortOptions& opts,
                         const Comparator& comp,
@@ -244,30 +250,24 @@ public:
         return _usedDisk;
     }
 
-    std::string getTempDir() const {
-        return _tempDir;
-    }
-
-    std::string getFileName() const {
-        return _fileName;
+    State getState() const {
+        return {_tempDir, _fileName, _getRangeInfos()};
     }
 
-    std::vector<SorterRangeInfo> getRangeInfos() const {
-        std::vector<SorterRangeInfo> ranges;
-        ranges.reserve(_iters.size());
-
-        std::transform(_iters.begin(), _iters.end(), std::back_inserter(ranges), [](const auto it) {
-            return it->getRangeInfo();
-        });
-
-        return ranges;
-    }
+    void persistDataForShutdown();
 
 protected:
     Sorter() {}  // can only be constructed as a base
 
+    virtual void spill() = 0;
+
+    std::vector<SorterRangeInfo> _getRangeInfos() const;
+
     bool _usedDisk{false};  // Keeps track of whether the sorter used disk or not
 
+    // Whether the files written by this Sorter should be kept on destruction.
+    bool _shouldKeepFilesOnDestruction = false;
+
     std::string _tempDir;
     std::string _fileName;
 
diff --git a/src/mongo/db/sorter/sorter_test.cpp b/src/mongo/db/sorter/sorter_test.cpp
index f9b56cb88cf..9e307fe2ba1 100644
--- a/src/mongo/db/sorter/sorter_test.cpp
+++ b/src/mongo/db/sorter/sorter_test.cpp
@@ -480,12 +480,13 @@ private:
     }
 
     void assertRangeInfo(unowned_ptr<IWSorter> sorter, const SortOptions& opts) {
+        auto state = sorter->getState();
         if (opts.extSortAllowed) {
-            ASSERT_EQ(sorter->getTempDir(), opts.tempDir);
-            ASSERT_NE(sorter->getFileName(), "");
+            ASSERT_EQ(state.tempDir, opts.tempDir);
+            ASSERT_NE(state.fileName, "");
         }
         if (auto numRanges = correctNumRanges()) {
-            ASSERT_EQ(sorter->getRangeInfos().size(), *numRanges);
+            ASSERT_EQ(state.ranges.size(), *numRanges);
         }
     }
 };