Async flash

build.yaml

@@ -21,13 +21,25 @@ machamp:
         parent: make-build
         # https://github.sc-corp.net/Snapchat/img/tree/master/keydb/ubuntu-20-04
         builder_image: us.gcr.io/snapchat-build-artifacts/prod/snapchat/img/keydb/keydb-ubuntu-20-04@sha256:cf869a3f5d1de1e1d976bb906689c37b7031938eb68661b844a38c532f27248c
-        command: ./runtest --clients 4 --verbose --tls
+        command: ./runtest --clients 4 --verbose --tls --skipunit unit/memefficiency
+      tests-with-flash:
+        type: cmd
+        parent: make-build
+        # https://github.sc-corp.net/Snapchat/img/tree/master/keydb/ubuntu-20-04
+        builder_image: us.gcr.io/snapchat-build-artifacts/prod/snapchat/img/keydb/keydb-ubuntu-20-04@sha256:cf869a3f5d1de1e1d976bb906689c37b7031938eb68661b844a38c532f27248c
+        command: ./runtest --verbose --skipunit unit/memefficiency --flash
       cluster-test:
         type: cmd
         parent: make-build
         # https://github.sc-corp.net/Snapchat/img/tree/master/keydb/ubuntu-20-04
         builder_image: us.gcr.io/snapchat-build-artifacts/prod/snapchat/img/keydb/keydb-ubuntu-20-04@sha256:cf869a3f5d1de1e1d976bb906689c37b7031938eb68661b844a38c532f27248c
         command: ./runtest-cluster --tls
+      cluster-test-flash:
+        type: cmd
+        parent: make-build
+        # https://github.sc-corp.net/Snapchat/img/tree/master/keydb/ubuntu-20-04
+        builder_image: us.gcr.io/snapchat-build-artifacts/prod/snapchat/img/keydb/keydb-ubuntu-20-04@sha256:cf869a3f5d1de1e1d976bb906689c37b7031938eb68661b844a38c532f27248c
+        command: ./runtest-cluster --tls --flash
       sentinel-test:
         type: cmd
         parent: make-build
@@ -59,7 +71,7 @@ machamp:
       # to ensure a clearer docker build env
       code-checkout:
         type: cmd
-        command: echo checkout
+        command: echo checkout && git submodule init && git submodule update
         # default machamp builder image does not work for multi arch
         builder_image: us.gcr.io/snapchat-build-artifacts/prod/snapchat/img/ubuntu/ubuntu-23-04@sha256:bd43177a80e6ce1c3583e8ea959b88a9081c0f56b765ec9c5a157c27a637c23b
       docker:

ci.yaml

@@ -26,6 +26,16 @@ on:
           # references a build defined in build.yaml
           build_name: keydb-docker-build
           arch_types: ["amd64", "arm64"]
+  # Doc: go/cool-guide
+  cool:
+    workflows:
+      - workflow_type: backend_workflow
+        build_name: keydb-build
+        arch_types: ["amd64", "arm64"]
+      - workflow_type: backend_workflow
+        # references a build defined in build.yaml
+        build_name: keydb-docker-build
+        arch_types: ["amd64", "arm64"]
 
 # below defines which branch is release branch / release tag
 machamp:

machamp_scripts/Dockerfile

@@ -24,7 +24,7 @@ RUN set -eux; \
         gosu nobody true
 # build KeyDB
 ARG MAKE_JOBS=""
-ARG ENABLE_FLASH=""
+ARG ENABLE_FLASH="yes"
 COPY . /tmp/keydb-internal
 RUN set -eux; \
         cd /tmp/keydb-internal; \
@@ -73,7 +73,25 @@ RUN set -eux; \
                | xargs -r apt-mark manual \
         ; \
         apt-get purge -y --auto-remove -o APT::AutoRemove::RecommendsImportant=false; \
