Fixed race condition and peer client reference issue

cache/lru.go

@@ -254,6 +254,8 @@ func (c *LRUCache) Size() int {
 
 // Returns stats about the current state of the cache
 func (c *LRUCache) Stats(clear bool) Stats {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
 	if clear {
 		defer func() {
 			c.stats = Stats{}

hash.go

@@ -61,9 +61,9 @@ func (ch *ConsistantHash) GetPeer(host string) *PeerClient {
 }
 
 // Given a key, return the peer that key is assigned too
-func (ch *ConsistantHash) Get(key string, peerInfo *PeerClient) error {
+func (ch *ConsistantHash) Get(key string) (*PeerClient, error) {
 	if ch.Size() == 0 {
-		return errors.New("unable to pick a peer; pool is empty")
+		return nil, errors.New("unable to pick a peer; pool is empty")
 	}
 
 	hash := int(ch.hashFunc([]byte(key)))
@@ -76,7 +76,5 @@ func (ch *ConsistantHash) Get(key string, peerInfo *PeerClient) error {
 		idx = 0
 	}
 
-	item := ch.peerMap[ch.peerKeys[idx]]
-	*peerInfo = *item
-	return nil
+	return ch.peerMap[ch.peerKeys[idx]], nil
 }

peers.go

@@ -2,6 +2,7 @@ package gubernator
 
 import (
 	"context"
+	"fmt"
 	"github.com/mailgun/gubernator/pb"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc"
@@ -10,7 +11,7 @@ import (
 type PeerPicker interface {
 	GetPeer(host string) *PeerClient
 	Peers() []*PeerClient
-	Get(string, *PeerClient) error
+	Get(string) (*PeerClient, error)
 	New() PeerPicker
 	Add(*PeerClient)
 	Size() int
@@ -58,6 +59,7 @@ func (c *PeerClient) GetPeerRateLimits(ctx context.Context, r *pb.RateLimitReque
 // Dial to a peer and initialize the GRPC client
 func (c *PeerClient) dialPeer() error {
 	// TODO: Allow TLS connections
+	fmt.Printf("Dial: %s\n", c.host)
 
 	var err error
 	c.conn, err = grpc.Dial(c.host, grpc.WithInsecure())

server.go

@@ -133,15 +133,14 @@ func (s *Server) GetRateLimits(ctx context.Context, reqs *pb.RateLimitRequestLis
 		globalKey := req.Namespace + "_" + req.UniqueKey
 
 		s.peerMutex.RLock()
-		var peer PeerClient
-		if err := s.conf.Picker.Get(globalKey, &peer); err != nil {
+		peer, err := s.conf.Picker.Get(globalKey)
+		if err != nil {
 			s.peerMutex.RUnlock()
 			return nil, errors.Wrapf(err, "while finding peer that owns key '%s'", globalKey)
 		}
 		s.peerMutex.RUnlock()
 
 		var resp *pb.RateLimitResponse
-		var err error
 
 		// If our server instance is the owner of this rate limit
 		if peer.isOwner {