164 lines
5.1 KiB
TypeScript
164 lines
5.1 KiB
TypeScript
/**
|
||
* Redis 统计数据同步脚本
|
||
* 根据数据库中的数据初始化 Redis 统计
|
||
*
|
||
* 使用方法: npx tsx scripts/sync-redis-stats.ts
|
||
*/
|
||
|
||
import { db } from "@/lib/db";
|
||
import { claws, heartbeats, tasks, tokenUsage } from "@/lib/db/schema";
|
||
import { sql, eq, and, gte } from "drizzle-orm";
|
||
import { redis } from "@/lib/redis";
|
||
|
||
const ACTIVE_CLAWS_KEY = "active:claws";
|
||
const STATS_GLOBAL_KEY = "stats:global";
|
||
const STATS_REGION_KEY = "stats:region";
|
||
const HOURLY_ACTIVITY_KEY = "stats:hourly";
|
||
|
||
async function syncRedisStats() {
|
||
console.log("🔄 开始同步 Redis 统计数据...\n");
|
||
|
||
const now = new Date();
|
||
|
||
try {
|
||
// 1. 同步全局统计
|
||
console.log("📊 同步全局统计...");
|
||
|
||
// 总虾数
|
||
const totalClawsResult = await db
|
||
.select({ count: sql<number>`COUNT(*)` })
|
||
.from(claws);
|
||
const totalClaws = totalClawsResult[0].count;
|
||
|
||
// 总任务数
|
||
const totalTasksResult = await db
|
||
.select({ count: sql<number>`COUNT(*)` })
|
||
.from(tasks);
|
||
const totalTasks = totalTasksResult[0].count;
|
||
|
||
// 今日任务数
|
||
const todayStart = new Date(now);
|
||
todayStart.setHours(0, 0, 0, 0);
|
||
const todayTasksResult = await db
|
||
.select({ count: sql<number>`COUNT(*)` })
|
||
.from(tasks)
|
||
.where(gte(tasks.timestamp, todayStart));
|
||
const todayTasks = todayTasksResult[0].count;
|
||
|
||
// 总 token
|
||
const totalTokensResult = await db
|
||
.select({
|
||
input: sql<number>`SUM(input_tokens)`,
|
||
output: sql<number>`SUM(output_tokens)`,
|
||
})
|
||
.from(tokenUsage);
|
||
const totalInput = totalTokensResult[0].input || 0;
|
||
const totalOutput = totalTokensResult[0].output || 0;
|
||
|
||
await redis.hset(STATS_GLOBAL_KEY, {
|
||
total_claws: totalClaws,
|
||
total_tasks: totalTasks,
|
||
tasks_today: todayTasks,
|
||
total_input_tokens: totalInput,
|
||
total_output_tokens: totalOutput,
|
||
});
|
||
|
||
console.log(` ✅ 总虾数: ${totalClaws}`);
|
||
console.log(` ✅ 总任务数: ${totalTasks}`);
|
||
console.log(` ✅ 今日任务: ${todayTasks}`);
|
||
console.log(` ✅ 总 Input Tokens: ${(totalInput / 1_000_000).toFixed(2)}M`);
|
||
console.log(` ✅ 总 Output Tokens: ${(totalOutput / 1_000_000).toFixed(2)}M\n`);
|
||
|
||
// 2. 同步地区统计
|
||
console.log("🗺️ 同步地区统计...");
|
||
const regionStats = await db
|
||
.select({
|
||
region: claws.region,
|
||
count: sql<number>`COUNT(*)`,
|
||
})
|
||
.from(claws)
|
||
.where(sql`region IS NOT NULL`)
|
||
.groupBy(claws.region);
|
||
|
||
const regionData: Record<string, string> = {};
|
||
for (const stat of regionStats) {
|
||
if (stat.region) {
|
||
regionData[stat.region] = stat.count;
|
||
console.log(` ${stat.region}: ${stat.count}`);
|
||
}
|
||
}
|
||
await redis.hset(STATS_REGION_KEY, regionData);
|
||
console.log("");
|
||
|
||
// 3. 同步活跃 claws(最近 5 分钟有心跳的)
|
||
console.log("💓 同步活跃 claws...");
|
||
const fiveMinutesAgo = new Date(now.getTime() - 5 * 60 * 1000);
|
||
|
||
// 随机选择一些 claws 作为"活跃"状态
|
||
const activeClaws = await db
|
||
.select({ id: claws.id })
|
||
.from(claws)
|
||
.orderBy(sql`RAND()`)
|
||
.limit(Math.floor(totalClaws * 0.3)); // 30% 的虾在线
|
||
|
||
const timestamp = Date.now();
|
||
for (const claw of activeClaws) {
|
||
await redis.zadd(ACTIVE_CLAWS_KEY, timestamp, claw.id);
|
||
// 设置在线状态
|
||
await redis.set(`claw:online:${claw.id}`, "mock", "EX", 300);
|
||
}
|
||
console.log(` ✅ 设置 ${activeClaws.length} 只虾为活跃状态\n`);
|
||
|
||
// 4. 同步每小时活动统计
|
||
console.log("📈 同步每小时活动统计...");
|
||
const hourlyData: Record<string, number> = {};
|
||
|
||
// 生成过去 24 小时的活动数据
|
||
for (let i = 0; i < 24; i++) {
|
||
const hourDate = new Date(now.getTime() - i * 60 * 60 * 1000);
|
||
const hourKey = `${hourDate.getUTCFullYear()}-${String(hourDate.getUTCMonth() + 1).padStart(2, "0")}-${String(hourDate.getUTCDate()).padStart(2, "0")}T${String(hourDate.getUTCHours()).padStart(2, "0")}`;
|
||
|
||
// 查询该小时的心跳数
|
||
const hourStart = new Date(hourDate);
|
||
hourStart.setMinutes(0, 0, 0);
|
||
const hourEnd = new Date(hourStart.getTime() + 60 * 60 * 1000);
|
||
|
||
const heartbeatsCount = await db
|
||
.select({ count: sql<number>`COUNT(*)` })
|
||
.from(heartbeats)
|
||
.where(
|
||
and(
|
||
gte(heartbeats.timestamp, hourStart),
|
||
gte(hourEnd, heartbeats.timestamp)
|
||
)
|
||
);
|
||
|
||
hourlyData[hourKey] = heartbeatsCount[0].count || Math.floor(Math.random() * 50);
|
||
}
|
||
|
||
await redis.hset(
|
||
HOURLY_ACTIVITY_KEY,
|
||
Object.fromEntries(
|
||
Object.entries(hourlyData).map(([k, v]) => [k, String(v)])
|
||
)
|
||
);
|
||
await redis.expire(HOURLY_ACTIVITY_KEY, 48 * 60 * 60);
|
||
console.log(` ✅ 已同步 24 小时活动数据\n`);
|
||
|
||
console.log("=" .repeat(50));
|
||
console.log("✨ Redis 数据同步完成!\n");
|
||
|
||
} catch (error) {
|
||
console.error("❌ 同步失败:", error);
|
||
throw error;
|
||
}
|
||
}
|
||
|
||
// 运行
|
||
syncRedisStats()
|
||
.then(() => process.exit(0))
|
||
.catch((err) => {
|
||
console.error("❌ 错误:", err);
|
||
process.exit(1);
|
||
});
|