Precise timeouts: fast exit for clientsHandleShortTimeout().

AOFRW on an empty stream created with MKSTREAM loads badkly

src/aof.c

@@ -1212,12 +1212,13 @@ int rewriteStreamObject(rio *r, robj *key, robj *o) {
         /* Use the XADD MAXLEN 0 trick to generate an empty stream if
          * the key we are serializing is an empty string, which is possible
          * for the Stream type. */
+        id.ms = 0; id.seq = 1; 
         if (rioWriteBulkCount(r,'*',7) == 0) return 0;
         if (rioWriteBulkString(r,"XADD",4) == 0) return 0;
         if (rioWriteBulkObject(r,key) == 0) return 0;
         if (rioWriteBulkString(r,"MAXLEN",6) == 0) return 0;
         if (rioWriteBulkString(r,"0",1) == 0) return 0;
-        if (rioWriteBulkStreamID(r,&s->last_id) == 0) return 0;
+        if (rioWriteBulkStreamID(r,&id) == 0) return 0;
         if (rioWriteBulkString(r,"x",1) == 0) return 0;
         if (rioWriteBulkString(r,"y",1) == 0) return 0;
     }

src/blocked.c

@@ -111,6 +111,7 @@ void blockClient(client *c, int btype) {
     c->btype = btype;
     server.blocked_clients++;
     server.blocked_clients_by_type[btype]++;
+    addClientToShortTimeoutTable(c);
 }
 
 /* This function is called in the beforeSleep() function of the event loop

src/server.c

@@ -1473,6 +1473,135 @@ int allPersistenceDisabled(void) {
     return server.saveparamslen == 0 && server.aof_state == AOF_OFF;
 }
 
+/* ========================== Clients timeouts ============================= */
+
+/* Check if this blocked client timedout (does nothing if the client is
+ * not blocked right now). If so send a reply, unblock it, and return 1.
+ * Otherwise 0 is returned and no operation is performed. */
+int checkBlockedClientTimeout(client *c, mstime_t now) {
+    if (c->flags & CLIENT_BLOCKED &&
+        c->bpop.timeout != 0
+        && c->bpop.timeout < now)
+    {
+        /* Handle blocking operation specific timeout. */
+        replyToBlockedClientTimedOut(c);
+        unblockClient(c);
+        return 1;
+    } else {
+        return 0;
+    }
+}
+
+/* Check for timeouts. Returns non-zero if the client was terminated.
+ * The function gets the current time in milliseconds as argument since
+ * it gets called multiple times in a loop, so calling gettimeofday() for
+ * each iteration would be costly without any actual gain. */
+int clientsCronHandleTimeout(client *c, mstime_t now_ms) {
+    time_t now = now_ms/1000;
+
+    if (server.maxidletime &&
+        /* This handles the idle clients connection timeout if set. */
+        !(c->flags & CLIENT_SLAVE) &&   /* No timeout for slaves and monitors */
+        !(c->flags & CLIENT_MASTER) &&  /* No timeout for masters */
+        !(c->flags & CLIENT_BLOCKED) && /* No timeout for BLPOP */
+        !(c->flags & CLIENT_PUBSUB) &&  /* No timeout for Pub/Sub clients */
+        (now - c->lastinteraction > server.maxidletime))
+    {
+        serverLog(LL_VERBOSE,"Closing idle client");
+        freeClient(c);
+        return 1;
+    } else if (c->flags & CLIENT_BLOCKED) {
+        /* Blocked OPS timeout is handled with milliseconds resolution.
+         * However note that the actual resolution is limited by
+         * server.hz. So for short timeouts (less than SERVER_SHORT_TIMEOUT
+         * milliseconds) we populate a Radix tree and handle such timeouts
+         * in clientsHandleShortTimeout(). */
+        if (checkBlockedClientTimeout(c,now_ms)) return 0;
+
+        /* Cluster: handle unblock & redirect of clients blocked
+         * into keys no longer served by this server. */
+        if (server.cluster_enabled) {
+            if (clusterRedirectBlockedClientIfNeeded(c))
+                unblockClient(c);
+        }
+    }
+    return 0;
+}
+
+/* For shor timeouts, less than < CLIENT_SHORT_TIMEOUT milliseconds, we
+ * populate a radix tree of 128 bit keys composed as such:
+ *
+ *  [8 byte big endian expire time]+[8 byte client ID]
+ *
+ * We don't do any cleanup in the Radix tree: when we run the clients that
+ * reached the timeout already, if they are no longer existing or no longer
+ * blocked with such timeout, we just go forward.
+ *
+ * Every time a client blocks with a short timeout, we add the client in
+ * the tree. In beforeSleep() we call clientsHandleShortTimeout() to run
+ * the tree and unblock the clients.
+ *
+ * Design hint: why we block only clients with short timeouts? For frugality:
+ * Clients blocking for 30 seconds usually don't need to be unblocked
+ * precisely, and anyway for the nature of Redis to *guarantee* unblock time
+ * precision is hard, so we can avoid putting a large number of clients in
+ * the radix tree without a good reason. This idea also has a role in memory
+ * usage as well given that we don't do cleanup, the shorter a client timeout,
+ * the less time it will stay in the radix tree. */
+
+#define CLIENT_ST_KEYLEN 16    /* 8 bytes mstime + 8 bytes client ID. */
+
+/* Given client ID and timeout, write the resulting radix tree key in buf. */
+void encodeTimeoutKey(unsigned char *buf, uint64_t timeout, uint64_t id) {
+    timeout = htonu64(timeout);
+    memcpy(buf,&timeout,sizeof(timeout));
+    memcpy(buf+8,&id,sizeof(id));
+}
+
+/* Given a key encoded with encodeTimeoutKey(), resolve the fields and write
+ * the timeout into *toptr and the client ID into *idptr. */
+void decodeTimeoutKey(unsigned char *buf, uint64_t *toptr, uint64_t *idptr) {
+    memcpy(toptr,buf,sizeof(*toptr));
+    *toptr = ntohu64(*toptr);
+    memcpy(idptr,buf+8,sizeof(*idptr));
+}
+
+/* Add the specified client id / timeout as a key in the radix tree we use
+ * to handle short timeouts. The client is not added to the list if its
+ * timeout is longer than CLIENT_SHORT_TIMEOUT milliseconds. */
+void addClientToShortTimeoutTable(client *c) {
+    if (c->bpop.timeout == 0 ||
+        c->bpop.timeout - mstime() > CLIENT_SHORT_TIMEOUT)
+    {
+        return;
+    }
+    uint64_t timeout = c->bpop.timeout;
+    uint64_t id = c->id;
+    unsigned char buf[CLIENT_ST_KEYLEN];
+    encodeTimeoutKey(buf,timeout,id);
+    raxTryInsert(server.clients_timeout_table,buf,sizeof(buf),NULL,NULL);
+}
+
+/* This function is called in beforeSleep() in order to unblock ASAP clients
+ * that are waiting in blocking operations with a short timeout set. */
+void clientsHandleShortTimeout(void) {
+    if (raxSize(server.clients_timeout_table) == 0) return;
+    uint64_t now = mstime();
+    raxIterator ri;
+    raxStart(&ri,server.clients_timeout_table);
+    raxSeek(&ri,"^",NULL,0);
+
+    while(raxNext(&ri)) {
+        uint64_t id, timeout;
+        decodeTimeoutKey(ri.key,&timeout,&id);
+        if (timeout >= now) break; /* All the timeouts are in the future. */
+        client *c = lookupClientByID(id);
+        if (c) checkBlockedClientTimeout(c,now);
+        raxRemove(server.clients_timeout_table,ri.key,ri.key_len,NULL);
+        raxSeek(&ri,"^",NULL,0);
+    }
+}
+
 /* ======================= Cron: called every 100 ms ======================== */
 
 /* Add a sample to the operations per second array of samples. */
