add memory tracking to workerd::mimetype

cloudflare · Oct 7, 2024 · a2415e0 · a2415e0
1 parent 7619848
commit a2415e0
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 14 deletions.
diff --git a/src/workerd/api/BUILD.bazel b/src/workerd/api/BUILD.bazel
@@ -98,6 +98,7 @@ wd_cc_library(
         "pyodide/pyodide.h",
         "//src/pyodide:generated/pyodide_extra.capnp.h",
     ],
+    implementation_deps = ["//src/workerd/util:string-buffer"],
     visibility = ["//visibility:public"],
     deps = [
         "//src/pyodide",
@@ -125,6 +126,7 @@ wd_cc_library(
     deps = [
         "//src/workerd/jsg",
         "//src/workerd/util",
+        "//src/workerd/util:mimetype",
         "@capnp-cpp//src/kj",
     ],
 )
@@ -154,6 +156,7 @@ wd_cc_library(
     deps = [
         "//src/workerd/jsg:url",
         "//src/workerd/util",
+        "//src/workerd/util:mimetype",
         "@capnp-cpp//src/kj",
     ],
 )

diff --git a/src/workerd/api/node/BUILD.bazel b/src/workerd/api/node/BUILD.bazel
@@ -56,6 +56,7 @@ wd_cc_library(
     deps = [
         "//src/workerd/io:compatibility-date_capnp",
         "//src/workerd/jsg",
+        "//src/workerd/util:mimetype",
     ],
 )
 

diff --git a/src/workerd/api/node/util.h b/src/workerd/api/node/util.h
@@ -88,7 +88,7 @@ class MIMEParams final: public jsg::Object {
 
 class MIMEType final: public jsg::Object {
 public:
-  MIMEType(MimeType inner);
+  explicit MIMEType(MimeType inner);
   ~MIMEType() noexcept(false);
   static jsg::Ref<MIMEType> constructor(kj::String input);
 
@@ -109,12 +109,6 @@ class MIMEType final: public jsg::Object {
     JSG_METHOD_NAMED(toJSON, toString);
   }
 
-  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
-    // TODO(cleanup): Better to have jsg::MimeType be a MemoryRetainer directly
-    tracker.trackFieldWithSize("mimeType", inner.toString().size());
-    tracker.trackField("params", params);
-  }
-
 private:
   workerd::MimeType inner;
   jsg::Ref<MIMEParams> params;

diff --git a/src/workerd/io/BUILD.bazel b/src/workerd/io/BUILD.bazel
@@ -74,6 +74,7 @@ wd_cc_library(
     }),
     implementation_deps = [
         "//src/workerd/util:perfetto",
+        "//src/workerd/util:string-buffer",
         "@capnp-cpp//src/kj/compat:kj-brotli",
         "@capnp-cpp//src/kj/compat:kj-gzip",
         "@nbytes",

diff --git a/src/workerd/tests/BUILD.bazel b/src/workerd/tests/BUILD.bazel
@@ -44,7 +44,7 @@ wd_cc_benchmark(
     name = "bench-mimetype",
     srcs = ["bench-mimetype.c++"],
     deps = [
-        "//src/workerd/util",
+        "//src/workerd/util:mimetype",
     ],
 )
 

diff --git a/src/workerd/util/BUILD.bazel b/src/workerd/util/BUILD.bazel
@@ -29,7 +29,6 @@ wd_cc_library(
 wd_cc_library(
     name = "util",
     srcs = [
-        "mimetype.c++",
         "stream-utils.c++",
         "wait-list.c++",
     ],
@@ -43,10 +42,7 @@ wd_cc_library(
         "color-util.h",
         "duration-exceeded-logger.h",
         "http-util.h",
-        "mimetype.h",
         "stream-utils.h",
-        "string-buffer.h",
-        "strings.h",
         "uncaught-exception-source.h",
         "wait-list.h",
         "weak-refs.h",
@@ -59,6 +55,36 @@ wd_cc_library(
     ],
 )
 
+wd_cc_library(
+    name = "mimetype",
+    srcs = ["mimetype.c++"],
+    hdrs = ["mimetype.h"],
+    implementation_deps = [
+        ":string-buffer",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//src/workerd/jsg:memory-tracker",
+        "@capnp-cpp//src/kj",
+        "@capnp-cpp//src/kj:kj-async",
+    ],
+)
+
+wd_cc_library(
+    name = "strings",
+    srcs = [],
+    hdrs = ["strings.h"],
+    visibility = ["//visibility:public"],
+    deps = [],
+)
+
+wd_cc_library(
+    name = "string-buffer",
+    hdrs = ["string-buffer.h"],
+    visibility = ["//visibility:public"],
+    deps = [":strings"],
+)
+
 wd_cc_library(
     name = "uuid",
     srcs = ["uuid.c++"],
@@ -184,13 +210,25 @@ exports_files(["autogate.h"])
     )
     for f in [
         "batch-queue-test.c++",
-        "mimetype-test.c++",
         "wait-list-test.c++",
         "duration-exceeded-logger-test.c++",
-        "string-buffer-test.c++",
     ]
 ]
 
+kj_test(
+    src = "string-buffer-test.c++",
+    deps = [
+        ":string-buffer",
+    ],
+)
+
+kj_test(
+    src = "mimetype-test.c++",
+    deps = [
+        ":mimetype",
+    ],
+)
+
 kj_test(
     src = "sqlite-test.c++",
     deps = [

diff --git a/src/workerd/util/mimetype.h b/src/workerd/util/mimetype.h
@@ -3,6 +3,8 @@
 //     https://opensource.org/licenses/Apache-2.0
 #pragma once
 
+#include <workerd/jsg/memory.h>
+
 #include <kj/common.h>
 #include <kj/exception.h>
 #include <kj/map.h>
@@ -102,6 +104,12 @@ class MimeType final {
   // https://fetch.spec.whatwg.org/#concept-header-extract-mime-type
   static kj::Maybe<MimeType> extract(kj::StringPtr input);
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackFieldWithSize("type", type_.size());
+    tracker.trackFieldWithSize("subtype", subtype_.size());
+    tracker.trackFieldWithSize("params", params_.size());
+  }
+
 private:
   kj::String type_;
   kj::String subtype_;