Cleanup old workers via daily gocron (#3663)

* cleanup zombie workers and query use index

* cleanup job rewrite

* retention period

* no advisory lock

Author: mnafees

internal/services/controllers/retention/controller.go

@@ -175,6 +175,23 @@ func (rc *RetentionControllerImpl) Start() (func() error, error) {
 		}
 	}
 
+	if rc.workerRetention {
+		workerInterval := 24 * time.Hour
+
+		_, err := rc.s.NewJob(
+			gocron.DurationJob(workerInterval),
+			gocron.NewTask(
+				rc.runCleanupOldWorkers(ctx),
+			),
+			gocron.WithSingletonMode(gocron.LimitModeReschedule),
+		)
+
+		if err != nil {
+			cancel()
+			return nil, fmt.Errorf("could not set up runCleanupOldWorkers: %w", err)
+		}
+	}
+
 	rc.s.Start()
 
 	cleanup := func() error {

internal/services/controllers/retention/worker.go

@@ -0,0 +1,48 @@
+package retention
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
+	"github.com/hatchet-dev/hatchet/pkg/telemetry"
+)
+
+func (rc *RetentionControllerImpl) runCleanupOldWorkers(ctx context.Context) func() {
+	return func() {
+		rc.l.Debug().Ctx(ctx).Msg("retention controller: cleaning up old workers")
+
+		ctx, cancel := context.WithTimeout(ctx, 30*time.Minute)
+		defer cancel()
+
+		if err := rc.ForTenants(ctx, rc.cleanupOldWorkersForTenant); err != nil {
+			rc.l.Err(err).Ctx(ctx).Msg("could not cleanup old workers")
+		}
+	}
+}
+
+func (rc *RetentionControllerImpl) cleanupOldWorkersForTenant(ctx context.Context, tenant sqlcv1.Tenant) error {
+	ctx, span := telemetry.NewSpan(ctx, "cleanup-old-workers-tenant")
+	defer span.End()
+
+	cutoff, err := GetDataRetentionExpiredTime(tenant.DataRetentionPeriod)
+	if err != nil {
+		return fmt.Errorf("could not get cutoff for tenant %s: %w", tenant.ID.String(), err)
+	}
+
+	shouldContinue := true
+
+	for shouldContinue {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
+		shouldContinue, err = rc.repo.Workers().CleanupOldWorkers(ctx, tenant.ID, cutoff)
+		if err != nil {
+			return fmt.Errorf("could not cleanup old workers for tenant %s: %w", tenant.ID.String(), err)
+		}
+	}
+
+	return nil
+}

pkg/repository/sqlcv1/tasks.sql

