From b023ebccfaf177135c91e79d924576c4d8c1b379 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Thu, 4 Jan 2024 12:11:49 +0100
Subject: [PATCH] Ensure race conditions handled

---
 .github/workflows/integration-tests.yml |  2 ++
 Makefile                                | 14 ++++++++++++++
 pkg/endpoints/endpoints.go              |  8 +++++++-
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml
index e2019d70..4d94880b 100644
--- a/.github/workflows/integration-tests.yml
+++ b/.github/workflows/integration-tests.yml
@@ -17,5 +17,7 @@ jobs:
           uses: actions/setup-go@v4
           with:
             go-version: '>=1.21.0'
+        - name: Run race tests
+          run: make test-race
         - name: Run integration tests
           run: make test-integration
diff --git a/Makefile b/Makefile
index 6054a733..ffe6eace 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,20 @@
 ENVTEST_K8S_VERSION = 1.27.1
 
 .PHONY: test
+test: test-unit
+
+.PHONY: test-all
+test-all: test-race test-integration
+
+.PHONY: test-unit
+test-unit:
+	go test -mod=readonly ./pkg/...
+
+.PHONY: test-race
+test-race:
+	go test -mod=readonly -race ./pkg/...
+
+.PHONY: test-integration
 test-integration: envtest
 	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) --bin-dir $(LOCALBIN) -p path)" go test ./tests/integration -v
 
diff --git a/pkg/endpoints/endpoints.go b/pkg/endpoints/endpoints.go
index e5112a92..0e51083a 100644
--- a/pkg/endpoints/endpoints.go
+++ b/pkg/endpoints/endpoints.go
@@ -45,7 +45,7 @@ func (e *endpointGroup) getBestHost(ctx context.Context, portName string) (strin
 	for len(e.endpoints) == 0 {
 		e.mtx.RUnlock()
 		select {
-		case <-e.bcast:
+		case <-e.awaitEndpoints():
 		case <-ctx.Done():
 			return "", ctx.Err()
 		}
@@ -65,6 +65,12 @@ func (e *endpointGroup) getBestHost(ctx context.Context, portName string) (strin
 	return fmt.Sprintf("%s:%v", bestIP, port), nil
 }
 
+func (e *endpointGroup) awaitEndpoints() chan struct{} {
+	e.bmtx.RLock()
+	defer e.bmtx.RUnlock()
+	return e.bcast
+}
+
 func (e *endpointGroup) getAllHosts(portName string) []string {
 	e.mtx.RLock()
 	defer e.mtx.RUnlock()