github.com/thanos-io/thanos@v0.32.5/internal/cortex/chunk/cache/memcached_client_selector.go (about)

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