Merge pull request #17507 from code-dot-org/redis-shard-factory

Internet Simulator: Test Redis shard selection

shared/middleware/helpers/sharded_redis_factory.rb

@@ -0,0 +1,64 @@
+require 'redis-slave-read'
+require 'consistent_hashing'
+
+# Helper for getting a correctly configured Redis client for a particular
+# shard key.  Can be configured for any number of shards, where each shard
+# can be a single Redis node or a replication group where reads will be
+# distributed to read-replicas.
+class ShardedRedisFactory
+  # @param [Hash<'master':String, 'read_replicas':String[]>[]] config
+  # The set of Redis node URLs to be used in the current environment, provided
+  # in the following format:
+  # [
+  #   {
+  #     'master': 'redis://master1',
+  #     'read_replicas': [
+  #       'redis://master1replica1',
+  #       'redis://master1replica2'
+  #     ]
+  #   },
+  #   {
+  #     'master': 'redis://master2',
+  #     'read_replicas': [
+  #       'redis://master2replica1',
+  #       'redis://master2replica2'
+  #     ]
+  #   }
+  # ]
+  #
+  # The read_replicas key is optional and you can specify any number of groups.
+  #
+  # @param [Proc] new_redis_proc (optional) is provided as a hook to allow
+  # us to construct fake Redis instances in test.
+  def initialize(shards, new_redis_proc = proc {|url| Redis.new(url: url)})
+    raise ArgumentError.new('Must provide at least one shard') if shards.empty?
+    @new_redis_proc = new_redis_proc
+    @ring = ConsistentHashing::Ring.new
+    shards.each {|shard_config| @ring.add shard_config}
+  end
+
+  # The set of Redis node URLs to be used for the given shard key.
+  #
+  # @param [String] shard_key
+  # @return [Hash<'master':String, 'read_replicas':String[]>]
+  def client_for_key(shard_key)
+    shard_config = @ring.node_for shard_key
+    client_for_shard_config shard_config
+  end
+
+  # Construct a Redis client for the given redis shard configuration.
+  #
+  # @param [Hash<'master':String, 'read_replicas':String[]>] shard_config
+  # @return [Redis]
+  def client_for_shard_config(shard_config)
+    master_url = shard_config['master']
+    replica_urls = shard_config['read_replicas'] || []
+
+    Redis::SlaveRead::Interface::Hiredis.new(
+      {
+        master: @new_redis_proc[master_url],
+        slaves: replica_urls.map {|url| @new_redis_proc[url]}
+      }
+    )
+  end
+end

shared/middleware/net_sim_api.rb

@@ -4,9 +4,8 @@
 require 'cdo/rack/request'
 require 'cgi'
 require 'csv'
-require 'redis-slave-read'
-require 'consistent_hashing'
 require_relative '../middleware/helpers/redis_table'
+require_relative '../middleware/helpers/sharded_redis_factory'
 require_relative '../middleware/channels_api'
 
 # NetSimApi implements a rest service for interacting with NetSim tables.
@@ -29,6 +28,8 @@ class NetSimApi < Sinatra::Base
     limit_reached: 'limit_reached'
   }
 
