Implement Add and Replace support (#28)

* add + replace #25

* soundness

* more test for add + replace

* build faiure fix?

Author: dkz2

Sources/SwiftMemcache/Extensions/ByteBuffer+SwiftMemcache.swift

@@ -88,15 +88,17 @@ extension ByteBuffer {
         }
 
         if let storageMode = flags.storageMode {
+            self.writeInteger(UInt8.whitespace)
+            self.writeInteger(UInt8.M)
             switch storageMode {
+            case .add:
+                self.writeInteger(UInt8.E)
             case .append:
-                self.writeInteger(UInt8.whitespace)
-                self.writeInteger(UInt8.M)
                 self.writeInteger(UInt8.A)
             case .prepend:
-                self.writeInteger(UInt8.whitespace)
-                self.writeInteger(UInt8.M)
                 self.writeInteger(UInt8.P)
+            case .replace:
+                self.writeInteger(UInt8.R)
             }
         }
     }

Sources/SwiftMemcache/Extensions/UInt8+Characters.swift

@@ -25,6 +25,8 @@ extension UInt8 {
     static var M: UInt8 = .init(ascii: "M")
     static var P: UInt8 = .init(ascii: "P")
     static var A: UInt8 = .init(ascii: "A")
+    static var E: UInt8 = .init(ascii: "E")
+    static var R: UInt8 = .init(ascii: "R")
     static var zero: UInt8 = .init(ascii: "0")
     static var nine: UInt8 = .init(ascii: "9")
 }

Sources/SwiftMemcache/MemcachedConnection.swift

@@ -61,6 +61,8 @@ public actor MemcachedConnection {
         case unexpectedNilResponse
         /// Indicates that the key was not found.
         case keyNotFound
+        /// Indicates that the key already exist
+        case keyExist
     }
 
     private var state: State
@@ -340,4 +342,86 @@ public actor MemcachedConnection {
             throw MemcachedConnectionError.connectionShutdown
         }
     }
+
+    // MARK: - Adding a Value
+
+    /// Adds a new key-value pair in the Memcached server.
+    /// The operation will fail if the key already exists.
+    ///
+    /// - Parameters:
+    ///   - key: The key to add the value to.
+    ///   - value: The `MemcachedValue` to add.
+    /// - Throws: A `MemcachedConnectionError.connectionShutdown` if the connection to the Memcached server is shut down.
+    /// - Throws: A `MemcachedConnectionError.keyExist` if the key already exists in the Memcached server.
+    /// - Throws: A `MemcachedConnectionError.unexpectedNilResponse` if an unexpected response code is returned.
+    public func add(_ key: String, value: some MemcachedValue) async throws {
+        switch self.state {
+        case .initial(_, let bufferAllocator, _, _),
+             .running(let bufferAllocator, _, _, _):
+
+            var buffer = bufferAllocator.buffer(capacity: 0)
+            value.writeToBuffer(&buffer)
+            var flags: MemcachedFlags
+
+            flags = MemcachedFlags()
+            flags.storageMode = .add
+
+            let command = MemcachedRequest.SetCommand(key: key, value: buffer, flags: flags)
+            let request = MemcachedRequest.set(command)
+
+            let response = try await sendRequest(request)
+
+            switch response.returnCode {
+            case .HD:
+                return
+            case .NS:
+                throw MemcachedConnectionError.keyExist
+            default:
+                throw MemcachedConnectionError.unexpectedNilResponse
+            }
+
+        case .finished:
+            throw MemcachedConnectionError.connectionShutdown
+        }
+    }
+
+    // MARK: - Replacing a Value
+
+    /// Replace the value for an existing key in the Memcache server.
+    /// The operation will fail if the key does not exist.
+    ///
+    /// - Parameters:
+    ///   - key: The key to replace the value for.
+    ///   - value: The `MemcachedValue` to replace.
+    /// - Throws: A `MemcachedConnectionError` if the connection to the Memcached server is shut down.
+    public func replace(_ key: String, value: some MemcachedValue) async throws {
+        switch self.state {
+        case .initial(_, let bufferAllocator, _, _),
+             .running(let bufferAllocator, _, _, _):
+
+            var buffer = bufferAllocator.buffer(capacity: 0)
+            value.writeToBuffer(&buffer)
+            var flags: MemcachedFlags
+
+            flags = MemcachedFlags()
+            flags.storageMode = .replace
+
+            let command = MemcachedRequest.SetCommand(key: key, value: buffer, flags: flags)
+            let request = MemcachedRequest.set(command)
+
+            let response = try await sendRequest(request)
+
+            switch response.returnCode {
+            case .HD:
+                return
+            case .NS:
+                throw MemcachedConnectionError.keyNotFound
+            default:
+                throw MemcachedConnectionError.unexpectedNilResponse
+            }
+
+        case .finished:
+            throw MemcachedConnectionError.connectionShutdown
+        }
+    }
 }

