coder · coadler · May 30, 2023 · May 26, 2023 · May 30, 2023 · May 30, 2023
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -129,7 +129,7 @@ type Options struct {
 	// AppSecurityKey is the crypto key used to sign and encrypt tokens related to
 	// workspace applications. It consists of both a signing and encryption key.
 	AppSecurityKey     workspaceapps.SecurityKey
-	HealthcheckFunc    func(ctx context.Context) (*healthcheck.Report, error)
+	HealthcheckFunc    func(ctx context.Context, apiKey string) (*healthcheck.Report, error)
 	HealthcheckTimeout time.Duration
 	HealthcheckRefresh time.Duration
 
@@ -256,10 +256,11 @@ func New(options *Options) *API {
 		options.TemplateScheduleStore.Store(&v)
 	}
 	if options.HealthcheckFunc == nil {
-		options.HealthcheckFunc = func(ctx context.Context) (*healthcheck.Report, error) {
+		options.HealthcheckFunc = func(ctx context.Context, apiKey string) (*healthcheck.Report, error) {
 			return healthcheck.Run(ctx, &healthcheck.ReportOptions{
 				AccessURL: options.AccessURL,
 				DERPMap:   options.DERPMap.Clone(),
+				APIKey:    apiKey,
 			})
 		}
 	}
@@ -787,6 +788,7 @@ func New(options *Options) *API {
 
 			r.Get("/coordinator", api.debugCoordinator)
 			r.Get("/health", api.debugDeploymentHealth)
+			r.Get("/ws", (&healthcheck.WebsocketEchoServer{}).ServeHTTP)
 		})
 	})
 
@@ -874,6 +876,7 @@ type API struct {
 	Experiments codersdk.Experiments
 
 	healthCheckGroup *singleflight.Group[string, *healthcheck.Report]
+	healthCheckCache atomic.Pointer[healthcheck.Report]
 }
 
 // Close waits for all WebSocket connections to drain before returning.

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -107,7 +107,7 @@ type Options struct {
 	TrialGenerator        func(context.Context, string) error
 	TemplateScheduleStore schedule.TemplateScheduleStore
 
-	HealthcheckFunc    func(ctx context.Context) (*healthcheck.Report, error)
+	HealthcheckFunc    func(ctx context.Context, apiKey string) (*healthcheck.Report, error)
 	HealthcheckTimeout time.Duration
 	HealthcheckRefresh time.Duration
 

diff --git a/coderd/debug.go b/coderd/debug.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/coder/coder/coderd/healthcheck"
 	"github.com/coder/coder/coderd/httpapi"
+	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/codersdk"
 )
 
@@ -29,11 +30,28 @@ func (api *API) debugCoordinator(rw http.ResponseWriter, r *http.Request) {
 // @Success 200 {object} healthcheck.Report
 // @Router /debug/health [get]
 func (api *API) debugDeploymentHealth(rw http.ResponseWriter, r *http.Request) {
+	apiKey := httpmw.APITokenFromRequest(r)
 	ctx, cancel := context.WithTimeout(r.Context(), api.HealthcheckTimeout)
 	defer cancel()
 
+	// Get cached report if it exists.
+	if report := api.healthCheckCache.Load(); report != nil {
+		if time.Since(report.Time) < api.HealthcheckRefresh {
+			httpapi.Write(ctx, rw, http.StatusOK, report)
+			return
+		}
+	}
+
 	resChan := api.healthCheckGroup.DoChan("", func() (*healthcheck.Report, error) {
-		return api.HealthcheckFunc(ctx)
+		// Create a new context not tied to the request.
+		ctx, cancel := context.WithTimeout(context.Background(), api.HealthcheckTimeout)
+		defer cancel()
+
+		report, err := api.HealthcheckFunc(ctx, apiKey)
+		if err == nil {
+			api.healthCheckCache.Store(report)
+		}
+		return report, err
 	})
 
 	select {
@@ -43,13 +61,19 @@ func (api *API) debugDeploymentHealth(rw http.ResponseWriter, r *http.Request) {
 		})
 		return
 	case res := <-resChan:
-		if time.Since(res.Val.Time) > api.HealthcheckRefresh {
-			api.healthCheckGroup.Forget("")
-			api.debugDeploymentHealth(rw, r)
-			return
-		}
-
 		httpapi.Write(ctx, rw, http.StatusOK, res.Val)
 		return
 	}
 }
+
+// For some reason the swagger docs need to be attached to a function.
+//
+// @Summary Debug Info Websocket Test
+// @ID debug-info-websocket-test
+// @Security CoderSessionToken
+// @Produce json
+// @Tags Debug
+// @Success 201 {object} codersdk.Response
+// @Router /debug/ws [get]
+// @x-apidocgen {"skip": true}
+func _debugws(http.ResponseWriter, *http.Request) {} //nolint:unused
diff --git a/coderd/debug_test.go b/coderd/debug_test.go
@@ -7,44 +7,48 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/coderd/healthcheck"
 	"github.com/coder/coder/testutil"
 )
 