-        rm -rf /var/lib/apt/lists/*; \
+        rm -rf /var/lib/apt/lists/*
+# build RedisBloom
+RUN set -eux; \
+        savedAptMark="$(apt-mark showmanual)"; \
+        apt-get update; \
+        DEBIAN_FRONTEND=noninteractive apt-get -o Dpkg::Options::="--force-confnew" install -qqy --no-install-recommends \
+                ca-certificates \
+                wget \
+                make \
+                pkg-config \
+                build-essential \
+                git; \
+        git clone --recursive https://github.com/RedisBloom/RedisBloom.git; \
+        cd RedisBloom; \
+        ./sbin/setup; \
+        bash -l; \
+        make; \
+        redisBloomPath=`make run -n | awk '{print $NF}'`; \
+        cp $redisBloomPath /usr/local/lib/
 # create working directories and organize files
 RUN \
         mkdir /data && chown keydb:keydb /data; \

pkg/docker/Dockerfile

@@ -53,7 +53,7 @@ RUN set -eux; \
         grep -E '^ *createBoolConfig[(]"protected-mode",.*, *1 *,.*[)],$' ./src/config.cpp; \
         sed -ri 's!^( *createBoolConfig[(]"protected-mode",.*, *)1( *,.*[)],)$!\10\2!' ./src/config.cpp; \
         grep -E '^ *createBoolConfig[(]"protected-mode",.*, *0 *,.*[)],$' ./src/config.cpp; \
-        make -j$(nproc) BUILD_TLS=yes; \
+        make -j$(nproc) BUILD_TLS=yes ENABLE_FLASH=yes; \
         cd src; \
         strip keydb-cli keydb-benchmark keydb-check-rdb keydb-check-aof keydb-diagnostic-tool keydb-sentinel keydb-server; \
         mv keydb-server keydb-cli keydb-benchmark keydb-check-rdb keydb-check-aof keydb-diagnostic-tool keydb-sentinel /usr/local/bin/; \
@@ -83,6 +83,7 @@ RUN \
         sed -i 's/^\(logfile .*\)$/# \1/' /etc/keydb/keydb.conf; \
         sed -i 's/protected-mode yes/protected-mode no/g' /etc/keydb/keydb.conf; \
         sed -i 's/^\(bind .*\)$/# \1/' /etc/keydb/keydb.conf; \
+        cd /usr/local/bin; \
         ln -s keydb-cli redis-cli; \
         cd /etc/keydb; \
         ln -s keydb.conf redis.conf; \

pkg/docker/Dockerfile_Alpine

@@ -5,7 +5,7 @@ RUN mkdir -p /etc/keydb
 ARG BRANCH
 RUN set -eux; \
         \
-        apk add --no-cache su-exec; \
+        apk add --no-cache su-exec tini; \
         apk add --no-cache --virtual .build-deps \
                 coreutils \
                 gcc \
@@ -32,7 +32,7 @@ RUN set -eux; \
         grep -E '^ *createBoolConfig[(]"protected-mode",.*, *1 *,.*[)],$' ./src/config.cpp; \
         sed -ri 's!^( *createBoolConfig[(]"protected-mode",.*, *)1( *,.*[)],)$!\10\2!' ./src/config.cpp; \
         grep -E '^ *createBoolConfig[(]"protected-mode",.*, *0 *,.*[)],$' ./src/config.cpp; \
-        make -j$(nproc) BUILD_TLS=yes; \
+        make -j$(nproc) BUILD_TLS=yes ENABLE_FLASH=yes; \
         cd src; \
         strip keydb-cli keydb-benchmark keydb-check-rdb keydb-check-aof keydb-diagnostic-tool keydb-sentinel keydb-server; \
         mv keydb-server keydb-cli keydb-benchmark keydb-check-rdb keydb-check-aof keydb-diagnostic-tool keydb-sentinel /usr/local/bin/; \
@@ -54,6 +54,7 @@ RUN set -eux; \
         sed -i 's/^\(logfile .*\)$/# \1/' /etc/keydb/keydb.conf; \
         sed -i 's/protected-mode yes/protected-mode no/g' /etc/keydb/keydb.conf; \
         sed -i 's/^\(bind .*\)$/# \1/' /etc/keydb/keydb.conf; \
+        cd /usr/local/bin; \
         ln -s keydb-cli redis-cli; \
         cd /etc/keydb; \
         ln -s keydb.conf redis.conf; \
@@ -76,6 +77,6 @@ RUN set -eux; \
         chmod +x /usr/local/bin/docker-entrypoint.sh
 VOLUME /data
 WORKDIR /data
-ENTRYPOINT ["docker-entrypoint.sh"]
+ENTRYPOINT ["tini", "--", "docker-entrypoint.sh"]
 EXPOSE 6379
 CMD ["keydb-server", "/etc/keydb/keydb.conf"]

src/IStorage.h

@@ -1,10 +1,25 @@
 #pragma once
 #include <functional>
+#include <unordered_set>
 #include "sds.h"
+#include "ae.h"
 #include <string>
 
 #define METADATA_DB_IDENTIFIER "c299fde0-6d42-4ec4-b939-34f680ffe39f"
 
+struct StorageToken {
+    enum class TokenType {
+        SingleRead,
+        SingleWrite,
+        Delete,
+        BatchWrite,
+    };
+    TokenType type;
+    std::unordered_set<struct client *> setc;
+    struct redisDbPersistentData *db;
+    virtual ~StorageToken() {}
+};
+
 class IStorageFactory
 {
 public:
@@ -36,6 +51,12 @@ class IStorage
     virtual bool enumerate_hashslot(callback fn, unsigned int hashslot) const = 0;
     virtual size_t count() const = 0;
 
+    virtual StorageToken *begin_retrieve(struct aeEventLoop *, aePostFunctionTokenProc, sds *, size_t) {return nullptr;};
+    virtual void complete_retrieve(StorageToken * /*tok*/, callbackSingle /*fn*/) {};
+
+    virtual StorageToken* begin_endWriteBatch(struct aeEventLoop *, aePostFunctionTokenProc*) { return nullptr; } // NOP
+    virtual void complete_endWriteBatch(StorageToken * /*tok*/) {};
+
     virtual void bulkInsert(char **rgkeys, size_t *rgcbkeys, char **rgvals, size_t *rgcbvals, size_t celem) {
         beginWriteBatch();
         for (size_t ielem = 0; ielem < celem; ++ielem) {

src/Makefile

@@ -73,10 +73,10 @@ ifneq ($(strip $(SANITIZE)),)
 endif
 
 ifeq ($(ENABLE_FLASH),yes)
+	FINAL_LIBS+= ../deps/rocksdb/librocksdb.a
 	FINAL_LIBS+= -lz -lcrypto -lbz2 -lzstd -llz4 -lsnappy
 	CXXFLAGS+= -I../deps/rocksdb/include/ -DENABLE_ROCKSDB
 	STORAGE_OBJ+= storage/rocksdb.o storage/rocksdbfactory.o
-	FINAL_LIBS+= ../deps/rocksdb/librocksdb.a 
 	DEPENDENCY_TARGETS+= rocksdb
 endif
 

src/SnapshotPayloadParseState.cpp

@@ -233,11 +233,6 @@ void SnapshotPayloadParseState::trimState() {
 
     if (stackParse.empty()) {
         flushQueuedKeys();
-        while (*insertsInFlight > 0) {
-            // TODO: ProcessEventsWhileBlocked
-            aeReleaseLock();
-            aeAcquireLock();
-        }
     }
 }
 

src/SnapshotPayloadParseState.h

@@ -63,4 +63,5 @@ class SnapshotPayloadParseState {
     void pushValue(const char *rgch, long long cch);
     void pushValue(long long value);
     bool shouldThrottle() const { return *insertsInFlight > (cserver.cthreads*4); }
+    bool hasIOInFlight() const { return *insertsInFlight > 0; }
 };

src/StorageCache.cpp

@@ -208,6 +208,30 @@ void StorageCache::retrieve(sds key, IStorage::callbackSingle fn) const
     m_spstorage->retrieve(key, sdslen(key), fn);
 }
 
+StorageToken *StorageCache::begin_retrieve(struct aeEventLoop *el, aePostFunctionTokenProc proc, sds *rgkey, size_t ckey) {
+    std::unique_lock<fastlock> ul(m_lock);
+    if (m_pdict != nullptr)
+    {
+        bool fAnyKey = false;
+        for (size_t ik = 0; ik < ckey; ++ik) {
+            uint64_t hash = dictSdsHash(rgkey[ik]);
+            dictEntry *de = dictFind(m_pdict, reinterpret_cast<void*>(hash));
+
+            if (de != nullptr)
+                fAnyKey = true;
+        }
+        if (!fAnyKey)
+            return nullptr; // All keys are missing - skip the io
+    }
+    ul.unlock();
+
+    return m_spstorage->begin_retrieve(el, proc, rgkey, ckey);
+}
+
+void StorageCache::complete_retrieve(StorageToken *tok, IStorage::callbackSingle fn) {
+    m_spstorage->complete_retrieve(tok, fn);
+}
+
 size_t StorageCache::count() const
 {
     std::unique_lock<fastlock> ul(m_lock, std::defer_lock);
@@ -233,4 +257,5 @@ void StorageCache::emergencyFreeCache() {
             dictRelease(d);
         });
     }
