Create HeadPoller for Multi-Node

common/client/head_poller.go

@@ -0,0 +1,93 @@
+package client
+
+import (
+	"sync"
+	"time"
+
+	"github.com/smartcontractkit/chainlink-common/pkg/services"
+
+	"github.com/smartcontractkit/chainlink/v2/common/types"
+)
+
+// HeadPoller is a component that polls a function at a given interval
+// and delivers the result to a channel. It is used to poll for new heads
+// and implements the Subscription interface.
+type HeadPoller[
+	HEAD Head,
+] struct {
+	services.StateMachine
+	pollingInterval time.Duration
+	pollingFunc     func() (HEAD, error)
+	channel         chan<- HEAD
+	errCh           chan error
+
+	stopCh chan struct{}
+	wg     sync.WaitGroup
+}
+
+// NewHeadPoller creates a new HeadPoller instance
+func NewHeadPoller[
+	HEAD Head,
+](pollingInterval time.Duration, pollingFunc func() (HEAD, error), channel chan<- HEAD) HeadPoller[HEAD] {
+	return HeadPoller[HEAD]{
+		pollingInterval: pollingInterval,
+		pollingFunc:     pollingFunc,
+		channel:         channel,
+		errCh:           make(chan error),
+		stopCh:          make(chan struct{}),
+	}
+}
+
+var _ types.Subscription = &HeadPoller[Head]{}
+
+func (p *HeadPoller[HEAD]) Start() error {
+	return p.StartOnce("HeadPoller", func() error {
+		p.wg.Add(1)
+		go p.pollingLoop()
+		return nil
+	})
+}
+
+// Unsubscribe cancels the sending of events to the data channel
+func (p *HeadPoller[HEAD]) Unsubscribe() {
+	_ = p.StopOnce("HeadPoller", func() error {
+		close(p.stopCh)
+		p.wg.Wait()
+		close(p.errCh)
+		return nil
+	})
+}
+
+func (p *HeadPoller[HEAD]) Err() <-chan error {
+	return p.errCh
+}
+
+func (p *HeadPoller[HEAD]) pollingLoop() {
+	defer p.wg.Done()
+
+	ticker := time.NewTicker(p.pollingInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-p.stopCh:
+			return
+		case <-ticker.C:
+			result, err := p.pollingFunc()
+			if err != nil {
+				select {
+				case p.errCh <- err:
+					continue
+				case <-p.stopCh:
+					return
+				}
+			}
+
+			select {
+			case p.channel <- result:
+			case <-p.stopCh:
+				return
+			}
+		}
+	}
+}

common/client/head_poller_test.go

@@ -0,0 +1,78 @@
+package client
+
+import (
+	"errors"
+	"math/big"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Poller(t *testing.T) {
+	// Mock polling function that returns a new value every time it's called
+	var pollNumber int
+	pollFunc := func() (Head, error) {
+		pollNumber++
+		h := head{
+			BlockNumber:     int64(pollNumber),
+			BlockDifficulty: big.NewInt(int64(pollNumber)),
+		}
+		return h.ToMockHead(t), nil
+	}
+
+	// data channel to receive updates from the poller
+	channel := make(chan Head, 1)
+
+	// Create poller and subscribe to receive data
+	poller := NewHeadPoller[Head](time.Millisecond, pollFunc, channel)
+
+	require.NoError(t, poller.Start())
+	defer poller.Unsubscribe()
+
+	// Create goroutine to receive updates from the poller
+	pollCount := 0
+	pollMax := 50
+	go func() {
+		for ; pollCount < pollMax; pollCount++ {
+			h := <-channel
+			assert.Equal(t, int64(pollNumber), h.BlockNumber())
+		}
+	}()
+
+	// Wait for a short duration to allow for some polling iterations
+	time.Sleep(100 * time.Millisecond)
+	require.Equal(t, pollMax, pollCount)
+}
+
+func Test_Poller_Error(t *testing.T) {
+	// Mock polling function that returns an error
+	pollFunc := func() (Head, error) {
+		return nil, errors.New("polling error")
+	}
+
+	// data channel to receive updates from the poller
+	channel := make(chan Head, 1)
+
+	// Create poller and subscribe to receive data
+	poller := NewHeadPoller[Head](time.Millisecond, pollFunc, channel)
+
+	require.NoError(t, poller.Start())
+	defer poller.Unsubscribe()
+
+	// Create goroutine to receive updates from the poller
+	pollCount := 0
+	pollMax := 50
+	go func() {
+		for ; pollCount < pollMax; pollCount++ {
+			err := <-poller.Err()
+			require.Error(t, err)
+			require.Equal(t, "polling error", err.Error())
+		}
+	}()
+
+	// Wait for a short duration to allow for some polling iterations
+	time.Sleep(100 * time.Millisecond)
+	require.Equal(t, pollMax, pollCount)
+}