swarm/pss: WaitTillSnapshotRecreated() func added

swarm/network/simulation/kademlia.go

@@ -18,12 +18,14 @@ package simulation
 
 import (
 	"context"
+	"encoding/binary"
 	"encoding/hex"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/simulations"
 	"github.com/ethereum/go-ethereum/swarm/network"
 )
 
@@ -96,3 +98,70 @@ func (s *Simulation) kademlias() (ks map[enode.ID]*network.Kademlia) {
 	}
 	return ks
 }
+
+func (s *Simulation) WaitTillSnapshotRecreated(ctx context.Context, snap simulations.Snapshot) error {
+	expected := listSnapshotConnections(snap.Conns)
+	ticker := time.NewTicker(256 * time.Millisecond) // todo: reduce
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
+			actual := listActualConnections(s.kademlias())
+			if isAllDeployed(expected, actual) {
+				return nil
+			}
+		}
+	}
+}
+
+func listActualConnections(kademlias map[enode.ID]*network.Kademlia) (res []uint64) {
+	for base, k := range kademlias {
+		k.EachConn(base[:], 256, func(p *network.Peer, _ int) bool {
+			res = append(res, getConnectionHash(base, p.ID()))
+			return true
+		})
+	}
+	return res
+}
+
+func listSnapshotConnections(conns []simulations.Conn) (res []uint64) {
+	for _, c := range conns {
+		res = append(res, getConnectionHash(c.One, c.Other))
+	}
+	return res
+}
+
+// returns an integer connection identifier (similar to 8-byte hash)
+func getConnectionHash(a, b enode.ID) uint64 {
+	var h [8]byte
+	for i := 0; i < 8; i++ {
+		h[i] = a[i] ^ b[i]
+	}
+	res := binary.LittleEndian.Uint64(h[:])
+	return res
+}
+
+// returns true if all connections in expected are listed in actual
+func isAllDeployed(expected []uint64, actual []uint64) bool {
+	exp := make([]uint64, len(expected))
+	copy(exp, expected)
+	if len(exp) > 0 {
+		for _, c := range actual {
+			// remove value c from exp
+			for i := 0; i < len(exp); i++ {
+				if exp[i] == c {
+					last := len(exp) - 1
+					if last == 0 {
+						return true
+					}
+					exp[i] = exp[last]
+					exp = exp[:last]
+				}
+			}
+		}
+	}
+	return len(exp) == 0
+}

swarm/pss/snapshot_test.go

@@ -30,29 +30,6 @@ import (
 	"github.com/ethereum/go-ethereum/swarm/state"
 )
 
-var (
-	runNodes    = flag.Int("nodes", 0, "nodes to start in the network")
-	runMessages = flag.Int("messages", 0, "messages to send during test")
-
-	topic         = BytesToTopic([]byte{0x00, 0x00, 0x06, 0x82})
-	mu            sync.Mutex // keeps handlerDonc in sync
-	kademlias     = make(map[enode.ID]*network.Kademlia)
-	nodeAddrs     = make(map[enode.ID][]byte)   // make predictable overlay addresses from the generated random enode ids
-	msgsToReceive int                           // total count of messages to receive, used for terminating the simulation run
-	recipients    = make(map[int][]enode.ID)    // for logging output only
-	msgs          [][]byte                      // recipient addresses of messages
-	expectedMsgs  = make(map[enode.ID][]uint64) // message serials we expect respective nodes to receive
-	senders       = make(map[int]enode.ID)      // originating nodes of the messages (intention is to choose as far as possible from the receiving neighborhood)
-	pof           = pot.DefaultPof(256)         // generate messages and index them
-	sim           *simulation.Simulation
-	handlerDone   bool                             // set to true on termination of the simulation run
-	handlerC      = make(chan handlerNotification) // passes message from pss message handler to simulation driver
-	doneC         = make(chan struct{})            // terminates the handler channel listener
-	errC          = make(chan error)               // error to pass to main sim thread
-	msgC          = make(chan handlerNotification) // message receipt notification to main sim thread
-	debugCnt      int
-)
-
 // needed to make the enode id of the receiving node available to the handler for triggers
 type handlerContextFunc func(*adapters.NodeConfig) *handler
 
@@ -65,6 +42,48 @@ type handlerNotification struct {
 	serial uint64
 }
 
