From b3e276ff9353e119d1e5391e713e6c698948f274 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcgstyle.net>
Date: Mon, 29 Aug 2016 13:39:24 -0700
Subject: [PATCH] Add hash map locking to proxy tests

Signed-off-by: Derek McGowan <derek@mcgstyle.net> (github: dmcgowan)
---
 registry/proxy/proxyblobstore_test.go      |  7 +++++++
 registry/proxy/scheduler/scheduler_test.go | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/registry/proxy/proxyblobstore_test.go b/registry/proxy/proxyblobstore_test.go
index 967dcd3d..8e3a0692 100644
--- a/registry/proxy/proxyblobstore_test.go
+++ b/registry/proxy/proxyblobstore_test.go
@@ -370,15 +370,20 @@ func testProxyStoreServe(t *testing.T, te *testEnv, numClients int) {
 	wg.Wait()
 
 	remoteBlobCount := len(te.inRemote)
+	sbsMu.Lock()
 	if (*localStats)["stat"] != remoteBlobCount*numClients && (*localStats)["create"] != te.numUnique {
+		sbsMu.Unlock()
 		t.Fatal("Expected: stat:", remoteBlobCount*numClients, "create:", remoteBlobCount)
 	}
+	sbsMu.Unlock()
 
 	// Wait for any async storage goroutines to finish
 	time.Sleep(3 * time.Second)
 
+	sbsMu.Lock()
 	remoteStatCount := (*remoteStats)["stat"]
 	remoteOpenCount := (*remoteStats)["open"]
+	sbsMu.Unlock()
 
 	// Serveblob - blobs come from local
 	for _, dr := range te.inRemote {
@@ -403,6 +408,8 @@ func testProxyStoreServe(t *testing.T, te *testEnv, numClients int) {
 	remoteStats = te.RemoteStats()
 
 	// Ensure remote unchanged
+	sbsMu.Lock()
+	defer sbsMu.Unlock()
 	if (*remoteStats)["stat"] != remoteStatCount && (*remoteStats)["open"] != remoteOpenCount {
 		t.Fatalf("unexpected remote stats: %#v", remoteStats)
 	}
diff --git a/registry/proxy/scheduler/scheduler_test.go b/registry/proxy/scheduler/scheduler_test.go
index 556f5204..b74ae4f9 100644
--- a/registry/proxy/scheduler/scheduler_test.go
+++ b/registry/proxy/scheduler/scheduler_test.go
@@ -2,6 +2,7 @@ package scheduler
 
 import (
 	"encoding/json"
+	"sync"
 	"testing"
 	"time"
 
@@ -38,6 +39,7 @@ func TestSchedule(t *testing.T) {
 		ref3.String(): true,
 	}
 
+	var mu sync.Mutex
 	s := New(context.Background(), inmemory.New(), "/ttl")
 	deleteFunc := func(repoName reference.Reference) error {
 		if len(remainingRepos) == 0 {
@@ -48,7 +50,9 @@ func TestSchedule(t *testing.T) {
 			t.Fatalf("Trying to remove nonexistent repo: %s", repoName)
 		}
 		t.Log("removing", repoName)
+		mu.Lock()
 		delete(remainingRepos, repoName.String())
+		mu.Unlock()
 
 		return nil
 	}
@@ -62,12 +66,17 @@ func TestSchedule(t *testing.T) {
 	s.add(ref2, 1*timeUnit, entryTypeBlob)
 
 	func() {
+		s.Lock()
 		s.add(ref3, 1*timeUnit, entryTypeBlob)
+		s.Unlock()
 
 	}()
 
 	// Ensure all repos are deleted
 	<-time.After(50 * timeUnit)
+
+	mu.Lock()
+	defer mu.Unlock()
 	if len(remainingRepos) != 0 {
 		t.Fatalf("Repositories remaining: %#v", remainingRepos)
 	}
@@ -80,6 +89,7 @@ func TestRestoreOld(t *testing.T) {
 		ref2.String(): true,
 	}
 
+	var mu sync.Mutex
 	deleteFunc := func(r reference.Reference) error {
 		if r.String() == ref1.String() && len(remainingRepos) == 2 {
 			t.Errorf("ref1 should be removed first")
@@ -88,7 +98,9 @@ func TestRestoreOld(t *testing.T) {
 		if !ok {
 			t.Fatalf("Trying to remove nonexistent repo: %s", r)
 		}
+		mu.Lock()
 		delete(remainingRepos, r.String())
+		mu.Unlock()
 		return nil
 	}
 
@@ -124,6 +136,8 @@ func TestRestoreOld(t *testing.T) {
 	}
 
 	<-time.After(50 * timeUnit)
+	mu.Lock()
+	defer mu.Unlock()
 	if len(remainingRepos) != 0 {
 		t.Fatalf("Repositories remaining: %#v", remainingRepos)
 	}
@@ -138,8 +152,11 @@ func TestStopRestore(t *testing.T) {
 		ref2.String(): true,
 	}
 
+	var mu sync.Mutex
 	deleteFunc := func(r reference.Reference) error {
+		mu.Lock()
 		delete(remainingRepos, r.String())
+		mu.Unlock()
 		return nil
 	}
 
@@ -169,6 +186,8 @@ func TestStopRestore(t *testing.T) {
 	}
 
 	<-time.After(500 * timeUnit)
+	mu.Lock()
+	defer mu.Unlock()
 	if len(remainingRepos) != 0 {
 		t.Fatalf("Repositories remaining: %#v", remainingRepos)
 	}