From 8322c2c62240155e1812f627125db34fe2170756 Mon Sep 17 00:00:00 2001 From: RainbowBird Date: Fri, 27 Mar 2026 16:40:44 +0800 Subject: [PATCH] refactor(server): split mq libs, organize billing services structure --- apps/server/src/app.ts | 21 +- apps/server/src/bin/run-billing-consumer.ts | 25 +- apps/server/src/libs/mq/index.ts | 13 + apps/server/src/libs/mq/stream.ts | 192 +++++++++++++++ .../mq/tests/worker.test.ts} | 10 +- apps/server/src/libs/mq/types.ts | 47 ++++ .../mq/worker.ts} | 33 ++- .../src/routes/tests/v1completions.test.ts | 7 +- .../billing/billing-consumer-handler.ts | 13 +- .../src/services/billing/billing-events.ts | 16 ++ .../server/src/services/billing/billing-mq.ts | 227 ------------------ .../src/services/billing/billing-service.ts | 42 ++-- .../services/billing/tests/billing-mq.test.ts | 28 +-- .../billing/tests/billing-service.test.ts | 12 +- 14 files changed, 363 insertions(+), 323 deletions(-) create mode 100644 apps/server/src/libs/mq/index.ts create mode 100644 apps/server/src/libs/mq/stream.ts rename apps/server/src/{services/billing/tests/billing-mq-worker.test.ts => libs/mq/tests/worker.test.ts} (92%) create mode 100644 apps/server/src/libs/mq/types.ts rename apps/server/src/{services/billing/billing-mq-worker.ts => libs/mq/worker.ts} (58%) delete mode 100644 apps/server/src/services/billing/billing-mq.ts diff --git a/apps/server/src/app.ts b/apps/server/src/app.ts index c5b650722..db334aaf2 100644 --- a/apps/server/src/app.ts +++ b/apps/server/src/app.ts @@ -1,8 +1,9 @@ import type Redis from 'ioredis' import type { Env } from './libs/env' +import type { MqService } from './libs/mq' import type { OtelInstance } from './libs/otel' -import type { BillingMqService } from './services/billing/billing-mq' +import type { BillingEvent } from './services/billing/billing-events' import type { BillingService } from './services/billing/billing-service' import type { CharacterService } from './services/characters' import type { ChatService } from './services/chats' @@ -40,7 +41,7 @@ import { createFluxRoutes } from './routes/flux' import { createProviderRoutes } from './routes/providers' import { createStripeRoutes } from './routes/stripe' import { createV1CompletionsRoutes } from './routes/v1completions' -import { createBillingMqService } from './services/billing/billing-mq' +import { createBillingMq } from './services/billing/billing-events' import { createBillingService } from './services/billing/billing-service' import { createCharacterService } from './services/characters' import { createChatService } from './services/chats' @@ -62,7 +63,7 @@ interface AppDeps { fluxAuditService: FluxAuditService stripeService: StripeService billingService: BillingService - billingMqService: BillingMqService + billingMq: MqService configKV: ConfigKVService redis: Redis env: Env @@ -161,7 +162,7 @@ function buildApp(deps: AppDeps) { /** * V1 routes for official provider. */ - .route('/api/v1', createV1CompletionsRoutes(deps.fluxService, deps.billingService, deps.configKV, deps.billingMqService, deps.otel?.llm)) + .route('/api/v1', createV1CompletionsRoutes(deps.fluxService, deps.billingService, deps.configKV, deps.billingMq, deps.otel?.llm)) /** * Flux routes. @@ -256,9 +257,9 @@ export async function createApp() { build: ({ dependsOn }) => createConfigKVService(dependsOn.redis), }) - const billingMqService = injeca.provide('services:billingMq', { + const billingMq = injeca.provide('services:billingMq', { dependsOn: { redis, env: parsedEnv }, - build: ({ dependsOn }) => createBillingMqService(dependsOn.redis, { + build: ({ dependsOn }) => createBillingMq(dependsOn.redis, { stream: dependsOn.env.BILLING_EVENTS_STREAM, }), }) @@ -304,8 +305,8 @@ export async function createApp() { }) const billingService = injeca.provide('services:billing', { - dependsOn: { db, redis, billingMqService, configKV, otel }, - build: ({ dependsOn }) => createBillingService(dependsOn.db, dependsOn.redis, dependsOn.billingMqService, dependsOn.configKV, dependsOn.otel?.revenue), + dependsOn: { db, redis, billingMq, configKV, otel }, + build: ({ dependsOn }) => createBillingService(dependsOn.db, dependsOn.redis, dependsOn.billingMq, dependsOn.configKV, dependsOn.otel?.revenue), }) await injeca.start() @@ -320,7 +321,7 @@ export async function createApp() { requestLogService, stripeService, billingService, - billingMqService, + billingMq, configKV, redis, env: parsedEnv, @@ -335,7 +336,7 @@ export async function createApp() { fluxAuditService: resolved.fluxAuditService, stripeService: resolved.stripeService, billingService: resolved.billingService, - billingMqService: resolved.billingMqService, + billingMq: resolved.billingMq, configKV: resolved.configKV, redis: resolved.redis, env: resolved.env, diff --git a/apps/server/src/bin/run-billing-consumer.ts b/apps/server/src/bin/run-billing-consumer.ts index c5c4d3804..78a65d3d8 100644 --- a/apps/server/src/bin/run-billing-consumer.ts +++ b/apps/server/src/bin/run-billing-consumer.ts @@ -5,19 +5,10 @@ import { initLogger, LoggerFormat, LoggerLevel, useLogger } from '@guiiai/logg' import { createDrizzle, migrateDatabase } from '../libs/db' import { parseEnv } from '../libs/env' import { initializeExternalDependency } from '../libs/external-dependency' +import { createMqWorker } from '../libs/mq' import { createRedis } from '../libs/redis' import { createBillingConsumerHandler } from '../services/billing/billing-consumer-handler' -import { createBillingMqService } from '../services/billing/billing-mq' -import { createBillingMqWorker } from '../services/billing/billing-mq-worker' - -function parsePositiveInteger(rawValue: string, envKey: string): number { - const parsed = Number(rawValue) - if (!Number.isInteger(parsed) || parsed <= 0) { - throw new Error(`${envKey} must be a positive integer`) - } - - return parsed -} +import { createBillingMq } from '../services/billing/billing-events' export async function runBillingConsumer(): Promise { initLogger(LoggerLevel.Debug, LoggerFormat.Pretty) @@ -28,7 +19,7 @@ export async function runBillingConsumer(): Promise { 'Database', logger, async (attempt) => { - const connection = createDrizzle(env.DATABASE_URL) + const connection = createDrizzle(env) try { await connection.db.execute('SELECT 1') @@ -77,20 +68,20 @@ export async function runBillingConsumer(): Promise { process.once('SIGTERM', () => shutdown('SIGTERM')) try { - const mq = createBillingMqService(redis, { + const mq = createBillingMq(redis, { stream: env.BILLING_EVENTS_STREAM, }) const handler = createBillingConsumerHandler(db) - const worker = createBillingMqWorker(mq) + const worker = createMqWorker(mq) await worker.run({ group: 'billing-consumer', consumer, signal: abortController.signal, - batchSize: parsePositiveInteger(env.BILLING_EVENTS_BATCH_SIZE, 'BILLING_EVENTS_BATCH_SIZE'), - blockMs: parsePositiveInteger(env.BILLING_EVENTS_BLOCK_MS, 'BILLING_EVENTS_BLOCK_MS'), - minIdleTimeMs: parsePositiveInteger(env.BILLING_EVENTS_MIN_IDLE_MS, 'BILLING_EVENTS_MIN_IDLE_MS'), + batchSize: env.BILLING_EVENTS_BATCH_SIZE, + blockMs: env.BILLING_EVENTS_BLOCK_MS, + minIdleTimeMs: env.BILLING_EVENTS_MIN_IDLE_MS, onMessage: message => handler.handleMessage(message), }) } diff --git a/apps/server/src/libs/mq/index.ts b/apps/server/src/libs/mq/index.ts new file mode 100644 index 000000000..362f156c6 --- /dev/null +++ b/apps/server/src/libs/mq/index.ts @@ -0,0 +1,13 @@ +export { createMqService } from './stream' +export type { MqService } from './stream' + +export type { + ClaimIdleOptions, + ConsumeOptions, + MqOptions, + RedisCommandClient, + StreamMessage, + WorkerOptions, +} from './types' +export { createMqWorker } from './worker' +export type { MqWorker } from './worker' diff --git a/apps/server/src/libs/mq/stream.ts b/apps/server/src/libs/mq/stream.ts new file mode 100644 index 000000000..41608c472 --- /dev/null +++ b/apps/server/src/libs/mq/stream.ts @@ -0,0 +1,192 @@ +import type { + ClaimIdleOptions, + ConsumeOptions, + MqOptions, + RedisArgument, + RedisCommandClient, + StreamMessage, +} from './types' + +import { useLogger } from '@guiiai/logg' + +type RedisStreamEntry = [streamMessageId: string, fieldValues: string[]] +type RedisReadGroupResponse = [stream: string, entries: RedisStreamEntry[]][] +type RedisAutoClaimResponse = [nextStartId: string, entries: RedisStreamEntry[], deletedIds?: string[]] + +const logger = useLogger('mq-stream').useGlobalConfig() + +/** + * Create a typed Redis Stream service. + * + * The caller supplies serialize/deserialize functions so this module + * stays domain-agnostic — it only knows how to talk to Redis Streams. + */ +export function createMqService(redis: RedisCommandClient, options: MqOptions) { + const { stream, serialize, deserialize } = options + + function parseEntry(entry: unknown): StreamMessage { + if (!Array.isArray(entry) || entry.length !== 2) { + throw new Error('Redis Stream entry has an invalid shape') + } + + const [streamMessageId, rawFieldValues] = entry + if (typeof streamMessageId !== 'string') { + throw new TypeError('Redis Stream entry is missing a valid message id') + } + + if (!Array.isArray(rawFieldValues)) { + throw new TypeError('Redis Stream entry fields are invalid') + } + + return { streamMessageId, event: deserialize(toFieldRecord(rawFieldValues)) } + } + + function parseReadGroupResponse(response: unknown): StreamMessage[] { + if (response == null) { + return [] + } + + if (!Array.isArray(response)) { + throw new TypeError('Redis XREADGROUP returned an invalid response') + } + + return response.flatMap((streamResponse) => { + if (!Array.isArray(streamResponse) || streamResponse.length !== 2) { + throw new Error('Redis XREADGROUP returned an invalid stream payload') + } + + const [, entries] = streamResponse as RedisReadGroupResponse[number] + return entries.map(parseEntry) + }) + } + + function parseAutoClaimResponse(response: unknown): StreamMessage[] { + if (response == null) { + return [] + } + + if (!Array.isArray(response) || response.length < 2) { + throw new Error('Redis XAUTOCLAIM returned an invalid response') + } + + const [, entries] = response as RedisAutoClaimResponse + if (!Array.isArray(entries)) { + throw new TypeError('Redis XAUTOCLAIM returned invalid entries') + } + + return entries.map(parseEntry) + } + + return { + stream, + + async publish(event: TEvent): Promise { + const fields = serialize(event) + const xaddArgs: RedisArgument[] = [stream] + + if (options.maxLength != null) { + xaddArgs.push('MAXLEN', '~', options.maxLength) + } + + xaddArgs.push('*', ...toRedisFieldArguments(fields)) + + const streamMessageId = await redis.call('XADD', ...xaddArgs) + if (typeof streamMessageId !== 'string') { + throw new TypeError('Redis XADD did not return a stream message id') + } + + logger.withFields({ stream, streamMessageId }).log('Published event to Redis Stream') + return streamMessageId + }, + + async ensureConsumerGroup(group: string, startId = '0'): Promise { + try { + await redis.call('XGROUP', 'CREATE', stream, group, startId, 'MKSTREAM') + return true + } + catch (error) { + if (error instanceof Error && error.message.includes('BUSYGROUP')) { + return false + } + + throw error + } + }, + + async consume(consumeOptions: ConsumeOptions): Promise[]> { + const response = await redis.call( + 'XREADGROUP', + 'GROUP', + consumeOptions.group, + consumeOptions.consumer, + 'COUNT', + consumeOptions.count ?? 10, + 'BLOCK', + consumeOptions.blockMs ?? 5_000, + 'STREAMS', + stream, + consumeOptions.startId ?? '>', + ) + + return parseReadGroupResponse(response) + }, + + async claimIdleMessages(claimOptions: ClaimIdleOptions): Promise[]> { + const response = await redis.call( + 'XAUTOCLAIM', + stream, + claimOptions.group, + claimOptions.consumer, + claimOptions.minIdleTimeMs, + claimOptions.startId ?? '0-0', + 'COUNT', + claimOptions.count ?? 10, + ) + + return parseAutoClaimResponse(response) + }, + + async ack(group: string, streamMessageIds: string | string[]): Promise { + const ids = Array.isArray(streamMessageIds) ? streamMessageIds : [streamMessageIds] + + if (ids.length === 0) { + return 0 + } + + const acked = await redis.call('XACK', stream, group, ...ids) + if (typeof acked !== 'number') { + throw new TypeError('Redis XACK did not return an acknowledgement count') + } + + return acked + }, + } +} + +function toRedisFieldArguments(fields: Record): RedisArgument[] { + return Object.entries(fields) + .filter(([, value]) => value !== undefined) + .flatMap(([key, value]) => [key, value as string]) +} + +function toFieldRecord(fieldValues: string[]): Record { + if (fieldValues.length % 2 !== 0) { + throw new Error('Redis Stream entry fields must be key/value pairs') + } + + const fields: Record = {} + for (let index = 0; index < fieldValues.length; index += 2) { + const key = fieldValues[index] + const value = fieldValues[index + 1] + + if (typeof key !== 'string' || typeof value !== 'string') { + throw new TypeError('Redis Stream entry contains non-string field data') + } + + fields[key] = value + } + + return fields +} + +export type MqService = ReturnType> diff --git a/apps/server/src/services/billing/tests/billing-mq-worker.test.ts b/apps/server/src/libs/mq/tests/worker.test.ts similarity index 92% rename from apps/server/src/services/billing/tests/billing-mq-worker.test.ts rename to apps/server/src/libs/mq/tests/worker.test.ts index 1489ec979..57aeac714 100644 --- a/apps/server/src/services/billing/tests/billing-mq-worker.test.ts +++ b/apps/server/src/libs/mq/tests/worker.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from 'vitest' -import { createBillingMqWorker } from '../billing-mq-worker' +import { createMqWorker } from '../worker' function createMessage() { return { @@ -22,7 +22,7 @@ function createMessage() { } } -describe('billingMqWorker', () => { +describe('mqWorker', () => { it('reclaims pending messages before reading new ones and acks after handling', async () => { const controller = new AbortController() const message = createMessage() @@ -34,7 +34,7 @@ describe('billingMqWorker', () => { ack: vi.fn(async () => 1), } - const worker = createBillingMqWorker(mq as any) + const worker = createMqWorker(mq as any) const handled: string[] = [] await worker.run({ @@ -70,7 +70,7 @@ describe('billingMqWorker', () => { ack: vi.fn(async () => 1), } - const worker = createBillingMqWorker(mq as any) + const worker = createMqWorker(mq as any) await worker.run({ group: 'billing', @@ -102,7 +102,7 @@ describe('billingMqWorker', () => { ack: vi.fn(async () => 1), } - const worker = createBillingMqWorker(mq as any) + const worker = createMqWorker(mq as any) await worker.run({ group: 'billing', diff --git a/apps/server/src/libs/mq/types.ts b/apps/server/src/libs/mq/types.ts new file mode 100644 index 000000000..555194ff3 --- /dev/null +++ b/apps/server/src/libs/mq/types.ts @@ -0,0 +1,47 @@ +export type RedisArgument = string | number + +export interface RedisCommandClient { + call: (command: string, ...args: RedisArgument[]) => Promise +} + +export interface MqOptions { + /** Redis Stream key name. */ + stream: string + /** Approximate max stream length (MAXLEN ~). Unbounded if omitted. */ + maxLength?: number + /** Convert a typed event into flat Redis field/value pairs. */ + serialize: (event: TEvent) => Record + /** Reconstruct a typed event from flat Redis field/value pairs. */ + deserialize: (fields: Record) => TEvent +} + +export interface StreamMessage { + streamMessageId: string + event: TEvent +} + +export interface ConsumeOptions { + group: string + consumer: string + count?: number + blockMs?: number + startId?: string +} + +export interface ClaimIdleOptions { + group: string + consumer: string + minIdleTimeMs: number + startId?: string + count?: number +} + +export interface WorkerOptions { + group: string + consumer: string + signal: AbortSignal + batchSize?: number + blockMs?: number + minIdleTimeMs?: number + onMessage: (message: StreamMessage) => Promise +} diff --git a/apps/server/src/services/billing/billing-mq-worker.ts b/apps/server/src/libs/mq/worker.ts similarity index 58% rename from apps/server/src/services/billing/billing-mq-worker.ts rename to apps/server/src/libs/mq/worker.ts index 8563ecabc..8f9edf14f 100644 --- a/apps/server/src/services/billing/billing-mq-worker.ts +++ b/apps/server/src/libs/mq/worker.ts @@ -1,22 +1,20 @@ -import type { BillingMqService, BillingStreamMessage } from './billing-mq' +import type { MqService } from './stream' +import type { StreamMessage, WorkerOptions } from './types' import { useLogger } from '@guiiai/logg' -export interface RunBillingMqWorkerOptions { - group: string - consumer: string - signal: AbortSignal - batchSize?: number - blockMs?: number - minIdleTimeMs?: number - onMessage: (message: BillingStreamMessage) => Promise -} +const logger = useLogger('mq-worker').useGlobalConfig() -const logger = useLogger('billing-mq-worker').useGlobalConfig() - -export function createBillingMqWorker(mq: BillingMqService) { +/** + * Create a consumer worker that processes messages from a Redis Stream. + * + * The loop first reclaims idle (possibly stalled) messages, then falls + * back to consuming new ones. Each message is passed to `onMessage`; + * on success it is acknowledged, on failure it stays pending for retry. + */ +export function createMqWorker(mq: MqService) { return { - async run(options: RunBillingMqWorkerOptions): Promise { + async run(options: WorkerOptions): Promise { await mq.ensureConsumerGroup(options.group) while (!options.signal.aborted) { @@ -27,7 +25,7 @@ export function createBillingMqWorker(mq: BillingMqService) { count: options.batchSize ?? 10, }) - const messages = reclaimedMessages.length > 0 + const messages: StreamMessage[] = reclaimedMessages.length > 0 ? reclaimedMessages : await mq.consume({ group: options.group, @@ -49,9 +47,8 @@ export function createBillingMqWorker(mq: BillingMqService) { logger.withError(error).withFields({ group: options.group, consumer: options.consumer, - eventId: message.event.eventId, streamMessageId: message.streamMessageId, - }).error('Billing MQ handler failed; leaving message pending') + }).error('MQ handler failed; leaving message pending') } } } @@ -59,4 +56,4 @@ export function createBillingMqWorker(mq: BillingMqService) { } } -export type BillingMqWorker = ReturnType +export type MqWorker = ReturnType> diff --git a/apps/server/src/routes/tests/v1completions.test.ts b/apps/server/src/routes/tests/v1completions.test.ts index 5a9ed7b2f..802356c7e 100644 --- a/apps/server/src/routes/tests/v1completions.test.ts +++ b/apps/server/src/routes/tests/v1completions.test.ts @@ -1,4 +1,5 @@ -import type { BillingMqService } from '../../services/billing/billing-mq' +import type { MqService } from '../../libs/mq' +import type { BillingEvent } from '../../services/billing/billing-events' import type { BillingService } from '../../services/billing/billing-service' import type { ConfigKVService } from '../../services/config-kv' import type { FluxService } from '../../services/flux' @@ -53,7 +54,7 @@ function createMockConfigKV(overrides: Record = {}): ConfigKVServic } as any } -function createMockBillingMq(): BillingMqService { +function createMockBillingMq(): MqService { return { stream: 'billing-events', publish: vi.fn(async () => '1-0'), @@ -68,7 +69,7 @@ function createTestApp( fluxService: FluxService, configKV: ConfigKVService, billingService?: BillingService, - billingMq?: BillingMqService, + billingMq?: MqService, ) { const routes = createV1CompletionsRoutes(fluxService, billingService ?? createMockBillingService(), configKV, billingMq ?? createMockBillingMq(), null) const app = new Hono() diff --git a/apps/server/src/services/billing/billing-consumer-handler.ts b/apps/server/src/services/billing/billing-consumer-handler.ts index 10bc7bbb2..56d753a6b 100644 --- a/apps/server/src/services/billing/billing-consumer-handler.ts +++ b/apps/server/src/services/billing/billing-consumer-handler.ts @@ -1,5 +1,6 @@ import type { Database } from '../../libs/db' -import type { BillingStreamMessage } from './billing-mq' +import type { StreamMessage } from '../../libs/mq' +import type { BillingEvent } from './billing-events' import { useLogger } from '@guiiai/logg' @@ -10,7 +11,7 @@ const logger = useLogger('billing-consumer-handler').useGlobalConfig() export function createBillingConsumerHandler(db: Database) { return { - async handleMessage(message: BillingStreamMessage): Promise { + async handleMessage(message: StreamMessage): Promise { const { event } = message switch (event.eventType) { @@ -19,6 +20,8 @@ export function createBillingConsumerHandler(db: Database) { ? event.payload.balanceAfter + event.payload.amount : 0 + // NOTICE: onConflictDoNothing handles redelivery after crash — + // the unique index (userId, requestId) prevents duplicate ledger entries. await db.insert(fluxLedgerSchema.fluxLedger).values({ userId: event.userId, type: 'debit', @@ -27,7 +30,7 @@ export function createBillingConsumerHandler(db: Database) { balanceAfter: event.payload.balanceAfter ?? balanceBefore - event.payload.amount, requestId: event.requestId, description: event.payload.source ?? 'LLM request', - }) + }).onConflictDoNothing() logger.withFields({ eventId: event.eventId, @@ -38,7 +41,9 @@ export function createBillingConsumerHandler(db: Database) { } case 'llm.request.log': { + // NOTICE: Use eventId as PK to make redelivery idempotent. await db.insert(llmRequestLogSchema.llmRequestLog).values({ + id: event.eventId, userId: event.userId, model: event.payload.model, status: event.payload.status, @@ -46,7 +51,7 @@ export function createBillingConsumerHandler(db: Database) { fluxConsumed: event.payload.fluxConsumed, promptTokens: event.payload.promptTokens, completionTokens: event.payload.completionTokens, - }) + }).onConflictDoNothing() logger.withFields({ eventId: event.eventId, diff --git a/apps/server/src/services/billing/billing-events.ts b/apps/server/src/services/billing/billing-events.ts index d1316895b..cb06ad337 100644 --- a/apps/server/src/services/billing/billing-events.ts +++ b/apps/server/src/services/billing/billing-events.ts @@ -1,5 +1,7 @@ import type { InferOutput } from 'valibot' +import type { RedisCommandClient } from '../../libs/mq' + import { literal, nonEmpty, @@ -13,6 +15,8 @@ import { unknown, } from 'valibot' +import { createMqService } from '../../libs/mq' + export const DEFAULT_BILLING_EVENTS_STREAM = 'billing-events' const BillingEventTypeSchema = union([ @@ -128,6 +132,18 @@ export function serializeBillingEvent(event: BillingEvent): SerializedBillingEve } } +/** + * Create a Redis Stream MQ service pre-configured for billing events. + */ +export function createBillingMq(redis: RedisCommandClient, options: { stream?: string, maxLength?: number } = {}) { + return createMqService(redis, { + stream: options.stream ?? DEFAULT_BILLING_EVENTS_STREAM, + maxLength: options.maxLength, + serialize: serializeBillingEvent, + deserialize: parseBillingEvent, + }) +} + export function parseBillingEvent(fields: Record): BillingEvent { const payload = fields.payload if (payload == null) { diff --git a/apps/server/src/services/billing/billing-mq.ts b/apps/server/src/services/billing/billing-mq.ts deleted file mode 100644 index ca22c53a5..000000000 --- a/apps/server/src/services/billing/billing-mq.ts +++ /dev/null @@ -1,227 +0,0 @@ -import type { BillingEvent } from './billing-events' - -import { useLogger } from '@guiiai/logg' - -import { - DEFAULT_BILLING_EVENTS_STREAM, - parseBillingEvent, - serializeBillingEvent, -} from './billing-events' - -type RedisArgument = string | number - -export interface RedisCommandClient { - call: (command: string, ...args: RedisArgument[]) => Promise -} - -export interface BillingMqOptions { - stream?: string - maxLength?: number -} - -export interface ConsumeBillingMessagesOptions { - group: string - consumer: string - count?: number - blockMs?: number - startId?: string -} - -export interface ClaimIdleBillingMessagesOptions { - group: string - consumer: string - minIdleTimeMs: number - startId?: string - count?: number -} - -export interface BillingStreamMessage { - streamMessageId: string - event: BillingEvent -} - -type RedisStreamEntry = [streamMessageId: string, fieldValues: string[]] -type RedisReadGroupResponse = [stream: string, entries: RedisStreamEntry[]][] -type RedisAutoClaimResponse = [nextStartId: string, entries: RedisStreamEntry[], deletedIds?: string[]] - -const logger = useLogger('billing-mq').useGlobalConfig() - -export function createBillingMqService(redis: RedisCommandClient, options: BillingMqOptions = {}) { - const stream = options.stream ?? DEFAULT_BILLING_EVENTS_STREAM - - return { - stream, - - async publish(event: BillingEvent): Promise { - const serializedFields = serializeBillingEvent(event) - const xaddArgs: RedisArgument[] = [stream] - - if (options.maxLength != null) { - xaddArgs.push('MAXLEN', '~', options.maxLength) - } - - xaddArgs.push('*', ...toRedisFieldArguments(serializedFields)) - - const streamMessageId = await redis.call('XADD', ...xaddArgs) - if (typeof streamMessageId !== 'string') { - throw new TypeError('Redis XADD did not return a stream message id') - } - - logger.withFields({ - stream, - eventId: event.eventId, - eventType: event.eventType, - streamMessageId, - }).log('Published billing event to Redis Stream') - - return streamMessageId - }, - - async ensureConsumerGroup(group: string, startId = '0'): Promise { - try { - await redis.call('XGROUP', 'CREATE', stream, group, startId, 'MKSTREAM') - return true - } - catch (error) { - if (error instanceof Error && error.message.includes('BUSYGROUP')) { - return false - } - - throw error - } - }, - - async consume(options: ConsumeBillingMessagesOptions): Promise { - const response = await redis.call( - 'XREADGROUP', - 'GROUP', - options.group, - options.consumer, - 'COUNT', - options.count ?? 10, - 'BLOCK', - options.blockMs ?? 5_000, - 'STREAMS', - stream, - options.startId ?? '>', - ) - - return parseReadGroupResponse(response) - }, - - async claimIdleMessages(options: ClaimIdleBillingMessagesOptions): Promise { - const response = await redis.call( - 'XAUTOCLAIM', - stream, - options.group, - options.consumer, - options.minIdleTimeMs, - options.startId ?? '0-0', - 'COUNT', - options.count ?? 10, - ) - - return parseAutoClaimResponse(response) - }, - - async ack(group: string, streamMessageIds: string | string[]): Promise { - const ids = Array.isArray(streamMessageIds) ? streamMessageIds : [streamMessageIds] - - if (ids.length === 0) { - return 0 - } - - const acked = await redis.call('XACK', stream, group, ...ids) - if (typeof acked !== 'number') { - throw new TypeError('Redis XACK did not return an acknowledgement count') - } - - return acked - }, - } -} - -function toRedisFieldArguments(fields: Record): RedisArgument[] { - return Object.entries(fields) - .filter(([, value]) => value !== undefined) - .flatMap(([key, value]) => [key, value as string]) -} - -function parseReadGroupResponse(response: unknown): BillingStreamMessage[] { - if (response == null) { - return [] - } - - if (!Array.isArray(response)) { - throw new TypeError('Redis XREADGROUP returned an invalid response') - } - - return response.flatMap((streamResponse) => { - if (!Array.isArray(streamResponse) || streamResponse.length !== 2) { - throw new Error('Redis XREADGROUP returned an invalid stream payload') - } - - const [, entries] = streamResponse as RedisReadGroupResponse[number] - return entries.map(parseRedisStreamEntry) - }) -} - -function parseAutoClaimResponse(response: unknown): BillingStreamMessage[] { - if (response == null) { - return [] - } - - if (!Array.isArray(response) || response.length < 2) { - throw new Error('Redis XAUTOCLAIM returned an invalid response') - } - - const [, entries] = response as RedisAutoClaimResponse - if (!Array.isArray(entries)) { - throw new TypeError('Redis XAUTOCLAIM returned invalid entries') - } - - return entries.map(parseRedisStreamEntry) -} - -function parseRedisStreamEntry(entry: unknown): BillingStreamMessage { - if (!Array.isArray(entry) || entry.length !== 2) { - throw new Error('Redis Stream entry has an invalid shape') - } - - const [streamMessageId, rawFieldValues] = entry - if (typeof streamMessageId !== 'string') { - throw new TypeError('Redis Stream entry is missing a valid message id') - } - - if (!Array.isArray(rawFieldValues)) { - throw new TypeError('Redis Stream entry fields are invalid') - } - - const event = parseBillingEvent(toFieldRecord(rawFieldValues)) - return { - streamMessageId, - event, - } -} - -function toFieldRecord(fieldValues: string[]): Record { - if (fieldValues.length % 2 !== 0) { - throw new Error('Redis Stream entry fields must be key/value pairs') - } - - const fields: Record = {} - for (let index = 0; index < fieldValues.length; index += 2) { - const key = fieldValues[index] - const value = fieldValues[index + 1] - - if (typeof key !== 'string' || typeof value !== 'string') { - throw new TypeError('Redis Stream entry contains non-string field data') - } - - fields[key] = value - } - - return fields -} - -export type BillingMqService = ReturnType diff --git a/apps/server/src/services/billing/billing-service.ts b/apps/server/src/services/billing/billing-service.ts index 42ee0963f..dabb9392a 100644 --- a/apps/server/src/services/billing/billing-service.ts +++ b/apps/server/src/services/billing/billing-service.ts @@ -1,13 +1,13 @@ import type Redis from 'ioredis' import type { Database } from '../../libs/db' +import type { MqService } from '../../libs/mq' import type { RevenueMetrics } from '../../libs/otel' import type { ConfigKVService } from '../config-kv' import type { BillingEvent } from './billing-events' -import type { BillingMqService } from './billing-mq' import { useLogger } from '@guiiai/logg' -import { eq } from 'drizzle-orm' +import { and, eq } from 'drizzle-orm' import { createPaymentRequiredError } from '../../utils/error' import { nanoid } from '../../utils/id' @@ -22,7 +22,7 @@ const logger = useLogger('billing-service') export function createBillingService( db: Database, redis: Redis, - billingMq: BillingMqService, + billingMq: MqService, _configKV: ConfigKVService, metrics?: RevenueMetrics | null, ) { @@ -205,11 +205,16 @@ export function createBillingService( fluxAmount: number }): Promise<{ applied: boolean, balanceAfter?: number }> { const txResult = await db.transaction(async (tx) => { - const record = await tx.query.stripeCheckoutSession.findFirst({ - where: (table, { eq }) => eq(table.stripeSessionId, input.stripeSessionId), - }) + // Atomic claim: set fluxCredited = true only if currently false + const [claimed] = await tx.update(stripeSchema.stripeCheckoutSession) + .set({ fluxCredited: true, updatedAt: new Date() }) + .where(and( + eq(stripeSchema.stripeCheckoutSession.stripeSessionId, input.stripeSessionId), + eq(stripeSchema.stripeCheckoutSession.fluxCredited, false), + )) + .returning() - if (!record || record.fluxCredited) { + if (!claimed) { return { applied: false } } @@ -233,11 +238,6 @@ export function createBillingService( .set({ flux: balanceAfter, updatedAt: new Date() }) .where(eq(fluxSchema.userFlux.userId, input.userId)) - // Mark checkout session as credited - await tx.update(stripeSchema.stripeCheckoutSession) - .set({ fluxCredited: true, updatedAt: new Date() }) - .where(eq(stripeSchema.stripeCheckoutSession.stripeSessionId, input.stripeSessionId)) - const description = `Stripe payment ${input.currency?.toUpperCase() ?? 'UNKNOWN'} ${(input.amountTotal / 100).toFixed(2)}` // Ledger entry @@ -313,11 +313,16 @@ export function createBillingService( fluxAmount: number }): Promise<{ applied: boolean, balanceAfter?: number }> { const txResult = await db.transaction(async (tx) => { - const record = await tx.query.stripeInvoice.findFirst({ - where: (table, { eq }) => eq(table.stripeInvoiceId, input.stripeInvoiceId), - }) + // Atomic claim: set fluxCredited = true only if currently false + const [claimed] = await tx.update(stripeSchema.stripeInvoice) + .set({ fluxCredited: true, updatedAt: new Date() }) + .where(and( + eq(stripeSchema.stripeInvoice.stripeInvoiceId, input.stripeInvoiceId), + eq(stripeSchema.stripeInvoice.fluxCredited, false), + )) + .returning() - if (!record || record.fluxCredited) { + if (!claimed) { return { applied: false } } @@ -341,11 +346,6 @@ export function createBillingService( .set({ flux: balanceAfter, updatedAt: new Date() }) .where(eq(fluxSchema.userFlux.userId, input.userId)) - // Mark invoice as credited - await tx.update(stripeSchema.stripeInvoice) - .set({ fluxCredited: true, updatedAt: new Date() }) - .where(eq(stripeSchema.stripeInvoice.stripeInvoiceId, input.stripeInvoiceId)) - const description = `Subscription invoice ${input.currency.toUpperCase()} ${(input.amountPaid / 100).toFixed(2)}` // Ledger entry diff --git a/apps/server/src/services/billing/tests/billing-mq.test.ts b/apps/server/src/services/billing/tests/billing-mq.test.ts index a0398def5..52b77ab4c 100644 --- a/apps/server/src/services/billing/tests/billing-mq.test.ts +++ b/apps/server/src/services/billing/tests/billing-mq.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from 'vitest' -import { createBillingMqService } from '../billing-mq' +import { createBillingMq } from '../billing-events' function createEvent() { return { @@ -25,7 +25,7 @@ describe('billingMqService', () => { call: vi.fn(async () => '1740000000000-0'), } - const mq = createBillingMqService(redis, { + const mq = createBillingMq(redis, { stream: 'billing-events-test', maxLength: 1_000, }) @@ -62,7 +62,7 @@ describe('billingMqService', () => { }) it('throws when publish does not return a stream message id', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => 123), }) @@ -70,7 +70,7 @@ describe('billingMqService', () => { }) it('creates a consumer group and returns true when the group is new', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => 'OK'), }) @@ -78,7 +78,7 @@ describe('billingMqService', () => { }) it('returns false when the consumer group already exists', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => { throw new Error('BUSYGROUP Consumer Group name already exists') }), @@ -88,7 +88,7 @@ describe('billingMqService', () => { }) it('rethrows non-BUSYGROUP errors when creating a consumer group', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => { throw new Error('NOAUTH') }), @@ -98,7 +98,7 @@ describe('billingMqService', () => { }) it('consumes stream entries from a consumer group', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => [[ 'billing-events', [[ @@ -141,7 +141,7 @@ describe('billingMqService', () => { }) it('returns an empty array when no messages are available', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => null), }) @@ -152,7 +152,7 @@ describe('billingMqService', () => { }) it('throws when xreadgroup returns an invalid payload', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => ['not-an-array-entry']), }) @@ -163,7 +163,7 @@ describe('billingMqService', () => { }) it('claims idle pending messages', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => [ '1740000000001-0', [[ @@ -206,7 +206,7 @@ describe('billingMqService', () => { }) it('throws when xautoclaim returns an invalid payload', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => ['1740000000001-0']), }) @@ -222,7 +222,7 @@ describe('billingMqService', () => { call: vi.fn(async () => 2), } - const mq = createBillingMqService(redis) + const mq = createBillingMq(redis) await expect(mq.ack('billing', ['1-0', '2-0'])).resolves.toBe(2) expect(redis.call).toHaveBeenCalledWith('XACK', 'billing-events', 'billing', '1-0', '2-0') }) @@ -232,13 +232,13 @@ describe('billingMqService', () => { call: vi.fn(), } - const mq = createBillingMqService(redis) + const mq = createBillingMq(redis) await expect(mq.ack('billing', [])).resolves.toBe(0) expect(redis.call).not.toHaveBeenCalled() }) it('throws when ack does not return a number', async () => { - const mq = createBillingMqService({ + const mq = createBillingMq({ call: vi.fn(async () => '2'), }) diff --git a/apps/server/src/services/billing/tests/billing-service.test.ts b/apps/server/src/services/billing/tests/billing-service.test.ts index 148dbab29..46535d57f 100644 --- a/apps/server/src/services/billing/tests/billing-service.test.ts +++ b/apps/server/src/services/billing/tests/billing-service.test.ts @@ -1,8 +1,9 @@ import type Redis from 'ioredis' import type { Database } from '../../../libs/db' +import type { MqService } from '../../../libs/mq' import type { createConfigKVService } from '../../config-kv' -import type { BillingMqService } from '../billing-mq' +import type { BillingEvent } from '../billing-events' import { eq } from 'drizzle-orm' import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' @@ -26,11 +27,14 @@ function createMockRedis(): Redis { const store = new Map() return { get: vi.fn(async (key: string) => store.get(key) ?? null), - set: vi.fn(async (key: string, value: string) => { store.set(key, value); return 'OK' }), + set: vi.fn(async (key: string, value: string) => { + store.set(key, value) + return 'OK' + }), } as unknown as Redis } -function createMockBillingMq(): BillingMqService { +function createMockBillingMq(): MqService { return { stream: 'billing-events', publish: vi.fn(async () => '1-0'), @@ -44,7 +48,7 @@ function createMockBillingMq(): BillingMqService { describe('billingService', () => { let db: Database let redis: Redis - let billingMq: BillingMqService + let billingMq: MqService let billingService: ReturnType beforeAll(async () => {