github.com/thanos-io/thanos@v0.32.5/pkg/cacheutil/memcached_server_selector.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package cacheutil
     5  
     6  import (
     7  	"net"
     8  	"sync"
     9  
    10  	"github.com/bradfitz/gomemcache/memcache"
    11  	"github.com/cespare/xxhash"
    12  	"github.com/facette/natsort"
    13  )
    14  
    15  var (
    16  	addrsPool = sync.Pool{
    17  		New: func() interface{} {
    18  			addrs := make([]net.Addr, 0, 64)
    19  			return &addrs
    20  		},
    21  	}
    22  )
    23  
    24  // MemcachedJumpHashSelector implements the memcache.ServerSelector
    25  // interface, utilizing a jump hash to distribute keys to servers.
    26  //
    27  // While adding or removing servers only requires 1/N keys to move,
    28  // servers are treated as a stack and can only be pushed/popped.
    29  // Therefore, MemcachedJumpHashSelector works best for servers
    30  // with consistent DNS names where the naturally sorted order
    31  // is predictable (ie. Kubernetes statefulsets).
    32  type MemcachedJumpHashSelector struct {
    33  	// To avoid copy and pasting all memcache server list logic,
    34  	// we embed it and implement our features on top of it.
    35  	servers memcache.ServerList
    36  }
    37  
    38  // SetServers changes a MemcachedJumpHashSelector's set of servers at
    39  // runtime and is safe for concurrent use by multiple goroutines.
    40  //
    41  // Each server is given equal weight. A server is given more weight
    42  // if it's listed multiple times.
    43  //
    44  // SetServers returns an error if any of the server names fail to
    45  // resolve. No attempt is made to connect to the server. If any
    46  // error occurs, no changes are made to the internal server list.
    47  //
    48  // To minimize the number of rehashes for keys when scaling the
    49  // number of servers in subsequent calls to SetServers, servers
    50  // are stored in natural sort order.
    51  func (s *MemcachedJumpHashSelector) SetServers(servers ...string) error {
    52  	sortedServers := make([]string, len(servers))
    53  	copy(sortedServers, servers)
    54  	natsort.Sort(sortedServers)
    55  
    56  	return s.servers.SetServers(sortedServers...)
    57  }
    58  
    59  // PickServer returns the server address that a given item
    60  // should be shared onto.
    61  func (s *MemcachedJumpHashSelector) PickServer(key string) (net.Addr, error) {
    62  	// Unfortunately we can't read the list of server addresses from
    63  	// the original implementation, so we use Each() to fetch all of them.
    64  	addrs := *(addrsPool.Get().(*[]net.Addr))
    65  	err := s.servers.Each(func(addr net.Addr) error {
    66  		addrs = append(addrs, addr)
    67  		return nil
    68  	})
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  
    73  	// No need of a jump hash in case of 0 or 1 servers.
    74  	if len(addrs) == 0 {
    75  		addrs = (addrs)[:0]
    76  		addrsPool.Put(&addrs)
    77  		return nil, memcache.ErrNoServers
    78  	}
    79  	if len(addrs) == 1 {
    80  		picked := addrs[0]
    81  
    82  		addrs = (addrs)[:0]
    83  		addrsPool.Put(&addrs)
    84  
    85  		return picked, nil
    86  	}
    87  
    88  	// Pick a server using the jump hash.
    89  	cs := xxhash.Sum64String(key)
    90  	idx := jumpHash(cs, len(addrs))
    91  	picked := (addrs)[idx]
    92  
    93  	addrs = (addrs)[:0]
    94  	addrsPool.Put(&addrs)
    95  
    96  	return picked, nil
    97  }
    98  
    99  // Each iterates over each server and calls the given function.
   100  // If f returns a non-nil error, iteration will stop and that
   101  // error will be returned.
   102  func (s *MemcachedJumpHashSelector) Each(f func(net.Addr) error) error {
   103  	return s.servers.Each(f)
   104  }