Replace DatabaseHealthcheckFn with interface

Signed-off-by: Danny Kopping <dannykopping@gmail.com>
coder · dannykopping · Apr 15, 2025 · Apr 11, 2025 · Apr 11, 2025 · Apr 11, 2025
commit 68867da2ecc58eb068b2cfd93e81c5be03084593
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -675,14 +675,11 @@ func New(options *Options) *API {
 	api.Auditor.Store(&options.Auditor)
 	api.TailnetCoordinator.Store(&options.TailnetCoordinator)
 	dialer := &InmemTailnetDialer{
-		CoordPtr: &api.TailnetCoordinator,
-		DERPFn:   api.DERPMap,
-		Logger:   options.Logger,
-		ClientID: uuid.New(),
-		DatabaseHealthcheckFn: func(ctx context.Context) error {
-			_, err := api.Database.Ping(ctx)
-			return err
-		},
+		CoordPtr:            &api.TailnetCoordinator,
+		DERPFn:              api.DERPMap,
+		Logger:              options.Logger,
+		ClientID:            uuid.New(),
+		DatabaseHealthCheck: api.Database,
 	}
 	stn, err := NewServerTailnet(api.ctx,
 		options.Logger,

diff --git a/coderd/tailnet.go b/coderd/tailnet.go
@@ -536,19 +536,24 @@ func NewMultiAgentController(ctx context.Context, logger slog.Logger, tracer tra
 	return m
 }
 
+type Pinger interface {
+	Ping(context.Context) (time.Duration, error)
+}
+
 // InmemTailnetDialer is a tailnet.ControlProtocolDialer that connects to a Coordinator and DERPMap
 // service running in the same memory space.
 type InmemTailnetDialer struct {
-	CoordPtr              *atomic.Pointer[tailnet.Coordinator]
-	DERPFn                func() *tailcfg.DERPMap
-	Logger                slog.Logger
-	ClientID              uuid.UUID
-	DatabaseHealthcheckFn func(ctx context.Context) error
+	CoordPtr *atomic.Pointer[tailnet.Coordinator]
+	DERPFn   func() *tailcfg.DERPMap
+	Logger   slog.Logger
+	ClientID uuid.UUID
+	// DatabaseHealthCheck is used to validate that the store is reachable.
+	DatabaseHealthCheck Pinger
 }
 
 func (a *InmemTailnetDialer) Dial(ctx context.Context, _ tailnet.ResumeTokenController) (tailnet.ControlProtocolClients, error) {
-	if a.DatabaseHealthcheckFn != nil {
-		if err := a.DatabaseHealthcheckFn(ctx); err != nil {
+	if a.DatabaseHealthCheck != nil {
+		if _, err := a.DatabaseHealthCheck.Ping(ctx); err != nil {
 			return tailnet.ControlProtocolClients{}, xerrors.Errorf("%w: %v", codersdk.ErrDatabaseNotReachable, err)
 		}
 	}

diff --git a/coderd/tailnet_test.go b/coderd/tailnet_test.go
@@ -11,6 +11,7 @@ import (
 	"strconv"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
@@ -365,6 +366,14 @@ func TestServerTailnet_ReverseProxy(t *testing.T) {
 	})
 }
 
+type fakePing struct {
+	err error
+}
+
+func (f *fakePing) Ping(context.Context) (time.Duration, error) {
+	return time.Duration(0), f.err
+}
+
 func TestServerTailnet_Healthcheck(t *testing.T) {
 	t.Parallel()
 
@@ -373,9 +382,8 @@ func TestServerTailnet_Healthcheck(t *testing.T) {
 		t.Parallel()
 
 		ctx := testutil.Context(t, testutil.WaitMedium)
-		fn := func(ctx context.Context) error { return nil }
 
-		agents, serverTailnet := setupServerTailnetAgent(t, 1, withHealthcheckFn(fn))
+		agents, serverTailnet := setupServerTailnetAgent(t, 1, withHealthChecker(&fakePing{}))
 
 		a := agents[0]
 		conn, release, err := serverTailnet.AgentConn(ctx, a.id)
@@ -391,9 +399,8 @@ func TestServerTailnet_Healthcheck(t *testing.T) {
 		t.Parallel()
 
 		ctx := testutil.Context(t, testutil.WaitMedium)
-		fn := func(ctx context.Context) error { return xerrors.Errorf("oops, db gone") }
 
-		agents, serverTailnet := setupServerTailnetAgent(t, 1, withHealthcheckFn(fn))
+		agents, serverTailnet := setupServerTailnetAgent(t, 1, withHealthChecker(&fakePing{err: xerrors.New("oops")}))
 
 		a := agents[0]
 		_, release, err := serverTailnet.AgentConn(ctx, a.id)
@@ -427,13 +434,13 @@ type agentWithID struct {
 }
 
 type serverOption struct {
-	HealthcheckFn      func(ctx context.Context) error
+	HealthCheck        coderd.Pinger
 	DERPAndStunOptions []tailnettest.DERPAndStunOption
 }
 
-func withHealthcheckFn(fn func(ctx context.Context) error) serverOption {
+func withHealthChecker(p coderd.Pinger) serverOption {
 	return serverOption{
-		HealthcheckFn: fn,
+		HealthCheck: p,
 	}
 }
 
@@ -446,12 +453,12 @@ func withDERPAndStunOptions(opts ...tailnettest.DERPAndStunOption) serverOption
 func setupServerTailnetAgent(t *testing.T, agentNum int, opts ...serverOption) ([]agentWithID, *coderd.ServerTailnet) {
 	logger := testutil.Logger(t)
 
-	var healthcheckFn func(ctx context.Context) error
+	var healthChecker coderd.Pinger
 	var derpAndStunOptions []tailnettest.DERPAndStunOption
 	for _, opt := range opts {
 		derpAndStunOptions = append(derpAndStunOptions, opt.DERPAndStunOptions...)
-		if opt.HealthcheckFn != nil {
-			healthcheckFn = opt.HealthcheckFn
+		if opt.HealthCheck != nil {
+			healthChecker = opt.HealthCheck
 		}
 	}
 
@@ -495,11 +502,11 @@ func setupServerTailnetAgent(t *testing.T, agentNum int, opts ...serverOption) (
 	}
 
 	dialer := &coderd.InmemTailnetDialer{
-		CoordPtr:              &coordPtr,
-		DERPFn:                func() *tailcfg.DERPMap { return derpMap },
-		Logger:                logger,
-		ClientID:              uuid.UUID{5},
-		DatabaseHealthcheckFn: healthcheckFn,
+		CoordPtr:            &coordPtr,
+		DERPFn:              func() *tailcfg.DERPMap { return derpMap },
+		Logger:              logger,
+		ClientID:            uuid.UUID{5},
+		DatabaseHealthCheck: healthChecker,
 	}
 	serverTailnet, err := coderd.NewServerTailnet(
 		context.Background(),