From 348530000f761f7e06e100a3882be2ee6716a08d Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Tue, 14 Mar 2023 15:14:47 +0200
Subject: [PATCH] fix(coderd): Ensure agent disconnect happens after timeout
 (#6600)

Fixes #6598
---
 coderd/coderd.go                      | 4 ++++
 coderd/workspaceagents.go             | 2 +-
 scaletest/reconnectingpty/run_test.go | 1 -
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/coderd/coderd.go b/coderd/coderd.go
index b17927203d..ed53b09575 100644
--- a/coderd/coderd.go
+++ b/coderd/coderd.go
@@ -175,6 +175,10 @@ func New(options *Options) *API {
 	if options.AgentInactiveDisconnectTimeout == 0 {
 		// Multiply the update by two to allow for some lag-time.
 		options.AgentInactiveDisconnectTimeout = options.AgentConnectionUpdateFrequency * 2
+		// Set a minimum timeout to avoid disconnecting too soon.
+		if options.AgentInactiveDisconnectTimeout < 2*time.Second {
+			options.AgentInactiveDisconnectTimeout = 2 * time.Second
+		}
 	}
 	if options.AgentStatsRefreshInterval == 0 {
 		options.AgentStatsRefreshInterval = 5 * time.Minute
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
index 041245825e..07784c2d45 100644
--- a/coderd/workspaceagents.go
+++ b/coderd/workspaceagents.go
@@ -615,7 +615,7 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
 
 	// We use a custom heartbeat routine here instead of `httpapi.Heartbeat`
 	// because we want to log the agent's last ping time.
-	var lastPing time.Time
+	lastPing := time.Now() // Since the agent initiated the request, assume it's alive.
 	var pingMu sync.Mutex
 	go pprof.Do(ctx, pprof.Labels("agent", workspaceAgent.ID.String()), func(ctx context.Context) {
 		// TODO(mafredri): Is this too frequent? Use separate ping disconnect timeout?
diff --git a/scaletest/reconnectingpty/run_test.go b/scaletest/reconnectingpty/run_test.go
index 09b7f3285c..f6f70bbf57 100644
--- a/scaletest/reconnectingpty/run_test.go
+++ b/scaletest/reconnectingpty/run_test.go
@@ -23,7 +23,6 @@ import (
 
 func Test_Runner(t *testing.T) {
 	t.Parallel()
-	t.Skip("https://github.com/coder/coder/issues/6598")
 
 	t.Run("OK", func(t *testing.T) {
 		t.Parallel()