SERVER-34944 Create a shim-override mechanism.

Modules for MongoDB need to be able to indicate their own overrides
of shimmed funcitons.  These overrides must be able to access the
underlying registered shim.  Further, the overrides must be installed
in the correct position in the initializer dependency graph.

1 files changed, 60 insertions, 6 deletions

diff --git a/src/mongo/base/shim.h b/src/mongo/base/shim.h
index a782109d24d..14de94df7c8 100644
--- a/src/mongo/base/shim.h
+++ b/src/mongo/base/shim.h
@@ -193,6 +193,15 @@ const bool checkShimsViaTUHook = false;
         }                                                                                         \
     }
 
+
+/**
+ * Evaluates to a string which represents the `MONGO_INITIALIZER` step name in which this specific
+ * shim is registered.  This can be useful to make sure that some initializers depend upon a shim's
+ * execution and presence in a binary.
+ */
+#define MONGO_SHIM_DEPENDENCY(...) MONGO_SHIM_EVIL_STRINGIFY_(__VA_ARGS__)
+#define MONGO_SHIM_EVIL_STRINGIFY_(args) #args
+
 /**
  * Define a shimmable function with name `SHIM_NAME`, returning a value of type `RETURN_TYPE`, with
  * any arguments.  This shim definition macro should go in the associated C++ file to the header
@@ -205,20 +214,25 @@ const bool checkShimsViaTUHook = false;
     namespace {                                                                               \
     namespace shim_namespace##LN {                                                            \
         using ShimType = decltype(__VA_ARGS__);                                               \
+        ::mongo::Status initializerGroupStartup(::mongo::InitializerContext*) {               \
+            return Status::OK();                                                              \
+        }                                                                                     \
+        ::mongo::GlobalInitializerRegisterer _mongoInitializerRegisterer(                     \
+            std::string(MONGO_SHIM_DEPENDENCY(__VA_ARGS__)),                                  \
+            {},                                                                               \
+            {MONGO_SHIM_DEPENDENTS},                                                          \
+            mongo::InitializerFunction(initializerGroupStartup));                             \
     } /*namespace shim_namespace*/                                                            \
     } /*namespace*/                                                                           \
     shim_namespace##LN::ShimType::MongoShimImplGuts::LibTUHookTypeBase::LibTUHookTypeBase() = \
         default;                                                                              \
     shim_namespace##LN::ShimType __VA_ARGS__{};
 
-#define MONGO_SHIM_EVIL_STRINGIFY_(args) #args
-
-
 /**
  * Define an implementation of a shimmable function with name `SHIM_NAME`.  The compiler will check
  * supplied parameters for correctness.  This shim registration macro should go in the associated
  * C++ implementation file to the header where a SHIM was defined.   Such a file would be a mock
- * implementation or a real implementation, for example
+ * implementation or a real implementation, for example.
  */
 #define MONGO_REGISTER_SHIM(/*SHIM_NAME*/...) MONGO_REGISTER_SHIM_1(__LINE__, __VA_ARGS__)
 #define MONGO_REGISTER_SHIM_1(LN, ...) MONGO_REGISTER_SHIM_2(LN, __VA_ARGS__)
@@ -240,9 +254,9 @@ const bool checkShimsViaTUHook = false;
         }                                                                                       \
                                                                                                 \
         const ::mongo::GlobalInitializerRegisterer registrationHook{                            \
-            std::string(MONGO_SHIM_EVIL_STRINGIFY_(__VA_ARGS__)),                               \
+            std::string(MONGO_SHIM_DEPENDENCY(__VA_ARGS__) "_registration"),                    \
             {},                                                                                 \
-            {MONGO_SHIM_DEPENDENTS},                                                            \
+            {MONGO_SHIM_DEPENDENCY(__VA_ARGS__), MONGO_SHIM_DEPENDENTS},                        \
             mongo::InitializerFunction(createInitializerRegistration)};                         \
     } /*namespace shim_namespace*/                                                              \
     } /*namespace*/                                                                             \
@@ -255,3 +269,43 @@ const bool checkShimsViaTUHook = false;
                                                                and return value (using arrow    \
                                                                notation).  Then they write the  \
                                                                body. */
+
+/**
+ * Define an overriding implementation of a shimmable function with `SHIM_NAME`.  The compiler will
+ * check the supplied parameters for correctness.  This shim override macro should go in the
+ * associated C++ implementation file to the header where a SHIM was defined.  Such a file
+ * specifying an override would be a C++ implementation used by a mongodb extension module.
+ * This creates a runtime dependency upon the original registration being linked in.
+ */
+#define MONGO_OVERRIDE_SHIM(/*SHIM_NAME*/...) MONGO_OVERRIDE_SHIM_1(__LINE__, __VA_ARGS__)
+#define MONGO_OVERRIDE_SHIM_1(LN, ...) MONGO_OVERRIDE_SHIM_2(LN, __VA_ARGS__)
+#define MONGO_OVERRIDE_SHIM_2(LN, ...)                                                           \
+    namespace {                                                                                  \
+    namespace shim_namespace##LN {                                                               \
+        using ShimType = decltype(__VA_ARGS__);                                                  \
+                                                                                                 \
+        class OverrideImplementation final : public ShimType::MongoShimImplGuts {                \
+            /* Some compilers don't work well with the trailing `override` in this kind of       \
+             * function declaration. */                                                          \
+            ShimType::MongoShimImplGuts::function_type implementation; /* override */            \
+        };                                                                                       \
+                                                                                                 \
+        ::mongo::Status createInitializerOverride(::mongo::InitializerContext* const) {          \
+            static OverrideImplementation overrideImpl;                                          \
+            ShimType::storage::data = &overrideImpl;                                             \
+            return Status::OK();                                                                 \
+        }                                                                                        \
+                                                                                                 \
+        const ::mongo::GlobalInitializerRegisterer overrideHook{                                 \
+            std::string(MONGO_SHIM_DEPENDENCY(__VA_ARGS__) "_override"),                         \
+            {MONGO_SHIM_DEPENDENCY(                                                              \
+                __VA_ARGS__) "_registration"},   /* Override happens after first registration */ \
+            {MONGO_SHIM_DEPENDENCY(__VA_ARGS__), /* Provides impl for this shim */               \
+             MONGO_SHIM_DEPENDENTS},             /* Still a shim registration */                 \
+            mongo::InitializerFunction(createInitializerOverride)};                              \
+    } /*namespace shim_namespace*/                                                               \
+    } /*namespace*/                                                                              \
+                                                                                                 \
+    auto shim_namespace##LN::OverrideImplementation::                                            \
+        implementation /* After this point someone just writes the signature's arguments and     \
+                          return value (using arrow notation).  Then they write the body. */