runner/poller/poller.go

package poller

import (
	"context"
	"errors"
	"sync"
	"time"

	runnerv1 "code.gitea.io/actions-proto-go/runner/v1"
	"codeberg.org/forgejo/runner/client"

	"github.com/bufbuild/connect-go"
	log "github.com/sirupsen/logrus"
)

var ErrDataLock = errors.New("Data Lock Error")

func New(cli client.Client, dispatch func(context.Context, *runnerv1.Task) error, workerNum int) *Poller {
	return &Poller{
		Client:       cli,
		Dispatch:     dispatch,
		routineGroup: newRoutineGroup(),
		metric:       &metric{},
		workerNum:    workerNum,
		ready:        make(chan struct{}, 1),
	}
}

type Poller struct {
	Client   client.Client
	Dispatch func(context.Context, *runnerv1.Task) error

	sync.Mutex
	routineGroup *routineGroup
	metric       *metric
	ready        chan struct{}
	workerNum    int
}

func (p *Poller) schedule() {
	p.Lock()
	defer p.Unlock()
	if int(p.metric.BusyWorkers()) >= p.workerNum {
		return
	}

	select {
	case p.ready <- struct{}{}:
	default:
	}
}

func (p *Poller) Wait() {
	p.routineGroup.Wait()
}

func (p *Poller) Poll(ctx context.Context) error {
	l := log.WithField("func", "Poll")

	for {
		// check worker number
		p.schedule()

		select {
		// wait worker ready
		case <-p.ready:
		case <-ctx.Done():
			return nil
		}
	LOOP:
		for {
			select {
			case <-ctx.Done():
				break LOOP
			default:
				task, err := p.pollTask(ctx)
				if task == nil || err != nil {
					if err != nil {
						l.Errorf("can't find the task: %v", err.Error())
					}
					time.Sleep(5 * time.Second)
					break
				}

				p.metric.IncBusyWorker()
				p.routineGroup.Run(func() {
					defer p.schedule()
					defer p.metric.DecBusyWorker()
					if err := p.dispatchTask(ctx, task); err != nil {
						l.Errorf("execute task: %v", err.Error())
					}
				})
				break LOOP
			}
		}
	}
}

func (p *Poller) pollTask(ctx context.Context) (*runnerv1.Task, error) {
	l := log.WithField("func", "pollTask")
	l.Info("poller: request stage from remote server")

	reqCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	// request a new build stage for execution from the central
	// build server.
	resp, err := p.Client.FetchTask(reqCtx, connect.NewRequest(&runnerv1.FetchTaskRequest{}))
	if err == context.Canceled || err == context.DeadlineExceeded {
		l.WithError(err).Trace("poller: no stage returned")
		return nil, nil
	}

	if err != nil && err == ErrDataLock {
		l.WithError(err).Info("task accepted by another runner")
		return nil, nil
	}

	if err != nil {
		l.WithError(err).Error("cannot accept task")
		return nil, err
	}

	// exit if a nil or empty stage is returned from the system
	// and allow the runner to retry.
	if resp.Msg.Task == nil || resp.Msg.Task.Id == 0 {
		return nil, nil
	}

	return resp.Msg.Task, nil
}

func (p *Poller) dispatchTask(ctx context.Context, task *runnerv1.Task) error {
	l := log.WithField("func", "dispatchTask")
	defer func() {
		e := recover()
		if e != nil {
			l.Errorf("panic error: %v", e)
		}
	}()

	runCtx, cancel := context.WithTimeout(ctx, time.Hour)
	defer cancel()

	return p.Dispatch(runCtx, task)
}