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

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) {