ValkeyKey storage (#116)

adam-fowler · web-flow · commit 0e3e17ff3abd · 2025-06-26T19:48:38.000+01:00
diff --git a/Benchmarks/ValkeyBenchmarks/ValkeyBenchmarks.swift b/Benchmarks/ValkeyBenchmarks/ValkeyBenchmarks.swift
@@ -104,7 +104,7 @@ let benchmarks: @Sendable () -> Void = {
         }
 
         Benchmark("ValkeyCommandEncoder – Simple MGET 15 keys", configuration: .init(metrics: defaultMetrics, scalingFactor: .kilo)) { benchmark in
-            let keys = (0..<15).map { ValkeyKey(rawValue: "foo-\($0)") }
+            let keys = (0..<15).map { ValkeyKey("foo-\($0)") }
             let command = MGET(key: keys)
             benchmark.startMeasurement()
 
@@ -135,8 +135,7 @@ let benchmarks: @Sendable () -> Void = {
         }
 
         Benchmark("HashSlot – {user}.whatever", configuration: .init(metrics: defaultMetrics, scalingFactor: .mega)) { benchmark in
-            let key = "{user}.whatever"
-
+            let key: ValkeyKey = "{user}.whatever"
             benchmark.startMeasurement()
             for _ in benchmark.scaledIterations {
                 blackHole(HashSlot(key: key))
diff --git a/Sources/Valkey/Cluster/HashSlot.swift b/Sources/Valkey/Cluster/HashSlot.swift
@@ -12,6 +12,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import NIOCore
+
 // This is a derived version from:
 // https://github.com/swift-server/RediStack/blob/2df32390e2366b58cc15c2612bb324b3fc37a190/Sources/RediStack/Cluster/RedisHashSlot.swift
 
@@ -130,7 +132,8 @@ extension HashSlot {
     ///
     /// - Parameter key: The key used in a Valkey command
     /// - Returns: A HashSlot representing where this key would be stored in the cluster
-    public init(key: String) {
+    @inlinable
+    public init(key: some BidirectionalCollection<UInt8>) {
         // Banging is safe because the modulo ensures we are in range
         self.init(rawValue: UInt16(HashSlot.crc16(HashSlot.hashTag(forKey: key)) % 16384))!
     }
@@ -139,8 +142,16 @@ extension HashSlot {
     ///
     /// - Parameter key: The Valkey key for which to calculate the hash slot
     /// - Returns: A HashSlot representing where this key would be stored in the cluster
+    @inlinable
     public init(key: ValkeyKey) {
-        self.init(key: key.rawValue)
+        switch key._storage {
+        case .string(let string):
+            self.init(key: string.utf8)
+        case .buffer(let buffer):
+            self = buffer.withUnsafeReadableBytes { bytes in
+                HashSlot(key: bytes)
+            }
+        }
     }
 
     /// Computes the portion of the key that should be used for hash slot calculation.
@@ -150,37 +161,36 @@ extension HashSlot {
     /// - If the pattern is empty "{}", or doesn't exist, the entire key is used
     /// - Only the first occurrence of "{...}" is considered
     ///
-    /// - Parameter key: The key for your operation
+    /// - Parameter keyUTF8View: The UTF8 view of key for your operation
     /// - Returns: A substring UTF8 view that will be used in the CRC16 computation
-    package static func hashTag(forKey key: String) -> Substring.UTF8View {
-        let utf8View = key.utf8
-
-        var firstOpenCurly: String.UTF8View.Index?
-        var index = utf8View.startIndex
+    @inlinable
+    package static func hashTag<Bytes: BidirectionalCollection<UInt8>>(forKey keyUTF8View: Bytes) -> Bytes.SubSequence {
+        var firstOpenCurly: Bytes.Index?
+        var index = keyUTF8View.startIndex
 
-        while index < utf8View.endIndex {
-            defer { index = utf8View.index(after: index) }
+        while index < keyUTF8View.endIndex {
+            defer { index = keyUTF8View.index(after: index) }
 
-            switch utf8View[index] {
+            switch keyUTF8View[index] {
             case UInt8(ascii: "{") where firstOpenCurly == nil:
                 firstOpenCurly = index
             case UInt8(ascii: "}"):
                 guard let firstOpenCurly = firstOpenCurly else {
                     continue
                 }
 
-                if firstOpenCurly == utf8View.index(before: index) {
+                if firstOpenCurly == keyUTF8View.index(before: index) {
                     // we had a `{}` combination... this means the complete key shall be used for hashing
-                    return utf8View[...]
+                    return keyUTF8View[...]
                 }
 
-                return utf8View[(utf8View.index(after: firstOpenCurly))..<index]
+                return keyUTF8View[(keyUTF8View.index(after: firstOpenCurly))..<index]
             default:
                 continue
             }
         }
 
-        return utf8View[...]
+        return keyUTF8View[...]
     }
 }
 
@@ -227,7 +237,8 @@ extension HashSlot {
  * Output for "123456789"     : 31C3
  */
 
-private let crc16tab: [UInt16] = [
+@usableFromInline
+/* private */ let crc16tab: [UInt16] = [
     0x0000, 0x1021, 0x2042, 0x3063, 0x4084, 0x50a5, 0x60c6, 0x70e7,
     0x8108, 0x9129, 0xa14a, 0xb16b, 0xc18c, 0xd1ad, 0xe1ce, 0xf1ef,
     0x1231, 0x0210, 0x3273, 0x2252, 0x52b5, 0x4294, 0x72f7, 0x62d6,
@@ -269,6 +280,7 @@ extension HashSlot {
     ///
     /// - Parameter bytes: A sequence of bytes to compute the CRC16 value for
     /// - Returns: The computed CRC16 value
+    @inlinable
     package static func crc16<Bytes: Sequence>(_ bytes: Bytes) -> UInt16 where Bytes.Element == UInt8 {
         var crc: UInt16 = 0
         for byte in bytes {
diff --git a/Sources/Valkey/Subscriptions/ValkeyConnection+subscribe.swift b/Sources/Valkey/Subscriptions/ValkeyConnection+subscribe.swift
@@ -192,7 +192,7 @@ extension ValkeyConnection {
         process: (AsyncMapSequence<ValkeySubscription, ValkeyKey>) async throws -> sending Value
     ) async throws -> sending Value {
         try await self.subscribe(to: [ValkeySubscriptions.invalidateChannel]) { subscription in
-            let keys = subscription.map { ValkeyKey(rawValue: $0.message) }
+            let keys = subscription.map { ValkeyKey($0.message) }
             return try await process(keys)
         }
     }
diff --git a/Sources/Valkey/Subscriptions/ValkeySubscriptions.swift b/Sources/Valkey/Subscriptions/ValkeySubscriptions.swift
@@ -87,7 +87,7 @@ struct ValkeySubscriptions {
             case .forwardMessage(let subscriptions):
                 for subscription in subscriptions {
                     for key in keys {
-                        subscription.sendMessage(.init(channel: Self.invalidateChannel, message: key.rawValue))
+                        subscription.sendMessage(.init(channel: Self.invalidateChannel, message: String(valkeyKey: key)))
                     }
                 }
             case .doNothing, .none:
diff --git a/Sources/Valkey/ValkeyKey.swift b/Sources/Valkey/ValkeyKey.swift
@@ -15,43 +15,117 @@
 import NIOCore
 
 /// Type representing a Valkey Key
-public struct ValkeyKey: RawRepresentable, Sendable, Equatable, Hashable {
-    public var rawValue: String
+public struct ValkeyKey: Sendable, Equatable, Hashable {
+    @usableFromInline
+    enum _Storage: Sendable {
+        case string(String)
+        case buffer(ByteBuffer)
+    }
+    @usableFromInline
+    let _storage: _Storage
+
+    /// Initialize ValkeyKey with String
+    /// - Parameter string: string
+    @inlinable
+    public init(_ string: String) {
+        self._storage = .string(string)
+    }
+
+    /// Initialize ValkeyKey with ByteBuffer
+    /// - Parameter buffer: ByteBuffer
+    @inlinable
+    public init(_ buffer: ByteBuffer) {
+        self._storage = .buffer(buffer)
+    }
+
+    @inlinable
+    static public func == (_ lhs: Self, _ rhs: Self) -> Bool {
+        switch (lhs._storage, rhs._storage) {
+        case (.string(let lhs), .string(let rhs)):
+            lhs == rhs
+        case (.buffer(let lhs), .buffer(let rhs)):
+            lhs == rhs
+        case (.string(let lhs), .buffer(let rhs)):
+            lhs.utf8.elementsEqual(rhs.readableBytesView)
+        case (.buffer(let lhs), .string(let rhs)):
+            rhs.utf8.elementsEqual(lhs.readableBytesView)
+        }
+    }
 
-    public init(rawValue: String) {
-        self.rawValue = rawValue
+    @inlinable
+    public func hash(into hasher: inout Hasher) {
+        switch self._storage {
+        case .string(let string): string.hash(into: &hasher)
+        case .buffer(let buffer): buffer.hash(into: &hasher)
+        }
     }
 }
 
 extension ValkeyKey: RESPTokenDecodable {
+    @inlinable
     public init(fromRESP token: RESPToken) throws {
         switch token.value {
         case .simpleString(let buffer), .bulkString(let buffer):
-            self.rawValue = String(buffer: buffer)
+            self._storage = .buffer(buffer)
         default:
             throw RESPParsingError(code: .unexpectedType, buffer: token.base)
         }
     }
 }
 
-extension ValkeyKey: CustomStringConvertible {
-    public var description: String { rawValue.description }
-}
-
 extension ValkeyKey: RESPRenderable {
-
     @inlinable
     public var respEntries: Int { 1 }
 
     @inlinable
     public func encode(into commandEncoder: inout ValkeyCommandEncoder) {
-        self.rawValue.encode(into: &commandEncoder)
+        switch self._storage {
+        case .string(let string):
+            string.encode(into: &commandEncoder)
+        case .buffer(let buffer):
+            buffer.encode(into: &commandEncoder)
+        }
+    }
+}
+
+extension ValkeyKey: RESPStringRenderable {}
+
+extension ValkeyKey: CustomStringConvertible {
+    public var description: String {
+        switch self._storage {
+        case .string(let string): string
+        case .buffer(let buffer): String(buffer: buffer)
+        }
     }
 }
 
 extension ValkeyKey: ExpressibleByStringLiteral {
     @inlinable
     public init(stringLiteral string: String) {
-        self.init(rawValue: string)
+        self.init(string)
+    }
+}
+
+extension String {
+    ///  Initialize String from ValkeyKey
+    /// - Parameter valkeyKey: key
+    @inlinable
+    public init(valkeyKey: ValkeyKey) {
+        switch valkeyKey._storage {
+        case .string(let string): self = string
+        case .buffer(let buffer): self = String(buffer: buffer)
+        }
+    }
+}
+
+extension ByteBuffer {
+    ///  Initialize ByteBuffer from ValkeyKey
+    /// - Parameter valkeyKey: key
+    @inlinable
+    public init(valkeyKey: ValkeyKey) {
+        switch valkeyKey._storage {
+        case .string(let string): self = ByteBuffer(string: string)
+        case .buffer(let buffer): self = buffer
+        }
     }
 }
diff --git a/Tests/ClusterIntegrationTests/ClusterIntegrationTests.swift b/Tests/ClusterIntegrationTests/ClusterIntegrationTests.swift
@@ -41,7 +41,7 @@ struct ClusterIntegrationTests {
         connection: some ValkeyConnectionProtocol,
         _ operation: (ValkeyKey) async throws -> Value
     ) async throws -> Value {
-        let key = ValkeyKey(rawValue: UUID().uuidString)
+        let key = ValkeyKey(UUID().uuidString)
         let result: Result<Value, any Error>
         do {
             result = try await .success(operation(key))
diff --git a/Tests/IntegrationTests/ValkeyTests.swift b/Tests/IntegrationTests/ValkeyTests.swift
@@ -23,7 +23,7 @@ import Valkey
 struct GeneratedCommands {
     let valkeyHostname = ProcessInfo.processInfo.environment["VALKEY_HOSTNAME"] ?? "localhost"
     func withKey<Value>(connection: some ValkeyConnectionProtocol, _ operation: (ValkeyKey) async throws -> Value) async throws -> Value {
-        let key = ValkeyKey(rawValue: UUID().uuidString)
+        let key = ValkeyKey(UUID().uuidString)
         let value: Value
         do {
             value = try await operation(key)
@@ -338,9 +338,9 @@ struct GeneratedCommands {
                 for _ in 0..<100 {
                     group.addTask {
                         try await withKey(connection: connection) { key in
-                            _ = try await connection.set(key: key, value: key.rawValue)
-                            let response = try await connection.get(key: key).map { String(buffer: $0) }
-                            #expect(response == key.rawValue)
+                            _ = try await connection.set(key: key, value: key)
+                            let response = try await connection.get(key: key).map { ValkeyKey($0) }
+                            #expect(response == key)
                         }
                     }
                 }
diff --git a/Tests/ValkeyTests/Cluster/HashSlotTests.swift b/Tests/ValkeyTests/Cluster/HashSlotTests.swift
@@ -30,6 +30,7 @@
 //===----------------------------------------------------------------------===//
 
 import Foundation
+import NIOCore
 import Testing
 import Valkey
 
@@ -88,20 +89,27 @@ struct HashSlotTests {
                 HashSlot.hashTag(forKey: "{user1000}.followers")
             )
         )
-        #expect(HashSlot.hashTag(forKey: "{user1000}.following").elementsEqual("user1000".utf8))
-        #expect(HashSlot.hashTag(forKey: "{user1000}.followers").elementsEqual("user1000".utf8))
-
-        #expect(HashSlot.hashTag(forKey: "foo{}{bar}").elementsEqual("foo{}{bar}".utf8))
-        #expect(HashSlot.hashTag(forKey: "foo{{bar}}zap").elementsEqual("{bar".utf8))
-        #expect(HashSlot.hashTag(forKey: "foo{bar}{zap}").elementsEqual("bar".utf8))
-        #expect(HashSlot.hashTag(forKey: "{}foo{bar}{zap}").elementsEqual("{}foo{bar}{zap}".utf8))
-        #expect(HashSlot.hashTag(forKey: "foo").elementsEqual("foo".utf8))
-        #expect(HashSlot.hashTag(forKey: "foo}").elementsEqual("foo}".utf8))
-        #expect(HashSlot.hashTag(forKey: "{foo}").elementsEqual("foo".utf8))
-        #expect(HashSlot.hashTag(forKey: "bar{foo}").elementsEqual("foo".utf8))
-        #expect(HashSlot.hashTag(forKey: "bar{}").elementsEqual("bar{}".utf8))
-        #expect(HashSlot.hashTag(forKey: "{}").elementsEqual("{}".utf8))
-        #expect(HashSlot.hashTag(forKey: "{}bar").elementsEqual("{}bar".utf8))
+        #expect(HashSlot.hashTag(forKey: "{user1000}.following").elementsEqual("user1000"))
+        #expect(HashSlot.hashTag(forKey: "{user1000}.followers").elementsEqual("user1000"))
+
+        #expect(HashSlot.hashTag(forKey: "foo{}{bar}").elementsEqual("foo{}{bar}"))
+        #expect(HashSlot.hashTag(forKey: "foo{{bar}}zap").elementsEqual("{bar"))
+        #expect(HashSlot.hashTag(forKey: "foo{bar}{zap}").elementsEqual("bar"))
+        #expect(HashSlot.hashTag(forKey: "{}foo{bar}{zap}").elementsEqual("{}foo{bar}{zap}"))
+        #expect(HashSlot.hashTag(forKey: "foo").elementsEqual("foo"))
+        #expect(HashSlot.hashTag(forKey: "foo}").elementsEqual("foo}"))
+        #expect(HashSlot.hashTag(forKey: "{foo}").elementsEqual("foo"))
+        #expect(HashSlot.hashTag(forKey: "bar{foo}").elementsEqual("foo"))
+        #expect(HashSlot.hashTag(forKey: "bar{}").elementsEqual("bar{}"))
+        #expect(HashSlot.hashTag(forKey: "{}").elementsEqual("{}"))
+        #expect(HashSlot.hashTag(forKey: "{}bar").elementsEqual("{}bar"))
+    }
+
+    @Test
+    func byteBufferHashTagComputation() {
+        #expect(HashSlot(key: ValkeyKey(ByteBuffer(string: "foo"))) == HashSlot(key: "foo"))
+        #expect(HashSlot(key: ValkeyKey(ByteBuffer(string: "foo{bar}"))) == HashSlot(key: "foo{bar}"))
+        #expect(HashSlot(key: ValkeyKey(ByteBuffer(string: "foo{bar}"))) == HashSlot(key: "bar"))
     }
 
     @Test
@@ -113,3 +121,12 @@ struct HashSlotTests {
         #expect("\(HashSlot(rawValue: 20)!)" == "20")
     }
 }
+
+extension HashSlot {
+    /// Computes the portion of the key that should be used for hash slot calculation.
+    ///
+    /// Helper for tests
+    static func hashTag(forKey key: String) -> Substring {
+        Substring(Self.hashTag(forKey: key.utf8))
+    }
+}
diff --git a/Tests/ValkeyTests/CommandTests.swift b/Tests/ValkeyTests/CommandTests.swift
@@ -236,11 +236,11 @@ struct CommandTests {
             try await withThrowingTaskGroup(of: Void.self) { group in
                 group.addTask {
                     var result = try await connection.zmpop(key: ["key", "key2"], where: .max)
-                    #expect(result?.key == ValkeyKey(rawValue: "key"))
+                    #expect(result?.key == "key")
                     #expect(result?.values[0].score == 3)
                     #expect((result?.values[0].value).map { String(buffer: $0) } == "three")
                     result = try await connection.zmpop(key: ["key", "key2"], where: .max, count: 2)
-                    #expect(result?.key == ValkeyKey(rawValue: "key2"))
+                    #expect(result?.key == ValkeyKey("key2"))
                     #expect(result?.values[0].score == 5)
                     #expect((result?.values[0].value).map { String(buffer: $0) } == "five")
                     #expect(result?.values[1].score == 4)
diff --git a/Tests/ValkeyTests/ValkeyKeyTests.swift b/Tests/ValkeyTests/ValkeyKeyTests.swift

Original file line number	Diff line number	Diff line change
`@@ -192,7 +192,7 @@ extension ValkeyConnection {`
`192`	`192`	`process: (AsyncMapSequence<ValkeySubscription, ValkeyKey>) async throws -> sending Value`
`193`	`193`	`) async throws -> sending Value {`
`194`	`194`	`try await self.subscribe(to: [ValkeySubscriptions.invalidateChannel]) { subscription in`
`195`		`- let keys = subscription.map { ValkeyKey(rawValue: $0.message) }`
	`195`	`+ let keys = subscription.map { ValkeyKey($0.message) }`
`196`	`196`	`return try await process(keys)`
`197`	`197`	`}`
`198`	`198`	`}`
Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ struct ValkeySubscriptions {`
`87`	`87`	`case .forwardMessage(let subscriptions):`
`88`	`88`	`for subscription in subscriptions {`
`89`	`89`	`for key in keys {`
`90`		`- subscription.sendMessage(.init(channel: Self.invalidateChannel, message: key.rawValue))`
	`90`	`+ subscription.sendMessage(.init(channel: Self.invalidateChannel, message: String(valkeyKey: key)))`
`91`	`91`	`}`
`92`	`92`	`}`
`93`	`93`	`case .doNothing, .none:`