utils.go raw

   1  package rpctest
   2  
   3  import (
   4  	"reflect"
   5  	"time"
   6  	
   7  	"github.com/p9c/p9/pkg/chainhash"
   8  	"github.com/p9c/p9/pkg/rpcclient"
   9  )
  10  
  11  // JoinType is an enum representing a particular type of "node join". A node
  12  // join is a synchronization tool used to wait until a subset of nodes have a
  13  // consistent state with respect to an attribute.
  14  type JoinType uint8
  15  
  16  const (
  17  	// Blocks is a JoinType which waits until all nodes share the same block
  18  	// height.
  19  	Blocks JoinType = iota
  20  	// Mempools is a JoinType which blocks until all nodes have identical mempool.
  21  	Mempools
  22  )
  23  
  24  // JoinNodes is a synchronization tool used to block until all passed nodes are
  25  // fully synced with respect to an attribute. This function will block for a
  26  // period of time, finally returning once all nodes are synced according to the
  27  // passed JoinType. This function be used to to ensure all active test harnesses
  28  // are at a consistent state before proceeding to an assertion or check within
  29  // rpc tests.
  30  func JoinNodes(nodes []*Harness, joinType JoinType) (e error) {
  31  	switch joinType {
  32  	case Blocks:
  33  		return syncBlocks(nodes)
  34  	case Mempools:
  35  		return syncMempools(nodes)
  36  	}
  37  	return nil
  38  }
  39  
  40  // syncMempools blocks until all nodes have identical mempools.
  41  func syncMempools(nodes []*Harness) (e error) {
  42  	poolsMatch := false
  43  retry:
  44  	for !poolsMatch {
  45  		firstPool, e := nodes[0].Node.GetRawMempool()
  46  		if e != nil {
  47  			return e
  48  		}
  49  		// If all nodes have an identical mempool with respect to the first node,
  50  		// then we're done. Otherwise drop back to the top of the loop and retry
  51  		// after a short wait period.
  52  		for _, node := range nodes[1:] {
  53  			nodePool, e := node.Node.GetRawMempool()
  54  			if e != nil {
  55  				return e
  56  			}
  57  			if !reflect.DeepEqual(firstPool, nodePool) {
  58  				time.Sleep(time.Millisecond * 100)
  59  				continue retry
  60  			}
  61  		}
  62  		poolsMatch = true
  63  	}
  64  	return nil
  65  }
  66  
  67  // syncBlocks blocks until all nodes report the same best chain.
  68  func syncBlocks(nodes []*Harness) (e error) {
  69  	blocksMatch := false
  70  retry:
  71  	for !blocksMatch {
  72  		var prevHash *chainhash.Hash
  73  		var prevHeight int32
  74  		
  75  		for _, node := range nodes {
  76  			blockHash, blockHeight, e := node.Node.GetBestBlock()
  77  			if e != nil {
  78  				return e
  79  			}
  80  			if prevHash != nil && (*blockHash != *prevHash ||
  81  				blockHeight != prevHeight) {
  82  				time.Sleep(time.Millisecond * 100)
  83  				continue retry
  84  			}
  85  			prevHash, prevHeight = blockHash, blockHeight
  86  		}
  87  		blocksMatch = true
  88  	}
  89  	return nil
  90  }
  91  
  92  // ConnectNode establishes a new peer-to-peer connection between the "from" harness and the "to" harness. The connection
  93  // made is flagged as persistent therefore in the case of disconnects, "from" will attempt to reestablish a connection
  94  // to the "to" harness.
  95  func ConnectNode(from *Harness, to *Harness) (e error) {
  96  	peerInfo, e := from.Node.GetPeerInfo()
  97  	if e != nil {
  98  		return e
  99  	}
 100  	numPeers := len(peerInfo)
 101  	targetAddr := to.node.config.listen
 102  	if e = from.Node.AddNode(targetAddr, rpcclient.ANAdd); E.Chk(e) {
 103  		return e
 104  	}
 105  	// Block until a new connection has been established.
 106  	peerInfo, e = from.Node.GetPeerInfo()
 107  	if e != nil {
 108  		return e
 109  	}
 110  	for len(peerInfo) <= numPeers {
 111  		peerInfo, e = from.Node.GetPeerInfo()
 112  		if e != nil {
 113  			return e
 114  		}
 115  	}
 116  	return nil
 117  }
 118  
 119  // TearDownAll tears down all active test harnesses.
 120  func TearDownAll() (e error) {
 121  	harnessStateMtx.Lock()
 122  	defer harnessStateMtx.Unlock()
 123  	for _, harness := range testInstances {
 124  		if e := harness.tearDown(); E.Chk(e) {
 125  			return e
 126  		}
 127  	}
 128  	return nil
 129  }
 130  
 131  // ActiveHarnesses returns a slice of all currently active test harnesses. A test harness if considered "active" if it
 132  // has been created, but not yet torn down.
 133  func ActiveHarnesses() []*Harness {
 134  	harnessStateMtx.RLock()
 135  	defer harnessStateMtx.RUnlock()
 136  	activeNodes := make([]*Harness, 0, len(testInstances))
 137  	for _, harness := range testInstances {
 138  		activeNodes = append(activeNodes, harness)
 139  	}
 140  	return activeNodes
 141  }
 142