From 9f5e9d6d9425d8464eaf7005c351a6ab9c903bcb Mon Sep 17 00:00:00 2001
From: Janos Guljas <janos@resenje.org>
Date: Thu, 21 Feb 2019 17:06:27 +0100
Subject: [PATCH] swarm/storage/localstore: protect chunk put with addr lock in
 tests

---
 .../localstore/subscription_pull_test.go      | 20 +++++++++++--------
 .../localstore/subscription_push_test.go      | 10 ++++++----
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/swarm/storage/localstore/subscription_pull_test.go b/swarm/storage/localstore/subscription_pull_test.go
index 958d71e50b..9800329eaa 100644
--- a/swarm/storage/localstore/subscription_pull_test.go
+++ b/swarm/storage/localstore/subscription_pull_test.go
@@ -151,6 +151,9 @@ func TestDB_SubscribePull_since(t *testing.T) {
 	})()
 
 	uploadRandomChunks := func(count int, wanted bool) (last map[uint8]ChunkDescriptor) {
+		addrsMu.Lock()
+		defer addrsMu.Unlock()
+
 		last = make(map[uint8]ChunkDescriptor)
 		for i := 0; i < count; i++ {
 			chunk := generateRandomChunk()
@@ -162,7 +165,6 @@ func TestDB_SubscribePull_since(t *testing.T) {
 
 			bin := db.po(chunk.Address())
 
-			addrsMu.Lock()
 			if _, ok := addrs[bin]; !ok {
 				addrs[bin] = make([]storage.Address, 0)
 			}
@@ -170,7 +172,6 @@ func TestDB_SubscribePull_since(t *testing.T) {
 				addrs[bin] = append(addrs[bin], chunk.Address())
 				wantedChunksCount++
 			}
-			addrsMu.Unlock()
 
 			lastTimestampMu.RLock()
 			storeTimestamp := lastTimestamp
@@ -241,6 +242,9 @@ func TestDB_SubscribePull_until(t *testing.T) {
 	})()
 
 	uploadRandomChunks := func(count int, wanted bool) (last map[uint8]ChunkDescriptor) {
+		addrsMu.Lock()
+		defer addrsMu.Unlock()
+
 		last = make(map[uint8]ChunkDescriptor)
 		for i := 0; i < count; i++ {
 			chunk := generateRandomChunk()
@@ -252,7 +256,6 @@ func TestDB_SubscribePull_until(t *testing.T) {
 
 			bin := db.po(chunk.Address())
 
-			addrsMu.Lock()
 			if _, ok := addrs[bin]; !ok {
 				addrs[bin] = make([]storage.Address, 0)
 			}
@@ -260,7 +263,6 @@ func TestDB_SubscribePull_until(t *testing.T) {
 				addrs[bin] = append(addrs[bin], chunk.Address())
 				wantedChunksCount++
 			}
-			addrsMu.Unlock()
 
 			lastTimestampMu.RLock()
 			storeTimestamp := lastTimestamp
@@ -330,6 +332,9 @@ func TestDB_SubscribePull_sinceAndUntil(t *testing.T) {
 	})()
 
 	uploadRandomChunks := func(count int, wanted bool) (last map[uint8]ChunkDescriptor) {
+		addrsMu.Lock()
+		defer addrsMu.Unlock()
+
 		last = make(map[uint8]ChunkDescriptor)
 		for i := 0; i < count; i++ {
 			chunk := generateRandomChunk()
@@ -341,7 +346,6 @@ func TestDB_SubscribePull_sinceAndUntil(t *testing.T) {
 
 			bin := db.po(chunk.Address())
 
-			addrsMu.Lock()
 			if _, ok := addrs[bin]; !ok {
 				addrs[bin] = make([]storage.Address, 0)
 			}
@@ -349,7 +353,6 @@ func TestDB_SubscribePull_sinceAndUntil(t *testing.T) {
 				addrs[bin] = append(addrs[bin], chunk.Address())
 				wantedChunksCount++
 			}
-			addrsMu.Unlock()
 
 			lastTimestampMu.RLock()
 			storeTimestamp := lastTimestamp
@@ -410,6 +413,9 @@ func TestDB_SubscribePull_sinceAndUntil(t *testing.T) {
 // uploadRandomChunksBin uploads random chunks to database and adds them to
 // the map of addresses ber bin.
 func uploadRandomChunksBin(t *testing.T, db *DB, uploader *Putter, addrs map[uint8][]storage.Address, addrsMu *sync.Mutex, wantedChunksCount *int, count int) {
+	addrsMu.Lock()
+	defer addrsMu.Unlock()
+
 	for i := 0; i < count; i++ {
 		chunk := generateRandomChunk()
 
@@ -418,13 +424,11 @@ func uploadRandomChunksBin(t *testing.T, db *DB, uploader *Putter, addrs map[uin
 			t.Fatal(err)
 		}
 
-		addrsMu.Lock()
 		bin := db.po(chunk.Address())
 		if _, ok := addrs[bin]; !ok {
 			addrs[bin] = make([]storage.Address, 0)
 		}
 		addrs[bin] = append(addrs[bin], chunk.Address())
-		addrsMu.Unlock()
 
 		*wantedChunksCount++
 	}
diff --git a/swarm/storage/localstore/subscription_push_test.go b/swarm/storage/localstore/subscription_push_test.go
index f1a0df7a34..0c8d7d0b9a 100644
--- a/swarm/storage/localstore/subscription_push_test.go
+++ b/swarm/storage/localstore/subscription_push_test.go
@@ -40,6 +40,9 @@ func TestDB_SubscribePush(t *testing.T) {
 	var chunksMu sync.Mutex
 
 	uploadRandomChunks := func(count int) {
+		chunksMu.Lock()
+		defer chunksMu.Unlock()
+
 		for i := 0; i < count; i++ {
 			chunk := generateRandomChunk()
 
@@ -48,9 +51,7 @@ func TestDB_SubscribePush(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			chunksMu.Lock()
 			chunks = append(chunks, chunk)
-			chunksMu.Unlock()
 		}
 	}
 
@@ -127,6 +128,9 @@ func TestDB_SubscribePush_multiple(t *testing.T) {
 	var addrsMu sync.Mutex
 
 	uploadRandomChunks := func(count int) {
+		addrsMu.Lock()
+		defer addrsMu.Unlock()
+
 		for i := 0; i < count; i++ {
 			chunk := generateRandomChunk()
 
@@ -135,9 +139,7 @@ func TestDB_SubscribePush_multiple(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			addrsMu.Lock()
 			addrs = append(addrs, chunk.Address())
-			addrsMu.Unlock()
 		}
 	}