update clusterStorageLoadCallback

* remove short circuit as it is unsafe

src/db.cpp

@@ -53,7 +53,7 @@ struct dbBackup {
  *----------------------------------------------------------------------------*/
 
 int expireIfNeeded(redisDb *db, robj *key, robj *o);
-void slotToKeyUpdateKeyCore(const char *key, size_t keylen, int add);
+void slotToKeyUpdateKeyCore(const char *key, size_t keylen, int add, int noModify);
 
 std::unique_ptr<expireEntry> deserializeExpire(sds key, const char *str, size_t cch, size_t *poffset);
 sds serializeStoredObjectAndExpire(redisDbPersistentData *db, const char *key, robj_roptr o);
@@ -2471,11 +2471,11 @@ int xreadGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult
  * a fast way a key that belongs to a specified hash slot. This is useful
  * while rehashing the cluster and in other conditions when we need to
  * understand if we have keys for a given hash slot. */
-void slotToKeyUpdateKey(sds key, int add) {
-    slotToKeyUpdateKeyCore(key, sdslen(key), add);
+void slotToKeyUpdateKey(sds key, int add, int noModify) {
+    slotToKeyUpdateKeyCore(key, sdslen(key), add, noModify);
 }
 
-void slotToKeyUpdateKeyCore(const char *key, size_t keylen, int add) {
+void slotToKeyUpdateKeyCore(const char *key, size_t keylen, int add, int noModify) {
     serverAssert(GlobalLocksAcquired());
 
     unsigned int hashslot = keyHashSlot(key,keylen);
@@ -2489,23 +2489,28 @@ void slotToKeyUpdateKeyCore(const char *key, size_t keylen, int add) {
     memcpy(indexed+2,key,keylen);
     int fModified = false;
     if (add) {
-        fModified = raxInsert(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL,NULL);
+        if (noModify) fModified = raxTryInsert(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL,NULL);
+        else fModified = raxInsert(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL,NULL);
     } else {
         fModified = raxRemove(g_pserver->cluster->slots_to_keys,indexed,keylen+2,NULL);
     }
     // This assert is disabled when a snapshot depth is >0 because prepOverwriteForSnapshot will add in a tombstone,
     //  this prevents ensure from adding the key to the dictionary which means the caller isn't aware we're already tracking
     //  the key.
-    serverAssert(fModified || g_pserver->db[0]->snapshot_depth() > 0);
+    serverAssert(fModified || g_pserver->db[0]->snapshot_depth() > 0 || noModify);
     if (indexed != buf) zfree(indexed);
 }
 
+void slotToKeyTryAdd(sds key) {
+    slotToKeyUpdateKey(key,1,1);
+}
+
 void slotToKeyAdd(sds key) {
-    slotToKeyUpdateKey(key,1);
+    slotToKeyUpdateKey(key,1,0);
 }
 
 void slotToKeyDel(sds key) {
-    slotToKeyUpdateKey(key,0);
+    slotToKeyUpdateKey(key,0,0);
 }
 
 /* Release the radix tree mapping Redis Cluster keys to slots. If 'async'
@@ -2604,7 +2609,7 @@ void redisDbPersistentData::endStorageProvider()
 
 void clusterStorageLoadCallback(const char *rgchkey, size_t cch, void *)
 {
-    slotToKeyUpdateKeyCore(rgchkey, cch, true /*add*/);
+    slotToKeyUpdateKeyCore(rgchkey, cch, true /*add*/, false /*noModify*/);
 }
 
 void moduleLoadCallback(const char * rgchKey, size_t cchKey, void *data) {
@@ -2858,6 +2863,9 @@ LNotFound:
             {
                 dictAdd(m_pdict, sdsNewKey, o);
                 o->SetFExpires(spexpire != nullptr);
+                if (g_pserver->cluster_enabled) {
+                    slotToKeyTryAdd(sdsNewKey);
+                }
 
                 if (spexpire != nullptr)
                 {
@@ -3265,7 +3273,7 @@ int dbnumFromDb(redisDb *db)
     serverPanic("invalid database pointer");
 }
 
-bool redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command, bool fExecOK)
+void redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command)
 {
     if (m_spstorage == nullptr) {
 #if defined(__x86_64__) || defined(__i386__)
@@ -3277,7 +3285,7 @@ bool redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command
             const char *cmd = szFromObj(command.argv[0]);
             if (!strcasecmp(cmd, "set") || !strcasecmp(cmd, "get")) {
                 if (c->db->m_spdbSnapshotHOLDER != nullptr)
-                    return false; // this is dangerous enough without a snapshot around
+                    return; // this is dangerous enough without a snapshot around
                 auto h = dictSdsHash(szFromObj(command.argv[1]));
                 for (int iht = 0; iht < 2; ++iht) {
                     auto hT = h & c->db->m_pdict->ht[iht].sizemask;
@@ -3297,7 +3305,7 @@ bool redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command
             }
         }
 #endif
-        return false;
+        return;
     }
 
     AeLocker lock;
@@ -3307,10 +3315,10 @@ bool redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command
     getKeysResult result = GETKEYS_RESULT_INIT;
     auto cmd = lookupCommand(szFromObj(command.argv[0]));
     if (cmd == nullptr)
-        return false; // Bad command? It's not for us to judge, just bail
+        return; // Bad command? It's not for us to judge, just bail
     
     if (command.argc < std::abs(cmd->arity))
-        return false; // Invalid number of args
+        return; // Invalid number of args
     
     int numkeys = getKeysFromCommand(cmd, command.argv, command.argc, &result);
     for (int ikey = 0; ikey < numkeys; ++ikey)
@@ -3343,7 +3351,6 @@ bool redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command
         }
     }
 
-    bool fNoInsert = false;
     if (!vecInserts.empty()) {
         lock.arm(c);
         for (auto &tuple : vecInserts)
@@ -3359,13 +3366,11 @@ bool redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command
                     // While unlocked this was already ensured
                     decrRefCount(o);
                     sdsfree(sharedKey);
-                    fNoInsert = true;
                 }
                 else
                 {
                     if (spexpire != nullptr) {
                         if (spexpire->when() < mstime()) {
-                            fNoInsert = true;
                             break;
                         }
                     }
@@ -3392,12 +3397,5 @@ bool redisDbPersistentData::prefetchKeysAsync(client *c, parsed_command &command
         lock.disarm();
     }
 
-    if (fExecOK && !fNoInsert && cmd->proc == getCommand && !vecInserts.empty()) {
-        robj *o = std::get<1>(vecInserts[0]);
-        if (o != nullptr) {
-            addReplyBulk(c, o);
-            return true;
-        }
-    }
-    return false;
-}
+    return;
+}

