Add worker pool-level gates

Introduces a pool-level admission control gating mechanism.
This allows worker pools to block/park incoming requests in-memory when
capacity is reached, preventing expensive broker nack-and-retry cycles.

Signed-off-by: Jacob Murry <jacobmurry@google.com>

Author: jtechapps

cmd/main.go

@@ -287,6 +287,19 @@ func main() {
 
 	dispatch := policy.MergeRequestChannels(impl.RequestChannels(), poolsMap)
 
+	poolGates := make(map[string]pipeline.Gate)
+	for poolID, pool := range poolsMap {
+		if pool.GateType != "" {
+			gate, err := gateFactory.CreateGate(pool.GateType, pool.GateParams)
+			if err != nil {
+				setupLog.Error(err, "Failed to create pool gate", "poolID", poolID, "gateType", pool.GateType)
+				os.Exit(1)
+			}
+			poolGates[poolID] = gate
+			setupLog.Info("Created pool gate", "poolID", poolID, "gateType", pool.GateType, "gateParams", pool.GateParams)
+		}
+	}
+
 	var wg sync.WaitGroup
 	for poolID, mergedChan := range dispatch.Channels {
 		pool, ok := poolsMap[poolID]
@@ -295,14 +308,15 @@ func main() {
 			os.Exit(1)
 		}
 		workersCount := pool.Workers
+		poolGate := poolGates[poolID]
 
-		setupLog.Info("Spawning workers for pool", "poolID", poolID, "workers", workersCount)
+		setupLog.Info("Spawning workers for pool", "poolID", poolID, "workers", workersCount, "hasGate", poolGate != nil)
 		for w := 1; w <= workersCount; w++ {
 			wg.Add(1)
-			go func(mergedChan chan pipeline.EmbelishedRequestMessage) {
+			go func(mergedChan chan pipeline.EmbelishedRequestMessage, poolGate pipeline.Gate) {
 				defer wg.Done()
-				asyncworker.Worker(signalCtx, drainCtx, impl.Characteristics(), inferenceClient, mergedChan, impl.RetryChannel(), impl.ResultChannel(), requestTimeout, transforms)
-			}(mergedChan)
+				asyncworker.WorkerWithGate(signalCtx, drainCtx, impl.Characteristics(), inferenceClient, mergedChan, impl.RetryChannel(), impl.ResultChannel(), requestTimeout, transforms, poolGate)
+			}(mergedChan, poolGate)
 		}
 	}
 

pipeline/worker_pool.go

@@ -4,13 +4,44 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"strconv"
 )
 
+// StringMap is a map[string]string that tolerates non-string JSON values
+// by converting them to their string representation during unmarshaling.
+type StringMap map[string]string
+
+func (m *StringMap) UnmarshalJSON(data []byte) error {
+	var raw map[string]interface{}
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+	result := make(map[string]string, len(raw))
+	for k, v := range raw {
+		switch val := v.(type) {
+		case string:
+			result[k] = val
+		case float64:
+			result[k] = strconv.FormatFloat(val, 'f', -1, 64)
+		case bool:
+			result[k] = strconv.FormatBool(val)
+		case nil:
+			result[k] = ""
+		default:
+			return fmt.Errorf("gate_params key %q: unsupported value type %T (only strings, numbers, and booleans are allowed)", k, v)
+		}
+	}
+	*m = result
+	return nil
+}
+
 // WorkerPoolConfig defines the configuration for a worker pool,
 // specifying the concurrency limit (number of workers) and its ID.
 type WorkerPoolConfig struct {
-	ID      string `json:"id"`
-	Workers int    `json:"workers"`
+	ID         string            `json:"id"`
+	Workers    int               `json:"workers"`
+	GateType   string            `json:"gate_type,omitempty"`
+	GateParams StringMap         `json:"gate_params,omitempty"`
 }
 
 // LoadWorkerPools loads and validates worker pool configurations from a JSON file.

pkg/async/inference/flowcontrol/gate_factory.go

