gistfile1.txt

package sharding

import (
	"sync"

	"github.com/golang/groupcache/consistenthash"
	"gopkg.in/redis.v1"
)

const replicas = 100

var noop = redis.NewTCPClient(&redis.Options{
	Addr: "noop",
})

type RedisShard struct {
	cons *consistenthash.Map

	mtx     sync.Mutex
	servers map[string]*redis.Options
	shards  map[string]*redis.Client
}

func NewRedisShard() *RedisShard {
	return &RedisShard{
		cons:    consistenthash.New(replicas, nil),
		servers: make(map[string]*redis.Options),
		shards:  make(map[string]*redis.Client),
	}
}

func (s *RedisShard) Add(server string, opt *redis.Options) {
	s.mtx.Lock()
	s.cons.Add(server)
	s.servers[server] = opt
	s.mtx.Unlock()
}

func (s *RedisShard) getShard(server string) *redis.Client {
	cl, ok := s.shards[server]
	if !ok {
		cl = redis.NewTCPClient(s.servers[server])
		s.shards[server] = cl
	}
	return cl
}

func (s *RedisShard) Get(key string) *redis.Client {
	s.mtx.Lock()
	server := s.cons.Get(key)
	if server == "" {
		return noop
	}
	shard := s.getShard(server)
	s.mtx.Unlock()
	return shard
}