Repository: honeycombio/refinery
Pull request #1171 (2024-05-23 14:55:55)

feat: add peer service for sharing peer information
<!--
Thank you for contributing to the project! 💜
Please make sure to:
- Chat with us first if this is a big change
  - Open a new issue (or comment on an existing one)
  - We want to make sure you don't spend time implementing something we might have to say No to
- Add unit tests
- Mention any relevant issues in the PR description (e.g. "Fixes #123")

Please see our [OSS process document](https://github.com/honeycombio/home/blob/main/honeycomb-oss-lifecycle-and-practices.md#) to get an idea of how we operate.
-->

## Which problem is this PR solving?

To share peer information in a centralized place so that we can take advantage of existing peer messages to manage active peers in a centralized place

## Short description of the changes

- create a new PeerStore that's responsible for gossiping peer information
- inject the PeerStore into stress relief and the sampler factory




Patch:
diff --git a/cmd/refinery/main.go b/cmd/refinery/main.go
index 228a031..6da89ce 100644
--- a/cmd/refinery/main.go
+++ b/cmd/refinery/main.go
@@ -30,6 +30,7 @@ import (
 	"github.com/honeycombio/refinery/internal/gossip"
 	"github.com/honeycombio/refinery/internal/health"
 	"github.com/honeycombio/refinery/internal/otelutil"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/redis"
@@ -230,6 +231,7 @@ func main() {
 		{Value: samplerFactory},
 		{Value: channels, Name: "gossip"},
 		{Value: stressRelief, Name: "stressRelief"},
+		{Value: &peer.PeerStore{}},
 		{Value: tracer, Name: "tracer"},
 		{Value: clockwork.NewRealClock()},
 		{Value: basicStore},
diff --git a/collect/stressRelief/stress_relief_redis.go b/collect/stressRelief/stress_relief_redis.go
index ab65d85..f8d47ae 100644
--- a/collect/stressRelief/stress_relief_redis.go
+++ b/collect/stressRelief/stress_relief_redis.go
@@ -2,18 +2,17 @@ package stressRelief
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"math"
-	"os"
 	"strconv"
 	"sync"
 	"time"
 
 	"github.com/dgryski/go-wyhash"
-	"github.com/gofrs/uuid/v5"
 	"github.com/honeycombio/refinery/config"
-	"github.com/honeycombio/refinery/internal/gossip"
 	"github.com/honeycombio/refinery/internal/health"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/jonboulle/clockwork"
@@ -22,14 +21,17 @@ import (
 
 var _ StressReliever = &StressRelief{}
 
-const stressReliefHealthSource = "stress_relief"
+const (
+	stressReliefHealthSource = "stress_relief"
+	peerMessagePrefix        = "sr"
+)
 
 var calculationInterval = 100 * time.Millisecond
 
 type StressRelief struct {
 	RefineryMetrics metrics.Metrics `inject:"genericMetrics"`
 	Logger          logger.Logger   `inject:""`
-	Gossip          gossip.Gossiper `inject:"gossip"`
+	Peer            peer.Peers      `inject:""`
 	Clock           clockwork.Clock `inject:""`
 	Health          health.Recorder `inject:""`
 	done            chan struct{}
@@ -45,8 +47,7 @@ type StressRelief struct {
 	stressed           bool
 	stayOnUntil        time.Time
 	minDuration        time.Duration
-	identification     string
-	stressGossipCh     chan []byte
+	peerChan           <-chan peer.PeerInfo
 
 	eg *errgroup.Group
 
@@ -90,18 +91,6 @@ func (s *StressRelief) Start() error {
 		//{Numerator: "redisstore_memory_used_total", }
 	}
 
-	// We need to identify ourselves to the cluster. We'll use the hostname if we can, but if we can't, we'll use a UUID.
-	if hostname, err := os.Hostname(); err == nil && hostname != "" {
-		s.identification = hostname
-	}
-	if s.identification == "" {
-		id, err := uuid.NewV7()
-		if err != nil {
-			panic("failed to generate a UUID for the StressRelief system")
-		}
-
-		s.identification = id.String()
-	}
 	s.stressLevels = make(map[string]stressReport)
 	s.done = make(chan struct{})
 
@@ -111,7 +100,7 @@ func (s *StressRelief) Start() error {
 	s.RefineryMetrics.Register("individual_stress_level", "gauge")
 	s.RefineryMetrics.Register("stress_relief_activated", "gauge")
 
-	s.stressGossipCh = s.Gossip.Subscribe(s.Gossip.GetChannel(gossip.ChannelStress), 20)
+	s.peerChan = s.Peer.Subscribe()
 	s.eg = &errgroup.Group{}
 	s.eg.Go(s.monitor)
 
@@ -121,38 +110,34 @@ func (s *StressRelief) Start() error {
 func (s *StressRelief) monitor() error {
 	tick := time.NewTicker(calculationInterval)
 	defer tick.Stop()
-	gossipchan := s.Gossip.GetChannel(gossip.ChannelStress)
 	for {
 		select {
 		case <-tick.C:
 			currentLevel := s.Recalc()
 			// publish the stress level to the rest of the cluster
-			msg := stressLevelMessage{
-				level: currentLevel,
-				id:    s.identification,
-			}
-			err := s.Gossip.Publish(gossipchan, msg.ToBytes())
+			err := s.Peer.PublishPeerInfo(peer.PeerInfo{
+				Data: peerMessageToBytes(currentLevel),
+			})
 			if err != nil {
 				s.Logger.Error().Logf("error publishing stress level: %s", err)
 			} else {
 				s.Health.Ready(stressReliefHealthSource, true)
 			}
 
-		case data := <-s.stressGossipCh:
-			msg, err := newMessageFromBytes(data)
+		case msg := <-s.peerChan:
+			level, err := newMessageFromBytes(msg.Data)
 			if err != nil {
 				s.Logger.Error().Logf("error parsing stress level message: %s", err)
 				continue
 			}
 
 			s.lock.Lock()
-			s.stressLevels[msg.id] = stressReport{
-				key:       msg.id,
-				level:     msg.level,
+			s.stressLevels[msg.ID()] = stressReport{
+				key:       msg.ID(),
+				level:     level,
 				timestamp: s.Clock.Now(),
 			}
 			s.lock.Unlock()
-
 		case <-s.done:
 			s.Logger.Debug().Logf("Stopping StressRelief system")
 			return nil
@@ -341,7 +326,7 @@ func (s *StressRelief) clusterStressLevel(level uint) uint {
 	// we need to calculate the stress level from the levels we've been given
 	// and then publish it to the cluster
 	report := stressReport{
-		key:       s.identification,
+		key:       s.Peer.HostID(),
 		level:     level,
 		timestamp: s.Clock.Now(),
 	}
@@ -349,7 +334,7 @@ func (s *StressRelief) clusterStressLevel(level uint) uint {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	s.stressLevels[s.identification] = report
+	s.stressLevels[report.key] = report
 	var total float64
 	availablePeers := 0
 	for _, report := range s.stressLevels {
@@ -373,28 +358,23 @@ func (s *StressRelief) clusterStressLevel(level uint) uint {
 	return uint(math.Sqrt(total / float64(availablePeers)))
 }
 
-// stressLevelMessage is used to communicate stress levels between refinery instances
-// it contains the stress level and the id of the instance that reported it.
-type stressLevelMessage struct {
-	level uint
-	id    string
-}
-
-func (s *stressLevelMessage) ToBytes() []byte {
-	return []byte(fmt.Sprintf("%s/%d", s.id, s.level))
+func peerMessageToBytes(level uint) []byte {
+	return []byte(peerMessagePrefix + "/" + strconv.Itoa(int(level)))
 }
 
-func newMessageFromBytes(b []byte) (stressLevelMessage, error) {
+func newMessageFromBytes(b []byte) (uint, error) {
 	parts := bytes.SplitN(b, []byte("/"), 2)
 	if len(parts) != 2 {
-		return stressLevelMessage{}, fmt.Errorf("invalid message format: %s", b)
+		return 0, fmt.Errorf("invalid message format: %s", b)
 	}
+
+	if string(parts[0]) != peerMessagePrefix {
+		return 0, errors.New("invalid message type")
+	}
+
 	level, err := strconv.Atoi(string(parts[1]))
 	if err != nil {
-		return stressLevelMessage{}, fmt.Errorf("invalid level: %s", parts[1])
+		return 0, fmt.Errorf("invalid level: %s", parts[1])
 	}
-	return stressLevelMessage{
-		id:    string(parts[0]),
-		level: uint(level),
-	}, nil
+	return uint(level), nil
 }
diff --git a/internal/gossip/gossip.go b/internal/gossip/gossip.go
index eaaf227..1c91c66 100644
--- a/internal/gossip/gossip.go
+++ b/internal/gossip/gossip.go
@@ -5,11 +5,11 @@ import (
 )
 
 const (
-	ChannelKeep   = "keep"
-	ChannelDrop   = "drop"
-	ChannelStress = "stress_health"
-	ChannelTest1  = "test1"
-	ChannelTest2  = "test2"
+	ChannelKeep  = "keep"
+	ChannelDrop  = "drop"
+	ChannelPeer  = "peer"
+	ChannelTest1 = "test1"
+	ChannelTest2 = "test2"
 )
 
 // Gossiper is an interface for broadcasting messages to all receivers
diff --git a/internal/gossip/gossip_inmem.go b/internal/gossip/gossip_inmem.go
index 10b75f4..5a4297a 100644
--- a/internal/gossip/gossip_inmem.go
+++ b/internal/gossip/gossip_inmem.go
@@ -24,6 +24,9 @@ var _ Gossiper = &InMemoryGossip{}
 
 func (g *InMemoryGossip) Publish(channel Channel, value []byte) error {
 	msg := NewMessage(channel, value)
+	if len(msg) == 0 {
+		return errors.New("empty message")
+	}
 	select {
 	case <-g.done:
 		return errors.New("gossip has been stopped")
@@ -59,6 +62,7 @@ func (g *InMemoryGossip) Start() error {
 			select {
 			case <-g.done:
 				return nil
+
 			case value := <-g.gossipCh:
 				msg := Message(value)
 				g.mut.RLock()
diff --git a/internal/gossip/gossip_redis.go b/internal/gossip/gossip_redis.go
index f31d742..4d841f1 100644
--- a/internal/gossip/gossip_redis.go
+++ b/internal/gossip/gossip_redis.go
@@ -50,9 +50,9 @@ func (g *GossipRedis) Start() error {
 	// This are all to avoid collisions with the ASCII values of the known channels and make it possible
 	// for tests to run in parallel.
 	g.channels = map[string]Channel{
-		ChannelKeep:   'k',
-		ChannelDrop:   'd',
-		ChannelStress: 's',
+		ChannelKeep: 'k',
+		ChannelDrop: 'd',
+		ChannelPeer: 'p',
 	}
 
 	g.Health.Register(gossipRedisHealth, redis.HealthCheckPeriod*2)
@@ -116,6 +116,11 @@ func (g *GossipRedis) Subscribe(channel Channel, depth int) chan []byte {
 
 // Publish sends a message to all subscribers of a given channel.
 func (g *GossipRedis) Publish(channel Channel, value []byte) error {
+	msg := NewMessage(channel, value)
+	if len(msg) == 0 {
+		return errors.New("empty message")
+	}
+
 	select {
 	case <-g.done:
 		return errors.New("gossip has been stopped")
@@ -125,8 +130,6 @@ func (g *GossipRedis) Publish(channel Channel, value []byte) error {
 	conn := g.Redis.GetPubSubConn()
 	defer conn.Close()
 
-	msg := NewMessage(channel, value)
-
 	return conn.Publish("refinery-gossip", msg.Bytes())
 }
 
diff --git a/internal/peer/peers_service.go b/internal/peer/peers_service.go
new file mode 100644
index 0000000..2280d20
--- /dev/null
+++ b/internal/peer/peers_service.go
@@ -0,0 +1,48 @@
+package peer
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/facebookgo/startstop"
+)
+
+// Peers provides information about all peers in a Refinery cluster.
+type Peers interface {
+	// Subscribe returns a channel that will receive updates from other active peers.
+	Subscribe() <-chan PeerInfo
+	// PublishPeerInfo sends information about the local peer to other peers.
+	PublishPeerInfo(PeerInfo) error
+	// GetPeerCount returns the number of active peers.
+	// This includes the local peer.
+	GetPeerCount() int
+	// HostID returns the unique identifier for the local peer.
+	HostID() string
+
+	startstop.Starter
+	startstop.Stopper
+}
+
+type PeerInfo struct {
+	id   string
+	Data []byte
+}
+
+func (p PeerInfo) ID() string {
+	return p.id
+}
+
+func (p PeerInfo) Bytes() []byte {
+	return append([]byte(p.id+"/"), p.Data...)
+}
+
+func newPeerInfoFromBytes(b []byte) (PeerInfo, error) {
+	parts := bytes.SplitN(b, []byte("/"), 2)
+	if len(parts) != 2 {
+		return PeerInfo{}, fmt.Errorf("invalid message format: %s", b)
+	}
+	return PeerInfo{
+		id:   string(parts[0]),
+		Data: parts[1],
+	}, nil
+}
diff --git a/internal/peer/store.go b/internal/peer/store.go
new file mode 100644
index 0000000..a782c87
--- /dev/null
+++ b/internal/peer/store.go
@@ -0,0 +1,217 @@
+package peer
+
+import (
+	"errors"
+	"maps"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/gofrs/uuid/v5"
+	"github.com/honeycombio/refinery/internal/gossip"
+	"github.com/jonboulle/clockwork"
+
+	"github.com/sirupsen/logrus"
+)
+
+const (
+	// keepAliveInterval is how frequently this host will announce itself
+	// with the rest of the cluster. This is a heartbeat to ensure that the host
+	// is still alive and participating in the cluster when no active peer info is
+	// being sent.
+	keepAliveInterval = 3 * time.Second
+
+	// peerEntryTimeout is how long to wait before expiring a peer that
+	// doesn't check in.
+	peerEntryTimeout = 10 * time.Second
+
+	defaultPeerInfoChannelSize = 100
+)
+
+var _ Peers = (*PeerStore)(nil)
+
+type PeerStore struct {
+	Gossip gossip.Gossiper `inject:"gossip"`
+	Clock  clockwork.Clock `inject:""`
+
+	done           chan struct{}
+	identification string
+
+	peerChan chan []byte
+
+	peerLock sync.RWMutex
+	peers    map[string]time.Time
+
+	publishChan chan PeerInfo
+
+	subscriptionLock sync.RWMutex
+	subscriptions    []chan PeerInfo
+
+	wg sync.WaitGroup
+}
+
+func (p *PeerStore) Start() error {
+	p.identification = p.HostID()
+	p.peerChan = p.Gossip.Subscribe(p.Gossip.GetChannel(gossip.ChannelPeer), defaultPeerInfoChannelSize)
+	p.publishChan = make(chan PeerInfo, defaultPeerInfoChannelSize)
+	p.done = make(chan struct{})
+	p.peers = make(map[string]time.Time)
+	p.wg = sync.WaitGroup{}
+	p.subscriptions = make([]chan PeerInfo, 0)
+
+	p.wg.Add(1)
+	go p.watchPeers()
+	p.wg.Add(1)
+	go p.publish()
+
+	return nil
+}
+
+func (p *PeerStore) Stop() error {
+	close(p.done)
+	close(p.publishChan)
+	p.wg.Wait()
+	return nil
+}
+
+// Subscribe adds a callback to be called when new peer info is received.
+func (p *PeerStore) Subscribe() <-chan PeerInfo {
+	ch := make(chan PeerInfo, defaultPeerInfoChannelSize)
+	p.subscriptionLock.Lock()
+	p.subscriptions = append(p.subscriptions, ch)
+	p.subscriptionLock.Unlock()
+	return ch
+}
+
+// PublishPeerInfo sends peer info to the rest of the cluster.
+func (p *PeerStore) PublishPeerInfo(info PeerInfo) error {
+	info.id = p.identification
+	select {
+	case <-p.done:
+		return errors.New("peer store has been stopped")
+	case p.publishChan <- info:
+		return nil
+	default:
+		return errors.New("failed to publish peer info")
+	}
+}
+
+// GetPeerCount returns the number of active peers in the cluster.
+// This includes the current host.
+func (p *PeerStore) GetPeerCount() int {
+	p.peerLock.Lock()
+	defer p.peerLock.Unlock()
+
+	maps.DeleteFunc(p.peers, func(id string, ts time.Time) bool {
+		return p.Clock.Since(ts) > peerEntryTimeout
+	})
+
+	// If we're the only peer, we won't be in the peer list, so we should return 1.
+	if len(p.peers) == 0 {
+		return 1
+	}
+
+	return len(p.peers) + 1
+}
+
+// HostID returns the unique identifier for this host.
+// If the host has a hostname, it will be used as part of the identifier.
+// Otherwise, a UUID will be generated.
+func (p *PeerStore) HostID() string {
+	if p.identification != "" {
+		return p.identification
+	}
+
+	var identification string
+	hostname, err := os.Hostname()
+	if err == nil && hostname != "" {
+		identification = hostname
+	}
+	id, err := uuid.NewV7()
+	if err != nil {
+		panic("failed to generate a UUID for the StressRelief system")
+	}
+
+	if identification == "" {
+		return id.String()
+	}
+
+	return identification + "-" + id.String()
+}
+
+// publish sends peer info from this host to the rest of the cluster.
+// If there is no new peer info to send, it will send a heartbeat.
+func (p *PeerStore) publish() {
+	defer p.wg.Done()
+
+	tk := p.Clock.NewTicker(keepAliveInterval)
+	defer tk.Stop()
+	peerChan := p.Gossip.GetChannel(gossip.ChannelPeer)
+
+	for {
+		var msg PeerInfo
+		select {
+		case <-p.done:
+			return
+		case data := <-p.publishChan:
+			msg = data
+			tk.Reset(keepAliveInterval)
+		case <-tk.Chan():
+			msg = PeerInfo{id: p.identification}
+		}
+
+		if err := p.Gossip.Publish(peerChan, msg.Bytes()); err != nil {
+			logrus.WithError(err).Error("failed to publish peer info")
+		}
+	}
+}
+
+// watchPeers listens for new peer info from the cluster and triggers callbacks.
+func (p *PeerStore) watchPeers() {
+	defer p.wg.Done()
+
+	for {
+		select {
+		case msg := <-p.peerChan:
+			info, err := newPeerInfoFromBytes(msg)
+			if err != nil {
+				logrus.WithError(err).Error("failed to decode peer info")
+				continue
+			}
+
+			// Ignore our own messages.
+			if info.id == p.identification {
+				continue
+			}
+
+			p.subscriptionLock.RLock()
+			for _, ch := range p.subscriptions {
+				ch <- info
+			}
+			p.subscriptionLock.RUnlock()
+
+			p.peerLock.Lock()
+			p.peers[info.id] = p.Clock.Now()
+			p.peerLock.Unlock()
+
+		case <-p.done:
+			return
+		}
+	}
+}
+
+type MockPeerStore struct {
+	PeerStore
+
+	Identification string
+	PeerCount      int
+}
+
+func (m *MockPeerStore) Start() error {
+	m.PeerStore.identification = m.Identification
+	return m.PeerStore.Start()
+}
+
+func (m *MockPeerStore) GetPeerCount() int {
+	return m.PeerCount
+}
diff --git a/sample/sample.go b/sample/sample.go
index 420c38c..298e56b 100644
--- a/sample/sample.go
+++ b/sample/sample.go
@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/honeycombio/refinery/config"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/types"
@@ -31,28 +32,23 @@ type ClusterSizer interface {
 
 // SamplerFactory is used to create new samplers with common (injected) resources
 type SamplerFactory struct {
-	Config    config.Config   `inject:""`
-	Logger    logger.Logger   `inject:""`
-	Metrics   metrics.Metrics `inject:"genericMetrics"`
-	peerCount int
-	samplers  []Sampler
+	Config   config.Config   `inject:""`
+	Logger   logger.Logger   `inject:""`
+	Metrics  metrics.Metrics `inject:"genericMetrics"`
+	Peer     peer.Peers      `inject:""`
+	samplers []Sampler
 }
 
 func (s *SamplerFactory) updatePeerCounts() {
 	// all the samplers who want it should use the stored count
 	for _, sampler := range s.samplers {
 		if clusterSizer, ok := sampler.(ClusterSizer); ok {
-			clusterSizer.SetClusterSize(s.peerCount)
+			peerCount := s.Peer.GetPeerCount()
+			clusterSizer.SetClusterSize(peerCount)
 		}
 	}
 }
 
-func (s *SamplerFactory) Start() error {
-	s.peerCount = 1
-	// TODO: register updatePeerCounts to be called whenever the peer count changes
-	return nil
-}
-
 // GetSamplerImplementationForKey returns the sampler implementation for the given
 // samplerKey (dataset for legacy keys, environment otherwise), or nil if it is not defined
 func (s *SamplerFactory) GetSamplerImplementationForKey(samplerKey string) Sampler {


Test patch:
diff --git a/app/app_test.go b/app/app_test.go
index ca4d57f..2d660fc 100644
--- a/app/app_test.go
+++ b/app/app_test.go
@@ -33,6 +33,7 @@ import (
 	"github.com/honeycombio/refinery/config"
 	"github.com/honeycombio/refinery/internal/gossip"
 	"github.com/honeycombio/refinery/internal/health"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/redis"
@@ -183,6 +184,7 @@ func newStartedApp(
 		&inject.Object{Value: samplerFactory},
 		&inject.Object{Value: &health.Health{}},
 		&inject.Object{Value: &stressRelief.StressRelief{}, Name: "stressRelief"},
+		&inject.Object{Value: &peer.PeerStore{}},
 		&inject.Object{Value: &a},
 	)
 	require.NoError(t, err)
diff --git a/centralstore/smartwrapper_test.go b/centralstore/smartwrapper_test.go
index bad244c..71a664f 100644
--- a/centralstore/smartwrapper_test.go
+++ b/centralstore/smartwrapper_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/facebookgo/startstop"
 	"github.com/honeycombio/refinery/collect/cache"
 	"github.com/honeycombio/refinery/config"
+	"github.com/honeycombio/refinery/generics"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/redis"
@@ -331,11 +332,11 @@ func TestSetTraceStatuses(t *testing.T) {
 
 			ctx := context.Background()
 			numberOfTraces := 5
-			traceids := make([]string, 0)
+			traceids := generics.NewSetWithCapacity[string](numberOfTraces)
 
 			for tr := 0; tr < numberOfTraces; tr++ {
 				tid := fmt.Sprintf("trace%02d", rand.Intn(1000))
-				traceids = append(traceids, tid)
+				traceids.Add(tid)
 				// write 9 child spans to the store
 				for s := 1; s < 10; s++ {
 					span := &CentralSpan{
@@ -355,15 +356,16 @@ func TestSetTraceStatuses(t *testing.T) {
 				require.NoError(t, err)
 			}
 
-			assert.Equal(t, numberOfTraces, len(traceids))
+			traces := traceids.Members()
+			assert.Equal(t, numberOfTraces, len(traces))
 			assert.Eventually(t, func() bool {
-				states, err := store.GetStatusForTraces(ctx, traceids, Collecting, DecisionDelay, ReadyToDecide)
+				states, err := store.GetStatusForTraces(ctx, traces, Collecting, DecisionDelay, ReadyToDecide)
 				return err == nil && len(states) == numberOfTraces
 			}, 1*time.Second, 100*time.Millisecond)
 
 			// wait for it to reach the Ready state
 			assert.EventuallyWithT(t, func(collect *assert.CollectT) {
-				states, err := store.GetStatusForTraces(ctx, traceids, ReadyToDecide)
+				states, err := store.GetStatusForTraces(ctx, traces, ReadyToDecide)
 				assert.NoError(collect, err)
 				assert.Equal(collect, numberOfTraces, len(states))
 				for _, state := range states {
@@ -377,7 +379,7 @@ func TestSetTraceStatuses(t *testing.T) {
 			assert.NoError(t, err)
 			assert.Equal(t, numberOfTraces, len(toDecide))
 			sort.Strings(toDecide)
-			expected := traceids[:numberOfTraces]
+			expected := traces[:numberOfTraces]
 			sort.Strings(expected)
 			assert.EqualValues(t, expected, toDecide)
 
@@ -397,7 +399,7 @@ func TestSetTraceStatuses(t *testing.T) {
 			}, 3*time.Second, 100*time.Millisecond)
 
 			for _, status := range statuses {
-				if status.TraceID == traceids[0] {
+				if status.TraceID == traces[0] {
 					status.State = DecisionKeep
 					status.KeepReason = "because"
 				} else {
@@ -410,11 +412,11 @@ func TestSetTraceStatuses(t *testing.T) {
 
 			// we need to give the dropped traces cache a chance to run or it might not process everything
 			time.Sleep(50 * time.Millisecond)
-			statuses, err = store.GetStatusForTraces(ctx, traceids, DecisionDrop, DecisionKeep)
+			statuses, err = store.GetStatusForTraces(ctx, traces, DecisionDrop, DecisionKeep)
 			assert.NoError(t, err)
 			assert.Equal(t, numberOfTraces, len(statuses))
 			for _, status := range statuses {
-				if status.TraceID == traceids[0] {
+				if status.TraceID == traces[0] {
 					assert.Equal(t, DecisionKeep, status.State)
 					assert.Equal(t, "because", status.KeepReason)
 				} else {
diff --git a/collect/central_collector_test.go b/collect/central_collector_test.go
index 4ddb59a..66417fa 100644
--- a/collect/central_collector_test.go
+++ b/collect/central_collector_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/honeycombio/refinery/generics"
 	"github.com/honeycombio/refinery/internal/gossip"
 	"github.com/honeycombio/refinery/internal/health"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/redis"
@@ -1444,6 +1445,7 @@ func startCollector(t *testing.T, cfg *config.MockConfig, collector *CentralColl
 		{Value: collector},
 		{Value: &health.Health{}},
 		{Value: &gossip.InMemoryGossip{}, Name: "gossip"},
+		{Value: &peer.PeerStore{}},
 	}
 	g := inject.Graph{}
 	require.NoError(t, g.Provide(objects...))
diff --git a/collect/stressRelief/stress_relief_test.go b/collect/stressRelief/stress_relief_test.go
index e9ad2da..da9e4ed 100644
--- a/collect/stressRelief/stress_relief_test.go
+++ b/collect/stressRelief/stress_relief_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/honeycombio/refinery/config"
 	"github.com/honeycombio/refinery/internal/gossip"
 	"github.com/honeycombio/refinery/internal/health"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/redis"
@@ -21,7 +22,7 @@ import (
 // for a given metrics data.
 func TestStressRelief_Monitor(t *testing.T) {
 	clock := clockwork.NewFakeClock()
-	sr, stop := newStressRelief(t, clock)
+	sr, stop := newStressRelief(t, clock, nil)
 	defer stop()
 	require.NoError(t, sr.Start())
 	defer sr.Stop()
@@ -70,7 +71,9 @@ func TestStressRelief_Monitor(t *testing.T) {
 // based on the stress level of a peer.
 func TestStressRelief_Peer(t *testing.T) {
 	clock := clockwork.NewFakeClock()
-	sr, stop := newStressRelief(t, clock)
+	channel := &gossip.InMemoryGossip{}
+
+	sr, stop := newStressRelief(t, clock, channel)
 	defer stop()
 	require.NoError(t, sr.Start())
 	defer sr.Stop()
@@ -101,13 +104,22 @@ func TestStressRelief_Peer(t *testing.T) {
 
 	// when a peer just started up, it should not affect the stress level of the
 	// cluster overall stress level
+	peer2 := &peer.MockPeerStore{
+		Identification: "peer2",
+		PeerStore: peer.PeerStore{
+			Gossip: channel,
+			Clock:  clock,
+		},
+	}
+	peer2.Start()
+	defer peer2.Stop()
+
 	require.Eventually(t, func() bool {
 		// pretend another refinery just started up
-		msg := stressLevelMessage{
-			level: 0,
-			id:    "peer",
+		msg := peer.PeerInfo{
+			Data: peerMessageToBytes(0),
 		}
-		require.NoError(t, sr.Gossip.Publish(sr.Gossip.GetChannel(gossip.ChannelStress), msg.ToBytes()))
+		require.NoError(t, peer2.PublishPeerInfo(msg))
 		clock.Advance(time.Second * 1)
 		return sr.Stressed()
 	}, 2*time.Second, 100*time.Millisecond, "stress relief should be false")
@@ -115,11 +127,10 @@ func TestStressRelief_Peer(t *testing.T) {
 	// now the peer has reported valid stress level
 	// it should be taken into account for the overall stress level
 	require.Eventually(t, func() bool {
-		msg := stressLevelMessage{
-			level: 5,
-			id:    "peer",
+		msg := peer.PeerInfo{
+			Data: peerMessageToBytes(5),
 		}
-		require.NoError(t, sr.Gossip.Publish(sr.Gossip.GetChannel(gossip.ChannelStress), msg.ToBytes()))
+		require.NoError(t, peer2.PublishPeerInfo(msg))
 
 		clock.Advance(time.Second * 1)
 		return !sr.Stressed()
@@ -168,7 +179,7 @@ func TestStressRelief_ShouldSampleDeterministically(t *testing.T) {
 	require.False(t, sr.ShouldSampleDeterministically(droppedTraceID))
 }
 
-func newStressRelief(t *testing.T, clock clockwork.Clock) (*StressRelief, func()) {
+func newStressRelief(t *testing.T, clock clockwork.Clock, channel gossip.Gossiper) (*StressRelief, func()) {
 	// Create a new StressRelief object
 	metric := &metrics.MockMetrics{}
 	metric.Start()
@@ -183,29 +194,37 @@ func newStressRelief(t *testing.T, clock clockwork.Clock) (*StressRelief, func()
 		},
 		Metrics: metric,
 	}
+	require.NoError(t, db.Start())
+
 	healthCheck := &health.Health{
 		Clock: clock,
 	}
 	require.NoError(t, healthCheck.Start())
-	channel := &gossip.GossipRedis{
-		Redis:  db,
-		Health: healthCheck,
+	if channel == nil {
+		channel = &gossip.GossipRedis{
+			Redis:  db,
+			Health: healthCheck,
+		}
 	}
+	peers := &peer.PeerStore{
+		Gossip: channel,
+		Clock:  clock,
+	}
+	require.NoError(t, channel.Start())
+	require.NoError(t, peers.Start())
 	sr := &StressRelief{
-		Gossip:          channel,
+		Peer:            peers,
 		Clock:           clock,
 		Logger:          &logger.NullLogger{},
 		RefineryMetrics: metric,
 		Health:          healthCheck,
 	}
 
-	require.NoError(t, db.Start())
-	require.NoError(t, channel.Start())
-
 	return sr, func() {
+		require.NoError(t, peers.Stop())
 		require.NoError(t, channel.Stop())
 		require.NoError(t, db.Stop())
-		require.NoError(t, healthCheck.Stop())
 		require.NoError(t, metric.Stop())
+		require.NoError(t, healthCheck.Stop())
 	}
 }
diff --git a/internal/peer/peers_service_test.go b/internal/peer/peers_service_test.go
new file mode 100644
index 0000000..20e045b
--- /dev/null
+++ b/internal/peer/peers_service_test.go
@@ -0,0 +1,92 @@
+package peer
+
+import (
+	"testing"
+	"time"
+
+	"github.com/honeycombio/refinery/internal/gossip"
+	"github.com/honeycombio/refinery/logger"
+	"github.com/jonboulle/clockwork"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPeers(t *testing.T) {
+	clock := clockwork.NewFakeClock()
+	gossip := &gossip.InMemoryGossip{}
+	gossip.Start()
+	peer1 := &PeerStore{
+		identification: "peer1",
+		Clock:          clock,
+		Gossip:         gossip,
+	}
+
+	peer2 := &PeerStore{
+		identification: "peer2",
+		Clock:          clock,
+		Gossip:         gossip,
+	}
+	require.NoError(t, peer1.Start())
+	require.NoError(t, peer2.Start())
+
+	peer1Chan := peer1.Subscribe()
+	peer2Chan := peer2.Subscribe()
+
+	peer1.PublishPeerInfo(PeerInfo{
+		Data: []byte("peer1"),
+	})
+
+	peer2.PublishPeerInfo(PeerInfo{
+		Data: []byte("peer2"),
+	})
+
+	for msg := range peer1Chan {
+		require.Equal(t, []byte("peer2"), msg.Data)
+		break
+	}
+
+	for msg := range peer2Chan {
+		require.Equal(t, []byte("peer1"), msg.Data)
+		break
+	}
+
+	require.Equal(t, 2, peer1.GetPeerCount())
+	require.Equal(t, 2, peer2.GetPeerCount())
+}
+
+func TestPeer_GetPeerCount(t *testing.T) {
+	// be able to get the count of peers
+	peer1Clock := clockwork.NewFakeClock()
+	gossip := &gossip.InMemoryGossip{
+		Logger: &logger.NullLogger{},
+	}
+	gossip.Start()
+	peer1 := &PeerStore{
+		identification: "peer1",
+		Clock:          peer1Clock,
+		Gossip:         gossip,
+	}
+
+	peer2Clock := clockwork.NewFakeClock()
+	peer2 := &PeerStore{
+		identification: "peer2",
+		Clock:          peer2Clock,
+		Gossip:         gossip,
+	}
+	require.NoError(t, peer1.Start())
+	require.NoError(t, peer2.Start())
+
+	require.EventuallyWithT(t, func(collect *assert.CollectT) {
+		peer1Clock.Advance(2 * keepAliveInterval)
+		peer2Clock.Advance(2 * keepAliveInterval)
+
+		assert.Equal(collect, 2, peer1.GetPeerCount())
+		assert.Equal(collect, 2, peer2.GetPeerCount())
+
+	}, 200*time.Millisecond, 10*time.Millisecond)
+
+	require.EventuallyWithT(t, func(collect *assert.CollectT) {
+		peer1Clock.Advance(10 * peerEntryTimeout)
+		require.Equal(t, 1, peer1.GetPeerCount())
+	}, 100*time.Millisecond, 10*time.Millisecond)
+}
diff --git a/route/route_test.go b/route/route_test.go
index 7345e46..598d194 100644
--- a/route/route_test.go
+++ b/route/route_test.go
@@ -20,6 +20,7 @@ import (
 	"github.com/honeycombio/refinery/config"
 	"github.com/honeycombio/refinery/internal/gossip"
 	"github.com/honeycombio/refinery/internal/health"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/redis"
@@ -502,6 +503,7 @@ func TestDependencyInjection(t *testing.T) {
 		&inject.Object{Value: &cache.CuckooSentCache{}},
 		&inject.Object{Value: &health.Health{}},
 		&inject.Object{Value: &gossip.InMemoryGossip{}, Name: "gossip"},
+		&inject.Object{Value: &peer.PeerStore{}},
 	)
 	if err != nil {
 		t.Error(err)
diff --git a/sample/sample_test.go b/sample/sample_test.go
index c1eeefb..7b1642e 100644
--- a/sample/sample_test.go
+++ b/sample/sample_test.go
@@ -7,9 +7,12 @@ import (
 
 	"github.com/facebookgo/inject"
 	"github.com/honeycombio/refinery/config"
+	"github.com/honeycombio/refinery/internal/gossip"
+	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
 	"github.com/honeycombio/refinery/types"
+	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/assert"
 	"gopkg.in/yaml.v3"
 )
@@ -74,7 +77,9 @@ func TestDependencyInjection(t *testing.T) {
 	var g inject.Graph
 	err := g.Provide(
 		&inject.Object{Value: &SamplerFactory{}},
-
+		&inject.Object{Value: clockwork.NewRealClock()},
+		&inject.Object{Value: &gossip.InMemoryGossip{}, Name: "gossip"},
+		&inject.Object{Value: &peer.PeerStore{}},
 		&inject.Object{Value: &config.MockConfig{}},
 		&inject.Object{Value: &logger.NullLogger{}},
 		&inject.Object{Value: &metrics.NullMetrics{}, Name: "genericMetrics"},
@@ -106,7 +111,6 @@ func TestDatasetPrefix(t *testing.T) {
 }
 
 func TestTotalThroughputClusterSize(t *testing.T) {
-	t.Skip()
 	cm := makeYAML(
 		"General/ConfigurationVersion", 2,
 	)
@@ -124,10 +128,10 @@ func TestTotalThroughputClusterSize(t *testing.T) {
 
 	factory := SamplerFactory{
 		Config:  c,
+		Peer:    &peer.MockPeerStore{PeerCount: 2},
 		Logger:  &logger.NullLogger{},
 		Metrics: &metrics.NullMetrics{},
 	}
-	factory.Start()
 	sampler := factory.GetSamplerImplementationForKey("production")
 	sampler.Start()
 	assert.NotNil(t, sampler)
@@ -137,7 +141,6 @@ func TestTotalThroughputClusterSize(t *testing.T) {
 }
 
 func TestEMAThroughputClusterSize(t *testing.T) {
-	t.Skip()
 	cm := makeYAML(
 		"General/ConfigurationVersion", 2,
 	)
@@ -155,10 +158,10 @@ func TestEMAThroughputClusterSize(t *testing.T) {
 
 	factory := SamplerFactory{
 		Config:  c,
+		Peer:    &peer.MockPeerStore{PeerCount: 2},
 		Logger:  &logger.NullLogger{},
 		Metrics: &metrics.NullMetrics{},
 	}
-	factory.Start()
 	sampler := factory.GetSamplerImplementationForKey("production")
 	sampler.Start()
 	assert.NotNil(t, sampler)
@@ -168,7 +171,6 @@ func TestEMAThroughputClusterSize(t *testing.T) {
 }
 
 func TestWindowedThroughputClusterSize(t *testing.T) {
-	t.Skip()
 	cm := makeYAML(
 		"General/ConfigurationVersion", 2,
 	)
@@ -186,10 +188,10 @@ func TestWindowedThroughputClusterSize(t *testing.T) {
 
 	factory := SamplerFactory{
 		Config:  c,
+		Peer:    &peer.MockPeerStore{PeerCount: 2},
 		Logger:  &logger.NullLogger{},
 		Metrics: &metrics.NullMetrics{},
 	}
-	factory.Start()
 	sampler := factory.GetSamplerImplementationForKey("production")
 	sampler.Start()
 	assert.NotNil(t, sampler)