fix: optimized queries / dedupe

This commit is contained in:
John Yeo
2026-04-22 12:09:33 +01:00
parent 0bffd9d249
commit e75965d922
30 changed files with 994 additions and 243 deletions

View File

@@ -33,7 +33,7 @@
"**/.cursor": true,
// "**/.github": true,
"**/.superset": true,
".claude": true,
// ".claude": true,
".codex": true,
".agents": true,
".cursor": true,

View File

@@ -206,7 +206,7 @@
},
"packages/autumn-js": {
"name": "autumn-js",
"version": "1.1.7",
"version": "1.2.10",
"dependencies": {
"query-string": "^9.2.2",
"rou3": "^0.6.1",

View File

@@ -17,8 +17,6 @@ export const prodTestOrgId = requireEnv({ key: "PROD_TEST_ORG_ID" });
export const prodTestCustomerId = requireEnv({
key: "PROD_TEST_CUSTOMER_ID",
});
export const prodTestEntityId = requireEnv({
key: "PROD_TEST_ENTITY_ID",
});
export const prodTestEntityId = process.env.PROD_TEST_ENTITY_ID || undefined;
export const { initDrizzle } = await import("../src/db/initDrizzle");

View File

@@ -5,6 +5,7 @@ import {
prodTestCustomerId,
prodTestOrgId,
} from "./experimentEnv";
import { OrgService } from "@/internal/orgs/OrgService";
const { getEntityAggregateForSync } = await import(
"../src/internal/customers/repos/getFullSubject/getEntityAggregateForSync"
@@ -19,6 +20,15 @@ const main = async () => {
const { db } = initDrizzle();
const org = await OrgService.getWithFeatures({
db,
orgId,
env,
});
const features = org?.features.filter((feature) => feature.id.toLowerCase().includes("credit"));
console.log("--- Running entity aggregate query ---");
const start = performance.now();
const result = await getEntityAggregateForSync({
@@ -26,6 +36,7 @@ const main = async () => {
orgId,
env,
customerId,
internalFeatureIds: features?.map((feature) => feature.internal_id),
});
const elapsed = performance.now() - start;
console.log(`Rows returned: ${result.length}`);
@@ -39,7 +50,10 @@ const main = async () => {
);
const statusFilter = sql`AND cp.status = ANY(ARRAY['active', 'past_due', 'scheduled'])`;
const entityFragments = getEntityAggregateFragments({ statusFilter });
const entityFragments = getEntityAggregateFragments({
statusFilter,
internalFeatureIds: features?.map((feature) => feature.internal_id),
});
const query = sql`
WITH subject_customer_records AS (

View File

@@ -0,0 +1,61 @@
import { AppEnv } from "@autumn/shared";
import { sql } from "drizzle-orm";
import {
initDrizzle,
prodTestCustomerId,
prodTestEntityId,
prodTestOrgId,
} from "./experimentEnv";
const { getFullSubjectQuery } = await import(
"../src/internal/customers/repos/getFullSubject/getFullSubjectQuery"
);
const { RELEVANT_STATUSES } = await import(
"../src/internal/customers/cusProducts/CusProductService"
);
// Run with:
// bun run experiments/explainGetFullSubject.ts
// Or scoped to an entity:
// PROD_TEST_ENTITY_ID=... bun run experiments/explainGetFullSubject.ts
const main = async () => {
const orgId = prodTestOrgId;
const env = AppEnv.Live;
const customerId = prodTestCustomerId;
const entityId = prodTestEntityId;
const { db } = initDrizzle();
const query = getFullSubjectQuery({
orgId,
env,
customerId,
entityId,
inStatuses: RELEVANT_STATUSES,
});
console.log(
`--- Running full subject query (customer=${customerId}${entityId ? `, entity=${entityId}` : ""}) ---`,
);
const start = performance.now();
const result = await db.execute(query);
const elapsed = performance.now() - start;
console.log(`Rows returned: ${result.length}`);
console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms`);
console.log("Result:", JSON.stringify(result, null, 2));
console.log();
console.log("--- EXPLAIN (ANALYZE, BUFFERS) ---\n");
const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${query}`;
const explainResult = await db.execute(explainQuery);
for (const row of explainResult) {
const line = (row as Record<string, unknown>)["QUERY PLAN"];
console.log(line);
}
process.exit(0);
};
await main();

View File

@@ -39,10 +39,12 @@ local function init_context(params)
missing_customer_entitlement_ids =
read_result.missing_customer_entitlement_ids or {},
logs = logs,
logger = {
logger = params.debug and {
log = function(fmt, ...)
table.insert(logs, string.format(fmt, ...))
end,
} or {
log = function() end,
},
}

View File

@@ -124,6 +124,7 @@ local context = init_context({
customer_id = customer_id,
customer_entitlement_deductions = customer_entitlement_deductions,
balance_keys_by_feature_id = params.balance_keys_by_feature_id,
debug = params.debug,
})
if #(context.missing_customer_entitlement_ids or {}) > 0 then

View File

@@ -0,0 +1,118 @@
-- ============================================================================
-- LOCK RECEIPT HELPERS (V2 — plain string storage)
-- V2 variant of deduction/lock/lockReceipt.lua. Stores the lock receipt as a
-- plain string via core Redis GET/SET instead of RedisJSON JSON.GET/JSON.SET.
--
-- Why: fewer Redis calls per lock save (combined SET NX EXAT), no RedisJSON
-- module parse/serialize overhead. V2 full-subject cache has no legacy
-- RedisJSON lock receipts so plain GET is safe.
--
-- Function names intentionally match deduction/lock/lockReceipt.lua so scripts
-- that import these helpers (deductFromSubjectBalances.lua, unwindLockV2.lua,
-- claimLockReceiptV2.lua) don't need call-site changes — only the bundle swaps
-- which helper file is included.
--
-- Depends on:
-- - is_nil (luaUtils.lua)
-- ============================================================================
-- ============================================================================
-- HELPER: Load lock receipt from Redis
-- Returns the decoded receipt table, or nil if the key does not exist / is
-- unparseable.
-- ============================================================================
local function load_lock_receipt(lock_receipt_key)
local raw = redis.call('GET', lock_receipt_key)
if is_nil(raw) or raw == false then
return nil
end
local ok, decoded = pcall(cjson.decode, raw)
if not ok or type(decoded) ~= 'table' then
return nil
end
return decoded
end
-- ============================================================================
-- HELPER: Store lock receipt in Redis (full overwrite, no TTL change)
-- Use `save_lock_receipt_from_updates` when you want to set a TTL atomically.
-- Use this when updating an existing receipt whose TTL should be preserved —
-- callers should pair with `SET ... KEEPTTL` instead of this helper if they
-- want to keep the existing TTL; plain SET here clears it.
-- ============================================================================
local function store_lock_receipt(lock_receipt_key, receipt)
redis.call('SET', lock_receipt_key, cjson.encode(receipt))
return receipt
end
-- ============================================================================
-- HELPER: Overwrite an existing lock receipt and preserve its TTL.
-- Used by the claim path (pending -> processing) where we mutate one field and
-- write the full receipt back without resetting the expiry.
-- ============================================================================
local function store_lock_receipt_keep_ttl(lock_receipt_key, receipt)
redis.call('SET', lock_receipt_key, cjson.encode(receipt), 'KEEPTTL')
return receipt
end
-- ============================================================================
-- HELPER: Save a lock receipt from deduction update objects.
--
-- Combines the previous `JSON.SET + EXPIREAT` into a single `SET ... EXAT` so
-- the whole operation is one Redis round trip. When `ttl_at` is nil the TTL is
-- omitted (the key will persist until explicitly deleted or overwritten).
--
-- params:
-- lock_receipt_key: string
-- receipt: table (base receipt metadata to persist)
-- mutation_logs: table | nil
-- ttl_at: number | nil (Unix seconds for EXAT)
-- ============================================================================
local function save_lock_receipt_from_updates(params)
local receipt = params.receipt or {}
local mutation_logs = params.mutation_logs or {}
receipt.items = #mutation_logs > 0 and mutation_logs or cjson.decode('[]')
local encoded = cjson.encode(receipt)
if not is_nil(params.ttl_at) then
redis.call('SET', params.lock_receipt_key, encoded, 'EXAT', params.ttl_at)
else
redis.call('SET', params.lock_receipt_key, encoded)
end
return receipt
end
-- ============================================================================
-- HELPER: Atomically create a lock receipt if no receipt exists at that key.
-- Returns true on success, false if a receipt is already present.
--
-- This is the Lua analogue of the TS `SET NX EXAT` path in saveLockReceipt.ts.
-- Use when you want a single-RT "create if absent" (replaces EXISTS + SET +
-- EXPIREAT).
--
-- params:
-- lock_receipt_key: string
-- receipt: table
-- mutation_logs: table | nil
-- ttl_at: number | nil (Unix seconds for EXAT)
-- ============================================================================
local function create_lock_receipt_if_absent(params)
local receipt = params.receipt or {}
local mutation_logs = params.mutation_logs or {}
receipt.items = #mutation_logs > 0 and mutation_logs or cjson.decode('[]')
local encoded = cjson.encode(receipt)
local result
if not is_nil(params.ttl_at) then
result = redis.call('SET', params.lock_receipt_key, encoded, 'NX', 'EXAT', params.ttl_at)
else
result = redis.call('SET', params.lock_receipt_key, encoded, 'NX')
end
return result == 'OK'
end

View File

@@ -55,6 +55,7 @@ export const initDrizzle = ({
const envDbUrl = replica
? process.env.DATABASE_REPLICA_URL
: process.env.DATABASE_URL;
const dbUrl = databaseUrl || envDbUrl || "";
const client = new pg.Pool({

View File

@@ -3,17 +3,12 @@ import { createHash } from "node:crypto";
const hash = (value: string) =>
createHash("sha256").update(value).digest("hex").slice(0, 12);
const mask = (value = "") =>
value.length > 6
? `${value.slice(0, 3)}***${value.slice(-3)}`
: `${value[0] ?? ""}***${value.slice(-1)}`;
const auth = (url: URL) => {
const username = decodeURIComponent(url.username);
const password = decodeURIComponent(url.password);
if (!username && !password) return "";
return `${mask(username)}${password ? `:${mask(password)}` : ""}@`;
return `${username}${password ? ":***" : ""}@`;
};
const formatUrl = (value: string) => {

View File

@@ -28,10 +28,6 @@ export const triggerAutoTopUp = async ({
featureId: relevantFeature.id,
});
console.log(
`resolved, feature ${relevantFeature.id}, balance below threshold: ${resolved?.balanceBelowThreshold}, customerEntitlement balance: ${resolved?.customerEntitlement.balance}`,
);
if (!resolved?.balanceBelowThreshold) continue;
// Enqueue the auto top-up job

View File

@@ -29,7 +29,7 @@ export const runCheckWithRollout = async ({
body,
requiredBalance,
}),
catch: (error) => error,
catch: (error: unknown) => error,
});
if (Result.isOk(result)) return result.value;

View File

@@ -160,6 +160,7 @@ export const executeRedisDeductionV2 = async ({
}
: null,
unwind_value: unwindValue ?? null,
debug: process.env.NODE_ENV !== "production",
};
const targetRedis = redisInstance ?? ctx.redisV2;
@@ -262,6 +263,9 @@ export const executeRedisDeductionV2 = async ({
});
}
} catch (error) {
// if (error.message?.includes("declined")) {
// return;
// }
if (error instanceof Error && !error?.message?.includes("declined")) {
ctx.logger.error(
`[executeRedisDeductionV2] Attempting rollback due to error: ${error}`,

View File

@@ -0,0 +1,196 @@
import type { AppEnv } from "@autumn/shared";
import { logger } from "@/external/logtail/logtailUtils.js";
import { currentRegion } from "@/external/redis/initRedis.js";
import { JobName } from "@/queue/JobName.js";
import { addTaskToQueue } from "@/queue/queueUtils.js";
import {
buildRefreshEntityAggregateDedupId,
REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS,
REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS,
} from "./queueRefreshEntityAggregate.js";
export type QueueRefreshEntityAggregatePayload = {
jobName: string;
payload: {
customerId: string;
orgId: string;
env: AppEnv;
region: string;
internalFeatureIds: string[];
};
messageGroupId: string;
messageDeduplicationId: string;
};
interface RefreshEntry {
customerId: string;
orgId: string;
env: AppEnv;
internalFeatureIds: Set<string>;
timer: NodeJS.Timeout | null;
}
/**
* Coalesces `RefreshEntityAggregate` enqueues per (org, env, customer).
*
* All `schedule()` calls that arrive inside the same 5s bucket share one
* timer aimed at the bucket's trailing edge (+ settle buffer). When the timer
* fires, the manager emits a single SQS message with the merged
* `internalFeatureIds`. Because we're in the same worker process that
* consumes the customer's sync-v4 stream (FIFO MessageGroupId already
* serializes that), this in-memory dedup is effectively per-customer global.
*/
export class RefreshEntityAggregateBatchingManager {
private entries: Map<string, RefreshEntry> = new Map();
private readonly bucketMs: number;
private readonly settleBufferMs: number;
private readonly _addTaskToQueue: (
args: QueueRefreshEntityAggregatePayload,
) => Promise<void>;
private readonly _now: () => number;
constructor({
addTaskToQueueFn,
bucketMs,
settleBufferMs,
now,
}: {
addTaskToQueueFn?: (
args: QueueRefreshEntityAggregatePayload,
) => Promise<void>;
bucketMs?: number;
settleBufferMs?: number;
now?: () => number;
} = {}) {
this._addTaskToQueue =
addTaskToQueueFn ??
(addTaskToQueue as unknown as (
args: QueueRefreshEntityAggregatePayload,
) => Promise<void>);
this.bucketMs = bucketMs ?? REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS;
this.settleBufferMs =
settleBufferMs ?? REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS;
this._now = now ?? Date.now;
}
schedule({
orgId,
env,
customerId,
internalFeatureIds,
}: {
orgId: string;
env: AppEnv;
customerId: string;
internalFeatureIds: string[];
}): void {
const key = this.buildKey({ orgId, env, customerId });
const existing = this.entries.get(key);
if (existing) {
for (const id of internalFeatureIds) {
existing.internalFeatureIds.add(id);
}
return;
}
const entry: RefreshEntry = {
customerId,
orgId,
env,
internalFeatureIds: new Set(internalFeatureIds),
timer: null,
};
this.entries.set(key, entry);
this.scheduleTimer({ key, entry });
}
async flush(): Promise<void> {
const keys = Array.from(this.entries.keys());
await Promise.all(keys.map((key) => this.fire({ key })));
}
getStats(): { totalPending: number } {
return { totalPending: this.entries.size };
}
private buildKey({
orgId,
env,
customerId,
}: {
orgId: string;
env: AppEnv;
customerId: string;
}): string {
return `${orgId}:${env}:${customerId}`;
}
private scheduleTimer({
key,
entry,
}: {
key: string;
entry: RefreshEntry;
}): void {
const nowMs = this._now();
const bucketEndMs =
(Math.floor(nowMs / this.bucketMs) + 1) * this.bucketMs;
const delayMs = bucketEndMs - nowMs + this.settleBufferMs;
entry.timer = setTimeout(() => {
this.fire({ key });
}, delayMs);
if (entry.timer.unref) entry.timer.unref();
}
private async fire({ key }: { key: string }): Promise<void> {
const entry = this.entries.get(key);
if (!entry) return;
if (entry.timer) {
clearTimeout(entry.timer);
entry.timer = null;
}
this.entries.delete(key);
if (entry.internalFeatureIds.size === 0) return;
const internalFeatureIds = Array.from(entry.internalFeatureIds).sort();
const messageDeduplicationId = buildRefreshEntityAggregateDedupId({
orgId: entry.orgId,
env: entry.env,
customerId: entry.customerId,
nowMs: this._now(),
bucketMs: this.bucketMs,
});
try {
await this._addTaskToQueue({
jobName: JobName.RefreshEntityAggregate,
payload: {
customerId: entry.customerId,
orgId: entry.orgId,
env: entry.env,
region: currentRegion,
internalFeatureIds,
},
messageGroupId: `refresh-agg:${entry.orgId}:${entry.env}:${entry.customerId}`,
messageDeduplicationId,
});
logger.info(
`[RefreshEntityAggregate] Queued refresh for ${entry.customerId}, ${internalFeatureIds.length} features`,
);
} catch (error) {
logger.error(
`[RefreshEntityAggregate] Failed to queue refresh for ${entry.customerId}: ${error}`,
);
}
}
}
export const globalRefreshEntityAggregateBatchingManager =
new RefreshEntityAggregateBatchingManager();

View File

@@ -0,0 +1,11 @@
export {
buildRefreshEntityAggregateDedupId,
REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS,
REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS,
} from "./queueRefreshEntityAggregate.js";
export { refreshEntityAggregateCache } from "./refreshEntityAggregateCache.js";
export {
globalRefreshEntityAggregateBatchingManager,
RefreshEntityAggregateBatchingManager,
type QueueRefreshEntityAggregatePayload,
} from "./RefreshEntityAggregateBatchingManager.js";

View File

@@ -0,0 +1,40 @@
import type { AppEnv } from "@autumn/shared";
import { JobName } from "@/queue/JobName.js";
export const REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS =
process.env.NODE_ENV === "development" ? 1000 : 5000;
/**
* Buffer added after the bucket boundary so the trailing enqueue fires *after*
* the final sync batch of the bucket has settled into Postgres.
* `SyncBatchingManagerV3` uses a 1s tumbling window, so 1.5s is enough.
*/
export const REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS = 1500;
/**
* Deterministic dedup id per (org, env, customer) within a time bucket.
* Rapid schedules inside the same bucket collapse to a single SQS message.
*/
export const buildRefreshEntityAggregateDedupId = ({
orgId,
env,
customerId,
nowMs,
bucketMs = REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS,
}: {
orgId: string;
env: AppEnv;
customerId: string;
nowMs: number;
bucketMs?: number;
}): string => {
const bucket = Math.floor(nowMs / bucketMs);
const key = JSON.stringify({
jobName: JobName.RefreshEntityAggregate,
orgId,
env,
customerId,
bucket,
});
return Bun.hash(key).toString();
};

View File

@@ -1,9 +1,8 @@
import type { AppEnv } from "@autumn/shared";
import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
import { buildSharedFullSubjectBalanceKey } from "@/internal/customers/cache/fullSubject/builders/buildSharedFullSubjectBalanceKey.js";
import { AGGREGATED_BALANCE_FIELD } from "@/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.js";
import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js";
import { getEntityAggregateForSync } from "@/internal/customers/repos/getFullSubject/getEntityAggregateForSync.js";
import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js";
/**
* After DB sync, recompute entity aggregation from the now-authoritative DB
@@ -12,33 +11,65 @@ import { getEntityAggregateForSync } from "@/internal/customers/repos/getFullSub
export const refreshEntityAggregateCache = async ({
ctx,
customerId,
orgId,
env,
featureIds,
internalFeatureIds,
}: {
ctx: AutumnContext;
customerId: string;
orgId: string;
env: AppEnv;
featureIds: string[];
internalFeatureIds: string[];
}): Promise<void> => {
try {
const orgId = ctx.org.id;
const env = ctx.env;
const internalIdSet = new Set(internalFeatureIds);
const featureIds = ctx.features
.filter((feature) => internalIdSet.has(feature.internal_id))
.map((feature) => feature.id);
if (featureIds.length === 0) return;
const { redisV2 } = ctx;
// Only refresh features whose balance hash already has `_aggregated`.
// If none of the hashes have it cached, there is nothing to refresh —
// avoid the expensive CTE entirely.
const existsPipeline = redisV2.pipeline();
for (const featureId of featureIds) {
const balanceKey = buildSharedFullSubjectBalanceKey({
orgId,
env,
customerId,
featureId,
});
existsPipeline.hexists(balanceKey, AGGREGATED_BALANCE_FIELD);
}
const existsResults = (await existsPipeline.exec()) ?? [];
const featuresWithAggregated = new Set<string>();
existsResults.forEach(([, exists], idx) => {
if (exists === 1) featuresWithAggregated.add(featureIds[idx]);
});
if (featuresWithAggregated.size === 0) {
ctx.logger.info(
`[SYNC V4] (${customerId}) No _aggregated fields cached — skipping refresh`,
);
return;
}
const aggregated = await getEntityAggregateForSync({
db: ctx.db,
orgId,
env,
customerId,
internalFeatureIds,
});
if (aggregated.length === 0) return;
const affectedFeatureIds = new Set(featureIds);
const { redisV2 } = ctx;
const pipeline = redisV2.pipeline();
let writeCount = 0;
for (const entry of aggregated) {
if (!affectedFeatureIds.has(entry.feature_id)) continue;
if (!featuresWithAggregated.has(entry.feature_id)) continue;
const balanceKey = buildSharedFullSubjectBalanceKey({
orgId,

View File

@@ -293,6 +293,7 @@ export class SyncBatchingManagerV3 {
entityId: context.entityId,
modifiedCusEntIdsByFeatureId: context.modifiedCusEntIdsByFeatureId,
},
messageGroupId: `sync-v4:${context.orgId}:${context.env}:${context.customerId}`,
messageDeduplicationId,
});

View File

@@ -7,9 +7,9 @@ import {
} from "@autumn/shared";
import { sql } from "drizzle-orm";
import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
import { deleteCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/deleteCachedFullCustomer.js";
import { getCachedFeatureBalance } from "@/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances.js";
import { refreshEntityAggregateCache } from "./refreshEntityAggregateCache.js";
import { deleteCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/deleteCachedFullCustomer.js";
import { globalRefreshEntityAggregateBatchingManager } from "../refreshEntityAggregate/index.js";
const SYNC_CONFLICT_CODES = {
ResetAtMismatch: "RESET_AT_MISMATCH",
@@ -135,8 +135,6 @@ export const syncItemV4 = async ({
const {
customerId,
entityId,
orgId,
env,
rolloverIds,
modifiedCusEntIdsByFeatureId,
} = payload;
@@ -244,12 +242,16 @@ export const syncItemV4 = async ({
(subjectBalance) => subjectBalance.isEntityLevel,
);
if (hasEntityLevel) {
await refreshEntityAggregateCache({
ctx,
const featureIds = Object.keys(modifiedCusEntIdsByFeatureId);
const internalFeatureIds = ctx.features
.filter((feature) => featureIds.includes(feature.id))
.map((feature) => feature.internal_id);
globalRefreshEntityAggregateBatchingManager.schedule({
orgId: ctx.org.id,
env: ctx.env,
customerId,
orgId,
env,
featureIds: Object.keys(modifiedCusEntIdsByFeatureId),
internalFeatureIds,
});
}
};

View File

@@ -20,12 +20,14 @@ export const getEntityAggregateForSync = async ({
orgId,
env,
customerId,
internalFeatureIds,
inStatuses = RELEVANT_STATUSES,
}: {
db: DrizzleCli;
orgId: string;
env: AppEnv;
customerId: string;
internalFeatureIds?: string[];
inStatuses?: CusProductStatus[];
}): Promise<AggregatedFeatureBalance[]> => {
const statusFilter =
@@ -38,6 +40,7 @@ export const getEntityAggregateForSync = async ({
const entityFragments = getEntityAggregateFragments({
statusFilter,
internalFeatureIds,
});
const query = sql`

View File

@@ -2,20 +2,27 @@ import { type SQL, sql } from "drizzle-orm";
import { getEntityOptionsAggregateFragments } from "./getEntityOptionsAggregateFragments.js";
/**
* Per-entity rollover rows sourced from the `rollovers` table, mirroring the
* four branches in `entity_balance_rows` (product-attached / loose, per-entity
* jsonb / top-level). Top-level rollovers are attributed to the owning entity
* via `cp.internal_entity_id` (product-attached) or `ce.internal_entity_id`
* (loose), matching main-balance behaviour. Also exposes per-entity and
* per-feature rollups used by the outer aggregate CTEs.
* Rollover CTEs driven from the shared `entity_product_cus_ents` and
* `entity_loose_cus_ents` base CTEs — avoids re-scanning `customer_entitlements`
* and `customer_products` per branch.
*/
const buildEntityRolloverCtes = ({
statusFilter,
}: {
statusFilter: SQL;
}) => sql`
const buildEntityRolloverCtes = () => sql`
entity_rollover_rows AS (
-- Product-attached cusEnt, per-entity rollover (rollovers.entities jsonb)
-- Top-level rollover: one row per (rollover × cus_ent).
-- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose.
SELECT
ce.internal_feature_id,
ce.internal_customer_id,
COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key,
r.balance::numeric AS rollover_balance,
COALESCE(r.usage, 0)::numeric AS rollover_usage
FROM rollovers r
JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id
WHERE (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
UNION ALL
-- Per-entity rollover from jsonb_each(r.entities).
SELECT
ce.internal_feature_id,
ce.internal_customer_id,
@@ -23,67 +30,10 @@ const buildEntityRolloverCtes = ({
COALESCE((kv.entity_value->>'balance')::numeric, 0) AS rollover_balance,
COALESCE((kv.entity_value->>'usage')::numeric, 0) AS rollover_usage
FROM rollovers r
JOIN customer_entitlements ce ON r.cus_ent_id = ce.id
JOIN customer_products cp ON ce.customer_product_id = cp.id
JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id
CROSS JOIN LATERAL jsonb_each(r.entities) AS kv(entity_key, entity_value)
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND cp.internal_entity_id IS NOT NULL
AND jsonb_typeof(r.entities) = 'object'
WHERE jsonb_typeof(r.entities) = 'object'
AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
${statusFilter}
UNION ALL
-- Product-attached cusEnt, top-level rollover (attributed to cp.internal_entity_id)
SELECT
ce.internal_feature_id,
ce.internal_customer_id,
cp.internal_entity_id AS entity_key,
r.balance::numeric AS rollover_balance,
COALESCE(r.usage, 0)::numeric AS rollover_usage
FROM rollovers r
JOIN customer_entitlements ce ON r.cus_ent_id = ce.id
JOIN customer_products cp ON ce.customer_product_id = cp.id
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND cp.internal_entity_id IS NOT NULL
AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
${statusFilter}
UNION ALL
-- Loose cusEnt (no customer_product), per-entity rollover
SELECT
ce.internal_feature_id,
ce.internal_customer_id,
kv.entity_key AS entity_key,
COALESCE((kv.entity_value->>'balance')::numeric, 0) AS rollover_balance,
COALESCE((kv.entity_value->>'usage')::numeric, 0) AS rollover_usage
FROM rollovers r
JOIN customer_entitlements ce ON r.cus_ent_id = ce.id
CROSS JOIN LATERAL jsonb_each(r.entities) AS kv(entity_key, entity_value)
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND ce.customer_product_id IS NULL
AND ce.internal_entity_id IS NOT NULL
AND jsonb_typeof(r.entities) = 'object'
AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
UNION ALL
-- Loose cusEnt, top-level rollover (attributed to ce.internal_entity_id)
SELECT
ce.internal_feature_id,
ce.internal_customer_id,
ce.internal_entity_id AS entity_key,
r.balance::numeric AS rollover_balance,
COALESCE(r.usage, 0)::numeric AS rollover_usage
FROM rollovers r
JOIN customer_entitlements ce ON r.cus_ent_id = ce.id
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND ce.customer_product_id IS NULL
AND ce.internal_entity_id IS NOT NULL
AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
),
entity_rollover_keys AS (
@@ -112,10 +62,20 @@ const buildEntityRolloverCtes = ({
export const getEntityAggregateFragments = ({
entityId,
statusFilter,
internalFeatureIds,
}: {
entityId?: string;
statusFilter: SQL;
internalFeatureIds?: string[];
}) => {
const featureFilter =
internalFeatureIds && internalFeatureIds.length > 0
? sql`AND ce.internal_feature_id = ANY(ARRAY[${sql.join(
internalFeatureIds.map((internalFeatureId) => sql`${internalFeatureId}`),
sql`, `,
)}])`
: sql``;
if (entityId) {
return {
ctes: sql``,
@@ -131,37 +91,50 @@ export const getEntityAggregateFragments = ({
const ctes = sql`,
entity_distinct_product_ids AS (
SELECT DISTINCT cp.internal_product_id, cp.internal_customer_id
FROM customer_products cp
JOIN subject_customer_records scr
ON cp.internal_customer_id = scr.internal_id
WHERE cp.internal_entity_id IS NOT NULL
${statusFilter}
),
entity_distinct_cus_products AS (
SELECT sub.*
FROM entity_distinct_product_ids edpi
JOIN LATERAL (
-- (A) Subject's entity-level customer_products, filtered once.
entity_cus_products AS (
SELECT cp.*
FROM customer_products cp
WHERE cp.internal_customer_id = edpi.internal_customer_id
AND cp.internal_product_id = edpi.internal_product_id
WHERE cp.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND cp.internal_entity_id IS NOT NULL
${statusFilter}
ORDER BY cp.created_at DESC
LIMIT 1
) sub ON true
),
entity_cus_products_for_options AS (
SELECT cp.*
FROM customer_products cp
JOIN subject_customer_records scr
ON cp.internal_customer_id = scr.internal_id
WHERE cp.internal_entity_id IS NOT NULL
${statusFilter}
-- (B) All entity-level cus_ents for the subject — product-attached (with
-- cp_entity_key set) and loose (cp_entity_key = NULL). Filtered once.
entity_level_cus_ents AS (
SELECT ce.*, cp.internal_entity_id AS cp_entity_key
FROM entity_cus_products cp
JOIN customer_entitlements ce ON ce.customer_product_id = cp.id
WHERE 1 = 1
${featureFilter}
UNION ALL
SELECT ce.*, NULL::text AS cp_entity_key
FROM customer_entitlements ce
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND ce.customer_product_id IS NULL
AND ce.internal_entity_id IS NOT NULL
AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
AND ce.balance != 0
${featureFilter}
),
-- Most-recent customer_product per (customer, product) — replaces the old
-- DISTINCT + LATERAL dance in entity_distinct_cus_products.
entity_distinct_cus_products AS (
SELECT *
FROM (
SELECT
cp.*,
ROW_NUMBER() OVER (
PARTITION BY cp.internal_customer_id, cp.internal_product_id
ORDER BY cp.created_at DESC
) AS rn
FROM entity_cus_products cp
) ranked
WHERE ranked.rn = 1
),
entity_cus_prices AS (
@@ -173,82 +146,8 @@ export const getEntityAggregateFragments = ({
${entityOptionsAggregateFragments.ctes},
entity_balance_rows AS (
SELECT
COALESCE(ce.external_id, ce.id) AS api_id,
ce.internal_feature_id,
ce.internal_customer_id,
ce.feature_id,
COALESCE(ent.allowance, 0)::numeric AS allowance,
ce.balance::numeric AS balance,
ce.adjustment::numeric AS adjustment,
COALESCE(ce.additional_balance, 0)::numeric AS additional_balance,
ce.unlimited,
ce.usage_allowed,
cp.internal_entity_id AS entity_key,
ce.balance::numeric AS entity_balance,
COALESCE(ce.adjustment, 0)::numeric AS entity_adjustment,
COALESCE(ce.additional_balance, 0)::numeric AS entity_additional_balance
FROM customer_entitlements ce
JOIN customer_products cp ON ce.customer_product_id = cp.id
JOIN entitlements ent ON ce.entitlement_id = ent.id
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND cp.internal_entity_id IS NOT NULL
${statusFilter}
UNION ALL
SELECT
COALESCE(ce.external_id, ce.id) AS api_id,
ce.internal_feature_id,
ce.internal_customer_id,
ce.feature_id,
COALESCE(ent.allowance, 0)::numeric AS allowance,
0::numeric AS balance,
0::numeric AS adjustment,
0::numeric AS additional_balance,
ce.unlimited,
ce.usage_allowed,
kv.entity_key AS entity_key,
(kv.entity_value->>'balance')::numeric AS entity_balance,
COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment,
COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance
FROM customer_entitlements ce
JOIN customer_products cp ON ce.customer_product_id = cp.id
JOIN entitlements ent ON ce.entitlement_id = ent.id
CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value)
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND cp.internal_entity_id IS NOT NULL
AND jsonb_typeof(ce.entities) = 'object'
${statusFilter}
UNION ALL
SELECT
COALESCE(ce.external_id, ce.id) AS api_id,
ce.internal_feature_id,
ce.internal_customer_id,
ce.feature_id,
COALESCE(ent.allowance, 0)::numeric AS allowance,
0::numeric AS balance,
0::numeric AS adjustment,
0::numeric AS additional_balance,
ce.unlimited,
ce.usage_allowed,
kv.entity_key AS entity_key,
(kv.entity_value->>'balance')::numeric AS entity_balance,
COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment,
COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance
FROM customer_entitlements ce
JOIN entitlements ent ON ce.entitlement_id = ent.id
CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value)
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND ce.customer_product_id IS NULL
AND ce.internal_entity_id IS NOT NULL
AND jsonb_typeof(ce.entities) = 'object'
AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
UNION ALL
-- Top-level: one row per cus_ent (product-attached or loose).
-- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose.
SELECT
COALESCE(ce.external_id, ce.id) AS api_id,
ce.internal_feature_id,
@@ -260,19 +159,39 @@ export const getEntityAggregateFragments = ({
COALESCE(ce.additional_balance, 0)::numeric AS additional_balance,
ce.unlimited,
ce.usage_allowed,
ce.internal_entity_id AS entity_key,
COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key,
ce.balance::numeric AS entity_balance,
COALESCE(ce.adjustment, 0)::numeric AS entity_adjustment,
COALESCE(ce.additional_balance, 0)::numeric AS entity_additional_balance
FROM customer_entitlements ce
FROM entity_level_cus_ents ce
JOIN entitlements ent ON ce.entitlement_id = ent.id
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND ce.customer_product_id IS NULL
AND ce.internal_entity_id IS NOT NULL
AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
UNION ALL
-- Per-entity: N rows per cus_ent from jsonb_each(ce.entities).
-- balance/adj/additional = 0 to avoid double-counting at the aggregate level.
SELECT
COALESCE(ce.external_id, ce.id) AS api_id,
ce.internal_feature_id,
ce.internal_customer_id,
ce.feature_id,
COALESCE(ent.allowance, 0)::numeric AS allowance,
0::numeric AS balance,
0::numeric AS adjustment,
0::numeric AS additional_balance,
ce.unlimited,
ce.usage_allowed,
kv.entity_key AS entity_key,
(kv.entity_value->>'balance')::numeric AS entity_balance,
COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment,
COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance
FROM entity_level_cus_ents ce
JOIN entitlements ent ON ce.entitlement_id = ent.id
CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value)
WHERE jsonb_typeof(ce.entities) = 'object'
),
${buildEntityRolloverCtes({ statusFilter })},
${buildEntityRolloverCtes()},
entity_balance_keys AS (
SELECT
@@ -389,21 +308,7 @@ export const getEntityAggregateFragments = ({
SELECT DISTINCT
ce.internal_customer_id,
ce.entitlement_id
FROM customer_entitlements ce
JOIN customer_products cp ON ce.customer_product_id = cp.id
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND cp.internal_entity_id IS NOT NULL
${statusFilter}
UNION
SELECT DISTINCT
ce.internal_customer_id,
ce.entitlement_id
FROM customer_entitlements ce
WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records)
AND ce.customer_product_id IS NULL
AND ce.internal_entity_id IS NOT NULL
AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
FROM entity_level_cus_ents ce
`;
const priceRefsUnion = sql`

View File

@@ -19,7 +19,7 @@ export const getEntityOptionsAggregateFragments = () => {
NULLIF(option_row.option_value->>'internal_feature_id', '') AS option_internal_feature_id,
NULLIF(option_row.option_value->>'feature_id', '') AS option_feature_id,
COALESCE((option_row.option_value->>'quantity')::numeric, 0) AS option_quantity
FROM entity_cus_products_for_options ecp
FROM entity_cus_products ecp
CROSS JOIN LATERAL unnest(
COALESCE(ecp.options, ARRAY[]::jsonb[])
) AS option_row(option_value)
@@ -35,7 +35,7 @@ export const getEntityOptionsAggregateFragments = () => {
* COALESCE((prepaid_price.config->>'billing_units')::numeric, 1)
AS prepaid_grant
FROM entity_option_rows eor
JOIN customer_entitlements ce
JOIN entity_level_cus_ents ce
ON ce.customer_product_id = eor.customer_product_id
JOIN entitlements ent
ON ent.id = ce.entitlement_id
@@ -83,4 +83,3 @@ export const getEntityOptionsAggregateFragments = () => {
ctes,
};
};

View File

@@ -232,7 +232,9 @@ export const getFullSubjectQuery = ({
ON ce.internal_customer_id = scr.internal_id
WHERE ce.customer_product_id IS NULL
AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000)
AND ce.balance != 0
${extraCustomerEntitlementEntityFilter}
LIMIT 20
),
all_cus_ent_ids AS (

View File

@@ -19,6 +19,7 @@ export enum JobName {
SyncBalanceBatchV2 = "sync-balance-batch-v2",
SyncBalanceBatchV3 = "sync-balance-batch-v3",
SyncBalanceBatchV4 = "sync-balance-batch-v4",
RefreshEntityAggregate = "refresh-entity-aggregate",
InsertEventBatch = "insert-event-batch",
Track = "track",

View File

@@ -10,6 +10,7 @@ import { runActionHandlerTask } from "@/internal/analytics/runActionHandlerTask.
import { autoTopup } from "@/internal/balances/autoTopUp/autoTopup.js";
import { runInsertEventBatch } from "@/internal/balances/events/runInsertEventBatch.js";
import { expireLock } from "@/internal/balances/finalizeLock/expireLock.js";
import { refreshEntityAggregateCache } from "@/internal/balances/utils/refreshEntityAggregate/index.js";
import { syncItemV3 } from "@/internal/balances/utils/sync/syncItemV3.js";
import { syncItemV4 } from "@/internal/balances/utils/sync/syncItemV4.js";
import { grantCheckoutReward } from "@/internal/billing/v2/workflows/grantCheckoutReward/grantCheckoutReward.js";
@@ -174,6 +175,20 @@ export const processMessage = async ({
return;
}
if (job.name === JobName.RefreshEntityAggregate) {
if (!ctx) {
workerLogger.error("No context found for refresh entity aggregate job");
return;
}
await refreshEntityAggregateCache({
ctx,
customerId: job.data.customerId,
internalFeatureIds: job.data.internalFeatureIds,
});
return;
}
if (job.name === JobName.InsertEventBatch) {
await runInsertEventBatch({
db,
@@ -278,7 +293,8 @@ export const processMessage = async ({
// won't fix on retry. DB errors (connection, timeout) will.
if (
(job.name === JobName.SyncBalanceBatchV3 ||
job.name === JobName.SyncBalanceBatchV4) &&
job.name === JobName.SyncBalanceBatchV4 ||
job.name === JobName.RefreshEntityAggregate) &&
isRetryableDbError({ error })
) {
Sentry.captureException(error);

View File

@@ -56,6 +56,13 @@ export interface Payloads {
entityId?: string;
modifiedCusEntIdsByFeatureId: Record<string, string[]>;
};
[JobName.RefreshEntityAggregate]: {
customerId: string;
orgId: string;
env: AppEnv;
region?: string;
internalFeatureIds: string[];
};
[JobName.InsertEventBatch]: {
events: EventInsert[];
};

View File

@@ -26,6 +26,7 @@ if (cluster.isPrimary) {
// await initHatchetWorker();
console.log(`Starting ${NUM_PROCESSES} worker processes`);
console.log(`SQS URL: ${process.env.SQS_QUEUE_URL}`);
// Fork workers
for (let i = 0; i < NUM_PROCESSES; i++) {

View File

@@ -1,7 +1,6 @@
import { expect, test } from "bun:test";
import type { ApiCustomerV3, TrackResponseV2 } from "@autumn/shared";
import { track } from "@tests/_groups/domains/balances/track";
import { TestFeature } from "@tests/setup/v2Features.js";
import { items } from "@tests/utils/fixtures/items.js";
import { products } from "@tests/utils/fixtures/products.js";

View File

@@ -0,0 +1,220 @@
import { describe, expect, test } from "bun:test";
import { AppEnv } from "@autumn/shared";
import chalk from "chalk";
import {
buildRefreshEntityAggregateDedupId,
type QueueRefreshEntityAggregatePayload,
REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS,
REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS,
RefreshEntityAggregateBatchingManager,
} from "@/internal/balances/utils/refreshEntityAggregate/index.js";
import { JobName } from "@/queue/JobName.js";
const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
const createMockQueue = () => {
const calls: QueueRefreshEntityAggregatePayload[] = [];
const fn = async (args: QueueRefreshEntityAggregatePayload) => {
calls.push(structuredClone(args));
};
return { fn, calls };
};
const baseArgs = {
orgId: "org-1",
env: AppEnv.Sandbox,
customerId: "cust-1",
internalFeatureIds: ["if-seats"],
};
describe("buildRefreshEntityAggregateDedupId", () => {
test(
`${chalk.yellowBright("dedup-id-1: same bucket → same id")}`,
() => {
const t0 = 1_700_000_000_000;
const a = buildRefreshEntityAggregateDedupId({
orgId: "org-1",
env: AppEnv.Sandbox,
customerId: "cust-1",
nowMs: t0,
});
const b = buildRefreshEntityAggregateDedupId({
orgId: "org-1",
env: AppEnv.Sandbox,
customerId: "cust-1",
nowMs: t0 + 4999,
});
expect(a).toBe(b);
},
);
test(
`${chalk.yellowBright("dedup-id-2: across bucket boundary → different ids")}`,
() => {
const t0 = 1_700_000_000_000;
const a = buildRefreshEntityAggregateDedupId({
orgId: "org-1",
env: AppEnv.Sandbox,
customerId: "cust-1",
nowMs: t0,
});
const b = buildRefreshEntityAggregateDedupId({
orgId: "org-1",
env: AppEnv.Sandbox,
customerId: "cust-1",
nowMs: t0 + REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS,
});
expect(a).not.toBe(b);
},
);
test(
`${chalk.yellowBright("dedup-id-3: different orgs/envs/customers → different ids")}`,
() => {
const nowMs = 1_700_000_000_000;
const ids = new Set([
buildRefreshEntityAggregateDedupId({
orgId: "org-1",
env: AppEnv.Sandbox,
customerId: "cust-1",
nowMs,
}),
buildRefreshEntityAggregateDedupId({
orgId: "org-2",
env: AppEnv.Sandbox,
customerId: "cust-1",
nowMs,
}),
buildRefreshEntityAggregateDedupId({
orgId: "org-1",
env: AppEnv.Live,
customerId: "cust-1",
nowMs,
}),
buildRefreshEntityAggregateDedupId({
orgId: "org-1",
env: AppEnv.Sandbox,
customerId: "cust-2",
nowMs,
}),
]);
expect(ids.size).toBe(4);
},
);
});
describe("RefreshEntityAggregateBatchingManager", () => {
test(
`${chalk.yellowBright("batch-1: rapid schedule() calls in one bucket → 1 enqueue with merged features")}`,
async () => {
const { fn, calls } = createMockQueue();
const manager = new RefreshEntityAggregateBatchingManager({
addTaskToQueueFn: fn,
bucketMs: 100,
settleBufferMs: 20,
});
for (let i = 0; i < 20; i++) {
manager.schedule({
...baseArgs,
internalFeatureIds: [`if-${i % 3}`],
});
}
expect(calls.length).toBe(0);
await wait(200);
expect(calls.length).toBe(1);
expect(calls[0].jobName).toBe(JobName.RefreshEntityAggregate);
expect(calls[0].payload.customerId).toBe(baseArgs.customerId);
expect(calls[0].payload.internalFeatureIds.sort()).toEqual([
"if-0",
"if-1",
"if-2",
]);
expect(calls[0].messageGroupId).toBe(
`refresh-agg:${baseArgs.orgId}:${baseArgs.env}:${baseArgs.customerId}`,
);
},
{ timeout: 5_000 },
);
test(
`${chalk.yellowBright("batch-2: schedules across a bucket boundary → 2 enqueues")}`,
async () => {
const { fn, calls } = createMockQueue();
const manager = new RefreshEntityAggregateBatchingManager({
addTaskToQueueFn: fn,
bucketMs: 100,
settleBufferMs: 20,
});
manager.schedule(baseArgs);
await wait(180);
manager.schedule(baseArgs);
await wait(180);
expect(calls.length).toBe(2);
expect(calls[0].messageDeduplicationId).not.toBe(
calls[1].messageDeduplicationId,
);
},
{ timeout: 5_000 },
);
test(
`${chalk.yellowBright("batch-3: different customers fire independently")}`,
async () => {
const { fn, calls } = createMockQueue();
const manager = new RefreshEntityAggregateBatchingManager({
addTaskToQueueFn: fn,
bucketMs: 100,
settleBufferMs: 20,
});
manager.schedule({ ...baseArgs, customerId: "cust-A" });
manager.schedule({ ...baseArgs, customerId: "cust-B" });
expect(manager.getStats().totalPending).toBe(2);
await wait(200);
expect(calls.length).toBe(2);
const customerIds = calls.map((c) => c.payload.customerId).sort();
expect(customerIds).toEqual(["cust-A", "cust-B"]);
},
{ timeout: 5_000 },
);
test(
`${chalk.yellowBright("batch-4: flush() drains all pending immediately")}`,
async () => {
const { fn, calls } = createMockQueue();
const manager = new RefreshEntityAggregateBatchingManager({
addTaskToQueueFn: fn,
bucketMs: 10_000, // Will never fire naturally
settleBufferMs: 1000,
});
manager.schedule({ ...baseArgs, customerId: "cust-A" });
manager.schedule({ ...baseArgs, customerId: "cust-B" });
expect(calls.length).toBe(0);
await manager.flush();
expect(calls.length).toBe(2);
expect(manager.getStats().totalPending).toBe(0);
},
{ timeout: 5_000 },
);
test(
`${chalk.yellowBright("batch-5: settle buffer honored")}`,
() => {
expect(REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS).toBe(5000);
expect(REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS).toBe(1500);
},
);
});

View File

@@ -0,0 +1,127 @@
/**
* Verifies that 20 concurrent tracks spread across ~3 seconds result in
* exactly ONE `RefreshEntityAggregate` enqueue to SQS at the end of the
* bucket, thanks to the producer-side `RefreshEntityAggregateBatchingManager`.
*/
import { expect, mock, test } from "bun:test";
import chalk from "chalk";
type QueueCall = {
jobName: string;
payload: Record<string, unknown>;
messageGroupId?: string;
messageDeduplicationId?: string;
};
const queueCalls: QueueCall[] = [];
mock.module("@/queue/queueUtils.js", () => ({
addTaskToQueue: async (args: QueueCall) => {
queueCalls.push(structuredClone(args));
},
}));
const { JobName } = await import("@/queue/JobName.js");
const { TestFeature } = await import("@tests/setup/v2Features.js");
const { items } = await import("@tests/utils/fixtures/items.js");
const { products } = await import("@tests/utils/fixtures/products.js");
const { timeout } = await import("@tests/utils/genUtils.js");
const { initScenario, s } = await import(
"@tests/utils/testInitUtils/initScenario.js"
);
const { globalRefreshEntityAggregateBatchingManager } = await import(
"@/internal/balances/utils/refreshEntityAggregate/index.js"
);
const { globalSyncBatchingManagerV3 } = await import(
"@/internal/balances/utils/sync/SyncBatchingManagerV3.js"
);
const { syncItemV4 } = await import(
"@/internal/balances/utils/sync/syncItemV4.js"
);
test(
`${chalk.yellowBright(
"refresh-dedup-track: 20 concurrent tracks across 3s → exactly 1 RefreshEntityAggregate enqueue",
)}`,
async () => {
const customerId = "refresh-agg-dedup-cus";
const perEntityMessages = items.monthlyMessages({
includedUsage: 500,
entityFeatureId: TestFeature.Users,
});
const prod = products.base({
id: "refresh-agg-dedup",
items: [perEntityMessages],
});
const { autumnV2_1, ctx, entities } = await initScenario({
customerId,
setup: [
s.customer({ testClock: false }),
s.products({ list: [prod] }),
s.entities({ count: 1, featureId: TestFeature.Users }),
],
actions: [s.attach({ productId: prod.id, entityIndex: 0 })],
});
// Discard enqueues produced during setup.
queueCalls.length = 0;
const trackPromises: Promise<unknown>[] = [];
for (let i = 0; i < 20; i++) {
trackPromises.push(
(async () => {
await timeout(i * 150); // 20 * 150ms ≈ 3s spread
return autumnV2_1.track({
customer_id: customerId,
entity_id: entities[0].id,
feature_id: TestFeature.Messages,
value: 1,
});
})(),
);
}
await Promise.all(trackPromises);
// Drain pending sync batches → SyncBalanceBatchV4 enqueues.
await globalSyncBatchingManagerV3.flush();
// Simulate the worker processing each sync-v4 job in this process
// (there's no live worker inside `bun test`). Each call schedules a
// refresh with the batching manager.
const syncJobs = queueCalls.filter(
(call) => call.jobName === JobName.SyncBalanceBatchV4,
);
expect(syncJobs.length).toBeGreaterThanOrEqual(1);
for (const job of syncJobs) {
await syncItemV4({
ctx: ctx as never,
payload: job.payload as never,
});
}
// Now drain the refresh batching manager — this is what the trailing
// timer would otherwise do at bucket end + settle buffer. Flushing
// directly keeps the test fast and deterministic.
await globalRefreshEntityAggregateBatchingManager.flush();
const refreshCalls = queueCalls.filter(
(call) => call.jobName === JobName.RefreshEntityAggregate,
);
expect(refreshCalls.length).toBe(1);
expect(refreshCalls[0].payload).toMatchObject({
customerId,
orgId: ctx.org.id,
env: ctx.env,
});
expect(refreshCalls[0].messageGroupId).toBe(
`refresh-agg:${ctx.org.id}:${ctx.env}:${customerId}`,
);
},
60_000,
);