+var (
+	runNodes    = flag.Int("nodes", 0, "nodes to start in the network")
+	runMessages = flag.Int("messages", 0, "messages to send during test")
+
+	pof           = pot.DefaultPof(256) // generate messages and index them
+	topic         = BytesToTopic([]byte{0x00, 0x00, 0x06, 0x82})
+	mu            sync.Mutex // keeps handlerDonc in sync
+	sim           *simulation.Simulation
+
+	handlerDone   bool // set to true on termination of the simulation run
+	msgsToReceive int  // total count of messages to receive, used for terminating the simulation run
+	debugCnt      int
+
+	kademlias     map[enode.ID]*network.Kademlia
+	nodeAddrs     map[enode.ID][]byte   // make predictable overlay addresses from the generated random enode ids
+	recipients    map[int][]enode.ID    // for logging output only
+	expectedMsgs  map[enode.ID][]uint64 // message serials we expect respective nodes to receive
+	senders       map[int]enode.ID      // originating nodes of the messages (intention is to choose as far as possible from the receiving neighborhood)
+	handlerC      chan handlerNotification // passes message from pss message handler to simulation driver
+	doneC         chan struct{}            // terminates the handler channel listener
+	errC          chan error               // error to pass to main sim thread
+	msgC          chan handlerNotification // message receipt notification to main sim thread
+	msgs          [][]byte                 // recipient addresses of messages
+)
+
+func resetTestVariables() {
+	handlerDone = false
+	msgsToReceive = 0
+	debugCnt = 0
+	msgs = nil
+
+	kademlias     = make(map[enode.ID]*network.Kademlia)
+	nodeAddrs     = make(map[enode.ID][]byte)
+	recipients    = make(map[int][]enode.ID)
+	expectedMsgs  = make(map[enode.ID][]uint64)
+	senders       = make(map[int]enode.ID)
+	handlerC      = make(chan handlerNotification)
+	doneC         = make(chan struct{})
+	errC          = make(chan error)
+	msgC          = make(chan handlerNotification)
+}
+
 func init() {
 	log.Root().SetHandler(log.LvlFilterHandler(log.LvlTrace, log.StreamHandler(os.Stderr, log.TerminalFormat(true))))
 }
@@ -111,7 +130,7 @@ func readSnapshot(t *testing.T, nodeCount int) simulations.Snapshot {
 	return snap
 }
 
-func initTestVariables(sim *simulation.Simulation, msgCount int) {
+func assingTestVariables(sim *simulation.Simulation, msgCount int) {
 	log.Debug("-------------------------------------------------------------------------")
 	var targets string
 	for _, nodeId := range sim.NodeIDs() {
@@ -169,6 +188,7 @@ func TestProxNetwork(t *testing.T) {
 // Upon sending the messages, it verifies that the respective message is passed to the message handlers of these recipients.
 // It will fail if a recipient handles a message it should not, or if after propagation not all expected messages are handled (timeout)
 func testProxNetwork(t *testing.T) {
+	resetTestVariables()
 	msgCount, nodeCount := getCmdParams(t)
 	handlerContextFuncs := make(map[Topic]handlerContextFunc)
 	handlerContextFuncs[topic] = nodeMsgHandler
@@ -180,10 +200,14 @@ func testProxNetwork(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*4)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
 	defer cancel()
-	waitTillSerenity(t, snap, sim, 1000)
+	err = sim.WaitTillSnapshotRecreated(ctx, snap)
-	initTestVariables(sim, msgCount)
+	if err != nil {
+		t.Fatalf("failed to recreate snapshot: %s", err)
+	}
+	//waitTillSerenity(t, snap, sim, 1000)
+	assingTestVariables(sim, msgCount)
 	result := sim.Run(ctx, runFunc)
 	if result.Error != nil {
 		log.Debug("--------------------------------------------------------------------------------", "rcv", debugCnt)
@@ -192,6 +216,7 @@ func testProxNetwork(t *testing.T) {
 	t.Logf("completed %d", result.Duration)
 }
 
+/*
 func waitTillSerenity(t *testing.T, snap simulations.Snapshot, sim *simulation.Simulation, timeout int) {
 	interval := 16
 	expected := listSnapConnections(snap.Conns)
@@ -203,7 +228,7 @@ func waitTillSerenity(t *testing.T, snap simulations.Snapshot, sim *simulation.S
 			time.Sleep(time.Millisecond * time.Duration(interval))
 		}
 	}
-	time.Sleep(time.Millisecond * 10) // todo: remove this later
+	time.Sleep(time.Millisecond * 16) // todo: remove this later
 }
 
 func listSnapConnections(conns []simulations.Conn) (res []uint64) {
@@ -251,6 +276,7 @@ func isSerenity(expected []uint64, actual []uint64) bool {
 	}
 	return len(exp) == 0
 }
+*/
 
 func sendAllMsgs(sim *simulation.Simulation, msgs [][]byte, senders map[int]enode.ID) {
 	for i, msg := range msgs {