Only wait for 5 seconds to report lifecycle at end

mafredri · mafredri · commit 5adc708e03da · 2023-02-10T12:22:24.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -1268,18 +1268,28 @@ func (a *agent) Close() error {
 	// Set final state and wait for it to be reported because context
 	// cancellation will stop the report loop.
 	a.setLifecycle(ctx, lifecycleState)
-	// TODO(mafredri): What if the agent token is revoked, build outdated, etc.?
-	for s := range a.lifecycleReported {
-		if s == lifecycleState {
-			break
-		}
-	}
 
 	if lifecycleState != codersdk.WorkspaceAgentLifecycleOff {
 		// TODO(mafredri): Delay shutdown, ensure debugging is possible.
 		_ = false
 	}
 
+	// Wait for the lifecycle to be reported, but don't wait forever so
+	// that we don't break user expectations.
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+lifecycleWaitLoop:
+	for {
+		select {
+		case <-ctx.Done():
+			break lifecycleWaitLoop
+		case s := <-a.lifecycleReported:
+			if s == lifecycleState {
+				break lifecycleWaitLoop
+			}
+		}
+	}
+
 	close(a.closed)
 	a.closeCancel()
 	_ = a.sshServer.Close()
