实现 Gitea assignment 调度循环

internal/taskscheduler/poller_test.go

@@ -0,0 +1,117 @@
+package taskscheduler
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	runnerv1 "gitea.dev/actionslib/runner/v1"
+	"google.golang.org/protobuf/types/known/structpb"
+)
+
+type fakePollClient struct {
+	mu        sync.Mutex
+	declared  int
+	fetches   int
+	responses []*runnerv1.FetchTaskResponse
+}
+
+func (c *fakePollClient) Declare(context.Context, string, []string) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.declared++
+	return nil
+}
+
+func (c *fakePollClient) FetchTask(ctx context.Context, _ int64) (*runnerv1.FetchTaskResponse, error) {
+	c.mu.Lock()
+	c.fetches++
+	if len(c.responses) > 0 {
+		response := c.responses[0]
+		c.responses = c.responses[1:]
+		c.mu.Unlock()
+		return response, nil
+	}
+	c.mu.Unlock()
+	<-ctx.Done()
+	return nil, ctx.Err()
+}
+
+type retryDispatcher struct {
+	mu               sync.Mutex
+	calls            int
+	failures         int
+	done             chan struct{}
+	fetches          func() int
+	fetchesAtSuccess int
+}
+
+func (d *retryDispatcher) Dispatch(context.Context, Assignment) error {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.calls++
+	if d.failures > 0 {
+		d.failures--
+		return errors.New("JetStream unavailable")
+	}
+	if d.fetches != nil {
+		d.fetchesAtSuccess = d.fetches()
+	}
+	select {
+	case <-d.done:
+	default:
+		close(d.done)
+	}
+	return nil
+}
+
+func pollTask(t *testing.T) *runnerv1.Task {
+	t.Helper()
+	fields, err := structpb.NewStruct(map[string]any{"repository": "owner/repo"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	return &runnerv1.Task{
+		Id: 42, Context: fields,
+		WorkflowPayload: []byte("jobs:\n  test:\n    runs-on: [self-hosted, pod]\n    steps: []\n"),
+	}
+}
+
+func TestPollerRetriesAssignedTaskBeforeFetchingAnother(t *testing.T) {
+	client := &fakePollClient{responses: []*runnerv1.FetchTaskResponse{{Task: pollTask(t), TasksVersion: 7}}}
+	dispatcher := &retryDispatcher{failures: 2, done: make(chan struct{})}
+	dispatcher.fetches = func() int {
+		client.mu.Lock()
+		defer client.mu.Unlock()
+		return client.fetches
+	}
+	poller := Poller{
+		Client:    client,
+		Scheduler: &Scheduler{TrustDomain: "ddupan.top", Dispatcher: dispatcher},
+		Config: PollerConfig{
+			Version: "dev", Labels: []string{"self-hosted:host", "pod:host", "vm:host"},
+			EmptyBackoff: time.Millisecond, ErrorBackoff: time.Millisecond,
+		},
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	finished := make(chan error, 1)
+	go func() { finished <- poller.Run(ctx) }()
+	select {
+	case <-dispatcher.done:
+	case <-time.After(time.Second):
+		t.Fatal("assignment was not dispatched")
+	}
+	cancel()
+	if err := <-finished; err != nil {
+		t.Fatal(err)
+	}
+	dispatcher.mu.Lock()
+	defer dispatcher.mu.Unlock()
+	client.mu.Lock()
+	defer client.mu.Unlock()
+	if dispatcher.calls != 3 || dispatcher.fetchesAtSuccess != 1 || client.declared != 1 {
+		t.Fatalf("dispatches=%d fetches-before-dispatch=%d declares=%d", dispatcher.calls, dispatcher.fetchesAtSuccess, client.declared)
+	}
+}