Sources/SwiftMemcache/MemcachedFlags.swift

@@ -49,11 +49,15 @@ public enum TimeToLive: Equatable, Hashable {
 }
 
 /// Enum representing the Memcached 'ms' (meta set) command modes (corresponding to the 'M' flag).
-public enum StorageMode: Equatable, Hashable {
+enum StorageMode: Equatable, Hashable {
+    /// The "add" command. If the item exists, LRU is bumped and NS is returned.
+    case add
     /// The 'append' command. If the item exists, append the new value to its data.
     case append
     /// The 'prepend' command. If the item exists, prepend the new value to its data.
     case prepend
+    /// The "replace" command. The new value is set only if the item already exists.
+    case replace
 }
 
 extension MemcachedFlags: Hashable {}

Tests/SwiftMemcacheTests/IntegrationTest/MemcachedIntegrationTests.swift

@@ -306,6 +306,134 @@ final class MemcachedIntegrationTest: XCTestCase {
         }
     }
 
+    func testAddValue() async throws {
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try! group.syncShutdownGracefully())
+        }
+        let memcachedConnection = MemcachedConnection(host: "memcached", port: 11211, eventLoopGroup: group)
+
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask { try await memcachedConnection.run() }
+
+            // Add a value to a key
+            let addValue = "foo"
+
+            // Attempt to delete the key, but ignore the error if it doesn't exist
+            do {
+                try await memcachedConnection.delete("adds")
+            } catch {
+                if "\(error)" != "keyNotFound" {
+                    throw error
+                }
+            }
+
+            // Proceed with adding the key-value pair
+            try await memcachedConnection.add("adds", value: addValue)
+
+            // Get value for the key after add operation
+            let addedValue: String? = try await memcachedConnection.get("adds")
+            XCTAssertEqual(addedValue, addValue, "Received value should be the same as the added value")
+
+            group.cancelAll()
+        }
+    }
+
+    func testAddValueKeyExists() async throws {
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try! group.syncShutdownGracefully())
+        }
+        let memcachedConnection = MemcachedConnection(host: "memcached", port: 11211, eventLoopGroup: group)
+
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask { try await memcachedConnection.run() }
+
+            // Add a value to a key
+            let initialValue = "foo"
+            let newValue = "bar"
+
+            // Attempt to delete the key, but ignore the error if it doesn't exist
+            do {
+                try await memcachedConnection.delete("adds")
+            } catch {
+                if "\(error)" != "keyNotFound" {
+                    throw error
+                }
+            }
+
+            // Set an initial value for the key
+            try await memcachedConnection.add("adds", value: initialValue)
+
+            do {
+                // Attempt to add a new value to the existing key
+                try await memcachedConnection.add("adds", value: newValue)
+                XCTFail("Expected an error indicating the key exists, but no error was thrown.")
+            } catch {
+                // Check if the error description or localized description matches the expected error
+                if "\(error)" != "keyExist" {
+                    XCTFail("Unexpected error: \(error)")
+                }
+            }
+
+            group.cancelAll()
+        }
+    }
+
+    func testReplaceValue() async throws {
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try! group.syncShutdownGracefully())
+        }
+        let memcachedConnection = MemcachedConnection(host: "memcached", port: 11211, eventLoopGroup: group)
+
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask { try await memcachedConnection.run() }
+
+            // Set key and initial value
+            let initialValue = "foo"
+            try await memcachedConnection.set("greet", value: initialValue)
+
+            // Replace value for the key
+            let replaceValue = "hi"
+            try await memcachedConnection.replace("greet", value: replaceValue)
+
+            // Get value for the key after replace operation
+            let replacedValue: String? = try await memcachedConnection.get("greet")
+            XCTAssertEqual(replacedValue, replaceValue, "Received value should be the same as the replaceValue")
+
+            group.cancelAll()
+        }
+    }
+
+    func testReplaceNonExistentKey() async throws {
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try! group.syncShutdownGracefully())
+        }
+        let memcachedConnection = MemcachedConnection(host: "memcached", port: 11211, eventLoopGroup: group)
+
+        await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask { try await memcachedConnection.run() }
+
+            do {
+                // Ensure the key is clean
+                try await memcachedConnection.delete("nonExistentKey")
+                // Attempt to replace value for a non-existent key
+                let replaceValue = "testValue"
+                try await memcachedConnection.replace("nonExistentKey", value: replaceValue)
+                XCTFail("Expected an error indicating the key was not found, but no error was thrown.")
+            } catch {
+                // Check if the error description or localized description matches the expected error
+                if "\(error)" != "keyNotFound" {
+                    XCTFail("Unexpected error: \(error)")
+                }
+            }
+
+            group.cancelAll()
+        }
+    }
+
     func testMemcachedConnectionWithUInt() async throws {
         let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         defer {

Tests/SwiftMemcacheTests/UnitTest/MemcachedFlagsTests.swift

@@ -38,6 +38,16 @@ final class MemcachedFlagsTests: XCTestCase {
         }
     }
 
+    func testStorageModeAdd() {
+        var flags = MemcachedFlags()
+        flags.storageMode = .add
+        if case .add? = flags.storageMode {
+            XCTAssertTrue(true)
+        } else {
+            XCTFail("Flag storageMode is not .add")
+        }
+    }
+
     func testStorageModeAppend() {
         var flags = MemcachedFlags()
         flags.storageMode = .append
@@ -57,4 +67,14 @@ final class MemcachedFlagsTests: XCTestCase {
             XCTFail("Flag storageMode is not .prepend")
         }
     }
+
+    func testStorageModeReplace() {
+        var flags = MemcachedFlags()
+        flags.storageMode = .replace
+        if case .replace? = flags.storageMode {
+            XCTAssertTrue(true)
+        } else {
+            XCTFail("Flag storageMode is not .replace")
+        }
+    }
 }

