Merge branch 'main' into snapshots

README.md

@@ -655,16 +655,10 @@ The network runner allows users to interact with an AvalancheGo network using th
 ```go
 // Network is an abstraction of an Avalanche network
 type Network interface {
-  // Returns a chan that is closed when
-  // all the nodes in the network are healthy.
-  // If an error is sent on this channel, at least 1
-  // node didn't become healthy before the timeout.
-  // If an error isn't sent on the channel before it
-  // closes, all the nodes are healthy.
+  // Returns nil if all the nodes in the network are healthy.
   // A stopped network is considered unhealthy.
   // Timeout is given by the context parameter.
-  // [ctx] must eventually be cancelled -- if it isn't, a goroutine is leaked.
-  Healthy(context.Context) chan error
+  Healthy(context.Context) error
   // Stop all the nodes.
   // Returns ErrStopped if Stop() was previously called.
   Stop(context.Context) error

examples/local/fivenodenetwork/main.go

@@ -89,9 +89,8 @@ func run(log logging.Logger, binaryPath string) error {
 	// Wait until the nodes in the network are ready
 	ctx, cancel := context.WithTimeout(context.Background(), healthyTimeout)
 	defer cancel()
-	healthyChan := nw.Healthy(ctx)
 	log.Info("waiting for all nodes to report healthy...")
-	if err := <-healthyChan; err != nil {
+	if err := nw.Healthy(ctx); err != nil {
 		return err
 	}
 

examples/local/indepth/main.go

@@ -93,10 +93,10 @@ func run(log logging.Logger, binaryPath string) error {
 	// Wait until the nodes in the network are ready
 	ctx, cancel := context.WithTimeout(context.Background(), healthyTimeout)
 	defer cancel()
-	healthyChan := nw.Healthy(ctx)
 	log.Info("waiting for all nodes to report healthy...")
-	if err := <-healthyChan; err != nil {
+	if err := nw.Healthy(ctx); err != nil {
 		return err
+
 	}
 
 	// Print the node names
@@ -150,9 +150,8 @@ func run(log logging.Logger, binaryPath string) error {
 	// Wait until the nodes in the updated network are ready
 	ctx, cancel = context.WithTimeout(context.Background(), healthyTimeout)
 	defer cancel()
-	healthyChan = nw.Healthy(ctx)
 	log.Info("waiting for updated network to report healthy...")
-	if err := <-healthyChan; err != nil {
+	if err := nw.Healthy(ctx); err != nil {
 		return err
 	}
 

local/network.go

@@ -73,8 +73,9 @@ type localNetwork struct {
 	newAPIClientF api.NewAPIClientF
 	// Used to create new node processes
 	nodeProcessCreator NodeProcessCreator
-	// Closed when network is done shutting down
-	closedOnStopCh chan struct{}
+	stopOnce           sync.Once
+	// Closed when Stop begins.
+	onStopCh chan struct{}
 	// For node name generation
 	nextNodeSuffix uint64
 	// Node Name --> Node
@@ -89,7 +90,8 @@ type localNetwork struct {
 	// directory where networks can be persistently saved
 	snapshotsDir string
 	// To keep track of network initialization
-	defined bool
+	definedLock sync.RWMutex
+	defined     bool
 }
 
 var (
@@ -188,7 +190,7 @@ func (npc *nodeProcessCreator) NewNodeProcess(config node.Config, args ...string
 	if config.RedirectStdout {
 		stdout, err := cmd.StdoutPipe()
 		if err != nil {
-			return nil, fmt.Errorf("Could not create stdout pipe: %s", err)
+			return nil, fmt.Errorf("couldn't create stdout pipe: %s", err)
 		}
 		// redirect stdout and assign a color to the text
 		utils.ColorAndPrepend(stdout, npc.stdout, config.Name, color)
@@ -196,7 +198,7 @@ func (npc *nodeProcessCreator) NewNodeProcess(config node.Config, args ...string
 	if config.RedirectStderr {
 		stderr, err := cmd.StderrPipe()
 		if err != nil {
-			return nil, fmt.Errorf("Could not create stderr pipe: %s", err)
+			return nil, fmt.Errorf("couldn't create stderr pipe: %s", err)
 		}
 		// redirect stderr and assign a color to the text
 		utils.ColorAndPrepend(stderr, npc.stderr, config.Name, color)
@@ -251,7 +253,7 @@ func newNetwork(
 	net := &localNetwork{
 		nextNodeSuffix:     1,
 		nodes:              map[string]*localNode{},
-		closedOnStopCh:     make(chan struct{}),
+		onStopCh:           make(chan struct{}),
 		log:                log,
 		bootstraps:         beacon.NewSet(),
 		newAPIClientF:      newAPIClientF,
@@ -325,13 +327,13 @@ func NewDefaultConfigNNodes(binaryPath string, numNodes uint32) (network.Config,
 func (ln *localNetwork) LoadConfig(ctx context.Context, networkConfig network.Config) error {
 	ln.lock.Lock()
 	defer ln.lock.Unlock()
+	if ln.wasDefined() {
+		return errors.New("configuration already loaded")
+	}
 	return ln.loadConfig(ctx, networkConfig)
 }
 
 func (ln *localNetwork) loadConfig(ctx context.Context, networkConfig network.Config) error {
-	if ln.defined {
-		return errors.New("configuration already loaded")
-	}
 	if err := networkConfig.Validate(); err != nil {
 		return fmt.Errorf("config failed validation: %w", err)
 	}
@@ -360,7 +362,10 @@ func (ln *localNetwork) loadConfig(ctx context.Context, networkConfig network.Co
 		}
 	}
 
+	ln.definedLock.Lock()
 	ln.defined = true
+	ln.definedLock.Unlock()
+
 	for _, nodeConfig := range nodeConfigs {
 		if _, err := ln.addNode(nodeConfig); err != nil {
 			if err := ln.stop(ctx); err != nil {
@@ -379,18 +384,18 @@ func (ln *localNetwork) AddNode(nodeConfig node.Config) (node.Node, error) {
 	ln.lock.Lock()
 	defer ln.lock.Unlock()
 
-	return ln.addNode(nodeConfig)
-}
-
-// Assumes [ln.lock] is held.
-func (ln *localNetwork) addNode(nodeConfig node.Config) (node.Node, error) {
-	if !ln.defined {
+	if !ln.wasDefined() {
 		return nil, network.ErrUndefined
 	}
-	if ln.isStopped() {
+	if ln.stopCalled() {
 		return nil, network.ErrStopped
 	}
 
+	return ln.addNode(nodeConfig)
+}
+
+// Assumes [ln.lock] is held and [ln.Stop] hasn't been called.
+func (ln *localNetwork) addNode(nodeConfig node.Config) (node.Node, error) {
 	if nodeConfig.Flags == nil {
 		nodeConfig.Flags = make(map[string]interface{})
 	}
@@ -461,57 +466,58 @@ func (ln *localNetwork) addNode(nodeConfig node.Config) (node.Node, error) {
 }
 
 // See network.Network
-func (ln *localNetwork) Healthy(ctx context.Context) chan error {
+func (ln *localNetwork) Healthy(ctx context.Context) error {
+	ln.lock.RLock()
+	defer ln.lock.RUnlock()
+
 	zap.L().Info("checking local network healthiness", zap.Int("nodes", len(ln.nodes)))
-	healthyChan := make(chan error, 1)
 
 	// Return unhealthy if the network is undefined
-	if !ln.defined {
-		healthyChan <- network.ErrUndefined
-		return healthyChan
+	if !ln.wasDefined() {
+		return network.ErrUndefined
 	}
 
 	// Return unhealthy if the network is stopped
-	if ln.isStopped() {
-		healthyChan <- network.ErrStopped
-		return healthyChan
+	if ln.stopCalled() {
+		return network.ErrStopped
 	}
 
-	go func() {
-		// TODO: This will block the network for the duration of the health call.
-		// Maybe a better solution can be found.
-		ln.lock.RLock()
-		defer ln.lock.RUnlock()
+	// Derive a new context that's cancelled when Stop is called,
+	// so that we calls to Healthy() below immediately return.
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	go func(ctx context.Context) {
+		// This goroutine runs until [ln.Stop] is called
+		// or this function returns.
+		select {
+		case <-ln.onStopCh:
+			cancel()
+		case <-ctx.Done():
+		}
+	}(ctx)
 
-		errGr, cctx := errgroup.WithContext(ctx)
-		for _, node := range ln.nodes {
-			node := node
-			errGr.Go(func() error {
-				// Every constants.HealthCheckInterval, query node for health status.
-				// Do this until ctx timeout
-				for {
-					select {
-					case <-ln.closedOnStopCh:
-						return network.ErrStopped
-					case <-cctx.Done():
-						return fmt.Errorf("node %q failed to become healthy within timeout", node.GetName())
-					case <-time.After(healthCheckFreq):
-					}
-					health, err := node.client.HealthAPI().Health(cctx)
-					if err == nil && health.Healthy {
-						ln.log.Debug("node %q became healthy", node.name)
-						return nil
-					}
+	errGr, ctx := errgroup.WithContext(ctx)
+	for _, node := range ln.nodes {
+		node := node
+		errGr.Go(func() error {
+			// Every [healthCheckFreq], query node for health status.
+			// Do this until ctx timeout or network closed.
+			for {
+				health, err := node.client.HealthAPI().Health(ctx)
+				if err == nil && health.Healthy {
+					ln.log.Debug("node %q became healthy", node.name)
+					return nil
 				}
-			})
-		}
-		// Wait until all nodes are ready or timeout
-		if err := errGr.Wait(); err != nil {
-			healthyChan <- err
-		}
-		close(healthyChan)
-	}()
-	return healthyChan
+				select {
+				case <-ctx.Done():
+					return fmt.Errorf("node %q failed to become healthy within timeout, or network stopped", node.GetName())
+				case <-time.After(healthCheckFreq):
+				}
+			}
+		})
+	}
+	// Wait until all nodes are ready or timeout
+	return errGr.Wait()
 }
 
 // See network.Network
@@ -519,11 +525,10 @@ func (ln *localNetwork) GetNode(nodeName string) (node.Node, error) {
 	ln.lock.RLock()
 	defer ln.lock.RUnlock()
 
-	if !ln.defined {
+	if !ln.wasDefined() {
 		return nil, network.ErrUndefined
 	}
-
-	if ln.isStopped() {
+	if ln.stopCalled() {
 		return nil, network.ErrStopped
 	}
 
@@ -539,11 +544,10 @@ func (ln *localNetwork) GetNodeNames() ([]string, error) {
 	ln.lock.RLock()
 	defer ln.lock.RUnlock()
 
-	if !ln.defined {
+	if !ln.wasDefined() {
 		return nil, network.ErrUndefined
 	}
-
-	if ln.isStopped() {
+	if ln.stopCalled() {
 		return nil, network.ErrStopped
 	}
 
@@ -561,11 +565,10 @@ func (ln *localNetwork) GetAllNodes() (map[string]node.Node, error) {
 	ln.lock.RLock()
 	defer ln.lock.RUnlock()
 
-	if !ln.defined {
+	if !ln.wasDefined() {
 		return nil, network.ErrUndefined
 	}
-
-	if ln.isStopped() {
+	if ln.stopCalled() {
 		return nil, network.ErrStopped
 	}
 
@@ -577,21 +580,26 @@ func (ln *localNetwork) GetAllNodes() (map[string]node.Node, error) {
 }
 
 func (ln *localNetwork) Stop(ctx context.Context) error {
-	ln.lock.Lock()
-	defer ln.lock.Unlock()
-
-	return ln.stop(ctx)
-}
-
-// Assumes [net.lock] is held
-func (ln *localNetwork) stop(ctx context.Context) error {
-	if !ln.defined {
+	if !ln.wasDefined() {
 		return network.ErrUndefined
 	}
-	if ln.isStopped() {
-		ln.log.Debug("stop() called multiple times")
-		return network.ErrStopped
-	}
+
+	err := network.ErrStopped
+	ln.stopOnce.Do(
+		func() {
+			close(ln.onStopCh)
+
+			ln.lock.Lock()
+			defer ln.lock.Unlock()
+
+			err = ln.stop(ctx)
+		},
+	)
+	return err
+}
+
+// Assumes [ln.lock] is held.
+func (ln *localNetwork) stop(ctx context.Context) error {
 	ctx, cancel := context.WithTimeout(ctx, stopTimeout)
 	defer cancel()
 	errs := wrappers.Errs{}
@@ -610,27 +618,26 @@ func (ln *localNetwork) stop(ctx context.Context) error {
 			errs.Add(err)
 		}
 	}
-	close(ln.closedOnStopCh)
 	ln.log.Info("done stopping network")
 	return errs.Err
 }
 
-// Sends a SIGTERM to the given node and removes it from this network
+// Sends a SIGTERM to the given node and removes it from this network.
 func (ln *localNetwork) RemoveNode(nodeName string) error {
 	ln.lock.Lock()
 	defer ln.lock.Unlock()
 
+	if !ln.wasDefined() {
+		return network.ErrUndefined
+	}
+	if ln.stopCalled() {
+		return network.ErrStopped
+	}
 	return ln.removeNode(nodeName)
 }
 
-// Assumes [net.lock] is held
+// Assumes [ln.lock] is held.
 func (ln *localNetwork) removeNode(nodeName string) error {
-	if !ln.defined {
-		return network.ErrUndefined
-	}
-	if ln.isStopped() {
-		return network.ErrStopped
-	}
 	ln.log.Debug("removing node %q", nodeName)
 	node, ok := ln.nodes[nodeName]
 	if !ok {
@@ -658,10 +665,10 @@ func (ln *localNetwork) removeNode(nodeName string) error {
 func (ln *localNetwork) SaveSnapshot(ctx context.Context, snapshotName string) error {
 	ln.lock.Lock()
 	defer ln.lock.Unlock()
-	if !ln.defined {
+	if !ln.wasDefined() {
 		return network.ErrUndefined
 	}
-	if ln.isStopped() {
+	if ln.stopCalled() {
 		return network.ErrStopped
 	}
 	if len(snapshotName) == 0 {
@@ -757,7 +764,7 @@ func (ln *localNetwork) SaveSnapshot(ctx context.Context, snapshotName string) e
 func (ln *localNetwork) LoadSnapshot(ctx context.Context, snapshotName string) error {
 	ln.lock.Lock()
 	defer ln.lock.Unlock()
-	if ln.defined {
+	if ln.wasDefined() {
 		return errors.New("configuration already loaded")
 	}
 	snapshotDir := filepath.Join(ln.snapshotsDir, snapshotPrefix+snapshotName)
@@ -818,16 +825,23 @@ func (ln *localNetwork) GetSnapshotNames() ([]string, error) {
 	return snapshots, nil
 }
 
-// Assumes [net.lock] is held
-func (ln *localNetwork) isStopped() bool {
+// Returns whether Stop has been called.
+func (ln *localNetwork) stopCalled() bool {
 	select {
-	case <-ln.closedOnStopCh:
+	case <-ln.onStopCh:
 		return true
 	default:
 		return false
 	}
 }
 
+// Returns whether the network has been defined
+func (ln *localNetwork) wasDefined() bool {
+	ln.definedLock.Lock()
+	defer ln.definedLock.Unlock()
+	return ln.defined
+}
+
 // createFileAndWrite creates a file with the given path and
 // writes the given contents
 func createFileAndWrite(path string, contents []byte) error {

local/network_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/ava-labs/avalanchego/snow/networking/router"
 	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/logging"
+	"github.com/ava-labs/avalanchego/utils/rpc"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 )
@@ -665,19 +666,19 @@ func TestStoppedNetwork(t *testing.T) {
 	assert.EqualValues(net.Stop(context.Background()), network.ErrStopped)
 	// AddNode failure
 	_, err = net.AddNode(networkConfig.NodeConfigs[1])
-	assert.EqualValues(err, network.ErrStopped)
+	assert.EqualValues(network.ErrStopped, err)
 	// GetNode failure
 	_, err = net.GetNode(networkConfig.NodeConfigs[0].Name)
 	assert.EqualValues(err, network.ErrStopped)
 	// second GetNodeNames should return no nodes
 	_, err = net.GetNodeNames()
-	assert.EqualValues(err, network.ErrStopped)
+	assert.EqualValues(network.ErrStopped, err)
 	// RemoveNode failure
-	assert.EqualValues(net.RemoveNode(networkConfig.NodeConfigs[0].Name), network.ErrStopped)
+	assert.EqualValues(network.ErrStopped, net.RemoveNode(networkConfig.NodeConfigs[0].Name))
 	// Healthy failure
 	assert.EqualValues(awaitNetworkHealthy(net, defaultHealthyTimeout), network.ErrStopped)
 	_, err = net.GetAllNodes()
-	assert.EqualValues(network.ErrStopped, err)
+	assert.EqualValues(err, network.ErrStopped)
 }
 
 func TestGetAllNodes(t *testing.T) {
@@ -938,8 +939,7 @@ func testNetworkConfig(t *testing.T) network.Config {
 func awaitNetworkHealthy(net network.Network, timeout time.Duration) error {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
-	healthyCh := net.Healthy(ctx)
-	return <-healthyCh
+	return net.Healthy(ctx)
 }
 
 func TestAddNetworkFlags(t *testing.T) {
@@ -1254,3 +1254,58 @@ func TestRemoveBeacon(t *testing.T) {
 	assert.NoError(err)
 	assert.Equal(0, net.bootstraps.Len())
 }
+
+// Returns an API client where:
+// * The Health API's Health method always returns an error after the
+//   given context is cancelled.
+// * The CChainEthAPI's Close method may be called
+// * Only the above 2 methods may be called
+func newMockAPIHealthyBlocks(ipAddr string, port uint16) api.Client {
+	healthClient := &healthmocks.Client{}
+	healthClient.On("Health", mock.MatchedBy(func(_ context.Context) bool { return true }), mock.Anything).Return(
+		func(ctx context.Context, _ ...rpc.Option) *health.APIHealthReply {
+			<-ctx.Done()
+			return nil
+		},
+		func(ctx context.Context, _ ...rpc.Option) error {
+			<-ctx.Done()
+			return ctx.Err()
+		},
+	)
+	// ethClient used when removing nodes, to close websocket connection
+	ethClient := &apimocks.EthClient{}
+	ethClient.On("Close").Return()
+	client := &apimocks.Client{}
+	client.On("HealthAPI").Return(healthClient)
+	client.On("CChainEthAPI").Return(ethClient)
+	return client
+}
+
+// Assert that if the network's Stop method is called while
+// a call to Healthy is ongoing, Healthy returns immediately.
+func TestHealthyDuringNetworkStop(t *testing.T) {
+	assert := assert.New(t)
+	networkConfig := testNetworkConfig(t)
+	// Calls to a node's Healthy() function blocks until context cancelled
+	net, err := newNetwork(logging.NoLog{}, newMockAPIHealthyBlocks, &localTestSuccessfulNodeProcessCreator{}, "", "")
+	assert.NoError(err)
+	err = net.LoadConfig(context.Background(), networkConfig)
+	assert.NoError(err)
+
+	healthyChan := make(chan error)
+	go func() {
+		healthyChan <- net.Healthy(context.Background())
+	}()
+	// Wait to make sure we're actually blocking on Health API call
+	time.Sleep(500 * time.Millisecond)
+	err = net.Stop(context.Background())
+	assert.NoError(err)
+	select {
+	case err := <-healthyChan:
+		assert.Error(err)
+	case <-time.After(1 * time.Second):
+		// Since [net.Stop] was called, [net.Healthy] should immediately return.
+		// We assume that it will do so within 1 second.
+		assert.Fail("Healthy should've returned immediately because network closed")
+	}
+}

local/node_test.go

@@ -17,7 +17,6 @@ import (
 	"github.com/ava-labs/avalanchego/network/peer"
 	"github.com/ava-labs/avalanchego/staking"
 	"github.com/ava-labs/avalanchego/utils"
-	avago_utils "github.com/ava-labs/avalanchego/utils"
 	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/wrappers"
 	"github.com/ava-labs/avalanchego/version"
@@ -66,7 +65,7 @@ func verifyProtocol(
 	// send the peer our version and peerlist
 
 	// create the version message
-	myIP := avago_utils.IPDesc{
+	myIP := utils.IPDesc{
 		IP:   net.IPv6zero,
 		Port: 0,
 	}

network/network.go

@@ -18,15 +18,10 @@ type Network interface {
 	// Initializes and starts network using the given snapshot
 	// To be executed after network creation. Enables the other calls.
 	LoadSnapshot(context.Context, string) error
-	// Returns a chan that is closed when
-	// all the nodes in the network are healthy.
-	// If an error is sent on this channel, at least 1
-	// node didn't become healthy before the timeout.
-	// If an error isn't sent on the channel before it
-	// closes, all the nodes are healthy.
+	// Returns nil if all the nodes in the network are healthy.
 	// A stopped network is considered unhealthy.
 	// Timeout is given by the context parameter.
-	Healthy(context.Context) chan error
+	Healthy(context.Context) error
 	// Stop all the nodes.
 	// Returns ErrStopped if Stop() was previously called.
 	Stop(context.Context) error

server/custom_vms.go

@@ -84,16 +84,8 @@ func (lc *localNetwork) waitForCustomVMsReady(ctx context.Context) error {
 	println()
 	color.Outf("{{blue}}{{bold}}waiting for custom VMs to report healthy...{{/}}\n")
 
-	hc := lc.nw.Healthy(ctx)
-	select {
-	case <-lc.stopc:
-		return errAborted
-	case <-ctx.Done():
-		return ctx.Err()
-	case err := <-hc:
-		if err != nil {
-			return err
-		}
+	if err := lc.nw.Healthy(ctx); err != nil {
+		return err
 	}
 
 	for nodeName, nodeInfo := range lc.nodeInfos {

server/network.go

@@ -352,16 +352,8 @@ var errAborted = errors.New("aborted")
 func (lc *localNetwork) waitForLocalClusterReady(ctx context.Context) error {
 	color.Outf("{{blue}}{{bold}}waiting for all nodes to report healthy...{{/}}\n")
 
-	hc := lc.nw.Healthy(ctx)
-	select {
-	case <-lc.stopc:
-		return errAborted
-	case <-ctx.Done():
-		return ctx.Err()
-	case err := <-hc:
-		if err != nil {
-			return err
-		}
+	if err := lc.nw.Healthy(ctx); err != nil {
+		return err
 	}
 
 	if err := lc.updateNodeInfo(); err != nil {