flvx/go-backend/internal/http/handler/jobs.go

package handler

import (
	"context"
	"database/sql"
	"time"
)

func (h *Handler) StartBackgroundJobs() {
	if h == nil || h.repo == nil || h.repo.DB() == nil {
		return
	}

	h.jobsMu.Lock()
	if h.jobsStarted {
		h.jobsMu.Unlock()
		return
	}
	ctx, cancel := context.WithCancel(context.Background())
	h.jobsCancel = cancel
	h.jobsStarted = true
	h.jobsWG.Add(2)
	h.jobsMu.Unlock()

	go h.runHourlyStatsLoop(ctx)
	go h.runDailyMaintenanceLoop(ctx)
}

func (h *Handler) StopBackgroundJobs() {
	if h == nil {
		return
	}

	h.jobsMu.Lock()
	if !h.jobsStarted {
		h.jobsMu.Unlock()
		return
	}
	cancel := h.jobsCancel
	h.jobsCancel = nil
	h.jobsStarted = false
	h.jobsMu.Unlock()

	if cancel != nil {
		cancel()
	}
	h.jobsWG.Wait()
}

func (h *Handler) runHourlyStatsLoop(ctx context.Context) {
	defer h.jobsWG.Done()

	for {
		wait := durationUntilNextHour(time.Now())
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return
		case <-timer.C:
			h.runStatisticsFlowJob(time.Now())
		}
	}
}

func (h *Handler) runDailyMaintenanceLoop(ctx context.Context) {
	defer h.jobsWG.Done()

	for {
		wait := durationUntilNextDailyMaintenance(time.Now())
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return
		case <-timer.C:
			h.runResetAndExpiryJob(time.Now())
		}
	}
}

func durationUntilNextHour(now time.Time) time.Duration {
	next := now.Truncate(time.Hour).Add(time.Hour)
	return next.Sub(now)
}

func durationUntilNextDailyMaintenance(now time.Time) time.Duration {
	next := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 5, 0, now.Location())
	if !next.After(now) {
		next = next.Add(24 * time.Hour)
	}
	return next.Sub(now)
}

func (h *Handler) runStatisticsFlowJob(now time.Time) {
	if h == nil || h.repo == nil || h.repo.DB() == nil {
		return
	}

	db := h.repo.DB()
	nowMs := now.UnixMilli()
	cutoffMs := nowMs - int64((48*time.Hour)/time.Millisecond)
	_, _ = db.Exec(`DELETE FROM statistics_flow WHERE created_time < ?`, cutoffMs)

	hourMark := now.Truncate(time.Hour)
	hourText := hourMark.Format("15:04")
	createdTime := hourMark.UnixMilli()

	rows, err := db.Query(`SELECT id, in_flow, out_flow FROM user ORDER BY id ASC`)
	if err != nil {
		return
	}
	type userFlowSnapshot struct {
		userID  int64
		inFlow  int64
		outFlow int64
	}
	users := make([]userFlowSnapshot, 0)

	for rows.Next() {
		var userID int64
		var inFlow int64
		var outFlow int64
		if err := rows.Scan(&userID, &inFlow, &outFlow); err != nil {
			continue
		}
		users = append(users, userFlowSnapshot{userID: userID, inFlow: inFlow, outFlow: outFlow})
	}
	_ = rows.Close()

	for _, user := range users {
		currentTotal := user.inFlow + user.outFlow
		increment := currentTotal

		var lastTotal sql.NullInt64
		err := db.QueryRow(`SELECT total_flow FROM statistics_flow WHERE user_id = ? ORDER BY id DESC LIMIT 1`, user.userID).Scan(&lastTotal)
		if err == nil && lastTotal.Valid {
			increment = currentTotal - lastTotal.Int64
			if increment < 0 {
				increment = currentTotal
			}
		}

		_, _ = db.Exec(`
			INSERT INTO statistics_flow(user_id, flow, total_flow, time, created_time)
			VALUES(?, ?, ?, ?, ?)
		`, user.userID, increment, currentTotal, hourText, createdTime)
	}
}

