Fix flaky mem tracker test (#1923)

src/dbnode/storage/mem_tracker.go

@@ -20,7 +20,9 @@
 
 package storage
 
-import "sync"
+import (
+	"sync"
+)
 
 type memoryTracker struct {
 	sync.Mutex

src/dbnode/storage/mem_tracker_test.go

@@ -23,7 +23,6 @@ package storage
 import (
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -121,58 +120,49 @@ func TestMemoryTrackerIncMarkAndDec(t *testing.T) {
 // or race conditions.
 func TestMemTrackerWaitForDec(t *testing.T) {
 	var (
-		memTracker = NewMemoryTracker(NewMemoryTrackerOptions(100))
-		doneCh     = make(chan struct{})
-		wg         sync.WaitGroup
+		numIterations = 1000
+		memTracker    = NewMemoryTracker(NewMemoryTrackerOptions(100))
+		doneCh        = make(chan struct{})
+		wg            sync.WaitGroup
 	)
 
 	// Start a goroutine to call MarkLoadedAsPending() in a loop.
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		for {
-			select {
-			case <-doneCh:
-				return
-			default:
-				memTracker.MarkLoadedAsPending()
-			}
+		for i := 0; i < numIterations; i++ {
+			memTracker.MarkLoadedAsPending()
 		}
 	}()
 
-	// Start a goroutine to call DecPendingLoadedBytes() in a loop.
+	// Start a goroutine to call WaitForDec() in a loop.
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		for {
-			select {
-			case <-doneCh:
-				return
-			default:
-				memTracker.DecPendingLoadedBytes()
-			}
+		for i := 0; i < numIterations; i++ {
+			memTracker.WaitForDec()
 		}
 	}()
 
-	// Start a goroutine to call WaitForDec() in a loop.
-	wg.Add(1)
+	// Start a goroutine to call DecPendingLoadedBytes() in a loop. Note that
+	// unlike the other two goroutines this one loops infinitely until the test
+	// is over. This is to prevent calls to WaitForDec() from getting stuck forever
+	// because a call to WaitForDec() was made after the goroutine that calls
+	// DecPendingLoadedBytes() had already shut down.
 	go func() {
-		defer wg.Done()
 		for {
 			select {
 			case <-doneCh:
 				return
 			default:
-				memTracker.WaitForDec()
+				memTracker.DecPendingLoadedBytes()
 			}
 		}
 	}()
 
-	time.Sleep(100 * time.Millisecond)
-
-	// Stop the background goroutines.
-	close(doneCh)
-
 	// Ensure all the goroutines exit cleanly (ensuring no deadlocks.)
 	wg.Wait()
+
+	// Stop the background goroutine calling DecPendingLoadedBytes().
+	close(doneCh)
 }