Go sharded client

clients/go/dkv/client.go

@@ -0,0 +1,277 @@
+package dkv
+
+import (
+	"fmt"
+	"github.com/OneOfOne/xxhash"
+	"github.com/dgraph-io/ristretto"
+	"github.com/flipkart-incubator/dkv/pkg/ctl"
+	"github.com/flipkart-incubator/dkv/pkg/serverpb"
+
+	"log"
+)
+
+//DKVServerRole defines the role of the DKV node
+type DKVServerRole string
+
+const (
+	noRole     DKVServerRole = "UNKNOWN"
+	masterRole               = "MASTER"
+	slaveRole                = "SLAVE"
+)
+
+func getNodeTypeByReadConsistency(rc serverpb.ReadConsistency) DKVServerRole {
+	if rc == serverpb.ReadConsistency_LINEARIZABLE {
+		return masterRole
+	}
+	return slaveRole
+}
+
+//DKVNode defines a DKV Node by a given Host & Port
+type DKVNode struct {
+	Host string `json:"host"`
+	Port int    `json:"port"`
+}
+
+//DKVNodeSet defines a group of DKV Node(s)
+type DKVNodeSet struct {
+	Name  string    `json:"name"`
+	Nodes []DKVNode `json:"nodes"`
+}
+
+//DKVShardKey defined a DKV ShardKey
+type DKVShardKey string
+
+//DKVShard defines a group of DKVNodeSet(s) along with their roles.
+type DKVShard struct {
+	Name     DKVShardKey                   `json:"name"`
+	Topology map[DKVServerRole]*DKVNodeSet `json:"topology"`
+}
+
+func (s DKVShard) getNodesByType(nodeType ...DKVServerRole) (*DKVNodeSet, error) {
+	for _, v := range nodeType {
+		if val, ok := s.Topology[v]; ok {
+			return val, nil
+		}
+	}
+	return nil, fmt.Errorf("valid DKV node type must be given")
+}
+
+// DKVKeyGroup Group of DKV Shard and associated keys.
+type DKVKeyGroup struct {
+	//Keys Set of byte Keys
+	Keys [][]byte
+	//Shard the DKV Shard associated with this keys
+	Shard *DKVShard
+}
+
+//ShardProvider Provides the ShardInformation for the given key(s)
+type ShardProvider interface {
+	ProvideShard(key []byte) (*DKVShard, error)
+	ProvideShards(keys ...[]byte) ([]DKVKeyGroup, error)
+}
+
+//KeyHashBasedShardProvider A xxhash based shared provider.
+type KeyHashBasedShardProvider struct {
+	shardConfiguration []DKVShard
+}
+
+//ProvideShard provides shardId based on input key
+func (p *KeyHashBasedShardProvider) ProvideShard(key []byte) (*DKVShard, error) {
+	shardId := p.getShardID(key)
+	return &p.shardConfiguration[shardId], nil
+}
+
+//ProvideShards provides list of pairs<shardId, keys> for given list of keys.
+func (p *KeyHashBasedShardProvider) ProvideShards(keys ...[]byte) ([]DKVKeyGroup, error) {
+	m := make(map[int][][]byte)
+	for _, key := range keys {
+		shardId := p.getShardID(key)
+		m[shardId] = append(m[shardId], key)
+	}
+
+	values := make([]DKVKeyGroup, 0, len(m))
+	for id, keys := range m {
+		values = append(values, DKVKeyGroup{
+			Keys:  keys,
+			Shard: &p.shardConfiguration[id],
+		})
+	}
+	return values, nil
+}
+
+func (p *KeyHashBasedShardProvider) getShardID(key []byte) int {
+	h := xxhash.New64()
+	h.Write(key)
+	hash := h.Sum64()
+	var id = (hash & 0xFFFF) % uint64(len(p.shardConfiguration))
+	return int(id)
+}
+
+type simpleDKVClient struct {
+	*ctl.DKVClient
+	addr string
+}
+
+// A ShardedDKVClient instance is used to communicate with a shared DKV cluster
+// over GRPC. It is an adapter to the underlying GRPC clients that
+// exposes a simpler API to its users without having to deal with timeouts,
+// contexts, sharding and other GRPC semantics.
+type ShardedDKVClient struct {
+	pool          *ristretto.Cache
+	shardProvider ShardProvider
+	connectOpts   ctl.ConnectOpts
+}
+
+// NewShardedDKVClient creates and returns a instance of ShardedDKVClient.
+func NewShardedDKVClient(shardProvider ShardProvider, opts ctl.ConnectOpts) (*ShardedDKVClient, error) {
+	if shardProvider == nil {
+		return nil, fmt.Errorf("shardProvider cannot be nil")
+	}
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1000,
+		MaxCost:     1000,
+		BufferItems: 64,
+		OnExit: func(val interface{}) {
+			client := val.(*simpleDKVClient)
+			log.Printf("[INFO] Closing Client of %s \n", client.addr)
+			client.Close()
+		},
+	})
+
+	if err != nil {
+		return nil, err
+	}
+	return &ShardedDKVClient{
+		shardProvider: shardProvider,
+		pool:          cache,
+		connectOpts:   opts,
+	}, nil
+}
+
+func (dkvClnt *ShardedDKVClient) getShardNode(shard *DKVShard, role ...DKVServerRole) (*DKVNodeSet, error) {
+	var nodeSet *DKVNodeSet
+	var err error
+	for _, serverRole := range role {
+		nodeSet, err = shard.getNodesByType(serverRole)
+		if err == nil {
+			return nodeSet, nil
+		}
+	}
+	return nodeSet, err
+}
+
+func (dkvClnt *ShardedDKVClient) getShardedClient(shard *DKVShard, role ...DKVServerRole) (*simpleDKVClient, error) {
+	var client *simpleDKVClient
+	nodeSet, err := dkvClnt.getShardNode(shard, role...)
+	if err != nil {
+		return nil, err
+	}
+	if nodeSet == nil {
+		return nil, fmt.Errorf("unable to get proper dkv shard endpoint")
+	}
+
+	//got a node
+	dkvNode := nodeSet.Nodes[0]
+	svcAddr := fmt.Sprintf("%s:%d", dkvNode.Host, dkvNode.Port)
+	if value, found := dkvClnt.pool.Get(svcAddr); !found {
+		log.Printf("[INFO] Creating new Client to : %s\n", svcAddr)
+		_client, err := ctl.NewInSecureDKVClient(svcAddr, nodeSet.Name, dkvClnt.connectOpts)
+		if err != nil {
+			return nil, err
+		}
+		client = &simpleDKVClient{_client, svcAddr}
+		dkvClnt.pool.Set(svcAddr, client, 1)
+	} else {
+		client = value.(*simpleDKVClient)
+	}
+	return client, nil
+}
+
+// Put takes the key and value as byte arrays, find the corresponding shard
+// and invokes the GRPC Put method.
+func (dkvClnt *ShardedDKVClient) Put(key []byte, value []byte) error {
+	dkvShard, err := dkvClnt.shardProvider.ProvideShard(key)
+	if err != nil {
+		return err
+	}
+	client, err := dkvClnt.getShardedClient(dkvShard, masterRole, noRole)
+	if err != nil {
+		return err
+	}
+	return client.Put(key, value)
+}
+
+// Delete takes the key and value as byte arrays, find the corresponding shard
+// and invokes the GRPC Delete method.
+func (dkvClnt *ShardedDKVClient) Delete(key []byte) error {
+	dkvShard, err := dkvClnt.shardProvider.ProvideShard(key)
+	if err != nil {
+		return err
+	}
+	client, err := dkvClnt.getShardedClient(dkvShard, masterRole, noRole)
+	if err != nil {
+		return err
+	}
+	return client.Delete(key)
+}
+
+// Get takes the key as byte array along with the consistency,
+// finds the corresponding shard and invokes the GRPC Get method.
+func (dkvClnt *ShardedDKVClient) Get(rc serverpb.ReadConsistency, key []byte) (*serverpb.GetResponse, error) {
+	dkvShard, err := dkvClnt.shardProvider.ProvideShard(key)
+	if err != nil {
+		return nil, err
+	}
+	nodeRole := getNodeTypeByReadConsistency(rc)
+	client, err := dkvClnt.getShardedClient(dkvShard, nodeRole, noRole)
+
+	if err != nil {
+		return nil, err
+	}
+	return client.Get(rc, key)
+}
+
+// MultiGet takes the keys as byte arrays along with the consistency,
+// finds the corresponding shard and invokes the GRPC MultiGet method.
+func (dkvClnt *ShardedDKVClient) MultiGet(rc serverpb.ReadConsistency, keys ...[]byte) ([]*serverpb.KVPair, error) {
+	var kv []*serverpb.KVPair
+	dkvShardIds, err := dkvClnt.shardProvider.ProvideShards(keys...)
+	if err != nil {
+		return nil, err
+	}
+	nodeType := getNodeTypeByReadConsistency(rc)
+	for _, keyGroup := range dkvShardIds {
+		_dkvClient, err := dkvClnt.getShardedClient(keyGroup.Shard, nodeType, noRole)
+		if err != nil {
+			return nil, err
+		}
+		kvs, err := _dkvClient.MultiGet(rc, keyGroup.Keys...)
+		if err != nil {
+			return nil, err
+		}
+		kv = append(kv, kvs...)
+	}
+	return kv, nil
+}
+
+// Iterate invokes the underlying GRPC method for iterating through the
+// entire keyspace in no particular order. `keyPrefix` can be used to
+// select only the keys matching the given prefix and `startKey` can
+// be used to set the lower bound for the iteration.
+func (dkvClnt *ShardedDKVClient) Iterate(keyPrefix, startKey []byte) (<-chan *ctl.KVPair, error) {
+	dkvShard, err := dkvClnt.shardProvider.ProvideShard(startKey)
+	if err != nil {
+		return nil, err
+	}
+	client, err := dkvClnt.getShardedClient(dkvShard, slaveRole, noRole)
+	if err != nil {
+		return nil, err
+	}
+	return client.Iterate(keyPrefix, startKey)
+}
+
+// Close closes the underlying GRPC client(s) connection to DKV service
+func (dkvClnt *ShardedDKVClient) Close() error {
+	dkvClnt.pool.Close()
+	return nil
+}