From fb1e912072f6922a8912e09f45dd55c09b0d16da Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Thu, 10 Sep 2026 14:24:15 +0300 Subject: [PATCH] feat: fixed calendar windows for day/week/month rate limits #1951 Converts day/week/month token/cost/request-count rate limit windows from a floating (rolling) model to fixed calendar windows anchored to a new deployment-wide rateLimitSchedule setting (timezone, week start day, daily reset time). Minute/hour limits keep their existing floating behavior. The bulk limits/usage endpoints now report a resetsAt instant for each fixed window. Pre-rollout floating-window records are implicitly treated as zero usage on first read, with no explicit migration step. --- .../com/epam/aidial/core/config/Config.java | 3 + .../aidial/core/config/RateLimitSchedule.java | 22 +++ .../com/epam/aidial/core/config/WeekDay.java | 23 +++ .../core/config/validation/ValidTimezone.java | 22 +++ .../validation/ValidTimezoneValidator.java | 20 +++ docs/dynamic-settings/roles.md | 37 ++++- docs/open_api_core.yaml | 60 ++++++- sample/aidial.config.json | 5 + .../com/epam/aidial/core/server/AiDial.java | 2 +- .../core/server/data/CostItemLimitStats.java | 10 +- .../core/server/data/ItemLimitStats.java | 8 + .../core/server/limiter/CalendarPeriod.java | 12 ++ .../limiter/CalendarWindowCalculator.java | 83 ++++++++++ .../server/limiter/CostFixedRateBucket.java | 46 ++++++ .../core/server/limiter/CostRateLimit.java | 78 +++++---- .../core/server/limiter/FixedRateBucket.java | 63 ++++++++ .../core/server/limiter/RateLimiter.java | 153 +++++++++++------- .../core/server/limiter/RateWindow.java | 5 +- .../core/server/limiter/RequestRateLimit.java | 23 +-- .../core/server/limiter/TokenRateLimit.java | 50 +++--- .../epam/aidial/core/server/LimitApiTest.java | 55 +++++-- .../limiter/CalendarWindowCalculatorTest.java | 148 +++++++++++++++++ .../limiter/CostFixedRateBucketTest.java | 62 +++++++ .../server/limiter/CostRateBucketTest.java | 120 +------------- .../server/limiter/CostRateLimitTest.java | 22 ++- .../server/limiter/FixedRateBucketTest.java | 88 ++++++++++ .../core/server/limiter/RateBucketTest.java | 91 ----------- .../core/server/limiter/RateLimiterTest.java | 22 ++- 28 files changed, 967 insertions(+), 366 deletions(-) create mode 100644 config/src/main/java/com/epam/aidial/core/config/RateLimitSchedule.java create mode 100644 config/src/main/java/com/epam/aidial/core/config/WeekDay.java create mode 100644 config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezone.java create mode 100644 config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezoneValidator.java create mode 100644 server/src/main/java/com/epam/aidial/core/server/limiter/CalendarPeriod.java create mode 100644 server/src/main/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculator.java create mode 100644 server/src/main/java/com/epam/aidial/core/server/limiter/CostFixedRateBucket.java create mode 100644 server/src/main/java/com/epam/aidial/core/server/limiter/FixedRateBucket.java create mode 100644 server/src/test/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculatorTest.java create mode 100644 server/src/test/java/com/epam/aidial/core/server/limiter/CostFixedRateBucketTest.java create mode 100644 server/src/test/java/com/epam/aidial/core/server/limiter/FixedRateBucketTest.java diff --git a/config/src/main/java/com/epam/aidial/core/config/Config.java b/config/src/main/java/com/epam/aidial/core/config/Config.java index 6a7f834f4..e61786f5e 100644 --- a/config/src/main/java/com/epam/aidial/core/config/Config.java +++ b/config/src/main/java/com/epam/aidial/core/config/Config.java @@ -67,6 +67,9 @@ public class Config { private List globalInterceptors = List.of(); + // deployment-wide anchor for the DAY/WEEK/MONTH fixed calendar rate-limit windows + private RateLimitSchedule rateLimitSchedule = new RateLimitSchedule(); + @JsonIgnore public Deployment selectDeployment(String deploymentId) { Application application = applications.get(deploymentId); diff --git a/config/src/main/java/com/epam/aidial/core/config/RateLimitSchedule.java b/config/src/main/java/com/epam/aidial/core/config/RateLimitSchedule.java new file mode 100644 index 000000000..63b0a06ea --- /dev/null +++ b/config/src/main/java/com/epam/aidial/core/config/RateLimitSchedule.java @@ -0,0 +1,22 @@ +package com.epam.aidial.core.config; + +import com.epam.aidial.core.config.validation.ValidTimezone; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import jakarta.validation.constraints.Pattern; +import lombok.Data; + +/** + * Deployment-wide anchor for the DAY/WEEK/MONTH fixed calendar rate-limit windows: what local time + * of day a period resets at, in which timezone, and - for WEEK - which day it starts on. Omitting + * this setting is equivalent to UTC/Mon/00:00, which reproduces the previous implicit UTC-midnight + * day/month behavior and only newly defines a week start day. + */ +@Data +@JsonIgnoreProperties(ignoreUnknown = true) +public class RateLimitSchedule { + @ValidTimezone + private String timezone = "UTC"; + private WeekDay weekStartDay = WeekDay.Mon; + @Pattern(regexp = "^([01][0-9]|2[0-3]):[0-5][0-9]$", message = "resetTime must be in HH:mm 24h format") + private String resetTime = "00:00"; +} diff --git a/config/src/main/java/com/epam/aidial/core/config/WeekDay.java b/config/src/main/java/com/epam/aidial/core/config/WeekDay.java new file mode 100644 index 000000000..f96c35df9 --- /dev/null +++ b/config/src/main/java/com/epam/aidial/core/config/WeekDay.java @@ -0,0 +1,23 @@ +package com.epam.aidial.core.config; + +import java.time.DayOfWeek; + +public enum WeekDay { + Mon(DayOfWeek.MONDAY), + Tue(DayOfWeek.TUESDAY), + Wed(DayOfWeek.WEDNESDAY), + Thu(DayOfWeek.THURSDAY), + Fri(DayOfWeek.FRIDAY), + Sat(DayOfWeek.SATURDAY), + Sun(DayOfWeek.SUNDAY); + + private final DayOfWeek dayOfWeek; + + WeekDay(DayOfWeek dayOfWeek) { + this.dayOfWeek = dayOfWeek; + } + + public DayOfWeek toDayOfWeek() { + return dayOfWeek; + } +} diff --git a/config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezone.java b/config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezone.java new file mode 100644 index 000000000..ba14b1fe4 --- /dev/null +++ b/config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezone.java @@ -0,0 +1,22 @@ +package com.epam.aidial.core.config.validation; + +import jakarta.validation.Constraint; +import jakarta.validation.ReportAsSingleViolation; + +import java.lang.annotation.Documented; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +import static java.lang.annotation.ElementType.FIELD; + +@Documented +@Constraint(validatedBy = { ValidTimezoneValidator.class }) +@Target({ FIELD }) +@Retention(RetentionPolicy.RUNTIME) +@ReportAsSingleViolation +public @interface ValidTimezone { + String message() default "Timezone must be a valid IANA timezone id, e.g. \"Europe/Warsaw\" or \"UTC\""; + Class[] groups() default {}; + Class[] payload() default {}; +} diff --git a/config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezoneValidator.java b/config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezoneValidator.java new file mode 100644 index 000000000..5496b46aa --- /dev/null +++ b/config/src/main/java/com/epam/aidial/core/config/validation/ValidTimezoneValidator.java @@ -0,0 +1,20 @@ +package com.epam.aidial.core.config.validation; + +import jakarta.validation.ConstraintValidator; +import jakarta.validation.ConstraintValidatorContext; + +import java.time.ZoneId; + +public class ValidTimezoneValidator implements ConstraintValidator { + + @Override + public boolean isValid(String value, ConstraintValidatorContext context) { + if (value == null) { + return true; + } + + // restricted to real IANA region ids - a fixed offset like "+02:00" would also parse via + // ZoneId.of, but the schedule must survive DST transitions, which a fixed offset cannot + return ZoneId.getAvailableZoneIds().contains(value); + } +} diff --git a/docs/dynamic-settings/roles.md b/docs/dynamic-settings/roles.md index 37d620d5a..23a67ba80 100644 --- a/docs/dynamic-settings/roles.md +++ b/docs/dynamic-settings/roles.md @@ -60,6 +60,31 @@ An object containing parameters for each [role](#roles). } ``` +## rateLimitSchedule + +A top-level (deployment-wide, not per-role) setting that anchors the fixed calendar windows used by the `day`/`week`/`month`/`requestDay` limits below - see [Fixed calendar windows](#rolesrole_namelimits). It is a single setting for the whole deployment: there is no per-role or per-user customization, and no self-service API for end users to change it. + +Available values: + +* `timezone`: An IANA timezone id (e.g. `"Europe/Warsaw"`, `"America/New_York"`, or `"UTC"`) the schedule is anchored to. Must be a real IANA region id, not a fixed UTC offset, since a fixed offset cannot express daylight-saving transitions. Default: `"UTC"`. +* `weekStartDay`: The day the `week` window starts on, using a 3-letter abbreviation (`Mon`, `Tue`, `Wed`, `Thu`, `Fri`, `Sat`, `Sun`). Default: `"Mon"`. +* `resetTime`: The 24h local time (`HH:mm`) at which `day`/`week`/`month`/`requestDay` periods reset. Default: `"00:00"`. + +Omitting `rateLimitSchedule` entirely is equivalent to `{"timezone": "UTC", "weekStartDay": "Mon", "resetTime": "00:00"}`, which reproduces the previous implicit UTC-midnight day/month behavior. + +**Example**: + +```json +{ + "rateLimitSchedule": { + "timezone": "Europe/Warsaw", + "weekStartDay": "Mon", + "resetTime": "09:00" + }, + "roles": {} +} +``` + #### roles..limits Use to define token usage limits for resources. @@ -68,12 +93,14 @@ Use to define token usage limits for resources. Available values: -* `requestHour`: Total requests per hour that can be sent to a specific resource. -* `requestDay`: Total requests per day that can be sent to a specific resource. +* `requestHour`: Total requests per hour that can be sent to a specific resource, managed via floating window approach for well-distributed rate limiting. +* `requestDay`: Total requests per day that can be sent to a specific resource, managed via a fixed calendar window - see below. * `minute`: Total tokens per minute that can be sent to a specific resource, managed via floating window approach for well-distributed rate limiting. -* `day`: Total tokens per day that can be sent to a specific resource, managed via floating window approach for balanced rate limiting. -* `week`: Total tokens per week that can be sent to a specific resource, managed via floating window approach for balanced rate limiting. -* `month`: Total tokens per month that can be sent to a specific resource, managed via floating window approach for balanced rate limiting. +* `day`: Total tokens per day that can be sent to a specific resource, managed via a fixed calendar window - see below. +* `week`: Total tokens per week that can be sent to a specific resource, managed via a fixed calendar window - see below. +* `month`: Total tokens per month that can be sent to a specific resource, managed via a fixed calendar window - see below. + +**Fixed calendar windows:** unlike `minute`/`requestHour`, which age usage out gradually as time advances, `day`/`week`/`month`/`requestDay` reset all at once at a deterministic boundary aligned to the calendar (e.g. "this month" is the 1st through the last day of the calendar month), so usage cannot exceed the limit until the next boundary regardless of when within the period it accrued. The boundary is anchored to the deployment-wide `rateLimitSchedule` top-level config setting (timezone, week start day, and daily reset time); omitting it defaults to UTC, Monday, 00:00, which reproduces the previous implicit UTC-midnight day/month behavior. The bulk limits/usage endpoints (`GET /v1/deployments/{deployment_id}/limits`, `GET /v1/user/limits`, `GET /v1/user/usage`) report this boundary as `resetsAt` on each fixed-window entry; `minute`/`requestHour` entries omit it, since a floating window has no single reset instant. **Requests served through a translator:** a model interface configured with `"mode": "translator"` does not add its tokens to these limits — the translator calls DIAL Core back to have the completion served, and that second call is what carries the usage, so it is counted once rather than twice. The limits are still checked before a translated request is forwarded, so an exhausted quota blocks it like any other request. `requestHour` and `requestDay` work the same way: the translated call is checked against them but spends no slot, so one client request never consumes two. Refer to [translators](translators.md#limits-and-a-translated-request). diff --git a/docs/open_api_core.yaml b/docs/open_api_core.yaml index 40e64b22c..00925d8cb 100644 --- a/docs/open_api_core.yaml +++ b/docs/open_api_core.yaml @@ -12791,7 +12791,7 @@ paths: summary: /v1/user/limits operationId: getUserLimits description: | - Returns limits and current rolling usage for every deployment available to the authenticated caller - + Returns limits and current usage for every deployment available to the authenticated caller - a JWT user or an API-key project. It replaces calling `/v1/deployments/{deployment_name}/limits` once per deployment, and one response labels an entire model picker, so switching model needs no refetch. @@ -12814,11 +12814,20 @@ paths: (`Long.MAX_VALUE`), meaning unlimited. That value exceeds JavaScript's `Number.MAX_SAFE_INTEGER` (`9007199254740991`), so treat any `total` at or above 2^53 as unlimited rather than rendering it as a `used / total` ratio. - 5. **Every window is trailing, not calendar-aligned.** `day` is the last 24 hours, `week` the last 7 - days, `month` the last 30 days - never "since midnight" or "since the 1st". As a result `used` - decreases on its own as older activity ages out; there is no refund and no periodic reset. - Eviction happens in steps on UTC boundaries at each window's granularity (1 hour for `day`, 1 day - for `week` and `month`). + 5. **`minute` and `hour` windows are trailing; `day`, `week` and `month` are calendar-aligned.** + `minuteTokenStats`/`minuteCostStats` and `hourRequestStats` are a trailing window - the last 60 + seconds or 60 minutes - so their `used` decreases on its own as older activity ages out, with no + periodic reset. `day`/`week`/`month` windows instead reset all at once at a deterministic + boundary (e.g. "since midnight" or "since the 1st") anchored to the deployment-wide + `rateLimitSchedule` config setting (default: UTC, Monday, 00:00) - `used` only ever grows within a + period, then drops to zero at the next boundary. + 6. **`resetsAt` names that boundary, only for calendar-aligned windows.** `dayTokenStats`, + `weekTokenStats`, `monthTokenStats`, `dayRequestStats`, `dayCostStats`, `weekCostStats` and + `monthCostStats` carry a `resetsAt` field - an absolute ISO-8601 instant (e.g. + `"2026-10-01T00:00:00+02:00"`), not a countdown - for when that window's usage resets next. + `minuteTokenStats`, `minuteCostStats` and `hourRequestStats` omit the field entirely rather than + serializing it as `null`: a trailing window has no single reset instant to report, so clients + should treat a missing `resetsAt` as a permanent, expected omission rather than poll for it. Per-deployment spend does not reconcile to the global figure. Attribution starts at rollout and does not back-fill, and a model without `pricing` never contributes while still consuming tokens, so the @@ -12846,30 +12855,37 @@ paths: dayTokenStats: total: 10000000 used: 42000 + resetsAt: 2026-09-10T00:00:00Z weekTokenStats: total: 9223372036854775807 used: 180000 + resetsAt: 2026-09-14T00:00:00Z monthTokenStats: total: 9223372036854775807 used: 640000 + resetsAt: 2026-10-01T00:00:00Z hourRequestStats: total: 9223372036854775807 used: 3 dayRequestStats: total: 9223372036854775807 used: 12 + resetsAt: 2026-09-10T00:00:00Z minuteCostStats: total: 9223372036854775807 used: 0.02 dayCostStats: total: 9223372036854775807 used: 1.85 + resetsAt: 2026-09-10T00:00:00Z weekCostStats: total: 9223372036854775807 used: 6.2 + resetsAt: 2026-09-14T00:00:00Z monthCostStats: total: 9223372036854775807 used: 21.4 + resetsAt: 2026-10-01T00:00:00Z gpt-4: minuteTokenStats: total: 100000 @@ -13062,30 +13078,37 @@ paths: dayTokenStats: total: 10000000 used: 42000 + resetsAt: 2026-09-10T00:00:00Z weekTokenStats: total: 9223372036854775807 used: 180000 + resetsAt: 2026-09-14T00:00:00Z monthTokenStats: total: 9223372036854775807 used: 640000 + resetsAt: 2026-10-01T00:00:00Z hourRequestStats: total: 9223372036854775807 used: 3 dayRequestStats: total: 9223372036854775807 used: 12 + resetsAt: 2026-09-10T00:00:00Z minuteCostStats: total: 9223372036854775807 used: 0.02 dayCostStats: total: 9223372036854775807 used: 1.85 + resetsAt: 2026-09-10T00:00:00Z weekCostStats: total: 9223372036854775807 used: 6.2 + resetsAt: 2026-09-14T00:00:00Z monthCostStats: total: 9223372036854775807 used: 21.4 + resetsAt: 2026-10-01T00:00:00Z minuteCostStats: total: 10.0 used: 0.4 @@ -15245,6 +15268,8 @@ components: type: string translators: $ref: "#/components/schemas/MapStringTranslator" + rateLimitSchedule: + $ref: "#/components/schemas/RateLimitSchedule" ConfigFileMigrateRequest: type: object properties: @@ -15370,6 +15395,8 @@ components: type: number used: type: number + resetsAt: + type: string CostLimit: type: object properties: @@ -16124,6 +16151,8 @@ components: type: integer used: type: integer + resetsAt: + type: string JsonNode: type: object Key: @@ -16772,6 +16801,15 @@ components: - PENDING - APPROVED - REJECTED + RateLimitSchedule: + type: object + properties: + resetTime: + type: string + timezone: + type: string + weekStartDay: + $ref: "#/components/schemas/WeekDay" RateRequest: required: - rate @@ -17819,6 +17857,16 @@ components: - SKIPPED ValueNode: type: object + WeekDay: + type: string + enum: + - Mon + - Tue + - Wed + - Thu + - Fri + - Sat + - Sun securitySchemes: ApiKeyAuth: type: apiKey diff --git a/sample/aidial.config.json b/sample/aidial.config.json index d204e63eb..fa1867134 100644 --- a/sample/aidial.config.json +++ b/sample/aidial.config.json @@ -250,6 +250,11 @@ "role": "default" } }, + "rateLimitSchedule": { + "timezone": "UTC", + "weekStartDay": "Mon", + "resetTime": "00:00" + }, "roles": { "default": { "limits": { diff --git a/server/src/main/java/com/epam/aidial/core/server/AiDial.java b/server/src/main/java/com/epam/aidial/core/server/AiDial.java index efba0a38e..b20b5a625 100644 --- a/server/src/main/java/com/epam/aidial/core/server/AiDial.java +++ b/server/src/main/java/com/epam/aidial/core/server/AiDial.java @@ -298,7 +298,7 @@ vertx, settings("config"), null, RuleService ruleService = new RuleService(resourceService); AccessService accessService = new AccessService(encryptionService, shareService, ruleService, applicationSchemaService, settings("access")); NotificationService notificationService = new NotificationService(resourceService, encryptionService); - RateLimiter rateLimiter = new RateLimiter(taskExecutor, resourceService); + RateLimiter rateLimiter = new RateLimiter(taskExecutor, resourceService, configStore); CodeInterpreterService codeInterpreterService = new CodeInterpreterService(vertx, taskExecutor, redis, resourceService, accessService, encryptionService, operatorService, generator, settings("codeInterpreter")); diff --git a/server/src/main/java/com/epam/aidial/core/server/data/CostItemLimitStats.java b/server/src/main/java/com/epam/aidial/core/server/data/CostItemLimitStats.java index 8a2322b2b..046c827f0 100644 --- a/server/src/main/java/com/epam/aidial/core/server/data/CostItemLimitStats.java +++ b/server/src/main/java/com/epam/aidial/core/server/data/CostItemLimitStats.java @@ -1,5 +1,6 @@ package com.epam.aidial.core.server.data; +import com.fasterxml.jackson.annotation.JsonInclude; import lombok.Data; import java.math.BigDecimal; @@ -11,4 +12,11 @@ public class CostItemLimitStats { private BigDecimal total = BigDecimal.ZERO; private BigDecimal used = BigDecimal.ZERO; -} \ No newline at end of file + /** + * Absolute instant this window's usage resets, as an ISO-8601 offset date-time - present only + * for the fixed calendar windows (DAY/WEEK/MONTH), omitted for the floating ones (MINUTE/HOUR), + * which have no single reset instant to report. + */ + @JsonInclude(JsonInclude.Include.NON_NULL) + private String resetsAt; +} diff --git a/server/src/main/java/com/epam/aidial/core/server/data/ItemLimitStats.java b/server/src/main/java/com/epam/aidial/core/server/data/ItemLimitStats.java index 43eb4150b..f89f49066 100644 --- a/server/src/main/java/com/epam/aidial/core/server/data/ItemLimitStats.java +++ b/server/src/main/java/com/epam/aidial/core/server/data/ItemLimitStats.java @@ -1,9 +1,17 @@ package com.epam.aidial.core.server.data; +import com.fasterxml.jackson.annotation.JsonInclude; import lombok.Data; @Data public class ItemLimitStats { private long total; private long used; + /** + * Absolute instant this window's usage resets, as an ISO-8601 offset date-time - present only + * for the fixed calendar windows (DAY/WEEK/MONTH), omitted for the floating ones (MINUTE/HOUR), + * which have no single reset instant to report. + */ + @JsonInclude(JsonInclude.Include.NON_NULL) + private String resetsAt; } diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/CalendarPeriod.java b/server/src/main/java/com/epam/aidial/core/server/limiter/CalendarPeriod.java new file mode 100644 index 000000000..83abc3165 --- /dev/null +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/CalendarPeriod.java @@ -0,0 +1,12 @@ +package com.epam.aidial.core.server.limiter; + +/** + * A fixed calendar rate-limit window: usage resets all at once at a deterministic boundary + * computed from the deployment-wide {@link com.epam.aidial.core.config.RateLimitSchedule}, unlike + * {@link RateWindow}'s MINUTE/HOUR, which keep aging usage out gradually. + */ +public enum CalendarPeriod { + DAY, + WEEK, + MONTH +} diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculator.java b/server/src/main/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculator.java new file mode 100644 index 000000000..751f2312a --- /dev/null +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculator.java @@ -0,0 +1,83 @@ +package com.epam.aidial.core.server.limiter; + +import com.epam.aidial.core.config.RateLimitSchedule; +import lombok.experimental.UtilityClass; + +import java.time.DayOfWeek; +import java.time.Instant; +import java.time.LocalTime; +import java.time.ZoneId; +import java.time.ZonedDateTime; + +/** + * Computes DAY/WEEK/MONTH period boundaries against a deployment-wide {@link RateLimitSchedule}, + * using {@link ZonedDateTime} arithmetic throughout so a "day"/"week"/"month" always advances by + * one calendar unit in local wall-clock time, regardless of how many real hours a DST transition + * makes that span (23h on a spring-forward day, 25h on a fall-back day). + */ +@UtilityClass +public class CalendarWindowCalculator { + + public long currentPeriodStart(CalendarPeriod period, long timestampMillis, RateLimitSchedule schedule) { + ZonedDateTime now = now(timestampMillis, schedule); + return periodStart(period, now, schedule).toInstant().toEpochMilli(); + } + + /** + * The instant the period containing {@code timestampMillis} ends - i.e. when the next period + * starts. This is both what {@link FixedRateBucket}/{@link CostFixedRateBucket} advance to on + * rollover, and what gets reported to callers as {@code resetsAt}. + */ + public long nextPeriodStart(CalendarPeriod period, long timestampMillis, RateLimitSchedule schedule) { + ZonedDateTime now = now(timestampMillis, schedule); + ZonedDateTime periodStart = periodStart(period, now, schedule); + return advance(period, periodStart).toInstant().toEpochMilli(); + } + + private ZonedDateTime now(long timestampMillis, RateLimitSchedule schedule) { + ZoneId zone = ZoneId.of(schedule.getTimezone()); + return Instant.ofEpochMilli(timestampMillis).atZone(zone); + } + + private ZonedDateTime periodStart(CalendarPeriod period, ZonedDateTime now, RateLimitSchedule schedule) { + return switch (period) { + case DAY -> dayPeriodStart(now, schedule); + case WEEK -> weekPeriodStart(now, schedule); + case MONTH -> monthPeriodStart(now, schedule); + }; + } + + private ZonedDateTime advance(CalendarPeriod period, ZonedDateTime periodStart) { + return switch (period) { + case DAY -> periodStart.plusDays(1); + case WEEK -> periodStart.plusWeeks(1); + case MONTH -> periodStart.plusMonths(1); + }; + } + + private ZonedDateTime dayPeriodStart(ZonedDateTime now, RateLimitSchedule schedule) { + ZonedDateTime candidate = todayReset(now, schedule); + return candidate.isAfter(now) ? candidate.minusDays(1) : candidate; + } + + private ZonedDateTime weekPeriodStart(ZonedDateTime now, RateLimitSchedule schedule) { + ZonedDateTime todayReset = todayReset(now, schedule); + DayOfWeek weekStartDay = schedule.getWeekStartDay().toDayOfWeek(); + int daysSinceWeekStart = (todayReset.getDayOfWeek().getValue() - weekStartDay.getValue() + 7) % 7; + ZonedDateTime candidate = todayReset.minusDays(daysSinceWeekStart); + return candidate.isAfter(now) ? candidate.minusWeeks(1) : candidate; + } + + private ZonedDateTime monthPeriodStart(ZonedDateTime now, RateLimitSchedule schedule) { + ZonedDateTime candidate = now.toLocalDate().withDayOfMonth(1).atTime(resetTime(schedule)).atZone(now.getZone()); + return candidate.isAfter(now) ? candidate.minusMonths(1) : candidate; + } + + private ZonedDateTime todayReset(ZonedDateTime now, RateLimitSchedule schedule) { + return now.toLocalDate().atTime(resetTime(schedule)).atZone(now.getZone()); + } + + private LocalTime resetTime(RateLimitSchedule schedule) { + return LocalTime.parse(schedule.getResetTime()); + } +} diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/CostFixedRateBucket.java b/server/src/main/java/com/epam/aidial/core/server/limiter/CostFixedRateBucket.java new file mode 100644 index 000000000..ee77ce81a --- /dev/null +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/CostFixedRateBucket.java @@ -0,0 +1,46 @@ +package com.epam.aidial.core.server.limiter; + +import com.epam.aidial.core.config.RateLimitSchedule; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import lombok.Data; +import lombok.NoArgsConstructor; + +import java.math.BigDecimal; +import java.util.concurrent.TimeUnit; + +/** + * {@link FixedRateBucket}'s BigDecimal analogue, for cost-based fixed calendar windows. See + * {@link FixedRateBucket} for why {@code ignoreUnknown = true} is load-bearing for the rollout. + */ +@Data +@NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +public class CostFixedRateBucket { + + private long periodStart = Long.MIN_VALUE; + private BigDecimal count = BigDecimal.ZERO; + + public BigDecimal reconcile(long timestamp, CalendarPeriod period, RateLimitSchedule schedule) { + long currentPeriodStart = CalendarWindowCalculator.currentPeriodStart(period, timestamp, schedule); + if (currentPeriodStart != periodStart) { + periodStart = currentPeriodStart; + count = BigDecimal.ZERO; + } + return count; + } + + public BigDecimal add(long timestamp, CalendarPeriod period, RateLimitSchedule schedule, BigDecimal amount) { + reconcile(timestamp, period, schedule); + count = count.add(amount); + return count; + } + + long retryAfterSeconds(long timestamp, CalendarPeriod period, RateLimitSchedule schedule) { + long resetsAt = CalendarWindowCalculator.nextPeriodStart(period, timestamp, schedule); + return TimeUnit.MILLISECONDS.toSeconds(resetsAt - timestamp); + } + + long resetsAtMillis(long timestamp, CalendarPeriod period, RateLimitSchedule schedule) { + return CalendarWindowCalculator.nextPeriodStart(period, timestamp, schedule); + } +} diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/CostRateLimit.java b/server/src/main/java/com/epam/aidial/core/server/limiter/CostRateLimit.java index 1ac6d2677..f35ce658d 100644 --- a/server/src/main/java/com/epam/aidial/core/server/limiter/CostRateLimit.java +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/CostRateLimit.java @@ -1,11 +1,11 @@ package com.epam.aidial.core.server.limiter; import com.epam.aidial.core.config.CostLimit; +import com.epam.aidial.core.config.RateLimitSchedule; import com.epam.aidial.core.server.data.LimitStats; import com.epam.aidial.core.storage.http.HttpStatus; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import lombok.Data; -import org.apache.commons.lang3.math.NumberUtils; import java.math.BigDecimal; import java.text.NumberFormat; @@ -23,61 +23,70 @@ public class CostRateLimit { private static final NumberFormat CURRENCY_FORMAT = NumberFormat.getCurrencyInstance(Locale.US); private final CostRateBucket minute = new CostRateBucket(RateWindow.MINUTE); - private final CostRateBucket day = new CostRateBucket(RateWindow.DAY); - private final CostRateBucket week = new CostRateBucket(RateWindow.WEEK); - private final CostRateBucket month = new CostRateBucket(RateWindow.MONTH); + private final CostFixedRateBucket day = new CostFixedRateBucket(); + private final CostFixedRateBucket week = new CostFixedRateBucket(); + private final CostFixedRateBucket month = new CostFixedRateBucket(); /** * Adds cost usage to all buckets. * * @param timestamp The current timestamp + * @param schedule The deployment-wide fixed-window schedule * @param cost The cost to add */ - public void add(long timestamp, BigDecimal cost) { + public void add(long timestamp, RateLimitSchedule schedule, BigDecimal cost) { if (cost == null || cost.compareTo(BigDecimal.ZERO) <= 0) { return; } - + minute.add(timestamp, cost); - day.add(timestamp, cost); - week.add(timestamp, cost); - month.add(timestamp, cost); + day.add(timestamp, CalendarPeriod.DAY, schedule, cost); + week.add(timestamp, CalendarPeriod.WEEK, schedule, cost); + month.add(timestamp, CalendarPeriod.MONTH, schedule, cost); } /** * Checks if any cost limit is exceeded. * * @param timestamp The current timestamp + * @param schedule The deployment-wide fixed-window schedule * @param costLimit The cost limits to check against * @return A RateLimitResult indicating success or failure */ - public RateLimitResult check(long timestamp, CostLimit costLimit) { + public RateLimitResult check(long timestamp, RateLimitSchedule schedule, CostLimit costLimit) { BigDecimal minuteTotal = minute.update(timestamp); - BigDecimal dayTotal = day.update(timestamp); - BigDecimal weekTotal = week.update(timestamp); - BigDecimal monthTotal = month.update(timestamp); + BigDecimal dayTotal = day.reconcile(timestamp, CalendarPeriod.DAY, schedule); + BigDecimal weekTotal = week.reconcile(timestamp, CalendarPeriod.WEEK, schedule); + BigDecimal monthTotal = month.reconcile(timestamp, CalendarPeriod.MONTH, schedule); - boolean result = minuteTotal.compareTo(costLimit.getMinute()) >= 0 + boolean result = minuteTotal.compareTo(costLimit.getMinute()) >= 0 || dayTotal.compareTo(costLimit.getDay()) >= 0 - || weekTotal.compareTo(costLimit.getWeek()) >= 0 + || weekTotal.compareTo(costLimit.getWeek()) >= 0 || monthTotal.compareTo(costLimit.getMonth()) >= 0; - + if (result) { String errorMsg = String.format( "Hit cost rate limit. Minute limit: %s / %s. Day limit: %s / %s. Week limit: %s / %s. Month limit: %s / %s.", format(minuteTotal), format(costLimit.getMinute()), format(dayTotal), format(costLimit.getDay()), format(weekTotal), format(costLimit.getWeek()), format(monthTotal), format(costLimit.getMonth())); - - long minuteRetryAfter = minute.retryAfter(costLimit.getMinute()); - long dayRetryAfter = day.retryAfter(costLimit.getDay()); - long weekRetryAfter = week.retryAfter(costLimit.getWeek()); - long monthRetryAfter = month.retryAfter(costLimit.getMonth()); - - long retryAfter = NumberUtils.max(minuteRetryAfter, dayRetryAfter, weekRetryAfter, monthRetryAfter); - + + long retryAfter = 0; + if (minuteTotal.compareTo(costLimit.getMinute()) >= 0) { + retryAfter = Math.max(retryAfter, minute.retryAfter(costLimit.getMinute())); + } + if (dayTotal.compareTo(costLimit.getDay()) >= 0) { + retryAfter = Math.max(retryAfter, day.retryAfterSeconds(timestamp, CalendarPeriod.DAY, schedule)); + } + if (weekTotal.compareTo(costLimit.getWeek()) >= 0) { + retryAfter = Math.max(retryAfter, week.retryAfterSeconds(timestamp, CalendarPeriod.WEEK, schedule)); + } + if (monthTotal.compareTo(costLimit.getMonth()) >= 0) { + retryAfter = Math.max(retryAfter, month.retryAfterSeconds(timestamp, CalendarPeriod.MONTH, schedule)); + } + List limits = new ArrayList<>(); StringBuilder displayError = new StringBuilder("You've exceeded your"); - + if (monthTotal.compareTo(costLimit.getMonth()) >= 0) { limits.add("monthly"); } @@ -90,7 +99,7 @@ public RateLimitResult check(long timestamp, CostLimit costLimit) { if (minuteTotal.compareTo(costLimit.getMinute()) >= 0) { limits.add("minute"); } - + for (int i = 0; i < limits.size(); i++) { if (i > 0) { if (i == limits.size() - 1) { @@ -102,12 +111,12 @@ public RateLimitResult check(long timestamp, CostLimit costLimit) { displayError.append(' '); displayError.append(limits.get(i)); } - + displayError.append(" cost limit"); if (limits.size() > 1) { displayError.append('s'); } - + return new RateLimitResult(HttpStatus.TOO_MANY_REQUESTS, errorMsg, displayError.toString(), retryAfter); } else { return RateLimitResult.SUCCESS; @@ -122,17 +131,18 @@ private static String format(BigDecimal n) { * Updates the limit statistics with the current usage. * * @param timestamp The current timestamp + * @param schedule The deployment-wide fixed-window schedule * @param limitStats The limit statistics to update */ - public void update(long timestamp, LimitStats limitStats) { + public void update(long timestamp, RateLimitSchedule schedule, LimitStats limitStats) { BigDecimal minuteTotal = minute.update(timestamp); - BigDecimal dayTotal = day.update(timestamp); - BigDecimal weekTotal = week.update(timestamp); - BigDecimal monthTotal = month.update(timestamp); - + BigDecimal dayTotal = day.reconcile(timestamp, CalendarPeriod.DAY, schedule); + BigDecimal weekTotal = week.reconcile(timestamp, CalendarPeriod.WEEK, schedule); + BigDecimal monthTotal = month.reconcile(timestamp, CalendarPeriod.MONTH, schedule); + limitStats.getMinuteCostStats().setUsed(minuteTotal); limitStats.getDayCostStats().setUsed(dayTotal); limitStats.getWeekCostStats().setUsed(weekTotal); limitStats.getMonthCostStats().setUsed(monthTotal); } -} \ No newline at end of file +} diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/FixedRateBucket.java b/server/src/main/java/com/epam/aidial/core/server/limiter/FixedRateBucket.java new file mode 100644 index 000000000..69ec795f4 --- /dev/null +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/FixedRateBucket.java @@ -0,0 +1,63 @@ +package com.epam.aidial.core.server.limiter; + +import com.epam.aidial.core.config.RateLimitSchedule; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import lombok.Data; +import lombok.NoArgsConstructor; + +import java.util.concurrent.TimeUnit; + +/** + * A fixed calendar-window counter: unlike {@link RateBucket}'s sliding sub-intervals, this only + * ever needs "how much has happened in the period we're currently in" - a read recomputes what + * {@code periodStart} should be for "now", and treats {@code count} as reset to zero if it doesn't + * match the stored value. + * + *

