diff --git a/spamoor/clientpool.go b/spamoor/clientpool.go index 09ef9a9..007595e 100644 --- a/spamoor/clientpool.go +++ b/spamoor/clientpool.go @@ -267,7 +267,9 @@ func (pool *ClientPool) watchClientStatus() error { goodClients = append(goodClients, client) } } + pool.selectionMutex.Lock() pool.goodClients = goodClients + pool.selectionMutex.Unlock() pool.logger.Infof("client check completed (%v good clients, %v bad clients)", len(goodClients), len(pool.allClients)-len(goodClients)) return nil @@ -345,19 +347,17 @@ func (pool *ClientPool) GetClient(options ...ClientSelectionOption) *Client { selectedIndex = opts.index % len(clientCandidates) case SelectClientRandom: selectedIndex = rand.Intn(len(clientCandidates)) - case SelectClientRoundRobin: - selectedIndex = pool.rrClientIdx - pool.rrClientIdx++ + default: + // Round-robin (also the fallback for any unrecognized mode). The cursor is + // shared across calls that can each see a differently sized candidate set + // (different group, different exclusions, clients toggling enabled), so it + // has to be clamped to the current candidate list before it is used, not + // just after advancing it. if pool.rrClientIdx >= len(clientCandidates) { pool.rrClientIdx = 0 } - default: - // Fallback to round-robin selectedIndex = pool.rrClientIdx pool.rrClientIdx++ - if pool.rrClientIdx >= len(clientCandidates) { - pool.rrClientIdx = 0 - } } return clientCandidates[selectedIndex] diff --git a/spamoor/clientpool_test.go b/spamoor/clientpool_test.go new file mode 100644 index 0000000..b4fed4e --- /dev/null +++ b/spamoor/clientpool_test.go @@ -0,0 +1,117 @@ +package spamoor + +import ( + "sync" + "testing" + "time" +) + +func newSelectionTestClient(groups ...string) *Client { + return &Client{enabled: true, clientGroups: groups} +} + +// The round-robin cursor is shared across calls that can see differently +// sized candidate sets (different group filters, or clients toggling +// enabled). Selecting from a smaller set than the cursor's current value +// must wrap instead of indexing out of range. +func TestGetClient_RoundRobinWrapsOnSmallerCandidateSet(t *testing.T) { + pool := &ClientPool{ + goodClients: []*Client{ + newSelectionTestClient("default"), + newSelectionTestClient("default"), + newSelectionTestClient("default"), + newSelectionTestClient("default"), + newSelectionTestClient("default"), + newSelectionTestClient("builders"), + }, + } + + // Advance the shared cursor past the size of the "builders" group. + pool.GetClient() + pool.GetClient() + + for i := 0; i < 10; i++ { + c := pool.GetClient(WithClientGroup("builders")) + if c == nil { + t.Fatalf("expected a client, got nil on iteration %d", i) + } + if !c.HasGroup("builders") { + t.Fatalf("expected a builders client, got one in groups %v", c.clientGroups) + } + } +} + +// A client set that shrinks between calls (a client goes unhealthy or +// disabled) must not leave the cursor pointing past the end of the next +// candidate list. +func TestGetClient_RoundRobinWrapsOnShrinkingPool(t *testing.T) { + c1 := newSelectionTestClient("default") + c2 := newSelectionTestClient("default") + c3 := newSelectionTestClient("default") + pool := &ClientPool{goodClients: []*Client{c1, c2, c3}} + + pool.GetClient() + pool.GetClient() + + pool.goodClients = []*Client{c1, c2} + + for i := 0; i < 10; i++ { + if c := pool.GetClient(); c == nil { + t.Fatalf("expected a client, got nil on iteration %d", i) + } + } +} + +// ByIndex selection was already bounds-safe; keep it that way. +func TestGetClient_ByIndexModeWrapsWithModulo(t *testing.T) { + pool := &ClientPool{goodClients: []*Client{ + newSelectionTestClient("builders"), + }} + + c := pool.GetClient(WithClientSelectionMode(SelectClientByIndex, 999), WithClientGroup("builders")) + if c == nil { + t.Fatal("expected a client, got nil") + } +} + +// The health-check loop replaces goodClients from a different goroutine than +// GetClient's readers. The write must be covered by the same lock the reader +// takes, or this fails under -race. +func TestClientPool_GoodClientsUpdateIsRaceFreeWithGetClient(t *testing.T) { + mk := func() *Client { return newSelectionTestClient("default") } + pool := &ClientPool{goodClients: []*Client{mk(), mk(), mk()}} + + stop := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + pool.selectionMutex.Lock() + pool.goodClients = []*Client{mk(), mk()} + pool.selectionMutex.Unlock() + } + }() + + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + pool.GetClient() + } + }() + + time.Sleep(150 * time.Millisecond) + close(stop) + wg.Wait() +}