-}
+}
+

src/StorageCache.h

@@ -43,6 +43,10 @@ class StorageCache
     void insert(sds key, const void *data, size_t cbdata, bool fOverwrite);
     void bulkInsert(char **rgkeys, size_t *rgcbkeys, char **rgvals, size_t *rgcbvals, size_t celem);
     void retrieve(sds key, IStorage::callbackSingle fn) const;
+    StorageToken *begin_retrieve(struct aeEventLoop *el, aePostFunctionTokenProc proc, sds *rgkey, size_t ckey);
+    void complete_retrieve(StorageToken *tok, IStorage::callbackSingle fn);
+    StorageToken* begin_endWriteBatch(struct aeEventLoop *el, aePostFunctionTokenProc* proc) {return m_spstorage->begin_endWriteBatch(el,proc);} // NOP
+    void complete_endWriteBatch(StorageToken *tok) {m_spstorage->complete_endWriteBatch(tok);};
     bool erase(sds key);
     void emergencyFreeCache();
     bool keycacheIsEnabled() const { return m_pdict != nullptr; }

src/ae.cpp

@@ -113,6 +113,8 @@ enum class AE_ASYNC_OP
     PostCppFunction,
     DeleteFileEvent,
     CreateFileEvent,
+    PostAsynDBFunction,
+
 };
 
 struct aeCommand
