mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-08 05:52:55 +00:00
2f0d245c2a
## Problem Proxy might overload the control plane. ## Summary of changes Implement rate limiter for proxy<->control plane connection. Resolves https://github.com/neondatabase/neon/issues/5707 Used implementation ideas from https://github.com/conradludgate/squeeze/
85 lines
2.7 KiB
Python
85 lines
2.7 KiB
Python
import asyncio
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
import pytest
|
|
from fixtures.neon_fixtures import (
|
|
PSQL,
|
|
NeonProxy,
|
|
)
|
|
from fixtures.port_distributor import PortDistributor
|
|
from pytest_httpserver import HTTPServer
|
|
from werkzeug.wrappers.response import Response
|
|
|
|
|
|
def waiting_handler(status_code: int) -> Response:
|
|
# wait more than timeout to make sure that both (two) connections are open.
|
|
# It would be better to use a barrier here, but I don't know how to do that together with pytest-httpserver.
|
|
time.sleep(2)
|
|
return Response(status=status_code)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def proxy_with_rate_limit(
|
|
port_distributor: PortDistributor,
|
|
neon_binpath: Path,
|
|
httpserver_listen_address,
|
|
test_output_dir: Path,
|
|
) -> Iterator[NeonProxy]:
|
|
"""Neon proxy that routes directly to vanilla postgres."""
|
|
|
|
proxy_port = port_distributor.get_port()
|
|
mgmt_port = port_distributor.get_port()
|
|
http_port = port_distributor.get_port()
|
|
external_http_port = port_distributor.get_port()
|
|
(host, port) = httpserver_listen_address
|
|
endpoint = f"http://{host}:{port}/billing/api/v1/usage_events"
|
|
|
|
with NeonProxy(
|
|
neon_binpath=neon_binpath,
|
|
test_output_dir=test_output_dir,
|
|
proxy_port=proxy_port,
|
|
http_port=http_port,
|
|
mgmt_port=mgmt_port,
|
|
external_http_port=external_http_port,
|
|
auth_backend=NeonProxy.Console(endpoint, fixed_rate_limit=5),
|
|
) as proxy:
|
|
proxy.start()
|
|
yield proxy
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_rate_limit(
|
|
httpserver: HTTPServer,
|
|
proxy_with_rate_limit: NeonProxy,
|
|
):
|
|
uri = "/billing/api/v1/usage_events/proxy_get_role_secret"
|
|
# mock control plane service
|
|
httpserver.expect_ordered_request(uri, method="GET").respond_with_handler(
|
|
lambda _: Response(status=200)
|
|
)
|
|
httpserver.expect_ordered_request(uri, method="GET").respond_with_handler(
|
|
lambda _: waiting_handler(429)
|
|
)
|
|
httpserver.expect_ordered_request(uri, method="GET").respond_with_handler(
|
|
lambda _: waiting_handler(500)
|
|
)
|
|
|
|
psql = PSQL(host=proxy_with_rate_limit.host, port=proxy_with_rate_limit.proxy_port)
|
|
f = await psql.run("select 42;")
|
|
await proxy_with_rate_limit.find_auth_link(uri, f)
|
|
# Limit should be 2.
|
|
|
|
# Run two queries in parallel.
|
|
f1, f2 = await asyncio.gather(psql.run("select 42;"), psql.run("select 42;"))
|
|
await proxy_with_rate_limit.find_auth_link(uri, f1)
|
|
await proxy_with_rate_limit.find_auth_link(uri, f2)
|
|
|
|
# Now limit should be 0.
|
|
f = await psql.run("select 42;")
|
|
await proxy_with_rate_limit.find_auth_link(uri, f)
|
|
|
|
# There last query shouldn't reach the http-server.
|
|
assert httpserver.assertions == []
|