sync2: fix race in the p2p test

Author: ivan4th

sync2/p2p_test.go

@@ -33,6 +33,7 @@ type fakeHandler struct {
 	mtx         *sync.Mutex
 	localPeerID p2p.Peer
 	synced      map[addedKey]struct{}
+	committed   map[string]struct{}
 }
 
 func (fh *fakeHandler) Receive(k rangesync.KeyBytes, peer p2p.Peer) (bool, error) {
@@ -50,11 +51,21 @@ func (fh *fakeHandler) Commit(peer p2p.Peer, base, new multipeer.OrderedSet) err
 	fh.mtx.Lock()
 	defer fh.mtx.Unlock()
 	for k := range fh.synced {
-		base.(*multipeer.DumbSet).AddUnchecked(rangesync.KeyBytes(k.key))
+		fh.committed[k.key] = struct{}{}
 	}
+	clear(fh.synced)
 	return nil
 }
 
+func (fh *fakeHandler) committedItems() (items []rangesync.KeyBytes) {
+	fh.mtx.Lock()
+	defer fh.mtx.Unlock()
+	for k := range fh.committed {
+		items = append(items, rangesync.KeyBytes(k))
+	}
+	return items
+}
+
 func TestP2P(t *testing.T) {
 	const (
 		numNodes  = 4
@@ -65,8 +76,8 @@ func TestP2P(t *testing.T) {
 	logger := zaptest.NewLogger(t)
 	mesh, err := mocknet.FullMeshConnected(numNodes)
 	require.NoError(t, err)
-	synced := make(map[addedKey]struct{})
 	hs := make([]*sync2.P2PHashSync, numNodes)
+	handlers := make([]*fakeHandler, numNodes)
 	initialSet := make([]rangesync.KeyBytes, numHashes)
 	for n := range initialSet {
 		initialSet[n] = rangesync.RandomKeyBytes(32)
@@ -85,10 +96,11 @@ func TestP2P(t *testing.T) {
 		cfg.EnableActiveSync = true
 		cfg.SyncInterval = 100 * time.Millisecond
 		host := mesh.Hosts()[n]
-		handler := &fakeHandler{
+		handlers[n] = &fakeHandler{
 			mtx:         &mtx,
 			localPeerID: host.ID(),
-			synced:      synced,
+			synced:      make(map[addedKey]struct{}),
+			committed:   make(map[string]struct{}),
 		}
 		os := multipeer.NewDumbHashSet()
 		d := rangesync.NewDispatcher(logger)
@@ -98,7 +110,7 @@ func TestP2P(t *testing.T) {
 		eg.Go(func() error { return srv.Run(ctx) })
 		hs[n] = sync2.NewP2PHashSync(
 			logger.Named(fmt.Sprintf("node%d", n)),
-			d, "test", os, keyLen, maxDepth, ps, handler, cfg, srv)
+			d, "test", os, keyLen, maxDepth, ps, handlers[n], cfg, srv)
 		require.NoError(t, hs[n].Load())
 		is := hs[n].Set().(*multipeer.DumbSet)
 		is.SetAllowMultiReceive(true)
@@ -112,12 +124,15 @@ func TestP2P(t *testing.T) {
 	}
 
 	require.Eventually(t, func() bool {
-		for _, hsync := range hs {
+		for n, hsync := range hs {
 			// use a snapshot to avoid races
 			if !hsync.Synced() {
 				return false
 			}
 			os := hsync.Set().Copy(false)
+			for _, k := range handlers[n].committedItems() {
+				os.(*multipeer.DumbSet).AddUnchecked(k)
+			}
 			empty, err := os.Empty()
 			require.NoError(t, err)
 			if empty {
@@ -134,9 +149,13 @@ func TestP2P(t *testing.T) {
 		return true
 	}, 30*time.Second, 300*time.Millisecond)
 
-	for _, hsync := range hs {
+	for n, hsync := range hs {
 		hsync.Stop()
-		actualItems, err := hsync.Set().Items().Collect()
+		os := hsync.Set().Copy(false)
+		for _, k := range handlers[n].committedItems() {
+			os.(*multipeer.DumbSet).AddUnchecked(k)
+		}
+		actualItems, err := os.Items().Collect()
 		require.NoError(t, err)
 		require.ElementsMatch(t, initialSet, actualItems)
 	}