github.com/drone/runner-go@v1.12.0/poller/poller.go (about)

     1  // Copyright 2019 Drone.IO Inc. All rights reserved.
     2  // Use of this source code is governed by the Polyform License
     3  // that can be found in the LICENSE file.
     4  
     5  package poller
     6  
     7  import (
     8  	"context"
     9  	"sync"
    10  
    11  	"github.com/drone/drone-go/drone"
    12  	"github.com/drone/runner-go/client"
    13  	"github.com/drone/runner-go/logger"
    14  )
    15  
    16  var noContext = context.Background()
    17  
    18  // Poller polls the server for pending stages and dispatches
    19  // for execution by the Runner.
    20  type Poller struct {
    21  	Client client.Client
    22  	Filter *client.Filter
    23  
    24  	// Dispatch is dispatches the resource for processing.
    25  	// It is invoked by the poller when a resource is
    26  	// received by the remote system.
    27  	Dispatch func(context.Context, *drone.Stage) error
    28  }
    29  
    30  // Poll opens N connections to the server to poll for pending
    31  // stages for execution. Pending stages are dispatched to a
    32  // Runner for execution.
    33  func (p *Poller) Poll(ctx context.Context, n int) {
    34  	var wg sync.WaitGroup
    35  	for i := 0; i < n; i++ {
    36  		wg.Add(1)
    37  		go func(i int) {
    38  			for {
    39  				select {
    40  				case <-ctx.Done():
    41  					wg.Done()
    42  					return
    43  				default:
    44  					p.poll(ctx, i+1)
    45  				}
    46  			}
    47  		}(i)
    48  	}
    49  
    50  	wg.Wait()
    51  }
    52  
    53  // poll requests a stage for execution from the server, and then
    54  // dispatches for execution.
    55  func (p *Poller) poll(ctx context.Context, thread int) error {
    56  	log := logger.FromContext(ctx).WithField("thread", thread)
    57  	log.WithField("thread", thread).Debug("poller: request stage from remote server")
    58  
    59  	// request a new build stage for execution from the central
    60  	// build server.
    61  	stage, err := p.Client.Request(ctx, p.Filter)
    62  	if err == context.Canceled || err == context.DeadlineExceeded {
    63  		log.WithError(err).Trace("poller: no stage returned")
    64  		return nil
    65  	}
    66  	if err != nil {
    67  		log.WithError(err).Error("poller: cannot request stage")
    68  		return err
    69  	}
    70  
    71  	// exit if a nil or empty stage is returned from the system
    72  	// and allow the runner to retry.
    73  	if stage == nil || stage.ID == 0 {
    74  		return nil
    75  	}
    76  
    77  	return p.Dispatch(
    78  		logger.WithContext(noContext, log), stage)
    79  }