feat: instrument RedisCluster and ClusterPipeline for record/replay (#91)

Author: sohankshirsagar

drift/instrumentation/redis/e2e-tests/docker-compose.yml

@@ -7,6 +7,46 @@ services:
       timeout: 5s
       retries: 5
 
+  redis-node-1:
+    image: redis:7-alpine
+    command: redis-server --cluster-enabled yes --cluster-config-file nodes.conf --cluster-node-timeout 5000 --cluster-announce-hostname redis-node-1 --appendonly yes
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+  redis-node-2:
+    image: redis:7-alpine
+    command: redis-server --cluster-enabled yes --cluster-config-file nodes.conf --cluster-node-timeout 5000 --cluster-announce-hostname redis-node-2 --appendonly yes
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+  redis-node-3:
+    image: redis:7-alpine
+    command: redis-server --cluster-enabled yes --cluster-config-file nodes.conf --cluster-node-timeout 5000 --cluster-announce-hostname redis-node-3 --appendonly yes
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+  redis-cluster-init:
+    image: redis:7-alpine
+    depends_on:
+      redis-node-1:
+        condition: service_healthy
+      redis-node-2:
+        condition: service_healthy
+      redis-node-3:
+        condition: service_healthy
+    command: >
+      sh -c "redis-cli --cluster create redis-node-1:6379 redis-node-2:6379 redis-node-3:6379 --cluster-replicas 0 --cluster-yes"
+    restart: "no"
+
   app:
     build:
       context: ../../../..
@@ -21,10 +61,14 @@ services:
     depends_on:
       redis:
         condition: service_healthy
+      redis-cluster-init:
+        condition: service_completed_successfully
     environment:
       - PORT=8000
       - REDIS_HOST=redis
       - REDIS_PORT=6379
+      - REDIS_CLUSTER_HOST=redis-node-1
+      - REDIS_CLUSTER_PORT=6379
       - TUSK_ANALYTICS_DISABLED=1
       - PYTHONUNBUFFERED=1
       - TUSK_USE_RUST_CORE=${TUSK_USE_RUST_CORE:-0}

drift/instrumentation/redis/e2e-tests/src/app.py

@@ -15,11 +15,27 @@
 
 app = Flask(__name__)
 
-# Initialize Redis client
+# Initialize Redis client (standalone)
 redis_client = redis.Redis(
     host=os.getenv("REDIS_HOST", "redis"), port=int(os.getenv("REDIS_PORT", "6379")), db=0, decode_responses=True
 )
 
+# Lazy-initialized RedisCluster client.
+# Initialized on first use so that cluster connection failures don't prevent
+# the rest of the app from starting (avoids breaking non-cluster tests).
+_cluster_client = None
+
+
+def get_cluster_client():
+    global _cluster_client
+    if _cluster_client is None:
+        _cluster_client = redis.RedisCluster(
+            host=os.getenv("REDIS_CLUSTER_HOST", "redis-node-1"),
+            port=int(os.getenv("REDIS_CLUSTER_PORT", "6379")),
+            decode_responses=True,
+        )
+    return _cluster_client
+
 
 @app.route("/health")
 def health():
@@ -245,6 +261,85 @@ def test_transaction_watch():
         return jsonify({"error": str(e)}), 500
 
 
+@app.route("/test/cluster-set-get", methods=["GET"])
+def test_cluster_set_get():
+    """Test SET/GET through RedisCluster.
+
+    RedisCluster.execute_command is a separate method from Redis.execute_command
+    """
+    try:
+        cluster = get_cluster_client()
+        cluster.set("test:cluster:key1", "cluster_value1")
+        cluster.set("test:cluster:key2", "cluster_value2")
+        val1 = cluster.get("test:cluster:key1")
+        val2 = cluster.get("test:cluster:key2")
+        cluster.delete("test:cluster:key1", "test:cluster:key2")
+        return jsonify({"success": True, "val1": val1, "val2": val2})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/test/cluster-incr", methods=["GET"])
+def test_cluster_incr():
+    """Test INCR through RedisCluster.
+
+    Exercises the cluster's execute_command path for a simple write operation.
+    """
+    try:
+        cluster = get_cluster_client()
+        cluster.set("test:cluster:counter", "0")
+        cluster.incr("test:cluster:counter")
+        cluster.incr("test:cluster:counter")
+        cluster.incr("test:cluster:counter")
+        val = cluster.get("test:cluster:counter")
+        cluster.delete("test:cluster:counter")
+        return jsonify({"success": True, "value": int(val)})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/test/cluster-pipeline", methods=["GET"])
+def test_cluster_pipeline():
+    """Test pipeline through RedisCluster (ClusterPipeline).
+
+    ClusterPipeline.execute_command is also a separate path.
+    All keys use the same hash tag {cp} to ensure they land on the same slot,
+    which is required for cluster pipelines.
+    """
+    try:
+        cluster = get_cluster_client()
+        pipe = cluster.pipeline()
+        pipe.set("{cp}:key1", "pipe_val1")
+        pipe.set("{cp}:key2", "pipe_val2")
+        pipe.get("{cp}:key1")
+        pipe.get("{cp}:key2")
+        results = pipe.execute()
+        cluster.delete("{cp}:key1", "{cp}:key2")
+        return jsonify({"success": True, "results": results})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/test/cluster-pipeline-transaction", methods=["GET"])
+def test_cluster_pipeline_transaction():
+    """Test ClusterPipeline with transaction mode.
+
+    Uses TransactionStrategy internally. All keys must be on the same slot.
+    """
+    try:
+        cluster = get_cluster_client()
+        pipe = cluster.pipeline(transaction=True)
+        pipe.set("{tx}:key1", "txval1")
+        pipe.set("{tx}:key2", "txval2")
+        pipe.get("{tx}:key1")
+        pipe.get("{tx}:key2")
+        results = pipe.execute()
+        cluster.delete("{tx}:key1", "{tx}:key2")
+        return jsonify({"success": True, "results": results})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
 if __name__ == "__main__":
     sdk.mark_app_as_ready()
     app.run(host="0.0.0.0", port=8000, debug=False)

drift/instrumentation/redis/e2e-tests/src/test_requests.py

@@ -44,4 +44,12 @@
 
     make_request("GET", "/test/transaction-watch")
 
+    # RedisCluster operations
+    make_request("GET", "/test/cluster-set-get")
+    make_request("GET", "/test/cluster-incr")
+    make_request("GET", "/test/cluster-pipeline")
+
+    # RedisCluster
+    make_request("GET", "/test/cluster-pipeline-transaction")
+
     print_request_summary()