lingniu-vehicle-ingest/go/vehicle-gateway/cmd/stats-backfill/customer_hydrogen_worker.go

package main

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"errors"
	"log/slog"
	"os"
	"os/signal"
	"syscall"
	"time"
)

func runCustomerHydrogenWorker(parent context.Context, mysqlDB, td *sql.DB, cfg config) error {
	ctx, stop := signal.NotifyContext(parent, os.Interrupt, syscall.SIGTERM)
	defer stop()
	maxJobs := envInt("CUSTOMER_HYDROGEN_MAX_JOBS", 0)
	completedJobs := 0
	for ctx.Err() == nil {
		date, owner, err := claimCustomerHydrogenDay(ctx, mysqlDB)
		if err != nil && !errors.Is(err, sql.ErrNoRows) {
			slog.Error("customer hydrogen claim failed", "error", err)
		}
		if err != nil {
			select {
			case <-ctx.Done():
				return nil
			case <-time.After(5 * time.Second):
				continue
			}
		}
		dayCtx, cancel := context.WithTimeout(ctx, 2*time.Hour)
		heartbeatDone := make(chan struct{})
		go func() {
			defer close(heartbeatDone)
			ticker := time.NewTicker(30 * time.Second)
			defer ticker.Stop()
			for {
				select {
				case <-dayCtx.Done():
					return
				case <-ticker.C:
					result, e := mysqlDB.ExecContext(dayCtx, `UPDATE platform_customer_hydrogen_job SET lease_until=DATE_ADD(NOW(3),INTERVAL 5 MINUTE) WHERE stat_date=? AND lease_owner=? AND status='running'`, date, owner)
					if e != nil {
						cancel()
						return
					}
					n, e := result.RowsAffected()
					if e != nil || n != 1 {
						cancel()
						return
					}
				}
			}
		}()
		task := cfg
		task.DateFrom = date
		task.DateTo = date
		task.Limit = 0
		task.DryRun = false
		task.Reset = false
		task.VINs = nil
		slog.Info("customer hydrogen job started", "date", date)
		err = runCustomerHydrogenBackfillWithLease(dayCtx, mysqlDB, td, task, owner)
		cancel()
		<-heartbeatDone
		status, message := "complete", ""
		if err != nil {
			status, message = "failed", "原始报文计算失败，请检查任务日志"
			slog.Error("customer hydrogen job failed", "date", date, "error", err)
		}
		// A cancelled/expired worker cannot acknowledge another worker's lease.
		finishCtx, finishCancel := context.WithTimeout(context.Background(), 5*time.Second)
		_, saveErr := mysqlDB.ExecContext(finishCtx, `UPDATE platform_customer_hydrogen_job SET status=?,error_message=?,completed_at=IF(?='complete',NOW(3),completed_at),lease_until=NULL,lease_owner='' WHERE stat_date=? AND lease_owner=?`, status, message, status, date, owner)
		finishCancel()
		if saveErr != nil {
			slog.Error("customer hydrogen job completion failed", "date", date, "error", saveErr)
		}
		completedJobs++
		if maxJobs > 0 && completedJobs >= maxJobs {
			return nil
		}
	}
	return nil
}

func claimCustomerHydrogenDay(ctx context.Context, db *sql.DB) (string, string, error) {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return "", "", err
	}
	defer tx.Rollback()
	var date string
	err = tx.QueryRowContext(ctx, `SELECT DATE_FORMAT(stat_date,'%Y-%m-%d') FROM platform_customer_hydrogen_job WHERE status='pending' OR (status='running' AND lease_until<NOW(3)) OR (status='failed' AND attempts<3 AND started_at<DATE_SUB(NOW(3),INTERVAL 1 MINUTE)) ORDER BY requested_at ASC,stat_date DESC LIMIT 1 FOR UPDATE`).Scan(&date)
	if err != nil {
		return "", "", err
	}
	var bytes [16]byte
	if _, err = rand.Read(bytes[:]); err != nil {
		return "", "", err
	}
	owner := hex.EncodeToString(bytes[:])
	_, err = tx.ExecContext(ctx, `UPDATE platform_customer_hydrogen_job SET status='running',started_at=NOW(3),lease_until=DATE_ADD(NOW(3),INTERVAL 5 MINUTE),lease_owner=?,attempts=attempts+1,error_message='' WHERE stat_date=?`, owner, date)
	if err != nil {
		return "", "", err
	}
	if err = tx.Commit(); err != nil {
		return "", "", err
	}
	return date, owner, nil
}