feat(coderd): insert provisioner daemons (#11207)

* Adds UpdateProvisionerDaemonLastSeenAt
* Adds heartbeat to provisioner daemons
* Inserts provisioner daemons to database upon start
* Ensures TagOwner is an empty string and not nil
* Adds COALESCE() in idx_provisioner_daemons_name_owner_key

coderd/provisionerdserver/provisionerdserver.go

@@ -44,9 +44,15 @@ import (
 	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
 )
 
-// DefaultAcquireJobLongPollDur is the time the (deprecated) AcquireJob rpc waits to try to obtain a job before
-// canceling and returning an empty job.
-const DefaultAcquireJobLongPollDur = time.Second * 5
+const (
+	// DefaultAcquireJobLongPollDur is the time the (deprecated) AcquireJob rpc waits to try to obtain a job before
+	// canceling and returning an empty job.
+	DefaultAcquireJobLongPollDur = time.Second * 5
+
+	// DefaultHeartbeatInterval is the interval at which the provisioner daemon
+	// will update its last seen at timestamp in the database.
+	DefaultHeartbeatInterval = time.Minute
+)
 
 type Options struct {
 	OIDCConfig          httpmw.OAuth2Config
@@ -56,6 +62,16 @@ type Options struct {
 
 	// AcquireJobLongPollDur is used in tests
 	AcquireJobLongPollDur time.Duration
+
+	// HeartbeatInterval is the interval at which the provisioner daemon
+	// will update its last seen at timestamp in the database.
+	HeartbeatInterval time.Duration
+
+	// HeartbeatFn is the function that will be called at the interval
+	// specified by HeartbeatInterval.
+	// The default function just calls UpdateProvisionerDaemonLastSeenAt.
+	// This is mainly used for testing.
+	HeartbeatFn func(context.Context) error
 }
 
 type server struct {
@@ -85,6 +101,9 @@ type server struct {
 	TimeNowFn func() time.Time
 
 	acquireJobLongPollDur time.Duration
+
+	heartbeatInterval time.Duration
+	heartbeatFn       func(ctx context.Context) error
 }
 
 // We use the null byte (0x00) in generating a canonical map key for tags, so
@@ -161,7 +180,21 @@ func NewServer(
 	if options.AcquireJobLongPollDur == 0 {
 		options.AcquireJobLongPollDur = DefaultAcquireJobLongPollDur
 	}
-	return &server{
+	if options.HeartbeatInterval == 0 {
+		options.HeartbeatInterval = DefaultHeartbeatInterval
+	}
+	// Avoid a nil check in s.heartbeat.
+	if options.HeartbeatFn == nil {
+		options.HeartbeatFn = func(hbCtx context.Context) error {
+			//nolint:gocritic // This is specifically for updating the last seen at timestamp.
+			return db.UpdateProvisionerDaemonLastSeenAt(dbauthz.AsSystemRestricted(hbCtx), database.UpdateProvisionerDaemonLastSeenAtParams{
+				ID:         id,
+				LastSeenAt: sql.NullTime{Time: time.Now(), Valid: true},
+			})
+		}
+	}
+
+	s := &server{
 		lifecycleCtx:                lifecycleCtx,
 		AccessURL:                   accessURL,
 		ID:                          id,
@@ -182,7 +215,12 @@ func NewServer(
 		OIDCConfig:                  options.OIDCConfig,
 		TimeNowFn:                   options.TimeNowFn,
 		acquireJobLongPollDur:       options.AcquireJobLongPollDur,
-	}, nil
+		heartbeatInterval:           options.HeartbeatInterval,
+		heartbeatFn:                 options.HeartbeatFn,
+	}
+
+	go s.heartbeatLoop()
+	return s, nil
 }
 
 // timeNow should be used when trying to get the current time for math
@@ -194,6 +232,50 @@ func (s *server) timeNow() time.Time {
 	return dbtime.Now()
 }
 
+// heartbeatLoop runs heartbeatOnce at the interval specified by HeartbeatInterval
+// until the lifecycle context is canceled.
+func (s *server) heartbeatLoop() {
+	tick := time.NewTicker(time.Nanosecond)
+	defer tick.Stop()
+	for {
+		select {
+		case <-s.lifecycleCtx.Done():
+			s.Logger.Debug(s.lifecycleCtx, "heartbeat loop canceled")
+			return
+		case <-tick.C:
+			if s.lifecycleCtx.Err() != nil {
+				return
+			}
+			start := s.timeNow()
+			hbCtx, hbCancel := context.WithTimeout(s.lifecycleCtx, s.heartbeatInterval)
+			if err := s.heartbeat(hbCtx); err != nil {
+				if !xerrors.Is(err, context.DeadlineExceeded) && !xerrors.Is(err, context.Canceled) {
+					s.Logger.Error(hbCtx, "heartbeat failed", slog.Error(err))
+				}
+			}
+			hbCancel()
+			elapsed := s.timeNow().Sub(start)
+			nextBeat := s.heartbeatInterval - elapsed
+			// avoid negative interval
+			if nextBeat <= 0 {
+				nextBeat = time.Nanosecond
+			}
+			tick.Reset(nextBeat)
+		}
+	}
+}
+
+// heartbeat updates the last seen at timestamp in the database.
+// If HeartbeatFn is set, it will be called instead.
+func (s *server) heartbeat(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return nil
+	default:
+		return s.heartbeatFn(ctx)
+	}
+}
+
 // AcquireJob queries the database to lock a job.
 //
 // Deprecated: This method is only available for back-level provisioner daemons.