From e1a3dd53bf09f0b6d702e4bc82f83141a8ca641d Mon Sep 17 00:00:00 2001
From: lash <dev@holbrook.no>
Date: Wed, 20 Mar 2019 15:40:34 +0100
Subject: [PATCH] swarm/network, swarm/pss: Simplify + correct lock in
 servicefunc sim

---
 swarm/network/simulation/simulation.go |  4 ++--
 swarm/pss/prox_test.go                 | 22 +++++++++++++++-------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/swarm/network/simulation/simulation.go b/swarm/network/simulation/simulation.go
index 722a01bdc5..5787cafdae 100644
--- a/swarm/network/simulation/simulation.go
+++ b/swarm/network/simulation/simulation.go
@@ -85,6 +85,8 @@ func New(services map[string]ServiceFunc) (s *Simulation) {
 		name, serviceFunc := name, serviceFunc
 		s.serviceNames = append(s.serviceNames, name)
 		adapterServices[name] = func(ctx *adapters.ServiceContext) (node.Service, error) {
+			s.mu.Lock()
+			defer s.mu.Unlock()
 			b, ok := s.buckets[ctx.Config.ID]
 			if !ok {
 				b = new(sync.Map)
@@ -93,8 +95,6 @@ func New(services map[string]ServiceFunc) (s *Simulation) {
 			if err != nil {
 				return nil, err
 			}
-			s.mu.Lock()
-			defer s.mu.Unlock()
 			if cleanup != nil {
 				s.cleanupFuncs = append(s.cleanupFuncs, cleanup)
 			}
diff --git a/swarm/pss/prox_test.go b/swarm/pss/prox_test.go
index 698e6553f1..ff53003950 100644
--- a/swarm/pss/prox_test.go
+++ b/swarm/pss/prox_test.go
@@ -145,17 +145,25 @@ func newTestData() *testData {
 	}
 }
 
+func (d *testData) getKademlia(nodeId *enode.ID) (*network.Kademlia, error) {
+	kadif, ok := d.sim.NodeItem(*nodeId, simulation.BucketKeyKademlia)
+	if !ok {
+		return nil, fmt.Errorf("no kademlia entry for %v", nodeId)
+	}
+	kad, ok := kadif.(*network.Kademlia)
+	if !ok {
+		return nil, fmt.Errorf("invalid kademlia entry for %v", nodeId)
+	}
+	return kad, nil
+}
+
 func (d *testData) init(msgCount int) error {
 	log.Debug("TestProxNetwork start")
 
 	for _, nodeId := range d.sim.NodeIDs() {
-		kadif, ok := d.sim.NodeItem(nodeId, simulation.BucketKeyKademlia)
-		if !ok {
-			return fmt.Errorf("no kademlia entry for %v", nodeId)
-		}
-		kad, ok := kadif.(*network.Kademlia)
-		if !ok {
-			return fmt.Errorf("invalid kademlia entry for %v", nodeId)
+		kad, err := d.getKademlia(&nodeId)
+		if err != nil {
+			return err
 		}
 		d.nodeAddrs[nodeId] = kad.BaseAddr()
 	}