RDB v9: Save Lua scripts state into RDB file.rdb9

This is currently needed in order to fix #4483, but this can be useful
in other contexts, so maybe later we may want to remove the conditionals
and always save/load scripts.

4 files changed, 56 insertions, 5 deletions

diff --git a/src/rdb.c b/src/rdb.c
index 00106cac4..481260902 100644
--- a/src/rdb.c
+++ b/src/rdb.c
@@ -943,6 +943,23 @@ int rdbSaveRio(rio *rdb, int *error, int flags, rdbSaveInfo *rsi) {
     }
     di = NULL; /* So that we don't release it again on error. */
 
+    /* If we are storing the replication information on disk, persist
+     * the script cache as well: on successful PSYNC after a restart, we need
+     * to be able to process any EVALSHA inside the replication backlog the
+     * master will send us. */
+    if (rsi && dictSize(server.lua_scripts)) {
+        di = dictGetIterator(server.lua_scripts);
+        while((de = dictNext(di)) != NULL) {
+            sds sha = dictGetKey(de);
+            robj *body = dictGetVal(de);
+            if (rdbSaveType(rdb,RDB_OPCODE_SCRIPT) == -1) goto werr;
+            if (rdbSaveRawString(rdb,(unsigned char*)sha,sdslen(sha)) == -1) goto werr;
+            if (rdbSaveRawString(rdb,body->ptr,sdslen(body->ptr)) == -1)
+                goto werr;
+        }
+        dictReleaseIterator(di);
+    }
+
     /* EOF opcode */
     if (rdbSaveType(rdb,RDB_OPCODE_EOF) == -1) goto werr;
 
@@ -1563,6 +1580,33 @@ int rdbLoadRio(rio *rdb, rdbSaveInfo *rsi) {
             dictExpand(db->dict,db_size);
             dictExpand(db->expires,expires_size);
             continue; /* Read type again. */
+        } else if (type == RDB_OPCODE_SCRIPT) {
+            sds sha;
+            robj *body;
+            char funcname[43];
+
+            if ((sha = rdbGenericLoadStringObject(rdb,RDB_LOAD_SDS,NULL))
+                == NULL) goto eoferr;
+            if (sdslen(sha) != 40) {
+                rdbExitReportCorruptRDB(
+                    "Lua script stored into the RDB file has invalid "
+                    "name length: '%s'", sha);
+            }
+            if ((body = rdbLoadStringObject(rdb)) == NULL) goto eoferr;
+
+            /* Compose the Lua function name, and load it. */
+            funcname[0] = 'f';
+            funcname[1] = '_';
+            memcpy(funcname+2,sha,41);
+            if (luaCreateFunction(NULL,server.lua,funcname,body) == C_ERR) {
+                rdbExitReportCorruptRDB(
+                    "Can't load Lua script from RDB file! "
+                    "Script SHA1: %s BODY: %d",
+                        sha, body->ptr);
+            }
+            sdsfree(sha);
+            decrRefCount(body);
+            continue; /* Read type again. */
         } else if (type == RDB_OPCODE_AUX) {
             /* AUX: generic string-string fields. Use to add state to RDB
              * which is backward compatible. Implementations of RDB loading
diff --git a/src/rdb.h b/src/rdb.h
index 62a13f444..2d1a5246f 100644
--- a/src/rdb.h
+++ b/src/rdb.h
@@ -95,6 +95,7 @@
 #define rdbIsObjectType(t) ((t >= 0 && t <= 7) || (t >= 9 && t <= 14))
 
 /* Special RDB opcodes (saved/loaded with rdbSaveType/rdbLoadType). */
+#define RDB_OPCODE_SCRIPT     249
 #define RDB_OPCODE_AUX        250
 #define RDB_OPCODE_RESIZEDB   251
 #define RDB_OPCODE_EXPIRETIME_MS 252
diff --git a/src/scripting.c b/src/scripting.c
index d9f954068..64de1edcd 100644
--- a/src/scripting.c
+++ b/src/scripting.c
@@ -1160,16 +1160,21 @@ int luaCreateFunction(client *c, lua_State *lua, char *funcname, robj *body) {
     funcdef = sdscatlen(funcdef,"\nend",4);
 
     if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"@user_script")) {
-        addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
-            lua_tostring(lua,-1));
+        if (c != NULL) {
+            addReplyErrorFormat(c,
+                "Error compiling script (new function): %s\n",
+                lua_tostring(lua,-1));
+        }
         lua_pop(lua,1);
         sdsfree(funcdef);
         return C_ERR;
     }
     sdsfree(funcdef);
     if (lua_pcall(lua,0,0,0)) {
-        addReplyErrorFormat(c,"Error running script (new function): %s\n",
-            lua_tostring(lua,-1));
+        if (c != NULL) {
+            addReplyErrorFormat(c,"Error running script (new function): %s\n",
+                lua_tostring(lua,-1));
+        }
         lua_pop(lua,1);
         return C_ERR;
     }
@@ -1180,7 +1185,7 @@ int luaCreateFunction(client *c, lua_State *lua, char *funcname, robj *body) {
     {
         int retval = dictAdd(server.lua_scripts,
                              sdsnewlen(funcname+2,40),body);
-        serverAssertWithInfo(c,NULL,retval == DICT_OK);
+        serverAssertWithInfo(c ? c : server.lua_client,NULL,retval == DICT_OK);
         incrRefCount(body);
     }
     return C_OK;
diff --git a/src/server.h b/src/server.h
index e3b56075a..a0c65355b 100644
--- a/src/server.h
+++ b/src/server.h
@@ -1781,6 +1781,7 @@ void scriptingInit(int setup);
 int ldbRemoveChild(pid_t pid);
 void ldbKillForkedSessions(void);
 int ldbPendingChildren(void);
+int luaCreateFunction(client *c, lua_State *lua, char *funcname, robj *body);
 
 /* Blocked clients */
 void processUnblockedClients(void);