fix: create result String/Buffer in native code

Refs: nodejs/node#54505 Refs: nodejs/node#53804
nxtedition · Aug 26, 2024 · 04e9ba4 · 04e9ba4
1 parent de668e1
commit 04e9ba4
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 45 deletions.
diff --git a/binding.cc b/binding.cc
@@ -968,6 +968,9 @@ NAPI_METHOD(db_get_many) {
   bool takeSnapshot = true;
   NAPI_STATUS_THROWS(GetProperty(env, options, "snapshot", takeSnapshot));
 
+  Encoding valueEncoding = Encoding::String;
+  NAPI_STATUS_THROWS(GetProperty(env, options, "valueEncoding", valueEncoding));
+
   auto callback = argv[3];
 
   std::shared_ptr<const rocksdb::Snapshot> snapshot;
@@ -1036,22 +1039,73 @@ NAPI_METHOD(db_get_many) {
         return rocksdb::Status::OK();
       },
       [=](auto& state, auto env, auto& argv) {
-        argv.resize(3);
+        argv.resize(2);
 
-        if (state.sizes.size() > 0) {
-          auto sizes = std::make_unique<std::vector<int32_t>>(std::move(state.sizes));
-          NAPI_STATUS_RETURN(napi_create_external_buffer(env, sizes->size() * 4, sizes->data(), Finalize<std::vector<int32_t>>, sizes.get(), &argv[1]));
-          sizes.release();
-        } else {
-          NAPI_STATUS_RETURN(napi_get_undefined(env, &argv[1]));
-        }
+        NAPI_STATUS_RETURN(napi_create_array_with_length(env, state.sizes.size(), &argv[1]));
+
+        if (valueEncoding == Encoding::Buffer) {
+          napi_value arraybuffer;
+          {
+            // We create one ArrayBuffer and then slice it into Buffer(s) to avoid the
+            // overhead of registering lots of finalizers which is super slow.
+            // https://github.com/nodejs/node/issues/53804
+            auto data = std::make_unique<std::vector<uint8_t>>(std::move(state.data));
+            NAPI_STATUS_RETURN(napi_create_external_arraybuffer(
+              env,
+              data->data(),
+              data->size(),
+              Finalize<std::vector<uint8_t>>,
+              data.get(),
+              &arraybuffer
+            ));
+            data.release();
+          }
 
-        if (state.data.size() > 0) {
-          auto data = std::make_unique<std::vector<uint8_t>>(std::move(state.data));
-          NAPI_STATUS_RETURN(napi_create_external_buffer(env, data->size(), data->data(), Finalize<std::vector<uint8_t>>, data.get(), &argv[2]));
-          data.release();
+          auto offset = 0;
+          for (auto n = 0; n < count; n++) {
+            const auto size = state.sizes[n];
+
+            // TODO (fix): We must create a Node::Buffere here...
+            // https://github.com/nodejs/node/pull/54505
+            napi_value row;
+            NAPI_STATUS_RETURN(napi_create_typedarray(
+              env,
+              napi_uint8_array,
+              size,
+              arraybuffer,
+              offset,
+              &row
+            ));
+
+            NAPI_STATUS_RETURN(napi_set_element(env, argv[1], n, row));
+
+            offset += size;
+            if (offset & 0x7) {
+              offset |= 0x7;
+              offset += 1;
+            }
+          }
         } else {
-          NAPI_STATUS_RETURN(napi_get_undefined(env, &argv[2]));
+          auto offset = 0;
+          auto ptr = reinterpret_cast<const char*>(state.data.data());
+          for (auto n = 0; n < count; n++) {
+            const auto size = state.sizes[n];
+
+            napi_value row;
+            NAPI_STATUS_RETURN(napi_create_string_utf8(
+              env,
+              ptr + offset,
+              size,
+              &row
+            ));
+            NAPI_STATUS_RETURN(napi_set_element(env, argv[1], n, row));
+
+            offset += size;
+            if (offset & 0x7) {
+              offset |= 0x7;
+              offset += 1;
+            }
+          }
         }
 
         return napi_ok;

diff --git a/index.js b/index.js
@@ -155,40 +155,10 @@ class RocksLevel extends AbstractLevel {
   _getMany (keys, options, callback) {
     callback = fromCallback(callback, kPromise)
 
-    const { valueEncoding } = options ?? EMPTY
     try {
       this[kRef]()
-      binding.db_get_many(this[kContext], keys, options ?? EMPTY, (err, sizes, data) => {
-        if (err) {
-          callback(err)
-        } else {
-          data ??= Buffer.alloc(0)
-          sizes ??= Buffer.alloc(0)
-
-          const rows = []
-          let offset = 0
-          const sizes32 = new Int32Array(sizes.buffer, sizes.byteOffset, sizes.byteLength / 4)
-          for (let n = 0; n < sizes32.length; n++) {
-            const size = sizes32[n]
-            const encoding = valueEncoding
-            if (size < 0) {
-              rows.push(undefined)
-            } else {
-              if (!encoding || encoding === 'buffer') {
-                rows.push(data.subarray(offset, offset + size))
-              } else if (encoding === 'slice') {
-                rows.push({ buffer: data, byteOffset: offset, byteLength: size })
-              } else {
-                rows.push(data.toString(encoding, offset, offset + size))
-              }
-              offset += size
-              if (offset & 0x7) {
-                offset = (offset | 0x7) + 1
-              }
-            }
-          }
-          callback(null, rows)
-        }
+      binding.db_get_many(this[kContext], keys, options ?? EMPTY, (err, rows) => {
+        callback(err, rows)
         this[kUnref]()
       })
     } catch (err) {