fix: fix tracker/head/head.go

tracker/head/head.go

@@ -63,33 +63,40 @@ func New(
 func (self *TrackerHead) Start() error {
 	level.Info(self.logger).Log("msg", "starting", "reorgWaitPeriod", self.reorgWaitPeriod)
 
-	src, subs := self.waitSubscribe()
-	defer func() {
-		if subs != nil {
-			subs.Unsubscribe()
-		}
-	}()
+	for self.ctx.Err() == nil {
+		func() {
+			src, subs := self.waitSubscribe()
+			// according to docs, Unsubscribe() should always be called
+			defer subs.Unsubscribe()
+
+			// chances are that the context was canceled while waiting for the subscription
+			if self.ctx.Err() != nil {
+				return
+			}
 
-	go self.listen(src)
+			ctx, cancel := context.WithCancel(self.ctx)
+			defer cancel()
 
-	for {
-		select {
-		case <-self.ctx.Done():
-			return nil
-		case err := <-subs.Err():
-			level.Error(self.logger).Log("msg", "subscription failed will try to resubscribe", "err", err)
-			src, subs = self.waitSubscribe()
-			self.listen(src)
-		}
+			// using child context to cancel the listener
+			go self.listen(ctx, src)
+
+			select {
+			case <-self.ctx.Done():
+			case err := <-subs.Err():
+				level.Error(self.logger).Log("msg", "subscription failed will try to resubscribe", "err", err)
+			}
+		}()
 	}
+
+	return nil
 }
 
-func (self *TrackerHead) listen(src chan *types.Header) {
+func (self *TrackerHead) listen(ctx context.Context, src chan *types.Header) {
 	level.Info(self.logger).Log("msg", "starting new subs listener")
 
 	for {
 		select {
-		case <-self.ctx.Done():
+		case <-ctx.Done():
 			level.Info(self.logger).Log("msg", "subscription listener canceled")
 			return
 		case event := <-src:
@@ -98,7 +105,7 @@ func (self *TrackerHead) listen(src chan *types.Header) {
 			level.Debug(logger).Log("msg", "new block")
 			if self.reorgWaitPeriod == 0 {
 				go func(event *types.Header) {
-					ctx, cncl := context.WithTimeout(self.ctx, time.Minute)
+					ctx, cncl := context.WithTimeout(ctx, time.Minute)
 					defer cncl()
 					block, err := self.client.BlockByNumber(ctx, event.Number)
 					if err != nil {
@@ -107,7 +114,7 @@ func (self *TrackerHead) listen(src chan *types.Header) {
 					}
 					select {
 					case self.dstChan <- block:
-					case <-self.ctx.Done():
+					case <-ctx.Done():
 						return
 					}
 				}(event)
@@ -122,11 +129,11 @@ func (self *TrackerHead) listen(src chan *types.Header) {
 
 				select {
 				case <-waitForReorg.C:
-				case <-self.ctx.Done():
+				case <-ctx.Done():
 					return
 				}
 
-				ctx, cncl := context.WithTimeout(self.ctx, 2*time.Minute)
+				ctx, cncl := context.WithTimeout(ctx, 2*time.Minute)
 				defer cncl()
 
 				// Duplicate event numbers will still return the same block when using this query.
@@ -149,7 +156,7 @@ func (self *TrackerHead) listen(src chan *types.Header) {
 				select {
 				case self.dstChan <- block:
 					return
-				case <-self.ctx.Done():
+				case <-ctx.Done():
 					return
 				}
 

tracker/head/head_test.go

@@ -0,0 +1,153 @@
+// Copyright (c) The Cryptorium Authors.
+// Licensed under the MIT License.
+
+package head
+
+import (
+	"context"
+	"math/big"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/cryptoriums/packages/logging"
+	"github.com/ethereum/go-ethereum"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/require"
+)
+
+type testChainReader struct {
+	subLocker sync.Locker
+	sub       *testSubscription
+}
+
+var _ ethereum.ChainReader = &testChainReader{}
+
+func (tcr *testChainReader) BlockByHash(ctx context.Context, hash common.Hash) (block *types.Block, err error) {
+	block = types.NewBlock(&types.Header{}, nil, nil, nil, nil)
+	return
+}
+
+func (tcr *testChainReader) BlockByNumber(ctx context.Context, number *big.Int) (block *types.Block, err error) {
+	block = types.NewBlock(&types.Header{}, nil, nil, nil, nil)
+	return
+}
+
+func (tcr *testChainReader) HeaderByHash(ctx context.Context, hash common.Hash) (header *types.Header, err error) {
+	header = &types.Header{}
+	return
+}
+
+func (tcr *testChainReader) HeaderByNumber(ctx context.Context, number *big.Int) (header *types.Header, err error) {
+	header = &types.Header{}
+	return
+}
+
+func (tcr *testChainReader) TransactionCount(ctx context.Context, blockHash common.Hash) (count uint, err error) {
+	count = 0
+	return
+}
+
+func (tcr *testChainReader) TransactionInBlock(ctx context.Context, blockHash common.Hash, index uint) (tx *types.Transaction, err error) {
+	return
+}
+
+func (tcr *testChainReader) SubscribeNewHead(ctx context.Context, ch chan<- *types.Header) (sub ethereum.Subscription, err error) {
+	tcr.subLocker.Lock()
+	defer tcr.subLocker.Unlock()
+	tcr.sub = &testSubscription{
+		ctx:       ctx,
+		chHeaders: ch,
+		chErrors:  make(chan error),
+	}
+	sub = tcr.sub
+	return
+}
+
+func (tcr *testChainReader) sendHeader() {
+	tcr.subLocker.Lock()
+	defer tcr.subLocker.Unlock()
+	tcr.sub.chHeaders <- &types.Header{}
+}
+
+func (tcr *testChainReader) sendError() {
+	tcr.subLocker.Lock()
+	defer tcr.subLocker.Unlock()
+	tcr.sub.chErrors <- errors.New("random error")
+}
+
+type testSubscription struct {
+	ctx       context.Context
+	chHeaders chan<- *types.Header
+	chErrors  chan error
+}
+
+var _ ethereum.Subscription = &testSubscription{}
+
+func (ts *testSubscription) Unsubscribe() {
+}
+
+func (ts *testSubscription) Err() <-chan error {
+	return ts.chErrors
+}
+
+func TestTrackerHead(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	logger := logging.NewLogger()
+	client := &testChainReader{
+		subLocker: &sync.Mutex{},
+	}
+
+	tracker, blocks, err := New(
+		ctx,
+		logger,
+		client,
+		0,
+	)
+	require.NoError(t, err)
+
+	go func() {
+		err := tracker.Start()
+		require.NoError(t, err)
+	}()
+
+	n := 7
+
+	go func() {
+		// since we are using moked client, we just need a bit of time to start the subscription
+		time.Sleep(time.Millisecond * 100)
+
+		for i := 0; i < n; i++ {
+			client.sendHeader()
+			time.Sleep(time.Millisecond * 100)
+
+			client.sendError()
+			time.Sleep(time.Millisecond * 100)
+		}
+
+		cancel()
+	}()
+
+	var output []*types.Block
+
+	func() {
+		timeout := time.After(time.Second * 10)
+
+		for {
+			select {
+			case <-timeout:
+				t.Fatal("deadlock")
+			case <-ctx.Done():
+				return
+			case blk := <-blocks:
+				output = append(output, blk)
+			}
+		}
+	}()
+
+	require.Len(t, output, n)
+}