Repository: honeycombio/refinery
Pull request #1268 (2024-08-12 18:30:23)

feat: redistribute traces on peer membership changes
<!--
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?

Draining incoming and peer queue on peer membership changes
#1258 

## Short description of the changes

A `RedistributeNotifier` is introduced in the collector in this PR. 
- It's initialized on refinery start up
- When a peer membership change is broadcasted, the notifier immediately signals the `collect()` goroutine to start the redistribution process
- When multiple peer membership changes happen in a row, it ignores the signal if there's already a redistribution process running




Patch:
diff --git a/collect/cache/cache.go b/collect/cache/cache.go
index 445b9aa..1c7c3ca 100644
--- a/collect/cache/cache.go
+++ b/collect/cache/cache.go
@@ -18,11 +18,19 @@ type Cache interface {
 	// GetAll is used during shutdown to get all in-flight traces to flush them
 	GetAll() []*types.Trace
 
+	// GetCacheCapacity returns the number of traces that can be stored in the cache
+	GetCacheCapacity() int
+
 	// Retrieve and remove all traces which are past their SendBy date.
 	// Does not check whether they've been sent.
 	TakeExpiredTraces(now time.Time) []*types.Trace
+
+	// RemoveTraces accepts a set of trace IDs and removes any matching ones from
+	RemoveTraces(toDelete generics.Set[string])
 }
 
+var _ Cache = (*DefaultInMemCache)(nil)
+
 // DefaultInMemCache keeps a bounded number of entries to avoid growing memory
 // forever. Traces are expunged from the cache in insertion order (not access
 // order) so it is important to have a cache larger than trace throughput *
@@ -68,7 +76,7 @@ func NewInMemCache(
 
 }
 
