diff --git a/.github/ansible/prod.us-west-2.hosts.yaml b/.github/ansible/prod.us-west-2.hosts.yaml index ff5d924a91..9cad79b986 100644 --- a/.github/ansible/prod.us-west-2.hosts.yaml +++ b/.github/ansible/prod.us-west-2.hosts.yaml @@ -29,6 +29,8 @@ storage: ansible_host: i-0c834be1dddba8b3f pageserver-2.us-west-2.aws.neon.tech: ansible_host: i-051642d372c0a4f32 + pageserver-3.us-west-2.aws.neon.tech: + ansible_host: i-00c3844beb9ad1c6b safekeepers: hosts: diff --git a/.github/ansible/staging.us-east-2.hosts.yaml b/.github/ansible/staging.us-east-2.hosts.yaml index b46e729e32..9a1a095282 100644 --- a/.github/ansible/staging.us-east-2.hosts.yaml +++ b/.github/ansible/staging.us-east-2.hosts.yaml @@ -31,6 +31,8 @@ storage: ansible_host: i-01e31cdf7e970586a pageserver-3.us-east-2.aws.neon.build: ansible_host: i-0602a0291365ef7cc + pageserver-99.us-east-2.aws.neon.build: + ansible_host: i-0c39491109bb88824 safekeepers: hosts: @@ -40,3 +42,5 @@ storage: ansible_host: i-0171efc3604a7b907 safekeeper-2.us-east-2.aws.neon.build: ansible_host: i-0de0b03a51676a6ce + safekeeper-99.us-east-2.aws.neon.build: + ansible_host: i-0d61b6a2ea32028d5 diff --git a/.github/workflows/deploy-prod.yml b/.github/workflows/deploy-prod.yml index b6800a8f7a..f4ce7e9afa 100644 --- a/.github/workflows/deploy-prod.yml +++ b/.github/workflows/deploy-prod.yml @@ -40,7 +40,9 @@ concurrency: jobs: deploy-prod-new: runs-on: prod - container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest + container: + image: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest + options: --user root --privileged if: inputs.deployStorage && inputs.disclamerAcknowledged defaults: run: @@ -168,7 +170,9 @@ jobs: deploy: runs-on: prod - container: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest + container: + image: 093970136003.dkr.ecr.eu-central-1.amazonaws.com/ansible:latest + options: --user root --privileged if: inputs.deployStorage && inputs.disclamerAcknowledged defaults: run: diff --git a/Cargo.lock b/Cargo.lock index 6be08d16b1..67e54d3833 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -917,6 +917,7 @@ dependencies = [ "reqwest", "safekeeper_api", "serde", + "serde_json", "serde_with", "storage_broker", "tar", @@ -2421,6 +2422,7 @@ dependencies = [ "crc32c", "criterion", "crossbeam-utils", + "either", "enum-map", "enumset", "fail", @@ -2484,6 +2486,7 @@ dependencies = [ "enum-map", "postgres_ffi", "serde", + "serde_json", "serde_with", "utils", "workspace_hack", @@ -2898,6 +2901,7 @@ dependencies = [ "serde_json", "sha2", "socket2", + "sync_wrapper", "thiserror", "tls-listener", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 9033671f55..4e4667f253 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ comfy-table = "6.1" const_format = "0.2" crc32c = "0.6" crossbeam-utils = "0.8.5" +either = "1.8" enum-map = "2.4.2" enumset = "1.0.12" fail = "0.5.0" @@ -68,7 +69,6 @@ once_cell = "1.13" opentelemetry = "0.18.0" opentelemetry-otlp = { version = "0.11.0", default_features=false, features = ["http-proto", "trace", "http", "reqwest-client"] } opentelemetry-semantic-conventions = "0.10.0" -tracing-opentelemetry = "0.18.0" parking_lot = "0.12" pin-project-lite = "0.2" prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency @@ -92,6 +92,7 @@ socket2 = "0.4.4" strum = "0.24" strum_macros = "0.24" svg_fmt = "0.4.1" +sync_wrapper = "0.1.2" tar = "0.4" thiserror = "1.0" tls-listener = { version = "0.6", features = ["rustls", "hyper-h1"] } @@ -104,6 +105,7 @@ toml = "0.5" toml_edit = { version = "0.17", features = ["easy"] } tonic = {version = "0.8", features = ["tls", "tls-roots"]} tracing = "0.1" +tracing-opentelemetry = "0.18.0" tracing-subscriber = { version = "0.3", features = ["env-filter"] } url = "2.2" uuid = { version = "1.2", features = ["v4", "serde"] } diff --git a/Dockerfile.compute-node b/Dockerfile.compute-node index 5a3110141c..5c58f4baaa 100644 --- a/Dockerfile.compute-node +++ b/Dockerfile.compute-node @@ -106,7 +106,9 @@ RUN wget https://github.com/plv8/plv8/archive/refs/tags/v3.1.4.tar.gz && \ sed -i 's/MemoryContextAlloc(/MemoryContextAllocZero(/' plv8.cc && \ make DOCKER=1 -j $(getconf _NPROCESSORS_ONLN) install && \ rm -rf /plv8-* && \ - echo 'trusted = true' >> /usr/local/pgsql/share/extension/plv8.control + echo 'trusted = true' >> /usr/local/pgsql/share/extension/plv8.control && \ + echo 'trusted = true' >> /usr/local/pgsql/share/extension/plcoffee.control && \ + echo 'trusted = true' >> /usr/local/pgsql/share/extension/plls.control ######################################################################################### # @@ -165,6 +167,21 @@ RUN wget https://github.com/df7cb/postgresql-unit/archive/refs/tags/7.7.tar.gz & find /usr/local/pgsql/share/extension/ -name "unit*.sql" -print0 | xargs -0 sed -i "s|pgsql/||g" && \ echo 'trusted = true' >> /usr/local/pgsql/share/extension/unit.control +######################################################################################### +# +# Layer "vector-pg-build" +# compile pgvector extension +# +######################################################################################### +FROM build-deps AS vector-pg-build +COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/ + +RUN git clone --branch v0.4.0 https://github.com/pgvector/pgvector.git && \ + cd pgvector && \ + make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \ + make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \ + echo 'trusted = true' >> /usr/local/pgsql/share/extension/vector.control + ######################################################################################### # # Layer "neon-pg-ext-build" @@ -178,6 +195,7 @@ COPY --from=plv8-build /usr/local/pgsql/ /usr/local/pgsql/ COPY --from=h3-pg-build /usr/local/pgsql/ /usr/local/pgsql/ COPY --from=h3-pg-build /h3/usr / COPY --from=unit-pg-build /usr/local/pgsql/ /usr/local/pgsql/ +COPY --from=vector-pg-build /usr/local/pgsql/ /usr/local/pgsql/ COPY pgxn/ pgxn/ RUN make -j $(getconf _NPROCESSORS_ONLN) \ @@ -228,7 +246,9 @@ RUN mkdir /var/db && useradd -m -d /var/db/postgres postgres && \ mkdir /var/db/postgres/compute && mkdir /var/db/postgres/specs && \ chown -R postgres:postgres /var/db/postgres && \ chmod 0750 /var/db/postgres/compute && \ - echo '/usr/local/lib' >> /etc/ld.so.conf && /sbin/ldconfig + echo '/usr/local/lib' >> /etc/ld.so.conf && /sbin/ldconfig && \ + # create folder for file cache + mkdir -p -m 777 /neon/cache COPY --from=postgres-cleanup-layer --chown=postgres /usr/local/pgsql /usr/local COPY --from=compute-tools --chown=postgres /home/nonroot/target/release-line-debug-size-lto/compute_ctl /usr/local/bin/compute_ctl diff --git a/control_plane/Cargo.toml b/control_plane/Cargo.toml index 0b2f561d39..309887e1fa 100644 --- a/control_plane/Cargo.toml +++ b/control_plane/Cargo.toml @@ -15,6 +15,7 @@ postgres.workspace = true regex.workspace = true reqwest = { workspace = true, features = ["blocking", "json"] } serde.workspace = true +serde_json.workspace = true serde_with.workspace = true tar.workspace = true thiserror.workspace = true diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs index 9cebe028e4..c49bd39f09 100644 --- a/control_plane/src/pageserver.rs +++ b/control_plane/src/pageserver.rs @@ -419,6 +419,11 @@ impl PageServerNode { .map(|x| x.parse::()) .transpose() .context("Failed to parse 'trace_read_requests' as bool")?, + eviction_policy: settings + .get("eviction_policy") + .map(|x| serde_json::from_str(x)) + .transpose() + .context("Failed to parse 'eviction_policy' json")?, }) .send()? .error_from_body()?; diff --git a/docs/synthetic-size.md b/docs/synthetic-size.md new file mode 100644 index 0000000000..8378efc842 --- /dev/null +++ b/docs/synthetic-size.md @@ -0,0 +1,335 @@ +# Synthetic size + +Neon storage has copy-on-write branching, which makes it difficult to +answer the question "how large is my database"? To give one reasonable +answer, we calculate _synthetic size_ for a project. + +The calculation is called "synthetic", because it is based purely on +the user-visible logical size, which is the size that you would see on +a standalone PostgreSQL installation, and the amount of WAL, which is +also the same as what you'd see on a standalone PostgreSQL, for the +same set of updates. + +The synthetic size does *not* depend on the actual physical size +consumed in the storage, or implementation details of the Neon storage +like garbage collection, compaction and compression. There is a +strong *correlation* between the physical size and the synthetic size, +but the synthetic size is designed to be independent of the +implementation details, so that any improvements we make in the +storage system simply reduce our COGS. And vice versa: any bugs or bad +implementation where we keep more data than we would need to, do not +change the synthetic size or incur any costs to the user. + +The synthetic size is calculated for the whole project. It is not +straighforward to attribute size to individual branches. See "What is +the size of an individual branch?" for discussion on those +difficulties. + +The synthetic size is designed to: + +- Take into account the copy-on-write nature of the storage. For + example, if you create a branch, it doesn't immediately add anything + to the synthetic size. It starts to affect the synthetic size only + as it diverges from the parent branch. + +- Be independent of any implementation details of the storage, like + garbage collection, remote storage, or compression. + +## Terms & assumptions + +- logical size is the size of a database *at a given point in + time*. It's the total size of all tables in all databases, as you + see with "\l+" in psql for example, plus the Postgres SLRUs and some + small amount of metadata. NOTE that currently, Neon does not include + the SLRUs and metadata in the logical size. See comment to `get_current_logical_size_non_incremental()`. + +- a "point in time" is defined as an LSN value. You can convert a + timestamp to an LSN, but the storage internally works with LSNs. + +- PITR horizon can be set per-branch. + +- PITR horizon can be set as a time interval, e.g. 5 days or hours, or + as amount of WAL, in bytes. If it's given as a time interval, it's + converted to an LSN for the calculation. + +- PITR horizon can be set to 0, if you don't want to retain any history. + +## Calculation + +Inputs to the calculation are: +- logical size of the database at different points in time, +- amount of WAL generated, and +- the PITR horizon settings + +The synthetic size is based on an idealistic model of the storage +system, where we pretend that the storage consists of two things: +- snapshots, containing a full snapshot of the database, at a given + point in time, and +- WAL. + +In the simple case that the project contains just one branch (main), +and a fixed PITR horizon, the synthetic size is the sum of: + +- the logical size of the branch *at the beginning of the PITR + horizon*, i.e. at the oldest point that you can still recover to, and +- the size of the WAL covering the PITR horizon. + +The snapshot allows you to recover to the beginning of the PITR +horizon, and the WAL allows you to recover from that point to any +point within the horizon. + +``` + WAL + -----------------------#########> + ^ + snapshot + +Legend: + ##### PITR horizon. This is the region that you can still access + with Point-in-time query and you can still create branches + from. + ----- history that has fallen out of the PITR horizon, and can no + longer be accessed +``` + +NOTE: This is not how the storage system actually works! The actual +implementation is also based on snapshots and WAL, but the snapshots +are taken for individual database pages and ranges of pages rather +than the whole database, and it is much more complicated. This model +is a reasonable approximation, however, to make the synthetic size a +useful proxy for the actual storage consumption. + + +## Example: Data is INSERTed + +For example, let's assume that your database contained 10 GB of data +at the beginning of the PITR horizon, and you have since then inserted +5 GB of additional data into it. The additional insertions of 5 GB of +data consume roughly 5 GB of WAL. In that case, the synthetic size is: + +> 10 GB (snapshot) + 5 GB (WAL) = 15 GB + +If you now set the PITR horizon on the project to 0, so that no +historical data is retained, then the beginning PITR horizon would be +at the end of the branch, so the size of the snapshot would be +calculated at the end of the branch, after the insertions. Then the +synthetic size is: + +> 15 GB (snapshot) + 0 GB (WAL) = 15 GB. + +In this case, the synthetic size is the same, regardless of the PITR horizon, +because all the history consists of inserts. The newly inserted data takes +up the same amount of space, whether it's stored as part of the logical +snapshot, or as WAL. (*) + +(*) This is a rough approximation. In reality, the WAL contains +headers and other overhead, and on the other hand, the logical +snapshot includes empty space on pages, so the size of insertions in +WAL can be smaller or greater than the size of the final table after +the insertions. But in most cases, it's in the same ballpark. + +## Example: Data is DELETEd + +Let's look at another example: + +Let's start again with a database that contains 10 GB of data. Then, +you DELETE 5 GB of the data, and run VACUUM to free up the space, so +that the logical size of the database is now only 5 GB. + +Let's assume that the WAL for the deletions and the vacuum take up +100 MB of space. In that case, the synthetic size of the project is: + +> 10 GB (snapshot) + 100 MB (WAL) = 10.1 GB + +This is much larger than the logical size of the database after the +deletions (5 GB). That's because the system still needs to retain the +deleted data, because it's still accessible to queries and branching +in the PITR window. + +If you now set the PITR horizon to 0 or just wait for time to pass so +that the data falls out of the PITR horizon, making the deleted data +inaccessible, the synthetic size shrinks: + +> 5 GB (snapshot) + 0 GB (WAL) = 5 GB + + +# Branching + +Things get more complicated with branching. Branches in Neon are +copy-on-write, which is also reflected in the synthetic size. + +When you create a branch, it doesn't immediately change the synthetic +size at all. The branch point is within the PITR horizon, and all the +data needed to recover to that point in time needs to be retained +anyway. + +However, if you make modifications on the branch, the system needs to +keep the WAL of those modifications. The WAL is included in the +synthetic size. + +## Example: branch and INSERT + +Let's assume that you again start with a 10 GB database. +On the main branch, you insert 2 GB of data. Then you create +a branch at that point, and insert another 3 GB of data on the +main branch, and 1 GB of data on the child branch + +``` + child +#####> + | + | WAL + main ---------###############> + ^ + snapshot +``` + +In this case, the synthetic size consists of: +- the snapshot at the beginning of the PITR horizon (10 GB) +- the WAL on the main branch (2 GB + 3 GB = 5 GB) +- the WAL on the child branch (1 GB) + +Total: 16 GB + +# Diverging branches + +If there is only a small amount of changes in the database on the +different branches, as in the previous example, the synthetic size +consists of a snapshot before the branch point, containing all the +shared data, and the WAL on both branches. However, if the branches +diverge a lot, it is more efficient to store a separate snapshot of +branches. + +## Example: diverging branches + +You start with a 10 GB database. You insert 5 GB of data on the main +branch. Then you create a branch, and immediately delete all the data +on the child branch and insert 5 GB of new data to it. Then you do the +same on the main branch. Let's assume +that the PITR horizon requires keeping the last 1 GB of WAL on the +both branches. + +``` + snapshot + v WAL + child +---------##############> + | + | + main -------------+---------##############> + ^ WAL + snapshot +``` + +In this case, the synthetic size consists of: +- snapshot at the beginning of the PITR horizon on the main branch (4 GB) +- WAL on the main branch (1 GB) +- snapshot at the beginning of the PITR horizon on the child branch (4 GB) +- last 1 GB of WAL on the child branch (1 GB) + +Total: 10 GB + +The alternative way to store this would be to take only one snapshot +at the beginning of branch point, and keep all the WAL on both +branches. However, the size with that method would be larger, as it +would require one 10 GB snapshot, and 5 GB + 5 GB of WAL. It depends +on the amount of changes (WAL) on both branches, and the logical size +at the branch point, which method would result in a smaller synthetic +size. On each branch point, the system performs the calculation with +both methods, and uses the method that is cheaper, i.e. the one that +results in a smaller synthetic size. + +One way to think about this is that when you create a branch, it +starts out as a thin branch that only stores the WAL since the branch +point. As you modify it, and the amount of WAL grows, at some point +it becomes cheaper to store a completely new snapshot of the branch +and truncate the WAL. + + +# What is the size of an individual branch? + +Synthetic size is calculated for the whole project, and includes all +branches. There is no such thing as the size of a branch, because it +is not straighforward to attribute the parts of size to individual +branches. + +## Example: attributing size to branches + +(copied from https://github.com/neondatabase/neon/pull/2884#discussion_r1029365278) + +Imagine that you create two branches, A and B, at the same point from +main branch, and do a couple of small updates on both branches. Then +six months pass, and during those six months the data on the main +branch churns over completely multiple times. The retention period is, +say 1 month. + +``` + +------> A + / +--------------------*-------------------------------> main + \ + +--------> B +``` + +In that situation, the synthetic tenant size would be calculated based +on a "logical snapshot" at the branch point, that is, the logical size +of the database at that point. Plus the WAL on branches A and B. Let's +say that the snapshot size is 10 GB, and the WAL is 1 MB on both +branches A and B. So the total synthetic storage size is 10002 +MB. (Let's ignore the main branch for now, that would be just added to +the sum) + +How would you break that down per branch? I can think of three +different ways to do it, and all of them have their own problems: + +### Subtraction method + +For each branch, calculate how much smaller the total synthetic size +would be, if that branch didn't exist. In other words, how much would +you save if you dropped the branch. With this method, the size of +branches A and B is 1 MB. + +With this method, the 10 GB shared logical snapshot is not included +for A nor B. So the size of all branches is not equal to the total +synthetic size of the tenant. If you drop branch A, you save 1 MB as +you'd expect, but also the size of B suddenly jumps from 1 MB to 10001 +MB, which might feel surprising. + +### Division method + +Divide the common parts evenly across all branches that need +them. With this method, the size of branches A and B would be 5001 MB. + +With this method, the sum of all branches adds up to the total +synthetic size. But it's surprising in other ways: if you drop branch +A, you might think that you save 5001 MB, but in reality you only save +1 MB, and the size of branch B suddenly grows from 5001 to 10001 MB. + +### Addition method + +For each branch, include all the snapshots and WAL that it depends on, +even if some of them are shared by other branches. With this method, +the size of branches A and B would be 10001 MB. + +The surprise with this method is that the sum of all the branches is +larger than the total synthetic size. And if you drop branch A, the +total synthetic size doesn't fall by 10001 MB as you might think. + +# Alternatives + +A sort of cop-out method would be to show the whole tree of branches +graphically, and for each section of WAL or logical snapshot, display +the size of that section. You can then see which branches depend on +which sections, which sections are shared etc. That would be good to +have in the UI anyway. + +Or perhaps calculate per-branch numbers using the subtraction method, +and in addition to that, one more number for "shared size" that +includes all the data that is needed by more than one branch. + +## Which is the right method? + +The bottom line is that it's not straightforward to attribute the +synthetic size to individual branches. There are things we can do, and +all of those methods are pretty straightforward to implement, but they +all have their own problems. What makes sense depends a lot on what +you want to do with the number, what question you are trying to +answer. diff --git a/libs/pageserver_api/Cargo.toml b/libs/pageserver_api/Cargo.toml index dafb246632..7709da1072 100644 --- a/libs/pageserver_api/Cargo.toml +++ b/libs/pageserver_api/Cargo.toml @@ -14,5 +14,6 @@ byteorder.workspace = true utils.workspace = true postgres_ffi.workspace = true enum-map.workspace = true +serde_json.workspace = true workspace_hack.workspace = true diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs index 9cdcf3a173..3ac7e31ec2 100644 --- a/libs/pageserver_api/src/models.rs +++ b/libs/pageserver_api/src/models.rs @@ -155,6 +155,11 @@ pub struct TenantConfigRequest { pub lagging_wal_timeout: Option, pub max_lsn_wal_lag: Option, pub trace_read_requests: Option, + // We defer the parsing of the eviction_policy field to the request handler. + // Otherwise we'd have to move the types for eviction policy into this package. + // We might do that once the eviction feature has stabilizied. + // For now, this field is not even documented in the openapi_spec.yml. + pub eviction_policy: Option, } impl TenantConfigRequest { @@ -174,6 +179,7 @@ impl TenantConfigRequest { lagging_wal_timeout: None, max_lsn_wal_lag: None, trace_read_requests: None, + eviction_policy: None, } } } @@ -263,11 +269,11 @@ pub struct LayerResidenceEvent { /// #[serde(rename = "timestamp_millis_since_epoch")] #[serde_as(as = "serde_with::TimestampMilliSeconds")] - timestamp: SystemTime, + pub timestamp: SystemTime, /// The new residence status of the layer. - status: LayerResidenceStatus, + pub status: LayerResidenceStatus, /// The reason why we had to record this event. - reason: LayerResidenceEventReason, + pub reason: LayerResidenceEventReason, } /// The reason for recording a given [`ResidenceEvent`]. diff --git a/pageserver/Cargo.toml b/pageserver/Cargo.toml index f3ad2c5de6..d2f0b84863 100644 --- a/pageserver/Cargo.toml +++ b/pageserver/Cargo.toml @@ -23,6 +23,7 @@ const_format.workspace = true consumption_metrics.workspace = true crc32c.workspace = true crossbeam-utils.workspace = true +either.workspace = true fail.workspace = true futures.workspace = true git-version.workspace = true @@ -51,7 +52,7 @@ thiserror.workspace = true tokio = { workspace = true, features = ["process", "sync", "fs", "rt", "io-util", "time"] } tokio-postgres.workspace = true tokio-util.workspace = true -toml_edit.workspace = true +toml_edit = { workspace = true, features = [ "serde" ] } tracing.workspace = true url.workspace = true walkdir.workspace = true diff --git a/pageserver/src/config.rs b/pageserver/src/config.rs index f88895a970..309e5367a4 100644 --- a/pageserver/src/config.rs +++ b/pageserver/src/config.rs @@ -731,6 +731,13 @@ impl PageServerConf { })?); } + if let Some(eviction_policy) = item.get("eviction_policy") { + t_conf.eviction_policy = Some( + toml_edit::de::from_item(eviction_policy.clone()) + .context("parse eviction_policy")?, + ); + } + Ok(t_conf) } diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index 229cf96ee3..6a9232e097 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -797,6 +797,14 @@ async fn update_tenant_config_handler( ); } + if let Some(eviction_policy) = request_data.eviction_policy { + tenant_conf.eviction_policy = Some( + serde_json::from_value(eviction_policy) + .context("parse field `eviction_policy`") + .map_err(ApiError::BadRequest)?, + ); + } + let state = get_state(&request); mgr::set_new_tenant_config(state.conf, tenant_conf, tenant_id) .instrument(info_span!("tenant_config", tenant = ?tenant_id)) diff --git a/pageserver/src/repository.rs b/pageserver/src/repository.rs index 092503b7c5..047fa761c3 100644 --- a/pageserver/src/repository.rs +++ b/pageserver/src/repository.rs @@ -7,11 +7,11 @@ use std::fmt; use std::ops::{AddAssign, Range}; use std::time::Duration; -#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Ord, PartialOrd, Serialize, Deserialize)] /// Key used in the Repository kv-store. /// /// The Repository treats this as an opaque struct, but see the code in pgdatadir_mapping.rs /// for what we actually store in these fields. +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Ord, PartialOrd, Serialize, Deserialize)] pub struct Key { pub field1: u8, pub field2: u32, diff --git a/pageserver/src/task_mgr.rs b/pageserver/src/task_mgr.rs index c4f213e755..2734031a09 100644 --- a/pageserver/src/task_mgr.rs +++ b/pageserver/src/task_mgr.rs @@ -231,6 +231,9 @@ pub enum TaskKind { // Compaction. One per tenant. Compaction, + // Eviction. One per timeline. + Eviction, + // Initial logical size calculation InitialLogicalSizeCalculation, diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index bc943372f8..23210b98d5 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -2757,6 +2757,7 @@ pub mod harness { lagging_wal_timeout: Some(tenant_conf.lagging_wal_timeout), max_lsn_wal_lag: Some(tenant_conf.max_lsn_wal_lag), trace_read_requests: Some(tenant_conf.trace_read_requests), + eviction_policy: Some(tenant_conf.eviction_policy), } } } diff --git a/pageserver/src/tenant/config.rs b/pageserver/src/tenant/config.rs index 087cff2537..fca08dd51a 100644 --- a/pageserver/src/tenant/config.rs +++ b/pageserver/src/tenant/config.rs @@ -91,6 +91,7 @@ pub struct TenantConf { /// to avoid eager reconnects. pub max_lsn_wal_lag: NonZeroU64, pub trace_read_requests: bool, + pub eviction_policy: EvictionPolicy, } /// Same as TenantConf, but this struct preserves the information about @@ -153,6 +154,34 @@ pub struct TenantConfOpt { #[serde(skip_serializing_if = "Option::is_none")] #[serde(default)] pub trace_read_requests: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] + pub eviction_policy: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind")] +pub enum EvictionPolicy { + NoEviction, + LayerAccessThreshold(EvictionPolicyLayerAccessThreshold), +} + +impl EvictionPolicy { + pub fn discriminant_str(&self) -> &'static str { + match self { + EvictionPolicy::NoEviction => "NoEviction", + EvictionPolicy::LayerAccessThreshold(_) => "LayerAccessThreshold", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct EvictionPolicyLayerAccessThreshold { + #[serde(with = "humantime_serde")] + pub period: Duration, + #[serde(with = "humantime_serde")] + pub threshold: Duration, } impl TenantConfOpt { @@ -189,6 +218,7 @@ impl TenantConfOpt { trace_read_requests: self .trace_read_requests .unwrap_or(global_conf.trace_read_requests), + eviction_policy: self.eviction_policy.unwrap_or(global_conf.eviction_policy), } } @@ -261,6 +291,7 @@ impl Default for TenantConf { max_lsn_wal_lag: NonZeroU64::new(DEFAULT_MAX_WALRECEIVER_LSN_WAL_LAG) .expect("cannot parse default max walreceiver Lsn wal lag"), trace_read_requests: false, + eviction_policy: EvictionPolicy::NoEviction, } } } diff --git a/pageserver/src/tenant/layer_map.rs b/pageserver/src/tenant/layer_map.rs index 59a358a355..e446e34f4e 100644 --- a/pageserver/src/tenant/layer_map.rs +++ b/pageserver/src/tenant/layer_map.rs @@ -731,16 +731,30 @@ where Ok(()) } + /// Similar to `Arc::ptr_eq`, but only compares the object pointers, not vtables. + /// + /// Returns `true` if the two `Arc` point to the same layer, false otherwise. #[inline(always)] - fn compare_arced_layers(left: &Arc, right: &Arc) -> bool { - // FIXME: ptr_eq might fail to return true for 'dyn' references because of multiple vtables - // can be created in compilation. Clippy complains about this. In practice it seems to - // work. + pub fn compare_arced_layers(left: &Arc, right: &Arc) -> bool { + // "dyn Trait" objects are "fat pointers" in that they have two components: + // - pointer to the object + // - pointer to the vtable // - // In future rust versions this might become Arc::as_ptr(left) as *const () == - // Arc::as_ptr(right) as *const (), we could change to that before. - #[allow(clippy::vtable_address_comparisons)] - Arc::ptr_eq(left, right) + // rust does not provide a guarantee that these vtables are unique, but however + // `Arc::ptr_eq` as of writing (at least up to 1.67) uses a comparison where both the + // pointer and the vtable need to be equal. + // + // See: https://github.com/rust-lang/rust/issues/103763 + // + // A future version of rust will most likely use this form below, where we cast each + // pointer into a pointer to unit, which drops the inaccessible vtable pointer, making it + // not affect the comparison. + // + // See: https://github.com/rust-lang/rust/pull/106450 + let left = Arc::as_ptr(left) as *const (); + let right = Arc::as_ptr(right) as *const (); + + left == right } } diff --git a/pageserver/src/tenant/storage_layer.rs b/pageserver/src/tenant/storage_layer.rs index e85359af16..9198cfd1df 100644 --- a/pageserver/src/tenant/storage_layer.rs +++ b/pageserver/src/tenant/storage_layer.rs @@ -13,6 +13,7 @@ use crate::task_mgr::TaskKind; use crate::walrecord::NeonWalRecord; use anyhow::Result; use bytes::Bytes; +use either::Either; use enum_map::EnumMap; use enumset::EnumSet; use pageserver_api::models::LayerAccessKind; @@ -92,7 +93,23 @@ pub enum ValueReconstructResult { } #[derive(Debug)] -pub struct LayerAccessStats(Mutex); +pub struct LayerAccessStats(Mutex); + +/// This struct holds two instances of [`LayerAccessStatsInner`]. +/// Accesses are recorded to both instances. +/// The `for_scraping_api`instance can be reset from the management API via [`LayerAccessStatsReset`]. +/// The `for_eviction_policy` is never reset. +#[derive(Debug, Default, Clone)] +struct LayerAccessStatsLocked { + for_scraping_api: LayerAccessStatsInner, + for_eviction_policy: LayerAccessStatsInner, +} + +impl LayerAccessStatsLocked { + fn iter_mut(&mut self) -> impl Iterator { + [&mut self.for_scraping_api, &mut self.for_eviction_policy].into_iter() + } +} #[derive(Debug, Default, Clone)] struct LayerAccessStatsInner { @@ -103,11 +120,11 @@ struct LayerAccessStatsInner { last_residence_changes: HistoryBufferWithDropCounter, } -#[derive(Debug, Clone)] -struct LayerAccessStatFullDetails { - when: SystemTime, - task_kind: TaskKind, - access_kind: LayerAccessKind, +#[derive(Debug, Clone, Copy)] +pub(super) struct LayerAccessStatFullDetails { + pub(super) when: SystemTime, + pub(super) task_kind: TaskKind, + pub(super) access_kind: LayerAccessKind, } #[derive(Clone, Copy, strum_macros::EnumString)] @@ -126,7 +143,7 @@ fn system_time_to_millis_since_epoch(ts: &SystemTime) -> u64 { } impl LayerAccessStatFullDetails { - fn to_api_model(&self) -> pageserver_api::models::LayerAccessStatFullDetails { + fn as_api_model(&self) -> pageserver_api::models::LayerAccessStatFullDetails { let Self { when, task_kind, @@ -142,13 +159,13 @@ impl LayerAccessStatFullDetails { impl LayerAccessStats { pub(crate) fn for_loading_layer(status: LayerResidenceStatus) -> Self { - let new = LayerAccessStats(Mutex::new(LayerAccessStatsInner::default())); + let new = LayerAccessStats(Mutex::new(LayerAccessStatsLocked::default())); new.record_residence_event(status, LayerResidenceEventReason::LayerLoad); new } pub(crate) fn for_new_layer_file() -> Self { - let new = LayerAccessStats(Mutex::new(LayerAccessStatsInner::default())); + let new = LayerAccessStats(Mutex::new(LayerAccessStatsLocked::default())); new.record_residence_event( LayerResidenceStatus::Resident, LayerResidenceEventReason::LayerCreate, @@ -176,38 +193,43 @@ impl LayerAccessStats { status: LayerResidenceStatus, reason: LayerResidenceEventReason, ) { - let mut inner = self.0.lock().unwrap(); - inner - .last_residence_changes - .write(LayerResidenceEvent::new(status, reason)); + let mut locked = self.0.lock().unwrap(); + locked.iter_mut().for_each(|inner| { + inner + .last_residence_changes + .write(LayerResidenceEvent::new(status, reason)) + }); } fn record_access(&self, access_kind: LayerAccessKind, task_kind: TaskKind) { - let mut inner = self.0.lock().unwrap(); let this_access = LayerAccessStatFullDetails { when: SystemTime::now(), task_kind, access_kind, }; - inner - .first_access - .get_or_insert_with(|| this_access.clone()); - inner.count_by_access_kind[access_kind] += 1; - inner.task_kind_flag |= task_kind; - inner.last_accesses.write(this_access); + + let mut locked = self.0.lock().unwrap(); + locked.iter_mut().for_each(|inner| { + inner.first_access.get_or_insert(this_access); + inner.count_by_access_kind[access_kind] += 1; + inner.task_kind_flag |= task_kind; + inner.last_accesses.write(this_access); + }) } - fn to_api_model( + + fn as_api_model( &self, reset: LayerAccessStatsReset, ) -> pageserver_api::models::LayerAccessStats { - let mut inner = self.0.lock().unwrap(); + let mut locked = self.0.lock().unwrap(); + let inner = &mut locked.for_scraping_api; let LayerAccessStatsInner { first_access, count_by_access_kind, task_kind_flag, last_accesses, last_residence_changes, - } = &*inner; + } = inner; let ret = pageserver_api::models::LayerAccessStats { access_count_by_access_kind: count_by_access_kind .iter() @@ -217,8 +239,8 @@ impl LayerAccessStats { .iter() .map(|task_kind| task_kind.into()) // into static str, powered by strum_macros .collect(), - first: first_access.as_ref().map(|a| a.to_api_model()), - accesses_history: last_accesses.map(|m| m.to_api_model()), + first: first_access.as_ref().map(|a| a.as_api_model()), + accesses_history: last_accesses.map(|m| m.as_api_model()), residence_events_history: last_residence_changes.clone(), }; match reset { @@ -232,6 +254,20 @@ impl LayerAccessStats { } ret } + + pub(super) fn most_recent_access_or_residence_event( + &self, + ) -> Either { + let locked = self.0.lock().unwrap(); + let inner = &locked.for_eviction_policy; + match inner.last_accesses.recent() { + Some(a) => Either::Left(*a), + None => match inner.last_residence_changes.recent() { + Some(e) => Either::Right(e.clone()), + None => unreachable!("constructors for LayerAccessStats ensure that there's always a residence change event"), + } + } + } } /// Supertrait of the [`Layer`] trait that captures the bare minimum interface @@ -449,3 +485,14 @@ enum PathOrConf { Path(PathBuf), Conf(&'static PageServerConf), } + +/// Range wrapping newtype, which uses display to render Debug. +/// +/// Useful with `Key`, which has too verbose `{:?}` for printing multiple layers. +struct RangeDisplayDebug<'a, T: std::fmt::Display>(&'a Range); + +impl<'a, T: std::fmt::Display> std::fmt::Debug for RangeDisplayDebug<'a, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}..{}", self.0.start, self.0.end) + } +} diff --git a/pageserver/src/tenant/storage_layer/delta_layer.rs b/pageserver/src/tenant/storage_layer/delta_layer.rs index 9b322faa65..4d1e08322d 100644 --- a/pageserver/src/tenant/storage_layer/delta_layer.rs +++ b/pageserver/src/tenant/storage_layer/delta_layer.rs @@ -194,8 +194,10 @@ pub struct DeltaLayer { impl std::fmt::Debug for DeltaLayer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use super::RangeDisplayDebug; + f.debug_struct("DeltaLayer") - .field("key_range", &self.key_range) + .field("key_range", &RangeDisplayDebug(&self.key_range)) .field("lsn_range", &self.lsn_range) .field("file_size", &self.file_size) .field("inner", &self.inner) @@ -450,7 +452,7 @@ impl PersistentLayer for DeltaLayer { let layer_file_name = self.filename().file_name(); let lsn_range = self.get_lsn_range(); - let access_stats = self.access_stats.to_api_model(reset); + let access_stats = self.access_stats.as_api_model(reset); HistoricLayerInfo::Delta { layer_file_name, diff --git a/pageserver/src/tenant/storage_layer/filename.rs b/pageserver/src/tenant/storage_layer/filename.rs index bd3d2c42c1..efd0769886 100644 --- a/pageserver/src/tenant/storage_layer/filename.rs +++ b/pageserver/src/tenant/storage_layer/filename.rs @@ -10,12 +10,23 @@ use std::str::FromStr; use utils::lsn::Lsn; // Note: Timeline::load_layer_map() relies on this sort order -#[derive(Debug, PartialEq, Eq, Clone, Hash)] +#[derive(PartialEq, Eq, Clone, Hash)] pub struct DeltaFileName { pub key_range: Range, pub lsn_range: Range, } +impl std::fmt::Debug for DeltaFileName { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + use super::RangeDisplayDebug; + + f.debug_struct("DeltaFileName") + .field("key_range", &RangeDisplayDebug(&self.key_range)) + .field("lsn_range", &self.lsn_range) + .finish() + } +} + impl PartialOrd for DeltaFileName { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) @@ -100,12 +111,23 @@ impl fmt::Display for DeltaFileName { } } -#[derive(Debug, PartialEq, Eq, Clone, Hash)] +#[derive(PartialEq, Eq, Clone, Hash)] pub struct ImageFileName { pub key_range: Range, pub lsn: Lsn, } +impl std::fmt::Debug for ImageFileName { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + use super::RangeDisplayDebug; + + f.debug_struct("ImageFileName") + .field("key_range", &RangeDisplayDebug(&self.key_range)) + .field("lsn", &self.lsn) + .finish() + } +} + impl PartialOrd for ImageFileName { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) diff --git a/pageserver/src/tenant/storage_layer/image_layer.rs b/pageserver/src/tenant/storage_layer/image_layer.rs index 86c1aee619..e48abd38dd 100644 --- a/pageserver/src/tenant/storage_layer/image_layer.rs +++ b/pageserver/src/tenant/storage_layer/image_layer.rs @@ -119,8 +119,10 @@ pub struct ImageLayer { impl std::fmt::Debug for ImageLayer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + use super::RangeDisplayDebug; + f.debug_struct("ImageLayer") - .field("key_range", &self.key_range) + .field("key_range", &RangeDisplayDebug(&self.key_range)) .field("file_size", &self.file_size) .field("lsn", &self.lsn) .field("inner", &self.inner) @@ -269,7 +271,7 @@ impl PersistentLayer for ImageLayer { layer_file_size: Some(self.file_size), lsn_start: lsn_range.start, remote: false, - access_stats: self.access_stats.to_api_model(reset), + access_stats: self.access_stats.as_api_model(reset), } } diff --git a/pageserver/src/tenant/storage_layer/remote_layer.rs b/pageserver/src/tenant/storage_layer/remote_layer.rs index 7391875d0c..51bb4dcc2a 100644 --- a/pageserver/src/tenant/storage_layer/remote_layer.rs +++ b/pageserver/src/tenant/storage_layer/remote_layer.rs @@ -171,7 +171,7 @@ impl PersistentLayer for RemoteLayer { lsn_start: lsn_range.start, lsn_end: lsn_range.end, remote: true, - access_stats: self.access_stats.to_api_model(reset), + access_stats: self.access_stats.as_api_model(reset), } } else { HistoricLayerInfo::Image { @@ -179,7 +179,7 @@ impl PersistentLayer for RemoteLayer { layer_file_size: self.layer_metadata.file_size(), lsn_start: lsn_range.start, remote: true, - access_stats: self.access_stats.to_api_model(reset), + access_stats: self.access_stats.as_api_model(reset), } } } diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index ff7b8c932c..5a829e42e5 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -1,5 +1,6 @@ //! +mod eviction_task; mod walreceiver; use anyhow::{anyhow, bail, ensure, Context}; @@ -47,7 +48,7 @@ use crate::metrics::TimelineMetrics; use crate::pgdatadir_mapping::LsnForTimestamp; use crate::pgdatadir_mapping::{is_rel_fsm_block_key, is_rel_vm_block_key}; use crate::pgdatadir_mapping::{BlockNumber, CalculateLogicalSizeError}; -use crate::tenant::config::TenantConfOpt; +use crate::tenant::config::{EvictionPolicy, TenantConfOpt}; use pageserver_api::reltag::RelTag; use postgres_connection::PgConnectionConfig; @@ -291,18 +292,9 @@ impl LogicalSize { // we change the type. match self.initial_logical_size.get() { Some(initial_size) => { - let absolute_size_increment = u64::try_from( - size_increment - .checked_abs() - .with_context(|| format!("Size added after initial {size_increment} is not expected to be i64::MIN"))?, - ).expect("casting nonnegative i64 to u64 should not fail"); - - if size_increment < 0 { - initial_size.checked_sub(absolute_size_increment) - } else { - initial_size.checked_add(absolute_size_increment) - }.with_context(|| format!("Overflow during logical size calculation, initial_size: {initial_size}, size_increment: {size_increment}")) - .map(CurrentLogicalSize::Exact) + initial_size.checked_add_signed(size_increment) + .with_context(|| format!("Overflow during logical size calculation, initial_size: {initial_size}, size_increment: {size_increment}")) + .map(CurrentLogicalSize::Exact) } None => { let non_negative_size_increment = u64::try_from(size_increment).unwrap_or(0); @@ -801,6 +793,7 @@ impl Timeline { pub fn activate(self: &Arc) { self.set_state(TimelineState::Active); self.launch_wal_receiver(); + self.launch_eviction_task(); } pub fn set_state(&self, new_state: TimelineState) { @@ -867,30 +860,101 @@ impl Timeline { Ok(Some(true)) } - /// Evicts one layer as in replaces a downloaded layer with a remote layer - /// - /// Returns: - /// - `Ok(Some(true))` when the layer was replaced - /// - `Ok(Some(false))` when the layer was found, but no changes were made - /// - evictee was not yet downloaded - /// - layermap replacement failed - /// - `Ok(None)` when the layer is not found + /// Like [`evict_layer_batch`], but for just one layer. + /// Additional case `Ok(None)` covers the case where the layer could not be found by its `layer_file_name`. pub async fn evict_layer(&self, layer_file_name: &str) -> anyhow::Result> { - use super::layer_map::Replacement; - let Some(local_layer) = self.find_layer(layer_file_name) else { return Ok(None) }; - if local_layer.is_remote_layer() { - return Ok(Some(false)); - } - - // ensure the current layer is uploaded for sure - self.remote_client + let remote_client = self + .remote_client .as_ref() - .ok_or_else(|| anyhow::anyhow!("remote storage not configured; cannot evict"))? + .ok_or_else(|| anyhow::anyhow!("remote storage not configured; cannot evict"))?; + + let cancel = CancellationToken::new(); + let results = self + .evict_layer_batch(remote_client, &[local_layer], cancel) + .await?; + assert_eq!(results.len(), 1); + let result: Option> = results.into_iter().next().unwrap(); + match result { + None => anyhow::bail!("task_mgr shutdown requested"), + Some(Ok(b)) => Ok(Some(b)), + Some(Err(e)) => Err(e), + } + } + + /// Evict multiple layers at once, continuing through errors. + /// + /// Try to evict the given `layers_to_evict` by + /// + /// 1. Replacing the given layer object in the layer map with a corresponding [`RemoteLayer`] object. + /// 2. Deleting the now unreferenced layer file from disk. + /// + /// The `remote_client` should be this timeline's `self.remote_client`. + /// We make the caller provide it so that they are responsible for handling the case + /// where someone wants to evict the layer but no remote storage is configured. + /// + /// Returns either `Err()` or `Ok(results)` where `results.len() == layers_to_evict.len()`. + /// If `Err()` is returned, no eviction was attempted. + /// Each position of `Ok(results)` corresponds to the layer in `layers_to_evict`. + /// Meaning of each `result[i]`: + /// - `Some(Err(...))` if layer replacement failed for an unexpected reason + /// - `Some(Ok(true))` if everything went well. + /// - `Some(Ok(false))` if there was an expected reason why the layer could not be replaced, e.g.: + /// - evictee was not yet downloaded + /// - replacement failed for an expectable reason (e.g., layer removed by GC before we grabbed all locks) + /// - `None` if no eviction attempt was made for the layer because `cancel.is_cancelled() == true`. + async fn evict_layer_batch( + &self, + remote_client: &Arc, + layers_to_evict: &[Arc], + cancel: CancellationToken, + ) -> anyhow::Result>>> { + // ensure that the layers have finished uploading + // (don't hold the layer_removal_cs while we do it, we're not removing anything yet) + remote_client .wait_completion() .await .context("wait for layer upload ops to complete")?; + // now lock out layer removal (compaction, gc, timeline deletion) + let layer_removal_guard = self.layer_removal_cs.lock().await; + + // start the batch update + let mut layer_map = self.layers.write().unwrap(); + let mut batch_updates = layer_map.batch_update(); + + let mut results = Vec::with_capacity(layers_to_evict.len()); + + for l in layers_to_evict.iter() { + let res = if cancel.is_cancelled() { + None + } else { + Some(self.evict_layer_batch_impl(&layer_removal_guard, l, &mut batch_updates)) + }; + results.push(res); + } + + // commit the updates & release locks + batch_updates.flush(); + drop(layer_map); + drop(layer_removal_guard); + + assert_eq!(results.len(), layers_to_evict.len()); + Ok(results) + } + + fn evict_layer_batch_impl( + &self, + _layer_removal_cs: &tokio::sync::MutexGuard<'_, ()>, + local_layer: &Arc, + batch_updates: &mut BatchedUpdates<'_, dyn PersistentLayer>, + ) -> anyhow::Result { + use super::layer_map::Replacement; + + if local_layer.is_remote_layer() { + return Ok(false); + } + let layer_metadata = LayerFileMetadata::new( local_layer .file_size() @@ -917,11 +981,7 @@ impl Timeline { ), }); - let gc_lock = self.layer_removal_cs.lock().await; - let mut layers = self.layers.write().unwrap(); - let mut updates = layers.batch_update(); - - let replaced = match updates.replace_historic(&local_layer, new_remote_layer)? { + let replaced = match batch_updates.replace_historic(local_layer, new_remote_layer)? { Replacement::Replaced { .. } => { let layer_size = local_layer.file_size(); @@ -944,7 +1004,7 @@ impl Timeline { } Replacement::Unexpected(other) => { error!( - local_layer.ptr=?Arc::as_ptr(&local_layer), + local_layer.ptr=?Arc::as_ptr(local_layer), other.ptr=?Arc::as_ptr(&other), ?other, "failed to replace"); @@ -952,11 +1012,7 @@ impl Timeline { } }; - updates.flush(); - drop(layers); - drop(gc_lock); - - Ok(Some(replaced)) + Ok(replaced) } } @@ -997,6 +1053,13 @@ impl Timeline { .unwrap_or(self.conf.default_tenant_conf.image_creation_threshold) } + fn get_eviction_policy(&self) -> EvictionPolicy { + let tenant_conf = self.tenant_conf.read().unwrap(); + tenant_conf + .eviction_policy + .unwrap_or(self.conf.default_tenant_conf.eviction_policy) + } + /// Open a Timeline handle. /// /// Loads the metadata for the timeline into memory, but not the layer map. @@ -1553,13 +1616,31 @@ impl Timeline { } x @ Err(_) => x.context("Failed to calculate logical size")?, }; + + // we cannot query current_logical_size.current_size() to know the current + // *negative* value, only truncated to u64. + let added = self_clone + .current_logical_size + .size_added_after_initial + .load(AtomicOrdering::Relaxed); + + let sum = calculated_size.saturating_add_signed(added); + + // set the gauge value before it can be set in `update_current_logical_size`. + self_clone.metrics.current_logical_size_gauge.set(sum); + match self_clone .current_logical_size .initial_logical_size .set(calculated_size) { Ok(()) => (), - Err(existing_size) => { + Err(_what_we_just_attempted_to_set) => { + let existing_size = self_clone + .current_logical_size + .initial_logical_size + .get() + .expect("once_cell set was lost, then get failed, impossible."); // This shouldn't happen because the semaphore is initialized with 1. // But if it happens, just complain & report success so there are no further retries. error!("Tried to update initial timeline size value to {calculated_size}, but the size was already set to {existing_size}, not changing") @@ -1742,10 +1823,15 @@ impl Timeline { // one value while current_logical_size is set to the // other. match logical_size.current_size() { - Ok(new_current_size) => self + Ok(CurrentLogicalSize::Exact(new_current_size)) => self .metrics .current_logical_size_gauge - .set(new_current_size.size()), + .set(new_current_size), + Ok(CurrentLogicalSize::Approximate(_)) => { + // don't update the gauge yet, this allows us not to update the gauge back and + // forth between the initial size calculation task. + } + // this is overflow Err(e) => error!("Failed to compute current logical size for metrics update: {e:?}"), } } @@ -2314,7 +2400,7 @@ impl Timeline { // Only one thread may call this function at a time (for this // timeline). If two threads tried to flush the same frozen // layer to disk at the same time, that would not work. - assert!(Arc::ptr_eq(&l.unwrap(), &frozen_layer)); + assert!(LayerMap::compare_arced_layers(&l.unwrap(), &frozen_layer)); // release lock on 'layers' } diff --git a/pageserver/src/tenant/timeline/eviction_task.rs b/pageserver/src/tenant/timeline/eviction_task.rs new file mode 100644 index 0000000000..e3e7ce4c9d --- /dev/null +++ b/pageserver/src/tenant/timeline/eviction_task.rs @@ -0,0 +1,199 @@ +//! The per-timeline layer eviction task. + +use std::{ + ops::ControlFlow, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use either::Either; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, info, instrument, warn}; + +use crate::{ + task_mgr::{self, TaskKind, BACKGROUND_RUNTIME}, + tenant::{ + config::{EvictionPolicy, EvictionPolicyLayerAccessThreshold}, + storage_layer::PersistentLayer, + }, +}; + +use super::Timeline; + +impl Timeline { + pub(super) fn launch_eviction_task(self: &Arc) { + let self_clone = Arc::clone(self); + task_mgr::spawn( + BACKGROUND_RUNTIME.handle(), + TaskKind::Eviction, + Some(self.tenant_id), + Some(self.timeline_id), + &format!("layer eviction for {}/{}", self.tenant_id, self.timeline_id), + false, + async move { + self_clone.eviction_task(task_mgr::shutdown_token()).await; + info!("eviction task finishing"); + Ok(()) + }, + ); + } + + #[instrument(skip_all, fields(tenant_id = %self.tenant_id, timeline_id = %self.timeline_id))] + async fn eviction_task(self: Arc, cancel: CancellationToken) { + loop { + let policy = self.get_eviction_policy(); + let cf = self.eviction_iteration(&policy, cancel.clone()).await; + match cf { + ControlFlow::Break(()) => break, + ControlFlow::Continue(sleep_until) => { + tokio::select! { + _ = cancel.cancelled() => { + info!("shutting down"); + break; + } + _ = tokio::time::sleep_until(sleep_until) => { } + } + } + } + } + } + + #[instrument(skip_all, fields(policy_kind = policy.discriminant_str()))] + async fn eviction_iteration( + self: &Arc, + policy: &EvictionPolicy, + cancel: CancellationToken, + ) -> ControlFlow<(), Instant> { + debug!("eviction iteration: {policy:?}"); + match policy { + EvictionPolicy::NoEviction => { + // check again in 10 seconds; XXX config watch mechanism + ControlFlow::Continue(Instant::now() + Duration::from_secs(10)) + } + EvictionPolicy::LayerAccessThreshold(p) => { + let start = Instant::now(); + match self.eviction_iteration_threshold(p, cancel).await { + ControlFlow::Break(()) => return ControlFlow::Break(()), + ControlFlow::Continue(()) => (), + } + let elapsed = start.elapsed(); + if elapsed > p.period { + warn!( + configured_period = %humantime::format_duration(p.period), + last_period = %humantime::format_duration(elapsed), + "this eviction period took longer than the configured period" + ); + } + ControlFlow::Continue(start + p.period) + } + } + } + + async fn eviction_iteration_threshold( + self: &Arc, + p: &EvictionPolicyLayerAccessThreshold, + cancel: CancellationToken, + ) -> ControlFlow<()> { + let now = SystemTime::now(); + + #[allow(dead_code)] + #[derive(Debug, Default)] + struct EvictionStats { + not_considered_due_to_clock_skew: usize, + candidates: usize, + evicted: usize, + errors: usize, + not_evictable: usize, + skipped_for_shutdown: usize, + } + let mut stats = EvictionStats::default(); + // Gather layers for eviction. + // NB: all the checks can be invalidated as soon as we release the layer map lock. + // We don't want to hold the layer map lock during eviction. + // So, we just need to deal with this. + let candidates: Vec> = { + let layers = self.layers.read().unwrap(); + let mut candidates = Vec::new(); + for hist_layer in layers.iter_historic_layers() { + if hist_layer.is_remote_layer() { + continue; + } + let last_activity_ts = match hist_layer + .access_stats() + .most_recent_access_or_residence_event() + { + Either::Left(mra) => mra.when, + Either::Right(re) => re.timestamp, + }; + let no_activity_for = match now.duration_since(last_activity_ts) { + Ok(d) => d, + Err(_e) => { + // NB: don't log the error. If there are many layers and the system clock + // is skewed, we'd be flooding the log. + stats.not_considered_due_to_clock_skew += 1; + continue; + } + }; + if no_activity_for > p.threshold { + candidates.push(hist_layer) + } + } + candidates + }; + stats.candidates = candidates.len(); + + let remote_client = match self.remote_client.as_ref() { + None => { + error!( + num_candidates = candidates.len(), + "no remote storage configured, cannot evict layers" + ); + return ControlFlow::Continue(()); + } + Some(c) => c, + }; + + let results = match self + .evict_layer_batch(remote_client, &candidates[..], cancel) + .await + { + Err(pre_err) => { + stats.errors += candidates.len(); + error!("could not do any evictions: {pre_err:#}"); + return ControlFlow::Continue(()); + } + Ok(results) => results, + }; + assert_eq!(results.len(), candidates.len()); + for (l, result) in candidates.iter().zip(results) { + match result { + None => { + stats.skipped_for_shutdown += 1; + } + Some(Ok(true)) => { + debug!("evicted layer {l:?}"); + stats.evicted += 1; + } + Some(Ok(false)) => { + debug!("layer is not evictable: {l:?}"); + stats.not_evictable += 1; + } + Some(Err(e)) => { + // This variant is the case where an unexpected error happened during eviction. + // Expected errors that result in non-eviction are `Some(Ok(false))`. + // So, dump Debug here to gather as much info as possible in this rare case. + warn!("failed to evict layer {l:?}: {e:?}"); + stats.errors += 1; + } + } + } + if stats.not_considered_due_to_clock_skew > 0 || stats.errors > 0 || stats.not_evictable > 0 + { + warn!(stats=?stats, "eviction iteration complete"); + } else { + info!(stats=?stats, "eviction iteration complete"); + } + ControlFlow::Continue(()) + } +} diff --git a/pageserver/src/tenant/timeline/walreceiver/connection_manager.rs b/pageserver/src/tenant/timeline/walreceiver/connection_manager.rs index cd7c7c51d2..64a79b6d1b 100644 --- a/pageserver/src/tenant/timeline/walreceiver/connection_manager.rs +++ b/pageserver/src/tenant/timeline/walreceiver/connection_manager.rs @@ -13,7 +13,7 @@ use std::{collections::HashMap, num::NonZeroU64, ops::ControlFlow, sync::Arc, ti use super::TaskStateUpdate; use crate::broker_client::get_broker_client; -use crate::context::RequestContext; +use crate::context::{DownloadBehavior, RequestContext}; use crate::task_mgr::WALRECEIVER_RUNTIME; use crate::task_mgr::{self, TaskKind}; use crate::tenant::Timeline; @@ -413,7 +413,7 @@ impl WalreceiverState { let timeline = Arc::clone(&self.timeline); let ctx = ctx.detached_child( TaskKind::WalReceiverConnectionHandler, - ctx.download_behavior(), + DownloadBehavior::Download, ); let connection_handle = TaskHandle::spawn(move |events_sender, cancellation| { async move { diff --git a/poetry.lock b/poetry.lock index f14c495556..7e80b1e10a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -865,50 +865,47 @@ files = [ [[package]] name = "cryptography" -version = "38.0.3" +version = "39.0.1" description = "cryptography is a package which provides cryptographic recipes and primitives to Python developers." category = "main" optional = false python-versions = ">=3.6" files = [ - {file = "cryptography-38.0.3-cp36-abi3-macosx_10_10_universal2.whl", hash = "sha256:984fe150f350a3c91e84de405fe49e688aa6092b3525f407a18b9646f6612320"}, - {file = "cryptography-38.0.3-cp36-abi3-macosx_10_10_x86_64.whl", hash = "sha256:ed7b00096790213e09eb11c97cc6e2b757f15f3d2f85833cd2d3ec3fe37c1722"}, - {file = "cryptography-38.0.3-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_24_aarch64.whl", hash = "sha256:bbf203f1a814007ce24bd4d51362991d5cb90ba0c177a9c08825f2cc304d871f"}, - {file = "cryptography-38.0.3-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:554bec92ee7d1e9d10ded2f7e92a5d70c1f74ba9524947c0ba0c850c7b011828"}, - {file = "cryptography-38.0.3-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b1b52c9e5f8aa2b802d48bd693190341fae201ea51c7a167d69fc48b60e8a959"}, - {file = "cryptography-38.0.3-cp36-abi3-manylinux_2_24_x86_64.whl", hash = "sha256:728f2694fa743a996d7784a6194da430f197d5c58e2f4e278612b359f455e4a2"}, - {file = "cryptography-38.0.3-cp36-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:dfb4f4dd568de1b6af9f4cda334adf7d72cf5bc052516e1b2608b683375dd95c"}, - {file = "cryptography-38.0.3-cp36-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:5419a127426084933076132d317911e3c6eb77568a1ce23c3ac1e12d111e61e0"}, - {file = "cryptography-38.0.3-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:9b24bcff7853ed18a63cfb0c2b008936a9554af24af2fb146e16d8e1aed75748"}, - {file = "cryptography-38.0.3-cp36-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:25c1d1f19729fb09d42e06b4bf9895212292cb27bb50229f5aa64d039ab29146"}, - {file = "cryptography-38.0.3-cp36-abi3-win32.whl", hash = "sha256:7f836217000342d448e1c9a342e9163149e45d5b5eca76a30e84503a5a96cab0"}, - {file = "cryptography-38.0.3-cp36-abi3-win_amd64.whl", hash = "sha256:c46837ea467ed1efea562bbeb543994c2d1f6e800785bd5a2c98bc096f5cb220"}, - {file = "cryptography-38.0.3-pp37-pypy37_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:06fc3cc7b6f6cca87bd56ec80a580c88f1da5306f505876a71c8cfa7050257dd"}, - {file = "cryptography-38.0.3-pp37-pypy37_pp73-manylinux_2_24_x86_64.whl", hash = "sha256:65535bc550b70bd6271984d9863a37741352b4aad6fb1b3344a54e6950249b55"}, - {file = "cryptography-38.0.3-pp37-pypy37_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:5e89468fbd2fcd733b5899333bc54d0d06c80e04cd23d8c6f3e0542358c6060b"}, - {file = "cryptography-38.0.3-pp38-pypy38_pp73-macosx_10_10_x86_64.whl", hash = "sha256:6ab9516b85bebe7aa83f309bacc5f44a61eeb90d0b4ec125d2d003ce41932d36"}, - {file = "cryptography-38.0.3-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:068147f32fa662c81aebab95c74679b401b12b57494872886eb5c1139250ec5d"}, - {file = "cryptography-38.0.3-pp38-pypy38_pp73-manylinux_2_24_x86_64.whl", hash = "sha256:402852a0aea73833d982cabb6d0c3bb582c15483d29fb7085ef2c42bfa7e38d7"}, - {file = "cryptography-38.0.3-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:b1b35d9d3a65542ed2e9d90115dfd16bbc027b3f07ee3304fc83580f26e43249"}, - {file = "cryptography-38.0.3-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:6addc3b6d593cd980989261dc1cce38263c76954d758c3c94de51f1e010c9a50"}, - {file = "cryptography-38.0.3-pp39-pypy39_pp73-macosx_10_10_x86_64.whl", hash = "sha256:be243c7e2bfcf6cc4cb350c0d5cdf15ca6383bbcb2a8ef51d3c9411a9d4386f0"}, - {file = "cryptography-38.0.3-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:78cf5eefac2b52c10398a42765bfa981ce2372cbc0457e6bf9658f41ec3c41d8"}, - {file = "cryptography-38.0.3-pp39-pypy39_pp73-manylinux_2_24_x86_64.whl", hash = "sha256:4e269dcd9b102c5a3d72be3c45d8ce20377b8076a43cbed6f660a1afe365e436"}, - {file = "cryptography-38.0.3-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:8d41a46251bf0634e21fac50ffd643216ccecfaf3701a063257fe0b2be1b6548"}, - {file = "cryptography-38.0.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:785e4056b5a8b28f05a533fab69febf5004458e20dad7e2e13a3120d8ecec75a"}, - {file = "cryptography-38.0.3.tar.gz", hash = "sha256:bfbe6ee19615b07a98b1d2287d6a6073f734735b49ee45b11324d85efc4d5cbd"}, + {file = "cryptography-39.0.1-cp36-abi3-macosx_10_12_universal2.whl", hash = "sha256:6687ef6d0a6497e2b58e7c5b852b53f62142cfa7cd1555795758934da363a965"}, + {file = "cryptography-39.0.1-cp36-abi3-macosx_10_12_x86_64.whl", hash = "sha256:706843b48f9a3f9b9911979761c91541e3d90db1ca905fd63fee540a217698bc"}, + {file = "cryptography-39.0.1-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_24_aarch64.whl", hash = "sha256:5d2d8b87a490bfcd407ed9d49093793d0f75198a35e6eb1a923ce1ee86c62b41"}, + {file = "cryptography-39.0.1-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:83e17b26de248c33f3acffb922748151d71827d6021d98c70e6c1a25ddd78505"}, + {file = "cryptography-39.0.1-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e124352fd3db36a9d4a21c1aa27fd5d051e621845cb87fb851c08f4f75ce8be6"}, + {file = "cryptography-39.0.1-cp36-abi3-manylinux_2_24_x86_64.whl", hash = "sha256:5aa67414fcdfa22cf052e640cb5ddc461924a045cacf325cd164e65312d99502"}, + {file = "cryptography-39.0.1-cp36-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:35f7c7d015d474f4011e859e93e789c87d21f6f4880ebdc29896a60403328f1f"}, + {file = "cryptography-39.0.1-cp36-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:f24077a3b5298a5a06a8e0536e3ea9ec60e4c7ac486755e5fb6e6ea9b3500106"}, + {file = "cryptography-39.0.1-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:f0c64d1bd842ca2633e74a1a28033d139368ad959872533b1bab8c80e8240a0c"}, + {file = "cryptography-39.0.1-cp36-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:0f8da300b5c8af9f98111ffd512910bc792b4c77392a9523624680f7956a99d4"}, + {file = "cryptography-39.0.1-cp36-abi3-win32.whl", hash = "sha256:fe913f20024eb2cb2f323e42a64bdf2911bb9738a15dba7d3cce48151034e3a8"}, + {file = "cryptography-39.0.1-cp36-abi3-win_amd64.whl", hash = "sha256:ced4e447ae29ca194449a3f1ce132ded8fcab06971ef5f618605aacaa612beac"}, + {file = "cryptography-39.0.1-pp38-pypy38_pp73-macosx_10_12_x86_64.whl", hash = "sha256:807ce09d4434881ca3a7594733669bd834f5b2c6d5c7e36f8c00f691887042ad"}, + {file = "cryptography-39.0.1-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:96f1157a7c08b5b189b16b47bc9db2332269d6680a196341bf30046330d15388"}, + {file = "cryptography-39.0.1-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:e422abdec8b5fa8462aa016786680720d78bdce7a30c652b7fadf83a4ba35336"}, + {file = "cryptography-39.0.1-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_24_aarch64.whl", hash = "sha256:b0afd054cd42f3d213bf82c629efb1ee5f22eba35bf0eec88ea9ea7304f511a2"}, + {file = "cryptography-39.0.1-pp39-pypy39_pp73-manylinux_2_24_x86_64.whl", hash = "sha256:6f8ba7f0328b79f08bdacc3e4e66fb4d7aab0c3584e0bd41328dce5262e26b2e"}, + {file = "cryptography-39.0.1-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:ef8b72fa70b348724ff1218267e7f7375b8de4e8194d1636ee60510aae104cd0"}, + {file = "cryptography-39.0.1-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:aec5a6c9864be7df2240c382740fcf3b96928c46604eaa7f3091f58b878c0bb6"}, + {file = "cryptography-39.0.1-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:fdd188c8a6ef8769f148f88f859884507b954cc64db6b52f66ef199bb9ad660a"}, + {file = "cryptography-39.0.1.tar.gz", hash = "sha256:d1f6198ee6d9148405e49887803907fe8962a23e6c6f83ea7d98f1c0de375695"}, ] [package.dependencies] cffi = ">=1.12" [package.extras] -docs = ["sphinx (>=1.6.5,!=1.8.0,!=3.1.0,!=3.1.1)", "sphinx-rtd-theme"] +docs = ["sphinx (>=5.3.0)", "sphinx-rtd-theme (>=1.1.1)"] docstest = ["pyenchant (>=1.6.11)", "sphinxcontrib-spelling (>=4.0.1)", "twine (>=1.12.0)"] -pep8test = ["black", "flake8", "flake8-import-order", "pep8-naming"] +pep8test = ["black", "check-manifest", "mypy", "ruff", "types-pytz", "types-requests"] sdist = ["setuptools-rust (>=0.11.4)"] ssh = ["bcrypt (>=3.1.5)"] -test = ["hypothesis (>=1.11.4,!=3.79.2)", "iso8601", "pretend", "pytest (>=6.2.0)", "pytest-benchmark", "pytest-cov", "pytest-subtests", "pytest-xdist", "pytz"] +test = ["hypothesis (>=1.11.4,!=3.79.2)", "iso8601", "pretend", "pytest (>=6.2.0)", "pytest-benchmark", "pytest-cov", "pytest-shard (>=0.1.2)", "pytest-subtests", "pytest-xdist", "pytz"] +test-randomorder = ["pytest-randomly"] +tox = ["tox"] [[package]] name = "docker" @@ -1241,6 +1238,7 @@ category = "main" optional = false python-versions = "*" files = [ + {file = "junit-xml-1.9.tar.gz", hash = "sha256:de16a051990d4e25a3982b2dd9e89d671067548718866416faec14d9de56db9f"}, {file = "junit_xml-1.9-py2.py3-none-any.whl", hash = "sha256:ec5ca1a55aefdd76d28fcc0b135251d156c7106fa979686a4b48d62b761b4732"}, ] @@ -1718,6 +1716,7 @@ python-versions = ">=3.6" files = [ {file = "psycopg2-binary-2.9.3.tar.gz", hash = "sha256:761df5313dc15da1502b21453642d7599d26be88bff659382f8f9747c7ebea4e"}, {file = "psycopg2_binary-2.9.3-cp310-cp310-macosx_10_14_x86_64.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", hash = "sha256:539b28661b71da7c0e428692438efbcd048ca21ea81af618d845e06ebfd29478"}, + {file = "psycopg2_binary-2.9.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2f2534ab7dc7e776a263b463a16e189eb30e85ec9bbe1bff9e78dae802608932"}, {file = "psycopg2_binary-2.9.3-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6e82d38390a03da28c7985b394ec3f56873174e2c88130e6966cb1c946508e65"}, {file = "psycopg2_binary-2.9.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:57804fc02ca3ce0dbfbef35c4b3a4a774da66d66ea20f4bda601294ad2ea6092"}, {file = "psycopg2_binary-2.9.3-cp310-cp310-manylinux_2_24_aarch64.whl", hash = "sha256:083a55275f09a62b8ca4902dd11f4b33075b743cf0d360419e2051a8a5d5ff76"}, @@ -1751,6 +1750,7 @@ files = [ {file = "psycopg2_binary-2.9.3-cp37-cp37m-win32.whl", hash = "sha256:adf20d9a67e0b6393eac162eb81fb10bc9130a80540f4df7e7355c2dd4af9fba"}, {file = "psycopg2_binary-2.9.3-cp37-cp37m-win_amd64.whl", hash = "sha256:2f9ffd643bc7349eeb664eba8864d9e01f057880f510e4681ba40a6532f93c71"}, {file = "psycopg2_binary-2.9.3-cp38-cp38-macosx_10_14_x86_64.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", hash = "sha256:def68d7c21984b0f8218e8a15d514f714d96904265164f75f8d3a70f9c295667"}, + {file = "psycopg2_binary-2.9.3-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:e6aa71ae45f952a2205377773e76f4e3f27951df38e69a4c95440c779e013560"}, {file = "psycopg2_binary-2.9.3-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:dffc08ca91c9ac09008870c9eb77b00a46b3378719584059c034b8945e26b272"}, {file = "psycopg2_binary-2.9.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:280b0bb5cbfe8039205c7981cceb006156a675362a00fe29b16fbc264e242834"}, {file = "psycopg2_binary-2.9.3-cp38-cp38-manylinux_2_24_aarch64.whl", hash = "sha256:af9813db73395fb1fc211bac696faea4ca9ef53f32dc0cfa27e4e7cf766dcf24"}, @@ -1762,6 +1762,7 @@ files = [ {file = "psycopg2_binary-2.9.3-cp38-cp38-win32.whl", hash = "sha256:6472a178e291b59e7f16ab49ec8b4f3bdada0a879c68d3817ff0963e722a82ce"}, {file = "psycopg2_binary-2.9.3-cp38-cp38-win_amd64.whl", hash = "sha256:35168209c9d51b145e459e05c31a9eaeffa9a6b0fd61689b48e07464ffd1a83e"}, {file = "psycopg2_binary-2.9.3-cp39-cp39-macosx_10_14_x86_64.macosx_10_9_intel.macosx_10_9_x86_64.macosx_10_10_intel.macosx_10_10_x86_64.whl", hash = "sha256:47133f3f872faf28c1e87d4357220e809dfd3fa7c64295a4a148bcd1e6e34ec9"}, + {file = "psycopg2_binary-2.9.3-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:b3a24a1982ae56461cc24f6680604fffa2c1b818e9dc55680da038792e004d18"}, {file = "psycopg2_binary-2.9.3-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:91920527dea30175cc02a1099f331aa8c1ba39bf8b7762b7b56cbf54bc5cce42"}, {file = "psycopg2_binary-2.9.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:887dd9aac71765ac0d0bac1d0d4b4f2c99d5f5c1382d8b770404f0f3d0ce8a39"}, {file = "psycopg2_binary-2.9.3-cp39-cp39-manylinux_2_24_aarch64.whl", hash = "sha256:1f14c8b0942714eb3c74e1e71700cbbcb415acbc311c730370e70c578a44a25c"}, @@ -1794,18 +1795,7 @@ category = "main" optional = false python-versions = "*" files = [ - {file = "pyasn1-0.4.8-py2.4.egg", hash = "sha256:fec3e9d8e36808a28efb59b489e4528c10ad0f480e57dcc32b4de5c9d8c9fdf3"}, - {file = "pyasn1-0.4.8-py2.5.egg", hash = "sha256:0458773cfe65b153891ac249bcf1b5f8f320b7c2ce462151f8fa74de8934becf"}, - {file = "pyasn1-0.4.8-py2.6.egg", hash = "sha256:5c9414dcfede6e441f7e8f81b43b34e834731003427e5b09e4e00e3172a10f00"}, - {file = "pyasn1-0.4.8-py2.7.egg", hash = "sha256:6e7545f1a61025a4e58bb336952c5061697da694db1cae97b116e9c46abcf7c8"}, {file = "pyasn1-0.4.8-py2.py3-none-any.whl", hash = "sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d"}, - {file = "pyasn1-0.4.8-py3.1.egg", hash = "sha256:78fa6da68ed2727915c4767bb386ab32cdba863caa7dbe473eaae45f9959da86"}, - {file = "pyasn1-0.4.8-py3.2.egg", hash = "sha256:08c3c53b75eaa48d71cf8c710312316392ed40899cb34710d092e96745a358b7"}, - {file = "pyasn1-0.4.8-py3.3.egg", hash = "sha256:03840c999ba71680a131cfaee6fab142e1ed9bbd9c693e285cc6aca0d555e576"}, - {file = "pyasn1-0.4.8-py3.4.egg", hash = "sha256:7ab8a544af125fb704feadb008c99a88805126fb525280b2270bb25cc1d78a12"}, - {file = "pyasn1-0.4.8-py3.5.egg", hash = "sha256:e89bf84b5437b532b0803ba5c9a5e054d21fec423a89952a74f87fa2c9b7bce2"}, - {file = "pyasn1-0.4.8-py3.6.egg", hash = "sha256:014c0e9976956a08139dc0712ae195324a75e142284d5f87f1a87ee1b068a359"}, - {file = "pyasn1-0.4.8-py3.7.egg", hash = "sha256:99fcc3c8d804d1bc6d9a099921e39d827026409a58f2a720dcdb89374ea0c776"}, {file = "pyasn1-0.4.8.tar.gz", hash = "sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba"}, ] @@ -2014,8 +2004,8 @@ files = [ [package.dependencies] pytest = [ - {version = ">=6.2.4", markers = "python_version >= \"3.10\""}, {version = ">=5.0", markers = "python_version < \"3.10\""}, + {version = ">=6.2.4", markers = "python_version >= \"3.10\""}, ] [[package]] @@ -2139,6 +2129,13 @@ files = [ {file = "PyYAML-6.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:f84fbc98b019fef2ee9a1cb3ce93e3187a6df0b2538a651bfb890254ba9f90b5"}, {file = "PyYAML-6.0-cp310-cp310-win32.whl", hash = "sha256:2cd5df3de48857ed0544b34e2d40e9fac445930039f3cfe4bcc592a1f836d513"}, {file = "PyYAML-6.0-cp310-cp310-win_amd64.whl", hash = "sha256:daf496c58a8c52083df09b80c860005194014c3698698d1a57cbcfa182142a3a"}, + {file = "PyYAML-6.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d4b0ba9512519522b118090257be113b9468d804b19d63c71dbcf4a48fa32358"}, + {file = "PyYAML-6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:81957921f441d50af23654aa6c5e5eaf9b06aba7f0a19c18a538dc7ef291c5a1"}, + {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afa17f5bc4d1b10afd4466fd3a44dc0e245382deca5b3c353d8b757f9e3ecb8d"}, + {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dbad0e9d368bb989f4515da330b88a057617d16b6a8245084f1b05400f24609f"}, + {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:432557aa2c09802be39460360ddffd48156e30721f5e8d917f01d31694216782"}, + {file = "PyYAML-6.0-cp311-cp311-win32.whl", hash = "sha256:bfaef573a63ba8923503d27530362590ff4f576c626d86a9fed95822a8255fd7"}, + {file = "PyYAML-6.0-cp311-cp311-win_amd64.whl", hash = "sha256:01b45c0191e6d66c470b6cf1b9531a771a83c1c4208272ead47a3ae4f2f603bf"}, {file = "PyYAML-6.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:897b80890765f037df3403d22bab41627ca8811ae55e9a722fd0392850ec4d86"}, {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:50602afada6d6cbfad699b0c7bb50d5ccffa7e46a3d738092afddc1f9758427f"}, {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:48c346915c114f5fdb3ead70312bd042a953a8ce5c7106d5bfb1a5254e47da92"}, diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index 1ff7eebd98..152c83e4a0 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -43,6 +43,7 @@ serde.workspace = true serde_json.workspace = true sha2.workspace = true socket2.workspace = true +sync_wrapper.workspace = true thiserror.workspace = true tls-listener.workspace = true tokio-postgres.workspace = true diff --git a/proxy/src/auth/backend/link.rs b/proxy/src/auth/backend/link.rs index ef92b1a444..5d0049c957 100644 --- a/proxy/src/auth/backend/link.rs +++ b/proxy/src/auth/backend/link.rs @@ -78,6 +78,8 @@ pub(super) async fn handle_user( client.write_message_noflush(&Be::NoticeResponse("Connecting to database."))?; + // This config should be self-contained, because we won't + // take username or dbname from client's startup message. let mut config = compute::ConnCfg::new(); config .host(&db_info.host) diff --git a/proxy/src/auth/credentials.rs b/proxy/src/auth/credentials.rs index 66ca8be73e..968104f058 100644 --- a/proxy/src/auth/credentials.rs +++ b/proxy/src/auth/credentials.rs @@ -32,7 +32,6 @@ impl UserFacingError for ClientCredsParseError {} #[derive(Debug, Clone, PartialEq, Eq)] pub struct ClientCredentials<'a> { pub user: &'a str, - pub dbname: &'a str, // TODO: this is a severe misnomer! We should think of a new name ASAP. pub project: Option>, /// If `True`, we'll use the old cleartext password flow. This is used for @@ -59,7 +58,6 @@ impl<'a> ClientCredentials<'a> { // Some parameters are stored in the startup message. let get_param = |key| params.get(key).ok_or(MissingKey(key)); let user = get_param("user")?; - let dbname = get_param("database")?; // Project name might be passed via PG's command-line options. let project_option = params.options_raw().and_then(|mut options| { @@ -100,7 +98,6 @@ impl<'a> ClientCredentials<'a> { info!( user = user, - dbname = dbname, project = project.as_deref(), use_cleartext_password_flow = use_cleartext_password_flow, "credentials" @@ -108,7 +105,6 @@ impl<'a> ClientCredentials<'a> { Ok(Self { user, - dbname, project, use_cleartext_password_flow, }) @@ -131,25 +127,27 @@ mod tests { use ClientCredsParseError::*; #[test] - #[ignore = "TODO: fix how database is handled"] fn parse_bare_minimum() -> anyhow::Result<()> { // According to postgresql, only `user` should be required. let options = StartupMessageParams::new([("user", "john_doe")]); - // TODO: check that `creds.dbname` is None. let creds = ClientCredentials::parse(&options, None, None, false)?; assert_eq!(creds.user, "john_doe"); + assert_eq!(creds.project, None); Ok(()) } #[test] - fn parse_missing_project() -> anyhow::Result<()> { - let options = StartupMessageParams::new([("user", "john_doe"), ("database", "world")]); + fn parse_excessive() -> anyhow::Result<()> { + let options = StartupMessageParams::new([ + ("user", "john_doe"), + ("database", "world"), // should be ignored + ("foo", "bar"), // should be ignored + ]); let creds = ClientCredentials::parse(&options, None, None, false)?; assert_eq!(creds.user, "john_doe"); - assert_eq!(creds.dbname, "world"); assert_eq!(creds.project, None); Ok(()) @@ -157,14 +155,13 @@ mod tests { #[test] fn parse_project_from_sni() -> anyhow::Result<()> { - let options = StartupMessageParams::new([("user", "john_doe"), ("database", "world")]); + let options = StartupMessageParams::new([("user", "john_doe")]); let sni = Some("foo.localhost"); let common_name = Some("localhost"); let creds = ClientCredentials::parse(&options, sni, common_name, false)?; assert_eq!(creds.user, "john_doe"); - assert_eq!(creds.dbname, "world"); assert_eq!(creds.project.as_deref(), Some("foo")); Ok(()) @@ -174,13 +171,11 @@ mod tests { fn parse_project_from_options() -> anyhow::Result<()> { let options = StartupMessageParams::new([ ("user", "john_doe"), - ("database", "world"), ("options", "-ckey=1 project=bar -c geqo=off"), ]); let creds = ClientCredentials::parse(&options, None, None, false)?; assert_eq!(creds.user, "john_doe"); - assert_eq!(creds.dbname, "world"); assert_eq!(creds.project.as_deref(), Some("bar")); Ok(()) @@ -188,18 +183,13 @@ mod tests { #[test] fn parse_projects_identical() -> anyhow::Result<()> { - let options = StartupMessageParams::new([ - ("user", "john_doe"), - ("database", "world"), - ("options", "project=baz"), - ]); + let options = StartupMessageParams::new([("user", "john_doe"), ("options", "project=baz")]); let sni = Some("baz.localhost"); let common_name = Some("localhost"); let creds = ClientCredentials::parse(&options, sni, common_name, false)?; assert_eq!(creds.user, "john_doe"); - assert_eq!(creds.dbname, "world"); assert_eq!(creds.project.as_deref(), Some("baz")); Ok(()) @@ -207,11 +197,8 @@ mod tests { #[test] fn parse_projects_different() { - let options = StartupMessageParams::new([ - ("user", "john_doe"), - ("database", "world"), - ("options", "project=first"), - ]); + let options = + StartupMessageParams::new([("user", "john_doe"), ("options", "project=first")]); let sni = Some("second.localhost"); let common_name = Some("localhost"); @@ -229,7 +216,7 @@ mod tests { #[test] fn parse_inconsistent_sni() { - let options = StartupMessageParams::new([("user", "john_doe"), ("database", "world")]); + let options = StartupMessageParams::new([("user", "john_doe")]); let sni = Some("project.localhost"); let common_name = Some("example.com"); diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index 0c0cbcde20..3f5eb3caff 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -65,6 +65,18 @@ impl ConnCfg { /// Apply startup message params to the connection config. pub fn set_startup_params(&mut self, params: &StartupMessageParams) { + // Only set `user` if it's not present in the config. + // Link auth flow takes username from the console's response. + if let (None, Some(user)) = (self.get_user(), params.get("user")) { + self.user(user); + } + + // Only set `dbname` if it's not present in the config. + // Link auth flow takes dbname from the console's response. + if let (None, Some(dbname)) = (self.get_dbname(), params.get("database")) { + self.dbname(dbname); + } + if let Some(options) = params.options_raw() { // We must drop all proxy-specific parameters. #[allow(unstable_name_collisions)] diff --git a/proxy/src/console/provider/mock.rs b/proxy/src/console/provider/mock.rs index 301c3be516..eaac9c06d9 100644 --- a/proxy/src/console/provider/mock.rs +++ b/proxy/src/console/provider/mock.rs @@ -82,16 +82,11 @@ impl Api { .await } - async fn do_wake_compute( - &self, - creds: &ClientCredentials<'_>, - ) -> Result { + async fn do_wake_compute(&self) -> Result { let mut config = compute::ConnCfg::new(); config .host(self.endpoint.host_str().unwrap_or("localhost")) - .port(self.endpoint.port().unwrap_or(5432)) - .dbname(creds.dbname) - .user(creds.user); + .port(self.endpoint.port().unwrap_or(5432)); let node = NodeInfo { config, @@ -117,9 +112,9 @@ impl super::Api for Api { async fn wake_compute( &self, _extra: &ConsoleReqExtra<'_>, - creds: &ClientCredentials<'_>, + _creds: &ClientCredentials<'_>, ) -> Result { - self.do_wake_compute(creds) + self.do_wake_compute() .map_ok(CachedNodeInfo::new_uncached) .await } diff --git a/proxy/src/console/provider/neon.rs b/proxy/src/console/provider/neon.rs index 00d3ca8352..4eca025d2d 100644 --- a/proxy/src/console/provider/neon.rs +++ b/proxy/src/console/provider/neon.rs @@ -97,12 +97,11 @@ impl Api { Some(x) => x, }; + // Don't set anything but host and port! This config will be cached. + // We'll set username and such later using the startup message. + // TODO: add more type safety (in progress). let mut config = compute::ConnCfg::new(); - config - .host(host) - .port(port) - .dbname(creds.dbname) - .user(creds.user); + config.host(host).port(port); let node = NodeInfo { config, diff --git a/proxy/src/http/websocket.rs b/proxy/src/http/websocket.rs index bedded7567..d4235c2c38 100644 --- a/proxy/src/http/websocket.rs +++ b/proxy/src/http/websocket.rs @@ -1,161 +1,136 @@ +use crate::{ + cancellation::CancelMap, config::ProxyConfig, error::io_error, proxy::handle_ws_client, +}; use bytes::{Buf, Bytes}; use futures::{Sink, Stream, StreamExt}; -use hyper::server::accept; -use hyper::server::conn::AddrIncoming; -use hyper::upgrade::Upgraded; -use hyper::{Body, Request, Response, StatusCode}; -use hyper_tungstenite::{tungstenite, WebSocketStream}; -use hyper_tungstenite::{tungstenite::Message, HyperWebsocket}; +use hyper::{ + server::{accept, conn::AddrIncoming}, + upgrade::Upgraded, + Body, Request, Response, StatusCode, +}; +use hyper_tungstenite::{tungstenite::Message, HyperWebsocket, WebSocketStream}; use pin_project_lite::pin_project; -use tokio::net::TcpListener; - -use std::convert::Infallible; -use std::future::ready; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{ready, Context, Poll}; +use std::{ + convert::Infallible, + future::ready, + pin::Pin, + sync::Arc, + task::{ready, Context, Poll}, +}; use tls_listener::TlsListener; - -use tokio::io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; - +use tokio::{ + io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}, + net::TcpListener, +}; use tracing::{error, info, info_span, warn, Instrument}; use utils::http::{error::ApiError, json::json_response}; -use crate::cancellation::CancelMap; -use crate::config::ProxyConfig; -use crate::proxy::handle_ws_client; +// TODO: use `std::sync::Exclusive` once it's stabilized. +// Tracking issue: https://github.com/rust-lang/rust/issues/98407. +use sync_wrapper::SyncWrapper; pin_project! { - /// This is a wrapper around a WebSocketStream that implements AsyncRead and AsyncWrite. - pub struct WebSocketRW { + /// This is a wrapper around a [`WebSocketStream`] that + /// implements [`AsyncRead`] and [`AsyncWrite`]. + pub struct WebSocketRw { #[pin] - stream: WebSocketStream, - chunk: Option, + stream: SyncWrapper>, + bytes: Bytes, } } -// FIXME: explain why this is safe or try to remove `unsafe impl`. -unsafe impl Sync for WebSocketRW {} - -impl WebSocketRW { +impl WebSocketRw { pub fn new(stream: WebSocketStream) -> Self { Self { - stream, - chunk: None, - } - } - - fn has_chunk(&self) -> bool { - if let Some(ref chunk) = self.chunk { - chunk.remaining() > 0 - } else { - false + stream: stream.into(), + bytes: Bytes::new(), } } } -fn ws_err_into(e: tungstenite::Error) -> io::Error { - io::Error::new(io::ErrorKind::Other, e.to_string()) -} - -impl AsyncWrite for WebSocketRW { +impl AsyncWrite for WebSocketRw { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], - ) -> Poll> { - let mut this = self.project(); - match this.stream.as_mut().poll_ready(cx) { - Poll::Ready(Ok(())) => { - if let Err(e) = this - .stream - .as_mut() - .start_send(Message::Binary(buf.to_vec())) - { - Poll::Ready(Err(ws_err_into(e))) - } else { - Poll::Ready(Ok(buf.len())) - } - } - Poll::Ready(Err(e)) => Poll::Ready(Err(ws_err_into(e))), - Poll::Pending => { - cx.waker().wake_by_ref(); - Poll::Pending - } + ) -> Poll> { + let mut stream = self.project().stream.get_pin_mut(); + + ready!(stream.as_mut().poll_ready(cx).map_err(io_error))?; + match stream.as_mut().start_send(Message::Binary(buf.into())) { + Ok(()) => Poll::Ready(Ok(buf.len())), + Err(e) => Poll::Ready(Err(io_error(e))), } } - fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.project().stream.poll_flush(cx).map_err(ws_err_into) + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let stream = self.project().stream.get_pin_mut(); + stream.poll_flush(cx).map_err(io_error) } - fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.project().stream.poll_close(cx).map_err(ws_err_into) + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let stream = self.project().stream.get_pin_mut(); + stream.poll_close(cx).map_err(io_error) } } -impl AsyncRead for WebSocketRW { +impl AsyncRead for WebSocketRw { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { - if buf.remaining() == 0 { - return Poll::Ready(Ok(())); + if buf.remaining() > 0 { + let bytes = ready!(self.as_mut().poll_fill_buf(cx))?; + let len = std::cmp::min(bytes.len(), buf.remaining()); + buf.put_slice(&bytes[..len]); + self.consume(len); } - let inner_buf = match ready!(self.as_mut().poll_fill_buf(cx)) { - Ok(buf) => buf, - Err(err) => return Poll::Ready(Err(err)), - }; - let len = std::cmp::min(inner_buf.len(), buf.remaining()); - buf.put_slice(&inner_buf[..len]); - - self.consume(len); Poll::Ready(Ok(())) } } -impl AsyncBufRead for WebSocketRW { - fn poll_fill_buf(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { +impl AsyncBufRead for WebSocketRw { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // Please refer to poll_fill_buf's documentation. + const EOF: Poll> = Poll::Ready(Ok(&[])); + + let mut this = self.project(); loop { - if self.as_mut().has_chunk() { - let buf = self.project().chunk.as_ref().unwrap().chunk(); - return Poll::Ready(Ok(buf)); - } else { - match ready!(self.as_mut().project().stream.poll_next(cx)) { - Some(Ok(message)) => match message { - Message::Text(_) => {} - Message::Binary(chunk) => { - *self.as_mut().project().chunk = Some(Bytes::from(chunk)); - } - Message::Ping(_) => { - // No need to send a reply: tungstenite takes care of this for you. - } - Message::Pong(_) => {} - Message::Close(_) => { - // No need to send a reply: tungstenite takes care of this for you. - return Poll::Ready(Ok(&[])); - } - Message::Frame(_) => { - unreachable!(); - } - }, - Some(Err(err)) => return Poll::Ready(Err(ws_err_into(err))), - None => return Poll::Ready(Ok(&[])), - } + if !this.bytes.chunk().is_empty() { + let chunk = (*this.bytes).chunk(); + return Poll::Ready(Ok(chunk)); + } + + let res = ready!(this.stream.as_mut().get_pin_mut().poll_next(cx)); + match res.transpose().map_err(io_error)? { + Some(message) => match message { + Message::Ping(_) => {} + Message::Pong(_) => {} + Message::Text(text) => { + // We expect to see only binary messages. + let error = "unexpected text message in the websocket"; + warn!(length = text.len(), error); + return Poll::Ready(Err(io_error(error))); + } + Message::Frame(_) => { + // This case is impossible according to Frame's doc. + panic!("unexpected raw frame in the websocket"); + } + Message::Binary(chunk) => { + assert!(this.bytes.is_empty()); + *this.bytes = Bytes::from(chunk); + } + Message::Close(_) => return EOF, + }, + None => return EOF, } } } - fn consume(self: Pin<&mut Self>, amt: usize) { - if amt > 0 { - self.project() - .chunk - .as_mut() - .expect("No chunk present") - .advance(amt); - } + fn consume(self: Pin<&mut Self>, amount: usize) { + self.project().bytes.advance(amount); } } @@ -171,7 +146,7 @@ async fn serve_websocket( config, cancel_map, session_id, - WebSocketRW::new(websocket), + WebSocketRw::new(websocket), hostname, ) .await?; @@ -199,7 +174,7 @@ async fn ws_handler( tokio::spawn(async move { if let Err(e) = serve_websocket(websocket, config, &cancel_map, session_id, host).await { - error!("error in websocket connection: {:?}", e); + error!("error in websocket connection: {e:?}"); } }); @@ -227,11 +202,12 @@ pub async fn task_main( } }; - let addr_incoming = AddrIncoming::from_listener(ws_listener)?; + let mut addr_incoming = AddrIncoming::from_listener(ws_listener)?; + let _ = addr_incoming.set_nodelay(true); let tls_listener = TlsListener::new(tls_acceptor, addr_incoming).filter(|conn| { if let Err(err) = conn { - error!("failed to accept TLS connection for websockets: {:?}", err); + error!("failed to accept TLS connection for websockets: {err:?}"); ready(false) } else { ready(true) diff --git a/proxy/src/main.rs b/proxy/src/main.rs index c96ca2a171..8812f77b62 100644 --- a/proxy/src/main.rs +++ b/proxy/src/main.rs @@ -41,12 +41,8 @@ async fn flatten_err( #[tokio::main] async fn main() -> anyhow::Result<()> { - tracing_subscriber::fmt() - .with_ansi(atty::is(atty::Stream::Stdout)) - .with_target(false) - .init(); - - // initialize sentry if SENTRY_DSN is provided + // First, initialize logging and troubleshooting subsystems. + init_tracing(); let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]); info!("Version: {GIT_VERSION}"); @@ -111,6 +107,21 @@ async fn main() -> anyhow::Result<()> { Ok(()) } +/// Tracing is used for logging and telemetry. +fn init_tracing() { + tracing_subscriber::fmt() + .with_env_filter({ + // This filter will examine the `RUST_LOG` env variable. + use tracing_subscriber::filter::{EnvFilter, LevelFilter}; + EnvFilter::builder() + .with_default_directive(LevelFilter::INFO.into()) + .from_env_lossy() + }) + .with_ansi(atty::is(atty::Stream::Stdout)) + .with_target(false) + .init(); +} + /// ProxyConfig is created at proxy startup, and lives forever. fn build_config(args: &clap::ArgMatches) -> anyhow::Result<&'static ProxyConfig> { let tls_config = match ( diff --git a/scripts/sk_cleanup_tenants/readme.md b/scripts/sk_cleanup_tenants/readme.md new file mode 100644 index 0000000000..f1bb2d540e --- /dev/null +++ b/scripts/sk_cleanup_tenants/readme.md @@ -0,0 +1,55 @@ +# Cleanup script for safekeeper + +This script can be used to remove tenant directories on safekeepers for projects which do not longer exist (deleted in console). + +To run this script you need to upload it to safekeeper (i.e. with SSH), and run it with python3. Ansible can be used to run this script on multiple safekeepers. + +NOTE: Console queries to check that project is deleted are slow and inefficient. +If you want to run this script on safekeeper with many tenants, consider +making PR to console repo to make projects search by tenant_id faster. + +## How to run on a single node + +``` +zsh nsh safekeeper-0.us-east-2.aws.neon.build + +ls /storage/safekeeper/data/ | grep -v safekeeper > tenants.txt + +mkdir -p /storage/neon-trash/2023-01-01--cleanup + + export CONSOLE_API_TOKEN= +python3 script.py --trash-dir /storage/neon-trash/2023-01-01--cleanup --safekeeper-id $(cat /storage/safekeeper/data/safekeeper.id) --safekeeper-host $HOSTNAME --dry-run + +cat tenants.txt | python3 script.py --trash-dir /storage/neon-trash/2023-01-01--cleanup --safekeeper-id $(cat /storage/safekeeper/data/safekeeper.id) --safekeeper-host $HOSTNAME --dry-run + +cat tenants.txt | python3 script.py --trash-dir /storage/neon-trash/2023-01-01--cleanup --safekeeper-id $(cat /storage/safekeeper/data/safekeeper.id) --safekeeper-host $HOSTNAME |& tee logs.txt +``` + +## How to use ansible (staging) + +``` +cd ~/neon/.github/ansible + +export AWS_DEFAULT_PROFILE=dev + +ansible-playbook -i staging.us-east-2.hosts.yaml -e @ssm_config ../../scripts/sk_cleanup_tenants/remote.yaml + +# add --extra-vars "api_token=" to set console api token +``` + +## How to use ansible (prod) + +- Change `endpoint` in `script.py` to "https://console.neon.tech/api" + +``` +cd ~/neon/.github/ansible + +export AWS_DEFAULT_PROFILE=prod + +ansible-playbook -i prod.us-east-2.hosts.yaml -e @ssm_config ../../scripts/sk_cleanup_tenants/remote.yaml + +# add --extra-vars "api_token=" to set console api token +``` + + +> Heavily inspired with script for pageserver cleanup: https://gist.github.com/problame/bafb6ca6334f0145757238e61380c3f1/9bef1845a8291ebfa1f3a51eb79c01d12498b2b5 \ No newline at end of file diff --git a/scripts/sk_cleanup_tenants/remote.yaml b/scripts/sk_cleanup_tenants/remote.yaml new file mode 100644 index 0000000000..c7eeb8516c --- /dev/null +++ b/scripts/sk_cleanup_tenants/remote.yaml @@ -0,0 +1,80 @@ +- name: Test safekeepers + hosts: safekeepers + gather_facts: False + remote_user: "{{ remote_user }}" + + vars: + script_dir: /storage/ansible_sk_cleanup + tenants_file: "{{ script_dir }}/tenants.txt" + trash_dir: /storage/neon-trash/2023-01-01--changeme + + tasks: + + - name: create script directory + file: + path: "{{ script_dir }}" + state: directory + mode: 0755 + tags: + - safekeeper + + - name: create trash dir + file: + path: "{{ trash_dir }}" + state: directory + mode: 0755 + tags: + - safekeeper + + - name: collect all tenant_ids to tenants.txt + shell: + cmd: ls /storage/safekeeper/data/ | grep -v safekeeper > {{ tenants_file }} + tags: + - safekeeper + + - name: count tenants + shell: + cmd: wc -l {{ tenants_file }} + register: tenants_count + tags: + - safekeeper + + - debug: msg="{{ tenants_count.stdout }}" + + - name: fetch safekeeper_id + shell: + cmd: cat /storage/safekeeper/data/safekeeper.id + register: safekeeper_id + tags: + - safekeeper + + - debug: msg="{{ safekeeper_id.stdout }}" + + - name: copy script.py to safekeeper + copy: + src: script.py + dest: "{{ script_dir }}" + mode: 0755 + tags: + - safekeeper + + - name: Run an async task + shell: + chdir: "{{ script_dir }}" + cmd: "cat tenants.txt | python3 script.py --trash-dir {{ trash_dir }} --safekeeper-id $(cat /storage/safekeeper/data/safekeeper.id) --safekeeper-host $HOSTNAME |& cat > {{ script_dir }}/run-`date +%Y-%m-%d-%H.%M.%S`.log" + args: + executable: /bin/bash + environment: + CONSOLE_API_TOKEN: "{{ api_token }}" + async: 30000 + poll: 0 + register: bg_async_task + + - name: Check on an async task + async_status: + jid: "{{ bg_async_task.ansible_job_id }}" + become: true + register: job_result + until: job_result.finished + retries: 3000 + delay: 10 diff --git a/scripts/sk_cleanup_tenants/script.py b/scripts/sk_cleanup_tenants/script.py new file mode 100644 index 0000000000..4d010d85ea --- /dev/null +++ b/scripts/sk_cleanup_tenants/script.py @@ -0,0 +1,133 @@ +import argparse +import logging +import os +import shutil +import sys +from pathlib import Path + +import requests + +level = logging.INFO +logging.basicConfig( + format="%(asctime)s,%(msecs)03d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s", + datefmt="%Y-%m-%d:%H:%M:%S", + level=level, +) + +parser = argparse.ArgumentParser() +parser.add_argument("--trash-dir", required=True, type=Path) +parser.add_argument("--dry-run", action="store_true") +parser.add_argument("--safekeeper-id", required=True, type=int) +parser.add_argument("--safekeeper-host", required=True, type=str) +args = parser.parse_args() + +access_key = os.getenv("CONSOLE_API_TOKEN") +endpoint: str = "https://console.stage.neon.tech/api" + +trash_dir: Path = args.trash_dir +dry_run: bool = args.dry_run +logging.info(f"dry_run={dry_run}") +sk_id: int = args.safekeeper_id +sk_host: str = args.safekeeper_host + +assert trash_dir.is_dir() + +### + + +def console_get(rel_url): + r = requests.get( + f"{endpoint}{rel_url}", + headers={ + "Authorization": f"Bearer {access_key}", + "Content-Type": "application/json", + "Accept": "application/json", + }, + ) + r.raise_for_status() + return r + + +def tenant_is_deleted_in_console(tenant_id): + r = console_get(f"/v1/admin/projects?search={tenant_id}&show_deleted=true") + r = r.json() + results = r["data"] + assert len(results) == 1, f"unexpected results len: {results}" + r = results[0] + assert r["tenant"] == tenant_id, f"tenant id doesn't match: {r}" + assert r["safekeepers"] is not None, f"safekeepers is None: {r}" + assert any(sk["id"] == sk_id for sk in r["safekeepers"]), f"safekeeper id not found: {r}" + assert "deleted" in r, f"{r}" + return r["deleted"] is True + + +def call_delete_tenant_api(tenant_id): + r = requests.delete(f"http://{sk_host}:7676/v1/tenant/{tenant_id}") + r.raise_for_status() + return r + + +def cleanup_tenant(tenant_id): + + tenant_dir = Path(f"/storage/safekeeper/data/{tenant_id}") + + if not tenant_dir.exists(): + logging.info("tenant directory doesn't exist, assuming it has been cleaned already") + return + + if not tenant_is_deleted_in_console(tenant_id): + logging.info("tenant is not deleted in console, skipping") + return + + logging.info("assertions passed") + + if dry_run: + return + + logging.info("deleting tenant") + + tenant_dir_in_trash = trash_dir / tenant_dir.relative_to("/") + tenant_dir_in_trash.parent.mkdir(parents=True, exist_ok=True) + + assert not tenant_dir_in_trash.exists(), f"{tenant_dir_in_trash}" + assert tenant_dir_in_trash.parent.exists(), f"{tenant_dir_in_trash}" + # double-check + assert tenant_dir.exists(), f"{tenant_dir}" + assert tenant_dir.is_dir(), f"{tenant_dir}" + + logging.info(f"copying {tenant_dir} to {tenant_dir_in_trash}") + shutil.copytree(src=tenant_dir, dst=tenant_dir_in_trash, symlinks=False, dirs_exist_ok=False) + + logging.info(f"deleting {tenant_dir}") + call_delete_tenant_api(tenant_id) + + logging.info("tenant is now deleted, checking that it's gone") + assert not tenant_dir.exists(), f"{tenant_dir}" + + +if os.path.exists("script.pid"): + logging.info( + f"script is already running, with pid={Path('script.pid').read_text()}. Terminate it first." + ) + exit(1) + +with open("script.pid", "w", encoding="utf-8") as f: + f.write(str(os.getpid())) + +logging.info(f"started script.py, pid={os.getpid()}") + +for line in sys.stdin: + tenant_id = line.strip() + try: + logging.info(f"start tenant {tenant_id}") + cleanup_tenant(tenant_id) + logging.info(f"done tenant {tenant_id}") + except KeyboardInterrupt: + print("KeyboardInterrupt exception is caught") + break + except: # noqa: E722 + logging.exception(f"failed to clean up tenant {tenant_id}") + +logging.info(f"finished script.py, pid={os.getpid()}") + +os.remove("script.pid")