diff --git a/swarm/network/stream/streamer_test.go b/swarm/network/stream/streamer_test.go index 5548c50bdc..01cbb426f3 100644 --- a/swarm/network/stream/streamer_test.go +++ b/swarm/network/stream/streamer_test.go @@ -23,15 +23,21 @@ import ( "fmt" "os" "strconv" + "strings" "sync" "testing" "time" "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/log" + "github.com/ethereum/go-ethereum/node" "github.com/ethereum/go-ethereum/p2p/enode" + "github.com/ethereum/go-ethereum/p2p/simulations/adapters" p2ptest "github.com/ethereum/go-ethereum/p2p/testing" "github.com/ethereum/go-ethereum/swarm/network" + "github.com/ethereum/go-ethereum/swarm/network/simulation" + "github.com/ethereum/go-ethereum/swarm/state" + "github.com/ethereum/go-ethereum/swarm/storage" "golang.org/x/crypto/sha3" ) @@ -1120,20 +1126,32 @@ func TestRequestPeerSubscriptions(t *testing.T) { TestGetSubscriptionsRPC sets up a simulation network of 16 nodes, starts the simulation, waits for SyncUpdateDelay in order to kick off stream registration, then tests that there are subscriptions. -If provided with the `-printstats = true` option, it will print -the information of who is subscribed to who to STDOUT */ func TestGetSubscriptionsRPC(t *testing.T) { - //arbitrarily set to 16 + // arbitrarily set to 16 nodeCount := 16 - //set the syncUpdateDelay for sync registrations to start + // set the syncUpdateDelay for sync registrations to start syncUpdateDelay := 500 * time.Millisecond - //we will later need the kad table for each node - bucketKeyKad := simulation.BucketKey("kademlia") - //holds the msg code for SubscribeMsg + // holds the msg code for SubscribeMsg var subscribeMsgCode uint64 var ok bool - //create a standard sim + var expectedMsgCount = 0 + + // this channel signalizes that the expected amount of subscriptiosn is done + allSubscriptionsDone := make(chan struct{}) + lock := sync.RWMutex{} + // after the test, we need to reset the subscriptionFunc to the default + defer func() { subscriptionFunc = doRequestSubscription }() + + // we use this subscriptionFunc for this test: just increases count and calls the actual subscription + subscriptionFunc = func(r *Registry, p *network.Peer, bin uint8, subs map[enode.ID]map[Stream]struct{}) bool { + lock.Lock() + expectedMsgCount++ + lock.Unlock() + doRequestSubscription(r, p, bin, subs) + return true + } + // create a standard sim sim := simulation.New(map[string]simulation.ServiceFunc{ "streamer": func(ctx *adapters.ServiceContext, bucket *sync.Map) (s node.Service, cleanup func(), err error) { n := ctx.Config.Node() @@ -1148,23 +1166,20 @@ func TestGetSubscriptionsRPC(t *testing.T) { return nil, nil, err } kad := network.NewKademlia(addr.Over(), network.NewKadParams()) - //store the kad table - bucket.Store(bucketKeyKad, kad) delivery := NewDelivery(kad, netStore) netStore.NewNetFetcherFunc = network.NewFetcherFactory(dummyRequestFromPeers, true).New - //configure so that sync registrations actually happen + // configure so that sync registrations actually happen r := NewRegistry(addr.ID(), delivery, netStore, state.NewInmemoryStore(), &RegistryOptions{ Retrieval: RetrievalEnabled, Syncing: SyncingAutoSubscribe, //enable sync registrations SyncUpdateDelay: syncUpdateDelay, }, nil) - //get the SubscribeMsg code + // get the SubscribeMsg code subscribeMsgCode, ok = r.GetSpec().GetCode(SubscribeMsg{}) if !ok { t.Fatal("Message code for SubscribeMsg not found") } - bucket.Store(bucketKeyRegistry, r) cleanup = func() { os.RemoveAll(datadir) netStore.Close() @@ -1180,64 +1195,52 @@ func TestGetSubscriptionsRPC(t *testing.T) { ctx, cancelSimRun := context.WithTimeout(context.Background(), 1*time.Minute) defer cancelSimRun() - //upload a snapshot + // upload a snapshot err := sim.UploadSnapshot(fmt.Sprintf("testing/snapshot_%d.json", nodeCount)) if err != nil { t.Fatal(err) } - //run the simulation - result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error { - log.Info("Simulation running") - nodes := sim.Net.Nodes - //setup the filter for SubscribeMsg - msgs := sim.PeerEvents( - context.Background(), - sim.NodeIDs(), - simulation.NewPeerEventsFilter().ReceivedMessages().Protocol("stream").MsgCode(subscribeMsgCode), - ) + // setup the filter for SubscribeMsg + msgs := sim.PeerEvents( + context.Background(), + sim.NodeIDs(), + simulation.NewPeerEventsFilter().ReceivedMessages().Protocol("stream").MsgCode(subscribeMsgCode), + ) - //setup the vars we need - msgCount := 0 - expectedMsgCount := 0 - allSubscriptionsDone := make(chan struct{}) + // strategy: listen to all SubscribeMsg events; after every event we wait + // if after 1 second no more messages are being received, we assume the + // subscription phase has terminated! - //in a loop, catch all SubscribeMsg and just add up - go func() { - for m := range msgs { + // the loop in this go routine will either wait for new message events + // or times out after 1 second, which signals that we are not receiving + // any new subscriptions any more + go func() { + for { + select { + case <-ctx.Done(): + return + case m := <-msgs: // just reset the loop if m.Error != nil { log.Error("stream message", "err", m.Error) continue } log.Trace("stream message", "node", m.NodeID, "peer", m.PeerID) - //add one - msgCount += 1 - if msgCount == expectedMsgCount { - //the expected amount is reached - allSubscriptionsDone <- struct{}{} - return - } - } - }() + case <-time.After(time.Second): + // one second passed, don't assume more subscriptions + allSubscriptionsDone <- struct{}{} + log.Info("All subscriptions received") + return - //first iterate all nodes to get the expected number of subscriptions from the kad table - for _, node := range nodes { - item, ok := sim.NodeItem(node.ID(), bucketKeyKad) - if !ok { - return fmt.Errorf("No kademlia") } - kad := item.(*network.Kademlia) - //define the function which should run for each connection - just count subscriptions - //this is not actually subscribing but iterating the same way as the subscriptions do, - //as we need just the number - eachBinFunc := func(p *network.Peer, bin int) bool { - expectedMsgCount += 1 - return true - } - //call the actual kademlia for the count - kad.EachBin(kad.BaseAddr(), pot.DefaultPof(kad.MaxProxDisplay), 0, eachBinFunc) } - log.Debug("Expected message count: ", "expectedMsgCount", expectedMsgCount) + }() + + //run the simulation + result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error { + log.Info("Simulation running") + nodes := sim.Net.Nodes + //wait until all subscriptions are done select { case <-allSubscriptionsDone: @@ -1245,8 +1248,9 @@ func TestGetSubscriptionsRPC(t *testing.T) { t.Fatal("Context timed out") } - log.Info("All subscriptions received") + log.Debug("Expected message count: ", "expectedMsgCount", expectedMsgCount) //now iterate again, this time we call each node via RPC to get its subscriptions + realCount := 0 for _, node := range nodes { //create rpc client client, err := node.Client() @@ -1254,11 +1258,6 @@ func TestGetSubscriptionsRPC(t *testing.T) { t.Fatalf("create node 1 rpc client fail: %v", err) } - item, ok := sim.NodeItem(node.ID(), bucketKeyRegistry) - if !ok { - return fmt.Errorf("No registry") - } - registry := item.(*Registry) //ask it for subscriptions pstreams := make(map[string][]string) err = client.Call(&pstreams, "stream_getPeerSubscriptions") @@ -1266,17 +1265,23 @@ func TestGetSubscriptionsRPC(t *testing.T) { t.Fatal(err) } //length of the subscriptions can not be smaller than number of peers - if len(pstreams) < len(registry.peers) { - t.Fatal("Subscription count is smaller than expected") - } log.Debug(fmt.Sprintf("node %s subscriptions:", node.String())) for p, ps := range pstreams { log.Debug(fmt.Sprintf("...with node %s: ", p)) for _, s := range ps { log.Debug(fmt.Sprintf("......%s", s)) + // each node also has subscriptions to RETRIEVE_REQUEST streams, + // we need to ignore those, we are only counting SYNC streams + if !strings.HasPrefix(s, "RETRIEVE_REQUEST") { + realCount++ + } } } } + // every node is mutually subscribed to each other, so the actual count is half of it + if realCount/2 != expectedMsgCount { + return errors.New(fmt.Sprintf("Real subscriptions and expected amount don't match; real: %d, expected: %d", realCount/2, expectedMsgCount)) + } return nil }) if result.Error != nil {