Make MemoryCache id optional, and include owner id on lookup

cloudflare · Feb 14, 2024 · b3a28c6 · b3a28c6
1 parent 7110195
commit b3a28c6
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 12 deletions.
diff --git a/src/workerd/api/memory-cache.c++ b/src/workerd/api/memory-cache.c++
@@ -402,7 +402,7 @@ public:
 
   // Gets an existing SharedMemoryCache instance or creates a new one if no
   // cache with the given id exists.
-  SharedMemoryCache& getInstance(kj::StringPtr cacheId) const override;
+  SharedMemoryCache& getInstance(kj::StringPtr cacheId, uint32_t ownerId) const override;
 
 private:
   using HashMap = kj::HashMap<kj::String, kj::Own<SharedMemoryCache>>;
@@ -416,17 +416,19 @@ private:
   // SharedMemoryCache::uuid() instead.
 };
 
-SharedMemoryCache& MemoryCacheMap::getInstance(kj::StringPtr cacheId) const {
+SharedMemoryCache& MemoryCacheMap::getInstance(
+    kj::StringPtr cacheId, uint32_t ownerId) const {
   auto lock = caches.lockExclusive();
-  return *lock->findOrCreate(cacheId, [this, &cacheId]() {
+  auto id = kj::str(cacheId, "::", ownerId);
+  return *lock->findOrCreate(id, [this, &id]() {
     auto handler = additionalResizeMemoryLimitHandler.map([](
             const SharedMemoryCache::AdditionalResizeMemoryLimitHandler& handler)
                 -> SharedMemoryCache::AdditionalResizeMemoryLimitHandler& {
       return const_cast<SharedMemoryCache::AdditionalResizeMemoryLimitHandler&>(handler);
     });
     return HashMap::Entry{
-      kj::str(cacheId),
-      kj::heap<SharedMemoryCache>(cacheId, handler)
+      kj::str(id),
+      kj::heap<SharedMemoryCache>(id, handler)
     };
   });
 }

diff --git a/src/workerd/api/memory-cache.h b/src/workerd/api/memory-cache.h
@@ -434,7 +434,8 @@ class MemoryCache: public jsg::Object {
 // uses a simple in-memory map to store the SharedMemoryCache instances.
 class MemoryCacheProvider {
 public:
-  virtual SharedMemoryCache& getInstance(kj::StringPtr cacheId) const = 0;
+  virtual ~MemoryCacheProvider() noexcept(false) = default;
+  virtual SharedMemoryCache& getInstance(kj::StringPtr cacheId, uint32_t ownerId) const = 0;
 
   static kj::Own<MemoryCacheProvider> createDefault(
     kj::Maybe<SharedMemoryCache::AdditionalResizeMemoryLimitHandler> additionalResizeMemoryLimitHandler = kj::none);

diff --git a/src/workerd/api/tests/memory-cache-test.wd-test b/src/workerd/api/tests/memory-cache-test.wd-test
@@ -11,23 +11,20 @@ const unitTests :Workerd.Config = (
         compatibilityFlags = ["nodejs_compat"],
         bindings = [
           (name = "CACHE", memoryCache = (
-            id = "abc123",
             limits = (
               maxKeys = 2,
               maxValueSize = 1024,
               maxTotalValueSize = 2056,
             ),
           )),
           (name = "CACHE2", memoryCache = (
-            id = "abc321",
             limits = (
               maxKeys = 2,
               maxValueSize = 1024,
               maxTotalValueSize = 2056,
             ),
           )),
           (name = "CACHE3", memoryCache = (
-            id = "xyz321",
             limits = (
               maxKeys = 2,
               maxValueSize = 500,

diff --git a/src/workerd/server/server.c++ b/src/workerd/server/server.c++
@@ -28,6 +28,7 @@
 #include <workerd/util/http-util.h>
 #include <workerd/api/actor-state.h>
 #include <workerd/util/mimetype.h>
+#include <workerd/util/uuid.h>
 #include <workerd/util/use-perfetto-categories.h>
 #include "workerd-api.h"
 #include "workerd/io/hibernation-manager.h"
@@ -2423,9 +2424,17 @@ static kj::Maybe<WorkerdApi::Global> createBinding(
     }
     case config::Worker::Binding::MEMORY_CACHE: {
       auto cache = binding.getMemoryCache();
-      KJ_REQUIRE(cache.hasId() && cache.hasLimits());
+      // TODO(cleanup): Should we have some reasonable default for these so they can
+      // be optional?
+      KJ_REQUIRE(cache.hasLimits());
       Global::MemoryCache cacheCopy;
-      cacheCopy.cacheId = kj::str(cache.getId());
+      // The id is optional. If provided, then multiple bindings with the same id will
+      // share the same cache. Otherwise, a unique id is generated for the cache.
+      if (cache.hasId()) {
+        cacheCopy.cacheId = kj::str(cache.getId());
+      } else {
+        cacheCopy.cacheId = randomUUID(kj::none);
+      }
       auto limits = cache.getLimits();
       cacheCopy.maxKeys = limits.getMaxKeys();
       cacheCopy.maxValueSize = limits.getMaxValueSize();

diff --git a/src/workerd/server/workerd-api.c++ b/src/workerd/server/workerd-api.c++
@@ -614,7 +614,8 @@ static v8::Local<v8::Value> createBindingValue(
       api::SharedMemoryCache::Limits limits = {.maxKeys = cache.maxKeys,
         .maxValueSize = cache.maxValueSize,
         .maxTotalValueSize = cache.maxTotalValueSize};
-      api::SharedMemoryCache& sharedCache = memoryCacheProvider.getInstance(cache.cacheId);
+      api::SharedMemoryCache& sharedCache =
+          memoryCacheProvider.getInstance(cache.cacheId, ownerId);
       api::SharedMemoryCache::Use cacheUse(sharedCache, limits);
       value = lock.wrap(context, jsg::alloc<api::MemoryCache>(kj::mv(cacheUse)));
     }