safe iterator tracking in hashtable. prevents invalid access if hashtable deleted from under safe iterator

Signed-off-by: Rain Valentine <[email protected]>

Author: rainsupreme

src/hashtable.c

@@ -284,6 +284,9 @@ typedef struct hashtableBucket {
 /* A key property is that the bucket size is one cache line. */
 static_assert(sizeof(bucket) == HASHTABLE_BUCKET_SIZE, "Bucket size mismatch");
 
+/* Forward declaration for iter type */
+typedef struct iter iter;
+
 struct hashtable {
     hashtableType *type;
     ssize_t rehash_idx;        /* -1 = rehashing not in progress. */
@@ -293,10 +296,11 @@ struct hashtable {
     int16_t pause_rehash;      /* Non-zero = rehashing is paused */
     int16_t pause_auto_shrink; /* Non-zero = automatic resizing disallowed. */
     size_t child_buckets[2];   /* Number of allocated child buckets. */
+    iter *safe_iterators; /* Head of linked list of safe iterators */
     void *metadata[];
 };
 
-typedef struct {
+struct iter {
     hashtable *hashtable;
     bucket *bucket;
     long index;
@@ -309,7 +313,8 @@ typedef struct {
         /* Safe iterator temporary storage for bucket chain compaction. */
         uint64_t last_seen_size;
     };
-} iter;
+    iter *next_safe_iter; /* Next safe iterator in hashtable's list */
+};
 
 /* The opaque hashtableIterator is defined as a blob of bytes. */
 static_assert(sizeof(hashtableIterator) >= sizeof(iter),
@@ -1084,6 +1089,41 @@ static inline int shouldPrefetchValues(iter *iter) {
     return (iter->flags & HASHTABLE_ITER_PREFETCH_VALUES);
 }
 
+/* Add a safe iterator to the hashtable's tracking list */
+static void trackSafeIterator(iter *it) {
+    hashtable *ht = it->hashtable;
+    it->next_safe_iter = ht->safe_iterators;
+    ht->safe_iterators = it;
+}
+
+/* Remove a safe iterator from the hashtable's tracking list */
+static void untrackSafeIterator(iter *it) {
+    hashtable *ht = it->hashtable;
+    if (ht->safe_iterators == it) {
+        ht->safe_iterators = it->next_safe_iter;
+    } else {
+        iter *prev = ht->safe_iterators;
+        while (prev && prev->next_safe_iter != it) {
+            prev = prev->next_safe_iter;
+        }
+        if (prev) {
+            prev->next_safe_iter = it->next_safe_iter;
+        }
+    }
+    it->next_safe_iter = NULL;
+    it->hashtable = NULL;
+}
+
+/* Invalidate all safe iterators by setting hashtable = NULL */
+static void invalidateAllSafeIterators(hashtable *ht) {
+    iter *it = ht->safe_iterators;
+    while (it) {
+        it->hashtable = NULL; /* Mark as invalid */
+        it = it->next_safe_iter;
+    }
+    ht->safe_iterators = NULL;
+}
+
 /* --- API functions --- */
 
 /* Allocates and initializes a new hashtable specified by the given type. */
@@ -1098,6 +1138,7 @@ hashtable *hashtableCreate(hashtableType *type) {
     ht->rehash_idx = -1;
     ht->pause_rehash = 0;
     ht->pause_auto_shrink = 0;
+    ht->safe_iterators = NULL;
     resetTable(ht, 0);
     resetTable(ht, 1);
     if (type->trackMemUsage) type->trackMemUsage(ht, alloc_size);
@@ -1153,6 +1194,7 @@ void hashtableEmpty(hashtable *ht, void(callback)(hashtable *)) {
 
 /* Deletes all the entries and frees the table. */
 void hashtableRelease(hashtable *ht) {
+    invalidateAllSafeIterators(ht);
     hashtableEmpty(ht, NULL);
     /* Call trackMemUsage before zfree, so trackMemUsage can access ht. */
     if (ht->type->trackMemUsage) {
@@ -1962,7 +2004,9 @@ size_t hashtableScanDefrag(hashtable *ht, size_t cursor, hashtableScanFunction f
  * rehashing to prevent entries from moving around. It's allowed to insert and
  * replace entries. Deleting entries is only allowed for the entry that was just
  * returned by hashtableNext. Deleting other entries is possible, but doing so
- * can cause internal fragmentation, so don't.
+ * can cause internal fragmentation, so don't. The hash table itself can be
+ * safely deleted while safe iterators exist - they will be invalidated and
+ * subsequent calls to hashtableNext will return false.
  *
  * Guarantees for safe iterators:
  *
@@ -1988,6 +2032,10 @@ void hashtableInitIterator(hashtableIterator *iterator, hashtable *ht, uint8_t f
     iter->table = 0;
     iter->index = -1;
     iter->flags = flags;
+    iter->next_safe_iter = NULL;
+    if (isSafe(iter) && ht != NULL) {
+        trackSafeIterator(iter);
+    }
 }
 
 /* Reinitializes the iterator for the provided hashtable while
@@ -2000,10 +2048,13 @@ void hashtableReinitIterator(hashtableIterator *iterator, hashtable *ht) {
 /* Resets a stack-allocated iterator. */
 void hashtableResetIterator(hashtableIterator *iterator) {
     iter *iter = iteratorFromOpaque(iterator);
+    if (iter->hashtable == NULL) return;
+
     if (!(iter->index == -1 && iter->table == 0)) {
         if (isSafe(iter)) {
             hashtableResumeRehashing(iter->hashtable);
             assert(iter->hashtable->pause_rehash >= 0);
+            untrackSafeIterator(iter);
         } else {
             assert(iter->fingerprint == hashtableFingerprint(iter->hashtable));
         }
@@ -2030,6 +2081,11 @@ void hashtableReleaseIterator(hashtableIterator *iterator) {
  * Returns false if there are no more entries. */
 bool hashtableNext(hashtableIterator *iterator, void **elemptr) {
     iter *iter = iteratorFromOpaque(iterator);
+    /* Check if iterator has been invalidated */
+    if (iter->hashtable == NULL) {
+        return false;
+    }
+    
     while (1) {
         if (iter->index == -1 && iter->table == 0) {
             /* It's the first call to next. */

src/hashtable.h

@@ -39,7 +39,7 @@ typedef struct hashtable hashtable;
 typedef struct hashtableStats hashtableStats;
 
 /* Can types that can be stack allocated. */
-typedef uint64_t hashtableIterator[5];
+typedef uint64_t hashtableIterator[6];
 typedef uint64_t hashtablePosition[2];
 typedef uint64_t hashtableIncrementalFindState[5];
 

src/unit/test_files.h

@@ -41,6 +41,11 @@ int test_compact_bucket_chain(int argc, char **argv, int flags);
 int test_random_entry(int argc, char **argv, int flags);
 int test_random_entry_with_long_chain(int argc, char **argv, int flags);
 int test_random_entry_sparse_table(int argc, char **argv, int flags);
+int test_safe_iterator_invalidation(int argc, char **argv, int flags);
+int test_safe_iterator_empty_no_invalidation(int argc, char **argv, int flags);
+int test_safe_iterator_reset_invalidation(int argc, char **argv, int flags);
+int test_safe_iterator_reset_untracking(int argc, char **argv, int flags);
+int test_safe_iterator_pause_resume_tracking(int argc, char **argv, int flags);
 int test_intsetValueEncodings(int argc, char **argv, int flags);
 int test_intsetBasicAdding(int argc, char **argv, int flags);
 int test_intsetLargeNumberRandomAdd(int argc, char **argv, int flags);
@@ -261,7 +266,7 @@ unitTest __test_crc64combine_c[] = {{"test_crc64combine", test_crc64combine}, {N
 unitTest __test_dict_c[] = {{"test_dictCreate", test_dictCreate}, {"test_dictAdd16Keys", test_dictAdd16Keys}, {"test_dictDisableResize", test_dictDisableResize}, {"test_dictAddOneKeyTriggerResize", test_dictAddOneKeyTriggerResize}, {"test_dictDeleteKeys", test_dictDeleteKeys}, {"test_dictDeleteOneKeyTriggerResize", test_dictDeleteOneKeyTriggerResize}, {"test_dictEmptyDirAdd128Keys", test_dictEmptyDirAdd128Keys}, {"test_dictDisableResizeReduceTo3", test_dictDisableResizeReduceTo3}, {"test_dictDeleteOneKeyTriggerResizeAgain", test_dictDeleteOneKeyTriggerResizeAgain}, {"test_dictBenchmark", test_dictBenchmark}, {NULL, NULL}};
 unitTest __test_endianconv_c[] = {{"test_endianconv", test_endianconv}, {NULL, NULL}};
 unitTest __test_entry_c[] = {{"test_entryCreate", test_entryCreate}, {"test_entryUpdate", test_entryUpdate}, {"test_entryHasexpiry_entrySetExpiry", test_entryHasexpiry_entrySetExpiry}, {"test_entryIsExpired", test_entryIsExpired}, {"test_entryMemUsage_entrySetExpiry_entrySetValue", test_entryMemUsage_entrySetExpiry_entrySetValue}, {NULL, NULL}};
-unitTest __test_hashtable_c[] = {{"test_cursor", test_cursor}, {"test_set_hash_function_seed", test_set_hash_function_seed}, {"test_add_find_delete", test_add_find_delete}, {"test_add_find_delete_avoid_resize", test_add_find_delete_avoid_resize}, {"test_instant_rehashing", test_instant_rehashing}, {"test_bucket_chain_length", test_bucket_chain_length}, {"test_two_phase_insert_and_pop", test_two_phase_insert_and_pop}, {"test_replace_reallocated_entry", test_replace_reallocated_entry}, {"test_incremental_find", test_incremental_find}, {"test_scan", test_scan}, {"test_iterator", test_iterator}, {"test_safe_iterator", test_safe_iterator}, {"test_compact_bucket_chain", test_compact_bucket_chain}, {"test_random_entry", test_random_entry}, {"test_random_entry_with_long_chain", test_random_entry_with_long_chain}, {"test_random_entry_sparse_table", test_random_entry_sparse_table}, {NULL, NULL}};
+unitTest __test_hashtable_c[] = {{"test_cursor", test_cursor}, {"test_set_hash_function_seed", test_set_hash_function_seed}, {"test_add_find_delete", test_add_find_delete}, {"test_add_find_delete_avoid_resize", test_add_find_delete_avoid_resize}, {"test_instant_rehashing", test_instant_rehashing}, {"test_bucket_chain_length", test_bucket_chain_length}, {"test_two_phase_insert_and_pop", test_two_phase_insert_and_pop}, {"test_replace_reallocated_entry", test_replace_reallocated_entry}, {"test_incremental_find", test_incremental_find}, {"test_scan", test_scan}, {"test_iterator", test_iterator}, {"test_safe_iterator", test_safe_iterator}, {"test_compact_bucket_chain", test_compact_bucket_chain}, {"test_random_entry", test_random_entry}, {"test_random_entry_with_long_chain", test_random_entry_with_long_chain}, {"test_random_entry_sparse_table", test_random_entry_sparse_table}, {"test_safe_iterator_invalidation", test_safe_iterator_invalidation}, {"test_safe_iterator_empty_no_invalidation", test_safe_iterator_empty_no_invalidation}, {"test_safe_iterator_reset_invalidation", test_safe_iterator_reset_invalidation}, {"test_safe_iterator_reset_untracking", test_safe_iterator_reset_untracking}, {"test_safe_iterator_pause_resume_tracking", test_safe_iterator_pause_resume_tracking}, {NULL, NULL}};
 unitTest __test_intset_c[] = {{"test_intsetValueEncodings", test_intsetValueEncodings}, {"test_intsetBasicAdding", test_intsetBasicAdding}, {"test_intsetLargeNumberRandomAdd", test_intsetLargeNumberRandomAdd}, {"test_intsetUpgradeFromint16Toint32", test_intsetUpgradeFromint16Toint32}, {"test_intsetUpgradeFromint16Toint64", test_intsetUpgradeFromint16Toint64}, {"test_intsetUpgradeFromint32Toint64", test_intsetUpgradeFromint32Toint64}, {"test_intsetStressLookups", test_intsetStressLookups}, {"test_intsetStressAddDelete", test_intsetStressAddDelete}, {NULL, NULL}};
 unitTest __test_kvstore_c[] = {{"test_kvstoreAdd16Keys", test_kvstoreAdd16Keys}, {"test_kvstoreIteratorRemoveAllKeysNoDeleteEmptyHashtable", test_kvstoreIteratorRemoveAllKeysNoDeleteEmptyHashtable}, {"test_kvstoreIteratorRemoveAllKeysDeleteEmptyHashtable", test_kvstoreIteratorRemoveAllKeysDeleteEmptyHashtable}, {"test_kvstoreHashtableIteratorRemoveAllKeysNoDeleteEmptyHashtable", test_kvstoreHashtableIteratorRemoveAllKeysNoDeleteEmptyHashtable}, {"test_kvstoreHashtableIteratorRemoveAllKeysDeleteEmptyHashtable", test_kvstoreHashtableIteratorRemoveAllKeysDeleteEmptyHashtable}, {"test_kvstoreHashtableExpand", test_kvstoreHashtableExpand}, {NULL, NULL}};
 unitTest __test_listpack_c[] = {{"test_listpackCreateIntList", test_listpackCreateIntList}, {"test_listpackCreateList", test_listpackCreateList}, {"test_listpackLpPrepend", test_listpackLpPrepend}, {"test_listpackLpPrependInteger", test_listpackLpPrependInteger}, {"test_listpackGetELementAtIndex", test_listpackGetELementAtIndex}, {"test_listpackPop", test_listpackPop}, {"test_listpackGetELementAtIndex2", test_listpackGetELementAtIndex2}, {"test_listpackIterate0toEnd", test_listpackIterate0toEnd}, {"test_listpackIterate1toEnd", test_listpackIterate1toEnd}, {"test_listpackIterate2toEnd", test_listpackIterate2toEnd}, {"test_listpackIterateBackToFront", test_listpackIterateBackToFront}, {"test_listpackIterateBackToFrontWithDelete", test_listpackIterateBackToFrontWithDelete}, {"test_listpackDeleteWhenNumIsMinusOne", test_listpackDeleteWhenNumIsMinusOne}, {"test_listpackDeleteWithNegativeIndex", test_listpackDeleteWithNegativeIndex}, {"test_listpackDeleteInclusiveRange0_0", test_listpackDeleteInclusiveRange0_0}, {"test_listpackDeleteInclusiveRange0_1", test_listpackDeleteInclusiveRange0_1}, {"test_listpackDeleteInclusiveRange1_2", test_listpackDeleteInclusiveRange1_2}, {"test_listpackDeleteWitStartIndexOutOfRange", test_listpackDeleteWitStartIndexOutOfRange}, {"test_listpackDeleteWitNumOverflow", test_listpackDeleteWitNumOverflow}, {"test_listpackBatchDelete", test_listpackBatchDelete}, {"test_listpackDeleteFooWhileIterating", test_listpackDeleteFooWhileIterating}, {"test_listpackReplaceWithSameSize", test_listpackReplaceWithSameSize}, {"test_listpackReplaceWithDifferentSize", test_listpackReplaceWithDifferentSize}, {"test_listpackRegressionGt255Bytes", test_listpackRegressionGt255Bytes}, {"test_listpackCreateLongListAndCheckIndices", test_listpackCreateLongListAndCheckIndices}, {"test_listpackCompareStrsWithLpEntries", test_listpackCompareStrsWithLpEntries}, {"test_listpackLpMergeEmptyLps", test_listpackLpMergeEmptyLps}, {"test_listpackLpMergeLp1Larger", test_listpackLpMergeLp1Larger}, {"test_listpackLpMergeLp2Larger", test_listpackLpMergeLp2Larger}, {"test_listpackLpNextRandom", test_listpackLpNextRandom}, {"test_listpackLpNextRandomCC", test_listpackLpNextRandomCC}, {"test_listpackRandomPairWithOneElement", test_listpackRandomPairWithOneElement}, {"test_listpackRandomPairWithManyElements", test_listpackRandomPairWithManyElements}, {"test_listpackRandomPairsWithOneElement", test_listpackRandomPairsWithOneElement}, {"test_listpackRandomPairsWithManyElements", test_listpackRandomPairsWithManyElements}, {"test_listpackRandomPairsUniqueWithOneElement", test_listpackRandomPairsUniqueWithOneElement}, {"test_listpackRandomPairsUniqueWithManyElements", test_listpackRandomPairsUniqueWithManyElements}, {"test_listpackPushVariousEncodings", test_listpackPushVariousEncodings}, {"test_listpackLpFind", test_listpackLpFind}, {"test_listpackLpValidateIntegrity", test_listpackLpValidateIntegrity}, {"test_listpackNumberOfElementsExceedsLP_HDR_NUMELE_UNKNOWN", test_listpackNumberOfElementsExceedsLP_HDR_NUMELE_UNKNOWN}, {"test_listpackStressWithRandom", test_listpackStressWithRandom}, {"test_listpackSTressWithVariableSize", test_listpackSTressWithVariableSize}, {"test_listpackBenchmarkInit", test_listpackBenchmarkInit}, {"test_listpackBenchmarkLpAppend", test_listpackBenchmarkLpAppend}, {"test_listpackBenchmarkLpFindString", test_listpackBenchmarkLpFindString}, {"test_listpackBenchmarkLpFindNumber", test_listpackBenchmarkLpFindNumber}, {"test_listpackBenchmarkLpSeek", test_listpackBenchmarkLpSeek}, {"test_listpackBenchmarkLpValidateIntegrity", test_listpackBenchmarkLpValidateIntegrity}, {"test_listpackBenchmarkLpCompareWithString", test_listpackBenchmarkLpCompareWithString}, {"test_listpackBenchmarkLpCompareWithNumber", test_listpackBenchmarkLpCompareWithNumber}, {"test_listpackBenchmarkFree", test_listpackBenchmarkFree}, {NULL, NULL}};