From d0af2178f1d1a18c2d7e61e6df199c18e4c373ac Mon Sep 17 00:00:00 2001 From: Thomas Lacroix Date: Mon, 19 Oct 2026 13:23:43 -0400 Subject: [PATCH 1/3] fix: processor allocator handler concurrency and reconnect issues Signed-off-by: Thomas Lacroix --- cmd/allocator/main.go | 3 + cmd/extensions/main.go | 3 + .../templates/extensions-deployment.yaml | 9 +++ .../processor/handler.go | 71 +++++++++++++------ .../processor/handler_test.go | 20 +++++- 5 files changed, 84 insertions(+), 22 deletions(-) diff --git a/cmd/allocator/main.go b/cmd/allocator/main.go index 4a420f16ff..f5e3d6f1df 100644 --- a/cmd/allocator/main.go +++ b/cmd/allocator/main.go @@ -166,6 +166,9 @@ func parseEnvFlags() config { runtime.Must(viper.BindEnv(maxListItemsFlag)) runtime.Must(viper.BindEnv(readinessShutdownDuration)) runtime.Must(viper.BindEnv(httpUnallocatedStatusCode)) + runtime.Must(viper.BindEnv(processorGRPCAddress)) + runtime.Must(viper.BindEnv(processorGRPCPort)) + runtime.Must(viper.BindEnv(processorMaxBatchSize)) runtime.Must(viper.BindPFlags(pflag.CommandLine)) runtime.Must(runtime.FeaturesBindEnv()) diff --git a/cmd/extensions/main.go b/cmd/extensions/main.go index 195831d948..fd466126cd 100644 --- a/cmd/extensions/main.go +++ b/cmd/extensions/main.go @@ -333,6 +333,9 @@ func parseEnvFlags() config { runtime.Must(viper.BindEnv(webhookPort)) runtime.Must(viper.BindEnv(allocationBatchWaitTime)) runtime.Must(viper.BindEnv(maxListItemsFlag)) + runtime.Must(viper.BindEnv(processorGRPCAddress)) + runtime.Must(viper.BindEnv(processorGRPCPort)) + runtime.Must(viper.BindEnv(processorMaxBatchSize)) runtime.Must(viper.BindPFlags(pflag.CommandLine)) runtime.Must(viper.BindEnv(readinessShutdownDuration)) runtime.Must(cloudproduct.BindEnv()) diff --git a/install/helm/agones/templates/extensions-deployment.yaml b/install/helm/agones/templates/extensions-deployment.yaml index a7e97723ca..7f281b7b8f 100644 --- a/install/helm/agones/templates/extensions-deployment.yaml +++ b/install/helm/agones/templates/extensions-deployment.yaml @@ -159,6 +159,15 @@ spec: value: {{ .Values.agones.extensions.webhooks.port | quote }} - name: HTTP_PORT value: {{ .Values.agones.extensions.http.port | quote }} +{{- $featureGates := include "agones.featureGates" . | fromYaml }} +{{- if $featureGates.ProcessorAllocator }} + - name: PROCESSOR_MAX_BATCH_SIZE + value: {{ .Values.agones.allocator.processor.maxBatchSize | quote }} + - name: PROCESSOR_GRPC_ADDRESS + value: {{ .Values.agones.allocator.processor.grpc.address | quote }} + - name: PROCESSOR_GRPC_PORT + value: {{ .Values.agones.allocator.processor.grpc.port | quote }} +{{- end }} {{- if .Values.agones.extensions.env }} {{- toYaml .Values.agones.extensions.env | nindent 8 }} {{- end }} diff --git a/pkg/gameserverallocations/processor/handler.go b/pkg/gameserverallocations/processor/handler.go index 97f90f4328..465c957adb 100644 --- a/pkg/gameserverallocations/processor/handler.go +++ b/pkg/gameserverallocations/processor/handler.go @@ -53,6 +53,22 @@ type allocationResult struct { error *rpcstatus.Status } +// lockedStream serialises the sends on a stream. gRPC does not allow concurrent Send calls on +// the same stream, but the pull request ticker and the batch goroutines (which each send their +// response once processed) can send at the same time, so the sends must take turns +type lockedStream struct { + allocationpb.Processor_StreamBatchesServer + mu sync.Mutex +} + +// Send sends the message, one at a time +func (s *lockedStream) Send(msg *allocationpb.ProcessorMessage) error { + s.mu.Lock() + defer s.mu.Unlock() + + return s.Processor_StreamBatchesServer.Send(msg) +} + // Handler is the gRPC server for processing allocation requests. type Handler struct { allocationpb.UnimplementedProcessorServer @@ -93,8 +109,16 @@ func (h *Handler) StreamBatches(stream allocationpb.Processor_StreamBatchesServe return status.Error(codes.InvalidArgument, "clientID is required") } + // batch responses and pull requests are sent from different goroutines, and a gRPC stream + // does not allow concurrent sends + stream = &lockedStream{Processor_StreamBatchesServer: stream} + h.addClient(clientID, stream) - defer h.removeClient(clientID) + defer h.removeClient(clientID, stream) + + // sends must be over before the handler returns + var batches sync.WaitGroup + defer batches.Wait() handlerLogger.WithField(logFieldClientID, clientID).Debug("Client registered") // Main loop: handle incoming messages @@ -131,24 +155,26 @@ func (h *Handler) StreamBatches(stream allocationpb.Processor_StreamBatchesServe logFieldRequestCount: len(requestWrappers), }).Debug("Received batch request") - // Submit batch for processing - response := h.submitBatch(stream.Context(), batchID, requestWrappers) + // Process the batch on its own, so that a slow request only delays the response of its own batch, + // and not the receipt and processing of the next ones. + batches.Go(func() { + response := h.submitBatch(stream.Context(), batchID, requestWrappers) - respMsg := &allocationpb.ProcessorMessage{ - ClientId: clientID, - Payload: &allocationpb.ProcessorMessage_BatchResponse{ - BatchResponse: response, - }, - } + respMsg := &allocationpb.ProcessorMessage{ + ClientId: clientID, + Payload: &allocationpb.ProcessorMessage_BatchResponse{ + BatchResponse: response, + }, + } - if err := stream.Send(respMsg); err != nil { - handlerLogger.WithFields(logrus.Fields{ - logFieldClientID: clientID, - logFieldBatchID: batchID, - logFieldRequestCount: len(requestWrappers), - }).WithError(err).Error("Failed to send response") - continue - } + if err := stream.Send(respMsg); err != nil { + handlerLogger.WithFields(logrus.Fields{ + logFieldClientID: clientID, + logFieldBatchID: batchID, + logFieldRequestCount: len(requestWrappers), + }).WithError(err).Error("Failed to send response") + } + }) } } @@ -189,7 +215,7 @@ func (h *Handler) sendPullRequestsToClients() { logFieldClientID: clientID, "error": err, }).Warn("Failed to send pull request, removing client") - h.removeClient(clientID) + h.removeClient(clientID, stream) } } } @@ -314,10 +340,13 @@ func (h *Handler) addClient(clientID string, stream allocationpb.Processor_Strea h.clients[clientID] = stream } -// removeClient unregisters a client from streaming allocation responses. -func (h *Handler) removeClient(clientID string) { +// removeClient unregisters the stream of a client from streaming allocation responses. It is a no-op if the +// client has registered another stream since, as a client reconnects with the same ID. +func (h *Handler) removeClient(clientID string, stream allocationpb.Processor_StreamBatchesServer) { h.mu.Lock() defer h.mu.Unlock() - delete(h.clients, clientID) + if h.clients[clientID] == stream { + delete(h.clients, clientID) + } } diff --git a/pkg/gameserverallocations/processor/handler_test.go b/pkg/gameserverallocations/processor/handler_test.go index 145fd1876f..52ee898303 100644 --- a/pkg/gameserverallocations/processor/handler_test.go +++ b/pkg/gameserverallocations/processor/handler_test.go @@ -134,13 +134,31 @@ func TestRemoveClient(t *testing.T) { h.addClient(id, stream) } - h.removeClient(tc.removeID) + h.removeClient(tc.removeID, stream) h.mu.RLock() defer h.mu.RUnlock() assert.Len(t, h.clients, tc.wantLen) }) } + + t.Run("removing a stale stream keeps the stream the client reconnected with", func(t *testing.T) { + t.Parallel() + + h := &Handler{ + clients: make(map[string]allocationpb.Processor_StreamBatchesServer), + } + stale := newMockServerStream(context.Background()) + current := newMockServerStream(context.Background()) + h.addClient("client-1", stale) + h.addClient("client-1", current) + + h.removeClient("client-1", stale) + + h.mu.RLock() + defer h.mu.RUnlock() + assert.Same(t, current, h.clients["client-1"]) + }) } func TestProcessAllocation(t *testing.T) { From f5c93fa9820f7758a88b1a9794f8738bc2e2d043 Mon Sep 17 00:00:00 2001 From: Thomas Lacroix Date: Mon, 19 Oct 2026 14:31:58 -0400 Subject: [PATCH 2/3] fix: improve context / errors on streams Signed-off-by: Thomas Lacroix --- pkg/gameserverallocations/allocator.go | 9 +- pkg/gameserverallocations/allocator_test.go | 14 + .../processor/handler.go | 39 ++- .../processor/handler_test.go | 251 ++++++++++++++++++ 4 files changed, 309 insertions(+), 4 deletions(-) diff --git a/pkg/gameserverallocations/allocator.go b/pkg/gameserverallocations/allocator.go index 79175c078c..55bebbe00d 100644 --- a/pkg/gameserverallocations/allocator.go +++ b/pkg/gameserverallocations/allocator.go @@ -475,8 +475,13 @@ func (c *Allocator) allocate(ctx context.Context, gsa *allocationv1.GameServerAl return nil, ErrTotalTimeoutExceeded } - // this pushes the request into the batching process - c.pendingRequests <- req + // this pushes the request into the batching process, giving up if the context is done while + // the queue is full + select { + case c.pendingRequests <- req: + case <-ctx.Done(): + return nil, ErrTotalTimeoutExceeded + } select { case res := <-req.response: // wait for the batch to be completed diff --git a/pkg/gameserverallocations/allocator_test.go b/pkg/gameserverallocations/allocator_test.go index f6a059fd89..d4bc83838f 100644 --- a/pkg/gameserverallocations/allocator_test.go +++ b/pkg/gameserverallocations/allocator_test.go @@ -1115,3 +1115,17 @@ func newFakeAllocator() (*Allocator, agtesting.Mocks) { return a, m } + +func TestAllocatorAllocateStopsWaitingOnFullQueue(t *testing.T) { + t.Parallel() + + // a full queue that nothing reads from, as when allocation is stalled + allocator := &Allocator{pendingRequests: make(chan request, 1)} + allocator.pendingRequests <- request{} + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + + _, err := allocator.allocate(ctx, &allocationv1.GameServerAllocation{}) + assert.ErrorIs(t, err, ErrTotalTimeoutExceeded) +} diff --git a/pkg/gameserverallocations/processor/handler.go b/pkg/gameserverallocations/processor/handler.go index 465c957adb..8071d8e24a 100644 --- a/pkg/gameserverallocations/processor/handler.go +++ b/pkg/gameserverallocations/processor/handler.go @@ -42,6 +42,10 @@ import ( var handlerLogger = runtime.NewLoggerWithSource("processorHandler") +// maxInFlightBatches is the maximum number of batches being processed at once for a single client stream +// Once reached, the handler stops receiving batches until one completes +const maxInFlightBatches = 10 + // GameServerAllocator represents the interface to allocate game servers. type GameServerAllocator interface { Allocate(ctx context.Context, gsa *allocationv1.GameServerAllocation) (k8sruntime.Object, error) @@ -59,6 +63,8 @@ type allocationResult struct { type lockedStream struct { allocationpb.Processor_StreamBatchesServer mu sync.Mutex + // inFlight holds one token per batch being processed, which bounds the work queued for the stream + inFlight chan struct{} } // Send sends the message, one at a time @@ -69,6 +75,17 @@ func (s *lockedStream) Send(msg *allocationpb.ProcessorMessage) error { return s.Processor_StreamBatchesServer.Send(msg) } +// sendPull sends a pull request, unless the stream is at capacity or busy sending something else, as a +// pull is not needed then and must not hold up the pulls to the other clients. +func (s *lockedStream) sendPull(msg *allocationpb.ProcessorMessage) error { + if len(s.inFlight) == cap(s.inFlight) || !s.mu.TryLock() { + return nil + } + defer s.mu.Unlock() + + return s.Processor_StreamBatchesServer.Send(msg) +} + // Handler is the gRPC server for processing allocation requests. type Handler struct { allocationpb.UnimplementedProcessorServer @@ -111,7 +128,11 @@ func (h *Handler) StreamBatches(stream allocationpb.Processor_StreamBatchesServe // batch responses and pull requests are sent from different goroutines, and a gRPC stream // does not allow concurrent sends - stream = &lockedStream{Processor_StreamBatchesServer: stream} + locked := &lockedStream{ + Processor_StreamBatchesServer: stream, + inFlight: make(chan struct{}, maxInFlightBatches), + } + stream = locked h.addClient(clientID, stream) defer h.removeClient(clientID, stream) @@ -155,9 +176,19 @@ func (h *Handler) StreamBatches(stream allocationpb.Processor_StreamBatchesServe logFieldRequestCount: len(requestWrappers), }).Debug("Received batch request") + // Wait for capacity before starting more work, so that a stalled allocator applies backpressure + // to the client instead of piling up goroutines and memory + select { + case locked.inFlight <- struct{}{}: + case <-stream.Context().Done(): + return stream.Context().Err() + } + // Process the batch on its own, so that a slow request only delays the response of its own batch, // and not the receipt and processing of the next ones. batches.Go(func() { + defer func() { <-locked.inFlight }() + response := h.submitBatch(stream.Context(), batchID, requestWrappers) respMsg := &allocationpb.ProcessorMessage{ @@ -210,7 +241,11 @@ func (h *Handler) sendPullRequestsToClients() { Pull: &allocationpb.PullRequest{Message: "pull"}, }, } - if err := stream.Send(pullMsg); err != nil { + send := stream.Send + if ls, ok := stream.(*lockedStream); ok { + send = ls.sendPull + } + if err := send(pullMsg); err != nil { handlerLogger.WithFields(logrus.Fields{ logFieldClientID: clientID, "error": err, diff --git a/pkg/gameserverallocations/processor/handler_test.go b/pkg/gameserverallocations/processor/handler_test.go index 52ee898303..80592a1398 100644 --- a/pkg/gameserverallocations/processor/handler_test.go +++ b/pkg/gameserverallocations/processor/handler_test.go @@ -20,6 +20,7 @@ import ( "fmt" "io" "net/http" + "sync/atomic" "testing" "time" @@ -784,3 +785,253 @@ type failingSendStream struct { func (f *failingSendStream) Send(_ *allocationpb.ProcessorMessage) error { return errors.New("send failed") } + +func allocatedGSA() k8sruntime.Object { + return &allocationv1.GameServerAllocation{ + Status: allocationv1.GameServerAllocationStatus{ + State: allocationv1.GameServerAllocationAllocated, + GameServerName: "gs-1", + }, + } +} + +// batchMessage builds a single request batch, whose namespace is the batch ID. +func batchMessage(batchID string) *allocationpb.ProcessorMessage { + return &allocationpb.ProcessorMessage{ + ClientId: "client-1", + Payload: &allocationpb.ProcessorMessage_BatchRequest{ + BatchRequest: &allocationpb.BatchRequest{ + BatchId: batchID, + Requests: []*allocationpb.RequestWrapper{ + {RequestId: "r-" + batchID, Request: &allocationpb.AllocationRequest{Namespace: batchID}}, + }, + }, + }, + } +} + +// runStreamBatches runs StreamBatches in the background and returns a channel with its result +func runStreamBatches(h *Handler, stream allocationpb.Processor_StreamBatchesServer) <-chan error { + done := make(chan error, 1) + go func() { done <- h.StreamBatches(stream) }() + return done +} + +// overlapDetectingStream records whether two sends were ever in progress at the same time +type overlapDetectingStream struct { + *mockServerStream + sending atomic.Int32 + overlap atomic.Bool +} + +func (o *overlapDetectingStream) Send(msg *allocationpb.ProcessorMessage) error { + if o.sending.Add(1) > 1 { + o.overlap.Store(true) + } + defer o.sending.Add(-1) + return o.mockServerStream.Send(msg) +} + +// blockingSendStream blocks every send until release is closed +type blockingSendStream struct { + *mockServerStream + release chan struct{} +} + +func (b *blockingSendStream) Send(msg *allocationpb.ProcessorMessage) error { + <-b.release + return b.mockServerStream.Send(msg) +} + +func TestStreamBatchesSerialisesSends(t *testing.T) { + t.Parallel() + ctx := t.Context() + + h := newTestHandler(ctx, func(context.Context, *allocationv1.GameServerAllocation) (k8sruntime.Object, error) { + return allocatedGSA(), nil + }) + stream := &overlapDetectingStream{mockServerStream: newMockServerStream(ctx)} + stream.sendChan = make(chan *allocationpb.ProcessorMessage, 100) + + const batches = 8 + stream.recvChan <- &allocationpb.ProcessorMessage{ClientId: "client-1"} + for i := range batches { + stream.recvChan <- batchMessage(fmt.Sprintf("b-%d", i)) + } + done := runStreamBatches(h, stream) + + // pull requests race with the batch responses + for range 50 { + h.sendPullRequestsToClients() + } + + responses := 0 + for responses < batches { + select { + case msg := <-stream.sendChan: + if msg.GetBatchResponse() != nil { + responses++ + } + case <-time.After(5 * time.Second): + t.Fatalf("timed out after %d responses", responses) + } + } + assert.False(t, stream.overlap.Load(), "sends on the stream must not overlap") + + close(stream.recvChan) + require.ErrorIs(t, <-done, io.EOF) +} + +func TestStreamBatchesBlockedBatchDoesNotDelayOthers(t *testing.T) { + t.Parallel() + + ctx := t.Context() + + release := make(chan struct{}) + h := newTestHandler(ctx, func(_ context.Context, gsa *allocationv1.GameServerAllocation) (k8sruntime.Object, error) { + if gsa.Namespace == "slow" { + <-release + } + return allocatedGSA(), nil + }) + stream := newMockServerStream(ctx) + stream.recvChan <- &allocationpb.ProcessorMessage{ClientId: "client-1"} + stream.recvChan <- batchMessage("slow") + stream.recvChan <- batchMessage("fast") + done := runStreamBatches(h, stream) + + select { + case msg := <-stream.sendChan: + assert.Equal(t, "fast", msg.GetBatchResponse().GetBatchId()) + case <-time.After(5 * time.Second): + t.Fatal("fast batch was held back by the slow one") + } + + close(release) + close(stream.recvChan) + require.ErrorIs(t, <-done, io.EOF) +} + +func TestStreamBatchesWaitsForResponsesBeforeReturning(t *testing.T) { + t.Parallel() + + ctx := t.Context() + + h := newTestHandler(ctx, func(context.Context, *allocationv1.GameServerAllocation) (k8sruntime.Object, error) { + return allocatedGSA(), nil + }) + stream := &blockingSendStream{ + mockServerStream: newMockServerStream(ctx), + release: make(chan struct{}), + } + stream.recvChan <- &allocationpb.ProcessorMessage{ClientId: "client-1"} + stream.recvChan <- batchMessage("b-1") + close(stream.recvChan) + done := runStreamBatches(h, stream) + + // the receive side is over, but the response is still stuck in Send + select { + case <-done: + t.Fatal("handler returned while a response was still being sent") + case <-time.After(50 * time.Millisecond): + } + + close(stream.release) + require.ErrorIs(t, <-done, io.EOF) + assert.Len(t, stream.sendChan, 1) +} + +func TestStreamBatchesInFlightLimit(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + batches int + wantStarted int // batches started while the allocator is blocked + wantUnread int // batches the handler has not read from the stream yet + }{ + { + name: "batches under the limit all start", + batches: maxInFlightBatches, + wantStarted: maxInFlightBatches, + }, + { + name: "batches over the limit wait for capacity", + batches: maxInFlightBatches + 2, + wantStarted: maxInFlightBatches, + wantUnread: 1, // the handler holds one received batch while waiting for capacity + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := t.Context() + release := make(chan struct{}) + var started atomic.Int32 + h := newTestHandler(ctx, func(context.Context, *allocationv1.GameServerAllocation) (k8sruntime.Object, error) { + started.Add(1) + <-release + return allocatedGSA(), nil + }) + stream := newMockServerStream(ctx) + stream.recvChan = make(chan *allocationpb.ProcessorMessage, tc.batches+1) + stream.sendChan = make(chan *allocationpb.ProcessorMessage, tc.batches) + + stream.recvChan <- &allocationpb.ProcessorMessage{ClientId: "client-1"} + for i := range tc.batches { + stream.recvChan <- batchMessage(fmt.Sprintf("b-%d", i)) + } + done := runStreamBatches(h, stream) + + require.Eventually(t, func() bool { return int(started.Load()) == tc.wantStarted }, 5*time.Second, time.Millisecond) + assert.Never(t, func() bool { return int(started.Load()) > tc.wantStarted }, 50*time.Millisecond, 5*time.Millisecond) + assert.Len(t, stream.recvChan, tc.wantUnread) + + // once the allocator is unblocked, the waiting batches get the capacity freed by finished ones + close(release) + close(stream.recvChan) + require.ErrorIs(t, <-done, io.EOF) + assert.EqualValues(t, tc.batches, started.Load()) + }) + } +} + +func TestLockedStreamSendPull(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + inFlight int // batches being processed + busy bool // another send is in progress + wantSent bool + }{ + {name: "idle stream gets the pull", wantSent: true}, + {name: "stream below capacity gets the pull", inFlight: maxInFlightBatches - 1, wantSent: true}, + {name: "stream at capacity is skipped", inFlight: maxInFlightBatches}, + {name: "stream busy sending is skipped", busy: true}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + mock := newMockServerStream(t.Context()) + stream := &lockedStream{ + Processor_StreamBatchesServer: mock, + inFlight: make(chan struct{}, maxInFlightBatches), + } + for range tc.inFlight { + stream.inFlight <- struct{}{} + } + if tc.busy { + stream.mu.Lock() + defer stream.mu.Unlock() + } + + require.NoError(t, stream.sendPull(&allocationpb.ProcessorMessage{})) + assert.Equal(t, tc.wantSent, len(mock.sendChan) == 1) + }) + } +} From 54132a3987e637fc0b76ec949a241265341343f3 Mon Sep 17 00:00:00 2001 From: Thomas Lacroix Date: Mon, 19 Oct 2026 14:44:57 -0400 Subject: [PATCH 3/3] fix: fix linter issues Signed-off-by: Thomas Lacroix --- pkg/gameserverallocations/processor/handler.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/gameserverallocations/processor/handler.go b/pkg/gameserverallocations/processor/handler.go index 8071d8e24a..9bf5a2662d 100644 --- a/pkg/gameserverallocations/processor/handler.go +++ b/pkg/gameserverallocations/processor/handler.go @@ -62,9 +62,9 @@ type allocationResult struct { // response once processed) can send at the same time, so the sends must take turns type lockedStream struct { allocationpb.Processor_StreamBatchesServer - mu sync.Mutex // inFlight holds one token per batch being processed, which bounds the work queued for the stream inFlight chan struct{} + mu sync.Mutex } // Send sends the message, one at a time