begin migration to kj/filesystem

cloudflare · Mar 19, 2024 · c04de8e · c04de8e
1 parent c53c908
commit c04de8e
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 60 deletions.
diff --git a/src/workerd/api/pyodide/pyodide.c++ b/src/workerd/api/pyodide/pyodide.c++
@@ -1,10 +1,6 @@
 #include "pyodide.h"
-#include <cstddef>
-#include <fstream>
-#include <iostream>
 #include <kj/string.h>
 #include <workerd/util/string-buffer.h>
-#include <filesystem>
 
 namespace workerd::api::pyodide {
 
@@ -114,67 +110,34 @@ jsg::Ref<PyodideMetadataReader> makePyodideMetadataReader(Worker::Reader conf) {
                                            false /* isTracing */);
 }
 
-namespace {
-// Checks that the given string contains no slashes, forward or backward
-bool hasNoSlashes(const kj::String &str) {
-  for (auto c : str) {
-    if (c == '/' || c == '\\') {
-      return false;
-    }
-  }
-  return true;
-}
-
-kj::String getDiskCachePath(kj::String &key) {
-  StringBuffer<64> fpath(64);
-  fpath.append("/home/garrett/.cache/workerd"); //TODO: do the thingy
-  fpath.append(key);
-  return fpath.toString();
-}
-} //namespace
-
 jsg::Optional<kj::Array<kj::byte>> DiskCache::get(jsg::Lock& js, kj::String key) {
-  if(!hasNoSlashes(key)) {
-    KJ_LOG(ERROR, "DiskCache key contains slashes: ", key);
-    return kj::none;
-  }
-
-  std::ifstream file;
-  kj::String fpath = getDiskCachePath(key);
-
-  file.open(fpath.cStr(), std::ios::binary | std::ios::ate);
-  if (!file.is_open()) {
-    KJ_LOG(INFO, "DiskCache get: file not found: ", fpath);
+  KJ_IF_SOME(root, cacheRoot) {
+    kj::Path path(key);
+    auto file = root->tryOpenFile(path);
+
+    KJ_IF_SOME(f, file) {
+      return f->readAllBytes();
+    } else {
+      return kj::none;
+    }
+  } else {
     return kj::none;
   }
-  std::size_t size = file.tellg();
-  file.seekg(0, std::ios::beg);
-
-  auto data = kj::heapArray<kj::byte>(size);
-  file.read(reinterpret_cast<char*>(data.begin()), size);
-
-  return data;
 }
 
 void DiskCache::put(jsg::Lock& js, kj::String key, kj::Array<kj::byte> data) {
-  if(!hasNoSlashes(key)) {
-    KJ_LOG(ERROR, "DiskCache key contains slashes: ", key);
-    return;
-  }
-  std::ofstream file;
-  kj::String fpath = getDiskCachePath(key);
-  file.open(fpath.cStr(), std::ios::binary | std::ios::trunc);
-
-  if(!file.is_open()) {
-    KJ_LOG(ERROR, "DiskCache put: failed to open file: ", fpath);
+  KJ_IF_SOME(root, cacheRoot) {
+    kj::Path path(key);
+    auto file = root->tryOpenFile(path, kj::WriteMode::CREATE | kj::WriteMode::MODIFY);
+
+    KJ_IF_SOME(f, file) {
+      f->writeAll(data);
+    } else {
+      KJ_LOG(ERROR, "DiskCache: Failed to open file", key);
+    }
+  } else {
     return;
   }
-
-  file.write(reinterpret_cast<const char*>(data.begin()), data.size());
-
-  if(!file.good()) {
-    KJ_LOG(ERROR, "DiskCache put: failed to write to file: ", fpath);
-  }
 }
 
 } // namespace workerd::api::pyodide
diff --git a/src/workerd/api/pyodide/pyodide.h b/src/workerd/api/pyodide/pyodide.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <kj/common.h>
+#include <kj/filesystem.h>
 #include <pyodide/generated/pyodide_extra.capnp.h>
 #include <pyodide/pyodide.capnp.h>
 #include <workerd/jsg/jsg.h>
@@ -169,8 +170,9 @@ class DisabledInternalJaeger : public jsg::Object {
 
 // This cache is used by Pyodide to store wheels fetched over the internet across workerd restarts in local dev only
 class DiskCache: public jsg::Object {
+  kj::Maybe<kj::Own<const kj::Directory>> cacheRoot;
 public:
-  DiskCache() = default;
+  DiskCache(kj::Maybe<kj::Own<const kj::Directory>> cacheRoot): cacheRoot(kj::mv(cacheRoot)) {};
 
   jsg::Optional<kj::Array<kj::byte>> get(jsg::Lock& js, kj::String key);
   void put(jsg::Lock& js, kj::String key, kj::Array<kj::byte> data);
@@ -198,8 +200,6 @@ template <class Registry> void registerPyodideModules(Registry& registry, auto f
     registry.addBuiltinBundle(PYODIDE_BUNDLE, kj::none);
     registry.template addBuiltinModule<PackagesTarReader>(
         "pyodide-internal:packages_tar_reader", workerd::jsg::ModuleRegistry::Type::INTERNAL);
-    registry.template addBuiltinModule<DiskCache>(
-        "pyodide-internal:disk_cache", workerd::jsg::ModuleRegistry::Type::INTERNAL); //TODO: ideally this should only happen in local dev
   }
 }
 

diff --git a/src/workerd/server/workerd-api.c++ b/src/workerd/server/workerd-api.c++
@@ -453,6 +453,24 @@ void WorkerdApi::compileModules(
             },
             jsg::ModuleRegistry::Type::INTERNAL);
       }
+
+      // Inject disk cache module
+      KJ_IF_SOME(pcr, pyodideCacheRoot) {
+        using ModuleInfo = jsg::ModuleRegistry::ModuleInfo;
+        using ObjectModuleInfo = jsg::ModuleRegistry::ObjectModuleInfo;
+        using ResolveMethod = jsg::ModuleRegistry::ResolveMethod;
+        auto specifier = "pyodide-internal:disk_cache";
+        auto diskCache = jsg::alloc<DiskCache>(pcr->clone());
+        modules->addBuiltinModule(
+          specifier,
+          [specifier = kj::str(specifier), diskCache = kj::mv(diskCache)](
+              jsg::Lock& js, ResolveMethod, kj::Maybe<const kj::Path&>&) mutable {
+            auto& wrapper = JsgWorkerdIsolate_TypeWrapper::from(js.v8Isolate);
+            auto wrap = wrapper.wrap(js.v8Context(), kj::none, kj::mv(diskCache));
+            return kj::Maybe(ModuleInfo(js, specifier, kj::none, ObjectModuleInfo(js, wrap)));
+          },
+          jsg::ModuleRegistry::Type::INTERNAL);
+      }
     }
 
     for (auto module: confModules) {

diff --git a/src/workerd/server/workerd-api.h b/src/workerd/server/workerd-api.h
@@ -223,6 +223,7 @@ class WorkerdApi final: public Worker::Api {
 private:
   struct Impl;
   kj::Own<Impl> impl;
+  kj::Maybe<kj::Own<const kj::Directory>> pyodideCacheRoot;
 
   kj::Array<Worker::Script::CompiledGlobal> compileScriptGlobals(
       jsg::Lock& lock,