-func TestDebug(t *testing.T) {
+func TestDebugHealth(t *testing.T) {
 	t.Parallel()
-	t.Run("Health/OK", func(t *testing.T) {
+	t.Run("OK", func(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx, cancel = context.WithTimeout(context.Background(), testutil.WaitShort)
-			client      = coderdtest.New(t, &coderdtest.Options{
-				HealthcheckFunc: func(context.Context) (*healthcheck.Report, error) {
+			ctx, cancel  = context.WithTimeout(context.Background(), testutil.WaitShort)
+			sessionToken string
+			client       = coderdtest.New(t, &coderdtest.Options{
+				HealthcheckFunc: func(_ context.Context, apiKey string) (*healthcheck.Report, error) {
+					assert.Equal(t, sessionToken, apiKey)
 					return &healthcheck.Report{}, nil
 				},
 			})
 			_ = coderdtest.CreateFirstUser(t, client)
 		)
 		defer cancel()
 
+		sessionToken = client.SessionToken()
 		res, err := client.Request(ctx, "GET", "/debug/health", nil)
 		require.NoError(t, err)
 		defer res.Body.Close()
 		_, _ = io.ReadAll(res.Body)
 		require.Equal(t, http.StatusOK, res.StatusCode)
 	})
 
-	t.Run("Health/Timeout", func(t *testing.T) {
+	t.Run("Timeout", func(t *testing.T) {
 		t.Parallel()
 
 		var (
 			ctx, cancel = context.WithTimeout(context.Background(), testutil.WaitShort)
 			client      = coderdtest.New(t, &coderdtest.Options{
 				HealthcheckTimeout: time.Microsecond,
-				HealthcheckFunc: func(context.Context) (*healthcheck.Report, error) {
+				HealthcheckFunc: func(context.Context, string) (*healthcheck.Report, error) {
 					t := time.NewTimer(time.Second)
 					defer t.Stop()
 
@@ -66,4 +70,48 @@ func TestDebug(t *testing.T) {
 		_, _ = io.ReadAll(res.Body)
 		require.Equal(t, http.StatusNotFound, res.StatusCode)
 	})
+
+	t.Run("Deduplicated", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			ctx, cancel = context.WithTimeout(context.Background(), testutil.WaitShort)
+			calls       int
+			client      = coderdtest.New(t, &coderdtest.Options{
+				HealthcheckRefresh: time.Hour,
+				HealthcheckTimeout: time.Hour,
+				HealthcheckFunc: func(context.Context, string) (*healthcheck.Report, error) {
+					calls++
+					return &healthcheck.Report{
+						Time: time.Now(),
+					}, nil
+				},
+			})
+			_ = coderdtest.CreateFirstUser(t, client)
+		)
+		defer cancel()
+
+		res, err := client.Request(ctx, "GET", "/api/v2/debug/health", nil)
+		require.NoError(t, err)
+		defer res.Body.Close()
+		_, _ = io.ReadAll(res.Body)
+
+		require.Equal(t, http.StatusOK, res.StatusCode)
+
+		res, err = client.Request(ctx, "GET", "/api/v2/debug/health", nil)
+		require.NoError(t, err)
+		defer res.Body.Close()
+		_, _ = io.ReadAll(res.Body)
+
+		require.Equal(t, http.StatusOK, res.StatusCode)
+		require.Equal(t, 1, calls)
+	})
+}
+
+func TestDebugWebsocket(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+	})
 }
diff --git a/coderd/healthcheck/healthcheck.go b/coderd/healthcheck/healthcheck.go
@@ -19,16 +19,15 @@ type Report struct {
 
 	DERP      DERPReport      `json:"derp"`
 	AccessURL AccessURLReport `json:"access_url"`
-
-	// TODO:
-	// Websocket WebsocketReport `json:"websocket"`
+	Websocket WebsocketReport `json:"websocket"`
 }
 
 type ReportOptions struct {
 	// TODO: support getting this over HTTP?
 	DERPMap   *tailcfg.DERPMap
 	AccessURL *url.URL
 	Client    *http.Client
+	APIKey    string
 }
 
 func Run(ctx context.Context, opts *ReportOptions) (*Report, error) {
@@ -65,11 +64,19 @@ func Run(ctx context.Context, opts *ReportOptions) (*Report, error) {
 		})
 	}()
 
-	// wg.Add(1)
-	// go func() {
-	// 	defer wg.Done()
-	// 	report.Websocket.Run(ctx, opts.AccessURL)
-	// }()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		defer func() {
+			if err := recover(); err != nil {
+				report.Websocket.Error = xerrors.Errorf("%v", err)
+			}
+		}()
+		report.Websocket.Run(ctx, &WebsocketReportOptions{
+			APIKey:    opts.APIKey,
+			AccessURL: opts.AccessURL,
+		})
+	}()
 
 	wg.Wait()
 	report.Time = time.Now()