add possibility to set replica weights explicitly (#28)

Author: Artem Klyukvin

rediscluster/cluster.go

@@ -64,6 +64,14 @@ const (
 	enabled  = 1
 )
 
+// WeightProvider explicitly provides weights for redis replicas
+type WeightProvider interface {
+	// GetWeightByHost provides weight by given name. If implementation does not have weight for a given host
+	// it must return `false` as the second return value. In this case scheduler will fallback to default
+	// scheduling strategy (either constant weights random or ping latency based random)
+	GetWeightByHost(host string) (uint32, bool)
+}
+
 // Opts holds the options for Cluster
 type Opts struct {
 	// HostOpts - per host options
@@ -98,8 +106,10 @@ type Opts struct {
 
 	// RoundRobinSeed - used to choose between master and replica.
 	RoundRobinSeed RoundRobinSeed
-	// LatencyOrientedRR - when MasterAndSlaves is used, prefer hosts with lower latency
+	// LatencyOrientedRR - when MasterAndSlaves is used, prefer hosts with lower latency (has lower priority than WeightProvider)
 	LatencyOrientedRR bool
+	// WeightProvider - enables to explicitly set weights of replicas (has higher priority than LatencyOrientedRR)
+	WeightProvider WeightProvider
 	// Enable connection with TLS
 	TLSEnabled bool
 	// Config for TLS connection

rediscluster/mapping.go

@@ -341,6 +341,25 @@ func (*Cluster) getHealthWeight(weights []uint32, health uint32) uint32 {
 }
 
 func (c *Cluster) weightsForPolicySlaves(policy ReplicaPolicyEnum, shard *shard) []uint32 {
+	if c.opts.WeightProvider == nil {
+		return c.weightsForPolicySlavesDefault(policy, shard)
+	}
+
+	weights := make([]uint32, 0, len(shard.addr))
+	for _, addr := range shard.addr {
+		weight, found := c.opts.WeightProvider.GetWeightByHost(addr)
+		if !found {
+			// there was some reconfiguration, so we fallback to default weights
+			return c.weightsForPolicySlavesDefault(policy, shard)
+		}
+
+		weights = append(weights, weight)
+	}
+
+	return weights
+}
+
+func (c *Cluster) weightsForPolicySlavesDefault(policy ReplicaPolicyEnum, shard *shard) []uint32 {
 	var ws []uint32
 	if atomic.LoadUint32(&c.latencyAwareness) == disabled {
 		ws = rr[:]