From cb65238a39dd1ed4f4abf49ee5dd1d2fae568a35 Mon Sep 17 00:00:00 2001
From: Yusef Napora <yusef@protocol.ai>
Date: Tue, 5 May 2020 14:13:26 -0400
Subject: [PATCH] fix race condition in rpc fragmentation test

---
 gossipsub_test.go | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/gossipsub_test.go b/gossipsub_test.go
index 3bbc0ad..daed76b 100644
--- a/gossipsub_test.go
+++ b/gossipsub_test.go
@@ -1752,6 +1752,8 @@ func TestGossipsubRPCFragmentation(t *testing.T) {
 
 	// wait a bit for them to be received via gossip by the fake peer
 	time.Sleep(5 * time.Second)
+	iwe.lk.Lock()
+	defer iwe.lk.Unlock()
 
 	// we should have received all the messages
 	if iwe.msgsReceived != nMessages {
@@ -1778,6 +1780,7 @@ func TestGossipsubRPCFragmentation(t *testing.T) {
 // test that large responses to IWANT requests are fragmented into multiple RPCs.
 type iwantEverything struct {
 	h                host.Host
+	lk               sync.Mutex
 	rpcsWithMessages int
 	msgsReceived     int
 	ihavesReceived   int
@@ -1816,6 +1819,7 @@ func (iwe *iwantEverything) handleStream(s network.Stream) {
 			return
 		}
 
+		iwe.lk.Lock()
 		if len(rpc.Publish) != 0 {
 			iwe.rpcsWithMessages++
 		}
@@ -1852,5 +1856,6 @@ func (iwe *iwantEverything) handleStream(s network.Stream) {
 				panic(err)
 			}
 		}
+		iwe.lk.Unlock()
 	}
 }