chore: fix flake in listening ports test (#10833)

deansheather · web-flow · commit a9c0c01629c3 · 2023-11-22T09:30:51.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -68,6 +68,7 @@ type Options struct {
 	EnvironmentVariables         map[string]string
 	Logger                       slog.Logger
 	IgnorePorts                  map[int]string
+	PortCacheDuration            time.Duration
 	SSHMaxTimeout                time.Duration
 	TailnetListenPort            uint16
 	Subsystems                   []codersdk.AgentSubsystem
@@ -126,6 +127,9 @@ func New(options Options) Agent {
 	if options.ServiceBannerRefreshInterval == 0 {
 		options.ServiceBannerRefreshInterval = 2 * time.Minute
 	}
+	if options.PortCacheDuration == 0 {
+		options.PortCacheDuration = 1 * time.Second
+	}
 
 	prometheusRegistry := options.PrometheusRegistry
 	if prometheusRegistry == nil {
@@ -153,6 +157,7 @@ func New(options Options) Agent {
 		lifecycleReported:            make(chan codersdk.WorkspaceAgentLifecycle, 1),
 		lifecycleStates:              []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},
 		ignorePorts:                  options.IgnorePorts,
+		portCacheDuration:            options.PortCacheDuration,
 		connStatsChan:                make(chan *agentsdk.Stats, 1),
 		reportMetadataInterval:       options.ReportMetadataInterval,
 		serviceBannerRefreshInterval: options.ServiceBannerRefreshInterval,
@@ -181,8 +186,9 @@ type agent struct {
 	// ignorePorts tells the api handler which ports to ignore when
 	// listing all listening ports. This is helpful to hide ports that
 	// are used by the agent, that the user does not care about.
-	ignorePorts map[int]string
-	subsystems  []codersdk.AgentSubsystem
+	ignorePorts       map[int]string
+	portCacheDuration time.Duration
+	subsystems        []codersdk.AgentSubsystem
 
 	reconnectingPTYs       sync.Map
 	reconnectingPTYTimeout time.Duration
diff --git a/agent/api.go b/agent/api.go
@@ -26,17 +26,27 @@ func (a *agent) apiHandler() http.Handler {
 		cpy[k] = b
 	}
 
-	lp := &listeningPortsHandler{ignorePorts: cpy}
+	cacheDuration := 1 * time.Second
+	if a.portCacheDuration > 0 {
+		cacheDuration = a.portCacheDuration
+	}
+
+	lp := &listeningPortsHandler{
+		ignorePorts:   cpy,
+		cacheDuration: cacheDuration,
+	}
 	r.Get("/api/v0/listening-ports", lp.handler)
 
 	return r
 }
 
 type listeningPortsHandler struct {
-	mut         sync.Mutex
-	ports       []codersdk.WorkspaceAgentListeningPort
-	mtime       time.Time
-	ignorePorts map[int]string
+	ignorePorts   map[int]string
+	cacheDuration time.Duration
+
+	mut   sync.Mutex
+	ports []codersdk.WorkspaceAgentListeningPort
+	mtime time.Time
 }
 
 // handler returns a list of listening ports. This is tested by coderd's
diff --git a/agent/ports_supported.go b/agent/ports_supported.go
@@ -15,7 +15,7 @@ func (lp *listeningPortsHandler) getListeningPorts() ([]codersdk.WorkspaceAgentL
 	lp.mut.Lock()
 	defer lp.mut.Unlock()
 
-	if time.Since(lp.mtime) < time.Second {
+	if time.Since(lp.mtime) < lp.cacheDuration {
 		// copy
 		ports := make([]codersdk.WorkspaceAgentListeningPort, len(lp.ports))
 		copy(ports, lp.ports)
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -21,6 +21,7 @@ import (
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/coderd"
 	"github.com/coder/coder/v2/coderd/coderdtest"
@@ -551,7 +552,9 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) {
 				},
 			}},
 		}).Do()
-		_ = agenttest.New(t, client.URL, authToken)
+		_ = agenttest.New(t, client.URL, authToken, func(o *agent.Options) {
+			o.PortCacheDuration = time.Millisecond
+		})
 		resources := coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
 		return client, uint16(coderdPort), resources[0].Agents[0].ID
 	}
@@ -670,15 +673,21 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) {
 
 			// Close the listener and check that the port is no longer in the response.
 			require.NoError(t, l.Close())
-			time.Sleep(2 * time.Second) // avoid cache
-			res, err = client.WorkspaceAgentListeningPorts(ctx, agentID)
-			require.NoError(t, err)
+			t.Log("checking for ports after listener close:")
+			require.Eventually(t, func() bool {
+				res, err = client.WorkspaceAgentListeningPorts(ctx, agentID)
+				if !assert.NoError(t, err) {
+					return false
+				}
 
-			for _, port := range res.Ports {
-				if port.Network == "tcp" && port.Port == lPort {
-					t.Fatalf("expected to not find TCP port %d in response", lPort)
+				for _, port := range res.Ports {
+					if port.Network == "tcp" && port.Port == lPort {
+						t.Logf("expected to not find TCP port %d in response", lPort)
+						return false
+					}
 				}
-			}
+				return true
+			}, testutil.WaitLong, testutil.IntervalMedium)
 		})
 
 		t.Run("Filter", func(t *testing.T) {
