Fix CuckooFilter losing elements when an insert or expansion fails

probables/cuckoo/countingcuckoo.py

@@ -245,10 +245,13 @@ def _insert_fingerprint_alt(
         # either move everything around or hit the maximum number of swaps
         idx = random.choice([idx_1, idx_2])
         prv_bin = CountingCuckooBin(fingerprint, 1)
+        orig_bin = prv_bin
+        swaps = []
         for _ in range(self.max_swaps):
             # select one element to be swapped out...
             swap_elm = random.randint(0, self.bucket_size - 1)
             swap_finger = self.buckets[idx][swap_elm]
+            swaps.append((idx, swap_elm, swap_finger))
             prv_bin, self.buckets[idx][swap_elm] = swap_finger, prv_bin
 
             # now find another place to put this fingerprint
@@ -261,8 +264,13 @@ def _insert_fingerprint_alt(
                 self.__unique_elements += 1
                 return None
 
+        # unable to place the new fingerprint; undo the swaps so that the
+        # failed insertion leaves the filter unchanged
+        for idx, swap_elm, elm in reversed(swaps):
+            self.buckets[idx][swap_elm] = elm
+
         # if we got here we have an error... we might need to know what is left
-        return prv_bin
+        return orig_bin
 
     def _check_if_present(self, idx_1: int, idx_2: int, fingerprint: int) -> int | None:
         """wrapper for checking if fingerprint is already inserted"""
@@ -304,16 +312,30 @@ def _parse_buckets(self, d: ByteString) -> None:
 
     def _expand_logic(self, extra_fingerprint: Union["CountingCuckooBin", None]) -> None:
         """the logic to acutally expand the cuckoo filter"""
-        # get all the fingerprints
-        fingerprints = self._setup_expand(extra_fingerprint)
-        self.__unique_elements = 0  # this needs to be reset!
-
-        for elm in fingerprints:
-            idx_1, idx_2 = self._indicies_from_fingerprint(elm.finger)
-            res = self._insert_fingerprint_alt(elm.finger, idx_1, idx_2, elm.count)
-            if res is not None:  # again, this *shouldn't* happen
-                msg = "The CountingCuckooFilter failed to expand"
-                raise CuckooFilterFullError(msg)
+        # snapshot the current state so that a failed expansion leaves the
+        # filter unchanged
+        old_capacity = self._cuckoo_capacity
+        old_buckets = self._buckets
+        old_inserted = self._inserted_elements
+        old_unique = self.__unique_elements
+
+        try:
+            # get all the fingerprints
+            fingerprints = self._setup_expand(extra_fingerprint)
+            self.__unique_elements = 0  # this needs to be reset!
+
+            for elm in fingerprints:
+                idx_1, idx_2 = self._indicies_from_fingerprint(elm.finger)
+                res = self._insert_fingerprint_alt(elm.finger, idx_1, idx_2, elm.count)
+                if res is not None:  # again, this *shouldn't* happen
+                    msg = "The CountingCuckooFilter failed to expand"
+                    raise CuckooFilterFullError(msg)
+        except Exception:
+            self._cuckoo_capacity = old_capacity
+            self._buckets = old_buckets
+            self._inserted_elements = old_inserted
+            self.__unique_elements = old_unique
+            raise
 
     def __insert_element(self, fingerprint, idx, count=1) -> bool:
         """insert an element"""

probables/cuckoo/cuckoo.py

@@ -371,12 +371,15 @@ def _insert_fingerprint(self, fingerprint, idx_1, idx_2):
         # and move things around to the other index, if possible, until we
         # either move everything around or hit the maximum number of swaps
         idx = random.choice([idx_1, idx_2])
+        orig_fingerprint = fingerprint
+        swaps = []
 
         for _ in range(self.max_swaps):
             # select one element to be swapped out...
             swap_elm = random.randint(0, self.bucket_size - 1)
 
             swb = self.buckets[idx][swap_elm]
+            swaps.append((idx, swap_elm, swb))
             fingerprint, self.buckets[idx][swap_elm] = swb, fingerprint
 
             # now find another place to put this fingerprint
@@ -388,8 +391,13 @@ def _insert_fingerprint(self, fingerprint, idx_1, idx_2):
                 self._inserted_elements += 1
                 return None
 
+        # unable to place the new fingerprint; undo the swaps so that the
+        # failed insertion leaves the filter unchanged
+        for idx, swap_elm, elm in reversed(swaps):
+            self.buckets[idx][swap_elm] = elm
+
         # if we got here we have an error... we might need to know what is left
-        return fingerprint
+        return orig_fingerprint
 
     def _load(self, file: Path | str | IOBase | mmap | bytes) -> None:
         """load a cuckoo filter from file"""
@@ -454,15 +462,27 @@ def __insert_element(self, fingerprint, idx) -> bool:
 
     def _expand_logic(self, extra_fingerprint):
         """the logic to acutally expand the cuckoo filter"""
-        # get all the fingerprints
-        fingerprints = self._setup_expand(extra_fingerprint)
-
-        for finger in fingerprints:
-            idx_1, idx_2 = self._indicies_from_fingerprint(finger)
-            res = self._insert_fingerprint(finger, idx_1, idx_2)
-            if res is not None:  # again, this *shouldn't* happen
-                msg = "The CuckooFilter failed to expand"
-                raise CuckooFilterFullError(msg)
+        # snapshot the current state so that a failed expansion leaves the
+        # filter unchanged
+        old_capacity = self._cuckoo_capacity
+        old_buckets = self._buckets
+        old_inserted = self._inserted_elements
+
+        try:
+            # get all the fingerprints
+            fingerprints = self._setup_expand(extra_fingerprint)
+
+            for finger in fingerprints:
+                idx_1, idx_2 = self._indicies_from_fingerprint(finger)
+                res = self._insert_fingerprint(finger, idx_1, idx_2)
+                if res is not None:  # again, this *shouldn't* happen
+                    msg = "The CuckooFilter failed to expand"
+                    raise CuckooFilterFullError(msg)
+        except Exception:
+            self._cuckoo_capacity = old_capacity
+            self._buckets = old_buckets
+            self._inserted_elements = old_inserted
+            raise
 
     def _setup_expand(self, extra_fingerprint):
         """setup this thing"""

tests/countingcuckoo_test.py

@@ -3,6 +3,7 @@
 
 import hashlib
 import os
+import random
 import sys
 import unittest
 from pathlib import Path
@@ -127,6 +128,33 @@ def test_c_cuckoo_full_msg(self):
         else:
             self.assertEqual(True, False)
 
+    def test_c_cuckoo_filter_full_preserves_elements(self):
+        """test that a failed insertion does not remove previously added elements"""
+        self.addCleanup(random.setstate, random.getstate())
+        random.seed(0)
+        cko = CountingCuckooFilter(capacity=4, bucket_size=1, max_swaps=1, auto_expand=False)
+        added = ["k54", "k24", "k48", "k56"]
+        for key in added:
+            cko.add(key)
+        self.assertRaises(CuckooFilterFullError, cko.add, "k36")
+        self.assertEqual(len(added), cko.elements_added)
+        for key in added:
+            self.assertTrue(cko.check(key))
+
+    def test_c_cuckoo_filter_failed_expand_preserves_elements(self):
+        """test that a failed expansion does not remove previously added elements"""
+        self.addCleanup(random.setstate, random.getstate())
+        random.seed(15)
+        cko = CountingCuckooFilter(capacity=8, bucket_size=1, max_swaps=2)
+        added = ["k{}".format(i) for i in range(18)]
+        for key in added:
+            cko.add(key)
+        self.assertEqual(32, cko.capacity)  # earlier expansions succeeded
+        self.assertRaises(CuckooFilterFullError, cko.add, "k18")
+        self.assertEqual(32, cko.capacity)  # the failed expansion is rolled back
+        for key in added:
+            self.assertTrue(cko.check(key))
+
     def test_c_cuckoo_idx(self):
         """test that the indexing works correctly for counting cuckoo filter
         swap"""

tests/cuckoo_test.py

@@ -3,6 +3,7 @@
 
 import hashlib
 import os
+import random
 import sys
 import unittest
 from pathlib import Path
@@ -185,6 +186,34 @@ def test_cuckoo_full_msg(self):
         else:
             self.assertEqual(True, False)
 
+    def test_cuckoo_filter_full_preserves_elements(self):
+        """test that a failed insertion does not remove previously added elements"""
+        self.addCleanup(random.setstate, random.getstate())
+        random.seed(0)
+        cko = CuckooFilter(capacity=4, bucket_size=1, max_swaps=1, auto_expand=False)
+        added = ["k24", "k26", "k2", "k16"]
+        for key in added:
+            cko.add(key)
+        self.assertRaises(CuckooFilterFullError, cko.add, "k32")
+        self.assertEqual(len(added), cko.elements_added)
+        for key in added:
+            self.assertTrue(cko.check(key))
+
+    def test_cuckoo_filter_failed_expand_preserves_elements(self):
+        """test that a failed expansion does not remove previously added elements"""
+        self.addCleanup(random.setstate, random.getstate())
+        random.seed(1)
+        cko = CuckooFilter(capacity=4, bucket_size=1, max_swaps=1)
+        added = ["k8", "k36", "k54", "k51", "k7", "k31"]
+        for key in added:
+            cko.add(key)
+        self.assertEqual(8, cko.capacity)  # an expansion already succeeded
+        self.assertRaises(CuckooFilterFullError, cko.add, "k48")
+        self.assertEqual(8, cko.capacity)  # the failed expansion is rolled back
+        self.assertEqual(len(added), cko.elements_added)
+        for key in added:
+            self.assertTrue(cko.check(key))
+
     def test_cuckoo_idx(self):
         """test that the indexing works correctly for cuckoo filter swap"""
         cko = CuckooFilter(capacity=100, bucket_size=2, max_swaps=5)