github.com/sequix/cortex@v1.1.6/pkg/chunk/cache/memcached_client_selector.go (about)

     1  package cache
     2  
     3  import (
     4  	"net"
     5  	"strings"
     6  	"sync"
     7  
     8  	"github.com/bradfitz/gomemcache/memcache"
     9  	"github.com/cespare/xxhash"
    10  	"github.com/facette/natsort"
    11  )
    12  
    13  // MemcachedJumpHashSelector implements the memcache.ServerSelector
    14  // interface. MemcachedJumpHashSelector utilizes a jump hash to
    15  // distribute keys to servers.
    16  //
    17  // While adding or removing servers only requires 1/N keys to move,
    18  // servers are treated as a stack and can only be pushed/popped.
    19  // Therefore, MemcachedJumpHashSelector works best for servers
    20  // with consistent DNS names where the naturally sorted order
    21  // is predictable.
    22  type MemcachedJumpHashSelector struct {
    23  	mu    sync.RWMutex
    24  	addrs []net.Addr
    25  }
    26  
    27  // staticAddr caches the Network() and String() values from
    28  // any net.Addr.
    29  //
    30  // Copied from github.com/bradfitz/gomemcache/selector.go.
    31  type staticAddr struct {
    32  	network, str string
    33  }
    34  
    35  func newStaticAddr(a net.Addr) net.Addr {
    36  	return &staticAddr{
    37  		network: a.Network(),
    38  		str:     a.String(),
    39  	}
    40  }
    41  
    42  func (a *staticAddr) Network() string { return a.network }
    43  func (a *staticAddr) String() string  { return a.str }
    44  
    45  // SetServers changes a MemcachedJumpHashSelector's set of servers at
    46  // runtime and is safe for concurrent use by multiple goroutines.
    47  //
    48  // Each server is given equal weight. A server is given more weight
    49  // if it's listed multiple times.
    50  //
    51  // SetServers returns an error if any of the server names fail to
    52  // resolve. No attempt is made to connect to the server. If any
    53  // error occurs, no changes are made to the internal server list.
    54  //
    55  // To minimize the number of rehashes for keys when scaling the
    56  // number of servers in subsequent calls to SetServers, servers
    57  // are stored in natural sort order.
    58  func (s *MemcachedJumpHashSelector) SetServers(servers ...string) error {
    59  	sortedServers := make([]string, len(servers))
    60  	copy(sortedServers, servers)
    61  	natsort.Sort(sortedServers)
    62  
    63  	naddrs := make([]net.Addr, len(sortedServers))
    64  	for i, server := range sortedServers {
    65  		if strings.Contains(server, "/") {
    66  			addr, err := net.ResolveUnixAddr("unix", server)
    67  			if err != nil {
    68  				return err
    69  			}
    70  			naddrs[i] = newStaticAddr(addr)
    71  		} else {
    72  			tcpAddr, err := net.ResolveTCPAddr("tcp", server)
    73  			if err != nil {
    74  				return err
    75  			}
    76  			naddrs[i] = newStaticAddr(tcpAddr)
    77  		}
    78  	}
    79  
    80  	s.mu.Lock()
    81  	defer s.mu.Unlock()
    82  	s.addrs = naddrs
    83  	return nil
    84  }
    85  
    86  // jumpHash consistently chooses a hash bucket number in the range [0, numBuckets) for the given key.
    87  // numBuckets must be >= 1.
    88  //
    89  // Copied from github.com/dgryski/go-jump/blob/master/jump.go
    90  func jumpHash(key uint64, numBuckets int) int32 {
    91  
    92  	var b int64 = -1
    93  	var j int64
    94  
    95  	for j < int64(numBuckets) {
    96  		b = j
    97  		key = key*2862933555777941757 + 1
    98  		j = int64(float64(b+1) * (float64(int64(1)<<31) / float64((key>>33)+1)))
    99  	}
   100  
   101  	return int32(b)
   102  }
   103  
   104  // PickServer returns the server address that a given item
   105  // should be shared onto.
   106  func (s *MemcachedJumpHashSelector) PickServer(key string) (net.Addr, error) {
   107  	s.mu.RLock()
   108  	defer s.mu.RUnlock()
   109  	if len(s.addrs) == 0 {
   110  		return nil, memcache.ErrNoServers
   111  	} else if len(s.addrs) == 1 {
   112  		return s.addrs[0], nil
   113  	}
   114  	cs := xxhash.Sum64String(key)
   115  	idx := jumpHash(cs, len(s.addrs))
   116  	return s.addrs[idx], nil
   117  }
   118  
   119  // Each iterates over each server and calls the given function.
   120  // If f returns a non-nil error, iteration will stop and that
   121  // error will be returned.
   122  func (s *MemcachedJumpHashSelector) Each(f func(net.Addr) error) error {
   123  	s.mu.RLock()
   124  	defer s.mu.RUnlock()
   125  	for _, def := range s.addrs {
   126  		if err := f(def); err != nil {
   127  			return err
   128  		}
   129  	}
   130  	return nil
   131  }