## Problem Tenant split test revealed another bug with PG backpressure throttling that under some cases PS may never report its progress back to SK (e.g., observed when aborting tenant shard where the old shard needs to re-establish SK connection and re-ingest WALs from a much older LSN). In this case, PG may get stuck forever. ## Summary of changes As a general precaution that PS feedback mechanism may not always be reliable, this PR uses the previously introduced WAL write rate limit mechanism to slow down write rates instead of completely pausing it. The idea is to introduce a new `databricks_effective_max_wal_bytes_per_second`, which is set to `databricks_max_wal_mb_per_second` when no PS back pressure and is set to `10KB` when there is back pressure. This way, PG can still write to SK, though at a very low speed. The PR also fixes the problem that the current WAL rate limiting mechanism is too coarse grained and cannot enforce limits < 1MB. This is because it always resets the rate limiter after 1 second, even if PG could have written more data in the past second. The fix is to introduce a `batch_end_time_us` which records the expected end time of the current batch. For example, if PG writes 10MB of data in a single batch, and max WAL write rate is set as `1MB/s`, then `batch_end_time_us` will be set as 10 seconds later. ## How is this tested? Tweaked the existing test, and also did manual testing on dev. I set `max_replication_flush_lag` as 1GB, and loaded 500GB pgbench tables. It's expected to see PG gets throttled periodically because PS will accumulate 4GB of data before flushing. Results: when PG is throttled: ``` 9500000 of 3300000000 tuples (0%) done (elapsed 10.36 s, remaining 3587.62 s) 9600000 of 3300000000 tuples (0%) done (elapsed 124.07 s, remaining 42523.59 s) 9700000 of 3300000000 tuples (0%) done (elapsed 255.79 s, remaining 86763.97 s) 9800000 of 3300000000 tuples (0%) done (elapsed 315.89 s, remaining 106056.52 s) 9900000 of 3300000000 tuples (0%) done (elapsed 412.75 s, remaining 137170.58 s) ``` when PS just flushed: ``` 18100000 of 3300000000 tuples (0%) done (elapsed 433.80 s, remaining 78655.96 s) 18200000 of 3300000000 tuples (0%) done (elapsed 433.85 s, remaining 78231.71 s) 18300000 of 3300000000 tuples (0%) done (elapsed 433.90 s, remaining 77810.62 s) 18400000 of 3300000000 tuples (0%) done (elapsed 433.96 s, remaining 77395.86 s) 18500000 of 3300000000 tuples (0%) done (elapsed 434.03 s, remaining 76987.27 s) 18600000 of 3300000000 tuples (0%) done (elapsed 434.08 s, remaining 76579.59 s) 18700000 of 3300000000 tuples (0%) done (elapsed 434.13 s, remaining 76177.12 s) 18800000 of 3300000000 tuples (0%) done (elapsed 434.19 s, remaining 75779.45 s) 18900000 of 3300000000 tuples (0%) done (elapsed 434.84 s, remaining 75489.40 s) 19000000 of 3300000000 tuples (0%) done (elapsed 434.89 s, remaining 75097.90 s) 19100000 of 3300000000 tuples (0%) done (elapsed 434.94 s, remaining 74712.56 s) 19200000 of 3300000000 tuples (0%) done (elapsed 498.93 s, remaining 85254.20 s) 19300000 of 3300000000 tuples (0%) done (elapsed 498.97 s, remaining 84817.95 s) 19400000 of 3300000000 tuples (0%) done (elapsed 623.80 s, remaining 105486.76 s) 19500000 of 3300000000 tuples (0%) done (elapsed 745.86 s, remaining 125476.51 s) ``` Co-authored-by: Chen Luo <chen.luo@databricks.com>
Neon test runner
This directory contains integration tests.
Prerequisites:
- Correctly configured Python, see
/docs/sourcetree.md - Neon and Postgres binaries
- See the root README.md for build directions
To run tests you need to add
--features testingto Rust code build commands. For convenience, repository cargo config containsbuild_testingalias, that serves as a subcommand, adding the required feature flags. Usage example:cargo build_testing --releaseis equivalent tocargo build --features testing --release - Tests can be run from the git tree; or see the environment variables below to run from other directories.
- See the root README.md for build directions
To run tests you need to add
- The neon git repo, including the postgres submodule
(for some tests, e.g.
pg_regress)
Test Organization
Regression tests are in the 'regress' directory. They can be run in parallel to minimize total runtime. Most regression test sets up their environment with its own pageservers and safekeepers.
'pg_clients' contains tests for connecting with various client libraries. Each client test uses a Dockerfile that pulls an image that contains the client, and connects to PostgreSQL with it. The client tests can be run against an existing PostgreSQL or Neon installation.
'performance' contains performance regression tests. Each test exercises a particular scenario or workload, and outputs measurements. They should be run serially, to avoid the tests interfering with the performance of each other. Some performance tests set up their own Neon environment, while others can be run against an existing PostgreSQL or Neon environment.
Running the tests
There is a wrapper script to invoke pytest: ./scripts/pytest.
It accepts all the arguments that are accepted by pytest.
Depending on your installation options pytest might be invoked directly.
Test state (postgres data, pageserver state, and log files) will
be stored under a directory test_output.
You can run all the tests with:
./scripts/pytest
If you want to run all the tests in a particular file:
./scripts/pytest test_pgbench.py
If you want to run all tests that have the string "bench" in their names:
./scripts/pytest -k bench
To run tests in parellel we utilize pytest-xdist plugin. By default everything runs single threaded. Number of workers can be specified with -n argument:
./scripts/pytest -n4
By default performance tests are excluded. To run them explicitly pass performance tests selection to the script:
./scripts/pytest test_runner/performance
Useful environment variables:
NEON_BIN: The directory where neon binaries can be found.
COMPATIBILITY_NEON_BIN: The directory where the previous version of Neon binaries can be found
POSTGRES_DISTRIB_DIR: The directory where postgres distribution can be found.
Since pageserver supports several postgres versions, POSTGRES_DISTRIB_DIR must contain
a subdirectory for each version with naming convention v{PG_VERSION}/.
Inside that dir, a bin/postgres binary should be present.
COMPATIBILITY_POSTGRES_DISTRIB_DIR: The directory where the prevoius version of postgres distribution can be found.
DEFAULT_PG_VERSION: The version of Postgres to use,
This is used to construct full path to the postgres binaries.
Format is 2-digit major version nubmer, i.e. DEFAULT_PG_VERSION=17
TEST_OUTPUT: Set the directory where test state and test output files
should go.
RUST_LOG: logging configuration to pass into Neon CLI
Useful parameters and commands:
--preserve-database-files to preserve pageserver (layer) and safekeer (segment) timeline files on disk
after running a test suite. Such files might be large, so removed by default; but might be useful for debugging or creation of svg images with layer file contents. If NeonEnvBuilder#preserve_database_files set to True for a particular test, the whole repo directory will be attached to Allure report (thus uploaded to S3) as everything.tar.zst for this test.
Let stdout, stderr and INFO log messages go to the terminal instead of capturing them:
./scripts/pytest -s --log-cli-level=INFO ...
(Note many tests capture subprocess outputs separately, so this may not
show much.)
Exit after the first test failure:
./scripts/pytest -x ...
(there are many more pytest options; run pytest -h to see them.)
Running Python tests against real S3 or S3-compatible services
Neon's libs/remote_storage supports multiple implementations of remote storage.
At the time of writing, that is
pub enum RemoteStorageKind {
/// Storage based on local file system.
/// Specify a root folder to place all stored files into.
LocalFs(Utf8PathBuf),
/// AWS S3 based storage, storing all files in the S3 bucket
/// specified by the config
AwsS3(S3Config),
/// Azure Blob based storage, storing all files in the container
/// specified by the config
AzureContainer(AzureConfig),
}
The test suite has a Python enum with equal name but different meaning:
@enum.unique
class RemoteStorageKind(StrEnum):
LOCAL_FS = "local_fs"
MOCK_S3 = "mock_s3"
REAL_S3 = "real_s3"
LOCAL_FS=>LocalFsMOCK_S3: startsmoto's S3 implementation, then configures Pageserver withAwsS3REAL_S3=> configureAwsS3as detailed below
When a test in the test suite needs an AwsS3, it is supposed to call remote_storage.s3_storage().
That function checks env var ENABLE_REAL_S3_REMOTE_STORAGE:
- If it is not set, use
MOCK_S3 - If it is set, use
REAL_S3.
For REAL_S3, the test suite creates the dict/toml representation of the RemoteStorageKind::AwsS3 based on env vars:
pub struct S3Config {
// test suite env var: REMOTE_STORAGE_S3_BUCKET
pub bucket_name: String,
// test suite env var: REMOTE_STORAGE_S3_REGION
pub bucket_region: String,
// test suite determines this
pub prefix_in_bucket: Option<String>,
// no env var exists; test suite sets it for MOCK_S3, because that's how moto works
pub endpoint: Option<String>,
...
}
Credentials are not part of the config, but discovered by the AWS SDK.
See the libs/remote_storage Rust code.
We're documenting two mechanism here:
The test suite supports two mechanisms (remote_storage.py):
Credential mechanism 1: env vars AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY.
Populate the env vars with AWS access keys that you created in IAM.
Our CI uses this mechanism.
However, it is not recommended for interactive use by developers (learn more).
Instead, use profiles (next section).
Credential mechanism 2: env var AWS_PROFILE.
This uses the AWS SDK's (and CLI's) profile mechanism.
Learn more about it in the official docs.
After configuring a profile (e.g. via the aws CLI), set the env var to its name.
In conclusion, the full command line is:
# with long-term AWS access keys
ENABLE_REAL_S3_REMOTE_STORAGE=true \
REMOTE_STORAGE_S3_BUCKET=mybucket \
REMOTE_STORAGE_S3_REGION=eu-central-1 \
AWS_ACCESS_KEY_ID=... \
AWS_SECRET_ACCESS_KEY=... \
./scripts/pytest
# with AWS PROFILE
ENABLE_REAL_S3_REMOTE_STORAGE=true \
REMOTE_STORAGE_S3_BUCKET=mybucket \
REMOTE_STORAGE_S3_REGION=eu-central-1 \
AWS_PROFILE=... \
./scripts/pytest
If you're using SSO, make sure to aws sso login --profile $AWS_PROFILE first.
Minio
If you want to run test without the cloud setup, we recommend minio.
# Start in Terminal 1
mkdir /tmp/minio_data
minio server /tmp/minio_data --console-address 127.0.0.1:9001 --address 127.0.0.1:9000
In another terminal, create an aws CLI profile for it:
# append to ~/.aws/config
[profile local-minio]
services = local-minio-services
[services local-minio-services]
s3 =
endpoint_url=http://127.0.0.1:9000/
Now configure the credentials (this is going to write ~/.aws/credentials for you).
It's an interactive prompt.
# Terminal 2
$ aws --profile local-minio configure
AWS Access Key ID [None]: minioadmin
AWS Secret Access Key [None]: minioadmin
Default region name [None]:
Default output format [None]:
Now create a bucket testbucket using the CLI.
# (don't forget to have AWS_PROFILE env var set; or use --profile)
aws --profile local-minio s3 mb s3://mybucket
(If it doesn't work, make sure you update your AWS CLI to a recent version. The service-specific endpoint feature that we're using is quite new.)
# with AWS PROFILE
ENABLE_REAL_S3_REMOTE_STORAGE=true \
REMOTE_STORAGE_S3_BUCKET=mybucket \
REMOTE_STORAGE_S3_REGION=doesntmatterforminio \
AWS_PROFILE=local-minio \
./scripts/pytest
NB: you can avoid the --profile by setting the AWS_PROFILE variable.
Just like the AWS SDKs, the aws CLI is sensible to it.
Running Rust tests against real S3 or S3-compatible services
We have some Rust tests that only run against real S3, e.g., here.
They use the same env vars as the Python test suite (see previous section) but interpret them on their own. However, at this time, the interpretation is identical.
So, above instructions apply to the Rust test as well.
Writing a test
Every test needs a Neon Environment, or NeonEnv to operate in. A Neon Environment is like a little cloud-in-a-box, and consists of a Pageserver, 0-N Safekeepers, and compute Postgres nodes. The connections between them can be configured to use JWT authentication tokens, and some other configuration options can be tweaked too.
The easiest way to get access to a Neon Environment is by using the neon_simple_env
fixture. For convenience, there is a branch called main in environments created with
'neon_simple_env', ready to be used in the test.
For more complicated cases, you can build a custom Neon Environment, with the neon_env
fixture:
def test_foobar(neon_env_builder: NeonEnvBuilder):
# Prescribe the environment.
# We want to have 3 safekeeper nodes, and use JWT authentication in the
# connections to the page server
neon_env_builder.num_safekeepers = 3
neon_env_builder.set_pageserver_auth(True)
# Now create the environment. This initializes the repository, and starts
# up the page server and the safekeepers
env = neon_env_builder.init_start()
# Run the test
...
The env includes a default tenant and timeline. Therefore, you do not need to create your own tenant/timeline for testing.
def test_foobar2(neon_env_builder: NeonEnvBuilder):
env = neon_env_builder.init_start() # Start the environment
with env.endpoints.create_start("main") as endpoint:
# Start the compute endpoint
client = env.pageserver.http_client() # Get the pageserver client
tenant_id = env.initial_tenant
timeline_id = env.initial_timeline
client.timeline_detail(tenant_id=tenant_id, timeline_id=timeline_id)
All the test which rely on NeonEnvBuilder, can check the various version combinations of the components. To do this yuo may want to add the parametrize decorator with the function fixtures.utils.allpairs_versions() E.g.
@pytest.mark.parametrize(**fixtures.utils.allpairs_versions())
def test_something(
...
For more information about pytest fixtures, see https://docs.pytest.org/en/stable/fixture.html
At the end of a test, all the nodes in the environment are automatically stopped, so you
don't need to worry about cleaning up. Logs and test data are preserved for the analysis,
in a directory under ../test_output/<testname>
Before submitting a patch
Ensure that you pass all obligatory checks.
Also consider:
- Writing a couple of docstrings to clarify the reasoning behind a new test.
- Adding more type hints to your code to avoid
Any, especially:- For fixture parameters, they are not automatically deduced.
- For function arguments and return values.