|
| 1 | +import type { MetricsV1Input } from "@internal/clickhouse"; |
| 2 | +import { env } from "~/env.server"; |
| 3 | +import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server"; |
| 4 | +import { singleton } from "~/utils/singleton"; |
| 5 | +import { meter } from "~/v3/tracer.server"; |
| 6 | +import { ApiRateLimitMetricsAggregator } from "./apiRateLimitMetricsAggregator.server"; |
| 7 | +import { |
| 8 | + apiRateLimitMetricsInsertSettings, |
| 9 | + exportApiRateLimitMetricRows, |
| 10 | +} from "./apiRateLimitMetricsExporter.server"; |
| 11 | +import type { |
| 12 | + RateLimitObservation, |
| 13 | + RateLimitTenant, |
| 14 | +} from "./authorizationRateLimitMiddleware.server"; |
| 15 | +import { logger } from "./logger.server"; |
| 16 | +import { signalsEmitter } from "./signals.server"; |
| 17 | + |
| 18 | +function enabledByEnv(): boolean { |
| 19 | + return env.API_RATE_LIMIT_METRICS_ENABLED !== "0"; |
| 20 | +} |
| 21 | + |
| 22 | +export function recordApiRateLimitObservation(observation: RateLimitObservation): void { |
| 23 | + if (!enabledByEnv()) { |
| 24 | + return; |
| 25 | + } |
| 26 | + getAggregator().record(observation); |
| 27 | +} |
| 28 | + |
| 29 | +/** |
| 30 | + * Builds the aggregator and starts its flush timer ahead of the first request. |
| 31 | + */ |
| 32 | +export function initApiRateLimitMetrics(): void { |
| 33 | + if (!enabledByEnv()) { |
| 34 | + return; |
| 35 | + } |
| 36 | + getAggregator(); |
| 37 | +} |
| 38 | + |
| 39 | +function getAggregator(): ApiRateLimitMetricsAggregator { |
| 40 | + return singleton("apiRateLimitMetricsAggregator", createAggregator); |
| 41 | +} |
| 42 | + |
| 43 | +/** |
| 44 | + * With the env value "allowlist" only tenants whose organization carries the |
| 45 | + * apiRateLimitMetricsEnabled feature flag are recorded. The flag travels with the cached |
| 46 | + * rate-limit resolution, so this costs nothing per request. Turning recording off is a redeploy. |
| 47 | + */ |
| 48 | +function createRuntimeGate(): (tenant: RateLimitTenant) => boolean { |
| 49 | + if (env.API_RATE_LIMIT_METRICS_ENABLED === "1") { |
| 50 | + return () => true; |
| 51 | + } |
| 52 | + return (tenant) => tenant.metricsEnabled; |
| 53 | +} |
| 54 | + |
| 55 | +function exportRows(rows: MetricsV1Input[], onInsertError: (rows: number) => void): Promise<void> { |
| 56 | + return exportApiRateLimitMetricRows(rows, { |
| 57 | + resolveClient: (organizationId) => |
| 58 | + clickhouseFactory.getClickhouseForOrganizationSync(organizationId, "events"), |
| 59 | + settings: apiRateLimitMetricsInsertSettings({ |
| 60 | + waitForAsyncInsert: env.API_RATE_LIMIT_METRICS_WAIT_FOR_ASYNC_INSERT === "1", |
| 61 | + busyTimeoutMs: env.API_RATE_LIMIT_METRICS_INSERT_BUSY_TIMEOUT_MS, |
| 62 | + }), |
| 63 | + onInsertError: (failedRows, error) => { |
| 64 | + onInsertError(failedRows); |
| 65 | + logger.error( |
| 66 | + "api rate limit metrics: clickhouse rejected the insert request, dropping rows", |
| 67 | + { |
| 68 | + rows: failedRows, |
| 69 | + error: error instanceof Error ? error.message : String(error), |
| 70 | + } |
| 71 | + ); |
| 72 | + }, |
| 73 | + }); |
| 74 | +} |
| 75 | + |
| 76 | +function createAggregator(): ApiRateLimitMetricsAggregator { |
| 77 | + const droppedCounter = meter.createCounter("api_rate_limit_metrics.dropped", { |
| 78 | + description: |
| 79 | + "API rate limit observations dropped, by reason: the in-process cap was reached, or shutdown came before data store routing was ready", |
| 80 | + }); |
| 81 | + const flushedCounter = meter.createCounter("api_rate_limit_metrics.rows_flushed", { |
| 82 | + description: "API rate limit metric rows sent to ClickHouse as async inserts", |
| 83 | + }); |
| 84 | + const insertFailedCounter = meter.createCounter("api_rate_limit_metrics.rows_insert_failed", { |
| 85 | + description: |
| 86 | + "API rate limit metric rows lost because ClickHouse rejected the insert request; failures while writing a queued async insert are not visible here", |
| 87 | + }); |
| 88 | + |
| 89 | + const aggregator = new ApiRateLimitMetricsAggregator({ |
| 90 | + bucketSeconds: env.API_RATE_LIMIT_METRICS_BUCKET_SECONDS, |
| 91 | + maxEntries: env.API_RATE_LIMIT_METRICS_MAX_ENTRIES, |
| 92 | + isEnabled: createRuntimeGate(), |
| 93 | + sink: (rows) => exportRows(rows, (failed) => insertFailedCounter.add(failed)), |
| 94 | + onDropped: (count) => droppedCounter.add(count, { reason: "cap" }), |
| 95 | + }); |
| 96 | + |
| 97 | + let routingReady = false; |
| 98 | + clickhouseFactory |
| 99 | + .isReady() |
| 100 | + .then(() => { |
| 101 | + routingReady = true; |
| 102 | + }) |
| 103 | + .catch((error) => { |
| 104 | + logger.error("api rate limit metrics: data store registry never became ready", { |
| 105 | + error: error instanceof Error ? error.message : String(error), |
| 106 | + }); |
| 107 | + }); |
| 108 | + |
| 109 | + let deferredWarned = false; |
| 110 | + const flush = () => { |
| 111 | + if (!routingReady) { |
| 112 | + if (aggregator.size > 0 && !deferredWarned) { |
| 113 | + deferredWarned = true; |
| 114 | + logger.warn("api rate limit metrics: flush deferred, data store routing not ready", { |
| 115 | + pendingEntries: aggregator.size, |
| 116 | + }); |
| 117 | + } |
| 118 | + return; |
| 119 | + } |
| 120 | + try { |
| 121 | + const flushed = aggregator.flush(); |
| 122 | + if (flushed > 0) { |
| 123 | + flushedCounter.add(flushed); |
| 124 | + } |
| 125 | + } catch (error) { |
| 126 | + logger.error("api rate limit metrics: flush failed", { |
| 127 | + error: error instanceof Error ? error.message : String(error), |
| 128 | + }); |
| 129 | + } |
| 130 | + }; |
| 131 | + |
| 132 | + const interval = setInterval(flush, env.API_RATE_LIMIT_METRICS_FLUSH_INTERVAL_MS); |
| 133 | + interval.unref(); |
| 134 | + |
| 135 | + const shutdown = () => { |
| 136 | + clearInterval(interval); |
| 137 | + if (!routingReady) { |
| 138 | + const dropped = aggregator.discard(); |
| 139 | + if (dropped > 0) { |
| 140 | + droppedCounter.add(dropped, { reason: "shutdown_before_ready" }); |
| 141 | + logger.warn("api rate limit metrics: shutting down before routing was ready, dropping", { |
| 142 | + droppedObservations: dropped, |
| 143 | + }); |
| 144 | + } |
| 145 | + return; |
| 146 | + } |
| 147 | + flush(); |
| 148 | + }; |
| 149 | + signalsEmitter.on("SIGTERM", shutdown); |
| 150 | + signalsEmitter.on("SIGINT", shutdown); |
| 151 | + |
| 152 | + return aggregator; |
| 153 | +} |
0 commit comments