swarm/network/stream: test improvements

swarm/network/stream/delivery.go

@@ -96,7 +96,12 @@ func (s *SwarmChunkServer) processDeliveries() {
 
 // SetNextBatch
 func (s *SwarmChunkServer) SetNextBatch(_, _ uint64) (hashes []byte, from uint64, to uint64, proof *HandoverProof, err error) {
-	hashes = <-s.batchC
+	select {
+	case hashes = <-s.batchC:
+	case <-s.quit:
+		return
+	}
+
 	from = s.currentLen
 	s.currentLen += uint64(len(hashes))
 	to = s.currentLen

swarm/network/stream/delivery_test.go

@@ -381,7 +381,7 @@ func testDeliveryFromNodes(t *testing.T, nodes, conns, chunkCount int, skipCheck
 		for j := 0; j < nodes-1; j++ {
 			id := sim.IDs[j]
 			err := sim.CallClient(id, func(client *rpc.Client) error {
-				err := streamTesting.WatchDisconnections(id, client, peerCount(id), errc, quitC)
+				err := streamTesting.WatchDisconnections(id, client, errc, quitC)
 				if err != nil {
 					return err
 				}
@@ -489,6 +489,7 @@ func BenchmarkDeliveryFromNodesWithCheck(b *testing.B) {
 func benchmarkDeliveryFromNodes(b *testing.B, nodes, conns, chunkCount int, skipCheck bool) {
 	defaultSkipCheck = skipCheck
 	toAddr = network.NewAddrFromNodeID
+
 	timeout := 300 * time.Second
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
@@ -527,9 +528,10 @@ func benchmarkDeliveryFromNodes(b *testing.B, nodes, conns, chunkCount int, skip
 
 	// channel to signal simulation initialisation with action call complete
 	// or node disconnections
-	simErrC := make(chan error)
+	disconnectC := make(chan error)
 	quitC := make(chan struct{})
-	defer close(quitC)
+
+	initC := make(chan error)
 
 	action := func(ctx context.Context) error {
 		// each node Subscribes to each other's swarmChunkServerStreamName
@@ -546,13 +548,13 @@ func benchmarkDeliveryFromNodes(b *testing.B, nodes, conns, chunkCount int, skip
 				break
 			}
 		}
-
+		var err error
 		// each node except the last one subscribes to the upstream swarm chunk server stream
 		// which responds to chunk retrieve requests
 		for j := 0; j < nodes-1; j++ {
 			id := sim.IDs[j]
-			simErrC <- sim.CallClient(id, func(client *rpc.Client) error {
+			err = sim.CallClient(id, func(client *rpc.Client) error {
-				err := streamTesting.WatchDisconnections(id, client, peerCount(id), simErrC, quitC)
+				err := streamTesting.WatchDisconnections(id, client, disconnectC, quitC)
 				if err != nil {
 					return err
 				}
@@ -561,23 +563,17 @@ func benchmarkDeliveryFromNodes(b *testing.B, nodes, conns, chunkCount int, skip
 				sid := sim.IDs[j+1] // the upstream peer's id
 				return client.CallContext(ctx, nil, "stream_subscribeStream", sid, swarmChunkServerStreamName, nil, 0, 0, Top, false)
 			})
+			if err != nil {
+				break
 			}
-		// signal to the benchmark that setup is complete
+		}
-		return err
+		initC <- err
+		return nil
 	}
 
 	// the check function is only triggered when the benchmark finishes
-	checkC := make(chan error)
 	trigger := make(chan discover.NodeID)
 	check := func(ctx context.Context, id discover.NodeID) (_ bool, err error) {
-		select {
-		case <-ctx.Done():
-			err = ctx.Err()
-		case err = <-checkC:
-		}
-		if err != nil {
-			return false, err
-		}
 		return true, nil
 	}
 
@@ -595,26 +591,15 @@ func benchmarkDeliveryFromNodes(b *testing.B, nodes, conns, chunkCount int, skip
 	errc := make(chan error)
 	go func() {
 		_, err := sim.Run(ctx, conf)
+		close(quitC)
 		errc <- err
 	}()
 
 	// wait for simulation action to complete stream subscriptions
-	err = <-simErrC
+	err = <-initC
 	if err != nil {
 		b.Fatalf("simulation failed to initialise. expected no error. got %v", err)
 	}
-	go func() {
-		for {
-			var err error
-			select {
-			case err = <-simErrC:
-			case <-quitC:
-				return
-			}
-			trigger <- sim.IDs[0]
-			checkC <- err
-		}
-	}()
 
 	// create a retriever dpa for the pivot node
 	// by now deliveries are set for each node by the streamer service
@@ -627,6 +612,7 @@ func benchmarkDeliveryFromNodes(b *testing.B, nodes, conns, chunkCount int, skip
 	// benchmark loop
 	b.ResetTimer()
 	b.StopTimer()
+Loop:
 	for i := 0; i < b.N; i++ {
 		// uploading chunkCount random chunks to the last node
 		hashes := make([]storage.Key, chunkCount)
@@ -666,12 +652,34 @@ func benchmarkDeliveryFromNodes(b *testing.B, nodes, conns, chunkCount int, skip
 			}
 		}
 		b.StopTimer()
+
+		select {
+		case err = <-disconnectC:
+			if err != nil {
+				break Loop
+			}
+		default:
+		}
+
 		if misses > 0 {
-			simErrC <- fmt.Errorf("%v chunk not found out of %v", misses, total)
+			err = fmt.Errorf("%v chunk not found out of %v", misses, total)
+			break Loop
 		}
 	}
-	// benchmark over, trigger the check function to conclude the simulation
+
+	select {
+	case <-quitC:
+	case trigger <- sim.IDs[0]:
+	}
+	if err == nil {
 		err = <-errc
+	} else {
+		if e := <-errc; e != nil {
+			b.Errorf("sim.Run function error: %v", e)
+		}
+	}
+
+	// benchmark over, trigger the check function to conclude the simulation
 	if err != nil {
 		b.Fatalf("expected no error. got %v", err)
 	}

swarm/network/stream/stream.go

@@ -198,7 +198,7 @@ func (r *Registry) run(p *protocols.Peer) error {
 	sp := NewPeer(p, r)
 	r.setPeer(sp)
 	defer r.deletePeer(sp)
-	// defer close(sp.quit
+	defer close(sp.quit)
 	defer sp.close()
 	return sp.Run(sp.HandleMsg)
 }

swarm/network/stream/syncer_test.go

@@ -159,7 +159,7 @@ func testSyncBetweenNodes(t *testing.T, nodes, conns, chunkCount int, skipCheck
 			id := sim.IDs[j]
 			err := sim.CallClient(id, func(client *rpc.Client) error {
 				// report disconnect events to the error channel cos peers should not disconnect
-				err := streamTesting.WatchDisconnections(id, client, peerCount(id), errc, quitC)
+				err := streamTesting.WatchDisconnections(id, client, errc, quitC)
 				if err != nil {
 					return err
 				}
@@ -224,6 +224,6 @@ func testSyncBetweenNodes(t *testing.T, nodes, conns, chunkCount int, skipCheck
 	}
 	if result.Error != nil {
 		t.Fatalf("Simulation failed: %s", result.Error)
+	}
 	streamTesting.CheckResult(t, result, startedAt, finishedAt)
 }
-}

swarm/network/stream/testing/testing.go

@@ -208,28 +208,25 @@ func (s *Simulation) Run(ctx context.Context, conf *RunConfig) (*simulations.Ste
 	return result, nil
 }
 
-func WatchDisconnections(id discover.NodeID, client *rpc.Client, expectedConnCount int, errc chan error, quitC chan struct{}) error {
+func WatchDisconnections(id discover.NodeID, client *rpc.Client, errc chan error, quitC chan struct{}) error {
 	events := make(chan *p2p.PeerEvent)
 	sub, err := client.Subscribe(context.Background(), "admin", events, "peerEvents")
 	if err != nil {
 		return fmt.Errorf("error getting peer events for node %v: %s", id, err)
 	}
 	go func() {
-		defer sub.Unsubscribe()
+		for {
 			select {
 			case <-quitC:
-			if expectedConnCount <= 0 {
 				return
-			}
 			case e := <-events:
-			expectedConnCount--
 				errc <- fmt.Errorf("peerEvent for node %v: %v", id, e)
 			case err := <-sub.Err():
-			expectedConnCount = 0
 				if err != nil {
 					errc <- fmt.Errorf("error getting peer events for node %v: %v", id, err)
 				}
 			}
+		}
 	}()
 	return nil
 }
@@ -237,6 +234,7 @@ func WatchDisconnections(id discover.NodeID, client *rpc.Client, expectedConnCou
 func Trigger(d time.Duration, quitC chan struct{}, ids ...discover.NodeID) chan discover.NodeID {
 	trigger := make(chan discover.NodeID)
 	go func() {
+		defer close(trigger)
 		ticker := time.NewTicker(d)
 		defer ticker.Stop()
 		// we are only testing the pivot node (net.Nodes[0])
@@ -245,6 +243,7 @@ func Trigger(d time.Duration, quitC chan struct{}, ids ...discover.NodeID) chan
 				select {
 				case trigger <- id:
 				case <-quitC:
+					return
 				}
 			}
 		}