{@code ignoreUnknown = true} is required, not cosmetic: it lets a pre-rollout record, still + * shaped like the old floating-window {@link RateBucket} ({@code window}/{@code sums}/{@code sum}/ + * {@code start}/{@code end}), deserialize here instead of throwing. Its unrecognized fields are + * dropped, {@code periodStart} is left at its sentinel default, and {@link #reconcile} then sees + * "not a match" on first touch and zeroes {@code count} - the implicit one-time reset the rollout + * relies on, without any explicit migration code. + */ +@Data +@NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +public class FixedRateBucket { + + private long periodStart = Long.MIN_VALUE; + private long count; + + /** + * Rolls the bucket over to the current calendar period if it has changed since the last touch, + * and returns the (possibly just-reset) running total for that period. + */ + public long reconcile(long timestamp, CalendarPeriod period, RateLimitSchedule schedule) { + long currentPeriodStart = CalendarWindowCalculator.currentPeriodStart(period, timestamp, schedule); + if (currentPeriodStart != periodStart) { + periodStart = currentPeriodStart; + count = 0; + } + return count; + } + + public long add(long timestamp, CalendarPeriod period, RateLimitSchedule schedule, long amount) { + reconcile(timestamp, period, schedule); + count += amount; + return count; + } + + /** + * Seconds until this window's next reset. Only meaningful when the caller already knows this + * window is over its limit - unlike {@link RateBucket#retryAfter}, there is no "walk forward + * and stop early" degenerate case that returns 0 when under limit. + */ + long retryAfterSeconds(long timestamp, CalendarPeriod period, RateLimitSchedule schedule) { + long resetsAt = CalendarWindowCalculator.nextPeriodStart(period, timestamp, schedule); + return TimeUnit.MILLISECONDS.toSeconds(resetsAt - timestamp); + } + + long resetsAtMillis(long timestamp, CalendarPeriod period, RateLimitSchedule schedule) { + return CalendarWindowCalculator.nextPeriodStart(period, timestamp, schedule); + } +} diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/RateLimiter.java b/server/src/main/java/com/epam/aidial/core/server/limiter/RateLimiter.java index 2cd1545d6..804416e67 100644 --- a/server/src/main/java/com/epam/aidial/core/server/limiter/RateLimiter.java +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/RateLimiter.java @@ -4,9 +4,11 @@ import com.epam.aidial.core.config.Deployment; import com.epam.aidial.core.config.InterfaceType; import com.epam.aidial.core.config.Limit; +import com.epam.aidial.core.config.RateLimitSchedule; import com.epam.aidial.core.config.Role; import com.epam.aidial.core.config.RoleBasedEntity; import com.epam.aidial.core.server.ProxyContext; +import com.epam.aidial.core.server.config.ConfigStore; import com.epam.aidial.core.server.data.CostItemLimitStats; import com.epam.aidial.core.server.data.ItemLimitStats; import com.epam.aidial.core.server.data.LimitStats; @@ -30,6 +32,10 @@ import org.apache.commons.lang3.tuple.Pair; import java.math.BigDecimal; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneId; +import java.time.format.DateTimeFormatter; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -43,11 +49,17 @@ public class RateLimiter { private static final Limit DEFAULT_LIMIT = new Limit(); private static final CostLimit DEFAULT_COST_LIMIT = new CostLimit(); private static final String DEFAULT_USER_ROLE = "default"; + // a safe "definitely stale, skip reading" upper bound for a calendar month (31 days) plus + // slack for a reset time near midnight in any configured zone; the authoritative correctness + // check is FixedRateBucket/CostFixedRateBucket#reconcile on actual read of a candidate record + private static final long WIDEST_WINDOW_MILLIS = Duration.ofDays(32).toMillis(); private final AsyncTaskExecutor taskExecutor; private final ResourceService resourceService; + private final ConfigStore configStore; + public Future increase( RoleBasedEntity roleBasedEntity, String bucket, TokenUsage usage, Buffer requestBody, Buffer responseBody, InterfaceType interfaceType, JsonNode liveUsageNode) { @@ -57,6 +69,7 @@ public Future increase( return Future.succeededFuture(); } + RateLimitSchedule schedule = configStore.get().getRateLimitSchedule(); BigDecimal cost = ModelCostCalculator.calculate( roleBasedEntity, usage, requestBody, responseBody, interfaceType, liveUsageNode); Future costFuture; @@ -66,7 +79,7 @@ public Future increase( usage.setAggCost(cost); } - costFuture = updateCostLimits(roleBasedEntity, bucket, cost); + costFuture = updateCostLimits(roleBasedEntity, bucket, cost, schedule); } else { costFuture = Future.succeededFuture(); } @@ -77,7 +90,7 @@ public Future increase( } else { String tokensPath = getPathToTokens(roleBasedEntity.getName()); ResourceDescriptor tokenResourceDescription = getResourceDescription(bucket, tokensPath); - tokenFuture = taskExecutor.submit(() -> updateTokenLimit(tokenResourceDescription, usage.getTotalTokens())); + tokenFuture = taskExecutor.submit(() -> updateTokenLimit(tokenResourceDescription, usage.getTotalTokens(), schedule)); } // Wait for every update to complete @@ -126,11 +139,12 @@ public Future getLimitStats(RoleBasedEntity roleBasedEntity, ProxyCo private LimitStats getLimitStats(ProxyContext context, Limit limit, String name) { CostLimit costLimit = getCostLimitByUser(context); - LimitStats limitStats = create(limit, costLimit); + RateLimitSchedule schedule = context.getConfig().getRateLimitSchedule(); long timestamp = System.currentTimeMillis(); - collectTokenLimitStats(context, limitStats, timestamp, name); - collectRequestLimitStats(context, limitStats, timestamp, name); - collectCostLimitStats(context, limitStats, timestamp); + LimitStats limitStats = create(limit, costLimit, timestamp, schedule); + collectTokenLimitStats(context, limitStats, timestamp, name, schedule); + collectRequestLimitStats(context, limitStats, timestamp, name, schedule); + collectCostLimitStats(context, limitStats, timestamp, schedule); return limitStats; } @@ -161,6 +175,8 @@ public Future getUserStats( private UserLimitStats collectUserStats( ProxyContext context, List deployments, boolean dropEmpty) { String bucketLocation = BucketBuilder.buildInitiatorBucket(context); + RateLimitSchedule schedule = context.getConfig().getRateLimitSchedule(); + long timestamp = System.currentTimeMillis(); UserLimitStats userLimitStats = new UserLimitStats(); Map statsByDeployment = userLimitStats.getDeployments(); @@ -170,7 +186,7 @@ private UserLimitStats collectUserStats( Limit limit = getLimitByUser(context, deployment); // DEFAULT_COST_LIMIT leaves every cost window at the unlimited sentinel: an entry reports the // deployment's attributed spend, and only the global budget can cap it - LimitStats limitStats = create(limit, DEFAULT_COST_LIMIT); + LimitStats limitStats = create(limit, DEFAULT_COST_LIMIT, timestamp, schedule); String tokensPath = getLimitAbsolutePath(bucketLocation, getPathToTokens(name)); String requestsPath = getLimitAbsolutePath(bucketLocation, getPathToRequests(name)); String costsPath = getLimitAbsolutePath(bucketLocation, getPathToDeploymentCosts(name)); @@ -184,7 +200,9 @@ private UserLimitStats collectUserStats( // record is a sibling of theirs, so a deployment named "costs" lands on "costs/costs" and cannot // collide with it CostLimit userCostLimit = getCostLimitByUser(context); - LimitStats costStats = create(DEFAULT_LIMIT, userCostLimit); + // token/request fields on this entry are discarded below - only its four cost pairs are copied + // out into UserLimitStats - so computing resetsAt for them here is harmless waste, not a bug + LimitStats costStats = create(DEFAULT_LIMIT, userCostLimit, timestamp, schedule); String userCostsPath = getLimitAbsolutePath(bucketLocation, getPathToCosts()); targetsByRecordPath.put(userCostsPath, new StatsTarget(costStats, LimitType.COSTS)); @@ -193,7 +211,7 @@ private UserLimitStats collectUserStats( for (Pair loaded : records) { String path = loaded.getKey().getDescriptor().getAbsoluteFilePath(); StatsTarget target = targetsByRecordPath.get(path); - target.type().collect(loaded.getValue(), target.stats()); + target.type().collect(loaded.getValue(), target.stats(), timestamp, schedule); } userLimitStats.setMinuteCostStats(costStats.getMinuteCostStats()); @@ -210,15 +228,14 @@ private UserLimitStats collectUserStats( /** * Lists the caller's {@code limits/} folder and loads the bodies of the records that belong to the - * response, ignoring any other name the listing turns up. A record last written before the widest - * window opened is left unread: {@link RateWindow#MONTH} keeps 30 one-day intervals, so nothing older - * can still project to a non-zero figure. The age comes from the listing entry, so skipping one costs - * nothing extra. + * response, ignoring any other name the listing turns up. A record last written before + * {@link #WIDEST_WINDOW_MILLIS} ago is left unread as a "definitely stale" pre-filter. The age comes + * from the listing entry, so skipping one costs nothing extra. */ private List> loadLimitRecords(String bucketLocation, Set wanted) { ResourceDescriptor folder = ResourceDescriptorFactory .fromEncoded(ResourceTypes.LIMIT, bucketLocation, bucketLocation, null); - long updatedAfter = System.currentTimeMillis() - RateWindow.MONTH.window(); + long updatedAfter = System.currentTimeMillis() - WIDEST_WINDOW_MILLIS; // the page's item list is immutable, so the unwanted records are filtered out by replacing it return resourceService.listResources(folder, page -> page.setItems(page.getItems().stream() .filter(item -> item instanceof ResourceItemMetadata metadata @@ -261,71 +278,74 @@ private record StatsTarget(LimitStats stats, LimitType type) { private enum LimitType { TOKENS { @Override - void collect(String json, LimitStats stats) { - collectTokenLimitStats(json, stats, System.currentTimeMillis()); + void collect(String json, LimitStats stats, long timestamp, RateLimitSchedule schedule) { + collectTokenLimitStats(json, stats, timestamp, schedule); } }, REQUESTS { @Override - void collect(String json, LimitStats stats) { - collectRequestLimitStats(json, stats, System.currentTimeMillis()); + void collect(String json, LimitStats stats, long timestamp, RateLimitSchedule schedule) { + collectRequestLimitStats(json, stats, timestamp, schedule); } }, COSTS { @Override - void collect(String json, LimitStats stats) { - collectCostLimitStats(json, stats, System.currentTimeMillis()); + void collect(String json, LimitStats stats, long timestamp, RateLimitSchedule schedule) { + collectCostLimitStats(json, stats, timestamp, schedule); } }; - abstract void collect(String json, LimitStats stats); + abstract void collect(String json, LimitStats stats, long timestamp, RateLimitSchedule schedule); } - private void collectTokenLimitStats(ProxyContext context, LimitStats limitStats, long timestamp, String name) { + private void collectTokenLimitStats( + ProxyContext context, LimitStats limitStats, long timestamp, String name, RateLimitSchedule schedule) { ResourceDescriptor resourceDescription = getResourceDescription(context, getPathToTokens(name)); - collectTokenLimitStats(resourceService.getResource(resourceDescription), limitStats, timestamp); + collectTokenLimitStats(resourceService.getResource(resourceDescription), limitStats, timestamp, schedule); } - private static void collectTokenLimitStats(String json, LimitStats limitStats, long timestamp) { + private static void collectTokenLimitStats(String json, LimitStats limitStats, long timestamp, RateLimitSchedule schedule) { TokenRateLimit rateLimit = ProxyUtil.convertToObject(json, TokenRateLimit.class); if (rateLimit == null) { return; } - rateLimit.update(timestamp, limitStats); + rateLimit.update(timestamp, schedule, limitStats); } - private void collectRequestLimitStats(ProxyContext context, LimitStats limitStats, long timestamp, String name) { + private void collectRequestLimitStats( + ProxyContext context, LimitStats limitStats, long timestamp, String name, RateLimitSchedule schedule) { ResourceDescriptor resourceDescription = getResourceDescription(context, getPathToRequests(name)); - collectRequestLimitStats(resourceService.getResource(resourceDescription), limitStats, timestamp); + collectRequestLimitStats(resourceService.getResource(resourceDescription), limitStats, timestamp, schedule); } - private static void collectRequestLimitStats(String json, LimitStats limitStats, long timestamp) { + private static void collectRequestLimitStats(String json, LimitStats limitStats, long timestamp, RateLimitSchedule schedule) { RequestRateLimit rateLimit = ProxyUtil.convertToObject(json, RequestRateLimit.class); if (rateLimit == null) { return; } - rateLimit.update(timestamp, limitStats); + rateLimit.update(timestamp, schedule, limitStats); } - private void collectCostLimitStats(ProxyContext context, LimitStats limitStats, long timestamp) { + private void collectCostLimitStats(ProxyContext context, LimitStats limitStats, long timestamp, RateLimitSchedule schedule) { ResourceDescriptor resourceDescription = getResourceDescription(context, getPathToCosts()); - collectCostLimitStats(resourceService.getResource(resourceDescription), limitStats, timestamp); + collectCostLimitStats(resourceService.getResource(resourceDescription), limitStats, timestamp, schedule); } - private static void collectCostLimitStats(String json, LimitStats limitStats, long timestamp) { + private static void collectCostLimitStats(String json, LimitStats limitStats, long timestamp, RateLimitSchedule schedule) { CostRateLimit rateLimit = ProxyUtil.convertToObject(json, CostRateLimit.class); if (rateLimit == null) { return; } - rateLimit.update(timestamp, limitStats); + rateLimit.update(timestamp, schedule, limitStats); } - private LimitStats create(Limit limit, CostLimit costLimit) { + private LimitStats create(Limit limit, CostLimit costLimit, long timestamp, RateLimitSchedule schedule) { LimitStats limitStats = new LimitStats(); // Token limits ItemLimitStats dayTokenStats = new ItemLimitStats(); dayTokenStats.setTotal(limit.getDay()); + dayTokenStats.setResetsAt(formatResetsAt(CalendarPeriod.DAY, timestamp, schedule)); limitStats.setDayTokenStats(dayTokenStats); ItemLimitStats minuteTokenStats = new ItemLimitStats(); @@ -334,10 +354,12 @@ private LimitStats create(Limit limit, CostLimit costLimit) { ItemLimitStats weekTokenStats = new ItemLimitStats(); weekTokenStats.setTotal(limit.getWeek()); + weekTokenStats.setResetsAt(formatResetsAt(CalendarPeriod.WEEK, timestamp, schedule)); limitStats.setWeekTokenStats(weekTokenStats); ItemLimitStats monthTokenStats = new ItemLimitStats(); monthTokenStats.setTotal(limit.getMonth()); + monthTokenStats.setResetsAt(formatResetsAt(CalendarPeriod.MONTH, timestamp, schedule)); limitStats.setMonthTokenStats(monthTokenStats); ItemLimitStats hourRequestStats = new ItemLimitStats(); @@ -346,6 +368,7 @@ private LimitStats create(Limit limit, CostLimit costLimit) { ItemLimitStats dayRequestStats = new ItemLimitStats(); dayRequestStats.setTotal(limit.getRequestDay()); + dayRequestStats.setResetsAt(formatResetsAt(CalendarPeriod.DAY, timestamp, schedule)); limitStats.setDayRequestStats(dayRequestStats); if (costLimit != null) { @@ -355,20 +378,35 @@ private LimitStats create(Limit limit, CostLimit costLimit) { CostItemLimitStats dayCostStats = new CostItemLimitStats(); dayCostStats.setTotal(costLimit.getDay()); + dayCostStats.setResetsAt(formatResetsAt(CalendarPeriod.DAY, timestamp, schedule)); limitStats.setDayCostStats(dayCostStats); CostItemLimitStats weekCostStats = new CostItemLimitStats(); weekCostStats.setTotal(costLimit.getWeek()); + weekCostStats.setResetsAt(formatResetsAt(CalendarPeriod.WEEK, timestamp, schedule)); limitStats.setWeekCostStats(weekCostStats); CostItemLimitStats monthCostStats = new CostItemLimitStats(); monthCostStats.setTotal(costLimit.getMonth()); + monthCostStats.setResetsAt(formatResetsAt(CalendarPeriod.MONTH, timestamp, schedule)); limitStats.setMonthCostStats(monthCostStats); } return limitStats; } + /** + * The absolute instant a fixed calendar window resets, formatted with the offset the configured + * timezone actually observes at that instant (so it can differ across two {@code resetsAt} values + * that share the same local wall-clock time, if a DST transition happened between them). Computable + * from "now" and the schedule alone, with no stored usage record involved. + */ + private static String formatResetsAt(CalendarPeriod period, long timestamp, RateLimitSchedule schedule) { + long resetsAtMillis = CalendarWindowCalculator.nextPeriodStart(period, timestamp, schedule); + return DateTimeFormatter.ISO_OFFSET_DATE_TIME.format( + Instant.ofEpochMilli(resetsAtMillis).atZone(ZoneId.of(schedule.getTimezone()))); + } + private ResourceDescriptor getResourceDescription(ProxyContext context, String path) { // use bucket location of request's initiator, // e.g. user -> core -> application -> core -> model, limits must be applied to the user by JWT @@ -390,25 +428,26 @@ private ResourceDescriptor getResourceDescription(String bucketLocation, String private RateLimitResult checkLimit(ProxyContext context, Limit limit, RoleBasedEntity roleBasedEntity) { long timestamp = System.currentTimeMillis(); + RateLimitSchedule schedule = context.getConfig().getRateLimitSchedule(); // Check token limits - RateLimitResult tokenResult = checkTokenLimit(context, limit, timestamp, roleBasedEntity); + RateLimitResult tokenResult = checkTokenLimit(context, limit, timestamp, schedule, roleBasedEntity); if (tokenResult.status() != HttpStatus.OK) { return tokenResult; } // Check request limits - RateLimitResult requestResult = checkRequestLimit(context, limit, timestamp, roleBasedEntity); + RateLimitResult requestResult = checkRequestLimit(context, limit, timestamp, schedule, roleBasedEntity); if (requestResult.status() != HttpStatus.OK) { return requestResult; } // Check cost limits CostLimit costLimit = getCostLimitByUser(context); - return checkCostLimit(context, costLimit, timestamp); + return checkCostLimit(context, costLimit, timestamp, schedule); } - private RateLimitResult checkCostLimit(ProxyContext context, CostLimit costLimit, long timestamp) { + private RateLimitResult checkCostLimit(ProxyContext context, CostLimit costLimit, long timestamp, RateLimitSchedule schedule) { String costsPath = getPathToCosts(); ResourceDescriptor resourceDescription = getResourceDescription(context, costsPath); String prevValue = resourceService.getResource(resourceDescription); @@ -416,10 +455,11 @@ private RateLimitResult checkCostLimit(ProxyContext context, CostLimit costLimit if (rateLimit == null) { return RateLimitResult.SUCCESS; } - return rateLimit.check(timestamp, costLimit); + return rateLimit.check(timestamp, schedule, costLimit); } - private RateLimitResult checkTokenLimit(ProxyContext context, Limit limit, long timestamp, RoleBasedEntity roleBasedEntity) { + private RateLimitResult checkTokenLimit( + ProxyContext context, Limit limit, long timestamp, RateLimitSchedule schedule, RoleBasedEntity roleBasedEntity) { String tokensPath = getPathToTokens(roleBasedEntity.getName()); ResourceDescriptor resourceDescription = getResourceDescription(context, tokensPath); String prevValue = resourceService.getResource(resourceDescription); @@ -427,43 +467,44 @@ private RateLimitResult checkTokenLimit(ProxyContext context, Limit limit, long if (rateLimit == null) { return RateLimitResult.SUCCESS; } - return rateLimit.update(timestamp, limit); + return rateLimit.update(timestamp, schedule, limit); } - private RateLimitResult checkRequestLimit(ProxyContext context, Limit limit, long timestamp, RoleBasedEntity roleBasedEntity) { + private RateLimitResult checkRequestLimit( + ProxyContext context, Limit limit, long timestamp, RateLimitSchedule schedule, RoleBasedEntity roleBasedEntity) { String tokensPath = getPathToRequests(roleBasedEntity.getName()); ResourceDescriptor resourceDescription = getResourceDescription(context, tokensPath); // pass array to hold rate limit result returned by the function to compute the resource RateLimitResult[] result = new RateLimitResult[1]; - resourceService.computeResource(resourceDescription, json -> updateRequestLimit(json, timestamp, limit, result)); + resourceService.computeResource(resourceDescription, json -> updateRequestLimit(json, timestamp, schedule, limit, result)); return result[0]; } - private String updateRequestLimit(String json, long timestamp, Limit limit, RateLimitResult[] result) { + private String updateRequestLimit(String json, long timestamp, RateLimitSchedule schedule, Limit limit, RateLimitResult[] result) { RequestRateLimit rateLimit = ProxyUtil.convertToObject(json, RequestRateLimit.class); if (rateLimit == null) { rateLimit = new RequestRateLimit(); } - result[0] = rateLimit.check(timestamp, limit, 1); + result[0] = rateLimit.check(timestamp, schedule, limit, 1); return ProxyUtil.convertToString(rateLimit); } - private Void updateTokenLimit(ResourceDescriptor resourceDescription, long totalUsedTokens) { - resourceService.computeResource(resourceDescription, json -> updateTokenLimit(json, totalUsedTokens)); + private Void updateTokenLimit(ResourceDescriptor resourceDescription, long totalUsedTokens, RateLimitSchedule schedule) { + resourceService.computeResource(resourceDescription, json -> updateTokenLimit(json, totalUsedTokens, schedule)); return null; } - private String updateTokenLimit(String json, long totalUsedTokens) { + private String updateTokenLimit(String json, long totalUsedTokens, RateLimitSchedule schedule) { TokenRateLimit rateLimit = ProxyUtil.convertToObject(json, TokenRateLimit.class); if (rateLimit == null) { rateLimit = new TokenRateLimit(); } long timestamp = System.currentTimeMillis(); - rateLimit.add(timestamp, totalUsedTokens); + rateLimit.add(timestamp, schedule, totalUsedTokens); return ProxyUtil.convertToString(rateLimit); } - private Future updateCostLimits(RoleBasedEntity roleBasedEntity, String bucket, BigDecimal cost) { + private Future updateCostLimits(RoleBasedEntity roleBasedEntity, String bucket, BigDecimal cost, RateLimitSchedule schedule) { ResourceDescriptor userCostDescriptor = getResourceDescription(bucket, getPathToCosts()); // the global document is what enforces; the deployment-scoped one only attributes the same // figure, so that a bulk report can break spend down without re-deriving it from stored @@ -472,23 +513,23 @@ private Future updateCostLimits(RoleBasedEntity roleBasedEntity, String bu getResourceDescription(bucket, getPathToDeploymentCosts(roleBasedEntity.getName())); // the enforcing document is written first, since the deployment-scoped one only reports return taskExecutor.submit(() -> { - updateCostLimit(userCostDescriptor, cost); - return updateCostLimit(deploymentCostDescriptor, cost); + updateCostLimit(userCostDescriptor, cost, schedule); + return updateCostLimit(deploymentCostDescriptor, cost, schedule); }); } - private Void updateCostLimit(ResourceDescriptor resourceDescription, BigDecimal cost) { - resourceService.computeResource(resourceDescription, json -> updateCostLimit(json, cost)); + private Void updateCostLimit(ResourceDescriptor resourceDescription, BigDecimal cost, RateLimitSchedule schedule) { + resourceService.computeResource(resourceDescription, json -> updateCostLimit(json, cost, schedule)); return null; } - private String updateCostLimit(String json, BigDecimal cost) { + private String updateCostLimit(String json, BigDecimal cost, RateLimitSchedule schedule) { CostRateLimit rateLimit = ProxyUtil.convertToObject(json, CostRateLimit.class); if (rateLimit == null) { rateLimit = new CostRateLimit(); } long timestamp = System.currentTimeMillis(); - rateLimit.add(timestamp, cost); + rateLimit.add(timestamp, schedule, cost); return ProxyUtil.convertToString(rateLimit); } diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/RateWindow.java b/server/src/main/java/com/epam/aidial/core/server/limiter/RateWindow.java index 95927a05c..0e4f10144 100644 --- a/server/src/main/java/com/epam/aidial/core/server/limiter/RateWindow.java +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/RateWindow.java @@ -7,10 +7,7 @@ @Accessors(fluent = true) public enum RateWindow { MINUTE(60L * 1000, 60), - HOUR(60 * 60 * 1000, 60), - DAY(24L * 60 * 60 * 1000, 24), - WEEK(7L * 24 * 60 * 60 * 1000, 7), - MONTH(30L * 24 * 60 * 60 * 1000, 30); + HOUR(60 * 60 * 1000, 60); private final long window; private final long interval; diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/RequestRateLimit.java b/server/src/main/java/com/epam/aidial/core/server/limiter/RequestRateLimit.java index c55a72d15..dc67fd8dc 100644 --- a/server/src/main/java/com/epam/aidial/core/server/limiter/RequestRateLimit.java +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/RequestRateLimit.java @@ -1,6 +1,7 @@ package com.epam.aidial.core.server.limiter; import com.epam.aidial.core.config.Limit; +import com.epam.aidial.core.config.RateLimitSchedule; import com.epam.aidial.core.server.data.LimitStats; import com.epam.aidial.core.storage.http.HttpStatus; import lombok.Data; @@ -11,19 +12,23 @@ @Data public class RequestRateLimit { private final RateBucket hour = new RateBucket(RateWindow.HOUR); - private final RateBucket day = new RateBucket(RateWindow.DAY); + private final FixedRateBucket day = new FixedRateBucket(); - public RateLimitResult check(long timestamp, Limit limit, long count) { + public RateLimitResult check(long timestamp, RateLimitSchedule schedule, Limit limit, long count) { long hourTotal = hour.update(timestamp); - long dayTotal = day.update(timestamp); + long dayTotal = day.reconcile(timestamp, CalendarPeriod.DAY, schedule); boolean result = hourTotal >= limit.getRequestHour() || dayTotal >= limit.getRequestDay(); if (result) { String errorMsg = String.format("Hit request rate limit. Hour limit: %d / %d requests. Day limit: %d / %d requests.", hourTotal, limit.getRequestHour(), dayTotal, limit.getRequestDay()); - long hourRetryAfter = hour.retryAfter(limit.getRequestHour()); - long dayRetryAfter = day.retryAfter(limit.getRequestDay()); - long retryAfter = Math.max(hourRetryAfter, dayRetryAfter); + long retryAfter = 0; + if (hourTotal >= limit.getRequestHour()) { + retryAfter = Math.max(retryAfter, hour.retryAfter(limit.getRequestHour())); + } + if (dayTotal >= limit.getRequestDay()) { + retryAfter = Math.max(retryAfter, day.retryAfterSeconds(timestamp, CalendarPeriod.DAY, schedule)); + } List limits = new ArrayList<>(); StringBuilder displayError = new StringBuilder("You've exceeded your"); if (dayTotal >= limit.getRequestDay()) { @@ -50,14 +55,14 @@ public RateLimitResult check(long timestamp, Limit limit, long count) { return new RateLimitResult(HttpStatus.TOO_MANY_REQUESTS, errorMsg, displayError.toString(), retryAfter); } else { hour.add(timestamp, count); - day.add(timestamp, count); + day.add(timestamp, CalendarPeriod.DAY, schedule, count); return RateLimitResult.SUCCESS; } } - public void update(long timestamp, LimitStats limitStats) { + public void update(long timestamp, RateLimitSchedule schedule, LimitStats limitStats) { long hourTotal = hour.update(timestamp); - long dayTotal = day.update(timestamp); + long dayTotal = day.reconcile(timestamp, CalendarPeriod.DAY, schedule); limitStats.getDayRequestStats().setUsed(dayTotal); limitStats.getHourRequestStats().setUsed(hourTotal); } diff --git a/server/src/main/java/com/epam/aidial/core/server/limiter/TokenRateLimit.java b/server/src/main/java/com/epam/aidial/core/server/limiter/TokenRateLimit.java index 5fef2a879..f08811178 100644 --- a/server/src/main/java/com/epam/aidial/core/server/limiter/TokenRateLimit.java +++ b/server/src/main/java/com/epam/aidial/core/server/limiter/TokenRateLimit.java @@ -1,11 +1,11 @@ package com.epam.aidial.core.server.limiter; import com.epam.aidial.core.config.Limit; +import com.epam.aidial.core.config.RateLimitSchedule; import com.epam.aidial.core.server.data.LimitStats; import com.epam.aidial.core.storage.http.HttpStatus; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import lombok.Data; -import org.apache.commons.lang3.math.NumberUtils; import java.util.ArrayList; import java.util.List; @@ -15,22 +15,22 @@ public class TokenRateLimit { private final RateBucket minute = new RateBucket(RateWindow.MINUTE); - private final RateBucket day = new RateBucket(RateWindow.DAY); - private final RateBucket week = new RateBucket(RateWindow.WEEK); - private final RateBucket month = new RateBucket(RateWindow.MONTH); + private final FixedRateBucket day = new FixedRateBucket(); + private final FixedRateBucket week = new FixedRateBucket(); + private final FixedRateBucket month = new FixedRateBucket(); - public void add(long timestamp, long count) { + public void add(long timestamp, RateLimitSchedule schedule, long count) { minute.add(timestamp, count); - day.add(timestamp, count); - week.add(timestamp, count); - month.add(timestamp, count); + day.add(timestamp, CalendarPeriod.DAY, schedule, count); + week.add(timestamp, CalendarPeriod.WEEK, schedule, count); + month.add(timestamp, CalendarPeriod.MONTH, schedule, count); } - public RateLimitResult update(long timestamp, Limit limit) { + public RateLimitResult update(long timestamp, RateLimitSchedule schedule, Limit limit) { long minuteTotal = minute.update(timestamp); - long dayTotal = day.update(timestamp); - long weekTotal = week.update(timestamp); - long monthTotal = month.update(timestamp); + long dayTotal = day.reconcile(timestamp, CalendarPeriod.DAY, schedule); + long weekTotal = week.reconcile(timestamp, CalendarPeriod.WEEK, schedule); + long monthTotal = month.reconcile(timestamp, CalendarPeriod.MONTH, schedule); boolean result = minuteTotal >= limit.getMinute() || dayTotal >= limit.getDay() || weekTotal >= limit.getWeek() || monthTotal >= limit.getMonth(); @@ -38,11 +38,19 @@ public RateLimitResult update(long timestamp, Limit limit) { String errorMsg = String.format( "Hit token rate limit. Minute limit: %d / %d tokens. Day limit: %d / %d tokens. Week limit: %d / %d tokens. Month limit: %d / %d tokens.", minuteTotal, limit.getMinute(), dayTotal, limit.getDay(), weekTotal, limit.getWeek(), monthTotal, limit.getMonth()); - long minuteRetryAfter = minute.retryAfter(limit.getMinute()); - long dayRetryAfter = day.retryAfter(limit.getDay()); - long weekRetryAfter = week.retryAfter(limit.getWeek()); - long monthRetryAfter = month.retryAfter(limit.getMonth()); - long retryAfter = NumberUtils.max(minuteRetryAfter, dayRetryAfter, weekRetryAfter, monthRetryAfter); + long retryAfter = 0; + if (minuteTotal >= limit.getMinute()) { + retryAfter = Math.max(retryAfter, minute.retryAfter(limit.getMinute())); + } + if (dayTotal >= limit.getDay()) { + retryAfter = Math.max(retryAfter, day.retryAfterSeconds(timestamp, CalendarPeriod.DAY, schedule)); + } + if (weekTotal >= limit.getWeek()) { + retryAfter = Math.max(retryAfter, week.retryAfterSeconds(timestamp, CalendarPeriod.WEEK, schedule)); + } + if (monthTotal >= limit.getMonth()) { + retryAfter = Math.max(retryAfter, month.retryAfterSeconds(timestamp, CalendarPeriod.MONTH, schedule)); + } List limits = new ArrayList<>(); StringBuilder displayError = new StringBuilder("You've exceeded your"); if (monthTotal >= limit.getMonth()) { @@ -78,11 +86,11 @@ public RateLimitResult update(long timestamp, Limit limit) { } } - public void update(long timestamp, LimitStats limitStats) { + public void update(long timestamp, RateLimitSchedule schedule, LimitStats limitStats) { long minuteTotal = minute.update(timestamp); - long dayTotal = day.update(timestamp); - long weekTotal = week.update(timestamp); - long monthTotal = month.update(timestamp); + long dayTotal = day.reconcile(timestamp, CalendarPeriod.DAY, schedule); + long weekTotal = week.reconcile(timestamp, CalendarPeriod.WEEK, schedule); + long monthTotal = month.reconcile(timestamp, CalendarPeriod.MONTH, schedule); limitStats.getDayTokenStats().setUsed(dayTotal); limitStats.getMinuteTokenStats().setUsed(minuteTotal); limitStats.getWeekTokenStats().setUsed(weekTotal); diff --git a/server/src/test/java/com/epam/aidial/core/server/LimitApiTest.java b/server/src/test/java/com/epam/aidial/core/server/LimitApiTest.java index 82b1a78ea..5da7fdc8f 100644 --- a/server/src/test/java/com/epam/aidial/core/server/LimitApiTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/LimitApiTest.java @@ -1,5 +1,8 @@ package com.epam.aidial.core.server; +import com.epam.aidial.core.config.RateLimitSchedule; +import com.epam.aidial.core.server.limiter.CalendarPeriod; +import com.epam.aidial.core.server.limiter.CalendarWindowCalculator; import com.epam.aidial.core.server.util.ProxyUtil; import com.fasterxml.jackson.databind.JsonNode; import io.vertx.core.http.HttpMethod; @@ -7,6 +10,9 @@ import org.junit.jupiter.api.Test; import java.math.BigDecimal; +import java.time.Instant; +import java.time.ZoneId; +import java.time.format.DateTimeFormatter; import java.util.ArrayList; import java.util.List; @@ -19,6 +25,15 @@ public class LimitApiTest extends ResourceBaseTest { @Test public void testGetLimitStats_Success() { + // the default (unconfigured) rateLimitSchedule - UTC, Monday, 00:00 - is what every dial-config + // fixture here relies on; resetsAt is a pure function of "now" against that schedule, computed + // the same way here and by the server, so it lands on the same instant unless the test happens + // to straddle a period boundary at the exact millisecond + RateLimitSchedule schedule = new RateLimitSchedule(); + String dayResetsAt = resetsAt(CalendarPeriod.DAY, schedule); + String weekResetsAt = resetsAt(CalendarPeriod.WEEK, schedule); + String monthResetsAt = resetsAt(CalendarPeriod.MONTH, schedule); + Response response = send(HttpMethod.GET, "/v1/deployments/test-model-v1/limits", null, null); verifyJson(response, 200, """ { @@ -28,15 +43,18 @@ public void testGetLimitStats_Success() { }, "dayTokenStats": { "total": %d, - "used": %d + "used": %d, + "resetsAt": "%s" }, "weekTokenStats": { "total": %d, - "used": %d + "used": %d, + "resetsAt": "%s" }, "monthTokenStats": { "total": %d, - "used": %d + "used": %d, + "resetsAt": "%s" }, "hourRequestStats": { "total": %d, @@ -44,7 +62,8 @@ public void testGetLimitStats_Success() { }, "dayRequestStats": { "total": %d, - "used": %d + "used": %d, + "resetsAt": "%s" }, "minuteCostStats": { "total": %d, @@ -52,21 +71,37 @@ public void testGetLimitStats_Success() { }, "dayCostStats": { "total": %d, - "used": %d + "used": %d, + "resetsAt": "%s" }, "weekCostStats": { "total": %d, - "used": %d + "used": %d, + "resetsAt": "%s" }, "monthCostStats": { "total": %d, - "used": %d + "used": %d, + "resetsAt": "%s" } } """.formatted( - Long.MAX_VALUE, 0, Long.MAX_VALUE, 0, Long.MAX_VALUE, 0, Long.MAX_VALUE, 0, - Long.MAX_VALUE, 0, Long.MAX_VALUE, 0, - Long.MAX_VALUE, 0, Long.MAX_VALUE, 0, Long.MAX_VALUE, 0, Long.MAX_VALUE, 0)); + Long.MAX_VALUE, 0, + Long.MAX_VALUE, 0, dayResetsAt, + Long.MAX_VALUE, 0, weekResetsAt, + Long.MAX_VALUE, 0, monthResetsAt, + Long.MAX_VALUE, 0, + Long.MAX_VALUE, 0, dayResetsAt, + Long.MAX_VALUE, 0, + Long.MAX_VALUE, 0, dayResetsAt, + Long.MAX_VALUE, 0, weekResetsAt, + Long.MAX_VALUE, 0, monthResetsAt)); + } + + private static String resetsAt(CalendarPeriod period, RateLimitSchedule schedule) { + long resetsAtMillis = CalendarWindowCalculator.nextPeriodStart(period, System.currentTimeMillis(), schedule); + return DateTimeFormatter.ISO_OFFSET_DATE_TIME.format( + Instant.ofEpochMilli(resetsAtMillis).atZone(ZoneId.of(schedule.getTimezone()))); } @Test diff --git a/server/src/test/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculatorTest.java b/server/src/test/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculatorTest.java new file mode 100644 index 000000000..794e1c162 --- /dev/null +++ b/server/src/test/java/com/epam/aidial/core/server/limiter/CalendarWindowCalculatorTest.java @@ -0,0 +1,148 @@ +package com.epam.aidial.core.server.limiter; + +import com.epam.aidial.core.config.RateLimitSchedule; +import com.epam.aidial.core.config.WeekDay; +import org.junit.jupiter.api.Test; + +import java.time.ZoneId; +import java.time.ZonedDateTime; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class CalendarWindowCalculatorTest { + + private static long instant(ZonedDateTime dateTime) { + return dateTime.toInstant().toEpochMilli(); + } + + private static RateLimitSchedule schedule(String timezone, WeekDay weekStartDay, String resetTime) { + RateLimitSchedule schedule = new RateLimitSchedule(); + schedule.setTimezone(timezone); + schedule.setWeekStartDay(weekStartDay); + schedule.setResetTime(resetTime); + return schedule; + } + + @Test + void defaultSchedule_isUtcMidnightMonday() { + RateLimitSchedule schedule = new RateLimitSchedule(); + assertEquals("UTC", schedule.getTimezone()); + assertEquals(WeekDay.Mon, schedule.getWeekStartDay()); + assertEquals("00:00", schedule.getResetTime()); + } + + @Test + void dayPeriod_beforeAndAfterResetTime() { + RateLimitSchedule schedule = schedule("UTC", WeekDay.Mon, "09:00"); + ZoneId utc = ZoneId.of("UTC"); + + // 08:00, before today's 09:00 reset -> still in yesterday's period + long beforeReset = instant(ZonedDateTime.of(2026, 9, 10, 8, 0, 0, 0, utc)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 9, 9, 0, 0, 0, utc)), + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.DAY, beforeReset, schedule)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 10, 9, 0, 0, 0, utc)), + CalendarWindowCalculator.nextPeriodStart(CalendarPeriod.DAY, beforeReset, schedule)); + + // 10:00, after today's 09:00 reset -> today's period + long afterReset = instant(ZonedDateTime.of(2026, 9, 10, 10, 0, 0, 0, utc)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 10, 9, 0, 0, 0, utc)), + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.DAY, afterReset, schedule)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 11, 9, 0, 0, 0, utc)), + CalendarWindowCalculator.nextPeriodStart(CalendarPeriod.DAY, afterReset, schedule)); + } + + @Test + void weekPeriod_rollsBackToConfiguredStartDay() { + RateLimitSchedule schedule = schedule("UTC", WeekDay.Mon, "09:00"); + ZoneId utc = ZoneId.of("UTC"); + + // Wednesday 2026-09-09 (before its own 09:00 reset does not matter for week boundary, + // only Monday 09:00 does) at 08:00 -> still within the week starting Monday 09:00 + long wednesdayMorning = instant(ZonedDateTime.of(2026, 9, 9, 8, 0, 0, 0, utc)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 7, 9, 0, 0, 0, utc)), // Monday 2026-09-07 + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.WEEK, wednesdayMorning, schedule)); + + // Monday 2026-09-07 at 08:00, before this Monday's 09:00 reset -> previous week + long mondayBeforeReset = instant(ZonedDateTime.of(2026, 9, 7, 8, 0, 0, 0, utc)); + assertEquals( + instant(ZonedDateTime.of(2026, 8, 31, 9, 0, 0, 0, utc)), // previous Monday + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.WEEK, mondayBeforeReset, schedule)); + + // Monday 2026-09-07 at 10:00, after reset -> this week + long mondayAfterReset = instant(ZonedDateTime.of(2026, 9, 7, 10, 0, 0, 0, utc)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 7, 9, 0, 0, 0, utc)), + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.WEEK, mondayAfterReset, schedule)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 14, 9, 0, 0, 0, utc)), + CalendarWindowCalculator.nextPeriodStart(CalendarPeriod.WEEK, mondayAfterReset, schedule)); + } + + @Test + void weekPeriod_supportsNonMondayStartDay() { + RateLimitSchedule schedule = schedule("UTC", WeekDay.Sun, "00:00"); + ZoneId utc = ZoneId.of("UTC"); + + // Wednesday 2026-09-09 -> week started Sunday 2026-09-06 + long wednesday = instant(ZonedDateTime.of(2026, 9, 9, 12, 0, 0, 0, utc)); + assertEquals( + instant(ZonedDateTime.of(2026, 9, 6, 0, 0, 0, 0, utc)), + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.WEEK, wednesday, schedule)); + } + + @Test + void monthPeriod_alwaysLandsOnTheFirst() { + RateLimitSchedule schedule = schedule("UTC", WeekDay.Mon, "00:00"); + ZoneId utc = ZoneId.of("UTC"); + + // mid-February (28-day month in 2026, non-leap) -> period started Feb 1 + long midFeb = instant(ZonedDateTime.of(2026, 2, 15, 12, 0, 0, 0, utc)); + assertEquals( + instant(ZonedDateTime.of(2026, 2, 1, 0, 0, 0, 0, utc)), + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.MONTH, midFeb, schedule)); + // next period start is March 1st regardless of February's length + assertEquals( + instant(ZonedDateTime.of(2026, 3, 1, 0, 0, 0, 0, utc)), + CalendarWindowCalculator.nextPeriodStart(CalendarPeriod.MONTH, midFeb, schedule)); + + // first of the month, before reset -> previous month + long firstOfMonthEarly = instant(ZonedDateTime.of(2026, 3, 1, 0, 0, 0, 0, utc).minusSeconds(1)); + assertEquals( + instant(ZonedDateTime.of(2026, 2, 1, 0, 0, 0, 0, utc)), + CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.MONTH, firstOfMonthEarly, schedule)); + } + + @Test + void dayPeriod_springForwardTransitionIsTwentyThreeHours() { + // Europe/Warsaw switches to DST on 2026-03-29, clocks jump 02:00 -> 03:00 + RateLimitSchedule schedule = schedule("Europe/Warsaw", WeekDay.Mon, "00:00"); + ZoneId warsaw = ZoneId.of("Europe/Warsaw"); + + long duringTransitionDay = instant(ZonedDateTime.of(2026, 3, 29, 12, 0, 0, 0, warsaw)); + long periodStart = CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.DAY, duringTransitionDay, schedule); + long nextStart = CalendarWindowCalculator.nextPeriodStart(CalendarPeriod.DAY, duringTransitionDay, schedule); + + assertEquals(instant(ZonedDateTime.of(2026, 3, 29, 0, 0, 0, 0, warsaw)), periodStart); + assertEquals(instant(ZonedDateTime.of(2026, 3, 30, 0, 0, 0, 0, warsaw)), nextStart); + // the transition day is only 23 real hours, not 24 + assertEquals(23 * 60 * 60 * 1000L, nextStart - periodStart); + } + + @Test + void dayPeriod_fallBackTransitionIsTwentyFiveHours() { + // Europe/Warsaw switches off DST on 2026-10-25, clocks fall back 03:00 -> 02:00 + RateLimitSchedule schedule = schedule("Europe/Warsaw", WeekDay.Mon, "00:00"); + ZoneId warsaw = ZoneId.of("Europe/Warsaw"); + + long duringTransitionDay = instant(ZonedDateTime.of(2026, 10, 25, 12, 0, 0, 0, warsaw)); + long periodStart = CalendarWindowCalculator.currentPeriodStart(CalendarPeriod.DAY, duringTransitionDay, schedule); + long nextStart = CalendarWindowCalculator.nextPeriodStart(CalendarPeriod.DAY, duringTransitionDay, schedule); + + assertEquals(25 * 60 * 60 * 1000L, nextStart - periodStart); + } +} diff --git a/server/src/test/java/com/epam/aidial/core/server/limiter/CostFixedRateBucketTest.java b/server/src/test/java/com/epam/aidial/core/server/limiter/CostFixedRateBucketTest.java new file mode 100644 index 000000000..632a7ba3d --- /dev/null +++ b/server/src/test/java/com/epam/aidial/core/server/limiter/CostFixedRateBucketTest.java @@ -0,0 +1,62 @@ +package com.epam.aidial.core.server.limiter; + +import com.epam.aidial.core.config.RateLimitSchedule; +import com.epam.aidial.core.config.WeekDay; +import com.epam.aidial.core.server.util.ProxyUtil; +import org.junit.jupiter.api.Test; + +import java.math.BigDecimal; +import java.time.ZoneId; +import java.time.ZonedDateTime; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class CostFixedRateBucketTest { + + private final RateLimitSchedule schedule = utcMidnightSchedule(); + + private static RateLimitSchedule utcMidnightSchedule() { + RateLimitSchedule schedule = new RateLimitSchedule(); + schedule.setTimezone("UTC"); + schedule.setWeekStartDay(WeekDay.Mon); + schedule.setResetTime("00:00"); + return schedule; + } + + private static long instant(ZonedDateTime dateTime) { + return dateTime.toInstant().toEpochMilli(); + } + + @Test + void addAccumulatesWithinTheSamePeriod() { + CostFixedRateBucket bucket = new CostFixedRateBucket(); + long day1 = instant(ZonedDateTime.of(2026, 9, 9, 10, 0, 0, 0, ZoneId.of("UTC"))); + long day1Later = instant(ZonedDateTime.of(2026, 9, 9, 20, 0, 0, 0, ZoneId.of("UTC"))); + + assertEquals(0, new BigDecimal("0.10").compareTo(bucket.add(day1, CalendarPeriod.DAY, schedule, new BigDecimal("0.10")))); + assertEquals(0, new BigDecimal("0.30").compareTo(bucket.add(day1Later, CalendarPeriod.DAY, schedule, new BigDecimal("0.20")))); + } + + @Test + void reconcileResetsCountOnPeriodRollover() { + CostFixedRateBucket bucket = new CostFixedRateBucket(); + long day1 = instant(ZonedDateTime.of(2026, 9, 9, 10, 0, 0, 0, ZoneId.of("UTC"))); + long day2 = instant(ZonedDateTime.of(2026, 9, 10, 10, 0, 0, 0, ZoneId.of("UTC"))); + + bucket.add(day1, CalendarPeriod.DAY, schedule, new BigDecimal("0.50")); + assertEquals(0, BigDecimal.ZERO.compareTo(bucket.reconcile(day2, CalendarPeriod.DAY, schedule))); + } + + @Test + void oldFloatingWindowJsonDeserializesWithoutThrowing() { + String oldJson = "{\"window\":\"DAY\",\"sums\":[\"0.1\",\"0.2\"],\"sum\":\"0.3\",\"start\":100,\"end\":124}"; + + CostFixedRateBucket bucket = ProxyUtil.convertToObject(oldJson, CostFixedRateBucket.class); + + assertTrue(bucket != null); + assertEquals(0, BigDecimal.ZERO.compareTo(bucket.getCount())); + long now = instant(ZonedDateTime.of(2026, 9, 9, 10, 0, 0, 0, ZoneId.of("UTC"))); + assertEquals(0, BigDecimal.ZERO.compareTo(bucket.reconcile(now, CalendarPeriod.DAY, schedule))); + } +} diff --git a/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateBucketTest.java b/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateBucketTest.java index 1ec827984..28ad42112 100644 --- a/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateBucketTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateBucketTest.java @@ -15,7 +15,7 @@ * - These tests verify: * - Constructor and initialization * - Adding costs and updating the window - * - Window sliding behavior for different time windows (minute, day, week, month) + * - Window sliding behavior for the floating windows (minute, hour) * - Retry calculation * - Edge cases and precision handling with BigDecimal */ @@ -42,20 +42,10 @@ void testConstructorAndInitialization() { assertEquals(BigDecimal.ZERO, bucket.getSum()); assertEquals(60, bucket.getSums().length); - bucket = new CostRateBucket(RateWindow.DAY); - assertEquals(RateWindow.DAY, bucket.getWindow()); + bucket = new CostRateBucket(RateWindow.HOUR); + assertEquals(RateWindow.HOUR, bucket.getWindow()); assertEquals(BigDecimal.ZERO, bucket.getSum()); - assertEquals(24, bucket.getSums().length); - - bucket = new CostRateBucket(RateWindow.WEEK); - assertEquals(RateWindow.WEEK, bucket.getWindow()); - assertEquals(BigDecimal.ZERO, bucket.getSum()); - assertEquals(7, bucket.getSums().length); - - bucket = new CostRateBucket(RateWindow.MONTH); - assertEquals(RateWindow.MONTH, bucket.getWindow()); - assertEquals(BigDecimal.ZERO, bucket.getSum()); - assertEquals(30, bucket.getSums().length); + assertEquals(60, bucket.getSums().length); } /** @@ -85,72 +75,6 @@ void testMinuteBucket() { update(121, "0.00"); } - @Test - void testDayBucket() { - bucket = new CostRateBucket(RateWindow.DAY); - - update(0, "0.00"); - add(0, "0.10", "0.10"); - add(0, "0.20", "0.30"); - update(0, "0.30"); - - add(1, "0.30", "0.60"); - add(23, "0.40", "1.00"); - update(23, "1.00"); - - add(24, "0.10", "0.80"); - update(24, "0.80"); - - add(25, "0.05", "0.55"); - update(25, "0.55"); - - update(49, "0.00"); - } - - @Test - void testWeekBucket() { - bucket = new CostRateBucket(RateWindow.WEEK); - - update(0, "0.00"); - add(0, "0.10", "0.10"); - add(0, "0.20", "0.30"); - update(0, "0.30"); - - add(1, "0.30", "0.60"); - add(6, "0.40", "1.00"); - update(6, "1.00"); - - add(7, "0.10", "0.80"); - update(7, "0.80"); - - add(8, "0.05", "0.55"); - update(8, "0.55"); - - update(15, "0.00"); - } - - @Test - void testMonthBucket() { - bucket = new CostRateBucket(RateWindow.MONTH); - - update(0, "0.00"); - add(0, "0.10", "0.10"); - add(0, "0.20", "0.30"); - update(0, "0.30"); - - add(1, "0.30", "0.60"); - add(29, "0.40", "1.00"); - update(29, "1.00"); - - add(30, "0.10", "0.80"); - update(30, "0.80"); - - add(31, "0.05", "0.55"); - update(31, "0.55"); - - update(61, "0.00"); - } - /** * Tests the retry calculation for the minute window. * Verifies that the retryAfter method correctly calculates how long to wait @@ -187,42 +111,6 @@ void testRetryAfterMinute() { assertTrue(retryTime3 < retryTime2, "Retry time should decrease after window slides"); } - /** - * Tests the retry calculation for the day window. - * Verifies that the retryAfter method correctly calculates how long to wait - * before making a retry request when the cost limit is exceeded with a day window. - * Also verifies that the retry time decreases as the window slides. - */ - @Test - void testRetryAfterDay() { - bucket = new CostRateBucket(RateWindow.DAY); - - update(0, "0.00"); - assertEquals(0, bucket.retryAfter(new BigDecimal("0.30"))); - add(0, "0.10", "0.10"); - - update(5, "0.10"); - assertEquals(0, bucket.retryAfter(new BigDecimal("0.30"))); - add(5, "0.20", "0.30"); - - update(10, "0.30"); - // When sum equals limit, retryAfter will return a non-zero value - // because of the >= comparison in the method - long retryTime1 = bucket.retryAfter(new BigDecimal("0.30")); - assertTrue(retryTime1 > 0, "Retry time should be greater than 0 when sum equals limit"); - add(10, "0.30", "0.60"); - - update(20, "0.60"); - long retryTime2 = bucket.retryAfter(new BigDecimal("0.30")); - assertTrue(retryTime2 > 0, "Retry time should be greater than 0 when sum exceeds limit"); - add(23, "0.10", "0.70"); - - update(24, "0.60"); - long retryTime3 = bucket.retryAfter(new BigDecimal("0.30")); - assertTrue(retryTime3 > 0, "Retry time should be greater than 0 when sum exceeds limit"); - assertTrue(retryTime3 < retryTime2, "Retry time should decrease after window slides"); - } - /** * Tests edge cases for the CostRateBucket. * Verifies behavior with: diff --git a/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateLimitTest.java b/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateLimitTest.java index bc7d77f18..bdcf426cb 100644 --- a/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateLimitTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/limiter/CostRateLimitTest.java @@ -56,6 +56,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -118,7 +119,11 @@ public void beforeEach() { LockService lockService = new LockService(redissonClient, null); ResourceService.Settings settings = new ResourceService.Settings(64 * 1048576, 1048576, 60000, 120000, 4096, 300000, 256); ResourceService resourceService = new ResourceService(mock(TimerService.class), redissonClient, blobStorage, lockService, settings, null); - rateLimiter = new RateLimiter(taskExecutor, resourceService); + // increase() only ever needs the schedule, and no test configures a non-default one; + // lenient() since not every test method calls increase() + ConfigStore configStore = mock(ConfigStore.class); + lenient().when(configStore.get()).thenReturn(new Config()); + rateLimiter = new RateLimiter(taskExecutor, resourceService, configStore); } private static Proxy mockProxy(Config config) { @@ -203,7 +208,8 @@ public void testCostLimit_User_LimitFound() { .thenReturn(new BigDecimal("0.05")); // First increase and limit check should succeed - Future increaseLimitFuture = rateLimiter.increase(model, bucketLocation, proxyContext.getTokenUsage(), null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); + Future increaseLimitFuture = rateLimiter.increase( + model, bucketLocation, proxyContext.getTokenUsage(), null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); assertNotNull(increaseLimitFuture); assertNull(increaseLimitFuture.cause()); @@ -217,7 +223,8 @@ public void testCostLimit_User_LimitFound() { .thenReturn(new BigDecimal("0.15")); // Second increase and limit check should fail due to cost limit - increaseLimitFuture = rateLimiter.increase(model, bucketLocation, proxyContext.getTokenUsage(), null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); + increaseLimitFuture = rateLimiter.increase( + model, bucketLocation, proxyContext.getTokenUsage(), null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); assertNotNull(increaseLimitFuture); assertNull(increaseLimitFuture.cause()); @@ -298,7 +305,8 @@ public void testGetLimitStats_WithCostLimits() { .thenReturn(new BigDecimal("0.05")); // Increase limit to record usage - Future increaseLimitFuture = rateLimiter.increase(model, bucketLocation, proxyContext.getTokenUsage(), null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); + Future increaseLimitFuture = rateLimiter.increase( + model, bucketLocation, proxyContext.getTokenUsage(), null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); assertNotNull(increaseLimitFuture); assertNull(increaseLimitFuture.cause()); @@ -405,12 +413,14 @@ public void testPerUserCostLimits() { .thenReturn(new BigDecimal("0.08")); // First user increases limit - Future increaseLimitFuture1 = rateLimiter.increase(model, bucketLocation1, tokenUsage1, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); + Future increaseLimitFuture1 = rateLimiter.increase( + model, bucketLocation1, tokenUsage1, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); assertNotNull(increaseLimitFuture1); assertNull(increaseLimitFuture1.cause()); // Second user increases limit - Future increaseLimitFuture2 = rateLimiter.increase(model, bucketLocation2, tokenUsage2, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); + Future increaseLimitFuture2 = rateLimiter.increase( + model, bucketLocation2, tokenUsage2, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null); assertNotNull(increaseLimitFuture2); assertNull(increaseLimitFuture2.cause()); diff --git a/server/src/test/java/com/epam/aidial/core/server/limiter/FixedRateBucketTest.java b/server/src/test/java/com/epam/aidial/core/server/limiter/FixedRateBucketTest.java new file mode 100644 index 000000000..ba9dc82c7 --- /dev/null +++ b/server/src/test/java/com/epam/aidial/core/server/limiter/FixedRateBucketTest.java @@ -0,0 +1,88 @@ +package com.epam.aidial.core.server.limiter; + +import com.epam.aidial.core.config.RateLimitSchedule; +import com.epam.aidial.core.config.WeekDay; +import com.epam.aidial.core.server.util.ProxyUtil; +import org.junit.jupiter.api.Test; + +import java.time.ZoneId; +import java.time.ZonedDateTime; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class FixedRateBucketTest { + + private final RateLimitSchedule schedule = utcMidnightSchedule(); + + private static RateLimitSchedule utcMidnightSchedule() { + RateLimitSchedule schedule = new RateLimitSchedule(); + schedule.setTimezone("UTC"); + schedule.setWeekStartDay(WeekDay.Mon); + schedule.setResetTime("00:00"); + return schedule; + } + + private static long instant(ZonedDateTime dateTime) { + return dateTime.toInstant().toEpochMilli(); + } + + @Test + void addAccumulatesWithinTheSamePeriod() { + FixedRateBucket bucket = new FixedRateBucket(); + long day1 = instant(ZonedDateTime.of(2026, 9, 9, 10, 0, 0, 0, ZoneId.of("UTC"))); + long day1Later = instant(ZonedDateTime.of(2026, 9, 9, 20, 0, 0, 0, ZoneId.of("UTC"))); + + assertEquals(10, bucket.add(day1, CalendarPeriod.DAY, schedule, 10)); + assertEquals(30, bucket.add(day1Later, CalendarPeriod.DAY, schedule, 20)); + } + + @Test + void reconcileResetsCountOnPeriodRollover() { + FixedRateBucket bucket = new FixedRateBucket(); + long day1 = instant(ZonedDateTime.of(2026, 9, 9, 10, 0, 0, 0, ZoneId.of("UTC"))); + long day2 = instant(ZonedDateTime.of(2026, 9, 10, 10, 0, 0, 0, ZoneId.of("UTC"))); + + bucket.add(day1, CalendarPeriod.DAY, schedule, 50); + assertEquals(50, bucket.reconcile(day1, CalendarPeriod.DAY, schedule)); + + // a new calendar day - the counter resets to zero before the new usage is recorded + assertEquals(0, bucket.reconcile(day2, CalendarPeriod.DAY, schedule)); + assertEquals(5, bucket.add(day2, CalendarPeriod.DAY, schedule, 5)); + } + + @Test + void retryAfterSecondsMatchesTimeUntilNextPeriod() { + FixedRateBucket bucket = new FixedRateBucket(); + long now = instant(ZonedDateTime.of(2026, 9, 9, 22, 0, 0, 0, ZoneId.of("UTC"))); + bucket.reconcile(now, CalendarPeriod.DAY, schedule); + + // 2 hours until UTC midnight + assertEquals(2 * 60 * 60, bucket.retryAfterSeconds(now, CalendarPeriod.DAY, schedule)); + } + + @Test + void resetsAtMillisIsTheNextPeriodStart() { + FixedRateBucket bucket = new FixedRateBucket(); + long now = instant(ZonedDateTime.of(2026, 9, 9, 22, 0, 0, 0, ZoneId.of("UTC"))); + + assertEquals( + instant(ZonedDateTime.of(2026, 9, 10, 0, 0, 0, 0, ZoneId.of("UTC"))), + bucket.resetsAtMillis(now, CalendarPeriod.DAY, schedule)); + } + + @Test + void oldFloatingWindowJsonDeserializesWithoutThrowing() { + // shape of a pre-rollout floating-window RateBucket record + String oldJson = "{\"window\":\"DAY\",\"sums\":[1,2,3],\"sum\":6,\"start\":100,\"end\":124}"; + + FixedRateBucket bucket = ProxyUtil.convertToObject(oldJson, FixedRateBucket.class); + + assertTrue(bucket != null); + assertEquals(0, bucket.getCount()); + // the sentinel default never matches a real computed period start, so the very first + // reconcile after rollout always resets - the implicit one-time reset the rollout relies on + long now = instant(ZonedDateTime.of(2026, 9, 9, 10, 0, 0, 0, ZoneId.of("UTC"))); + assertEquals(0, bucket.reconcile(now, CalendarPeriod.DAY, schedule)); + } +} diff --git a/server/src/test/java/com/epam/aidial/core/server/limiter/RateBucketTest.java b/server/src/test/java/com/epam/aidial/core/server/limiter/RateBucketTest.java index 676805593..d26b03f6f 100644 --- a/server/src/test/java/com/epam/aidial/core/server/limiter/RateBucketTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/limiter/RateBucketTest.java @@ -32,28 +32,6 @@ void testMinuteBucket() { update(121, 0); } - @Test - void testDayBucket() { - bucket = new RateBucket(RateWindow.DAY); - - update(0, 0); - add(0, 10, 10); - add(0, 20, 30); - update(0, 30); - - add(1, 30, 60); - add(23, 40, 100); - update(23, 100); - - add(24, 10, 80); - update(24, 80); - - add(25, 5, 55); - update(25, 55); - - update(49, 0); - } - @Test public void testRetryAfterMinute() { bucket = new RateBucket(RateWindow.MINUTE); @@ -77,75 +55,6 @@ public void testRetryAfterMinute() { assertEquals(15, bucket.retryAfter(30)); } - @Test - void testWeekBucket() { - bucket = new RateBucket(RateWindow.WEEK); - - update(0, 0); - add(0, 10, 10); - add(0, 20, 30); - update(0, 30); - - add(1, 30, 60); - add(6, 40, 100); - update(6, 100); - - add(7, 10, 80); - update(7, 80); - - add(8, 5, 55); - update(8, 55); - - update(15, 0); - } - - @Test - void testMonthBucket() { - bucket = new RateBucket(RateWindow.MONTH); - - update(0, 0); - add(0, 10, 10); - add(0, 20, 30); - update(0, 30); - - add(1, 30, 60); - add(29, 40, 100); - update(29, 100); - - add(30, 10, 80); - update(30, 80); - - add(31, 5, 55); - update(31, 55); - - update(61, 0); - } - - @Test - public void testRetryAfterDay() { - bucket = new RateBucket(RateWindow.DAY); - - update(0, 0); - assertEquals(0, bucket.retryAfter(30)); - add(0, 10, 10); - - update(5, 10); - assertEquals(0, bucket.retryAfter(30)); - add(5, 20, 30); - - update(10, 30); - // need to wait 14 hours - assertEquals(14 * 60 * 60, bucket.retryAfter(30)); - add(10, 30, 60); - - update(20, 60); - add(23, 10, 70); - - update(24, 60); - // need to wait 10 hours - assertEquals(10 * 60 * 60, bucket.retryAfter(30)); - } - private void add(long interval, long count, long expected) { RateWindow window = bucket.getWindow(); long whole = interval * window.interval(); diff --git a/server/src/test/java/com/epam/aidial/core/server/limiter/RateLimiterTest.java b/server/src/test/java/com/epam/aidial/core/server/limiter/RateLimiterTest.java index 484bcf69e..69395fdda 100644 --- a/server/src/test/java/com/epam/aidial/core/server/limiter/RateLimiterTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/limiter/RateLimiterTest.java @@ -66,6 +66,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -154,7 +155,11 @@ public void beforeEach() throws IOException { ResourceService.Settings settings = new ResourceService.Settings(64 * 1048576, 1048576, 60000, 120000, 4096, 300000, 256); resourceService = new ResourceService(mock(TimerService.class), redissonClient, blobStorage, lockService, settings, null); - rateLimiter = new RateLimiter(taskExecutor, resourceService); + // increase() only ever needs the schedule, and no test configures a non-default one; + // lenient() since not every test method calls increase() + ConfigStore configStore = mock(ConfigStore.class); + lenient().when(configStore.get()).thenReturn(new Config()); + rateLimiter = new RateLimiter(taskExecutor, resourceService, configStore); } @AfterEach @@ -556,7 +561,8 @@ public void testGetUserLimitStats_MultipleDeploymentsAndEmptyUsage() { TokenUsage tokenUsage = new TokenUsage(); tokenUsage.setTotalTokens(90); - assertNull(rateLimiter.increase(usedModel, BucketBuilder.buildInitiatorBucket(proxyContext), tokenUsage, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null).cause()); + assertNull(rateLimiter.increase( + usedModel, BucketBuilder.buildInitiatorBucket(proxyContext), tokenUsage, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null).cause()); UserLimitStats stats = rateLimiter.getUserStats(proxyContext, List.of(usedModel, unusedModel), false).result(); @@ -616,7 +622,8 @@ public void testGetUserStats_CostIsAttributedPerDeployment() { tokenUsage.setPromptTokens(1000); tokenUsage.setCompletionTokens(2000); tokenUsage.setTotalTokens(3000); - assertNull(rateLimiter.increase(model, BucketBuilder.buildInitiatorBucket(proxyContext), tokenUsage, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null).cause()); + assertNull(rateLimiter.increase( + model, BucketBuilder.buildInitiatorBucket(proxyContext), tokenUsage, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null).cause()); UserLimitStats stats = rateLimiter.getUserStats(proxyContext, List.of(model), false).result(); @@ -679,7 +686,8 @@ public void testGetUserStats_DeploymentNamedCostsDoesNotShadowTheGlobalDocument( TokenUsage tokenUsage = new TokenUsage(); tokenUsage.setTotalTokens(11); - assertNull(rateLimiter.increase(model, BucketBuilder.buildInitiatorBucket(proxyContext), tokenUsage, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null).cause()); + assertNull(rateLimiter.increase( + model, BucketBuilder.buildInitiatorBucket(proxyContext), tokenUsage, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null).cause()); UserLimitStats stats = rateLimiter.getUserStats(proxyContext, List.of(model), false).result(); @@ -852,7 +860,8 @@ public void testGetUserStats_SkipsRecordsOlderThanWidestWindow() { assertEquals(HttpStatus.OK, rateLimiter.limit(proxyContext, model).result().status()); assertNull(rateLimiter.increase(model, bucket, tokenUsage, null, null, InterfaceType.OPENAI_CHAT_COMPLETIONS, null).cause()); - // age the token record's listing entry past RateWindow.MONTH, leaving its counter untouched + // age the token record's listing entry past RateLimiter's widest-window pre-filter (32 days), + // leaving its counter untouched ResourceDescriptor tokens = ResourceDescriptorFactory .fromEncoded(ResourceTypes.LIMIT, bucket, bucket, "aged-model/tokens"); long agedOut = System.currentTimeMillis() - Duration.ofDays(60).toMillis(); @@ -867,7 +876,8 @@ public void testGetUserStats_SkipsRecordsOlderThanWidestWindow() { } return page; }).when(listingWithAgedRecord).getFolderMetadata(any(), any(), anyInt(), anyBoolean()); - RateLimiter limiter = new RateLimiter(taskExecutor, listingWithAgedRecord); + // this instance never calls increase(), so its ConfigStore is never actually read + RateLimiter limiter = new RateLimiter(taskExecutor, listingWithAgedRecord, mock(ConfigStore.class)); Future future = limiter.getUserStats(proxyContext, List.of(model), false);