SERVER-21220 ESE - mongo side temp file coverage (fixed unit tests)

10 files changed, 180 insertions, 15 deletions

diff --git a/buildscripts/resmokeconfig/suites/aggregation_ese.yml b/buildscripts/resmokeconfig/suites/aggregation_ese.yml
new file mode 100644
index 00000000000..fe670ba48dc
--- /dev/null
+++ b/buildscripts/resmokeconfig/suites/aggregation_ese.yml
@@ -0,0 +1,29 @@
+# Section that is ignored by resmoke.py.
+config_variables:
+- &keyFile src/mongo/db/modules/enterprise/jstests/encryptdb/libs/ekf2
+
+selector:
+  js_test:
+    roots:
+    - jstests/aggregation/*.js
+    - jstests/aggregation/bugs/*.js
+    - src/mongo/db/modules/*/jstests/aggregation/*.js
+    exclude_files:
+    # Skip any tests that run with auth explicitly.
+    - jstests/aggregation/*[aA]uth*.js
+
+executor:
+  js_test:
+    config:
+      shell_options:
+        readMode: commands
+    hooks:
+    - class: CleanEveryN
+      n: 20
+    fixture:
+      class: MongoDFixture
+      mongod_options:
+        enableEncryption: ''
+        encryptionKeyFile: *keyFile
+        set_parameters:
+          enableTestCommands: 1
diff --git a/etc/evergreen.yml b/etc/evergreen.yml
index c48e178d0b4..8924655ba06 100644
--- a/etc/evergreen.yml
+++ b/etc/evergreen.yml
@@ -636,6 +636,17 @@ tasks:
       run_multiple_jobs: true
 
 - <<: *task_template
+  name: aggregation_WT_ese
+  depends_on:
+  - name: jsCore_WT
+  commands:
+  - func: "do setup"
+  - func: "run tests"
+    vars:
+      resmoke_args: --suites=aggregation_ese --storageEngine=wiredTiger
+      run_multiple_jobs: true
+
+- <<: *task_template
   name: aggregation_auth
   depends_on:
   - name: aggregation_WT
@@ -3541,6 +3552,7 @@ buildvariants:
   - name: lint
   - name: aggregation
   - name: aggregation_WT
+  - name: aggregation_WT_ese
   - name: aggregation_auth
   - name: audit
   - name: audit_WT
diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index bf7faa17531..c8d379f88ce 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -121,8 +121,9 @@ docSourceEnv.Library(
         'document_value',
         'expression',
         '$BUILD_DIR/mongo/client/clientdriver',
-        '$BUILD_DIR/mongo/db/service_context',
         '$BUILD_DIR/mongo/db/matcher/expressions',
+        '$BUILD_DIR/mongo/db/service_context',
+        '$BUILD_DIR/mongo/db/storage/wiredtiger/storage_wiredtiger_customization_hooks',
         '$BUILD_DIR/third_party/shim_snappy',
     ],
     LIBDEPS_TAGS=[
diff --git a/src/mongo/db/pipeline/document_source_test.cpp b/src/mongo/db/pipeline/document_source_test.cpp
index 9400aae3d45..bdf9ac79a39 100644
--- a/src/mongo/db/pipeline/document_source_test.cpp
+++ b/src/mongo/db/pipeline/document_source_test.cpp
@@ -28,12 +28,14 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/base/init.h"
 #include "mongo/db/operation_context_noop.h"
 #include "mongo/db/pipeline/dependencies.h"
 #include "mongo/db/pipeline/document_source.h"
 #include "mongo/db/pipeline/expression_context.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/service_context_noop.h"
+#include "mongo/stdx/memory.h"
 #include "mongo/db/storage/storage_options.h"
 #include "mongo/dbtests/dbtests.h"
 #include "mongo/unittest/temp_dir.h"
@@ -44,6 +46,12 @@ bool isMongos() {
 }
 }
 
+// Stub to avoid including the server environment library.
+MONGO_INITIALIZER(SetGlobalEnvironment)(InitializerContext* context) {
+    setGlobalServiceContext(stdx::make_unique<ServiceContextNoop>());
+    return Status::OK();
+}
+
 namespace DocumentSourceTests {
 
 using boost::intrusive_ptr;
diff --git a/src/mongo/db/sorter/SConscript b/src/mongo/db/sorter/SConscript
index 38ac55fafb8..3bcab3fb843 100644
--- a/src/mongo/db/sorter/SConscript
+++ b/src/mongo/db/sorter/SConscript
@@ -2,4 +2,8 @@ Import("env")
 
 sorterEnv = env.Clone()
 sorterEnv.InjectThirdPartyIncludePaths(libraries=['snappy'])
-sorterEnv.CppUnitTest('sorter_test', 'sorter_test.cpp', LIBDEPS=['$BUILD_DIR/third_party/shim_snappy'])
+sorterEnv.CppUnitTest('sorter_test',
+                      'sorter_test.cpp',
+                       LIBDEPS=['$BUILD_DIR/mongo/db/service_context',
+                                '$BUILD_DIR/mongo/db/storage/wiredtiger/storage_wiredtiger_customization_hooks',
+                                '$BUILD_DIR/third_party/shim_snappy'])
diff --git a/src/mongo/db/sorter/sorter.cpp b/src/mongo/db/sorter/sorter.cpp
index d393b450992..11a8fce129c 100644
--- a/src/mongo/db/sorter/sorter.cpp
+++ b/src/mongo/db/sorter/sorter.cpp
@@ -49,11 +49,14 @@
 
 #include <boost/filesystem/operations.hpp>
 #include <snappy.h>
+#include <vector>
 
 #include "mongo/base/string_data.h"
 #include "mongo/config.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/db/service_context.h"
 #include "mongo/db/storage/storage_options.h"
+#include "mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h"
 #include "mongo/platform/atomic_word.h"
 #include "mongo/s/mongos_options.h"
 #include "mongo/util/assert_util.h"
@@ -217,12 +220,28 @@ private:
 
         // negative size means compressed
         const bool compressed = rawSize < 0;
-        const int32_t blockSize = std::abs(rawSize);
+        int32_t blockSize = std::abs(rawSize);
 
         _buffer.reset(new char[blockSize]);
         read(_buffer.get(), blockSize);
         massert(16816, "file too short?", !_done);
 
+        auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
+        if (hooks->enabled()) {
+            std::unique_ptr<char[]> out(new char[blockSize]);
+            size_t outLen;
+            Status status = hooks->unprotectTmpData(reinterpret_cast<uint8_t*>(_buffer.get()),
+                                                    blockSize,
+                                                    reinterpret_cast<uint8_t*>(out.get()),
+                                                    blockSize,
+                                                    &outLen);
+            massert(28841,
+                    str::stream() << "Failed to unprotect data: " << status.toString(),
+                    status.isOK());
+            blockSize = outLen;
+            _buffer.swap(out);
+        }
+
         if (!compressed) {
             _reader.reset(new BufReader(_buffer.get(), blockSize));
             return;
@@ -841,23 +860,46 @@ template <typename Key, typename Value>
 void SortedFileWriter<Key, Value>::spill() {
     namespace str = mongoutils::str;
 
-    if (_buffer.len() == 0)
+    int32_t size = _buffer.len();
+    char* outBuffer = _buffer.buf();
+
+    if (size == 0)
         return;
 
     std::string compressed;
-    snappy::Compress(_buffer.buf(), _buffer.len(), &compressed);
+    snappy::Compress(outBuffer, size, &compressed);
     verify(compressed.size() <= size_t(std::numeric_limits<int32_t>::max()));
 
+    const bool shouldCompress = compressed.size() < size_t(_buffer.len() / 10 * 9);
+    if (shouldCompress) {
+        size = compressed.size();
+        outBuffer = const_cast<char*>(compressed.data());
+    }
+
+    std::unique_ptr<char[]> out;
+    auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
+    if (hooks->enabled()) {
+        size_t protectedSizeMax = size + hooks->additionalBytesForProtectedBuffer();
+        out.reset(new char[protectedSizeMax]);
+        size_t resultLen;
+        Status status = hooks->protectTmpData(reinterpret_cast<const uint8_t*>(outBuffer),
+                                              size,
+                                              reinterpret_cast<uint8_t*>(out.get()),
+                                              protectedSizeMax,
+                                              &resultLen);
+        massert(28842,
+                str::stream() << "Failed to compress data: " << status.toString(),
+                status.isOK());
+        outBuffer = out.get();
+        size = resultLen;
+    }
+
+    // negative size means compressed
+    size = shouldCompress ? -size : size;
     try {
-        if (compressed.size() < size_t(_buffer.len() / 10 * 9)) {
-            const int32_t size = -int32_t(compressed.size());  // negative means compressed
-            _file.write(reinterpret_cast<const char*>(&size), sizeof(size));
-            _file.write(compressed.data(), compressed.size());
-        } else {
-            const int32_t size = _buffer.len();
-            _file.write(reinterpret_cast<const char*>(&size), sizeof(size));
-            _file.write(_buffer.buf(), _buffer.len());
-        }
+        _file.write(reinterpret_cast<const char*>(&size), sizeof(size));
+        _file.write(outBuffer, std::abs(size));
+
     } catch (const std::exception&) {
         msgasserted(16821,
                     str::stream() << "error writing to file \"" << _fileName
diff --git a/src/mongo/db/sorter/sorter_test.cpp b/src/mongo/db/sorter/sorter_test.cpp
index 7c485c365a1..e53108bff6b 100644
--- a/src/mongo/db/sorter/sorter_test.cpp
+++ b/src/mongo/db/sorter/sorter_test.cpp
@@ -33,6 +33,10 @@
 #include <boost/filesystem.hpp>
 
 #include "mongo/config.h"
+#include "mongo/base/init.h"
+#include "mongo/db/service_context.h"
+#include "mongo/db/service_context_noop.h"
+#include "mongo/stdx/memory.h"
 #include "mongo/stdx/thread.h"
 #include "mongo/unittest/temp_dir.h"
 #include "mongo/unittest/unittest.h"
@@ -52,6 +56,12 @@ bool isMongos() {
     return false;
 }
 
+// Stub to avoid including the server environment library.
+MONGO_INITIALIZER(SetGlobalEnvironment)(InitializerContext* context) {
+    setGlobalServiceContext(stdx::make_unique<ServiceContextNoop>());
+    return Status::OK();
+}
+
 //
 // Sorter framework testing utilities
 //
diff --git a/src/mongo/db/storage/wiredtiger/SConscript b/src/mongo/db/storage/wiredtiger/SConscript
index e79d41f16e6..9904c8b6226 100644
--- a/src/mongo/db/storage/wiredtiger/SConscript
+++ b/src/mongo/db/storage/wiredtiger/SConscript
@@ -1,6 +1,17 @@
 Import("env")
 Import("wiredtiger")
 
+env.Library(
+    target='storage_wiredtiger_customization_hooks',
+    source= ['wiredtiger_customization_hooks.cpp'],
+    LIBDEPS= ['$BUILD_DIR/mongo/base',
+              '$BUILD_DIR/mongo/db/service_context'],
+    PROGDEPS_DEPENDENTS=[
+        '$BUILD_DIR/mongo/mongod',
+        '$BUILD_DIR/mongo/mongos',
+    ],
+    )
+
 if wiredtiger:
     wtEnv = env.Clone()
     wtEnv.InjectThirdPartyIncludePaths(libraries=['wiredtiger'])
@@ -10,7 +21,6 @@ if wiredtiger:
     wtEnv.Library(
         target='storage_wiredtiger_core',
         source= [
-            'wiredtiger_customization_hooks.cpp',
             'wiredtiger_global_options.cpp',
             'wiredtiger_index.cpp',
             'wiredtiger_kv_engine.cpp',
@@ -22,6 +32,7 @@ if wiredtiger:
             'wiredtiger_util.cpp',
             ],
         LIBDEPS= [
+            'storage_wiredtiger_customization_hooks',
             '$BUILD_DIR/mongo/base',
             '$BUILD_DIR/mongo/db/namespace_string',
             '$BUILD_DIR/mongo/db/catalog/collection_options',
@@ -57,6 +68,7 @@ if wiredtiger:
             'wiredtiger_server_status.cpp',
             ],
         LIBDEPS=['storage_wiredtiger_core',
+                 'storage_wiredtiger_customization_hooks',
                  '$BUILD_DIR/mongo/db/storage/kv/kv_engine',
                  '$BUILD_DIR/mongo/db/storage/storage_engine_lock_file',
                  '$BUILD_DIR/mongo/db/storage/storage_engine_metadata',
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp
index b2e51b5f9d1..48f135e1bae 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp
@@ -65,6 +65,10 @@ WiredTigerCustomizationHooks* WiredTigerCustomizationHooks::get(ServiceContext*
 
 EmptyWiredTigerCustomizationHooks::~EmptyWiredTigerCustomizationHooks() {}
 
+bool EmptyWiredTigerCustomizationHooks::enabled() const {
+    return false;
+}
+
 bool EmptyWiredTigerCustomizationHooks::restartRequired() {
     return false;
 }
@@ -73,4 +77,15 @@ std::string EmptyWiredTigerCustomizationHooks::getOpenConfig(StringData tableNam
     return "";
 }
 
+Status EmptyWiredTigerCustomizationHooks::protectTmpData(
+    const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) {
+    return Status(ErrorCodes::InternalError,
+                  "Customization hooks must be enabled to use preprocessTmpData.");
+}
+
+Status EmptyWiredTigerCustomizationHooks::unprotectTmpData(
+    const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) {
+    return Status(ErrorCodes::InternalError,
+                  "Customization hooks must be enabled to use postprocessTmpData.");
+}
 }  // namespace mongo
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h
index 2a12fad4d6b..c3bc0e1d08b 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h
@@ -49,6 +49,11 @@ public:
     virtual ~WiredTigerCustomizationHooks() = default;
 
     /**
+     * Returns true if the customization hooks are enabled.
+     */
+    virtual bool enabled() const = 0;
+
+    /**
      * Perform any encryption engine initialization/sanity checking that needs to happen after
      * storage engine initialization but before the server starts accepting incoming connections.
      *
@@ -61,6 +66,25 @@ public:
      *  provided table name
      */
     virtual std::string getOpenConfig(StringData tableName) = 0;
+
+    /**
+     * Returns the maximum size addition when doing transforming temp data.
+     */
+    size_t additionalBytesForProtectedBuffer() {
+        return 32;
+    }
+
+    /**
+     * Transform temp data to non-readable form before writing it to disk.
+     */
+    virtual Status protectTmpData(
+        const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) = 0;
+
+    /**
+     * Tranforms temp data back to readable form, after reading from disk.
+     */
+    virtual Status unprotectTmpData(
+        const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) = 0;
 };
 
 // Empty default implementation of the abstract class WiredTigerCustomizationHooks
@@ -68,8 +92,16 @@ class EmptyWiredTigerCustomizationHooks : public WiredTigerCustomizationHooks {
 public:
     ~EmptyWiredTigerCustomizationHooks() override;
 
+    bool enabled() const override;
+
     bool restartRequired() override;
 
     std::string getOpenConfig(StringData tableName) override;
+
+    Status protectTmpData(
+        const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) override;
+
+    Status unprotectTmpData(
+        const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) override;
 };
 }  // namespace mongo