-func (d *DefaultInMemCache) GetCacheSize() int {
+func (d *DefaultInMemCache) GetCacheCapacity() int {
 	return len(d.traceBuffer)
 }
 
diff --git a/collect/collect.go b/collect/collect.go
index 197021e..0cbe6cd 100644
--- a/collect/collect.go
+++ b/collect/collect.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
+	"math/rand"
 	"os"
 	"runtime"
 	"sort"
@@ -17,6 +19,7 @@ import (
 	"github.com/honeycombio/refinery/generics"
 	"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/sample"
@@ -68,6 +71,7 @@ type InMemCollector struct {
 	Metrics        metrics.Metrics        `inject:"genericMetrics"`
 	SamplerFactory *sample.SamplerFactory `inject:""`
 	StressRelief   StressReliever         `inject:"stressRelief"`
+	Peers          peer.Peers             `inject:""`
 
 	// For test use only
 	BlockOnAddSpan bool
@@ -75,16 +79,17 @@ type InMemCollector struct {
 	// mutex must be held whenever non-channel internal fields are accessed.
 	// This exists to avoid data races in tests and startup/shutdown.
 	mutex sync.RWMutex
+	cache cache.Cache
 
-	cache           cache.Cache
 	datasetSamplers map[string]sample.Sampler
 
 	sampleTraceCache cache.TraceSentCache
 
-	incoming chan *types.Span
-	fromPeer chan *types.Span
-	reload   chan struct{}
-	done     chan struct{}
+	incoming          chan *types.Span
+	fromPeer          chan *types.Span
+	reload            chan struct{}
+	done              chan struct{}
+	redistributeTimer *redistributeNotifier
 
 	hostname string
 }
@@ -118,6 +123,9 @@ func (i *InMemCollector) Start() error {
 	i.Metrics.Register("trace_send_dropped", "counter")
 	i.Metrics.Register("trace_send_has_root", "counter")
 	i.Metrics.Register("trace_send_no_root", "counter")
+	i.Metrics.Register("trace_forwarded_on_peer_change", "gauge")
+	i.Metrics.Register("trace_send_on_shutdown", "gauge")
+	i.Metrics.Register("trace_forwarded_on_shutdown", "gauge")
 
 	i.Metrics.Register(TraceSendGotRoot, "counter")
 	i.Metrics.Register(TraceSendExpired, "counter")
@@ -143,6 +151,8 @@ func (i *InMemCollector) Start() error {
 	i.reload = make(chan struct{}, 1)
 	i.done = make(chan struct{})
 	i.datasetSamplers = make(map[string]sample.Sampler)
+	i.done = make(chan struct{})
+	i.redistributeTimer = newRedistributeNotifier(i.Logger, i.Clock)
 
 	if i.Config.GetAddHostMetadataToTrace() {
 		if hostname, err := os.Hostname(); err == nil && hostname != "" {
@@ -150,6 +160,10 @@ func (i *InMemCollector) Start() error {
 		}
 	}
 
+	if !i.Config.GetCollectionConfig().DisableRedistribution {
+		i.Peers.RegisterUpdatedPeersCallback(i.redistributeTimer.Reset)
+	}
+
 	// spin up one collector because this is a single threaded collector
 	go i.collect()
 
@@ -171,28 +185,24 @@ func (i *InMemCollector) reloadConfigs() {
 	i.Logger.Debug().Logf("reloading in-mem collect config")
 	imcConfig := i.Config.GetCollectionConfig()
 
-	if existingCache, ok := i.cache.(*cache.DefaultInMemCache); ok {
-		if imcConfig.CacheCapacity != existingCache.GetCacheSize() {
-			i.Logger.Debug().WithField("cache_size.previous", existingCache.GetCacheSize()).WithField("cache_size.new", imcConfig.CacheCapacity).Logf("refreshing the cache because it changed size")
-			c := cache.NewInMemCache(imcConfig.CacheCapacity, i.Metrics, i.Logger)
-			// pull the old cache contents into the new cache
-			for j, trace := range existingCache.GetAll() {
-				if j >= imcConfig.CacheCapacity {
-					i.sendTrace(trace, TraceSendEjectedFull)
-					continue
-				}
-				c.Set(trace)
+	if imcConfig.CacheCapacity != i.cache.GetCacheCapacity() {
+		i.Logger.Debug().WithField("cache_size.previous", i.cache.GetCacheCapacity()).WithField("cache_size.new", imcConfig.CacheCapacity).Logf("refreshing the cache because it changed size")
+		c := cache.NewInMemCache(imcConfig.CacheCapacity, i.Metrics, i.Logger)
+		// pull the old cache contents into the new cache
+		for j, trace := range i.cache.GetAll() {
+			if j >= imcConfig.CacheCapacity {
+				i.send(trace, TraceSendEjectedFull)
+				continue
 			}
-			i.cache = c
-		} else {
-			i.Logger.Debug().Logf("skipping reloading the in-memory cache on config reload because it hasn't changed capacity")
+			c.Set(trace)
 		}
-
-		i.sampleTraceCache.Resize(i.Config.GetSampleCacheConfig())
+		i.cache = c
 	} else {
-		i.Logger.Error().WithField("cache", i.cache.(*cache.DefaultInMemCache)).Logf("skipping reloading the cache on config reload because it's not an in-memory cache")
+		i.Logger.Debug().Logf("skipping reloading the in-memory cache on config reload because it hasn't changed capacity")
 	}
 
+	i.sampleTraceCache.Resize(i.Config.GetSampleCacheConfig())
+
 	i.StressRelief.UpdateFromConfig(i.Config.GetStressReliefConfig())
 
 	// clear out any samplers that we have previously created
@@ -223,14 +233,7 @@ func (i *InMemCollector) checkAlloc() {
 	// remove the traces from the cache that have had the most impact on allocation.
 	// To do this, we sort the traces by their CacheImpact value and then remove traces
 	// until the total size is less than the amount to which we want to shrink.
-	existingCache, ok := i.cache.(*cache.DefaultInMemCache)
-	if !ok {
-		i.Logger.Error().WithField("alloc", mem.Alloc).Logf(
-			"total allocation exceeds limit, but unable to control cache",
-		)
-		return
-	}
-	allTraces := existingCache.GetAll()
+	allTraces := i.cache.GetAll()
 	timeout := i.Config.GetTracesConfig().GetTraceTimeout()
 	if timeout == 0 {
 		timeout = 60 * time.Second
@@ -243,7 +246,7 @@ func (i *InMemCollector) checkAlloc() {
 	// successive traces until we've crossed the totalToRemove threshold
 	// or just run out of traces to delete.
 
-	cap := existingCache.GetCacheSize()
+	cap := i.cache.GetCacheCapacity()
 	i.Metrics.Gauge("collector_cache_size", cap)
 
 	totalDataSizeSent := 0
@@ -252,12 +255,12 @@ func (i *InMemCollector) checkAlloc() {
 	for _, trace := range allTraces {
 		tracesSent.Add(trace.TraceID)
 		totalDataSizeSent += trace.DataSize
-		i.sendTrace(trace, TraceSendEjectedMemsize)
+		i.send(trace, TraceSendEjectedMemsize)
 		if totalDataSizeSent > int(totalToRemove) {
 			break
 		}
 	}
-	existingCache.RemoveTraces(tracesSent)
+	i.cache.RemoveTraces(tracesSent)
 
 	// Treat any MaxAlloc overage as an error so we know it's happening
 	i.Logger.Error().
@@ -265,7 +268,7 @@ func (i *InMemCollector) checkAlloc() {
 		WithField("alloc", mem.Alloc).
 		WithField("num_traces_sent", len(tracesSent)).
 		WithField("datasize_sent", totalDataSizeSent).
-		WithField("new_trace_count", existingCache.GetCacheSize()).
+		WithField("new_trace_count", i.cache.GetCacheCapacity()).
 		Logf("evicting large traces early due to memory overage")
 
 	// Manually GC here - without this we can easily end up evicting more than we
@@ -340,6 +343,8 @@ func (i *InMemCollector) collect() {
 		select {
 		case <-i.done:
 			return
+		case <-i.redistributeTimer.Notify():
+			i.redistributeTraces()
 		case sp, ok := <-i.fromPeer:
 			if !ok {
 				// channel's been closed; we should shut down.
@@ -351,13 +356,19 @@ func (i *InMemCollector) collect() {
 			case <-i.done:
 				return
 			case <-ticker.C:
-				i.sendExpiredTracesInCache(i.Clock.Now())
-				i.checkAlloc()
-
-				// Briefly unlock the cache, to allow test access.
-				i.mutex.Unlock()
-				runtime.Gosched()
-				i.mutex.Lock()
+				select {
+				case <-i.done:
+				default:
+					i.sendExpiredTracesInCache(i.Clock.Now())
+					i.checkAlloc()
+
+					// Briefly unlock the cache, to allow test access.
+					i.mutex.Unlock()
+					runtime.Gosched()
+					i.mutex.Lock()
+				}
+			case <-i.redistributeTimer.Notify():
+				i.redistributeTraces()
 			case sp, ok := <-i.incoming:
 				if !ok {
 					// channel's been closed; we should shut down.
@@ -379,17 +390,75 @@ func (i *InMemCollector) collect() {
 	}
 }
 
+func (i *InMemCollector) redistributeTraces() {
+	_, span := otelutil.StartSpan(context.Background(), i.Tracer, "redistributeTraces")
+	defer span.End()
+	// loop through eveything in the cache of live traces
+	// if it doesn't belong to this peer, we should forward it to the correct peer
+	peers, err := i.Peers.GetPeers()
+	if err != nil {
+		i.Logger.Error().Logf("unable to get peer list with error %s", err.Error())
+		return
+	}
+	numOfPeers := len(peers)
+	if numOfPeers <= 1 {
+		return
+	}
+
+	traces := i.cache.GetAll()
+	forwardedTraces := generics.NewSetWithCapacity[string](len(traces) / numOfPeers)
+	for _, trace := range traces {
+		if trace == nil {
+			continue
+		}
+
+		newTarget := i.Sharder.WhichShard(trace.TraceID)
+
+		if newTarget.Equals(i.Sharder.MyShard()) {
+			continue
+		}
+
+		for _, sp := range trace.GetSpans() {
+			sp.APIHost = newTarget.GetAddress()
+
+			if sp.Data == nil {
+				sp.Data = make(map[string]interface{})
+			}
+			if v, ok := sp.Data["meta.refinery.forwarded"]; ok {
+				sp.Data["meta.refinery.forwarded"] = fmt.Sprintf("%s,%s", v, i.hostname)
+			} else {
+				sp.Data["meta.refinery.forwarded"] = i.hostname
+			}
+
+			i.Transmission.EnqueueSpan(sp)
+		}
+
+		forwardedTraces.Add(trace.TraceID)
+	}
+
+	otelutil.AddSpanFields(span, map[string]interface{}{
+		"forwarded_trace_count": len(forwardedTraces.Members()),
+		"total_trace_count":     len(traces),
+		"hostname":              i.hostname,
+	})
+
+	i.Metrics.Gauge("trace_forwarded_on_peer_change", len(forwardedTraces))
+	if len(forwardedTraces) > 0 {
+		i.cache.RemoveTraces(forwardedTraces)
+	}
+}
+
 func (i *InMemCollector) sendExpiredTracesInCache(now time.Time) {
 	traces := i.cache.TakeExpiredTraces(now)
 	spanLimit := uint32(i.Config.GetTracesConfig().SpanLimit)
 	for _, t := range traces {
 		if t.RootSpan != nil {
-			i.sendTrace(t, TraceSendGotRoot)
+			i.send(t, TraceSendGotRoot)
 		} else {
 			if t.SpanCount() > spanLimit {
-				i.sendTrace(t, TraceSendSpanLimit)
+				i.send(t, TraceSendSpanLimit)
 			} else {
-				i.sendTrace(t, TraceSendExpired)
+				i.send(t, TraceSendExpired)
 			}
 		}
 	}
@@ -439,7 +508,7 @@ func (i *InMemCollector) processSpan(sp *types.Span) {
 		// push this into the cache and if we eject an unsent trace, send it ASAP
 		ejectedTrace := i.cache.Set(trace)
 		if ejectedTrace != nil {
-			i.sendTrace(ejectedTrace, TraceSendEjectedFull)
+			i.send(ejectedTrace, TraceSendEjectedFull)
 		}
 	}
 	// if the trace we got back from the cache has already been sent, deal with the
@@ -653,18 +722,7 @@ func (i *InMemCollector) isRootSpan(sp *types.Span) bool {
 	return true
 }
 
-// drainTrace is called when a trace should be sent to honeycomb during shutdown.
-func (i *InMemCollector) drainTrace(trace *types.Trace, sendReason string) {
-	i.send(trace, sendReason, true)
-}
-
-// sendTrace is called when a trace is ready to be sent to honeycomb during normal operation.
-func (i *InMemCollector) sendTrace(trace *types.Trace, sendReason string) {
-	i.send(trace, sendReason, false)
-}
-
-// send is the actual workhorse of sending a trace to honeycomb. It's called by both drainTrace and sendTrace.
-func (i *InMemCollector) send(trace *types.Trace, sendReason string, onShutdown bool) {
+func (i *InMemCollector) send(trace *types.Trace, sendReason string) {
 	if trace.Sent {
 		// someone else already sent this so we shouldn't also send it.
 		i.Logger.Debug().
@@ -776,9 +834,6 @@ func (i *InMemCollector) send(trace *types.Trace, sendReason string, onShutdown
 		if i.hostname != "" {
 			sp.Data["meta.refinery.local_hostname"] = i.hostname
 		}
-		if onShutdown {
-			sp.Data["meta.refinery.shutdown.send"] = true
-		}
 		mergeTraceAndSpanSampleRates(sp, trace.SampleRate(), isDryRun)
 		i.addAdditionalAttributes(sp)
 		i.Transmission.EnqueueSpan(sp)
@@ -786,6 +841,7 @@ func (i *InMemCollector) send(trace *types.Trace, sendReason string, onShutdown
 }
 
 func (i *InMemCollector) Stop() error {
+	i.redistributeTimer.Stop()
 	close(i.done)
 	// signal the health system to not be ready
 	// so that no new traces are accepted
@@ -794,7 +850,13 @@ func (i *InMemCollector) Stop() error {
 	i.mutex.Lock()
 
 	if !i.Config.GetCollectionConfig().DisableRedistribution {
-		i.sendTracesOnShutdown()
+		peers, err := i.Peers.GetPeers()
+		if err != nil {
+			i.Logger.Error().Logf("unable to get peer list with error %s", err.Error())
+		}
+		if len(peers) > 1 {
+			i.sendTracesOnShutdown()
+		}
 	}
 
 	if i.Transmission != nil {
@@ -893,50 +955,58 @@ func (i *InMemCollector) sendTracesOnShutdown() {
 }
 
 // distributeSpansInCache takes a list of spans and sends them to the appropriate channel based on the state of the trace.
-func (i *InMemCollector) distributeSpansOnShutdown(sentTraceChan chan sentRecord, forwardTraceChan chan *types.Span, spans ...*types.Span) {
+func (i *InMemCollector) distributeSpansOnShutdown(sentSpanChan chan sentRecord, forwardSpanChan chan *types.Span, spans ...*types.Span) {
 	for _, sp := range spans {
 		if sp != nil {
 
 			// first check if there's a trace decision
-			record, reason, found := i.sampleTraceCache.CheckTrace(sp.TraceID)
+			record, reason, found := i.sampleTraceCache.CheckSpan(sp)
 			if found {
-				sentTraceChan <- sentRecord{sp, record, reason}
+				sentSpanChan <- sentRecord{sp, record, reason}
 				continue
 			}
 
 			// if there's no trace decision, then we need to forward the trace to its new home
-			forwardTraceChan <- sp
+			forwardSpanChan <- sp
 		}
 	}
 }
 
 // sendSpansOnShutdown is a helper function that sends span to their final destination
 // on shutdown.
-func (i *InMemCollector) sendSpansOnShutdown(ctx context.Context, sentTraceChan <-chan sentRecord, forwardTraceChan <-chan *types.Span) {
+func (i *InMemCollector) sendSpansOnShutdown(ctx context.Context, sentSpanChan <-chan sentRecord, forwardSpanChan <-chan *types.Span) {
+	sentTraces := make(map[string]struct{})
+	forwardedTraces := make(map[string]struct{})
+	defer func() {
+		i.Metrics.Gauge("trace_send_on_shutdown", len(sentTraces))
+		i.Metrics.Gauge("trace_forwarded_on_shutdown", len(forwardedTraces))
+	}()
+
 	for {
 		select {
 		case <-ctx.Done():
 			i.Logger.Info().Logf("Timed out waiting for traces to send")
 			return
 
-		case r, ok := <-sentTraceChan:
+		case r, ok := <-sentSpanChan:
 			if !ok {
 				return
 			}
 
-			ctx, span := otelutil.StartSpanMulti(ctx, i.Tracer, "shutdown_sent_trace", map[string]interface{}{"trace_id": r.span.TraceID, "hostname": i.hostname})
+			ctx, span := otelutil.StartSpanMulti(ctx, i.Tracer, "shutdown_sent_span", map[string]interface{}{"trace_id": r.span.TraceID, "hostname": i.hostname})
 			r.span.Data["meta.refinery.shutdown.send"] = true
 
 			i.dealWithSentTrace(ctx, r.record, r.reason, r.span)
+			sentTraces[r.span.TraceID] = struct{}{}
 
 			span.End()
 
-		case sp, ok := <-forwardTraceChan:
+		case sp, ok := <-forwardSpanChan:
 			if !ok {
 				return
 			}
 
-			_, span := otelutil.StartSpanMulti(ctx, i.Tracer, "shutdown_forwarded_trace", map[string]interface{}{"trace_id": sp.TraceID, "hostname": i.hostname})
+			_, span := otelutil.StartSpanMulti(ctx, i.Tracer, "shutdown_forwarded_span", map[string]interface{}{"trace_id": sp.TraceID, "hostname": i.hostname})
 
 			targetShard := i.Sharder.WhichShard(sp.TraceID)
 			url := targetShard.GetAddress()
@@ -947,8 +1017,18 @@ func (i *InMemCollector) sendSpansOnShutdown(ctx context.Context, sentTraceChan
 			// the downstream consumers can make decisions based on the metadata without having
 			// to restart the TraceTimeout or SendDelay
 			sp.APIHost = url
-			sp.Data["meta.refinery.shutdown.send"] = false
+
+			if sp.Data == nil {
+				sp.Data = make(map[string]interface{})
+			}
+			if v, ok := sp.Data["meta.refinery.forwarded"]; ok {
+				sp.Data["meta.refinery.forwarded"] = fmt.Sprintf("%s,%s", v, i.hostname)
+			} else {
+				sp.Data["meta.refinery.forwarded"] = i.hostname
+			}
+
 			i.Transmission.EnqueueSpan(sp)
+			forwardedTraces[sp.TraceID] = struct{}{}
 			span.End()
 		}
 
@@ -957,9 +1037,8 @@ func (i *InMemCollector) sendSpansOnShutdown(ctx context.Context, sentTraceChan
 
 // Convenience method for tests.
 func (i *InMemCollector) getFromCache(traceID string) *types.Trace {
-	i.mutex.RLock()
-	defer i.mutex.RUnlock()
-
+	i.mutex.Lock()
+	defer i.mutex.Unlock()
 	return i.cache.Get(traceID)
 }
 
@@ -968,3 +1047,100 @@ func (i *InMemCollector) addAdditionalAttributes(sp *types.Span) {
 		sp.Data[k] = v
 	}
 }
+
+func newRedistributeNotifier(logger logger.Logger, clock clockwork.Clock) *redistributeNotifier {
+	r := &redistributeNotifier{
+		initialDelay: 3 * time.Second,
+		maxDelay:     30 * time.Second,
+		maxAttempts:  5,
+		done:         make(chan struct{}),
+		clock:        clock,
+		logger:       logger,
+		triggered:    make(chan struct{}),
+		reset:        make(chan struct{}),
+	}
+
+	return r
+}
+
+type redistributeNotifier struct {
+	clock        clockwork.Clock
+	logger       logger.Logger
+	initialDelay time.Duration
+	maxAttempts  int
+	maxDelay     time.Duration
+
+	reset     chan struct{}
+	done      chan struct{}
+	triggered chan struct{}
+	once      sync.Once
+}
+
+func (r *redistributeNotifier) Notify() <-chan struct{} {
+	return r.triggered
+}
+
+func (r *redistributeNotifier) Reset() {
+	var started bool
+	r.once.Do(func() {
+		go r.run()
+		started = true
+	})
+
+	if started {
+		return
+	}
+
+	select {
+	case r.reset <- struct{}{}:
+	case <-r.done:
+		return
+	default:
+		r.logger.Debug().Logf("A trace redistribution is ongoing. Ignoring reset.")
+	}
+}
+
+func (r *redistributeNotifier) Stop() {
+	close(r.done)
+}
+
+func (r *redistributeNotifier) run() {
+	var attempts int
+	lastBackoff := r.initialDelay
+	for {
+		// if we've reached the max attempts, reset the backoff and attempts
+		// only when the reset signal is received.
+		if attempts >= r.maxAttempts {
+			<-r.reset
+			lastBackoff = r.initialDelay
+			attempts = 0
+		}
+		select {
+		case <-r.done:
+			return
+		case r.triggered <- struct{}{}:
+		}
+
+		attempts++
+
+		// Calculate the backoff interval using exponential backoff with a base time.
+		backoff := time.Duration(math.Min(float64(lastBackoff)*2, float64(r.maxDelay)))
+		// Add jitter to the backoff to avoid retry collisions.
+		jitter := time.Duration(rand.Float64() * float64(backoff) * 0.5)
+		nextBackoff := backoff + jitter
+		lastBackoff = nextBackoff
+
+		timer := r.clock.NewTimer(nextBackoff)
+		select {
+		case <-timer.Chan():
+			timer.Stop()
+		case <-r.reset:
+			lastBackoff = r.initialDelay
+			attempts = 0
+			timer.Stop()
+		case <-r.done:
+			timer.Stop()
+			return
+		}
+	}
+}
diff --git a/sharder/mock.go b/sharder/mock.go
new file mode 100644
index 0000000..c3c90b6
--- /dev/null
+++ b/sharder/mock.go
@@ -0,0 +1,23 @@
+package sharder
+
+type MockSharder struct {
+	Self  *TestShard
+	Other *TestShard
+}
+
+func (s *MockSharder) MyShard() Shard { return s.Self }
+
+func (s *MockSharder) WhichShard(traceID string) Shard {
+	if s.Other != nil {
+		return s.Other
+	}
+
+	return s.Self
+}
+
+type TestShard struct {
+	Addr string
+}
+
+func (s *TestShard) Equals(other Shard) bool { return s.Addr == other.GetAddress() }
+func (s *TestShard) GetAddress() string      { return s.Addr }


Test patch:
diff --git a/app/app_test.go b/app/app_test.go
index 7ed3679..3b3904c 100644
--- a/app/app_test.go
+++ b/app/app_test.go
@@ -109,7 +109,7 @@ func newStartedApp(
 		GetListenAddrVal:         "127.0.0.1:" + strconv.Itoa(basePort),
 		GetPeerListenAddrVal:     "127.0.0.1:" + strconv.Itoa(basePort+1),
 		GetHoneycombAPIVal:       "http://api.honeycomb.io",
-		GetCollectionConfigVal:   config.CollectionConfig{CacheCapacity: 10000},
+		GetCollectionConfigVal:   config.CollectionConfig{CacheCapacity: 10000, ShutdownDelay: config.Duration(1 * time.Second)},
 		AddHostMetadataToTrace:   enableHostMetadata,
 		TraceIdFieldNames:        []string{"trace.trace_id"},
 		ParentIdFieldNames:       []string{"trace.parent_id"},
@@ -445,8 +445,11 @@ func TestHostMetadataSpanAdditions(t *testing.T) {
 
 	time.Sleep(5 * app.Config.GetTracesConfig().GetSendTickerValue())
 
-	events := sender.Events()
-	require.Len(t, events, 1)
+	var events []*transmission.Event
+	require.Eventually(t, func() bool {
+		events = sender.Events()
+		return len(events) == 1
+	}, 2*time.Second, 10*time.Millisecond)
 
 	assert.Equal(t, "dataset", events[0].Dataset)
 	assert.Equal(t, "bar", events[0].Data["foo"])
diff --git a/collect/cache/cache_test.go b/collect/cache/cache_test.go
index 8ce6651..4ca35e8 100644
--- a/collect/cache/cache_test.go
+++ b/collect/cache/cache_test.go
@@ -1,6 +1,7 @@
 package cache
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
@@ -140,3 +141,88 @@ func TestSkipOldUnsentTraces(t *testing.T) {
 	// make sure we kicked out #4
 	assert.Equal(t, traces[3], prev)
 }
+
+// Benchamark the cache's Set method
+func BenchmarkCache_Set(b *testing.B) {
+	s := &metrics.MockMetrics{}
+	s.Start()
+	c := NewInMemCache(100000, s, &logger.NullLogger{})
+	now := time.Now()
+	traces := make([]*types.Trace, 0, b.N)
+	for i := 0; i < b.N; i++ {
+		traces = append(traces, &types.Trace{
+			TraceID: "trace" + fmt.Sprint(i),
+			SendBy:  now.Add(time.Duration(i) * time.Second),
+		})
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		c.Set(traces[i])
+	}
+}
+
+// Benchmark the cache's Get method
+func BenchmarkCache_Get(b *testing.B) {
+	s := &metrics.MockMetrics{}
+	s.Start()
+	c := NewInMemCache(100000, s, &logger.NullLogger{})
+	now := time.Now()
+	traces := make([]*types.Trace, 0, b.N)
+	for i := 0; i < b.N; i++ {
+		traces = append(traces, &types.Trace{
+			TraceID: "trace" + fmt.Sprint(i),
+			SendBy:  now.Add(time.Duration(i) * time.Second),
+		})
+		c.Set(traces[i])
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		c.Get(traces[i].TraceID)
+	}
+}
+
+// Benchmark the cache's TakeExpiredTraces method
+func BenchmarkCache_TakeExpiredTraces(b *testing.B) {
+	s := &metrics.MockMetrics{}
+	s.Start()
+	c := NewInMemCache(100000, s, &logger.NullLogger{})
+	now := time.Now()
+	traces := make([]*types.Trace, 0, b.N)
+	for i := 0; i < b.N; i++ {
+		traces = append(traces, &types.Trace{
+			TraceID: "trace" + fmt.Sprint(i),
+			SendBy:  now.Add(time.Duration(i) * time.Second),
+		})
+		c.Set(traces[i])
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		c.TakeExpiredTraces(now.Add(time.Duration(i) * time.Second))
+	}
+}
+
+// Benchmark the cache's RemoveTraces method
+func BenchmarkCache_RemoveTraces(b *testing.B) {
+	s := &metrics.MockMetrics{}
+	s.Start()
+	c := NewInMemCache(100000, s, &logger.NullLogger{})
+	now := time.Now()
+	traces := make([]*types.Trace, 0, b.N)
+	for i := 0; i < b.N; i++ {
+		traces = append(traces, &types.Trace{
+			TraceID: "trace" + fmt.Sprint(i),
+			SendBy:  now.Add(time.Duration(i) * time.Second),
+		})
+		c.Set(traces[i])
+	}
+
+	deletes := generics.NewSetWithCapacity[string](b.N / 2)
+	for i := 0; i < b.N/2; i++ {
+		deletes.Add("trace" + fmt.Sprint(i))
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		c.RemoveTraces(deletes)
+	}
+}
diff --git a/collect/collect_test.go b/collect/collect_test.go
index 26912a0..e05110c 100644
--- a/collect/collect_test.go
+++ b/collect/collect_test.go
@@ -50,14 +50,11 @@ func newTestCollector(conf config.Config, transmission transmit.Transmission) *I
 	healthReporter.Start()
 
 	return &InMemCollector{
-		Config: conf,
-		Clock:  clock,
-		Logger: &logger.NullLogger{},
-		Tracer: noop.NewTracerProvider().Tracer("test"),
-		Health: healthReporter,
-		Sharder: &sharder.SingleServerSharder{
-			Logger: &logger.NullLogger{},
-		},
+		Config:       conf,
+		Clock:        clock,
+		Logger:       &logger.NullLogger{},
+		Tracer:       noop.NewTracerProvider().Tracer("test"),
+		Health:       healthReporter,
 		Transmission: transmission,
 		Metrics:      &metrics.NullMetrics{},
 		StressRelief: &MockStressReliever{},
@@ -67,6 +64,15 @@ func newTestCollector(conf config.Config, transmission transmit.Transmission) *I
 			Logger:  &logger.NullLogger{},
 		},
 		done: make(chan struct{}),
+		Peers: &peer.MockPeers{
+			Peers: []string{"api1", "api2"},
+		},
+		Sharder: &sharder.MockSharder{
+			Self: &sharder.TestShard{
+				Addr: "api1",
+			},
+		},
+		redistributeTimer: newRedistributeNotifier(&logger.NullLogger{}, clock),
 	}
 }
 
@@ -351,8 +357,9 @@ func TestAddSpan(t *testing.T) {
 	}
 	coll.AddSpanFromPeer(span)
 	time.Sleep(conf.GetTracesConfig().GetSendTickerValue() * 2)
-
-	assert.Equal(t, traceID, coll.getFromCache(traceID).TraceID, "after adding the span, we should have a trace in the cache with the right trace ID")
+	trace := coll.getFromCache(traceID)
+	require.NotNil(t, trace)
+	assert.Equal(t, traceID, trace.TraceID, "after adding the span, we should have a trace in the cache with the right trace ID")
 	assert.Equal(t, 0, len(transmission.Events), "adding a non-root span should not yet send the span")
 	// ok now let's add the root span and verify that both got sent
 	rootSpan := &types.Span{
@@ -532,6 +539,7 @@ func TestCacheSizeReload(t *testing.T) {
 	transmission := &transmit.MockTransmission{}
 	transmission.Start()
 	coll := newTestCollector(conf, transmission)
+	coll.Peers = &peer.MockPeers{}
 
 	err := coll.Start()
 	assert.NoError(t, err)
@@ -568,8 +576,7 @@ func TestCacheSizeReload(t *testing.T) {
 	assert.Eventually(t, func() bool {
 		coll.mutex.RLock()
 		defer coll.mutex.RUnlock()
-
-		return coll.cache.(*cache.DefaultInMemCache).GetCacheSize() == 2
+		return coll.cache.GetCacheCapacity() == 2
 	}, 60*wait, wait, "cache size to change")
 
 	err = coll.AddSpan(&types.Span{TraceID: "3", Event: event})
@@ -626,9 +633,9 @@ func TestSampleConfigReload(t *testing.T) {
 
 	assert.Eventually(t, func() bool {
 		coll.mutex.Lock()
-		defer coll.mutex.Unlock()
-
 		_, ok := coll.datasetSamplers[dataset]
+		coll.mutex.Unlock()
+
 		return ok
 	}, conf.GetTracesConfig().GetTraceTimeout()*2, conf.GetTracesConfig().GetSendTickerValue())
 
@@ -636,9 +643,8 @@ func TestSampleConfigReload(t *testing.T) {
 
 	assert.Eventually(t, func() bool {
 		coll.mutex.Lock()
-		defer coll.mutex.Unlock()
-
 		_, ok := coll.datasetSamplers[dataset]
+		coll.mutex.Unlock()
 		return !ok
 	}, conf.GetTracesConfig().GetTraceTimeout()*2, conf.GetTracesConfig().GetSendTickerValue())
 
@@ -654,9 +660,8 @@ func TestSampleConfigReload(t *testing.T) {
 
 	assert.Eventually(t, func() bool {
 		coll.mutex.Lock()
-		defer coll.mutex.Unlock()
-
 		_, ok := coll.datasetSamplers[dataset]
+		coll.mutex.Unlock()
 		return ok
 	}, conf.GetTracesConfig().GetTraceTimeout()*2, conf.GetTracesConfig().GetSendTickerValue())
 }
@@ -730,9 +735,7 @@ func TestStableMaxAlloc(t *testing.T) {
 	runtime.ReadMemStats(&mem)
 	// Set MaxAlloc, which should cause cache evictions.
 	conf.GetCollectionConfigVal.MaxAlloc = config.MemorySize(mem.Alloc * 99 / 100)
-
 	coll.mutex.Unlock()
-
 	// wait for the cache to take some action
 	var traces []*types.Trace
 	for {
@@ -746,7 +749,7 @@ func TestStableMaxAlloc(t *testing.T) {
 		time.Sleep(conf.GetTracesConfig().GetSendTickerValue())
 	}
 
-	assert.Equal(t, 1000, coll.cache.(*cache.DefaultInMemCache).GetCacheSize(), "cache size shouldn't change")
+	assert.Equal(t, 1000, coll.cache.GetCacheCapacity(), "cache size shouldn't change")
 
 	tracesLeft := len(traces)
 	assert.Less(t, tracesLeft, 480, "should have sent some traces")
@@ -1136,6 +1139,7 @@ func TestLateRootGetsSpanCount(t *testing.T) {
 	trace := coll.getFromCache(traceID)
 	assert.Nil(t, trace, "trace should have been sent although the root span hasn't arrived")
 	assert.Equal(t, 1, len(transmission.Events), "adding a non-root span and waiting should send the span")
+
 	// now we add the root span and verify that both got sent and that the root span had the span count
 	rootSpan := &types.Span{
 		TraceID: traceID,
@@ -1149,12 +1153,14 @@ func TestLateRootGetsSpanCount(t *testing.T) {
 	time.Sleep(conf.GetTracesConfig().GetSendTickerValue() * 2)
 
 	assert.Nil(t, coll.getFromCache(traceID), "after adding a leaf and root span, it should be removed from the cache")
-	transmission.Mux.RLock()
-	assert.Equal(t, 2, len(transmission.Events), "adding a root span should send all spans in the trace")
-	assert.Equal(t, nil, transmission.Events[0].Data["meta.span_count"], "child span metadata should NOT be populated with span count")
-	assert.Equal(t, int64(2), transmission.Events[1].Data["meta.span_count"], "root span metadata should be populated with span count")
-	assert.Equal(t, "deterministic/always - late arriving span", transmission.Events[1].Data["meta.refinery.reason"], "late spans should have meta.refinery.reason set to late.")
-	transmission.Mux.RUnlock()
+	assert.EventuallyWithT(t, func(collect *assert.CollectT) {
+		transmission.Mux.RLock()
+		assert.Equal(collect, 2, len(transmission.Events), "adding a root span should send all spans in the trace")
+		assert.Equal(collect, nil, transmission.Events[0].Data["meta.span_count"], "child span metadata should NOT be populated with span count")
+		assert.Equal(collect, int64(2), transmission.Events[1].Data["meta.span_count"], "root span metadata should be populated with span count")
+		assert.Equal(collect, "deterministic/always - late arriving span", transmission.Events[1].Data["meta.refinery.reason"], "late spans should have meta.refinery.reason set to late.")
+		transmission.Mux.RUnlock()
+	}, 2*conf.GetTracesConfig().GetSendTickerValue(), 1*time.Millisecond)
 }
 
 // TestLateRootNotDecorated tests that spans do not get decorated with 'meta.refinery.reason' meta field
@@ -1567,6 +1573,90 @@ func TestIsRootSpan(t *testing.T) {
 	}
 }
 
+func TestRedistributeTraces(t *testing.T) {
+	conf := &config.MockConfig{
+		GetTracesConfigVal: config.TracesConfig{
+			SendDelay:    config.Duration(1 * time.Millisecond),
+			TraceTimeout: config.Duration(1 * time.Second),
+			SendTicker:   config.Duration(2 * time.Millisecond),
+		},
+		GetSamplerTypeVal:      &config.DeterministicSamplerConfig{SampleRate: 1},
+		ParentIdFieldNames:     []string{"trace.parent_id", "parentId"},
+		GetCollectionConfigVal: config.CollectionConfig{CacheCapacity: 10},
+		SampleCache: config.SampleCacheConfig{
+			KeptSize:          100,
+			DroppedSize:       100,
+			SizeCheckInterval: config.Duration(1 * time.Second),
+		},
+	}
+
+	transmission := &transmit.MockTransmission{}
+	transmission.Start()
+	coll := newTestCollector(conf, transmission)
+	s := &sharder.MockSharder{
+		Self: &sharder.TestShard{Addr: "api1"},
+	}
+
+	coll.Sharder = s
+
+	err := coll.Start()
+	assert.NoError(t, err)
+	defer coll.Stop()
+
+	dataset := "aoeu"
+
+	span := &types.Span{
+		TraceID: "1",
+		Event: types.Event{
+			Dataset: dataset,
+			APIKey:  legacyAPIKey,
+			APIHost: "api1",
+			Data:    make(map[string]interface{}),
+		},
+	}
+
+	coll.AddSpan(span)
+
+	assert.Eventually(t, func() bool {
+		transmission.Mux.Lock()
+		defer transmission.Mux.Unlock()
+
+		return len(transmission.Events) == 1 && transmission.Events[0].APIHost == "api1"
+	}, conf.GetTracesConfig().GetTraceTimeout()*2, conf.GetTracesConfig().GetSendTickerValue())
+	transmission.Flush()
+
+	s.Other = &sharder.TestShard{Addr: "api2"}
+	span = &types.Span{
+		TraceID: "11",
+		Event: types.Event{
+			Dataset: dataset,
+			APIKey:  legacyAPIKey,
+			Data:    make(map[string]interface{}),
+		},
+	}
+	trace := &types.Trace{
+		TraceID: span.TraceID,
+		Dataset: dataset,
+		SendBy:  coll.Clock.Now().Add(5 * time.Second),
+	}
+	trace.AddSpan(span)
+
+	coll.mutex.Lock()
+	coll.cache.Set(trace)
+	coll.mutex.Unlock()
+	coll.Peers.RegisterUpdatedPeersCallback(coll.redistributeTimer.Reset)
+
+	assert.Eventually(t, func() bool {
+		transmission.Mux.Lock()
+		defer transmission.Mux.Unlock()
+		if len(transmission.Events) == 0 {
+			return false
+		}
+
+		return len(transmission.Events) == 1 && transmission.Events[0].APIHost == "api2"
+	}, conf.GetTracesConfig().GetTraceTimeout()*2, conf.GetTracesConfig().GetSendTickerValue())
+}
+
 func TestDrainTracesOnShutdown(t *testing.T) {
 	// set up the trace cache
 	conf := &config.MockConfig{
@@ -1587,6 +1677,11 @@ func TestDrainTracesOnShutdown(t *testing.T) {
 	transmission.Start()
 	coll := newTestCollector(conf, transmission)
 	coll.hostname = "host123"
+	coll.Sharder = &sharder.MockSharder{
+		Self:  &sharder.TestShard{Addr: "api1"},
+		Other: &sharder.TestShard{Addr: "api2"},
+	}
+
 	c := cache.NewInMemCache(3, &metrics.NullMetrics{}, &logger.NullLogger{})
 	coll.cache = c
 	stc, err := newCache()
@@ -1656,11 +1751,10 @@ func TestDrainTracesOnShutdown(t *testing.T) {
 		transmission.Mux.Lock()
 		require.Len(collect, transmission.Events, 1)
 		require.Equal(collect, span2.Dataset, transmission.Events[0].Dataset)
-		require.Equal(collect, "http://self", transmission.Events[0].APIHost)
+		require.Equal(collect, "api2", transmission.Events[0].APIHost)
 		transmission.Mux.Unlock()
 	}, 2*time.Second, 100*time.Millisecond)
 	cancel2()
-
 }
 
 func TestBigTracesGoEarly(t *testing.T) {
diff --git a/route/route_test.go b/route/route_test.go
index e2f87f6..1649df1 100644
--- a/route/route_test.go
+++ b/route/route_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/honeycombio/refinery/internal/peer"
 	"github.com/honeycombio/refinery/logger"
 	"github.com/honeycombio/refinery/metrics"
+	"github.com/honeycombio/refinery/sharder"
 	"github.com/honeycombio/refinery/transmit"
 	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/assert"
@@ -30,7 +31,6 @@ import (
 	trace "go.opentelemetry.io/proto/otlp/trace/v1"
 
 	"github.com/gorilla/mux"
-	"github.com/honeycombio/refinery/sharder"
 	"github.com/klauspost/compress/zstd"
 	"github.com/vmihailenco/msgpack/v5"
 	"google.golang.org/grpc/health/grpc_health_v1"
@@ -321,7 +321,10 @@ func TestDebugTrace(t *testing.T) {
 
 	rr := httptest.NewRecorder()
 	router := &Router{
-		Sharder: &TestSharder{},
+		Sharder: &sharder.MockSharder{
+			Self:  &sharder.TestShard{Addr: "http://localhost:12345"},
+			Other: &sharder.TestShard{Addr: "http://localhost:12345"},
+		},
 	}
 
 	router.debugTrace(rr, req)
@@ -481,7 +484,7 @@ func TestDependencyInjection(t *testing.T) {
 		&inject.Object{Value: http.DefaultTransport, Name: "upstreamTransport"},
 		&inject.Object{Value: &transmit.MockTransmission{}, Name: "upstreamTransmission"},
 		&inject.Object{Value: &transmit.MockTransmission{}, Name: "peerTransmission"},
-		&inject.Object{Value: &TestSharder{}},
+		&inject.Object{Value: &sharder.MockSharder{}},
 		&inject.Object{Value: &collect.InMemCollector{}},
 		&inject.Object{Value: &metrics.NullMetrics{}, Name: "metrics"},
 		&inject.Object{Value: &metrics.NullMetrics{}, Name: "genericMetrics"},
@@ -498,23 +501,6 @@ func TestDependencyInjection(t *testing.T) {
 	}
 }
 
-type TestSharder struct{}
-
-func (s *TestSharder) MyShard() sharder.Shard { return nil }
-
-func (s *TestSharder) WhichShard(string) sharder.Shard {
-	return &TestShard{
-		addr: "http://localhost:12345",
-	}
-}
-
-type TestShard struct {
-	addr string
-}
-
-func (s *TestShard) Equals(other sharder.Shard) bool { return true }
-func (s *TestShard) GetAddress() string              { return s.addr }
-
 func TestEnvironmentCache(t *testing.T) {
 	t.Run("calls getFn on cache miss", func(t *testing.T) {
 		cache := newEnvironmentCache(time.Second, func(key string) (string, error) {