Test ClusterLuaScript against a real Redis cluster.

service/pom.xml

@@ -172,6 +172,13 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.signal</groupId>
+            <artifactId>embedded-redis</artifactId>
+            <version>0.8.1</version>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>com.fasterxml.uuid</groupId>
             <artifactId>java-uuid-generator</artifactId>

service/src/main/java/org/whispersystems/textsecuregcm/redis/ClusterLuaScript.java

@@ -1,5 +1,6 @@
 package org.whispersystems.textsecuregcm.redis;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.lettuce.core.RedisNoScriptException;
 import io.lettuce.core.ScriptOutputType;
 import io.lettuce.core.api.sync.RedisCommands;
@@ -36,7 +37,8 @@ public class ClusterLuaScript {
         }
     }
 
-    private ClusterLuaScript(final FaultTolerantRedisCluster redisCluster, final String script, final ScriptOutputType scriptOutputType) {
+    @VisibleForTesting
+    ClusterLuaScript(final FaultTolerantRedisCluster redisCluster, final String script, final ScriptOutputType scriptOutputType) {
         this.redisCluster     = redisCluster;
         this.scriptOutputType = scriptOutputType;
         this.script           = script;

service/src/test/java/org/whispersystems/textsecuregcm/redis/AbstractRedisClusterTest.java

@@ -0,0 +1,152 @@
+package org.whispersystems.textsecuregcm.redis;
+
+import io.lettuce.core.RedisClient;
+import io.lettuce.core.RedisURI;
+import io.lettuce.core.api.StatefulRedisConnection;
+import io.lettuce.core.api.sync.RedisCommands;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
+import redis.embedded.RedisServer;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.ServerSocket;
+import java.net.URISyntaxException;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static org.junit.Assume.assumeFalse;
+
+/**
+ * An abstract base class that assembles a real (local!) Redis cluster and provides a client to that cluster for
+ * subclasses.
+ */
+public abstract class AbstractRedisClusterTest {
+
+    private static final int MAX_SLOT   = 16384;
+    private static final int NODE_COUNT = 2;
+
+    private static RedisServer[] clusterNodes;
+
+    private FaultTolerantRedisCluster redisCluster;
+
+    @BeforeClass
+    public static void setUpBeforeClass() throws IOException, URISyntaxException, InterruptedException {
+        assumeFalse(System.getProperty("os.name").equalsIgnoreCase("windows"));
+
+        clusterNodes = new RedisServer[NODE_COUNT];
+
+        for (int i = 0; i < NODE_COUNT; i++) {
+            clusterNodes[i] = buildClusterNode(getNextPort());
+            clusterNodes[i].start();
+        }
+
+        assembleCluster(clusterNodes);
+    }
+
+    @Before
+    public void setUp() {
+        final List<String> urls = Arrays.stream(clusterNodes)
+                                        .map(node -> String.format("redis://127.0.0.1:%d", node.ports().get(0)))
+                                        .collect(Collectors.toList());
+
+        redisCluster = new FaultTolerantRedisCluster("test-cluster", urls, Duration.ofSeconds(2), new CircuitBreakerConfiguration());
+    }
+
+    protected FaultTolerantRedisCluster getRedisCluster() {
+        return redisCluster;
+    }
+
+    @After
+    public void tearDown() {
+        redisCluster.stop();
+    }
+
+    @AfterClass
+    public static void tearDownAfterClass() {
+        for (final RedisServer node : clusterNodes) {
+            node.stop();
+        }
+    }
+
+    private static RedisServer buildClusterNode(final int port) throws IOException, URISyntaxException {
+        final File clusterConfigFile     = File.createTempFile("redis", ".conf");
+        final File rdbFile               = File.createTempFile("redis", ".rdb");
+
+        // Redis struggles with existing-but-empty RDB files
+        rdbFile.delete();
+        rdbFile.deleteOnExit();
+        clusterConfigFile.deleteOnExit();
+
+        return RedisServer.builder()
+                .setting("cluster-enabled yes")
+                .setting("cluster-config-file " + clusterConfigFile.getAbsolutePath())
+                .setting("cluster-node-timeout 5000")
+                .setting("dir " + System.getProperty("java.io.tmpdir"))
+                .setting("dbfilename " + rdbFile.getName())
+                .port(port)
+                .build();
+    }
+
+    private static void assembleCluster(final RedisServer... nodes) throws InterruptedException {
+        final RedisClient meetClient = RedisClient.create(RedisURI.create("127.0.0.1", nodes[0].ports().get(0)));
+
+        try {
+            final StatefulRedisConnection<String, String> connection = meetClient.connect();
+            final RedisCommands<String, String>           commands   = connection.sync();
+
+            for (int i = 1; i < nodes.length; i++) {
+                commands.clusterMeet("127.0.0.1", nodes[i].ports().get(0));
+            }
+        } finally {
+            meetClient.shutdown();
+        }
+
+        final int slotsPerNode = MAX_SLOT / nodes.length;
+
+        for (int i = 0; i < nodes.length; i++) {
+            final int startInclusive = i * slotsPerNode;
+            final int endExclusive   = i == nodes.length - 1 ? MAX_SLOT : (i + 1) * slotsPerNode;
+
+            final RedisClient assignSlotClient = RedisClient.create(RedisURI.create("127.0.0.1", nodes[0].ports().get(0)));
+
+            try {
+                final int[] slots = new int[endExclusive - startInclusive];
+
+                for (int s = startInclusive; s < endExclusive; s++) {
+                    slots[s - startInclusive] = s;
+                }
+
+                assignSlotClient.connect().sync().clusterAddSlots(slots);
+            } finally {
+                assignSlotClient.shutdown();
+            }
+        }
+
+        final RedisClient                             waitClient = RedisClient.create(RedisURI.create("127.0.0.1", nodes[0].ports().get(0)));
+        final StatefulRedisConnection<String, String> connection = waitClient.connect();
+
+        try {
+            // CLUSTER INFO gives us a big blob of key-value pairs, but the one we're interested in is `cluster_state`.
+            // According to https://redis.io/commands/cluster-info, `cluster_state:ok` means that the node is ready to
+            // receive queries, all slots are assigned, and a majority of master nodes are reachable.
+            while (!connection.sync().clusterInfo().contains("cluster_state:ok")) {
+                Thread.sleep(500);
+            }
+        } finally {
+            waitClient.shutdown();
+        }
+    }
+
+    private static int getNextPort() throws IOException {
+        try (ServerSocket socket = new ServerSocket(0)) {
+            socket.setReuseAddress(false);
+            return socket.getLocalPort();
+        }
+    }
+}

service/src/test/java/org/whispersystems/textsecuregcm/redis/ClusterLuaScriptTest.java

@@ -0,0 +1,106 @@
+package org.whispersystems.textsecuregcm.redis;
+
+import io.lettuce.core.RedisNoScriptException;
+import io.lettuce.core.ScriptOutputType;
+import io.lettuce.core.api.sync.RedisCommands;
+import io.lettuce.core.cluster.SlotHash;
+import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
+import io.lettuce.core.cluster.models.partitions.RedisClusterNode;
+import org.junit.Test;
+import org.whispersystems.textsecuregcm.tests.util.RedisClusterHelper;
+
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class ClusterLuaScriptTest extends AbstractRedisClusterTest {
+
+    @Test
+    public void testExecuteMovedKey() {
+        final String key   = "key";
+        final String value = "value";
+
+        final FaultTolerantRedisCluster redisCluster = getRedisCluster();
+
+        final ClusterLuaScript script = new ClusterLuaScript(redisCluster, "return redis.call(\"SET\", KEYS[1], ARGV[1])", ScriptOutputType.VALUE);
+
+        assertEquals("OK", script.execute(List.of(key), List.of("value")));
+        assertEquals(value, redisCluster.withReadCluster(connection -> connection.sync().get(key)));
+
+        final int    slot            = SlotHash.getSlot(key);
+
+        final int                           sourcePort          = redisCluster.withWriteCluster(connection -> connection.sync().nodes(node -> node.hasSlot(slot) && node.is(RedisClusterNode.NodeFlag.MASTER)).node(0).getUri().getPort());
+        final RedisCommands<String, String> sourceCommands      = redisCluster.withWriteCluster(connection -> connection.sync().nodes(node -> node.hasSlot(slot) && node.is(RedisClusterNode.NodeFlag.MASTER)).commands(0));
+        final RedisCommands<String, String> destinationCommands = redisCluster.withWriteCluster(connection -> connection.sync().nodes(node -> !node.hasSlot(slot) && node.is(RedisClusterNode.NodeFlag.MASTER)).commands(0));
+
+        destinationCommands.clusterSetSlotImporting(slot, sourceCommands.clusterMyId());
+
+        assertEquals("OK", script.execute(List.of(key), List.of("value")));
+        assertEquals(value, redisCluster.withReadCluster(connection -> connection.sync().get(key)));
+
+        sourceCommands.clusterSetSlotMigrating(slot, destinationCommands.clusterMyId());
+
+        assertEquals("OK", script.execute(List.of(key), List.of("value")));
+        assertEquals(value, redisCluster.withReadCluster(connection -> connection.sync().get(key)));
+
+        for (final String migrateKey : sourceCommands.clusterGetKeysInSlot(slot, Integer.MAX_VALUE)) {
+            destinationCommands.migrate("127.0.0.1", sourcePort, migrateKey, 0, 1000);
+        }
+
+        assertEquals("OK", script.execute(List.of(key), List.of("value")));
+        assertEquals(value, redisCluster.withReadCluster(connection -> connection.sync().get(key)));
+
+        destinationCommands.clusterSetSlotNode(slot, destinationCommands.clusterMyId());
+
+        assertEquals("OK", script.execute(List.of(key), List.of("value")));
+        assertEquals(value, redisCluster.withReadCluster(connection -> connection.sync().get(key)));
+    }
+
+    @Test
+    public void testExecute() {
+        final RedisAdvancedClusterCommands<String, String> commands    = mock(RedisAdvancedClusterCommands.class);
+        final FaultTolerantRedisCluster                    mockCluster = RedisClusterHelper.buildMockRedisCluster(commands);
+
+        final String           script           = "return redis.call(\"SET\", KEYS[1], ARGV[1])";
+        final String           sha              = "abc123";
+        final ScriptOutputType scriptOutputType = ScriptOutputType.VALUE;
+        final List<String>     keys             = List.of("key");
+        final List<String>     values           = List.of("value");
+
+        when(commands.scriptLoad(script)).thenReturn(sha);
+        when(commands.evalsha(any(), any(), any(), any())).thenReturn("OK");
+
+        new ClusterLuaScript(mockCluster, script, scriptOutputType).execute(keys, values);
+
+        verify(commands).scriptLoad(script);
+        verify(commands).evalsha(sha, scriptOutputType, keys.toArray(new String[0]), values.toArray(new String[0]));
+    }
+
+    @Test
+    public void testExecuteNoScriptException() {
+        final RedisAdvancedClusterCommands<String, String> commands    = mock(RedisAdvancedClusterCommands.class);
+        final FaultTolerantRedisCluster                    mockCluster = RedisClusterHelper.buildMockRedisCluster(commands);
+
+        final String           script           = "return redis.call(\"SET\", KEYS[1], ARGV[1])";
+        final String           sha              = "abc123";
+        final ScriptOutputType scriptOutputType = ScriptOutputType.VALUE;
+        final List<String>     keys             = List.of("key");
+        final List<String>     values           = List.of("value");
+
+        when(commands.scriptLoad(script)).thenReturn(sha);
+        when(commands.evalsha(any(), any(), any(), any()))
+            .thenThrow(new RedisNoScriptException("OH NO"))
+            .thenReturn("OK");
+
+        new ClusterLuaScript(mockCluster, script, scriptOutputType).execute(keys, values);
+
+        verify(commands, times(2)).scriptLoad(script);
+        verify(commands, times(2)).evalsha(sha, scriptOutputType, keys.toArray(new String[0]), values.toArray(new String[0]));
+    }
+}

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/RedisClusterHelper.java

@@ -0,0 +1,44 @@
+package org.whispersystems.textsecuregcm.tests.util;
+
+import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
+import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
+
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class RedisClusterHelper {
+
+    @SuppressWarnings("unchecked")
+    public static FaultTolerantRedisCluster buildMockRedisCluster(final RedisAdvancedClusterCommands<String, String> commands) {
+        final FaultTolerantRedisCluster                      cluster    = mock(FaultTolerantRedisCluster.class);
+        final StatefulRedisClusterConnection<String, String> connection = mock(StatefulRedisClusterConnection.class);
+
+        when(connection.sync()).thenReturn(commands);
+
+        when(cluster.withReadCluster(any(Function.class))).thenAnswer(invocation -> {
+            return invocation.getArgument(0, Function.class).apply(connection);
+        });
+        
+        doAnswer(invocation -> {
+            invocation.getArgument(0, Consumer.class).accept(connection);
+            return null;
+        }).when(cluster).useReadCluster(any(Consumer.class));
+
+        when(cluster.withWriteCluster(any(Function.class))).thenAnswer(invocation -> {
+            return invocation.getArgument(0, Function.class).apply(connection);
+        });
+
+        doAnswer(invocation -> {
+            invocation.getArgument(0, Consumer.class).accept(connection);
+            return null;
+        }).when(cluster).useWriteCluster(any(Consumer.class));
+        
+        return cluster;
+    }
+}