src/networking.cpp

@@ -2595,9 +2595,7 @@ void parseClientCommandBuffer(client *c) {
             (g_pserver->m_pstorageFactory || aeLockContested(cserver.cthreads/2) || cserver.cthreads == 1) && !GlobalLocksAcquired()) {
             auto &query = c->vecqueuedcmd.back();
             if (query.argc > 0 && query.argc == query.argcMax) {
-                if (c->db->prefetchKeysAsync(c, query, c->vecqueuedcmd.size() == 1)) {
-                    c->vecqueuedcmd.erase(c->vecqueuedcmd.begin());
-                }
+                c->db->prefetchKeysAsync(c, query);
             }
         }
         c->reqtype = 0;

src/server.cpp

@@ -2219,6 +2219,7 @@ void checkChildrenDone(void) {
             g_pserver->rdbThreadVars.fRdbThreadCancel = false;
             g_pserver->rdbThreadVars.fDone = false;
             if (exitcode == 0) receiveChildInfo();
+            closeChildInfoPipe();
         }
     }
     else if ((pid = waitpid(-1, &statloc, WNOHANG)) != 0) {

src/server.h

@@ -1204,7 +1204,7 @@ public:
     void disableKeyCache();
     bool keycacheIsEnabled();
 
-    bool prefetchKeysAsync(client *c, struct parsed_command &command, bool fExecOK);
+    void prefetchKeysAsync(client *c, struct parsed_command &command);
 
     bool FSnapshot() const { return m_spdbSnapshotHOLDER != nullptr; }