@@ -425,6 +425,7 @@ WITH tasks_on_inactive_workers AS (
         v1_task_runtime runtime ON w."id" = runtime.worker_id
     WHERE
         w."tenantId" = @tenantId::uuid
+        AND w."tenantId" = runtime.tenant_id
         AND w."lastHeartbeatAt" < NOW() - INTERVAL '30 seconds'
         -- evicted tasks are not eligible for re-assignment
         AND runtime.evicted_at IS NULL

pkg/repository/sqlcv1/tasks.sql.go

@@ -379,36 +379,16 @@ SET
     "strValue" = sqlc.narg('strValue')::text
 RETURNING *;
 
--- name: DeleteOldWorkers :one
-WITH for_delete AS (
-    SELECT
-        "id"
-    FROM "Worker" w
-    WHERE
-        w."tenantId" = @tenantId::uuid AND
-        w."lastHeartbeatAt" < @lastHeartbeatBefore::timestamp
-    LIMIT sqlc.arg('limit') + 1
-), expired_with_limit AS (
-    SELECT
-        for_delete."id" as "id"
-    FROM for_delete
-    LIMIT sqlc.arg('limit')
-), has_more AS (
-    SELECT
-        CASE
-            WHEN COUNT(*) > sqlc.arg('limit') THEN TRUE
-            ELSE FALSE
-        END as has_more
-    FROM for_delete
-), delete_events AS (
-    DELETE FROM "WorkerAssignEvent" wae
-    WHERE wae."workerId" IN (SELECT "id" FROM expired_with_limit)
-    RETURNING wae."id"
+-- name: CleanupOldWorkers :execresult
+WITH old_workers AS (
+    SELECT "id"
+    FROM "Worker"
+    WHERE "tenantId" = @tenantId::uuid
+      AND "lastHeartbeatAt" < @lastHeartbeatBefore::timestamp
+    LIMIT @batchSize::int
 )
-DELETE FROM "Worker" w
-WHERE w."id" IN (SELECT "id" FROM expired_with_limit)
-RETURNING
-    (SELECT has_more FROM has_more) as has_more;
+DELETE FROM "Worker"
+WHERE "id" IN (SELECT "id" FROM old_workers);
 
 -- name: ListDispatcherIdsForWorkers :many
 SELECT

pkg/repository/sqlcv1/workers.sql

@@ -124,7 +124,7 @@ type WorkerRepository interface {
 
 	UpsertWorkerLabels(ctx context.Context, workerId uuid.UUID, opts []UpsertWorkerLabelOpts) ([]*sqlcv1.WorkerLabel, error)
 
-	DeleteOldWorkers(ctx context.Context, tenantId uuid.UUID, lastHeartbeatBefore time.Time) (bool, error)
+	CleanupOldWorkers(ctx context.Context, tenantId uuid.UUID, lastHeartbeatBefore time.Time) (bool, error)
 
 	GetDispatcherIdsForWorkers(ctx context.Context, tenantId uuid.UUID, workerIds []uuid.UUID) (map[uuid.UUID]uuid.UUID, map[uuid.UUID]struct{}, error)
 
@@ -723,22 +723,30 @@ func (w *workerRepository) UpsertWorkerLabels(ctx context.Context, workerId uuid
 	return affinities, nil
 }
 
-func (w *workerRepository) DeleteOldWorkers(ctx context.Context, tenantId uuid.UUID, lastHeartbeatBefore time.Time) (bool, error) {
-	hasMore, err := w.queries.DeleteOldWorkers(ctx, w.pool, sqlcv1.DeleteOldWorkersParams{
+func (w *workerRepository) CleanupOldWorkers(ctx context.Context, tenantId uuid.UUID, lastHeartbeatBefore time.Time) (bool, error) {
+	const timeout = 1000 * 60 * 3 // 3 minutes
+	const batchSize int32 = 10000
+
+	tx, commit, rollback, err := sqlchelpers.PrepareTxWithStatementTimeout(ctx, w.pool, w.l, timeout)
+	if err != nil {
+		return false, fmt.Errorf("error beginning transaction: %w", err)
+	}
+	defer rollback()
+
+	result, err := w.queries.CleanupOldWorkers(ctx, tx, sqlcv1.CleanupOldWorkersParams{
 		Tenantid:            tenantId,
 		Lastheartbeatbefore: sqlchelpers.TimestampFromTime(lastHeartbeatBefore),
-		Limit:               20,
+		Batchsize:           batchSize,
 	})
-
 	if err != nil {
-		if errors.Is(err, pgx.ErrNoRows) {
-			return false, nil
-		}
+		return false, fmt.Errorf("error cleaning up old workers: %w", err)
+	}
 
-		return false, err
+	if err := commit(ctx); err != nil {
+		return false, fmt.Errorf("error committing transaction: %w", err)
 	}
 
-	return hasMore, nil
+	return result.RowsAffected() == int64(batchSize), nil
 }
 
 func (w *workerRepository) GetDispatcherIdsForWorkers(ctx context.Context, tenantId uuid.UUID, workerIds []uuid.UUID) (map[uuid.UUID]uuid.UUID, map[uuid.UUID]struct{}, error) {