diff --git a/BUILD b/BUILD index f29f4de8..8b330199 100644 --- a/BUILD +++ b/BUILD @@ -76,6 +76,7 @@ cc_library( "src/signature_util.cc", "src/vm_id_handle.cc", "src/wasm.cc", + "src/wasm.h", ], hdrs = [ "include/proxy-wasm/bytecode_util.h", diff --git a/src/wasm.cc b/src/wasm.cc index 9d2566fc..cb1dd9b3 100644 --- a/src/wasm.cc +++ b/src/wasm.cc @@ -670,4 +670,24 @@ void clearWasmCachesForTesting() { } } +std::vector staleLocalPluginsKeysForTesting() { + std::vector keys; + for (const auto &kv : local_plugins) { + if (kv.second.expired()) { + keys.push_back(kv.first); + } + } + return keys; +} + +std::vector staleLocalWasmsKeysForTesting() { + std::vector keys; + for (const auto &kv : local_wasms) { + if (kv.second.expired()) { + keys.push_back(kv.first); + } + } + return keys; +} + } // namespace proxy_wasm diff --git a/src/wasm.h b/src/wasm.h new file mode 100644 index 00000000..882dbcd9 --- /dev/null +++ b/src/wasm.h @@ -0,0 +1,25 @@ +// Copyright 2023 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace proxy_wasm { + +std::vector staleLocalWasmsKeysForTesting(); +std::vector staleLocalPluginsKeysForTesting(); + +} // namespace proxy_wasm diff --git a/test/wasm_test.cc b/test/wasm_test.cc index b8873269..2a3e60cc 100644 --- a/test/wasm_test.cc +++ b/test/wasm_test.cc @@ -20,6 +20,8 @@ #include "test/utility.h" +#include "src/wasm.h" + namespace proxy_wasm { INSTANTIATE_TEST_SUITE_P(WasmEngines, TestVm, testing::ValuesIn(getWasmEngines()), @@ -242,4 +244,84 @@ TEST_P(TestVm, AlwaysApplyCanary) { } } +// Check that there are no stale thread-local cache keys (eventually) +TEST_P(TestVm, CleanupThreadLocalCacheKeys) { + const auto *const plugin_name = "plugin_name"; + const auto *const root_id = "root_id"; + const auto *const vm_id = "vm_id"; + const auto *const vm_config = "vm_config"; + const auto *const plugin_config = "plugin_config"; + const auto fail_open = false; + + WasmHandleFactory wasm_handle_factory = + [this, vm_id, vm_config](std::string_view vm_key) -> std::shared_ptr { + auto base_wasm = std::make_shared(makeVm(engine_), vm_id, vm_config, vm_key, + std::unordered_map{}, + AllowedCapabilitiesMap{}); + return std::make_shared(base_wasm); + }; + + WasmHandleCloneFactory wasm_handle_clone_factory = + [this](const std::shared_ptr &base_wasm_handle) + -> std::shared_ptr { + auto wasm = std::make_shared( + base_wasm_handle, [this]() -> std::unique_ptr { return makeVm(engine_); }); + return std::make_shared(wasm); + }; + + PluginHandleFactory plugin_handle_factory = + [](const std::shared_ptr &base_wasm, + const std::shared_ptr &plugin) -> std::shared_ptr { + return std::make_shared(base_wasm, plugin); + }; + + // Read the minimal loadable binary. + auto source = readTestWasmFile("abi_export.wasm"); + + // Simulate a plugin lifetime. + const auto plugin1 = std::make_shared(plugin_name, root_id, vm_id, engine_, + plugin_config, fail_open, "plugin_1"); + auto base_wasm_handle1 = + createWasm("vm_1", source, plugin1, wasm_handle_factory, wasm_handle_clone_factory, false); + ASSERT_TRUE(base_wasm_handle1 && base_wasm_handle1->wasm()); + + auto local_plugin1 = getOrCreateThreadLocalPlugin( + base_wasm_handle1, plugin1, wasm_handle_clone_factory, plugin_handle_factory); + ASSERT_TRUE(local_plugin1 && local_plugin1->plugin()); + local_plugin1.reset(); + + auto stale_plugins_keys = staleLocalPluginsKeysForTesting(); + EXPECT_EQ(1, stale_plugins_keys.size()); + + // Now we create another plugin with a slightly different key and expect that there are no stale + // thread-local cache entries. + const auto plugin2 = std::make_shared(plugin_name, root_id, vm_id, engine_, + plugin_config, fail_open, "plugin_2"); + auto local_plugin2 = getOrCreateThreadLocalPlugin( + base_wasm_handle1, plugin2, wasm_handle_clone_factory, plugin_handle_factory); + ASSERT_TRUE(local_plugin2 && local_plugin2->plugin()); + + stale_plugins_keys = staleLocalPluginsKeysForTesting(); + EXPECT_TRUE(stale_plugins_keys.empty()); + + // Trigger deletion of the thread-local WasmVM cloned from base_wasm_handle1 by freeing objects + // referencing it. + local_plugin2.reset(); + + auto stale_wasms_keys = staleLocalWasmsKeysForTesting(); + EXPECT_EQ(1, stale_wasms_keys.size()); + + // Create another base WASM handle and invoke WASM thread-local cache key cleanup. + auto base_wasm_handle2 = + createWasm("vm_2", source, plugin2, wasm_handle_factory, wasm_handle_clone_factory, false); + ASSERT_TRUE(base_wasm_handle2 && base_wasm_handle2->wasm()); + + auto local_plugin3 = getOrCreateThreadLocalPlugin( + base_wasm_handle2, plugin2, wasm_handle_clone_factory, plugin_handle_factory); + ASSERT_TRUE(local_plugin3 && local_plugin3->plugin()); + + stale_wasms_keys = staleLocalWasmsKeysForTesting(); + EXPECT_TRUE(stale_wasms_keys.empty()); +} + } // namespace proxy_wasm