@@ -346,7 +346,15 @@ func (f *GateFactory) CreateGate(gateType string, params map[string]string) (pip
 		if limit <= 0 {
 			return nil, fmt.Errorf("local-max-concurrency limit must be greater than 0, got %d", limit)
 		}
-		return NewLocalConcurrencyGate(limit), nil
+		gate := NewLocalConcurrencyGate(limit)
+		gatingMode := params["gating_mode"]
+		if gatingMode != "" {
+			if gatingMode != string(GatingModeBlocking) && gatingMode != string(GatingModeClassifying) {
+				return nil, fmt.Errorf("local-max-concurrency gating_mode must be either 'blocking' or 'classifying', got %q", gatingMode)
+			}
+			gate.WithGatingMode(GatingMode(gatingMode))
+		}
+		return gate, nil
 
 	default:
 		// Unknown gate types default to open gate

pkg/async/inference/flowcontrol/local_concurrency_gate.go

@@ -10,21 +10,42 @@ import (
 
 var _ pipeline.Gate = (*LocalConcurrencyGate)(nil)
 
+type GatingMode string
+
+const (
+	GatingModeBlocking    GatingMode = "blocking"
+	GatingModeClassifying GatingMode = "classifying"
+)
+
 // LocalConcurrencyGate limits the number of concurrent in-flight requests
 // processed from a single queue locally.
 type LocalConcurrencyGate struct {
-	mu       sync.Mutex
-	limit    int
-	inFlight int
+	mu         sync.Mutex
+	limit      int
+	inFlight   int
+	gatingMode GatingMode
+	sem        chan struct{}
 }
 
 // NewLocalConcurrencyGate creates a new LocalConcurrencyGate with the specified limit.
 func NewLocalConcurrencyGate(limit int) *LocalConcurrencyGate {
 	return &LocalConcurrencyGate{
-		limit: limit,
+		limit:      limit,
+		gatingMode: GatingModeClassifying,
 	}
 }
 
+// WithGatingMode configures the gating mode (blocking or classifying).
+func (g *LocalConcurrencyGate) WithGatingMode(mode GatingMode) *LocalConcurrencyGate {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	g.gatingMode = mode
+	if mode == GatingModeBlocking && g.limit > 0 {
+		g.sem = make(chan struct{}, g.limit)
+	}
+	return g
+}
+
 // Budget implements pipeline.Gate.
 // Returns the fraction of available capacity in [0.0, 1.0].
 func (g *LocalConcurrencyGate) Budget(ctx context.Context) float64 {
@@ -44,17 +65,43 @@ func (g *LocalConcurrencyGate) Budget(ctx context.Context) float64 {
 // Returns VerdictContinue if request fits in budget, VerdictRefuse with redeliver otherwise.
 func (g *LocalConcurrencyGate) Apply(ctx context.Context, msg *api.InternalRequest) (pipeline.Verdict, error) {
 	g.mu.Lock()
-	defer g.mu.Unlock()
 
 	if g.limit <= 0 {
+		g.mu.Unlock()
 		return pipeline.Refuse(), nil
 	}
 
+	if g.gatingMode == GatingModeBlocking {
+		sem := g.sem
+		g.mu.Unlock()
+
+		select {
+		case sem <- struct{}{}:
+			g.mu.Lock()
+			g.inFlight++
+			g.mu.Unlock()
+
+			msg.AttachRelease(func() {
+				<-sem
+				g.mu.Lock()
+				g.inFlight--
+				g.mu.Unlock()
+			})
+			return pipeline.Continue(), nil
+		case <-ctx.Done():
+			return pipeline.Verdict{}, ctx.Err()
+		}
+	}
+
+	// Classifying/non-blocking mode
 	if g.inFlight >= g.limit {
+		g.mu.Unlock()
 		return pipeline.Refuse(), nil
 	}
 
 	g.inFlight++
+	g.mu.Unlock()
+
 	msg.AttachRelease(func() {
 		g.mu.Lock()
 		g.inFlight--

pkg/async/inference/flowcontrol/local_concurrency_gate_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/llm-d-incubation/llm-d-async/api"
 	"github.com/llm-d-incubation/llm-d-async/pipeline"
@@ -131,3 +132,93 @@ func TestLocalConcurrencyGate_Concurrency(t *testing.T) {
 	// Budget should be fully recovered to 1.0
 	assert.Equal(t, 1.0, gate.Budget(ctx))
 }
+
+func TestLocalConcurrencyGate_BlockingMode(t *testing.T) {
+	gate := NewLocalConcurrencyGate(2).WithGatingMode(GatingModeBlocking)
+	ctx := context.Background()
+
+	// 1. Initial Budget is 1.0
+	assert.Equal(t, 1.0, gate.Budget(ctx))
+
+	// 2. Admit 2 requests
+	r1 := api.NewInternalRequest(api.InternalRouting{}, &api.RequestMessage{})
+	v1, err := gate.Apply(ctx, r1)
+	require.NoError(t, err)
+	assert.Equal(t, pipeline.ActionContinue, v1.Action)
+
+	r2 := api.NewInternalRequest(api.InternalRouting{}, &api.RequestMessage{})
+	v2, err := gate.Apply(ctx, r2)
+	require.NoError(t, err)
+	assert.Equal(t, pipeline.ActionContinue, v2.Action)
+
+	// Budget should now be 0.0
+	assert.Equal(t, 0.0, gate.Budget(ctx))
+
+	// 3. Attempt to admit 3rd request in a separate goroutine (should block)
+	r3 := api.NewInternalRequest(api.InternalRouting{}, &api.RequestMessage{})
+	blockedCh := make(chan struct{})
+	resultCh := make(chan pipeline.Verdict, 1)
+	errCh := make(chan error, 1)
+
+	go func() {
+		close(blockedCh)
+		v, e := gate.Apply(ctx, r3)
+		resultCh <- v
+		errCh <- e
+	}()
+
+	<-blockedCh
+	// Sleep briefly to ensure the goroutine is indeed parked waiting on semaphore
+	time.Sleep(50 * time.Millisecond)
+
+	select {
+	case <-resultCh:
+		t.Fatal("Apply should have blocked, but returned result immediately")
+	default:
+		// Passed: goroutine is blocked
+	}
+
+	// 4. Release request 1
+	r1.Release()
+
+	// 5. Goroutine should unblock and the request should be admitted
+	select {
+	case verdict := <-resultCh:
+		require.NoError(t, <-errCh)
+		assert.Equal(t, pipeline.ActionContinue, verdict.Action)
+	case <-time.After(1 * time.Second):
+		t.Fatal("timed out waiting for 3rd request to unblock")
+	}
+
+	// Budget should be back to 0.0
+	assert.Equal(t, 0.0, gate.Budget(ctx))
+
+	// Clean up
+	r2.Release()
+	r3.Release()
+	assert.Equal(t, 1.0, gate.Budget(ctx))
+}
+
+func TestLocalConcurrencyGate_BlockingModeCancel(t *testing.T) {
+	gate := NewLocalConcurrencyGate(1).WithGatingMode(GatingModeBlocking)
+	ctx := context.Background()
+
+	// Admit 1 request to exhaust capacity
+	r1 := api.NewInternalRequest(api.InternalRouting{}, &api.RequestMessage{})
+	v1, err := gate.Apply(ctx, r1)
+	require.NoError(t, err)
+	assert.Equal(t, pipeline.ActionContinue, v1.Action)
+
+	// Try to admit 2nd request with a cancelled context
+	r2 := api.NewInternalRequest(api.InternalRouting{}, &api.RequestMessage{})
+	cancelCtx, cancel := context.WithCancel(ctx)
+	cancel() // cancel immediately
+
+	_, err = gate.Apply(cancelCtx, r2)
+	assert.ErrorIs(t, err, context.Canceled)
+
+	// Clean up
+	r1.Release()
+	assert.Equal(t, 1.0, gate.Budget(ctx))
+}
+

pkg/asyncworker/worker.go

@@ -30,6 +30,11 @@ const (
 
 func Worker(consumeCtx, requestCtx context.Context, characteristics pipeline.Characteristics, client asyncapi.InferenceClient, requestChannel chan pipeline.EmbelishedRequestMessage,
 	retryChannel chan pipeline.RetryMessage, resultChannel chan asyncapi.ResultMessage, requestTimeout time.Duration, transforms *transform.Chain) {
+	WorkerWithGate(consumeCtx, requestCtx, characteristics, client, requestChannel, retryChannel, resultChannel, requestTimeout, transforms, nil)
+}
+
+func WorkerWithGate(consumeCtx, requestCtx context.Context, characteristics pipeline.Characteristics, client asyncapi.InferenceClient, requestChannel chan pipeline.EmbelishedRequestMessage,
+	retryChannel chan pipeline.RetryMessage, resultChannel chan asyncapi.ResultMessage, requestTimeout time.Duration, transforms *transform.Chain, poolGate pipeline.Gate) {
 
 	logger := log.FromContext(requestCtx)
 	for {
@@ -75,9 +80,6 @@ func Worker(consumeCtx, requestCtx context.Context, characteristics pipeline.Cha
 			}
 
 			processMessage := func() {
-				metrics.IncInflight(queueID, queueName, msg.WorkerPoolID)
-				defer metrics.DecInflight(queueID, queueName, msg.WorkerPoolID)
-
 				if msg.RetryCount == 0 {
 					metrics.RecordAsyncReq(queueID, queueName, msg.WorkerPoolID)
 				}
@@ -87,6 +89,75 @@ func Worker(consumeCtx, requestCtx context.Context, characteristics pipeline.Cha
 					return
 				}
 
+				snapshot := len(msg.Releases())
+				defer msg.RollbackReleases(snapshot)
+
+				if poolGate != nil {
+					reqDeadline := time.Now().Add(requestTimeout)
+					if dline := msg.PublicRequest.ReqDeadline(); dline > 0 {
+						if msgDeadline := time.Unix(dline, 0); msgDeadline.Before(reqDeadline) {
+							reqDeadline = msgDeadline
+						}
+					}
+					gateCtx, cancelGate := context.WithDeadline(requestCtx, reqDeadline)
+					defer cancelGate()
+
+					var verdict pipeline.Verdict
+					var err error
+					for {
+						verdict, err = poolGate.Apply(gateCtx, msg.InternalRequest)
+						if err != nil {
+							if errors.Is(err, context.DeadlineExceeded) || gateCtx.Err() != nil {
+								metrics.RecordExceededDeadlineReq(queueID, queueName, msg.WorkerPoolID)
+								select {
+								case resultChannel <- CreateDeadlineExceededResultMessage(msg.PublicRequest, msg.InternalRouting):
+								case <-requestCtx.Done():
+								}
+								return
+							}
+							select {
+							case resultChannel <- CreateErrorResultMessage(msg.PublicRequest, msg.InternalRouting, fmt.Sprintf("Pool gating error: %s", err.Error())):
+							case <-requestCtx.Done():
+							}
+							return
+						}
+
+						if verdict.Action == pipeline.ActionContinue {
+							break
+						}
+
+						if verdict.Action == pipeline.ActionDrop {
+							var resultMsg asyncapi.ResultMessage
+							if verdict.Result != nil {
+								resultMsg = *verdict.Result
+							} else {
+								resultMsg = CreateErrorResultMessage(msg.PublicRequest, msg.InternalRouting, "Pool gating dropped request")
+							}
+							select {
+							case resultChannel <- resultMsg:
+							case <-requestCtx.Done():
+							}
+							return
+						}
+
+						// ActionRefuse: park/wait and retry
+						select {
+						case <-gateCtx.Done():
+							metrics.RecordExceededDeadlineReq(queueID, queueName, msg.WorkerPoolID)
+							select {
+							case resultChannel <- CreateDeadlineExceededResultMessage(msg.PublicRequest, msg.InternalRouting):
+							case <-requestCtx.Done():
+							}
+							return
+						case <-time.After(50 * time.Millisecond):
+							// poll again
+						}
+					}
+				}
+
+				metrics.IncInflight(queueID, queueName, msg.WorkerPoolID)
+				defer metrics.DecInflight(queueID, queueName, msg.WorkerPoolID)
+
 				sendInferenceRequest := func() {
 					reqCtx := requestCtx
 					if md := msg.PublicRequest.ReqMetadata(); len(md) > 0 {