Repository: grafana/grafana-app-sdk
Pull request #729: [MemcachedStore] Allow for specifying a ServerSelector and automatic server refresh
Author: login: IfSentient
Labels: name: area/controller
Created: 2025-04-02T19:28:53Z | Merged: 2025-04-11T14:18:41Z

Allow for specifying a `ServerSelector` in the `MemcachedStore` and automatically refresh the servers when an EOF or timeout is encountered.

A default implementation of the `MemcachedServerSelector` using `memcache.ServerList` has been added to the `simple` package.



Diff:
diff --git a/operator/cache_memcached.go b/operator/cache_memcached.go
index a1a768f03..d62898130 100644
--- a/operator/cache_memcached.go
+++ b/operator/cache_memcached.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/bradfitz/gomemcache/memcache"
@@ -26,19 +27,29 @@ const (
 	keysCacheKey = "%s-keys"
 )
 
+// MemcachedServerSelector enhances memcache.ServerSelector with a method to refresh the server list used.
+type MemcachedServerSelector interface {
+	memcache.ServerSelector
+	// RefreshServers refreshes the underlying list of servers, re-resolving any hostnames.
+	// This is used by the MemcachedStore when an operation results in an EOF or timeout error to ensure that
+	// the resolved addresses are still valid.
+	RefreshServers() error
+}
+
 // MemcachedStore implements cache.Store using memcached as the store for objects.
 // It should be instantiated with NewMemcachedStore.
 type MemcachedStore struct {
-	client       *memcache.Client
-	keyFunc      func(any) (string, error)
-	kind         resource.Kind
-	readLatency  *prometheus.HistogramVec
-	writeLatency *prometheus.HistogramVec
-	keys         sync.Map
-	trackKeys    bool
-	syncTicker   *time.Ticker
-	pageSize     int
-	cacheKeysKey string
+	client          *memcache.Client
+	keyFunc         func(any) (string, error)
+	kind            resource.Kind
+	readLatency     *prometheus.HistogramVec
+	writeLatency    *prometheus.HistogramVec
+	keys            sync.Map
+	trackKeys       bool
+	syncTicker      *time.Ticker
+	pageSize        int
+	cacheKeysKey    string
+	serverRefresher *memcachedServerRefresher
 }
 
 // MemcachedStoreConfig is a collection of config values for a MemcachedStore
