gitea-dynamic-runner/internal/taskscheduler/poller_test.go

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

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