Add test for module diskless short reads

2 files changed, 73 insertions, 2 deletions

diff --git a/tests/modules/testrdb.c b/tests/modules/testrdb.c
index 415497a2f..d73c8bfd3 100644
--- a/tests/modules/testrdb.c
+++ b/tests/modules/testrdb.c
@@ -15,6 +15,8 @@ RedisModuleString *after_str = NULL;
 
 void *testrdb_type_load(RedisModuleIO *rdb, int encver) {
     int count = RedisModule_LoadSigned(rdb);
+    if (RedisModule_IsIOError(rdb))
+        return NULL;
     assert(count==1);
     assert(encver==1);
     RedisModuleString *str = RedisModule_LoadString(rdb);
@@ -57,6 +59,8 @@ int testrdb_aux_load(RedisModuleIO *rdb, int encver, int when) {
             RedisModule_FreeString(ctx, before_str);
         before_str = NULL;
         int count = RedisModule_LoadSigned(rdb);
+        if (RedisModule_IsIOError(rdb))
+            return REDISMODULE_ERR;
         if (count)
             before_str = RedisModule_LoadString(rdb);
     } else {
@@ -64,14 +68,19 @@ int testrdb_aux_load(RedisModuleIO *rdb, int encver, int when) {
             RedisModule_FreeString(ctx, after_str);
         after_str = NULL;
         int count = RedisModule_LoadSigned(rdb);
+        if (RedisModule_IsIOError(rdb))
+            return REDISMODULE_ERR;
         if (count)
             after_str = RedisModule_LoadString(rdb);
     }
+    if (RedisModule_IsIOError(rdb))
+        return REDISMODULE_ERR;
     return REDISMODULE_OK;
 }
 
 void testrdb_type_free(void *value) {
-    RedisModule_FreeString(NULL, (RedisModuleString*)value);
+    if (value)
+        RedisModule_FreeString(NULL, (RedisModuleString*)value);
 }
 
 int testrdb_set_before(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
@@ -171,6 +180,8 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
     if (RedisModule_Init(ctx,"testrdb",1,REDISMODULE_APIVER_1) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
+    RedisModule_SetModuleOptions(ctx, REDISMODULE_OPTIONS_HANDLE_IO_ERRORS);
+
     if (argc > 0)
         RedisModule_StringToLongLong(argv[0], &conf_aux_count);
 
diff --git a/tests/unit/moduleapi/testrdb.tcl b/tests/unit/moduleapi/testrdb.tcl
index 22201a08e..c72570002 100644
--- a/tests/unit/moduleapi/testrdb.tcl
+++ b/tests/unit/moduleapi/testrdb.tcl
@@ -56,7 +56,67 @@ tags "modules" {
         }
     }
 
+    tags {repl} {
+        test {diskless loading short read with module} {
+            start_server [list overrides [list loadmodule "$testmodule"]] {
+                set replica [srv 0 client]
+                set replica_host [srv 0 host]
+                set replica_port [srv 0 port]
+                start_server [list overrides [list loadmodule "$testmodule"]] {
+                    set master [srv 0 client]
+                    set master_host [srv 0 host]
+                    set master_port [srv 0 port]
 
-    # TODO: test short read handling
+                    # Set master and replica to use diskless replication
+                    $master config set repl-diskless-sync yes
+                    $master config set rdbcompression no
+                    $replica config set repl-diskless-load swapdb
+                    for {set k 0} {$k < 30} {incr k} {
+                        r testrdb.set.key key$k [string repeat A [expr {int(rand()*1000000)}]]
+                    }
 
+                    # Start the replication process...
+                    $master config set repl-diskless-sync-delay 0
+                    $replica replicaof $master_host $master_port
+
+                    # kill the replication at various points
+                    set attempts 3
+                    if {$::accurate} { set attempts 10 }
+                    for {set i 0} {$i < $attempts} {incr i} {
+                        # wait for the replica to start reading the rdb
+                        # using the log file since the replica only responds to INFO once in 2mb
+                        wait_for_log_message -1 "*Loading DB in memory*" 5 2000 1
+
+                        # add some additional random sleep so that we kill the master on a different place each time
+                        after [expr {int(rand()*100)}]
+
+                        # kill the replica connection on the master
+                        set killed [$master client kill type replica]
+
+                        if {[catch {
+                            set res [wait_for_log_message -1 "*Internal error in RDB*" 5 100 10]
+                            if {$::verbose} {
+                                puts $res
+                            }
+                        }]} {
+                            puts "failed triggering short read"
+                            # force the replica to try another full sync
+                            $master client kill type replica
+                            $master set asdf asdf
+                            # the side effect of resizing the backlog is that it is flushed (16k is the min size)
+                            $master config set repl-backlog-size [expr {16384 + $i}]
+                        }
+                        # wait for loading to stop (fail)
+                        wait_for_condition 100 10 {
+                            [s -1 loading] eq 0
+                        } else {
+                            fail "Replica didn't disconnect"
+                        }
+                    }
+                    # enable fast shutdown
+                    $master config set rdb-key-save-delay 0
+                }
+            }
+        }
+    }
 }