fix bug due to data racing on VidMap (#3606)

1 files changed, 40 insertions, 1 deletions

diff --git a/weed/wdclient/vid_map_test.go b/weed/wdclient/vid_map_test.go
index eae456c9f..980e5bd8c 100644
--- a/weed/wdclient/vid_map_test.go
+++ b/weed/wdclient/vid_map_test.go
@@ -1,15 +1,17 @@
 package wdclient
 
 import (
+	"context"
 	"fmt"
 	"google.golang.org/grpc"
 	"strconv"
 	"sync"
 	"testing"
+	"time"
 )
 
 func TestLocationIndex(t *testing.T) {
-	vm := vidMap{}
+	vm := &vidMap{}
 	// test must be failed
 	mustFailed := func(length int) {
 		_, err := vm.getLocationIndex(length)
@@ -132,6 +134,43 @@ func TestLookupFileId(t *testing.T) {
 	wg.Wait()
 }
 
+func TestConcurrentGetLocations(t *testing.T) {
+	mc := NewMasterClient(grpc.EmptyDialOption{}, "", "", "", "", "", nil)
+	location := Location{Url: "TestDataRacing"}
+	mc.addLocation(1, location)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	wg := sync.WaitGroup{}
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					_, found := mc.GetLocations(1)
+					if !found {
+						cancel()
+						t.Error("vid map invalid due to data racing. ")
+						return
+					}
+				}
+			}
+		}()
+	}
+
+	//Simulate vidmap reset with cache when leader changes
+	for i := 0; i < 100; i++ {
+		mc.resetVidMap()
+		mc.addLocation(1, location)
+		time.Sleep(1 * time.Microsecond)
+	}
+	cancel()
+	wg.Wait()
+}
+
 func BenchmarkLocationIndex(b *testing.B) {
 	b.SetParallelism(8)
 	vm := vidMap{