@@ -125,6 +127,8 @@ struct aeCommand
         aePostFunctionProc *proc;
         aeFileProc *fproc;
         std::function<void()> *pfn;
+        aePostFunctionTokenProc* tproc;
+
     };
     void *clientData;
 };
@@ -175,6 +179,17 @@ void aeProcessCmd(aeEventLoop *eventLoop, int fd, void *, int )
             (*cmd.pfn)();
 
             delete cmd.pfn;
+            break;
+        }
+        case AE_ASYNC_OP::PostAsynDBFunction:
+        {   //added to support async api IStorage
+           if (cmd.fLock && !ulock.owns_lock()) {
+                g_forkLock.releaseRead();
+                ulock.lock();
+                g_forkLock.acquireRead();
+            }  
+            ((aePostFunctionTokenProc*)cmd.tproc)(eventLoop,(StorageToken*)cmd.clientData);
+            break; 
         }
             break;
         }
@@ -256,6 +271,20 @@ int aePostFunction(aeEventLoop *eventLoop, aePostFunctionProc *proc, void *arg)
     return AE_OK;
 }
 
+int aePostFunction(aeEventLoop *eventLoop, aePostFunctionTokenProc *proc, StorageToken *token)
+{
+  //added to support async api IStorage
+    aeCommand cmd = {};
+    cmd.op = AE_ASYNC_OP::PostAsynDBFunction; 
+    cmd.tproc = proc;
+    cmd.clientData = (void*)token; 
+    cmd.fLock = false;
+    auto size = write(eventLoop->fdCmdWrite, &cmd, sizeof(cmd));
+    if (size != sizeof(cmd))
+        return AE_ERR;
+    return AE_OK;
+}
+
 int aePostFunction(aeEventLoop *eventLoop, std::function<void()> fn, bool fLock, bool fForceQueue)
 {
     if (eventLoop == g_eventLoopThisThread && !fForceQueue)

src/ae.h

@@ -72,13 +72,15 @@ extern "C" {
 #define AE_NOTUSED(V) ((void) V)
 
 struct aeEventLoop;
-
 /* Types and data structures */
 typedef void aeFileProc(struct aeEventLoop *eventLoop, int fd, void *clientData, int mask);
 typedef int aeTimeProc(struct aeEventLoop *eventLoop, long long id, void *clientData);
 typedef void aeEventFinalizerProc(struct aeEventLoop *eventLoop, void *clientData);
 typedef void aeBeforeSleepProc(struct aeEventLoop *eventLoop);
 typedef void aePostFunctionProc(void *pvArgs);
+//added to support async api IStorage
+struct StorageToken;
+typedef void aePostFunctionTokenProc(struct aeEventLoop *el, struct StorageToken *token);
 
 /* File event structure */
 typedef struct aeFileEvent {
@@ -134,6 +136,8 @@ int aePostFunction(aeEventLoop *eventLoop, aePostFunctionProc *proc, void *arg);
 #ifdef __cplusplus
 }   // EXTERN C
 int aePostFunction(aeEventLoop *eventLoop, std::function<void()> fn, bool fLock = true, bool fForceQueue = false);
+//added to support async api IStorage
+int aePostFunction(aeEventLoop *eventLoop, aePostFunctionTokenProc *proc, StorageToken *token);
 extern "C" {
 #endif
 void aeDeleteEventLoop(aeEventLoop *eventLoop);

src/blocked.cpp

@@ -92,7 +92,8 @@ void blockClient(client *c, int btype) {
     /* Master client should never be blocked unless pause or module */
     serverAssert(!(c->flags & CLIENT_MASTER &&
                    btype != BLOCKED_MODULE &&
-                   btype != BLOCKED_PAUSE));
+                   btype != BLOCKED_PAUSE &&
+                   btype != BLOCKED_STORAGE));
 
     c->flags |= CLIENT_BLOCKED;
     c->btype = btype;
@@ -207,6 +208,8 @@ void unblockClient(client *c) {
     } else if (c->btype == BLOCKED_PAUSE) {
         listDelNode(g_pserver->paused_clients,c->paused_list_node);
         c->paused_list_node = NULL;
+    } else if (c->btype == BLOCKED_STORAGE) {
+        serverTL->vecclientsProcess.push_back(c);
     } else {
         serverPanic("Unknown btype in unblockClient().");
     }