i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-01-02 14:22:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: migrate DynamoDB quota to composite key schema (#426)

Browse Source 
- Change from single key (PK only) to composite key (PK + SK)
- PK = user ID, SK = date for per-day history tracking
- Remove two-step daily reset logic (SK handles day separation)
- Rename dailyReqCount/dailyTokenCount to reqCount/tokenCount
- Remove TTL (data never expires per user request)
- Simplify checkAndIncrementRequest to single atomic update
- Fix recordTokenUsage to handle new items explicitly

New table: next-ai-drawio-quota-v2
This commit is contained in:
Dayuan Jiang
2025-12-27 10:24:43 +09:00
committed by GitHub
parent 467561df47
commit 18ab1bffa0
1 changed files with 44 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

128
lib/dynamo-quota-manager.ts
View File

@@ -27,6 +27,7 @@ try {
/** /**
* Get today's date string in the configured timezone (YYYY-MM-DD format) * Get today's date string in the configured timezone (YYYY-MM-DD format)
* This is used as the Sort Key (SK) for per-day tracking
*/ */
function getTodayInTimezone(): string { function getTodayInTimezone(): string {
return new Intl.DateTimeFormat("en-CA", { return new Intl.DateTimeFormat("en-CA", {
@@ -61,8 +62,8 @@ interface QuotaCheckResult {
/** /**
* Check all quotas and increment request count atomically. * Check all quotas and increment request count atomically.
* Uses ConditionExpression to prevent race conditions. * Uses composite key (PK=user, SK=date) for per-day tracking.
* Returns which limit was exceeded if any. * Each day automatically gets a new item - no explicit reset needed.
*/ */
export async function checkAndIncrementRequest( export async function checkAndIncrementRequest(
ip: string, ip: string,
@@ -73,77 +74,33 @@ export async function checkAndIncrementRequest(
return { allowed: true } return { allowed: true }
} }
const today = getTodayInTimezone() const pk = ip // User identifier (base64 IP)
const sk = getTodayInTimezone() // Date as sort key (YYYY-MM-DD)
const currentMinute = Math.floor(Date.now() / 60000).toString() const currentMinute = Math.floor(Date.now() / 60000).toString()
const ttl = Math.floor(Date.now() / 1000) + 7 * 24 * 60 * 60
try { try {
// First, try to reset counts if it's a new day (atomic day reset) // Single atomic update - handles creation AND increment
// This will succeed only if lastResetDate < today or doesn't exist // New day automatically creates new item (different SK)
try { // Note: lastMinute/tpmCount are managed by recordTokenUsage only
await client.send( await client.send(
new UpdateItemCommand({ new UpdateItemCommand({
TableName: TABLE, TableName: TABLE,
Key: { PK: { S: `IP#${ip}` } }, Key: {
// Reset all counts to 1/0 for the new day PK: { S: pk },
UpdateExpression: ` SK: { S: sk },
SET lastResetDate = :today,
dailyReqCount = :one,
dailyTokenCount = :zero,
lastMinute = :minute,
tpmCount = :zero,
#ttl = :ttl
`,
// Only succeed if it's a new day (or new item)
ConditionExpression: `
attribute_not_exists(lastResetDate) OR lastResetDate < :today
`,
ExpressionAttributeNames: { "#ttl": "ttl" },
ExpressionAttributeValues: {
":today": { S: today },
":zero": { N: "0" },
":one": { N: "1" },
":minute": { S: currentMinute },
":ttl": { N: String(ttl) },
}, },
}), UpdateExpression: "ADD reqCount :one",
)
// New day reset successful
return { allowed: true }
} catch (resetError: any) {
// If condition failed, it's the same day - continue to increment logic
if (!(resetError instanceof ConditionalCheckFailedException)) {
throw resetError // Re-throw unexpected errors
}
}
// Same day - increment request count with limit checks
await client.send(
new UpdateItemCommand({
TableName: TABLE,
Key: { PK: { S: `IP#${ip}` } },
// Increment request count, handle minute boundary for TPM
UpdateExpression: `
SET lastMinute = :minute,
tpmCount = if_not_exists(tpmCount, :zero),
#ttl = :ttl
ADD dailyReqCount :one
`,
// Check all limits before allowing increment // Check all limits before allowing increment
// TPM check: allow if new minute OR under limit
ConditionExpression: ` ConditionExpression: `
lastResetDate = :today AND (attribute_not_exists(reqCount) OR reqCount < :reqLimit) AND
(attribute_not_exists(dailyReqCount) OR dailyReqCount < :reqLimit) AND (attribute_not_exists(tokenCount) OR tokenCount < :tokenLimit) AND
(attribute_not_exists(dailyTokenCount) OR dailyTokenCount < :tokenLimit) AND
(attribute_not_exists(lastMinute) OR lastMinute <> :minute OR (attribute_not_exists(lastMinute) OR lastMinute <> :minute OR
attribute_not_exists(tpmCount) OR tpmCount < :tpmLimit) attribute_not_exists(tpmCount) OR tpmCount < :tpmLimit)
`, `,
ExpressionAttributeNames: { "#ttl": "ttl" },
ExpressionAttributeValues: { ExpressionAttributeValues: {
":today": { S: today },
":zero": { N: "0" },
":one": { N: "1" }, ":one": { N: "1" },
":minute": { S: currentMinute }, ":minute": { S: currentMinute },
":ttl": { N: String(ttl) },
":reqLimit": { N: String(limits.requests || 999999) }, ":reqLimit": { N: String(limits.requests || 999999) },
":tokenLimit": { N: String(limits.tokens || 999999) }, ":tokenLimit": { N: String(limits.tokens || 999999) },
":tpmLimit": { N: String(limits.tpm || 999999) }, ":tpmLimit": { N: String(limits.tpm || 999999) },
@@ -160,42 +117,39 @@ export async function checkAndIncrementRequest(
const getResult = await client.send( const getResult = await client.send(
new GetItemCommand({ new GetItemCommand({
TableName: TABLE, TableName: TABLE,
Key: { PK: { S: `IP#${ip}` } }, Key: {
PK: { S: pk },
SK: { S: sk },
},
}), }),
) )
const item = getResult.Item const item = getResult.Item
const storedDate = item?.lastResetDate?.S
const storedMinute = item?.lastMinute?.S const storedMinute = item?.lastMinute?.S
const isNewDay = !storedDate || storedDate < today
const dailyReqCount = isNewDay const reqCount = Number(item?.reqCount?.N || 0)
? 0 const tokenCount = Number(item?.tokenCount?.N || 0)
: Number(item?.dailyReqCount?.N || 0)
const dailyTokenCount = isNewDay
? 0
: Number(item?.dailyTokenCount?.N || 0)
const tpmCount = const tpmCount =
storedMinute !== currentMinute storedMinute !== currentMinute
? 0 ? 0
: Number(item?.tpmCount?.N || 0) : Number(item?.tpmCount?.N || 0)
// Determine which limit was exceeded // Determine which limit was exceeded
if (limits.requests > 0 && dailyReqCount >= limits.requests) { if (limits.requests > 0 && reqCount >= limits.requests) {
return { return {
allowed: false, allowed: false,
type: "request", type: "request",
error: "Daily request limit exceeded", error: "Daily request limit exceeded",
used: dailyReqCount, used: reqCount,
limit: limits.requests, limit: limits.requests,
} }
} }
if (limits.tokens > 0 && dailyTokenCount >= limits.tokens) { if (limits.tokens > 0 && tokenCount >= limits.tokens) {
return { return {
allowed: false, allowed: false,
type: "token", type: "token",
error: "Daily token limit exceeded", error: "Daily token limit exceeded",
used: dailyTokenCount, used: tokenCount,
limit: limits.tokens, limit: limits.tokens,
} }
} }
@@ -210,7 +164,7 @@ export async function checkAndIncrementRequest(
} }
// Condition failed but no limit clearly exceeded - race condition edge case // Condition failed but no limit clearly exceeded - race condition edge case
// Fail safe by allowing (could be a reset race) // Fail safe by allowing (could be a TPM reset race)
console.warn( console.warn(
`[quota] Condition failed but no limit exceeded for IP prefix: ${ip.slice(0, 8)}...`, `[quota] Condition failed but no limit exceeded for IP prefix: ${ip.slice(0, 8)}...`,
) )
@@ -233,7 +187,7 @@ export async function checkAndIncrementRequest(
/** /**
* Record token usage after response completes. * Record token usage after response completes.
* Uses atomic operations to update both daily token count and TPM count. * Uses composite key (PK=user, SK=date) for per-day tracking.
* Handles minute boundaries atomically to prevent race conditions. * Handles minute boundaries atomically to prevent race conditions.
*/ */
export async function recordTokenUsage( export async function recordTokenUsage(
@@ -244,24 +198,27 @@ export async function recordTokenUsage(
if (!client || !TABLE) return if (!client || !TABLE) return
if (!Number.isFinite(tokens) || tokens <= 0) return if (!Number.isFinite(tokens) || tokens <= 0) return
const pk = ip // User identifier (base64 IP)
const sk = getTodayInTimezone() // Date as sort key (YYYY-MM-DD)
const currentMinute = Math.floor(Date.now() / 60000).toString() const currentMinute = Math.floor(Date.now() / 60000).toString()
const ttl = Math.floor(Date.now() / 1000) + 7 * 24 * 60 * 60
try { try {
// Try to update assuming same minute (most common case) // Try to update for same minute OR new item (most common cases)
// Uses condition to ensure we're in the same minute // Handles: 1) new item (no lastMinute), 2) same minute (lastMinute matches)
await client.send( await client.send(
new UpdateItemCommand({ new UpdateItemCommand({
TableName: TABLE, TableName: TABLE,
Key: { PK: { S: `IP#${ip}` } }, Key: {
PK: { S: pk },
SK: { S: sk },
},
UpdateExpression: UpdateExpression:
"SET #ttl = :ttl ADD dailyTokenCount :tokens, tpmCount :tokens", "SET lastMinute = if_not_exists(lastMinute, :minute) ADD tokenCount :tokens, tpmCount :tokens",
ConditionExpression: "lastMinute = :minute", ConditionExpression:
ExpressionAttributeNames: { "#ttl": "ttl" }, "attribute_not_exists(lastMinute) OR lastMinute = :minute",
ExpressionAttributeValues: { ExpressionAttributeValues: {
":minute": { S: currentMinute }, ":minute": { S: currentMinute },
":tokens": { N: String(tokens) }, ":tokens": { N: String(tokens) },
":ttl": { N: String(ttl) },
}, },
}), }),
) )
@@ -272,14 +229,15 @@ export async function recordTokenUsage(
await client.send( await client.send(
new UpdateItemCommand({ new UpdateItemCommand({
TableName: TABLE, TableName: TABLE,
Key: { PK: { S: `IP#${ip}` } }, Key: {
PK: { S: pk },
SK: { S: sk },
},
UpdateExpression: UpdateExpression:
"SET lastMinute = :minute, tpmCount = :tokens, #ttl = :ttl ADD dailyTokenCount :tokens", "SET lastMinute = :minute, tpmCount = :tokens ADD tokenCount :tokens",
ExpressionAttributeNames: { "#ttl": "ttl" },
ExpressionAttributeValues: { ExpressionAttributeValues: {
":minute": { S: currentMinute }, ":minute": { S: currentMinute },
":tokens": { N: String(tokens) }, ":tokens": { N: String(tokens) },
":ttl": { N: String(ttl) },
}, },
}), }),
) )