From 333e6f1cb95aced491d96e18cbb45ac7c096d93a Mon Sep 17 00:00:00 2001
From: apolcyn <apolcyn@google.com>
Date: Sun, 19 May 2024 15:58:17 -0700
Subject: [PATCH] [probes.external] Fix a race condition in server mode (#744)

---
 probes/external/external_server.go | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/probes/external/external_server.go b/probes/external/external_server.go
index 982485a96e..d0c7a0a5cf 100644
--- a/probes/external/external_server.go
+++ b/probes/external/external_server.go
@@ -198,7 +198,6 @@ func (p *Probe) runServerProbe(ctx, startCtx context.Context) {
 
 	outstandingReqs := make(map[int32]requestInfo)
 	var outstandingReqsMu sync.RWMutex
-	sendDoneCh := make(chan struct{})
 
 	if err := p.startCmdIfNotRunning(startCtx); err != nil {
 		p.l.Error(err.Error())
@@ -212,6 +211,7 @@ func (p *Probe) runServerProbe(ctx, startCtx context.Context) {
 
 		// Read probe replies until we have no outstanding requests or context has
 		// run out.
+		expectedRepliesReceived := 0
 		for {
 			select {
 			case <-ctx.Done():
@@ -221,6 +221,7 @@ func (p *Probe) runServerProbe(ctx, startCtx context.Context) {
 				outstandingReqsMu.Lock()
 				reqInfo, ok := outstandingReqs[rep.GetRequestId()]
 				if ok {
+					expectedRepliesReceived++
 					delete(outstandingReqs, rep.GetRequestId())
 				}
 				outstandingReqsMu.Unlock()
@@ -243,14 +244,10 @@ func (p *Probe) runServerProbe(ctx, startCtx context.Context) {
 				p.processProbeResult(ps, p.results[reqInfo.target.Key()])
 			}
 
-			// If we are done sending requests, we can exit if we have no
-			// outstanding requests.
-			select {
-			case <-sendDoneCh:
-				if len(outstandingReqs) == 0 {
-					return
-				}
-			default:
+			// We send a total if len(p.targets) requests. We can exit if we've
+			// seen replies for all of them.
+			if expectedRepliesReceived == len(p.targets) {
+				return
 			}
 		}
 	}()
@@ -269,9 +266,6 @@ func (p *Probe) runServerProbe(ctx, startCtx context.Context) {
 		time.Sleep(TimeBetweenRequests)
 	}
 
-	// Send signal to receiver loop that we are done sending requests.
-	close(sendDoneCh)
-
 	// Wait for receiver goroutine to exit.
 	wg.Wait()