diff --git a/api/mod.ts b/api/mod.ts index 66e4120..471f36d 100644 --- a/api/mod.ts +++ b/api/mod.ts @@ -14,7 +14,7 @@ export async function serveUi() { aliasMap: { "/utils/*": "../utils/", }, - quiet: true, + log: (_request, response) => response.status >= 400, }), })); diff --git a/api/serveApi.ts b/api/serveApi.ts index 0cbecec..9373870 100644 --- a/api/serveApi.ts +++ b/api/serveApi.ts @@ -1,14 +1,19 @@ -import { createPathFilter } from "t_rest/server"; +import { createLoggerMiddleware, createPathFilter } from "t_rest/server"; import { jobsRoute } from "./jobsRoute.ts"; -import { sessionsRoute } from "./sessionsRoute.ts"; -import { usersRoute } from "./usersRoute.ts"; import { paramsRoute } from "./paramsRoute.ts"; +import { sessionsRoute } from "./sessionsRoute.ts"; +import { statsRoute } from "./statsRoute.ts"; +import { usersRoute } from "./usersRoute.ts"; -export const serveApi = createPathFilter({ - "jobs": jobsRoute, - "sessions": sessionsRoute, - "users": usersRoute, - "settings/params": paramsRoute, -}); +export const serveApi = createLoggerMiddleware( + createPathFilter({ + "jobs": jobsRoute, + "sessions": sessionsRoute, + "users": usersRoute, + "settings/params": paramsRoute, + "stats": statsRoute, + }), + { filterStatus: (status) => status >= 400 }, +); export type ApiHandler = typeof serveApi; diff --git a/api/statsRoute.ts b/api/statsRoute.ts new file mode 100644 index 0000000..170bb0a --- /dev/null +++ b/api/statsRoute.ts @@ -0,0 +1,46 @@ +import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server"; +import { liveGlobalStats } from "../app/globalStatsStore.ts"; +import { getDailyStats } from "../app/dailyStatsStore.ts"; + +export const statsRoute = createPathFilter({ + "global": createMethodFilter({ + GET: createEndpoint( + { query: null, body: null }, + async () => { + return { + status: 200, + body: { + type: "application/json", + data: liveGlobalStats, + }, + }; + }, + ), + }), + "daily/{year}/{month}/{day}": createMethodFilter({ + GET: createEndpoint( + { query: null, body: null }, + async ({ params }) => { + const year = Number(params.year); + const month = Number(params.month); + const day = Number(params.day); + const minDate = new Date("2023-01-01"); + const maxDate = new Date(); + const date = new Date(Date.UTC(year, month - 1, day)); + if (date < minDate || date > maxDate) { + return { + status: 404, + body: { type: "text/plain", data: "Not found" }, + }; + } + return { + status: 200, + body: { + type: "application/json", + data: await getDailyStats(year, month, day), + }, + }; + }, + ), + }), +}); diff --git a/app/dailyStatsStore.ts b/app/dailyStatsStore.ts new file mode 100644 index 0000000..8dc4353 --- /dev/null +++ b/app/dailyStatsStore.ts @@ -0,0 +1,65 @@ +import { UTCDateMini } from "@date-fns/utc"; +import { hoursToMilliseconds, isSameDay, minutesToMilliseconds } from "date-fns"; +import { getLogger } from "std/log/mod.ts"; +import { JsonSchema, jsonType } from "t_rest/server"; +import { db } from "./db.ts"; +import { generationStore } from "./generationStore.ts"; +import { kvMemoize } from "./kvMemoize.ts"; + +const logger = () => getLogger(); + +export const dailyStatsSchema = { + type: "object", + properties: { + userIds: { type: "array", items: { type: "number" } }, + imageCount: { type: "number" }, + pixelCount: { type: "number" }, + timestamp: { type: "number" }, + }, + required: ["userIds", "imageCount", "pixelCount", "timestamp"], +} as const satisfies JsonSchema; + +export type DailyStats = jsonType; + +export const getDailyStats = kvMemoize( + db, + ["dailyStats"], + async (year: number, month: number, day: number): Promise => { + const userIdSet = new Set(); + let imageCount = 0; + let pixelCount = 0; + + const after = new Date(Date.UTC(year, month - 1, day)); + const before = new Date(Date.UTC(year, month - 1, day + 1)); + + logger().info(`Calculating daily stats for ${year}-${month}-${day}`); + + for await ( + const generation of generationStore.listAll({ after, before }) + ) { + userIdSet.add(generation.value.from.id); + imageCount++; + pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0); + } + + return { + userIds: [...userIdSet], + imageCount, + pixelCount, + timestamp: Date.now(), + }; + }, + { + // expire in 1 minute if was calculated on the same day, otherwise 7-14 days. + expireIn: (result, year, month, day) => { + const requestDate = new UTCDateMini(year, month - 1, day); + const calculatedDate = new UTCDateMini(result.timestamp); + return isSameDay(requestDate, calculatedDate) + ? minutesToMilliseconds(1) + : hoursToMilliseconds(24 * 7 + Math.random() * 24 * 7); + }, + // should cache if the stats are non-zero + shouldCache: (result) => + result.userIds.length > 0 || result.imageCount > 0 || result.pixelCount > 0, + }, +); diff --git a/app/globalStatsStore.ts b/app/globalStatsStore.ts new file mode 100644 index 0000000..a71136d --- /dev/null +++ b/app/globalStatsStore.ts @@ -0,0 +1,56 @@ +import { addDays } from "date-fns"; +import { JsonSchema, jsonType } from "t_rest/server"; +import { decodeTime } from "ulid"; +import { getDailyStats } from "./dailyStatsStore.ts"; +import { generationStore } from "./generationStore.ts"; + +export const globalStatsSchema = { + type: "object", + properties: { + userIds: { type: "array", items: { type: "number" } }, + imageCount: { type: "number" }, + pixelCount: { type: "number" }, + timestamp: { type: "number" }, + }, + required: ["userIds", "imageCount", "pixelCount", "timestamp"], +} as const satisfies JsonSchema; + +export type GlobalStats = jsonType; + +export const liveGlobalStats: GlobalStats = await getGlobalStats(); + +export async function getGlobalStats(): Promise { + // find the year/month/day of the first generation + const startDate = await generationStore.getAll({}, { limit: 1 }) + .then((generations) => generations[0]?.id) + .then((generationId) => generationId ? new Date(decodeTime(generationId)) : new Date()); + + // iterate to today and sum up stats + const userIdSet = new Set(); + let imageCount = 0; + let pixelCount = 0; + + const tomorrow = addDays(new Date(), 1); + + for ( + let date = startDate; + date < tomorrow; + date = addDays(date, 1) + ) { + const dailyStats = await getDailyStats( + date.getUTCFullYear(), + date.getUTCMonth() + 1, + date.getUTCDate(), + ); + for (const userId of dailyStats.userIds) userIdSet.add(userId); + imageCount += dailyStats.imageCount; + pixelCount += dailyStats.pixelCount; + } + + return { + userIds: [...userIdSet], + imageCount, + pixelCount, + timestamp: Date.now(), + }; +} diff --git a/app/kvMemoize.ts b/app/kvMemoize.ts new file mode 100644 index 0000000..0b74d4d --- /dev/null +++ b/app/kvMemoize.ts @@ -0,0 +1,45 @@ +/** + * Memoizes the function result in KV storage. + */ +export function kvMemoize( + db: Deno.Kv, + key: Deno.KvKey, + fn: (...args: A) => Promise, + options?: { + expireIn?: number | ((result: R, ...args: A) => number); + shouldRecalculate?: (result: R, ...args: A) => boolean; + shouldCache?: (result: R, ...args: A) => boolean; + override?: { + set: (key: Deno.KvKey, args: A, value: R, options: { expireIn?: number }) => Promise; + get: (key: Deno.KvKey, args: A) => Promise; + }; + }, +): (...args: A) => Promise { + return async (...args) => { + const cachedResult = options?.override?.get + ? await options.override.get(key, args) + : (await db.get([...key, ...args])).value; + + if (cachedResult != null) { + if (!options?.shouldRecalculate?.(cachedResult, ...args)) { + return cachedResult; + } + } + + const result = await fn(...args); + + const expireIn = typeof options?.expireIn === "function" + ? options.expireIn(result, ...args) + : options?.expireIn; + + if (options?.shouldCache?.(result, ...args) ?? (result != null)) { + if (options?.override?.set) { + await options.override.set(key, args, result, { expireIn }); + } else { + await db.set([...key, ...args], result, { expireIn }); + } + } + + return result; + }; +} diff --git a/app/uploadQueue.ts b/app/uploadQueue.ts index e49ae07..5fd747e 100644 --- a/app/uploadQueue.ts +++ b/app/uploadQueue.ts @@ -9,6 +9,7 @@ import { bot } from "../bot/mod.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; import { db, fs } from "./db.ts"; import { generationStore, SdGenerationInfo } from "./generationStore.ts"; +import { liveGlobalStats } from "./globalStatsStore.ts"; const logger = () => getLogger(); @@ -109,6 +110,15 @@ export async function processUploadQueue() { info: state.info, }); + // update live stats + { + liveGlobalStats.imageCount++; + liveGlobalStats.pixelCount += state.info.width * state.info.height; + const userIdSet = new Set(liveGlobalStats.userIds); + userIdSet.add(state.from.id); + liveGlobalStats.userIds = [...userIdSet]; + } + // delete the status message await bot.api.deleteMessage(state.replyMessage.chat.id, state.replyMessage.message_id) .catch(() => undefined); diff --git a/app/userStatsStore.ts b/app/userStatsStore.ts new file mode 100644 index 0000000..961c523 --- /dev/null +++ b/app/userStatsStore.ts @@ -0,0 +1,100 @@ +import { minutesToMilliseconds } from "date-fns"; +import { Store } from "indexed_kv"; +import { getLogger } from "std/log/mod.ts"; +import { JsonSchema, jsonType } from "t_rest/server"; +import { db } from "./db.ts"; +import { generationStore } from "./generationStore.ts"; +import { kvMemoize } from "./kvMemoize.ts"; + +const logger = () => getLogger(); + +export const userStatsSchema = { + type: "object", + properties: { + userId: { type: "number" }, + imageCount: { type: "number" }, + pixelCount: { type: "number" }, + tagsCount: { + type: "object", + additionalProperties: { type: "number" }, + }, + timestamp: { type: "number" }, + }, + required: ["userId", "imageCount", "pixelCount", "tagsCount", "timestamp"], +} as const satisfies JsonSchema; + +export type UserStats = jsonType; + +type UserStatsIndices = { + userId: number; + imageCount: number; + pixelCount: number; +}; + +const userStatsStore = new Store( + db, + "userStats", + { + indices: { + userId: { getValue: (item) => item.userId }, + imageCount: { getValue: (item) => item.imageCount }, + pixelCount: { getValue: (item) => item.pixelCount }, + }, + }, +); + +export const getUserStats = kvMemoize( + db, + ["userStats"], + async (userId: number): Promise => { + let imageCount = 0; + let pixelCount = 0; + const tagsCount: Record = {}; + + logger().info(`Calculating user stats for ${userId}`); + + for await ( + const generation of generationStore.listBy("fromId", { value: userId }) + ) { + imageCount++; + pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0); + const tags = generation.value.info?.prompt.split(/[,;.\n]/) + .map((tag) => tag.trim()) + .filter((tag) => tag.length > 0) + .map((tag) => tag.toLowerCase()) + .map((tag) => tag.replace(/[()[\]]/, "")) + .map((tag) => tag.replace(/:[\d.]/g, "")) + .map((tag) => tag.replace(/ +/g, " ")) ?? []; + for (const tag of tags) { + tagsCount[tag] = (tagsCount[tag] ?? 0) + 1; + } + } + + return { + userId, + imageCount, + pixelCount, + tagsCount, + timestamp: Date.now(), + }; + }, + { + // expire in random time between 5-10 minutes + expireIn: () => minutesToMilliseconds(5 + Math.random() * 5), + // override default set/get behavior to use userStatsStore + override: { + get: async (_key, [userId]) => { + const items = await userStatsStore.getBy("userId", { value: userId }, { reverse: true }); + return items[0]?.value; + }, + set: async (_key, [userId], value, options) => { + // delete old stats + for await (const item of userStatsStore.listBy("userId", { value: userId })) { + await item.delete(); + } + // set new stats + await userStatsStore.create(value, options); + }, + }, + }, +); diff --git a/deno.json b/deno.json index fe295f8..ae61dd7 100644 --- a/deno.json +++ b/deno.json @@ -19,10 +19,10 @@ "async": "https://deno.land/x/async@v2.0.2/mod.ts", "ulid": "https://deno.land/x/ulid@v0.3.0/mod.ts", - "indexed_kv": "https://deno.land/x/indexed_kv@v0.4.0/mod.ts", + "indexed_kv": "https://deno.land/x/indexed_kv@v0.5.0/mod.ts", "kvmq": "https://deno.land/x/kvmq@v0.3.0/mod.ts", "kvfs": "https://deno.land/x/kvfs@v0.1.0/mod.ts", - "serve_spa": "https://deno.land/x/serve_spa@v0.1.0/mod.ts", + "serve_spa": "https://deno.land/x/serve_spa@v0.2.0/mod.ts", "reroute": "https://deno.land/x/reroute@v0.1.0/mod.ts", "grammy": "https://lib.deno.dev/x/grammy@1/mod.ts", "grammy_types": "https://lib.deno.dev/x/grammy_types@3/mod.ts", @@ -30,13 +30,15 @@ "grammy_parse_mode": "https://lib.deno.dev/x/grammy_parse_mode@1/mod.ts", "grammy_stateless_question": "https://lib.deno.dev/x/grammy_stateless_question_alpha@3/mod.ts", "grammy_files": "https://lib.deno.dev/x/grammy_files@1/mod.ts", + "date-fns": "https://cdn.skypack.dev/date-fns@2.30.0?dts", + "@date-fns/utc": "https://cdn.skypack.dev/@date-fns/utc@1.1.0?dts", "file_type": "https://esm.sh/file-type@18.5.0", "png_chunks_extract": "https://esm.sh/png-chunks-extract@1.0.0", "png_chunk_text": "https://esm.sh/png-chunk-text@1.0.0", "openapi_fetch": "https://esm.sh/openapi-fetch@0.7.6", - "t_rest/server": "https://esm.sh/ty-rest@0.3.1/server?dev", + "t_rest/server": "https://esm.sh/ty-rest@0.3.2/server?dev", - "t_rest/client": "https://esm.sh/ty-rest@0.3.1/client?dev", + "t_rest/client": "https://esm.sh/ty-rest@0.3.2/client?dev", "react": "https://esm.sh/react@18.2.0?dev", "react-dom/client": "https://esm.sh/react-dom@18.2.0/client?dev", "react-router-dom": "https://esm.sh/react-router-dom@6.16.0?dev",