实现 Gitea assignment 调度循环

2026-09-20 19:19:11 +00:00
parent a856176c56
commit b77245e4b9
3 changed files with 232 additions and 1 deletions
@@ -0,0 +1,111 @@
+package taskscheduler
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	runnerv1 "gitea.dev/actionslib/runner/v1"
+)
+
+type PollClient interface {
+	Declare(context.Context, string, []string) error
+	FetchTask(context.Context, int64) (*runnerv1.FetchTaskResponse, error)
+}
+
+type PollerConfig struct {
+	Version      string
+	Labels       []string
+	EmptyBackoff time.Duration
+	ErrorBackoff time.Duration
+}
+
+// Poller is the scheduler component. Once Gitea assigns a task, it never
+// fetches another one until the current assignment is durably dispatched.
+type Poller struct {
+	Client    PollClient
+	Scheduler *Scheduler
+	Config    PollerConfig
+	OnError   func(error)
+}
+
+func (p Poller) Run(ctx context.Context) error {
+	if p.Client == nil || p.Scheduler == nil {
+		return errors.New("Gitea poll client and task scheduler are required")
+	}
+	if p.Config.Version == "" || len(p.Config.Labels) == 0 {
+		return errors.New("runner version and labels are required")
+	}
+	if err := p.Client.Declare(ctx, p.Config.Version, p.Config.Labels); err != nil {
+		return fmt.Errorf("declare scheduler labels: %w", err)
+	}
+
+	emptyBackoff := p.Config.EmptyBackoff
+	if emptyBackoff <= 0 {
+		emptyBackoff = time.Second
+	}
+	errorBackoff := p.Config.ErrorBackoff
+	if errorBackoff <= 0 {
+		errorBackoff = 5 * time.Second
+	}
+	var tasksVersion int64
+	for {
+		response, err := p.Client.FetchTask(ctx, tasksVersion)
+		if err != nil {
+			if ctx.Err() != nil {
+				return nil
+			}
+			p.report(fmt.Errorf("fetch Gitea task: %w", err))
+			if !wait(ctx, errorBackoff) {
+				return nil
+			}
+			continue
+		}
+		if response == nil {
+			p.report(errors.New("fetch Gitea task returned an empty response"))
+			if !wait(ctx, errorBackoff) {
+				return nil
+			}
+			continue
+		}
+		tasksVersion = response.GetTasksVersion()
+		task := response.GetTask()
+		if task == nil {
+			if !wait(ctx, emptyBackoff) {
+				return nil
+			}
+			continue
+		}
+		for {
+			if err := p.Scheduler.Run(ctx, task); err == nil {
+				break
+			} else {
+				if ctx.Err() != nil {
+					return nil
+				}
+				p.report(fmt.Errorf("dispatch Gitea task %d: %w", task.GetId(), err))
+			}
+			if !wait(ctx, errorBackoff) {
+				return nil
+			}
+		}
+	}
+}
+
+func (p Poller) report(err error) {
+	if p.OnError != nil {
+		p.OnError(err)
+	}
+}
+
+func wait(ctx context.Context, duration time.Duration) bool {
+	timer := time.NewTimer(duration)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return false
+	case <-timer.C:
+		return true
+	}
+}
@@ -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)
+	}
+}