mirror of
https://github.com/fosrl/pangolin.git
synced 2026-02-01 07:39:09 +00:00
689 lines
23 KiB
TypeScript
689 lines
23 KiB
TypeScript
import { eq, sql, and } from "drizzle-orm";
|
|
import { v4 as uuidv4 } from "uuid";
|
|
import { PutObjectCommand } from "@aws-sdk/client-s3";
|
|
import {
|
|
db,
|
|
usage,
|
|
customers,
|
|
sites,
|
|
newts,
|
|
limits,
|
|
Usage,
|
|
Limit,
|
|
Transaction
|
|
} from "@server/db";
|
|
import { FeatureId, getFeatureMeterId } from "./features";
|
|
import logger from "@server/logger";
|
|
import { sendToClient } from "#dynamic/routers/ws";
|
|
import { build } from "@server/build";
|
|
import { s3Client } from "@server/lib/s3";
|
|
import cache from "@server/lib/cache";
|
|
|
|
interface StripeEvent {
|
|
identifier?: string;
|
|
timestamp: number;
|
|
event_name: string;
|
|
payload: {
|
|
value: number;
|
|
stripe_customer_id: string;
|
|
};
|
|
}
|
|
|
|
export function noop() {
|
|
if (
|
|
build !== "saas" ||
|
|
!process.env.S3_BUCKET
|
|
) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
export class UsageService {
|
|
private bucketName: string | undefined;
|
|
private events: StripeEvent[] = [];
|
|
private lastUploadTime: number = Date.now();
|
|
private isUploading: boolean = false;
|
|
|
|
constructor() {
|
|
if (noop()) {
|
|
return;
|
|
}
|
|
this.bucketName = process.env.S3_BUCKET || undefined;
|
|
|
|
// Periodically check and upload events
|
|
setInterval(() => {
|
|
this.checkAndUploadEvents().catch((err) => {
|
|
logger.error("Error in periodic event upload:", err);
|
|
});
|
|
}, 30000); // every 30 seconds
|
|
|
|
// Handle graceful shutdown on SIGTERM
|
|
process.on("SIGTERM", async () => {
|
|
logger.info("SIGTERM received, uploading events before shutdown...");
|
|
await this.forceUpload();
|
|
logger.info("Events uploaded, proceeding with shutdown");
|
|
});
|
|
|
|
// Handle SIGINT as well (Ctrl+C)
|
|
process.on("SIGINT", async () => {
|
|
logger.info("SIGINT received, uploading events before shutdown...");
|
|
await this.forceUpload();
|
|
logger.info("Events uploaded, proceeding with shutdown");
|
|
process.exit(0);
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Truncate a number to 11 decimal places to prevent precision issues
|
|
*/
|
|
private truncateValue(value: number): number {
|
|
return Math.round(value * 100000000000) / 100000000000; // 11 decimal places
|
|
}
|
|
|
|
public async add(
|
|
orgId: string,
|
|
featureId: FeatureId,
|
|
value: number,
|
|
transaction: any = null
|
|
): Promise<Usage | null> {
|
|
if (noop()) {
|
|
return null;
|
|
}
|
|
|
|
// Truncate value to 11 decimal places
|
|
value = this.truncateValue(value);
|
|
|
|
// Implement retry logic for deadlock handling
|
|
const maxRetries = 3;
|
|
let attempt = 0;
|
|
|
|
while (attempt <= maxRetries) {
|
|
try {
|
|
// Get subscription data for this org (with caching)
|
|
const customerId = await this.getCustomerId(orgId, featureId);
|
|
|
|
if (!customerId) {
|
|
logger.warn(
|
|
`No subscription data found for org ${orgId} and feature ${featureId}`
|
|
);
|
|
return null;
|
|
}
|
|
|
|
let usage;
|
|
if (transaction) {
|
|
usage = await this.internalAddUsage(
|
|
orgId,
|
|
featureId,
|
|
value,
|
|
transaction
|
|
);
|
|
} else {
|
|
await db.transaction(async (trx) => {
|
|
usage = await this.internalAddUsage(
|
|
orgId,
|
|
featureId,
|
|
value,
|
|
trx
|
|
);
|
|
});
|
|
}
|
|
|
|
// Log event for Stripe
|
|
await this.logStripeEvent(featureId, value, customerId);
|
|
|
|
return usage || null;
|
|
} catch (error: any) {
|
|
// Check if this is a deadlock error
|
|
const isDeadlock =
|
|
error?.code === "40P01" ||
|
|
error?.cause?.code === "40P01" ||
|
|
(error?.message && error.message.includes("deadlock"));
|
|
|
|
if (isDeadlock && attempt < maxRetries) {
|
|
attempt++;
|
|
// Exponential backoff with jitter: 50-150ms, 100-300ms, 200-600ms
|
|
const baseDelay = Math.pow(2, attempt - 1) * 50;
|
|
const jitter = Math.random() * baseDelay;
|
|
const delay = baseDelay + jitter;
|
|
|
|
logger.warn(
|
|
`Deadlock detected for ${orgId}/${featureId}, retrying attempt ${attempt}/${maxRetries} after ${delay.toFixed(0)}ms`
|
|
);
|
|
|
|
await new Promise((resolve) => setTimeout(resolve, delay));
|
|
continue;
|
|
}
|
|
|
|
logger.error(
|
|
`Failed to add usage for ${orgId}/${featureId} after ${attempt} attempts:`,
|
|
error
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
private async internalAddUsage(
|
|
orgId: string,
|
|
featureId: FeatureId,
|
|
value: number,
|
|
trx: Transaction
|
|
): Promise<Usage> {
|
|
// Truncate value to 11 decimal places
|
|
value = this.truncateValue(value);
|
|
|
|
const usageId = `${orgId}-${featureId}`;
|
|
const meterId = getFeatureMeterId(featureId);
|
|
|
|
// Use upsert: insert if not exists, otherwise increment
|
|
const [returnUsage] = await trx
|
|
.insert(usage)
|
|
.values({
|
|
usageId,
|
|
featureId,
|
|
orgId,
|
|
meterId,
|
|
latestValue: value,
|
|
updatedAt: Math.floor(Date.now() / 1000)
|
|
})
|
|
.onConflictDoUpdate({
|
|
target: usage.usageId,
|
|
set: {
|
|
latestValue: sql`${usage.latestValue} + ${value}`
|
|
}
|
|
})
|
|
.returning();
|
|
|
|
return returnUsage;
|
|
}
|
|
|
|
// Helper function to get today's date as string (YYYY-MM-DD)
|
|
getTodayDateString(): string {
|
|
return new Date().toISOString().split("T")[0];
|
|
}
|
|
|
|
// Helper function to get date string from Date object
|
|
getDateString(date: number): string {
|
|
return new Date(date * 1000).toISOString().split("T")[0];
|
|
}
|
|
|
|
async updateDaily(
|
|
orgId: string,
|
|
featureId: FeatureId,
|
|
value?: number,
|
|
customerId?: string
|
|
): Promise<void> {
|
|
if (noop()) {
|
|
return;
|
|
}
|
|
try {
|
|
if (!customerId) {
|
|
customerId =
|
|
(await this.getCustomerId(orgId, featureId)) || undefined;
|
|
if (!customerId) {
|
|
logger.warn(
|
|
`No subscription data found for org ${orgId} and feature ${featureId}`
|
|
);
|
|
return;
|
|
}
|
|
}
|
|
|
|
// Truncate value to 11 decimal places if provided
|
|
if (value !== undefined && value !== null) {
|
|
value = this.truncateValue(value);
|
|
}
|
|
|
|
const today = this.getTodayDateString();
|
|
|
|
let currentUsage: Usage | null = null;
|
|
|
|
await db.transaction(async (trx) => {
|
|
// Get existing meter record
|
|
const usageId = `${orgId}-${featureId}`;
|
|
// Get current usage record
|
|
[currentUsage] = await trx
|
|
.select()
|
|
.from(usage)
|
|
.where(eq(usage.usageId, usageId))
|
|
.limit(1);
|
|
|
|
if (currentUsage) {
|
|
const lastUpdateDate = this.getDateString(
|
|
currentUsage.updatedAt
|
|
);
|
|
const currentRunningTotal = currentUsage.latestValue;
|
|
const lastDailyValue = currentUsage.instantaneousValue || 0;
|
|
|
|
if (value == undefined || value === null) {
|
|
value = currentUsage.instantaneousValue || 0;
|
|
}
|
|
|
|
if (lastUpdateDate === today) {
|
|
// Same day update: replace the daily value
|
|
// Remove old daily value from running total, add new value
|
|
const newRunningTotal = this.truncateValue(
|
|
currentRunningTotal - lastDailyValue + value
|
|
);
|
|
|
|
await trx
|
|
.update(usage)
|
|
.set({
|
|
latestValue: newRunningTotal,
|
|
instantaneousValue: value,
|
|
updatedAt: Math.floor(Date.now() / 1000)
|
|
})
|
|
.where(eq(usage.usageId, usageId));
|
|
} else {
|
|
// New day: add to running total
|
|
const newRunningTotal = this.truncateValue(
|
|
currentRunningTotal + value
|
|
);
|
|
|
|
await trx
|
|
.update(usage)
|
|
.set({
|
|
latestValue: newRunningTotal,
|
|
instantaneousValue: value,
|
|
updatedAt: Math.floor(Date.now() / 1000)
|
|
})
|
|
.where(eq(usage.usageId, usageId));
|
|
}
|
|
} else {
|
|
// First record for this meter
|
|
const meterId = getFeatureMeterId(featureId);
|
|
const truncatedValue = this.truncateValue(value || 0);
|
|
await trx.insert(usage).values({
|
|
usageId,
|
|
featureId,
|
|
orgId,
|
|
meterId,
|
|
instantaneousValue: truncatedValue,
|
|
latestValue: truncatedValue,
|
|
updatedAt: Math.floor(Date.now() / 1000)
|
|
});
|
|
}
|
|
});
|
|
|
|
await this.logStripeEvent(featureId, value || 0, customerId);
|
|
} catch (error) {
|
|
logger.error(
|
|
`Failed to update daily usage for ${orgId}/${featureId}:`,
|
|
error
|
|
);
|
|
}
|
|
}
|
|
|
|
private async getCustomerId(
|
|
orgId: string,
|
|
featureId: FeatureId
|
|
): Promise<string | null> {
|
|
const cacheKey = `customer_${orgId}_${featureId}`;
|
|
const cached = cache.get<string>(cacheKey);
|
|
|
|
if (cached) {
|
|
return cached;
|
|
}
|
|
|
|
try {
|
|
// Query subscription data
|
|
const [customer] = await db
|
|
.select({
|
|
customerId: customers.customerId
|
|
})
|
|
.from(customers)
|
|
.where(eq(customers.orgId, orgId))
|
|
.limit(1);
|
|
|
|
if (!customer) {
|
|
return null;
|
|
}
|
|
|
|
const customerId = customer.customerId;
|
|
|
|
// Cache the result
|
|
cache.set(cacheKey, customerId, 300); // 5 minute TTL
|
|
|
|
return customerId;
|
|
} catch (error) {
|
|
logger.error(
|
|
`Failed to get subscription data for ${orgId}/${featureId}:`,
|
|
error
|
|
);
|
|
return null;
|
|
}
|
|
}
|
|
|
|
private async logStripeEvent(
|
|
featureId: FeatureId,
|
|
value: number,
|
|
customerId: string
|
|
): Promise<void> {
|
|
// Truncate value to 11 decimal places before sending to Stripe
|
|
const truncatedValue = this.truncateValue(value);
|
|
|
|
const event: StripeEvent = {
|
|
identifier: uuidv4(),
|
|
timestamp: Math.floor(new Date().getTime() / 1000),
|
|
event_name: featureId,
|
|
payload: {
|
|
value: truncatedValue,
|
|
stripe_customer_id: customerId
|
|
}
|
|
};
|
|
|
|
this.addEventToMemory(event);
|
|
await this.checkAndUploadEvents();
|
|
}
|
|
|
|
private addEventToMemory(event: StripeEvent): void {
|
|
if (!this.bucketName) {
|
|
logger.warn(
|
|
"S3 bucket name is not configured, skipping event storage."
|
|
);
|
|
return;
|
|
}
|
|
this.events.push(event);
|
|
}
|
|
|
|
private async checkAndUploadEvents(): Promise<void> {
|
|
const now = Date.now();
|
|
const timeSinceLastUpload = now - this.lastUploadTime;
|
|
|
|
// Check if at least 1 minute has passed since last upload
|
|
if (timeSinceLastUpload >= 60000 && this.events.length > 0) {
|
|
await this.uploadEventsToS3();
|
|
}
|
|
}
|
|
|
|
private async uploadEventsToS3(): Promise<void> {
|
|
if (!this.bucketName) {
|
|
logger.warn(
|
|
"S3 bucket name is not configured, skipping S3 upload."
|
|
);
|
|
return;
|
|
}
|
|
|
|
if (this.events.length === 0) {
|
|
return;
|
|
}
|
|
|
|
// Check if already uploading
|
|
if (this.isUploading) {
|
|
logger.debug("Already uploading events, skipping");
|
|
return;
|
|
}
|
|
|
|
this.isUploading = true;
|
|
|
|
try {
|
|
// Take a snapshot of current events and clear the array
|
|
const eventsToUpload = [...this.events];
|
|
this.events = [];
|
|
this.lastUploadTime = Date.now();
|
|
|
|
const fileName = this.generateEventFileName();
|
|
const fileContent = JSON.stringify(eventsToUpload, null, 2);
|
|
|
|
// Upload to S3
|
|
const uploadCommand = new PutObjectCommand({
|
|
Bucket: this.bucketName,
|
|
Key: fileName,
|
|
Body: fileContent,
|
|
ContentType: "application/json"
|
|
});
|
|
|
|
await s3Client.send(uploadCommand);
|
|
|
|
logger.info(
|
|
`Uploaded ${fileName} to S3 with ${eventsToUpload.length} events`
|
|
);
|
|
} catch (error) {
|
|
logger.error("Failed to upload events to S3:", error);
|
|
// Note: Events are lost if upload fails. In a production system,
|
|
// you might want to add the events back to the array or implement retry logic
|
|
} finally {
|
|
this.isUploading = false;
|
|
}
|
|
}
|
|
|
|
private generateEventFileName(): string {
|
|
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
|
|
const uuid = uuidv4().substring(0, 8);
|
|
return `events-${timestamp}-${uuid}.json`;
|
|
}
|
|
|
|
public async getUsage(
|
|
orgId: string,
|
|
featureId: FeatureId,
|
|
trx: Transaction | typeof db = db
|
|
): Promise<Usage | null> {
|
|
if (noop()) {
|
|
return null;
|
|
}
|
|
|
|
const usageId = `${orgId}-${featureId}`;
|
|
|
|
try {
|
|
const [result] = await trx
|
|
.select()
|
|
.from(usage)
|
|
.where(eq(usage.usageId, usageId))
|
|
.limit(1);
|
|
|
|
if (!result) {
|
|
// Lets create one if it doesn't exist using upsert to handle race conditions
|
|
logger.info(
|
|
`Creating new usage record for ${orgId}/${featureId}`
|
|
);
|
|
const meterId = getFeatureMeterId(featureId);
|
|
|
|
try {
|
|
const [newUsage] = await trx
|
|
.insert(usage)
|
|
.values({
|
|
usageId,
|
|
featureId,
|
|
orgId,
|
|
meterId,
|
|
latestValue: 0,
|
|
updatedAt: Math.floor(Date.now() / 1000)
|
|
})
|
|
.onConflictDoNothing()
|
|
.returning();
|
|
|
|
if (newUsage) {
|
|
return newUsage;
|
|
} else {
|
|
// Record was created by another process, fetch it
|
|
const [existingUsage] = await trx
|
|
.select()
|
|
.from(usage)
|
|
.where(eq(usage.usageId, usageId))
|
|
.limit(1);
|
|
return existingUsage || null;
|
|
}
|
|
} catch (insertError) {
|
|
// Fallback: try to fetch existing record in case of any insert issues
|
|
logger.warn(
|
|
`Insert failed for ${orgId}/${featureId}, attempting to fetch existing record:`,
|
|
insertError
|
|
);
|
|
const [existingUsage] = await trx
|
|
.select()
|
|
.from(usage)
|
|
.where(eq(usage.usageId, usageId))
|
|
.limit(1);
|
|
return existingUsage || null;
|
|
}
|
|
}
|
|
|
|
return result;
|
|
} catch (error) {
|
|
logger.error(
|
|
`Failed to get usage for ${orgId}/${featureId}:`,
|
|
error
|
|
);
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
public async getUsageDaily(
|
|
orgId: string,
|
|
featureId: FeatureId
|
|
): Promise<Usage | null> {
|
|
if (noop()) {
|
|
return null;
|
|
}
|
|
await this.updateDaily(orgId, featureId); // Ensure daily usage is updated
|
|
return this.getUsage(orgId, featureId);
|
|
}
|
|
|
|
public async forceUpload(): Promise<void> {
|
|
if (this.events.length > 0) {
|
|
// Force upload regardless of time
|
|
this.lastUploadTime = 0; // Reset to force upload
|
|
await this.uploadEventsToS3();
|
|
}
|
|
}
|
|
|
|
public async checkLimitSet(
|
|
orgId: string,
|
|
kickSites = false,
|
|
featureId?: FeatureId,
|
|
usage?: Usage,
|
|
trx: Transaction | typeof db = db
|
|
): Promise<boolean> {
|
|
if (noop()) {
|
|
return false;
|
|
}
|
|
// This method should check the current usage against the limits set for the organization
|
|
// and kick out all of the sites on the org
|
|
let hasExceededLimits = false;
|
|
|
|
try {
|
|
let orgLimits: Limit[] = [];
|
|
if (featureId) {
|
|
// Get all limits set for this organization
|
|
orgLimits = await trx
|
|
.select()
|
|
.from(limits)
|
|
.where(
|
|
and(
|
|
eq(limits.orgId, orgId),
|
|
eq(limits.featureId, featureId)
|
|
)
|
|
);
|
|
} else {
|
|
// Get all limits set for this organization
|
|
orgLimits = await trx
|
|
.select()
|
|
.from(limits)
|
|
.where(eq(limits.orgId, orgId));
|
|
}
|
|
|
|
if (orgLimits.length === 0) {
|
|
logger.debug(`No limits set for org ${orgId}`);
|
|
return false;
|
|
}
|
|
|
|
// Check each limit against current usage
|
|
for (const limit of orgLimits) {
|
|
let currentUsage: Usage | null;
|
|
if (usage) {
|
|
currentUsage = usage;
|
|
} else {
|
|
currentUsage = await this.getUsage(
|
|
orgId,
|
|
limit.featureId as FeatureId,
|
|
trx
|
|
);
|
|
}
|
|
|
|
const usageValue =
|
|
currentUsage?.instantaneousValue ||
|
|
currentUsage?.latestValue ||
|
|
0;
|
|
logger.debug(
|
|
`Current usage for org ${orgId} on feature ${limit.featureId}: ${usageValue}`
|
|
);
|
|
logger.debug(
|
|
`Limit for org ${orgId} on feature ${limit.featureId}: ${limit.value}`
|
|
);
|
|
if (
|
|
currentUsage &&
|
|
limit.value !== null &&
|
|
usageValue > limit.value
|
|
) {
|
|
logger.debug(
|
|
`Org ${orgId} has exceeded limit for ${limit.featureId}: ` +
|
|
`${usageValue} > ${limit.value}`
|
|
);
|
|
hasExceededLimits = true;
|
|
break; // Exit early if any limit is exceeded
|
|
}
|
|
}
|
|
|
|
// If any limits are exceeded, disconnect all sites for this organization
|
|
if (hasExceededLimits && kickSites) {
|
|
logger.warn(
|
|
`Disconnecting all sites for org ${orgId} due to exceeded limits`
|
|
);
|
|
|
|
// Get all sites for this organization
|
|
const orgSites = await trx
|
|
.select()
|
|
.from(sites)
|
|
.where(eq(sites.orgId, orgId));
|
|
|
|
// Mark all sites as offline and send termination messages
|
|
const siteUpdates = orgSites.map((site) => site.siteId);
|
|
|
|
if (siteUpdates.length > 0) {
|
|
// Send termination messages to newt sites
|
|
for (const site of orgSites) {
|
|
if (site.type === "newt") {
|
|
const [newt] = await trx
|
|
.select()
|
|
.from(newts)
|
|
.where(eq(newts.siteId, site.siteId))
|
|
.limit(1);
|
|
|
|
if (newt) {
|
|
const payload = {
|
|
type: `newt/wg/terminate`,
|
|
data: {
|
|
reason: "Usage limits exceeded"
|
|
}
|
|
};
|
|
|
|
// Don't await to prevent blocking
|
|
await sendToClient(newt.newtId, payload).catch(
|
|
(error: any) => {
|
|
logger.error(
|
|
`Failed to send termination message to newt ${newt.newtId}:`,
|
|
error
|
|
);
|
|
}
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
logger.info(
|
|
`Disconnected ${orgSites.length} sites for org ${orgId} due to exceeded limits`
|
|
);
|
|
}
|
|
}
|
|
} catch (error) {
|
|
logger.error(`Error checking limits for org ${orgId}:`, error);
|
|
}
|
|
|
|
return hasExceededLimits;
|
|
}
|
|
}
|
|
|
|
export const usageService = new UsageService();
|