+  DEFAULT_LOCAL_REDIS = 'redis://localhost:6379'
+
   helpers do
     %w{
       core.rb
@@ -370,33 +371,7 @@ def self.override_redis_for_test(override_redis)
   # @return [Redis]
   def get_redis_client(shard_id)
     return @@overridden_redis unless @@overridden_redis.nil?
-    new_redis_client_for_group redis_group_for_shard shard_id
-  end
-
-  # Construct a Redis client for the given redis group definition.
-  #
-  # @param [Hash<'master':String, 'read_replicas':String[]>]
-  # @return [Redis]
-  def new_redis_client_for_group(redis_group)
-    master_url = redis_group['master']
-    replica_urls = redis_group['read_replicas'] || []
-
-    Redis::SlaveRead::Interface::Hiredis.new(
-      {
-        master: Redis.new(url: master_url),
-        slaves: replica_urls.map {|url| Redis.new(url: url)}
-      }
-    )
-  end
-
-  # The set of Redis node URLs to be used for the given shard id.
-  #
-  # @param [String] shard_id
-  # @return [Hash<'master':String, 'read_replicas':String[]>]
-  def redis_group_for_shard(shard_id)
-    ring = ConsistentHashing::Ring.new
-    redis_groups.each {|group| ring.add group}
-    ring.node_for shard_id
+    ShardedRedisFactory.new(redis_groups).client_for_key(shard_id)
   end
 
   # The set of Redis node URLs to be used in the current environment.
@@ -420,7 +395,7 @@ def redis_group_for_shard(shard_id)
   #
   # @return [Hash<'master':String, 'read_replicas':String[]>[]]
   def redis_groups
-    CDO.netsim_redis_groups || [{'master': 'redis://localhost:6379'}]
+    CDO.netsim_redis_groups || [{'master' => DEFAULT_LOCAL_REDIS}]
   end
 
   # Get the Pub/Sub API interface for the current configuration

shared/test/middleware/helpers/test_sharded_redis_factory.rb

@@ -0,0 +1,117 @@
+# Unit tests for ShardedRedisFactory
+# This uses a fake in-memory Redis service that does not actually support
+# replication.
+
+require_relative '../../test_helper'
+require 'fakeredis'
+require 'helpers/sharded_redis_factory'
+
+# Monkeypatch Redis client so we can store the URL we used when creating it;
+# FakeRedis always lists a localhost URL when connecting to an in-memory store,
+# which isn't helpful for checking that we passed the right configuration.
+class Redis
+  attr_accessor :url_according_to_test
+end
+
+# Provide an alternate Redis construction Proc that will use FakeRedis (because
+# we imported it above) and store the URL we wanted to pass so we can check it
+# in our tests.
+TEST_NEW_REDIS_PROC = proc do |url|
+  Redis.new(url: url).tap do |fake_redis|
+    fake_redis.url_according_to_test = url
+  end
+end
+
+class ShardedRedisFactoryTest < MiniTest::Test
+  include SetupTest
+
+  def test_raises_if_constructed_with_empty_shards
+    assert_raises ArgumentError do
+      ShardedRedisFactory.new []
+    end
+  end
+
+  def test_single_node_config
+    factory = ShardedRedisFactory.new(
+      [{'master' => 'redis://master'}],
+      TEST_NEW_REDIS_PROC
+    )
+    client = factory.client_for_key('any old key')
+    assert_equal 'redis://master', client.master.url_according_to_test
+  end
+
+  def test_single_replication_group_config
+    factory = ShardedRedisFactory.new(
+      [
+        {
+          'master' => 'redis://master',
+          'read_replicas' => [
+            'redis://replica1',
+            'redis://replica2'
+          ]
+        }
+      ],
+      TEST_NEW_REDIS_PROC
+    )
+    client = factory.client_for_key('does not matter')
+    assert_equal 'redis://master', client.master.url_according_to_test
+    assert_equal 'redis://replica1', client.slaves[0].url_according_to_test
+    assert_equal 'redis://replica2', client.slaves[1].url_according_to_test
+  end
+
+  def test_three_master_config
+    factory = ShardedRedisFactory.new(
+      [
+        {'master' => 'redis://master1'},
+        {'master' => 'redis://master2'},
+        {'master' => 'redis://master3'}
+      ],
+      TEST_NEW_REDIS_PROC
+    )
+
+    # Shard selection is deterministic, so we just found appropriate keys
+    # for this test without pinning any randomness.
+    client = factory.client_for_key('a shard key')
+    assert_equal 'redis://master1', client.master.url_according_to_test
+
+    client = factory.client_for_key('alternate shard key')
+    assert_equal 'redis://master2', client.master.url_according_to_test
+
+    client = factory.client_for_key('a different shard key')
+    assert_equal 'redis://master3', client.master.url_according_to_test
+  end
+
+  def test_two_replication_group_config
+    factory = ShardedRedisFactory.new(
+      [
+        {
+          'master' => 'redis://master1',
+          'read_replicas' => [
+            'redis://replica1_1',
+            'redis://replica1_2'
+          ]
+        },
+        {
+          'master' => 'redis://master2',
+          'read_replicas' => [
+            'redis://replica2_1',
+            'redis://replica2_2'
+          ]
+        }
+      ],
+      TEST_NEW_REDIS_PROC
+    )
+
+    # Shard selection is deterministic, so we just found appropriate keys
+    # for this test without pinning any randomness.
+    client = factory.client_for_key('shard 1')
+    assert_equal 'redis://master1', client.master.url_according_to_test
+    assert_equal 'redis://replica1_1', client.slaves[0].url_according_to_test
+    assert_equal 'redis://replica1_2', client.slaves[1].url_according_to_test
+
+    client = factory.client_for_key('second shard')
+    assert_equal 'redis://master2', client.master.url_according_to_test
+    assert_equal 'redis://replica2_1', client.slaves[0].url_according_to_test
+    assert_equal 'redis://replica2_2', client.slaves[1].url_according_to_test
+  end
+end

shared/test/test_net_sim_api.rb → shared/test/middleware/test_net_sim_api.rb

@@ -5,11 +5,11 @@
 # Caution: This test is destructive, clears the whole Redis instance between
 # tests, so be careful when using real redis.
 
-require_relative 'test_helper'
+require_relative '../test_helper'
 require 'fakeredis' unless use_real_redis?
 require 'redis-slave-read'
 require 'net_sim_api'
-require_relative 'spy_pub_sub_api'
+require_relative '../spy_pub_sub_api'
 
 class NetSimApiTest < Minitest::Test
   include SetupTest