func (h *Handler) runResetAndExpiryJob(now time.Time) {
	if h == nil || h.repo == nil || h.repo.DB() == nil {
		return
	}

	h.resetMonthlyFlow(now)
	h.disableExpiredUsers(now.UnixMilli())
	h.disableExpiredUserTunnels(now.UnixMilli())
}

func (h *Handler) resetMonthlyFlow(now time.Time) {
	db := h.repo.DB()
	currentDay := now.Day()
	lastDay := time.Date(now.Year(), now.Month()+1, 0, 0, 0, 0, 0, now.Location()).Day()

	if currentDay == lastDay {
		_, _ = db.Exec(`
			UPDATE user
			SET in_flow = 0, out_flow = 0
			WHERE flow_reset_time != 0
			  AND (flow_reset_time = ? OR flow_reset_time > ?)
		`, currentDay, lastDay)
		_, _ = db.Exec(`
			UPDATE user_tunnel
			SET in_flow = 0, out_flow = 0
			WHERE flow_reset_time != 0
			  AND (flow_reset_time = ? OR flow_reset_time > ?)
		`, currentDay, lastDay)
		return
	}

	_, _ = db.Exec(`
		UPDATE user
		SET in_flow = 0, out_flow = 0
		WHERE flow_reset_time != 0
		  AND flow_reset_time = ?
	`, currentDay)
	_, _ = db.Exec(`
		UPDATE user_tunnel
		SET in_flow = 0, out_flow = 0
		WHERE flow_reset_time != 0
		  AND flow_reset_time = ?
	`, currentDay)
}

func (h *Handler) disableExpiredUsers(nowMs int64) {
	db := h.repo.DB()
	rows, err := db.Query(`
		SELECT id
		FROM user
		WHERE role_id != 0
		  AND status = 1
		  AND exp_time IS NOT NULL
		  AND exp_time < ?
	`, nowMs)
	if err != nil {
		return
	}
	userIDs := make([]int64, 0)

	for rows.Next() {
		var userID int64
		if err := rows.Scan(&userID); err != nil {
			continue
		}
		userIDs = append(userIDs, userID)
	}
	_ = rows.Close()

	for _, userID := range userIDs {
		forwards, err := h.listActiveForwardsByUser(userID)
		if err == nil {
			h.pauseForwardRecords(forwards, nowMs)
		}
		_, _ = db.Exec(`UPDATE user SET status = 0 WHERE id = ?`, userID)
	}
}

func (h *Handler) disableExpiredUserTunnels(nowMs int64) {
	db := h.repo.DB()
	rows, err := db.Query(`
		SELECT id, user_id, tunnel_id
		FROM user_tunnel
		WHERE status = 1
		  AND exp_time IS NOT NULL
		  AND exp_time < ?
	`, nowMs)
	if err != nil {
		return
	}
	type expiredUserTunnel struct {
		userTunnelID int64
		userID       int64
		tunnelID     int64
	}
	items := make([]expiredUserTunnel, 0)

	for rows.Next() {
		var userTunnelID int64
		var userID int64
		var tunnelID int64
		if err := rows.Scan(&userTunnelID, &userID, &tunnelID); err != nil {
			continue
		}
		items = append(items, expiredUserTunnel{userTunnelID: userTunnelID, userID: userID, tunnelID: tunnelID})
	}
	_ = rows.Close()

	for _, item := range items {
		forwards, err := h.listActiveForwardsByUserTunnel(item.userID, item.tunnelID)
		if err == nil {
			h.pauseForwardRecords(forwards, nowMs)
		}
		_, _ = db.Exec(`UPDATE user_tunnel SET status = 0 WHERE id = ?`, item.userTunnelID)
	}
}