Skip to content

Commit 262b5b8

Browse files
Merge pull request #937 from Killerjunior/fix-815-container-cpu-metrics
Use container-aware CPU metrics for quota limits
2 parents 474a9ee + 33a4039 commit 262b5b8

6 files changed

Lines changed: 286 additions & 17 deletions

src/rate-limiting/rate-limiting.module.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { QuotaDefinitionService } from './services/quota-definition.service';
1111
import { QuotaTrackingService } from './services/quota-tracking.service';
1212
import { QuotaResetScheduler } from './services/quota-reset.scheduler';
1313
import { AdaptiveRateLimitingService } from './services/adaptive-rate-limiting.service';
14+
import { ContainerCpuMetricsService } from './services/container-cpu-metrics.service';
1415

1516
// Guard & Decorator
1617
import { QuotaGuard } from './guards/quota.guard';
@@ -28,6 +29,7 @@ import { UserQuotaController } from './controllers/user-quota.controller';
2829
QuotaManagementService,
2930
QuotaResetScheduler,
3031
AdaptiveRateLimitingService,
32+
ContainerCpuMetricsService,
3133
QuotaGuard,
3234
],
3335
exports: [QuotaManagementService, QuotaDefinitionService, QuotaTrackingService, QuotaGuard],
Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,18 @@
11
import { Injectable } from '@nestjs/common';
2-
import * as os from 'os';
2+
import { ContainerCpuMetricsService } from './container-cpu-metrics.service';
33