@@ -47,6 +58,9 @@ type MemcachedStoreConfig struct {
 	KeyFunc func(any) (string, error)
 	// Addrs is a list of addresses (including ports) to connect to
 	Addrs []string
+	// ServerSelector is a server selector for the memcached client.
+	// If present, it overrides Addrs and is used to determine the memcached servers to connect to.
+	ServerSelector MemcachedServerSelector
 	// Metrics is metrics configuration
 	Metrics metrics.Config
 	// KeySyncInterval is the interval at which keys stored in the in-memory map will be pushed to memcached.
@@ -79,7 +93,12 @@ func NewMemcachedStore(kind resource.Kind, cfg MemcachedStoreConfig) (*Memcached
 	if cfg.KeyFunc != nil {
 		keyFunc = cfg.KeyFunc
 	}
-	client := memcache.New(cfg.Addrs...)
+	var client *memcache.Client
+	if cfg.ServerSelector != nil {
+		client = memcache.NewFromSelector(cfg.ServerSelector)
+	} else {
+		client = memcache.New(cfg.Addrs...)
+	}
 	client.Timeout = cfg.Timeout
 	client.MaxIdleConns = cfg.MaxIdleConns
 	store := &MemcachedStore{
@@ -111,6 +130,11 @@ func NewMemcachedStore(kind resource.Kind, cfg MemcachedStoreConfig) (*Memcached
 		pageSize:     MemcachedStoreDefaultPageSize,
 		cacheKeysKey: fmt.Sprintf(keysCacheKey, kind.Plural()),
 	}
+	if cfg.ServerSelector != nil {
+		store.serverRefresher = &memcachedServerRefresher{
+			refresher: cfg.ServerSelector,
+		}
+	}
 	if cfg.PageSize > 0 {
 		store.pageSize = cfg.PageSize
 	}
@@ -173,9 +197,11 @@ func (m *MemcachedStore) Update(obj any) error {
 		return err
 	}
 	start := time.Now()
-	err = m.client.Replace(&memcache.Item{
-		Key:   key,
-		Value: o,
+	err = m.attemptWithRefreshOnTimeout(func() error {
+		return m.client.Replace(&memcache.Item{
+			Key:   key,
+			Value: o,
+		})
 	})
 	m.writeLatency.WithLabelValues(m.kind.Kind()).Observe(time.Since(start).Seconds())
 	return err
@@ -186,7 +212,9 @@ func (m *MemcachedStore) Delete(obj any) error {
 		return err
 	}
 	start := time.Now()
-	err = m.client.Delete(key)
+	err = m.attemptWithRefreshOnTimeout(func() error {
+		return m.client.Delete(key)
+	})
 	m.writeLatency.WithLabelValues(m.kind.Kind()).Observe(time.Since(start).Seconds())
 	if err == nil && m.trackKeys {
 		m.keys.Delete(trackKey)
@@ -209,7 +237,12 @@ func (m *MemcachedStore) List() []any {
 		for j := 0; j < len(fetchKeys); j++ {
 			fetchKeys[j] = fmt.Sprintf("%s/%s", m.kind.Plural(), fetchKeys[j])
 		}
-		res, err := m.client.GetMulti(fetchKeys)
+		var res map[string]*memcache.Item
+		err := m.attemptWithRefreshOnTimeout(func() error {
+			var err error
+			res, err = m.client.GetMulti(fetchKeys)
+			return err
+		})
 		if err != nil {
 			// TODO: ???
 			return nil
@@ -253,7 +286,11 @@ func (m *MemcachedStore) GetByKey(key string) (item any, exists bool, err error)
 
 func (m *MemcachedStore) getByKey(key string) (item any, exists bool, err error) {
 	start := time.Now()
-	fromCache, err := m.client.Get(key)
+	var fromCache *memcache.Item
+	err = m.attemptWithRefreshOnTimeout(func() error {
+		fromCache, err = m.client.Get(key)
+		return err
+	})
 	m.readLatency.WithLabelValues(m.kind.Kind()).Observe(time.Since(start).Seconds())
 	if err != nil && !errors.Is(err, memcache.ErrCacheMiss) {
 		return nil, false, err
@@ -287,7 +324,12 @@ func (m *MemcachedStore) getKey(obj any) (prefixedKey string, externalKey string
 }
 
 func (m *MemcachedStore) setKeysFromCache() error {
-	item, err := m.client.Get(m.cacheKeysKey)
+	var item *memcache.Item
+	err := m.attemptWithRefreshOnTimeout(func() error {
+		var err error
+		item, err = m.client.Get(m.cacheKeysKey)
+		return err
+	})
 	if err != nil {
 		if errors.Is(err, memcache.ErrCacheMiss) {
 			return nil
@@ -306,19 +348,22 @@ func (m *MemcachedStore) setKeysFromCache() error {
 }
 
 func (m *MemcachedStore) syncKeys() error {
-	_, err := m.client.Get(m.cacheKeysKey)
-	if err != nil {
-		if !errors.Is(err, memcache.ErrCacheMiss) {
-			return err
-		}
-		item := &memcache.Item{
-			Key:   m.cacheKeysKey,
-			Value: []byte("[]"),
-		}
-		err = m.client.Add(item)
+	err := m.attemptWithRefreshOnTimeout(func() error {
+		_, err := m.client.Get(m.cacheKeysKey)
 		if err != nil {
-			return err
+			if !errors.Is(err, memcache.ErrCacheMiss) {
+				return err
+			}
+			item := &memcache.Item{
+				Key:   m.cacheKeysKey,
+				Value: []byte("[]"),
+			}
+			return m.client.Add(item)
 		}
+		return nil
+	})
+	if err != nil {
+		return err
 	}
 	externalKeys := make([]string, 0)
 	m.keys.Range(func(key, _ any) bool {
@@ -329,8 +374,51 @@ func (m *MemcachedStore) syncKeys() error {
 	if err != nil {
 		return err
 	}
-	return m.client.Replace(&memcache.Item{
-		Key:   m.cacheKeysKey,
-		Value: externalKeysJSON,
+	return m.attemptWithRefreshOnTimeout(func() error {
+		return m.client.Replace(&memcache.Item{
+			Key:   m.cacheKeysKey,
+			Value: externalKeysJSON,
+		})
 	})
 }
+
+func (m *MemcachedStore) attemptWithRefreshOnTimeout(f func() error) error {
+	if err := f(); err != nil {
+		if m.serverRefresher != nil {
+			var e *memcache.ConnectTimeoutError
+			if errors.As(err, &e) || err.Error() == "EOF" {
+				refreshErr := m.serverRefresher.refresh()
+				if refreshErr != nil {
+					return fmt.Errorf("%w: could not refresh server list: %w", err, refreshErr)
+				}
+				return f()
+			}
+		}
+		return err
+	}
+	return nil
+}
+
+type memcachedServerRefresher struct {
+	refresher  MemcachedServerSelector
+	mux        sync.Mutex
+	inProgress atomic.Bool
+}
+
+func (m *memcachedServerRefresher) refresh() error {
+	if !m.inProgress.CompareAndSwap(false, true) {
+		// Block until the ongoing refresh is finished,
+		// just attempt to lock the mutex so we wait until it's freed by the existing process
+		m.mux.Lock()
+		defer m.mux.Unlock()
+		return nil
+	}
+	if logging.DefaultLogger != nil {
+		logging.DefaultLogger.Info("Refreshing memcached servers")
+	}
+	m.mux.Lock()
+	defer m.mux.Unlock()
+	defer m.inProgress.Store(false)
+
+	return m.refresher.RefreshServers()
+}
diff --git a/operator/informer_customcache.go b/operator/informer_customcache.go
index 80026d731..7248f4ced 100644
--- a/operator/informer_customcache.go
+++ b/operator/informer_customcache.go
@@ -47,7 +47,10 @@ type CustomCacheInformer struct {
 }
 
 type MemcachedInformerOptions struct {
-	Addrs            []string
+	Addrs []string
+	// ServerSelector is a server selector for the memcached client.
+	// If present, it overrides Addrs and is used to determine the memcached servers to connect to.
+	ServerSelector   MemcachedServerSelector
 	ListWatchOptions ListWatchOptions
 }
 
@@ -56,7 +59,8 @@ type MemcachedInformerOptions struct {
 // To set additional memcached options, use NewCustomCacheInformer and NewMemcachedStore.
 func NewMemcachedInformer(kind resource.Kind, client ListWatchClient, opts MemcachedInformerOptions) (*CustomCacheInformer, error) {
 	c, err := NewMemcachedStore(kind, MemcachedStoreConfig{
-		Addrs: opts.Addrs,
+		Addrs:          opts.Addrs,
+		ServerSelector: opts.ServerSelector,
 	})
 	if err != nil {
 		return nil, err
diff --git a/simple/memcached.go b/simple/memcached.go
new file mode 100644
index 000000000..19fcbe7f5
--- /dev/null
+++ b/simple/memcached.go
@@ -0,0 +1,48 @@
+package simple
+
+import (
+	"net"
+	"sync"
+
+	"github.com/bradfitz/gomemcache/memcache"
+
+	"github.com/grafana/grafana-app-sdk/operator"
+)
+
+type MemcachedHostList struct {
+	dnsMux     sync.Mutex
+	endpoints  []string
+	serverList *memcache.ServerList
+}
+
+// NewMemcachedHostList takes a list of endpoints and returns a MemcachedHostList which will resolve the hosts
+// in the endpoints slice, which should be in the format [host]:[port].
+// RefreshServers will re-resolve the endpoints for the MemcachedHostList.
+func NewMemcachedHostList(endpoints []string) (*MemcachedHostList, error) {
+	s := &MemcachedHostList{
+		endpoints:  endpoints,
+		serverList: &memcache.ServerList{},
+	}
+	// Resolve the initial list
+	err := s.serverList.SetServers(endpoints...)
+	return s, err
+}
+
+func (m *MemcachedHostList) PickServer(key string) (net.Addr, error) {
+	return m.serverList.PickServer(key)
+}
+
+func (m *MemcachedHostList) Each(f func(net.Addr) error) error {
+	return m.serverList.Each(f)
+}
+
+// RefreshServers re-resolves the endpoints for this MemcachedHostList.
+// This can be used if the memcached client runs into a timeout or EOF due to the host IP changing.
+func (m *MemcachedHostList) RefreshServers() error {
+	m.dnsMux.Lock()
+	defer m.dnsMux.Unlock()
+	// Setting the ServerList with SetServers re-resolves the hosts to inet addresses
+	return m.serverList.SetServers(m.endpoints...)
+}
+
+var _ operator.MemcachedServerSelector = &MemcachedHostList{}


Reviews:
user: login: IfSentient
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T19:30:44Z
user: login: ferruvich
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T08:34:47Z
user: login: IfSentient
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T12:18:24Z
user: login: ferruvich
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T14:15:43Z

Review comments:
diff_hunk: @@ -0,0 +1,48 @@
+package simple
path: simple/memcached.go
user: login: IfSentient
body: I am on the fence about this being in the `simple` package, but it didn't feel like it belonged in the `operator` package either.
created_at: 2025-04-02T19:30:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -253,7 +281,11 @@ func (m *MemcachedStore) GetByKey(key string) (item any, exists bool, err error)
 
 func (m *MemcachedStore) getByKey(key string) (item any, exists bool, err error) {
 	start := time.Now()
-	fromCache, err := m.client.Get(key)
+	var fromCache *memcache.Item
+	err = m.attemptWithRefreshOnTimeout(func() error {
path: operator/cache_memcached.go
user: login: ferruvich
body: one question: should the `attemptWithRefreshOnTimeout` be used **each** time the `m.client` is used? Because I can see some "plain" calls to `m.client` in this file 🤔 
created_at: 2025-04-04T08:34:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -253,7 +281,11 @@ func (m *MemcachedStore) GetByKey(key string) (item any, exists bool, err error)
 
 func (m *MemcachedStore) getByKey(key string) (item any, exists bool, err error) {
 	start := time.Now()
-	fromCache, err := m.client.Get(key)
+	var fromCache *memcache.Item
+	err = m.attemptWithRefreshOnTimeout(func() error {
path: operator/cache_memcached.go
user: login: IfSentient
body: Ah, you're right, I missed doing this with the key sync stuff for handling cache resyncs. Let me push a fix for that.
created_at: 2025-04-07T12:18:24Z
author_association: CONTRIBUTOR

PR comments:
user: login: IfSentient
created_at: 2025-04-03T04:27:33Z
body: CC @ferruvich 
author_association: CONTRIBUTOR
user: login: ferruvich
created_at: 2025-04-08T13:59:02Z
body: @IfSentient is this ready for review again?
author_association: CONTRIBUTOR
user: login: IfSentient
created_at: 2025-04-10T14:24:12Z
body: > @IfSentient is this ready for review again?

@ferruvich Ah, yup, sorry that I missed this ping
author_association: CONTRIBUTOR