simplify BlockMiner.

itests/kit/blockminer.go

@@ -2,7 +2,6 @@ package kit
 
 import (
 	"context"
-	"fmt"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -11,42 +10,40 @@ import (
 	"github.com/filecoin-project/lotus/miner"
 )
 
-// BlockMiner is a utility that makes a test Miner Mine blocks on a timer.
+// BlockMiner is a utility that makes a test miner Mine blocks on a timer.
 type BlockMiner struct {
-	ctx       context.Context
-	t         *testing.T
-	miner     TestStorageNode
-	blocktime time.Duration
-	done      chan struct{}
+	t     *testing.T
+	miner TestStorageNode
 
-	Mine  int64
-	Nulls int64
+	nextNulls int64
+	stopCh    chan chan struct{}
 }
 
-func NewBlockMiner(ctx context.Context, t *testing.T, miner TestStorageNode, blocktime time.Duration) *BlockMiner {
+func NewBlockMiner(t *testing.T, miner TestStorageNode) *BlockMiner {
 	return &BlockMiner{
-		ctx:       ctx,
-		t:         t,
-		miner:     miner,
-		blocktime: blocktime,
-		Mine:      int64(1),
-		done:      make(chan struct{}),
+		t:      t,
+		miner:  miner,
+		stopCh: make(chan chan struct{}),
 	}
 }
 
-func (bm *BlockMiner) MineBlocks() {
+func (bm *BlockMiner) MineBlocks(ctx context.Context, blocktime time.Duration) {
 	time.Sleep(time.Second)
+
 	go func() {
-		defer close(bm.done)
-		for atomic.LoadInt64(&bm.Mine) == 1 {
+		for {
 			select {
-			case <-bm.ctx.Done():
+			case <-time.After(blocktime):
+			case <-ctx.Done():
+				return
+			case ch := <-bm.stopCh:
+				close(ch)
+				close(bm.stopCh)
 				return
-			case <-time.After(bm.blocktime):
 			}
 
-			nulls := atomic.SwapInt64(&bm.Nulls, 0)
-			if err := bm.miner.MineOne(bm.ctx, miner.MineReq{
+			nulls := atomic.SwapInt64(&bm.nextNulls, 0)
+			if err := bm.miner.MineOne(ctx, miner.MineReq{
 				InjectNulls: abi.ChainEpoch(nulls),
 				Done:        func(bool, abi.ChainEpoch, error) {},
 			}); err != nil {
@@ -56,8 +53,20 @@ func (bm *BlockMiner) MineBlocks() {
 	}()
 }
 
-func (bm *BlockMiner) Stop() {
-	atomic.AddInt64(&bm.Mine, -1)
-	fmt.Println("shutting down mining")
-	<-bm.done
+// InjectNulls injects the specified amount of null rounds in the next
+// mining rounds.
+func (bm *BlockMiner) InjectNulls(rounds abi.ChainEpoch) {
+	atomic.AddInt64(&bm.nextNulls, int64(rounds))
+}
+
+// Stop stops the block miner.
+func (bm *BlockMiner) Stop() {
+	bm.t.Log("shutting down mining")
+	if _, ok := <-bm.stopCh; ok {
+		// already stopped
+		return
+	}
+	ch := make(chan struct{})
+	bm.stopCh <- ch
+	<-ch
 }

itests/kit/deals.go

@@ -279,8 +279,8 @@ func ConnectAndStartMining(t *testing.T, blocktime time.Duration, client *impl.F
 	}
 	time.Sleep(time.Second)
 
-	blockMiner := NewBlockMiner(ctx, t, miner, blocktime)
-	blockMiner.MineBlocks()
+	blockMiner := NewBlockMiner(t, miner)
+	blockMiner.MineBlocks(ctx, blocktime)
 
 	return &DealsScaffold{
 		Ctx:        ctx,

itests/kit/net.go

@@ -28,8 +28,8 @@ func StartOneNodeOneMiner(ctx context.Context, t *testing.T, blocktime time.Dura
 	}
 
 	// Start mining blocks
-	bm := NewBlockMiner(ctx, t, miner, blocktime)
-	bm.MineBlocks()
+	bm := NewBlockMiner(t, miner)
+	bm.MineBlocks(ctx, blocktime)
 	t.Cleanup(bm.Stop)
 
 	// Get the full node's wallet address
@@ -64,8 +64,8 @@ func StartTwoNodesOneMiner(ctx context.Context, t *testing.T, blocktime time.Dur
 	}
 
 	// Start mining blocks
-	bm := NewBlockMiner(ctx, t, miner, blocktime)
-	bm.MineBlocks()
+	bm := NewBlockMiner(t, miner)
+	bm.MineBlocks(ctx, blocktime)
 	t.Cleanup(bm.Stop)
 
 	// Send some funds to register the second node

itests/paych_api_test.go

@@ -3,7 +3,6 @@ package itests
 import (
 	"context"
 	"fmt"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -52,8 +51,9 @@ func TestPaymentChannelsAPI(t *testing.T) {
 	}
 
 	// start mining blocks
-	bm := kit.NewBlockMiner(ctx, t, miner, 5*time.Millisecond)
-	bm.MineBlocks()
+	bm := kit.NewBlockMiner(t, miner)
+	bm.MineBlocks(ctx, 5*time.Millisecond)
+	t.Cleanup(bm.Stop)
 
 	// send some funds to register the receiver
 	receiverAddr, err := paymentReceiver.WalletNew(ctx, types.KTSecp256k1)
@@ -278,7 +278,7 @@ func waitForBlocks(ctx context.Context, t *testing.T, bm *kit.BlockMiner, paymen
 		}
 
 		// Add a batch of null blocks
-		atomic.StoreInt64(&bm.Nulls, int64(size-1))
+		bm.InjectNulls(abi.ChainEpoch(size - 1))
 
 		// Add a real block
 		m, err := paymentReceiver.MpoolPushMessage(ctx, &types.Message{