97 lines
2.5 KiB
Go
97 lines
2.5 KiB
Go
|
package cache
|
||
|
|
|
||
|
|
import (
|
||
|
|
"io"
|
||
|
|
"sort"
|
||
|
|
"sync"
|
||
|
|
"sync/atomic"
|
||
|
|
|
||
|
|
"github.com/go-kit/kit/endpoint"
|
||
|
|
"github.com/go-kit/kit/log"
|
||
|
|
"github.com/go-kit/kit/sd"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Cache collects the most recent set of endpoints from a service discovery
|
||
|
|
// system via a subscriber, and makes them available to consumers. Cache is
|
||
|
|
// meant to be embedded inside of a concrete subscriber, and can serve Service
|
||
|
|
// invocations directly.
|
||
|
|
type Cache struct {
|
||
|
|
mtx sync.RWMutex
|
||
|
|
factory sd.Factory
|
||
|
|
cache map[string]endpointCloser
|
||
|
|
slice atomic.Value // []endpoint.Endpoint
|
||
|
|
logger log.Logger
|
||
|
|
}
|
||
|
|
|
||
|
|
type endpointCloser struct {
|
||
|
|
endpoint.Endpoint
|
||
|
|
io.Closer
|
||
|
|
}
|
||
|
|
|
||
|
|
// New returns a new, empty endpoint cache.
|
||
|
|
func New(factory sd.Factory, logger log.Logger) *Cache {
|
||
|
|
return &Cache{
|
||
|
|
factory: factory,
|
||
|
|
cache: map[string]endpointCloser{},
|
||
|
|
logger: logger,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Update should be invoked by clients with a complete set of current instance
|
||
|
|
// strings whenever that set changes. The cache manufactures new endpoints via
|
||
|
|
// the factory, closes old endpoints when they disappear, and persists existing
|
||
|
|
// endpoints if they survive through an update.
|
||
|
|
func (c *Cache) Update(instances []string) {
|
||
|
|
c.mtx.Lock()
|
||
|
|
defer c.mtx.Unlock()
|
||
|
|
|
||
|
|
// Deterministic order (for later).
|
||
|
|
sort.Strings(instances)
|
||
|
|
|
||
|
|
// Produce the current set of services.
|
||
|
|
cache := make(map[string]endpointCloser, len(instances))
|
||
|
|
for _, instance := range instances {
|
||
|
|
// If it already exists, just copy it over.
|
||
|
|
if sc, ok := c.cache[instance]; ok {
|
||
|
|
cache[instance] = sc
|
||
|
|
delete(c.cache, instance)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
|
||
|
|
// If it doesn't exist, create it.
|
||
|
|
service, closer, err := c.factory(instance)
|
||
|
|
if err != nil {
|
||
|
|
c.logger.Log("instance", instance, "err", err)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
cache[instance] = endpointCloser{service, closer}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Close any leftover endpoints.
|
||
|
|
for _, sc := range c.cache {
|
||
|
|
if sc.Closer != nil {
|
||
|
|
sc.Closer.Close()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Populate the slice of endpoints.
|
||
|
|
slice := make([]endpoint.Endpoint, 0, len(cache))
|
||
|
|
for _, instance := range instances {
|
||
|
|
// A bad factory may mean an instance is not present.
|
||
|
|
if _, ok := cache[instance]; !ok {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
slice = append(slice, cache[instance].Endpoint)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Swap and trigger GC for old copies.
|
||
|
|
c.slice.Store(slice)
|
||
|
|
c.cache = cache
|
||
|
|
}
|
||
|
|
|
||
|
|
// Endpoints yields the current set of (presumably identical) endpoints, ordered
|
||
|
|
// lexicographically by the corresponding instance string.
|
||
|
|
func (c *Cache) Endpoints() []endpoint.Endpoint {
|
||
|
|
return c.slice.Load().([]endpoint.Endpoint)
|
||
|
|
}
|