Skip to content

Commit

Permalink
fix(HMS-1783): allow 2 reservations in one second
Browse files Browse the repository at this point in the history
  • Loading branch information
lzap committed Jul 12, 2023
1 parent 10c8213 commit 67663b0
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 6 deletions.
6 changes: 3 additions & 3 deletions api/openapi.gen.json
Original file line number Diff line number Diff line change
Expand Up @@ -1280,7 +1280,7 @@
},
"/reservations/aws": {
"post": {
"description": "A reservation is a way to activate a job, keeps all data needed for a job to start. An AWS reservation is a reservation created for an AWS job. Image Builder UUID image is required, the service will also launch any AMI image prefixed with \"ami-\". Optionally, AWS EC2 launch template ID can be provided. All flags set through this endpoint override template values. Public key must exist prior calling this endpoint and ID must be provided, even when AWS EC2 launch template provides ssh-keys. Public key will be always be overwritten.\n",
"description": "A reservation is a way to activate a job, keeps all data needed for a job to start. An AWS reservation is a reservation created for an AWS job. Image Builder UUID image is required, the service will also launch any AMI image prefixed with \"ami-\". Optionally, AWS EC2 launch template ID can be provided. All flags set through this endpoint override template values. Public key must exist prior calling this endpoint and ID must be provided, even when AWS EC2 launch template provides ssh-keys. Public key will be always be overwritten. A single account can create maximum of 2 reservations per second.\n",
"operationId": "createAwsReservation",
"requestBody": {
"content": {
Expand Down Expand Up @@ -1367,7 +1367,7 @@
},
"/reservations/azure": {
"post": {
"description": "A reservation is a way to activate a job, keeps all data needed for a job to start. An Azure reservation is a reservation created for an Azure job. Image Builder UUID image is required and needs to be stored under same account as provided by SourceID.\n",
"description": "A reservation is a way to activate a job, keeps all data needed for a job to start. An Azure reservation is a reservation created for an Azure job. Image Builder UUID image is required and needs to be stored under same account as provided by SourceID. A single account can create maximum of 2 reservations per second.\n",
"operationId": "createAzureReservation",
"requestBody": {
"content": {
Expand Down Expand Up @@ -1454,7 +1454,7 @@
},
"/reservations/gcp": {
"post": {
"description": "A reservation is a way to activate a job, keeps all data needed for a job to start. A GCP reservation is a reservation created for a GCP job. Image Builder UUID image is required and needs to be shared with the service account. Furthermore, by specifying the name pattern for example as \"instance\", instances names will be created in the format: \"instance-#####\".\n",
"description": "A reservation is a way to activate a job, keeps all data needed for a job to start. A GCP reservation is a reservation created for a GCP job. Image Builder UUID image is required and needs to be shared with the service account. Furthermore, by specifying the name pattern for example as \"instance\", instances names will be created in the format: \"instance-#####\". A single account can create maximum of 2 reservations per second.\n",
"operationId": "createGCPReservation",
"requestBody": {
"content": {
Expand Down
6 changes: 3 additions & 3 deletions api/openapi.gen.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -927,7 +927,7 @@ paths:
tags:
- Reservation
description: |
A reservation is a way to activate a job, keeps all data needed for a job to start. An AWS reservation is a reservation created for an AWS job. Image Builder UUID image is required, the service will also launch any AMI image prefixed with "ami-". Optionally, AWS EC2 launch template ID can be provided. All flags set through this endpoint override template values. Public key must exist prior calling this endpoint and ID must be provided, even when AWS EC2 launch template provides ssh-keys. Public key will be always be overwritten.
A reservation is a way to activate a job, keeps all data needed for a job to start. An AWS reservation is a reservation created for an AWS job. Image Builder UUID image is required, the service will also launch any AMI image prefixed with "ami-". Optionally, AWS EC2 launch template ID can be provided. All flags set through this endpoint override template values. Public key must exist prior calling this endpoint and ID must be provided, even when AWS EC2 launch template provides ssh-keys. Public key will be always be overwritten. A single account can create maximum of 2 reservations per second.
operationId: createAwsReservation
requestBody:
description: aws request body
Expand Down Expand Up @@ -983,7 +983,7 @@ paths:
tags:
- Reservation
description: |
A reservation is a way to activate a job, keeps all data needed for a job to start. An Azure reservation is a reservation created for an Azure job. Image Builder UUID image is required and needs to be stored under same account as provided by SourceID.
A reservation is a way to activate a job, keeps all data needed for a job to start. An Azure reservation is a reservation created for an Azure job. Image Builder UUID image is required and needs to be stored under same account as provided by SourceID. A single account can create maximum of 2 reservations per second.
operationId: createAzureReservation
requestBody:
description: azure request body
Expand Down Expand Up @@ -1039,7 +1039,7 @@ paths:
tags:
- Reservation
description: |
A reservation is a way to activate a job, keeps all data needed for a job to start. A GCP reservation is a reservation created for a GCP job. Image Builder UUID image is required and needs to be shared with the service account. Furthermore, by specifying the name pattern for example as "instance", instances names will be created in the format: "instance-#####".
A reservation is a way to activate a job, keeps all data needed for a job to start. A GCP reservation is a reservation created for a GCP job. Image Builder UUID image is required and needs to be shared with the service account. Furthermore, by specifying the name pattern for example as "instance", instances names will be created in the format: "instance-#####". A single account can create maximum of 2 reservations per second.
operationId: createGCPReservation
requestBody:
description: gcp request body
Expand Down
3 changes: 3 additions & 0 deletions cmd/spec/path.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,7 @@ paths:
endpoint override template values.
Public key must exist prior calling this endpoint and ID must be provided, even when
AWS EC2 launch template provides ssh-keys. Public key will be always be overwritten.
A single account can create maximum of 2 reservations per second.
requestBody:
content:
application/json:
Expand Down Expand Up @@ -446,6 +447,7 @@ paths:
A reservation is a way to activate a job, keeps all data needed for a job to start.
An Azure reservation is a reservation created for an Azure job. Image Builder UUID image
is required and needs to be stored under same account as provided by SourceID.
A single account can create maximum of 2 reservations per second.
requestBody:
content:
application/json:
Expand Down Expand Up @@ -476,6 +478,7 @@ paths:
is required and needs to be shared with the service account.
Furthermore, by specifying the name pattern for example as "instance",
instances names will be created in the format: "instance-#####".
A single account can create maximum of 2 reservations per second.
requestBody:
content:
application/json:
Expand Down
45 changes: 45 additions & 0 deletions internal/migrations/sql/020_reservations_rate_limit.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
--
-- Function (and trigger) for throttling reservations. AWS does not allow more launches per account that 2 per second
-- (with the initial bucket set to 5). This prevents from launching more than 2 instances per second for an individual
-- account. It is applied on all hyperscalers despite other clouds might have different limits.
--
-- Enforcement on the database level seems to be the most efficient way of throttling down clients, in-app solution
-- would need to have a shared place (e.g. Redis) for storing state. This solution is also transactional safe.
-- Error generated by this function is returned in the error payload to the user or UI.
--
-- A side effect of this solution is that there can be gaps in the primary key sequence which is not a problem at all.
--

-- Actual reservation limit per account and provider type (returns constant number)
CREATE OR REPLACE FUNCTION reservations_rate_limit() RETURNS INTEGER AS
$reservations_rate_limit$
BEGIN
RETURN 5;
END;
$reservations_rate_limit$ LANGUAGE plpgsql;

-- Rate limiting function (throws exception when exceeded)
CREATE FUNCTION reservations_rate() RETURNS TRIGGER AS
$reservations_rate$
DECLARE
maximum INTEGER := reservations_rate_limit();
last_rec RECORD;
BEGIN
FOR last_rec IN SELECT COUNT(*) FROM reservations WHERE account_id = NEW.account_id AND provider = NEW.provider AND success IS NULL AND created_at >= now() - INTERVAL '1 second'
LOOP
IF last_rec.count >= maximum THEN
RAISE EXCEPTION 'too many pending reservations (%) for this provider (maximum % per second)', last_rec.count, maximum;
END IF;
END LOOP;

RETURN NEW;
END;
$reservations_rate$ LANGUAGE plpgsql;

-- Rate limiting trigger
CREATE TRIGGER reservations_rate_trigger
BEFORE INSERT ON reservations
FOR EACH ROW
EXECUTE FUNCTION reservations_rate();

CREATE INDEX reservations_created_at_idx ON reservations(created_at);

0 comments on commit 67663b0

Please sign in to comment.