Merge pull request #7207 from HuntercodeT/fix/group-usage-summary-seq-scan

fix: 分组用量汇总不再全表扫 usage_logs（48.7s → 34ms）
2026-10-14 12:02:27 +08:00 · 2026-09-17 16:37:21 +08:00
parent 44831bb854 18d483c2a1
commit 7e4164c1cf
2 changed files with 147 additions and 33 deletions
@@ -10,6 +10,62 @@ import (
 	"github.com/Wei-Shaw/sub2api/internal/service"
 )

+// groupUsageRollupSnapshot 是一次汇总水位的快照。
+// closedBefore / retainedFrom 在水位无效时退化为 1970-01-01，
+// 与此前 SQL 里 CASE WHEN valid 的语义一一对应。
+type groupUsageRollupSnapshot struct {
+	valid        bool
+	closedBefore string    // date，历史日桶的右开边界
+	retainedDate string    // retained_from 在服务端时区内的日期，历史日桶的左闭边界
+	tailStart    time.Time // 尾段（还没进日桶的那部分 usage_logs）的起点
+}
+
+// readGroupUsageRollupSnapshot 读一次汇总水位并在 Go 侧判定有效性。
+//
+// 判定条件与原先 SQL 里的 state_values CTE 完全一致：
+// 恰好一行、时区名与当前服务端配置相同、且水位不在未来。
+func (r *usageLogRepository) readGroupUsageRollupSnapshot(ctx context.Context, timezoneName, todayDate string) (groupUsageRollupSnapshot, error) {
+	epoch := time.Unix(0, 0).UTC()
+	invalid := groupUsageRollupSnapshot{
+		closedBefore: "1970-01-01",
+		retainedDate: "1970-01-01",
+		tailStart:    epoch,
+	}
+
+	var rowCount int
+	var closedBefore sql.NullString
+	var retainedFrom sql.NullTime
+	var stateTimezone sql.NullString
+	if err := scanSingleRow(ctx, r.sql, `
+		SELECT
+			COUNT(*),
+			MAX(closed_before)::text,
+			MAX(retained_from),
+			MAX(timezone_name)
+		FROM usage_group_rollup_state
+		WHERE id = 1
+	`, nil, &rowCount, &closedBefore, &retainedFrom, &stateTimezone); err != nil {
+		return groupUsageRollupSnapshot{}, fmt.Errorf("读取分组用量汇总水位: %w", err)
+	}
+	if rowCount != 1 || !closedBefore.Valid || !retainedFrom.Valid ||
+		!stateTimezone.Valid || stateTimezone.String != timezoneName ||
+		closedBefore.String > todayDate {
+		return invalid, nil
+	}
+
+	tailStart, err := service.ParseGroupUsageDate(closedBefore.String)
+	if err != nil {
+		// 水位值解析不出来，按无效处理：宁可多扫一次，也不要算错钱。
+		return invalid, nil //nolint:nilerr // 与上面的 valid 判定同语义，均降级为全量重算
+	}
+	return groupUsageRollupSnapshot{
+		valid:        true,
+		closedBefore: closedBefore.String,
+		retainedDate: service.GroupUsageDate(retainedFrom.Time),
+		tailStart:    tailStart.UTC(),
+	}, nil
+}
+
 func (r *usageLogRepository) getAllGroupUsageSummaryFromRollups(ctx context.Context, todayStart time.Time) (results []usagestats.GroupUsageSummary, err error) {
 	todayStart = service.GroupUsageTodayStart(todayStart)
 	yesterdayStart := service.GroupUsageYesterdayStart(todayStart)
@@ -17,40 +73,34 @@ func (r *usageLogRepository) getAllGroupUsageSummaryFromRollups(ctx context.Cont
 	todayDate := service.GroupUsageDate(todayStart)
 	yesterdayDate := service.GroupUsageDate(yesterdayStart)

+	state, err := r.readGroupUsageRollupSnapshot(ctx, timezoneName, todayDate)
+	if err != nil {
+		return nil, err
+	}
+
+	// 尾段起点必须以**查询参数**的形式给进来。
+	//
+	// 此前它是在同一条 SQL 里由 state CTE 算出、再 CROSS JOIN 给 tail 用的，
+	// 于是 created_at 的下界对 planner 来说是个运行期才知道的值：既进不了
+	// index cond，也估不准选择性，只能退化成 usage_logs 全表扫。
+	// 生产实测（1721 万行 / PostgreSQL 17）：
+	//   Seq Scan on usage_logs … rows=17212870, Rows Removed by Join Filter: 17178082
+	//   Execution Time: 48720 ms
+	// 换成参数之后走 idx_usage_logs_created_at：
+	//   Index Scan … Index Cond: (created_at >= $7)
+	//   Execution Time: 34 ms
 	const query = `
-		WITH state_values AS (
-			SELECT
-				COUNT(*) = 1
-					AND MAX(timezone_name) = $3
-					AND MAX(closed_before) <= $4::date AS valid,
-				MAX(closed_before) AS closed_before,
-				MAX(retained_from) AS retained_from
-			FROM usage_group_rollup_state
-			WHERE id = 1
-		),
-		state AS (
-			SELECT
-				CASE WHEN valid THEN closed_before ELSE DATE '1970-01-01' END AS closed_before,
-				CASE WHEN valid THEN retained_from ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00' END AS retained_from,
-				CASE
-					WHEN valid THEN closed_before::timestamp AT TIME ZONE $3::text
-					ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00'
-				END AS tail_start,
-				valid
-			FROM state_values
-		),
-		historical AS (
+		WITH historical AS (
 			SELECT
 				rollup.group_id,
 				COALESCE(SUM(rollup.actual_cost), 0) AS actual_cost,
 				COALESCE(SUM(rollup.actual_cost) FILTER (
-					WHERE rollup.bucket_date = $5::date
+					WHERE rollup.bucket_date = $3::date
 				), 0) AS yesterday_cost
 			FROM usage_group_daily_rollups rollup
-			CROSS JOIN state
-			WHERE state.valid
-				AND rollup.bucket_date >= (state.retained_from AT TIME ZONE $3::text)::date
-				AND rollup.bucket_date < state.closed_before
+			WHERE $4::boolean
+				AND rollup.bucket_date >= $5::date
+				AND rollup.bucket_date < $6::date
 			GROUP BY rollup.group_id
 		),
 		tail AS (
@@ -63,8 +113,7 @@ func (r *usageLogRepository) getAllGroupUsageSummaryFromRollups(ctx context.Cont
 						AND ul.created_at < $1
 				), 0) AS yesterday_cost
 			FROM usage_logs ul
-			CROSS JOIN state
-			WHERE ul.created_at >= state.tail_start
+			WHERE ul.created_at >= $7
 			GROUP BY ul.group_id
 		)
 		SELECT
@@ -83,9 +132,11 @@ func (r *usageLogRepository) getAllGroupUsageSummaryFromRollups(ctx context.Cont
 		query,
 		todayStart,
 		yesterdayStart,
-		timezoneName,
-		todayDate,
 		yesterdayDate,
+		state.valid,
+		state.retainedDate,
+		state.closedBefore,
+		state.tailStart,
 	)
 	if err != nil {
 		return nil, err
@@ -18,8 +18,23 @@ func TestUsageLogRepositoryGetAllGroupUsageSummaryUsesRollupTail(t *testing.T) {
 	todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
 	yesterdayStart := time.Date(2026, 3, 8, 5, 0, 0, 0, time.UTC)

-	mock.ExpectQuery(`(?s)usage_group_rollup_state.*usage_group_daily_rollups.*created_at >= state\.tail_start`).
-		WithArgs(todayStart, yesterdayStart, "America/New_York", "2026-03-09", "2026-03-08").
+	// 水位有效：closed_before = 2026-03-07（≤ today），时区与配置一致。
+	mock.ExpectQuery(`(?s)COUNT\(\*\).*usage_group_rollup_state.*WHERE id = 1`).
+		WillReturnRows(sqlmock.NewRows([]string{"count", "closed_before", "retained_from", "timezone_name"}).
+			AddRow(1, "2026-03-07", time.Date(2026, 2, 1, 5, 0, 0, 0, time.UTC), "America/New_York"))
+
+	// 尾段起点必须以参数形式进 WHERE —— 这正是把 usage_logs 全表扫换回
+	// idx_usage_logs_created_at 索引扫的那一步。
+	mock.ExpectQuery(`(?s)usage_group_daily_rollups.*FROM usage_logs ul\s+WHERE ul\.created_at >= \$7`).
+		WithArgs(
+			todayStart,
+			yesterdayStart,
+			"2026-03-08", // yesterdayDate
+			true,         // 水位有效
+			"2026-02-01", // retained_from 在配置时区内的日期
+			"2026-03-07", // closed_before
+			time.Date(2026, 3, 7, 5, 0, 0, 0, time.UTC), // tail_start = closed_before 在配置时区的零点
+		).
 		WillReturnRows(sqlmock.NewRows([]string{"group_id", "total_cost", "today_cost", "yesterday_cost"}).
 			AddRow(int64(7), 12.5, 1.25, 2.5))

@@ -31,3 +46,51 @@ func TestUsageLogRepositoryGetAllGroupUsageSummaryUsesRollupTail(t *testing.T) {
 	require.InDelta(t, 2.5, result[0].YesterdayCost, 0.0000001)
 	require.NoError(t, mock.ExpectationsWereMet())
 }
+
+// 水位无效（时区变了 / 水位在未来 / 行缺失）时必须退化成"从 epoch 起全量重算"，
+// 与改动前 SQL 里 CASE WHEN valid 的 ELSE 分支逐字对应：
+// 历史日桶整段不参与，尾段起点回到 1970-01-01。
+func TestUsageLogRepositoryGetAllGroupUsageSummaryFallsBackWhenWatermarkInvalid(t *testing.T) {
+	epoch := time.Unix(0, 0).UTC()
+	todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
+	yesterdayStart := time.Date(2026, 3, 8, 5, 0, 0, 0, time.UTC)
+
+	cases := []struct {
+		name string
+		rows *sqlmock.Rows
+	}{
+		{
+			name: "水位行缺失",
+			rows: sqlmock.NewRows([]string{"count", "closed_before", "retained_from", "timezone_name"}).
+				AddRow(0, nil, nil, nil),
+		},
+		{
+			name: "时区与当前配置不一致",
+			rows: sqlmock.NewRows([]string{"count", "closed_before", "retained_from", "timezone_name"}).
+				AddRow(1, "2026-03-07", time.Date(2026, 2, 1, 5, 0, 0, 0, time.UTC), "Asia/Shanghai"),
+		},
+		{
+			name: "水位位于未来",
+			rows: sqlmock.NewRows([]string{"count", "closed_before", "retained_from", "timezone_name"}).
+				AddRow(1, "2026-03-10", time.Date(2026, 2, 1, 5, 0, 0, 0, time.UTC), "America/New_York"),
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			db, mock := newSQLMock(t)
+			repo := newUsageLogRepositoryWithSQL(nil, db)
+			useGroupUsageRepositoryTestTimezone(t, "America/New_York")
+
+			mock.ExpectQuery(`(?s)COUNT\(\*\).*usage_group_rollup_state.*WHERE id = 1`).
+				WillReturnRows(tc.rows)
+			mock.ExpectQuery(`(?s)usage_group_daily_rollups.*FROM usage_logs ul\s+WHERE ul\.created_at >= \$7`).
+				WithArgs(todayStart, yesterdayStart, "2026-03-08", false, "1970-01-01", "1970-01-01", epoch).
+				WillReturnRows(sqlmock.NewRows([]string{"group_id", "total_cost", "today_cost", "yesterday_cost"}))
+
+			_, err := repo.GetAllGroupUsageSummary(context.Background(), todayStart)
+			require.NoError(t, err)
+			require.NoError(t, mock.ExpectationsWereMet())
+		})
+	}
+}