Fix config_server_test flakiness (#5175)

Author: pjanotti

internal/configconverter/config_server.go

@@ -38,7 +38,8 @@ import (
 const (
 	configServerEnabledEnvVar   = "SPLUNK_DEBUG_CONFIG_SERVER"
 	configServerPortEnvVar      = "SPLUNK_DEBUG_CONFIG_SERVER_PORT"
-	defaultConfigServerEndpoint = "localhost:55554"
+	defaultConfigServerPort     = "55554"
+	defaultConfigServerEndpoint = "localhost:" + defaultConfigServerPort
 	effectivePath               = "/debug/configz/effective"
 	initialPath                 = "/debug/configz/initial"
 )
@@ -58,7 +59,6 @@ type ConfigServer struct {
 	initial        map[string]any
 	effective      map[string]any
 	server         *http.Server
-	doneCh         chan struct{}
 	initialMutex   sync.RWMutex
 	effectiveMutex sync.RWMutex
 	wg             sync.WaitGroup
@@ -73,7 +73,6 @@ func NewConfigServer() *ConfigServer {
 		effectiveMutex: sync.RWMutex{},
 		wg:             sync.WaitGroup{},
 		once:           sync.Once{},
-		doneCh:         make(chan struct{}),
 	}
 
 	mux := http.NewServeMux()
@@ -156,8 +155,6 @@ func (cs *ConfigServer) start() {
 			}
 
 			go func() {
-				defer close(cs.doneCh)
-
 				httpErr := cs.server.Serve(listener)
 				if httpErr != http.ErrServerClosed {
 					log.Print(fmt.Errorf("config server error: %w", httpErr).Error())
@@ -166,15 +163,17 @@ func (cs *ConfigServer) start() {
 
 			go func() {
 				cs.wg.Wait()
-				if cs.server != nil {
-					_ = cs.server.Close()
-				}
+				_ = cs.server.Close()
 
 			}()
 
 		})
 }
 
+// OnShutdown doesn't guarantee that server is down. The server is stopped when
+// the number of OnShutdown calls match the number of OnNew calls. The actual
+// shutdown time is non-deterministic since it happens in a goroutine that
+// was waiting for the last OnShutdown call.
 func (cs *ConfigServer) OnShutdown() {
 	cs.wg.Done()
 }

internal/configconverter/config_server_test.go

@@ -18,10 +18,12 @@ package configconverter
 import (
 	"context"
 	"io"
+	"net"
 	"net/http"
 	"os"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -39,7 +41,9 @@ func TestConfigServer_RequireEnvVar(t *testing.T) {
 	cs := NewConfigServer()
 	require.NotNil(t, cs)
 	cs.OnNew()
-	t.Cleanup(cs.OnShutdown)
+	t.Cleanup(func() {
+		shutdownAndWaitForPort(t, cs, defaultConfigServerPort)
+	})
 	require.NoError(t, cs.Convert(context.Background(), confmap.NewFromStringMap(initial)))
 
 	client := &http.Client{}
@@ -90,18 +94,16 @@ func TestConfigServer_EnvVar(t *testing.T) {
 				}()
 			}
 
+			actualServerPort := defaultConfigServerPort
+			if tt.portEnvVar != "" {
+				actualServerPort = tt.portEnvVar
+			}
 			cs := NewConfigServer()
 			require.NotNil(t, cs)
 			cs.OnNew()
 
 			require.NoError(t, cs.Convert(context.Background(), confmap.NewFromStringMap(initial)))
-			defer func() {
-				cs.OnShutdown() // Call shutdown even if the server is not up.
-				if !tt.serverDown {
-					// Wait for the server to shutdown. Otherwise the port might be still in use.
-					<-cs.doneCh
-				}
-			}()
+			defer shutdownAndWaitForPort(t, cs, actualServerPort)
 
 			endpoint := tt.endpoint
 			if endpoint == "" {
@@ -154,7 +156,9 @@ func TestConfigServer_Serve(t *testing.T) {
 	cs := NewConfigServer()
 	require.NotNil(t, cs)
 	cs.OnNew()
-	t.Cleanup(cs.OnShutdown)
+	t.Cleanup(func() {
+		shutdownAndWaitForPort(t, cs, defaultConfigServerPort)
+	})
 
 	cs.OnRetrieve("scheme", initial)
 	require.NoError(t, cs.Convert(context.Background(), confmap.NewFromStringMap(initial)))
@@ -225,3 +229,15 @@ func TestSimpleRedact(t *testing.T) {
 		"X-SF-Token": "<redacted>",
 	}, result)
 }
+
+func shutdownAndWaitForPort(t *testing.T, cs *ConfigServer, port string) {
+	cs.OnShutdown()
+	require.Eventually(t, func() bool {
+		ln, err := net.Listen("tcp", "localhost:"+port)
+		if err == nil {
+			ln.Close()
+			return true
+		}
+		return false
+	}, 5*time.Second, 200*time.Millisecond)
+}