1 files changed, 50 insertions, 0 deletions
diff --git a/src/dict.c b/src/dict.c
index 9ed461d62..7fecb2daa 100644
--- a/src/dict.c
+++ b/src/dict.c
@@ -541,6 +541,56 @@ void *dictFetchValue(dict *d, const void *key) {
     return he ? dictGetVal(he) : NULL;
 }
 
+/* Find an element from the table, also get the plink of the entry. The entry
+ * is returned if the element is found, and the user should later call
+ * `dictTwoPhaseUnlinkFree` with it in order to unlink and release it. Otherwise if
+ * the key is not found, NULL is returned. These two functions should be used in pair.
+ * `dictTwoPhaseUnlinkFind` pauses rehash and `dictTwoPhaseUnlinkFree` resumes rehash.
+ *
+ * We can use like this:
+ *
+ * dictEntry *de = dictTwoPhaseUnlinkFind(db->dict,key->ptr,&plink, &table);
+ * // Do something, but we can't modify the dict
+ * dictTwoPhaseUnlinkFree(db->dict,de,plink,table); // We don't need to lookup again
+ *
+ * If we want to find an entry before delete this entry, this an optimization to avoid
+ * dictFind followed by dictDelete. i.e. the first API is a find, and it gives some info
+ * to the second one to avoid repeating the lookup
+ */
+dictEntry *dictTwoPhaseUnlinkFind(dict *d, const void *key, dictEntry ***plink, int *table_index) {
+    uint64_t h, idx, table;
+
+    if (dictSize(d) == 0) return NULL; /* dict is empty */
+    if (dictIsRehashing(d)) _dictRehashStep(d);
+    h = dictHashKey(d, key);
+
+    for (table = 0; table <= 1; table++) {
+        idx = h & DICTHT_SIZE_MASK(d->ht_size_exp[table]);
+        dictEntry **ref = &d->ht_table[table][idx];
+        while(*ref) {
+            if (key==(*ref)->key || dictCompareKeys(d, key, (*ref)->key)) {
+                *table_index = table;
+                *plink = ref;
+                dictPauseRehashing(d);
+                return *ref;
+            }
+            ref = &(*ref)->next;
+        }
+        if (!dictIsRehashing(d)) return NULL;
+    }
+    return NULL;
+}
+
+void dictTwoPhaseUnlinkFree(dict *d, dictEntry *he, dictEntry **plink, int table_index) {
+    if (he == NULL) return;
+    d->ht_used[table_index]--;
+    *plink = he->next;
+    dictFreeKey(d, he);
+    dictFreeVal(d, he);
+    zfree(he);
+    dictResumeRehashing(d);
+}
+
 /* A fingerprint is a 64 bit number that represents the state of the dictionary
  * at a given time, it's just a few dict properties xored together.
  * When an unsafe iterator is initialized, we get the dict fingerprint, and check