Fix disk cache cost tracking and purge edge cases

Sources/Arsenal/Implementations/DiskArsenal.swift

@@ -109,15 +109,25 @@ import os
         ArsenalActor.assertIsolated()
         await ensureCostCalculated()
 
-        guard let url = urlProvider.url(for: key) else {
+        guard var url = urlProvider.url(for: key) else {
             return
         }
         if let data = value?.toData() {
             do {
+                // Get old file size before overwriting
+                let oldSize = (try? url.resourceValues(forKeys: [.fileSizeKey]).fileSize) ?? 0
+
                 try data.write(to: url)
-                let size = try url.resourceValues(forKeys: [.fileSizeKey]).fileSize ?? 0
-                if size > 0 {
-                    cost += UInt64(size)
+
+                // Write succeeded - adjust cost (subtract old, add new)
+                if oldSize > 0 {
+                    cost -= UInt64(oldSize)
+                }
+                // Clear cached resource values to get accurate new file size
+                url.removeCachedResourceValue(forKey: .fileSizeKey)
+                let newSize = try url.resourceValues(forKeys: [.fileSizeKey]).fileSize ?? 0
+                if newSize > 0 {
+                    cost += UInt64(newSize)
                 }
                 await purge()
             } catch {
@@ -196,19 +206,26 @@ import os
         // Purge based on date
         if maxStaleness > 0 {
             let now = Date()
+            var itemsWithoutDates: [URL] = []
             while let url = sortedUrls.popLast() {
                 guard let date = try? url.resourceValues(forKeys: [.contentModificationDateKey]).contentModificationDate else {
+                    // Keep items without valid dates for cost-based purge
+                    itemsWithoutDates.append(url)
                     continue
                 }
 
                 // Item is too young, since we're sorted we can bail
                 guard now.timeIntervalSince1970 - date.timeIntervalSince1970 > maxStaleness else {
+                    // Put the item back so it can be considered for cost-based purge
+                    sortedUrls.append(url)
                     break
                 }
 
                 // We now know we need to delete the item
                 deleteItem(at: url)
             }
+            // Add back items without dates so they can be considered for cost-based purge
+            sortedUrls.append(contentsOf: itemsWithoutDates)
         }
 
         // Purge based on cost
@@ -244,8 +261,9 @@ import os
     private func deleteItem(at url: URL) {
         ArsenalActor.assertIsolated()
 
+        // Read size separately so we can still adjust cost even if this fails
+        let size = (try? url.resourceValues(forKeys: [.fileSizeKey]).fileSize) ?? 0
         do {
-            let size = try url.resourceValues(forKeys: [.fileSizeKey]).fileSize ?? 0
             try urlProvider.fileManager.removeItem(at: url)
             if size > 0 {
                 cost -= UInt64(size)

Sources/Arsenal/Implementations/MemoryArsenal.swift

@@ -148,7 +148,7 @@ import os
 
         // check our limits again in case we're
         // good after removing non-referenced items.
-        guard costLimit > 0, cost >= costLimit else {
+        guard costLimit > 0, cost > costLimit else {
             return
         }
 
@@ -157,7 +157,7 @@ import os
             item1.timestamp.compare(item2.timestamp) == .orderedAscending
         }
 
-        while !sorted.isEmpty, cost >= costLimit {
+        while !sorted.isEmpty, cost > costLimit {
             guard let item = sorted.first else {
                 break
             }

Tests/ArsenalTests/ArsenalTests.swift

@@ -83,7 +83,9 @@ class ArsenalTests: XCTestCase {
 
     func testMemoryPurgeOnLimitExceed() async {
         // Setting items until the cache exceeds its limit
-        for i in 0 ..< 500 { // Each item is 1024 bytes, limit is 512 KB
+        // Each item is 1024 bytes, limit is 512000 bytes (1024 * 500)
+        // Adding 501 items (513024 bytes) exceeds the limit
+        for i in 0 ..< 501 {
             let item = TestItem(data: Data(repeating: UInt8(i % 256), count: 1024), cost: 1024)
             await memoryCache.set(item, key: "key\(i)")
         }
@@ -297,6 +299,60 @@ class ArsenalTests: XCTestCase {
         XCTAssertEqual(retrieved?.toData(), item2.toData(), "Retrieved item should be the replacement.")
     }
 
+    func testOverwriteExistingKeyOnDisk() async {
+        let key = "diskOverwriteKey"
+        let item1 = TestItem(data: Data(repeating: 1, count: 100), cost: 100)
+        let item2 = TestItem(data: Data(repeating: 2, count: 200), cost: 200)
+
+        await diskCache.set(item1, key: key)
+        let cost1 = await diskCache.diskResourceCost
+        XCTAssertEqual(cost1, 100, "Disk cost should reflect first item.")
+
+        await diskCache.set(item2, key: key)
+        let cost2 = await diskCache.diskResourceCost
+        XCTAssertEqual(cost2, 200, "Disk cost should reflect replaced item, not sum.")
+
+        let retrieved = await diskCache.value(for: key)
+        XCTAssertEqual(retrieved?.toData(), item2.toData(), "Retrieved item should be the replacement.")
+    }
+
+    func testOverwriteExistingKeyOnCombinedCache() async {
+        let key = "combinedOverwriteKey"
+        let item1 = TestItem(data: Data(repeating: 1, count: 100), cost: 100)
+        let item2 = TestItem(data: Data(repeating: 2, count: 200), cost: 200)
+
+        await combinedCache.set(item1, key: key)
+        let memoryCost1 = await combinedCache.memoryResourceCost
+        let diskCost1 = await combinedCache.diskResourceCost
+        XCTAssertEqual(memoryCost1, 100, "Memory cost should reflect first item.")
+        XCTAssertEqual(diskCost1, 100, "Disk cost should reflect first item.")
+
+        await combinedCache.set(item2, key: key)
+        let memoryCost2 = await combinedCache.memoryResourceCost
+        let diskCost2 = await combinedCache.diskResourceCost
+        XCTAssertEqual(memoryCost2, 200, "Memory cost should reflect replaced item, not sum.")
+        XCTAssertEqual(diskCost2, 200, "Disk cost should reflect replaced item, not sum.")
+
+        let retrieved = await combinedCache.value(for: key)
+        XCTAssertEqual(retrieved?.toData(), item2.toData(), "Retrieved item should be the replacement.")
+    }
+
+    func testMultipleOverwritesOnDisk() async {
+        let key = "multiOverwriteKey"
+
+        // Overwrite the same key multiple times
+        for i in 1 ... 5 {
+            let item = TestItem(data: Data(repeating: UInt8(i), count: 100), cost: 100)
+            await diskCache.set(item, key: key)
+        }
+
+        let finalCost = await diskCache.diskResourceCost
+        XCTAssertEqual(finalCost, 100, "Cost should only reflect the single item, not accumulated from overwrites.")
+
+        let retrieved = await diskCache.value(for: key)
+        XCTAssertEqual(retrieved?.toData(), Data(repeating: 5, count: 100), "Retrieved item should be the last written value.")
+    }
+
     func testEmptyCachePurge() async {
         // Purging empty cache should not crash
         await memoryCache.purge([.memory])