refactor: memory usage improvements

Author: izadoesdev

apps/basket/src/index.ts

@@ -2,17 +2,37 @@ import "./polyfills/compression";
 
 import { Elysia } from "elysia";
 import { logger } from "./lib/logger";
-import { getProducerStats } from "./lib/producer";
+import { disconnectProducer, getProducerStats } from "./lib/producer";
 import {
 	endRequestSpan,
 	initTracing,
+	shutdownTracing,
 	startRequestSpan,
 } from "./lib/tracing";
 import basketRouter from "./routes/basket";
 import emailRouter from "./routes/email";
+import { closeGeoIPReader } from "./utils/ip-geo";
 
 initTracing();
 
+process.on("SIGTERM", async () => {
+	logger.info("SIGTERM received, shutting down gracefully...");
+	await Promise.all([disconnectProducer(), shutdownTracing()]).catch((error) =>
+		logger.error({ error }, "Shutdown error")
+	);
+	closeGeoIPReader();
+	process.exit(0);
+});
+
+process.on("SIGINT", async () => {
+	logger.info("SIGINT received, shutting down gracefully...");
+	await Promise.all([disconnectProducer(), shutdownTracing()]).catch((error) =>
+		logger.error({ error }, "Shutdown error")
+	);
+	closeGeoIPReader();
+	process.exit(0);
+});
+
 function getKafkaHealth() {
 	const stats = getProducerStats();
 
@@ -99,7 +119,7 @@ const app = new Elysia()
 
 const port = process.env.PORT || 4000;
 
-console.log(`Starting basket service on port ${port}`);
+logger.info(`Starting basket service on port ${port}`);
 
 export default {
 	fetch: app.fetch,

apps/basket/src/lib/blocked-traffic.ts

@@ -24,11 +24,11 @@ async function _logBlockedTrafficAsync(
 			) || "";
 
 
-		const [geo, ua, now] = await Promise.all([
+		const [geo, ua] = await Promise.all([
 			getGeo(ip),
 			parseUserAgent(userAgent),
-			Date.now(),
 		]);
+		const now = Date.now();
 		const { anonymizedIP, country, region, city } = geo;
 		const { browserName, browserVersion, osName, osVersion, deviceType } = ua;
 

apps/basket/src/lib/event-service.ts

@@ -59,7 +59,7 @@ export function insertError(
 
 		const { anonymizedIP, country, region } = geoData;
 		const { browserName, browserVersion, osName, osVersion, deviceType } =
-			parseUserAgent(userAgent);
+			await parseUserAgent(userAgent);
 
 		const errorEvent: ErrorEvent = {
 			id: randomUUID(),
@@ -141,7 +141,7 @@ export async function insertWebVitals(
 
 	const { country, region } = await getGeo(ip);
 	const { browserName, browserVersion, osName, osVersion, deviceType } =
-		parseUserAgent(userAgent);
+		await parseUserAgent(userAgent);
 
 	const webVitalsEvent: WebVitalsEvent = {
 		id: randomUUID(),
@@ -331,7 +331,7 @@ export function insertTrackEvent(
 			deviceType,
 			deviceBrand,
 			deviceModel,
-		} = parseUserAgent(userAgent);
+		} = await parseUserAgent(userAgent);
 		const now = Date.now();
 
 		const trackEvent: AnalyticsEvent = {

apps/basket/src/lib/producer.ts

@@ -193,7 +193,8 @@ export class EventProducer {
 		}
 
 		this.flushing = true;
-		const items = this.buffer.splice(0);
+		const batchSize = Math.min(this.buffer.length, this.config.bufferMax);
+		const items = this.buffer.splice(0, batchSize);
 
 		try {
 			const grouped = items.reduce(
@@ -238,7 +239,7 @@ export class EventProducer {
 
 						for (const { event, retries, timestamp } of items) {
 							const age = Date.now() - timestamp;
-							if (retries < this.config.maxRetries && age < 300_000) {
+							if (retries < this.config.maxRetries && age < 300_000 && this.buffer.length < this.config.bufferHardMax) {
 								this.buffer.push({
 									table,
 									event,
@@ -249,7 +250,7 @@ export class EventProducer {
 								this.stats.dropped += 1;
 								logger.error(
 									{ error },
-									`Dropped event (retries: ${retries}, age: ${age}ms)`,
+									`Dropped event (retries: ${retries}, age: ${age}ms, buffer: ${this.buffer.length})`,
 									{ table, eventId: (event as { event_id?: string }).event_id }
 								);
 							}
@@ -611,17 +612,3 @@ export const disconnectProducer = async (): Promise<void> => {
 };
 
 export const getProducerStats = () => getDefaultProducer().getStats();
-
-process.on("SIGTERM", async () => {
-	await disconnectProducer().catch((error) =>
-		logger.error({ error }, "SIGTERM error")
-	);
-	process.exit(0);
-});
-
-process.on("SIGINT", async () => {
-	await disconnectProducer().catch((error) =>
-		logger.error({ error }, "SIGINT error")
-	);
-	process.exit(0);
-});

apps/basket/src/lib/tracing.ts

@@ -19,37 +19,37 @@ export function initTracing(): void {
         return;
     }
 
+    const exporter = new OTLPTraceExporter({
+        url: "https://api.axiom.co/v1/traces",
+        headers: {
+            Authorization: `Bearer ${process.env.AXIOM_TOKEN}`,
+            "X-Axiom-Dataset": process.env.AXIOM_DATASET ?? "basket",
+        },
+    });
+
     sdk = new NodeSDK({
         resource: resourceFromAttributes({
             [ATTR_SERVICE_NAME]: "basket",
             [ATTR_SERVICE_VERSION]: pkg.version,
         }),
-        traceExporter: new OTLPTraceExporter({
-            url: "https://api.axiom.co/v1/traces",
-            headers: {
-                Authorization: `Bearer ${process.env.AXIOM_TOKEN}`,
-                "X-Axiom-Dataset": process.env.AXIOM_DATASET ?? "basket",
-            },
+        spanProcessor: new BatchSpanProcessor(exporter, {
+            scheduledDelayMillis: 1000,
+            exportTimeoutMillis: 30_000,
+            maxExportBatchSize: 512,
+            maxQueueSize: 2048,
         }),
-        spanProcessor: new BatchSpanProcessor(
-            new OTLPTraceExporter({
-                url: "https://api.axiom.co/v1/traces",
-                headers: {
-                    Authorization: `Bearer ${process.env.AXIOM_TOKEN}`,
-                    "X-Axiom-Dataset": process.env.AXIOM_DATASET ?? "basket",
-                },
-            }),
-            {
-                scheduledDelayMillis: 1000,
-                exportTimeoutMillis: 30_000,
-                maxExportBatchSize: 512,
-            }
-        ),
     });
 
     sdk.start();
 }
 
+export async function shutdownTracing(): Promise<void> {
+    if (sdk) {
+        await sdk.shutdown();
+        sdk = null;
+    }
+}
+
 /**
  * Get tracer
  */

apps/basket/src/utils/ip-geo.ts

@@ -19,6 +19,7 @@ let reader: GeoIPReader | null = null;
 let isLoading = false;
 let loadPromise: Promise<void> | null = null;
 let loadError: Error | null = null;
+let dbBuffer: Buffer | null = null;
 
 async function loadDatabaseFromCdn(): Promise<Buffer> {
 	try {
@@ -61,12 +62,13 @@ function loadDatabase() {
 	isLoading = true;
 	loadPromise = (async () => {
 		try {
-			const dbBuffer = await loadDatabaseFromCdn();
+			dbBuffer = await loadDatabaseFromCdn();
 			reader = Reader.openBuffer(dbBuffer) as GeoIPReader;
 		} catch (error) {
 			logger.error({ error }, "Failed to load GeoIP database");
 			loadError = error as Error;
 			reader = null;
+			dbBuffer = null;
 		} finally {
 			isLoading = false;
 		}
@@ -206,3 +208,15 @@ export function extractIpFromRequest(request: Request): string {
 
 	return "";
 }
+
+export function closeGeoIPReader(): void {
+	if (reader) {
+		reader = null;
+	}
+	if (dbBuffer) {
+		dbBuffer = null;
+	}
+	loadPromise = null;
+	loadError = null;
+	isLoading = false;
+}

apps/basket/src/utils/user-agent.ts

@@ -5,6 +5,7 @@
  * and platform identification.
  */
 
+import { cacheable } from "@databuddy/redis";
 import { bots } from "@databuddy/shared/lists/bots";
 import { UAParser } from "ua-parser-js";
 import { logger } from "../lib/logger";
@@ -21,9 +22,9 @@ export type UserAgentInfo = {
 };
 
 /**
- * Parse user agent to extract useful information
+ * Parse user agent to extract useful information (uncached)
  */
-export function parseUserAgent(userAgent: string): {
+function _parseUserAgent(userAgent: string): {
 	browserName?: string;
 	browserVersion?: string;
 	osName?: string;
@@ -58,10 +59,7 @@ export function parseUserAgent(userAgent: string): {
 			deviceModel: result.device.model || undefined,
 		};
 	} catch (error) {
-		logger.error(
-			{ error, userAgent },
-			"Failed to parse user agent"
-		);
+		logger.error({ error, userAgent }, "Failed to parse user agent");
 		return {
 			browserName: undefined,
 			browserVersion: undefined,
@@ -74,6 +72,16 @@ export function parseUserAgent(userAgent: string): {
 	}
 }
 
+/**
+ * Parse user agent to extract useful information (cached)
+ */
+export const parseUserAgent = cacheable(_parseUserAgent, {
+	expireInSec: 3600,
+	prefix: "ua_parse",
+	staleWhileRevalidate: true,
+	staleTime: 1800,
+});
+
 export function detectBot(
 	userAgent: string,
 	request: Request

packages/redis/cacheable.ts

@@ -2,9 +2,11 @@ import { getRedisCache } from "./redis";
 
 const logger = console;
 
+const activeRevalidations = new Map<string, Promise<void>>();
+
 const stringifyRegex = /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}.*Z$/;
 
-interface CacheOptions {
+type CacheOptions = {
 	expireInSec: number;
 	prefix?: string;
 	serialize?: (data: unknown) => string;
@@ -47,9 +49,9 @@ export async function getCache<T>(
 
 				if (staleWhileRevalidate) {
 					const ttl = await redis.ttl(key);
-					if (ttl < staleTime) {
+					if (ttl < staleTime && !activeRevalidations.has(key)) {
 						// Return stale data and revalidate in background
-						fn()
+						const revalidationPromise = fn()
 							.then(async (freshData: T) => {
 								if (freshData !== undefined && freshData !== null) {
 									const redis = getRedisCache();
@@ -61,7 +63,11 @@ export async function getCache<T>(
 									`Background revalidation failed for key ${key}:`,
 									error
 								);
+							})
+							.finally(() => {
+								activeRevalidations.delete(key);
 							});
+						activeRevalidations.set(key, revalidationPromise);
 					}
 				}
 
@@ -74,7 +80,7 @@ export async function getCache<T>(
 			}
 			return data;
 		} catch (error: unknown) {
-			retries++;
+			retries += 1;
 			if (retries === maxRetries) {
 				logger.error(
 					`Cache error for key ${key} after ${maxRetries} retries:`,
@@ -157,9 +163,9 @@ export function cacheable<T extends (...args: any) => any>(
 
 					if (staleWhileRevalidate) {
 						const ttl = await redis.ttl(key);
-						if (ttl < staleTime) {
+						if (ttl < staleTime && !activeRevalidations.has(key)) {
 							// Return stale data and revalidate in background
-							fn(...args)
+							const revalidationPromise = fn(...args)
 								.then(async (freshData: Awaited<ReturnType<T>>) => {
 									if (freshData !== undefined && freshData !== null) {
 										const redis = getRedisCache();
@@ -171,7 +177,11 @@ export function cacheable<T extends (...args: any) => any>(
 										`Background revalidation failed for function ${fn.name}:`,
 										error
 									);
+								})
+								.finally(() => {
+									activeRevalidations.delete(key);
 								});
+							activeRevalidations.set(key, revalidationPromise);
 						}
 					}
 
@@ -184,7 +194,7 @@ export function cacheable<T extends (...args: any) => any>(
 				}
 				return result;
 			} catch (error: unknown) {
-				retries++;
+				retries += 1;
 				if (retries === maxRetries) {
 					logger.error(
 						`Cache error for function ${fn.name} after ${maxRetries} retries:`,