fix(lib/babe): add --babe-lead flag, update epoch handling logic

lib/babe/babe.go

@@ -45,7 +45,7 @@ type Service struct {
 	dev       bool
 	// lead is used when setting up a new network from genesis.
 	// the "lead" node is the node that is designated to build block 1, after which the rest of the nodes
-	// will sync it and determine the first slot of the network based on that block
+	// will sync block 1 and determine the first slot of the network based on it
 	lead bool
 
 	// Storage interfaces
@@ -206,10 +206,7 @@ func (b *Service) Start() error {
 		}
 	}
 
-	go func() {
-		b.initiate()
-	}()
-
+	go b.initiate()
 	return nil
 }
 
@@ -219,22 +216,24 @@ func (b *Service) waitForFirstBlock() error {
 
 	const firstBlockTimeout = time.Minute
 	timeout := time.NewTimer(firstBlockTimeout)
-	defer func() {
+	cleanup := func() {
 		if !timeout.Stop() {
 			<-timeout.C
 		}
-	}()
+	}
 
 	// loop until block 1
 	for {
 		select {
-		case block := <-ch:
-			if block != nil && block.Header.Number.Int64() > 0 {
+		case block, ok := <-ch:
+			if ok && block.Header.Number.Int64() > 0 {
+				cleanup()
 				return nil
 			}
 		case <-timeout.C:
 			return errFirstBlockTimeout
 		case <-b.ctx.Done():
+			cleanup()
 			return b.ctx.Err()
 		}
 	}
@@ -412,11 +411,11 @@ func (b *Service) invokeBlockAuthoring() error {
 
 		nextEpochStartTime := getSlotStartTime(nextEpochStart, b.slotDuration)
 		epochTimer := time.NewTimer(time.Until(nextEpochStartTime))
-		defer func() {
+		cleanup := func() {
 			if !epochTimer.Stop() {
 				<-epochTimer.C
 			}
-		}()
+		}
 
 		slotDone := make([]<-chan time.Time, b.epochLength-intoEpoch)
 		for i := 0; i < int(b.epochLength-intoEpoch); i++ {
@@ -428,8 +427,10 @@ func (b *Service) invokeBlockAuthoring() error {
 
 			select {
 			case <-b.ctx.Done():
+				cleanup()
 				return nil
 			case <-b.pause:
+				cleanup()
 				return nil
 			case <-slotDone[i]:
 				slotNum := startSlot + uint64(i)

tests/stress/grandpa_test.go

@@ -135,7 +135,7 @@ func TestStress_Grandpa_CatchUp(t *testing.T) {
 
 	time.Sleep(time.Second * 70) // let some rounds run
 	//nolint
-	node, err := utils.RunGossamer(t, numNodes-1, utils.TestDir(t, utils.KeyList[numNodes-1]), utils.GenesisSixAuths, utils.ConfigDefault, false)
+	node, err := utils.RunGossamer(t, numNodes-1, utils.TestDir(t, utils.KeyList[numNodes-1]), utils.GenesisSixAuths, utils.ConfigDefault, false, false)
 	require.NoError(t, err)
 	nodes = append(nodes, node)
 

tests/stress/stress_test.go

@@ -102,7 +102,7 @@ func TestSync_SingleBlockProducer(t *testing.T) {
 
 	// start block producing node first
 	//nolint
-	node, err := utils.RunGossamer(t, numNodes-1, utils.TestDir(t, utils.KeyList[numNodes-1]), utils.GenesisDev, utils.ConfigNoGrandpa, false)
+	node, err := utils.RunGossamer(t, numNodes-1, utils.TestDir(t, utils.KeyList[numNodes-1]), utils.GenesisDev, utils.ConfigNoGrandpa, false, true)
 	require.NoError(t, err)
 
 	// wait and start rest of nodes - if they all start at the same time the first round usually doesn't complete since
@@ -182,40 +182,40 @@ func TestSync_MultipleEpoch(t *testing.T) {
 	}
 }
 
-func TestSync_SingleSyncingNode(t *testing.T) {
-	// TODO: Fix this test and enable it.
-	t.Skip("skipping TestSync_SingleSyncingNode")
-	utils.SetLogLevel(log.LvlInfo)
-
-	// start block producing node
-	alice, err := utils.RunGossamer(t, 0, utils.TestDir(t, utils.KeyList[0]), utils.GenesisDev, utils.ConfigDefault, false)
-	require.NoError(t, err)
-	time.Sleep(time.Second * 15)
-
-	// start syncing node
-	bob, err := utils.RunGossamer(t, 1, utils.TestDir(t, utils.KeyList[1]), utils.GenesisDev, utils.ConfigNoBABE, false)
-	require.NoError(t, err)
-
-	nodes := []*utils.Node{alice, bob}
-	defer func() {
-		errList := utils.StopNodes(t, nodes)
-		require.Len(t, errList, 0)
-	}()
-
-	numCmps := 100
-	for i := 0; i < numCmps; i++ {
-		t.Log("comparing...", i)
-		_, err = compareBlocksByNumberWithRetry(t, nodes, strconv.Itoa(i))
-		require.NoError(t, err, i)
-	}
-}
+// func TestSync_SingleSyncingNode(t *testing.T) {
+// 	// TODO: Fix this test and enable it.
+// 	t.Skip("skipping TestSync_SingleSyncingNode")
+// 	utils.SetLogLevel(log.LvlInfo)
+
+// 	// start block producing node
+// 	alice, err := utils.RunGossamer(t, 0, utils.TestDir(t, utils.KeyList[0]), utils.GenesisDev, utils.ConfigDefault, false)
+// 	require.NoError(t, err)
+// 	time.Sleep(time.Second * 15)
+
+// 	// start syncing node
+// 	bob, err := utils.RunGossamer(t, 1, utils.TestDir(t, utils.KeyList[1]), utils.GenesisDev, utils.ConfigNoBABE, false)
+// 	require.NoError(t, err)
+
+// 	nodes := []*utils.Node{alice, bob}
+// 	defer func() {
+// 		errList := utils.StopNodes(t, nodes)
+// 		require.Len(t, errList, 0)
+// 	}()
+
+// 	numCmps := 100
+// 	for i := 0; i < numCmps; i++ {
+// 		t.Log("comparing...", i)
+// 		_, err = compareBlocksByNumberWithRetry(t, nodes, strconv.Itoa(i))
+// 		require.NoError(t, err, i)
+// 	}
+// }
 
 func TestSync_Bench(t *testing.T) {
 	utils.SetLogLevel(log.LvlInfo)
 	numBlocks := 64
 
 	// start block producing node
-	alice, err := utils.RunGossamer(t, 0, utils.TestDir(t, utils.KeyList[1]), utils.GenesisDev, utils.ConfigNoGrandpa, false)
+	alice, err := utils.RunGossamer(t, 0, utils.TestDir(t, utils.KeyList[1]), utils.GenesisDev, utils.ConfigNoGrandpa, false, true)
 	require.NoError(t, err)
 
 	for {
@@ -236,7 +236,7 @@ func TestSync_Bench(t *testing.T) {
 	t.Log("BABE paused")
 
 	// start syncing node
-	bob, err := utils.RunGossamer(t, 1, utils.TestDir(t, utils.KeyList[0]), utils.GenesisDev, utils.ConfigNotAuthority, false)
+	bob, err := utils.RunGossamer(t, 1, utils.TestDir(t, utils.KeyList[0]), utils.GenesisDev, utils.ConfigNotAuthority, false, true)
 	require.NoError(t, err)
 
 	nodes := []*utils.Node{alice, bob}
@@ -290,7 +290,7 @@ func TestSync_Restart(t *testing.T) {
 
 	// start block producing node first
 	//nolint
-	node, err := utils.RunGossamer(t, numNodes-1, utils.TestDir(t, utils.KeyList[numNodes-1]), utils.GenesisDefault, utils.ConfigDefault, false)
+	node, err := utils.RunGossamer(t, numNodes-1, utils.TestDir(t, utils.KeyList[numNodes-1]), utils.GenesisDefault, utils.ConfigDefault, false, true)
 	require.NoError(t, err)
 
 	// wait and start rest of nodes
@@ -343,15 +343,15 @@ func TestSync_SubmitExtrinsic(t *testing.T) {
 	idx := 0 // TODO: randomise this
 
 	// start block producing node first
-	node, err := utils.RunGossamer(t, 0, utils.TestDir(t, utils.KeyList[0]), utils.GenesisDev, utils.ConfigNoGrandpa, false)
+	node, err := utils.RunGossamer(t, 0, utils.TestDir(t, utils.KeyList[0]), utils.GenesisDev, utils.ConfigNoGrandpa, false, true)
 	require.NoError(t, err)
 	nodes := []*utils.Node{node}
 
 	// Start rest of nodes
-	node, err = utils.RunGossamer(t, 1, utils.TestDir(t, utils.KeyList[1]), utils.GenesisDev, utils.ConfigNotAuthority, false)
+	node, err = utils.RunGossamer(t, 1, utils.TestDir(t, utils.KeyList[1]), utils.GenesisDev, utils.ConfigNotAuthority, false, false)
 	require.NoError(t, err)
 	nodes = append(nodes, node)
-	node, err = utils.RunGossamer(t, 2, utils.TestDir(t, utils.KeyList[2]), utils.GenesisDev, utils.ConfigNotAuthority, false)
+	node, err = utils.RunGossamer(t, 2, utils.TestDir(t, utils.KeyList[2]), utils.GenesisDev, utils.ConfigNotAuthority, false, false)
 	require.NoError(t, err)
 	nodes = append(nodes, node)
 

tests/utils/gossamer_utils.go

@@ -220,14 +220,14 @@ func StartGossamer(t *testing.T, node *Node, websocket bool) error {
 }
 
 // RunGossamer will initialise and start a gossamer instance
-func RunGossamer(t *testing.T, idx int, basepath, genesis, config string, websocket bool) (*Node, error) {
+func RunGossamer(t *testing.T, idx int, basepath, genesis, config string, websocket, babeLead bool) (*Node, error) {
 	node, err := InitGossamer(idx, basepath, genesis, config)
 	if err != nil {
 		logger.Crit("could not initialise gossamer", "error", err)
 		os.Exit(1)
 	}
 
-	if idx == 0 {
+	if idx == 0 || babeLead {
 		node.BABELead = true
 	}
 
@@ -316,7 +316,7 @@ func InitializeAndStartNodes(t *testing.T, num int, genesis, config string) ([]*
 			if i < len(KeyList) {
 				name = KeyList[i]
 			}
-			node, err := RunGossamer(t, i, TestDir(t, name), genesis, config, false)
+			node, err := RunGossamer(t, i, TestDir(t, name), genesis, config, false, false)
 			if err != nil {
 				logger.Error("failed to run gossamer", "i", i)
 			}
@@ -346,7 +346,7 @@ func InitializeAndStartNodesWebsocket(t *testing.T, num int, genesis, config str
 			if i < len(KeyList) {
 				name = KeyList[i]
 			}
-			node, err := RunGossamer(t, i, TestDir(t, name), genesis, config, true)
+			node, err := RunGossamer(t, i, TestDir(t, name), genesis, config, true, false)
 			if err != nil {
 				logger.Error("failed to run gossamer", "i", i)
 			}