@@ -1502,42 +1631,6 @@ long long getInstantaneousMetric(int metric) {
     return sum / STATS_METRIC_SAMPLES;
 }
 
-/* Check for timeouts. Returns non-zero if the client was terminated.
- * The function gets the current time in milliseconds as argument since
- * it gets called multiple times in a loop, so calling gettimeofday() for
- * each iteration would be costly without any actual gain. */
-int clientsCronHandleTimeout(client *c, mstime_t now_ms) {
-    time_t now = now_ms/1000;
-
-    if (server.maxidletime &&
-        !(c->flags & CLIENT_SLAVE) &&    /* no timeout for slaves and monitors */
-        !(c->flags & CLIENT_MASTER) &&   /* no timeout for masters */
-        !(c->flags & CLIENT_BLOCKED) &&  /* no timeout for BLPOP */
-        !(c->flags & CLIENT_PUBSUB) &&   /* no timeout for Pub/Sub clients */
-        (now - c->lastinteraction > server.maxidletime))
-    {
-        serverLog(LL_VERBOSE,"Closing idle client");
-        freeClient(c);
-        return 1;
-    } else if (c->flags & CLIENT_BLOCKED) {
-        /* Blocked OPS timeout is handled with milliseconds resolution.
-         * However note that the actual resolution is limited by
-         * server.hz. */
-
-        if (c->bpop.timeout != 0 && c->bpop.timeout < now_ms) {
-            /* Handle blocking operation specific timeout. */
-            replyToBlockedClientTimedOut(c);
-            unblockClient(c);
-        } else if (server.cluster_enabled) {
-            /* Cluster: handle unblock & redirect of clients blocked
-             * into keys no longer served by this server. */
-            if (clusterRedirectBlockedClientIfNeeded(c))
-                unblockClient(c);
-        }
-    }
-    return 0;
-}
-
 /* The client query buffer is an sds.c string that can end with a lot of
  * free space not used, this function reclaims space if needed.
  *
@@ -1654,6 +1747,9 @@ void getExpansiveClientsInfo(size_t *in_usage, size_t *out_usage) {
  */
 #define CLIENTS_CRON_MIN_ITERATIONS 5
 void clientsCron(void) {
+    /* Unblock short timeout clients ASAP. */
+    clientsHandleShortTimeout();
+
     /* Try to process at least numclients/server.hz of clients
      * per call. Since normally (if there are no big latency events) this
      * function is called server.hz times per second, in the average case we
@@ -2092,11 +2188,15 @@ int serverCron(struct aeEventLoop *eventLoop, long long id, void *clientData) {
 void beforeSleep(struct aeEventLoop *eventLoop) {
     UNUSED(eventLoop);
 
+    /* Handle precise timeouts of blocked clients. */
+    clientsHandleShortTimeout();
+
     /* We should handle pending reads clients ASAP after event loop. */
     handleClientsWithPendingReadsUsingThreads();
 
     /* Handle TLS pending data. (must be done before flushAppendOnlyFile) */
     tlsProcessPendingData();
+
     /* If tls still has pending unread data don't sleep at all. */
     aeSetDontWait(server.el, tlsHasPendingData());
 
@@ -2721,6 +2821,7 @@ void initServer(void) {
     server.monitors = listCreate();
     server.clients_pending_write = listCreate();
     server.clients_pending_read = listCreate();
+    server.clients_timeout_table = raxNew();
     server.slaveseldb = -1; /* Force to emit the first SELECT command. */
     server.unblocked_clients = listCreate();
     server.ready_keys = listCreate();

src/server.h

@@ -277,6 +277,9 @@ typedef long long ustime_t; /* microsecond time type. */
                                     buffer configuration. Just the first
                                     three: normal, slave, pubsub. */
 
+/* Other client related defines. */
+#define CLIENT_SHORT_TIMEOUT 2000  /* See clientsHandleShortTimeout(). */
+
 /* Slave replication state. Used in server.repl_state for slaves to remember
  * what to do next. */
 #define REPL_STATE_NONE 0 /* No active replication */
@@ -1067,6 +1070,7 @@ struct redisServer {
     list *clients_pending_read;  /* Client has pending read socket buffers. */
     list *slaves, *monitors;    /* List of slaves and MONITORs */
     client *current_client;     /* Current client executing the command. */
+    rax *clients_timeout_table; /* Radix tree for clients with short timeout. */
     long fixed_time_expire;     /* If > 0, expire keys against server.mstime. */
     rax *clients_index;         /* Active clients dictionary by client ID. */
     int clients_paused;         /* True if clients are currently paused */
@@ -2136,6 +2140,7 @@ void disconnectAllBlockedClients(void);
 void handleClientsBlockedOnKeys(void);
 void signalKeyAsReady(redisDb *db, robj *key);
 void blockForKeys(client *c, int btype, robj **keys, int numkeys, mstime_t timeout, robj *target, streamID *ids);
+void addClientToShortTimeoutTable(client *c);
 
 /* expire.c -- Handling of expired keys */
 void activeExpireCycle(int type);

tests/unit/type/stream-cgroups.tcl

@@ -311,4 +311,17 @@ start_server {
             }
         }
     }
+
+    start_server {tags {"stream"} overrides {appendonly yes aof-use-rdb-preamble no}} {
+        test {Empty stream with no lastid can be rewrite into AOF correctly} {
+            r XGROUP CREATE mystream group-name $ MKSTREAM
+            assert {[dict get [r xinfo stream mystream] length] == 0}
+            set grpinfo [r xinfo groups mystream]
+            r bgrewriteaof
+            waitForBgrewriteaof r
+            r debug loadaof
+            assert {[dict get [r xinfo stream mystream] length] == 0}
+            assert {[r xinfo groups mystream] == $grpinfo}
+        }
+    }
 }