Tests/SwiftMemcacheTests/UnitTest/MemcachedRequestEncoderTests.swift

@@ -48,38 +48,37 @@ final class MemcachedRequestEncoderTests: XCTestCase {
         XCTAssertEqual(outBuffer.getString(at: 0, length: outBuffer.readableBytes), expectedEncodedData)
     }
 
-    func testEncodeAppendRequest() {
+    func testEncodeStorageRequest(withMode mode: StorageMode, expectedEncodedData: String) {
         // Prepare a MemcachedRequest
         var buffer = ByteBufferAllocator().buffer(capacity: 2)
         buffer.writeString("hi")
 
         var flags = MemcachedFlags()
-        flags.storageMode = .append
+        flags.storageMode = mode
         let command = MemcachedRequest.SetCommand(key: "foo", value: buffer, flags: flags)
         let request = MemcachedRequest.set(command)
 
         // pass our request through the encoder
         let outBuffer = self.encodeRequest(request)
 
-        let expectedEncodedData = "ms foo 2 MA\r\nhi\r\n"
+        // assert the encoded request
         XCTAssertEqual(outBuffer.getString(at: 0, length: outBuffer.readableBytes), expectedEncodedData)
     }
 
-    func testEncodePrependRequest() {
-        // Prepare a MemcachedRequest
-        var buffer = ByteBufferAllocator().buffer(capacity: 2)
-        buffer.writeString("hi")
+    func testEncodeAppendRequest() {
+        self.testEncodeStorageRequest(withMode: .append, expectedEncodedData: "ms foo 2 MA\r\nhi\r\n")
+    }
 
-        var flags = MemcachedFlags()
-        flags.storageMode = .prepend
-        let command = MemcachedRequest.SetCommand(key: "foo", value: buffer, flags: flags)
-        let request = MemcachedRequest.set(command)
+    func testEncodePrependRequest() {
+        self.testEncodeStorageRequest(withMode: .prepend, expectedEncodedData: "ms foo 2 MP\r\nhi\r\n")
+    }
 
-        // pass our request through the encoder
-        let outBuffer = self.encodeRequest(request)
+    func testEncodeAddRequest() {
+        self.testEncodeStorageRequest(withMode: .add, expectedEncodedData: "ms foo 2 ME\r\nhi\r\n")
+    }
 
-        let expectedEncodedData = "ms foo 2 MP\r\nhi\r\n"
-        XCTAssertEqual(outBuffer.getString(at: 0, length: outBuffer.readableBytes), expectedEncodedData)
+    func testEncodeReplaceRequest() {
+        self.testEncodeStorageRequest(withMode: .replace, expectedEncodedData: "ms foo 2 MR\r\nhi\r\n")
     }
 
     func testEncodeTouchRequest() {