refactor: decouple foreign functions from WasmBase. (#190)

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

include/proxy-wasm/exports.h

@@ -30,6 +30,36 @@ ::proxy_wasm::ContextBase *contextOrEffectiveContext();
 
 extern thread_local ContextBase *current_context_;
 
+/**
+ * WasmForeignFunction is used for registering host-specific host functions.
+ * A foreign function can be registered via RegisterForeignFunction and available
+ * to Wasm modules via proxy_call_foreign_function.
+ * @param wasm is the WasmBase which the Wasm module is running on.
+ * @param argument is the view to the argument to the function passed by the module.
+ * @param alloc_result is used to allocate the result data of this foreign function.
+ */
+using WasmForeignFunction = std::function<WasmResult(
+    WasmBase &wasm, std::string_view argument, std::function<void *(size_t size)> alloc_result)>;
+
+/**
+ * Used to get the foreign function registered via RegisterForeignFunction for a given name.
+ * @param function_name is the name used to lookup the foreign function table.
+ * @return a WasmForeignFunction if registered.
+ */
+WasmForeignFunction getForeignFunction(std::string_view function_name);
+
+/**
+ * RegisterForeignFunction is used to register a foreign function in the lookup table
+ * used internally in getForeignFunction.
+ */
+struct RegisterForeignFunction {
+  /**
+   * @param function_name is the key for this foreign function.
+   * @param f is the function instance.
+   */
+  RegisterForeignFunction(std::string function_name, WasmForeignFunction f);
+};
+
 namespace exports {
 
 template <typename Pairs> size_t pairsSize(const Pairs &result) {

include/proxy-wasm/wasm.h

@@ -34,11 +34,8 @@ namespace proxy_wasm {
 #include "proxy_wasm_common.h"
 
 class ContextBase;
-class WasmBase;
 class WasmHandleBase;
 
-using WasmForeignFunction =
-    std::function<WasmResult(WasmBase &, std::string_view, std::function<void *(size_t size)>)>;
 using WasmVmFactory = std::function<std::unique_ptr<WasmVm>()>;
 using CallOnThreadFunction = std::function<void(std::function<void()>)>;
 
@@ -129,8 +126,6 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
   bool copyToPointerSize(std::string_view s, uint64_t ptr_ptr, uint64_t size_ptr);
   template <typename T> bool setDatatype(uint64_t ptr, const T &t);
 
-  WasmForeignFunction getForeignFunction(std::string_view function_name);
-
   void fail(FailState fail_state, std::string_view message) {
     error(message);
     failed_ = fail_state;
@@ -439,8 +434,4 @@ template <typename T> inline bool WasmBase::setDatatype(uint64_t ptr, const T &t
   return wasm_vm_->setMemory(ptr, sizeof(T), &t);
 }
 
-struct RegisterForeignFunction {
-  RegisterForeignFunction(std::string name, WasmForeignFunction f);
-};
-
 } // namespace proxy_wasm

src/exports.cc

@@ -36,6 +36,24 @@ ContextBase *contextOrEffectiveContext() {
 // of current_context_.
 extern thread_local uint32_t effective_context_id_;
 
+std::unordered_map<std::string, WasmForeignFunction> &foreignFunctions() {
+  static auto ptr = new std::unordered_map<std::string, WasmForeignFunction>;
+  return *ptr;
+}
+
+WasmForeignFunction getForeignFunction(std::string_view function_name) {
+  auto foreign_functions = foreignFunctions();
+  auto it = foreign_functions.find(std::string(function_name));
+  if (it != foreign_functions.end()) {
+    return it->second;
+  }
+  return nullptr;
+}
+
+RegisterForeignFunction::RegisterForeignFunction(std::string name, WasmForeignFunction f) {
+  foreignFunctions()[name] = f;
+}
+
 namespace exports {
 
 namespace {
@@ -220,7 +238,7 @@ Word call_foreign_function(Word function_name, Word function_name_size, Word arg
   if (!args_opt) {
     return WasmResult::InvalidMemoryAccess;
   }
-  auto f = context->wasm()->getForeignFunction(function.value());
+  auto f = getForeignFunction(function.value());
   if (!f) {
     return WasmResult::NotFound;
   }

src/wasm.cc

@@ -45,7 +45,6 @@ thread_local std::unordered_map<std::string, std::weak_ptr<PluginHandleBase>> lo
 // Map from Wasm Key to the base Wasm instance, using a pointer to avoid the initialization fiasco.
 std::mutex base_wasms_mutex;
 std::unordered_map<std::string, std::weak_ptr<WasmHandleBase>> *base_wasms = nullptr;
-std::unordered_map<std::string, WasmForeignFunction> *foreign_functions = nullptr;
 
 std::vector<uint8_t> Sha256(const std::vector<std::string_view> parts) {
   uint8_t sha256[SHA256_DIGEST_LENGTH];
@@ -85,13 +84,6 @@ class WasmBase::ShutdownHandle {
   std::shared_ptr<WasmBase> wasm_;
 };
 
-RegisterForeignFunction::RegisterForeignFunction(std::string name, WasmForeignFunction f) {
-  if (!foreign_functions) {
-    foreign_functions = new std::remove_reference<decltype(*foreign_functions)>::type;
-  }
-  (*foreign_functions)[name] = f;
-}
-
 void WasmBase::registerCallbacks() {
 #define _REGISTER(_fn)                                                                             \
   wasm_vm_->registerCallback(                                                                      \
@@ -454,14 +446,6 @@ void WasmBase::finishShutdown() {
   }
 }
 
-WasmForeignFunction WasmBase::getForeignFunction(std::string_view function_name) {
-  auto it = foreign_functions->find(std::string(function_name));
-  if (it != foreign_functions->end()) {
-    return it->second;
-  }
-  return nullptr;
-}
-
 std::shared_ptr<WasmHandleBase> createWasm(std::string vm_key, std::string code,
                                            std::shared_ptr<PluginBase> plugin,
                                            WasmHandleFactory factory,