44
/**
55
* Provides adaptive Rate Limiting operations.
66
*/
77
@Injectable()
88
export class AdaptiveRateLimitingService {
9+
constructor(private readonly cpuMetrics: ContainerCpuMetricsService) {}
910
/**
1011
* Retrieves system Load Factor.
1112
* @returns The calculated numeric value.
1213
*/
13-
getSystemLoadFactor(): number {
14-
const load = os.loadavg()[0]; // 1-minute average
15-
const cpuCount = os.cpus().length;
16-
17-
const loadPercentage = load / cpuCount;
14+
async getSystemLoadFactor(): Promise<number> {
15+
const loadPercentage = await this.cpuMetrics.getCpuLoadRatio();
1816

1917
if (loadPercentage > 0.9) return 0.5; // reduce limits by 50%
2018
if (loadPercentage > 0.7) return 0.7;
@@ -26,8 +24,8 @@ export class AdaptiveRateLimitingService {
2624
* @param baseLimit The maximum number of results.
2725
* @returns The calculated numeric value.
2826
*/
29-
adjustLimit(baseLimit: number): number {
30-
const factor = this.getSystemLoadFactor();
27+
async adjustLimit(baseLimit: number): Promise<number> {
28+
const factor = await this.getSystemLoadFactor();
3129
return Math.floor(baseLimit * factor);
3230
}
3331
}
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
jest.mock('fs', () => ({
2+
readFileSync: jest.fn(),
3+
}));
4+
5+
jest.mock('os', () => ({
6+
cpus: jest.fn(),
7+
loadavg: jest.fn(),
8+
}));
9+
10+
import { readFileSync } from 'fs';
11+
import * as os from 'os';
12+
import { ContainerCpuMetricsService } from './container-cpu-metrics.service';
13+
14+
describe('ContainerCpuMetricsService', () => {
15+
const mockedReadFileSync = readFileSync as jest.MockedFunction<typeof readFileSync>;
16+
const mockedCpus = os.cpus as jest.MockedFunction<typeof os.cpus>;
17+
const mockedLoadavg = os.loadavg as jest.MockedFunction<typeof os.loadavg>;
18+
19+
const originalFetch = global.fetch;
20+
const originalPrometheusUrl = process.env.PROMETHEUS_METRICS_URL;
21+
22+
beforeEach(() => {
23+
jest.clearAllMocks();
24+
process.env.PROMETHEUS_METRICS_URL = undefined;
25+
});
26+
27+
afterEach(() => {
28+
global.fetch = originalFetch;
29+
process.env.PROMETHEUS_METRICS_URL = originalPrometheusUrl;
30+
});
31+
32+
it('reads cgroup v2 cpu.stat and returns throttling ratio', async () => {
33+
mockedCpus.mockReturnValue([{ model: 'cpu', speed: 1000, times: {} as never }] as never);
34+
mockedReadFileSync.mockReturnValue(
35+
'usage_usec 100000\nnr_periods 100\nnr_throttled 80\n' as never,
36+
);
37+
38+
const service = new ContainerCpuMetricsService();
39+
const ratio = await service.getCpuLoadRatio();
40+
41+
expect(ratio).toBeCloseTo(0.8, 4);
42+
});
43+
44+
it('falls back to os.loadavg when cgroup is unavailable', async () => {
45+
mockedReadFileSync.mockImplementation(() => {
46+
throw new Error('missing');
47+
});
48+
mockedLoadavg.mockReturnValue([2, 1, 1]);
49+
mockedCpus.mockReturnValue([
50+
{ model: 'cpu-1', speed: 1000, times: {} as never },
51+
{ model: 'cpu-2', speed: 1000, times: {} as never },
52+
{ model: 'cpu-3', speed: 1000, times: {} as never },
53+
{ model: 'cpu-4', speed: 1000, times: {} as never },
54+
] as never);
55+
56+
const service = new ContainerCpuMetricsService();
57+
const ratio = await service.getCpuLoadRatio();
58+
59+
expect(ratio).toBeCloseTo(0.5, 4);
60+
});
61+
62+
it('falls back to Prometheus process_cpu_seconds_total when cgroup is unavailable', async () => {
63+
process.env.PROMETHEUS_METRICS_URL = 'http://127.0.0.1:3000/metrics';
64+
65+
mockedReadFileSync.mockImplementation(() => {
66+
throw new Error('missing');
67+
});
68+
mockedCpus.mockReturnValue([
69+
{ model: 'cpu-1', speed: 1000, times: {} as never },
70+
{ model: 'cpu-2', speed: 1000, times: {} as never },
71+
] as never);
72+
mockedLoadavg.mockReturnValue([0.4, 0, 0]);
73+
74+
jest.spyOn(Date, 'now').mockReturnValueOnce(1000).mockReturnValueOnce(2000);
75+
76+
global.fetch = jest
77+
.fn()
78+
.mockResolvedValueOnce({
79+
ok: true,
80+
text: async () => '# HELP process_cpu_seconds_total\nprocess_cpu_seconds_total 10\n',
81+
} as Response)
82+
.mockResolvedValueOnce({
83+
ok: true,
84+
text: async () => '# HELP process_cpu_seconds_total\nprocess_cpu_seconds_total 11\n',
85+
} as Response);
86+
87+
const service = new ContainerCpuMetricsService();
88+
89+
const firstRatio = await service.getCpuLoadRatio();
90+
const secondRatio = await service.getCpuLoadRatio();
91+
92+
expect(firstRatio).toBeCloseTo(0.2, 4);
93+
expect(secondRatio).toBeCloseTo(0.5, 4);
94+
});
95+
});
Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
import { Injectable, Logger } from '@nestjs/common';
2+
import { readFileSync } from 'fs';
3+
import * as os from 'os';
4+
5+
interface CpuStatSnapshot {
6+
usageUsec: number;
7+
nrPeriods: number;
8+
nrThrottled: number;
9+
timestampMs: number;
10+
}
11+
12+
interface ProcessCpuSnapshot {
13+
cpuSecondsTotal: number;
14+
timestampMs: number;
15+
}
16+
17+
/**
18+
* Reads container-aware CPU metrics with ordered fallbacks:
19+
* 1) cgroups v2 cpu.stat (container-level)
20+
* 2) Prometheus process_cpu_seconds_total scrape (optional)
21+
* 3) host loadavg (last resort)
22+
*/
23+
@Injectable()
24+
export class ContainerCpuMetricsService {
25+
private readonly logger = new Logger(ContainerCpuMetricsService.name);
26+
private readonly cpuStatPath = '/sys/fs/cgroup/cpu.stat';
27+
private previousCpuStatSnapshot?: CpuStatSnapshot;
28+
private previousProcessCpuSnapshot?: ProcessCpuSnapshot;
29+
30+
async getCpuLoadRatio(): Promise<number> {
31+
const cgroupRatio = this.getCgroupCpuLoadRatio();
32+
if (cgroupRatio !== null) {
33+
return cgroupRatio;
34+
}
35+
36+
const prometheusRatio = await this.getPrometheusCpuLoadRatio();
37+
if (prometheusRatio !== null) {
38+
return prometheusRatio;
39+
}
40+
41+
return this.getLoadAvgRatio();
42+
}
43+
44+
private getCgroupCpuLoadRatio(): number | null {
45+
try {
46+
const cpuStat = readFileSync(this.cpuStatPath, 'utf8');
47+
const usageUsec = this.readMetric(cpuStat, 'usage_usec');
48+
const nrPeriods = this.readMetric(cpuStat, 'nr_periods');
49+
const nrThrottled = this.readMetric(cpuStat, 'nr_throttled');
50+
51+
if (usageUsec === null || nrPeriods === null || nrThrottled === null) {
52+
return null;
53+
}
54+
55+
const current: CpuStatSnapshot = {
56+
usageUsec,
57+
nrPeriods,
58+
nrThrottled,
59+
timestampMs: Date.now(),
60+
};
61+
62+
const throttleRatio = nrPeriods > 0 ? this.clamp(nrThrottled / nrPeriods) : 0;
63+
const cpuCount = Math.max(1, os.cpus().length);
64+
65+
if (!this.previousCpuStatSnapshot) {
66+
this.previousCpuStatSnapshot = current;
67+
return throttleRatio;
68+
}
69+
70+
const elapsedUsec = (current.timestampMs - this.previousCpuStatSnapshot.timestampMs) * 1000;
71+
const deltaUsageUsec = current.usageUsec - this.previousCpuStatSnapshot.usageUsec;
72+
this.previousCpuStatSnapshot = current;
73+
74+
if (elapsedUsec <= 0 || deltaUsageUsec < 0) {
75+
return throttleRatio;
76+
}
77+
78+
const usageRatio = this.clamp(deltaUsageUsec / elapsedUsec / cpuCount);
79+
return Math.max(usageRatio, throttleRatio);
80+
} catch {
81+
return null;
82+
}
83+
}
84+
85+
private async getPrometheusCpuLoadRatio(): Promise<number | null> {
86+
const url = process.env.PROMETHEUS_METRICS_URL;
87+
if (!url) {
88+
return null;
89+
}
90+
91+
try {
92+
const response = await fetch(url);
93+
if (!response.ok) {
94+
this.logger.warn(`Prometheus metrics scrape failed: status=${response.status}`);
95+
return null;
96+
}
97+
98+
const metricsText = await response.text();
99+
const cpuSecondsTotal = this.readPrometheusProcessCpuSeconds(metricsText);
100+
if (cpuSecondsTotal === null) {
101+
return null;
102+
}
103+
104+
const current: ProcessCpuSnapshot = { cpuSecondsTotal, timestampMs: Date.now() };
105+
const cpuCount = Math.max(1, os.cpus().length);
106+
107+
if (!this.previousProcessCpuSnapshot) {
108+
this.previousProcessCpuSnapshot = current;
109+
return null;
110+
}
111+
112+
const elapsedSeconds =
113+
(current.timestampMs - this.previousProcessCpuSnapshot.timestampMs) / 1000;
114+
const deltaCpuSeconds =
115+
current.cpuSecondsTotal - this.previousProcessCpuSnapshot.cpuSecondsTotal;
116+
this.previousProcessCpuSnapshot = current;
117+
118+
if (elapsedSeconds <= 0 || deltaCpuSeconds < 0) {
119+
return null;
120+
}
121+
122+
return this.clamp(deltaCpuSeconds / elapsedSeconds / cpuCount);
123+
} catch {
124+
return null;
125+
}
126+
}
127+
128+
private getLoadAvgRatio(): number {
129+
const load = os.loadavg()[0];
130+
const cpuCount = Math.max(1, os.cpus().length);
131+
return this.clamp(load / cpuCount);
132+
}
133+
134+
private readMetric(content: string, key: string): number | null {
135+
const match = content.match(new RegExp(`^${key}\\s+(\\d+)`, 'm'));
136+
if (!match) return null;
137+
138+
const parsed = Number(match[1]);
139+
return Number.isFinite(parsed) ? parsed : null;
140+
}
141+
142+
private readPrometheusProcessCpuSeconds(metricsText: string): number | null {
143+
const match = metricsText.match(
144+
/^process_cpu_seconds_total(?:\{[^}]*\})?\s+([0-9]+(?:\.[0-9]+)?)$/m,
145+
);
146+
if (!match) return null;
147+
148+
const parsed = Number(match[1]);
149+
return Number.isFinite(parsed) ? parsed : null;
150+
}
151+
152+
private clamp(value: number): number {
153+
if (!Number.isFinite(value)) return 0;
154+
if (value < 0) return 0;
155+
if (value > 1) return 1;
156+
return value;
157+
}
158+
}

src/rate-limiting/services/quota-tracking.service.ts

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,15 @@ export class QuotaTrackingService {
3131
*/
3232
async checkAndIncrement(userId: string, tier: UserTier): Promise<QuotaCheckResult> {
3333
const baseLimits = await this.definitionService.resolveForUser(userId, tier);
34+
const [requestsPerMinute, requestsPerHour, requestsPerDay] = await Promise.all([
35+
this.adaptive.adjustLimit(baseLimits.requestsPerMinute),
36+
this.adaptive.adjustLimit(baseLimits.requestsPerHour),
37+
this.adaptive.adjustLimit(baseLimits.requestsPerDay),
38+
]);
3439
const limits = {
35-
requestsPerMinute: this.adaptive.adjustLimit(baseLimits.requestsPerMinute),
36-
requestsPerHour: this.adaptive.adjustLimit(baseLimits.requestsPerHour),
37-
requestsPerDay: this.adaptive.adjustLimit(baseLimits.requestsPerDay),
40+
requestsPerMinute,
41+
requestsPerHour,
42+
requestsPerDay,
3843
};
3944
const now = new Date();
4045

@@ -86,10 +91,15 @@ export class QuotaTrackingService {
8691
/** Get quota status without incrementing (for status endpoint). */
8792
async getStatus(userId: string, tier: UserTier): Promise<QuotaStatusDto> {
8893
const baseLimits = await this.definitionService.resolveForUser(userId, tier);
94+
const [requestsPerMinute, requestsPerHour, requestsPerDay] = await Promise.all([
95+
this.adaptive.adjustLimit(baseLimits.requestsPerMinute),
96+
this.adaptive.adjustLimit(baseLimits.requestsPerHour),
97+
this.adaptive.adjustLimit(baseLimits.requestsPerDay),
98+
]);
8999
const limits = {
90-
requestsPerMinute: this.adaptive.adjustLimit(baseLimits.requestsPerMinute),
91-
requestsPerHour: this.adaptive.adjustLimit(baseLimits.requestsPerHour),
92-
requestsPerDay: this.adaptive.adjustLimit(baseLimits.requestsPerDay),
100+
requestsPerMinute,
101+
requestsPerHour,
102+
requestsPerDay,
93103
};
94104
const now = new Date();
95105

src/rate-limiting/services/quota.service.ts

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,16 @@ export class QuotaManagementService {
2121
/** Resolve and return the effective quota limits for a user. */
2222
async getQuotaForUser(userId: string, tier: UserTier) {
2323
const base = await this.definitions.resolveForUser(userId, tier);
24+
const [requestsPerMinute, requestsPerHour, requestsPerDay] = await Promise.all([
25+
this.adaptive.adjustLimit(base.requestsPerMinute),
26+
this.adaptive.adjustLimit(base.requestsPerHour),
27+
this.adaptive.adjustLimit(base.requestsPerDay),
28+
]);
29+
2430
return {
25-
requestsPerMinute: this.adaptive.adjustLimit(base.requestsPerMinute),
26-
requestsPerHour: this.adaptive.adjustLimit(base.requestsPerHour),
27-
requestsPerDay: this.adaptive.adjustLimit(base.requestsPerDay),
31+
requestsPerMinute,
32+
requestsPerHour,
33+
requestsPerDay,
2834
};
2935
}
3036

0 commit comments

Comments
 (0)