feat: add periodic cleanup of PG Coordinator state (#8142)

* PG Coordinator cleans orphaned state

Signed-off-by: Spike Curtis <spike@coder.com>

* Don't need pubsub

Signed-off-by: Spike Curtis <spike@coder.com>

---------

Signed-off-by: Spike Curtis <spike@coder.com>

coderd/database/dbauthz/dbauthz.go

@@ -683,6 +683,13 @@ func (q *querier) AcquireProvisionerJob(ctx context.Context, arg database.Acquir
 	return q.db.AcquireProvisionerJob(ctx, arg)
 }
 
+func (q *querier) CleanTailnetCoordinators(ctx context.Context) error {
+	if err := q.authorizeContext(ctx, rbac.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
+		return err
+	}
+	return q.db.CleanTailnetCoordinators(ctx)
+}
+
 func (q *querier) DeleteAPIKeyByID(ctx context.Context, id string) error {
 	return deleteQ(q.log, q.auth, q.db.GetAPIKeyByID, q.db.DeleteAPIKeyByID)(ctx, id)
 }

coderd/database/dbfake/dbfake.go

@@ -1035,6 +1035,10 @@ func (q *fakeQuerier) AcquireProvisionerJob(_ context.Context, arg database.Acqu
 	return database.ProvisionerJob{}, sql.ErrNoRows
 }
 
+func (*fakeQuerier) CleanTailnetCoordinators(_ context.Context) error {
+	return ErrUnimplemented
+}
+
 func (q *fakeQuerier) DeleteAPIKeyByID(_ context.Context, id string) error {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()

coderd/database/dbmetrics/dbmetrics.go

@@ -122,6 +122,13 @@ func (m metricsStore) AcquireProvisionerJob(ctx context.Context, arg database.Ac
 	return provisionerJob, err
 }
 
+func (m metricsStore) CleanTailnetCoordinators(ctx context.Context) error {
+	start := time.Now()
+	err := m.s.CleanTailnetCoordinators(ctx)
+	m.queryLatencies.WithLabelValues("CleanTailnetCoordinators").Observe(time.Since(start).Seconds())
+	return err
+}
+
 func (m metricsStore) DeleteAPIKeyByID(ctx context.Context, id string) error {
 	start := time.Now()
 	err := m.s.DeleteAPIKeyByID(ctx, id)

coderd/database/dbmock/dbmock.go

@@ -68,6 +68,20 @@ func (mr *MockStoreMockRecorder) AcquireProvisionerJob(arg0, arg1 interface{}) *
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcquireProvisionerJob", reflect.TypeOf((*MockStore)(nil).AcquireProvisionerJob), arg0, arg1)
 }
 
+// CleanTailnetCoordinators mocks base method.
+func (m *MockStore) CleanTailnetCoordinators(arg0 context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CleanTailnetCoordinators", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// CleanTailnetCoordinators indicates an expected call of CleanTailnetCoordinators.
+func (mr *MockStoreMockRecorder) CleanTailnetCoordinators(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CleanTailnetCoordinators", reflect.TypeOf((*MockStore)(nil).CleanTailnetCoordinators), arg0)
+}
+
 // DeleteAPIKeyByID mocks base method.
 func (m *MockStore) DeleteAPIKeyByID(arg0 context.Context, arg1 string) error {
 	m.ctrl.T.Helper()

coderd/database/querier.go

@@ -26,6 +26,7 @@ type sqlcQuerier interface {
 	// multiple provisioners from acquiring the same jobs. See:
 	// https://www.postgresql.org/docs/9.5/sql-select.html#SQL-FOR-UPDATE-SHARE
 	AcquireProvisionerJob(ctx context.Context, arg AcquireProvisionerJobParams) (ProvisionerJob, error)
+	CleanTailnetCoordinators(ctx context.Context) error
 	DeleteAPIKeyByID(ctx context.Context, id string) error
 	DeleteAPIKeysByUserID(ctx context.Context, userID uuid.UUID) error
 	DeleteApplicationConnectAPIKeysByUserID(ctx context.Context, userID uuid.UUID) error

coderd/database/queries.sql.go

@@ -3261,6 +3261,17 @@ func (q *sqlQuerier) UpsertServiceBanner(ctx context.Context, value string) erro
 	return err
 }
 
+const cleanTailnetCoordinators = `-- name: CleanTailnetCoordinators :exec
+DELETE
+FROM tailnet_coordinators
+WHERE heartbeat_at < now() - INTERVAL '24 HOURS'
+`
+
+func (q *sqlQuerier) CleanTailnetCoordinators(ctx context.Context) error {
+	_, err := q.db.ExecContext(ctx, cleanTailnetCoordinators)
+	return err
+}
+
 const deleteCoordinator = `-- name: DeleteCoordinator :exec
 DELETE
 FROM tailnet_coordinators

coderd/database/queries/tailnet.sql

@@ -77,3 +77,8 @@ DO UPDATE SET
   id = $1,
   heartbeat_at = now() at time zone 'utc'
 RETURNING *;
+
+-- name: CleanTailnetCoordinators :exec
+DELETE
+FROM tailnet_coordinators
+WHERE heartbeat_at < now() - INTERVAL '24 HOURS';

enterprise/tailnet/pgcoord.go

@@ -31,6 +31,7 @@ const (
 	numQuerierWorkers = 10
 	numBinderWorkers  = 10
 	dbMaxBackoff      = 10 * time.Second
+	cleanupPeriod     = time.Hour
 )
 
 // pgCoord is a postgres-backed coordinator
@@ -1041,6 +1042,9 @@ type heartbeats struct {
 	lock         sync.RWMutex
 	coordinators map[uuid.UUID]time.Time
 	timer        *time.Timer
+
+	// overwritten in tests, but otherwise constant
+	cleanupPeriod time.Duration
 }
 
 func newHeartbeats(
@@ -1058,9 +1062,11 @@ func newHeartbeats(
 		update:         update,
 		firstHeartbeat: firstHeartbeat,
 		coordinators:   make(map[uuid.UUID]time.Time),
+		cleanupPeriod:  cleanupPeriod,
 	}
 	go h.subscribe()
 	go h.sendBeats()
+	go h.cleanupLoop()
 	return h
 }
 
@@ -1211,3 +1217,31 @@ func (h *heartbeats) sendDelete() {
 	}
 	h.logger.Debug(h.ctx, "deleted coordinator")
 }
+
+func (h *heartbeats) cleanupLoop() {
+	h.cleanup()
+	tkr := time.NewTicker(h.cleanupPeriod)
+	defer tkr.Stop()
+	for {
+		select {
+		case <-h.ctx.Done():
+			h.logger.Debug(h.ctx, "ending cleanupLoop", slog.Error(h.ctx.Err()))
+			return
+		case <-tkr.C:
+			h.cleanup()
+		}
+	}
+}
+
+// cleanup issues a DB command to clean out any old expired coordinators state.  The cleanup is idempotent, so no need
+// to synchronize with other coordinators.
+func (h *heartbeats) cleanup() {
+	err := h.store.CleanTailnetCoordinators(h.ctx)
+	if err != nil {
+		// the records we are attempting to clean up do no serious harm other than
+		// accumulating in the tables, so we don't bother retrying if it fails.
+		h.logger.Error(h.ctx, "failed to cleanup old coordinators", slog.Error(err))
+		return
+	}
+	h.logger.Debug(h.ctx, "cleaned up old coordinators")
+}

enterprise/tailnet/pgcoord_internal_test.go

@@ -0,0 +1,52 @@
+package tailnet
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+
+	"github.com/coder/coder/coderd/database/dbmock"
+	"github.com/coder/coder/testutil"
+)
+
+// TestHeartbeat_Cleanup is internal so that we can overwrite the cleanup period and not wait an hour for the timed
+// cleanup.
+func TestHeartbeat_Cleanup(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+	mStore := dbmock.NewMockStore(ctrl)
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+
+	waitForCleanup := make(chan struct{})
+	mStore.EXPECT().CleanTailnetCoordinators(gomock.Any()).MinTimes(2).DoAndReturn(func(_ context.Context) error {
+		<-waitForCleanup
+		return nil
+	})
+
+	uut := &heartbeats{
+		ctx:           ctx,
+		logger:        logger,
+		store:         mStore,
+		cleanupPeriod: time.Millisecond,
+	}
+	go uut.cleanupLoop()
+
+	for i := 0; i < 2; i++ {
+		select {
+		case <-ctx.Done():
+			t.Fatal("timeout")
+		case waitForCleanup <- struct{}{}:
+			// ok
+		}
+	}
+	close(waitForCleanup)
+}