Fix flaky ping test for remote clusters service (#26821)

* fix flaky ping test for remote clusters service

* refactor ping test

* Update server/platform/services/remotecluster/ping_test.go

server/platform/services/remotecluster/ping_test.go

@@ -9,7 +9,6 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"sync"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -28,16 +27,14 @@ func TestPing(t *testing.T) {
 	disablePing = false
 
 	t.Run("No error", func(t *testing.T) {
-		var countWebReq int32
 		merr := merror.New()
 
-		wg := &sync.WaitGroup{}
-		wg.Add(NumRemotes)
+		var remotes []*model.RemoteCluster
+		pingsReceived := make(map[string]struct{})
+		var mux sync.Mutex
 
 		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			defer wg.Done()
 			defer w.WriteHeader(200)
-			atomic.AddInt32(&countWebReq, 1)
 
 			var frame model.RemoteClusterFrame
 			err := json.NewDecoder(r.Body).Decode(&frame)
@@ -50,12 +47,19 @@ func TestPing(t *testing.T) {
 				return
 			}
 
+			// Make sure ping is from a remote that was added for this test.
+			if !hasRemoteID(frame.RemoteId, remotes) {
+				merr.Append(fmt.Errorf("RemoteID not in list of remotes for this test; remote_id=%s", frame.RemoteId))
+				return
+			}
+
 			var ping model.RemoteClusterPing
 			err = json.Unmarshal(frame.Msg.Payload, &ping)
 			if err != nil {
 				merr.Append(err)
 				return
 			}
+
 			if !checkRecent(ping.SentAt, Recent) {
 				merr.Append(fmt.Errorf("timestamp out of range, got %d", ping.SentAt))
 				return
@@ -64,10 +68,15 @@ func TestPing(t *testing.T) {
 				merr.Append(fmt.Errorf("timestamp should be 0, got %d", ping.RecvAt))
 				return
 			}
+
+			mux.Lock()
+			defer mux.Unlock()
+			pingsReceived[frame.RemoteId] = struct{}{}
 		}))
 		defer ts.Close()
 
-		mockServer := newMockServer(t, makeRemoteClusters(NumRemotes, ts.URL, false))
+		remotes = makeRemoteClusters(NumRemotes, ts.URL, false)
+		mockServer := newMockServer(t, remotes)
 		mockApp := newMockApp(t, nil)
 
 		service, err := NewRemoteClusterService(mockServer, mockApp)
@@ -77,31 +86,37 @@ func TestPing(t *testing.T) {
 		require.NoError(t, err)
 		defer service.Shutdown()
 
-		wg.Wait()
+		// wait up to 10 seconds for all remotes to get pinged. This will normally take less than 1 second
+		// unless the server is very busy.
+		assert.Eventually(t, func() bool {
+			mux.Lock()
+			defer mux.Unlock()
+			return len(pingsReceived) == NumRemotes
+		}, time.Second*10, time.Millisecond*50, "all remotes must get pinged")
 
 		assert.NoError(t, merr.ErrorOrNil())
-
-		assert.Equal(t, int32(NumRemotes), atomic.LoadInt32(&countWebReq))
-		t.Logf("%d web requests counted;  %d expected",
-			atomic.LoadInt32(&countWebReq), NumRemotes)
 	})
 
 	t.Run("HTTP errors", func(t *testing.T) {
-		var countWebReq int32
 		merr := merror.New()
 
-		wg := &sync.WaitGroup{}
-		wg.Add(NumRemotes)
+		var remotes []*model.RemoteCluster
+		pingsReceived := make(map[string]struct{})
+		var mux sync.Mutex
 
 		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			defer wg.Done()
-			atomic.AddInt32(&countWebReq, 1)
-
 			var frame model.RemoteClusterFrame
 			err := json.NewDecoder(r.Body).Decode(&frame)
 			if err != nil {
 				merr.Append(err)
 			}
+
+			// Make sure ping is from a remote that was added for this test.
+			if !hasRemoteID(frame.RemoteId, remotes) {
+				merr.Append(fmt.Errorf("RemoteID not in list of remotes for this test; remote_id=%s", frame.RemoteId))
+				return
+			}
+
 			var ping model.RemoteClusterPing
 			err = json.Unmarshal(frame.Msg.Payload, &ping)
 			if err != nil {
@@ -113,11 +128,17 @@ func TestPing(t *testing.T) {
 			if ping.RecvAt != 0 {
 				merr.Append(fmt.Errorf("timestamp should be 0, got %d", ping.RecvAt))
 			}
+
 			w.WriteHeader(500)
+
+			mux.Lock()
+			defer mux.Unlock()
+			pingsReceived[frame.RemoteId] = struct{}{}
 		}))
 		defer ts.Close()
 
-		mockServer := newMockServer(t, makeRemoteClusters(NumRemotes, ts.URL, false))
+		remotes = makeRemoteClusters(NumRemotes, ts.URL, false)
+		mockServer := newMockServer(t, remotes)
 		mockApp := newMockApp(t, nil)
 
 		service, err := NewRemoteClusterService(mockServer, mockApp)
@@ -127,13 +148,15 @@ func TestPing(t *testing.T) {
 		require.NoError(t, err)
 		defer service.Shutdown()
 
-		wg.Wait()
+		// wait up to 10 seconds for all remotes to get pinged. This will normally take less than 1 second
+		// until the server is very busy.
+		assert.Eventually(t, func() bool {
+			mux.Lock()
+			defer mux.Unlock()
+			return len(pingsReceived) == NumRemotes
+		}, time.Second*10, time.Millisecond*50, "all remotes must get pinged")
 
 		assert.NoError(t, merr.ErrorOrNil())
-
-		assert.Equal(t, int32(NumRemotes), atomic.LoadInt32(&countWebReq))
-		t.Logf("%d web requests counted;  %d expected",
-			atomic.LoadInt32(&countWebReq), NumRemotes)
 	})
 
 	t.Run("Plugin ping", func(t *testing.T) {
@@ -169,3 +192,12 @@ func checkRecent(millis int64, within int64) bool {
 	now := model.GetMillis()
 	return millis > now-within && millis < now+within
 }
+
+func hasRemoteID(remoteID string, remotes []*model.RemoteCluster) bool {
+	for _, r := range remotes {
+		if r.RemoteId == remoteID {
+			return true
+		}
+	}
+	return false
+}