feat: 为执行后端增加独立容量池

internal/assignmentqueue/jetstream.go

@@ -57,6 +57,11 @@ type Claims interface {
 	WaitClaimed(context.Context, string) error
 }
 
+type Admission interface {
+	Acquire(string) bool
+	Release(string)
+}
+
 // Message is the subset of jetstream.Msg needed by one reconciliation.
 type Message interface {
 	Data() []byte
@@ -70,13 +75,14 @@ type Processor struct {
 	TrustDomain  string
 	Accepter     Accepter
 	Claims       Claims
+	Admission    Admission
 	RetryDelay   time.Duration
 	ClaimTimeout time.Duration
 }
 
 func (p Processor) Process(ctx context.Context, message Message) error {
-	if p.Accepter == nil || p.Claims == nil {
-		return errors.New("assignment accepter and claim registry are required")
+	if p.Accepter == nil || p.Claims == nil || p.Admission == nil {
+		return errors.New("assignment accepter, claim registry, and backend admission pool are required")
 	}
 	assignment, err := taskassignment.Unmarshal(message.Data(), p.TrustDomain)
 	if err != nil {
@@ -85,8 +91,16 @@ func (p Processor) Process(ctx context.Context, message Message) error {
 	if _, err := p.Claims.Offer(assignment); err != nil {
 		return errors.Join(err, message.TermWithReason("conflicting assignment"))
 	}
+	if !p.Admission.Acquire(assignment.ID) {
+		delay := p.RetryDelay
+		if delay <= 0 {
+			delay = 2 * time.Second
+		}
+		return message.NakWithDelay(delay)
+	}
 	accepted, err := p.Accepter.Accept(ctx, assignment)
 	if err != nil {
+		p.Admission.Release(assignment.ID)
 		delay := p.RetryDelay
 		if delay <= 0 {
 			delay = 15 * time.Second

internal/assignmentqueue/jetstream_test.go

@@ -61,6 +61,28 @@ type fakeClaims struct {
 	claimed bool
 }
 
+type fakeAdmission struct {
+	allowed  bool
+	active   map[string]bool
+	released int
+}
+
+func (a *fakeAdmission) Acquire(assignmentID string) bool {
+	if !a.allowed {
+		return false
+	}
+	if a.active == nil {
+		a.active = make(map[string]bool)
+	}
+	a.active[assignmentID] = true
+	return true
+}
+
+func (a *fakeAdmission) Release(assignmentID string) {
+	delete(a.active, assignmentID)
+	a.released++
+}
+
 func (c *fakeClaims) Offer(taskassignment.Assignment) (<-chan struct{}, error) {
 	ready := make(chan struct{})
 	if c.claimed {
@@ -104,7 +126,7 @@ func encodedAssignment(t *testing.T) []byte {
 
 func TestProcessorAcknowledgesPersistedHandoff(t *testing.T) {
 	message := &fakeMessage{data: encodedAssignment(t)}
-	processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{accepted: true}, Claims: &fakeClaims{claimed: true}}
+	processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{accepted: true}, Claims: &fakeClaims{claimed: true}, Admission: &fakeAdmission{allowed: true}}
 	if err := processor.Process(context.Background(), message); err != nil {
 		t.Fatal(err)
 	}
@@ -115,7 +137,7 @@ func TestProcessorAcknowledgesPersistedHandoff(t *testing.T) {
 
 func TestProcessorRetriesUntilBackendHandoffIsDurable(t *testing.T) {
 	message := &fakeMessage{data: encodedAssignment(t)}
-	processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{}, Claims: &fakeClaims{}, RetryDelay: 2 * time.Second}
+	processor := Processor{TrustDomain: "ddupan.top", Accepter: &fakeAccepter{}, Claims: &fakeClaims{}, Admission: &fakeAdmission{allowed: true}, RetryDelay: 2 * time.Second}
 	if err := processor.Process(context.Background(), message); err != nil {
 		t.Fatal(err)
 	}
@@ -126,10 +148,12 @@ func TestProcessorRetriesUntilBackendHandoffIsDurable(t *testing.T) {
 
 func TestProcessorRetriesBackendFailureAndTerminatesPoisonMessage(t *testing.T) {
 	retry := &fakeMessage{data: encodedAssignment(t)}
+	admission := &fakeAdmission{allowed: true}
 	processor := Processor{
 		TrustDomain: "ddupan.top",
 		Accepter:    &fakeAccepter{err: errors.New("backend unavailable")},
 		Claims:      &fakeClaims{},
+		Admission:   admission,
 		RetryDelay:  time.Minute,
 	}
 	if err := processor.Process(context.Background(), retry); err == nil {
@@ -138,6 +162,9 @@ func TestProcessorRetriesBackendFailureAndTerminatesPoisonMessage(t *testing.T)
 	if retry.nacked != time.Minute {
 		t.Fatalf("retry delay = %s", retry.nacked)
 	}
+	if admission.released != 1 {
+		t.Fatalf("released slots = %d", admission.released)
+	}
 
 	poison := &fakeMessage{data: []byte("not-json")}
 	if err := processor.Process(context.Background(), poison); err == nil {
@@ -148,6 +175,24 @@ func TestProcessorRetriesBackendFailureAndTerminatesPoisonMessage(t *testing.T)
 	}
 }
 
+func TestProcessorLeavesAssignmentPendingWhenBackendPoolIsFull(t *testing.T) {
+	message := &fakeMessage{data: encodedAssignment(t)}
+	accepter := &fakeAccepter{accepted: true}
+	processor := Processor{
+		TrustDomain: "ddupan.top",
+		Accepter:    accepter,
+		Claims:      &fakeClaims{},
+		Admission:   &fakeAdmission{},
+		RetryDelay:  3 * time.Second,
+	}
+	if err := processor.Process(context.Background(), message); err != nil {
+		t.Fatal(err)
+	}
+	if message.acked != 0 || message.nacked != 3*time.Second {
+		t.Fatalf("message = %#v", message)
+	}
+}
+
 type fakeConsumerManager struct{ config jetstream.ConsumerConfig }
 
 func (m *fakeConsumerManager) CreateOrUpdateConsumer(_ context.Context, _ string, config jetstream.ConsumerConfig) (jetstream.Consumer, error) {

internal/backendpool/pool.go

@@ -0,0 +1,50 @@
+// Package backendpool manages runtime capacity independently for each executor backend.
+package backendpool
+
+import "sync"
+
+type Pool struct {
+	mu       sync.Mutex
+	capacity int
+	active   map[string]struct{}
+}
+
+func New(capacity int) *Pool {
+	return &Pool{capacity: capacity, active: make(map[string]struct{})}
+}
+
+// Acquire reserves a backend slot without blocking. Redelivery of the same
+// assignment is idempotent and succeeds even while the pool is full.
+func (p *Pool) Acquire(assignmentID string) bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if _, exists := p.active[assignmentID]; exists {
+		return true
+	}
+	if assignmentID == "" || len(p.active) >= p.capacity {
+		return false
+	}
+	p.active[assignmentID] = struct{}{}
+	return true
+}
+
+func (p *Pool) Restore(assignmentID string) {
+	if assignmentID == "" {
+		return
+	}
+	p.mu.Lock()
+	p.active[assignmentID] = struct{}{}
+	p.mu.Unlock()
+}
+
+func (p *Pool) Release(assignmentID string) {
+	p.mu.Lock()
+	delete(p.active, assignmentID)
+	p.mu.Unlock()
+}
+
+func (p *Pool) Active() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return len(p.active)
+}

internal/backendpool/pool_test.go

@@ -0,0 +1,26 @@
+package backendpool
+
+import "testing"
+
+func TestPoolSeparatesRuntimeCapacityFromDeliveries(t *testing.T) {
+	pool := New(2)
+	if !pool.Acquire("one") || !pool.Acquire("two") || pool.Acquire("three") {
+		t.Fatal("capacity was not enforced")
+	}
+	if !pool.Acquire("one") {
+		t.Fatal("redelivery must be idempotent")
+	}
+	pool.Release("one")
+	if !pool.Acquire("three") || pool.Active() != 2 {
+		t.Fatalf("active=%d", pool.Active())
+	}
+}
+
+func TestRestoreMayTemporarilyExceedReducedCapacity(t *testing.T) {
+	pool := New(1)
+	pool.Restore("one")
+	pool.Restore("two")
+	if pool.Active() != 2 || pool.Acquire("three") {
+		t.Fatalf("active=%d", pool.Active())
+	}
+}

internal/taskscheduler/poller.go

@@ -4,8 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync/atomic"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+
 	runnerv1 "gitea.dev/actionslib/runner/v1"
 )
 
@@ -19,10 +22,12 @@ type PollerConfig struct {
 	Labels       []string
 	EmptyBackoff time.Duration
 	ErrorBackoff time.Duration
+	Capacity     int
 }
 
-// Poller is the scheduler component. Once Gitea assigns a task, it never
-// fetches another one until the current assignment is durably dispatched.
+// Poller is the scheduler component. Each fetcher keeps its assigned task
+// until that assignment is durably dispatched; all fetchers share one runner
+// declaration and a monotonic tasks version.
 type Poller struct {
 	Client    PollClient
 	Scheduler *Scheduler
@@ -49,9 +54,21 @@ func (p Poller) Run(ctx context.Context) error {
 	if errorBackoff <= 0 {
 		errorBackoff = 5 * time.Second
 	}
-	var tasksVersion int64
+	capacity := p.Config.Capacity
+	if capacity < 1 {
+		capacity = 1
+	}
+	var tasksVersion atomic.Int64
+	group, groupContext := errgroup.WithContext(ctx)
+	for range capacity {
+		group.Go(func() error { return p.runFetcher(groupContext, &tasksVersion, emptyBackoff, errorBackoff) })
+	}
+	return group.Wait()
+}
+
+func (p Poller) runFetcher(ctx context.Context, tasksVersion *atomic.Int64, emptyBackoff, errorBackoff time.Duration) error {
 	for {
-		response, err := p.Client.FetchTask(ctx, tasksVersion)
+		response, err := p.Client.FetchTask(ctx, tasksVersion.Load())
 		if err != nil {
 			if ctx.Err() != nil {
 				return nil
@@ -69,7 +86,7 @@ func (p Poller) Run(ctx context.Context) error {
 			}
 			continue
 		}
-		tasksVersion = response.GetTasksVersion()
+		storeMaximum(tasksVersion, response.GetTasksVersion())
 		task := response.GetTask()
 		if task == nil {
 			if !wait(ctx, emptyBackoff) {
@@ -93,6 +110,14 @@ func (p Poller) Run(ctx context.Context) error {
 	}
 }
 
+func storeMaximum(value *atomic.Int64, candidate int64) {
+	for current := value.Load(); candidate > current; current = value.Load() {
+		if value.CompareAndSwap(current, candidate) {
+			return
+		}
+	}
+}
+
 func (p Poller) report(err error) {
 	if p.OnError != nil {
 		p.OnError(err)

internal/taskscheduler/poller_test.go

@@ -115,3 +115,52 @@ func TestPollerRetriesAssignedTaskBeforeFetchingAnother(t *testing.T) {
 		t.Fatalf("dispatches=%d fetches-before-dispatch=%d declares=%d", dispatcher.calls, dispatcher.fetchesAtSuccess, client.declared)
 	}
 }
+
+type blockingPollClient struct {
+	mu       sync.Mutex
+	declared int
+	started  chan struct{}
+}
+
+func (c *blockingPollClient) Declare(context.Context, string, []string) error {
+	c.mu.Lock()
+	c.declared++
+	c.mu.Unlock()
+	return nil
+}
+
+func (c *blockingPollClient) FetchTask(ctx context.Context, _ int64) (*runnerv1.FetchTaskResponse, error) {
+	c.started <- struct{}{}
+	<-ctx.Done()
+	return nil, ctx.Err()
+}
+
+func TestPollerStartsConfiguredNumberOfFetchersAfterOneDeclare(t *testing.T) {
+	client := &blockingPollClient{started: make(chan struct{}, 3)}
+	poller := Poller{
+		Client:    client,
+		Scheduler: &Scheduler{TrustDomain: "ddupan.top", Dispatcher: &retryDispatcher{done: make(chan struct{})}},
+		Config: PollerConfig{
+			Version: "dev", Labels: []string{"self-hosted:host", "pod:host", "vm:host"}, Capacity: 3,
+		},
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	finished := make(chan error, 1)
+	go func() { finished <- poller.Run(ctx) }()
+	for range 3 {
+		select {
+		case <-client.started:
+		case <-time.After(time.Second):
+			t.Fatal("configured fetchers did not start")
+		}
+	}
+	cancel()
+	if err := <-finished; err != nil {
+		t.Fatal(err)
+	}
+	client.mu.Lock()
+	defer client.mu.Unlock()
+	if client.declared != 1 {
+		t.Fatalf("declares = %d", client.declared)
+	}
+}