Coordinate retry attempts with backend permits

internal/capacity/client.go

@@ -28,13 +28,13 @@ func (c *client) Generate(
 	ctx context.Context,
 	req domain.GenerateRequest,
 ) (*domain.GenerateResponse, error) {
-	pool := c.manager.getPool(req.Target.BackendID)
-	if pool == nil {
+	if c.manager.getPool(req.Target.BackendID) == nil {
 		return c.next.Generate(ctx, req)
 	}
-	if err := pool.acquire(ctx); err != nil {
+	release, err := c.manager.Acquire(ctx, req.Target.BackendID)
+	if err != nil {
 		return nil, err
 	}
-	defer pool.releaseActive()
+	defer release()
 	return c.next.Generate(ctx, req)
 }

internal/capacity/manager.go

@@ -107,6 +107,23 @@ func (m *Manager) getPool(backendID string) *pool {
 	return m.pools[backendID]
 }
 
+// Acquire obtains one active-generation permit for a backend. Its release
+// function is idempotent. Backends without a configured pool are unlimited.
+func (m *Manager) Acquire(ctx context.Context, backendID string) (func(), error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	pool := m.getPool(backendID)
+	if pool == nil {
+		return releaseNothing, nil
+	}
+	if err := pool.acquire(ctx); err != nil {
+		return nil, err
+	}
+	var once sync.Once
+	return func() { once.Do(pool.releaseActive) }, nil
+}
+
 func (p *pool) acquire(ctx context.Context) error {
 	p.mu.Lock()
 	if err := ctx.Err(); err != nil {

internal/capacity/manager_test.go

@@ -60,6 +60,26 @@ func TestNewManagerRejectsInvalidPolicies(t *testing.T) {
 	}
 }
 
+func TestManagerAcquireReleaseIsIdempotent(t *testing.T) {
+	manager, err := NewManager(map[string]domain.BackendCapacityPolicy{
+		"limited": {ConcurrencyLimit: 1},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	release, err := manager.Acquire(context.Background(), "limited")
+	if err != nil {
+		t.Fatal(err)
+	}
+	release()
+	release()
+	other, err := manager.Acquire(context.Background(), "limited")
+	if err != nil {
+		t.Fatal(err)
+	}
+	other()
+}
+
 func TestManagerAdmissionIsBoundedAndReleaseIsIdempotent(t *testing.T) {
 	policies := map[string]domain.BackendCapacityPolicy{
 		"limited": {