@@ -18,6 +18,7 @@ import (
1818 "github.com/stretchr/testify/assert"
1919 "github.com/stretchr/testify/require"
2020 "go.opentelemetry.io/otel/trace"
21+ "golang.org/x/xerrors"
2122 "tailscale.com/tailcfg"
2223
2324 "github.com/coder/coder/v2/agent"
@@ -56,8 +57,7 @@ func TestServerTailnet_AgentConn_NoSTUN(t *testing.T) {
5657 defer cancel ()
5758
5859 // Connect through the ServerTailnet
59- agents , serverTailnet := setupServerTailnetAgent (t , 1 ,
60- tailnettest .DisableSTUN , tailnettest .DERPIsEmbedded )
60+ agents , serverTailnet := setupServerTailnetAgent (t , 1 , withDERPAndStunOptions (tailnettest .DisableSTUN , tailnettest .DERPIsEmbedded ))
6161 a := agents [0 ]
6262
6363 conn , release , err := serverTailnet .AgentConn (ctx , a .id )
@@ -340,7 +340,7 @@ func TestServerTailnet_ReverseProxy(t *testing.T) {
340340 ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
341341 defer cancel ()
342342
343- agents , serverTailnet := setupServerTailnetAgent (t , 1 , tailnettest .DisableSTUN )
343+ agents , serverTailnet := setupServerTailnetAgent (t , 1 , withDERPAndStunOptions ( tailnettest .DisableSTUN ) )
344344 a := agents [0 ]
345345
346346 require .True (t , serverTailnet .Conn ().GetBlockEndpoints (), "expected BlockEndpoints to be set" )
@@ -365,6 +365,43 @@ func TestServerTailnet_ReverseProxy(t *testing.T) {
365365 })
366366}
367367
368+ func TestServerTailnet_Healthcheck (t * testing.T ) {
369+ t .Parallel ()
370+
371+ // Verifies that a non-nil healthcheck which returns a non-error response behaves as expected.
372+ t .Run ("Passing" , func (t * testing.T ) {
373+ t .Parallel ()
374+
375+ ctx := testutil .Context (t , testutil .WaitMedium )
376+ fn := func (ctx context.Context ) error { return nil }
377+
378+ agents , serverTailnet := setupServerTailnetAgent (t , 1 , withHealthcheckFn (fn ))
379+
380+ a := agents [0 ]
381+ conn , release , err := serverTailnet .AgentConn (ctx , a .id )
382+ t .Cleanup (release )
383+ require .NoError (t , err )
384+ assert .True (t , conn .AwaitReachable (ctx ))
385+ })
386+
387+ // If the healthcheck fails, we have no insight into this at this level.
388+ // The dial against the control plane is retried, so we wait for the context to timeout as an indication that the
389+ // healthcheck is performing as expected.
390+ t .Run ("Failing" , func (t * testing.T ) {
391+ t .Parallel ()
392+
393+ ctx := testutil .Context (t , testutil .WaitMedium )
394+ fn := func (ctx context.Context ) error { return xerrors .Errorf ("oops, db gone" ) }
395+
396+ agents , serverTailnet := setupServerTailnetAgent (t , 1 , withHealthcheckFn (fn ))
397+
398+ a := agents [0 ]
399+ _ , release , err := serverTailnet .AgentConn (ctx , a .id )
400+ require .Nil (t , release )
401+ require .ErrorContains (t , err , "agent is unreachable" )
402+ })
403+ }
404+
368405type wrappedListener struct {
369406 net.Listener
370407 dials int32
@@ -389,9 +426,36 @@ type agentWithID struct {
389426 agent.Agent
390427}
391428
392- func setupServerTailnetAgent (t * testing.T , agentNum int , opts ... tailnettest.DERPAndStunOption ) ([]agentWithID , * coderd.ServerTailnet ) {
429+ type serverOption struct {
430+ HealthcheckFn func (ctx context.Context ) error
431+ DERPAndStunOptions []tailnettest.DERPAndStunOption
432+ }
433+
434+ func withHealthcheckFn (fn func (ctx context.Context ) error ) serverOption {
435+ return serverOption {
436+ HealthcheckFn : fn ,
437+ }
438+ }
439+
440+ func withDERPAndStunOptions (opts ... tailnettest.DERPAndStunOption ) serverOption {
441+ return serverOption {
442+ DERPAndStunOptions : opts ,
443+ }
444+ }
445+
446+ func setupServerTailnetAgent (t * testing.T , agentNum int , opts ... serverOption ) ([]agentWithID , * coderd.ServerTailnet ) {
393447 logger := testutil .Logger (t )
394- derpMap , derpServer := tailnettest .RunDERPAndSTUN (t , opts ... )
448+
449+ var healthcheckFn func (ctx context.Context ) error
450+ var derpAndStunOptions []tailnettest.DERPAndStunOption
451+ for _ , opt := range opts {
452+ derpAndStunOptions = append (derpAndStunOptions , opt .DERPAndStunOptions ... )
453+ if opt .HealthcheckFn != nil {
454+ healthcheckFn = opt .HealthcheckFn
455+ }
456+ }
457+
458+ derpMap , derpServer := tailnettest .RunDERPAndSTUN (t , derpAndStunOptions ... )
395459
396460 coord := tailnet .NewCoordinator (logger )
397461 t .Cleanup (func () {
@@ -431,10 +495,11 @@ func setupServerTailnetAgent(t *testing.T, agentNum int, opts ...tailnettest.DER
431495 }
432496
433497 dialer := & coderd.InmemTailnetDialer {
434- CoordPtr : & coordPtr ,
435- DERPFn : func () * tailcfg.DERPMap { return derpMap },
436- Logger : logger ,
437- ClientID : uuid.UUID {5 },
498+ CoordPtr : & coordPtr ,
499+ DERPFn : func () * tailcfg.DERPMap { return derpMap },
500+ Logger : logger ,
501+ ClientID : uuid.UUID {5 },
502+ DatabaseHealthcheckFn : healthcheckFn ,
438503 }
439504 serverTailnet , err := coderd .NewServerTailnet (
440505 context .Background (),
0 commit comments