runner-go/poller/poller.go

// Copyright 2019 Drone.IO Inc. All rights reserved.
// Use of this source code is governed by the Polyform License
// that can be found in the LICENSE file.

package poller

import (
	"context"
	"sync"

	"git.awesome-for.me/liuzhiguo/drone-go/drone"
	"git.awesome-for.me/liuzhiguo/runner-go/client"
	"git.awesome-for.me/liuzhiguo/runner-go/logger"
)

var noContext = context.Background()

// Poller polls the server for pending stages and dispatches
// for execution by the Runner.
type Poller struct {
	Client client.Client
	Filter *client.Filter

	// Dispatch is dispatches the resource for processing.
	// It is invoked by the poller when a resource is
	// received by the remote system.
	Dispatch func(context.Context, *drone.Stage) error
}

// Poll opens N connections to the server to poll for pending
// stages for execution. Pending stages are dispatched to a
// Runner for execution.
func (p *Poller) Poll(ctx context.Context, n int) {
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(i int) {
			for {
				select {
				case <-ctx.Done():
					wg.Done()
					return
				default:
					p.poll(ctx, i+1)
				}
			}
		}(i)
	}

	wg.Wait()
}

// poll requests a stage for execution from the server, and then
// dispatches for execution.
func (p *Poller) poll(ctx context.Context, thread int) error {
	log := logger.FromContext(ctx).WithField("thread", thread)
	log.WithField("thread", thread).Debug("poller: request stage from remote server")

	// request a new build stage for execution from the central
	// build server.
	stage, err := p.Client.Request(ctx, p.Filter)
	if err == context.Canceled || err == context.DeadlineExceeded {
		log.WithError(err).Trace("poller: no stage returned")
		return nil
	}
	if err != nil {
		log.WithError(err).Error("poller: cannot request stage")
		return err
	}

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

	return p.Dispatch(
		logger.WithContext(noContext, log), stage)
}