From d09cbdfae9f0fe8b2c11828170b2bf66018352b6 Mon Sep 17 00:00:00 2001 From: Alistair Smith Date: Fri, 9 May 2025 21:25:40 -0700 Subject: [PATCH] Add a RedisClient#getBuffer method (#19567) --- .vscode/settings.json | 3 + packages/bun-types/redis.d.ts | 9 +- src/bun.js/api/valkey.classes.ts | 4 + src/valkey/ValkeyCommand.zig | 9 +- src/valkey/js_valkey.zig | 1 + src/valkey/js_valkey_functions.zig | 20 +++ src/valkey/valkey_protocol.zig | 39 ++++-- test/js/node/no-addons.test.ts | 2 +- test/js/valkey/unit/buffer-operations.test.ts | 115 ++++++++++++++++++ 9 files changed, 188 insertions(+), 14 deletions(-) create mode 100644 test/js/valkey/unit/buffer-operations.test.ts diff --git a/.vscode/settings.json b/.vscode/settings.json index 244a941e4e..167a601132 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -43,6 +43,9 @@ "editor.tabSize": 4, "editor.useTabStops": false, "editor.defaultFormatter": "ziglang.vscode-zig", + "editor.codeActionsOnSave": { + "source.organizeImports": "never", + }, }, // lldb diff --git a/packages/bun-types/redis.d.ts b/packages/bun-types/redis.d.ts index c4631f5c1a..dccc6f5852 100644 --- a/packages/bun-types/redis.d.ts +++ b/packages/bun-types/redis.d.ts @@ -110,10 +110,17 @@ declare module "bun" { /** * Get the value of a key * @param key The key to get - * @returns Promise that resolves with the key's value, or null if the key doesn't exist + * @returns Promise that resolves with the key's value as a string, or null if the key doesn't exist */ get(key: string | ArrayBufferView | Blob): Promise; + /** + * Get the value of a key as a Uint8Array + * @param key The key to get + * @returns Promise that resolves with the key's value as a Uint8Array, or null if the key doesn't exist + */ + getBuffer(key: string | ArrayBufferView | Blob): Promise | null>; + /** * Set key to hold the string value * @param key The key to set diff --git a/src/bun.js/api/valkey.classes.ts b/src/bun.js/api/valkey.classes.ts index c035bab32f..9a7af095ff 100644 --- a/src/bun.js/api/valkey.classes.ts +++ b/src/bun.js/api/valkey.classes.ts @@ -31,6 +31,10 @@ export default [ fn: "get", length: 1, }, + getBuffer: { + fn: "getBuffer", + length: 1, + }, set: { fn: "set", length: 2, diff --git a/src/valkey/ValkeyCommand.zig b/src/valkey/ValkeyCommand.zig index 13a3df96a6..47d8885071 100644 --- a/src/valkey/ValkeyCommand.zig +++ b/src/valkey/ValkeyCommand.zig @@ -81,7 +81,8 @@ pub fn deinit(_: *Command) void { pub const Meta = packed struct(u8) { return_as_bool: bool = false, supports_auto_pipelining: bool = true, - _padding: u6 = 0, + return_as_buffer: bool = false, + _padding: u5 = 0, const not_allowed_autopipeline_commands = bun.ComptimeStringMap(void, .{ .{"AUTH"}, @@ -123,7 +124,11 @@ pub const Promise = struct { } pub fn resolve(self: *Promise, globalObject: *JSC.JSGlobalObject, value: *protocol.RESPValue) void { - const js_value = value.toJS(globalObject) catch |err| { + const options = protocol.RESPValue.ToJSOptions{ + .return_as_buffer = self.meta.return_as_buffer, + }; + + const js_value = value.toJSWithOptions(globalObject, options) catch |err| { self.reject(globalObject, globalObject.takeError(err)); return; }; diff --git a/src/valkey/js_valkey.zig b/src/valkey/js_valkey.zig index 135cc2eb97..6193b6bf47 100644 --- a/src/valkey/js_valkey.zig +++ b/src/valkey/js_valkey.zig @@ -648,6 +648,7 @@ pub const JSValkeyClient = struct { pub const expire = fns.expire; pub const expiretime = fns.expiretime; pub const get = fns.get; + pub const getBuffer = fns.getBuffer; pub const getdel = fns.getdel; pub const getex = fns.getex; pub const getset = fns.getset; diff --git a/src/valkey/js_valkey_functions.zig b/src/valkey/js_valkey_functions.zig index d3f11768d6..7cc629724a 100644 --- a/src/valkey/js_valkey_functions.zig +++ b/src/valkey/js_valkey_functions.zig @@ -60,6 +60,26 @@ pub fn get(this: *JSValkeyClient, globalObject: *JSC.JSGlobalObject, callframe: return promise.asValue(globalObject); } +pub fn getBuffer(this: *JSValkeyClient, globalObject: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) bun.JSError!JSValue { + const key = (try fromJS(globalObject, callframe.argument(0))) orelse { + return globalObject.throwInvalidArgumentType("getBuffer", "key", "string or buffer"); + }; + defer key.deinit(); + + const promise = this.send( + globalObject, + callframe.this(), + &.{ + .command = "GET", + .args = .{ .args = &.{key} }, + .meta = .{ .return_as_buffer = true }, + }, + ) catch |err| { + return protocol.valkeyErrorToJS(globalObject, "Failed to send GET command", err); + }; + return promise.asValue(globalObject); +} + pub fn set(this: *JSValkeyClient, globalObject: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) bun.JSError!JSValue { const args_view = callframe.arguments(); var stack_fallback = std.heap.stackFallback(512, bun.default_allocator); diff --git a/src/valkey/valkey_protocol.zig b/src/valkey/valkey_protocol.zig index e9b67e8120..528255356a 100644 --- a/src/valkey/valkey_protocol.zig +++ b/src/valkey/valkey_protocol.zig @@ -1,7 +1,9 @@ const std = @import("std"); + const bun = @import("bun"); const JSC = bun.JSC; const String = bun.String; + const debug = bun.Output.scoped(.Redis, false); pub const RedisError = error{ @@ -244,15 +246,32 @@ pub const RESPValue = union(RESPType) { } } - // Convert RESPValue to JSValue pub fn toJS(self: *RESPValue, globalObject: *JSC.JSGlobalObject) bun.JSError!JSC.JSValue { + return self.toJSWithOptions(globalObject, .{}); + } + + pub const ToJSOptions = struct { + return_as_buffer: bool = false, + }; + + fn valkeyStrToJSValue(globalObject: *JSC.JSGlobalObject, str: []const u8, options: *const ToJSOptions) bun.JSError!JSC.JSValue { + if (options.return_as_buffer) { + // TODO: handle values > 4.7 GB + const buf = JSC.ArrayBuffer.createBuffer(globalObject, str); + return buf.toJS(globalObject); + } else { + return bun.String.createUTF8ForJS(globalObject, str); + } + } + + pub fn toJSWithOptions(self: *RESPValue, globalObject: *JSC.JSGlobalObject, options: ToJSOptions) bun.JSError!JSC.JSValue { switch (self.*) { - .SimpleString => |str| return bun.String.createUTF8ForJS(globalObject, str), + .SimpleString => |str| return valkeyStrToJSValue(globalObject, str, &options), .Error => |str| return valkeyErrorToJS(globalObject, str, RedisError.InvalidResponse), .Integer => |int| return JSC.JSValue.jsNumber(int), .BulkString => |maybe_str| { if (maybe_str) |str| { - return bun.String.createUTF8ForJS(globalObject, str); + return valkeyStrToJSValue(globalObject, str, &options); } else { return JSC.JSValue.jsNull(); } @@ -260,7 +279,7 @@ pub const RESPValue = union(RESPType) { .Array => |array| { var js_array = JSC.JSValue.createEmptyArray(globalObject, array.len); for (array, 0..) |*item, i| { - const js_item = try item.toJS(globalObject); + const js_item = try item.toJSWithOptions(globalObject, options); js_array.putIndex(globalObject, @intCast(i), js_item); } return js_array; @@ -269,14 +288,14 @@ pub const RESPValue = union(RESPType) { .Double => |d| return JSC.JSValue.jsNumber(d), .Boolean => |b| return JSC.JSValue.jsBoolean(b), .BlobError => |str| return valkeyErrorToJS(globalObject, str, RedisError.InvalidBlobError), - .VerbatimString => |verbatim| return bun.String.createUTF8ForJS(globalObject, verbatim.content), + .VerbatimString => |verbatim| return valkeyStrToJSValue(globalObject, verbatim.content, &options), .Map => |entries| { var js_obj = JSC.JSValue.createEmptyObjectWithNullPrototype(globalObject); for (entries) |*entry| { - const js_key = try entry.key.toJS(globalObject); + const js_key = try entry.key.toJSWithOptions(globalObject, .{}); var key_str = try js_key.toBunString(globalObject); defer key_str.deref(); - const js_value = try entry.value.toJS(globalObject); + const js_value = try entry.value.toJSWithOptions(globalObject, options); js_obj.putMayBeIndex(globalObject, &key_str, js_value); } @@ -285,7 +304,7 @@ pub const RESPValue = union(RESPType) { .Set => |set| { var js_array = JSC.JSValue.createEmptyArray(globalObject, set.len); for (set, 0..) |*item, i| { - const js_item = try item.toJS(globalObject); + const js_item = try item.toJSWithOptions(globalObject, options); js_array.putIndex(globalObject, @intCast(i), js_item); } return js_array; @@ -293,7 +312,7 @@ pub const RESPValue = union(RESPType) { .Attribute => |attribute| { // For now, we just return the value and ignore attributes // In the future, we could attach the attributes as a hidden property - return try attribute.value.toJS(globalObject); + return try attribute.value.toJSWithOptions(globalObject, options); }, .Push => |push| { var js_obj = JSC.JSValue.createEmptyObjectWithNullPrototype(globalObject); @@ -305,7 +324,7 @@ pub const RESPValue = union(RESPType) { // Add the data as an array var data_array = JSC.JSValue.createEmptyArray(globalObject, push.data.len); for (push.data, 0..) |*item, i| { - const js_item = try item.toJS(globalObject); + const js_item = try item.toJSWithOptions(globalObject, options); data_array.putIndex(globalObject, @intCast(i), js_item); } js_obj.put(globalObject, "data", data_array); diff --git a/test/js/node/no-addons.test.ts b/test/js/node/no-addons.test.ts index f507e7c18f..0df8cf5155 100644 --- a/test/js/node/no-addons.test.ts +++ b/test/js/node/no-addons.test.ts @@ -1,5 +1,5 @@ -import { test, expect } from "bun:test"; import { spawnSync } from "bun"; +import { expect, test } from "bun:test"; import { bunExe, bunEnv as env } from "harness"; test("--no-addons throws an error on process.dlopen", () => { diff --git a/test/js/valkey/unit/buffer-operations.test.ts b/test/js/valkey/unit/buffer-operations.test.ts new file mode 100644 index 0000000000..d87f2de128 --- /dev/null +++ b/test/js/valkey/unit/buffer-operations.test.ts @@ -0,0 +1,115 @@ +import { beforeEach, describe, expect, test } from "bun:test"; +import { ConnectionType, createClient, ctx, isEnabled } from "../test-utils"; + +describe.skipIf(!isEnabled)("Valkey: Buffer Operations", () => { + beforeEach(() => { + if (ctx.redis?.connected) { + ctx.redis.close?.(); + } + ctx.redis = createClient(ConnectionType.TCP); + }); + + test("getBuffer returns binary data as Uint8Array", async () => { + const key = ctx.generateKey("buffer-test"); + + const binaryData = new Uint8Array([0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x20, 0x77, 0x6f, 0x72, 0x6c, 0x64]); + await ctx.redis.set(key, binaryData); + + const asString = await ctx.redis.get(key); + const asBuffer = await ctx.redis.getBuffer(key); + + expectAssert(asString); + expectAssert(asBuffer); + + expect(asBuffer.buffer).toBeInstanceOf(ArrayBuffer); + expect(asBuffer).toBeInstanceOf(Uint8Array); + expect(asBuffer.length).toBe(binaryData.length); + expect(asBuffer).toStrictEqual(binaryData); + + for (let i = 0; i < binaryData.length; i++) { + expect(asBuffer[i]).toBe(binaryData[i]); + } + + const stringBuffer = Buffer.from(asString); + expect(stringBuffer.length).toBe(binaryData.length); + }); + + test("getBuffer for non-existent key returns null", async () => { + const key = ctx.generateKey("non-existent"); + const result = await ctx.redis.getBuffer(key); + expect(result).toBeNull(); + }); + + test("Really long buffer", async () => { + const key = ctx.generateKey("long-buffer"); + const binaryData = new Uint8Array(1000000); + await ctx.redis.set(key, binaryData); + const result = await ctx.redis.getBuffer(key); + expect(result).toBeInstanceOf(Uint8Array); + }); + + test("Buffer with no bytes", async () => { + const key = ctx.generateKey("empty-buffer"); + const binaryData = new Uint8Array(0); + await ctx.redis.set(key, binaryData); + const result = await ctx.redis.getBuffer(key); + expectAssert(result); + expect(result).toBeInstanceOf(Uint8Array); + expect(result.length).toBe(0); + }); + + test("Buffer with null bytes", async () => { + const key = ctx.generateKey("null-bytes"); + const binaryData = new Uint8Array([0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09]); + await ctx.redis.set(key, binaryData); + const result = await ctx.redis.getBuffer(key); + expectAssert(result); + expect(result).toBeInstanceOf(Uint8Array); + expect(result.length).toBe(binaryData.length); + for (let i = 0; i < binaryData.length; i++) { + expect(result[i]).toBe(binaryData[i]); + } + }); + + test("concurrent getBuffer against large blob", async () => { + const key = ctx.generateKey("concurrent"); + const big = new Uint8Array(500_000).map((_, i) => i % 256); + await ctx.redis.set(key, big); + const readers = Array.from({ length: 20 }, () => ctx.redis.getBuffer(key)); + const results = await Promise.all(readers); + for (const r of results) expect(r).toStrictEqual(big); + }); + + test("set and getBuffer with ArrayBufferView key", async () => { + const keyBytes = new Uint8Array([0x6b, 0x65, 0x79, 0x21]); // "key!" + const value = new Uint8Array([0x01, 0x02, 0x03]); + await ctx.redis.set(keyBytes, value); + const out = await ctx.redis.getBuffer(keyBytes); + expect(out).toBeInstanceOf(Uint8Array); + expect(out).toStrictEqual(value); + }); + + test("set and getBuffer with ArrayBuffer key", async () => { + const keyBuffer = new Uint8Array([0x62, 0x75, 0x6e, 0x21]).buffer; // "bun!" + expect(keyBuffer).toBeInstanceOf(ArrayBuffer); + const value = new Uint8Array([0x0a, 0x0b]); + await ctx.redis.set(keyBuffer, value); + const out = await ctx.redis.getBuffer(keyBuffer); + expect(out).toBeInstanceOf(Uint8Array); + expect(out).toStrictEqual(value); + }); + + test("set and getBuffer with Blob key", async () => { + const keyBytes = new Uint8Array([0x74, 0x65, 0x73, 0x74]); // "test" + const keyBlob = new Blob([keyBytes]); + const value = new Uint8Array([0xff, 0xee, 0xdd]); + await ctx.redis.set(keyBlob, value); + const out = await ctx.redis.getBuffer(keyBlob); + expect(out).toBeInstanceOf(Uint8Array); + expect(out).toStrictEqual(value); + }); +}); + +function expectAssert(value: unknown): asserts value { + expect(value).toBeTruthy(); +}