From 66393535682c906b780b47d13e0eaa32b252dab9 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Tue, 12 Apr 2016 14:40:22 +0100
Subject: [PATCH] Lock round TestRouter map accesses, and copy the set where we
 may block.

---
 testing/gossip/mocks.go | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/testing/gossip/mocks.go b/testing/gossip/mocks.go
index 124afe30..8cafd89d 100644
--- a/testing/gossip/mocks.go
+++ b/testing/gossip/mocks.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"math"
 	"math/rand"
+	"sync"
 	"time"
 
 	"github.com/weaveworks/mesh"
@@ -32,22 +33,34 @@ type flushMessage struct {
 }
 
 type TestRouter struct {
+	sync.Mutex
 	gossipChans map[mesh.PeerName]chan interface{}
 	loss        float32 // 0.0 means no loss
 }
 
 func NewTestRouter(loss float32) *TestRouter {
-	return &TestRouter{make(map[mesh.PeerName]chan interface{}, 100), loss}
+	return &TestRouter{gossipChans: make(map[mesh.PeerName]chan interface{}, 100), loss: loss}
+}
+
+// Copy so we can access outside of a lock
+func (grouter *TestRouter) copyGossipChans() map[mesh.PeerName]chan interface{} {
+	ret := make(map[mesh.PeerName]chan interface{})
+	grouter.Lock()
+	defer grouter.Unlock()
+	for p, c := range grouter.gossipChans {
+		ret[p] = c
+	}
+	return ret
 }
 
 func (grouter *TestRouter) Stop() {
-	for peer := range grouter.gossipChans {
+	for peer := range grouter.copyGossipChans() {
 		grouter.RemovePeer(peer)
 	}
 }
 
 func (grouter *TestRouter) gossipBroadcast(sender mesh.PeerName, update mesh.GossipData) {
-	for _, gossipChan := range grouter.gossipChans {
+	for _, gossipChan := range grouter.copyGossipChans() {
 		select {
 		case gossipChan <- broadcastMessage{sender: sender, data: update}:
 		default: // drop the message if we cannot send it
@@ -57,8 +70,9 @@ func (grouter *TestRouter) gossipBroadcast(sender mesh.PeerName, update mesh.Gos
 }
 
 func (grouter *TestRouter) gossip(sender mesh.PeerName, update mesh.GossipData) error {
-	count := int(math.Log2(float64(len(grouter.gossipChans))))
-	for dest, gossipChan := range grouter.gossipChans {
+	gossipChans := grouter.copyGossipChans()
+	count := int(math.Log2(float64(len(gossipChans))))
+	for dest, gossipChan := range gossipChans {
 		if dest == sender {
 			continue
 		}
@@ -76,7 +90,7 @@ func (grouter *TestRouter) gossip(sender mesh.PeerName, update mesh.GossipData)
 }
 
 func (grouter *TestRouter) Flush() {
-	for _, gossipChan := range grouter.gossipChans {
+	for _, gossipChan := range grouter.copyGossipChans() {
 		flushChan := make(chan struct{})
 		gossipChan <- flushMessage{flushChan: flushChan}
 		<-flushChan
@@ -84,11 +98,15 @@ func (grouter *TestRouter) Flush() {
 }
 
 func (grouter *TestRouter) RemovePeer(peer mesh.PeerName) {
+	grouter.Lock()
 	gossipChan := grouter.gossipChans[peer]
+	grouter.Unlock()
 	resultChan := make(chan struct{})
 	gossipChan <- exitMessage{exitChan: resultChan}
 	<-resultChan
+	grouter.Lock()
 	delete(grouter.gossipChans, peer)
+	grouter.Unlock()
 }
 
 type TestRouterClient struct {
@@ -160,13 +178,18 @@ func (grouter *TestRouter) Connect(sender mesh.PeerName, gossiper mesh.Gossiper)
 
 	go grouter.run(sender, gossiper, gossipChan)
 
+	grouter.Lock()
 	grouter.gossipChans[sender] = gossipChan
+	grouter.Unlock()
 	return TestRouterClient{grouter, sender}
 }
 
 func (client TestRouterClient) GossipUnicast(dstPeerName mesh.PeerName, buf []byte) error {
+	client.router.Lock()
+	gossipChan := client.router.gossipChans[dstPeerName]
+	client.router.Unlock()
 	select {
-	case client.router.gossipChans[dstPeerName] <- unicastMessage{sender: client.sender, buf: buf}:
+	case gossipChan <- unicastMessage{sender: client.sender, buf: buf}:
 	default: // drop the message if we cannot send it
 		common.Log.Errorf("Dropping message")
 	}