diff --git a/bitchat/Services/MessageDeduplicationService.swift b/bitchat/Services/MessageDeduplicationService.swift index 6eba2e59..89de0fb0 100644 --- a/bitchat/Services/MessageDeduplicationService.swift +++ b/bitchat/Services/MessageDeduplicationService.swift @@ -43,6 +43,10 @@ final class LRUDeduplicationCache { map.count } + var orderStorageCountForTesting: Int { + order.count + } + /// Checks if a key exists in the cache func contains(_ key: String) -> Bool { map[key] != nil @@ -65,11 +69,13 @@ final class LRUDeduplicationCache { } trimIfNeeded() + compactIfNeeded() } /// Removes a specific key from the cache func remove(_ key: String) { - map.removeValue(forKey: key) + guard map.removeValue(forKey: key) != nil else { return } + compactIfNeeded() } /// Clears all entries from the cache @@ -93,13 +99,6 @@ final class LRUDeduplicationCache { let candidate = order[head] head += 1 - // Periodically compact the consumed prefix, matching the previous - // cache behavior without scanning unconsumed order nodes. - if head >= 32 && head * 2 >= order.count { - order.removeFirst(head) - head = 0 - } - guard let liveEntry = map[candidate.key], liveEntry.generation == candidate.generation else { continue @@ -110,6 +109,24 @@ final class LRUDeduplicationCache { } return false } + + private func compactIfNeeded() { + let unconsumedCount = order.count - head + let staleNodeCount = unconsumedCount - map.count + let shouldDropConsumedPrefix = head >= 32 && head * 2 >= order.count + let shouldRemoveStaleNodes = staleNodeCount >= 32 && staleNodeCount * 2 >= unconsumedCount + + guard shouldDropConsumedPrefix || shouldRemoveStaleNodes else { return } + + if shouldRemoveStaleNodes { + order = order[head...].filter { candidate in + map[candidate.key]?.generation == candidate.generation + } + } else { + order.removeFirst(head) + } + head = 0 + } } // MARK: - Content Normalizer diff --git a/bitchatTests/MessageDeduplicationServiceTests.swift b/bitchatTests/MessageDeduplicationServiceTests.swift index 5028b614..9e988cda 100644 --- a/bitchatTests/MessageDeduplicationServiceTests.swift +++ b/bitchatTests/MessageDeduplicationServiceTests.swift @@ -166,6 +166,29 @@ struct LRUDeduplicationCacheTests { #expect(cache.contains("c")) } + @Test func repeatedRemovalAndReinsertionCompactsStaleNodes() { + let cache = LRUDeduplicationCache(capacity: 3) + cache.record("a", value: 0) + + for value in 1...1_001 { + cache.remove("a") + cache.record("a", value: value) + } + + #expect(cache.count == 1) + #expect(cache.value(for: "a") == 1_001) + #expect(cache.orderStorageCountForTesting <= 32) + + cache.record("b", value: 1) + cache.record("c", value: 2) + cache.record("d", value: 3) + + #expect(!cache.contains("a")) + #expect(cache.contains("b")) + #expect(cache.contains("c")) + #expect(cache.contains("d")) + } + // MARK: - Edge Cases @Test func emptyKey_works() {