From 70f646ffe2fe9829316f1ed02a5a1529bc296fd6 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Mon, 5 Feb 2024 09:34:03 +0200 Subject: [PATCH 01/72] More logging fixes (#6584) I was on-call this week, these would had made me understand more/faster of the system: - move stray attaching start logging inside the span it starts, add generation - log ancestor timeline_id or bootstrapping in the beginning of timeline creation --- pageserver/src/http/routes.rs | 6 +++++ pageserver/src/tenant.rs | 28 ++++++++++------------- pageserver/src/tenant/config.rs | 4 ++-- pageserver/src/tenant/mgr.rs | 7 ------ test_runner/regress/test_timeline_size.py | 2 +- 5 files changed, 21 insertions(+), 26 deletions(-) diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index 5735489742..b97e272c86 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -489,6 +489,12 @@ async fn timeline_create_handler( tenant.wait_to_become_active(ACTIVE_TENANT_TIMEOUT).await?; + if let Some(ancestor_id) = request_data.ancestor_timeline_id.as_ref() { + tracing::info!(%ancestor_id, "starting to branch"); + } else { + tracing::info!("bootstrapping"); + } + match tenant.create_timeline( new_timeline_id, request_data.ancestor_timeline_id.map(TimelineId::from), diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index 58af80238d..dd4f9107f9 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -205,7 +205,7 @@ impl AttachedTenantConf { match &location_conf.mode { LocationMode::Attached(attach_conf) => Ok(Self { tenant_conf: location_conf.tenant_conf, - location: attach_conf.clone(), + location: *attach_conf, }), LocationMode::Secondary(_) => { anyhow::bail!("Attempted to construct AttachedTenantConf from a LocationConf in secondary mode") @@ -625,6 +625,9 @@ impl Tenant { deletion_queue_client, } = resources; + let attach_mode = attached_conf.location.attach_mode; + let generation = attached_conf.location.generation; + let tenant = Arc::new(Tenant::new( TenantState::Attaching, conf, @@ -654,6 +657,12 @@ impl Tenant { "attach tenant", false, async move { + + info!( + ?attach_mode, + "Attaching tenant" + ); + let _gate_guard = attach_gate_guard; // Is this tenant being spawned as part of process startup? @@ -865,7 +874,7 @@ impl Tenant { Ok(()) } .instrument({ - let span = tracing::info_span!(parent: None, "attach", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug()); + let span = tracing::info_span!(parent: None, "attach", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), gen=?generation); span.follows_from(Span::current()); span }), @@ -2354,12 +2363,7 @@ impl Tenant { } pub(crate) fn get_attach_mode(&self) -> AttachmentMode { - self.tenant_conf - .read() - .unwrap() - .location - .attach_mode - .clone() + self.tenant_conf.read().unwrap().location.attach_mode } /// For API access: generate a LocationConfig equivalent to the one that would be used to @@ -3225,8 +3229,6 @@ impl Tenant { .context("branch initial metadata upload")?; } - info!("branched timeline {dst_id} from {src_id} at {start_lsn}"); - Ok(new_timeline) } @@ -3444,12 +3446,6 @@ impl Tenant { // All done! let timeline = raw_timeline.finish_creation()?; - info!( - "created root timeline {} timeline.lsn {}", - timeline_id, - timeline.get_last_record_lsn() - ); - Ok(timeline) } diff --git a/pageserver/src/tenant/config.rs b/pageserver/src/tenant/config.rs index 63bd56cf5f..563887088d 100644 --- a/pageserver/src/tenant/config.rs +++ b/pageserver/src/tenant/config.rs @@ -51,7 +51,7 @@ pub mod defaults { pub const DEFAULT_INGEST_BATCH_SIZE: u64 = 100; } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq, Eq)] pub(crate) enum AttachmentMode { /// Our generation is current as far as we know, and as far as we know we are the only attached /// pageserver. This is the "normal" attachment mode. @@ -66,7 +66,7 @@ pub(crate) enum AttachmentMode { Stale, } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq, Eq)] pub(crate) struct AttachedLocationConfig { pub(crate) generation: Generation, pub(crate) attach_mode: AttachmentMode, diff --git a/pageserver/src/tenant/mgr.rs b/pageserver/src/tenant/mgr.rs index 64fd709386..de0b636d47 100644 --- a/pageserver/src/tenant/mgr.rs +++ b/pageserver/src/tenant/mgr.rs @@ -607,13 +607,6 @@ pub(crate) fn tenant_spawn( "Cannot load tenant, ignore mark found at {tenant_ignore_mark:?}" ); - info!( - tenant_id = %tenant_shard_id.tenant_id, - shard_id = %tenant_shard_id.shard_slug(), - generation = ?location_conf.location.generation, - attach_mode = ?location_conf.location.attach_mode, - "Attaching tenant" - ); let tenant = match Tenant::spawn( conf, tenant_shard_id, diff --git a/test_runner/regress/test_timeline_size.py b/test_runner/regress/test_timeline_size.py index 303aabb58d..cd7203bba6 100644 --- a/test_runner/regress/test_timeline_size.py +++ b/test_runner/regress/test_timeline_size.py @@ -883,7 +883,7 @@ def test_ondemand_activation(neon_env_builder: NeonEnvBuilder): # Deletion itself won't complete due to our failpoint: Tenant::shutdown can't complete while calculating # logical size is paused in a failpoint. So instead we will use a log observation to check that # on-demand activation was triggered by the tenant deletion - log_match = f".*attach{{tenant_id={delete_tenant_id} shard_id=0000}}: Activating tenant \\(on-demand\\).*" + log_match = f".*attach{{tenant_id={delete_tenant_id} shard_id=0000 gen=[0-9a-f]+}}: Activating tenant \\(on-demand\\).*" def activated_on_demand(): assert env.pageserver.log_contains(log_match) is not None From df7bee7cfaba8f2129fd9ea88976da5d079684a5 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Sat, 3 Feb 2024 00:02:33 +0200 Subject: [PATCH 02/72] Fix compilation with recent glibc headers with close_range(2). MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit I was getting an error: /home/heikki/git-sandbox/neon//pgxn/neon_walredo/walredoproc.c:161:5: error: conflicting types for ‘close_range’; have ‘int(unsigned int, unsigned int, unsigned int)’ 161 | int close_range(unsigned int start_fd, unsigned int count, unsigned int flags) { | ^~~~~~~~~~~ In file included from /usr/include/x86_64-linux-gnu/bits/sigstksz.h:24, from /usr/include/signal.h:328, from /home/heikki/git-sandbox/neon//pgxn/neon_walredo/walredoproc.c:50: /usr/include/unistd.h:1208:12: note: previous declaration of ‘close_range’ with type ‘int(unsigned int, unsigned int, int)’ 1208 | extern int close_range (unsigned int __fd, unsigned int __max_fd, | ^~~~~~~~~~~ The discrepancy is in the 3rd argument. Apparently in the glibc wrapper it's signed. As a quick fix, rename our close_range() function, the one that calls syscall() directly, to avoid the clash with the glibc wrapper. In the long term, an autoconf test would be nice, and some equivalent on macOS, see issue #6580. --- pgxn/neon_walredo/walredoproc.c | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pgxn/neon_walredo/walredoproc.c b/pgxn/neon_walredo/walredoproc.c index 6ca0b2a274..1fdd3801c6 100644 --- a/pgxn/neon_walredo/walredoproc.c +++ b/pgxn/neon_walredo/walredoproc.c @@ -158,7 +158,10 @@ static XLogReaderState *reader_state; #include #include #include -int close_range(unsigned int start_fd, unsigned int count, unsigned int flags) { + +static int +close_range_syscall(unsigned int start_fd, unsigned int count, unsigned int flags) +{ return syscall(__NR_close_range, start_fd, count, flags); } @@ -172,7 +175,7 @@ enter_seccomp_mode(void) * wal records. See the comment in the Rust code that launches this process. */ int err; - if (err = close_range(3, ~0U, 0)) { + if (err = close_range_syscall(3, ~0U, 0)) { ereport(FATAL, (errcode(ERRCODE_SYSTEM_ERROR), errmsg("seccomp: could not close files >= fd 3"))); } From 56cf3604395125b9283ba643cfbb98efd926ff49 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Mon, 5 Feb 2024 10:53:37 +0100 Subject: [PATCH 03/72] Don't preserve temp files on creation errors of delta layers (#6612) There is currently no cleanup done after a delta layer creation error, so delta layers can accumulate. The problem gets worse as the operation gets retried and delta layers accumulate on the disk. Therefore, delete them from disk (if something has been written to disk). --- pageserver/src/tenant/storage_layer/delta_layer.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/pageserver/src/tenant/storage_layer/delta_layer.rs b/pageserver/src/tenant/storage_layer/delta_layer.rs index ec031d6089..2a51884c0b 100644 --- a/pageserver/src/tenant/storage_layer/delta_layer.rs +++ b/pageserver/src/tenant/storage_layer/delta_layer.rs @@ -609,7 +609,19 @@ impl DeltaLayerWriter { key_end: Key, timeline: &Arc, ) -> anyhow::Result { - self.inner.take().unwrap().finish(key_end, timeline).await + let inner = self.inner.take().unwrap(); + let temp_path = inner.path.clone(); + let result = inner.finish(key_end, timeline).await; + // The delta layer files can sometimes be really large. Clean them up. + if result.is_err() { + tracing::warn!( + "Cleaning up temporary delta file {temp_path} after error during writing" + ); + if let Err(e) = std::fs::remove_file(&temp_path) { + tracing::warn!("Error cleaning up temporary delta layer file {temp_path}: {e:?}") + } + } + result } } From 01c57ec547cb701f2253c8c445931644cc9f60b9 Mon Sep 17 00:00:00 2001 From: Abhijeet Patil Date: Mon, 5 Feb 2024 10:08:20 +0000 Subject: [PATCH 04/72] Removed Uploading of perf result to git repo 'zenith-perf-data' (#6590) ## Problem We were archiving the pref benchmarks to - neon DB - git repo `zenith-perf-data` As the pref batch ran in parallel when the uploading of results to zenith-perf-data` git repo resulted in merge conflicts. Which made the run flaky and as a side effect the build started failing . The problem is been expressed in https://github.com/neondatabase/neon/issues/5160 ## Summary of changes As the results were not used from the git repo it was redundant hence in this PR cleaning up the results uploading of of perf results to git repo The shell script `generate_and_push_perf_report.sh` was using a py script [git-upload](https://github.com/neondatabase/neon/compare/remove-perf-benchmark-git-upload?expand=1#diff-c6d938e7f060e487367d9dc8055245c82b51a73c1f97956111a495a8a86e9a33) and [scripts/generate_perf_report_page.py](https://github.com/neondatabase/neon/pull/6590/files#diff-81af2147e72d07e4cf8ee4395632596d805d6168ba75c71cab58db2659956ef8) which are not used anywhere else in repo hence also cleaning that up ## Checklist before requesting a review - [ ] I have performed a self-review of my code. - [ ] If it is a core feature, I have added thorough tests. - [ ] Do we need to implement analytics? if so did you add the relevant metrics to the dashboard? - [ ] If this PR requires public announcement, mark it with /release-notes label and add several sentences in this section. ## Checklist before merging - [ ] Do not forget to reformat the commit message to not include the above checklist --- scripts/generate_and_push_perf_report.sh | 14 -- scripts/generate_perf_report_page.py | 219 ----------------------- scripts/git-upload | 170 ------------------ 3 files changed, 403 deletions(-) delete mode 100755 scripts/generate_perf_report_page.py delete mode 100755 scripts/git-upload diff --git a/scripts/generate_and_push_perf_report.sh b/scripts/generate_and_push_perf_report.sh index 9e03302b0f..178c570b13 100755 --- a/scripts/generate_and_push_perf_report.sh +++ b/scripts/generate_and_push_perf_report.sh @@ -8,17 +8,3 @@ SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) echo "Uploading perf report to neon pg" # ingest per test results data into neon backed postgres running in staging to build grafana reports on that data DATABASE_URL="$PERF_TEST_RESULT_CONNSTR" poetry run python "$SCRIPT_DIR"/ingest_perf_test_result.py --ingest "$REPORT_FROM" - -# Activate poetry's venv. Needed because git upload does not run in a project dir (it uses tmp to store the repository) -# so the problem occurs because poetry cannot find pyproject.toml in temp dir created by git upload -# shellcheck source=/dev/null -. "$(poetry env info --path)"/bin/activate - -echo "Uploading perf result to zenith-perf-data" -scripts/git-upload \ - --repo=https://"$VIP_VAP_ACCESS_TOKEN"@github.com/neondatabase/zenith-perf-data.git \ - --message="add performance test result for $GITHUB_SHA neon revision" \ - --branch=master \ - copy "$REPORT_FROM" "data/$REPORT_TO" `# COPY FROM TO_RELATIVE`\ - --merge \ - --run-cmd "python $SCRIPT_DIR/generate_perf_report_page.py --input-dir data/$REPORT_TO --out reports/$REPORT_TO.html" diff --git a/scripts/generate_perf_report_page.py b/scripts/generate_perf_report_page.py deleted file mode 100755 index b5b49bb600..0000000000 --- a/scripts/generate_perf_report_page.py +++ /dev/null @@ -1,219 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import json -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple, cast - -from jinja2 import Template - -# skip 'input' columns. They are included in the header and just blow the table -EXCLUDE_COLUMNS = frozenset( - { - "scale", - "duration", - "number_of_clients", - "number_of_threads", - "init_start_timestamp", - "init_end_timestamp", - "run_start_timestamp", - "run_end_timestamp", - } -) - -KEY_EXCLUDE_FIELDS = frozenset( - { - "init_start_timestamp", - "init_end_timestamp", - "run_start_timestamp", - "run_end_timestamp", - } -) -NEGATIVE_COLOR = "negative" -POSITIVE_COLOR = "positive" -EPS = 1e-6 - - -@dataclass -class SuitRun: - revision: str - values: Dict[str, Any] - - -@dataclass -class SuitRuns: - platform: str - suit: str - common_columns: List[Tuple[str, str]] - value_columns: List[str] - runs: List[SuitRun] - - -@dataclass -class RowValue: - value: str - color: str - ratio: str - - -def get_columns(values: List[Dict[Any, Any]]) -> Tuple[List[Tuple[str, str]], List[str]]: - value_columns = [] - common_columns = [] - for item in values: - if item["name"] in KEY_EXCLUDE_FIELDS: - continue - if item["report"] != "test_param": - value_columns.append(cast(str, item["name"])) - else: - common_columns.append((cast(str, item["name"]), cast(str, item["value"]))) - value_columns.sort() - common_columns.sort(key=lambda x: x[0]) # sort by name - return common_columns, value_columns - - -def format_ratio(ratio: float, report: str) -> Tuple[str, str]: - color = "" - sign = "+" if ratio > 0 else "" - if abs(ratio) < 0.05: - return f" ({sign}{ratio:.2f})", color - - if report not in {"test_param", "higher_is_better", "lower_is_better"}: - raise ValueError(f"Unknown report type: {report}") - - if report == "test_param": - return f"{ratio:.2f}", color - - if ratio > 0: - if report == "higher_is_better": - color = POSITIVE_COLOR - elif report == "lower_is_better": - color = NEGATIVE_COLOR - elif ratio < 0: - if report == "higher_is_better": - color = NEGATIVE_COLOR - elif report == "lower_is_better": - color = POSITIVE_COLOR - - return f" ({sign}{ratio:.2f})", color - - -def extract_value(name: str, suit_run: SuitRun) -> Optional[Dict[str, Any]]: - for item in suit_run.values["data"]: - if item["name"] == name: - return cast(Dict[str, Any], item) - return None - - -def get_row_values( - columns: List[str], run_result: SuitRun, prev_result: Optional[SuitRun] -) -> List[RowValue]: - row_values = [] - for column in columns: - current_value = extract_value(column, run_result) - if current_value is None: - # should never happen - raise ValueError(f"{column} not found in {run_result.values}") - - value = current_value["value"] - if isinstance(value, float): - value = f"{value:.2f}" - - if prev_result is None: - row_values.append(RowValue(value, "", "")) - continue - - prev_value = extract_value(column, prev_result) - if prev_value is None: - # this might happen when new metric is added and there is no value for it in previous run - # let this be here, TODO add proper handling when this actually happens - raise ValueError(f"{column} not found in previous result") - # adding `EPS` to each term to avoid ZeroDivisionError when the denominator is zero - ratio = (float(value) + EPS) / (float(prev_value["value"]) + EPS) - 1 - ratio_display, color = format_ratio(ratio, current_value["report"]) - row_values.append(RowValue(value, color, ratio_display)) - return row_values - - -@dataclass -class SuiteRunTableRow: - revision: str - values: List[RowValue] - - -def prepare_rows_from_runs(value_columns: List[str], runs: List[SuitRun]) -> List[SuiteRunTableRow]: - rows = [] - prev_run = None - for run in runs: - rows.append( - SuiteRunTableRow( - revision=run.revision, values=get_row_values(value_columns, run, prev_run) - ) - ) - prev_run = run - - return rows - - -def main(args: argparse.Namespace) -> None: - input_dir = Path(args.input_dir) - grouped_runs: Dict[str, SuitRuns] = {} - # we have files in form: _.json - # fill them in the hashmap so we have grouped items for the - # same run configuration (scale, duration etc.) ordered by counter. - for item in sorted(input_dir.iterdir(), key=lambda x: int(x.name.split("_")[0])): - run_data = json.loads(item.read_text()) - revision = run_data["revision"] - - for suit_result in run_data["result"]: - key = "{}{}".format(run_data["platform"], suit_result["suit"]) - # pack total duration as a synthetic value - total_duration = suit_result["total_duration"] - suit_result["data"].append( - { - "name": "total_duration", - "value": total_duration, - "unit": "s", - "report": "lower_is_better", - } - ) - common_columns, value_columns = get_columns(suit_result["data"]) - - grouped_runs.setdefault( - key, - SuitRuns( - platform=run_data["platform"], - suit=suit_result["suit"], - common_columns=common_columns, - value_columns=value_columns, - runs=[], - ), - ) - - grouped_runs[key].runs.append(SuitRun(revision=revision, values=suit_result)) - context = {} - for result in grouped_runs.values(): - suit = result.suit - context[suit] = { - "common_columns": result.common_columns, - "value_columns": result.value_columns, - "platform": result.platform, - # reverse the order so newest results are on top of the table - "rows": reversed(prepare_rows_from_runs(result.value_columns, result.runs)), - } - - template = Template((Path(__file__).parent / "perf_report_template.html").read_text()) - - Path(args.out).write_text(template.render(context=context)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--input-dir", - dest="input_dir", - required=True, - help="Directory with jsons generated by the test suite", - ) - parser.add_argument("--out", required=True, help="Output html file path") - args = parser.parse_args() - main(args) diff --git a/scripts/git-upload b/scripts/git-upload deleted file mode 100755 index d56c0f8e94..0000000000 --- a/scripts/git-upload +++ /dev/null @@ -1,170 +0,0 @@ -#!/usr/bin/env python3 - -import argparse -import os -import shlex -import shutil -import subprocess -import sys -import textwrap -from contextlib import contextmanager -from distutils.dir_util import copy_tree -from pathlib import Path -from tempfile import TemporaryDirectory -from typing import Optional - - -def absolute_path(path): - return Path(path).resolve() - - -def relative_path(path): - path = Path(path) - if path.is_absolute(): - raise Exception(f'path `{path}` must be relative!') - return path - - -@contextmanager -def chdir(cwd: Path): - old = os.getcwd() - os.chdir(cwd) - try: - yield cwd - finally: - os.chdir(old) - - -def run(cmd, *args, **kwargs): - print('$', ' '.join(cmd)) - subprocess.check_call(cmd, *args, **kwargs) - - -class GitRepo: - def __init__(self, url, branch: Optional[str] = None): - self.url = url - self.cwd = TemporaryDirectory() - self.branch = branch - - args = [ - 'git', - 'clone', - '--single-branch', - ] - if self.branch: - args.extend(['--branch', self.branch]) - - subprocess.check_call([ - *args, - str(url), - self.cwd.name, - ]) - - def is_dirty(self): - res = subprocess.check_output(['git', 'status', '--porcelain'], text=True).strip() - return bool(res) - - def update(self, message, action, branch=None): - with chdir(self.cwd.name): - if not branch: - cmd = ['git', 'branch', '--show-current'] - branch = subprocess.check_output(cmd, text=True).strip() - - # Run action in repo's directory - action() - - run(['git', 'add', '.']) - - if not self.is_dirty(): - print('No changes detected, quitting') - return - - git_with_user = [ - 'git', - '-c', - 'user.name=vipvap', - '-c', - 'user.email=vipvap@zenith.tech', - ] - run(git_with_user + [ - 'commit', - '--author="vipvap "', - f'--message={message}', - ]) - - for _ in range(5): - try: - run(['git', 'fetch', 'origin', branch]) - run(git_with_user + ['rebase', f'origin/{branch}']) - run(['git', 'push', 'origin', branch]) - return - - except subprocess.CalledProcessError as e: - print(f'failed to update branch `{branch}`: {e}', file=sys.stderr) - - raise Exception(f'failed to update branch `{branch}`') - - -def do_copy(args): - src = args.src - dst = args.dst - - if args.forbid_overwrite and dst.exists(): - raise FileExistsError(f"File exists: '{dst}'") - - if src.is_dir(): - if not args.merge: - shutil.rmtree(dst, ignore_errors=True) - # distutils is deprecated, but this is a temporary workaround before python version bump - # here we need dir_exists_ok=True from shutil.copytree which is available in python 3.8+ - copy_tree(str(src), str(dst)) - else: - shutil.copy(src, dst) - - if args.run_cmd: - run(shlex.split(args.run_cmd)) - - -def main(): - parser = argparse.ArgumentParser(description='Git upload tool') - parser.add_argument('--repo', type=str, metavar='URL', required=True, help='git repo url') - parser.add_argument('--message', type=str, metavar='TEXT', help='commit message') - parser.add_argument('--branch', type=str, metavar='TEXT', help='target git repo branch') - - commands = parser.add_subparsers(title='commands', dest='subparser_name') - - p_copy = commands.add_parser( - 'copy', - help='copy file into the repo', - formatter_class=argparse.RawTextHelpFormatter, - ) - p_copy.add_argument('src', type=absolute_path, help='source path') - p_copy.add_argument('dst', type=relative_path, help='relative dest path') - p_copy.add_argument('--forbid-overwrite', action='store_true', help='do not allow overwrites') - p_copy.add_argument( - '--merge', - action='store_true', - help='when copying a directory do not delete existing data, but add new files') - p_copy.add_argument('--run-cmd', - help=textwrap.dedent('''\ - run arbitrary cmd on top of copied files, - example usage is static content generation - based on current repository state\ - ''')) - - args = parser.parse_args() - - commands = { - 'copy': do_copy, - } - - action = commands.get(args.subparser_name) - if action: - message = args.message or 'update' - GitRepo(args.repo, args.branch).update(message, lambda: action(args)) - else: - parser.print_usage() - - -if __name__ == '__main__': - main() From db89b13aaa45266227b89884490c11e10abb8054 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Mon, 5 Feb 2024 14:10:08 +0200 Subject: [PATCH 05/72] fix: use the shared constant download buffer size (#6620) Noticed that we had forgotten to use `remote_timeline_client.rs::BUFFER_SIZE` in one instance. --- pageserver/src/tenant/remote_timeline_client/download.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pageserver/src/tenant/remote_timeline_client/download.rs b/pageserver/src/tenant/remote_timeline_client/download.rs index 4309c683e2..b84b5ca33b 100644 --- a/pageserver/src/tenant/remote_timeline_client/download.rs +++ b/pageserver/src/tenant/remote_timeline_client/download.rs @@ -471,7 +471,7 @@ pub(crate) async fn download_initdb_tar_zst( Err(other) => Err(other)?, }; let mut download = tokio_util::io::StreamReader::new(download.download_stream); - let mut writer = tokio::io::BufWriter::with_capacity(8 * 1024, file); + let mut writer = tokio::io::BufWriter::with_capacity(super::BUFFER_SIZE, file); // TODO: this consumption of the response body should be subject to timeout + cancellation, but // not without thinking carefully about how to recover safely from cancelling a write to From 5e8deca26862f190e6f38b31ccea5f0a22c36c69 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Mon, 5 Feb 2024 14:49:35 +0200 Subject: [PATCH 06/72] metrics: remove broken tenants (#6586) Before tenant migration it made sense to leak broken tenants in the metrics until restart. Nowdays it makes less sense because on cancellations we set the tenant broken. The set metric still allows filterable alerting. Fixes: #6507 --- pageserver/src/tenant.rs | 45 +++++++++++------------ test_runner/fixtures/metrics.py | 2 +- test_runner/regress/test_tenant_detach.py | 39 +++++--------------- 3 files changed, 32 insertions(+), 54 deletions(-) diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index dd4f9107f9..b801347c06 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -67,7 +67,9 @@ use crate::deletion_queue::DeletionQueueError; use crate::import_datadir; use crate::is_uninit_mark; use crate::metrics::TENANT; -use crate::metrics::{remove_tenant_metrics, TENANT_STATE_METRIC, TENANT_SYNTHETIC_SIZE_METRIC}; +use crate::metrics::{ + remove_tenant_metrics, BROKEN_TENANTS_SET, TENANT_STATE_METRIC, TENANT_SYNTHETIC_SIZE_METRIC, +}; use crate::repository::GcResult; use crate::task_mgr; use crate::task_mgr::TaskKind; @@ -2637,9 +2639,16 @@ impl Tenant { let (state, mut rx) = watch::channel(state); tokio::spawn(async move { - // Strings for metric labels + // reflect tenant state in metrics: + // - global per tenant state: TENANT_STATE_METRIC + // - "set" of broken tenants: BROKEN_TENANTS_SET + // + // set of broken tenants should not have zero counts so that it remains accessible for + // alerting. + let tid = tenant_shard_id.to_string(); - let shard_id_str = format!("{}", tenant_shard_id.shard_slug()); + let shard_id = tenant_shard_id.shard_slug().to_string(); + let set_key = &[tid.as_str(), shard_id.as_str()][..]; fn inspect_state(state: &TenantState) -> ([&'static str; 1], bool) { ([state.into()], matches!(state, TenantState::Broken { .. })) @@ -2648,21 +2657,13 @@ impl Tenant { let mut tuple = inspect_state(&rx.borrow_and_update()); let is_broken = tuple.1; - let mut counted_broken = if !is_broken { - // the tenant might be ignored and reloaded, so first remove any previous set - // element. it most likely has already been scraped, as these are manual operations - // right now. most likely we will add it back very soon. - drop( - crate::metrics::BROKEN_TENANTS_SET.remove_label_values(&[&tid, &shard_id_str]), - ); - false - } else { + let mut counted_broken = if is_broken { // add the id to the set right away, there should not be any updates on the channel - // after - crate::metrics::BROKEN_TENANTS_SET - .with_label_values(&[&tid, &shard_id_str]) - .set(1); + // after before tenant is removed, if ever + BROKEN_TENANTS_SET.with_label_values(set_key).set(1); true + } else { + false }; loop { @@ -2671,10 +2672,9 @@ impl Tenant { current.inc(); if rx.changed().await.is_err() { - // tenant has been dropped; decrement the counter because a tenant with that - // state is no longer in tenant map, but allow any broken set item to exist - // still. + // tenant has been dropped current.dec(); + drop(BROKEN_TENANTS_SET.remove_label_values(set_key)); break; } @@ -2684,10 +2684,9 @@ impl Tenant { let is_broken = tuple.1; if is_broken && !counted_broken { counted_broken = true; - // insert the tenant_id (back) into the set - crate::metrics::BROKEN_TENANTS_SET - .with_label_values(&[&tid, &shard_id_str]) - .inc(); + // insert the tenant_id (back) into the set while avoiding needless counter + // access + BROKEN_TENANTS_SET.with_label_values(set_key).set(1); } } }); diff --git a/test_runner/fixtures/metrics.py b/test_runner/fixtures/metrics.py index 7c489bda67..ef41774289 100644 --- a/test_runner/fixtures/metrics.py +++ b/test_runner/fixtures/metrics.py @@ -96,5 +96,5 @@ PAGESERVER_PER_TENANT_METRICS: Tuple[str, ...] = ( "pageserver_evictions_total", "pageserver_evictions_with_low_residence_duration_total", *PAGESERVER_PER_TENANT_REMOTE_TIMELINE_CLIENT_METRICS, - # pageserver_broken_tenants_count is a leaked "metric" which is "cleared" on restart or reload + # "pageserver_broken_tenants_count" -- used only for broken ) diff --git a/test_runner/regress/test_tenant_detach.py b/test_runner/regress/test_tenant_detach.py index 8d5ef4e3c4..4752699abb 100644 --- a/test_runner/regress/test_tenant_detach.py +++ b/test_runner/regress/test_tenant_detach.py @@ -742,8 +742,6 @@ def ensure_test_data(data_id: int, data: str, endpoint: Endpoint): def test_metrics_while_ignoring_broken_tenant_and_reloading( neon_env_builder: NeonEnvBuilder, ): - neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.LOCAL_FS) - env = neon_env_builder.init_start() client = env.pageserver.http_client() @@ -761,56 +759,37 @@ def test_metrics_while_ignoring_broken_tenant_and_reloading( client.tenant_break(env.initial_tenant) - found_broken = False - active, broken, broken_set = ([], [], []) - for _ in range(10): + def found_broken(): m = client.get_metrics() active = m.query_all("pageserver_tenant_states_count", {"state": "Active"}) broken = m.query_all("pageserver_tenant_states_count", {"state": "Broken"}) broken_set = m.query_all( "pageserver_broken_tenants_count", {"tenant_id": str(env.initial_tenant)} ) - found_broken = only_int(active) == 0 and only_int(broken) == 1 and only_int(broken_set) == 1 + assert only_int(active) == 0 and only_int(broken) == 1 and only_int(broken_set) == 1 - if found_broken: - break - log.info(f"active: {active}, broken: {broken}, broken_set: {broken_set}") - time.sleep(0.5) - assert ( - found_broken - ), f"tenant shows up as broken; active={active}, broken={broken}, broken_set={broken_set}" + wait_until(10, 0.5, found_broken) client.tenant_ignore(env.initial_tenant) - found_broken = False - broken, broken_set = ([], []) - for _ in range(10): + def found_cleaned_up(): m = client.get_metrics() broken = m.query_all("pageserver_tenant_states_count", {"state": "Broken"}) broken_set = m.query_all( "pageserver_broken_tenants_count", {"tenant_id": str(env.initial_tenant)} ) - found_broken = only_int(broken) == 0 and only_int(broken_set) == 1 + assert only_int(broken) == 0 and len(broken_set) == 0 - if found_broken: - break - time.sleep(0.5) - assert found_broken, f"broken should still be in set, but it is not in the tenant state count: broken={broken}, broken_set={broken_set}" + wait_until(10, 0.5, found_cleaned_up) env.pageserver.tenant_load(env.initial_tenant) - found_active = False - active, broken_set = ([], []) - for _ in range(10): + def found_active(): m = client.get_metrics() active = m.query_all("pageserver_tenant_states_count", {"state": "Active"}) broken_set = m.query_all( "pageserver_broken_tenants_count", {"tenant_id": str(env.initial_tenant)} ) - found_active = only_int(active) == 1 and len(broken_set) == 0 + assert only_int(active) == 1 and len(broken_set) == 0 - if found_active: - break - time.sleep(0.5) - - assert found_active, f"reloaded tenant should be active, and broken tenant set item removed: active={active}, broken_set={broken_set}" + wait_until(10, 0.5, found_active) From 74c5e3d9b877ae006c0c5c4b4ea176ed36f647c1 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Mon, 5 Feb 2024 14:27:25 +0000 Subject: [PATCH 07/72] use string interner for project cache (#6578) ## Problem Running some memory profiling with high concurrent request rate shows seemingly some memory fragmentation. ## Summary of changes Eventually, we will want to separate global memory (caches) from local memory (per connection handshake and per passthrough). Using a string interner for project info cache helps reduce some of the fragmentation of the global cache by having a single heap dedicated to project strings, and not scattering them throughout all a requests. At the same time, the interned key is 4 bytes vs the 24 bytes that `SmolStr` offers. Important: we should only store verified strings in the interner because there's no way to remove them afterwards. Good for caching responses from console. --- Cargo.lock | 13 ++ Cargo.toml | 1 + proxy/Cargo.toml | 2 + proxy/src/cache/project_info.rs | 84 ++++++----- proxy/src/intern.rs | 237 +++++++++++++++++++++++++++++++ proxy/src/lib.rs | 1 + proxy/src/redis/notifications.rs | 31 ++-- workspace_hack/Cargo.toml | 5 +- 8 files changed, 321 insertions(+), 53 deletions(-) create mode 100644 proxy/src/intern.rs diff --git a/Cargo.lock b/Cargo.lock index 02450709d1..c16331636a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2718,6 +2718,16 @@ dependencies = [ "libc", ] +[[package]] +name = "lasso" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4644821e1c3d7a560fe13d842d13f587c07348a1a05d3a797152d41c90c56df2" +dependencies = [ + "dashmap", + "hashbrown 0.13.2", +] + [[package]] name = "lazy_static" version = "1.4.0" @@ -4075,6 +4085,7 @@ dependencies = [ "hyper-tungstenite", "ipnet", "itertools", + "lasso", "md5", "metrics", "native-tls", @@ -4091,6 +4102,7 @@ dependencies = [ "pq_proto", "prometheus", "rand 0.8.5", + "rand_distr", "rcgen", "redis", "regex", @@ -6803,6 +6815,7 @@ dependencies = [ "futures-sink", "futures-util", "getrandom 0.2.11", + "hashbrown 0.13.2", "hashbrown 0.14.0", "hex", "hmac", diff --git a/Cargo.toml b/Cargo.toml index 0cfe522ff9..271edee742 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -95,6 +95,7 @@ inotify = "0.10.2" ipnet = "2.9.0" itertools = "0.10" jsonwebtoken = "9" +lasso = "0.7" libc = "0.2" md5 = "0.7.0" memoffset = "0.8" diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index 79abe639ed..1247f08ee6 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -31,6 +31,7 @@ hyper-tungstenite.workspace = true hyper.workspace = true ipnet.workspace = true itertools.workspace = true +lasso = { workspace = true, features = ["multi-threaded"] } md5.workspace = true metrics.workspace = true once_cell.workspace = true @@ -92,3 +93,4 @@ rcgen.workspace = true rstest.workspace = true tokio-postgres-rustls.workspace = true walkdir.workspace = true +rand_distr = "0.4" diff --git a/proxy/src/cache/project_info.rs b/proxy/src/cache/project_info.rs index 6f37868a8c..62015312a9 100644 --- a/proxy/src/cache/project_info.rs +++ b/proxy/src/cache/project_info.rs @@ -12,15 +12,18 @@ use tokio::time::Instant; use tracing::{debug, info}; use crate::{ - auth::IpPattern, config::ProjectInfoCacheOptions, console::AuthSecret, EndpointId, ProjectId, - RoleName, + auth::IpPattern, + config::ProjectInfoCacheOptions, + console::AuthSecret, + intern::{EndpointIdInt, ProjectIdInt, RoleNameInt}, + EndpointId, ProjectId, RoleName, }; use super::{Cache, Cached}; pub trait ProjectInfoCache { - fn invalidate_allowed_ips_for_project(&self, project_id: &ProjectId); - fn invalidate_role_secret_for_project(&self, project_id: &ProjectId, role_name: &RoleName); + fn invalidate_allowed_ips_for_project(&self, project_id: ProjectIdInt); + fn invalidate_role_secret_for_project(&self, project_id: ProjectIdInt, role_name: RoleNameInt); fn enable_ttl(&self); fn disable_ttl(&self); } @@ -47,7 +50,7 @@ impl From for Entry { #[derive(Default)] struct EndpointInfo { - secret: std::collections::HashMap>>, + secret: std::collections::HashMap>>, allowed_ips: Option>>>, } @@ -60,11 +63,11 @@ impl EndpointInfo { } pub fn get_role_secret( &self, - role_name: &RoleName, + role_name: RoleNameInt, valid_since: Instant, ignore_cache_since: Option, ) -> Option<(Option, bool)> { - if let Some(secret) = self.secret.get(role_name) { + if let Some(secret) = self.secret.get(&role_name) { if valid_since < secret.created_at { return Some(( secret.value.clone(), @@ -93,8 +96,8 @@ impl EndpointInfo { pub fn invalidate_allowed_ips(&mut self) { self.allowed_ips = None; } - pub fn invalidate_role_secret(&mut self, role_name: &RoleName) { - self.secret.remove(role_name); + pub fn invalidate_role_secret(&mut self, role_name: RoleNameInt) { + self.secret.remove(&role_name); } } @@ -106,9 +109,9 @@ impl EndpointInfo { /// One may ask, why the data is stored per project, when on the user request there is only data about the endpoint available? /// On the cplane side updates are done per project (or per branch), so it's easier to invalidate the whole project cache. pub struct ProjectInfoCacheImpl { - cache: DashMap, + cache: DashMap, - project2ep: DashMap>, + project2ep: DashMap>, config: ProjectInfoCacheOptions, start_time: Instant, @@ -116,11 +119,11 @@ pub struct ProjectInfoCacheImpl { } impl ProjectInfoCache for ProjectInfoCacheImpl { - fn invalidate_allowed_ips_for_project(&self, project_id: &ProjectId) { + fn invalidate_allowed_ips_for_project(&self, project_id: ProjectIdInt) { info!("invalidating allowed ips for project `{}`", project_id); let endpoints = self .project2ep - .get(project_id) + .get(&project_id) .map(|kv| kv.value().clone()) .unwrap_or_default(); for endpoint_id in endpoints { @@ -129,14 +132,14 @@ impl ProjectInfoCache for ProjectInfoCacheImpl { } } } - fn invalidate_role_secret_for_project(&self, project_id: &ProjectId, role_name: &RoleName) { + fn invalidate_role_secret_for_project(&self, project_id: ProjectIdInt, role_name: RoleNameInt) { info!( "invalidating role secret for project_id `{}` and role_name `{}`", - project_id, role_name + project_id, role_name, ); let endpoints = self .project2ep - .get(project_id) + .get(&project_id) .map(|kv| kv.value().clone()) .unwrap_or_default(); for endpoint_id in endpoints { @@ -173,15 +176,17 @@ impl ProjectInfoCacheImpl { endpoint_id: &EndpointId, role_name: &RoleName, ) -> Option>> { + let endpoint_id = EndpointIdInt::get(endpoint_id)?; + let role_name = RoleNameInt::get(role_name)?; let (valid_since, ignore_cache_since) = self.get_cache_times(); - let endpoint_info = self.cache.get(endpoint_id)?; + let endpoint_info = self.cache.get(&endpoint_id)?; let (value, ignore_cache) = endpoint_info.get_role_secret(role_name, valid_since, ignore_cache_since)?; if !ignore_cache { let cached = Cached { token: Some(( self, - CachedLookupInfo::new_role_secret(endpoint_id.clone(), role_name.clone()), + CachedLookupInfo::new_role_secret(endpoint_id, role_name), )), value, }; @@ -193,13 +198,14 @@ impl ProjectInfoCacheImpl { &self, endpoint_id: &EndpointId, ) -> Option>>> { + let endpoint_id = EndpointIdInt::get(endpoint_id)?; let (valid_since, ignore_cache_since) = self.get_cache_times(); - let endpoint_info = self.cache.get(endpoint_id)?; + let endpoint_info = self.cache.get(&endpoint_id)?; let value = endpoint_info.get_allowed_ips(valid_since, ignore_cache_since); let (value, ignore_cache) = value?; if !ignore_cache { let cached = Cached { - token: Some((self, CachedLookupInfo::new_allowed_ips(endpoint_id.clone()))), + token: Some((self, CachedLookupInfo::new_allowed_ips(endpoint_id))), value, }; return Some(cached); @@ -213,14 +219,17 @@ impl ProjectInfoCacheImpl { role_name: &RoleName, secret: Option, ) { + let project_id = ProjectIdInt::from(project_id); + let endpoint_id = EndpointIdInt::from(endpoint_id); + let role_name = RoleNameInt::from(role_name); if self.cache.len() >= self.config.size { // If there are too many entries, wait until the next gc cycle. return; } - self.inser_project2endpoint(project_id, endpoint_id); - let mut entry = self.cache.entry(endpoint_id.clone()).or_default(); + self.insert_project2endpoint(project_id, endpoint_id); + let mut entry = self.cache.entry(endpoint_id).or_default(); if entry.secret.len() < self.config.max_roles { - entry.secret.insert(role_name.clone(), secret.into()); + entry.secret.insert(role_name, secret.into()); } } pub fn insert_allowed_ips( @@ -229,22 +238,21 @@ impl ProjectInfoCacheImpl { endpoint_id: &EndpointId, allowed_ips: Arc>, ) { + let project_id = ProjectIdInt::from(project_id); + let endpoint_id = EndpointIdInt::from(endpoint_id); if self.cache.len() >= self.config.size { // If there are too many entries, wait until the next gc cycle. return; } - self.inser_project2endpoint(project_id, endpoint_id); - self.cache - .entry(endpoint_id.clone()) - .or_default() - .allowed_ips = Some(allowed_ips.into()); + self.insert_project2endpoint(project_id, endpoint_id); + self.cache.entry(endpoint_id).or_default().allowed_ips = Some(allowed_ips.into()); } - fn inser_project2endpoint(&self, project_id: &ProjectId, endpoint_id: &EndpointId) { - if let Some(mut endpoints) = self.project2ep.get_mut(project_id) { - endpoints.insert(endpoint_id.clone()); + fn insert_project2endpoint(&self, project_id: ProjectIdInt, endpoint_id: EndpointIdInt) { + if let Some(mut endpoints) = self.project2ep.get_mut(&project_id) { + endpoints.insert(endpoint_id); } else { self.project2ep - .insert(project_id.clone(), HashSet::from([endpoint_id.clone()])); + .insert(project_id, HashSet::from([endpoint_id])); } } fn get_cache_times(&self) -> (Instant, Option) { @@ -300,18 +308,18 @@ impl ProjectInfoCacheImpl { /// This is used to invalidate cache entries. pub struct CachedLookupInfo { /// Search by this key. - endpoint_id: EndpointId, + endpoint_id: EndpointIdInt, lookup_type: LookupType, } impl CachedLookupInfo { - pub(self) fn new_role_secret(endpoint_id: EndpointId, role_name: RoleName) -> Self { + pub(self) fn new_role_secret(endpoint_id: EndpointIdInt, role_name: RoleNameInt) -> Self { Self { endpoint_id, lookup_type: LookupType::RoleSecret(role_name), } } - pub(self) fn new_allowed_ips(endpoint_id: EndpointId) -> Self { + pub(self) fn new_allowed_ips(endpoint_id: EndpointIdInt) -> Self { Self { endpoint_id, lookup_type: LookupType::AllowedIps, @@ -320,7 +328,7 @@ impl CachedLookupInfo { } enum LookupType { - RoleSecret(RoleName), + RoleSecret(RoleNameInt), AllowedIps, } @@ -335,7 +343,7 @@ impl Cache for ProjectInfoCacheImpl { match &key.lookup_type { LookupType::RoleSecret(role_name) => { if let Some(mut endpoint_info) = self.cache.get_mut(&key.endpoint_id) { - endpoint_info.invalidate_role_secret(role_name); + endpoint_info.invalidate_role_secret(*role_name); } } LookupType::AllowedIps => { @@ -457,7 +465,7 @@ mod tests { assert_eq!(cached.value, secret2); // The only way to invalidate this value is to invalidate via the api. - cache.invalidate_role_secret_for_project(&project_id, &user2); + cache.invalidate_role_secret_for_project((&project_id).into(), (&user2).into()); assert!(cache.get_role_secret(&endpoint_id, &user2).is_none()); let cached = cache.get_allowed_ips(&endpoint_id).unwrap(); diff --git a/proxy/src/intern.rs b/proxy/src/intern.rs new file mode 100644 index 0000000000..a6519bdff9 --- /dev/null +++ b/proxy/src/intern.rs @@ -0,0 +1,237 @@ +use std::{ + hash::BuildHasherDefault, marker::PhantomData, num::NonZeroUsize, ops::Index, sync::OnceLock, +}; + +use lasso::{Capacity, MemoryLimits, Spur, ThreadedRodeo}; +use rustc_hash::FxHasher; + +use crate::{BranchId, EndpointId, ProjectId, RoleName}; + +pub trait InternId: Sized + 'static { + fn get_interner() -> &'static StringInterner; +} + +pub struct StringInterner { + inner: ThreadedRodeo>, + _id: PhantomData, +} + +#[derive(PartialEq, Debug, Clone, Copy, Eq, Hash)] +pub struct InternedString { + inner: Spur, + _id: PhantomData, +} + +impl std::fmt::Display for InternedString { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.as_str().fmt(f) + } +} + +impl InternedString { + pub fn as_str(&self) -> &'static str { + Id::get_interner().inner.resolve(&self.inner) + } + pub fn get(s: &str) -> Option { + Id::get_interner().get(s) + } +} + +impl AsRef for InternedString { + fn as_ref(&self) -> &str { + self.as_str() + } +} + +impl std::ops::Deref for InternedString { + type Target = str; + fn deref(&self) -> &str { + self.as_str() + } +} + +impl<'de, Id: InternId> serde::de::Deserialize<'de> for InternedString { + fn deserialize>(d: D) -> Result { + struct Visitor(PhantomData); + impl<'de, Id: InternId> serde::de::Visitor<'de> for Visitor { + type Value = InternedString; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("a string") + } + + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + Ok(Id::get_interner().get_or_intern(v)) + } + } + d.deserialize_str(Visitor::(PhantomData)) + } +} + +impl serde::Serialize for InternedString { + fn serialize(&self, s: S) -> Result { + self.as_str().serialize(s) + } +} + +impl StringInterner { + pub fn new() -> Self { + StringInterner { + inner: ThreadedRodeo::with_capacity_memory_limits_and_hasher( + Capacity::new(2500, NonZeroUsize::new(1 << 16).unwrap()), + // unbounded + MemoryLimits::for_memory_usage(usize::MAX), + BuildHasherDefault::::default(), + ), + _id: PhantomData, + } + } + + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + pub fn len(&self) -> usize { + self.inner.len() + } + + pub fn current_memory_usage(&self) -> usize { + self.inner.current_memory_usage() + } + + pub fn get_or_intern(&self, s: &str) -> InternedString { + InternedString { + inner: self.inner.get_or_intern(s), + _id: PhantomData, + } + } + + pub fn get(&self, s: &str) -> Option> { + Some(InternedString { + inner: self.inner.get(s)?, + _id: PhantomData, + }) + } +} + +impl Index> for StringInterner { + type Output = str; + + fn index(&self, index: InternedString) -> &Self::Output { + self.inner.resolve(&index.inner) + } +} + +impl Default for StringInterner { + fn default() -> Self { + Self::new() + } +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +pub struct RoleNameTag; +impl InternId for RoleNameTag { + fn get_interner() -> &'static StringInterner { + pub static ROLE_NAMES: OnceLock> = OnceLock::new(); + ROLE_NAMES.get_or_init(Default::default) + } +} +pub type RoleNameInt = InternedString; +impl From<&RoleName> for RoleNameInt { + fn from(value: &RoleName) -> Self { + RoleNameTag::get_interner().get_or_intern(value) + } +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +pub struct EndpointIdTag; +impl InternId for EndpointIdTag { + fn get_interner() -> &'static StringInterner { + pub static ROLE_NAMES: OnceLock> = OnceLock::new(); + ROLE_NAMES.get_or_init(Default::default) + } +} +pub type EndpointIdInt = InternedString; +impl From<&EndpointId> for EndpointIdInt { + fn from(value: &EndpointId) -> Self { + EndpointIdTag::get_interner().get_or_intern(value) + } +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +pub struct BranchIdTag; +impl InternId for BranchIdTag { + fn get_interner() -> &'static StringInterner { + pub static ROLE_NAMES: OnceLock> = OnceLock::new(); + ROLE_NAMES.get_or_init(Default::default) + } +} +pub type BranchIdInt = InternedString; +impl From<&BranchId> for BranchIdInt { + fn from(value: &BranchId) -> Self { + BranchIdTag::get_interner().get_or_intern(value) + } +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +pub struct ProjectIdTag; +impl InternId for ProjectIdTag { + fn get_interner() -> &'static StringInterner { + pub static ROLE_NAMES: OnceLock> = OnceLock::new(); + ROLE_NAMES.get_or_init(Default::default) + } +} +pub type ProjectIdInt = InternedString; +impl From<&ProjectId> for ProjectIdInt { + fn from(value: &ProjectId) -> Self { + ProjectIdTag::get_interner().get_or_intern(value) + } +} + +#[cfg(test)] +mod tests { + use std::sync::OnceLock; + + use crate::intern::StringInterner; + + use super::InternId; + + struct MyId; + impl InternId for MyId { + fn get_interner() -> &'static StringInterner { + pub static ROLE_NAMES: OnceLock> = OnceLock::new(); + ROLE_NAMES.get_or_init(Default::default) + } + } + + #[test] + fn push_many_strings() { + use rand::{rngs::StdRng, Rng, SeedableRng}; + use rand_distr::Zipf; + + let endpoint_dist = Zipf::new(500000, 0.8).unwrap(); + let endpoints = StdRng::seed_from_u64(272488357).sample_iter(endpoint_dist); + + let interner = MyId::get_interner(); + + const N: usize = 100_000; + let mut verify = Vec::with_capacity(N); + for endpoint in endpoints.take(N) { + let endpoint = format!("ep-string-interning-{endpoint}"); + let key = interner.get_or_intern(&endpoint); + verify.push((endpoint, key)); + } + + for (s, key) in verify { + assert_eq!(interner[key], s); + } + + // 2031616/59861 = 34 bytes per string + assert_eq!(interner.len(), 59_861); + // will have other overhead for the internal hashmaps that are not accounted for. + assert_eq!(interner.current_memory_usage(), 2_031_616); + } +} diff --git a/proxy/src/lib.rs b/proxy/src/lib.rs index db6256d611..da7c7f3ed2 100644 --- a/proxy/src/lib.rs +++ b/proxy/src/lib.rs @@ -16,6 +16,7 @@ pub mod console; pub mod context; pub mod error; pub mod http; +pub mod intern; pub mod jemalloc; pub mod logging; pub mod metrics; diff --git a/proxy/src/redis/notifications.rs b/proxy/src/redis/notifications.rs index 9cd70b109b..158884aa17 100644 --- a/proxy/src/redis/notifications.rs +++ b/proxy/src/redis/notifications.rs @@ -4,7 +4,10 @@ use futures::StreamExt; use redis::aio::PubSub; use serde::Deserialize; -use crate::{cache::project_info::ProjectInfoCache, ProjectId, RoleName}; +use crate::{ + cache::project_info::ProjectInfoCache, + intern::{ProjectIdInt, RoleNameInt}, +}; const CHANNEL_NAME: &str = "neondb-proxy-ws-updates"; const RECONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); @@ -45,12 +48,12 @@ enum Notification { } #[derive(Clone, Debug, Deserialize, Eq, PartialEq)] struct AllowedIpsUpdate { - project_id: ProjectId, + project_id: ProjectIdInt, } #[derive(Clone, Debug, Deserialize, Eq, PartialEq)] struct PasswordUpdate { - project_id: ProjectId, - role_name: RoleName, + project_id: ProjectIdInt, + role_name: RoleNameInt, } fn deserialize_json_string<'de, D, T>(deserializer: D) -> Result where @@ -65,11 +68,11 @@ fn invalidate_cache(cache: Arc, msg: Notification) { use Notification::*; match msg { AllowedIpsUpdate { allowed_ips_update } => { - cache.invalidate_allowed_ips_for_project(&allowed_ips_update.project_id) + cache.invalidate_allowed_ips_for_project(allowed_ips_update.project_id) } PasswordUpdate { password_update } => cache.invalidate_role_secret_for_project( - &password_update.project_id, - &password_update.role_name, + password_update.project_id, + password_update.role_name, ), } } @@ -141,12 +144,14 @@ where #[cfg(test)] mod tests { + use crate::{ProjectId, RoleName}; + use super::*; use serde_json::json; #[test] fn parse_allowed_ips() -> anyhow::Result<()> { - let project_id = "new_project".to_string(); + let project_id: ProjectId = "new_project".into(); let data = format!("{{\"project_id\": \"{project_id}\"}}"); let text = json!({ "type": "message", @@ -161,7 +166,7 @@ mod tests { result, Notification::AllowedIpsUpdate { allowed_ips_update: AllowedIpsUpdate { - project_id: project_id.into() + project_id: (&project_id).into() } } ); @@ -171,8 +176,8 @@ mod tests { #[test] fn parse_password_updated() -> anyhow::Result<()> { - let project_id = "new_project".to_string(); - let role_name = "new_role".to_string(); + let project_id: ProjectId = "new_project".into(); + let role_name: RoleName = "new_role".into(); let data = format!("{{\"project_id\": \"{project_id}\", \"role_name\": \"{role_name}\"}}"); let text = json!({ "type": "message", @@ -187,8 +192,8 @@ mod tests { result, Notification::PasswordUpdate { password_update: PasswordUpdate { - project_id: project_id.into(), - role_name: role_name.into() + project_id: (&project_id).into(), + role_name: (&role_name).into(), } } ); diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index f58b912a77..74464dd4c8 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -39,7 +39,8 @@ futures-io = { version = "0.3" } futures-sink = { version = "0.3" } futures-util = { version = "0.3", features = ["channel", "io", "sink"] } getrandom = { version = "0.2", default-features = false, features = ["std"] } -hashbrown = { version = "0.14", default-features = false, features = ["raw"] } +hashbrown-582f2526e08bb6a0 = { package = "hashbrown", version = "0.14", default-features = false, features = ["raw"] } +hashbrown-594e8ee84c453af0 = { package = "hashbrown", version = "0.13", features = ["raw"] } hex = { version = "0.4", features = ["serde"] } hmac = { version = "0.12", default-features = false, features = ["reset"] } hyper = { version = "0.14", features = ["full"] } @@ -91,7 +92,7 @@ cc = { version = "1", default-features = false, features = ["parallel"] } chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "wasmbind"] } either = { version = "1" } getrandom = { version = "0.2", default-features = false, features = ["std"] } -hashbrown = { version = "0.14", default-features = false, features = ["raw"] } +hashbrown-582f2526e08bb6a0 = { package = "hashbrown", version = "0.14", default-features = false, features = ["raw"] } indexmap = { version = "1", default-features = false, features = ["std"] } itertools = { version = "0.10" } libc = { version = "0.2", features = ["extra_traits", "use_std"] } From cb7c89332f25c652fa7dd06a9be7d984f8cc3989 Mon Sep 17 00:00:00 2001 From: John Spray Date: Mon, 5 Feb 2024 14:29:05 +0000 Subject: [PATCH 08/72] control_plane: fix tenant GET, clean up endpoints (#6553) Cleanups from https://github.com/neondatabase/neon/pull/6394 - There was a rogue `*` breaking the `GET /tenant/:tenant_id`, which passes through to shard zero - There was a duplicate migrate endpoint - There are un-prefixed API endpoints that were only needed for compat tests and can now be removed. --- control_plane/attachment_service/src/http.rs | 10 +--------- test_runner/regress/test_sharding_service.py | 7 +++++++ 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/control_plane/attachment_service/src/http.rs b/control_plane/attachment_service/src/http.rs index aa8c73c493..049e66fddf 100644 --- a/control_plane/attachment_service/src/http.rs +++ b/control_plane/attachment_service/src/http.rs @@ -403,10 +403,6 @@ pub fn make_router( .put("/v1/tenant/:tenant_id/location_config", |r| { tenant_service_handler(r, handle_tenant_location_config) }) - // Tenant Shard operations (low level/maintenance) - .put("/tenant/:tenant_shard_id/migrate", |r| { - tenant_service_handler(r, handle_tenant_shard_migrate) - }) // Timeline operations .delete("/v1/tenant/:tenant_id/timeline/:timeline_id", |r| { tenant_service_handler(r, handle_tenant_timeline_delete) @@ -415,7 +411,7 @@ pub fn make_router( tenant_service_handler(r, handle_tenant_timeline_create) }) // Tenant detail GET passthrough to shard zero - .get("/v1/tenant/:tenant_id*", |r| { + .get("/v1/tenant/:tenant_id", |r| { tenant_service_handler(r, handle_tenant_timeline_passthrough) }) // Timeline GET passthrough to shard zero. Note that the `*` in the URL is a wildcard: any future @@ -423,8 +419,4 @@ pub fn make_router( .get("/v1/tenant/:tenant_id/timeline*", |r| { tenant_service_handler(r, handle_tenant_timeline_passthrough) }) - // Path aliases for tests_forward_compatibility - // TODO: remove these in future PR - .post("/re-attach", |r| request_span(r, handle_re_attach)) - .post("/validate", |r| request_span(r, handle_validate)) } diff --git a/test_runner/regress/test_sharding_service.py b/test_runner/regress/test_sharding_service.py index 346df708de..5c70378ab0 100644 --- a/test_runner/regress/test_sharding_service.py +++ b/test_runner/regress/test_sharding_service.py @@ -140,6 +140,13 @@ def test_sharding_service_passthrough( timelines = client.timeline_list(tenant_id=env.initial_tenant) assert len(timelines) == 1 + status = client.tenant_status(env.initial_tenant) + assert TenantId(status["id"]) == env.initial_tenant + assert set(TimelineId(t) for t in status["timelines"]) == { + env.initial_timeline, + } + assert status["state"]["slug"] == "Active" + def test_sharding_service_restart(neon_env_builder: NeonEnvBuilder): env = neon_env_builder.init_start() From 8e114bd6101dee117e1125ea68dfbdbbc59c965f Mon Sep 17 00:00:00 2001 From: John Spray Date: Mon, 5 Feb 2024 19:31:55 +0000 Subject: [PATCH 09/72] control_plane/attachment_service: make --database-url optional (#6636) ## Problem This change was left out of #6585 accidentally -- just forgot to push the very last version of my branch. Now that we can load database url from Secrets Manager, we don't always need it on the CLI any more. We should let the user omit it instead of passing `--database-url ""` ## Summary of changes - Make `--database-url` optional --- control_plane/attachment_service/src/main.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/control_plane/attachment_service/src/main.rs b/control_plane/attachment_service/src/main.rs index eda9c7aad6..37b06c4090 100644 --- a/control_plane/attachment_service/src/main.rs +++ b/control_plane/attachment_service/src/main.rs @@ -53,7 +53,7 @@ struct Cli { /// URL to connect to postgres, like postgresql://localhost:1234/attachment_service #[arg(long)] - database_url: String, + database_url: Option, } /// Secrets may either be provided on the command line (for testing), or loaded from AWS SecretManager: this @@ -74,10 +74,9 @@ impl Secrets { const PUBLIC_KEY_SECRET: &'static str = "neon-storage-controller-public-key"; async fn load(args: &Cli) -> anyhow::Result { - if args.database_url.is_empty() { - Self::load_aws_sm().await - } else { - Self::load_cli(args) + match &args.database_url { + Some(url) => Self::load_cli(url, args), + None => Self::load_aws_sm().await, } } @@ -153,13 +152,13 @@ impl Secrets { }) } - fn load_cli(args: &Cli) -> anyhow::Result { + fn load_cli(database_url: &str, args: &Cli) -> anyhow::Result { let public_key = match &args.public_key { None => None, Some(key_path) => Some(JwtAuth::from_key_path(key_path)?), }; Ok(Self { - database_url: args.database_url.clone(), + database_url: database_url.to_owned(), public_key, jwt_token: args.jwt_token.clone(), control_plane_jwt_token: args.control_plane_jwt_token.clone(), From 947165788dc2447b17b8cd163568d10b8c4ddeaa Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Tue, 6 Feb 2024 09:39:06 +0200 Subject: [PATCH 10/72] refactor: needless cancellation token cloning (#6618) The solution we ended up for `backoff::retry` requires always cloning of cancellation tokens even though there is just `.await`. Fix that, and also turn the return type into `Option>` avoiding the need for the `E::cancelled()` fn passed in. Cc: #6096 --- .../attachment_service/src/compute_hook.rs | 4 +- libs/remote_storage/src/azure_blob.rs | 2 +- libs/remote_storage/src/lib.rs | 4 +- libs/remote_storage/src/local_fs.rs | 2 +- libs/remote_storage/src/s3_bucket.rs | 14 ++- libs/remote_storage/src/simulate_failures.rs | 2 +- libs/remote_storage/tests/test_real_s3.rs | 11 ++- libs/utils/src/backoff.rs | 92 ++++++++----------- pageserver/src/consumption_metrics/upload.rs | 56 +++++------ pageserver/src/control_plane_client.rs | 35 ++----- pageserver/src/deletion_queue/deleter.rs | 4 +- pageserver/src/tenant.rs | 8 +- pageserver/src/tenant/delete.rs | 8 +- .../src/tenant/remote_timeline_client.rs | 16 +++- .../tenant/remote_timeline_client/download.rs | 19 ++-- .../tenant/remote_timeline_client/upload.rs | 8 +- pageserver/src/tenant/secondary/downloader.rs | 8 +- .../src/tenant/secondary/heatmap_uploader.rs | 4 +- proxy/src/context/parquet.rs | 4 +- safekeeper/src/wal_backup.rs | 9 +- 20 files changed, 156 insertions(+), 154 deletions(-) diff --git a/control_plane/attachment_service/src/compute_hook.rs b/control_plane/attachment_service/src/compute_hook.rs index 9c1185f259..4ca26431ca 100644 --- a/control_plane/attachment_service/src/compute_hook.rs +++ b/control_plane/attachment_service/src/compute_hook.rs @@ -244,9 +244,11 @@ impl ComputeHook { 3, 10, "Send compute notification", - backoff::Cancel::new(cancel.clone(), || NotifyError::ShuttingDown), + cancel, ) .await + .ok_or_else(|| NotifyError::ShuttingDown) + .and_then(|x| x) } /// Call this to notify the compute (postgres) tier of new pageservers to use diff --git a/libs/remote_storage/src/azure_blob.rs b/libs/remote_storage/src/azure_blob.rs index 57c57a2b70..c6d5224706 100644 --- a/libs/remote_storage/src/azure_blob.rs +++ b/libs/remote_storage/src/azure_blob.rs @@ -379,7 +379,7 @@ impl RemoteStorage for AzureBlobStorage { _prefix: Option<&RemotePath>, _timestamp: SystemTime, _done_if_after: SystemTime, - _cancel: CancellationToken, + _cancel: &CancellationToken, ) -> Result<(), TimeTravelError> { // TODO use Azure point in time recovery feature for this // https://learn.microsoft.com/en-us/azure/storage/blobs/point-in-time-restore-overview diff --git a/libs/remote_storage/src/lib.rs b/libs/remote_storage/src/lib.rs index 4aeaee70b1..e64b1de6f9 100644 --- a/libs/remote_storage/src/lib.rs +++ b/libs/remote_storage/src/lib.rs @@ -218,7 +218,7 @@ pub trait RemoteStorage: Send + Sync + 'static { prefix: Option<&RemotePath>, timestamp: SystemTime, done_if_after: SystemTime, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result<(), TimeTravelError>; } @@ -442,7 +442,7 @@ impl GenericRemoteStorage> { prefix: Option<&RemotePath>, timestamp: SystemTime, done_if_after: SystemTime, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result<(), TimeTravelError> { match self { Self::LocalFs(s) => { diff --git a/libs/remote_storage/src/local_fs.rs b/libs/remote_storage/src/local_fs.rs index d47fa75b37..36ec15e1b1 100644 --- a/libs/remote_storage/src/local_fs.rs +++ b/libs/remote_storage/src/local_fs.rs @@ -431,7 +431,7 @@ impl RemoteStorage for LocalFs { _prefix: Option<&RemotePath>, _timestamp: SystemTime, _done_if_after: SystemTime, - _cancel: CancellationToken, + _cancel: &CancellationToken, ) -> Result<(), TimeTravelError> { Err(TimeTravelError::Unimplemented) } diff --git a/libs/remote_storage/src/s3_bucket.rs b/libs/remote_storage/src/s3_bucket.rs index 4d6564cba6..c9ad9ef225 100644 --- a/libs/remote_storage/src/s3_bucket.rs +++ b/libs/remote_storage/src/s3_bucket.rs @@ -638,7 +638,7 @@ impl RemoteStorage for S3Bucket { prefix: Option<&RemotePath>, timestamp: SystemTime, done_if_after: SystemTime, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result<(), TimeTravelError> { let kind = RequestKind::TimeTravel; let _guard = self.permit(kind).await; @@ -678,9 +678,11 @@ impl RemoteStorage for S3Bucket { warn_threshold, max_retries, "listing object versions for time_travel_recover", - backoff::Cancel::new(cancel.clone(), || TimeTravelError::Cancelled), + cancel, ) - .await?; + .await + .ok_or_else(|| TimeTravelError::Cancelled) + .and_then(|x| x)?; tracing::trace!( " Got List response version_id_marker={:?}, key_marker={:?}", @@ -805,9 +807,11 @@ impl RemoteStorage for S3Bucket { warn_threshold, max_retries, "copying object version for time_travel_recover", - backoff::Cancel::new(cancel.clone(), || TimeTravelError::Cancelled), + cancel, ) - .await?; + .await + .ok_or_else(|| TimeTravelError::Cancelled) + .and_then(|x| x)?; tracing::info!(%version_id, %key, "Copied old version in S3"); } VerOrDelete { diff --git a/libs/remote_storage/src/simulate_failures.rs b/libs/remote_storage/src/simulate_failures.rs index ee9792232a..82d5a61fda 100644 --- a/libs/remote_storage/src/simulate_failures.rs +++ b/libs/remote_storage/src/simulate_failures.rs @@ -190,7 +190,7 @@ impl RemoteStorage for UnreliableWrapper { prefix: Option<&RemotePath>, timestamp: SystemTime, done_if_after: SystemTime, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result<(), TimeTravelError> { self.attempt(RemoteOp::TimeTravelRecover(prefix.map(|p| p.to_owned()))) .map_err(|e| TimeTravelError::Other(anyhow::Error::new(e)))?; diff --git a/libs/remote_storage/tests/test_real_s3.rs b/libs/remote_storage/tests/test_real_s3.rs index 679be66bf7..fc52dabc36 100644 --- a/libs/remote_storage/tests/test_real_s3.rs +++ b/libs/remote_storage/tests/test_real_s3.rs @@ -56,9 +56,10 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow: warn_threshold, max_retries, "test retry", - backoff::Cancel::new(CancellationToken::new(), || unreachable!()), + &CancellationToken::new(), ) .await + .expect("never cancelled") } async fn time_point() -> SystemTime { @@ -76,6 +77,8 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow: .collect::>()) } + let cancel = CancellationToken::new(); + let path1 = RemotePath::new(Utf8Path::new(format!("{}/path1", ctx.base_prefix).as_str())) .with_context(|| "RemotePath conversion")?; @@ -142,7 +145,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow: // No changes after recovery to t2 (no-op) let t_final = time_point().await; ctx.client - .time_travel_recover(None, t2, t_final, CancellationToken::new()) + .time_travel_recover(None, t2, t_final, &cancel) .await?; let t2_files_recovered = list_files(&ctx.client).await?; println!("after recovery to t2: {t2_files_recovered:?}"); @@ -153,7 +156,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow: // after recovery to t1: path1 is back, path2 has the old content let t_final = time_point().await; ctx.client - .time_travel_recover(None, t1, t_final, CancellationToken::new()) + .time_travel_recover(None, t1, t_final, &cancel) .await?; let t1_files_recovered = list_files(&ctx.client).await?; println!("after recovery to t1: {t1_files_recovered:?}"); @@ -164,7 +167,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow: // after recovery to t0: everything is gone except for path1 let t_final = time_point().await; ctx.client - .time_travel_recover(None, t0, t_final, CancellationToken::new()) + .time_travel_recover(None, t0, t_final, &cancel) .await?; let t0_files_recovered = list_files(&ctx.client).await?; println!("after recovery to t0: {t0_files_recovered:?}"); diff --git a/libs/utils/src/backoff.rs b/libs/utils/src/backoff.rs index d50ad39585..096c7e5854 100644 --- a/libs/utils/src/backoff.rs +++ b/libs/utils/src/backoff.rs @@ -37,69 +37,53 @@ pub fn exponential_backoff_duration_seconds(n: u32, base_increment: f64, max_sec } } -/// Configure cancellation for a retried operation: when to cancel (the token), and -/// what kind of error to return on cancellation -pub struct Cancel -where - E: Display + Debug + 'static, - CF: Fn() -> E, -{ - token: CancellationToken, - on_cancel: CF, -} - -impl Cancel -where - E: Display + Debug + 'static, - CF: Fn() -> E, -{ - pub fn new(token: CancellationToken, on_cancel: CF) -> Self { - Self { token, on_cancel } - } -} - -/// retries passed operation until one of the following conditions are met: -/// Encountered error is considered as permanent (non-retryable) -/// Retries have been exhausted. -/// `is_permanent` closure should be used to provide distinction between permanent/non-permanent errors -/// When attempts cross `warn_threshold` function starts to emit log warnings. +/// Retries passed operation until one of the following conditions are met: +/// - encountered error is considered as permanent (non-retryable) +/// - retries have been exhausted +/// - cancellation token has been cancelled +/// +/// `is_permanent` closure should be used to provide distinction between permanent/non-permanent +/// errors. When attempts cross `warn_threshold` function starts to emit log warnings. /// `description` argument is added to log messages. Its value should identify the `op` is doing -/// `cancel` argument is required: any time we are looping on retry, we should be using a CancellationToken -/// to drop out promptly on shutdown. -pub async fn retry( +/// `cancel` cancels new attempts and the backoff sleep. +/// +/// If attempts fail, they are being logged with `{:#}` which works for anyhow, but does not work +/// for any other error type. Final failed attempt is logged with `{:?}`. +/// +/// Returns `None` if cancellation was noticed during backoff or the terminal result. +pub async fn retry( mut op: O, is_permanent: impl Fn(&E) -> bool, warn_threshold: u32, max_retries: u32, description: &str, - cancel: Cancel, -) -> Result + cancel: &CancellationToken, +) -> Option> where // Not std::error::Error because anyhow::Error doesnt implement it. // For context see https://github.com/dtolnay/anyhow/issues/63 E: Display + Debug + 'static, O: FnMut() -> F, F: Future>, - CF: Fn() -> E, { let mut attempts = 0; loop { - if cancel.token.is_cancelled() { - return Err((cancel.on_cancel)()); + if cancel.is_cancelled() { + return None; } let result = op().await; - match result { + match &result { Ok(_) => { if attempts > 0 { tracing::info!("{description} succeeded after {attempts} retries"); } - return result; + return Some(result); } // These are "permanent" errors that should not be retried. - Err(ref e) if is_permanent(e) => { - return result; + Err(e) if is_permanent(e) => { + return Some(result); } // Assume that any other failure might be transient, and the operation might // succeed if we just keep trying. @@ -109,12 +93,12 @@ where Err(err) if attempts < max_retries => { tracing::warn!("{description} failed, will retry (attempt {attempts}): {err:#}"); } - Err(ref err) => { + Err(err) => { // Operation failed `max_attempts` times. Time to give up. tracing::warn!( "{description} still failed after {attempts} retries, giving up: {err:?}" ); - return result; + return Some(result); } } // sleep and retry @@ -122,7 +106,7 @@ where attempts, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS, - &cancel.token, + cancel, ) .await; attempts += 1; @@ -131,11 +115,9 @@ where #[cfg(test)] mod tests { - use std::io; - - use tokio::sync::Mutex; - use super::*; + use std::io; + use tokio::sync::Mutex; #[test] fn backoff_defaults_produce_growing_backoff_sequence() { @@ -166,7 +148,7 @@ mod tests { #[tokio::test(start_paused = true)] async fn retry_always_error() { let count = Mutex::new(0); - let err_result = retry( + retry( || async { *count.lock().await += 1; Result::<(), io::Error>::Err(io::Error::from(io::ErrorKind::Other)) @@ -175,11 +157,11 @@ mod tests { 1, 1, "work", - Cancel::new(CancellationToken::new(), || -> io::Error { unreachable!() }), + &CancellationToken::new(), ) - .await; - - assert!(err_result.is_err()); + .await + .expect("not cancelled") + .expect_err("it can only fail"); assert_eq!(*count.lock().await, 2); } @@ -201,10 +183,11 @@ mod tests { 2, 2, "work", - Cancel::new(CancellationToken::new(), || -> io::Error { unreachable!() }), + &CancellationToken::new(), ) .await - .unwrap(); + .expect("not cancelled") + .expect("success on second try"); } #[tokio::test(start_paused = true)] @@ -224,10 +207,11 @@ mod tests { 2, 2, "work", - Cancel::new(CancellationToken::new(), || -> io::Error { unreachable!() }), + &CancellationToken::new(), ) .await - .unwrap_err(); + .expect("was not cancellation") + .expect_err("it was permanent error"); assert_eq!(*count.lock().await, 1); } diff --git a/pageserver/src/consumption_metrics/upload.rs b/pageserver/src/consumption_metrics/upload.rs index 322ed95cc8..6b840a3136 100644 --- a/pageserver/src/consumption_metrics/upload.rs +++ b/pageserver/src/consumption_metrics/upload.rs @@ -262,35 +262,33 @@ async fn upload( ) -> Result<(), UploadError> { let warn_after = 3; let max_attempts = 10; + + // this is used only with tests so far + let last_value = if is_last { "true" } else { "false" }; + let res = utils::backoff::retry( - move || { - let body = body.clone(); - async move { - let res = client - .post(metric_collection_endpoint.clone()) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .header( - LAST_IN_BATCH.clone(), - if is_last { "true" } else { "false" }, - ) - .body(body) - .send() - .await; + || async { + let res = client + .post(metric_collection_endpoint.clone()) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .header(LAST_IN_BATCH.clone(), last_value) + .body(body.clone()) + .send() + .await; - let res = res.and_then(|res| res.error_for_status()); + let res = res.and_then(|res| res.error_for_status()); - // 10 redirects are normally allowed, so we don't need worry about 3xx - match res { - Ok(_response) => Ok(()), - Err(e) => { - let status = e.status().filter(|s| s.is_client_error()); - if let Some(status) = status { - // rejection used to be a thing when the server could reject a - // whole batch of metrics if one metric was bad. - Err(UploadError::Rejected(status)) - } else { - Err(UploadError::Reqwest(e)) - } + // 10 redirects are normally allowed, so we don't need worry about 3xx + match res { + Ok(_response) => Ok(()), + Err(e) => { + let status = e.status().filter(|s| s.is_client_error()); + if let Some(status) = status { + // rejection used to be a thing when the server could reject a + // whole batch of metrics if one metric was bad. + Err(UploadError::Rejected(status)) + } else { + Err(UploadError::Reqwest(e)) } } } @@ -299,9 +297,11 @@ async fn upload( warn_after, max_attempts, "upload consumption_metrics", - utils::backoff::Cancel::new(cancel.clone(), || UploadError::Cancelled), + cancel, ) - .await; + .await + .ok_or_else(|| UploadError::Cancelled) + .and_then(|x| x); match &res { Ok(_) => {} diff --git a/pageserver/src/control_plane_client.rs b/pageserver/src/control_plane_client.rs index 950791ea48..61c7d03408 100644 --- a/pageserver/src/control_plane_client.rs +++ b/pageserver/src/control_plane_client.rs @@ -82,46 +82,29 @@ impl ControlPlaneClient { R: Serialize, T: DeserializeOwned, { - #[derive(thiserror::Error, Debug)] - enum RemoteAttemptError { - #[error("shutdown")] - Shutdown, - #[error("remote: {0}")] - Remote(reqwest::Error), - } - - match backoff::retry( + let res = backoff::retry( || async { let response = self .http_client .post(url.clone()) .json(&request) .send() - .await - .map_err(RemoteAttemptError::Remote)?; + .await?; - response - .error_for_status_ref() - .map_err(RemoteAttemptError::Remote)?; - response - .json::() - .await - .map_err(RemoteAttemptError::Remote) + response.error_for_status_ref()?; + response.json::().await }, |_| false, 3, u32::MAX, "calling control plane generation validation API", - backoff::Cancel::new(self.cancel.clone(), || RemoteAttemptError::Shutdown), + &self.cancel, ) .await - { - Err(RemoteAttemptError::Shutdown) => Err(RetryForeverError::ShuttingDown), - Err(RemoteAttemptError::Remote(_)) => { - panic!("We retry forever, this should never be reached"); - } - Ok(r) => Ok(r), - } + .ok_or(RetryForeverError::ShuttingDown)? + .expect("We retry forever, this should never be reached"); + + Ok(res) } } diff --git a/pageserver/src/deletion_queue/deleter.rs b/pageserver/src/deletion_queue/deleter.rs index 57421b1547..a75c73f2b1 100644 --- a/pageserver/src/deletion_queue/deleter.rs +++ b/pageserver/src/deletion_queue/deleter.rs @@ -77,9 +77,11 @@ impl Deleter { 3, 10, "executing deletion batch", - backoff::Cancel::new(self.cancel.clone(), || anyhow::anyhow!("Shutting down")), + &self.cancel, ) .await + .ok_or_else(|| anyhow::anyhow!("Shutting down")) + .and_then(|x| x) } /// Block until everything in accumulator has been executed diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index b801347c06..624c3e365f 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -3294,11 +3294,11 @@ impl Tenant { 3, u32::MAX, "persist_initdb_tar_zst", - backoff::Cancel::new(self.cancel.clone(), || anyhow::anyhow!("Cancelled")), + &self.cancel, ) - .await?; - - Ok(()) + .await + .ok_or_else(|| anyhow::anyhow!("Cancelled")) + .and_then(|x| x) } /// - run initdb to init temporary instance and get bootstrap data diff --git a/pageserver/src/tenant/delete.rs b/pageserver/src/tenant/delete.rs index 0dbaa3ec93..7c35914b61 100644 --- a/pageserver/src/tenant/delete.rs +++ b/pageserver/src/tenant/delete.rs @@ -91,9 +91,11 @@ async fn create_remote_delete_mark( FAILED_UPLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, "mark_upload", - backoff::Cancel::new(cancel.clone(), || anyhow::anyhow!("Cancelled")), + cancel, ) .await + .ok_or_else(|| anyhow::anyhow!("Cancelled")) + .and_then(|x| x) .context("mark_upload")?; Ok(()) @@ -187,9 +189,11 @@ async fn remove_tenant_remote_delete_mark( FAILED_UPLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, "remove_tenant_remote_delete_mark", - backoff::Cancel::new(cancel.clone(), || anyhow::anyhow!("Cancelled")), + cancel, ) .await + .ok_or_else(|| anyhow::anyhow!("Cancelled")) + .and_then(|x| x) .context("remove_tenant_remote_delete_mark")?; } Ok(()) diff --git a/pageserver/src/tenant/remote_timeline_client.rs b/pageserver/src/tenant/remote_timeline_client.rs index 2e429ee9bc..831a073d17 100644 --- a/pageserver/src/tenant/remote_timeline_client.rs +++ b/pageserver/src/tenant/remote_timeline_client.rs @@ -1046,9 +1046,11 @@ impl RemoteTimelineClient { // when executed as part of tenant deletion this happens in the background 2, "persist_index_part_with_deleted_flag", - backoff::Cancel::new(self.cancel.clone(), || anyhow::anyhow!("Cancelled")), + &self.cancel, ) - .await?; + .await + .ok_or_else(|| anyhow::anyhow!("Cancelled")) + .and_then(|x| x)?; // all good, disarm the guard and mark as success ScopeGuard::into_inner(undo_deleted_at); @@ -1083,9 +1085,11 @@ impl RemoteTimelineClient { FAILED_DOWNLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, "preserve_initdb_tar_zst", - backoff::Cancel::new(cancel.clone(), || anyhow::anyhow!("Cancelled!")), + &cancel.clone(), ) .await + .ok_or_else(|| anyhow::anyhow!("Cancellled")) + .and_then(|x| x) .context("backing up initdb archive")?; Ok(()) } @@ -1141,6 +1145,8 @@ impl RemoteTimelineClient { // taking the burden of listing all the layers that we already know we should delete. self.deletion_queue_client.flush_immediate().await?; + let cancel = shutdown_token(); + let remaining = backoff::retry( || async { self.storage_impl @@ -1151,9 +1157,11 @@ impl RemoteTimelineClient { FAILED_DOWNLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, "list_prefixes", - backoff::Cancel::new(shutdown_token(), || anyhow::anyhow!("Cancelled!")), + &cancel, ) .await + .ok_or_else(|| anyhow::anyhow!("Cancelled!")) + .and_then(|x| x) .context("list prefixes")?; // We will delete the current index_part object last, since it acts as a deletion diff --git a/pageserver/src/tenant/remote_timeline_client/download.rs b/pageserver/src/tenant/remote_timeline_client/download.rs index b84b5ca33b..2c50726b43 100644 --- a/pageserver/src/tenant/remote_timeline_client/download.rs +++ b/pageserver/src/tenant/remote_timeline_client/download.rs @@ -76,7 +76,6 @@ pub async fn download_layer_file<'a>( // If pageserver crashes the temp file will be deleted on startup and re-downloaded. let temp_file_path = path_with_suffix_extension(&local_path, TEMP_DOWNLOAD_EXTENSION); - let cancel_inner = cancel.clone(); let (mut destination_file, bytes_amount) = download_retry( || async { let destination_file = tokio::fs::File::create(&temp_file_path) @@ -87,7 +86,7 @@ pub async fn download_layer_file<'a>( // Cancellation safety: it is safe to cancel this future, because it isn't writing to a local // file: the write to local file doesn't start until after the request header is returned // and we start draining the body stream below - let download = download_cancellable(&cancel_inner, storage.download(&remote_path)) + let download = download_cancellable(cancel, storage.download(&remote_path)) .await .with_context(|| { format!( @@ -107,7 +106,7 @@ pub async fn download_layer_file<'a>( // we will imminiently try and write to again. let bytes_amount: u64 = match timeout_cancellable( DOWNLOAD_TIMEOUT, - &cancel_inner, + cancel, tokio::io::copy_buf(&mut reader, &mut destination_file), ) .await @@ -386,9 +385,11 @@ pub(super) async fn download_index_part( FAILED_DOWNLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, "listing index_part files", - backoff::Cancel::new(cancel.clone(), || anyhow::anyhow!("Cancelled")), + &cancel, ) .await + .ok_or_else(|| anyhow::anyhow!("Cancelled")) + .and_then(|x| x) .map_err(DownloadError::Other)?; // General case logic for which index to use: the latest index whose generation @@ -510,7 +511,7 @@ pub(crate) async fn download_initdb_tar_zst( /// Helper function to handle retries for a download operation. /// -/// Remote operations can fail due to rate limits (IAM, S3), spurious network +/// Remote operations can fail due to rate limits (S3), spurious network /// problems, or other external reasons. Retry FAILED_DOWNLOAD_RETRIES times, /// with backoff. /// @@ -530,9 +531,11 @@ where FAILED_DOWNLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, description, - backoff::Cancel::new(cancel.clone(), || DownloadError::Cancelled), + cancel, ) .await + .ok_or_else(|| DownloadError::Cancelled) + .and_then(|x| x) } async fn download_retry_forever( @@ -550,7 +553,9 @@ where FAILED_DOWNLOAD_WARN_THRESHOLD, u32::MAX, description, - backoff::Cancel::new(cancel, || DownloadError::Cancelled), + &cancel, ) .await + .ok_or_else(|| DownloadError::Cancelled) + .and_then(|x| x) } diff --git a/pageserver/src/tenant/remote_timeline_client/upload.rs b/pageserver/src/tenant/remote_timeline_client/upload.rs index 76df9ba5c4..e8ba1d3d6e 100644 --- a/pageserver/src/tenant/remote_timeline_client/upload.rs +++ b/pageserver/src/tenant/remote_timeline_client/upload.rs @@ -188,16 +188,18 @@ pub(crate) async fn time_travel_recover_tenant( backoff::retry( || async { storage - .time_travel_recover(Some(prefix), timestamp, done_if_after, cancel.clone()) + .time_travel_recover(Some(prefix), timestamp, done_if_after, cancel) .await }, |e| !matches!(e, TimeTravelError::Other(_)), warn_after, max_attempts, "time travel recovery of tenant prefix", - backoff::Cancel::new(cancel.clone(), || TimeTravelError::Cancelled), + cancel, ) - .await?; + .await + .ok_or_else(|| TimeTravelError::Cancelled) + .and_then(|x| x)?; } Ok(()) } diff --git a/pageserver/src/tenant/secondary/downloader.rs b/pageserver/src/tenant/secondary/downloader.rs index 702c0b1ec1..55af4f9f2b 100644 --- a/pageserver/src/tenant/secondary/downloader.rs +++ b/pageserver/src/tenant/secondary/downloader.rs @@ -537,11 +537,11 @@ impl<'a> TenantDownloader<'a> { FAILED_DOWNLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, "download heatmap", - backoff::Cancel::new(self.secondary_state.cancel.clone(), || { - UpdateError::Cancelled - }), + &self.secondary_state.cancel, ) - .await?; + .await + .ok_or_else(|| UpdateError::Cancelled) + .and_then(|x| x)?; SECONDARY_MODE.download_heatmap.inc(); diff --git a/pageserver/src/tenant/secondary/heatmap_uploader.rs b/pageserver/src/tenant/secondary/heatmap_uploader.rs index df865658a4..fff29b2487 100644 --- a/pageserver/src/tenant/secondary/heatmap_uploader.rs +++ b/pageserver/src/tenant/secondary/heatmap_uploader.rs @@ -426,9 +426,11 @@ async fn upload_tenant_heatmap( 3, u32::MAX, "Uploading heatmap", - backoff::Cancel::new(tenant_cancel.clone(), || anyhow::anyhow!("Shutting down")), + &tenant_cancel, ) .await + .ok_or_else(|| anyhow::anyhow!("Shutting down")) + .and_then(|x| x) { if tenant_cancel.is_cancelled() { return Err(UploadHeatmapError::Cancelled); diff --git a/proxy/src/context/parquet.rs b/proxy/src/context/parquet.rs index e920d7be01..8510c5c586 100644 --- a/proxy/src/context/parquet.rs +++ b/proxy/src/context/parquet.rs @@ -315,9 +315,11 @@ async fn upload_parquet( FAILED_UPLOAD_MAX_RETRIES, "request_data_upload", // we don't want cancellation to interrupt here, so we make a dummy cancel token - backoff::Cancel::new(CancellationToken::new(), || anyhow::anyhow!("Cancelled")), + &CancellationToken::new(), ) .await + .ok_or_else(|| anyhow::anyhow!("Cancelled")) + .and_then(|x| x) .context("request_data_upload")?; Ok(buffer.writer()) diff --git a/safekeeper/src/wal_backup.rs b/safekeeper/src/wal_backup.rs index c47381351d..df99244770 100644 --- a/safekeeper/src/wal_backup.rs +++ b/safekeeper/src/wal_backup.rs @@ -558,16 +558,17 @@ pub async fn delete_timeline(ttid: &TenantTimelineId) -> Result<()> { backoff::retry( || async { let files = storage.list_files(Some(&remote_path)).await?; - storage.delete_objects(&files).await?; - Ok(()) + storage.delete_objects(&files).await }, |_| false, 3, 10, "executing WAL segments deletion batch", - backoff::Cancel::new(token, || anyhow::anyhow!("canceled")), + &token, ) - .await?; + .await + .ok_or_else(|| anyhow::anyhow!("canceled")) + .and_then(|x| x)?; Ok(()) } From e196d974cc585341ee38f8fd6b54c257a3ad78a4 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Tue, 6 Feb 2024 10:34:16 +0100 Subject: [PATCH 11/72] pagebench: actually implement `--num_clients` (#6640) Will need this to validate per-tenant throttling in https://github.com/neondatabase/neon/issues/5899 --- .../pagebench/src/cmd/getpage_latest_lsn.rs | 139 ++++++++++-------- 1 file changed, 78 insertions(+), 61 deletions(-) diff --git a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs index 400b5476b7..aa809d8d26 100644 --- a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs +++ b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs @@ -79,6 +79,12 @@ impl KeyRange { } } +#[derive(PartialEq, Eq, Hash, Copy, Clone)] +struct WorkerId { + timeline: TenantTimelineId, + num_client: usize, // from 0..args.num_clients +} + #[derive(serde::Serialize)] struct Output { total: request_stats::Output, @@ -206,7 +212,7 @@ async fn main_impl( let live_stats = Arc::new(LiveStats::default()); - let num_client_tasks = timelines.len(); + let num_client_tasks = args.num_clients.get() * timelines.len(); let num_live_stats_dump = 1; let num_work_sender_tasks = 1; let num_main_impl = 1; @@ -235,19 +241,25 @@ async fn main_impl( let cancel = CancellationToken::new(); - let mut work_senders: HashMap = HashMap::new(); + let mut work_senders: HashMap = HashMap::new(); let mut tasks = Vec::new(); - for tl in &timelines { - let (sender, receiver) = tokio::sync::mpsc::channel(10); // TODO: not sure what the implications of this are - work_senders.insert(*tl, sender); - tasks.push(tokio::spawn(client( - args, - *tl, - Arc::clone(&start_work_barrier), - receiver, - Arc::clone(&live_stats), - cancel.clone(), - ))); + for timeline in timelines.iter().cloned() { + for num_client in 0..args.num_clients.get() { + let (sender, receiver) = tokio::sync::mpsc::channel(10); // TODO: not sure what the implications of this are + let worker_id = WorkerId { + timeline, + num_client, + }; + work_senders.insert(worker_id, sender); + tasks.push(tokio::spawn(client( + args, + worker_id, + Arc::clone(&start_work_barrier), + receiver, + Arc::clone(&live_stats), + cancel.clone(), + ))); + } } let work_sender: Pin>> = { @@ -271,7 +283,10 @@ async fn main_impl( let (rel_tag, block_no) = key_to_rel_block(key).expect("we filter non-rel-block keys out above"); ( - r.timeline, + WorkerId { + timeline: r.timeline, + num_client: rng.gen_range(0..args.num_clients.get()), + }, PagestreamGetPageRequest { latest: rng.gen_bool(args.req_latest_probability), lsn: r.timeline_lsn, @@ -289,56 +304,54 @@ async fn main_impl( }), Some(rps_limit) => Box::pin(async move { let period = Duration::from_secs_f64(1.0 / (rps_limit as f64)); - let make_timeline_task: &dyn Fn( - TenantTimelineId, - ) - -> Pin>> = &|timeline| { - let sender = work_senders.get(&timeline).unwrap(); - let ranges: Vec = all_ranges - .iter() - .filter(|r| r.timeline == timeline) - .cloned() - .collect(); - let weights = rand::distributions::weighted::WeightedIndex::new( - ranges.iter().map(|v| v.len()), - ) - .unwrap(); + let make_task: &dyn Fn(WorkerId) -> Pin>> = + &|worker_id| { + let sender = work_senders.get(&worker_id).unwrap(); + let ranges: Vec = all_ranges + .iter() + .filter(|r| r.timeline == worker_id.timeline) + .cloned() + .collect(); + let weights = rand::distributions::weighted::WeightedIndex::new( + ranges.iter().map(|v| v.len()), + ) + .unwrap(); - let cancel = cancel.clone(); - Box::pin(async move { - let mut ticker = tokio::time::interval(period); - ticker.set_missed_tick_behavior( - /* TODO review this choice */ - tokio::time::MissedTickBehavior::Burst, - ); - while !cancel.is_cancelled() { - ticker.tick().await; - let req = { - let mut rng = rand::thread_rng(); - let r = &ranges[weights.sample(&mut rng)]; - let key: i128 = rng.gen_range(r.start..r.end); - let key = Key::from_i128(key); - assert!(is_rel_block_key(&key)); - let (rel_tag, block_no) = key_to_rel_block(key) - .expect("we filter non-rel-block keys out above"); - PagestreamGetPageRequest { - latest: rng.gen_bool(args.req_latest_probability), - lsn: r.timeline_lsn, - rel: rel_tag, - blkno: block_no, + let cancel = cancel.clone(); + Box::pin(async move { + let mut ticker = tokio::time::interval(period); + ticker.set_missed_tick_behavior( + /* TODO review this choice */ + tokio::time::MissedTickBehavior::Burst, + ); + while !cancel.is_cancelled() { + ticker.tick().await; + let req = { + let mut rng = rand::thread_rng(); + let r = &ranges[weights.sample(&mut rng)]; + let key: i128 = rng.gen_range(r.start..r.end); + let key = Key::from_i128(key); + assert!(is_rel_block_key(&key)); + let (rel_tag, block_no) = key_to_rel_block(key) + .expect("we filter non-rel-block keys out above"); + PagestreamGetPageRequest { + latest: rng.gen_bool(args.req_latest_probability), + lsn: r.timeline_lsn, + rel: rel_tag, + blkno: block_no, + } + }; + if sender.send(req).await.is_err() { + assert!( + cancel.is_cancelled(), + "client has gone away unexpectedly" + ); } - }; - if sender.send(req).await.is_err() { - assert!(cancel.is_cancelled(), "client has gone away unexpectedly"); } - } - }) - }; + }) + }; - let tasks: Vec<_> = work_senders - .keys() - .map(|tl| make_timeline_task(*tl)) - .collect(); + let tasks: Vec<_> = work_senders.keys().map(|tl| make_task(*tl)).collect(); start_work_barrier.wait().await; @@ -390,12 +403,16 @@ async fn main_impl( #[instrument(skip_all)] async fn client( args: &'static Args, - timeline: TenantTimelineId, + id: WorkerId, start_work_barrier: Arc, mut work: tokio::sync::mpsc::Receiver, live_stats: Arc, cancel: CancellationToken, ) { + let WorkerId { + timeline, + num_client: _, + } = id; let client = pageserver_client::page_service::Client::new(args.page_service_connstring.clone()) .await .unwrap(); From edcde05c1cdf75f9bd5f0669b95ef61946d25549 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Tue, 6 Feb 2024 10:44:49 +0100 Subject: [PATCH 12/72] refactor(walredo): split up the massive `walredo.rs` (#6583) Part of https://github.com/neondatabase/neon/issues/6581 --- pageserver/src/walredo.rs | 825 +----------------- pageserver/src/walredo/apply_neon.rs | 235 +++++ pageserver/src/walredo/process.rs | 406 +++++++++ .../src/walredo/process/no_leak_child.rs | 126 +++ pageserver/src/walredo/process/protocol.rs | 57 ++ 5 files changed, 848 insertions(+), 801 deletions(-) create mode 100644 pageserver/src/walredo/apply_neon.rs create mode 100644 pageserver/src/walredo/process.rs create mode 100644 pageserver/src/walredo/process/no_leak_child.rs create mode 100644 pageserver/src/walredo/process/protocol.rs diff --git a/pageserver/src/walredo.rs b/pageserver/src/walredo.rs index 5bc897b730..773e5fc051 100644 --- a/pageserver/src/walredo.rs +++ b/pageserver/src/walredo.rs @@ -17,71 +17,30 @@ //! records. It achieves it by dropping privileges before replaying //! any WAL records, so that even if an attacker hijacks the Postgres //! process, he cannot escape out of it. -//! -use anyhow::Context; -use byteorder::{ByteOrder, LittleEndian}; -use bytes::{BufMut, Bytes, BytesMut}; -use nix::poll::*; -use pageserver_api::models::WalRedoManagerStatus; -use pageserver_api::shard::TenantShardId; -use serde::Serialize; -use std::collections::VecDeque; -use std::io; -use std::io::prelude::*; -use std::ops::{Deref, DerefMut}; -use std::os::unix::io::AsRawFd; -use std::process::Stdio; -use std::process::{Child, ChildStdin, ChildStdout, Command}; -use std::sync::{Arc, Mutex, MutexGuard, RwLock}; -use std::time::Duration; -use std::time::Instant; -use tracing::*; -use utils::{bin_ser::BeSer, lsn::Lsn, nonblock::set_nonblock}; -#[cfg(feature = "testing")] -use std::sync::atomic::{AtomicUsize, Ordering}; +/// Process lifecycle and abstracction for the IPC protocol. +mod process; + +/// Code to apply [`NeonWalRecord`]s. +mod apply_neon; use crate::config::PageServerConf; use crate::metrics::{ - WalRedoKillCause, WAL_REDO_BYTES_HISTOGRAM, WAL_REDO_PROCESS_COUNTERS, - WAL_REDO_PROCESS_LAUNCH_DURATION_HISTOGRAM, WAL_REDO_RECORDS_HISTOGRAM, - WAL_REDO_RECORD_COUNTER, WAL_REDO_TIME, + WAL_REDO_BYTES_HISTOGRAM, WAL_REDO_PROCESS_LAUNCH_DURATION_HISTOGRAM, + WAL_REDO_RECORDS_HISTOGRAM, WAL_REDO_TIME, }; use crate::repository::Key; use crate::walrecord::NeonWalRecord; - -use pageserver_api::key::{key_to_rel_block, key_to_slru_block}; -use pageserver_api::reltag::{RelTag, SlruKind}; -use postgres_ffi::pg_constants; -use postgres_ffi::relfile_utils::VISIBILITYMAP_FORKNUM; -use postgres_ffi::v14::nonrelfile_utils::{ - mx_offset_to_flags_bitshift, mx_offset_to_flags_offset, mx_offset_to_member_offset, - transaction_id_set_status, -}; -use postgres_ffi::BLCKSZ; - -/// -/// `RelTag` + block number (`blknum`) gives us a unique id of the page in the cluster. -/// -/// In Postgres `BufferTag` structure is used for exactly the same purpose. -/// [See more related comments here](https://github.com/postgres/postgres/blob/99c5852e20a0987eca1c38ba0c09329d4076b6a0/src/include/storage/buf_internals.h#L91). -/// -#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Serialize)] -pub(crate) struct BufferTag { - pub rel: RelTag, - pub blknum: u32, -} - -struct ProcessInput { - stdin: ChildStdin, - n_requests: usize, -} - -struct ProcessOutput { - stdout: ChildStdout, - pending_responses: VecDeque>, - n_processed_responses: usize, -} +use anyhow::Context; +use bytes::{Bytes, BytesMut}; +use pageserver_api::key::key_to_rel_block; +use pageserver_api::models::WalRedoManagerStatus; +use pageserver_api::shard::TenantShardId; +use std::sync::{Arc, RwLock}; +use std::time::Duration; +use std::time::Instant; +use tracing::*; +use utils::lsn::Lsn; /// /// This is the real implementation that uses a Postgres process to @@ -94,22 +53,7 @@ pub struct PostgresRedoManager { tenant_shard_id: TenantShardId, conf: &'static PageServerConf, last_redo_at: std::sync::Mutex>, - redo_process: RwLock>>, -} - -/// Can this request be served by neon redo functions -/// or we need to pass it to wal-redo postgres process? -fn can_apply_in_neon(rec: &NeonWalRecord) -> bool { - // Currently, we don't have bespoken Rust code to replay any - // Postgres WAL records. But everything else is handled in neon. - #[allow(clippy::match_like_matches_macro)] - match rec { - NeonWalRecord::Postgres { - will_init: _, - rec: _, - } => false, - _ => true, - } + redo_process: RwLock>>, } /// @@ -139,10 +83,10 @@ impl PostgresRedoManager { let base_img_lsn = base_img.as_ref().map(|p| p.0).unwrap_or(Lsn::INVALID); let mut img = base_img.map(|p| p.1); - let mut batch_neon = can_apply_in_neon(&records[0].1); + let mut batch_neon = apply_neon::can_apply_in_neon(&records[0].1); let mut batch_start = 0; for (i, record) in records.iter().enumerate().skip(1) { - let rec_neon = can_apply_in_neon(&record.1); + let rec_neon = apply_neon::can_apply_in_neon(&record.1); if rec_neon != batch_neon { let result = if batch_neon { @@ -248,7 +192,7 @@ impl PostgresRedoManager { let mut n_attempts = 0u32; loop { // launch the WAL redo process on first use - let proc: Arc = { + let proc: Arc = { let proc_guard = self.redo_process.read().unwrap(); match &*proc_guard { None => { @@ -259,7 +203,7 @@ impl PostgresRedoManager { None => { let start = Instant::now(); let proc = Arc::new( - WalRedoProcess::launch( + process::WalRedoProcess::launch( self.conf, self.tenant_shard_id, pg_version, @@ -287,9 +231,8 @@ impl PostgresRedoManager { let started_at = std::time::Instant::now(); // Relational WAL records are applied using wal-redo-postgres - let buf_tag = BufferTag { rel, blknum }; let result = proc - .apply_wal_records(buf_tag, &base_img, records, wal_redo_timeout) + .apply_wal_records(rel, blknum, &base_img, records, wal_redo_timeout) .context("apply_wal_records"); let duration = started_at.elapsed(); @@ -416,732 +359,12 @@ impl PostgresRedoManager { _record_lsn: Lsn, record: &NeonWalRecord, ) -> anyhow::Result<()> { - match record { - NeonWalRecord::Postgres { - will_init: _, - rec: _, - } => { - anyhow::bail!("tried to pass postgres wal record to neon WAL redo"); - } - NeonWalRecord::ClearVisibilityMapFlags { - new_heap_blkno, - old_heap_blkno, - flags, - } => { - // sanity check that this is modifying the correct relation - let (rel, blknum) = key_to_rel_block(key).context("invalid record")?; - assert!( - rel.forknum == VISIBILITYMAP_FORKNUM, - "ClearVisibilityMapFlags record on unexpected rel {}", - rel - ); - if let Some(heap_blkno) = *new_heap_blkno { - // Calculate the VM block and offset that corresponds to the heap block. - let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno); - let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno); - let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno); - - // Check that we're modifying the correct VM block. - assert!(map_block == blknum); - - // equivalent to PageGetContents(page) - let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..]; - - map[map_byte as usize] &= !(flags << map_offset); - } - - // Repeat for 'old_heap_blkno', if any - if let Some(heap_blkno) = *old_heap_blkno { - let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno); - let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno); - let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno); - - assert!(map_block == blknum); - - let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..]; - - map[map_byte as usize] &= !(flags << map_offset); - } - } - // Non-relational WAL records are handled here, with custom code that has the - // same effects as the corresponding Postgres WAL redo function. - NeonWalRecord::ClogSetCommitted { xids, timestamp } => { - let (slru_kind, segno, blknum) = - key_to_slru_block(key).context("invalid record")?; - assert_eq!( - slru_kind, - SlruKind::Clog, - "ClogSetCommitted record with unexpected key {}", - key - ); - for &xid in xids { - let pageno = xid / pg_constants::CLOG_XACTS_PER_PAGE; - let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; - let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; - - // Check that we're modifying the correct CLOG block. - assert!( - segno == expected_segno, - "ClogSetCommitted record for XID {} with unexpected key {}", - xid, - key - ); - assert!( - blknum == expected_blknum, - "ClogSetCommitted record for XID {} with unexpected key {}", - xid, - key - ); - - transaction_id_set_status( - xid, - pg_constants::TRANSACTION_STATUS_COMMITTED, - page, - ); - } - - // Append the timestamp - if page.len() == BLCKSZ as usize + 8 { - page.truncate(BLCKSZ as usize); - } - if page.len() == BLCKSZ as usize { - page.extend_from_slice(×tamp.to_be_bytes()); - } else { - warn!( - "CLOG blk {} in seg {} has invalid size {}", - blknum, - segno, - page.len() - ); - } - } - NeonWalRecord::ClogSetAborted { xids } => { - let (slru_kind, segno, blknum) = - key_to_slru_block(key).context("invalid record")?; - assert_eq!( - slru_kind, - SlruKind::Clog, - "ClogSetAborted record with unexpected key {}", - key - ); - for &xid in xids { - let pageno = xid / pg_constants::CLOG_XACTS_PER_PAGE; - let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; - let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; - - // Check that we're modifying the correct CLOG block. - assert!( - segno == expected_segno, - "ClogSetAborted record for XID {} with unexpected key {}", - xid, - key - ); - assert!( - blknum == expected_blknum, - "ClogSetAborted record for XID {} with unexpected key {}", - xid, - key - ); - - transaction_id_set_status(xid, pg_constants::TRANSACTION_STATUS_ABORTED, page); - } - } - NeonWalRecord::MultixactOffsetCreate { mid, moff } => { - let (slru_kind, segno, blknum) = - key_to_slru_block(key).context("invalid record")?; - assert_eq!( - slru_kind, - SlruKind::MultiXactOffsets, - "MultixactOffsetCreate record with unexpected key {}", - key - ); - // Compute the block and offset to modify. - // See RecordNewMultiXact in PostgreSQL sources. - let pageno = mid / pg_constants::MULTIXACT_OFFSETS_PER_PAGE as u32; - let entryno = mid % pg_constants::MULTIXACT_OFFSETS_PER_PAGE as u32; - let offset = (entryno * 4) as usize; - - // Check that we're modifying the correct multixact-offsets block. - let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; - let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; - assert!( - segno == expected_segno, - "MultiXactOffsetsCreate record for multi-xid {} with unexpected key {}", - mid, - key - ); - assert!( - blknum == expected_blknum, - "MultiXactOffsetsCreate record for multi-xid {} with unexpected key {}", - mid, - key - ); - - LittleEndian::write_u32(&mut page[offset..offset + 4], *moff); - } - NeonWalRecord::MultixactMembersCreate { moff, members } => { - let (slru_kind, segno, blknum) = - key_to_slru_block(key).context("invalid record")?; - assert_eq!( - slru_kind, - SlruKind::MultiXactMembers, - "MultixactMembersCreate record with unexpected key {}", - key - ); - for (i, member) in members.iter().enumerate() { - let offset = moff + i as u32; - - // Compute the block and offset to modify. - // See RecordNewMultiXact in PostgreSQL sources. - let pageno = offset / pg_constants::MULTIXACT_MEMBERS_PER_PAGE as u32; - let memberoff = mx_offset_to_member_offset(offset); - let flagsoff = mx_offset_to_flags_offset(offset); - let bshift = mx_offset_to_flags_bitshift(offset); - - // Check that we're modifying the correct multixact-members block. - let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; - let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; - assert!( - segno == expected_segno, - "MultiXactMembersCreate record for offset {} with unexpected key {}", - moff, - key - ); - assert!( - blknum == expected_blknum, - "MultiXactMembersCreate record for offset {} with unexpected key {}", - moff, - key - ); - - let mut flagsval = LittleEndian::read_u32(&page[flagsoff..flagsoff + 4]); - flagsval &= !(((1 << pg_constants::MXACT_MEMBER_BITS_PER_XACT) - 1) << bshift); - flagsval |= member.status << bshift; - LittleEndian::write_u32(&mut page[flagsoff..flagsoff + 4], flagsval); - LittleEndian::write_u32(&mut page[memberoff..memberoff + 4], member.xid); - } - } - } + apply_neon::apply_in_neon(record, key, page)?; Ok(()) } } -struct WalRedoProcess { - #[allow(dead_code)] - conf: &'static PageServerConf, - tenant_shard_id: TenantShardId, - // Some() on construction, only becomes None on Drop. - child: Option, - stdout: Mutex, - stdin: Mutex, - /// Counter to separate same sized walredo inputs failing at the same millisecond. - #[cfg(feature = "testing")] - dump_sequence: AtomicUsize, -} - -impl WalRedoProcess { - // - // Start postgres binary in special WAL redo mode. - // - #[instrument(skip_all,fields(tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), pg_version=pg_version))] - fn launch( - conf: &'static PageServerConf, - tenant_shard_id: TenantShardId, - pg_version: u32, - ) -> anyhow::Result { - let pg_bin_dir_path = conf.pg_bin_dir(pg_version).context("pg_bin_dir")?; // TODO these should be infallible. - let pg_lib_dir_path = conf.pg_lib_dir(pg_version).context("pg_lib_dir")?; - - // Start postgres itself - let child = Command::new(pg_bin_dir_path.join("postgres")) - // the first arg must be --wal-redo so the child process enters into walredo mode - .arg("--wal-redo") - // the child doesn't process this arg, but, having it in the argv helps indentify the - // walredo process for a particular tenant when debugging a pagserver - .args(["--tenant-shard-id", &format!("{tenant_shard_id}")]) - .stdin(Stdio::piped()) - .stderr(Stdio::piped()) - .stdout(Stdio::piped()) - .env_clear() - .env("LD_LIBRARY_PATH", &pg_lib_dir_path) - .env("DYLD_LIBRARY_PATH", &pg_lib_dir_path) - // NB: The redo process is not trusted after we sent it the first - // walredo work. Before that, it is trusted. Specifically, we trust - // it to - // 1. close all file descriptors except stdin, stdout, stderr because - // pageserver might not be 100% diligent in setting FD_CLOEXEC on all - // the files it opens, and - // 2. to use seccomp to sandbox itself before processing the first - // walredo request. - .spawn_no_leak_child(tenant_shard_id) - .context("spawn process")?; - WAL_REDO_PROCESS_COUNTERS.started.inc(); - let mut child = scopeguard::guard(child, |child| { - error!("killing wal-redo-postgres process due to a problem during launch"); - child.kill_and_wait(WalRedoKillCause::Startup); - }); - - let stdin = child.stdin.take().unwrap(); - let stdout = child.stdout.take().unwrap(); - let stderr = child.stderr.take().unwrap(); - let stderr = tokio::process::ChildStderr::from_std(stderr) - .context("convert to tokio::ChildStderr")?; - macro_rules! set_nonblock_or_log_err { - ($file:ident) => {{ - let res = set_nonblock($file.as_raw_fd()); - if let Err(e) = &res { - error!(error = %e, file = stringify!($file), pid = child.id(), "set_nonblock failed"); - } - res - }}; - } - set_nonblock_or_log_err!(stdin)?; - set_nonblock_or_log_err!(stdout)?; - - // all fallible operations post-spawn are complete, so get rid of the guard - let child = scopeguard::ScopeGuard::into_inner(child); - - tokio::spawn( - async move { - scopeguard::defer! { - debug!("wal-redo-postgres stderr_logger_task finished"); - crate::metrics::WAL_REDO_PROCESS_COUNTERS.active_stderr_logger_tasks_finished.inc(); - } - debug!("wal-redo-postgres stderr_logger_task started"); - crate::metrics::WAL_REDO_PROCESS_COUNTERS.active_stderr_logger_tasks_started.inc(); - - use tokio::io::AsyncBufReadExt; - let mut stderr_lines = tokio::io::BufReader::new(stderr); - let mut buf = Vec::new(); - let res = loop { - buf.clear(); - // TODO we don't trust the process to cap its stderr length. - // Currently it can do unbounded Vec allocation. - match stderr_lines.read_until(b'\n', &mut buf).await { - Ok(0) => break Ok(()), // eof - Ok(num_bytes) => { - let output = String::from_utf8_lossy(&buf[..num_bytes]); - error!(%output, "received output"); - } - Err(e) => { - break Err(e); - } - } - }; - match res { - Ok(()) => (), - Err(e) => { - error!(error=?e, "failed to read from walredo stderr"); - } - } - }.instrument(tracing::info_span!(parent: None, "wal-redo-postgres-stderr", pid = child.id(), tenant_id = %tenant_shard_id.tenant_id, shard_id = %tenant_shard_id.shard_slug(), %pg_version)) - ); - - Ok(Self { - conf, - tenant_shard_id, - child: Some(child), - stdin: Mutex::new(ProcessInput { - stdin, - n_requests: 0, - }), - stdout: Mutex::new(ProcessOutput { - stdout, - pending_responses: VecDeque::new(), - n_processed_responses: 0, - }), - #[cfg(feature = "testing")] - dump_sequence: AtomicUsize::default(), - }) - } - - fn id(&self) -> u32 { - self.child - .as_ref() - .expect("must not call this during Drop") - .id() - } - - // Apply given WAL records ('records') over an old page image. Returns - // new page image. - // - #[instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), pid=%self.id()))] - fn apply_wal_records( - &self, - tag: BufferTag, - base_img: &Option, - records: &[(Lsn, NeonWalRecord)], - wal_redo_timeout: Duration, - ) -> anyhow::Result { - let input = self.stdin.lock().unwrap(); - - // Serialize all the messages to send the WAL redo process first. - // - // This could be problematic if there are millions of records to replay, - // but in practice the number of records is usually so small that it doesn't - // matter, and it's better to keep this code simple. - // - // Most requests start with a before-image with BLCKSZ bytes, followed by - // by some other WAL records. Start with a buffer that can hold that - // comfortably. - let mut writebuf: Vec = Vec::with_capacity((BLCKSZ as usize) * 3); - build_begin_redo_for_block_msg(tag, &mut writebuf); - if let Some(img) = base_img { - build_push_page_msg(tag, img, &mut writebuf); - } - for (lsn, rec) in records.iter() { - if let NeonWalRecord::Postgres { - will_init: _, - rec: postgres_rec, - } = rec - { - build_apply_record_msg(*lsn, postgres_rec, &mut writebuf); - } else { - anyhow::bail!("tried to pass neon wal record to postgres WAL redo"); - } - } - build_get_page_msg(tag, &mut writebuf); - WAL_REDO_RECORD_COUNTER.inc_by(records.len() as u64); - - let res = self.apply_wal_records0(&writebuf, input, wal_redo_timeout); - - if res.is_err() { - // not all of these can be caused by this particular input, however these are so rare - // in tests so capture all. - self.record_and_log(&writebuf); - } - - res - } - - fn apply_wal_records0( - &self, - writebuf: &[u8], - input: MutexGuard, - wal_redo_timeout: Duration, - ) -> anyhow::Result { - let mut proc = { input }; // TODO: remove this legacy rename, but this keep the patch small. - let mut nwrite = 0usize; - - while nwrite < writebuf.len() { - let mut stdin_pollfds = [PollFd::new(&proc.stdin, PollFlags::POLLOUT)]; - let n = loop { - match nix::poll::poll(&mut stdin_pollfds[..], wal_redo_timeout.as_millis() as i32) { - Err(nix::errno::Errno::EINTR) => continue, - res => break res, - } - }?; - - if n == 0 { - anyhow::bail!("WAL redo timed out"); - } - - // If 'stdin' is writeable, do write. - let in_revents = stdin_pollfds[0].revents().unwrap(); - if in_revents & (PollFlags::POLLERR | PollFlags::POLLOUT) != PollFlags::empty() { - nwrite += proc.stdin.write(&writebuf[nwrite..])?; - } - if in_revents.contains(PollFlags::POLLHUP) { - // We still have more data to write, but the process closed the pipe. - anyhow::bail!("WAL redo process closed its stdin unexpectedly"); - } - } - let request_no = proc.n_requests; - proc.n_requests += 1; - drop(proc); - - // To improve walredo performance we separate sending requests and receiving - // responses. Them are protected by different mutexes (output and input). - // If thread T1, T2, T3 send requests D1, D2, D3 to walredo process - // then there is not warranty that T1 will first granted output mutex lock. - // To address this issue we maintain number of sent requests, number of processed - // responses and ring buffer with pending responses. After sending response - // (under input mutex), threads remembers request number. Then it releases - // input mutex, locks output mutex and fetch in ring buffer all responses until - // its stored request number. The it takes correspondent element from - // pending responses ring buffer and truncate all empty elements from the front, - // advancing processed responses number. - - let mut output = self.stdout.lock().unwrap(); - let n_processed_responses = output.n_processed_responses; - while n_processed_responses + output.pending_responses.len() <= request_no { - // We expect the WAL redo process to respond with an 8k page image. We read it - // into this buffer. - let mut resultbuf = vec![0; BLCKSZ.into()]; - let mut nresult: usize = 0; // # of bytes read into 'resultbuf' so far - while nresult < BLCKSZ.into() { - let mut stdout_pollfds = [PollFd::new(&output.stdout, PollFlags::POLLIN)]; - // We do two things simultaneously: reading response from stdout - // and forward any logging information that the child writes to its stderr to the page server's log. - let n = loop { - match nix::poll::poll( - &mut stdout_pollfds[..], - wal_redo_timeout.as_millis() as i32, - ) { - Err(nix::errno::Errno::EINTR) => continue, - res => break res, - } - }?; - - if n == 0 { - anyhow::bail!("WAL redo timed out"); - } - - // If we have some data in stdout, read it to the result buffer. - let out_revents = stdout_pollfds[0].revents().unwrap(); - if out_revents & (PollFlags::POLLERR | PollFlags::POLLIN) != PollFlags::empty() { - nresult += output.stdout.read(&mut resultbuf[nresult..])?; - } - if out_revents.contains(PollFlags::POLLHUP) { - anyhow::bail!("WAL redo process closed its stdout unexpectedly"); - } - } - output - .pending_responses - .push_back(Some(Bytes::from(resultbuf))); - } - // Replace our request's response with None in `pending_responses`. - // Then make space in the ring buffer by clearing out any seqence of contiguous - // `None`'s from the front of `pending_responses`. - // NB: We can't pop_front() because other requests' responses because another - // requester might have grabbed the output mutex before us: - // T1: grab input mutex - // T1: send request_no 23 - // T1: release input mutex - // T2: grab input mutex - // T2: send request_no 24 - // T2: release input mutex - // T2: grab output mutex - // T2: n_processed_responses + output.pending_responses.len() <= request_no - // 23 0 24 - // T2: enters poll loop that reads stdout - // T2: put response for 23 into pending_responses - // T2: put response for 24 into pending_resposnes - // pending_responses now looks like this: Front Some(response_23) Some(response_24) Back - // T2: takes its response_24 - // pending_responses now looks like this: Front Some(response_23) None Back - // T2: does the while loop below - // pending_responses now looks like this: Front Some(response_23) None Back - // T2: releases output mutex - // T1: grabs output mutex - // T1: n_processed_responses + output.pending_responses.len() > request_no - // 23 2 23 - // T1: skips poll loop that reads stdout - // T1: takes its response_23 - // pending_responses now looks like this: Front None None Back - // T2: does the while loop below - // pending_responses now looks like this: Front Back - // n_processed_responses now has value 25 - let res = output.pending_responses[request_no - n_processed_responses] - .take() - .expect("we own this request_no, nobody else is supposed to take it"); - while let Some(front) = output.pending_responses.front() { - if front.is_none() { - output.pending_responses.pop_front(); - output.n_processed_responses += 1; - } else { - break; - } - } - Ok(res) - } - - #[cfg(feature = "testing")] - fn record_and_log(&self, writebuf: &[u8]) { - let millis = std::time::SystemTime::now() - .duration_since(std::time::SystemTime::UNIX_EPOCH) - .unwrap() - .as_millis(); - - let seq = self.dump_sequence.fetch_add(1, Ordering::Relaxed); - - // these files will be collected to an allure report - let filename = format!("walredo-{millis}-{}-{seq}.walredo", writebuf.len()); - - let path = self.conf.tenant_path(&self.tenant_shard_id).join(&filename); - - let res = std::fs::OpenOptions::new() - .write(true) - .create_new(true) - .read(true) - .open(path) - .and_then(|mut f| f.write_all(writebuf)); - - // trip up allowed_errors - if let Err(e) = res { - tracing::error!(target=%filename, length=writebuf.len(), "failed to write out the walredo errored input: {e}"); - } else { - tracing::error!(filename, "erroring walredo input saved"); - } - } - - #[cfg(not(feature = "testing"))] - fn record_and_log(&self, _: &[u8]) {} -} - -impl Drop for WalRedoProcess { - fn drop(&mut self) { - self.child - .take() - .expect("we only do this once") - .kill_and_wait(WalRedoKillCause::WalRedoProcessDrop); - // no way to wait for stderr_logger_task from Drop because that is async only - } -} - -/// Wrapper type around `std::process::Child` which guarantees that the child -/// will be killed and waited-for by this process before being dropped. -struct NoLeakChild { - tenant_id: TenantShardId, - child: Option, -} - -impl Deref for NoLeakChild { - type Target = Child; - - fn deref(&self) -> &Self::Target { - self.child.as_ref().expect("must not use from drop") - } -} - -impl DerefMut for NoLeakChild { - fn deref_mut(&mut self) -> &mut Self::Target { - self.child.as_mut().expect("must not use from drop") - } -} - -impl NoLeakChild { - fn spawn(tenant_id: TenantShardId, command: &mut Command) -> io::Result { - let child = command.spawn()?; - Ok(NoLeakChild { - tenant_id, - child: Some(child), - }) - } - - fn kill_and_wait(mut self, cause: WalRedoKillCause) { - let child = match self.child.take() { - Some(child) => child, - None => return, - }; - Self::kill_and_wait_impl(child, cause); - } - - #[instrument(skip_all, fields(pid=child.id(), ?cause))] - fn kill_and_wait_impl(mut child: Child, cause: WalRedoKillCause) { - scopeguard::defer! { - WAL_REDO_PROCESS_COUNTERS.killed_by_cause[cause].inc(); - } - let res = child.kill(); - if let Err(e) = res { - // This branch is very unlikely because: - // - We (= pageserver) spawned this process successfully, so, we're allowed to kill it. - // - This is the only place that calls .kill() - // - We consume `self`, so, .kill() can't be called twice. - // - If the process exited by itself or was killed by someone else, - // .kill() will still succeed because we haven't wait()'ed yet. - // - // So, if we arrive here, we have really no idea what happened, - // whether the PID stored in self.child is still valid, etc. - // If this function were fallible, we'd return an error, but - // since it isn't, all we can do is log an error and proceed - // with the wait(). - error!(error = %e, "failed to SIGKILL; subsequent wait() might fail or wait for wrong process"); - } - - match child.wait() { - Ok(exit_status) => { - info!(exit_status = %exit_status, "wait successful"); - } - Err(e) => { - error!(error = %e, "wait error; might leak the child process; it will show as zombie (defunct)"); - } - } - } -} - -impl Drop for NoLeakChild { - fn drop(&mut self) { - let child = match self.child.take() { - Some(child) => child, - None => return, - }; - let tenant_shard_id = self.tenant_id; - // Offload the kill+wait of the child process into the background. - // If someone stops the runtime, we'll leak the child process. - // We can ignore that case because we only stop the runtime on pageserver exit. - tokio::runtime::Handle::current().spawn(async move { - tokio::task::spawn_blocking(move || { - // Intentionally don't inherit the tracing context from whoever is dropping us. - // This thread here is going to outlive of our dropper. - let span = tracing::info_span!( - "walredo", - tenant_id = %tenant_shard_id.tenant_id, - shard_id = %tenant_shard_id.shard_slug() - ); - let _entered = span.enter(); - Self::kill_and_wait_impl(child, WalRedoKillCause::NoLeakChildDrop); - }) - .await - }); - } -} - -trait NoLeakChildCommandExt { - fn spawn_no_leak_child(&mut self, tenant_id: TenantShardId) -> io::Result; -} - -impl NoLeakChildCommandExt for Command { - fn spawn_no_leak_child(&mut self, tenant_id: TenantShardId) -> io::Result { - NoLeakChild::spawn(tenant_id, self) - } -} - -// Functions for constructing messages to send to the postgres WAL redo -// process. See pgxn/neon_walredo/walredoproc.c for -// explanation of the protocol. - -fn build_begin_redo_for_block_msg(tag: BufferTag, buf: &mut Vec) { - let len = 4 + 1 + 4 * 4; - - buf.put_u8(b'B'); - buf.put_u32(len as u32); - - tag.ser_into(buf) - .expect("serialize BufferTag should always succeed"); -} - -fn build_push_page_msg(tag: BufferTag, base_img: &[u8], buf: &mut Vec) { - assert!(base_img.len() == 8192); - - let len = 4 + 1 + 4 * 4 + base_img.len(); - - buf.put_u8(b'P'); - buf.put_u32(len as u32); - tag.ser_into(buf) - .expect("serialize BufferTag should always succeed"); - buf.put(base_img); -} - -fn build_apply_record_msg(endlsn: Lsn, rec: &[u8], buf: &mut Vec) { - let len = 4 + 8 + rec.len(); - - buf.put_u8(b'A'); - buf.put_u32(len as u32); - buf.put_u64(endlsn.0); - buf.put(rec); -} - -fn build_get_page_msg(tag: BufferTag, buf: &mut Vec) { - let len = 4 + 1 + 4 * 4; - - buf.put_u8(b'G'); - buf.put_u32(len as u32); - tag.ser_into(buf) - .expect("serialize BufferTag should always succeed"); -} - #[cfg(test)] mod tests { use super::PostgresRedoManager; diff --git a/pageserver/src/walredo/apply_neon.rs b/pageserver/src/walredo/apply_neon.rs new file mode 100644 index 0000000000..52899349c4 --- /dev/null +++ b/pageserver/src/walredo/apply_neon.rs @@ -0,0 +1,235 @@ +use crate::walrecord::NeonWalRecord; +use anyhow::Context; +use byteorder::{ByteOrder, LittleEndian}; +use bytes::BytesMut; +use pageserver_api::key::{key_to_rel_block, key_to_slru_block, Key}; +use pageserver_api::reltag::SlruKind; +use postgres_ffi::pg_constants; +use postgres_ffi::relfile_utils::VISIBILITYMAP_FORKNUM; +use postgres_ffi::v14::nonrelfile_utils::{ + mx_offset_to_flags_bitshift, mx_offset_to_flags_offset, mx_offset_to_member_offset, + transaction_id_set_status, +}; +use postgres_ffi::BLCKSZ; +use tracing::*; + +/// Can this request be served by neon redo functions +/// or we need to pass it to wal-redo postgres process? +pub(crate) fn can_apply_in_neon(rec: &NeonWalRecord) -> bool { + // Currently, we don't have bespoken Rust code to replay any + // Postgres WAL records. But everything else is handled in neon. + #[allow(clippy::match_like_matches_macro)] + match rec { + NeonWalRecord::Postgres { + will_init: _, + rec: _, + } => false, + _ => true, + } +} + +pub(crate) fn apply_in_neon( + record: &NeonWalRecord, + key: Key, + page: &mut BytesMut, +) -> Result<(), anyhow::Error> { + match record { + NeonWalRecord::Postgres { + will_init: _, + rec: _, + } => { + anyhow::bail!("tried to pass postgres wal record to neon WAL redo"); + } + NeonWalRecord::ClearVisibilityMapFlags { + new_heap_blkno, + old_heap_blkno, + flags, + } => { + // sanity check that this is modifying the correct relation + let (rel, blknum) = key_to_rel_block(key).context("invalid record")?; + assert!( + rel.forknum == VISIBILITYMAP_FORKNUM, + "ClearVisibilityMapFlags record on unexpected rel {}", + rel + ); + if let Some(heap_blkno) = *new_heap_blkno { + // Calculate the VM block and offset that corresponds to the heap block. + let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno); + let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno); + let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno); + + // Check that we're modifying the correct VM block. + assert!(map_block == blknum); + + // equivalent to PageGetContents(page) + let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..]; + + map[map_byte as usize] &= !(flags << map_offset); + } + + // Repeat for 'old_heap_blkno', if any + if let Some(heap_blkno) = *old_heap_blkno { + let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno); + let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno); + let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno); + + assert!(map_block == blknum); + + let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..]; + + map[map_byte as usize] &= !(flags << map_offset); + } + } + // Non-relational WAL records are handled here, with custom code that has the + // same effects as the corresponding Postgres WAL redo function. + NeonWalRecord::ClogSetCommitted { xids, timestamp } => { + let (slru_kind, segno, blknum) = key_to_slru_block(key).context("invalid record")?; + assert_eq!( + slru_kind, + SlruKind::Clog, + "ClogSetCommitted record with unexpected key {}", + key + ); + for &xid in xids { + let pageno = xid / pg_constants::CLOG_XACTS_PER_PAGE; + let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; + let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; + + // Check that we're modifying the correct CLOG block. + assert!( + segno == expected_segno, + "ClogSetCommitted record for XID {} with unexpected key {}", + xid, + key + ); + assert!( + blknum == expected_blknum, + "ClogSetCommitted record for XID {} with unexpected key {}", + xid, + key + ); + + transaction_id_set_status(xid, pg_constants::TRANSACTION_STATUS_COMMITTED, page); + } + + // Append the timestamp + if page.len() == BLCKSZ as usize + 8 { + page.truncate(BLCKSZ as usize); + } + if page.len() == BLCKSZ as usize { + page.extend_from_slice(×tamp.to_be_bytes()); + } else { + warn!( + "CLOG blk {} in seg {} has invalid size {}", + blknum, + segno, + page.len() + ); + } + } + NeonWalRecord::ClogSetAborted { xids } => { + let (slru_kind, segno, blknum) = key_to_slru_block(key).context("invalid record")?; + assert_eq!( + slru_kind, + SlruKind::Clog, + "ClogSetAborted record with unexpected key {}", + key + ); + for &xid in xids { + let pageno = xid / pg_constants::CLOG_XACTS_PER_PAGE; + let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; + let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; + + // Check that we're modifying the correct CLOG block. + assert!( + segno == expected_segno, + "ClogSetAborted record for XID {} with unexpected key {}", + xid, + key + ); + assert!( + blknum == expected_blknum, + "ClogSetAborted record for XID {} with unexpected key {}", + xid, + key + ); + + transaction_id_set_status(xid, pg_constants::TRANSACTION_STATUS_ABORTED, page); + } + } + NeonWalRecord::MultixactOffsetCreate { mid, moff } => { + let (slru_kind, segno, blknum) = key_to_slru_block(key).context("invalid record")?; + assert_eq!( + slru_kind, + SlruKind::MultiXactOffsets, + "MultixactOffsetCreate record with unexpected key {}", + key + ); + // Compute the block and offset to modify. + // See RecordNewMultiXact in PostgreSQL sources. + let pageno = mid / pg_constants::MULTIXACT_OFFSETS_PER_PAGE as u32; + let entryno = mid % pg_constants::MULTIXACT_OFFSETS_PER_PAGE as u32; + let offset = (entryno * 4) as usize; + + // Check that we're modifying the correct multixact-offsets block. + let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; + let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; + assert!( + segno == expected_segno, + "MultiXactOffsetsCreate record for multi-xid {} with unexpected key {}", + mid, + key + ); + assert!( + blknum == expected_blknum, + "MultiXactOffsetsCreate record for multi-xid {} with unexpected key {}", + mid, + key + ); + + LittleEndian::write_u32(&mut page[offset..offset + 4], *moff); + } + NeonWalRecord::MultixactMembersCreate { moff, members } => { + let (slru_kind, segno, blknum) = key_to_slru_block(key).context("invalid record")?; + assert_eq!( + slru_kind, + SlruKind::MultiXactMembers, + "MultixactMembersCreate record with unexpected key {}", + key + ); + for (i, member) in members.iter().enumerate() { + let offset = moff + i as u32; + + // Compute the block and offset to modify. + // See RecordNewMultiXact in PostgreSQL sources. + let pageno = offset / pg_constants::MULTIXACT_MEMBERS_PER_PAGE as u32; + let memberoff = mx_offset_to_member_offset(offset); + let flagsoff = mx_offset_to_flags_offset(offset); + let bshift = mx_offset_to_flags_bitshift(offset); + + // Check that we're modifying the correct multixact-members block. + let expected_segno = pageno / pg_constants::SLRU_PAGES_PER_SEGMENT; + let expected_blknum = pageno % pg_constants::SLRU_PAGES_PER_SEGMENT; + assert!( + segno == expected_segno, + "MultiXactMembersCreate record for offset {} with unexpected key {}", + moff, + key + ); + assert!( + blknum == expected_blknum, + "MultiXactMembersCreate record for offset {} with unexpected key {}", + moff, + key + ); + + let mut flagsval = LittleEndian::read_u32(&page[flagsoff..flagsoff + 4]); + flagsval &= !(((1 << pg_constants::MXACT_MEMBER_BITS_PER_XACT) - 1) << bshift); + flagsval |= member.status << bshift; + LittleEndian::write_u32(&mut page[flagsoff..flagsoff + 4], flagsval); + LittleEndian::write_u32(&mut page[memberoff..memberoff + 4], member.xid); + } + } + } + Ok(()) +} diff --git a/pageserver/src/walredo/process.rs b/pageserver/src/walredo/process.rs new file mode 100644 index 0000000000..85db3b4a4a --- /dev/null +++ b/pageserver/src/walredo/process.rs @@ -0,0 +1,406 @@ +use self::no_leak_child::NoLeakChild; +use crate::{ + config::PageServerConf, + metrics::{WalRedoKillCause, WAL_REDO_PROCESS_COUNTERS, WAL_REDO_RECORD_COUNTER}, + walrecord::NeonWalRecord, +}; +use anyhow::Context; +use bytes::Bytes; +use nix::poll::{PollFd, PollFlags}; +use pageserver_api::{reltag::RelTag, shard::TenantShardId}; +use postgres_ffi::BLCKSZ; +use std::os::fd::AsRawFd; +#[cfg(feature = "testing")] +use std::sync::atomic::AtomicUsize; +use std::{ + collections::VecDeque, + io::{Read, Write}, + process::{ChildStdin, ChildStdout, Command, Stdio}, + sync::{Mutex, MutexGuard}, + time::Duration, +}; +use tracing::{debug, error, instrument, Instrument}; +use utils::{lsn::Lsn, nonblock::set_nonblock}; + +mod no_leak_child; +/// The IPC protocol that pageserver and walredo process speak over their shared pipe. +mod protocol; + +pub struct WalRedoProcess { + #[allow(dead_code)] + conf: &'static PageServerConf, + tenant_shard_id: TenantShardId, + // Some() on construction, only becomes None on Drop. + child: Option, + stdout: Mutex, + stdin: Mutex, + /// Counter to separate same sized walredo inputs failing at the same millisecond. + #[cfg(feature = "testing")] + dump_sequence: AtomicUsize, +} + +struct ProcessInput { + stdin: ChildStdin, + n_requests: usize, +} + +struct ProcessOutput { + stdout: ChildStdout, + pending_responses: VecDeque>, + n_processed_responses: usize, +} + +impl WalRedoProcess { + // + // Start postgres binary in special WAL redo mode. + // + #[instrument(skip_all,fields(tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), pg_version=pg_version))] + pub(crate) fn launch( + conf: &'static PageServerConf, + tenant_shard_id: TenantShardId, + pg_version: u32, + ) -> anyhow::Result { + let pg_bin_dir_path = conf.pg_bin_dir(pg_version).context("pg_bin_dir")?; // TODO these should be infallible. + let pg_lib_dir_path = conf.pg_lib_dir(pg_version).context("pg_lib_dir")?; + + use no_leak_child::NoLeakChildCommandExt; + // Start postgres itself + let child = Command::new(pg_bin_dir_path.join("postgres")) + // the first arg must be --wal-redo so the child process enters into walredo mode + .arg("--wal-redo") + // the child doesn't process this arg, but, having it in the argv helps indentify the + // walredo process for a particular tenant when debugging a pagserver + .args(["--tenant-shard-id", &format!("{tenant_shard_id}")]) + .stdin(Stdio::piped()) + .stderr(Stdio::piped()) + .stdout(Stdio::piped()) + .env_clear() + .env("LD_LIBRARY_PATH", &pg_lib_dir_path) + .env("DYLD_LIBRARY_PATH", &pg_lib_dir_path) + // NB: The redo process is not trusted after we sent it the first + // walredo work. Before that, it is trusted. Specifically, we trust + // it to + // 1. close all file descriptors except stdin, stdout, stderr because + // pageserver might not be 100% diligent in setting FD_CLOEXEC on all + // the files it opens, and + // 2. to use seccomp to sandbox itself before processing the first + // walredo request. + .spawn_no_leak_child(tenant_shard_id) + .context("spawn process")?; + WAL_REDO_PROCESS_COUNTERS.started.inc(); + let mut child = scopeguard::guard(child, |child| { + error!("killing wal-redo-postgres process due to a problem during launch"); + child.kill_and_wait(WalRedoKillCause::Startup); + }); + + let stdin = child.stdin.take().unwrap(); + let stdout = child.stdout.take().unwrap(); + let stderr = child.stderr.take().unwrap(); + let stderr = tokio::process::ChildStderr::from_std(stderr) + .context("convert to tokio::ChildStderr")?; + macro_rules! set_nonblock_or_log_err { + ($file:ident) => {{ + let res = set_nonblock($file.as_raw_fd()); + if let Err(e) = &res { + error!(error = %e, file = stringify!($file), pid = child.id(), "set_nonblock failed"); + } + res + }}; + } + set_nonblock_or_log_err!(stdin)?; + set_nonblock_or_log_err!(stdout)?; + + // all fallible operations post-spawn are complete, so get rid of the guard + let child = scopeguard::ScopeGuard::into_inner(child); + + tokio::spawn( + async move { + scopeguard::defer! { + debug!("wal-redo-postgres stderr_logger_task finished"); + crate::metrics::WAL_REDO_PROCESS_COUNTERS.active_stderr_logger_tasks_finished.inc(); + } + debug!("wal-redo-postgres stderr_logger_task started"); + crate::metrics::WAL_REDO_PROCESS_COUNTERS.active_stderr_logger_tasks_started.inc(); + + use tokio::io::AsyncBufReadExt; + let mut stderr_lines = tokio::io::BufReader::new(stderr); + let mut buf = Vec::new(); + let res = loop { + buf.clear(); + // TODO we don't trust the process to cap its stderr length. + // Currently it can do unbounded Vec allocation. + match stderr_lines.read_until(b'\n', &mut buf).await { + Ok(0) => break Ok(()), // eof + Ok(num_bytes) => { + let output = String::from_utf8_lossy(&buf[..num_bytes]); + error!(%output, "received output"); + } + Err(e) => { + break Err(e); + } + } + }; + match res { + Ok(()) => (), + Err(e) => { + error!(error=?e, "failed to read from walredo stderr"); + } + } + }.instrument(tracing::info_span!(parent: None, "wal-redo-postgres-stderr", pid = child.id(), tenant_id = %tenant_shard_id.tenant_id, shard_id = %tenant_shard_id.shard_slug(), %pg_version)) + ); + + Ok(Self { + conf, + tenant_shard_id, + child: Some(child), + stdin: Mutex::new(ProcessInput { + stdin, + n_requests: 0, + }), + stdout: Mutex::new(ProcessOutput { + stdout, + pending_responses: VecDeque::new(), + n_processed_responses: 0, + }), + #[cfg(feature = "testing")] + dump_sequence: AtomicUsize::default(), + }) + } + + pub(crate) fn id(&self) -> u32 { + self.child + .as_ref() + .expect("must not call this during Drop") + .id() + } + + // Apply given WAL records ('records') over an old page image. Returns + // new page image. + // + #[instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), pid=%self.id()))] + pub(crate) fn apply_wal_records( + &self, + rel: RelTag, + blknum: u32, + base_img: &Option, + records: &[(Lsn, NeonWalRecord)], + wal_redo_timeout: Duration, + ) -> anyhow::Result { + let tag = protocol::BufferTag { rel, blknum }; + let input = self.stdin.lock().unwrap(); + + // Serialize all the messages to send the WAL redo process first. + // + // This could be problematic if there are millions of records to replay, + // but in practice the number of records is usually so small that it doesn't + // matter, and it's better to keep this code simple. + // + // Most requests start with a before-image with BLCKSZ bytes, followed by + // by some other WAL records. Start with a buffer that can hold that + // comfortably. + let mut writebuf: Vec = Vec::with_capacity((BLCKSZ as usize) * 3); + protocol::build_begin_redo_for_block_msg(tag, &mut writebuf); + if let Some(img) = base_img { + protocol::build_push_page_msg(tag, img, &mut writebuf); + } + for (lsn, rec) in records.iter() { + if let NeonWalRecord::Postgres { + will_init: _, + rec: postgres_rec, + } = rec + { + protocol::build_apply_record_msg(*lsn, postgres_rec, &mut writebuf); + } else { + anyhow::bail!("tried to pass neon wal record to postgres WAL redo"); + } + } + protocol::build_get_page_msg(tag, &mut writebuf); + WAL_REDO_RECORD_COUNTER.inc_by(records.len() as u64); + + let res = self.apply_wal_records0(&writebuf, input, wal_redo_timeout); + + if res.is_err() { + // not all of these can be caused by this particular input, however these are so rare + // in tests so capture all. + self.record_and_log(&writebuf); + } + + res + } + + fn apply_wal_records0( + &self, + writebuf: &[u8], + input: MutexGuard, + wal_redo_timeout: Duration, + ) -> anyhow::Result { + let mut proc = { input }; // TODO: remove this legacy rename, but this keep the patch small. + let mut nwrite = 0usize; + + while nwrite < writebuf.len() { + let mut stdin_pollfds = [PollFd::new(&proc.stdin, PollFlags::POLLOUT)]; + let n = loop { + match nix::poll::poll(&mut stdin_pollfds[..], wal_redo_timeout.as_millis() as i32) { + Err(nix::errno::Errno::EINTR) => continue, + res => break res, + } + }?; + + if n == 0 { + anyhow::bail!("WAL redo timed out"); + } + + // If 'stdin' is writeable, do write. + let in_revents = stdin_pollfds[0].revents().unwrap(); + if in_revents & (PollFlags::POLLERR | PollFlags::POLLOUT) != PollFlags::empty() { + nwrite += proc.stdin.write(&writebuf[nwrite..])?; + } + if in_revents.contains(PollFlags::POLLHUP) { + // We still have more data to write, but the process closed the pipe. + anyhow::bail!("WAL redo process closed its stdin unexpectedly"); + } + } + let request_no = proc.n_requests; + proc.n_requests += 1; + drop(proc); + + // To improve walredo performance we separate sending requests and receiving + // responses. Them are protected by different mutexes (output and input). + // If thread T1, T2, T3 send requests D1, D2, D3 to walredo process + // then there is not warranty that T1 will first granted output mutex lock. + // To address this issue we maintain number of sent requests, number of processed + // responses and ring buffer with pending responses. After sending response + // (under input mutex), threads remembers request number. Then it releases + // input mutex, locks output mutex and fetch in ring buffer all responses until + // its stored request number. The it takes correspondent element from + // pending responses ring buffer and truncate all empty elements from the front, + // advancing processed responses number. + + let mut output = self.stdout.lock().unwrap(); + let n_processed_responses = output.n_processed_responses; + while n_processed_responses + output.pending_responses.len() <= request_no { + // We expect the WAL redo process to respond with an 8k page image. We read it + // into this buffer. + let mut resultbuf = vec![0; BLCKSZ.into()]; + let mut nresult: usize = 0; // # of bytes read into 'resultbuf' so far + while nresult < BLCKSZ.into() { + let mut stdout_pollfds = [PollFd::new(&output.stdout, PollFlags::POLLIN)]; + // We do two things simultaneously: reading response from stdout + // and forward any logging information that the child writes to its stderr to the page server's log. + let n = loop { + match nix::poll::poll( + &mut stdout_pollfds[..], + wal_redo_timeout.as_millis() as i32, + ) { + Err(nix::errno::Errno::EINTR) => continue, + res => break res, + } + }?; + + if n == 0 { + anyhow::bail!("WAL redo timed out"); + } + + // If we have some data in stdout, read it to the result buffer. + let out_revents = stdout_pollfds[0].revents().unwrap(); + if out_revents & (PollFlags::POLLERR | PollFlags::POLLIN) != PollFlags::empty() { + nresult += output.stdout.read(&mut resultbuf[nresult..])?; + } + if out_revents.contains(PollFlags::POLLHUP) { + anyhow::bail!("WAL redo process closed its stdout unexpectedly"); + } + } + output + .pending_responses + .push_back(Some(Bytes::from(resultbuf))); + } + // Replace our request's response with None in `pending_responses`. + // Then make space in the ring buffer by clearing out any seqence of contiguous + // `None`'s from the front of `pending_responses`. + // NB: We can't pop_front() because other requests' responses because another + // requester might have grabbed the output mutex before us: + // T1: grab input mutex + // T1: send request_no 23 + // T1: release input mutex + // T2: grab input mutex + // T2: send request_no 24 + // T2: release input mutex + // T2: grab output mutex + // T2: n_processed_responses + output.pending_responses.len() <= request_no + // 23 0 24 + // T2: enters poll loop that reads stdout + // T2: put response for 23 into pending_responses + // T2: put response for 24 into pending_resposnes + // pending_responses now looks like this: Front Some(response_23) Some(response_24) Back + // T2: takes its response_24 + // pending_responses now looks like this: Front Some(response_23) None Back + // T2: does the while loop below + // pending_responses now looks like this: Front Some(response_23) None Back + // T2: releases output mutex + // T1: grabs output mutex + // T1: n_processed_responses + output.pending_responses.len() > request_no + // 23 2 23 + // T1: skips poll loop that reads stdout + // T1: takes its response_23 + // pending_responses now looks like this: Front None None Back + // T2: does the while loop below + // pending_responses now looks like this: Front Back + // n_processed_responses now has value 25 + let res = output.pending_responses[request_no - n_processed_responses] + .take() + .expect("we own this request_no, nobody else is supposed to take it"); + while let Some(front) = output.pending_responses.front() { + if front.is_none() { + output.pending_responses.pop_front(); + output.n_processed_responses += 1; + } else { + break; + } + } + Ok(res) + } + + #[cfg(feature = "testing")] + fn record_and_log(&self, writebuf: &[u8]) { + use std::sync::atomic::Ordering; + + let millis = std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH) + .unwrap() + .as_millis(); + + let seq = self.dump_sequence.fetch_add(1, Ordering::Relaxed); + + // these files will be collected to an allure report + let filename = format!("walredo-{millis}-{}-{seq}.walredo", writebuf.len()); + + let path = self.conf.tenant_path(&self.tenant_shard_id).join(&filename); + + let res = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .read(true) + .open(path) + .and_then(|mut f| f.write_all(writebuf)); + + // trip up allowed_errors + if let Err(e) = res { + tracing::error!(target=%filename, length=writebuf.len(), "failed to write out the walredo errored input: {e}"); + } else { + tracing::error!(filename, "erroring walredo input saved"); + } + } + + #[cfg(not(feature = "testing"))] + fn record_and_log(&self, _: &[u8]) {} +} + +impl Drop for WalRedoProcess { + fn drop(&mut self) { + self.child + .take() + .expect("we only do this once") + .kill_and_wait(WalRedoKillCause::WalRedoProcessDrop); + // no way to wait for stderr_logger_task from Drop because that is async only + } +} diff --git a/pageserver/src/walredo/process/no_leak_child.rs b/pageserver/src/walredo/process/no_leak_child.rs new file mode 100644 index 0000000000..ca016408e6 --- /dev/null +++ b/pageserver/src/walredo/process/no_leak_child.rs @@ -0,0 +1,126 @@ +use tracing; +use tracing::error; +use tracing::info; +use tracing::instrument; + +use crate::metrics::WalRedoKillCause; +use crate::metrics::WAL_REDO_PROCESS_COUNTERS; + +use std::io; +use std::process::Command; + +use std::ops::DerefMut; + +use std::ops::Deref; + +use std::process::Child; + +use pageserver_api::shard::TenantShardId; + +/// Wrapper type around `std::process::Child` which guarantees that the child +/// will be killed and waited-for by this process before being dropped. +pub(crate) struct NoLeakChild { + pub(crate) tenant_id: TenantShardId, + pub(crate) child: Option, +} + +impl Deref for NoLeakChild { + type Target = Child; + + fn deref(&self) -> &Self::Target { + self.child.as_ref().expect("must not use from drop") + } +} + +impl DerefMut for NoLeakChild { + fn deref_mut(&mut self) -> &mut Self::Target { + self.child.as_mut().expect("must not use from drop") + } +} + +impl NoLeakChild { + pub(crate) fn spawn(tenant_id: TenantShardId, command: &mut Command) -> io::Result { + let child = command.spawn()?; + Ok(NoLeakChild { + tenant_id, + child: Some(child), + }) + } + + pub(crate) fn kill_and_wait(mut self, cause: WalRedoKillCause) { + let child = match self.child.take() { + Some(child) => child, + None => return, + }; + Self::kill_and_wait_impl(child, cause); + } + + #[instrument(skip_all, fields(pid=child.id(), ?cause))] + pub(crate) fn kill_and_wait_impl(mut child: Child, cause: WalRedoKillCause) { + scopeguard::defer! { + WAL_REDO_PROCESS_COUNTERS.killed_by_cause[cause].inc(); + } + let res = child.kill(); + if let Err(e) = res { + // This branch is very unlikely because: + // - We (= pageserver) spawned this process successfully, so, we're allowed to kill it. + // - This is the only place that calls .kill() + // - We consume `self`, so, .kill() can't be called twice. + // - If the process exited by itself or was killed by someone else, + // .kill() will still succeed because we haven't wait()'ed yet. + // + // So, if we arrive here, we have really no idea what happened, + // whether the PID stored in self.child is still valid, etc. + // If this function were fallible, we'd return an error, but + // since it isn't, all we can do is log an error and proceed + // with the wait(). + error!(error = %e, "failed to SIGKILL; subsequent wait() might fail or wait for wrong process"); + } + + match child.wait() { + Ok(exit_status) => { + info!(exit_status = %exit_status, "wait successful"); + } + Err(e) => { + error!(error = %e, "wait error; might leak the child process; it will show as zombie (defunct)"); + } + } + } +} + +impl Drop for NoLeakChild { + fn drop(&mut self) { + let child = match self.child.take() { + Some(child) => child, + None => return, + }; + let tenant_shard_id = self.tenant_id; + // Offload the kill+wait of the child process into the background. + // If someone stops the runtime, we'll leak the child process. + // We can ignore that case because we only stop the runtime on pageserver exit. + tokio::runtime::Handle::current().spawn(async move { + tokio::task::spawn_blocking(move || { + // Intentionally don't inherit the tracing context from whoever is dropping us. + // This thread here is going to outlive of our dropper. + let span = tracing::info_span!( + "walredo", + tenant_id = %tenant_shard_id.tenant_id, + shard_id = %tenant_shard_id.shard_slug() + ); + let _entered = span.enter(); + Self::kill_and_wait_impl(child, WalRedoKillCause::NoLeakChildDrop); + }) + .await + }); + } +} + +pub(crate) trait NoLeakChildCommandExt { + fn spawn_no_leak_child(&mut self, tenant_id: TenantShardId) -> io::Result; +} + +impl NoLeakChildCommandExt for Command { + fn spawn_no_leak_child(&mut self, tenant_id: TenantShardId) -> io::Result { + NoLeakChild::spawn(tenant_id, self) + } +} diff --git a/pageserver/src/walredo/process/protocol.rs b/pageserver/src/walredo/process/protocol.rs new file mode 100644 index 0000000000..b703344cc8 --- /dev/null +++ b/pageserver/src/walredo/process/protocol.rs @@ -0,0 +1,57 @@ +use bytes::BufMut; +use pageserver_api::reltag::RelTag; +use serde::Serialize; +use utils::bin_ser::BeSer; +use utils::lsn::Lsn; + +/// +/// `RelTag` + block number (`blknum`) gives us a unique id of the page in the cluster. +/// +/// In Postgres `BufferTag` structure is used for exactly the same purpose. +/// [See more related comments here](https://github.com/postgres/postgres/blob/99c5852e20a0987eca1c38ba0c09329d4076b6a0/src/include/storage/buf_internals.h#L91). +/// +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Serialize)] +pub(crate) struct BufferTag { + pub rel: RelTag, + pub blknum: u32, +} + +pub(crate) fn build_begin_redo_for_block_msg(tag: BufferTag, buf: &mut Vec) { + let len = 4 + 1 + 4 * 4; + + buf.put_u8(b'B'); + buf.put_u32(len as u32); + + tag.ser_into(buf) + .expect("serialize BufferTag should always succeed"); +} + +pub(crate) fn build_push_page_msg(tag: BufferTag, base_img: &[u8], buf: &mut Vec) { + assert!(base_img.len() == 8192); + + let len = 4 + 1 + 4 * 4 + base_img.len(); + + buf.put_u8(b'P'); + buf.put_u32(len as u32); + tag.ser_into(buf) + .expect("serialize BufferTag should always succeed"); + buf.put(base_img); +} + +pub(crate) fn build_apply_record_msg(endlsn: Lsn, rec: &[u8], buf: &mut Vec) { + let len = 4 + 8 + rec.len(); + + buf.put_u8(b'A'); + buf.put_u32(len as u32); + buf.put_u64(endlsn.0); + buf.put(rec); +} + +pub(crate) fn build_get_page_msg(tag: BufferTag, buf: &mut Vec) { + let len = 4 + 1 + 4 * 4; + + buf.put_u8(b'G'); + buf.put_u32(len as u32); + tag.ser_into(buf) + .expect("serialize BufferTag should always succeed"); +} From 431f4234d43f3fe42fbda441e601a89d2421b52e Mon Sep 17 00:00:00 2001 From: John Spray Date: Tue, 6 Feb 2024 10:07:10 +0000 Subject: [PATCH 13/72] storage controller: embed database migrations in binary (#6637) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem We don't have a neat way to carry around migration .sql files during deploy, and in any case would prefer to avoid depending on diesel CLI to deploy. ## Summary of changes - Use `diesel_migrations` crate to embed migrations in our binary - Run migrations on startup - Drop the diesel dependency in the `neon_local` binary, as the attachment_service binary just needs the database to exist. Do database creation with a simple `createdb`. Co-authored-by: Arpad Müller --- Cargo.lock | 1 + control_plane/attachment_service/Cargo.toml | 1 + control_plane/attachment_service/src/main.rs | 24 +++++++++++++++++++- 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index c16331636a..b2b2777408 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -281,6 +281,7 @@ dependencies = [ "clap", "control_plane", "diesel", + "diesel_migrations", "futures", "git-version", "hyper", diff --git a/control_plane/attachment_service/Cargo.toml b/control_plane/attachment_service/Cargo.toml index d3c62d74d2..3a65153c41 100644 --- a/control_plane/attachment_service/Cargo.toml +++ b/control_plane/attachment_service/Cargo.toml @@ -25,6 +25,7 @@ tokio-util.workspace = true tracing.workspace = true diesel = { version = "2.1.4", features = ["serde_json", "postgres"] } +diesel_migrations = { version = "2.1.0" } utils = { path = "../../libs/utils/" } metrics = { path = "../../libs/metrics/" } diff --git a/control_plane/attachment_service/src/main.rs b/control_plane/attachment_service/src/main.rs index 37b06c4090..7ac5918244 100644 --- a/control_plane/attachment_service/src/main.rs +++ b/control_plane/attachment_service/src/main.rs @@ -4,13 +4,14 @@ /// This enables running & testing pageservers without a full-blown /// deployment of the Neon cloud platform. /// -use anyhow::anyhow; +use anyhow::{anyhow, Context}; use attachment_service::http::make_router; use attachment_service::persistence::Persistence; use attachment_service::service::{Config, Service}; use aws_config::{self, BehaviorVersion, Region}; use camino::Utf8PathBuf; use clap::Parser; +use diesel::Connection; use metrics::launch_timestamp::LaunchTimestamp; use std::sync::Arc; use tokio::signal::unix::SignalKind; @@ -22,6 +23,9 @@ use utils::{project_build_tag, project_git_version, tcp_listener}; project_git_version!(GIT_VERSION); project_build_tag!(BUILD_TAG); +use diesel_migrations::{embed_migrations, EmbeddedMigrations}; +pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); + #[derive(Parser)] #[command(author, version, about, long_about = None)] #[command(arg_required_else_help(true))] @@ -166,6 +170,19 @@ impl Secrets { } } +async fn migration_run(database_url: &str) -> anyhow::Result<()> { + use diesel::PgConnection; + use diesel_migrations::{HarnessWithOutput, MigrationHarness}; + let mut conn = PgConnection::establish(database_url)?; + + HarnessWithOutput::write_to_stdout(&mut conn) + .run_pending_migrations(MIGRATIONS) + .map(|_| ()) + .map_err(|e| anyhow::anyhow!(e))?; + + Ok(()) +} + #[tokio::main] async fn main() -> anyhow::Result<()> { let launch_ts = Box::leak(Box::new(LaunchTimestamp::generate())); @@ -194,6 +211,11 @@ async fn main() -> anyhow::Result<()> { compute_hook_url: args.compute_hook_url, }; + // After loading secrets & config, but before starting anything else, apply database migrations + migration_run(&secrets.database_url) + .await + .context("Running database migrations")?; + let json_path = args.path; let persistence = Arc::new(Persistence::new(secrets.database_url, json_path.clone())); From 53743991decd9f1d13fd5063a8e840a38cbda383 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Tue, 6 Feb 2024 13:34:13 +0200 Subject: [PATCH 14/72] uploader: avoid cloning vecs just to get Bytes (#6645) Fix cloning the serialized heatmap on every attempt by just turning it into `bytes::Bytes` before clone so it will be a refcounted instead of refcounting a vec clone later on. Also fixes one cancellation token cloning I had missed in #6618. Cc: #6096 --- .../src/tenant/secondary/heatmap_uploader.rs | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/pageserver/src/tenant/secondary/heatmap_uploader.rs b/pageserver/src/tenant/secondary/heatmap_uploader.rs index fff29b2487..806e3fb0e8 100644 --- a/pageserver/src/tenant/secondary/heatmap_uploader.rs +++ b/pageserver/src/tenant/secondary/heatmap_uploader.rs @@ -371,8 +371,6 @@ async fn upload_tenant_heatmap( }; let timelines = tenant.timelines.lock().unwrap().clone(); - let tenant_cancel = tenant.cancel.clone(); - // Ensure that Tenant::shutdown waits for any upload in flight: this is needed because otherwise // when we delete a tenant, we might race with an upload in flight and end up leaving a heatmap behind // in remote storage. @@ -401,6 +399,7 @@ async fn upload_tenant_heatmap( // Serialize the heatmap let bytes = serde_json::to_vec(&heatmap).map_err(|e| anyhow::anyhow!(e))?; + let bytes = bytes::Bytes::from(bytes); let size = bytes.len(); // Drop out early if nothing changed since our last upload @@ -411,13 +410,12 @@ async fn upload_tenant_heatmap( let path = remote_heatmap_path(tenant.get_tenant_shard_id()); - // Write the heatmap. + let cancel = &tenant.cancel; + tracing::debug!("Uploading {size} byte heatmap to {path}"); if let Err(e) = backoff::retry( || async { - let bytes = futures::stream::once(futures::future::ready(Ok(bytes::Bytes::from( - bytes.clone(), - )))); + let bytes = futures::stream::once(futures::future::ready(Ok(bytes.clone()))); remote_storage .upload_storage_object(bytes, size, &path) .await @@ -426,13 +424,13 @@ async fn upload_tenant_heatmap( 3, u32::MAX, "Uploading heatmap", - &tenant_cancel, + cancel, ) .await .ok_or_else(|| anyhow::anyhow!("Shutting down")) .and_then(|x| x) { - if tenant_cancel.is_cancelled() { + if cancel.is_cancelled() { return Err(UploadHeatmapError::Cancelled); } else { return Err(e.into()); From 0de46fd6f265e1ef0d27b0ab0f51fb7da2e52705 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Tue, 6 Feb 2024 13:04:15 +0100 Subject: [PATCH 15/72] heavier_once_cell: switch to tokio::sync::RwLock (#6589) Using the RwLock reduces contention on the hot path. Co-authored-by: Joonas Koivunen --- libs/utils/src/sync/heavier_once_cell.rs | 153 ++++++++++++++----- pageserver/src/tenant/storage_layer/layer.rs | 24 +-- pageserver/src/tenant/timeline.rs | 2 +- 3 files changed, 127 insertions(+), 52 deletions(-) diff --git a/libs/utils/src/sync/heavier_once_cell.rs b/libs/utils/src/sync/heavier_once_cell.rs index 0ccaf4e716..f733d107f1 100644 --- a/libs/utils/src/sync/heavier_once_cell.rs +++ b/libs/utils/src/sync/heavier_once_cell.rs @@ -1,6 +1,6 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, - Arc, Mutex, MutexGuard, + Arc, }; use tokio::sync::Semaphore; @@ -12,7 +12,7 @@ use tokio::sync::Semaphore; /// /// [`OwnedSemaphorePermit`]: tokio::sync::OwnedSemaphorePermit pub struct OnceCell { - inner: Mutex>, + inner: tokio::sync::RwLock>, initializers: AtomicUsize, } @@ -50,7 +50,7 @@ impl OnceCell { let sem = Semaphore::new(1); sem.close(); Self { - inner: Mutex::new(Inner { + inner: tokio::sync::RwLock::new(Inner { init_semaphore: Arc::new(sem), value: Some(value), }), @@ -61,18 +61,18 @@ impl OnceCell { /// Returns a guard to an existing initialized value, or uniquely initializes the value before /// returning the guard. /// - /// Initializing might wait on any existing [`Guard::take_and_deinit`] deinitialization. + /// Initializing might wait on any existing [`GuardMut::take_and_deinit`] deinitialization. /// /// Initialization is panic-safe and cancellation-safe. - pub async fn get_or_init(&self, factory: F) -> Result, E> + pub async fn get_mut_or_init(&self, factory: F) -> Result, E> where F: FnOnce(InitPermit) -> Fut, Fut: std::future::Future>, { let sem = { - let guard = self.inner.lock().unwrap(); + let guard = self.inner.write().await; if guard.value.is_some() { - return Ok(Guard(guard)); + return Ok(GuardMut(guard)); } guard.init_semaphore.clone() }; @@ -88,29 +88,72 @@ impl OnceCell { let permit = InitPermit(permit); let (value, _permit) = factory(permit).await?; - let guard = self.inner.lock().unwrap(); + let guard = self.inner.write().await; Ok(Self::set0(value, guard)) } Err(_closed) => { - let guard = self.inner.lock().unwrap(); + let guard = self.inner.write().await; assert!( guard.value.is_some(), "semaphore got closed, must be initialized" ); - return Ok(Guard(guard)); + return Ok(GuardMut(guard)); } } } - /// Assuming a permit is held after previous call to [`Guard::take_and_deinit`], it can be used + /// Returns a guard to an existing initialized value, or uniquely initializes the value before + /// returning the guard. + /// + /// Initialization is panic-safe and cancellation-safe. + pub async fn get_or_init(&self, factory: F) -> Result, E> + where + F: FnOnce(InitPermit) -> Fut, + Fut: std::future::Future>, + { + let sem = { + let guard = self.inner.read().await; + if guard.value.is_some() { + return Ok(GuardRef(guard)); + } + guard.init_semaphore.clone() + }; + + let permit = { + // increment the count for the duration of queued + let _guard = CountWaitingInitializers::start(self); + sem.acquire_owned().await + }; + + match permit { + Ok(permit) => { + let permit = InitPermit(permit); + let (value, _permit) = factory(permit).await?; + + let guard = self.inner.write().await; + + Ok(Self::set0(value, guard).downgrade()) + } + Err(_closed) => { + let guard = self.inner.read().await; + assert!( + guard.value.is_some(), + "semaphore got closed, must be initialized" + ); + return Ok(GuardRef(guard)); + } + } + } + + /// Assuming a permit is held after previous call to [`GuardMut::take_and_deinit`], it can be used /// to complete initializing the inner value. /// /// # Panics /// /// If the inner has already been initialized. - pub fn set(&self, value: T, _permit: InitPermit) -> Guard<'_, T> { - let guard = self.inner.lock().unwrap(); + pub async fn set(&self, value: T, _permit: InitPermit) -> GuardMut<'_, T> { + let guard = self.inner.write().await; // cannot assert that this permit is for self.inner.semaphore, but we can assert it cannot // give more permits right now. @@ -122,21 +165,31 @@ impl OnceCell { Self::set0(value, guard) } - fn set0(value: T, mut guard: std::sync::MutexGuard<'_, Inner>) -> Guard<'_, T> { + fn set0(value: T, mut guard: tokio::sync::RwLockWriteGuard<'_, Inner>) -> GuardMut<'_, T> { if guard.value.is_some() { drop(guard); unreachable!("we won permit, must not be initialized"); } guard.value = Some(value); guard.init_semaphore.close(); - Guard(guard) + GuardMut(guard) } /// Returns a guard to an existing initialized value, if any. - pub fn get(&self) -> Option> { - let guard = self.inner.lock().unwrap(); + pub async fn get_mut(&self) -> Option> { + let guard = self.inner.write().await; if guard.value.is_some() { - Some(Guard(guard)) + Some(GuardMut(guard)) + } else { + None + } + } + + /// Returns a guard to an existing initialized value, if any. + pub async fn get(&self) -> Option> { + let guard = self.inner.read().await; + if guard.value.is_some() { + Some(GuardRef(guard)) } else { None } @@ -168,9 +221,9 @@ impl<'a, T> Drop for CountWaitingInitializers<'a, T> { /// Uninteresting guard object to allow short-lived access to inspect or clone the held, /// initialized value. #[derive(Debug)] -pub struct Guard<'a, T>(MutexGuard<'a, Inner>); +pub struct GuardMut<'a, T>(tokio::sync::RwLockWriteGuard<'a, Inner>); -impl std::ops::Deref for Guard<'_, T> { +impl std::ops::Deref for GuardMut<'_, T> { type Target = T; fn deref(&self) -> &Self::Target { @@ -181,7 +234,7 @@ impl std::ops::Deref for Guard<'_, T> { } } -impl std::ops::DerefMut for Guard<'_, T> { +impl std::ops::DerefMut for GuardMut<'_, T> { fn deref_mut(&mut self) -> &mut Self::Target { self.0 .value @@ -190,7 +243,7 @@ impl std::ops::DerefMut for Guard<'_, T> { } } -impl<'a, T> Guard<'a, T> { +impl<'a, T> GuardMut<'a, T> { /// Take the current value, and a new permit for it's deinitialization. /// /// The permit will be on a semaphore part of the new internal value, and any following @@ -208,6 +261,24 @@ impl<'a, T> Guard<'a, T> { .map(|v| (v, InitPermit(permit))) .expect("guard is not created unless value has been initialized") } + + pub fn downgrade(self) -> GuardRef<'a, T> { + GuardRef(self.0.downgrade()) + } +} + +#[derive(Debug)] +pub struct GuardRef<'a, T>(tokio::sync::RwLockReadGuard<'a, Inner>); + +impl std::ops::Deref for GuardRef<'_, T> { + type Target = T; + + fn deref(&self) -> &Self::Target { + self.0 + .value + .as_ref() + .expect("guard is not created unless value has been initialized") + } } /// Type held by OnceCell (de)initializing task. @@ -248,7 +319,7 @@ mod tests { barrier.wait().await; let won = { let g = cell - .get_or_init(|permit| { + .get_mut_or_init(|permit| { counters.factory_got_to_run.fetch_add(1, Ordering::Relaxed); async { counters.future_polled.fetch_add(1, Ordering::Relaxed); @@ -295,7 +366,11 @@ mod tests { let cell = cell.clone(); let deinitialization_started = deinitialization_started.clone(); async move { - let (answer, _permit) = cell.get().expect("initialized to value").take_and_deinit(); + let (answer, _permit) = cell + .get_mut() + .await + .expect("initialized to value") + .take_and_deinit(); assert_eq!(answer, initial); deinitialization_started.wait().await; @@ -306,7 +381,7 @@ mod tests { deinitialization_started.wait().await; let started_at = tokio::time::Instant::now(); - cell.get_or_init(|permit| async { Ok::<_, Infallible>((reinit, permit)) }) + cell.get_mut_or_init(|permit| async { Ok::<_, Infallible>((reinit, permit)) }) .await .unwrap(); @@ -318,21 +393,21 @@ mod tests { jh.await.unwrap(); - assert_eq!(*cell.get().unwrap(), reinit); + assert_eq!(*cell.get_mut().await.unwrap(), reinit); } - #[test] - fn reinit_with_deinit_permit() { + #[tokio::test] + async fn reinit_with_deinit_permit() { let cell = Arc::new(OnceCell::new(42)); - let (mol, permit) = cell.get().unwrap().take_and_deinit(); - cell.set(5, permit); - assert_eq!(*cell.get().unwrap(), 5); + let (mol, permit) = cell.get_mut().await.unwrap().take_and_deinit(); + cell.set(5, permit).await; + assert_eq!(*cell.get_mut().await.unwrap(), 5); - let (five, permit) = cell.get().unwrap().take_and_deinit(); + let (five, permit) = cell.get_mut().await.unwrap().take_and_deinit(); assert_eq!(5, five); - cell.set(mol, permit); - assert_eq!(*cell.get().unwrap(), 42); + cell.set(mol, permit).await; + assert_eq!(*cell.get_mut().await.unwrap(), 42); } #[tokio::test] @@ -340,13 +415,13 @@ mod tests { let cell = OnceCell::default(); for _ in 0..10 { - cell.get_or_init(|_permit| async { Err("whatever error") }) + cell.get_mut_or_init(|_permit| async { Err("whatever error") }) .await .unwrap_err(); } let g = cell - .get_or_init(|permit| async { Ok::<_, Infallible>(("finally success", permit)) }) + .get_mut_or_init(|permit| async { Ok::<_, Infallible>(("finally success", permit)) }) .await .unwrap(); assert_eq!(*g, "finally success"); @@ -358,7 +433,7 @@ mod tests { let barrier = tokio::sync::Barrier::new(2); - let initializer = cell.get_or_init(|permit| async { + let initializer = cell.get_mut_or_init(|permit| async { barrier.wait().await; futures::future::pending::<()>().await; @@ -372,10 +447,10 @@ mod tests { // now initializer is dropped - assert!(cell.get().is_none()); + assert!(cell.get_mut().await.is_none()); let g = cell - .get_or_init(|permit| async { Ok::<_, Infallible>(("now initialized", permit)) }) + .get_mut_or_init(|permit| async { Ok::<_, Infallible>(("now initialized", permit)) }) .await .unwrap(); assert_eq!(*g, "now initialized"); diff --git a/pageserver/src/tenant/storage_layer/layer.rs b/pageserver/src/tenant/storage_layer/layer.rs index 12af866810..1f337adf53 100644 --- a/pageserver/src/tenant/storage_layer/layer.rs +++ b/pageserver/src/tenant/storage_layer/layer.rs @@ -299,8 +299,8 @@ impl Layer { }) } - pub(crate) fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { - self.0.info(reset) + pub(crate) async fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { + self.0.info(reset).await } pub(crate) fn access_stats(&self) -> &LayerAccessStats { @@ -611,10 +611,10 @@ impl LayerInner { let mut rx = self.status.subscribe(); let strong = { - match self.inner.get() { + match self.inner.get_mut().await { Some(mut either) => { self.wanted_evicted.store(true, Ordering::Relaxed); - either.downgrade() + ResidentOrWantedEvicted::downgrade(&mut either) } None => return Err(EvictionError::NotFound), } @@ -640,7 +640,7 @@ impl LayerInner { // use however late (compared to the initial expressing of wanted) as the // "outcome" now LAYER_IMPL_METRICS.inc_broadcast_lagged(); - match self.inner.get() { + match self.inner.get_mut().await { Some(_) => Err(EvictionError::Downloaded), None => Ok(()), } @@ -758,7 +758,7 @@ impl LayerInner { // use the already held initialization permit because it is impossible to hit the // below paths anymore essentially limiting the max loop iterations to 2. let (value, init_permit) = download(init_permit).await?; - let mut guard = self.inner.set(value, init_permit); + let mut guard = self.inner.set(value, init_permit).await; let (strong, _upgraded) = guard .get_and_upgrade() .expect("init creates strong reference, we held the init permit"); @@ -766,7 +766,7 @@ impl LayerInner { } let (weak, permit) = { - let mut locked = self.inner.get_or_init(download).await?; + let mut locked = self.inner.get_mut_or_init(download).await?; if let Some((strong, upgraded)) = locked.get_and_upgrade() { if upgraded { @@ -986,12 +986,12 @@ impl LayerInner { } } - fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { + async fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { let layer_file_name = self.desc.filename().file_name(); // this is not accurate: we could have the file locally but there was a cancellation // and now we are not in sync, or we are currently downloading it. - let remote = self.inner.get().is_none(); + let remote = self.inner.get_mut().await.is_none(); let access_stats = self.access_stats.as_api_model(reset); @@ -1050,7 +1050,7 @@ impl LayerInner { LAYER_IMPL_METRICS.inc_eviction_cancelled(EvictionCancelled::LayerGone); return; }; - match this.evict_blocking(version) { + match tokio::runtime::Handle::current().block_on(this.evict_blocking(version)) { Ok(()) => LAYER_IMPL_METRICS.inc_completed_evictions(), Err(reason) => LAYER_IMPL_METRICS.inc_eviction_cancelled(reason), } @@ -1058,7 +1058,7 @@ impl LayerInner { } } - fn evict_blocking(&self, only_version: usize) -> Result<(), EvictionCancelled> { + async fn evict_blocking(&self, only_version: usize) -> Result<(), EvictionCancelled> { // deleted or detached timeline, don't do anything. let Some(timeline) = self.timeline.upgrade() else { return Err(EvictionCancelled::TimelineGone); @@ -1067,7 +1067,7 @@ impl LayerInner { // to avoid starting a new download while we evict, keep holding on to the // permit. let _permit = { - let maybe_downloaded = self.inner.get(); + let maybe_downloaded = self.inner.get_mut().await; let (_weak, permit) = match maybe_downloaded { Some(mut guard) => { diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 0ba3fe728a..50ffc4d265 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -1268,7 +1268,7 @@ impl Timeline { let mut historic_layers = Vec::new(); for historic_layer in layer_map.iter_historic_layers() { let historic_layer = guard.get_from_desc(&historic_layer); - historic_layers.push(historic_layer.info(reset)); + historic_layers.push(historic_layer.info(reset).await); } LayerMapInfo { From dae56ef60ca33643b3d80b4d2497fb6902620db0 Mon Sep 17 00:00:00 2001 From: Vadim Kharitonov Date: Tue, 6 Feb 2024 13:15:42 +0100 Subject: [PATCH 16/72] Do not suspend compute if there is an active logical replication subscription. (#6570) ## Problem the idea is to keep compute up and running if there are any active logical replication subscriptions. ### Rationale Rationale: - The Write-Ahead Logging (WAL) files, which contain the data changes, will need to be retained on the publisher side until the subscriber is able to connect again and apply these changes. This could potentially lead to increased disk usage on the publisher - and we do not want to disrupt the source - I think it is more pain for our customer to resolve storage issues on the source than to pay for the compute at the target. - Upon resuming the compute resources, the subscriber will start consuming and applying the changes from the retained WAL files. The time taken to catch up will depend on the volume of changes and the configured vCPUs. we can avoid explaining complex situations where we lag behind (in extreme cases we could lag behind hours, days or even months) - I think an important use case for logical replication from a source is a one-time migration or release upgrade. In this case the customer would not mind if we are not suspended for the duration of the migration. We need to document this in the release notes and the documentation in the context of logical replication where Neon is the target (subscriber) ### See internal discussion here https://neondb.slack.com/archives/C04DGM6SMTM/p1706793400746539?thread_ts=1706792628.701279&cid=C04DGM6SMTM --- compute_tools/src/monitor.rs | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/compute_tools/src/monitor.rs b/compute_tools/src/monitor.rs index f09bd02664..872a3f7750 100644 --- a/compute_tools/src/monitor.rs +++ b/compute_tools/src/monitor.rs @@ -138,6 +138,34 @@ fn watch_compute_activity(compute: &ComputeNode) { } } // + // Don't suspend compute if there is an active logical replication subscription + // + // `where pid is not null` – to filter out read only computes and subscription on branches + // + let logical_subscriptions_query = + "select count(*) from pg_stat_subscription where pid is not null;"; + match cli.query_one(logical_subscriptions_query, &[]) { + Ok(row) => match row.try_get::<&str, i64>("count") { + Ok(num_subscribers) => { + if num_subscribers > 0 { + compute.update_last_active(Some(Utc::now())); + continue; + } + } + Err(e) => { + warn!("failed to parse `pg_stat_subscription` count: {:?}", e); + continue; + } + }, + Err(e) => { + warn!( + "failed to get list of active logical replication subscriptions: {:?}", + e + ); + continue; + } + } + // // Do not suspend compute if autovacuum is running // let autovacuum_count_query = "select count(*) from pg_stat_activity where backend_type = 'autovacuum worker'"; From 62978433176ca6a9679baea769aa751c48fa037d Mon Sep 17 00:00:00 2001 From: John Spray Date: Tue, 6 Feb 2024 12:49:41 +0000 Subject: [PATCH 17/72] tests: flakiness fixes in pageserver tests (#6632) Fix several test flakes: - test_sharding_service_smoke had log failures on "Dropped LSN updates" - test_emergency_mode had log failures on a deletion queue shutdown check, where the check was incorrect because it was expecting channel receiver to stay alive after cancellation token was fired. - test_secondary_mode_eviction had racing heatmap uploads because the test was using a live migration hook to set up locations, where that migration was itself uploading heatmaps and generally making the situation more complex than it needed to be. These are the failure modes that I saw when spot checking the last few failures of each test. This will mostly/completely address #6511, but I'll leave that ticket open for a couple days and then check if either of the tests named in that ticket are flaky. Related #6511 --- pageserver/src/deletion_queue.rs | 6 ++-- test_runner/fixtures/neon_fixtures.py | 3 +- .../regress/test_disk_usage_eviction.py | 30 ++++++++++--------- test_runner/regress/test_sharding_service.py | 5 ++++ 4 files changed, 27 insertions(+), 17 deletions(-) diff --git a/pageserver/src/deletion_queue.rs b/pageserver/src/deletion_queue.rs index 6a820e1bdc..da1da9331a 100644 --- a/pageserver/src/deletion_queue.rs +++ b/pageserver/src/deletion_queue.rs @@ -700,8 +700,6 @@ impl DeletionQueue { } pub async fn shutdown(&mut self, timeout: Duration) { - self.cancel.cancel(); - match tokio::time::timeout(timeout, self.client.flush()).await { Ok(Ok(())) => { tracing::info!("Deletion queue flushed successfully on shutdown") @@ -715,6 +713,10 @@ impl DeletionQueue { tracing::warn!("Timed out flushing deletion queue on shutdown") } } + + // We only cancel _after_ flushing: otherwise we would be shutting down the + // components that do the flush. + self.cancel.cancel(); } } diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 5ce2fca820..bf7c6ccc14 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1162,7 +1162,8 @@ class NeonEnv: to the attachment service. """ meta = self.attachment_service.inspect(tenant_id) - assert meta is not None, f"{tenant_id} attachment location not found" + if meta is None: + return None pageserver_id = meta[1] return self.get_pageserver(pageserver_id) diff --git a/test_runner/regress/test_disk_usage_eviction.py b/test_runner/regress/test_disk_usage_eviction.py index dcbf8a5025..061c57c88b 100644 --- a/test_runner/regress/test_disk_usage_eviction.py +++ b/test_runner/regress/test_disk_usage_eviction.py @@ -17,7 +17,7 @@ from fixtures.neon_fixtures import ( from fixtures.pageserver.http import PageserverHttpClient from fixtures.pageserver.utils import wait_for_upload_queue_empty from fixtures.remote_storage import RemoteStorageKind -from fixtures.types import Lsn, TenantId, TenantShardId, TimelineId +from fixtures.types import Lsn, TenantId, TimelineId from fixtures.utils import wait_until GLOBAL_LRU_LOG_LINE = "tenant_min_resident_size-respecting LRU would not relieve pressure, evicting more following global LRU policy" @@ -194,8 +194,10 @@ class EvictionEnv: # we now do initial logical size calculation on startup, which on debug builds can fight with disk usage based eviction for tenant_id, timeline_id in self.timelines: - pageserver_http = self.neon_env.get_tenant_pageserver(tenant_id).http_client() - pageserver_http.timeline_wait_logical_size(tenant_id, timeline_id) + tenant_ps = self.neon_env.get_tenant_pageserver(tenant_id) + # Pageserver may be none if we are currently not attached anywhere, e.g. during secondary eviction test + if tenant_ps is not None: + tenant_ps.http_client().timeline_wait_logical_size(tenant_id, timeline_id) def statvfs_called(): assert pageserver.log_contains(".*running mocked statvfs.*") @@ -864,18 +866,18 @@ def test_secondary_mode_eviction(eviction_env_ha: EvictionEnv): # Set up a situation where one pageserver _only_ has secondary locations on it, # so that when we release space we are sure it is via secondary locations. - - log.info("Setting up secondary location...") - ps_attached = env.neon_env.pageservers[0] + log.info("Setting up secondary locations...") ps_secondary = env.neon_env.pageservers[1] for tenant_id in tenant_ids: - # Migrate all attached tenants to the same pageserver, so that all the secondaries - # will run on the other pageserver. This is necessary because when we create tenants, - # they are spread over pageservers by default. - env.neon_env.attachment_service.tenant_shard_migrate( - TenantShardId(tenant_id, 0, 0), ps_attached.id - ) + # Find where it is attached + pageserver = env.neon_env.get_tenant_pageserver(tenant_id) + pageserver.http_client().tenant_heatmap_upload(tenant_id) + # Detach it + pageserver.tenant_detach(tenant_id) + + # Create a secondary mode location for the tenant, all tenants on one pageserver that will only + # contain secondary locations: this is the one where we will exercise disk usage eviction ps_secondary.tenant_location_configure( tenant_id, { @@ -887,8 +889,8 @@ def test_secondary_mode_eviction(eviction_env_ha: EvictionEnv): readback_conf = ps_secondary.read_tenant_location_conf(tenant_id) log.info(f"Read back conf: {readback_conf}") - # Request secondary location to download all layers that the attached location has - ps_attached.http_client().tenant_heatmap_upload(tenant_id) + # Request secondary location to download all layers that the attached location indicated + # in its heatmap ps_secondary.http_client().tenant_secondary_download(tenant_id) # Configure the secondary pageserver to have a phony small disk size diff --git a/test_runner/regress/test_sharding_service.py b/test_runner/regress/test_sharding_service.py index 5c70378ab0..ee57fcb2cf 100644 --- a/test_runner/regress/test_sharding_service.py +++ b/test_runner/regress/test_sharding_service.py @@ -35,6 +35,11 @@ def test_sharding_service_smoke( neon_env_builder.num_pageservers = 3 env = neon_env_builder.init_configs() + for pageserver in env.pageservers: + # This test detaches tenants during migration, which can race with deletion queue operations, + # during detach we only do an advisory flush, we don't wait for it. + pageserver.allowed_errors.extend([".*Dropped remote consistent LSN updates.*"]) + # Start services by hand so that we can skip a pageserver (this will start + register later) env.broker.try_start() env.attachment_service.start() From 27a3c9ecbe8fd09f35bbe534c0628831f29d0a1f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 6 Feb 2024 13:15:07 +0000 Subject: [PATCH 18/72] build(deps): bump cryptography from 41.0.6 to 42.0.0 (#6643) --- poetry.lock | 65 ++++++++++++++++++++++++++++++----------------------- 1 file changed, 37 insertions(+), 28 deletions(-) diff --git a/poetry.lock b/poetry.lock index 2904e2872e..e18cd4a74d 100644 --- a/poetry.lock +++ b/poetry.lock @@ -836,47 +836,56 @@ files = [ [[package]] name = "cryptography" -version = "41.0.6" +version = "42.0.0" description = "cryptography is a package which provides cryptographic recipes and primitives to Python developers." optional = false python-versions = ">=3.7" files = [ - {file = "cryptography-41.0.6-cp37-abi3-macosx_10_12_universal2.whl", hash = "sha256:0f27acb55a4e77b9be8d550d762b0513ef3fc658cd3eb15110ebbcbd626db12c"}, - {file = "cryptography-41.0.6-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ae236bb8760c1e55b7a39b6d4d32d2279bc6c7c8500b7d5a13b6fb9fc97be35b"}, - {file = "cryptography-41.0.6-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afda76d84b053923c27ede5edc1ed7d53e3c9f475ebaf63c68e69f1403c405a8"}, - {file = "cryptography-41.0.6-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:da46e2b5df770070412c46f87bac0849b8d685c5f2679771de277a422c7d0b86"}, - {file = "cryptography-41.0.6-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:ff369dd19e8fe0528b02e8df9f2aeb2479f89b1270d90f96a63500afe9af5cae"}, - {file = "cryptography-41.0.6-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b648fe2a45e426aaee684ddca2632f62ec4613ef362f4d681a9a6283d10e079d"}, - {file = "cryptography-41.0.6-cp37-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:5daeb18e7886a358064a68dbcaf441c036cbdb7da52ae744e7b9207b04d3908c"}, - {file = "cryptography-41.0.6-cp37-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:068bc551698c234742c40049e46840843f3d98ad7ce265fd2bd4ec0d11306596"}, - {file = "cryptography-41.0.6-cp37-abi3-win32.whl", hash = "sha256:2132d5865eea673fe6712c2ed5fb4fa49dba10768bb4cc798345748380ee3660"}, - {file = "cryptography-41.0.6-cp37-abi3-win_amd64.whl", hash = "sha256:48783b7e2bef51224020efb61b42704207dde583d7e371ef8fc2a5fb6c0aabc7"}, - {file = "cryptography-41.0.6-pp310-pypy310_pp73-macosx_10_12_x86_64.whl", hash = "sha256:8efb2af8d4ba9dbc9c9dd8f04d19a7abb5b49eab1f3694e7b5a16a5fc2856f5c"}, - {file = "cryptography-41.0.6-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:c5a550dc7a3b50b116323e3d376241829fd326ac47bc195e04eb33a8170902a9"}, - {file = "cryptography-41.0.6-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:85abd057699b98fce40b41737afb234fef05c67e116f6f3650782c10862c43da"}, - {file = "cryptography-41.0.6-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:f39812f70fc5c71a15aa3c97b2bbe213c3f2a460b79bd21c40d033bb34a9bf36"}, - {file = "cryptography-41.0.6-pp38-pypy38_pp73-macosx_10_12_x86_64.whl", hash = "sha256:742ae5e9a2310e9dade7932f9576606836ed174da3c7d26bc3d3ab4bd49b9f65"}, - {file = "cryptography-41.0.6-pp38-pypy38_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:35f3f288e83c3f6f10752467c48919a7a94b7d88cc00b0668372a0d2ad4f8ead"}, - {file = "cryptography-41.0.6-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:4d03186af98b1c01a4eda396b137f29e4e3fb0173e30f885e27acec8823c1b09"}, - {file = "cryptography-41.0.6-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:b27a7fd4229abef715e064269d98a7e2909ebf92eb6912a9603c7e14c181928c"}, - {file = "cryptography-41.0.6-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:398ae1fc711b5eb78e977daa3cbf47cec20f2c08c5da129b7a296055fbb22aed"}, - {file = "cryptography-41.0.6-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:7e00fb556bda398b99b0da289ce7053639d33b572847181d6483ad89835115f6"}, - {file = "cryptography-41.0.6-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:60e746b11b937911dc70d164060d28d273e31853bb359e2b2033c9e93e6f3c43"}, - {file = "cryptography-41.0.6-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:3288acccef021e3c3c10d58933f44e8602cf04dba96d9796d70d537bb2f4bbc4"}, - {file = "cryptography-41.0.6.tar.gz", hash = "sha256:422e3e31d63743855e43e5a6fcc8b4acab860f560f9321b0ee6269cc7ed70cc3"}, + {file = "cryptography-42.0.0-cp37-abi3-macosx_10_12_universal2.whl", hash = "sha256:c640b0ef54138fde761ec99a6c7dc4ce05e80420262c20fa239e694ca371d434"}, + {file = "cryptography-42.0.0-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:678cfa0d1e72ef41d48993a7be75a76b0725d29b820ff3cfd606a5b2b33fda01"}, + {file = "cryptography-42.0.0-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:146e971e92a6dd042214b537a726c9750496128453146ab0ee8971a0299dc9bd"}, + {file = "cryptography-42.0.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:87086eae86a700307b544625e3ba11cc600c3c0ef8ab97b0fda0705d6db3d4e3"}, + {file = "cryptography-42.0.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0a68bfcf57a6887818307600c3c0ebc3f62fbb6ccad2240aa21887cda1f8df1b"}, + {file = "cryptography-42.0.0-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:5a217bca51f3b91971400890905a9323ad805838ca3fa1e202a01844f485ee87"}, + {file = "cryptography-42.0.0-cp37-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:ca20550bb590db16223eb9ccc5852335b48b8f597e2f6f0878bbfd9e7314eb17"}, + {file = "cryptography-42.0.0-cp37-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:33588310b5c886dfb87dba5f013b8d27df7ffd31dc753775342a1e5ab139e59d"}, + {file = "cryptography-42.0.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9515ea7f596c8092fdc9902627e51b23a75daa2c7815ed5aa8cf4f07469212ec"}, + {file = "cryptography-42.0.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:35cf6ed4c38f054478a9df14f03c1169bb14bd98f0b1705751079b25e1cb58bc"}, + {file = "cryptography-42.0.0-cp37-abi3-win32.whl", hash = "sha256:8814722cffcfd1fbd91edd9f3451b88a8f26a5fd41b28c1c9193949d1c689dc4"}, + {file = "cryptography-42.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:a2a8d873667e4fd2f34aedab02ba500b824692c6542e017075a2efc38f60a4c0"}, + {file = "cryptography-42.0.0-cp39-abi3-macosx_10_12_universal2.whl", hash = "sha256:8fedec73d590fd30c4e3f0d0f4bc961aeca8390c72f3eaa1a0874d180e868ddf"}, + {file = "cryptography-42.0.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be41b0c7366e5549265adf2145135dca107718fa44b6e418dc7499cfff6b4689"}, + {file = "cryptography-42.0.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3ca482ea80626048975360c8e62be3ceb0f11803180b73163acd24bf014133a0"}, + {file = "cryptography-42.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:c58115384bdcfe9c7f644c72f10f6f42bed7cf59f7b52fe1bf7ae0a622b3a139"}, + {file = "cryptography-42.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:56ce0c106d5c3fec1038c3cca3d55ac320a5be1b44bf15116732d0bc716979a2"}, + {file = "cryptography-42.0.0-cp39-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:324721d93b998cb7367f1e6897370644751e5580ff9b370c0a50dc60a2003513"}, + {file = "cryptography-42.0.0-cp39-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:d97aae66b7de41cdf5b12087b5509e4e9805ed6f562406dfcf60e8481a9a28f8"}, + {file = "cryptography-42.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:85f759ed59ffd1d0baad296e72780aa62ff8a71f94dc1ab340386a1207d0ea81"}, + {file = "cryptography-42.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:206aaf42e031b93f86ad60f9f5d9da1b09164f25488238ac1dc488334eb5e221"}, + {file = "cryptography-42.0.0-cp39-abi3-win32.whl", hash = "sha256:74f18a4c8ca04134d2052a140322002fef535c99cdbc2a6afc18a8024d5c9d5b"}, + {file = "cryptography-42.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:14e4b909373bc5bf1095311fa0f7fcabf2d1a160ca13f1e9e467be1ac4cbdf94"}, + {file = "cryptography-42.0.0-pp310-pypy310_pp73-macosx_10_12_x86_64.whl", hash = "sha256:3005166a39b70c8b94455fdbe78d87a444da31ff70de3331cdec2c568cf25b7e"}, + {file = "cryptography-42.0.0-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:be14b31eb3a293fc6e6aa2807c8a3224c71426f7c4e3639ccf1a2f3ffd6df8c3"}, + {file = "cryptography-42.0.0-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:bd7cf7a8d9f34cc67220f1195884151426ce616fdc8285df9054bfa10135925f"}, + {file = "cryptography-42.0.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:c310767268d88803b653fffe6d6f2f17bb9d49ffceb8d70aed50ad45ea49ab08"}, + {file = "cryptography-42.0.0-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:bdce70e562c69bb089523e75ef1d9625b7417c6297a76ac27b1b8b1eb51b7d0f"}, + {file = "cryptography-42.0.0-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:e9326ca78111e4c645f7e49cbce4ed2f3f85e17b61a563328c85a5208cf34440"}, + {file = "cryptography-42.0.0-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:69fd009a325cad6fbfd5b04c711a4da563c6c4854fc4c9544bff3088387c77c0"}, + {file = "cryptography-42.0.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:988b738f56c665366b1e4bfd9045c3efae89ee366ca3839cd5af53eaa1401bce"}, + {file = "cryptography-42.0.0.tar.gz", hash = "sha256:6cf9b76d6e93c62114bd19485e5cb003115c134cf9ce91f8ac924c44f8c8c3f4"}, ] [package.dependencies] -cffi = ">=1.12" +cffi = {version = ">=1.12", markers = "platform_python_implementation != \"PyPy\""} [package.extras] 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)"] +docstest = ["pyenchant (>=1.6.11)", "readme-renderer", "sphinxcontrib-spelling (>=4.0.1)"] nox = ["nox"] -pep8test = ["black", "check-sdist", "mypy", "ruff"] +pep8test = ["check-sdist", "click", "mypy", "ruff"] sdist = ["build"] ssh = ["bcrypt (>=3.1.5)"] -test = ["pretend", "pytest (>=6.2.0)", "pytest-benchmark", "pytest-cov", "pytest-xdist"] +test = ["certifi", "pretend", "pytest (>=6.2.0)", "pytest-benchmark", "pytest-cov", "pytest-xdist"] test-randomorder = ["pytest-randomly"] [[package]] From 53a3ed0a7e26ddba5a6a70b2a5176ee7d5491283 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Tue, 6 Feb 2024 15:43:33 +0100 Subject: [PATCH 19/72] debug_assert presence of `shard_id` tracing field (#6572) also: fixes https://github.com/neondatabase/neon/issues/6638 --- libs/utils/src/tracing_span_assert.rs | 51 +++++++------- pageserver/src/http/routes.rs | 10 +-- pageserver/src/lib.rs | 1 + pageserver/src/page_service.rs | 68 +++++++++++++------ pageserver/src/pgdatadir_mapping.rs | 3 +- pageserver/src/span.rs | 43 ++++++++++++ pageserver/src/tenant.rs | 16 +++-- pageserver/src/tenant/mgr.rs | 22 +++--- .../src/tenant/remote_timeline_client.rs | 1 + .../tenant/remote_timeline_client/download.rs | 2 +- pageserver/src/tenant/span.rs | 17 ----- pageserver/src/tenant/storage_layer/layer.rs | 3 + pageserver/src/tenant/timeline.rs | 8 +-- pageserver/src/tenant/timeline/span.rs | 19 ------ pageserver/src/walredo.rs | 9 +++ pageserver/src/walredo/process.rs | 4 +- 16 files changed, 165 insertions(+), 112 deletions(-) create mode 100644 pageserver/src/span.rs delete mode 100644 pageserver/src/tenant/span.rs diff --git a/libs/utils/src/tracing_span_assert.rs b/libs/utils/src/tracing_span_assert.rs index db17f7d8cd..d24c81ad0b 100644 --- a/libs/utils/src/tracing_span_assert.rs +++ b/libs/utils/src/tracing_span_assert.rs @@ -20,13 +20,13 @@ //! //! // Then, in the main code: //! -//! let span = tracing::info_span!("TestSpan", test_id = 1); +//! let span = tracing::info_span!("TestSpan", tenant_id = 1); //! let _guard = span.enter(); //! //! // ... down the call stack //! -//! use utils::tracing_span_assert::{check_fields_present, MultiNameExtractor}; -//! let extractor = MultiNameExtractor::new("TestExtractor", ["test", "test_id"]); +//! use utils::tracing_span_assert::{check_fields_present, ConstExtractor}; +//! let extractor = ConstExtractor::new("tenant_id"); //! if let Err(missing) = check_fields_present!([&extractor]) { //! // if you copypaste this to a custom assert method, remember to add #[track_caller] //! // to get the "user" code location for the panic. @@ -45,27 +45,26 @@ pub enum ExtractionResult { } pub trait Extractor: Send + Sync + std::fmt::Debug { - fn name(&self) -> &str; + fn id(&self) -> &str; fn extract(&self, fields: &tracing::field::FieldSet) -> ExtractionResult; } #[derive(Debug)] -pub struct MultiNameExtractor { - name: &'static str, - field_names: [&'static str; L], +pub struct ConstExtractor { + field_name: &'static str, } -impl MultiNameExtractor { - pub fn new(name: &'static str, field_names: [&'static str; L]) -> MultiNameExtractor { - MultiNameExtractor { name, field_names } +impl ConstExtractor { + pub const fn new(field_name: &'static str) -> ConstExtractor { + ConstExtractor { field_name } } } -impl Extractor for MultiNameExtractor { - fn name(&self) -> &str { - self.name +impl Extractor for ConstExtractor { + fn id(&self) -> &str { + self.field_name } fn extract(&self, fields: &tracing::field::FieldSet) -> ExtractionResult { - if fields.iter().any(|f| self.field_names.contains(&f.name())) { + if fields.iter().any(|f| f.name() == self.field_name) { ExtractionResult::Present } else { ExtractionResult::Absent @@ -203,19 +202,19 @@ mod tests { } impl<'a> fmt::Debug for MemoryIdentity<'a> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{:p}: {}", self.as_ptr(), self.0.name()) + write!(f, "{:p}: {}", self.as_ptr(), self.0.id()) } } struct Setup { _current_thread_subscriber_guard: tracing::subscriber::DefaultGuard, - tenant_extractor: MultiNameExtractor<2>, - timeline_extractor: MultiNameExtractor<2>, + tenant_extractor: ConstExtractor, + timeline_extractor: ConstExtractor, } fn setup_current_thread() -> Setup { - let tenant_extractor = MultiNameExtractor::new("TenantId", ["tenant_id", "tenant"]); - let timeline_extractor = MultiNameExtractor::new("TimelineId", ["timeline_id", "timeline"]); + let tenant_extractor = ConstExtractor::new("tenant_id"); + let timeline_extractor = ConstExtractor::new("timeline_id"); let registry = tracing_subscriber::registry() .with(tracing_subscriber::fmt::layer()) @@ -343,12 +342,12 @@ mod tests { let span = tracing::info_span!("foo", e = "some value"); let _guard = span.enter(); - let extractor = MultiNameExtractor::new("E", ["e"]); + let extractor = ConstExtractor::new("e"); let res = check_fields_present0([&extractor]); assert!(matches!(res, Ok(Summary::Unconfigured)), "{res:?}"); // similarly for a not found key - let extractor = MultiNameExtractor::new("F", ["foobar"]); + let extractor = ConstExtractor::new("foobar"); let res = check_fields_present0([&extractor]); assert!(matches!(res, Ok(Summary::Unconfigured)), "{res:?}"); } @@ -368,16 +367,14 @@ mod tests { // normally this would work, but without any tracing-subscriber configured, both // check_field_present find nothing let _guard = subspan.enter(); - let extractors: [&dyn Extractor; 2] = [ - &MultiNameExtractor::new("E", ["e"]), - &MultiNameExtractor::new("F", ["f"]), - ]; + let extractors: [&dyn Extractor; 2] = + [&ConstExtractor::new("e"), &ConstExtractor::new("f")]; let res = check_fields_present0(extractors); assert!(matches!(res, Ok(Summary::Unconfigured)), "{res:?}"); // similarly for a not found key - let extractor = MultiNameExtractor::new("G", ["g"]); + let extractor = ConstExtractor::new("g"); let res = check_fields_present0([&extractor]); assert!(matches!(res, Ok(Summary::Unconfigured)), "{res:?}"); } @@ -410,7 +407,7 @@ mod tests { let span = tracing::info_span!("foo", e = "some value"); let _guard = span.enter(); - let extractors: [&dyn Extractor; 1] = [&MultiNameExtractor::new("E", ["e"])]; + let extractors: [&dyn Extractor; 1] = [&ConstExtractor::new("e")]; if span.is_disabled() { // the tests are running single threaded, or we got lucky and no other tests subscriber diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index b97e272c86..792089ebe7 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -535,7 +535,7 @@ async fn timeline_create_handler( } .instrument(info_span!("timeline_create", tenant_id = %tenant_shard_id.tenant_id, - shard = %tenant_shard_id.shard_slug(), + shard_id = %tenant_shard_id.shard_slug(), timeline_id = %new_timeline_id, lsn=?request_data.ancestor_start_lsn, pg_version=?request_data.pg_version)) .await } @@ -831,7 +831,7 @@ async fn timeline_delete_handler( } })?; tenant.wait_to_become_active(ACTIVE_TENANT_TIMEOUT).await?; - tenant.delete_timeline(timeline_id).instrument(info_span!("timeline_delete", tenant_id=%tenant_shard_id.tenant_id, shard=%tenant_shard_id.shard_slug(), %timeline_id)) + tenant.delete_timeline(timeline_id).instrument(info_span!("timeline_delete", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), %timeline_id)) .await?; json_response(StatusCode::ACCEPTED, ()) @@ -856,7 +856,7 @@ async fn tenant_detach_handler( detach_ignored.unwrap_or(false), &state.deletion_queue_client, ) - .instrument(info_span!("tenant_detach", %tenant_id)) + .instrument(info_span!("tenant_detach", %tenant_id, shard_id=%tenant_shard_id.shard_slug())) .await?; json_response(StatusCode::OK, ()) @@ -1007,7 +1007,7 @@ async fn tenant_delete_handler( .delete_tenant(tenant_shard_id, ACTIVE_TENANT_TIMEOUT) .instrument(info_span!("tenant_delete_handler", tenant_id = %tenant_shard_id.tenant_id, - shard = %tenant_shard_id.shard_slug() + shard_id = %tenant_shard_id.shard_slug() )) .await?; @@ -1363,7 +1363,7 @@ async fn put_tenant_location_config_handler( mgr::detach_tenant(conf, tenant_shard_id, true, &state.deletion_queue_client) .instrument(info_span!("tenant_detach", tenant_id = %tenant_shard_id.tenant_id, - shard = %tenant_shard_id.shard_slug() + shard_id = %tenant_shard_id.shard_slug() )) .await { diff --git a/pageserver/src/lib.rs b/pageserver/src/lib.rs index bcde1166b7..c3f35142ec 100644 --- a/pageserver/src/lib.rs +++ b/pageserver/src/lib.rs @@ -17,6 +17,7 @@ pub mod page_cache; pub mod page_service; pub mod pgdatadir_mapping; pub mod repository; +pub mod span; pub(crate) mod statvfs; pub mod task_mgr; pub mod tenant; diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs index 754c021c88..6fc38a76d4 100644 --- a/pageserver/src/page_service.rs +++ b/pageserver/src/page_service.rs @@ -63,9 +63,10 @@ use crate::import_datadir::import_wal_from_tar; use crate::metrics; use crate::metrics::LIVE_CONNECTIONS_COUNT; use crate::pgdatadir_mapping::Version; +use crate::span::debug_assert_current_span_has_tenant_and_timeline_id; +use crate::span::debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id; use crate::task_mgr; use crate::task_mgr::TaskKind; -use crate::tenant::debug_assert_current_span_has_tenant_and_timeline_id; use crate::tenant::mgr; use crate::tenant::mgr::get_active_tenant_with_timeout; use crate::tenant::mgr::GetActiveTenantError; @@ -549,7 +550,7 @@ impl PageServerHandler { where IO: AsyncRead + AsyncWrite + Send + Sync + Unpin, { - debug_assert_current_span_has_tenant_and_timeline_id(); + debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id(); let tenant = mgr::get_active_tenant_with_timeout( tenant_id, @@ -631,6 +632,7 @@ impl PageServerHandler { ) } PagestreamFeMessage::GetPage(req) => { + // shard_id is filled in by the handler let span = tracing::info_span!("handle_get_page_at_lsn_request", rel = %req.rel, blkno = %req.blkno, req_lsn = %req.lsn); ( self.handle_get_page_at_lsn_request(tenant_id, timeline_id, &req, &ctx) @@ -719,7 +721,7 @@ impl PageServerHandler { where IO: AsyncRead + AsyncWrite + Send + Sync + Unpin, { - debug_assert_current_span_has_tenant_and_timeline_id(); + debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id(); // Create empty timeline info!("creating new timeline"); @@ -772,7 +774,7 @@ impl PageServerHandler { Ok(()) } - #[instrument(skip_all, fields(%start_lsn, %end_lsn))] + #[instrument(skip_all, fields(shard_id, %start_lsn, %end_lsn))] async fn handle_import_wal( &self, pgb: &mut PostgresBackend, @@ -785,8 +787,6 @@ impl PageServerHandler { where IO: AsyncRead + AsyncWrite + Send + Sync + Unpin, { - debug_assert_current_span_has_tenant_and_timeline_id(); - let timeline = self .get_active_tenant_timeline(tenant_id, timeline_id, ShardSelector::Zero) .await?; @@ -893,6 +893,7 @@ impl PageServerHandler { Ok(lsn) } + #[instrument(skip_all, fields(shard_id))] async fn handle_get_rel_exists_request( &mut self, tenant_id: TenantId, @@ -919,6 +920,7 @@ impl PageServerHandler { })) } + #[instrument(skip_all, fields(shard_id))] async fn handle_get_nblocks_request( &mut self, tenant_id: TenantId, @@ -946,6 +948,7 @@ impl PageServerHandler { })) } + #[instrument(skip_all, fields(shard_id))] async fn handle_db_size_request( &mut self, tenant_id: TenantId, @@ -1096,6 +1099,7 @@ impl PageServerHandler { } } + #[instrument(skip_all, fields(shard_id))] async fn handle_get_page_at_lsn_request( &mut self, tenant_id: TenantId, @@ -1129,6 +1133,9 @@ impl PageServerHandler { } }; + // load_timeline_for_page sets shard_id, but get_cached_timeline_for_page doesn't + set_tracing_field_shard_id(timeline); + let _timer = timeline .query_metrics .start_timer(metrics::SmgrQueryType::GetPageAtLsn); @@ -1147,6 +1154,7 @@ impl PageServerHandler { })) } + #[instrument(skip_all, fields(shard_id))] async fn handle_get_slru_segment_request( &mut self, tenant_id: TenantId, @@ -1175,7 +1183,7 @@ impl PageServerHandler { } #[allow(clippy::too_many_arguments)] - #[instrument(skip_all, fields(?lsn, ?prev_lsn, %full_backup))] + #[instrument(skip_all, fields(shard_id, ?lsn, ?prev_lsn, %full_backup))] async fn handle_basebackup_request( &mut self, pgb: &mut PostgresBackend, @@ -1190,8 +1198,6 @@ impl PageServerHandler { where IO: AsyncRead + AsyncWrite + Send + Sync + Unpin, { - debug_assert_current_span_has_tenant_and_timeline_id(); - let started = std::time::Instant::now(); // check that the timeline exists @@ -1313,6 +1319,7 @@ impl PageServerHandler { .await .map_err(GetActiveTimelineError::Tenant)?; let timeline = tenant.get_timeline(timeline_id, true)?; + set_tracing_field_shard_id(&timeline); Ok(timeline) } } @@ -1477,21 +1484,29 @@ where .record("timeline_id", field::display(timeline_id)); self.check_permission(Some(tenant_id))?; - let timeline = self - .get_active_tenant_timeline(tenant_id, timeline_id, ShardSelector::Zero) - .await?; + async { + let timeline = self + .get_active_tenant_timeline(tenant_id, timeline_id, ShardSelector::Zero) + .await?; - let end_of_timeline = timeline.get_last_record_rlsn(); + let end_of_timeline = timeline.get_last_record_rlsn(); - pgb.write_message_noflush(&BeMessage::RowDescription(&[ - RowDescriptor::text_col(b"prev_lsn"), - RowDescriptor::text_col(b"last_lsn"), - ]))? - .write_message_noflush(&BeMessage::DataRow(&[ - Some(end_of_timeline.prev.to_string().as_bytes()), - Some(end_of_timeline.last.to_string().as_bytes()), - ]))? - .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?; + pgb.write_message_noflush(&BeMessage::RowDescription(&[ + RowDescriptor::text_col(b"prev_lsn"), + RowDescriptor::text_col(b"last_lsn"), + ]))? + .write_message_noflush(&BeMessage::DataRow(&[ + Some(end_of_timeline.prev.to_string().as_bytes()), + Some(end_of_timeline.last.to_string().as_bytes()), + ]))? + .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?; + anyhow::Ok(()) + } + .instrument(info_span!( + "handle_get_last_record_lsn", + shard_id = tracing::field::Empty + )) + .await?; } // same as basebackup, but result includes relational data as well else if query_string.starts_with("fullbackup ") { @@ -1748,3 +1763,12 @@ impl From for QueryError { } } } + +fn set_tracing_field_shard_id(timeline: &Timeline) { + debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id(); + tracing::Span::current().record( + "shard_id", + tracing::field::display(timeline.tenant_shard_id.shard_slug()), + ); + debug_assert_current_span_has_tenant_and_timeline_id(); +} diff --git a/pageserver/src/pgdatadir_mapping.rs b/pageserver/src/pgdatadir_mapping.rs index a36785a69f..f1d18c0146 100644 --- a/pageserver/src/pgdatadir_mapping.rs +++ b/pageserver/src/pgdatadir_mapping.rs @@ -10,6 +10,7 @@ use super::tenant::{PageReconstructError, Timeline}; use crate::context::RequestContext; use crate::keyspace::{KeySpace, KeySpaceAccum}; use crate::repository::*; +use crate::span::debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id; use crate::walrecord::NeonWalRecord; use anyhow::{ensure, Context}; use bytes::{Buf, Bytes, BytesMut}; @@ -699,7 +700,7 @@ impl Timeline { lsn: Lsn, ctx: &RequestContext, ) -> Result { - crate::tenant::debug_assert_current_span_has_tenant_and_timeline_id(); + debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id(); // Fetch list of database dirs and iterate them let buf = self.get(DBDIR_KEY, lsn, ctx).await?; diff --git a/pageserver/src/span.rs b/pageserver/src/span.rs new file mode 100644 index 0000000000..91fee50514 --- /dev/null +++ b/pageserver/src/span.rs @@ -0,0 +1,43 @@ +use utils::tracing_span_assert::check_fields_present; + +mod extractors { + use utils::tracing_span_assert::ConstExtractor; + + pub(super) const TENANT_ID: ConstExtractor = ConstExtractor::new("tenant_id"); + pub(super) const SHARD_ID: ConstExtractor = ConstExtractor::new("shard_id"); + pub(super) const TIMELINE_ID: ConstExtractor = ConstExtractor::new("timeline_id"); +} + +#[track_caller] +pub(crate) fn debug_assert_current_span_has_tenant_id() { + if cfg!(debug_assertions) { + if let Err(missing) = check_fields_present!([&extractors::TENANT_ID, &extractors::SHARD_ID]) + { + panic!("missing extractors: {missing:?}") + } + } +} + +#[track_caller] +pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id() { + if cfg!(debug_assertions) { + if let Err(missing) = check_fields_present!([ + &extractors::TENANT_ID, + &extractors::SHARD_ID, + &extractors::TIMELINE_ID, + ]) { + panic!("missing extractors: {missing:?}") + } + } +} + +#[track_caller] +pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id() { + if cfg!(debug_assertions) { + if let Err(missing) = + check_fields_present!([&extractors::TENANT_ID, &extractors::TIMELINE_ID,]) + { + panic!("missing extractors: {missing:?}") + } + } +} diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index 624c3e365f..fe85cf9753 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -100,6 +100,7 @@ use std::sync::Arc; use std::sync::{Mutex, RwLock}; use std::time::{Duration, Instant}; +use crate::span; use crate::tenant::timeline::delete::DeleteTimelineFlow; use crate::tenant::timeline::uninit::cleanup_timeline_directory; use crate::virtual_file::VirtualFile; @@ -150,7 +151,6 @@ pub mod block_io; pub mod disk_btree; pub(crate) mod ephemeral_file; pub mod layer_map; -mod span; pub mod metadata; mod par_fsync; @@ -168,7 +168,7 @@ pub(crate) mod timeline; pub mod size; -pub(crate) use timeline::span::debug_assert_current_span_has_tenant_and_timeline_id; +pub(crate) use crate::span::debug_assert_current_span_has_tenant_and_timeline_id; pub(crate) use timeline::{LogicalSizeCalculationCause, PageReconstructError, Timeline}; // re-export for use in remote_timeline_client.rs @@ -3998,6 +3998,10 @@ pub(crate) mod harness { }) } + pub fn span(&self) -> tracing::Span { + info_span!("TenantHarness", tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug()) + } + pub async fn load(&self) -> (Arc, RequestContext) { let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error); ( @@ -4602,7 +4606,7 @@ mod tests { // so that all uploads finish & we can call harness.load() below again tenant .shutdown(Default::default(), true) - .instrument(info_span!("test_shutdown", tenant_id=%tenant.tenant_shard_id)) + .instrument(harness.span()) .await .ok() .unwrap(); @@ -4643,7 +4647,7 @@ mod tests { // so that all uploads finish & we can call harness.load() below again tenant .shutdown(Default::default(), true) - .instrument(info_span!("test_shutdown", tenant_id=%tenant.tenant_shard_id)) + .instrument(harness.span()) .await .ok() .unwrap(); @@ -4705,7 +4709,7 @@ mod tests { // so that all uploads finish & we can call harness.try_load() below again tenant .shutdown(Default::default(), true) - .instrument(info_span!("test_shutdown", tenant_id=%tenant.tenant_shard_id)) + .instrument(harness.span()) .await .ok() .unwrap(); @@ -5238,7 +5242,7 @@ mod tests { let raw_tline = tline.raw_timeline().unwrap(); raw_tline .shutdown() - .instrument(info_span!("test_shutdown", tenant_id=%raw_tline.tenant_shard_id, timeline_id=%TIMELINE_ID)) + .instrument(info_span!("test_shutdown", tenant_id=%raw_tline.tenant_shard_id, shard_id=%raw_tline.tenant_shard_id.shard_slug(), timeline_id=%TIMELINE_ID)) .await; std::mem::forget(tline); } diff --git a/pageserver/src/tenant/mgr.rs b/pageserver/src/tenant/mgr.rs index de0b636d47..5ec910ca3e 100644 --- a/pageserver/src/tenant/mgr.rs +++ b/pageserver/src/tenant/mgr.rs @@ -684,7 +684,7 @@ async fn shutdown_all_tenants0(tenants: &std::sync::RwLock) { // going to log too many lines debug!("tenant successfully stopped"); } - .instrument(info_span!("shutdown", tenant_id=%tenant_shard_id.tenant_id, shard=%tenant_shard_id.shard_slug())), + .instrument(info_span!("shutdown", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug())), ); total_attached += 1; @@ -1720,6 +1720,7 @@ pub(crate) async fn ignore_tenant( ignore_tenant0(conf, &TENANTS, tenant_id).await } +#[instrument(skip_all, fields(shard_id))] async fn ignore_tenant0( conf: &'static PageServerConf, tenants: &std::sync::RwLock, @@ -1727,6 +1728,10 @@ async fn ignore_tenant0( ) -> Result<(), TenantStateError> { // This is a legacy API (replaced by `/location_conf`). It does not support sharding let tenant_shard_id = TenantShardId::unsharded(tenant_id); + tracing::Span::current().record( + "shard_id", + tracing::field::display(tenant_shard_id.shard_slug()), + ); remove_tenant_from_memory(tenants, tenant_shard_id, async { let ignore_mark_file = conf.tenant_ignore_mark_file_path(&tenant_shard_id); @@ -2122,7 +2127,7 @@ fn tenant_map_acquire_slot_impl( METRICS.tenant_slot_writes.inc(); let mut locked = tenants.write().unwrap(); - let span = tracing::info_span!("acquire_slot", tenant_id=%tenant_shard_id.tenant_id, shard = %tenant_shard_id.shard_slug()); + let span = tracing::info_span!("acquire_slot", tenant_id=%tenant_shard_id.tenant_id, shard_id = %tenant_shard_id.shard_slug()); let _guard = span.enter(); let m = match &mut *locked { @@ -2358,7 +2363,7 @@ pub(crate) async fn immediate_gc( mod tests { use std::collections::BTreeMap; use std::sync::Arc; - use tracing::{info_span, Instrument}; + use tracing::Instrument; use crate::tenant::mgr::TenantSlot; @@ -2369,17 +2374,16 @@ mod tests { // Test that if an InProgress tenant is in the map during shutdown, the shutdown will gracefully // wait for it to complete before proceeding. - let (t, _ctx) = TenantHarness::create("shutdown_awaits_in_progress_tenant") - .unwrap() - .load() - .await; + let h = TenantHarness::create("shutdown_awaits_in_progress_tenant").unwrap(); + let (t, _ctx) = h.load().await; // harness loads it to active, which is forced and nothing is running on the tenant let id = t.tenant_shard_id(); // tenant harness configures the logging and we cannot escape it - let _e = info_span!("testing", tenant_id = %id).entered(); + let span = h.span(); + let _e = span.enter(); let tenants = BTreeMap::from([(id, TenantSlot::Attached(t.clone()))]); let tenants = Arc::new(std::sync::RwLock::new(TenantsMap::Open(tenants))); @@ -2400,7 +2404,7 @@ mod tests { }; super::remove_tenant_from_memory(&tenants, id, cleanup).await } - .instrument(info_span!("foobar", tenant_id = %id)) + .instrument(h.span()) }); // now the long cleanup should be in place, with the stopping state diff --git a/pageserver/src/tenant/remote_timeline_client.rs b/pageserver/src/tenant/remote_timeline_client.rs index 831a073d17..152c9a2b7d 100644 --- a/pageserver/src/tenant/remote_timeline_client.rs +++ b/pageserver/src/tenant/remote_timeline_client.rs @@ -1952,6 +1952,7 @@ mod tests { tracing::info_span!( "test", tenant_id = %self.harness.tenant_shard_id.tenant_id, + shard_id = %self.harness.tenant_shard_id.shard_slug(), timeline_id = %TIMELINE_ID ) } diff --git a/pageserver/src/tenant/remote_timeline_client/download.rs b/pageserver/src/tenant/remote_timeline_client/download.rs index 2c50726b43..6c1125746b 100644 --- a/pageserver/src/tenant/remote_timeline_client/download.rs +++ b/pageserver/src/tenant/remote_timeline_client/download.rs @@ -17,11 +17,11 @@ use utils::timeout::timeout_cancellable; use utils::{backoff, crashsafe}; use crate::config::PageServerConf; +use crate::span::debug_assert_current_span_has_tenant_and_timeline_id; use crate::tenant::remote_timeline_client::{ download_cancellable, remote_layer_path, remote_timelines_path, DOWNLOAD_TIMEOUT, }; use crate::tenant::storage_layer::LayerFileName; -use crate::tenant::timeline::span::debug_assert_current_span_has_tenant_and_timeline_id; use crate::tenant::Generation; use crate::virtual_file::on_fatal_io_error; use crate::TEMP_FILE_SUFFIX; diff --git a/pageserver/src/tenant/span.rs b/pageserver/src/tenant/span.rs deleted file mode 100644 index 04e92f4096..0000000000 --- a/pageserver/src/tenant/span.rs +++ /dev/null @@ -1,17 +0,0 @@ -#[cfg(debug_assertions)] -use utils::tracing_span_assert::{check_fields_present, MultiNameExtractor}; - -#[cfg(not(debug_assertions))] -pub(crate) fn debug_assert_current_span_has_tenant_id() {} - -#[cfg(debug_assertions)] -pub(crate) static TENANT_ID_EXTRACTOR: once_cell::sync::Lazy> = - once_cell::sync::Lazy::new(|| MultiNameExtractor::new("TenantId", ["tenant_id"])); - -#[cfg(debug_assertions)] -#[track_caller] -pub(crate) fn debug_assert_current_span_has_tenant_id() { - if let Err(missing) = check_fields_present!([&*TENANT_ID_EXTRACTOR]) { - panic!("missing extractors: {missing:?}") - } -} diff --git a/pageserver/src/tenant/storage_layer/layer.rs b/pageserver/src/tenant/storage_layer/layer.rs index 1f337adf53..52c0f8abdc 100644 --- a/pageserver/src/tenant/storage_layer/layer.rs +++ b/pageserver/src/tenant/storage_layer/layer.rs @@ -15,6 +15,7 @@ use utils::sync::heavier_once_cell; use crate::config::PageServerConf; use crate::context::RequestContext; use crate::repository::Key; +use crate::span::debug_assert_current_span_has_tenant_and_timeline_id; use crate::tenant::{remote_timeline_client::LayerFileMetadata, Timeline}; use super::delta_layer::{self, DeltaEntry}; @@ -836,6 +837,8 @@ impl LayerInner { timeline: Arc, permit: heavier_once_cell::InitPermit, ) -> Result { + debug_assert_current_span_has_tenant_and_timeline_id(); + let task_name = format!("download layer {}", self); let (tx, rx) = tokio::sync::oneshot::channel(); diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 50ffc4d265..43aa178ab5 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -1138,7 +1138,7 @@ impl Timeline { /// Shut down immediately, without waiting for any open layers to flush to disk. This is a subset of /// the graceful [`Timeline::flush_and_shutdown`] function. pub(crate) async fn shutdown(&self) { - span::debug_assert_current_span_has_tenant_and_timeline_id(); + debug_assert_current_span_has_tenant_and_timeline_id(); // Signal any subscribers to our cancellation token to drop out tracing::debug!("Cancelling CancellationToken"); @@ -1964,7 +1964,7 @@ impl Timeline { .await; Ok(()) } - .instrument(info_span!(parent: None, "initial_size_calculation", tenant_id=%self.tenant_shard_id.tenant_id, timeline_id=%self.timeline_id)), + .instrument(info_span!(parent: None, "initial_size_calculation", tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), timeline_id=%self.timeline_id)), ); } @@ -2151,7 +2151,7 @@ impl Timeline { cause: LogicalSizeCalculationCause, ctx: &RequestContext, ) -> Result { - span::debug_assert_current_span_has_tenant_and_timeline_id(); + crate::span::debug_assert_current_span_has_tenant_and_timeline_id(); // We should never be calculating logical sizes on shard !=0, because these shards do not have // accurate relation sizes, and they do not emit consumption metrics. debug_assert!(self.tenant_shard_id.is_zero()); @@ -2849,7 +2849,7 @@ impl Timeline { frozen_layer: Arc, ctx: &RequestContext, ) -> Result<(), FlushLayerError> { - span::debug_assert_current_span_has_tenant_and_timeline_id(); + debug_assert_current_span_has_tenant_and_timeline_id(); // As a special case, when we have just imported an image into the repository, // instead of writing out a L0 delta layer, we directly write out image layer // files instead. This is possible as long as *all* the data imported into the diff --git a/pageserver/src/tenant/timeline/span.rs b/pageserver/src/tenant/timeline/span.rs index 3b580c9d1b..8b13789179 100644 --- a/pageserver/src/tenant/timeline/span.rs +++ b/pageserver/src/tenant/timeline/span.rs @@ -1,20 +1 @@ -#[cfg(debug_assertions)] -use utils::tracing_span_assert::{check_fields_present, Extractor, MultiNameExtractor}; -#[cfg(not(debug_assertions))] -pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id() {} - -#[cfg(debug_assertions)] -#[track_caller] -pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id() { - static TIMELINE_ID_EXTRACTOR: once_cell::sync::Lazy> = - once_cell::sync::Lazy::new(|| MultiNameExtractor::new("TimelineId", ["timeline_id"])); - - let fields: [&dyn Extractor; 2] = [ - &*crate::tenant::span::TENANT_ID_EXTRACTOR, - &*TIMELINE_ID_EXTRACTOR, - ]; - if let Err(missing) = check_fields_present!(fields) { - panic!("missing extractors: {missing:?}") - } -} diff --git a/pageserver/src/walredo.rs b/pageserver/src/walredo.rs index 773e5fc051..98a6a0bb6c 100644 --- a/pageserver/src/walredo.rs +++ b/pageserver/src/walredo.rs @@ -373,6 +373,7 @@ mod tests { use bytes::Bytes; use pageserver_api::shard::TenantShardId; use std::str::FromStr; + use tracing::Instrument; use utils::{id::TenantId, lsn::Lsn}; #[tokio::test] @@ -397,6 +398,7 @@ mod tests { short_records(), 14, ) + .instrument(h.span()) .await .unwrap(); @@ -424,6 +426,7 @@ mod tests { short_records(), 14, ) + .instrument(h.span()) .await .unwrap(); @@ -444,6 +447,7 @@ mod tests { short_records(), 16, /* 16 currently produces stderr output on startup, which adds a nice extra edge */ ) + .instrument(h.span()) .await .unwrap_err(); } @@ -472,6 +476,7 @@ mod tests { // underscored because unused, except for removal at drop _repo_dir: camino_tempfile::Utf8TempDir, manager: PostgresRedoManager, + tenant_shard_id: TenantShardId, } impl RedoHarness { @@ -488,7 +493,11 @@ mod tests { Ok(RedoHarness { _repo_dir: repo_dir, manager, + tenant_shard_id, }) } + fn span(&self) -> tracing::Span { + tracing::info_span!("RedoHarness", tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug()) + } } } diff --git a/pageserver/src/walredo/process.rs b/pageserver/src/walredo/process.rs index 85db3b4a4a..bcbb263663 100644 --- a/pageserver/src/walredo/process.rs +++ b/pageserver/src/walredo/process.rs @@ -54,12 +54,14 @@ impl WalRedoProcess { // // Start postgres binary in special WAL redo mode. // - #[instrument(skip_all,fields(tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), pg_version=pg_version))] + #[instrument(skip_all,fields(pg_version=pg_version))] pub(crate) fn launch( conf: &'static PageServerConf, tenant_shard_id: TenantShardId, pg_version: u32, ) -> anyhow::Result { + crate::span::debug_assert_current_span_has_tenant_id(); + let pg_bin_dir_path = conf.pg_bin_dir(pg_version).context("pg_bin_dir")?; // TODO these should be infallible. let pg_lib_dir_path = conf.pg_lib_dir(pg_version).context("pg_lib_dir")?; From d7b29aace7eec730af45e7f12fbe5620545b48aa Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Tue, 6 Feb 2024 16:20:02 +0100 Subject: [PATCH 20/72] refactor(walredo): don't create WalRedoManager for broken tenants (#6597) When we'll later introduce a global pool of pre-spawned walredo processes (https://github.com/neondatabase/neon/issues/6581), this refactoring avoids plumbing through the reference to the pool to all the places where we create a broken tenant. Builds atop the refactoring in #6583 --- pageserver/src/tenant.rs | 18 +++++++----------- pageserver/src/tenant/tasks.rs | 4 +++- pageserver/src/tenant/timeline.rs | 9 ++++++--- 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index fe85cf9753..f704f8c0dd 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -278,7 +278,7 @@ pub struct Tenant { // with timelines, which in turn may cause dropping replication connection, expiration of wait_for_lsn // timeout... gc_cs: tokio::sync::Mutex<()>, - walredo_mgr: Arc, + walredo_mgr: Option>, // provides access to timeline data sitting in the remote storage pub(crate) remote_storage: Option, @@ -635,7 +635,7 @@ impl Tenant { conf, attached_conf, shard_identity, - wal_redo_manager, + Some(wal_redo_manager), tenant_shard_id, remote_storage.clone(), deletion_queue_client, @@ -1195,10 +1195,6 @@ impl Tenant { tenant_shard_id: TenantShardId, reason: String, ) -> Arc { - let wal_redo_manager = Arc::new(WalRedoManager::from(PostgresRedoManager::new( - conf, - tenant_shard_id, - ))); Arc::new(Tenant::new( TenantState::Broken { reason, @@ -1209,7 +1205,7 @@ impl Tenant { // Shard identity isn't meaningful for a broken tenant: it's just a placeholder // to occupy the slot for this TenantShardId. ShardIdentity::broken(tenant_shard_id.shard_number, tenant_shard_id.shard_count), - wal_redo_manager, + None, tenant_shard_id, None, DeletionQueueClient::broken(), @@ -1978,7 +1974,7 @@ impl Tenant { } pub(crate) fn wal_redo_manager_status(&self) -> Option { - self.walredo_mgr.status() + self.walredo_mgr.as_ref().and_then(|mgr| mgr.status()) } /// Changes tenant status to active, unless shutdown was already requested. @@ -2613,7 +2609,7 @@ impl Tenant { self.tenant_shard_id, self.generation, self.shard_identity, - Arc::clone(&self.walredo_mgr), + self.walredo_mgr.as_ref().map(Arc::clone), resources, pg_version, state, @@ -2631,7 +2627,7 @@ impl Tenant { conf: &'static PageServerConf, attached_conf: AttachedTenantConf, shard_identity: ShardIdentity, - walredo_mgr: Arc, + walredo_mgr: Option>, tenant_shard_id: TenantShardId, remote_storage: Option, deletion_queue_client: DeletionQueueClient, @@ -4055,7 +4051,7 @@ pub(crate) mod harness { .unwrap(), // This is a legacy/test code path: sharding isn't supported here. ShardIdentity::unsharded(), - walredo_mgr, + Some(walredo_mgr), self.tenant_shard_id, Some(self.remote_storage.clone()), self.deletion_queue.new_client(), diff --git a/pageserver/src/tenant/tasks.rs b/pageserver/src/tenant/tasks.rs index 5f39c46a84..950cc46e71 100644 --- a/pageserver/src/tenant/tasks.rs +++ b/pageserver/src/tenant/tasks.rs @@ -199,7 +199,9 @@ async fn compaction_loop(tenant: Arc, cancel: CancellationToken) { // Perhaps we did no work and the walredo process has been idle for some time: // give it a chance to shut down to avoid leaving walredo process running indefinitely. - tenant.walredo_mgr.maybe_quiesce(period * 10); + if let Some(walredo_mgr) = &tenant.walredo_mgr { + walredo_mgr.maybe_quiesce(period * 10); + } // Sleep if tokio::time::timeout(sleep_duration, cancel.cancelled()) diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 43aa178ab5..735b8003b4 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -215,8 +215,8 @@ pub struct Timeline { // Atomic would be more appropriate here. last_freeze_ts: RwLock, - // WAL redo manager - walredo_mgr: Arc, + // WAL redo manager. `None` only for broken tenants. + walredo_mgr: Option>, /// Remote storage client. /// See [`remote_timeline_client`](super::remote_timeline_client) module comment for details. @@ -1427,7 +1427,7 @@ impl Timeline { tenant_shard_id: TenantShardId, generation: Generation, shard_identity: ShardIdentity, - walredo_mgr: Arc, + walredo_mgr: Option>, resources: TimelineResources, pg_version: u32, state: TimelineState, @@ -4457,6 +4457,9 @@ impl Timeline { let img = match self .walredo_mgr + .as_ref() + .context("timeline has no walredo manager") + .map_err(PageReconstructError::WalRedo)? .request_redo(key, request_lsn, data.img, data.records, self.pg_version) .await .context("reconstruct a page image") From bb9272116816690f806b8932af037a8b69e10aa2 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Tue, 6 Feb 2024 17:53:04 +0200 Subject: [PATCH 21/72] build: migrate check-style-rust to small runners (#6588) We have more small runners than large runners, and often a shortage of large runners. Migrate `check-style-rust` to run on small runners. --- .github/workflows/build_and_test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 2d7edf2e22..9fe9636d67 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -132,7 +132,7 @@ jobs: check-codestyle-rust: needs: [ check-permissions, build-buildtools-image ] - runs-on: [ self-hosted, gen3, large ] + runs-on: [ self-hosted, gen3, small ] container: image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/build-tools:${{ needs.build-buildtools-image.outputs.build-tools-tag }} options: --init From e65f0fe874aa4762d5d4702349647677ea2c352e Mon Sep 17 00:00:00 2001 From: Alexander Bayandin Date: Tue, 6 Feb 2024 17:00:55 +0000 Subject: [PATCH 22/72] CI(benchmarks): make job split consistent across reruns (#6614) ## Problem We've got several issues with the current `benchmarks` job setup: - `benchmark_durations.json` file (that we generate in runtime to split tests into several jobs[0]) is not consistent between these jobs (and very not consistent with the file if we rerun the job). I.e. test selection for each job can be different, which could end up in missed tests in a test run. - `scripts/benchmark_durations` doesn't fetch all tests from the database (it doesn't expect any extra directories inside `test_runner/performance`) - For some reason, currently split into 4 groups ends up with the 4th group has no tests to run, which fails the job[1] - [0] https://github.com/neondatabase/neon/pull/4683 - [1] https://github.com/neondatabase/neon/issues/6629 ## Summary of changes - Generate `benchmark_durations.json` file once before we start `benchmarks` jobs (this makes it consistent across the jobs) and pass the file content through the GitHub Actions input (this makes it consistent for reruns) - `scripts/benchmark_durations` fix SQL query for getting all required tests - Split benchmarks into 5 jobs instead of 4 jobs. --- .../actions/run-python-test-set/action.yml | 6 +- .github/workflows/build_and_test.yml | 39 ++++- scripts/benchmark_durations.py | 133 +++++++++--------- 3 files changed, 111 insertions(+), 67 deletions(-) diff --git a/.github/actions/run-python-test-set/action.yml b/.github/actions/run-python-test-set/action.yml index 8dfa6c465f..7a88e4f73b 100644 --- a/.github/actions/run-python-test-set/action.yml +++ b/.github/actions/run-python-test-set/action.yml @@ -44,6 +44,10 @@ inputs: description: 'Postgres version to use for tests' required: false default: 'v14' + benchmark_durations: + description: 'benchmark durations JSON' + required: false + default: '{}' runs: using: "composite" @@ -160,7 +164,7 @@ runs: # We use pytest-split plugin to run benchmarks in parallel on different CI runners if [ "${TEST_SELECTION}" = "test_runner/performance" ] && [ "${{ inputs.build_type }}" != "remote" ]; then mkdir -p $TEST_OUTPUT - poetry run ./scripts/benchmark_durations.py "${TEST_RESULT_CONNSTR}" --days 10 --output "$TEST_OUTPUT/benchmark_durations.json" + echo '${{ inputs.benchmark_durations || '{}' }}' > $TEST_OUTPUT/benchmark_durations.json EXTRA_PARAMS="--durations-path $TEST_OUTPUT/benchmark_durations.json $EXTRA_PARAMS" fi diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 9fe9636d67..066f4a21eb 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -478,8 +478,40 @@ jobs: if: matrix.build_type == 'debug' && matrix.pg_version == 'v14' uses: ./.github/actions/save-coverage-data + get-benchmarks-durations: + outputs: + json: ${{ steps.get-benchmark-durations.outputs.json }} + needs: [ check-permissions, build-buildtools-image ] + runs-on: [ self-hosted, gen3, small ] + container: + image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/build-tools:${{ needs.build-buildtools-image.outputs.build-tools-tag }} + options: --init + if: github.ref_name == 'main' || contains(github.event.pull_request.labels.*.name, 'run-benchmarks') + steps: + - name: Checkout + uses: actions/checkout@v3 + + - name: Cache poetry deps + uses: actions/cache@v3 + with: + path: ~/.cache/pypoetry/virtualenvs + key: v1-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }} + + - name: Install Python deps + run: ./scripts/pysync + + - name: get benchmark durations + id: get-benchmark-durations + env: + TEST_RESULT_CONNSTR: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }} + run: | + poetry run ./scripts/benchmark_durations.py "${TEST_RESULT_CONNSTR}" \ + --days 10 \ + --output /tmp/benchmark_durations.json + echo "json=$(jq --compact-output '.' /tmp/benchmark_durations.json)" >> $GITHUB_OUTPUT + benchmarks: - needs: [ check-permissions, build-neon, build-buildtools-image ] + needs: [ check-permissions, build-neon, build-buildtools-image, get-benchmarks-durations ] runs-on: [ self-hosted, gen3, small ] container: image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/build-tools:${{ needs.build-buildtools-image.outputs.build-tools-tag }} @@ -490,7 +522,7 @@ jobs: fail-fast: false matrix: # the amount of groups (N) should be reflected in `extra_params: --splits N ...` - pytest_split_group: [ 1, 2, 3, 4 ] + pytest_split_group: [ 1, 2, 3, 4, 5 ] build_type: [ release ] steps: - name: Checkout @@ -503,7 +535,8 @@ jobs: test_selection: performance run_in_parallel: false save_perf_report: ${{ github.ref_name == 'main' }} - extra_params: --splits 4 --group ${{ matrix.pytest_split_group }} + extra_params: --splits 5 --group ${{ matrix.pytest_split_group }} + benchmark_durations: ${{ needs.get-benchmarks-durations.outputs.json }} env: VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}" PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}" diff --git a/scripts/benchmark_durations.py b/scripts/benchmark_durations.py index 7f05d72a03..01f34a1b96 100755 --- a/scripts/benchmark_durations.py +++ b/scripts/benchmark_durations.py @@ -20,7 +20,7 @@ BENCHMARKS_DURATION_QUERY = """ FROM results WHERE started_at > CURRENT_DATE - INTERVAL '%s' day - AND parent_suite = 'test_runner.performance' + AND starts_with(parent_suite, 'test_runner.performance') AND status = 'passed' GROUP BY parent_suite, suite, name @@ -31,68 +31,75 @@ BENCHMARKS_DURATION_QUERY = """ # the total duration varies from 8 to 40 minutes. # We use some pre-collected durations as a fallback to have a better distribution. FALLBACK_DURATION = { - "test_runner/performance/test_branch_creation.py::test_branch_creation_heavy_write[20]": 62.144, - "test_runner/performance/test_branch_creation.py::test_branch_creation_many[1024]": 90.941, - "test_runner/performance/test_branch_creation.py::test_branch_creation_many_relations": 26.053, - "test_runner/performance/test_branching.py::test_compare_child_and_root_pgbench_perf": 25.67, - "test_runner/performance/test_branching.py::test_compare_child_and_root_read_perf": 14.497, - "test_runner/performance/test_branching.py::test_compare_child_and_root_write_perf": 18.852, - "test_runner/performance/test_bulk_insert.py::test_bulk_insert[neon]": 26.572, - "test_runner/performance/test_bulk_insert.py::test_bulk_insert[vanilla]": 6.259, - "test_runner/performance/test_bulk_tenant_create.py::test_bulk_tenant_create[10]": 21.206, - "test_runner/performance/test_bulk_tenant_create.py::test_bulk_tenant_create[1]": 3.474, - "test_runner/performance/test_bulk_tenant_create.py::test_bulk_tenant_create[5]": 11.262, - "test_runner/performance/test_bulk_update.py::test_bulk_update[100]": 94.225, - "test_runner/performance/test_bulk_update.py::test_bulk_update[10]": 68.159, - "test_runner/performance/test_bulk_update.py::test_bulk_update[50]": 76.719, - "test_runner/performance/test_compaction.py::test_compaction": 110.222, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_ro_with_pgbench_select_only[neon-5-10-100]": 10.743, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_ro_with_pgbench_select_only[vanilla-5-10-100]": 16.541, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_rw_with_pgbench_default[neon-5-10-100]": 11.109, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_rw_with_pgbench_default[vanilla-5-10-100]": 18.121, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wal_with_pgbench_default[neon-5-10-100]": 11.3, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wal_with_pgbench_default[vanilla-5-10-100]": 16.086, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[neon-10-10]": 12.024, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[neon-10-1]": 11.14, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[vanilla-10-10]": 10.375, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[vanilla-10-1]": 10.075, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_pgbench_simple_update[neon-5-10-100]": 11.147, - "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_pgbench_simple_update[vanilla-5-10-100]": 16.321, - "test_runner/performance/test_copy.py::test_copy[neon]": 16.579, - "test_runner/performance/test_copy.py::test_copy[vanilla]": 10.094, - "test_runner/performance/test_gc_feedback.py::test_gc_feedback": 590.157, - "test_runner/performance/test_gist_build.py::test_gist_buffering_build[neon]": 14.102, - "test_runner/performance/test_gist_build.py::test_gist_buffering_build[vanilla]": 8.677, - "test_runner/performance/test_latency.py::test_measure_read_latency_heavy_write_workload[neon-1]": 31.079, - "test_runner/performance/test_latency.py::test_measure_read_latency_heavy_write_workload[vanilla-1]": 38.119, - "test_runner/performance/test_layer_map.py::test_layer_map": 24.784, - "test_runner/performance/test_logical_replication.py::test_logical_replication": 117.707, - "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_different_tables[neon]": 21.194, - "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_different_tables[vanilla]": 59.068, - "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_same_table[neon]": 73.235, - "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_same_table[vanilla]": 82.586, - "test_runner/performance/test_perf_pgbench.py::test_pgbench[neon-45-10]": 106.536, - "test_runner/performance/test_perf_pgbench.py::test_pgbench[vanilla-45-10]": 98.753, - "test_runner/performance/test_random_writes.py::test_random_writes[neon]": 6.975, - "test_runner/performance/test_random_writes.py::test_random_writes[vanilla]": 3.69, - "test_runner/performance/test_seqscans.py::test_seqscans[neon-100000-100-0]": 3.529, - "test_runner/performance/test_seqscans.py::test_seqscans[neon-10000000-1-0]": 64.522, - "test_runner/performance/test_seqscans.py::test_seqscans[neon-10000000-1-4]": 40.964, - "test_runner/performance/test_seqscans.py::test_seqscans[vanilla-100000-100-0]": 0.55, - "test_runner/performance/test_seqscans.py::test_seqscans[vanilla-10000000-1-0]": 12.189, - "test_runner/performance/test_seqscans.py::test_seqscans[vanilla-10000000-1-4]": 13.899, - "test_runner/performance/test_startup.py::test_startup_simple": 2.51, - "test_runner/performance/test_wal_backpressure.py::test_heavy_write_workload[neon_off-10-5-5]": 527.245, - "test_runner/performance/test_wal_backpressure.py::test_heavy_write_workload[neon_on-10-5-5]": 583.46, - "test_runner/performance/test_wal_backpressure.py::test_heavy_write_workload[vanilla-10-5-5]": 113.653, - "test_runner/performance/test_wal_backpressure.py::test_pgbench_intensive_init_workload[neon_off-1000]": 233.728, - "test_runner/performance/test_wal_backpressure.py::test_pgbench_intensive_init_workload[neon_on-1000]": 419.093, - "test_runner/performance/test_wal_backpressure.py::test_pgbench_intensive_init_workload[vanilla-1000]": 982.461, - "test_runner/performance/test_wal_backpressure.py::test_pgbench_simple_update_workload[neon_off-45-100]": 116.522, - "test_runner/performance/test_wal_backpressure.py::test_pgbench_simple_update_workload[neon_on-45-100]": 115.583, - "test_runner/performance/test_wal_backpressure.py::test_pgbench_simple_update_workload[vanilla-45-100]": 155.282, - "test_runner/performance/test_write_amplification.py::test_write_amplification[neon]": 26.704, - "test_runner/performance/test_write_amplification.py::test_write_amplification[vanilla]": 16.088, + "test_runner/performance/pageserver/pagebench/test_pageserver_max_throughput_getpage_at_latest_lsn.py::test_pageserver_max_throughput_getpage_at_latest_lsn[1-13-30]": 400.15, + "test_runner/performance/pageserver/pagebench/test_pageserver_max_throughput_getpage_at_latest_lsn.py::test_pageserver_max_throughput_getpage_at_latest_lsn[1-6-30]": 372.521, + "test_runner/performance/pageserver/pagebench/test_pageserver_max_throughput_getpage_at_latest_lsn.py::test_pageserver_max_throughput_getpage_at_latest_lsn[10-13-30]": 420.017, + "test_runner/performance/pageserver/pagebench/test_pageserver_max_throughput_getpage_at_latest_lsn.py::test_pageserver_max_throughput_getpage_at_latest_lsn[10-6-30]": 373.769, + "test_runner/performance/pageserver/pagebench/test_pageserver_max_throughput_getpage_at_latest_lsn.py::test_pageserver_max_throughput_getpage_at_latest_lsn[100-13-30]": 678.742, + "test_runner/performance/pageserver/pagebench/test_pageserver_max_throughput_getpage_at_latest_lsn.py::test_pageserver_max_throughput_getpage_at_latest_lsn[100-6-30]": 512.135, + "test_runner/performance/test_branch_creation.py::test_branch_creation_heavy_write[20]": 58.036, + "test_runner/performance/test_branch_creation.py::test_branch_creation_many_relations": 22.104, + "test_runner/performance/test_branch_creation.py::test_branch_creation_many[1024]": 126.073, + "test_runner/performance/test_branching.py::test_compare_child_and_root_pgbench_perf": 25.759, + "test_runner/performance/test_branching.py::test_compare_child_and_root_read_perf": 6.885, + "test_runner/performance/test_branching.py::test_compare_child_and_root_write_perf": 8.758, + "test_runner/performance/test_bulk_insert.py::test_bulk_insert[neon]": 18.275, + "test_runner/performance/test_bulk_insert.py::test_bulk_insert[vanilla]": 9.533, + "test_runner/performance/test_bulk_tenant_create.py::test_bulk_tenant_create[1]": 12.09, + "test_runner/performance/test_bulk_tenant_create.py::test_bulk_tenant_create[10]": 35.145, + "test_runner/performance/test_bulk_tenant_create.py::test_bulk_tenant_create[5]": 22.28, + "test_runner/performance/test_bulk_update.py::test_bulk_update[10]": 66.353, + "test_runner/performance/test_bulk_update.py::test_bulk_update[100]": 75.487, + "test_runner/performance/test_bulk_update.py::test_bulk_update[50]": 54.142, + "test_runner/performance/test_compaction.py::test_compaction": 110.715, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_ro_with_pgbench_select_only[neon-5-10-100]": 11.68, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_ro_with_pgbench_select_only[vanilla-5-10-100]": 16.384, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_rw_with_pgbench_default[neon-5-10-100]": 11.315, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_rw_with_pgbench_default[vanilla-5-10-100]": 18.783, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wal_with_pgbench_default[neon-5-10-100]": 11.647, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wal_with_pgbench_default[vanilla-5-10-100]": 17.04, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[neon-10-1]": 11.01, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[neon-10-10]": 11.902, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[vanilla-10-1]": 10.077, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_heavy_write[vanilla-10-10]": 10.4, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_pgbench_simple_update[neon-5-10-100]": 11.33, + "test_runner/performance/test_compare_pg_stats.py::test_compare_pg_stats_wo_with_pgbench_simple_update[vanilla-5-10-100]": 16.434, + "test_runner/performance/test_copy.py::test_copy[neon]": 13.817, + "test_runner/performance/test_copy.py::test_copy[vanilla]": 11.736, + "test_runner/performance/test_gc_feedback.py::test_gc_feedback": 575.735, + "test_runner/performance/test_gist_build.py::test_gist_buffering_build[neon]": 14.868, + "test_runner/performance/test_gist_build.py::test_gist_buffering_build[vanilla]": 14.393, + "test_runner/performance/test_latency.py::test_measure_read_latency_heavy_write_workload[neon-1]": 20.588, + "test_runner/performance/test_latency.py::test_measure_read_latency_heavy_write_workload[vanilla-1]": 30.849, + "test_runner/performance/test_layer_map.py::test_layer_map": 39.378, + "test_runner/performance/test_lazy_startup.py::test_lazy_startup": 2848.938, + "test_runner/performance/test_logical_replication.py::test_logical_replication": 120.952, + "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_different_tables[neon]": 35.552, + "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_different_tables[vanilla]": 66.762, + "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_same_table[neon]": 85.177, + "test_runner/performance/test_parallel_copy_to.py::test_parallel_copy_same_table[vanilla]": 92.12, + "test_runner/performance/test_perf_pgbench.py::test_pgbench[neon-45-10]": 107.009, + "test_runner/performance/test_perf_pgbench.py::test_pgbench[vanilla-45-10]": 99.582, + "test_runner/performance/test_random_writes.py::test_random_writes[neon]": 4.737, + "test_runner/performance/test_random_writes.py::test_random_writes[vanilla]": 2.686, + "test_runner/performance/test_seqscans.py::test_seqscans[neon-100000-100-0]": 3.271, + "test_runner/performance/test_seqscans.py::test_seqscans[neon-10000000-1-0]": 50.719, + "test_runner/performance/test_seqscans.py::test_seqscans[neon-10000000-1-4]": 15.992, + "test_runner/performance/test_seqscans.py::test_seqscans[vanilla-100000-100-0]": 0.566, + "test_runner/performance/test_seqscans.py::test_seqscans[vanilla-10000000-1-0]": 13.542, + "test_runner/performance/test_seqscans.py::test_seqscans[vanilla-10000000-1-4]": 13.35, + "test_runner/performance/test_startup.py::test_startup_simple": 13.043, + "test_runner/performance/test_wal_backpressure.py::test_heavy_write_workload[neon_off-10-5-5]": 194.841, + "test_runner/performance/test_wal_backpressure.py::test_heavy_write_workload[neon_on-10-5-5]": 286.667, + "test_runner/performance/test_wal_backpressure.py::test_heavy_write_workload[vanilla-10-5-5]": 85.577, + "test_runner/performance/test_wal_backpressure.py::test_pgbench_intensive_init_workload[neon_off-1000]": 297.626, + "test_runner/performance/test_wal_backpressure.py::test_pgbench_intensive_init_workload[neon_on-1000]": 646.187, + "test_runner/performance/test_wal_backpressure.py::test_pgbench_intensive_init_workload[vanilla-1000]": 989.776, + "test_runner/performance/test_wal_backpressure.py::test_pgbench_simple_update_workload[neon_off-45-100]": 125.638, + "test_runner/performance/test_wal_backpressure.py::test_pgbench_simple_update_workload[neon_on-45-100]": 123.554, + "test_runner/performance/test_wal_backpressure.py::test_pgbench_simple_update_workload[vanilla-45-100]": 190.083, + "test_runner/performance/test_write_amplification.py::test_write_amplification[neon]": 21.016, + "test_runner/performance/test_write_amplification.py::test_write_amplification[vanilla]": 23.028, } From dc811d19231273ff9ce3e235d34c45c0fd0d443a Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Tue, 6 Feb 2024 20:37:35 +0200 Subject: [PATCH 23/72] Add a span to 'create_neon_superuser' for better OpenTelemetry traces (#6644) create_neon_superuser runs the first queries in the database after cold start. Traces suggest that those first queries can make up a significant fraction of the cold start time. Make it more visible by adding an explict tracing span to it; currently you just have to deduce it by looking at the time spent in the parent 'apply_config' span subtracted by all the other child spans. --- compute_tools/src/compute.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 1976299e93..098e06cca9 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -207,6 +207,7 @@ fn maybe_cgexec(cmd: &str) -> Command { /// Create special neon_superuser role, that's a slightly nerfed version of a real superuser /// that we give to customers +#[instrument(skip_all)] fn create_neon_superuser(spec: &ComputeSpec, client: &mut Client) -> Result<()> { let roles = spec .cluster From 4f57dc6cc6ac69d9d342b8eb566237907dcff85b Mon Sep 17 00:00:00 2001 From: John Spray Date: Tue, 6 Feb 2024 19:08:39 +0000 Subject: [PATCH 24/72] control_plane/attachment_service: take public key as value (#6651) It's awkward to point to a file when doing some kinds of ad-hoc deployment (like right now, when I'm hacking a helm chart having not quite hooked up secrets properly yet). We take all the rest of the secrets as CLI args directly, so let's do the same for public key. --- control_plane/attachment_service/src/main.rs | 6 ++-- control_plane/src/attachment_service.rs | 31 ++++++++++++++++---- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/control_plane/attachment_service/src/main.rs b/control_plane/attachment_service/src/main.rs index 7ac5918244..bc8a8786c2 100644 --- a/control_plane/attachment_service/src/main.rs +++ b/control_plane/attachment_service/src/main.rs @@ -34,9 +34,9 @@ struct Cli { #[arg(short, long)] listen: std::net::SocketAddr, - /// Path to public key for JWT authentication of clients + /// Public key for JWT authentication of clients #[arg(long)] - public_key: Option, + public_key: Option, /// Token for authenticating this service with the pageservers it controls #[arg(long)] @@ -159,7 +159,7 @@ impl Secrets { fn load_cli(database_url: &str, args: &Cli) -> anyhow::Result { let public_key = match &args.public_key { None => None, - Some(key_path) => Some(JwtAuth::from_key_path(key_path)?), + Some(key) => Some(JwtAuth::from_key(key.clone()).context("Loading public key")?), }; Ok(Self { database_url: database_url.to_owned(), diff --git a/control_plane/src/attachment_service.rs b/control_plane/src/attachment_service.rs index 140e5c4e34..a3f832036c 100644 --- a/control_plane/src/attachment_service.rs +++ b/control_plane/src/attachment_service.rs @@ -28,7 +28,7 @@ pub struct AttachmentService { listen: String, path: Utf8PathBuf, jwt_token: Option, - public_key_path: Option, + public_key: Option, postgres_port: u16, client: reqwest::Client, } @@ -207,7 +207,7 @@ impl AttachmentService { .pageservers .first() .expect("Config is validated to contain at least one pageserver"); - let (jwt_token, public_key_path) = match ps_conf.http_auth_type { + let (jwt_token, public_key) = match ps_conf.http_auth_type { AuthType::Trust => (None, None), AuthType::NeonJWT => { let jwt_token = env @@ -219,7 +219,26 @@ impl AttachmentService { let public_key_path = camino::Utf8PathBuf::try_from(env.base_data_dir.join("auth_public_key.pem")) .unwrap(); - (Some(jwt_token), Some(public_key_path)) + + // This service takes keys as a string rather than as a path to a file/dir: read the key into memory. + let public_key = if std::fs::metadata(&public_key_path) + .expect("Can't stat public key") + .is_dir() + { + // Our config may specify a directory: this is for the pageserver's ability to handle multiple + // keys. We only use one key at a time, so, arbitrarily load the first one in the directory. + let mut dir = + std::fs::read_dir(&public_key_path).expect("Can't readdir public key path"); + let dent = dir + .next() + .expect("Empty key dir") + .expect("Error reading key dir"); + + std::fs::read_to_string(dent.path()).expect("Can't read public key") + } else { + std::fs::read_to_string(&public_key_path).expect("Can't read public key") + }; + (Some(jwt_token), Some(public_key)) } }; @@ -228,7 +247,7 @@ impl AttachmentService { path, listen, jwt_token, - public_key_path, + public_key, postgres_port, client: reqwest::ClientBuilder::new() .build() @@ -453,8 +472,8 @@ impl AttachmentService { args.push(format!("--jwt-token={jwt_token}")); } - if let Some(public_key_path) = &self.public_key_path { - args.push(format!("--public-key={public_key_path}")); + if let Some(public_key) = &self.public_key { + args.push(format!("--public-key=\"{public_key}\"")); } if let Some(control_plane_compute_hook_api) = &self.env.control_plane_compute_hook_api { From f4cc7cae1412c14e49a795dc6a8d0ca21413affd Mon Sep 17 00:00:00 2001 From: Alexander Bayandin Date: Tue, 6 Feb 2024 20:30:43 +0000 Subject: [PATCH 25/72] CI(build-tools): Update Python from 3.9.2 to 3.9.18 (#6615) ## Problem We use an outdated version of Python (3.9.2) ## Summary of changes - Update Python to the latest patch version (3.9.18) - Unify the usage of python caches where possible --- .github/actions/allure-report-generate/action.yml | 6 ++++++ .github/actions/run-python-test-set/action.yml | 3 +-- .github/workflows/build_and_test.yml | 3 +-- .github/workflows/pg_clients.yml | 3 +-- Dockerfile.buildtools | 2 +- 5 files changed, 10 insertions(+), 7 deletions(-) diff --git a/.github/actions/allure-report-generate/action.yml b/.github/actions/allure-report-generate/action.yml index a33adf8bdd..f474dd3444 100644 --- a/.github/actions/allure-report-generate/action.yml +++ b/.github/actions/allure-report-generate/action.yml @@ -179,6 +179,12 @@ runs: aws s3 rm "s3://${BUCKET}/${LOCK_FILE}" fi + - name: Cache poetry deps + uses: actions/cache@v3 + with: + path: ~/.cache/pypoetry/virtualenvs + key: v2-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }} + - name: Store Allure test stat in the DB (new) if: ${{ !cancelled() && inputs.store-test-results-into-db == 'true' }} shell: bash -euxo pipefail {0} diff --git a/.github/actions/run-python-test-set/action.yml b/.github/actions/run-python-test-set/action.yml index 7a88e4f73b..8852a28da9 100644 --- a/.github/actions/run-python-test-set/action.yml +++ b/.github/actions/run-python-test-set/action.yml @@ -86,11 +86,10 @@ runs: fetch-depth: 1 - name: Cache poetry deps - id: cache_poetry uses: actions/cache@v3 with: path: ~/.cache/pypoetry/virtualenvs - key: v1-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }} + key: v2-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }} - name: Install Python deps shell: bash -euxo pipefail {0} diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 066f4a21eb..f12f020634 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -112,11 +112,10 @@ jobs: fetch-depth: 1 - name: Cache poetry deps - id: cache_poetry uses: actions/cache@v3 with: path: ~/.cache/pypoetry/virtualenvs - key: v1-codestyle-python-deps-${{ hashFiles('poetry.lock') }} + key: v2-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }} - name: Install Python deps run: ./scripts/pysync diff --git a/.github/workflows/pg_clients.yml b/.github/workflows/pg_clients.yml index 224b7b4a6d..28016cadb1 100644 --- a/.github/workflows/pg_clients.yml +++ b/.github/workflows/pg_clients.yml @@ -38,11 +38,10 @@ jobs: uses: snok/install-poetry@v1 - name: Cache poetry deps - id: cache_poetry uses: actions/cache@v3 with: path: ~/.cache/pypoetry/virtualenvs - key: v1-${{ runner.os }}-python-deps-${{ hashFiles('poetry.lock') }} + key: v2-${{ runner.os }}-python-deps-ubunutu-latest-${{ hashFiles('poetry.lock') }} - name: Install Python deps shell: bash -euxo pipefail {0} diff --git a/Dockerfile.buildtools b/Dockerfile.buildtools index 213aed1679..220e995d64 100644 --- a/Dockerfile.buildtools +++ b/Dockerfile.buildtools @@ -111,7 +111,7 @@ USER nonroot:nonroot WORKDIR /home/nonroot # Python -ENV PYTHON_VERSION=3.9.2 \ +ENV PYTHON_VERSION=3.9.18 \ PYENV_ROOT=/home/nonroot/.pyenv \ PATH=/home/nonroot/.pyenv/shims:/home/nonroot/.pyenv/bin:/home/nonroot/.poetry/bin:$PATH RUN set -e \ From 9f75da7c0ac483e612b7382b0b050588c5587584 Mon Sep 17 00:00:00 2001 From: Alexander Bayandin Date: Wed, 7 Feb 2024 00:31:26 +0000 Subject: [PATCH 26/72] test_lazy_startup: fix statement_timeout setting (#6654) ## Problem Test `test_lazy_startup` is flaky[0], sometimes (pretty frequently) it fails with `canceling statement due to statement timeout`. - [0] https://neon-github-public-dev.s3.amazonaws.com/reports/main/7803316870/index.html#suites/355b1a7a5b1e740b23ea53728913b4fa/7263782d30986c50/history ## Summary of changes - Fix setting `statement_timeout` setting by reusing a connection for all queries. - Also fix label (`lazy`, `eager`) assignment - Split `test_lazy_startup` into two, by `slru` laziness and make tests smaller --- test_runner/performance/test_lazy_startup.py | 143 +++++++++---------- 1 file changed, 69 insertions(+), 74 deletions(-) diff --git a/test_runner/performance/test_lazy_startup.py b/test_runner/performance/test_lazy_startup.py index 1a431e272e..e929bd4d05 100644 --- a/test_runner/performance/test_lazy_startup.py +++ b/test_runner/performance/test_lazy_startup.py @@ -26,86 +26,81 @@ from fixtures.neon_fixtures import NeonEnvBuilder # apply during config step, like more users, databases, or extensions. By default # we load extensions 'neon,pg_stat_statements,timescaledb,pg_cron', but in this # test we only load neon. -@pytest.mark.timeout(1000) -def test_lazy_startup(neon_env_builder: NeonEnvBuilder, zenbenchmark: NeonBenchmarker): +@pytest.mark.timeout(1800) +@pytest.mark.parametrize("slru", ["lazy", "eager"]) +def test_lazy_startup(slru: str, neon_env_builder: NeonEnvBuilder, zenbenchmark: NeonBenchmarker): neon_env_builder.num_safekeepers = 3 env = neon_env_builder.init_start() - lazy_tenant, _ = env.neon_cli.create_tenant( + lazy_slru_download = "true" if slru == "lazy" else "false" + tenant, _ = env.neon_cli.create_tenant( conf={ - "lazy_slru_download": "true", + "lazy_slru_download": lazy_slru_download, } ) - eager_tenant, _ = env.neon_cli.create_tenant( - conf={ - "lazy_slru_download": "false", - } - ) - tenants = [lazy_tenant, eager_tenant] - slru = "lazy" - for tenant in tenants: - endpoint = env.endpoints.create_start("main", tenant_id=tenant) - endpoint.safe_psql("CREATE TABLE t (pk integer PRIMARY KEY, x integer)") - endpoint.safe_psql("ALTER TABLE t SET (autovacuum_enabled = false)") - endpoint.safe_psql("INSERT INTO t VALUES (1, 0)") - endpoint.safe_psql( - """ - CREATE PROCEDURE updating() as - $$ - DECLARE - i integer; - BEGIN - FOR i IN 1..10000000 LOOP - UPDATE t SET x = x + 1 WHERE pk=1; - COMMIT; - END LOOP; - END - $$ LANGUAGE plpgsql - """ - ) - endpoint.safe_psql("SET statement_timeout=0") - endpoint.safe_psql("call updating()") + endpoint = env.endpoints.create_start("main", tenant_id=tenant) + with endpoint.cursor() as cur: + cur.execute("CREATE TABLE t (pk integer PRIMARY KEY, x integer)") + cur.execute("ALTER TABLE t SET (autovacuum_enabled = false)") + cur.execute("INSERT INTO t VALUES (1, 0)") + cur.execute( + """ + CREATE PROCEDURE updating() as + $$ + DECLARE + i integer; + BEGIN + FOR i IN 1..1000000 LOOP + UPDATE t SET x = x + 1 WHERE pk=1; + COMMIT; + END LOOP; + END + $$ LANGUAGE plpgsql + """ + ) + cur.execute("SET statement_timeout=0") + cur.execute("call updating()") + + endpoint.stop() + + # We do two iterations so we can see if the second startup is faster. It should + # be because the compute node should already be configured with roles, databases, + # extensions, etc from the first run. + for i in range(2): + # Start + with zenbenchmark.record_duration(f"{slru}_{i}_start"): + endpoint.start() + + with zenbenchmark.record_duration(f"{slru}_{i}_select"): + sum = endpoint.safe_psql("select sum(x) from t")[0][0] + assert sum == 1000000 + + # Get metrics + metrics = requests.get(f"http://localhost:{endpoint.http_port}/metrics.json").json() + durations = { + "wait_for_spec_ms": f"{slru}_{i}_wait_for_spec", + "sync_safekeepers_ms": f"{slru}_{i}_sync_safekeepers", + "sync_sk_check_ms": f"{slru}_{i}_sync_sk_check", + "basebackup_ms": f"{slru}_{i}_basebackup", + "start_postgres_ms": f"{slru}_{i}_start_postgres", + "config_ms": f"{slru}_{i}_config", + "total_startup_ms": f"{slru}_{i}_total_startup", + } + for key, name in durations.items(): + value = metrics[key] + zenbenchmark.record(name, value, "ms", report=MetricReport.LOWER_IS_BETTER) + + basebackup_bytes = metrics["basebackup_bytes"] + zenbenchmark.record( + f"{slru}_{i}_basebackup_bytes", + basebackup_bytes, + "bytes", + report=MetricReport.LOWER_IS_BETTER, + ) + + # Stop so we can restart endpoint.stop() - # We do two iterations so we can see if the second startup is faster. It should - # be because the compute node should already be configured with roles, databases, - # extensions, etc from the first run. - for i in range(2): - # Start - with zenbenchmark.record_duration(f"{slru}_{i}_start"): - endpoint.start() - - with zenbenchmark.record_duration(f"{slru}_{i}_select"): - sum = endpoint.safe_psql("select sum(x) from t")[0][0] - assert sum == 10000000 - - # Get metrics - metrics = requests.get(f"http://localhost:{endpoint.http_port}/metrics.json").json() - durations = { - "wait_for_spec_ms": f"{slru}_{i}_wait_for_spec", - "sync_safekeepers_ms": f"{slru}_{i}_sync_safekeepers", - "sync_sk_check_ms": f"{slru}_{i}_sync_sk_check", - "basebackup_ms": f"{slru}_{i}_basebackup", - "start_postgres_ms": f"{slru}_{i}_start_postgres", - "config_ms": f"{slru}_{i}_config", - "total_startup_ms": f"{slru}_{i}_total_startup", - } - for key, name in durations.items(): - value = metrics[key] - zenbenchmark.record(name, value, "ms", report=MetricReport.LOWER_IS_BETTER) - - basebackup_bytes = metrics["basebackup_bytes"] - zenbenchmark.record( - f"{slru}_{i}_basebackup_bytes", - basebackup_bytes, - "bytes", - report=MetricReport.LOWER_IS_BETTER, - ) - - # Stop so we can restart - endpoint.stop() - - # Imitate optimizations that console would do for the second start - endpoint.respec(skip_pg_catalog_updates=True) - slru = "eager" + # Imitate optimizations that console would do for the second start + endpoint.respec(skip_pg_catalog_updates=True) From f3d7d2380566948d5bf7250c32c1e11ef5099ab3 Mon Sep 17 00:00:00 2001 From: Konstantin Knizhnik Date: Wed, 7 Feb 2024 08:47:19 +0200 Subject: [PATCH 27/72] Some small WAL records can write a lot of data to KV storage, so perform checkpoint check more frequently (#6639) ## Problem See https://neondb.slack.com/archives/C04DGM6SMTM/p1707149618314539?thread_ts=1707081520.140049&cid=C04DGM6SMTM ## Summary of changes Perform checkpoint check after processing `ingest_batch_size` (default 100) WAL records. ## Checklist before requesting a review - [ ] I have performed a self-review of my code. - [ ] If it is a core feature, I have added thorough tests. - [ ] Do we need to implement analytics? if so did you add the relevant metrics to the dashboard? - [ ] If this PR requires public announcement, mark it with /release-notes label and add several sentences in this section. ## Checklist before merging - [ ] Do not forget to reformat commit message to not include the above checklist --------- Co-authored-by: Konstantin Knizhnik --- .../walreceiver/walreceiver_connection.rs | 17 +++++ test_runner/regress/test_layer_bloating.py | 66 +++++++++++++++++++ 2 files changed, 83 insertions(+) create mode 100644 test_runner/regress/test_layer_bloating.py diff --git a/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs b/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs index 73eb42bb30..9cb53f46d1 100644 --- a/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs +++ b/pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs @@ -343,6 +343,23 @@ pub(super) async fn handle_walreceiver_connection( modification.commit(&ctx).await?; uncommitted_records = 0; filtered_records = 0; + + // + // We should check checkpoint distance after appending each ingest_batch_size bytes because otherwise + // layer size can become much larger than `checkpoint_distance`. + // It can append because wal-sender is sending WAL using 125kb chucks and some WAL records can cause writing large + // amount of data to key-value storage. So performing this check only after processing + // all WAL records in the chunk, can cause huge L0 layer files. + // + timeline + .check_checkpoint_distance() + .await + .with_context(|| { + format!( + "Failed to check checkpoint distance for timeline {}", + timeline.timeline_id + ) + })?; } } diff --git a/test_runner/regress/test_layer_bloating.py b/test_runner/regress/test_layer_bloating.py new file mode 100644 index 0000000000..70b115ad61 --- /dev/null +++ b/test_runner/regress/test_layer_bloating.py @@ -0,0 +1,66 @@ +import os +import time + +import pytest +from fixtures.log_helper import log +from fixtures.neon_fixtures import ( + NeonEnv, + logical_replication_sync, +) +from fixtures.pg_version import PgVersion + + +def test_layer_bloating(neon_simple_env: NeonEnv, vanilla_pg): + env = neon_simple_env + + if env.pg_version != PgVersion.V16: + pytest.skip("pg_log_standby_snapshot() function is available only in PG16") + + timeline = env.neon_cli.create_branch("test_logical_replication", "empty") + endpoint = env.endpoints.create_start( + "test_logical_replication", config_lines=["log_statement=all"] + ) + + log.info("postgres is running on 'test_logical_replication' branch") + pg_conn = endpoint.connect() + cur = pg_conn.cursor() + + # create table... + cur.execute("create table t(pk integer primary key)") + cur.execute("create publication pub1 for table t") + # Create slot to hold WAL + cur.execute("select pg_create_logical_replication_slot('my_slot', 'pgoutput')") + + # now start subscriber + vanilla_pg.start() + vanilla_pg.safe_psql("create table t(pk integer primary key)") + + connstr = endpoint.connstr().replace("'", "''") + log.info(f"ep connstr is {endpoint.connstr()}, subscriber connstr {vanilla_pg.connstr()}") + vanilla_pg.safe_psql(f"create subscription sub1 connection '{connstr}' publication pub1") + + cur.execute( + """create or replace function create_snapshots(n integer) returns void as $$ + declare + i integer; + begin + for i in 1..n loop + perform pg_log_standby_snapshot(); + end loop; + end; $$ language plpgsql""" + ) + cur.execute("set statement_timeout=0") + cur.execute("select create_snapshots(10000)") + # Wait logical replication to sync + logical_replication_sync(vanilla_pg, endpoint) + time.sleep(10) + + # Check layer file sizes + timeline_path = "{}/tenants/{}/timelines/{}/".format( + env.pageserver.workdir, env.initial_tenant, timeline + ) + log.info(f"Check {timeline_path}") + for filename in os.listdir(timeline_path): + if filename.startswith("00000"): + log.info(f"layer {filename} size is {os.path.getsize(timeline_path + filename)}") + assert os.path.getsize(timeline_path + filename) < 512_000_000 From f7516df6c155162aa2d935adadf95524379e0a58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Wed, 7 Feb 2024 12:56:53 +0100 Subject: [PATCH 28/72] Pass timestamp as a datetime (#6656) This saves some repetition. I did this in #6533 for `tenant_time_travel_remote_storage` already. --- test_runner/fixtures/pageserver/http.py | 4 ++-- test_runner/regress/test_lsn_mapping.py | 16 ++++------------ 2 files changed, 6 insertions(+), 14 deletions(-) diff --git a/test_runner/fixtures/pageserver/http.py b/test_runner/fixtures/pageserver/http.py index 92e5027a9f..adea9ca764 100644 --- a/test_runner/fixtures/pageserver/http.py +++ b/test_runner/fixtures/pageserver/http.py @@ -563,13 +563,13 @@ class PageserverHttpClient(requests.Session): self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, - timestamp, + timestamp: datetime, ): log.info( f"Requesting lsn by timestamp {timestamp}, tenant {tenant_id}, timeline {timeline_id}" ) res = self.get( - f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/get_lsn_by_timestamp?timestamp={timestamp}", + f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/get_lsn_by_timestamp?timestamp={timestamp.isoformat()}Z", ) self.verbose_error(res) res_json = res.json() diff --git a/test_runner/regress/test_lsn_mapping.py b/test_runner/regress/test_lsn_mapping.py index 9788e8c0d7..50d7c74af0 100644 --- a/test_runner/regress/test_lsn_mapping.py +++ b/test_runner/regress/test_lsn_mapping.py @@ -64,18 +64,14 @@ def test_lsn_mapping(neon_env_builder: NeonEnvBuilder): # Check edge cases # Timestamp is in the future probe_timestamp = tbl[-1][1] + timedelta(hours=1) - result = client.timeline_get_lsn_by_timestamp( - tenant_id, timeline_id, f"{probe_timestamp.isoformat()}Z" - ) + result = client.timeline_get_lsn_by_timestamp(tenant_id, timeline_id, probe_timestamp) assert result["kind"] == "future" # make sure that we return a well advanced lsn here assert Lsn(result["lsn"]) > start_lsn # Timestamp is in the unreachable past probe_timestamp = tbl[0][1] - timedelta(hours=10) - result = client.timeline_get_lsn_by_timestamp( - tenant_id, timeline_id, f"{probe_timestamp.isoformat()}Z" - ) + result = client.timeline_get_lsn_by_timestamp(tenant_id, timeline_id, probe_timestamp) assert result["kind"] == "past" # make sure that we return the minimum lsn here at the start of the range assert Lsn(result["lsn"]) < start_lsn @@ -83,9 +79,7 @@ def test_lsn_mapping(neon_env_builder: NeonEnvBuilder): # Probe a bunch of timestamps in the valid range for i in range(1, len(tbl), 100): probe_timestamp = tbl[i][1] - result = client.timeline_get_lsn_by_timestamp( - tenant_id, timeline_id, f"{probe_timestamp.isoformat()}Z" - ) + result = client.timeline_get_lsn_by_timestamp(tenant_id, timeline_id, probe_timestamp) assert result["kind"] not in ["past", "nodata"] lsn = result["lsn"] # Call get_lsn_by_timestamp to get the LSN @@ -108,9 +102,7 @@ def test_lsn_mapping(neon_env_builder: NeonEnvBuilder): # Timestamp is in the unreachable past probe_timestamp = tbl[0][1] - timedelta(hours=10) - result = client.timeline_get_lsn_by_timestamp( - tenant_id, timeline_id_child, f"{probe_timestamp.isoformat()}Z" - ) + result = client.timeline_get_lsn_by_timestamp(tenant_id, timeline_id_child, probe_timestamp) assert result["kind"] == "past" # make sure that we return the minimum lsn here at the start of the range assert Lsn(result["lsn"]) >= last_flush_lsn From 3d4fe205ba260c6cd878bf8d0c19623d45920e4f Mon Sep 17 00:00:00 2001 From: John Spray Date: Wed, 7 Feb 2024 13:08:09 +0000 Subject: [PATCH 29/72] control_plane/attachment_service: database connection pool (#6622) ## Problem This is mainly to limit our concurrency, rather than to speed up requests (I was doing some sanity checks on performance of the service with thousands of shards) ## Summary of changes - Enable the `diesel:r2d2` feature, which provides an async connection pool - Acquire a connection before entering spawn_blocking for a database transaction (recall that diesel's interface is sync) - Set a connection pool size of 99 to fit within default postgres limit (100) - Also set the tokio blocking thread count to accomodate the same number of blocking tasks (the only thing we use spawn_blocking for is database calls). --- Cargo.lock | 23 +++++++++++ control_plane/attachment_service/Cargo.toml | 3 +- control_plane/attachment_service/src/main.rs | 15 ++++++- .../attachment_service/src/persistence.rs | 41 ++++++++++++++----- .../attachment_service/src/service.rs | 4 +- workspace_hack/Cargo.toml | 3 +- 6 files changed, 74 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b2b2777408..a25725f90d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -289,6 +289,7 @@ dependencies = [ "pageserver_api", "pageserver_client", "postgres_connection", + "r2d2", "reqwest", "serde", "serde_json", @@ -1651,6 +1652,7 @@ dependencies = [ "diesel_derives", "itoa", "pq-sys", + "r2d2", "serde_json", ] @@ -4166,6 +4168,17 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot 0.12.1", + "scheduled-thread-pool", +] + [[package]] name = "rand" version = "0.7.3" @@ -4879,6 +4892,15 @@ dependencies = [ "windows-sys 0.42.0", ] +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot 0.12.1", +] + [[package]] name = "scopeguard" version = "1.1.0" @@ -6807,6 +6829,7 @@ dependencies = [ "clap_builder", "crossbeam-utils", "diesel", + "diesel_derives", "either", "fail", "futures-channel", diff --git a/control_plane/attachment_service/Cargo.toml b/control_plane/attachment_service/Cargo.toml index 3a65153c41..0b93211dbc 100644 --- a/control_plane/attachment_service/Cargo.toml +++ b/control_plane/attachment_service/Cargo.toml @@ -24,8 +24,9 @@ tokio.workspace = true tokio-util.workspace = true tracing.workspace = true -diesel = { version = "2.1.4", features = ["serde_json", "postgres"] } +diesel = { version = "2.1.4", features = ["serde_json", "postgres", "r2d2"] } diesel_migrations = { version = "2.1.0" } +r2d2 = { version = "0.8.10" } utils = { path = "../../libs/utils/" } metrics = { path = "../../libs/metrics/" } diff --git a/control_plane/attachment_service/src/main.rs b/control_plane/attachment_service/src/main.rs index bc8a8786c2..7229a2517b 100644 --- a/control_plane/attachment_service/src/main.rs +++ b/control_plane/attachment_service/src/main.rs @@ -170,6 +170,7 @@ impl Secrets { } } +/// Execute the diesel migrations that are built into this binary async fn migration_run(database_url: &str) -> anyhow::Result<()> { use diesel::PgConnection; use diesel_migrations::{HarnessWithOutput, MigrationHarness}; @@ -183,8 +184,18 @@ async fn migration_run(database_url: &str) -> anyhow::Result<()> { Ok(()) } -#[tokio::main] -async fn main() -> anyhow::Result<()> { +fn main() -> anyhow::Result<()> { + tokio::runtime::Builder::new_current_thread() + // We use spawn_blocking for database operations, so require approximately + // as many blocking threads as we will open database connections. + .max_blocking_threads(Persistence::MAX_CONNECTIONS as usize) + .enable_all() + .build() + .unwrap() + .block_on(async_main()) +} + +async fn async_main() -> anyhow::Result<()> { let launch_ts = Box::leak(Box::new(LaunchTimestamp::generate())); logging::init( diff --git a/control_plane/attachment_service/src/persistence.rs b/control_plane/attachment_service/src/persistence.rs index 574441c409..db487bcec6 100644 --- a/control_plane/attachment_service/src/persistence.rs +++ b/control_plane/attachment_service/src/persistence.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::str::FromStr; +use std::time::Duration; use camino::Utf8Path; use camino::Utf8PathBuf; @@ -44,7 +45,7 @@ use crate::PlacementPolicy; /// updated, and reads of nodes are always from memory, not the database. We only require that /// we can UPDATE a node's scheduling mode reasonably quickly to mark a bad node offline. pub struct Persistence { - database_url: String, + connection_pool: diesel::r2d2::Pool>, // In test environments, we support loading+saving a JSON file. This is temporary, for the benefit of // test_compatibility.py, so that we don't have to commit to making the database contents fully backward/forward @@ -64,6 +65,8 @@ pub(crate) enum DatabaseError { Query(#[from] diesel::result::Error), #[error(transparent)] Connection(#[from] diesel::result::ConnectionError), + #[error(transparent)] + ConnectionPool(#[from] r2d2::Error), #[error("Logical error: {0}")] Logical(String), } @@ -71,9 +74,31 @@ pub(crate) enum DatabaseError { pub(crate) type DatabaseResult = Result; impl Persistence { + // The default postgres connection limit is 100. We use up to 99, to leave one free for a human admin under + // normal circumstances. This assumes we have exclusive use of the database cluster to which we connect. + pub const MAX_CONNECTIONS: u32 = 99; + + // We don't want to keep a lot of connections alive: close them down promptly if they aren't being used. + const IDLE_CONNECTION_TIMEOUT: Duration = Duration::from_secs(10); + const MAX_CONNECTION_LIFETIME: Duration = Duration::from_secs(60); + pub fn new(database_url: String, json_path: Option) -> Self { + let manager = diesel::r2d2::ConnectionManager::::new(database_url); + + // We will use a connection pool: this is primarily to _limit_ our connection count, rather than to optimize time + // to execute queries (database queries are not generally on latency-sensitive paths). + let connection_pool = diesel::r2d2::Pool::builder() + .max_size(Self::MAX_CONNECTIONS) + .max_lifetime(Some(Self::MAX_CONNECTION_LIFETIME)) + .idle_timeout(Some(Self::IDLE_CONNECTION_TIMEOUT)) + // Always keep at least one connection ready to go + .min_idle(Some(1)) + .test_on_check_out(true) + .build(manager) + .expect("Could not build connection pool"); + Self { - database_url, + connection_pool, json_path, } } @@ -84,14 +109,10 @@ impl Persistence { F: Fn(&mut PgConnection) -> DatabaseResult + Send + 'static, R: Send + 'static, { - let database_url = self.database_url.clone(); - tokio::task::spawn_blocking(move || -> DatabaseResult { - // TODO: connection pooling, such as via diesel::r2d2 - let mut conn = PgConnection::establish(&database_url)?; - func(&mut conn) - }) - .await - .expect("Task panic") + let mut conn = self.connection_pool.get()?; + tokio::task::spawn_blocking(move || -> DatabaseResult { func(&mut conn) }) + .await + .expect("Task panic") } /// When a node is first registered, persist it before using it for anything diff --git a/control_plane/attachment_service/src/service.rs b/control_plane/attachment_service/src/service.rs index 6f0e3ebb74..febee1aa0d 100644 --- a/control_plane/attachment_service/src/service.rs +++ b/control_plane/attachment_service/src/service.rs @@ -103,7 +103,9 @@ impl From for ApiError { match err { DatabaseError::Query(e) => ApiError::InternalServerError(e.into()), // FIXME: ApiError doesn't have an Unavailable variant, but ShuttingDown maps to 503. - DatabaseError::Connection(_e) => ApiError::ShuttingDown, + DatabaseError::Connection(_) | DatabaseError::ConnectionPool(_) => { + ApiError::ShuttingDown + } DatabaseError::Logical(reason) => { ApiError::InternalServerError(anyhow::anyhow!(reason)) } diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 74464dd4c8..70b238913d 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -29,7 +29,7 @@ chrono = { version = "0.4", default-features = false, features = ["clock", "serd clap = { version = "4", features = ["derive", "string"] } clap_builder = { version = "4", default-features = false, features = ["color", "help", "std", "string", "suggestions", "usage"] } crossbeam-utils = { version = "0.8" } -diesel = { version = "2", features = ["postgres", "serde_json"] } +diesel = { version = "2", features = ["postgres", "r2d2", "serde_json"] } either = { version = "1" } fail = { version = "0.5", default-features = false, features = ["failpoints"] } futures-channel = { version = "0.3", features = ["sink"] } @@ -90,6 +90,7 @@ anyhow = { version = "1", features = ["backtrace"] } bytes = { version = "1", features = ["serde"] } cc = { version = "1", default-features = false, features = ["parallel"] } chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "wasmbind"] } +diesel_derives = { version = "2", features = ["32-column-tables", "postgres", "r2d2", "with-deprecated"] } either = { version = "1" } getrandom = { version = "0.2", default-features = false, features = ["std"] } hashbrown-582f2526e08bb6a0 = { package = "hashbrown", version = "0.14", default-features = false, features = ["raw"] } From 090a789408e4bd95656132248bdbcbdba0fd3c4a Mon Sep 17 00:00:00 2001 From: John Spray Date: Wed, 7 Feb 2024 13:24:10 +0000 Subject: [PATCH 30/72] storage controller: use PUT instead of POST (#6659) This was a typo, the server expects PUT. --- control_plane/attachment_service/src/compute_hook.rs | 2 +- test_runner/regress/test_sharding_service.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/control_plane/attachment_service/src/compute_hook.rs b/control_plane/attachment_service/src/compute_hook.rs index 4ca26431ca..0d3610aafa 100644 --- a/control_plane/attachment_service/src/compute_hook.rs +++ b/control_plane/attachment_service/src/compute_hook.rs @@ -170,7 +170,7 @@ impl ComputeHook { reconfigure_request: &ComputeHookNotifyRequest, cancel: &CancellationToken, ) -> Result<(), NotifyError> { - let req = client.request(Method::POST, url); + let req = client.request(Method::PUT, url); let req = if let Some(value) = &self.authorization_header { req.header(reqwest::header::AUTHORIZATION, value) } else { diff --git a/test_runner/regress/test_sharding_service.py b/test_runner/regress/test_sharding_service.py index ee57fcb2cf..fd811a9d02 100644 --- a/test_runner/regress/test_sharding_service.py +++ b/test_runner/regress/test_sharding_service.py @@ -310,7 +310,7 @@ def test_sharding_service_compute_hook( notifications.append(request.json) return Response(status=200) - httpserver.expect_request("/notify", method="POST").respond_with_handler(handler) + httpserver.expect_request("/notify", method="PUT").respond_with_handler(handler) # Start running env = neon_env_builder.init_start() From 75f1a01d4aba488012c9fd86b56b6dcf46726c92 Mon Sep 17 00:00:00 2001 From: Abhijeet Patil Date: Wed, 7 Feb 2024 16:14:10 +0000 Subject: [PATCH 31/72] Optimise e2e run (#6513) ## Problem We have finite amount of runners and intermediate results are often wanted before a PR is ready for merging. Currently all PRs get e2e tests run and this creates a lot of throwaway e2e results which may or may not get to start or complete before a new push. ## Summary of changes 1. Skip e2e test when PR is in draft mode 2. Run e2e when PR status changes from draft to ready for review (change this to having its trigger in below PR and update results of build and test) 3. Abstract e2e test in a Separate workflow and call it from the main workflow for the e2e test 5. Add a label, if that label is present run e2e test in draft (run-e2e-test-in-draft) 6. Auto add a label(approve to ci) so that all the external contributors PR , e2e run in draft 7. Document the new label changes and the above behaviour Draft PR : https://github.com/neondatabase/neon/actions/runs/7729128470 Ready To Review : https://github.com/neondatabase/neon/actions/runs/7733779916 Draft PR with label : https://github.com/neondatabase/neon/actions/runs/7725691012/job/21062432342 and https://github.com/neondatabase/neon/actions/runs/7733854028 ## Checklist before requesting a review - [x] I have performed a self-review of my code. - [ ] If it is a core feature, I have added thorough tests. - [ ] Do we need to implement analytics? if so did you add the relevant metrics to the dashboard? - [ ] If this PR requires public announcement, mark it with /release-notes label and add several sentences in this section. ## Checklist before merging - [ ] Do not forget to reformat commit message to not include the above checklist --------- Co-authored-by: Alexander Bayandin --- .github/workflows/approved-for-ci-run.yml | 1 + .github/workflows/build_and_test.yml | 48 +-------- .github/workflows/trigger-e2e-tests.yml | 118 ++++++++++++++++++++++ CONTRIBUTING.md | 3 + 4 files changed, 126 insertions(+), 44 deletions(-) create mode 100644 .github/workflows/trigger-e2e-tests.yml diff --git a/.github/workflows/approved-for-ci-run.yml b/.github/workflows/approved-for-ci-run.yml index 5b21011b83..ae2f173b47 100644 --- a/.github/workflows/approved-for-ci-run.yml +++ b/.github/workflows/approved-for-ci-run.yml @@ -93,6 +93,7 @@ jobs: --body-file "body.md" \ --head "${BRANCH}" \ --base "main" \ + --label "run-e2e-tests-in-draft" \ --draft fi diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index f12f020634..078916e1ea 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -22,7 +22,7 @@ env: AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_DEV }} AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_KEY_DEV }} # A concurrency group that we use for e2e-tests runs, matches `concurrency.group` above with `github.repository` as a prefix - E2E_CONCURRENCY_GROUP: ${{ github.repository }}-${{ github.workflow }}-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }} + E2E_CONCURRENCY_GROUP: ${{ github.repository }}-e2e-tests-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }} jobs: check-permissions: @@ -692,50 +692,10 @@ jobs: }) trigger-e2e-tests: + if: ${{ !github.event.pull_request.draft || contains( github.event.pull_request.labels.*.name, 'run-e2e-tests-in-draft') || github.ref_name == 'main' || github.ref_name == 'release' }} needs: [ check-permissions, promote-images, tag ] - runs-on: [ self-hosted, gen3, small ] - container: - image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/base:pinned - options: --init - steps: - - name: Set PR's status to pending and request a remote CI test - run: | - # For pull requests, GH Actions set "github.sha" variable to point at a fake merge commit - # but we need to use a real sha of a latest commit in the PR's branch for the e2e job, - # to place a job run status update later. - COMMIT_SHA=${{ github.event.pull_request.head.sha }} - # For non-PR kinds of runs, the above will produce an empty variable, pick the original sha value for those - COMMIT_SHA=${COMMIT_SHA:-${{ github.sha }}} - - REMOTE_REPO="${{ github.repository_owner }}/cloud" - - curl -f -X POST \ - https://api.github.com/repos/${{ github.repository }}/statuses/$COMMIT_SHA \ - -H "Accept: application/vnd.github.v3+json" \ - --user "${{ secrets.CI_ACCESS_TOKEN }}" \ - --data \ - "{ - \"state\": \"pending\", - \"context\": \"neon-cloud-e2e\", - \"description\": \"[$REMOTE_REPO] Remote CI job is about to start\" - }" - - curl -f -X POST \ - https://api.github.com/repos/$REMOTE_REPO/actions/workflows/testing.yml/dispatches \ - -H "Accept: application/vnd.github.v3+json" \ - --user "${{ secrets.CI_ACCESS_TOKEN }}" \ - --data \ - "{ - \"ref\": \"main\", - \"inputs\": { - \"ci_job_name\": \"neon-cloud-e2e\", - \"commit_hash\": \"$COMMIT_SHA\", - \"remote_repo\": \"${{ github.repository }}\", - \"storage_image_tag\": \"${{ needs.tag.outputs.build-tag }}\", - \"compute_image_tag\": \"${{ needs.tag.outputs.build-tag }}\", - \"concurrency_group\": \"${{ env.E2E_CONCURRENCY_GROUP }}\" - } - }" + uses: ./.github/workflows/trigger-e2e-tests.yml + secrets: inherit neon-image: needs: [ check-permissions, build-buildtools-image, tag ] diff --git a/.github/workflows/trigger-e2e-tests.yml b/.github/workflows/trigger-e2e-tests.yml new file mode 100644 index 0000000000..2776033805 --- /dev/null +++ b/.github/workflows/trigger-e2e-tests.yml @@ -0,0 +1,118 @@ +name: Trigger E2E Tests + +on: + pull_request: + types: + - ready_for_review + workflow_call: + +defaults: + run: + shell: bash -euxo pipefail {0} + +env: + # A concurrency group that we use for e2e-tests runs, matches `concurrency.group` above with `github.repository` as a prefix + E2E_CONCURRENCY_GROUP: ${{ github.repository }}-e2e-tests-${{ github.ref_name }}-${{ github.ref_name == 'main' && github.sha || 'anysha' }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_DEV }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_KEY_DEV }} + +jobs: + cancel-previous-e2e-tests: + if: github.event_name == 'pull_request' + runs-on: ubuntu-latest + + steps: + - name: Cancel previous e2e-tests runs for this PR + env: + GH_TOKEN: ${{ secrets.CI_ACCESS_TOKEN }} + run: | + gh workflow --repo neondatabase/cloud \ + run cancel-previous-in-concurrency-group.yml \ + --field concurrency_group="${{ env.E2E_CONCURRENCY_GROUP }}" + + tag: + runs-on: [ ubuntu-latest ] + outputs: + build-tag: ${{ steps.build-tag.outputs.tag }} + + steps: + - name: Checkout + uses: actions/checkout@v3 + with: + fetch-depth: 0 + + - name: Get build tag + env: + GH_TOKEN: ${{ secrets.CI_ACCESS_TOKEN }} + CURRENT_BRANCH: ${{ github.head_ref || github.ref_name }} + CURRENT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + run: | + if [[ "$GITHUB_REF_NAME" == "main" ]]; then + echo "tag=$(git rev-list --count HEAD)" | tee -a $GITHUB_OUTPUT + elif [[ "$GITHUB_REF_NAME" == "release" ]]; then + echo "tag=release-$(git rev-list --count HEAD)" | tee -a $GITHUB_OUTPUT + else + echo "GITHUB_REF_NAME (value '$GITHUB_REF_NAME') is not set to either 'main' or 'release'" + BUILD_AND_TEST_RUN_ID=$(gh run list -b $CURRENT_BRANCH -c $CURRENT_SHA -w 'Build and Test' -L 1 --json databaseId --jq '.[].databaseId') + echo "tag=$BUILD_AND_TEST_RUN_ID" | tee -a $GITHUB_OUTPUT + fi + id: build-tag + + trigger-e2e-tests: + needs: [ tag ] + runs-on: [ self-hosted, gen3, small ] + env: + TAG: ${{ needs.tag.outputs.build-tag }} + container: + image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/base:pinned + options: --init + steps: + - name: check if ecr image are present + run: | + for REPO in neon compute-tools compute-node-v14 vm-compute-node-v14 compute-node-v15 vm-compute-node-v15 compute-node-v16 vm-compute-node-v16; do + OUTPUT=$(aws ecr describe-images --repository-name ${REPO} --region eu-central-1 --query "imageDetails[?imageTags[?contains(@, '${TAG}')]]" --output text) + if [ "$OUTPUT" == "" ]; then + echo "$REPO with image tag $TAG not found" >> $GITHUB_OUTPUT + exit 1 + fi + done + + - name: Set PR's status to pending and request a remote CI test + run: | + # For pull requests, GH Actions set "github.sha" variable to point at a fake merge commit + # but we need to use a real sha of a latest commit in the PR's branch for the e2e job, + # to place a job run status update later. + COMMIT_SHA=${{ github.event.pull_request.head.sha }} + # For non-PR kinds of runs, the above will produce an empty variable, pick the original sha value for those + COMMIT_SHA=${COMMIT_SHA:-${{ github.sha }}} + + REMOTE_REPO="${{ github.repository_owner }}/cloud" + + curl -f -X POST \ + https://api.github.com/repos/${{ github.repository }}/statuses/$COMMIT_SHA \ + -H "Accept: application/vnd.github.v3+json" \ + --user "${{ secrets.CI_ACCESS_TOKEN }}" \ + --data \ + "{ + \"state\": \"pending\", + \"context\": \"neon-cloud-e2e\", + \"description\": \"[$REMOTE_REPO] Remote CI job is about to start\" + }" + + curl -f -X POST \ + https://api.github.com/repos/$REMOTE_REPO/actions/workflows/testing.yml/dispatches \ + -H "Accept: application/vnd.github.v3+json" \ + --user "${{ secrets.CI_ACCESS_TOKEN }}" \ + --data \ + "{ + \"ref\": \"main\", + \"inputs\": { + \"ci_job_name\": \"neon-cloud-e2e\", + \"commit_hash\": \"$COMMIT_SHA\", + \"remote_repo\": \"${{ github.repository }}\", + \"storage_image_tag\": \"${TAG}\", + \"compute_image_tag\": \"${TAG}\", + \"concurrency_group\": \"${{ env.E2E_CONCURRENCY_GROUP }}\" + } + }" + \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7e177693fa..2e447fba47 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -54,6 +54,9 @@ _An instruction for maintainers_ - If and only if it looks **safe** (i.e. it doesn't contain any malicious code which could expose secrets or harm the CI), then: - Press the "Approve and run" button in GitHub UI - Add the `approved-for-ci-run` label to the PR + - Currently draft PR will skip e2e test (only for internal contributors). After turning the PR 'Ready to Review' CI will trigger e2e test + - Add `run-e2e-tests-in-draft` label to run e2e test in draft PR (override above behaviour) + - The `approved-for-ci-run` workflow will add `run-e2e-tests-in-draft` automatically to run e2e test for external contributors Repeat all steps after any change to the PR. - When the changes are ready to get merged — merge the original PR (not the internal one) From 7b49e5e5c334bc8d07232f385d08e370ba85fb5a Mon Sep 17 00:00:00 2001 From: Sasha Krassovsky Date: Wed, 7 Feb 2024 07:55:55 -0900 Subject: [PATCH 32/72] Remove compute migrations feature flag (#6653) --- compute_tools/src/compute.rs | 11 +++++------ libs/compute_api/src/spec.rs | 3 --- test_runner/fixtures/neon_fixtures.py | 5 +---- test_runner/regress/test_migrations.py | 2 +- test_runner/regress/test_neon_superuser.py | 4 ++-- 5 files changed, 9 insertions(+), 16 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 098e06cca9..0ca1a47fbf 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -773,12 +773,11 @@ impl ComputeNode { // 'Close' connection drop(client); - if self.has_feature(ComputeFeature::Migrations) { - thread::spawn(move || { - let mut client = Client::connect(connstr.as_str(), NoTls)?; - handle_migrations(&mut client) - }); - } + // Run migrations separately to not hold up cold starts + thread::spawn(move || { + let mut client = Client::connect(connstr.as_str(), NoTls)?; + handle_migrations(&mut client) + }); Ok(()) } diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index 5361d14004..13ac18e0c5 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -90,9 +90,6 @@ pub enum ComputeFeature { /// track short-lived connections as user activity. ActivityMonitorExperimental, - /// Enable running migrations - Migrations, - /// This is a special feature flag that is used to represent unknown feature flags. /// Basically all unknown to enum flags are represented as this one. See unit test /// `parse_unknown_features()` for more details. diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index bf7c6ccc14..4491655aeb 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -3131,10 +3131,7 @@ class Endpoint(PgProtocol): log.info(json.dumps(dict(data_dict, **kwargs))) json.dump(dict(data_dict, **kwargs), file, indent=4) - # Please note: if you didn't respec this endpoint to have the `migrations` - # feature, this function will probably fail because neon_migration.migration_id - # won't exist. This is temporary - soon we'll get rid of the feature flag and - # migrations will be enabled for everyone. + # Please note: Migrations only run if pg_skip_catalog_updates is false def wait_for_migrations(self): with self.cursor() as cur: diff --git a/test_runner/regress/test_migrations.py b/test_runner/regress/test_migrations.py index 30dd54a8c1..8954810451 100644 --- a/test_runner/regress/test_migrations.py +++ b/test_runner/regress/test_migrations.py @@ -10,7 +10,7 @@ def test_migrations(neon_simple_env: NeonEnv): endpoint = env.endpoints.create("test_migrations") log_path = endpoint.endpoint_path() / "compute.log" - endpoint.respec(skip_pg_catalog_updates=False, features=["migrations"]) + endpoint.respec(skip_pg_catalog_updates=False) endpoint.start() endpoint.wait_for_migrations() diff --git a/test_runner/regress/test_neon_superuser.py b/test_runner/regress/test_neon_superuser.py index eff2cadabf..34f1e64b34 100644 --- a/test_runner/regress/test_neon_superuser.py +++ b/test_runner/regress/test_neon_superuser.py @@ -12,10 +12,10 @@ def test_neon_superuser(neon_simple_env: NeonEnv, pg_version: PgVersion): env.neon_cli.create_branch("test_neon_superuser_subscriber") sub = env.endpoints.create("test_neon_superuser_subscriber") - pub.respec(skip_pg_catalog_updates=False, features=["migrations"]) + pub.respec(skip_pg_catalog_updates=False) pub.start() - sub.respec(skip_pg_catalog_updates=False, features=["migrations"]) + sub.respec(skip_pg_catalog_updates=False) sub.start() pub.wait_for_migrations() From 51f9385b1bd60f3152a580332ba4b19ec131f89a Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Wed, 7 Feb 2024 18:47:55 +0100 Subject: [PATCH 33/72] live-reconfigurable virtual_file::IoEngine (#6552) This PR adds an API to live-reconfigure the VirtualFile io engine. It also adds a flag to `pagebench get-page-latest-lsn`, which is where I found this functionality to be useful: it helps compare the io engines in a benchmark without re-compiling a release build, which took ~50s on the i3en.3xlarge where I was doing the benchmark. Switching the IO engine is completely safe at runtime. --- libs/pageserver_api/src/models.rs | 21 +++ pageserver/client/src/mgmt_api.rs | 12 ++ pageserver/ctl/src/layer_map_analyzer.rs | 2 +- pageserver/ctl/src/layers.rs | 4 +- pageserver/ctl/src/main.rs | 2 +- .../pagebench/src/cmd/getpage_latest_lsn.rs | 8 ++ pageserver/src/http/routes.rs | 10 ++ pageserver/src/virtual_file.rs | 5 +- pageserver/src/virtual_file/io_engine.rs | 130 +++++++++++------- pageserver/src/virtual_file/open_options.rs | 7 +- 10 files changed, 144 insertions(+), 57 deletions(-) diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs index 5a638df9cc..c08cacb822 100644 --- a/libs/pageserver_api/src/models.rs +++ b/libs/pageserver_api/src/models.rs @@ -649,6 +649,27 @@ pub struct WalRedoManagerStatus { pub pid: Option, } +pub mod virtual_file { + #[derive( + Copy, + Clone, + PartialEq, + Eq, + Hash, + strum_macros::EnumString, + strum_macros::Display, + serde_with::DeserializeFromStr, + serde_with::SerializeDisplay, + Debug, + )] + #[strum(serialize_all = "kebab-case")] + pub enum IoEngineKind { + StdFs, + #[cfg(target_os = "linux")] + TokioEpollUring, + } +} + // Wrapped in libpq CopyData #[derive(PartialEq, Eq, Debug)] pub enum PagestreamFeMessage { diff --git a/pageserver/client/src/mgmt_api.rs b/pageserver/client/src/mgmt_api.rs index 91b9afa026..8abe58e1a2 100644 --- a/pageserver/client/src/mgmt_api.rs +++ b/pageserver/client/src/mgmt_api.rs @@ -339,4 +339,16 @@ impl Client { .await .map_err(Error::ReceiveBody) } + + pub async fn put_io_engine( + &self, + engine: &pageserver_api::models::virtual_file::IoEngineKind, + ) -> Result<()> { + let uri = format!("{}/v1/io_engine", self.mgmt_api_endpoint); + self.request(Method::PUT, uri, engine) + .await? + .json() + .await + .map_err(Error::ReceiveBody) + } } diff --git a/pageserver/ctl/src/layer_map_analyzer.rs b/pageserver/ctl/src/layer_map_analyzer.rs index eb5c3f15cf..42c4e9ff48 100644 --- a/pageserver/ctl/src/layer_map_analyzer.rs +++ b/pageserver/ctl/src/layer_map_analyzer.rs @@ -142,7 +142,7 @@ pub(crate) async fn main(cmd: &AnalyzeLayerMapCmd) -> Result<()> { let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error); // Initialize virtual_file (file desriptor cache) and page cache which are needed to access layer persistent B-Tree. - pageserver::virtual_file::init(10, virtual_file::IoEngineKind::StdFs); + pageserver::virtual_file::init(10, virtual_file::api::IoEngineKind::StdFs); pageserver::page_cache::init(100); let mut total_delta_layers = 0usize; diff --git a/pageserver/ctl/src/layers.rs b/pageserver/ctl/src/layers.rs index dbbcfedac0..27efa6d028 100644 --- a/pageserver/ctl/src/layers.rs +++ b/pageserver/ctl/src/layers.rs @@ -59,7 +59,7 @@ pub(crate) enum LayerCmd { async fn read_delta_file(path: impl AsRef, ctx: &RequestContext) -> Result<()> { let path = Utf8Path::from_path(path.as_ref()).expect("non-Unicode path"); - virtual_file::init(10, virtual_file::IoEngineKind::StdFs); + virtual_file::init(10, virtual_file::api::IoEngineKind::StdFs); page_cache::init(100); let file = FileBlockReader::new(VirtualFile::open(path).await?); let summary_blk = file.read_blk(0, ctx).await?; @@ -187,7 +187,7 @@ pub(crate) async fn main(cmd: &LayerCmd) -> Result<()> { new_tenant_id, new_timeline_id, } => { - pageserver::virtual_file::init(10, virtual_file::IoEngineKind::StdFs); + pageserver::virtual_file::init(10, virtual_file::api::IoEngineKind::StdFs); pageserver::page_cache::init(100); let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error); diff --git a/pageserver/ctl/src/main.rs b/pageserver/ctl/src/main.rs index 3c90933fe9..e73d961e36 100644 --- a/pageserver/ctl/src/main.rs +++ b/pageserver/ctl/src/main.rs @@ -123,7 +123,7 @@ fn read_pg_control_file(control_file_path: &Utf8Path) -> anyhow::Result<()> { async fn print_layerfile(path: &Utf8Path) -> anyhow::Result<()> { // Basic initialization of things that don't change after startup - virtual_file::init(10, virtual_file::IoEngineKind::StdFs); + virtual_file::init(10, virtual_file::api::IoEngineKind::StdFs); page_cache::init(100); let ctx = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error); dump_layerfile_from_path(path, true, &ctx).await diff --git a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs index aa809d8d26..647f571e59 100644 --- a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs +++ b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs @@ -51,6 +51,10 @@ pub(crate) struct Args { /// It doesn't get invalidated if the keyspace changes under the hood, e.g., due to new ingested data or compaction. #[clap(long)] keyspace_cache: Option, + /// Before starting the benchmark, live-reconfigure the pageserver to use the given + /// [`pageserver_api::models::virtual_file::IoEngineKind`]. + #[clap(long)] + set_io_engine: Option, targets: Option>, } @@ -109,6 +113,10 @@ async fn main_impl( args.pageserver_jwt.as_deref(), )); + if let Some(engine_str) = &args.set_io_engine { + mgmt_api_client.put_io_engine(engine_str).await?; + } + // discover targets let timelines: Vec = crate::util::cli::targets::discover( &mgmt_api_client, diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index 792089ebe7..ebcb27fa08 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -1908,6 +1908,15 @@ async fn post_tracing_event_handler( json_response(StatusCode::OK, ()) } +async fn put_io_engine_handler( + mut r: Request, + _cancel: CancellationToken, +) -> Result, ApiError> { + let kind: crate::virtual_file::IoEngineKind = json_request(&mut r).await?; + crate::virtual_file::io_engine::set(kind); + json_response(StatusCode::OK, ()) +} + /// Common functionality of all the HTTP API handlers. /// /// - Adds a tracing span to each request (by `request_span`) @@ -2165,5 +2174,6 @@ pub fn make_router( "/v1/tenant/:tenant_shard_id/timeline/:timeline_id/keyspace", |r| testing_api_handler("read out the keyspace", r, timeline_collect_keyspace), ) + .put("/v1/io_engine", |r| api_handler(r, put_io_engine_handler)) .any(handler_404)) } diff --git a/pageserver/src/virtual_file.rs b/pageserver/src/virtual_file.rs index 066f06c88f..059a6596d3 100644 --- a/pageserver/src/virtual_file.rs +++ b/pageserver/src/virtual_file.rs @@ -28,9 +28,10 @@ use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; use tokio::time::Instant; use utils::fs_ext; -mod io_engine; +pub use pageserver_api::models::virtual_file as api; +pub(crate) mod io_engine; mod open_options; -pub use io_engine::IoEngineKind; +pub(crate) use io_engine::IoEngineKind; pub(crate) use open_options::*; /// diff --git a/pageserver/src/virtual_file/io_engine.rs b/pageserver/src/virtual_file/io_engine.rs index f7b46fe653..892affa326 100644 --- a/pageserver/src/virtual_file/io_engine.rs +++ b/pageserver/src/virtual_file/io_engine.rs @@ -7,67 +7,100 @@ //! //! Then use [`get`] and [`super::OpenOptions`]. -#[derive( - Copy, - Clone, - PartialEq, - Eq, - Hash, - strum_macros::EnumString, - strum_macros::Display, - serde_with::DeserializeFromStr, - serde_with::SerializeDisplay, - Debug, -)] -#[strum(serialize_all = "kebab-case")] -pub enum IoEngineKind { +pub(crate) use super::api::IoEngineKind; +#[derive(Clone, Copy)] +#[repr(u8)] +pub(crate) enum IoEngine { + NotSet, StdFs, #[cfg(target_os = "linux")] TokioEpollUring, } -static IO_ENGINE: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); - -#[cfg(not(test))] -pub(super) fn init(engine: IoEngineKind) { - if IO_ENGINE.set(engine).is_err() { - panic!("called twice"); +impl From for IoEngine { + fn from(value: IoEngineKind) -> Self { + match value { + IoEngineKind::StdFs => IoEngine::StdFs, + #[cfg(target_os = "linux")] + IoEngineKind::TokioEpollUring => IoEngine::TokioEpollUring, + } } - crate::metrics::virtual_file_io_engine::KIND - .with_label_values(&[&format!("{engine}")]) - .set(1); } -pub(super) fn get() -> &'static IoEngineKind { - #[cfg(test)] - { - let env_var_name = "NEON_PAGESERVER_UNIT_TEST_VIRTUAL_FILE_IOENGINE"; - IO_ENGINE.get_or_init(|| match std::env::var(env_var_name) { - Ok(v) => match v.parse::() { - Ok(engine_kind) => engine_kind, - Err(e) => { - panic!("invalid VirtualFile io engine for env var {env_var_name}: {e:#}: {v:?}") - } - }, - Err(std::env::VarError::NotPresent) => { - crate::config::defaults::DEFAULT_VIRTUAL_FILE_IO_ENGINE - .parse() - .unwrap() - } - Err(std::env::VarError::NotUnicode(_)) => { - panic!("env var {env_var_name} is not unicode"); - } +impl TryFrom for IoEngine { + type Error = u8; + + fn try_from(value: u8) -> Result { + Ok(match value { + v if v == (IoEngine::NotSet as u8) => IoEngine::NotSet, + v if v == (IoEngine::StdFs as u8) => IoEngine::StdFs, + #[cfg(target_os = "linux")] + v if v == (IoEngine::TokioEpollUring as u8) => IoEngine::TokioEpollUring, + x => return Err(x), }) } - #[cfg(not(test))] - IO_ENGINE.get().unwrap() } -use std::os::unix::prelude::FileExt; +static IO_ENGINE: AtomicU8 = AtomicU8::new(IoEngine::NotSet as u8); + +pub(crate) fn set(engine_kind: IoEngineKind) { + let engine: IoEngine = engine_kind.into(); + IO_ENGINE.store(engine as u8, std::sync::atomic::Ordering::Relaxed); + #[cfg(not(test))] + { + let metric = &crate::metrics::virtual_file_io_engine::KIND; + metric.reset(); + metric + .with_label_values(&[&format!("{engine_kind}")]) + .set(1); + } +} + +#[cfg(not(test))] +pub(super) fn init(engine_kind: IoEngineKind) { + set(engine_kind); +} + +pub(super) fn get() -> IoEngine { + let cur = IoEngine::try_from(IO_ENGINE.load(Ordering::Relaxed)).unwrap(); + if cfg!(test) { + let env_var_name = "NEON_PAGESERVER_UNIT_TEST_VIRTUAL_FILE_IOENGINE"; + match cur { + IoEngine::NotSet => { + let kind = match std::env::var(env_var_name) { + Ok(v) => match v.parse::() { + Ok(engine_kind) => engine_kind, + Err(e) => { + panic!("invalid VirtualFile io engine for env var {env_var_name}: {e:#}: {v:?}") + } + }, + Err(std::env::VarError::NotPresent) => { + crate::config::defaults::DEFAULT_VIRTUAL_FILE_IO_ENGINE + .parse() + .unwrap() + } + Err(std::env::VarError::NotUnicode(_)) => { + panic!("env var {env_var_name} is not unicode"); + } + }; + self::set(kind); + self::get() + } + x => x, + } + } else { + cur + } +} + +use std::{ + os::unix::prelude::FileExt, + sync::atomic::{AtomicU8, Ordering}, +}; use super::FileGuard; -impl IoEngineKind { +impl IoEngine { pub(super) async fn read_at( &self, file_guard: FileGuard, @@ -78,7 +111,8 @@ impl IoEngineKind { B: tokio_epoll_uring::BoundedBufMut + Send, { match self { - IoEngineKind::StdFs => { + IoEngine::NotSet => panic!("not initialized"), + IoEngine::StdFs => { // SAFETY: `dst` only lives at most as long as this match arm, during which buf remains valid memory. let dst = unsafe { std::slice::from_raw_parts_mut(buf.stable_mut_ptr(), buf.bytes_total()) @@ -96,7 +130,7 @@ impl IoEngineKind { ((file_guard, buf), res) } #[cfg(target_os = "linux")] - IoEngineKind::TokioEpollUring => { + IoEngine::TokioEpollUring => { let system = tokio_epoll_uring::thread_local_system().await; let (resources, res) = system.read(file_guard, offset, buf).await; ( diff --git a/pageserver/src/virtual_file/open_options.rs b/pageserver/src/virtual_file/open_options.rs index 1e5ffe15cc..f75edb0bac 100644 --- a/pageserver/src/virtual_file/open_options.rs +++ b/pageserver/src/virtual_file/open_options.rs @@ -1,6 +1,6 @@ //! Enum-dispatch to the `OpenOptions` type of the respective [`super::IoEngineKind`]; -use super::IoEngineKind; +use super::io_engine::IoEngine; use std::{os::fd::OwnedFd, path::Path}; #[derive(Debug, Clone)] @@ -13,9 +13,10 @@ pub enum OpenOptions { impl Default for OpenOptions { fn default() -> Self { match super::io_engine::get() { - IoEngineKind::StdFs => Self::StdFs(std::fs::OpenOptions::new()), + IoEngine::NotSet => panic!("io engine not set"), + IoEngine::StdFs => Self::StdFs(std::fs::OpenOptions::new()), #[cfg(target_os = "linux")] - IoEngineKind::TokioEpollUring => { + IoEngine::TokioEpollUring => { Self::TokioEpollUring(tokio_epoll_uring::ops::open_at::OpenOptions::new()) } } From 2e9b1f7aaf61d5886f312628d4fb54a1526317f2 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Tue, 6 Feb 2024 14:34:20 -0600 Subject: [PATCH 34/72] Update Postgres 14 to 14.11 --- vendor/postgres-v14 | 2 +- vendor/revisions.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vendor/postgres-v14 b/vendor/postgres-v14 index be7a65fe67..018fb05201 160000 --- a/vendor/postgres-v14 +++ b/vendor/postgres-v14 @@ -1 +1 @@ -Subproject commit be7a65fe67dc81d85bbcbebb13e00d94715f4b88 +Subproject commit 018fb052011081dc2733d3118d12e5c36df6eba1 diff --git a/vendor/revisions.json b/vendor/revisions.json index 80699839ba..c2f9244116 100644 --- a/vendor/revisions.json +++ b/vendor/revisions.json @@ -1,5 +1,5 @@ { "postgres-v16": "f7ea954989a2e7901f858779cff55259f203479a", "postgres-v15": "81e16cd537053f49e175d4a08ab7c8aec3d9b535", - "postgres-v14": "be7a65fe67dc81d85bbcbebb13e00d94715f4b88" + "postgres-v14": "018fb052011081dc2733d3118d12e5c36df6eba1" } From 5541244dc4736208e802dd60d6f9861392d9b743 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Tue, 6 Feb 2024 14:35:37 -0600 Subject: [PATCH 35/72] Update Postgres 15 to 15.6 --- vendor/postgres-v15 | 2 +- vendor/revisions.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vendor/postgres-v15 b/vendor/postgres-v15 index 81e16cd537..6ee78a3c29 160000 --- a/vendor/postgres-v15 +++ b/vendor/postgres-v15 @@ -1 +1 @@ -Subproject commit 81e16cd537053f49e175d4a08ab7c8aec3d9b535 +Subproject commit 6ee78a3c29e33cafd85ba09568b6b5eb031d29b9 diff --git a/vendor/revisions.json b/vendor/revisions.json index c2f9244116..c7076231e5 100644 --- a/vendor/revisions.json +++ b/vendor/revisions.json @@ -1,5 +1,5 @@ { "postgres-v16": "f7ea954989a2e7901f858779cff55259f203479a", - "postgres-v15": "81e16cd537053f49e175d4a08ab7c8aec3d9b535", + "postgres-v15": "6ee78a3c29e33cafd85ba09568b6b5eb031d29b9", "postgres-v14": "018fb052011081dc2733d3118d12e5c36df6eba1" } From 128fae70548f06ebc8ac44c38576c993ae6cba52 Mon Sep 17 00:00:00 2001 From: Tristan Partin Date: Tue, 6 Feb 2024 14:37:21 -0600 Subject: [PATCH 36/72] Update Postgres 16 to 16.2 --- libs/walproposer/src/walproposer.rs | 7 +++++-- vendor/postgres-v16 | 2 +- vendor/revisions.json | 2 +- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/libs/walproposer/src/walproposer.rs b/libs/walproposer/src/walproposer.rs index 7251545792..8ab8fb1a07 100644 --- a/libs/walproposer/src/walproposer.rs +++ b/libs/walproposer/src/walproposer.rs @@ -453,9 +453,12 @@ mod tests { event_mask: 0, }), expected_messages: vec![ - // Greeting(ProposerGreeting { protocol_version: 2, pg_version: 160001, proposer_id: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], system_id: 0, timeline_id: 9e4c8f36063c6c6e93bc20d65a820f3d, tenant_id: 9e4c8f36063c6c6e93bc20d65a820f3d, tli: 1, wal_seg_size: 16777216 }) + // TODO: When updating Postgres versions, this test will cause + // problems. Postgres version in message needs updating. + // + // Greeting(ProposerGreeting { protocol_version: 2, pg_version: 160002, proposer_id: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], system_id: 0, timeline_id: 9e4c8f36063c6c6e93bc20d65a820f3d, tenant_id: 9e4c8f36063c6c6e93bc20d65a820f3d, tli: 1, wal_seg_size: 16777216 }) vec![ - 103, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 1, 113, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 103, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 2, 113, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 158, 76, 143, 54, 6, 60, 108, 110, 147, 188, 32, 214, 90, 130, 15, 61, 158, 76, 143, 54, 6, 60, 108, 110, 147, 188, 32, 214, 90, 130, 15, 61, 1, 0, 0, 0, 0, 0, 0, 1, diff --git a/vendor/postgres-v16 b/vendor/postgres-v16 index f7ea954989..550cdd26d4 160000 --- a/vendor/postgres-v16 +++ b/vendor/postgres-v16 @@ -1 +1 @@ -Subproject commit f7ea954989a2e7901f858779cff55259f203479a +Subproject commit 550cdd26d445afdd26b15aa93c8c2f3dc52f8361 diff --git a/vendor/revisions.json b/vendor/revisions.json index c7076231e5..91ebb8cb34 100644 --- a/vendor/revisions.json +++ b/vendor/revisions.json @@ -1,5 +1,5 @@ { - "postgres-v16": "f7ea954989a2e7901f858779cff55259f203479a", + "postgres-v16": "550cdd26d445afdd26b15aa93c8c2f3dc52f8361", "postgres-v15": "6ee78a3c29e33cafd85ba09568b6b5eb031d29b9", "postgres-v14": "018fb052011081dc2733d3118d12e5c36df6eba1" } From 3bd2a4fd56803b0aabb87e9076872ceff0147a77 Mon Sep 17 00:00:00 2001 From: John Spray Date: Wed, 7 Feb 2024 19:14:18 +0000 Subject: [PATCH 37/72] control_plane: avoid feedback loop with /location_config if compute hook fails. (#6668) ## Problem The existing behavior isn't exactly incorrect, but is operationally risky: if the control plane compute hook breaks, then all the control plane operations trying to call /location_config will end up retrying forever, which could put more load on the system. ## Summary of changes - Treat 404s as fatal errors to do fewer retries: a 404 either indicates we have the wrong URL, or some control plane bug is failing to recognize our tenant ID as existing. - Do not return an error on reconcilation errors in a non-creating /location_config response: this allows the control plane to finish its Operation (and we will eventually retry the compute notification later) --- control_plane/attachment_service/src/compute_hook.rs | 2 +- control_plane/attachment_service/src/service.rs | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/control_plane/attachment_service/src/compute_hook.rs b/control_plane/attachment_service/src/compute_hook.rs index 0d3610aafa..5bd1b6bf09 100644 --- a/control_plane/attachment_service/src/compute_hook.rs +++ b/control_plane/attachment_service/src/compute_hook.rs @@ -240,7 +240,7 @@ impl ComputeHook { let client = reqwest::Client::new(); backoff::retry( || self.do_notify_iteration(&client, url, &reconfigure_request, cancel), - |e| matches!(e, NotifyError::Fatal(_)), + |e| matches!(e, NotifyError::Fatal(_) | NotifyError::Unexpected(_)), 3, 10, "Send compute notification", diff --git a/control_plane/attachment_service/src/service.rs b/control_plane/attachment_service/src/service.rs index febee1aa0d..1db1906df8 100644 --- a/control_plane/attachment_service/src/service.rs +++ b/control_plane/attachment_service/src/service.rs @@ -989,7 +989,15 @@ impl Service { .collect(); } else { // This was an update, wait for reconciliation - self.await_waiters(waiters).await?; + if let Err(e) = self.await_waiters(waiters).await { + // Do not treat a reconcile error as fatal: we have already applied any requested + // Intent changes, and the reconcile can fail for external reasons like unavailable + // compute notification API. In these cases, it is important that we do not + // cause the cloud control plane to retry forever on this API. + tracing::warn!( + "Failed to reconcile after /location_config: {e}, returning success anyway" + ); + } } Ok(result) From c561ad4e2e900409141e8c6c9963bab90288fd12 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Wed, 7 Feb 2024 20:39:52 +0100 Subject: [PATCH 38/72] feat: expose locked memory in pageserver `/metrics` (#6669) context: https://github.com/neondatabase/neon/issues/6667 --- Cargo.lock | 3 ++ Cargo.toml | 1 + libs/metrics/Cargo.toml | 3 ++ libs/metrics/src/lib.rs | 2 + libs/metrics/src/more_process_metrics.rs | 54 ++++++++++++++++++++++++ pageserver/src/bin/pageserver.rs | 2 + 6 files changed, 65 insertions(+) create mode 100644 libs/metrics/src/more_process_metrics.rs diff --git a/Cargo.lock b/Cargo.lock index a25725f90d..bf1ecfa89d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2869,6 +2869,7 @@ dependencies = [ "chrono", "libc", "once_cell", + "procfs", "prometheus", "rand 0.8.5", "rand_distr", @@ -3986,6 +3987,8 @@ checksum = "b1de8dacb0873f77e6aefc6d71e044761fcc68060290f5b1089fcdf84626bb69" dependencies = [ "bitflags 1.3.2", "byteorder", + "chrono", + "flate2", "hex", "lazy_static", "rustix 0.36.16", diff --git a/Cargo.toml b/Cargo.toml index 271edee742..6a2c3fa563 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,6 +113,7 @@ parquet = { version = "49.0.0", default-features = false, features = ["zstd"] } parquet_derive = "49.0.0" pbkdf2 = { version = "0.12.1", features = ["simple", "std"] } pin-project-lite = "0.2" +procfs = "0.14" prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency prost = "0.11" rand = "0.8" diff --git a/libs/metrics/Cargo.toml b/libs/metrics/Cargo.toml index a547d492df..f6a49a0166 100644 --- a/libs/metrics/Cargo.toml +++ b/libs/metrics/Cargo.toml @@ -13,6 +13,9 @@ twox-hash.workspace = true workspace_hack.workspace = true +[target.'cfg(target_os = "linux")'.dependencies] +procfs.workspace = true + [dev-dependencies] rand = "0.8" rand_distr = "0.4.3" diff --git a/libs/metrics/src/lib.rs b/libs/metrics/src/lib.rs index cb9914e5de..b57fd9f33b 100644 --- a/libs/metrics/src/lib.rs +++ b/libs/metrics/src/lib.rs @@ -31,6 +31,8 @@ pub use wrappers::{CountedReader, CountedWriter}; mod hll; pub mod metric_vec_duration; pub use hll::{HyperLogLog, HyperLogLogVec}; +#[cfg(target_os = "linux")] +pub mod more_process_metrics; pub type UIntGauge = GenericGauge; pub type UIntGaugeVec = GenericGaugeVec; diff --git a/libs/metrics/src/more_process_metrics.rs b/libs/metrics/src/more_process_metrics.rs new file mode 100644 index 0000000000..920724fdec --- /dev/null +++ b/libs/metrics/src/more_process_metrics.rs @@ -0,0 +1,54 @@ +//! process metrics that the [`::prometheus`] crate doesn't provide. + +// This module has heavy inspiration from the prometheus crate's `process_collector.rs`. + +use crate::UIntGauge; + +pub struct Collector { + descs: Vec, + vmlck: crate::UIntGauge, +} + +const NMETRICS: usize = 1; + +impl prometheus::core::Collector for Collector { + fn desc(&self) -> Vec<&prometheus::core::Desc> { + self.descs.iter().collect() + } + + fn collect(&self) -> Vec { + let Ok(myself) = procfs::process::Process::myself() else { + return vec![]; + }; + let mut mfs = Vec::with_capacity(NMETRICS); + if let Ok(status) = myself.status() { + if let Some(vmlck) = status.vmlck { + self.vmlck.set(vmlck); + mfs.extend(self.vmlck.collect()) + } + } + mfs + } +} + +impl Collector { + pub fn new() -> Self { + let mut descs = Vec::new(); + + let vmlck = + UIntGauge::new("libmetrics_process_status_vmlck", "/proc/self/status vmlck").unwrap(); + descs.extend( + prometheus::core::Collector::desc(&vmlck) + .into_iter() + .cloned(), + ); + + Self { descs, vmlck } + } +} + +impl Default for Collector { + fn default() -> Self { + Self::new() + } +} diff --git a/pageserver/src/bin/pageserver.rs b/pageserver/src/bin/pageserver.rs index eaddcb4607..7a93830c14 100644 --- a/pageserver/src/bin/pageserver.rs +++ b/pageserver/src/bin/pageserver.rs @@ -272,6 +272,8 @@ fn start_pageserver( ); set_build_info_metric(GIT_VERSION, BUILD_TAG); set_launch_timestamp_metric(launch_ts); + #[cfg(target_os = "linux")] + metrics::register_internal(Box::new(metrics::more_process_metrics::Collector::new())).unwrap(); pageserver::preinitialize_metrics(); // If any failpoints were set from FAILPOINTS environment variable, From 9a017778a9f89d5adfb6869a883ee2532dcaf13a Mon Sep 17 00:00:00 2001 From: Andreas Scherbaum Date: Thu, 8 Feb 2024 00:48:31 +0100 Subject: [PATCH 39/72] Update copyright notice, set it to current year (#6671) ## Problem Copyright notice is outdated ## Summary of changes Replace the initial year `2022` with `2022 - 2024`, after brief discussion with Stas about the format Co-authored-by: Andreas Scherbaum --- NOTICE | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NOTICE b/NOTICE index c13dc2f0b3..52fc751c41 100644 --- a/NOTICE +++ b/NOTICE @@ -1,5 +1,5 @@ Neon -Copyright 2022 Neon Inc. +Copyright 2022 - 2024 Neon Inc. The PostgreSQL submodules in vendor/ are licensed under the PostgreSQL license. See vendor/postgres-vX/COPYRIGHT for details. From c52495774d5151db63059515a524621660236f75 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Thu, 8 Feb 2024 00:58:54 +0100 Subject: [PATCH 40/72] tokio-epoll-uring: expose its metrics in pageserver's `/metrics` (#6672) context: https://github.com/neondatabase/neon/issues/6667 --- Cargo.lock | 4 +- pageserver/src/bin/pageserver.rs | 4 ++ pageserver/src/metrics.rs | 66 ++++++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bf1ecfa89d..30e233ecc1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5739,7 +5739,7 @@ dependencies = [ [[package]] name = "tokio-epoll-uring" version = "0.1.0" -source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#0e1af4ccddf2f01805cfc9eaefa97ee13c04b52d" +source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#d6a1c93442fb6b3a5bec490204961134e54925dc" dependencies = [ "futures", "nix 0.26.4", @@ -6264,7 +6264,7 @@ dependencies = [ [[package]] name = "uring-common" version = "0.1.0" -source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#0e1af4ccddf2f01805cfc9eaefa97ee13c04b52d" +source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#d6a1c93442fb6b3a5bec490204961134e54925dc" dependencies = [ "io-uring", "libc", diff --git a/pageserver/src/bin/pageserver.rs b/pageserver/src/bin/pageserver.rs index 7a93830c14..2f172bd384 100644 --- a/pageserver/src/bin/pageserver.rs +++ b/pageserver/src/bin/pageserver.rs @@ -274,6 +274,10 @@ fn start_pageserver( set_launch_timestamp_metric(launch_ts); #[cfg(target_os = "linux")] metrics::register_internal(Box::new(metrics::more_process_metrics::Collector::new())).unwrap(); + metrics::register_internal(Box::new( + pageserver::metrics::tokio_epoll_uring::Collector::new(), + )) + .unwrap(); pageserver::preinitialize_metrics(); // If any failpoints were set from FAILPOINTS environment variable, diff --git a/pageserver/src/metrics.rs b/pageserver/src/metrics.rs index 489ec58e62..98c98ef6e7 100644 --- a/pageserver/src/metrics.rs +++ b/pageserver/src/metrics.rs @@ -2400,6 +2400,72 @@ impl>, O, E> Future for MeasuredRemoteOp { } } +pub mod tokio_epoll_uring { + use metrics::UIntGauge; + + pub struct Collector { + descs: Vec, + systems_created: UIntGauge, + systems_destroyed: UIntGauge, + } + + const NMETRICS: usize = 2; + + impl metrics::core::Collector for Collector { + fn desc(&self) -> Vec<&metrics::core::Desc> { + self.descs.iter().collect() + } + + fn collect(&self) -> Vec { + let mut mfs = Vec::with_capacity(NMETRICS); + let tokio_epoll_uring::metrics::Metrics { + systems_created, + systems_destroyed, + } = tokio_epoll_uring::metrics::global(); + self.systems_created.set(systems_created); + mfs.extend(self.systems_created.collect()); + self.systems_destroyed.set(systems_destroyed); + mfs.extend(self.systems_destroyed.collect()); + mfs + } + } + + impl Collector { + #[allow(clippy::new_without_default)] + pub fn new() -> Self { + let mut descs = Vec::new(); + + let systems_created = UIntGauge::new( + "pageserver_tokio_epoll_uring_systems_created", + "counter of tokio-epoll-uring systems that were created", + ) + .unwrap(); + descs.extend( + metrics::core::Collector::desc(&systems_created) + .into_iter() + .cloned(), + ); + + let systems_destroyed = UIntGauge::new( + "pageserver_tokio_epoll_uring_systems_destroyed", + "counter of tokio-epoll-uring systems that were destroyed", + ) + .unwrap(); + descs.extend( + metrics::core::Collector::desc(&systems_destroyed) + .into_iter() + .cloned(), + ); + + Self { + descs, + systems_created, + systems_destroyed, + } + } + } +} + pub fn preinitialize_metrics() { // Python tests need these and on some we do alerting. // From c63e3e7e84c2dd9c9792619cc4fee15b07cfe7d7 Mon Sep 17 00:00:00 2001 From: Anna Khanova <32508607+khanova@users.noreply.github.com> Date: Thu, 8 Feb 2024 12:57:05 +0100 Subject: [PATCH 41/72] Proxy: improve http-pool (#6577) ## Problem The password check logic for the sql-over-http is a bit non-intuitive. ## Summary of changes 1. Perform scram auth using the same logic as for websocket cleartext password. 2. Split establish connection logic and connection pool. 3. Parallelize param parsing logic with authentication + wake compute. 4. Limit the total number of clients --- Cargo.lock | 1 + proxy/Cargo.toml | 1 + proxy/src/auth/backend.rs | 12 + proxy/src/auth/flow.rs | 2 +- proxy/src/bin/proxy.rs | 5 + proxy/src/console/provider/neon.rs | 2 + proxy/src/context.rs | 4 + proxy/src/metrics.rs | 44 +- proxy/src/proxy/connect_compute.rs | 22 +- proxy/src/proxy/tests.rs | 3 + proxy/src/serverless.rs | 41 +- proxy/src/serverless/backend.rs | 157 +++++ proxy/src/serverless/conn_pool.rs | 797 +++++++++++++------------- proxy/src/serverless/json.rs | 28 +- proxy/src/serverless/sql_over_http.rs | 92 ++- test_runner/regress/test_proxy.py | 20 +- 16 files changed, 753 insertions(+), 478 deletions(-) create mode 100644 proxy/src/serverless/backend.rs diff --git a/Cargo.lock b/Cargo.lock index 30e233ecc1..c0c319cd89 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4079,6 +4079,7 @@ dependencies = [ "clap", "consumption_metrics", "dashmap", + "env_logger", "futures", "git-version", "hashbrown 0.13.2", diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index 1247f08ee6..83cab381b3 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -19,6 +19,7 @@ chrono.workspace = true clap.workspace = true consumption_metrics.workspace = true dashmap.workspace = true +env_logger.workspace = true futures.workspace = true git-version.workspace = true hashbrown.workspace = true diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs index 236567163e..fa2782bee3 100644 --- a/proxy/src/auth/backend.rs +++ b/proxy/src/auth/backend.rs @@ -68,6 +68,7 @@ pub trait TestBackend: Send + Sync + 'static { fn get_allowed_ips_and_secret( &self, ) -> Result<(CachedAllowedIps, Option), console::errors::GetAuthInfoError>; + fn get_role_secret(&self) -> Result; } impl std::fmt::Display for BackendType<'_, ()> { @@ -358,6 +359,17 @@ impl<'a> BackendType<'a, ComputeUserInfoMaybeEndpoint> { } impl BackendType<'_, ComputeUserInfo> { + pub async fn get_role_secret( + &self, + ctx: &mut RequestMonitoring, + ) -> Result { + use BackendType::*; + match self { + Console(api, user_info) => api.get_role_secret(ctx, user_info).await, + Link(_) => Ok(Cached::new_uncached(None)), + } + } + pub async fn get_allowed_ips_and_secret( &self, ctx: &mut RequestMonitoring, diff --git a/proxy/src/auth/flow.rs b/proxy/src/auth/flow.rs index 077178d107..c2783e236c 100644 --- a/proxy/src/auth/flow.rs +++ b/proxy/src/auth/flow.rs @@ -167,7 +167,7 @@ impl AuthFlow<'_, S, Scram<'_>> { } } -pub(super) fn validate_password_and_exchange( +pub(crate) fn validate_password_and_exchange( password: &[u8], secret: AuthSecret, ) -> super::Result> { diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index 3bbb87808d..6974f1a274 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -165,6 +165,10 @@ struct SqlOverHttpArgs { #[clap(long, default_value_t = 20)] sql_over_http_pool_max_conns_per_endpoint: usize, + /// How many connections to pool for each endpoint. Excess connections are discarded + #[clap(long, default_value_t = 20000)] + sql_over_http_pool_max_total_conns: usize, + /// How long pooled connections should remain idle for before closing #[clap(long, default_value = "5m", value_parser = humantime::parse_duration)] sql_over_http_idle_timeout: tokio::time::Duration, @@ -387,6 +391,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> { pool_shards: args.sql_over_http.sql_over_http_pool_shards, idle_timeout: args.sql_over_http.sql_over_http_idle_timeout, opt_in: args.sql_over_http.sql_over_http_pool_opt_in, + max_total_conns: args.sql_over_http.sql_over_http_pool_max_total_conns, }, }; let authentication_config = AuthenticationConfig { diff --git a/proxy/src/console/provider/neon.rs b/proxy/src/console/provider/neon.rs index 0785419790..71b34cb676 100644 --- a/proxy/src/console/provider/neon.rs +++ b/proxy/src/console/provider/neon.rs @@ -188,6 +188,7 @@ impl super::Api for Api { ep, Arc::new(auth_info.allowed_ips), ); + ctx.set_project_id(project_id); } // When we just got a secret, we don't need to invalidate it. Ok(Cached::new_uncached(auth_info.secret)) @@ -221,6 +222,7 @@ impl super::Api for Api { self.caches .project_info .insert_allowed_ips(&project_id, ep, allowed_ips.clone()); + ctx.set_project_id(project_id); } Ok(( Cached::new_uncached(allowed_ips), diff --git a/proxy/src/context.rs b/proxy/src/context.rs index e2b0294cd3..fe204534b7 100644 --- a/proxy/src/context.rs +++ b/proxy/src/context.rs @@ -89,6 +89,10 @@ impl RequestMonitoring { self.project = Some(x.project_id); } + pub fn set_project_id(&mut self, project_id: ProjectId) { + self.project = Some(project_id); + } + pub fn set_endpoint_id(&mut self, endpoint_id: EndpointId) { crate::metrics::CONNECTING_ENDPOINTS .with_label_values(&[self.protocol]) diff --git a/proxy/src/metrics.rs b/proxy/src/metrics.rs index fa663d8ff6..e2d96a9c27 100644 --- a/proxy/src/metrics.rs +++ b/proxy/src/metrics.rs @@ -1,8 +1,10 @@ use ::metrics::{ exponential_buckets, register_histogram, register_histogram_vec, register_hll_vec, - register_int_counter_pair_vec, register_int_counter_vec, register_int_gauge_vec, Histogram, - HistogramVec, HyperLogLogVec, IntCounterPairVec, IntCounterVec, IntGaugeVec, + register_int_counter_pair_vec, register_int_counter_vec, register_int_gauge, + register_int_gauge_vec, Histogram, HistogramVec, HyperLogLogVec, IntCounterPairVec, + IntCounterVec, IntGauge, IntGaugeVec, }; +use metrics::{register_int_counter_pair, IntCounterPair}; use once_cell::sync::Lazy; use tokio::time; @@ -112,6 +114,44 @@ pub static ALLOWED_IPS_NUMBER: Lazy = Lazy::new(|| { .unwrap() }); +pub static HTTP_CONTENT_LENGTH: Lazy = Lazy::new(|| { + register_histogram!( + "proxy_http_conn_content_length_bytes", + "Time it took for proxy to establish a connection to the compute endpoint", + // largest bucket = 3^16 * 0.05ms = 2.15s + exponential_buckets(8.0, 2.0, 20).unwrap() + ) + .unwrap() +}); + +pub static GC_LATENCY: Lazy = Lazy::new(|| { + register_histogram!( + "proxy_http_pool_reclaimation_lag_seconds", + "Time it takes to reclaim unused connection pools", + // 1us -> 65ms + exponential_buckets(1e-6, 2.0, 16).unwrap(), + ) + .unwrap() +}); + +pub static ENDPOINT_POOLS: Lazy = Lazy::new(|| { + register_int_counter_pair!( + "proxy_http_pool_endpoints_registered_total", + "Number of endpoints we have registered pools for", + "proxy_http_pool_endpoints_unregistered_total", + "Number of endpoints we have unregistered pools for", + ) + .unwrap() +}); + +pub static NUM_OPEN_CLIENTS_IN_HTTP_POOL: Lazy = Lazy::new(|| { + register_int_gauge!( + "proxy_http_pool_opened_connections", + "Number of opened connections to a database.", + ) + .unwrap() +}); + #[derive(Clone)] pub struct LatencyTimer { // time since the stopwatch was started diff --git a/proxy/src/proxy/connect_compute.rs b/proxy/src/proxy/connect_compute.rs index 58c59dba36..b9346aa743 100644 --- a/proxy/src/proxy/connect_compute.rs +++ b/proxy/src/proxy/connect_compute.rs @@ -34,21 +34,6 @@ pub fn invalidate_cache(node_info: console::CachedNodeInfo) -> compute::ConnCfg node_info.invalidate().config } -/// Try to connect to the compute node once. -#[tracing::instrument(name = "connect_once", fields(pid = tracing::field::Empty), skip_all)] -async fn connect_to_compute_once( - ctx: &mut RequestMonitoring, - node_info: &console::CachedNodeInfo, - timeout: time::Duration, -) -> Result { - let allow_self_signed_compute = node_info.allow_self_signed_compute; - - node_info - .config - .connect(ctx, allow_self_signed_compute, timeout) - .await -} - #[async_trait] pub trait ConnectMechanism { type Connection; @@ -75,13 +60,18 @@ impl ConnectMechanism for TcpMechanism<'_> { type ConnectError = compute::ConnectionError; type Error = compute::ConnectionError; + #[tracing::instrument(fields(pid = tracing::field::Empty), skip_all)] async fn connect_once( &self, ctx: &mut RequestMonitoring, node_info: &console::CachedNodeInfo, timeout: time::Duration, ) -> Result { - connect_to_compute_once(ctx, node_info, timeout).await + let allow_self_signed_compute = node_info.allow_self_signed_compute; + node_info + .config + .connect(ctx, allow_self_signed_compute, timeout) + .await } fn update_connect_config(&self, config: &mut compute::ConnCfg) { diff --git a/proxy/src/proxy/tests.rs b/proxy/src/proxy/tests.rs index 2000774224..656cabac75 100644 --- a/proxy/src/proxy/tests.rs +++ b/proxy/src/proxy/tests.rs @@ -478,6 +478,9 @@ impl TestBackend for TestConnectMechanism { { unimplemented!("not used in tests") } + fn get_role_secret(&self) -> Result { + unimplemented!("not used in tests") + } } fn helper_create_cached_node_info() -> CachedNodeInfo { diff --git a/proxy/src/serverless.rs b/proxy/src/serverless.rs index 7ff93b23b8..58aa925a6a 100644 --- a/proxy/src/serverless.rs +++ b/proxy/src/serverless.rs @@ -2,6 +2,7 @@ //! //! Handles both SQL over HTTP and SQL over Websockets. +mod backend; mod conn_pool; mod json; mod sql_over_http; @@ -18,11 +19,11 @@ pub use reqwest_middleware::{ClientWithMiddleware, Error}; pub use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware}; use tokio_util::task::TaskTracker; -use crate::config::TlsConfig; use crate::context::RequestMonitoring; use crate::metrics::NUM_CLIENT_CONNECTION_GAUGE; use crate::protocol2::{ProxyProtocolAccept, WithClientIp}; use crate::rate_limiter::EndpointRateLimiter; +use crate::serverless::backend::PoolingBackend; use crate::{cancellation::CancelMap, config::ProxyConfig}; use futures::StreamExt; use hyper::{ @@ -54,12 +55,13 @@ pub async fn task_main( info!("websocket server has shut down"); } - let conn_pool = conn_pool::GlobalConnPool::new(config); - - let conn_pool2 = Arc::clone(&conn_pool); - tokio::spawn(async move { - conn_pool2.gc_worker(StdRng::from_entropy()).await; - }); + let conn_pool = conn_pool::GlobalConnPool::new(&config.http_config); + { + let conn_pool = Arc::clone(&conn_pool); + tokio::spawn(async move { + conn_pool.gc_worker(StdRng::from_entropy()).await; + }); + } // shutdown the connection pool tokio::spawn({ @@ -73,6 +75,11 @@ pub async fn task_main( } }); + let backend = Arc::new(PoolingBackend { + pool: Arc::clone(&conn_pool), + config, + }); + let tls_config = match config.tls_config.as_ref() { Some(config) => config, None => { @@ -106,7 +113,7 @@ pub async fn task_main( let client_addr = io.client_addr(); let remote_addr = io.inner.remote_addr(); let sni_name = tls.server_name().map(|s| s.to_string()); - let conn_pool = conn_pool.clone(); + let backend = backend.clone(); let ws_connections = ws_connections.clone(); let endpoint_rate_limiter = endpoint_rate_limiter.clone(); @@ -119,7 +126,7 @@ pub async fn task_main( Ok(MetricService::new(hyper::service::service_fn( move |req: Request| { let sni_name = sni_name.clone(); - let conn_pool = conn_pool.clone(); + let backend = backend.clone(); let ws_connections = ws_connections.clone(); let endpoint_rate_limiter = endpoint_rate_limiter.clone(); @@ -130,8 +137,7 @@ pub async fn task_main( request_handler( req, config, - tls_config, - conn_pool, + backend, ws_connections, cancel_map, session_id, @@ -200,8 +206,7 @@ where async fn request_handler( mut request: Request, config: &'static ProxyConfig, - tls: &'static TlsConfig, - conn_pool: Arc, + backend: Arc, ws_connections: TaskTracker, cancel_map: Arc, session_id: uuid::Uuid, @@ -248,15 +253,7 @@ async fn request_handler( } else if request.uri().path() == "/sql" && request.method() == Method::POST { let mut ctx = RequestMonitoring::new(session_id, peer_addr, "http", &config.region); - sql_over_http::handle( - tls, - &config.http_config, - &mut ctx, - request, - sni_hostname, - conn_pool, - ) - .await + sql_over_http::handle(config, &mut ctx, request, sni_hostname, backend).await } else if request.uri().path() == "/sql" && request.method() == Method::OPTIONS { Response::builder() .header("Allow", "OPTIONS, POST") diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs new file mode 100644 index 0000000000..466a74f0ea --- /dev/null +++ b/proxy/src/serverless/backend.rs @@ -0,0 +1,157 @@ +use std::{sync::Arc, time::Duration}; + +use anyhow::Context; +use async_trait::async_trait; +use tracing::info; + +use crate::{ + auth::{backend::ComputeCredentialKeys, check_peer_addr_is_in_list, AuthError}, + compute, + config::ProxyConfig, + console::CachedNodeInfo, + context::RequestMonitoring, + proxy::connect_compute::ConnectMechanism, +}; + +use super::conn_pool::{poll_client, Client, ConnInfo, GlobalConnPool, APP_NAME}; + +pub struct PoolingBackend { + pub pool: Arc>, + pub config: &'static ProxyConfig, +} + +impl PoolingBackend { + pub async fn authenticate( + &self, + ctx: &mut RequestMonitoring, + conn_info: &ConnInfo, + ) -> Result { + let user_info = conn_info.user_info.clone(); + let backend = self.config.auth_backend.as_ref().map(|_| user_info.clone()); + let (allowed_ips, maybe_secret) = backend.get_allowed_ips_and_secret(ctx).await?; + if !check_peer_addr_is_in_list(&ctx.peer_addr, &allowed_ips) { + return Err(AuthError::ip_address_not_allowed()); + } + let cached_secret = match maybe_secret { + Some(secret) => secret, + None => backend.get_role_secret(ctx).await?, + }; + + let secret = match cached_secret.value.clone() { + Some(secret) => secret, + None => { + // If we don't have an authentication secret, for the http flow we can just return an error. + info!("authentication info not found"); + return Err(AuthError::auth_failed(&*user_info.user)); + } + }; + let auth_outcome = + crate::auth::validate_password_and_exchange(conn_info.password.as_bytes(), secret)?; + match auth_outcome { + crate::sasl::Outcome::Success(key) => Ok(key), + crate::sasl::Outcome::Failure(reason) => { + info!("auth backend failed with an error: {reason}"); + Err(AuthError::auth_failed(&*conn_info.user_info.user)) + } + } + } + + // Wake up the destination if needed. Code here is a bit involved because + // we reuse the code from the usual proxy and we need to prepare few structures + // that this code expects. + #[tracing::instrument(fields(pid = tracing::field::Empty), skip_all)] + pub async fn connect_to_compute( + &self, + ctx: &mut RequestMonitoring, + conn_info: ConnInfo, + keys: ComputeCredentialKeys, + force_new: bool, + ) -> anyhow::Result> { + let maybe_client = if !force_new { + info!("pool: looking for an existing connection"); + self.pool.get(ctx, &conn_info).await? + } else { + info!("pool: pool is disabled"); + None + }; + + if let Some(client) = maybe_client { + return Ok(client); + } + let conn_id = uuid::Uuid::new_v4(); + info!(%conn_id, "pool: opening a new connection '{conn_info}'"); + ctx.set_application(Some(APP_NAME)); + let backend = self + .config + .auth_backend + .as_ref() + .map(|_| conn_info.user_info.clone()); + + let mut node_info = backend + .wake_compute(ctx) + .await? + .context("missing cache entry from wake_compute")?; + + match keys { + #[cfg(any(test, feature = "testing"))] + ComputeCredentialKeys::Password(password) => node_info.config.password(password), + ComputeCredentialKeys::AuthKeys(auth_keys) => node_info.config.auth_keys(auth_keys), + }; + + ctx.set_project(node_info.aux.clone()); + + crate::proxy::connect_compute::connect_to_compute( + ctx, + &TokioMechanism { + conn_id, + conn_info, + pool: self.pool.clone(), + }, + node_info, + &backend, + ) + .await + } +} + +struct TokioMechanism { + pool: Arc>, + conn_info: ConnInfo, + conn_id: uuid::Uuid, +} + +#[async_trait] +impl ConnectMechanism for TokioMechanism { + type Connection = Client; + type ConnectError = tokio_postgres::Error; + type Error = anyhow::Error; + + async fn connect_once( + &self, + ctx: &mut RequestMonitoring, + node_info: &CachedNodeInfo, + timeout: Duration, + ) -> Result { + let mut config = (*node_info.config).clone(); + let config = config + .user(&self.conn_info.user_info.user) + .password(&*self.conn_info.password) + .dbname(&self.conn_info.dbname) + .connect_timeout(timeout); + + let (client, connection) = config.connect(tokio_postgres::NoTls).await?; + + tracing::Span::current().record("pid", &tracing::field::display(client.get_process_id())); + Ok(poll_client( + self.pool.clone(), + ctx, + self.conn_info.clone(), + client, + connection, + self.conn_id, + node_info.aux.clone(), + )) + } + + fn update_connect_config(&self, _config: &mut compute::ConnCfg) {} +} diff --git a/proxy/src/serverless/conn_pool.rs b/proxy/src/serverless/conn_pool.rs index 312fa2b36f..a7b2c532d2 100644 --- a/proxy/src/serverless/conn_pool.rs +++ b/proxy/src/serverless/conn_pool.rs @@ -1,15 +1,7 @@ -use anyhow::Context; -use async_trait::async_trait; use dashmap::DashMap; use futures::{future::poll_fn, Future}; -use metrics::{register_int_counter_pair, IntCounterPair, IntCounterPairGuard}; -use once_cell::sync::Lazy; +use metrics::IntCounterPairGuard; use parking_lot::RwLock; -use pbkdf2::{ - password_hash::{PasswordHashString, PasswordHasher, PasswordVerifier, SaltString}, - Params, Pbkdf2, -}; -use prometheus::{exponential_buckets, register_histogram, Histogram}; use rand::Rng; use smol_str::SmolStr; use std::{collections::HashMap, pin::pin, sync::Arc, sync::Weak, time::Duration}; @@ -21,19 +13,17 @@ use std::{ ops::Deref, sync::atomic::{self, AtomicUsize}, }; -use tokio::time::{self, Instant}; -use tokio_postgres::{AsyncMessage, ReadyForQueryStatus}; +use tokio::time::Instant; +use tokio_postgres::tls::NoTlsStream; +use tokio_postgres::{AsyncMessage, ReadyForQueryStatus, Socket}; +use crate::console::messages::MetricsAuxInfo; +use crate::metrics::{ENDPOINT_POOLS, GC_LATENCY, NUM_OPEN_CLIENTS_IN_HTTP_POOL}; +use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS}; use crate::{ - auth::{self, backend::ComputeUserInfo, check_peer_addr_is_in_list}, - console::{self, messages::MetricsAuxInfo}, - context::RequestMonitoring, - metrics::NUM_DB_CONNECTIONS_GAUGE, - proxy::connect_compute::ConnectMechanism, - usage_metrics::{Ids, MetricCounter, USAGE_METRICS}, + auth::backend::ComputeUserInfo, context::RequestMonitoring, metrics::NUM_DB_CONNECTIONS_GAUGE, DbName, EndpointCacheKey, RoleName, }; -use crate::{compute, config}; use tracing::{debug, error, warn, Span}; use tracing::{info, info_span, Instrument}; @@ -72,39 +62,51 @@ impl fmt::Display for ConnInfo { } } -struct ConnPoolEntry { - conn: ClientInner, +struct ConnPoolEntry { + conn: ClientInner, _last_access: std::time::Instant, } // Per-endpoint connection pool, (dbname, username) -> DbUserConnPool // Number of open connections is limited by the `max_conns_per_endpoint`. -pub struct EndpointConnPool { - pools: HashMap<(DbName, RoleName), DbUserConnPool>, +pub struct EndpointConnPool { + pools: HashMap<(DbName, RoleName), DbUserConnPool>, total_conns: usize, max_conns: usize, _guard: IntCounterPairGuard, + global_connections_count: Arc, + global_pool_size_max_conns: usize, } -impl EndpointConnPool { - fn get_conn_entry(&mut self, db_user: (DbName, RoleName)) -> Option { +impl EndpointConnPool { + fn get_conn_entry(&mut self, db_user: (DbName, RoleName)) -> Option> { let Self { - pools, total_conns, .. + pools, + total_conns, + global_connections_count, + .. } = self; - pools - .get_mut(&db_user) - .and_then(|pool_entries| pool_entries.get_conn_entry(total_conns)) + pools.get_mut(&db_user).and_then(|pool_entries| { + pool_entries.get_conn_entry(total_conns, global_connections_count.clone()) + }) } fn remove_client(&mut self, db_user: (DbName, RoleName), conn_id: uuid::Uuid) -> bool { let Self { - pools, total_conns, .. + pools, + total_conns, + global_connections_count, + .. } = self; if let Some(pool) = pools.get_mut(&db_user) { let old_len = pool.conns.len(); pool.conns.retain(|conn| conn.conn.conn_id != conn_id); let new_len = pool.conns.len(); let removed = old_len - new_len; + if removed > 0 { + global_connections_count.fetch_sub(removed, atomic::Ordering::Relaxed); + NUM_OPEN_CLIENTS_IN_HTTP_POOL.sub(removed as i64); + } *total_conns -= removed; removed > 0 } else { @@ -112,13 +114,27 @@ impl EndpointConnPool { } } - fn put(pool: &RwLock, conn_info: &ConnInfo, client: ClientInner) -> anyhow::Result<()> { + fn put( + pool: &RwLock, + conn_info: &ConnInfo, + client: ClientInner, + ) -> anyhow::Result<()> { let conn_id = client.conn_id; - if client.inner.is_closed() { + if client.is_closed() { info!(%conn_id, "pool: throwing away connection '{conn_info}' because connection is closed"); return Ok(()); } + let global_max_conn = pool.read().global_pool_size_max_conns; + if pool + .read() + .global_connections_count + .load(atomic::Ordering::Relaxed) + >= global_max_conn + { + info!(%conn_id, "pool: throwing away connection '{conn_info}' because pool is full"); + return Ok(()); + } // return connection to the pool let mut returned = false; @@ -127,18 +143,19 @@ impl EndpointConnPool { let mut pool = pool.write(); if pool.total_conns < pool.max_conns { - // we create this db-user entry in get, so it should not be None - if let Some(pool_entries) = pool.pools.get_mut(&conn_info.db_and_user()) { - pool_entries.conns.push(ConnPoolEntry { - conn: client, - _last_access: std::time::Instant::now(), - }); + let pool_entries = pool.pools.entry(conn_info.db_and_user()).or_default(); + pool_entries.conns.push(ConnPoolEntry { + conn: client, + _last_access: std::time::Instant::now(), + }); - returned = true; - per_db_size = pool_entries.conns.len(); + returned = true; + per_db_size = pool_entries.conns.len(); - pool.total_conns += 1; - } + pool.total_conns += 1; + pool.global_connections_count + .fetch_add(1, atomic::Ordering::Relaxed); + NUM_OPEN_CLIENTS_IN_HTTP_POOL.inc(); } pool.total_conns @@ -155,49 +172,61 @@ impl EndpointConnPool { } } -/// 4096 is the number of rounds that SCRAM-SHA-256 recommends. -/// It's not the 600,000 that OWASP recommends... but our passwords are high entropy anyway. -/// -/// Still takes 1.4ms to hash on my hardware. -/// We don't want to ruin the latency improvements of using the pool by making password verification take too long -const PARAMS: Params = Params { - rounds: 4096, - output_length: 32, -}; - -#[derive(Default)] -pub struct DbUserConnPool { - conns: Vec, - password_hash: Option, +impl Drop for EndpointConnPool { + fn drop(&mut self) { + if self.total_conns > 0 { + self.global_connections_count + .fetch_sub(self.total_conns, atomic::Ordering::Relaxed); + NUM_OPEN_CLIENTS_IN_HTTP_POOL.sub(self.total_conns as i64); + } + } } -impl DbUserConnPool { - fn clear_closed_clients(&mut self, conns: &mut usize) { +pub struct DbUserConnPool { + conns: Vec>, +} + +impl Default for DbUserConnPool { + fn default() -> Self { + Self { conns: Vec::new() } + } +} + +impl DbUserConnPool { + fn clear_closed_clients(&mut self, conns: &mut usize) -> usize { let old_len = self.conns.len(); - self.conns.retain(|conn| !conn.conn.inner.is_closed()); + self.conns.retain(|conn| !conn.conn.is_closed()); let new_len = self.conns.len(); let removed = old_len - new_len; *conns -= removed; + removed } - fn get_conn_entry(&mut self, conns: &mut usize) -> Option { - self.clear_closed_clients(conns); + fn get_conn_entry( + &mut self, + conns: &mut usize, + global_connections_count: Arc, + ) -> Option> { + let mut removed = self.clear_closed_clients(conns); let conn = self.conns.pop(); if conn.is_some() { *conns -= 1; + removed += 1; } + global_connections_count.fetch_sub(removed, atomic::Ordering::Relaxed); + NUM_OPEN_CLIENTS_IN_HTTP_POOL.sub(removed as i64); conn } } -pub struct GlobalConnPool { +pub struct GlobalConnPool { // endpoint -> per-endpoint connection pool // // That should be a fairly conteded map, so return reference to the per-endpoint // pool as early as possible and release the lock. - global_pool: DashMap>>, + global_pool: DashMap>>>, /// Number of endpoint-connection pools /// @@ -206,7 +235,10 @@ pub struct GlobalConnPool { /// It's only used for diagnostics. global_pool_size: AtomicUsize, - proxy_config: &'static crate::config::ProxyConfig, + /// Total number of connections in the pool + global_connections_count: Arc, + + config: &'static crate::config::HttpConfig, } #[derive(Debug, Clone, Copy)] @@ -224,45 +256,39 @@ pub struct GlobalConnPoolOptions { pub idle_timeout: Duration, pub opt_in: bool, + + // Total number of connections in the pool. + pub max_total_conns: usize, } -pub static GC_LATENCY: Lazy = Lazy::new(|| { - register_histogram!( - "proxy_http_pool_reclaimation_lag_seconds", - "Time it takes to reclaim unused connection pools", - // 1us -> 65ms - exponential_buckets(1e-6, 2.0, 16).unwrap(), - ) - .unwrap() -}); - -pub static ENDPOINT_POOLS: Lazy = Lazy::new(|| { - register_int_counter_pair!( - "proxy_http_pool_endpoints_registered_total", - "Number of endpoints we have registered pools for", - "proxy_http_pool_endpoints_unregistered_total", - "Number of endpoints we have unregistered pools for", - ) - .unwrap() -}); - -impl GlobalConnPool { - pub fn new(config: &'static crate::config::ProxyConfig) -> Arc { - let shards = config.http_config.pool_options.pool_shards; +impl GlobalConnPool { + pub fn new(config: &'static crate::config::HttpConfig) -> Arc { + let shards = config.pool_options.pool_shards; Arc::new(Self { global_pool: DashMap::with_shard_amount(shards), global_pool_size: AtomicUsize::new(0), - proxy_config: config, + config, + global_connections_count: Arc::new(AtomicUsize::new(0)), }) } + #[cfg(test)] + pub fn get_global_connections_count(&self) -> usize { + self.global_connections_count + .load(atomic::Ordering::Relaxed) + } + + pub fn get_idle_timeout(&self) -> Duration { + self.config.pool_options.idle_timeout + } + pub fn shutdown(&self) { // drops all strong references to endpoint-pools self.global_pool.clear(); } pub async fn gc_worker(&self, mut rng: impl Rng) { - let epoch = self.proxy_config.http_config.pool_options.gc_epoch; + let epoch = self.config.pool_options.gc_epoch; let mut interval = tokio::time::interval(epoch / (self.global_pool.shards().len()) as u32); loop { interval.tick().await; @@ -280,6 +306,7 @@ impl GlobalConnPool { let timer = GC_LATENCY.start_timer(); let current_len = shard.len(); + let mut clients_removed = 0; shard.retain(|endpoint, x| { // if the current endpoint pool is unique (no other strong or weak references) // then it is currently not in use by any connections. @@ -289,9 +316,9 @@ impl GlobalConnPool { } = pool.get_mut(); // ensure that closed clients are removed - pools - .iter_mut() - .for_each(|(_, db_pool)| db_pool.clear_closed_clients(total_conns)); + pools.iter_mut().for_each(|(_, db_pool)| { + clients_removed += db_pool.clear_closed_clients(total_conns); + }); // we only remove this pool if it has no active connections if *total_conns == 0 { @@ -302,10 +329,20 @@ impl GlobalConnPool { true }); + let new_len = shard.len(); drop(shard); timer.observe_duration(); + // Do logging outside of the lock. + if clients_removed > 0 { + let size = self + .global_connections_count + .fetch_sub(clients_removed, atomic::Ordering::Relaxed) + - clients_removed; + NUM_OPEN_CLIENTS_IN_HTTP_POOL.sub(clients_removed as i64); + info!("pool: performed global pool gc. removed {clients_removed} clients, total number of clients in pool is {size}"); + } let removed = current_len - new_len; if removed > 0 { @@ -320,61 +357,24 @@ impl GlobalConnPool { pub async fn get( self: &Arc, ctx: &mut RequestMonitoring, - conn_info: ConnInfo, - force_new: bool, - ) -> anyhow::Result { - let mut client: Option = None; + conn_info: &ConnInfo, + ) -> anyhow::Result>> { + let mut client: Option> = None; - let mut hash_valid = false; - let mut endpoint_pool = Weak::new(); - if !force_new { - let pool = self.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key()); - endpoint_pool = Arc::downgrade(&pool); - let mut hash = None; - - // find a pool entry by (dbname, username) if exists - { - let pool = pool.read(); - if let Some(pool_entries) = pool.pools.get(&conn_info.db_and_user()) { - if !pool_entries.conns.is_empty() { - hash = pool_entries.password_hash.clone(); - } - } - } - - // a connection exists in the pool, verify the password hash - if let Some(hash) = hash { - let pw = conn_info.password.clone(); - let validate = tokio::task::spawn_blocking(move || { - Pbkdf2.verify_password(pw.as_bytes(), &hash.password_hash()) - }) - .await?; - - // if the hash is invalid, don't error - // we will continue with the regular connection flow - if validate.is_ok() { - hash_valid = true; - if let Some(entry) = pool.write().get_conn_entry(conn_info.db_and_user()) { - client = Some(entry.conn) - } - } - } + let endpoint_pool = self.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key()); + if let Some(entry) = endpoint_pool + .write() + .get_conn_entry(conn_info.db_and_user()) + { + client = Some(entry.conn) } + let endpoint_pool = Arc::downgrade(&endpoint_pool); // ok return cached connection if found and establish a new one otherwise - let new_client = if let Some(client) = client { - ctx.set_project(client.aux.clone()); - if client.inner.is_closed() { - let conn_id = uuid::Uuid::new_v4(); - info!(%conn_id, "pool: cached connection '{conn_info}' is closed, opening a new one"); - connect_to_compute( - self.proxy_config, - ctx, - &conn_info, - conn_id, - endpoint_pool.clone(), - ) - .await + if let Some(client) = client { + if client.is_closed() { + info!("pool: cached connection '{conn_info}' is closed, opening a new one"); + return Ok(None); } else { info!("pool: reusing connection '{conn_info}'"); client.session.send(ctx.session_id)?; @@ -384,67 +384,16 @@ impl GlobalConnPool { ); ctx.latency_timer.pool_hit(); ctx.latency_timer.success(); - return Ok(Client::new(client, conn_info, endpoint_pool).await); + return Ok(Some(Client::new(client, conn_info.clone(), endpoint_pool))); } - } else { - let conn_id = uuid::Uuid::new_v4(); - info!(%conn_id, "pool: opening a new connection '{conn_info}'"); - connect_to_compute( - self.proxy_config, - ctx, - &conn_info, - conn_id, - endpoint_pool.clone(), - ) - .await - }; - if let Ok(client) = &new_client { - tracing::Span::current().record( - "pid", - &tracing::field::display(client.inner.get_process_id()), - ); } - - match &new_client { - // clear the hash. it's no longer valid - // TODO: update tokio-postgres fork to allow access to this error kind directly - Err(err) - if hash_valid && err.to_string().contains("password authentication failed") => - { - let pool = self.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key()); - let mut pool = pool.write(); - if let Some(entry) = pool.pools.get_mut(&conn_info.db_and_user()) { - entry.password_hash = None; - } - } - // new password is valid and we should insert/update it - Ok(_) if !force_new && !hash_valid => { - let pw = conn_info.password.clone(); - let new_hash = tokio::task::spawn_blocking(move || { - let salt = SaltString::generate(rand::rngs::OsRng); - Pbkdf2 - .hash_password_customized(pw.as_bytes(), None, None, PARAMS, &salt) - .map(|s| s.serialize()) - }) - .await??; - - let pool = self.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key()); - let mut pool = pool.write(); - pool.pools - .entry(conn_info.db_and_user()) - .or_default() - .password_hash = Some(new_hash); - } - _ => {} - } - let new_client = new_client?; - Ok(Client::new(new_client, conn_info, endpoint_pool).await) + Ok(None) } fn get_or_create_endpoint_pool( - &self, + self: &Arc, endpoint: &EndpointCacheKey, - ) -> Arc> { + ) -> Arc>> { // fast path if let Some(pool) = self.global_pool.get(endpoint) { return pool.clone(); @@ -454,12 +403,10 @@ impl GlobalConnPool { let new_pool = Arc::new(RwLock::new(EndpointConnPool { pools: HashMap::new(), total_conns: 0, - max_conns: self - .proxy_config - .http_config - .pool_options - .max_conns_per_endpoint, + max_conns: self.config.pool_options.max_conns_per_endpoint, _guard: ENDPOINT_POOLS.guard(), + global_connections_count: self.global_connections_count.clone(), + global_pool_size_max_conns: self.config.pool_options.max_total_conns, })); // find or create a pool for this endpoint @@ -488,196 +435,128 @@ impl GlobalConnPool { } } -struct TokioMechanism<'a> { - pool: Weak>, - conn_info: &'a ConnInfo, - conn_id: uuid::Uuid, - idle: Duration, -} - -#[async_trait] -impl ConnectMechanism for TokioMechanism<'_> { - type Connection = ClientInner; - type ConnectError = tokio_postgres::Error; - type Error = anyhow::Error; - - async fn connect_once( - &self, - ctx: &mut RequestMonitoring, - node_info: &console::CachedNodeInfo, - timeout: time::Duration, - ) -> Result { - connect_to_compute_once( - ctx, - node_info, - self.conn_info, - timeout, - self.conn_id, - self.pool.clone(), - self.idle, - ) - .await - } - - fn update_connect_config(&self, _config: &mut compute::ConnCfg) {} -} - -// Wake up the destination if needed. Code here is a bit involved because -// we reuse the code from the usual proxy and we need to prepare few structures -// that this code expects. -#[tracing::instrument(fields(pid = tracing::field::Empty), skip_all)] -async fn connect_to_compute( - config: &config::ProxyConfig, +pub fn poll_client( + global_pool: Arc>, ctx: &mut RequestMonitoring, - conn_info: &ConnInfo, + conn_info: ConnInfo, + client: C, + mut connection: tokio_postgres::Connection, conn_id: uuid::Uuid, - pool: Weak>, -) -> anyhow::Result { - ctx.set_application(Some(APP_NAME)); - let backend = config - .auth_backend - .as_ref() - .map(|_| conn_info.user_info.clone()); - - if !config.disable_ip_check_for_http { - let (allowed_ips, _) = backend.get_allowed_ips_and_secret(ctx).await?; - if !check_peer_addr_is_in_list(&ctx.peer_addr, &allowed_ips) { - return Err(auth::AuthError::ip_address_not_allowed().into()); - } - } - let node_info = backend - .wake_compute(ctx) - .await? - .context("missing cache entry from wake_compute")?; - - ctx.set_project(node_info.aux.clone()); - - crate::proxy::connect_compute::connect_to_compute( - ctx, - &TokioMechanism { - conn_id, - conn_info, - pool, - idle: config.http_config.pool_options.idle_timeout, - }, - node_info, - &backend, - ) - .await -} - -async fn connect_to_compute_once( - ctx: &mut RequestMonitoring, - node_info: &console::CachedNodeInfo, - conn_info: &ConnInfo, - timeout: time::Duration, - conn_id: uuid::Uuid, - pool: Weak>, - idle: Duration, -) -> Result { - let mut config = (*node_info.config).clone(); - let mut session = ctx.session_id; - - let (client, mut connection) = config - .user(&conn_info.user_info.user) - .password(&*conn_info.password) - .dbname(&conn_info.dbname) - .connect_timeout(timeout) - .connect(tokio_postgres::NoTls) - .await?; - + aux: MetricsAuxInfo, +) -> Client { let conn_gauge = NUM_DB_CONNECTIONS_GAUGE .with_label_values(&[ctx.protocol]) .guard(); - - tracing::Span::current().record("pid", &tracing::field::display(client.get_process_id())); - - let (tx, mut rx) = tokio::sync::watch::channel(session); + let mut session_id = ctx.session_id; + let (tx, mut rx) = tokio::sync::watch::channel(session_id); let span = info_span!(parent: None, "connection", %conn_id); span.in_scope(|| { - info!(%conn_info, %session, "new connection"); + info!(%conn_info, %session_id, "new connection"); }); + let pool = + Arc::downgrade(&global_pool.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key())); + let pool_clone = pool.clone(); let db_user = conn_info.db_and_user(); + let idle = global_pool.get_idle_timeout(); tokio::spawn( - async move { - let _conn_gauge = conn_gauge; - let mut idle_timeout = pin!(tokio::time::sleep(idle)); - poll_fn(move |cx| { - if matches!(rx.has_changed(), Ok(true)) { - session = *rx.borrow_and_update(); - info!(%session, "changed session"); - idle_timeout.as_mut().reset(Instant::now() + idle); - } + async move { + let _conn_gauge = conn_gauge; + let mut idle_timeout = pin!(tokio::time::sleep(idle)); + poll_fn(move |cx| { + if matches!(rx.has_changed(), Ok(true)) { + session_id = *rx.borrow_and_update(); + info!(%session_id, "changed session"); + idle_timeout.as_mut().reset(Instant::now() + idle); + } - // 5 minute idle connection timeout - if idle_timeout.as_mut().poll(cx).is_ready() { - idle_timeout.as_mut().reset(Instant::now() + idle); - info!("connection idle"); - if let Some(pool) = pool.clone().upgrade() { - // remove client from pool - should close the connection if it's idle. - // does nothing if the client is currently checked-out and in-use - if pool.write().remove_client(db_user.clone(), conn_id) { - info!("idle connection removed"); - } - } - } - - loop { - let message = ready!(connection.poll_message(cx)); - - match message { - Some(Ok(AsyncMessage::Notice(notice))) => { - info!(%session, "notice: {}", notice); - } - Some(Ok(AsyncMessage::Notification(notif))) => { - warn!(%session, pid = notif.process_id(), channel = notif.channel(), "notification received"); - } - Some(Ok(_)) => { - warn!(%session, "unknown message"); - } - Some(Err(e)) => { - error!(%session, "connection error: {}", e); - break - } - None => { - info!("connection closed"); - break - } - } - } - - // remove from connection pool + // 5 minute idle connection timeout + if idle_timeout.as_mut().poll(cx).is_ready() { + idle_timeout.as_mut().reset(Instant::now() + idle); + info!("connection idle"); if let Some(pool) = pool.clone().upgrade() { + // remove client from pool - should close the connection if it's idle. + // does nothing if the client is currently checked-out and in-use if pool.write().remove_client(db_user.clone(), conn_id) { - info!("closed connection removed"); + info!("idle connection removed"); } } + } - Poll::Ready(()) - }).await; + loop { + let message = ready!(connection.poll_message(cx)); - } - .instrument(span) - ); + match message { + Some(Ok(AsyncMessage::Notice(notice))) => { + info!(%session_id, "notice: {}", notice); + } + Some(Ok(AsyncMessage::Notification(notif))) => { + warn!(%session_id, pid = notif.process_id(), channel = notif.channel(), "notification received"); + } + Some(Ok(_)) => { + warn!(%session_id, "unknown message"); + } + Some(Err(e)) => { + error!(%session_id, "connection error: {}", e); + break + } + None => { + info!("connection closed"); + break + } + } + } - Ok(ClientInner { + // remove from connection pool + if let Some(pool) = pool.clone().upgrade() { + if pool.write().remove_client(db_user.clone(), conn_id) { + info!("closed connection removed"); + } + } + + Poll::Ready(()) + }).await; + + } + .instrument(span)); + let inner = ClientInner { inner: client, session: tx, - aux: node_info.aux.clone(), + aux, conn_id, - }) + }; + Client::new(inner, conn_info, pool_clone) } -struct ClientInner { - inner: tokio_postgres::Client, +struct ClientInner { + inner: C, session: tokio::sync::watch::Sender, aux: MetricsAuxInfo, conn_id: uuid::Uuid, } -impl Client { +pub trait ClientInnerExt: Sync + Send + 'static { + fn is_closed(&self) -> bool; + fn get_process_id(&self) -> i32; +} + +impl ClientInnerExt for tokio_postgres::Client { + fn is_closed(&self) -> bool { + self.is_closed() + } + fn get_process_id(&self) -> i32 { + self.get_process_id() + } +} + +impl ClientInner { + pub fn is_closed(&self) -> bool { + self.inner.is_closed() + } +} + +impl Client { pub fn metrics(&self) -> Arc { let aux = &self.inner.as_ref().unwrap().aux; USAGE_METRICS.register(Ids { @@ -687,51 +566,46 @@ impl Client { } } -pub struct Client { - conn_id: uuid::Uuid, +pub struct Client { span: Span, - inner: Option, + inner: Option>, conn_info: ConnInfo, - pool: Weak>, + pool: Weak>>, } -pub struct Discard<'a> { +pub struct Discard<'a, C: ClientInnerExt> { conn_id: uuid::Uuid, conn_info: &'a ConnInfo, - pool: &'a mut Weak>, + pool: &'a mut Weak>>, } -impl Client { - pub(self) async fn new( - inner: ClientInner, +impl Client { + pub(self) fn new( + inner: ClientInner, conn_info: ConnInfo, - pool: Weak>, + pool: Weak>>, ) -> Self { Self { - conn_id: inner.conn_id, inner: Some(inner), span: Span::current(), conn_info, pool, } } - pub fn inner(&mut self) -> (&mut tokio_postgres::Client, Discard<'_>) { + pub fn inner(&mut self) -> (&mut C, Discard<'_, C>) { let Self { inner, pool, - conn_id, conn_info, span: _, } = self; + let inner = inner.as_mut().expect("client inner should not be removed"); ( - &mut inner - .as_mut() - .expect("client inner should not be removed") - .inner, + &mut inner.inner, Discard { pool, conn_info, - conn_id: *conn_id, + conn_id: inner.conn_id, }, ) } @@ -744,7 +618,7 @@ impl Client { } } -impl Discard<'_> { +impl Discard<'_, C> { pub fn check_idle(&mut self, status: ReadyForQueryStatus) { let conn_info = &self.conn_info; if status != ReadyForQueryStatus::Idle && std::mem::take(self.pool).strong_count() > 0 { @@ -759,8 +633,8 @@ impl Discard<'_> { } } -impl Deref for Client { - type Target = tokio_postgres::Client; +impl Deref for Client { + type Target = C; fn deref(&self) -> &Self::Target { &self @@ -771,8 +645,8 @@ impl Deref for Client { } } -impl Drop for Client { - fn drop(&mut self) { +impl Client { + fn do_drop(&mut self) -> Option { let conn_info = self.conn_info.clone(); let client = self .inner @@ -781,10 +655,161 @@ impl Drop for Client { if let Some(conn_pool) = std::mem::take(&mut self.pool).upgrade() { let current_span = self.span.clone(); // return connection to the pool - tokio::task::spawn_blocking(move || { + return Some(move || { let _span = current_span.enter(); let _ = EndpointConnPool::put(&conn_pool, &conn_info, client); }); } + None + } +} + +impl Drop for Client { + fn drop(&mut self) { + if let Some(drop) = self.do_drop() { + tokio::task::spawn_blocking(drop); + } + } +} + +#[cfg(test)] +mod tests { + use env_logger; + use std::{mem, sync::atomic::AtomicBool}; + + use super::*; + + struct MockClient(Arc); + impl MockClient { + fn new(is_closed: bool) -> Self { + MockClient(Arc::new(is_closed.into())) + } + } + impl ClientInnerExt for MockClient { + fn is_closed(&self) -> bool { + self.0.load(atomic::Ordering::Relaxed) + } + fn get_process_id(&self) -> i32 { + 0 + } + } + + fn create_inner() -> ClientInner { + create_inner_with(MockClient::new(false)) + } + + fn create_inner_with(client: MockClient) -> ClientInner { + ClientInner { + inner: client, + session: tokio::sync::watch::Sender::new(uuid::Uuid::new_v4()), + aux: Default::default(), + conn_id: uuid::Uuid::new_v4(), + } + } + + #[tokio::test] + async fn test_pool() { + let _ = env_logger::try_init(); + let config = Box::leak(Box::new(crate::config::HttpConfig { + pool_options: GlobalConnPoolOptions { + max_conns_per_endpoint: 2, + gc_epoch: Duration::from_secs(1), + pool_shards: 2, + idle_timeout: Duration::from_secs(1), + opt_in: false, + max_total_conns: 3, + }, + request_timeout: Duration::from_secs(1), + })); + let pool = GlobalConnPool::new(config); + let conn_info = ConnInfo { + user_info: ComputeUserInfo { + user: "user".into(), + endpoint: "endpoint".into(), + options: Default::default(), + }, + dbname: "dbname".into(), + password: "password".into(), + }; + let ep_pool = + Arc::downgrade(&pool.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key())); + { + let mut client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone()); + assert_eq!(0, pool.get_global_connections_count()); + client.discard(); + // Discard should not add the connection from the pool. + assert_eq!(0, pool.get_global_connections_count()); + } + { + let mut client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone()); + client.do_drop().unwrap()(); + mem::forget(client); // drop the client + assert_eq!(1, pool.get_global_connections_count()); + } + { + let mut closed_client = Client::new( + create_inner_with(MockClient::new(true)), + conn_info.clone(), + ep_pool.clone(), + ); + closed_client.do_drop().unwrap()(); + mem::forget(closed_client); // drop the client + // The closed client shouldn't be added to the pool. + assert_eq!(1, pool.get_global_connections_count()); + } + let is_closed: Arc = Arc::new(false.into()); + { + let mut client = Client::new( + create_inner_with(MockClient(is_closed.clone())), + conn_info.clone(), + ep_pool.clone(), + ); + client.do_drop().unwrap()(); + mem::forget(client); // drop the client + + // The client should be added to the pool. + assert_eq!(2, pool.get_global_connections_count()); + } + { + let mut client = Client::new(create_inner(), conn_info, ep_pool); + client.do_drop().unwrap()(); + mem::forget(client); // drop the client + + // The client shouldn't be added to the pool. Because the ep-pool is full. + assert_eq!(2, pool.get_global_connections_count()); + } + + let conn_info = ConnInfo { + user_info: ComputeUserInfo { + user: "user".into(), + endpoint: "endpoint-2".into(), + options: Default::default(), + }, + dbname: "dbname".into(), + password: "password".into(), + }; + let ep_pool = + Arc::downgrade(&pool.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key())); + { + let mut client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone()); + client.do_drop().unwrap()(); + mem::forget(client); // drop the client + assert_eq!(3, pool.get_global_connections_count()); + } + { + let mut client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone()); + client.do_drop().unwrap()(); + mem::forget(client); // drop the client + + // The client shouldn't be added to the pool. Because the global pool is full. + assert_eq!(3, pool.get_global_connections_count()); + } + + is_closed.store(true, atomic::Ordering::Relaxed); + // Do gc for all shards. + pool.gc(0); + pool.gc(1); + // Closed client should be removed from the pool. + assert_eq!(2, pool.get_global_connections_count()); } } diff --git a/proxy/src/serverless/json.rs b/proxy/src/serverless/json.rs index 05835b23ce..a089d34040 100644 --- a/proxy/src/serverless/json.rs +++ b/proxy/src/serverless/json.rs @@ -9,23 +9,23 @@ use tokio_postgres::Row; // as parameters. // pub fn json_to_pg_text(json: Vec) -> Vec> { - json.iter() - .map(|value| { - match value { - // special care for nulls - Value::Null => None, + json.iter().map(json_value_to_pg_text).collect() +} - // convert to text with escaping - v @ (Value::Bool(_) | Value::Number(_) | Value::Object(_)) => Some(v.to_string()), +fn json_value_to_pg_text(value: &Value) -> Option { + match value { + // special care for nulls + Value::Null => None, - // avoid escaping here, as we pass this as a parameter - Value::String(s) => Some(s.to_string()), + // convert to text with escaping + v @ (Value::Bool(_) | Value::Number(_) | Value::Object(_)) => Some(v.to_string()), - // special care for arrays - Value::Array(_) => json_array_to_pg_array(value), - } - }) - .collect() + // avoid escaping here, as we pass this as a parameter + Value::String(s) => Some(s.to_string()), + + // special care for arrays + Value::Array(_) => json_array_to_pg_array(value), + } } // diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 96bf39c915..7092b65f03 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -13,6 +13,7 @@ use hyper::StatusCode; use hyper::{Body, HeaderMap, Request}; use serde_json::json; use serde_json::Value; +use tokio::join; use tokio_postgres::error::DbError; use tokio_postgres::error::ErrorPosition; use tokio_postgres::GenericClient; @@ -20,6 +21,7 @@ use tokio_postgres::IsolationLevel; use tokio_postgres::ReadyForQueryStatus; use tokio_postgres::Transaction; use tracing::error; +use tracing::info; use tracing::instrument; use url::Url; use utils::http::error::ApiError; @@ -27,22 +29,25 @@ use utils::http::json::json_response; use crate::auth::backend::ComputeUserInfo; use crate::auth::endpoint_sni; -use crate::config::HttpConfig; +use crate::config::ProxyConfig; use crate::config::TlsConfig; use crate::context::RequestMonitoring; +use crate::metrics::HTTP_CONTENT_LENGTH; use crate::metrics::NUM_CONNECTION_REQUESTS_GAUGE; use crate::proxy::NeonOptions; use crate::RoleName; +use super::backend::PoolingBackend; use super::conn_pool::ConnInfo; -use super::conn_pool::GlobalConnPool; -use super::json::{json_to_pg_text, pg_text_row_to_json}; +use super::json::json_to_pg_text; +use super::json::pg_text_row_to_json; use super::SERVERLESS_DRIVER_SNI; #[derive(serde::Deserialize)] struct QueryData { query: String, - params: Vec, + #[serde(deserialize_with = "bytes_to_pg_text")] + params: Vec>, } #[derive(serde::Deserialize)] @@ -69,6 +74,15 @@ static TXN_DEFERRABLE: HeaderName = HeaderName::from_static("neon-batch-deferrab static HEADER_VALUE_TRUE: HeaderValue = HeaderValue::from_static("true"); +fn bytes_to_pg_text<'de, D>(deserializer: D) -> Result>, D::Error> +where + D: serde::de::Deserializer<'de>, +{ + // TODO: consider avoiding the allocation here. + let json: Vec = serde::de::Deserialize::deserialize(deserializer)?; + Ok(json_to_pg_text(json)) +} + fn get_conn_info( ctx: &mut RequestMonitoring, headers: &HeaderMap, @@ -171,16 +185,15 @@ fn check_matches(sni_hostname: &str, hostname: &str) -> Result, sni_hostname: Option, - conn_pool: Arc, + backend: Arc, ) -> Result, ApiError> { let result = tokio::time::timeout( - config.request_timeout, - handle_inner(tls, config, ctx, request, sni_hostname, conn_pool), + config.http_config.request_timeout, + handle_inner(config, ctx, request, sni_hostname, backend), ) .await; let mut response = match result { @@ -265,7 +278,7 @@ pub async fn handle( Err(_) => { let message = format!( "HTTP-Connection timed out, execution time exeeded {} seconds", - config.request_timeout.as_secs() + config.http_config.request_timeout.as_secs() ); error!(message); json_response( @@ -283,22 +296,36 @@ pub async fn handle( #[instrument(name = "sql-over-http", fields(pid = tracing::field::Empty), skip_all)] async fn handle_inner( - tls: &'static TlsConfig, - config: &'static HttpConfig, + config: &'static ProxyConfig, ctx: &mut RequestMonitoring, request: Request, sni_hostname: Option, - conn_pool: Arc, + backend: Arc, ) -> anyhow::Result> { let _request_gauge = NUM_CONNECTION_REQUESTS_GAUGE - .with_label_values(&["http"]) + .with_label_values(&[ctx.protocol]) .guard(); + info!( + protocol = ctx.protocol, + "handling interactive connection from client" + ); // // Determine the destination and connection params // let headers = request.headers(); - let conn_info = get_conn_info(ctx, headers, sni_hostname, tls)?; + // TLS config should be there. + let conn_info = get_conn_info( + ctx, + headers, + sni_hostname, + config.tls_config.as_ref().unwrap(), + )?; + info!( + user = conn_info.user_info.user.as_str(), + project = conn_info.user_info.endpoint.as_str(), + "credentials" + ); // Determine the output options. Default behaviour is 'false'. Anything that is not // strictly 'true' assumed to be false. @@ -307,8 +334,8 @@ async fn handle_inner( // Allow connection pooling only if explicitly requested // or if we have decided that http pool is no longer opt-in - let allow_pool = - !config.pool_options.opt_in || headers.get(&ALLOW_POOL) == Some(&HEADER_VALUE_TRUE); + let allow_pool = !config.http_config.pool_options.opt_in + || headers.get(&ALLOW_POOL) == Some(&HEADER_VALUE_TRUE); // isolation level, read only and deferrable @@ -333,6 +360,8 @@ async fn handle_inner( None => MAX_REQUEST_SIZE + 1, }; drop(paused); + info!(request_content_length, "request size in bytes"); + HTTP_CONTENT_LENGTH.observe(request_content_length as f64); // we don't have a streaming request support yet so this is to prevent OOM // from a malicious user sending an extremely large request body @@ -342,13 +371,28 @@ async fn handle_inner( )); } - // - // Read the query and query params from the request body - // - let body = hyper::body::to_bytes(request.into_body()).await?; - let payload: Payload = serde_json::from_slice(&body)?; + let fetch_and_process_request = async { + let body = hyper::body::to_bytes(request.into_body()) + .await + .map_err(anyhow::Error::from)?; + let payload: Payload = serde_json::from_slice(&body)?; + Ok::(payload) // Adjust error type accordingly + }; - let mut client = conn_pool.get(ctx, conn_info, !allow_pool).await?; + let authenticate_and_connect = async { + let keys = backend.authenticate(ctx, &conn_info).await?; + backend + .connect_to_compute(ctx, conn_info, keys, !allow_pool) + .await + }; + + // Run both operations in parallel + let (payload_result, auth_and_connect_result) = + join!(fetch_and_process_request, authenticate_and_connect,); + + // Handle the results + let payload = payload_result?; // Handle errors appropriately + let mut client = auth_and_connect_result?; // Handle errors appropriately let mut response = Response::builder() .status(StatusCode::OK) @@ -482,7 +526,7 @@ async fn query_to_json( raw_output: bool, array_mode: bool, ) -> anyhow::Result<(ReadyForQueryStatus, Value)> { - let query_params = json_to_pg_text(data.params); + let query_params = data.params; let row_stream = client.query_raw_txt(&data.query, query_params).await?; // Manually drain the stream into a vector to leave row_stream hanging diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py index 1d62f09840..b3b35e446d 100644 --- a/test_runner/regress/test_proxy.py +++ b/test_runner/regress/test_proxy.py @@ -393,11 +393,11 @@ def test_sql_over_http_batch(static_proxy: NeonProxy): def test_sql_over_http_pool(static_proxy: NeonProxy): static_proxy.safe_psql("create user http_auth with password 'http' superuser") - def get_pid(status: int, pw: str) -> Any: + def get_pid(status: int, pw: str, user="http_auth") -> Any: return static_proxy.http_query( GET_CONNECTION_PID_QUERY, [], - user="http_auth", + user=user, password=pw, expected_code=status, ) @@ -418,20 +418,14 @@ def test_sql_over_http_pool(static_proxy: NeonProxy): static_proxy.safe_psql("alter user http_auth with password 'http2'") - # after password change, should open a new connection to verify it - pid2 = get_pid(200, "http2")["rows"][0]["pid"] - assert pid1 != pid2 + # after password change, shouldn't open a new connection because it checks password in proxy. + rows = get_pid(200, "http2")["rows"] + assert rows == [{"pid": pid1}] time.sleep(0.02) - # query should be on an existing connection - pid = get_pid(200, "http2")["rows"][0]["pid"] - assert pid in [pid1, pid2] - - time.sleep(0.02) - - # old password should not work - res = get_pid(400, "http") + # incorrect user shouldn't reveal that the user doesn't exists + res = get_pid(400, "http", user="http_auth2") assert "password authentication failed for user" in res["message"] From 6c34d4cd147eb3704d8e54b434afee35b7d08704 Mon Sep 17 00:00:00 2001 From: Anna Khanova <32508607+khanova@users.noreply.github.com> Date: Thu, 8 Feb 2024 14:52:04 +0100 Subject: [PATCH 42/72] Proxy: set timeout on establishing connection (#6679) ## Problem There is no timeout on the handshake. ## Summary of changes Set the timeout on the establishing connection. --- proxy/src/bin/proxy.rs | 4 ++++ proxy/src/config.rs | 1 + proxy/src/proxy.rs | 9 +++++---- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index 6974f1a274..8fbcb56758 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -88,6 +88,9 @@ struct ProxyCliArgs { /// path to directory with TLS certificates for client postgres connections #[clap(long)] certs_dir: Option, + /// timeout for the TLS handshake + #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)] + handshake_timeout: tokio::time::Duration, /// http endpoint to receive periodic metric updates #[clap(long)] metric_collection_endpoint: Option, @@ -411,6 +414,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> { require_client_ip: args.require_client_ip, disable_ip_check_for_http: args.disable_ip_check_for_http, endpoint_rps_limit, + handshake_timeout: args.handshake_timeout, // TODO: add this argument region: args.region.clone(), })); diff --git a/proxy/src/config.rs b/proxy/src/config.rs index 2c46458a49..31c9228b35 100644 --- a/proxy/src/config.rs +++ b/proxy/src/config.rs @@ -22,6 +22,7 @@ pub struct ProxyConfig { pub disable_ip_check_for_http: bool, pub endpoint_rps_limit: Vec, pub region: String, + pub handshake_timeout: Duration, } #[derive(Debug)] diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs index b68fb26e42..b3b221d3e2 100644 --- a/proxy/src/proxy.rs +++ b/proxy/src/proxy.rs @@ -194,10 +194,11 @@ pub async fn handle_client( let pause = ctx.latency_timer.pause(); let do_handshake = handshake(stream, mode.handshake_tls(tls), &cancel_map); - let (mut stream, params) = match do_handshake.await? { - Some(x) => x, - None => return Ok(()), // it's a cancellation request - }; + let (mut stream, params) = + match tokio::time::timeout(config.handshake_timeout, do_handshake).await?? { + Some(x) => x, + None => return Ok(()), // it's a cancellation request + }; drop(pause); let hostname = mode.hostname(stream.get_ref()); From 43eae17f0d2e84b0c88e34f3fff6bfe515008b89 Mon Sep 17 00:00:00 2001 From: Konstantin Knizhnik Date: Thu, 8 Feb 2024 17:31:15 +0200 Subject: [PATCH 43/72] Drop unused replication slots (#6655) ## Problem See #6626 If there is inactive replication slot then Postgres will not bw able to shrink WAL and delete unused snapshots. If she other active subscription is present, then snapshots created each 15 seconds will overflow AUX_DIR. Setting `max_slot_wal_keep_size` doesn't solve the problem, because even small WAL segment will be enough to overflow AUX_DIR if there is no other activity on the system. ## Summary of changes If there are active subscriptions and some logical replication slots are not used during `neon.logical_replication_max_time_lag` interval, then unused slot is dropped. ## Checklist before requesting a review - [ ] I have performed a self-review of my code. - [ ] If it is a core feature, I have added thorough tests. - [ ] Do we need to implement analytics? if so did you add the relevant metrics to the dashboard? - [ ] If this PR requires public announcement, mark it with /release-notes label and add several sentences in this section. ## Checklist before merging - [ ] Do not forget to reformat commit message to not include the above checklist Co-authored-by: Konstantin Knizhnik --- pgxn/neon/neon.c | 133 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 133 insertions(+) diff --git a/pgxn/neon/neon.c b/pgxn/neon/neon.c index b930fdb3ca..799f88751c 100644 --- a/pgxn/neon/neon.c +++ b/pgxn/neon/neon.c @@ -11,16 +11,23 @@ #include "postgres.h" #include "fmgr.h" +#include "miscadmin.h" #include "access/xact.h" #include "access/xlog.h" #include "storage/buf_internals.h" #include "storage/bufmgr.h" #include "catalog/pg_type.h" +#include "postmaster/bgworker.h" +#include "postmaster/interrupt.h" +#include "replication/slot.h" #include "replication/walsender.h" +#include "storage/procsignal.h" +#include "tcop/tcopprot.h" #include "funcapi.h" #include "access/htup_details.h" #include "utils/pg_lsn.h" #include "utils/guc.h" +#include "utils/wait_event.h" #include "neon.h" #include "walproposer.h" @@ -30,6 +37,130 @@ PG_MODULE_MAGIC; void _PG_init(void); +static int logical_replication_max_time_lag = 3600; + +static void +InitLogicalReplicationMonitor(void) +{ + BackgroundWorker bgw; + + DefineCustomIntVariable( + "neon.logical_replication_max_time_lag", + "Threshold for dropping unused logical replication slots", + NULL, + &logical_replication_max_time_lag, + 3600, 0, INT_MAX, + PGC_SIGHUP, + GUC_UNIT_S, + NULL, NULL, NULL); + + memset(&bgw, 0, sizeof(bgw)); + bgw.bgw_flags = BGWORKER_SHMEM_ACCESS; + bgw.bgw_start_time = BgWorkerStart_RecoveryFinished; + snprintf(bgw.bgw_library_name, BGW_MAXLEN, "neon"); + snprintf(bgw.bgw_function_name, BGW_MAXLEN, "LogicalSlotsMonitorMain"); + snprintf(bgw.bgw_name, BGW_MAXLEN, "Logical replication monitor"); + snprintf(bgw.bgw_type, BGW_MAXLEN, "Logical replication monitor"); + bgw.bgw_restart_time = 5; + bgw.bgw_notify_pid = 0; + bgw.bgw_main_arg = (Datum) 0; + + RegisterBackgroundWorker(&bgw); +} + +typedef struct +{ + NameData name; + bool dropped; + XLogRecPtr confirmed_flush_lsn; + TimestampTz last_updated; +} SlotStatus; + +/* + * Unused logical replication slots pins WAL and prevents deletion of snapshots. + */ +PGDLLEXPORT void +LogicalSlotsMonitorMain(Datum main_arg) +{ + SlotStatus* slots; + TimestampTz now, last_checked; + + /* Establish signal handlers. */ + pqsignal(SIGUSR1, procsignal_sigusr1_handler); + pqsignal(SIGHUP, SignalHandlerForConfigReload); + pqsignal(SIGTERM, die); + + BackgroundWorkerUnblockSignals(); + + slots = (SlotStatus*)calloc(max_replication_slots, sizeof(SlotStatus)); + last_checked = GetCurrentTimestamp(); + + for (;;) + { + (void) WaitLatch(MyLatch, + WL_LATCH_SET | WL_EXIT_ON_PM_DEATH | WL_TIMEOUT, + logical_replication_max_time_lag*1000/2, + PG_WAIT_EXTENSION); + ResetLatch(MyLatch); + CHECK_FOR_INTERRUPTS(); + + now = GetCurrentTimestamp(); + + if (now - last_checked > logical_replication_max_time_lag*USECS_PER_SEC) + { + int n_active_slots = 0; + last_checked = now; + + LWLockAcquire(ReplicationSlotControlLock, LW_SHARED); + for (int i = 0; i < max_replication_slots; i++) + { + ReplicationSlot *s = &ReplicationSlotCtl->replication_slots[i]; + + /* Consider only logical repliction slots */ + if (!s->in_use || !SlotIsLogical(s)) + continue; + + if (s->active_pid != 0) + { + n_active_slots += 1; + continue; + } + + /* Check if there was some activity with the slot since last check */ + if (s->data.confirmed_flush != slots[i].confirmed_flush_lsn) + { + slots[i].confirmed_flush_lsn = s->data.confirmed_flush; + slots[i].last_updated = now; + } + else if (now - slots[i].last_updated > logical_replication_max_time_lag*USECS_PER_SEC) + { + slots[i].name = s->data.name; + slots[i].dropped = true; + } + } + LWLockRelease(ReplicationSlotControlLock); + + /* + * If there are no active subscriptions, then no new snapshots are generated + * and so no need to force slot deletion. + */ + if (n_active_slots != 0) + { + for (int i = 0; i < max_replication_slots; i++) + { + if (slots[i].dropped) + { + elog(LOG, "Drop logical replication slot because it was not update more than %ld seconds", + (now - slots[i].last_updated)/USECS_PER_SEC); + ReplicationSlotDrop(slots[i].name.data, true); + slots[i].dropped = false; + } + } + } + } + } +} + void _PG_init(void) { @@ -44,6 +175,8 @@ _PG_init(void) pg_init_libpagestore(); pg_init_walproposer(); + InitLogicalReplicationMonitor(); + InitControlPlaneConnector(); pg_init_extension_server(); From af91a28936eef0b1e5149dc71d92394a89410372 Mon Sep 17 00:00:00 2001 From: John Spray Date: Thu, 8 Feb 2024 15:35:13 +0000 Subject: [PATCH 44/72] pageserver: shard splitting (#6379) ## Problem One doesn't know at tenant creation time how large the tenant will grow. We need to be able to dynamically adjust the shard count at runtime. This is implemented as "splitting" of shards into smaller child shards, which cover a subset of the keyspace that the parent covered. Refer to RFC: https://github.com/neondatabase/neon/pull/6358 Part of epic: #6278 ## Summary of changes This PR implements the happy path (does not cleanly recover from a crash mid-split, although won't lose any data), without any optimizations (e.g. child shards re-download their own copies of layers that the parent shard already had on local disk) - Add `/v1/tenant/:tenant_shard_id/shard_split` API to pageserver: this copies the shard's index to the child shards' paths, instantiates child `Tenant` object, and tears down parent `Tenant` object. - Add `splitting` column to `tenant_shards` table. This is written into an existing migration because we haven't deployed yet, so don't need to cleanly upgrade. - Add `/control/v1/tenant/:tenant_id/shard_split` API to attachment_service, - Add `test_sharding_split_smoke` test. This covers the happy path: future PRs will add tests that exercise failure cases. --- Dockerfile | 5 + .../up.sql | 1 + control_plane/attachment_service/src/http.rs | 19 +- .../attachment_service/src/persistence.rs | 102 +++++- .../src/persistence/split_state.rs | 46 +++ .../attachment_service/src/schema.rs | 1 + .../attachment_service/src/service.rs | 333 +++++++++++++++++- .../attachment_service/src/tenant_state.rs | 10 + control_plane/src/attachment_service.rs | 21 +- control_plane/src/bin/neon_local.rs | 25 ++ libs/pageserver_api/src/models.rs | 10 + libs/pageserver_api/src/shard.rs | 128 +++++++ pageserver/client/src/mgmt_api.rs | 16 + pageserver/src/http/routes.rs | 27 +- pageserver/src/tenant.rs | 66 ++++ pageserver/src/tenant/mgr.rs | 169 ++++++++- .../tenant/remote_timeline_client/upload.rs | 2 +- test_runner/fixtures/neon_fixtures.py | 2 +- test_runner/regress/test_sharding.py | 129 ++++++- 19 files changed, 1088 insertions(+), 24 deletions(-) create mode 100644 control_plane/attachment_service/src/persistence/split_state.rs diff --git a/Dockerfile b/Dockerfile index bb926643dc..c37f94b981 100644 --- a/Dockerfile +++ b/Dockerfile @@ -100,6 +100,11 @@ RUN mkdir -p /data/.neon/ && chown -R neon:neon /data/.neon/ \ -c "listen_pg_addr='0.0.0.0:6400'" \ -c "listen_http_addr='0.0.0.0:9898'" +# When running a binary that links with libpq, default to using our most recent postgres version. Binaries +# that want a particular postgres version will select it explicitly: this is just a default. +ENV LD_LIBRARY_PATH /usr/local/v16/lib + + VOLUME ["/data"] USER neon EXPOSE 6400 diff --git a/control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/up.sql b/control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/up.sql index 585dbc79a0..2ffdae6287 100644 --- a/control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/up.sql +++ b/control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/up.sql @@ -7,6 +7,7 @@ CREATE TABLE tenant_shards ( generation INTEGER NOT NULL, generation_pageserver BIGINT NOT NULL, placement_policy VARCHAR NOT NULL, + splitting SMALLINT NOT NULL, -- config is JSON encoded, opaque to the database. config TEXT NOT NULL ); \ No newline at end of file diff --git a/control_plane/attachment_service/src/http.rs b/control_plane/attachment_service/src/http.rs index 049e66fddf..38eecaf7ef 100644 --- a/control_plane/attachment_service/src/http.rs +++ b/control_plane/attachment_service/src/http.rs @@ -3,7 +3,8 @@ use crate::service::{Service, STARTUP_RECONCILE_TIMEOUT}; use hyper::{Body, Request, Response}; use hyper::{StatusCode, Uri}; use pageserver_api::models::{ - TenantCreateRequest, TenantLocationConfigRequest, TimelineCreateRequest, + TenantCreateRequest, TenantLocationConfigRequest, TenantShardSplitRequest, + TimelineCreateRequest, }; use pageserver_api::shard::TenantShardId; use pageserver_client::mgmt_api; @@ -292,6 +293,19 @@ async fn handle_node_configure(mut req: Request) -> Result, json_response(StatusCode::OK, state.service.node_configure(config_req)?) } +async fn handle_tenant_shard_split( + service: Arc, + mut req: Request, +) -> Result, ApiError> { + let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?; + let split_req = json_request::(&mut req).await?; + + json_response( + StatusCode::OK, + service.tenant_shard_split(tenant_id, split_req).await?, + ) +} + async fn handle_tenant_shard_migrate( service: Arc, mut req: Request, @@ -391,6 +405,9 @@ pub fn make_router( .put("/control/v1/tenant/:tenant_shard_id/migrate", |r| { tenant_service_handler(r, handle_tenant_shard_migrate) }) + .put("/control/v1/tenant/:tenant_id/shard_split", |r| { + tenant_service_handler(r, handle_tenant_shard_split) + }) // Tenant operations // The ^/v1/ endpoints act as a "Virtual Pageserver", enabling shard-naive clients to call into // this service to manage tenants that actually consist of many tenant shards, as if they are a single entity. diff --git a/control_plane/attachment_service/src/persistence.rs b/control_plane/attachment_service/src/persistence.rs index db487bcec6..cead540058 100644 --- a/control_plane/attachment_service/src/persistence.rs +++ b/control_plane/attachment_service/src/persistence.rs @@ -1,7 +1,9 @@ +pub(crate) mod split_state; use std::collections::HashMap; use std::str::FromStr; use std::time::Duration; +use self::split_state::SplitState; use camino::Utf8Path; use camino::Utf8PathBuf; use control_plane::attachment_service::{NodeAvailability, NodeSchedulingPolicy}; @@ -363,19 +365,101 @@ impl Persistence { Ok(()) } - // TODO: when we start shard splitting, we must durably mark the tenant so that - // on restart, we know that we must go through recovery (list shards that exist - // and pick up where we left off and/or revert to parent shards). + // When we start shard splitting, we must durably mark the tenant so that + // on restart, we know that we must go through recovery. + // + // We create the child shards here, so that they will be available for increment_generation calls + // if some pageserver holding a child shard needs to restart before the overall tenant split is complete. #[allow(dead_code)] - pub(crate) async fn begin_shard_split(&self, _tenant_id: TenantId) -> anyhow::Result<()> { - todo!(); + pub(crate) async fn begin_shard_split( + &self, + old_shard_count: ShardCount, + split_tenant_id: TenantId, + parent_to_children: Vec<(TenantShardId, Vec)>, + ) -> DatabaseResult<()> { + use crate::schema::tenant_shards::dsl::*; + self.with_conn(move |conn| -> DatabaseResult<()> { + conn.transaction(|conn| -> DatabaseResult<()> { + // Mark parent shards as splitting + let updated = diesel::update(tenant_shards) + .filter(tenant_id.eq(split_tenant_id.to_string())) + .filter(shard_count.eq(old_shard_count.0 as i32)) + .set((splitting.eq(1),)) + .execute(conn)?; + if ShardCount(updated.try_into().map_err(|_| DatabaseError::Logical(format!("Overflow existing shard count {} while splitting", updated)))?) != old_shard_count { + // Perhaps a deletion or another split raced with this attempt to split, mutating + // the parent shards that we intend to split. In this case the split request should fail. + return Err(DatabaseError::Logical( + format!("Unexpected existing shard count {updated} when preparing tenant for split (expected {old_shard_count:?})") + )); + } + + // FIXME: spurious clone to sidestep closure move rules + let parent_to_children = parent_to_children.clone(); + + // Insert child shards + for (parent_shard_id, children) in parent_to_children { + let mut parent = crate::schema::tenant_shards::table + .filter(tenant_id.eq(parent_shard_id.tenant_id.to_string())) + .filter(shard_number.eq(parent_shard_id.shard_number.0 as i32)) + .filter(shard_count.eq(parent_shard_id.shard_count.0 as i32)) + .load::(conn)?; + let parent = if parent.len() != 1 { + return Err(DatabaseError::Logical(format!( + "Parent shard {parent_shard_id} not found" + ))); + } else { + parent.pop().unwrap() + }; + for mut shard in children { + // Carry the parent's generation into the child + shard.generation = parent.generation; + + debug_assert!(shard.splitting == SplitState::Splitting); + diesel::insert_into(tenant_shards) + .values(shard) + .execute(conn)?; + } + } + + Ok(()) + })?; + + Ok(()) + }) + .await } - // TODO: when we finish shard splitting, we must atomically clean up the old shards + // When we finish shard splitting, we must atomically clean up the old shards // and insert the new shards, and clear the splitting marker. #[allow(dead_code)] - pub(crate) async fn complete_shard_split(&self, _tenant_id: TenantId) -> anyhow::Result<()> { - todo!(); + pub(crate) async fn complete_shard_split( + &self, + split_tenant_id: TenantId, + old_shard_count: ShardCount, + ) -> DatabaseResult<()> { + use crate::schema::tenant_shards::dsl::*; + self.with_conn(move |conn| -> DatabaseResult<()> { + conn.transaction(|conn| -> QueryResult<()> { + // Drop parent shards + diesel::delete(tenant_shards) + .filter(tenant_id.eq(split_tenant_id.to_string())) + .filter(shard_count.eq(old_shard_count.0 as i32)) + .execute(conn)?; + + // Clear sharding flag + let updated = diesel::update(tenant_shards) + .filter(tenant_id.eq(split_tenant_id.to_string())) + .set((splitting.eq(0),)) + .execute(conn)?; + debug_assert!(updated > 0); + + Ok(()) + })?; + + Ok(()) + }) + .await } } @@ -403,6 +487,8 @@ pub(crate) struct TenantShardPersistence { #[serde(default)] pub(crate) placement_policy: String, #[serde(default)] + pub(crate) splitting: SplitState, + #[serde(default)] pub(crate) config: String, } diff --git a/control_plane/attachment_service/src/persistence/split_state.rs b/control_plane/attachment_service/src/persistence/split_state.rs new file mode 100644 index 0000000000..bce1a75843 --- /dev/null +++ b/control_plane/attachment_service/src/persistence/split_state.rs @@ -0,0 +1,46 @@ +use diesel::pg::{Pg, PgValue}; +use diesel::{ + deserialize::FromSql, deserialize::FromSqlRow, expression::AsExpression, serialize::ToSql, + sql_types::Int2, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, FromSqlRow, AsExpression)] +#[diesel(sql_type = SplitStateSQLRepr)] +#[derive(Deserialize, Serialize)] +pub enum SplitState { + Idle = 0, + Splitting = 1, +} + +impl Default for SplitState { + fn default() -> Self { + Self::Idle + } +} + +type SplitStateSQLRepr = Int2; + +impl ToSql for SplitState { + fn to_sql<'a>( + &'a self, + out: &'a mut diesel::serialize::Output, + ) -> diesel::serialize::Result { + let raw_value: i16 = *self as i16; + let mut new_out = out.reborrow(); + ToSql::::to_sql(&raw_value, &mut new_out) + } +} + +impl FromSql for SplitState { + fn from_sql(pg_value: PgValue) -> diesel::deserialize::Result { + match FromSql::::from_sql(pg_value).map(|v| match v { + 0 => Some(Self::Idle), + 1 => Some(Self::Splitting), + _ => None, + })? { + Some(v) => Ok(v), + None => Err(format!("Invalid SplitState value, was: {:?}", pg_value.as_bytes()).into()), + } + } +} diff --git a/control_plane/attachment_service/src/schema.rs b/control_plane/attachment_service/src/schema.rs index de80fc8f64..db5a957443 100644 --- a/control_plane/attachment_service/src/schema.rs +++ b/control_plane/attachment_service/src/schema.rs @@ -20,6 +20,7 @@ diesel::table! { generation -> Int4, generation_pageserver -> Int8, placement_policy -> Varchar, + splitting -> Int2, config -> Text, } } diff --git a/control_plane/attachment_service/src/service.rs b/control_plane/attachment_service/src/service.rs index 1db1906df8..0ec2b9dc4c 100644 --- a/control_plane/attachment_service/src/service.rs +++ b/control_plane/attachment_service/src/service.rs @@ -1,4 +1,5 @@ use std::{ + cmp::Ordering, collections::{BTreeMap, HashMap}, str::FromStr, sync::Arc, @@ -23,7 +24,7 @@ use pageserver_api::{ models::{ LocationConfig, LocationConfigMode, ShardParameters, TenantConfig, TenantCreateRequest, TenantLocationConfigRequest, TenantLocationConfigResponse, TenantShardLocation, - TimelineCreateRequest, TimelineInfo, + TenantShardSplitRequest, TenantShardSplitResponse, TimelineCreateRequest, TimelineInfo, }, shard::{ShardCount, ShardIdentity, ShardNumber, ShardStripeSize, TenantShardId}, }; @@ -40,7 +41,11 @@ use utils::{ use crate::{ compute_hook::{self, ComputeHook}, node::Node, - persistence::{DatabaseError, NodePersistence, Persistence, TenantShardPersistence}, + persistence::{ + split_state::SplitState, DatabaseError, NodePersistence, Persistence, + TenantShardPersistence, + }, + reconciler::attached_location_conf, scheduler::Scheduler, tenant_state::{ IntentState, ObservedState, ObservedStateLocation, ReconcileResult, ReconcileWaitError, @@ -476,6 +481,7 @@ impl Service { generation_pageserver: i64::MAX, placement_policy: serde_json::to_string(&PlacementPolicy::default()).unwrap(), config: serde_json::to_string(&TenantConfig::default()).unwrap(), + splitting: SplitState::default(), }; match self.persistence.insert_tenant_shards(vec![tsp]).await { @@ -718,6 +724,7 @@ impl Service { generation_pageserver: i64::MAX, placement_policy: serde_json::to_string(&placement_policy).unwrap(), config: serde_json::to_string(&create_req.config).unwrap(), + splitting: SplitState::default(), }) .collect(); self.persistence @@ -1100,6 +1107,7 @@ impl Service { self.ensure_attached_wait(tenant_id).await?; // TODO: refuse to do this if shard splitting is in progress + // (https://github.com/neondatabase/neon/issues/6676) let targets = { let locked = self.inner.read().unwrap(); let mut targets = Vec::new(); @@ -1180,6 +1188,7 @@ impl Service { self.ensure_attached_wait(tenant_id).await?; // TODO: refuse to do this if shard splitting is in progress + // (https://github.com/neondatabase/neon/issues/6676) let targets = { let locked = self.inner.read().unwrap(); let mut targets = Vec::new(); @@ -1352,6 +1361,326 @@ impl Service { }) } + pub(crate) async fn tenant_shard_split( + &self, + tenant_id: TenantId, + split_req: TenantShardSplitRequest, + ) -> Result { + let mut policy = None; + let mut shard_ident = None; + + // TODO: put a cancellation token on Service for clean shutdown + let cancel = CancellationToken::new(); + + // A parent shard which will be split + struct SplitTarget { + parent_id: TenantShardId, + node: Node, + child_ids: Vec, + } + + // Validate input, and calculate which shards we will create + let (old_shard_count, targets, compute_hook) = { + let locked = self.inner.read().unwrap(); + + let pageservers = locked.nodes.clone(); + + let mut targets = Vec::new(); + + // In case this is a retry, count how many already-split shards we found + let mut children_found = Vec::new(); + let mut old_shard_count = None; + + for (tenant_shard_id, shard) in + locked.tenants.range(TenantShardId::tenant_range(tenant_id)) + { + match shard.shard.count.0.cmp(&split_req.new_shard_count) { + Ordering::Equal => { + // Already split this + children_found.push(*tenant_shard_id); + continue; + } + Ordering::Greater => { + return Err(ApiError::BadRequest(anyhow::anyhow!( + "Requested count {} but already have shards at count {}", + split_req.new_shard_count, + shard.shard.count.0 + ))); + } + Ordering::Less => { + // Fall through: this shard has lower count than requested, + // is a candidate for splitting. + } + } + + match old_shard_count { + None => old_shard_count = Some(shard.shard.count), + Some(old_shard_count) => { + if old_shard_count != shard.shard.count { + // We may hit this case if a caller asked for two splits to + // different sizes, before the first one is complete. + // e.g. 1->2, 2->4, where the 4 call comes while we have a mixture + // of shard_count=1 and shard_count=2 shards in the map. + return Err(ApiError::Conflict( + "Cannot split, currently mid-split".to_string(), + )); + } + } + } + if policy.is_none() { + policy = Some(shard.policy.clone()); + } + if shard_ident.is_none() { + shard_ident = Some(shard.shard); + } + + if tenant_shard_id.shard_count == ShardCount(split_req.new_shard_count) { + tracing::info!( + "Tenant shard {} already has shard count {}", + tenant_shard_id, + split_req.new_shard_count + ); + continue; + } + + let node_id = + shard + .intent + .attached + .ok_or(ApiError::BadRequest(anyhow::anyhow!( + "Cannot split a tenant that is not attached" + )))?; + + let node = pageservers + .get(&node_id) + .expect("Pageservers may not be deleted while referenced"); + + // TODO: if any reconciliation is currently in progress for this shard, wait for it. + + targets.push(SplitTarget { + parent_id: *tenant_shard_id, + node: node.clone(), + child_ids: tenant_shard_id.split(ShardCount(split_req.new_shard_count)), + }); + } + + if targets.is_empty() { + if children_found.len() == split_req.new_shard_count as usize { + return Ok(TenantShardSplitResponse { + new_shards: children_found, + }); + } else { + // No shards found to split, and no existing children found: the + // tenant doesn't exist at all. + return Err(ApiError::NotFound( + anyhow::anyhow!("Tenant {} not found", tenant_id).into(), + )); + } + } + + (old_shard_count, targets, locked.compute_hook.clone()) + }; + + // unwrap safety: we would have returned above if we didn't find at least one shard to split + let old_shard_count = old_shard_count.unwrap(); + let shard_ident = shard_ident.unwrap(); + let policy = policy.unwrap(); + + // FIXME: we have dropped self.inner lock, and not yet written anything to the database: another + // request could occur here, deleting or mutating the tenant. begin_shard_split checks that the + // parent shards exist as expected, but it would be neater to do the above pre-checks within the + // same database transaction rather than pre-check in-memory and then maybe-fail the database write. + // (https://github.com/neondatabase/neon/issues/6676) + + // Before creating any new child shards in memory or on the pageservers, persist them: this + // enables us to ensure that we will always be able to clean up if something goes wrong. This also + // acts as the protection against two concurrent attempts to split: one of them will get a database + // error trying to insert the child shards. + let mut child_tsps = Vec::new(); + for target in &targets { + let mut this_child_tsps = Vec::new(); + for child in &target.child_ids { + let mut child_shard = shard_ident; + child_shard.number = child.shard_number; + child_shard.count = child.shard_count; + + this_child_tsps.push(TenantShardPersistence { + tenant_id: child.tenant_id.to_string(), + shard_number: child.shard_number.0 as i32, + shard_count: child.shard_count.0 as i32, + shard_stripe_size: shard_ident.stripe_size.0 as i32, + // Note: this generation is a placeholder, [`Persistence::begin_shard_split`] will + // populate the correct generation as part of its transaction, to protect us + // against racing with changes in the state of the parent. + generation: 0, + generation_pageserver: target.node.id.0 as i64, + placement_policy: serde_json::to_string(&policy).unwrap(), + // TODO: get the config out of the map + config: serde_json::to_string(&TenantConfig::default()).unwrap(), + splitting: SplitState::Splitting, + }); + } + + child_tsps.push((target.parent_id, this_child_tsps)); + } + + if let Err(e) = self + .persistence + .begin_shard_split(old_shard_count, tenant_id, child_tsps) + .await + { + match e { + DatabaseError::Query(diesel::result::Error::DatabaseError( + DatabaseErrorKind::UniqueViolation, + _, + )) => { + // Inserting a child shard violated a unique constraint: we raced with another call to + // this function + tracing::warn!("Conflicting attempt to split {tenant_id}: {e}"); + return Err(ApiError::Conflict("Tenant is already splitting".into())); + } + _ => return Err(ApiError::InternalServerError(e.into())), + } + } + + // FIXME: we have now committed the shard split state to the database, so any subsequent + // failure needs to roll it back. We will later wrap this function in logic to roll back + // the split if it fails. + // (https://github.com/neondatabase/neon/issues/6676) + + // TODO: issue split calls concurrently (this only matters once we're splitting + // N>1 shards into M shards -- initially we're usually splitting 1 shard into N). + + for target in &targets { + let SplitTarget { + parent_id, + node, + child_ids, + } = target; + let client = mgmt_api::Client::new(node.base_url(), self.config.jwt_token.as_deref()); + let response = client + .tenant_shard_split( + *parent_id, + TenantShardSplitRequest { + new_shard_count: split_req.new_shard_count, + }, + ) + .await + .map_err(|e| ApiError::Conflict(format!("Failed to split {}: {}", parent_id, e)))?; + + tracing::info!( + "Split {} into {}", + parent_id, + response + .new_shards + .iter() + .map(|s| format!("{:?}", s)) + .collect::>() + .join(",") + ); + + if &response.new_shards != child_ids { + // This should never happen: the pageserver should agree with us on how shard splits work. + return Err(ApiError::InternalServerError(anyhow::anyhow!( + "Splitting shard {} resulted in unexpected IDs: {:?} (expected {:?})", + parent_id, + response.new_shards, + child_ids + ))); + } + } + + // TODO: if the pageserver restarted concurrently with our split API call, + // the actual generation of the child shard might differ from the generation + // we expect it to have. In order for our in-database generation to end up + // correct, we should carry the child generation back in the response and apply it here + // in complete_shard_split (and apply the correct generation in memory) + // (or, we can carry generation in the request and reject the request if + // it doesn't match, but that requires more retry logic on this side) + + self.persistence + .complete_shard_split(tenant_id, old_shard_count) + .await?; + + // Replace all the shards we just split with their children + let mut response = TenantShardSplitResponse { + new_shards: Vec::new(), + }; + let mut child_locations = Vec::new(); + { + let mut locked = self.inner.write().unwrap(); + for target in targets { + let SplitTarget { + parent_id, + node: _node, + child_ids, + } = target; + let (pageserver, generation, config) = { + let old_state = locked + .tenants + .remove(&parent_id) + .expect("It was present, we just split it"); + ( + old_state.intent.attached.unwrap(), + old_state.generation, + old_state.config.clone(), + ) + }; + + locked.tenants.remove(&parent_id); + + for child in child_ids { + let mut child_shard = shard_ident; + child_shard.number = child.shard_number; + child_shard.count = child.shard_count; + + let mut child_observed: HashMap = HashMap::new(); + child_observed.insert( + pageserver, + ObservedStateLocation { + conf: Some(attached_location_conf(generation, &child_shard, &config)), + }, + ); + + let mut child_state = TenantState::new(child, child_shard, policy.clone()); + child_state.intent = IntentState::single(Some(pageserver)); + child_state.observed = ObservedState { + locations: child_observed, + }; + child_state.generation = generation; + child_state.config = config.clone(); + + child_locations.push((child, pageserver)); + + locked.tenants.insert(child, child_state); + response.new_shards.push(child); + } + } + } + + // Send compute notifications for all the new shards + let mut failed_notifications = Vec::new(); + for (child_id, child_ps) in child_locations { + if let Err(e) = compute_hook.notify(child_id, child_ps, &cancel).await { + tracing::warn!("Failed to update compute of {}->{} during split, proceeding anyway to complete split ({e})", + child_id, child_ps); + failed_notifications.push(child_id); + } + } + + // If we failed any compute notifications, make a note to retry later. + if !failed_notifications.is_empty() { + let mut locked = self.inner.write().unwrap(); + for failed in failed_notifications { + if let Some(shard) = locked.tenants.get_mut(&failed) { + shard.pending_compute_notification = true; + } + } + } + + Ok(response) + } + pub(crate) async fn tenant_shard_migrate( &self, tenant_shard_id: TenantShardId, diff --git a/control_plane/attachment_service/src/tenant_state.rs b/control_plane/attachment_service/src/tenant_state.rs index a358e1ff7b..c0ab076a55 100644 --- a/control_plane/attachment_service/src/tenant_state.rs +++ b/control_plane/attachment_service/src/tenant_state.rs @@ -193,6 +193,13 @@ impl IntentState { result } + pub(crate) fn single(node_id: Option) -> Self { + Self { + attached: node_id, + secondary: vec![], + } + } + /// When a node goes offline, we update intents to avoid using it /// as their attached pageserver. /// @@ -286,6 +293,9 @@ impl TenantState { // self.intent refers to pageservers that are offline, and pick other // pageservers if so. + // TODO: respect the splitting bit on tenants: if they are currently splitting then we may not + // change their attach location. + // Build the set of pageservers already in use by this tenant, to avoid scheduling // more work on the same pageservers we're already using. let mut used_pageservers = self.intent.all_pageservers(); diff --git a/control_plane/src/attachment_service.rs b/control_plane/src/attachment_service.rs index a3f832036c..c3e071aa71 100644 --- a/control_plane/src/attachment_service.rs +++ b/control_plane/src/attachment_service.rs @@ -8,7 +8,10 @@ use diesel::{ use diesel_migrations::{HarnessWithOutput, MigrationHarness}; use hyper::Method; use pageserver_api::{ - models::{ShardParameters, TenantCreateRequest, TimelineCreateRequest, TimelineInfo}, + models::{ + ShardParameters, TenantCreateRequest, TenantShardSplitRequest, TenantShardSplitResponse, + TimelineCreateRequest, TimelineInfo, + }, shard::TenantShardId, }; use pageserver_client::mgmt_api::ResponseErrorMessageExt; @@ -648,7 +651,7 @@ impl AttachmentService { ) -> anyhow::Result { self.dispatch( Method::PUT, - format!("tenant/{tenant_shard_id}/migrate"), + format!("control/v1/tenant/{tenant_shard_id}/migrate"), Some(TenantShardMigrateRequest { tenant_shard_id, node_id, @@ -657,6 +660,20 @@ impl AttachmentService { .await } + #[instrument(skip(self), fields(%tenant_id, %new_shard_count))] + pub async fn tenant_split( + &self, + tenant_id: TenantId, + new_shard_count: u8, + ) -> anyhow::Result { + self.dispatch( + Method::PUT, + format!("control/v1/tenant/{tenant_id}/shard_split"), + Some(TenantShardSplitRequest { new_shard_count }), + ) + .await + } + #[instrument(skip_all, fields(node_id=%req.node_id))] pub async fn node_register(&self, req: NodeRegisterRequest) -> anyhow::Result<()> { self.dispatch::<_, ()>(Method::POST, "control/v1/node".to_string(), Some(req)) diff --git a/control_plane/src/bin/neon_local.rs b/control_plane/src/bin/neon_local.rs index e56007dd20..b9af467fdf 100644 --- a/control_plane/src/bin/neon_local.rs +++ b/control_plane/src/bin/neon_local.rs @@ -575,6 +575,26 @@ async fn handle_tenant( println!("{tenant_table}"); println!("{shard_table}"); } + Some(("shard-split", matches)) => { + let tenant_id = get_tenant_id(matches, env)?; + let shard_count: u8 = matches.get_one::("shard-count").cloned().unwrap_or(0); + + let attachment_service = AttachmentService::from_env(env); + let result = attachment_service + .tenant_split(tenant_id, shard_count) + .await?; + println!( + "Split tenant {} into shards {}", + tenant_id, + result + .new_shards + .iter() + .map(|s| format!("{:?}", s)) + .collect::>() + .join(",") + ); + } + Some((sub_name, _)) => bail!("Unexpected tenant subcommand '{}'", sub_name), None => bail!("no tenant subcommand provided"), } @@ -1524,6 +1544,11 @@ fn cli() -> Command { .subcommand(Command::new("status") .about("Human readable summary of the tenant's shards and attachment locations") .arg(tenant_id_arg.clone())) + .subcommand(Command::new("shard-split") + .about("Increase the number of shards in the tenant") + .arg(tenant_id_arg.clone()) + .arg(Arg::new("shard-count").value_parser(value_parser!(u8)).long("shard-count").action(ArgAction::Set).help("Number of shards in the new tenant (default 1)")) + ) ) .subcommand( Command::new("pageserver") diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs index c08cacb822..46324efd43 100644 --- a/libs/pageserver_api/src/models.rs +++ b/libs/pageserver_api/src/models.rs @@ -192,6 +192,16 @@ pub struct TimelineCreateRequest { pub pg_version: Option, } +#[derive(Serialize, Deserialize)] +pub struct TenantShardSplitRequest { + pub new_shard_count: u8, +} + +#[derive(Serialize, Deserialize)] +pub struct TenantShardSplitResponse { + pub new_shards: Vec, +} + /// Parameters that apply to all shards in a tenant. Used during tenant creation. #[derive(Serialize, Deserialize, Debug)] #[serde(deny_unknown_fields)] diff --git a/libs/pageserver_api/src/shard.rs b/libs/pageserver_api/src/shard.rs index e27aad8156..322b6c642e 100644 --- a/libs/pageserver_api/src/shard.rs +++ b/libs/pageserver_api/src/shard.rs @@ -88,12 +88,36 @@ impl TenantShardId { pub fn is_unsharded(&self) -> bool { self.shard_number == ShardNumber(0) && self.shard_count == ShardCount(0) } + + /// Convenience for dropping the tenant_id and just getting the ShardIndex: this + /// is useful when logging from code that is already in a span that includes tenant ID, to + /// keep messages reasonably terse. pub fn to_index(&self) -> ShardIndex { ShardIndex { shard_number: self.shard_number, shard_count: self.shard_count, } } + + /// Calculate the children of this TenantShardId when splitting the overall tenant into + /// the given number of shards. + pub fn split(&self, new_shard_count: ShardCount) -> Vec { + let effective_old_shard_count = std::cmp::max(self.shard_count.0, 1); + let mut child_shards = Vec::new(); + for shard_number in 0..ShardNumber(new_shard_count.0).0 { + // Key mapping is based on a round robin mapping of key hash modulo shard count, + // so our child shards are the ones which the same keys would map to. + if shard_number % effective_old_shard_count == self.shard_number.0 { + child_shards.push(TenantShardId { + tenant_id: self.tenant_id, + shard_number: ShardNumber(shard_number), + shard_count: new_shard_count, + }) + } + } + + child_shards + } } /// Formatting helper @@ -793,4 +817,108 @@ mod tests { let shard = key_to_shard_number(ShardCount(10), DEFAULT_STRIPE_SIZE, &key); assert_eq!(shard, ShardNumber(8)); } + + #[test] + fn shard_id_split() { + let tenant_id = TenantId::generate(); + let parent = TenantShardId::unsharded(tenant_id); + + // Unsharded into 2 + assert_eq!( + parent.split(ShardCount(2)), + vec![ + TenantShardId { + tenant_id, + shard_count: ShardCount(2), + shard_number: ShardNumber(0) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(2), + shard_number: ShardNumber(1) + } + ] + ); + + // Unsharded into 4 + assert_eq!( + parent.split(ShardCount(4)), + vec![ + TenantShardId { + tenant_id, + shard_count: ShardCount(4), + shard_number: ShardNumber(0) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(4), + shard_number: ShardNumber(1) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(4), + shard_number: ShardNumber(2) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(4), + shard_number: ShardNumber(3) + } + ] + ); + + // count=1 into 2 (check this works the same as unsharded.) + let parent = TenantShardId { + tenant_id, + shard_count: ShardCount(1), + shard_number: ShardNumber(0), + }; + assert_eq!( + parent.split(ShardCount(2)), + vec![ + TenantShardId { + tenant_id, + shard_count: ShardCount(2), + shard_number: ShardNumber(0) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(2), + shard_number: ShardNumber(1) + } + ] + ); + + // count=2 into count=8 + let parent = TenantShardId { + tenant_id, + shard_count: ShardCount(2), + shard_number: ShardNumber(1), + }; + assert_eq!( + parent.split(ShardCount(8)), + vec![ + TenantShardId { + tenant_id, + shard_count: ShardCount(8), + shard_number: ShardNumber(1) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(8), + shard_number: ShardNumber(3) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(8), + shard_number: ShardNumber(5) + }, + TenantShardId { + tenant_id, + shard_count: ShardCount(8), + shard_number: ShardNumber(7) + }, + ] + ); + } } diff --git a/pageserver/client/src/mgmt_api.rs b/pageserver/client/src/mgmt_api.rs index 8abe58e1a2..200369df90 100644 --- a/pageserver/client/src/mgmt_api.rs +++ b/pageserver/client/src/mgmt_api.rs @@ -310,6 +310,22 @@ impl Client { .map_err(Error::ReceiveBody) } + pub async fn tenant_shard_split( + &self, + tenant_shard_id: TenantShardId, + req: TenantShardSplitRequest, + ) -> Result { + let uri = format!( + "{}/v1/tenant/{}/shard_split", + self.mgmt_api_endpoint, tenant_shard_id + ); + self.request(Method::PUT, &uri, req) + .await? + .json() + .await + .map_err(Error::ReceiveBody) + } + pub async fn timeline_list( &self, tenant_shard_id: &TenantShardId, diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index ebcb27fa08..af9a3c7301 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -19,11 +19,14 @@ use pageserver_api::models::ShardParameters; use pageserver_api::models::TenantDetails; use pageserver_api::models::TenantLocationConfigResponse; use pageserver_api::models::TenantShardLocation; +use pageserver_api::models::TenantShardSplitRequest; +use pageserver_api::models::TenantShardSplitResponse; use pageserver_api::models::TenantState; use pageserver_api::models::{ DownloadRemoteLayersTaskSpawnRequest, LocationConfigMode, TenantAttachRequest, TenantLoadRequest, TenantLocationConfigRequest, }; +use pageserver_api::shard::ShardCount; use pageserver_api::shard::TenantShardId; use remote_storage::GenericRemoteStorage; use remote_storage::TimeTravelError; @@ -875,7 +878,7 @@ async fn tenant_reset_handler( let state = get_state(&request); state .tenant_manager - .reset_tenant(tenant_shard_id, drop_cache.unwrap_or(false), ctx) + .reset_tenant(tenant_shard_id, drop_cache.unwrap_or(false), &ctx) .await .map_err(ApiError::InternalServerError)?; @@ -1104,6 +1107,25 @@ async fn tenant_size_handler( ) } +async fn tenant_shard_split_handler( + mut request: Request, + _cancel: CancellationToken, +) -> Result, ApiError> { + let req: TenantShardSplitRequest = json_request(&mut request).await?; + + let tenant_shard_id: TenantShardId = parse_request_param(&request, "tenant_shard_id")?; + let state = get_state(&request); + let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Warn); + + let new_shards = state + .tenant_manager + .shard_split(tenant_shard_id, ShardCount(req.new_shard_count), &ctx) + .await + .map_err(ApiError::InternalServerError)?; + + json_response(StatusCode::OK, TenantShardSplitResponse { new_shards }) +} + async fn layer_map_info_handler( request: Request, _cancel: CancellationToken, @@ -2063,6 +2085,9 @@ pub fn make_router( .put("/v1/tenant/config", |r| { api_handler(r, update_tenant_config_handler) }) + .put("/v1/tenant/:tenant_shard_id/shard_split", |r| { + api_handler(r, tenant_shard_split_handler) + }) .get("/v1/tenant/:tenant_shard_id/config", |r| { api_handler(r, get_tenant_config_handler) }) diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index f704f8c0dd..f086f46213 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -53,6 +53,7 @@ use self::metadata::TimelineMetadata; use self::mgr::GetActiveTenantError; use self::mgr::GetTenantError; use self::mgr::TenantsMap; +use self::remote_timeline_client::upload::upload_index_part; use self::remote_timeline_client::RemoteTimelineClient; use self::timeline::uninit::TimelineExclusionError; use self::timeline::uninit::TimelineUninitMark; @@ -2397,6 +2398,67 @@ impl Tenant { pub(crate) fn get_generation(&self) -> Generation { self.generation } + + /// This function partially shuts down the tenant (it shuts down the Timelines) and is fallible, + /// and can leave the tenant in a bad state if it fails. The caller is responsible for + /// resetting this tenant to a valid state if we fail. + pub(crate) async fn split_prepare( + &self, + child_shards: &Vec, + ) -> anyhow::Result<()> { + let timelines = self.timelines.lock().unwrap().clone(); + for timeline in timelines.values() { + let Some(tl_client) = &timeline.remote_client else { + anyhow::bail!("Remote storage is mandatory"); + }; + + let Some(remote_storage) = &self.remote_storage else { + anyhow::bail!("Remote storage is mandatory"); + }; + + // We do not block timeline creation/deletion during splits inside the pageserver: it is up to higher levels + // to ensure that they do not start a split if currently in the process of doing these. + + // Upload an index from the parent: this is partly to provide freshness for the + // child tenants that will copy it, and partly for general ease-of-debugging: there will + // always be a parent shard index in the same generation as we wrote the child shard index. + tl_client.schedule_index_upload_for_file_changes()?; + tl_client.wait_completion().await?; + + // Shut down the timeline's remote client: this means that the indices we write + // for child shards will not be invalidated by the parent shard deleting layers. + tl_client.shutdown().await?; + + // Download methods can still be used after shutdown, as they don't flow through the remote client's + // queue. In principal the RemoteTimelineClient could provide this without downloading it, but this + // operation is rare, so it's simpler to just download it (and robustly guarantees that the index + // we use here really is the remotely persistent one). + let result = tl_client + .download_index_file(self.cancel.clone()) + .instrument(info_span!("download_index_file", tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), timeline_id=%timeline.timeline_id)) + .await?; + let index_part = match result { + MaybeDeletedIndexPart::Deleted(_) => { + anyhow::bail!("Timeline deletion happened concurrently with split") + } + MaybeDeletedIndexPart::IndexPart(p) => p, + }; + + for child_shard in child_shards { + upload_index_part( + remote_storage, + child_shard, + &timeline.timeline_id, + self.generation, + &index_part, + &self.cancel, + ) + .await?; + } + } + + Ok(()) + } } /// Given a Vec of timelines and their ancestors (timeline_id, ancestor_id), @@ -3732,6 +3794,10 @@ impl Tenant { Ok(()) } + + pub(crate) fn get_tenant_conf(&self) -> TenantConfOpt { + self.tenant_conf.read().unwrap().tenant_conf + } } fn remove_timeline_and_uninit_mark( diff --git a/pageserver/src/tenant/mgr.rs b/pageserver/src/tenant/mgr.rs index 5ec910ca3e..9aee39bd35 100644 --- a/pageserver/src/tenant/mgr.rs +++ b/pageserver/src/tenant/mgr.rs @@ -2,6 +2,7 @@ //! page server. use camino::{Utf8DirEntry, Utf8Path, Utf8PathBuf}; +use itertools::Itertools; use pageserver_api::key::Key; use pageserver_api::models::ShardParameters; use pageserver_api::shard::{ShardCount, ShardIdentity, ShardNumber, TenantShardId}; @@ -22,7 +23,7 @@ use tokio_util::sync::CancellationToken; use tracing::*; use remote_storage::GenericRemoteStorage; -use utils::crashsafe; +use utils::{completion, crashsafe}; use crate::config::PageServerConf; use crate::context::{DownloadBehavior, RequestContext}; @@ -644,8 +645,6 @@ pub(crate) async fn shutdown_all_tenants() { } async fn shutdown_all_tenants0(tenants: &std::sync::RwLock) { - use utils::completion; - let mut join_set = JoinSet::new(); // Atomically, 1. create the shutdown tasks and 2. prevent creation of new tenants. @@ -1200,7 +1199,7 @@ impl TenantManager { &self, tenant_shard_id: TenantShardId, drop_cache: bool, - ctx: RequestContext, + ctx: &RequestContext, ) -> anyhow::Result<()> { let mut slot_guard = tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?; let Some(old_slot) = slot_guard.get_old_value() else { @@ -1253,7 +1252,7 @@ impl TenantManager { None, self.tenants, SpawnMode::Normal, - &ctx, + ctx, )?; slot_guard.upsert(TenantSlot::Attached(tenant))?; @@ -1375,6 +1374,164 @@ impl TenantManager { slot_guard.revert(); result } + + #[instrument(skip_all, fields(tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), new_shard_count=%new_shard_count.0))] + pub(crate) async fn shard_split( + &self, + tenant_shard_id: TenantShardId, + new_shard_count: ShardCount, + ctx: &RequestContext, + ) -> anyhow::Result> { + let tenant = get_tenant(tenant_shard_id, true)?; + + // Plan: identify what the new child shards will be + let effective_old_shard_count = std::cmp::max(tenant_shard_id.shard_count.0, 1); + if new_shard_count <= ShardCount(effective_old_shard_count) { + anyhow::bail!("Requested shard count is not an increase"); + } + let expansion_factor = new_shard_count.0 / effective_old_shard_count; + if !expansion_factor.is_power_of_two() { + anyhow::bail!("Requested split is not a power of two"); + } + + let parent_shard_identity = tenant.shard_identity; + let parent_tenant_conf = tenant.get_tenant_conf(); + let parent_generation = tenant.generation; + + let child_shards = tenant_shard_id.split(new_shard_count); + tracing::info!( + "Shard {} splits into: {}", + tenant_shard_id.to_index(), + child_shards + .iter() + .map(|id| format!("{}", id.to_index())) + .join(",") + ); + + // Phase 1: Write out child shards' remote index files, in the parent tenant's current generation + if let Err(e) = tenant.split_prepare(&child_shards).await { + // If [`Tenant::split_prepare`] fails, we must reload the tenant, because it might + // have been left in a partially-shut-down state. + tracing::warn!("Failed to prepare for split: {e}, reloading Tenant before returning"); + self.reset_tenant(tenant_shard_id, false, ctx).await?; + return Err(e); + } + + self.resources.deletion_queue_client.flush_advisory(); + + // Phase 2: Put the parent shard to InProgress and grab a reference to the parent Tenant + drop(tenant); + let mut parent_slot_guard = + tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?; + let parent = match parent_slot_guard.get_old_value() { + Some(TenantSlot::Attached(t)) => t, + Some(TenantSlot::Secondary(_)) => anyhow::bail!("Tenant location in secondary mode"), + Some(TenantSlot::InProgress(_)) => { + // tenant_map_acquire_slot never returns InProgress, if a slot was InProgress + // it would return an error. + unreachable!() + } + None => { + // We don't actually need the parent shard to still be attached to do our work, but it's + // a weird enough situation that the caller probably didn't want us to continue working + // if they had detached the tenant they requested the split on. + anyhow::bail!("Detached parent shard in the middle of split!") + } + }; + + // TODO: hardlink layers from the parent into the child shard directories so that they don't immediately re-download + // TODO: erase the dentries from the parent + + // Take a snapshot of where the parent's WAL ingest had got to: we will wait for + // child shards to reach this point. + let mut target_lsns = HashMap::new(); + for timeline in parent.timelines.lock().unwrap().clone().values() { + target_lsns.insert(timeline.timeline_id, timeline.get_last_record_lsn()); + } + + // TODO: we should have the parent shard stop its WAL ingest here, it's a waste of resources + // and could slow down the children trying to catch up. + + // Phase 3: Spawn the child shards + for child_shard in &child_shards { + let mut child_shard_identity = parent_shard_identity; + child_shard_identity.count = child_shard.shard_count; + child_shard_identity.number = child_shard.shard_number; + + let child_location_conf = LocationConf { + mode: LocationMode::Attached(AttachedLocationConfig { + generation: parent_generation, + attach_mode: AttachmentMode::Single, + }), + shard: child_shard_identity, + tenant_conf: parent_tenant_conf, + }; + + self.upsert_location( + *child_shard, + child_location_conf, + None, + SpawnMode::Normal, + ctx, + ) + .await?; + } + + // Phase 4: wait for child chards WAL ingest to catch up to target LSN + for child_shard_id in &child_shards { + let child_shard = { + let locked = TENANTS.read().unwrap(); + let peek_slot = + tenant_map_peek_slot(&locked, child_shard_id, TenantSlotPeekMode::Read)?; + peek_slot.and_then(|s| s.get_attached()).cloned() + }; + if let Some(t) = child_shard { + let timelines = t.timelines.lock().unwrap().clone(); + for timeline in timelines.values() { + let Some(target_lsn) = target_lsns.get(&timeline.timeline_id) else { + continue; + }; + + tracing::info!( + "Waiting for child shard {}/{} to reach target lsn {}...", + child_shard_id, + timeline.timeline_id, + target_lsn + ); + if let Err(e) = timeline.wait_lsn(*target_lsn, ctx).await { + // Failure here might mean shutdown, in any case this part is an optimization + // and we shouldn't hold up the split operation. + tracing::warn!( + "Failed to wait for timeline {} to reach lsn {target_lsn}: {e}", + timeline.timeline_id + ); + } else { + tracing::info!( + "Child shard {}/{} reached target lsn {}", + child_shard_id, + timeline.timeline_id, + target_lsn + ); + } + } + } + } + + // Phase 5: Shut down the parent shard. + let (_guard, progress) = completion::channel(); + match parent.shutdown(progress, false).await { + Ok(()) => {} + Err(other) => { + other.wait().await; + } + } + parent_slot_guard.drop_old_value()?; + + // Phase 6: Release the InProgress on the parent shard + drop(parent_slot_guard); + + Ok(child_shards) + } } #[derive(Debug, thiserror::Error)] @@ -2209,8 +2366,6 @@ async fn remove_tenant_from_memory( where F: std::future::Future>, { - use utils::completion; - let mut slot_guard = tenant_map_acquire_slot_impl(&tenant_shard_id, tenants, TenantSlotAcquireMode::MustExist)?; diff --git a/pageserver/src/tenant/remote_timeline_client/upload.rs b/pageserver/src/tenant/remote_timeline_client/upload.rs index e8ba1d3d6e..c17e27b446 100644 --- a/pageserver/src/tenant/remote_timeline_client/upload.rs +++ b/pageserver/src/tenant/remote_timeline_client/upload.rs @@ -27,7 +27,7 @@ use super::index::LayerFileMetadata; use tracing::info; /// Serializes and uploads the given index part data to the remote storage. -pub(super) async fn upload_index_part<'a>( +pub(crate) async fn upload_index_part<'a>( storage: &'a GenericRemoteStorage, tenant_shard_id: &TenantShardId, timeline_id: &TimelineId, diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 4491655aeb..3d2549a8c3 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -4054,7 +4054,7 @@ def logical_replication_sync(subscriber: VanillaPostgres, publisher: Endpoint) - def tenant_get_shards( - env: NeonEnv, tenant_id: TenantId, pageserver_id: Optional[int] + env: NeonEnv, tenant_id: TenantId, pageserver_id: Optional[int] = None ) -> list[tuple[TenantShardId, NeonPageserver]]: """ Helper for when you want to talk to one or more pageservers, and the diff --git a/test_runner/regress/test_sharding.py b/test_runner/regress/test_sharding.py index c16bfc2ec6..805eaa34b0 100644 --- a/test_runner/regress/test_sharding.py +++ b/test_runner/regress/test_sharding.py @@ -1,6 +1,7 @@ from fixtures.log_helper import log from fixtures.neon_fixtures import ( NeonEnvBuilder, + tenant_get_shards, ) from fixtures.remote_storage import s3_storage from fixtures.types import TimelineId @@ -82,4 +83,130 @@ def test_sharding_smoke( ) assert timelines == {env.initial_timeline, timeline_b} - # TODO: test timeline deletion and tenant deletion (depends on change in attachment_service) + +def test_sharding_split_smoke( + neon_env_builder: NeonEnvBuilder, +): + """ + Test the basics of shard splitting: + - The API results in more shards than we started with + - The tenant's data remains readable + + """ + + # We will start with 4 shards and split into 8, then migrate all those + # 8 shards onto separate pageservers + shard_count = 4 + split_shard_count = 8 + neon_env_builder.num_pageservers = split_shard_count + + # 1MiB stripes: enable getting some meaningful data distribution without + # writing large quantities of data in this test. The stripe size is given + # in number of 8KiB pages. + stripe_size = 128 + + # Use S3-compatible remote storage so that we can scrub: this test validates + # that the scrubber doesn't barf when it sees a sharded tenant. + neon_env_builder.enable_pageserver_remote_storage(s3_storage()) + neon_env_builder.enable_scrub_on_exit() + + neon_env_builder.preserve_database_files = True + + env = neon_env_builder.init_start( + initial_tenant_shard_count=shard_count, initial_tenant_shard_stripe_size=stripe_size + ) + tenant_id = env.initial_tenant + timeline_id = env.initial_timeline + workload = Workload(env, tenant_id, timeline_id, branch_name="main") + workload.init() + + # Initial data + workload.write_rows(256) + + # Note which pageservers initially hold a shard after tenant creation + pre_split_pageserver_ids = [loc["node_id"] for loc in env.attachment_service.locate(tenant_id)] + + # For pageservers holding a shard, validate their ingest statistics + # reflect a proper splitting of the WAL. + for pageserver in env.pageservers: + if pageserver.id not in pre_split_pageserver_ids: + continue + + metrics = pageserver.http_client().get_metrics_values( + [ + "pageserver_wal_ingest_records_received_total", + "pageserver_wal_ingest_records_committed_total", + "pageserver_wal_ingest_records_filtered_total", + ] + ) + + log.info(f"Pageserver {pageserver.id} metrics: {metrics}") + + # Not everything received was committed + assert ( + metrics["pageserver_wal_ingest_records_received_total"] + > metrics["pageserver_wal_ingest_records_committed_total"] + ) + + # Something was committed + assert metrics["pageserver_wal_ingest_records_committed_total"] > 0 + + # Counts are self consistent + assert ( + metrics["pageserver_wal_ingest_records_received_total"] + == metrics["pageserver_wal_ingest_records_committed_total"] + + metrics["pageserver_wal_ingest_records_filtered_total"] + ) + + # TODO: validate that shards have different sizes + + workload.validate() + + assert len(pre_split_pageserver_ids) == 4 + + env.attachment_service.tenant_shard_split(tenant_id, shard_count=split_shard_count) + + post_split_pageserver_ids = [loc["node_id"] for loc in env.attachment_service.locate(tenant_id)] + # We should have split into 8 shards, on the same 4 pageservers we started on. + assert len(post_split_pageserver_ids) == split_shard_count + assert len(set(post_split_pageserver_ids)) == shard_count + assert set(post_split_pageserver_ids) == set(pre_split_pageserver_ids) + + workload.validate() + + workload.churn_rows(256) + + workload.validate() + + # Run GC on all new shards, to check they don't barf or delete anything that breaks reads + # (compaction was already run as part of churn_rows) + all_shards = tenant_get_shards(env, tenant_id) + for tenant_shard_id, pageserver in all_shards: + pageserver.http_client().timeline_gc(tenant_shard_id, timeline_id, None) + + # Restart all nodes, to check that the newly created shards are durable + for ps in env.pageservers: + ps.restart() + + workload.validate() + + migrate_to_pageserver_ids = list( + set(p.id for p in env.pageservers) - set(pre_split_pageserver_ids) + ) + assert len(migrate_to_pageserver_ids) == split_shard_count - shard_count + + # Migrate shards away from the node where the split happened + for ps_id in pre_split_pageserver_ids: + shards_here = [ + tenant_shard_id + for (tenant_shard_id, pageserver) in all_shards + if pageserver.id == ps_id + ] + assert len(shards_here) == 2 + migrate_shard = shards_here[0] + destination = migrate_to_pageserver_ids.pop() + + log.info(f"Migrating shard {migrate_shard} from {ps_id} to {destination}") + env.neon_cli.tenant_migrate(migrate_shard, destination, timeout_secs=10) + + workload.validate() From e8d2843df63ba05cd74baa8017736a903f9a322a Mon Sep 17 00:00:00 2001 From: John Spray Date: Thu, 8 Feb 2024 18:00:53 +0000 Subject: [PATCH 45/72] storage controller: improved handling of node availability on restart (#6658) - Automatically set a node's availability to Active if it is responsive in startup_reconcile - Impose a 5s timeout of HTTP request to list location conf, so that an unresponsive node can't hang it for minutes - Do several retries if the request fails with a retryable error, to be tolerant of concurrent pageserver & storage controller restarts - Add a readiness hook for use with k8s so that we can tell when the startup reconciliaton is done and the service is fully ready to do work. - Add /metrics to the list of un-authenticated endpoints (this is unrelated but we're touching the line in this PR already, and it fixes auth error spam in deployed container.) - A test for the above. Closes: #6670 --- control_plane/attachment_service/src/http.rs | 14 ++- .../attachment_service/src/service.rs | 107 +++++++++++++----- libs/utils/src/completion.rs | 5 + pageserver/client/src/mgmt_api.rs | 10 +- test_runner/fixtures/neon_fixtures.py | 9 ++ test_runner/regress/test_sharding_service.py | 32 ++++++ 6 files changed, 149 insertions(+), 28 deletions(-) diff --git a/control_plane/attachment_service/src/http.rs b/control_plane/attachment_service/src/http.rs index 38eecaf7ef..8501e4980f 100644 --- a/control_plane/attachment_service/src/http.rs +++ b/control_plane/attachment_service/src/http.rs @@ -42,7 +42,7 @@ pub struct HttpState { impl HttpState { pub fn new(service: Arc, auth: Option>) -> Self { - let allowlist_routes = ["/status"] + let allowlist_routes = ["/status", "/ready", "/metrics"] .iter() .map(|v| v.parse().unwrap()) .collect::>(); @@ -325,6 +325,17 @@ async fn handle_status(_req: Request) -> Result, ApiError> json_response(StatusCode::OK, ()) } +/// Readiness endpoint indicates when we're done doing startup I/O (e.g. reconciling +/// with remote pageserver nodes). This is intended for use as a kubernetes readiness probe. +async fn handle_ready(req: Request) -> Result, ApiError> { + let state = get_state(&req); + if state.service.startup_complete.is_ready() { + json_response(StatusCode::OK, ()) + } else { + json_response(StatusCode::SERVICE_UNAVAILABLE, ()) + } +} + impl From for ApiError { fn from(value: ReconcileError) -> Self { ApiError::Conflict(format!("Reconciliation error: {}", value)) @@ -380,6 +391,7 @@ pub fn make_router( .data(Arc::new(HttpState::new(service, auth))) // Non-prefixed generic endpoints (status, metrics) .get("/status", |r| request_span(r, handle_status)) + .get("/ready", |r| request_span(r, handle_ready)) // Upcalls for the pageserver: point the pageserver's `control_plane_api` config to this prefix .post("/upcall/v1/re-attach", |r| { request_span(r, handle_re_attach) diff --git a/control_plane/attachment_service/src/service.rs b/control_plane/attachment_service/src/service.rs index 0ec2b9dc4c..0331087e0d 100644 --- a/control_plane/attachment_service/src/service.rs +++ b/control_plane/attachment_service/src/service.rs @@ -1,6 +1,6 @@ use std::{ cmp::Ordering, - collections::{BTreeMap, HashMap}, + collections::{BTreeMap, HashMap, HashSet}, str::FromStr, sync::Arc, time::{Duration, Instant}, @@ -31,6 +31,7 @@ use pageserver_api::{ use pageserver_client::mgmt_api; use tokio_util::sync::CancellationToken; use utils::{ + backoff, completion::Barrier, generation::Generation, http::error::ApiError, @@ -150,31 +151,71 @@ impl Service { // indeterminate, same as in [`ObservedStateLocation`]) let mut observed = HashMap::new(); - let nodes = { - let locked = self.inner.read().unwrap(); - locked.nodes.clone() - }; + let mut nodes_online = HashSet::new(); + + // TODO: give Service a cancellation token for clean shutdown + let cancel = CancellationToken::new(); // TODO: issue these requests concurrently - for node in nodes.values() { - let client = mgmt_api::Client::new(node.base_url(), self.config.jwt_token.as_deref()); + { + let nodes = { + let locked = self.inner.read().unwrap(); + locked.nodes.clone() + }; + for node in nodes.values() { + let http_client = reqwest::ClientBuilder::new() + .timeout(Duration::from_secs(5)) + .build() + .expect("Failed to construct HTTP client"); + let client = mgmt_api::Client::from_client( + http_client, + node.base_url(), + self.config.jwt_token.as_deref(), + ); - tracing::info!("Scanning shards on node {}...", node.id); - match client.list_location_config().await { - Err(e) => { - tracing::warn!("Could not contact pageserver {} ({e})", node.id); - // TODO: be more tolerant, apply a generous 5-10 second timeout with retries, in case - // pageserver is being restarted at the same time as we are + fn is_fatal(e: &mgmt_api::Error) -> bool { + use mgmt_api::Error::*; + match e { + ReceiveBody(_) | ReceiveErrorBody(_) => false, + ApiError(StatusCode::SERVICE_UNAVAILABLE, _) + | ApiError(StatusCode::GATEWAY_TIMEOUT, _) + | ApiError(StatusCode::REQUEST_TIMEOUT, _) => false, + ApiError(_, _) => true, + } } - Ok(listing) => { - tracing::info!( - "Received {} shard statuses from pageserver {}, setting it to Active", - listing.tenant_shards.len(), - node.id - ); - for (tenant_shard_id, conf_opt) in listing.tenant_shards { - observed.insert(tenant_shard_id, (node.id, conf_opt)); + let list_response = backoff::retry( + || client.list_location_config(), + is_fatal, + 1, + 5, + "Location config listing", + &cancel, + ) + .await; + let Some(list_response) = list_response else { + tracing::info!("Shutdown during startup_reconcile"); + return; + }; + + tracing::info!("Scanning shards on node {}...", node.id); + match list_response { + Err(e) => { + tracing::warn!("Could not contact pageserver {} ({e})", node.id); + // TODO: be more tolerant, do some retries, in case + // pageserver is being restarted at the same time as we are + } + Ok(listing) => { + tracing::info!( + "Received {} shard statuses from pageserver {}, setting it to Active", + listing.tenant_shards.len(), + node.id + ); + nodes_online.insert(node.id); + + for (tenant_shard_id, conf_opt) in listing.tenant_shards { + observed.insert(tenant_shard_id, (node.id, conf_opt)); + } } } } @@ -185,8 +226,19 @@ impl Service { let mut compute_notifications = Vec::new(); // Populate intent and observed states for all tenants, based on reported state on pageservers - let shard_count = { + let (shard_count, nodes) = { let mut locked = self.inner.write().unwrap(); + + // Mark nodes online if they responded to us: nodes are offline by default after a restart. + let mut nodes = (*locked.nodes).clone(); + for (node_id, node) in nodes.iter_mut() { + if nodes_online.contains(node_id) { + node.availability = NodeAvailability::Active; + } + } + locked.nodes = Arc::new(nodes); + let nodes = locked.nodes.clone(); + for (tenant_shard_id, (node_id, observed_loc)) in observed { let Some(tenant_state) = locked.tenants.get_mut(&tenant_shard_id) else { cleanup.push((tenant_shard_id, node_id)); @@ -218,7 +270,7 @@ impl Service { } } - locked.tenants.len() + (locked.tenants.len(), nodes) }; // TODO: if any tenant's intent now differs from its loaded generation_pageserver, we should clear that @@ -279,9 +331,8 @@ impl Service { let stream = futures::stream::iter(compute_notifications.into_iter()) .map(|(tenant_shard_id, node_id)| { let compute_hook = compute_hook.clone(); + let cancel = cancel.clone(); async move { - // TODO: give Service a cancellation token for clean shutdown - let cancel = CancellationToken::new(); if let Err(e) = compute_hook.notify(tenant_shard_id, node_id, &cancel).await { tracing::error!( tenant_shard_id=%tenant_shard_id, @@ -387,7 +438,7 @@ impl Service { ))), config, persistence, - startup_complete, + startup_complete: startup_complete.clone(), }); let result_task_this = this.clone(); @@ -984,6 +1035,10 @@ impl Service { } }; + // TODO: if we timeout/fail on reconcile, we should still succeed this request, + // because otherwise a broken compute hook causes a feedback loop where + // location_config returns 500 and gets retried forever. + if let Some(create_req) = maybe_create { let create_resp = self.tenant_create(create_req).await?; result.shards = create_resp diff --git a/libs/utils/src/completion.rs b/libs/utils/src/completion.rs index ca6827c9b8..ea05cf54b1 100644 --- a/libs/utils/src/completion.rs +++ b/libs/utils/src/completion.rs @@ -27,6 +27,11 @@ impl Barrier { b.wait().await } } + + /// Return true if a call to wait() would complete immediately + pub fn is_ready(&self) -> bool { + futures::future::FutureExt::now_or_never(self.0.wait()).is_some() + } } impl PartialEq for Barrier { diff --git a/pageserver/client/src/mgmt_api.rs b/pageserver/client/src/mgmt_api.rs index 200369df90..baea747d3c 100644 --- a/pageserver/client/src/mgmt_api.rs +++ b/pageserver/client/src/mgmt_api.rs @@ -56,10 +56,18 @@ pub enum ForceAwaitLogicalSize { impl Client { pub fn new(mgmt_api_endpoint: String, jwt: Option<&str>) -> Self { + Self::from_client(reqwest::Client::new(), mgmt_api_endpoint, jwt) + } + + pub fn from_client( + client: reqwest::Client, + mgmt_api_endpoint: String, + jwt: Option<&str>, + ) -> Self { Self { mgmt_api_endpoint, authorization_header: jwt.map(|jwt| format!("Bearer {jwt}")), - client: reqwest::Client::new(), + client, } } diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 3d2549a8c3..0af8098cad 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1949,6 +1949,15 @@ class NeonAttachmentService: return headers + def ready(self) -> bool: + resp = self.request("GET", f"{self.env.attachment_service_api}/ready") + if resp.status_code == 503: + return False + elif resp.status_code == 200: + return True + else: + raise RuntimeError(f"Unexpected status {resp.status_code} from readiness endpoint") + def attach_hook_issue( self, tenant_shard_id: Union[TenantId, TenantShardId], pageserver_id: int ) -> int: diff --git a/test_runner/regress/test_sharding_service.py b/test_runner/regress/test_sharding_service.py index fd811a9d02..babb0d261c 100644 --- a/test_runner/regress/test_sharding_service.py +++ b/test_runner/regress/test_sharding_service.py @@ -128,6 +128,38 @@ def test_sharding_service_smoke( assert counts[env.pageservers[2].id] == tenant_shard_count // 2 +def test_node_status_after_restart( + neon_env_builder: NeonEnvBuilder, +): + neon_env_builder.num_pageservers = 2 + env = neon_env_builder.init_start() + + # Initially we have two online pageservers + nodes = env.attachment_service.node_list() + assert len(nodes) == 2 + + env.pageservers[1].stop() + + env.attachment_service.stop() + env.attachment_service.start() + + # Initially readiness check should fail because we're trying to connect to the offline node + assert env.attachment_service.ready() is False + + def is_ready(): + assert env.attachment_service.ready() is True + + wait_until(30, 1, is_ready) + + # We loaded nodes from database on restart + nodes = env.attachment_service.node_list() + assert len(nodes) == 2 + + # We should still be able to create a tenant, because the pageserver which is still online + # should have had its availabilty state set to Active. + env.attachment_service.tenant_create(TenantId.generate()) + + def test_sharding_service_passthrough( neon_env_builder: NeonEnvBuilder, ): From c0e0fc8151f2c00d45ebb8e39ef3c271c65a38f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Thu, 8 Feb 2024 19:57:02 +0100 Subject: [PATCH 46/72] Update Rust to 1.76.0 (#6683) [Release notes](https://github.com/rust-lang/rust/releases/tag/1.75.0). --- Dockerfile.buildtools | 2 +- compute_tools/src/pg_helpers.rs | 5 +++-- control_plane/src/background_process.rs | 1 - rust-toolchain.toml | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Dockerfile.buildtools b/Dockerfile.buildtools index 220e995d64..3a452fec32 100644 --- a/Dockerfile.buildtools +++ b/Dockerfile.buildtools @@ -135,7 +135,7 @@ WORKDIR /home/nonroot # Rust # Please keep the version of llvm (installed above) in sync with rust llvm (`rustc --version --verbose | grep LLVM`) -ENV RUSTC_VERSION=1.75.0 +ENV RUSTC_VERSION=1.76.0 ENV RUSTUP_HOME="/home/nonroot/.rustup" ENV PATH="/home/nonroot/.cargo/bin:${PATH}" RUN curl -sSO https://static.rust-lang.org/rustup/dist/$(uname -m)-unknown-linux-gnu/rustup-init && whoami && \ diff --git a/compute_tools/src/pg_helpers.rs b/compute_tools/src/pg_helpers.rs index ce704385c6..5deb50d6b7 100644 --- a/compute_tools/src/pg_helpers.rs +++ b/compute_tools/src/pg_helpers.rs @@ -264,9 +264,10 @@ pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> { // case we miss some events for some reason. Not strictly necessary, but // better safe than sorry. let (tx, rx) = std::sync::mpsc::channel(); - let (mut watcher, rx): (Box, _) = match notify::recommended_watcher(move |res| { + let watcher_res = notify::recommended_watcher(move |res| { let _ = tx.send(res); - }) { + }); + let (mut watcher, rx): (Box, _) = match watcher_res { Ok(watcher) => (Box::new(watcher), rx), Err(e) => { match e.kind { diff --git a/control_plane/src/background_process.rs b/control_plane/src/background_process.rs index 364cc01c39..0e59b28230 100644 --- a/control_plane/src/background_process.rs +++ b/control_plane/src/background_process.rs @@ -72,7 +72,6 @@ where let log_path = datadir.join(format!("{process_name}.log")); let process_log_file = fs::OpenOptions::new() .create(true) - .write(true) .append(true) .open(&log_path) .with_context(|| { diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 9b5a965f7d..b0949c32b1 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,5 +1,5 @@ [toolchain] -channel = "1.75.0" +channel = "1.76.0" profile = "default" # The default profile includes rustc, rust-std, cargo, rust-docs, rustfmt and clippy. # https://rust-lang.github.io/rustup/concepts/profiles.html From 9a31311990d19eb607e087e0e12d4369bfab8b6c Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Thu, 8 Feb 2024 22:40:14 +0200 Subject: [PATCH 47/72] fix(heavier_once_cell): assertion failure can be hit (#6652) @problame noticed that the `tokio::sync::AcquireError` branch assertion can be hit like in the first commit. We haven't seen this yet in production, but I'd prefer not to see it there. There `take_and_deinit` is being used, but this race must be quite timing sensitive. --- libs/utils/src/sync/heavier_once_cell.rs | 241 +++++++++++++++++------ 1 file changed, 176 insertions(+), 65 deletions(-) diff --git a/libs/utils/src/sync/heavier_once_cell.rs b/libs/utils/src/sync/heavier_once_cell.rs index f733d107f1..81625b907e 100644 --- a/libs/utils/src/sync/heavier_once_cell.rs +++ b/libs/utils/src/sync/heavier_once_cell.rs @@ -69,37 +69,44 @@ impl OnceCell { F: FnOnce(InitPermit) -> Fut, Fut: std::future::Future>, { - let sem = { + loop { + let sem = { + let guard = self.inner.write().await; + if guard.value.is_some() { + return Ok(GuardMut(guard)); + } + guard.init_semaphore.clone() + }; + + { + let permit = { + // increment the count for the duration of queued + let _guard = CountWaitingInitializers::start(self); + sem.acquire().await + }; + + let Ok(permit) = permit else { + let guard = self.inner.write().await; + if !Arc::ptr_eq(&sem, &guard.init_semaphore) { + // there was a take_and_deinit in between + continue; + } + assert!( + guard.value.is_some(), + "semaphore got closed, must be initialized" + ); + return Ok(GuardMut(guard)); + }; + + permit.forget(); + } + + let permit = InitPermit(sem); + let (value, _permit) = factory(permit).await?; + let guard = self.inner.write().await; - if guard.value.is_some() { - return Ok(GuardMut(guard)); - } - guard.init_semaphore.clone() - }; - let permit = { - // increment the count for the duration of queued - let _guard = CountWaitingInitializers::start(self); - sem.acquire_owned().await - }; - - match permit { - Ok(permit) => { - let permit = InitPermit(permit); - let (value, _permit) = factory(permit).await?; - - let guard = self.inner.write().await; - - Ok(Self::set0(value, guard)) - } - Err(_closed) => { - let guard = self.inner.write().await; - assert!( - guard.value.is_some(), - "semaphore got closed, must be initialized" - ); - return Ok(GuardMut(guard)); - } + return Ok(Self::set0(value, guard)); } } @@ -112,37 +119,44 @@ impl OnceCell { F: FnOnce(InitPermit) -> Fut, Fut: std::future::Future>, { - let sem = { - let guard = self.inner.read().await; - if guard.value.is_some() { - return Ok(GuardRef(guard)); - } - guard.init_semaphore.clone() - }; - - let permit = { - // increment the count for the duration of queued - let _guard = CountWaitingInitializers::start(self); - sem.acquire_owned().await - }; - - match permit { - Ok(permit) => { - let permit = InitPermit(permit); - let (value, _permit) = factory(permit).await?; - - let guard = self.inner.write().await; - - Ok(Self::set0(value, guard).downgrade()) - } - Err(_closed) => { + loop { + let sem = { let guard = self.inner.read().await; - assert!( - guard.value.is_some(), - "semaphore got closed, must be initialized" - ); - return Ok(GuardRef(guard)); + if guard.value.is_some() { + return Ok(GuardRef(guard)); + } + guard.init_semaphore.clone() + }; + + { + let permit = { + // increment the count for the duration of queued + let _guard = CountWaitingInitializers::start(self); + sem.acquire().await + }; + + let Ok(permit) = permit else { + let guard = self.inner.read().await; + if !Arc::ptr_eq(&sem, &guard.init_semaphore) { + // there was a take_and_deinit in between + continue; + } + assert!( + guard.value.is_some(), + "semaphore got closed, must be initialized" + ); + return Ok(GuardRef(guard)); + }; + + permit.forget(); } + + let permit = InitPermit(sem); + let (value, _permit) = factory(permit).await?; + + let guard = self.inner.write().await; + + return Ok(Self::set0(value, guard).downgrade()); } } @@ -250,15 +264,12 @@ impl<'a, T> GuardMut<'a, T> { /// [`OnceCell::get_or_init`] will wait on it to complete. pub fn take_and_deinit(&mut self) -> (T, InitPermit) { let mut swapped = Inner::default(); - let permit = swapped - .init_semaphore - .clone() - .try_acquire_owned() - .expect("we just created this"); + let sem = swapped.init_semaphore.clone(); + sem.try_acquire().expect("we just created this").forget(); std::mem::swap(&mut *self.0, &mut swapped); swapped .value - .map(|v| (v, InitPermit(permit))) + .map(|v| (v, InitPermit(sem))) .expect("guard is not created unless value has been initialized") } @@ -282,13 +293,23 @@ impl std::ops::Deref for GuardRef<'_, T> { } /// Type held by OnceCell (de)initializing task. -pub struct InitPermit(tokio::sync::OwnedSemaphorePermit); +pub struct InitPermit(Arc); + +impl Drop for InitPermit { + fn drop(&mut self) { + debug_assert_eq!(self.0.available_permits(), 0); + self.0.add_permits(1); + } +} #[cfg(test)] mod tests { + use futures::Future; + use super::*; use std::{ convert::Infallible, + pin::{pin, Pin}, sync::atomic::{AtomicUsize, Ordering}, time::Duration, }; @@ -455,4 +476,94 @@ mod tests { .unwrap(); assert_eq!(*g, "now initialized"); } + + #[tokio::test(start_paused = true)] + async fn reproduce_init_take_deinit_race() { + init_take_deinit_scenario(|cell, factory| { + Box::pin(async { + cell.get_or_init(factory).await.unwrap(); + }) + }) + .await; + } + + #[tokio::test(start_paused = true)] + async fn reproduce_init_take_deinit_race_mut() { + init_take_deinit_scenario(|cell, factory| { + Box::pin(async { + cell.get_mut_or_init(factory).await.unwrap(); + }) + }) + .await; + } + + type BoxedInitFuture = Pin>>>; + type BoxedInitFunction = Box BoxedInitFuture>; + + /// Reproduce an assertion failure with both initialization methods. + /// + /// This has interesting generics to be generic between `get_or_init` and `get_mut_or_init`. + /// Alternative would be a macro_rules! but that is the last resort. + async fn init_take_deinit_scenario(init_way: F) + where + F: for<'a> Fn( + &'a OnceCell<&'static str>, + BoxedInitFunction<&'static str, Infallible>, + ) -> Pin + 'a>>, + { + let cell = OnceCell::default(); + + // acquire the init_semaphore only permit to drive initializing tasks in order to waiting + // on the same semaphore. + let permit = cell + .inner + .read() + .await + .init_semaphore + .clone() + .try_acquire_owned() + .unwrap(); + + let mut t1 = pin!(init_way( + &cell, + Box::new(|permit| Box::pin(async move { Ok(("t1", permit)) })), + )); + + let mut t2 = pin!(init_way( + &cell, + Box::new(|permit| Box::pin(async move { Ok(("t2", permit)) })), + )); + + // drive t2 first to the init_semaphore + tokio::select! { + _ = &mut t2 => unreachable!("it cannot get permit"), + _ = tokio::time::sleep(Duration::from_secs(3600 * 24 * 7 * 365)) => {} + } + + // followed by t1 in the init_semaphore + tokio::select! { + _ = &mut t1 => unreachable!("it cannot get permit"), + _ = tokio::time::sleep(Duration::from_secs(3600 * 24 * 7 * 365)) => {} + } + + // now let t2 proceed and initialize + drop(permit); + t2.await; + + let (s, permit) = { cell.get_mut().await.unwrap().take_and_deinit() }; + assert_eq!("t2", s); + + // now originally t1 would see the semaphore it has as closed. it cannot yet get a permit from + // the new one. + tokio::select! { + _ = &mut t1 => unreachable!("it cannot get permit"), + _ = tokio::time::sleep(Duration::from_secs(3600 * 24 * 7 * 365)) => {} + } + + // only now we get to initialize it + drop(permit); + t1.await; + + assert_eq!("t1", *cell.get().await.unwrap()); + } } From c09993396ea026758bfda83c477361d656a5b647 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Fri, 9 Feb 2024 00:37:57 +0200 Subject: [PATCH 48/72] fix: secondary tenant relative order eviction (#6491) Calculate the `relative_last_activity` using the total evicted and resident layers similar to what we originally planned. Cc: #5331 --- pageserver/src/disk_usage_eviction_task.rs | 73 +++++++++++++------ pageserver/src/tenant/secondary.rs | 2 +- pageserver/src/tenant/secondary/downloader.rs | 27 ++++--- 3 files changed, 67 insertions(+), 35 deletions(-) diff --git a/pageserver/src/disk_usage_eviction_task.rs b/pageserver/src/disk_usage_eviction_task.rs index 1f0525b045..d5f5a20683 100644 --- a/pageserver/src/disk_usage_eviction_task.rs +++ b/pageserver/src/disk_usage_eviction_task.rs @@ -623,6 +623,7 @@ impl std::fmt::Display for EvictionLayer { } } +#[derive(Default)] pub(crate) struct DiskUsageEvictionInfo { /// Timeline's largest layer (remote or resident) pub max_layer_size: Option, @@ -854,19 +855,27 @@ async fn collect_eviction_candidates( let total = tenant_candidates.len(); - for (i, mut candidate) in tenant_candidates.into_iter().enumerate() { - // as we iterate this reverse sorted list, the most recently accessed layer will always - // be 1.0; this is for us to evict it last. - candidate.relative_last_activity = eviction_order.relative_last_activity(total, i); + let tenant_candidates = + tenant_candidates + .into_iter() + .enumerate() + .map(|(i, mut candidate)| { + // as we iterate this reverse sorted list, the most recently accessed layer will always + // be 1.0; this is for us to evict it last. + candidate.relative_last_activity = + eviction_order.relative_last_activity(total, i); - let partition = if cumsum > min_resident_size as i128 { - MinResidentSizePartition::Above - } else { - MinResidentSizePartition::Below - }; - cumsum += i128::from(candidate.layer.get_file_size()); - candidates.push((partition, candidate)); - } + let partition = if cumsum > min_resident_size as i128 { + MinResidentSizePartition::Above + } else { + MinResidentSizePartition::Below + }; + cumsum += i128::from(candidate.layer.get_file_size()); + + (partition, candidate) + }); + + candidates.extend(tenant_candidates); } // Note: the same tenant ID might be hit twice, if it transitions from attached to @@ -882,21 +891,41 @@ async fn collect_eviction_candidates( ); for secondary_tenant in secondary_tenants { - let mut layer_info = secondary_tenant.get_layers_for_eviction(); + // for secondary tenants we use a sum of on_disk layers and already evicted layers. this is + // to prevent repeated disk usage based evictions from completely draining less often + // updating secondaries. + let (mut layer_info, total_layers) = secondary_tenant.get_layers_for_eviction(); + + debug_assert!( + total_layers >= layer_info.resident_layers.len(), + "total_layers ({total_layers}) must be at least the resident_layers.len() ({})", + layer_info.resident_layers.len() + ); layer_info .resident_layers .sort_unstable_by_key(|layer_info| std::cmp::Reverse(layer_info.last_activity_ts)); - candidates.extend(layer_info.resident_layers.into_iter().map(|candidate| { - ( - // Secondary locations' layers are always considered above the min resident size, - // i.e. secondary locations are permitted to be trimmed to zero layers if all - // the layers have sufficiently old access times. - MinResidentSizePartition::Above, - candidate, - ) - })); + let tenant_candidates = + layer_info + .resident_layers + .into_iter() + .enumerate() + .map(|(i, mut candidate)| { + candidate.relative_last_activity = + eviction_order.relative_last_activity(total_layers, i); + ( + // Secondary locations' layers are always considered above the min resident size, + // i.e. secondary locations are permitted to be trimmed to zero layers if all + // the layers have sufficiently old access times. + MinResidentSizePartition::Above, + candidate, + ) + }); + + candidates.extend(tenant_candidates); + + tokio::task::yield_now().await; } debug_assert!(MinResidentSizePartition::Above < MinResidentSizePartition::Below, diff --git a/pageserver/src/tenant/secondary.rs b/pageserver/src/tenant/secondary.rs index 4269e1dec1..926cd0302b 100644 --- a/pageserver/src/tenant/secondary.rs +++ b/pageserver/src/tenant/secondary.rs @@ -160,7 +160,7 @@ impl SecondaryTenant { &self.tenant_shard_id } - pub(crate) fn get_layers_for_eviction(self: &Arc) -> DiskUsageEvictionInfo { + pub(crate) fn get_layers_for_eviction(self: &Arc) -> (DiskUsageEvictionInfo, usize) { self.detail.lock().unwrap().get_layers_for_eviction(self) } diff --git a/pageserver/src/tenant/secondary/downloader.rs b/pageserver/src/tenant/secondary/downloader.rs index 55af4f9f2b..9330edf946 100644 --- a/pageserver/src/tenant/secondary/downloader.rs +++ b/pageserver/src/tenant/secondary/downloader.rs @@ -146,14 +146,15 @@ impl SecondaryDetail { } } + /// Additionally returns the total number of layers, used for more stable relative access time + /// based eviction. pub(super) fn get_layers_for_eviction( &self, parent: &Arc, - ) -> DiskUsageEvictionInfo { - let mut result = DiskUsageEvictionInfo { - max_layer_size: None, - resident_layers: Vec::new(), - }; + ) -> (DiskUsageEvictionInfo, usize) { + let mut result = DiskUsageEvictionInfo::default(); + let mut total_layers = 0; + for (timeline_id, timeline_detail) in &self.timelines { result .resident_layers @@ -169,6 +170,10 @@ impl SecondaryDetail { relative_last_activity: finite_f32::FiniteF32::ZERO, } })); + + // total might be missing currently downloading layers, but as a lower than actual + // value it is good enough approximation. + total_layers += timeline_detail.on_disk_layers.len() + timeline_detail.evicted_at.len(); } result.max_layer_size = result .resident_layers @@ -183,7 +188,7 @@ impl SecondaryDetail { result.resident_layers.len() ); - result + (result, total_layers) } } @@ -312,9 +317,7 @@ impl JobGenerator Date: Fri, 9 Feb 2024 08:14:41 +0200 Subject: [PATCH 49/72] Increment generation which LFC is disabled by assigning 0 to neon.file_cache_size_limit (#6692) ## Problem test_lfc_resize sometimes filed with assertion failure when require lock in write operation: ``` if (lfc_ctl->generation == generation) { Assert(LFC_ENABLED()); ``` ## Summary of changes Increment generation when 0 is assigned to neon.file_cache_size_limit ## Checklist before requesting a review - [ ] I have performed a self-review of my code. - [ ] If it is a core feature, I have added thorough tests. - [ ] Do we need to implement analytics? if so did you add the relevant metrics to the dashboard? - [ ] If this PR requires public announcement, mark it with /release-notes label and add several sentences in this section. ## Checklist before merging - [ ] Do not forget to reformat commit message to not include the above checklist Co-authored-by: Konstantin Knizhnik --- pgxn/neon/file_cache.c | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pgxn/neon/file_cache.c b/pgxn/neon/file_cache.c index 21db666caa..448b9263f3 100644 --- a/pgxn/neon/file_cache.c +++ b/pgxn/neon/file_cache.c @@ -314,6 +314,9 @@ lfc_change_limit_hook(int newval, void *extra) lfc_ctl->used -= 1; } lfc_ctl->limit = new_size; + if (new_size == 0) { + lfc_ctl->generation += 1; + } neon_log(DEBUG1, "set local file cache limit to %d", new_size); LWLockRelease(lfc_lock); From a18aa14754fc44f7b38970bc546e4340386c32c9 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Fri, 9 Feb 2024 11:01:07 +0200 Subject: [PATCH 50/72] test: shutdown endpoints before deletion (#6619) this avoids a page_service error in the log sometimes. keeping the endpoint running while deleting has no function for this test. --- test_runner/regress/test_timeline_delete.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/test_runner/regress/test_timeline_delete.py b/test_runner/regress/test_timeline_delete.py index 352b82d525..5fda5aa569 100644 --- a/test_runner/regress/test_timeline_delete.py +++ b/test_runner/regress/test_timeline_delete.py @@ -651,9 +651,7 @@ def test_timeline_delete_works_for_remote_smoke( timeline_ids = [env.initial_timeline] for i in range(2): branch_timeline_id = env.neon_cli.create_branch(f"new{i}", "main") - pg = env.endpoints.create_start(f"new{i}") - - with pg.cursor() as cur: + with env.endpoints.create_start(f"new{i}") as pg, pg.cursor() as cur: cur.execute("CREATE TABLE f (i integer);") cur.execute("INSERT INTO f VALUES (generate_series(1,1000));") current_lsn = Lsn(query_scalar(cur, "SELECT pg_current_wal_flush_lsn()")) From 568f91420a9c677e77aeb736cb3f995a85f0b106 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Fri, 9 Feb 2024 11:34:15 +0200 Subject: [PATCH 51/72] tests: try to make restored-datadir comparison tests not flaky (#6666) This test occasionally fails with a difference in "pg_xact/0000" file between the local and restored datadirs. My hypothesis is that something changed in the database between the last explicit checkpoint and the shutdown. I suspect autovacuum, it could certainly create transactions. To fix, be more precise about the point in time that we compare. Shut down the endpoint first, then read the last LSN (i.e. the shutdown checkpoint's LSN), from the local disk with pg_controldata. And use exactly that LSN in the basebackup. Closes #559. I'm proposing this as an alternative to https://github.com/neondatabase/neon/pull/6662. --- test_runner/fixtures/neon_fixtures.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 0af8098cad..a6aff77ddf 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -3964,24 +3964,27 @@ def list_files_to_compare(pgdata_dir: Path) -> List[str]: # pg is the existing and running compute node, that we want to compare with a basebackup def check_restored_datadir_content(test_output_dir: Path, env: NeonEnv, endpoint: Endpoint): + pg_bin = PgBin(test_output_dir, env.pg_distrib_dir, env.pg_version) + # Get the timeline ID. We need it for the 'basebackup' command timeline_id = TimelineId(endpoint.safe_psql("SHOW neon.timeline_id")[0][0]) - # many tests already checkpoint, but do it just in case - with closing(endpoint.connect()) as conn: - with conn.cursor() as cur: - cur.execute("CHECKPOINT") - - # wait for pageserver to catch up - wait_for_last_flush_lsn(env, endpoint, endpoint.tenant_id, timeline_id) # stop postgres to ensure that files won't change endpoint.stop() + # Read the shutdown checkpoint's LSN + pg_controldata_path = os.path.join(pg_bin.pg_bin_path, "pg_controldata") + cmd = f"{pg_controldata_path} -D {endpoint.pgdata_dir}" + result = subprocess.run(cmd, capture_output=True, text=True, shell=True) + checkpoint_lsn = re.findall( + "Latest checkpoint location:\\s+([0-9A-F]+/[0-9A-F]+)", result.stdout + )[0] + log.debug(f"last checkpoint at {checkpoint_lsn}") + # Take a basebackup from pageserver restored_dir_path = env.repo_dir / f"{endpoint.endpoint_id}_restored_datadir" restored_dir_path.mkdir(exist_ok=True) - pg_bin = PgBin(test_output_dir, env.pg_distrib_dir, env.pg_version) psql_path = os.path.join(pg_bin.pg_bin_path, "psql") pageserver_id = env.attachment_service.locate(endpoint.tenant_id)[0]["node_id"] @@ -3989,7 +3992,7 @@ def check_restored_datadir_content(test_output_dir: Path, env: NeonEnv, endpoint {psql_path} \ --no-psqlrc \ postgres://localhost:{env.get_pageserver(pageserver_id).service_port.pg} \ - -c 'basebackup {endpoint.tenant_id} {timeline_id}' \ + -c 'basebackup {endpoint.tenant_id} {timeline_id} {checkpoint_lsn}' \ | tar -x -C {restored_dir_path} """ From 951c9bf4cad6a651f9531f3c4e1e58d90c27910e Mon Sep 17 00:00:00 2001 From: John Spray Date: Fri, 9 Feb 2024 10:12:40 +0000 Subject: [PATCH 52/72] control_plane: fix shard splitting on unsharded tenant (#6689) ## Problem Previous test started with a new-style TenantShardId with a non-zero ShardCount. We also need to handle the case of a ShardCount() (aka `unsharded`) parent shard. **A followup PR will refactor ShardCount to make its inner value private and thereby make this kind of mistake harder** ## Summary of changes - Fix a place we were incorrectly treating a ShardCount as a number of shards rather than as thing that can be zero or the number of shards. - Add a test for this case. --- .../attachment_service/src/persistence.rs | 10 ++++-- test_runner/regress/test_sharding.py | 31 ++++++++++++++++++- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/control_plane/attachment_service/src/persistence.rs b/control_plane/attachment_service/src/persistence.rs index cead540058..623d625767 100644 --- a/control_plane/attachment_service/src/persistence.rs +++ b/control_plane/attachment_service/src/persistence.rs @@ -381,16 +381,22 @@ impl Persistence { self.with_conn(move |conn| -> DatabaseResult<()> { conn.transaction(|conn| -> DatabaseResult<()> { // Mark parent shards as splitting + + let expect_parent_records = std::cmp::max(1, old_shard_count.0); + let updated = diesel::update(tenant_shards) .filter(tenant_id.eq(split_tenant_id.to_string())) .filter(shard_count.eq(old_shard_count.0 as i32)) .set((splitting.eq(1),)) .execute(conn)?; - if ShardCount(updated.try_into().map_err(|_| DatabaseError::Logical(format!("Overflow existing shard count {} while splitting", updated)))?) != old_shard_count { + if u8::try_from(updated) + .map_err(|_| DatabaseError::Logical( + format!("Overflow existing shard count {} while splitting", updated)) + )? != expect_parent_records { // Perhaps a deletion or another split raced with this attempt to split, mutating // the parent shards that we intend to split. In this case the split request should fail. return Err(DatabaseError::Logical( - format!("Unexpected existing shard count {updated} when preparing tenant for split (expected {old_shard_count:?})") + format!("Unexpected existing shard count {updated} when preparing tenant for split (expected {expect_parent_records})") )); } diff --git a/test_runner/regress/test_sharding.py b/test_runner/regress/test_sharding.py index 805eaa34b0..27d1cf2f34 100644 --- a/test_runner/regress/test_sharding.py +++ b/test_runner/regress/test_sharding.py @@ -4,7 +4,7 @@ from fixtures.neon_fixtures import ( tenant_get_shards, ) from fixtures.remote_storage import s3_storage -from fixtures.types import TimelineId +from fixtures.types import TenantShardId, TimelineId from fixtures.workload import Workload @@ -84,6 +84,35 @@ def test_sharding_smoke( assert timelines == {env.initial_timeline, timeline_b} +def test_sharding_split_unsharded( + neon_env_builder: NeonEnvBuilder, +): + """ + Test that shard splitting works on a tenant created as unsharded (i.e. with + ShardCount(0)). + """ + env = neon_env_builder.init_start() + tenant_id = env.initial_tenant + timeline_id = env.initial_timeline + + workload = Workload(env, tenant_id, timeline_id, branch_name="main") + workload.init() + workload.write_rows(256) + + # Check that we created with an unsharded TenantShardId: this is the default, + # but check it in case we change the default in future + assert env.attachment_service.inspect(TenantShardId(tenant_id, 0, 0)) is not None + + # Split one shard into two + env.attachment_service.tenant_shard_split(tenant_id, shard_count=2) + + # Check we got the shard IDs we expected + assert env.attachment_service.inspect(TenantShardId(tenant_id, 0, 2)) is not None + assert env.attachment_service.inspect(TenantShardId(tenant_id, 1, 2)) is not None + + workload.validate() + + def test_sharding_split_smoke( neon_env_builder: NeonEnvBuilder, ): From ea089dc97700732788f2d9f0ea44e10fb59c2f6f Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Fri, 9 Feb 2024 10:29:20 +0000 Subject: [PATCH 53/72] proxy: add per query array mode flag (#6678) ## Problem Drizzle needs to be able to configure the array_mode flag per query. ## Summary of changes Adds an array_mode flag to the query data json that will otherwise default to the header flag. --- proxy/src/serverless/sql_over_http.rs | 163 ++++++++++++++------------ test_runner/regress/test_proxy.py | 33 ++++++ 2 files changed, 119 insertions(+), 77 deletions(-) diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 7092b65f03..25e8813625 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -44,10 +44,13 @@ use super::json::pg_text_row_to_json; use super::SERVERLESS_DRIVER_SNI; #[derive(serde::Deserialize)] +#[serde(rename_all = "camelCase")] struct QueryData { query: String, #[serde(deserialize_with = "bytes_to_pg_text")] params: Vec>, + #[serde(default)] + array_mode: Option, } #[derive(serde::Deserialize)] @@ -330,7 +333,7 @@ async fn handle_inner( // Determine the output options. Default behaviour is 'false'. Anything that is not // strictly 'true' assumed to be false. let raw_output = headers.get(&RAW_TEXT_OUTPUT) == Some(&HEADER_VALUE_TRUE); - let array_mode = headers.get(&ARRAY_MODE) == Some(&HEADER_VALUE_TRUE); + let default_array_mode = headers.get(&ARRAY_MODE) == Some(&HEADER_VALUE_TRUE); // Allow connection pooling only if explicitly requested // or if we have decided that http pool is no longer opt-in @@ -402,83 +405,87 @@ async fn handle_inner( // Now execute the query and return the result // let mut size = 0; - let result = - match payload { - Payload::Single(stmt) => { - let (status, results) = - query_to_json(&*client, stmt, &mut 0, raw_output, array_mode) - .await - .map_err(|e| { - client.discard(); - e - })?; - client.check_idle(status); - results + let result = match payload { + Payload::Single(stmt) => { + let (status, results) = + query_to_json(&*client, stmt, &mut 0, raw_output, default_array_mode) + .await + .map_err(|e| { + client.discard(); + e + })?; + client.check_idle(status); + results + } + Payload::Batch(statements) => { + let (inner, mut discard) = client.inner(); + let mut builder = inner.build_transaction(); + if let Some(isolation_level) = txn_isolation_level { + builder = builder.isolation_level(isolation_level); } - Payload::Batch(statements) => { - let (inner, mut discard) = client.inner(); - let mut builder = inner.build_transaction(); - if let Some(isolation_level) = txn_isolation_level { - builder = builder.isolation_level(isolation_level); - } - if txn_read_only { - builder = builder.read_only(true); - } - if txn_deferrable { - builder = builder.deferrable(true); - } - - let transaction = builder.start().await.map_err(|e| { - // if we cannot start a transaction, we should return immediately - // and not return to the pool. connection is clearly broken - discard.discard(); - e - })?; - - let results = - match query_batch(&transaction, statements, &mut size, raw_output, array_mode) - .await - { - Ok(results) => { - let status = transaction.commit().await.map_err(|e| { - // if we cannot commit - for now don't return connection to pool - // TODO: get a query status from the error - discard.discard(); - e - })?; - discard.check_idle(status); - results - } - Err(err) => { - let status = transaction.rollback().await.map_err(|e| { - // if we cannot rollback - for now don't return connection to pool - // TODO: get a query status from the error - discard.discard(); - e - })?; - discard.check_idle(status); - return Err(err); - } - }; - - if txn_read_only { - response = response.header( - TXN_READ_ONLY.clone(), - HeaderValue::try_from(txn_read_only.to_string())?, - ); - } - if txn_deferrable { - response = response.header( - TXN_DEFERRABLE.clone(), - HeaderValue::try_from(txn_deferrable.to_string())?, - ); - } - if let Some(txn_isolation_level) = txn_isolation_level_raw { - response = response.header(TXN_ISOLATION_LEVEL.clone(), txn_isolation_level); - } - json!({ "results": results }) + if txn_read_only { + builder = builder.read_only(true); } - }; + if txn_deferrable { + builder = builder.deferrable(true); + } + + let transaction = builder.start().await.map_err(|e| { + // if we cannot start a transaction, we should return immediately + // and not return to the pool. connection is clearly broken + discard.discard(); + e + })?; + + let results = match query_batch( + &transaction, + statements, + &mut size, + raw_output, + default_array_mode, + ) + .await + { + Ok(results) => { + let status = transaction.commit().await.map_err(|e| { + // if we cannot commit - for now don't return connection to pool + // TODO: get a query status from the error + discard.discard(); + e + })?; + discard.check_idle(status); + results + } + Err(err) => { + let status = transaction.rollback().await.map_err(|e| { + // if we cannot rollback - for now don't return connection to pool + // TODO: get a query status from the error + discard.discard(); + e + })?; + discard.check_idle(status); + return Err(err); + } + }; + + if txn_read_only { + response = response.header( + TXN_READ_ONLY.clone(), + HeaderValue::try_from(txn_read_only.to_string())?, + ); + } + if txn_deferrable { + response = response.header( + TXN_DEFERRABLE.clone(), + HeaderValue::try_from(txn_deferrable.to_string())?, + ); + } + if let Some(txn_isolation_level) = txn_isolation_level_raw { + response = response.header(TXN_ISOLATION_LEVEL.clone(), txn_isolation_level); + } + json!({ "results": results }) + } + }; ctx.set_success(); ctx.log(); @@ -524,7 +531,7 @@ async fn query_to_json( data: QueryData, current_size: &mut usize, raw_output: bool, - array_mode: bool, + default_array_mode: bool, ) -> anyhow::Result<(ReadyForQueryStatus, Value)> { let query_params = data.params; let row_stream = client.query_raw_txt(&data.query, query_params).await?; @@ -578,6 +585,8 @@ async fn query_to_json( columns.push(client.get_type(c.type_oid()).await?); } + let array_mode = data.array_mode.unwrap_or(default_array_mode); + // convert rows to JSON let rows = rows .iter() diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py index b3b35e446d..49a0450f0c 100644 --- a/test_runner/regress/test_proxy.py +++ b/test_runner/regress/test_proxy.py @@ -390,6 +390,39 @@ def test_sql_over_http_batch(static_proxy: NeonProxy): assert result[0]["rows"] == [{"answer": 42}] +def test_sql_over_http_batch_output_options(static_proxy: NeonProxy): + static_proxy.safe_psql("create role http with login password 'http' superuser") + + connstr = f"postgresql://http:http@{static_proxy.domain}:{static_proxy.proxy_port}/postgres" + response = requests.post( + f"https://{static_proxy.domain}:{static_proxy.external_http_port}/sql", + data=json.dumps( + { + "queries": [ + {"query": "select $1 as answer", "params": [42], "arrayMode": True}, + {"query": "select $1 as answer", "params": [42], "arrayMode": False}, + ] + } + ), + headers={ + "Content-Type": "application/sql", + "Neon-Connection-String": connstr, + "Neon-Batch-Isolation-Level": "Serializable", + "Neon-Batch-Read-Only": "false", + "Neon-Batch-Deferrable": "false", + }, + verify=str(static_proxy.test_output_dir / "proxy.crt"), + ) + assert response.status_code == 200 + results = response.json()["results"] + + assert results[0]["rowAsArray"] + assert results[0]["rows"] == [["42"]] + + assert not results[1]["rowAsArray"] + assert results[1]["rows"] == [{"answer": "42"}] + + def test_sql_over_http_pool(static_proxy: NeonProxy): static_proxy.safe_psql("create user http_auth with password 'http' superuser") From eec1e1a19223750e16401962c978fdeee2a305c8 Mon Sep 17 00:00:00 2001 From: Anastasia Lubennikova Date: Thu, 4 Jan 2024 12:34:15 +0000 Subject: [PATCH 54/72] Pre-install anon extension from compute_ctl if anon is in shared_preload_libraries. Users cannot install it themselves, because superuser is required. GRANT all priveleged needed to use it to db_owner We use the neon fork of the extension, because small change to sql file is needed to allow db_owner to use it. This feature is behind a feature flag AnonExtension, so it is not enabled by default. --- Dockerfile.compute-node | 5 +- compute_tools/src/compute.rs | 14 +++- compute_tools/src/spec.rs | 132 ++++++++++++++++++++++++++++++++++- libs/compute_api/src/spec.rs | 3 + 4 files changed, 149 insertions(+), 5 deletions(-) diff --git a/Dockerfile.compute-node b/Dockerfile.compute-node index d91c7cfd72..cc7a110008 100644 --- a/Dockerfile.compute-node +++ b/Dockerfile.compute-node @@ -639,8 +639,8 @@ FROM build-deps AS pg-anon-pg-build COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/ ENV PATH "/usr/local/pgsql/bin/:$PATH" -RUN wget https://gitlab.com/dalibo/postgresql_anonymizer/-/archive/1.1.0/postgresql_anonymizer-1.1.0.tar.gz -O pg_anon.tar.gz && \ - echo "08b09d2ff9b962f96c60db7e6f8e79cf7253eb8772516998fc35ece08633d3ad pg_anon.tar.gz" | sha256sum --check && \ +RUN wget https://github.com/neondatabase/postgresql_anonymizer/archive/refs/tags/neon_1.1.1.tar.gz -O pg_anon.tar.gz && \ + echo "321ea8d5c1648880aafde850a2c576e4a9e7b9933a34ce272efc839328999fa9 pg_anon.tar.gz" | sha256sum --check && \ mkdir pg_anon-src && cd pg_anon-src && tar xvzf ../pg_anon.tar.gz --strip-components=1 -C . && \ find /usr/local/pgsql -type f | sed 's|^/usr/local/pgsql/||' > /before.txt &&\ make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \ @@ -809,6 +809,7 @@ COPY --from=pg-roaringbitmap-pg-build /usr/local/pgsql/ /usr/local/pgsql/ COPY --from=pg-semver-pg-build /usr/local/pgsql/ /usr/local/pgsql/ COPY --from=pg-embedding-pg-build /usr/local/pgsql/ /usr/local/pgsql/ COPY --from=wal2json-pg-build /usr/local/pgsql /usr/local/pgsql +COPY --from=pg-anon-pg-build /usr/local/pgsql/ /usr/local/pgsql/ COPY pgxn/ pgxn/ RUN make -j $(getconf _NPROCESSORS_ONLN) \ diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 0ca1a47fbf..993b5725a4 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -765,7 +765,12 @@ impl ComputeNode { handle_roles(spec, &mut client)?; handle_databases(spec, &mut client)?; handle_role_deletions(spec, connstr.as_str(), &mut client)?; - handle_grants(spec, &mut client, connstr.as_str())?; + handle_grants( + spec, + &mut client, + connstr.as_str(), + self.has_feature(ComputeFeature::AnonExtension), + )?; handle_extensions(spec, &mut client)?; handle_extension_neon(&mut client)?; create_availability_check_data(&mut client)?; @@ -839,7 +844,12 @@ impl ComputeNode { handle_roles(&spec, &mut client)?; handle_databases(&spec, &mut client)?; handle_role_deletions(&spec, self.connstr.as_str(), &mut client)?; - handle_grants(&spec, &mut client, self.connstr.as_str())?; + handle_grants( + &spec, + &mut client, + self.connstr.as_str(), + self.has_feature(ComputeFeature::AnonExtension), + )?; handle_extensions(&spec, &mut client)?; handle_extension_neon(&mut client)?; // We can skip handle_migrations here because a new migration can only appear diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs index 2b1bff75fe..3df5f10e23 100644 --- a/compute_tools/src/spec.rs +++ b/compute_tools/src/spec.rs @@ -581,7 +581,12 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> { /// Grant CREATE ON DATABASE to the database owner and do some other alters and grants /// to allow users creating trusted extensions and re-creating `public` schema, for example. #[instrument(skip_all)] -pub fn handle_grants(spec: &ComputeSpec, client: &mut Client, connstr: &str) -> Result<()> { +pub fn handle_grants( + spec: &ComputeSpec, + client: &mut Client, + connstr: &str, + enable_anon_extension: bool, +) -> Result<()> { info!("modifying database permissions"); let existing_dbs = get_existing_dbs(client)?; @@ -678,6 +683,11 @@ pub fn handle_grants(spec: &ComputeSpec, client: &mut Client, connstr: &str) -> inlinify(&grant_query) ); db_client.simple_query(&grant_query)?; + + // it is important to run this after all grants + if enable_anon_extension { + handle_extension_anon(spec, &db.owner, &mut db_client, false)?; + } } Ok(()) @@ -809,5 +819,125 @@ $$;"#, "Ran {} migrations", (migrations.len() - starting_migration_id) ); + + Ok(()) +} + +/// Connect to the database as superuser and pre-create anon extension +/// if it is present in shared_preload_libraries +#[instrument(skip_all)] +pub fn handle_extension_anon( + spec: &ComputeSpec, + db_owner: &str, + db_client: &mut Client, + grants_only: bool, +) -> Result<()> { + info!("handle extension anon"); + + if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") { + if libs.contains("anon") { + if !grants_only { + // check if extension is already initialized using anon.is_initialized() + let query = "SELECT anon.is_initialized()"; + match db_client.query(query, &[]) { + Ok(rows) => { + if !rows.is_empty() { + let is_initialized: bool = rows[0].get(0); + if is_initialized { + info!("anon extension is already initialized"); + return Ok(()); + } + } + } + Err(e) => { + warn!( + "anon extension is_installed check failed with expected error: {}", + e + ); + } + }; + + // Create anon extension if this compute needs it + // Users cannot create it themselves, because superuser is required. + let mut query = "CREATE EXTENSION IF NOT EXISTS anon CASCADE"; + info!("creating anon extension with query: {}", query); + match db_client.query(query, &[]) { + Ok(_) => {} + Err(e) => { + error!("anon extension creation failed with error: {}", e); + return Ok(()); + } + } + + // check that extension is installed + query = "SELECT extname FROM pg_extension WHERE extname = 'anon'"; + let rows = db_client.query(query, &[])?; + if rows.is_empty() { + error!("anon extension is not installed"); + return Ok(()); + } + + // Initialize anon extension + // This also requires superuser privileges, so users cannot do it themselves. + query = "SELECT anon.init()"; + match db_client.query(query, &[]) { + Ok(_) => {} + Err(e) => { + error!("anon.init() failed with error: {}", e); + return Ok(()); + } + } + } + + // check that extension is installed, if not bail early + let query = "SELECT extname FROM pg_extension WHERE extname = 'anon'"; + match db_client.query(query, &[]) { + Ok(rows) => { + if rows.is_empty() { + error!("anon extension is not installed"); + return Ok(()); + } + } + Err(e) => { + error!("anon extension check failed with error: {}", e); + return Ok(()); + } + }; + + let query = format!("GRANT ALL ON SCHEMA anon TO {}", db_owner); + info!("granting anon extension permissions with query: {}", query); + db_client.simple_query(&query)?; + + // Grant permissions to db_owner to use anon extension functions + let query = format!("GRANT ALL ON ALL FUNCTIONS IN SCHEMA anon TO {}", db_owner); + info!("granting anon extension permissions with query: {}", query); + db_client.simple_query(&query)?; + + // This is needed, because some functions are defined as SECURITY DEFINER. + // In Postgres SECURITY DEFINER functions are executed with the privileges + // of the owner. + // In anon extension this it is needed to access some GUCs, which are only accessible to + // superuser. But we've patched postgres to allow db_owner to access them as well. + // So we need to change owner of these functions to db_owner. + let query = format!(" + SELECT 'ALTER FUNCTION '||nsp.nspname||'.'||p.proname||'('||pg_get_function_identity_arguments(p.oid)||') OWNER TO {};' + from pg_proc p + join pg_namespace nsp ON p.pronamespace = nsp.oid + where nsp.nspname = 'anon';", db_owner); + + info!("change anon extension functions owner to db owner"); + db_client.simple_query(&query)?; + + // affects views as well + let query = format!("GRANT ALL ON ALL TABLES IN SCHEMA anon TO {}", db_owner); + info!("granting anon extension permissions with query: {}", query); + db_client.simple_query(&query)?; + + let query = format!("GRANT ALL ON ALL SEQUENCES IN SCHEMA anon TO {}", db_owner); + info!("granting anon extension permissions with query: {}", query); + db_client.simple_query(&query)?; + } + } + Ok(()) } diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index 13ac18e0c5..2f412b61a3 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -90,6 +90,9 @@ pub enum ComputeFeature { /// track short-lived connections as user activity. ActivityMonitorExperimental, + /// Pre-install and initialize anon extension for every database in the cluster + AnonExtension, + /// This is a special feature flag that is used to represent unknown feature flags. /// Basically all unknown to enum flags are represented as this one. See unit test /// `parse_unknown_features()` for more details. From eb919cab88b8a28eb423b33eb07a858acbd61eab Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Fri, 9 Feb 2024 14:52:58 +0200 Subject: [PATCH 55/72] prepare to move timeouts and cancellation handling to remote_storage (#6696) This PR is preliminary cleanups and refactoring around `remote_storage` for next PR which will move the timeouts and cancellation into `remote_storage`. Summary: - smaller drive-by fixes - code simplification - refactor common parts like `DownloadError::is_permanent` - align error types with `RemoteStorage::list_*` to use more `download_retry` helper Cc: #6096 --- libs/remote_storage/src/lib.rs | 26 ++++++- libs/remote_storage/src/local_fs.rs | 50 ++++++++---- libs/remote_storage/src/s3_bucket.rs | 77 ++++++------------- libs/remote_storage/src/simulate_failures.rs | 28 ++++--- libs/remote_storage/src/support.rs | 33 ++++++++ pageserver/src/task_mgr.rs | 4 +- pageserver/src/tenant.rs | 4 +- .../src/tenant/remote_timeline_client.rs | 35 ++++----- .../tenant/remote_timeline_client/download.rs | 59 +++++--------- pageserver/src/tenant/secondary/downloader.rs | 2 +- 10 files changed, 175 insertions(+), 143 deletions(-) create mode 100644 libs/remote_storage/src/support.rs diff --git a/libs/remote_storage/src/lib.rs b/libs/remote_storage/src/lib.rs index e64b1de6f9..b6648931ac 100644 --- a/libs/remote_storage/src/lib.rs +++ b/libs/remote_storage/src/lib.rs @@ -13,6 +13,7 @@ mod azure_blob; mod local_fs; mod s3_bucket; mod simulate_failures; +mod support; use std::{ collections::HashMap, fmt::Debug, num::NonZeroUsize, pin::Pin, sync::Arc, time::SystemTime, @@ -170,7 +171,10 @@ pub trait RemoteStorage: Send + Sync + 'static { /// whereas, /// list_prefixes("foo/bar/") = ["cat", "dog"] /// See `test_real_s3.rs` for more details. - async fn list_files(&self, prefix: Option<&RemotePath>) -> anyhow::Result> { + async fn list_files( + &self, + prefix: Option<&RemotePath>, + ) -> Result, DownloadError> { let result = self.list(prefix, ListingMode::NoDelimiter).await?.keys; Ok(result) } @@ -179,7 +183,7 @@ pub trait RemoteStorage: Send + Sync + 'static { &self, prefix: Option<&RemotePath>, _mode: ListingMode, - ) -> anyhow::Result; + ) -> Result; /// Streams the local file contents into remote into the remote storage entry. async fn upload( @@ -269,6 +273,19 @@ impl std::fmt::Display for DownloadError { impl std::error::Error for DownloadError {} +impl DownloadError { + /// Returns true if the error should not be retried with backoff + pub fn is_permanent(&self) -> bool { + use DownloadError::*; + match self { + BadInput(_) => true, + NotFound => true, + Cancelled => true, + Other(_) => false, + } + } +} + #[derive(Debug)] pub enum TimeTravelError { /// Validation or other error happened due to user input. @@ -336,7 +353,10 @@ impl GenericRemoteStorage> { // A function for listing all the files in a "directory" // Example: // list_files("foo/bar") = ["foo/bar/a.txt", "foo/bar/b.txt"] - pub async fn list_files(&self, folder: Option<&RemotePath>) -> anyhow::Result> { + pub async fn list_files( + &self, + folder: Option<&RemotePath>, + ) -> Result, DownloadError> { match self { Self::LocalFs(s) => s.list_files(folder).await, Self::AwsS3(s) => s.list_files(folder).await, diff --git a/libs/remote_storage/src/local_fs.rs b/libs/remote_storage/src/local_fs.rs index 36ec15e1b1..3ebea76181 100644 --- a/libs/remote_storage/src/local_fs.rs +++ b/libs/remote_storage/src/local_fs.rs @@ -18,9 +18,7 @@ use tokio_util::{io::ReaderStream, sync::CancellationToken}; use tracing::*; use utils::{crashsafe::path_with_suffix_extension, fs_ext::is_directory_empty}; -use crate::{ - Download, DownloadError, DownloadStream, Listing, ListingMode, RemotePath, TimeTravelError, -}; +use crate::{Download, DownloadError, Listing, ListingMode, RemotePath, TimeTravelError}; use super::{RemoteStorage, StorageMetadata}; @@ -365,27 +363,33 @@ impl RemoteStorage for LocalFs { format!("Failed to open source file {target_path:?} to use in the download") }) .map_err(DownloadError::Other)?; + + let len = source + .metadata() + .await + .context("query file length") + .map_err(DownloadError::Other)? + .len(); + source .seek(io::SeekFrom::Start(start_inclusive)) .await .context("Failed to seek to the range start in a local storage file") .map_err(DownloadError::Other)?; + let metadata = self .read_storage_metadata(&target_path) .await .map_err(DownloadError::Other)?; - let download_stream: DownloadStream = match end_exclusive { - Some(end_exclusive) => Box::pin(ReaderStream::new( - source.take(end_exclusive - start_inclusive), - )), - None => Box::pin(ReaderStream::new(source)), - }; + let source = source.take(end_exclusive.unwrap_or(len) - start_inclusive); + let source = ReaderStream::new(source); + Ok(Download { metadata, last_modified: None, etag: None, - download_stream, + download_stream: Box::pin(source), }) } else { Err(DownloadError::NotFound) @@ -514,10 +518,8 @@ mod fs_tests { use futures_util::Stream; use std::{collections::HashMap, io::Write}; - async fn read_and_assert_remote_file_contents( + async fn read_and_check_metadata( storage: &LocalFs, - #[allow(clippy::ptr_arg)] - // have to use &Utf8PathBuf due to `storage.local_path` parameter requirements remote_storage_path: &RemotePath, expected_metadata: Option<&StorageMetadata>, ) -> anyhow::Result { @@ -596,7 +598,7 @@ mod fs_tests { let upload_name = "upload_1"; let upload_target = upload_dummy_file(&storage, upload_name, None).await?; - let contents = read_and_assert_remote_file_contents(&storage, &upload_target, None).await?; + let contents = read_and_check_metadata(&storage, &upload_target, None).await?; assert_eq!( dummy_contents(upload_name), contents, @@ -618,7 +620,7 @@ mod fs_tests { let upload_target = upload_dummy_file(&storage, upload_name, None).await?; let full_range_download_contents = - read_and_assert_remote_file_contents(&storage, &upload_target, None).await?; + read_and_check_metadata(&storage, &upload_target, None).await?; assert_eq!( dummy_contents(upload_name), full_range_download_contents, @@ -660,6 +662,22 @@ mod fs_tests { "Second part bytes should be returned when requested" ); + let suffix_bytes = storage + .download_byte_range(&upload_target, 13, None) + .await? + .download_stream; + let suffix_bytes = aggregate(suffix_bytes).await?; + let suffix = std::str::from_utf8(&suffix_bytes)?; + assert_eq!(upload_name, suffix); + + let all_bytes = storage + .download_byte_range(&upload_target, 0, None) + .await? + .download_stream; + let all_bytes = aggregate(all_bytes).await?; + let all_bytes = std::str::from_utf8(&all_bytes)?; + assert_eq!(dummy_contents("upload_1"), all_bytes); + Ok(()) } @@ -736,7 +754,7 @@ mod fs_tests { upload_dummy_file(&storage, upload_name, Some(metadata.clone())).await?; let full_range_download_contents = - read_and_assert_remote_file_contents(&storage, &upload_target, Some(&metadata)).await?; + read_and_check_metadata(&storage, &upload_target, Some(&metadata)).await?; assert_eq!( dummy_contents(upload_name), full_range_download_contents, diff --git a/libs/remote_storage/src/s3_bucket.rs b/libs/remote_storage/src/s3_bucket.rs index c9ad9ef225..2b33a6ffd1 100644 --- a/libs/remote_storage/src/s3_bucket.rs +++ b/libs/remote_storage/src/s3_bucket.rs @@ -45,8 +45,9 @@ use utils::backoff; use super::StorageMetadata; use crate::{ - ConcurrencyLimiter, Download, DownloadError, Listing, ListingMode, RemotePath, RemoteStorage, - S3Config, TimeTravelError, MAX_KEYS_PER_DELETE, REMOTE_STORAGE_PREFIX_SEPARATOR, + support::PermitCarrying, ConcurrencyLimiter, Download, DownloadError, Listing, ListingMode, + RemotePath, RemoteStorage, S3Config, TimeTravelError, MAX_KEYS_PER_DELETE, + REMOTE_STORAGE_PREFIX_SEPARATOR, }; pub(super) mod metrics; @@ -63,7 +64,6 @@ pub struct S3Bucket { concurrency_limiter: ConcurrencyLimiter, } -#[derive(Default)] struct GetObjectRequest { bucket: String, key: String, @@ -232,24 +232,8 @@ impl S3Bucket { let started_at = ScopeGuard::into_inner(started_at); - match get_object { - Ok(object_output) => { - let metadata = object_output.metadata().cloned().map(StorageMetadata); - let etag = object_output.e_tag.clone(); - let last_modified = object_output.last_modified.and_then(|t| t.try_into().ok()); - - let body = object_output.body; - let body = ByteStreamAsStream::from(body); - let body = PermitCarrying::new(permit, body); - let body = TimedDownload::new(started_at, body); - - Ok(Download { - metadata, - etag, - last_modified, - download_stream: Box::pin(body), - }) - } + let object_output = match get_object { + Ok(object_output) => object_output, Err(SdkError::ServiceError(e)) if matches!(e.err(), GetObjectError::NoSuchKey(_)) => { // Count this in the AttemptOutcome::Ok bucket, because 404 is not // an error: we expect to sometimes fetch an object and find it missing, @@ -259,7 +243,7 @@ impl S3Bucket { AttemptOutcome::Ok, started_at, ); - Err(DownloadError::NotFound) + return Err(DownloadError::NotFound); } Err(e) => { metrics::BUCKET_METRICS.req_seconds.observe_elapsed( @@ -268,11 +252,27 @@ impl S3Bucket { started_at, ); - Err(DownloadError::Other( + return Err(DownloadError::Other( anyhow::Error::new(e).context("download s3 object"), - )) + )); } - } + }; + + let metadata = object_output.metadata().cloned().map(StorageMetadata); + let etag = object_output.e_tag; + let last_modified = object_output.last_modified.and_then(|t| t.try_into().ok()); + + let body = object_output.body; + let body = ByteStreamAsStream::from(body); + let body = PermitCarrying::new(permit, body); + let body = TimedDownload::new(started_at, body); + + Ok(Download { + metadata, + etag, + last_modified, + download_stream: Box::pin(body), + }) } async fn delete_oids( @@ -354,33 +354,6 @@ impl Stream for ByteStreamAsStream { // sense and Stream::size_hint does not really } -pin_project_lite::pin_project! { - /// An `AsyncRead` adapter which carries a permit for the lifetime of the value. - struct PermitCarrying { - permit: tokio::sync::OwnedSemaphorePermit, - #[pin] - inner: S, - } -} - -impl PermitCarrying { - fn new(permit: tokio::sync::OwnedSemaphorePermit, inner: S) -> Self { - Self { permit, inner } - } -} - -impl>> Stream for PermitCarrying { - type Item = ::Item; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.project().inner.poll_next(cx) - } - - fn size_hint(&self) -> (usize, Option) { - self.inner.size_hint() - } -} - pin_project_lite::pin_project! { /// Times and tracks the outcome of the request. struct TimedDownload { diff --git a/libs/remote_storage/src/simulate_failures.rs b/libs/remote_storage/src/simulate_failures.rs index 82d5a61fda..14bdb5ed4d 100644 --- a/libs/remote_storage/src/simulate_failures.rs +++ b/libs/remote_storage/src/simulate_failures.rs @@ -60,7 +60,7 @@ impl UnreliableWrapper { /// On the first attempts of this operation, return an error. After 'attempts_to_fail' /// attempts, let the operation go ahead, and clear the counter. /// - fn attempt(&self, op: RemoteOp) -> Result { + fn attempt(&self, op: RemoteOp) -> anyhow::Result { let mut attempts = self.attempts.lock().unwrap(); match attempts.entry(op) { @@ -78,13 +78,13 @@ impl UnreliableWrapper { } else { let error = anyhow::anyhow!("simulated failure of remote operation {:?}", e.key()); - Err(DownloadError::Other(error)) + Err(error) } } Entry::Vacant(e) => { let error = anyhow::anyhow!("simulated failure of remote operation {:?}", e.key()); e.insert(1); - Err(DownloadError::Other(error)) + Err(error) } } } @@ -105,12 +105,17 @@ impl RemoteStorage for UnreliableWrapper { &self, prefix: Option<&RemotePath>, ) -> Result, DownloadError> { - self.attempt(RemoteOp::ListPrefixes(prefix.cloned()))?; + self.attempt(RemoteOp::ListPrefixes(prefix.cloned())) + .map_err(DownloadError::Other)?; self.inner.list_prefixes(prefix).await } - async fn list_files(&self, folder: Option<&RemotePath>) -> anyhow::Result> { - self.attempt(RemoteOp::ListPrefixes(folder.cloned()))?; + async fn list_files( + &self, + folder: Option<&RemotePath>, + ) -> Result, DownloadError> { + self.attempt(RemoteOp::ListPrefixes(folder.cloned())) + .map_err(DownloadError::Other)?; self.inner.list_files(folder).await } @@ -119,7 +124,8 @@ impl RemoteStorage for UnreliableWrapper { prefix: Option<&RemotePath>, mode: ListingMode, ) -> Result { - self.attempt(RemoteOp::ListPrefixes(prefix.cloned()))?; + self.attempt(RemoteOp::ListPrefixes(prefix.cloned())) + .map_err(DownloadError::Other)?; self.inner.list(prefix, mode).await } @@ -137,7 +143,8 @@ impl RemoteStorage for UnreliableWrapper { } async fn download(&self, from: &RemotePath) -> Result { - self.attempt(RemoteOp::Download(from.clone()))?; + self.attempt(RemoteOp::Download(from.clone())) + .map_err(DownloadError::Other)?; self.inner.download(from).await } @@ -150,7 +157,8 @@ impl RemoteStorage for UnreliableWrapper { // Note: We treat any download_byte_range as an "attempt" of the same // operation. We don't pay attention to the ranges. That's good enough // for now. - self.attempt(RemoteOp::Download(from.clone()))?; + self.attempt(RemoteOp::Download(from.clone())) + .map_err(DownloadError::Other)?; self.inner .download_byte_range(from, start_inclusive, end_exclusive) .await @@ -193,7 +201,7 @@ impl RemoteStorage for UnreliableWrapper { cancel: &CancellationToken, ) -> Result<(), TimeTravelError> { self.attempt(RemoteOp::TimeTravelRecover(prefix.map(|p| p.to_owned()))) - .map_err(|e| TimeTravelError::Other(anyhow::Error::new(e)))?; + .map_err(TimeTravelError::Other)?; self.inner .time_travel_recover(prefix, timestamp, done_if_after, cancel) .await diff --git a/libs/remote_storage/src/support.rs b/libs/remote_storage/src/support.rs new file mode 100644 index 0000000000..4688a484a5 --- /dev/null +++ b/libs/remote_storage/src/support.rs @@ -0,0 +1,33 @@ +use std::{ + pin::Pin, + task::{Context, Poll}, +}; + +use futures_util::Stream; + +pin_project_lite::pin_project! { + /// An `AsyncRead` adapter which carries a permit for the lifetime of the value. + pub(crate) struct PermitCarrying { + permit: tokio::sync::OwnedSemaphorePermit, + #[pin] + inner: S, + } +} + +impl PermitCarrying { + pub(crate) fn new(permit: tokio::sync::OwnedSemaphorePermit, inner: S) -> Self { + Self { permit, inner } + } +} + +impl Stream for PermitCarrying { + type Item = ::Item; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} diff --git a/pageserver/src/task_mgr.rs b/pageserver/src/task_mgr.rs index 5a06a97525..3cec5fa850 100644 --- a/pageserver/src/task_mgr.rs +++ b/pageserver/src/task_mgr.rs @@ -576,8 +576,8 @@ pub fn shutdown_token() -> CancellationToken { /// Has the current task been requested to shut down? pub fn is_shutdown_requested() -> bool { - if let Ok(cancel) = SHUTDOWN_TOKEN.try_with(|t| t.clone()) { - cancel.is_cancelled() + if let Ok(true_or_false) = SHUTDOWN_TOKEN.try_with(|t| t.is_cancelled()) { + true_or_false } else { if !cfg!(test) { warn!("is_shutdown_requested() called in an unexpected task or thread"); diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index f086f46213..4446c410b0 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -1377,7 +1377,7 @@ impl Tenant { async move { debug!("starting index part download"); - let index_part = client.download_index_file(cancel_clone).await; + let index_part = client.download_index_file(&cancel_clone).await; debug!("finished index part download"); @@ -2434,7 +2434,7 @@ impl Tenant { // operation is rare, so it's simpler to just download it (and robustly guarantees that the index // we use here really is the remotely persistent one). let result = tl_client - .download_index_file(self.cancel.clone()) + .download_index_file(&self.cancel) .instrument(info_span!("download_index_file", tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), timeline_id=%timeline.timeline_id)) .await?; let index_part = match result { diff --git a/pageserver/src/tenant/remote_timeline_client.rs b/pageserver/src/tenant/remote_timeline_client.rs index 152c9a2b7d..0c7dd68c3f 100644 --- a/pageserver/src/tenant/remote_timeline_client.rs +++ b/pageserver/src/tenant/remote_timeline_client.rs @@ -217,6 +217,7 @@ use crate::metrics::{ }; use crate::task_mgr::shutdown_token; use crate::tenant::debug_assert_current_span_has_tenant_and_timeline_id; +use crate::tenant::remote_timeline_client::download::download_retry; use crate::tenant::storage_layer::AsLayerDesc; use crate::tenant::upload_queue::Delete; use crate::tenant::TIMELINES_SEGMENT_NAME; @@ -262,6 +263,11 @@ pub(crate) const INITDB_PRESERVED_PATH: &str = "initdb-preserved.tar.zst"; /// Default buffer size when interfacing with [`tokio::fs::File`]. pub(crate) const BUFFER_SIZE: usize = 32 * 1024; +/// This timeout is intended to deal with hangs in lower layers, e.g. stuck TCP flows. It is not +/// intended to be snappy enough for prompt shutdown, as we have a CancellationToken for that. +pub(crate) const UPLOAD_TIMEOUT: Duration = Duration::from_secs(120); +pub(crate) const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(120); + pub enum MaybeDeletedIndexPart { IndexPart(IndexPart), Deleted(IndexPart), @@ -325,11 +331,6 @@ pub struct RemoteTimelineClient { cancel: CancellationToken, } -/// This timeout is intended to deal with hangs in lower layers, e.g. stuck TCP flows. It is not -/// intended to be snappy enough for prompt shutdown, as we have a CancellationToken for that. -const UPLOAD_TIMEOUT: Duration = Duration::from_secs(120); -const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(120); - /// Wrapper for timeout_cancellable that flattens result and converts TimeoutCancellableError to anyhow. /// /// This is a convenience for the various upload functions. In future @@ -506,7 +507,7 @@ impl RemoteTimelineClient { /// Download index file pub async fn download_index_file( &self, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result { let _unfinished_gauge_guard = self.metrics.call_begin( &RemoteOpFileKind::Index, @@ -1147,22 +1148,17 @@ impl RemoteTimelineClient { let cancel = shutdown_token(); - let remaining = backoff::retry( + let remaining = download_retry( || async { self.storage_impl .list_files(Some(&timeline_storage_path)) .await }, - |_e| false, - FAILED_DOWNLOAD_WARN_THRESHOLD, - FAILED_REMOTE_OP_RETRIES, - "list_prefixes", + "list remaining files", &cancel, ) .await - .ok_or_else(|| anyhow::anyhow!("Cancelled!")) - .and_then(|x| x) - .context("list prefixes")?; + .context("list files remaining files")?; // We will delete the current index_part object last, since it acts as a deletion // marker via its deleted_at attribute @@ -1351,6 +1347,7 @@ impl RemoteTimelineClient { /// queue. /// async fn perform_upload_task(self: &Arc, task: Arc) { + let cancel = shutdown_token(); // Loop to retry until it completes. loop { // If we're requested to shut down, close up shop and exit. @@ -1362,7 +1359,7 @@ impl RemoteTimelineClient { // the Future, but we're not 100% sure if the remote storage library // is cancellation safe, so we don't dare to do that. Hopefully, the // upload finishes or times out soon enough. - if task_mgr::is_shutdown_requested() { + if cancel.is_cancelled() { info!("upload task cancelled by shutdown request"); match self.stop() { Ok(()) => {} @@ -1473,7 +1470,7 @@ impl RemoteTimelineClient { retries, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS, - &shutdown_token(), + &cancel, ) .await; } @@ -1990,7 +1987,7 @@ mod tests { // Download back the index.json, and check that the list of files is correct let initial_index_part = match client - .download_index_file(CancellationToken::new()) + .download_index_file(&CancellationToken::new()) .await .unwrap() { @@ -2084,7 +2081,7 @@ mod tests { // Download back the index.json, and check that the list of files is correct let index_part = match client - .download_index_file(CancellationToken::new()) + .download_index_file(&CancellationToken::new()) .await .unwrap() { @@ -2286,7 +2283,7 @@ mod tests { let client = test_state.build_client(get_generation); let download_r = client - .download_index_file(CancellationToken::new()) + .download_index_file(&CancellationToken::new()) .await .expect("download should always succeed"); assert!(matches!(download_r, MaybeDeletedIndexPart::IndexPart(_))); diff --git a/pageserver/src/tenant/remote_timeline_client/download.rs b/pageserver/src/tenant/remote_timeline_client/download.rs index 6c1125746b..33287fc8f4 100644 --- a/pageserver/src/tenant/remote_timeline_client/download.rs +++ b/pageserver/src/tenant/remote_timeline_client/download.rs @@ -216,16 +216,15 @@ pub async fn list_remote_timelines( anyhow::bail!("storage-sync-list-remote-timelines"); }); - let cancel_inner = cancel.clone(); let listing = download_retry_forever( || { download_cancellable( - &cancel_inner, + &cancel, storage.list(Some(&remote_path), ListingMode::WithDelimiter), ) }, &format!("list timelines for {tenant_shard_id}"), - cancel, + &cancel, ) .await?; @@ -258,19 +257,18 @@ async fn do_download_index_part( tenant_shard_id: &TenantShardId, timeline_id: &TimelineId, index_generation: Generation, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result { use futures::stream::StreamExt; let remote_path = remote_index_path(tenant_shard_id, timeline_id, index_generation); - let cancel_inner = cancel.clone(); let index_part_bytes = download_retry_forever( || async { // Cancellation: if is safe to cancel this future because we're just downloading into // a memory buffer, not touching local disk. let index_part_download = - download_cancellable(&cancel_inner, storage.download(&remote_path)).await?; + download_cancellable(cancel, storage.download(&remote_path)).await?; let mut index_part_bytes = Vec::new(); let mut stream = std::pin::pin!(index_part_download.download_stream); @@ -288,7 +286,7 @@ async fn do_download_index_part( .await?; let index_part: IndexPart = serde_json::from_slice(&index_part_bytes) - .with_context(|| format!("download index part file at {remote_path:?}")) + .with_context(|| format!("deserialize index part file at {remote_path:?}")) .map_err(DownloadError::Other)?; Ok(index_part) @@ -305,7 +303,7 @@ pub(super) async fn download_index_part( tenant_shard_id: &TenantShardId, timeline_id: &TimelineId, my_generation: Generation, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result { debug_assert_current_span_has_tenant_and_timeline_id(); @@ -325,14 +323,8 @@ pub(super) async fn download_index_part( // index in our generation. // // This is an optimization to avoid doing the listing for the general case below. - let res = do_download_index_part( - storage, - tenant_shard_id, - timeline_id, - my_generation, - cancel.clone(), - ) - .await; + let res = + do_download_index_part(storage, tenant_shard_id, timeline_id, my_generation, cancel).await; match res { Ok(index_part) => { tracing::debug!( @@ -357,7 +349,7 @@ pub(super) async fn download_index_part( tenant_shard_id, timeline_id, my_generation.previous(), - cancel.clone(), + cancel, ) .await; match res { @@ -379,18 +371,13 @@ pub(super) async fn download_index_part( // objects, and select the highest one with a generation <= my_generation. Constructing the prefix is equivalent // to constructing a full index path with no generation, because the generation is a suffix. let index_prefix = remote_index_path(tenant_shard_id, timeline_id, Generation::none()); - let indices = backoff::retry( + + let indices = download_retry( || async { storage.list_files(Some(&index_prefix)).await }, - |_| false, - FAILED_DOWNLOAD_WARN_THRESHOLD, - FAILED_REMOTE_OP_RETRIES, - "listing index_part files", - &cancel, + "list index_part files", + cancel, ) - .await - .ok_or_else(|| anyhow::anyhow!("Cancelled")) - .and_then(|x| x) - .map_err(DownloadError::Other)?; + .await?; // General case logic for which index to use: the latest index whose generation // is <= our own. See "Finding the remote indices for timelines" in docs/rfcs/025-generation-numbers.md @@ -447,8 +434,6 @@ pub(crate) async fn download_initdb_tar_zst( "{INITDB_PATH}.download-{timeline_id}.{TEMP_FILE_SUFFIX}" )); - let cancel_inner = cancel.clone(); - let file = download_retry( || async { let file = OpenOptions::new() @@ -461,13 +446,11 @@ pub(crate) async fn download_initdb_tar_zst( .with_context(|| format!("tempfile creation {temp_path}")) .map_err(DownloadError::Other)?; - let download = match download_cancellable(&cancel_inner, storage.download(&remote_path)) - .await + let download = match download_cancellable(cancel, storage.download(&remote_path)).await { Ok(dl) => dl, Err(DownloadError::NotFound) => { - download_cancellable(&cancel_inner, storage.download(&remote_preserved_path)) - .await? + download_cancellable(cancel, storage.download(&remote_preserved_path)).await? } Err(other) => Err(other)?, }; @@ -516,7 +499,7 @@ pub(crate) async fn download_initdb_tar_zst( /// with backoff. /// /// (See similar logic for uploads in `perform_upload_task`) -async fn download_retry( +pub(super) async fn download_retry( op: O, description: &str, cancel: &CancellationToken, @@ -527,7 +510,7 @@ where { backoff::retry( op, - |e| matches!(e, DownloadError::BadInput(_) | DownloadError::NotFound), + DownloadError::is_permanent, FAILED_DOWNLOAD_WARN_THRESHOLD, FAILED_REMOTE_OP_RETRIES, description, @@ -541,7 +524,7 @@ where async fn download_retry_forever( op: O, description: &str, - cancel: CancellationToken, + cancel: &CancellationToken, ) -> Result where O: FnMut() -> F, @@ -549,11 +532,11 @@ where { backoff::retry( op, - |e| matches!(e, DownloadError::BadInput(_) | DownloadError::NotFound), + DownloadError::is_permanent, FAILED_DOWNLOAD_WARN_THRESHOLD, u32::MAX, description, - &cancel, + cancel, ) .await .ok_or_else(|| DownloadError::Cancelled) diff --git a/pageserver/src/tenant/secondary/downloader.rs b/pageserver/src/tenant/secondary/downloader.rs index 9330edf946..0666e104f8 100644 --- a/pageserver/src/tenant/secondary/downloader.rs +++ b/pageserver/src/tenant/secondary/downloader.rs @@ -533,7 +533,7 @@ impl<'a> TenantDownloader<'a> { .map_err(UpdateError::from)?; let mut heatmap_bytes = Vec::new(); let mut body = tokio_util::io::StreamReader::new(download.download_stream); - let _size = tokio::io::copy(&mut body, &mut heatmap_bytes).await?; + let _size = tokio::io::copy_buf(&mut body, &mut heatmap_bytes).await?; Ok(heatmap_bytes) }, |e| matches!(e, UpdateError::NoData | UpdateError::Cancelled), From 8d98981fe580fcdfb7066a5698c2448af0cbc61d Mon Sep 17 00:00:00 2001 From: John Spray Date: Fri, 9 Feb 2024 13:20:04 +0000 Subject: [PATCH 56/72] tests: deflake test_sharding_split_unsharded (#6699) ## Problem This test was a subset of the larger sharding test, and it missed the validate() call on workload that was implicitly waiting for a tenant to become active before trying to split it. It could therefore fail to split due to tenant not yet being active. ## Summary of changes - Insert .validate() call, and move the Workload setup to after the check of shard ID (as the shard ID check should pass immediately) --- test_runner/regress/test_sharding.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/test_runner/regress/test_sharding.py b/test_runner/regress/test_sharding.py index 27d1cf2f34..fa40219d0e 100644 --- a/test_runner/regress/test_sharding.py +++ b/test_runner/regress/test_sharding.py @@ -95,14 +95,15 @@ def test_sharding_split_unsharded( tenant_id = env.initial_tenant timeline_id = env.initial_timeline - workload = Workload(env, tenant_id, timeline_id, branch_name="main") - workload.init() - workload.write_rows(256) - # Check that we created with an unsharded TenantShardId: this is the default, # but check it in case we change the default in future assert env.attachment_service.inspect(TenantShardId(tenant_id, 0, 0)) is not None + workload = Workload(env, tenant_id, timeline_id, branch_name="main") + workload.init() + workload.write_rows(256) + workload.validate() + # Split one shard into two env.attachment_service.tenant_shard_split(tenant_id, shard_count=2) From 84a0e7b022e37b041004e7d9299060a3777c63eb Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Fri, 9 Feb 2024 11:07:42 +0200 Subject: [PATCH 57/72] tests: Allow setting shutdown mode separately from 'destroy' flag In neon_local, the default mode is now always 'fast', regardless of 'destroy'. You can override it with the "neon_local endpoint stop --mode=immediate" flag. In python tests, we still default to 'immediate' mode when using the stop_and_destroy() function, and 'fast' with plain stop(). I kept that to avoid changing behavior in existing tests. I don't think existing tests depend on it, but I wasn't 100% certain. --- control_plane/src/bin/neon_local.rs | 16 +++++++++++++--- control_plane/src/endpoint.rs | 18 ++---------------- test_runner/fixtures/neon_fixtures.py | 11 +++++++---- 3 files changed, 22 insertions(+), 23 deletions(-) diff --git a/control_plane/src/bin/neon_local.rs b/control_plane/src/bin/neon_local.rs index b9af467fdf..d71cdf02c0 100644 --- a/control_plane/src/bin/neon_local.rs +++ b/control_plane/src/bin/neon_local.rs @@ -1014,12 +1014,13 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re .get_one::("endpoint_id") .ok_or_else(|| anyhow!("No endpoint ID was provided to stop"))?; let destroy = sub_args.get_flag("destroy"); + let mode = sub_args.get_one::("mode").expect("has a default"); let endpoint = cplane .endpoints .get(endpoint_id.as_str()) .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?; - endpoint.stop(destroy)?; + endpoint.stop(mode, destroy)?; } _ => bail!("Unexpected endpoint subcommand '{sub_name}'"), @@ -1303,7 +1304,7 @@ async fn try_stop_all(env: &local_env::LocalEnv, immediate: bool) { match ComputeControlPlane::load(env.clone()) { Ok(cplane) => { for (_k, node) in cplane.endpoints { - if let Err(e) = node.stop(false) { + if let Err(e) = node.stop(if immediate { "immediate" } else { "fast " }, false) { eprintln!("postgres stop failed: {e:#}"); } } @@ -1652,7 +1653,16 @@ fn cli() -> Command { .long("destroy") .action(ArgAction::SetTrue) .required(false) - ) + ) + .arg( + Arg::new("mode") + .help("Postgres shutdown mode, passed to \"pg_ctl -m \"") + .long("mode") + .action(ArgAction::Set) + .required(false) + .value_parser(["smart", "fast", "immediate"]) + .default_value("fast") + ) ) ) diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs index b19a6a1a18..f1fe12e05f 100644 --- a/control_plane/src/endpoint.rs +++ b/control_plane/src/endpoint.rs @@ -761,22 +761,8 @@ impl Endpoint { } } - pub fn stop(&self, destroy: bool) -> Result<()> { - // If we are going to destroy data directory, - // use immediate shutdown mode, otherwise, - // shutdown gracefully to leave the data directory sane. - // - // Postgres is always started from scratch, so stop - // without destroy only used for testing and debugging. - // - self.pg_ctl( - if destroy { - &["-m", "immediate", "stop"] - } else { - &["stop"] - }, - &None, - )?; + pub fn stop(&self, mode: &str, destroy: bool) -> Result<()> { + self.pg_ctl(&["-m", mode, "stop"], &None)?; // Also wait for the compute_ctl process to die. It might have some // cleanup work to do after postgres stops, like syncing safekeepers, diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index a6aff77ddf..9996853525 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1816,6 +1816,7 @@ class NeonCli(AbstractNeonCli): endpoint_id: str, destroy=False, check_return_code=True, + mode: Optional[str] = None, ) -> "subprocess.CompletedProcess[str]": args = [ "endpoint", @@ -1823,6 +1824,8 @@ class NeonCli(AbstractNeonCli): ] if destroy: args.append("--destroy") + if mode is not None: + args.append(f"--mode={mode}") if endpoint_id is not None: args.append(endpoint_id) @@ -3162,7 +3165,7 @@ class Endpoint(PgProtocol): with open(remote_extensions_spec_path, "w") as file: json.dump(spec, file, indent=4) - def stop(self) -> "Endpoint": + def stop(self, mode: str = "fast") -> "Endpoint": """ Stop the Postgres instance if it's running. Returns self. @@ -3171,13 +3174,13 @@ class Endpoint(PgProtocol): if self.running: assert self.endpoint_id is not None self.env.neon_cli.endpoint_stop( - self.endpoint_id, check_return_code=self.check_stop_result + self.endpoint_id, check_return_code=self.check_stop_result, mode=mode ) self.running = False return self - def stop_and_destroy(self) -> "Endpoint": + def stop_and_destroy(self, mode: str = "immediate") -> "Endpoint": """ Stop the Postgres instance, then destroy the endpoint. Returns self. @@ -3185,7 +3188,7 @@ class Endpoint(PgProtocol): assert self.endpoint_id is not None self.env.neon_cli.endpoint_stop( - self.endpoint_id, True, check_return_code=self.check_stop_result + self.endpoint_id, True, check_return_code=self.check_stop_result, mode=mode ) self.endpoint_id = None self.running = False From 5239cdc29fdfe8458798cefad51f8871108f9811 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Fri, 9 Feb 2024 11:07:47 +0200 Subject: [PATCH 58/72] Fix test_vm_bit_clear_on_heap_lock test The test was supposed to reproduce the bug fixed in commit 66fa176cc8, i.e. that the clearing of the VM bit was not replayed in the pageserver on HEAP_LOCK records. But it was broken in many ways and failed to reproduce the original problem if you reverted the fix: - The comparison of XIDs was broken. The test read the XID in to a variable in python, but it was treated as a string rather than an integer. As a result, e.g. "999" > "1000". - The test accessed the locked tuple too early, in the loop. Accessing it early, before the pg_xact page had been removed, set the hint bits. That masked the problem on subsequent accesses. - The on-demand SLRU download that was introduced in commit 9a9d9beaee hid the issue. Even though an SLRU segment was removed by Postgres, when it later tried to access it, it could still download it from the pageserver. To ensure that doesn't happen, shorten the GC period and compact and GC aggressively in the test. I also added a more direct check that the VM page is updated, using the get_page_at_lsn() debugging function. Right after locking the row, we now fetch the VM page from pageserver and directly compare it with the VM page in the page cache. They should match. That assertion is more robust to things like on-demand SLRU download that could mask the bug. --- test_runner/regress/test_vm_bits.py | 118 +++++++++++++++++----------- 1 file changed, 72 insertions(+), 46 deletions(-) diff --git a/test_runner/regress/test_vm_bits.py b/test_runner/regress/test_vm_bits.py index 415f086bd3..06c30b8d81 100644 --- a/test_runner/regress/test_vm_bits.py +++ b/test_runner/regress/test_vm_bits.py @@ -1,6 +1,7 @@ -import pytest +import time + from fixtures.log_helper import log -from fixtures.neon_fixtures import NeonEnv, fork_at_current_lsn +from fixtures.neon_fixtures import NeonEnv, NeonEnvBuilder, fork_at_current_lsn # @@ -118,12 +119,20 @@ def test_vm_bit_clear(neon_simple_env: NeonEnv): # Test that the ALL_FROZEN VM bit is cleared correctly at a HEAP_LOCK # record. # -# FIXME: This test is broken -@pytest.mark.skip("See https://github.com/neondatabase/neon/pull/6412#issuecomment-1902072541") -def test_vm_bit_clear_on_heap_lock(neon_simple_env: NeonEnv): - env = neon_simple_env +def test_vm_bit_clear_on_heap_lock(neon_env_builder: NeonEnvBuilder): + tenant_conf = { + "checkpoint_distance": f"{128 * 1024}", + "compaction_target_size": f"{128 * 1024}", + "compaction_threshold": "1", + # create image layers eagerly, so that GC can remove some layers + "image_creation_threshold": "1", + # set PITR interval to be small, so we can do GC + "pitr_interval": "0 s", + } + env = neon_env_builder.init_start(initial_tenant_conf=tenant_conf) - env.neon_cli.create_branch("test_vm_bit_clear_on_heap_lock", "empty") + tenant_id = env.initial_tenant + timeline_id = env.neon_cli.create_branch("test_vm_bit_clear_on_heap_lock") endpoint = env.endpoints.create_start( "test_vm_bit_clear_on_heap_lock", config_lines=[ @@ -139,72 +148,88 @@ def test_vm_bit_clear_on_heap_lock(neon_simple_env: NeonEnv): # Install extension containing function needed for test cur.execute("CREATE EXTENSION neon_test_utils") - - cur.execute("SELECT pg_switch_wal()") + cur.execute("CREATE EXTENSION pageinspect") # Create a test table and freeze it to set the all-frozen VM bit on all pages. cur.execute("CREATE TABLE vmtest_lock (id integer PRIMARY KEY)") cur.execute("INSERT INTO vmtest_lock SELECT g FROM generate_series(1, 50000) g") - cur.execute("VACUUM FREEZE vmtest_lock") + + cur.execute("VACUUM (FREEZE, DISABLE_PAGE_SKIPPING true) vmtest_lock") # Lock a row. This clears the all-frozen VM bit for that page. + cur.execute("BEGIN") cur.execute("SELECT * FROM vmtest_lock WHERE id = 40000 FOR UPDATE") # Remember the XID. We will use it later to verify that we have consumed a lot of # XIDs after this. cur.execute("select pg_current_xact_id()") - locking_xid = cur.fetchall()[0][0] + locking_xid = int(cur.fetchall()[0][0]) - # Stop and restart postgres, to clear the buffer cache. + cur.execute("COMMIT") + + # The VM page in shared buffer cache, and the same page as reconstructed + # by the pageserver, should be equal. + cur.execute("select get_raw_page( 'vmtest_lock', 'vm', 0 )") + vm_page_in_cache = (cur.fetchall()[0][0])[:100].hex() + cur.execute("select get_raw_page_at_lsn( 'vmtest_lock', 'vm', 0, pg_current_wal_insert_lsn() )") + vm_page_at_pageserver = (cur.fetchall()[0][0])[:100].hex() + + assert vm_page_at_pageserver == vm_page_in_cache + + # The above assert is enough to verify the bug that was fixed in + # commit 66fa176cc8. But for good measure, we also reproduce the + # original problem that the missing VM page update caused. The + # rest of the test does that. + + # Kill and restart postgres, to clear the buffer cache. # # NOTE: clear_buffer_cache() will not do, because it evicts the dirty pages # in a "clean" way. Our neon extension will write a full-page image of the VM - # page, and we want to avoid that. - endpoint.stop() + # page, and we want to avoid that. A clean shutdown will also not do, for the + # same reason. + endpoint.stop(mode="immediate") + endpoint.start() pg_conn = endpoint.connect() cur = pg_conn.cursor() - cur.execute("select xmin, xmax, * from vmtest_lock where id = 40000 ") - tup = cur.fetchall() - xmax_before = tup[0][1] - # Consume a lot of XIDs, so that anti-wraparound autovacuum kicks # in and the clog gets truncated. We set autovacuum_freeze_max_age to a very # low value, so it doesn't take all that many XIDs for autovacuum to kick in. - for i in range(1000): + # + # We could use test_consume_xids() to consume XIDs much faster, + # but it wouldn't speed up the overall test, because we'd still + # need to wait for autovacuum to run. + for _ in range(1000): + cur.execute("select test_consume_xids(10000);") + for _ in range(1000): cur.execute( - """ - CREATE TEMP TABLE othertable (i int) ON COMMIT DROP; - do $$ - begin - for i in 1..100000 loop - -- Use a begin-exception block to generate a new subtransaction on each iteration - begin - insert into othertable values (i); - exception when others then - raise 'not expected %', sqlerrm; - end; - end loop; - end; - $$; - """ + "select get_raw_page_at_lsn( 'vmtest_lock', 'vm', 0, pg_current_wal_insert_lsn() )" ) - cur.execute("select xmin, xmax, * from vmtest_lock where id = 40000 ") - tup = cur.fetchall() - log.info(f"tuple = {tup}") - xmax = tup[0][1] - assert xmax == xmax_before + page = (cur.fetchall()[0][0])[:100].hex() + log.info(f"VM page contents: {page}") - if i % 50 == 0: - cur.execute("select datfrozenxid from pg_database where datname='postgres'") - datfrozenxid = cur.fetchall()[0][0] - if datfrozenxid > locking_xid: - break + cur.execute("select get_raw_page( 'vmtest_lock', 'vm', 0 )") + page = (cur.fetchall()[0][0])[:100].hex() + log.info(f"VM page contents in cache: {page}") + + cur.execute("select min(datfrozenxid::text::int) from pg_database") + datfrozenxid = int(cur.fetchall()[0][0]) + log.info(f"datfrozenxid {datfrozenxid} locking_xid: {locking_xid}") + if datfrozenxid > locking_xid + 3000000: + break + time.sleep(0.5) cur.execute("select pg_current_xact_id()") - curr_xid = cur.fetchall()[0][0] - assert int(curr_xid) - int(locking_xid) >= 100000 + curr_xid = int(cur.fetchall()[0][0]) + assert curr_xid - locking_xid >= 100000 + + # Perform GC in the pageserver. Otherwise the compute might still + # be able to download the already-deleted SLRU segment from the + # pageserver. That masks the original bug. + env.pageserver.http_client().timeline_checkpoint(tenant_id, timeline_id) + env.pageserver.http_client().timeline_compact(tenant_id, timeline_id) + env.pageserver.http_client().timeline_gc(tenant_id, timeline_id, 0) # Now, if the VM all-frozen bit was not correctly cleared on # replay, we will try to fetch the status of the XID that was @@ -214,3 +239,4 @@ def test_vm_bit_clear_on_heap_lock(neon_simple_env: NeonEnv): cur.execute("select xmin, xmax, * from vmtest_lock where id = 40000 for update") tup = cur.fetchall() log.info(f"tuple = {tup}") + cur.execute("commit transaction") From 89a5c654bfc688babcdfa6c9dcda68876c0d6f98 Mon Sep 17 00:00:00 2001 From: John Spray Date: Fri, 9 Feb 2024 14:26:50 +0000 Subject: [PATCH 59/72] control_plane: follow up for embedded migrations (#6647) ## Problem In https://github.com/neondatabase/neon/pull/6637, we remove the need to run migrations externally, but for compat tests to work we can't remove those invocations from the neon_local binary. Once that previous PR merges, we can make the followup changes without upsetting compat tests. --- Cargo.lock | 4 - control_plane/Cargo.toml | 2 - control_plane/src/attachment_service.rs | 118 +++++------------------- workspace_hack/Cargo.toml | 2 - 4 files changed, 22 insertions(+), 104 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c0c319cd89..a2939e6c75 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1329,8 +1329,6 @@ dependencies = [ "clap", "comfy-table", "compute_api", - "diesel", - "diesel_migrations", "futures", "git-version", "hex", @@ -6832,8 +6830,6 @@ dependencies = [ "clap", "clap_builder", "crossbeam-utils", - "diesel", - "diesel_derives", "either", "fail", "futures-channel", diff --git a/control_plane/Cargo.toml b/control_plane/Cargo.toml index 09c171f1d3..75e5dcb7f8 100644 --- a/control_plane/Cargo.toml +++ b/control_plane/Cargo.toml @@ -10,8 +10,6 @@ async-trait.workspace = true camino.workspace = true clap.workspace = true comfy-table.workspace = true -diesel = { version = "2.1.4", features = ["postgres"]} -diesel_migrations = { version = "2.1.0", features = ["postgres"]} futures.workspace = true git-version.workspace = true nix.workspace = true diff --git a/control_plane/src/attachment_service.rs b/control_plane/src/attachment_service.rs index c3e071aa71..14bfda47c3 100644 --- a/control_plane/src/attachment_service.rs +++ b/control_plane/src/attachment_service.rs @@ -1,11 +1,5 @@ use crate::{background_process, local_env::LocalEnv}; use camino::{Utf8Path, Utf8PathBuf}; -use diesel::{ - backend::Backend, - query_builder::{AstPass, QueryFragment, QueryId}, - Connection, PgConnection, QueryResult, RunQueryDsl, -}; -use diesel_migrations::{HarnessWithOutput, MigrationHarness}; use hyper::Method; use pageserver_api::{ models::{ @@ -17,7 +11,7 @@ use pageserver_api::{ use pageserver_client::mgmt_api::ResponseErrorMessageExt; use postgres_backend::AuthType; use serde::{de::DeserializeOwned, Deserialize, Serialize}; -use std::{env, str::FromStr}; +use std::str::FromStr; use tokio::process::Command; use tracing::instrument; use url::Url; @@ -273,37 +267,6 @@ impl AttachmentService { .expect("non-Unicode path") } - /// In order to access database migrations, we need to find the Neon source tree - async fn find_source_root(&self) -> anyhow::Result { - // We assume that either prd or our binary is in the source tree. The former is usually - // true for automated test runners, the latter is usually true for developer workstations. Often - // both are true, which is fine. - let candidate_start_points = [ - // Current working directory - Utf8PathBuf::from_path_buf(std::env::current_dir()?).unwrap(), - // Directory containing the binary we're running inside - Utf8PathBuf::from_path_buf(env::current_exe()?.parent().unwrap().to_owned()).unwrap(), - ]; - - // For each candidate start point, search through ancestors looking for a neon.git source tree root - for start_point in &candidate_start_points { - // Start from the build dir: assumes we are running out of a built neon source tree - for path in start_point.ancestors() { - // A crude approximation: the root of the source tree is whatever contains a "control_plane" - // subdirectory. - let control_plane = path.join("control_plane"); - if tokio::fs::try_exists(&control_plane).await? { - return Ok(path.to_owned()); - } - } - } - - // Fall-through - Err(anyhow::anyhow!( - "Could not find control_plane src dir, after searching ancestors of {candidate_start_points:?}" - )) - } - /// Find the directory containing postgres binaries, such as `initdb` and `pg_ctl` /// /// This usually uses ATTACHMENT_SERVICE_POSTGRES_VERSION of postgres, but will fall back @@ -343,69 +306,32 @@ impl AttachmentService { /// /// Returns the database url pub async fn setup_database(&self) -> anyhow::Result { - let database_url = format!( - "postgresql://localhost:{}/attachment_service", - self.postgres_port - ); - println!("Running attachment service database setup..."); - fn change_database_of_url(database_url: &str, default_database: &str) -> (String, String) { - let base = ::url::Url::parse(database_url).unwrap(); - let database = base.path_segments().unwrap().last().unwrap().to_owned(); - let mut new_url = base.join(default_database).unwrap(); - new_url.set_query(base.query()); - (database, new_url.into()) - } + const DB_NAME: &str = "attachment_service"; + let database_url = format!("postgresql://localhost:{}/{DB_NAME}", self.postgres_port); - #[derive(Debug, Clone)] - pub struct CreateDatabaseStatement { - db_name: String, - } + let pg_bin_dir = self.get_pg_bin_dir().await?; + let createdb_path = pg_bin_dir.join("createdb"); + let output = Command::new(&createdb_path) + .args([ + "-h", + "localhost", + "-p", + &format!("{}", self.postgres_port), + &DB_NAME, + ]) + .output() + .await + .expect("Failed to spawn createdb"); - impl CreateDatabaseStatement { - pub fn new(db_name: &str) -> Self { - CreateDatabaseStatement { - db_name: db_name.to_owned(), - } + if !output.status.success() { + let stderr = String::from_utf8(output.stderr).expect("Non-UTF8 output from createdb"); + if stderr.contains("already exists") { + tracing::info!("Database {DB_NAME} already exists"); + } else { + anyhow::bail!("createdb failed with status {}: {stderr}", output.status); } } - impl QueryFragment for CreateDatabaseStatement { - fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, DB>) -> QueryResult<()> { - out.push_sql("CREATE DATABASE "); - out.push_identifier(&self.db_name)?; - Ok(()) - } - } - - impl RunQueryDsl for CreateDatabaseStatement {} - - impl QueryId for CreateDatabaseStatement { - type QueryId = (); - - const HAS_STATIC_QUERY_ID: bool = false; - } - if PgConnection::establish(&database_url).is_err() { - let (database, postgres_url) = change_database_of_url(&database_url, "postgres"); - println!("Creating database: {database}"); - let mut conn = PgConnection::establish(&postgres_url)?; - CreateDatabaseStatement::new(&database).execute(&mut conn)?; - } - let mut conn = PgConnection::establish(&database_url)?; - - let migrations_dir = self - .find_source_root() - .await? - .join("control_plane/attachment_service/migrations"); - - let migrations = diesel_migrations::FileBasedMigrations::from_path(migrations_dir)?; - println!("Running migrations in {}", migrations.path().display()); - HarnessWithOutput::write_to_stdout(&mut conn) - .run_pending_migrations(migrations) - .map(|_| ()) - .map_err(|e| anyhow::anyhow!(e))?; - - println!("Migrations complete"); - Ok(database_url) } diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 70b238913d..8e9cc43152 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -29,7 +29,6 @@ chrono = { version = "0.4", default-features = false, features = ["clock", "serd clap = { version = "4", features = ["derive", "string"] } clap_builder = { version = "4", default-features = false, features = ["color", "help", "std", "string", "suggestions", "usage"] } crossbeam-utils = { version = "0.8" } -diesel = { version = "2", features = ["postgres", "r2d2", "serde_json"] } either = { version = "1" } fail = { version = "0.5", default-features = false, features = ["failpoints"] } futures-channel = { version = "0.3", features = ["sink"] } @@ -90,7 +89,6 @@ anyhow = { version = "1", features = ["backtrace"] } bytes = { version = "1", features = ["serde"] } cc = { version = "1", default-features = false, features = ["parallel"] } chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "wasmbind"] } -diesel_derives = { version = "2", features = ["32-column-tables", "postgres", "r2d2", "with-deprecated"] } either = { version = "1" } getrandom = { version = "0.2", default-features = false, features = ["std"] } hashbrown-582f2526e08bb6a0 = { package = "hashbrown", version = "0.14", default-features = false, features = ["raw"] } From 96d89cde5108850d1f0f41c23ff175552297ab9d Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Fri, 9 Feb 2024 15:50:51 +0000 Subject: [PATCH 60/72] Proxy error reworking (#6453) ## Problem Taking my ideas from https://github.com/neondatabase/neon/pull/6283 and doing a bit less radical changes. smaller commits. We currently don't report error classifications in proxy as the current error handling made it hard to do so. ## Summary of changes 1. Add a `ReportableError` trait that all errors will implement. This provides the error classification functionality. 2. Handle Client requests a strongly typed error * this error is a `ReportableError` and is logged appropriately 3. The handle client error only has a few possible error types, to account for the fact that at this point errors should be returned to the user. --- proxy/src/auth.rs | 37 ++++++++- proxy/src/auth/backend/classic.rs | 4 +- proxy/src/auth/backend/link.rs | 18 ++-- proxy/src/auth/credentials.rs | 14 +++- proxy/src/bin/pg_sni_router.rs | 11 ++- proxy/src/cancellation.rs | 37 +++++++-- proxy/src/compute.rs | 22 ++++- proxy/src/console/provider.rs | 31 ++++++- proxy/src/context.rs | 18 +++- proxy/src/context/parquet.rs | 2 +- proxy/src/error.rs | 38 +++++++-- proxy/src/metrics.rs | 19 +++++ proxy/src/proxy.rs | 95 ++++++++++++++++++---- proxy/src/proxy/handshake.rs | 76 +++++++++++++---- proxy/src/proxy/passthrough.rs | 23 ++++-- proxy/src/proxy/tests.rs | 8 +- proxy/src/proxy/tests/mitm.rs | 10 +-- proxy/src/sasl.rs | 14 +++- proxy/src/serverless.rs | 14 ++-- proxy/src/serverless/backend.rs | 29 +++++-- proxy/src/serverless/conn_pool.rs | 4 +- proxy/src/serverless/json.rs | 32 ++++++-- proxy/src/serverless/sql_over_http.rs | 113 ++++++++++++-------------- proxy/src/serverless/websocket.rs | 30 +++++-- proxy/src/stream.rs | 75 ++++++++++++++--- 25 files changed, 588 insertions(+), 186 deletions(-) diff --git a/proxy/src/auth.rs b/proxy/src/auth.rs index 8d1b861a66..48de4e2353 100644 --- a/proxy/src/auth.rs +++ b/proxy/src/auth.rs @@ -5,7 +5,8 @@ pub use backend::BackendType; mod credentials; pub use credentials::{ - check_peer_addr_is_in_list, endpoint_sni, ComputeUserInfoMaybeEndpoint, IpPattern, + check_peer_addr_is_in_list, endpoint_sni, ComputeUserInfoMaybeEndpoint, + ComputeUserInfoParseError, IpPattern, }; mod password_hack; @@ -14,8 +15,12 @@ use password_hack::PasswordHackPayload; mod flow; pub use flow::*; +use tokio::time::error::Elapsed; -use crate::{console, error::UserFacingError}; +use crate::{ + console, + error::{ReportableError, UserFacingError}, +}; use std::io; use thiserror::Error; @@ -67,6 +72,9 @@ pub enum AuthErrorImpl { #[error("Too many connections to this endpoint. Please try again later.")] TooManyConnections, + + #[error("Authentication timed out")] + UserTimeout(Elapsed), } #[derive(Debug, Error)] @@ -93,6 +101,10 @@ impl AuthError { pub fn is_auth_failed(&self) -> bool { matches!(self.0.as_ref(), AuthErrorImpl::AuthFailed(_)) } + + pub fn user_timeout(elapsed: Elapsed) -> Self { + AuthErrorImpl::UserTimeout(elapsed).into() + } } impl> From for AuthError { @@ -116,6 +128,27 @@ impl UserFacingError for AuthError { Io(_) => "Internal error".to_string(), IpAddressNotAllowed => self.to_string(), TooManyConnections => self.to_string(), + UserTimeout(_) => self.to_string(), + } + } +} + +impl ReportableError for AuthError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + use AuthErrorImpl::*; + match self.0.as_ref() { + Link(e) => e.get_error_kind(), + GetAuthInfo(e) => e.get_error_kind(), + WakeCompute(e) => e.get_error_kind(), + Sasl(e) => e.get_error_kind(), + AuthFailed(_) => crate::error::ErrorKind::User, + BadAuthMethod(_) => crate::error::ErrorKind::User, + MalformedPassword(_) => crate::error::ErrorKind::User, + MissingEndpointName => crate::error::ErrorKind::User, + Io(_) => crate::error::ErrorKind::ClientDisconnect, + IpAddressNotAllowed => crate::error::ErrorKind::User, + TooManyConnections => crate::error::ErrorKind::RateLimit, + UserTimeout(_) => crate::error::ErrorKind::User, } } } diff --git a/proxy/src/auth/backend/classic.rs b/proxy/src/auth/backend/classic.rs index 384063ceae..745dd75107 100644 --- a/proxy/src/auth/backend/classic.rs +++ b/proxy/src/auth/backend/classic.rs @@ -45,9 +45,9 @@ pub(super) async fn authenticate( } ) .await - .map_err(|error| { + .map_err(|e| { warn!("error processing scram messages error = authentication timed out, execution time exeeded {} seconds", config.scram_protocol_timeout.as_secs()); - auth::io::Error::new(auth::io::ErrorKind::TimedOut, error) + auth::AuthError::user_timeout(e) })??; let client_key = match auth_outcome { diff --git a/proxy/src/auth/backend/link.rs b/proxy/src/auth/backend/link.rs index d8ae362c03..c71637dd1a 100644 --- a/proxy/src/auth/backend/link.rs +++ b/proxy/src/auth/backend/link.rs @@ -2,7 +2,7 @@ use crate::{ auth, compute, console::{self, provider::NodeInfo}, context::RequestMonitoring, - error::UserFacingError, + error::{ReportableError, UserFacingError}, stream::PqStream, waiters, }; @@ -14,10 +14,6 @@ use tracing::{info, info_span}; #[derive(Debug, Error)] pub enum LinkAuthError { - /// Authentication error reported by the console. - #[error("Authentication failed: {0}")] - AuthFailed(String), - #[error(transparent)] WaiterRegister(#[from] waiters::RegisterError), @@ -30,10 +26,16 @@ pub enum LinkAuthError { impl UserFacingError for LinkAuthError { fn to_string_client(&self) -> String { - use LinkAuthError::*; + "Internal error".to_string() + } +} + +impl ReportableError for LinkAuthError { + fn get_error_kind(&self) -> crate::error::ErrorKind { match self { - AuthFailed(_) => self.to_string(), - _ => "Internal error".to_string(), + LinkAuthError::WaiterRegister(_) => crate::error::ErrorKind::Service, + LinkAuthError::WaiterWait(_) => crate::error::ErrorKind::Service, + LinkAuthError::Io(_) => crate::error::ErrorKind::ClientDisconnect, } } } diff --git a/proxy/src/auth/credentials.rs b/proxy/src/auth/credentials.rs index 875baaec47..d32609e44c 100644 --- a/proxy/src/auth/credentials.rs +++ b/proxy/src/auth/credentials.rs @@ -1,8 +1,12 @@ //! User credentials used in authentication. use crate::{ - auth::password_hack::parse_endpoint_param, context::RequestMonitoring, error::UserFacingError, - metrics::NUM_CONNECTION_ACCEPTED_BY_SNI, proxy::NeonOptions, serverless::SERVERLESS_DRIVER_SNI, + auth::password_hack::parse_endpoint_param, + context::RequestMonitoring, + error::{ReportableError, UserFacingError}, + metrics::NUM_CONNECTION_ACCEPTED_BY_SNI, + proxy::NeonOptions, + serverless::SERVERLESS_DRIVER_SNI, EndpointId, RoleName, }; use itertools::Itertools; @@ -39,6 +43,12 @@ pub enum ComputeUserInfoParseError { impl UserFacingError for ComputeUserInfoParseError {} +impl ReportableError for ComputeUserInfoParseError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + crate::error::ErrorKind::User + } +} + /// Various client credentials which we use for authentication. /// Note that we don't store any kind of client key or password here. #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/proxy/src/bin/pg_sni_router.rs b/proxy/src/bin/pg_sni_router.rs index 471be7af25..43b805e8a1 100644 --- a/proxy/src/bin/pg_sni_router.rs +++ b/proxy/src/bin/pg_sni_router.rs @@ -240,7 +240,9 @@ async fn ssl_handshake( ?unexpected, "unexpected startup packet, rejecting connection" ); - stream.throw_error_str(ERR_INSECURE_CONNECTION).await? + stream + .throw_error_str(ERR_INSECURE_CONNECTION, proxy::error::ErrorKind::User) + .await? } } } @@ -272,5 +274,10 @@ async fn handle_client( let client = tokio::net::TcpStream::connect(destination).await?; let metrics_aux: MetricsAuxInfo = Default::default(); - proxy::proxy::passthrough::proxy_pass(ctx, tls_stream, client, metrics_aux).await + + // doesn't yet matter as pg-sni-router doesn't report analytics logs + ctx.set_success(); + ctx.log(); + + proxy::proxy::passthrough::proxy_pass(tls_stream, client, metrics_aux).await } diff --git a/proxy/src/cancellation.rs b/proxy/src/cancellation.rs index d4ee657144..fe614628d8 100644 --- a/proxy/src/cancellation.rs +++ b/proxy/src/cancellation.rs @@ -1,24 +1,45 @@ -use anyhow::Context; use dashmap::DashMap; use pq_proto::CancelKeyData; use std::{net::SocketAddr, sync::Arc}; +use thiserror::Error; use tokio::net::TcpStream; use tokio_postgres::{CancelToken, NoTls}; use tracing::info; +use crate::error::ReportableError; + /// Enables serving `CancelRequest`s. #[derive(Default)] pub struct CancelMap(DashMap>); +#[derive(Debug, Error)] +pub enum CancelError { + #[error("{0}")] + IO(#[from] std::io::Error), + #[error("{0}")] + Postgres(#[from] tokio_postgres::Error), +} + +impl ReportableError for CancelError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + CancelError::IO(_) => crate::error::ErrorKind::Compute, + CancelError::Postgres(e) if e.as_db_error().is_some() => { + crate::error::ErrorKind::Postgres + } + CancelError::Postgres(_) => crate::error::ErrorKind::Compute, + } + } +} + impl CancelMap { /// Cancel a running query for the corresponding connection. - pub async fn cancel_session(&self, key: CancelKeyData) -> anyhow::Result<()> { + pub async fn cancel_session(&self, key: CancelKeyData) -> Result<(), CancelError> { // NB: we should immediately release the lock after cloning the token. - let cancel_closure = self - .0 - .get(&key) - .and_then(|x| x.clone()) - .with_context(|| format!("query cancellation key not found: {key}"))?; + let Some(cancel_closure) = self.0.get(&key).and_then(|x| x.clone()) else { + tracing::warn!("query cancellation key not found: {key}"); + return Ok(()); + }; info!("cancelling query per user's request using key {key}"); cancel_closure.try_cancel_query().await @@ -81,7 +102,7 @@ impl CancelClosure { } /// Cancels the query running on user's compute node. - pub async fn try_cancel_query(self) -> anyhow::Result<()> { + async fn try_cancel_query(self) -> Result<(), CancelError> { let socket = TcpStream::connect(self.socket_addr).await?; self.cancel_token.cancel_query_raw(socket, NoTls).await?; diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index aef1aab733..83940d80ec 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -1,6 +1,10 @@ use crate::{ - auth::parse_endpoint_param, cancellation::CancelClosure, console::errors::WakeComputeError, - context::RequestMonitoring, error::UserFacingError, metrics::NUM_DB_CONNECTIONS_GAUGE, + auth::parse_endpoint_param, + cancellation::CancelClosure, + console::errors::WakeComputeError, + context::RequestMonitoring, + error::{ReportableError, UserFacingError}, + metrics::NUM_DB_CONNECTIONS_GAUGE, proxy::neon_option, }; use futures::{FutureExt, TryFutureExt}; @@ -58,6 +62,20 @@ impl UserFacingError for ConnectionError { } } +impl ReportableError for ConnectionError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + ConnectionError::Postgres(e) if e.as_db_error().is_some() => { + crate::error::ErrorKind::Postgres + } + ConnectionError::Postgres(_) => crate::error::ErrorKind::Compute, + ConnectionError::CouldNotConnect(_) => crate::error::ErrorKind::Compute, + ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute, + ConnectionError::WakeComputeError(e) => e.get_error_kind(), + } + } +} + /// A pair of `ClientKey` & `ServerKey` for `SCRAM-SHA-256`. pub type ScramKeys = tokio_postgres::config::ScramKeys<32>; diff --git a/proxy/src/console/provider.rs b/proxy/src/console/provider.rs index c53d929470..e5cad42753 100644 --- a/proxy/src/console/provider.rs +++ b/proxy/src/console/provider.rs @@ -20,7 +20,7 @@ use tracing::info; pub mod errors { use crate::{ - error::{io_error, UserFacingError}, + error::{io_error, ReportableError, UserFacingError}, http, proxy::retry::ShouldRetry, }; @@ -81,6 +81,15 @@ pub mod errors { } } + impl ReportableError for ApiError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + ApiError::Console { .. } => crate::error::ErrorKind::ControlPlane, + ApiError::Transport(_) => crate::error::ErrorKind::ControlPlane, + } + } + } + impl ShouldRetry for ApiError { fn could_retry(&self) -> bool { match self { @@ -150,6 +159,16 @@ pub mod errors { } } } + + impl ReportableError for GetAuthInfoError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + GetAuthInfoError::BadSecret => crate::error::ErrorKind::ControlPlane, + GetAuthInfoError::ApiError(_) => crate::error::ErrorKind::ControlPlane, + } + } + } + #[derive(Debug, Error)] pub enum WakeComputeError { #[error("Console responded with a malformed compute address: {0}")] @@ -194,6 +213,16 @@ pub mod errors { } } } + + impl ReportableError for WakeComputeError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + WakeComputeError::BadComputeAddress(_) => crate::error::ErrorKind::ControlPlane, + WakeComputeError::ApiError(e) => e.get_error_kind(), + WakeComputeError::TimeoutError => crate::error::ErrorKind::RateLimit, + } + } + } } /// Auth secret which is managed by the cloud. diff --git a/proxy/src/context.rs b/proxy/src/context.rs index fe204534b7..d2bf3f68d3 100644 --- a/proxy/src/context.rs +++ b/proxy/src/context.rs @@ -8,8 +8,10 @@ use tokio::sync::mpsc; use uuid::Uuid; use crate::{ - console::messages::MetricsAuxInfo, error::ErrorKind, metrics::LatencyTimer, BranchId, - EndpointId, ProjectId, RoleName, + console::messages::MetricsAuxInfo, + error::ErrorKind, + metrics::{LatencyTimer, ENDPOINT_ERRORS_BY_KIND, ERROR_BY_KIND}, + BranchId, EndpointId, ProjectId, RoleName, }; pub mod parquet; @@ -108,6 +110,18 @@ impl RequestMonitoring { self.user = Some(user); } + pub fn set_error_kind(&mut self, kind: ErrorKind) { + ERROR_BY_KIND + .with_label_values(&[kind.to_metric_label()]) + .inc(); + if let Some(ep) = &self.endpoint_id { + ENDPOINT_ERRORS_BY_KIND + .with_label_values(&[kind.to_metric_label()]) + .measure(ep); + } + self.error_kind = Some(kind); + } + pub fn set_success(&mut self) { self.success = true; } diff --git a/proxy/src/context/parquet.rs b/proxy/src/context/parquet.rs index 8510c5c586..0fe46915bc 100644 --- a/proxy/src/context/parquet.rs +++ b/proxy/src/context/parquet.rs @@ -108,7 +108,7 @@ impl From for RequestData { branch: value.branch.as_deref().map(String::from), protocol: value.protocol, region: value.region, - error: value.error_kind.as_ref().map(|e| e.to_str()), + error: value.error_kind.as_ref().map(|e| e.to_metric_label()), success: value.success, duration_us: SystemTime::from(value.first_packet) .elapsed() diff --git a/proxy/src/error.rs b/proxy/src/error.rs index 5b2dd7ecfd..eafe92bf48 100644 --- a/proxy/src/error.rs +++ b/proxy/src/error.rs @@ -17,7 +17,7 @@ pub fn log_error(e: E) -> E { /// NOTE: This trait should not be implemented for [`anyhow::Error`], since it /// is way too convenient and tends to proliferate all across the codebase, /// ultimately leading to accidental leaks of sensitive data. -pub trait UserFacingError: fmt::Display { +pub trait UserFacingError: ReportableError { /// Format the error for client, stripping all sensitive info. /// /// Although this might be a no-op for many types, it's highly @@ -29,13 +29,13 @@ pub trait UserFacingError: fmt::Display { } } -#[derive(Clone)] +#[derive(Copy, Clone, Debug)] pub enum ErrorKind { /// Wrong password, unknown endpoint, protocol violation, etc... User, /// Network error between user and proxy. Not necessarily user error - Disconnect, + ClientDisconnect, /// Proxy self-imposed rate limits RateLimit, @@ -46,6 +46,9 @@ pub enum ErrorKind { /// Error communicating with control plane ControlPlane, + /// Postgres error + Postgres, + /// Error communicating with compute Compute, } @@ -54,11 +57,36 @@ impl ErrorKind { pub fn to_str(&self) -> &'static str { match self { ErrorKind::User => "request failed due to user error", - ErrorKind::Disconnect => "client disconnected", + ErrorKind::ClientDisconnect => "client disconnected", ErrorKind::RateLimit => "request cancelled due to rate limit", ErrorKind::Service => "internal service error", ErrorKind::ControlPlane => "non-retryable control plane error", - ErrorKind::Compute => "non-retryable compute error (or exhausted retry capacity)", + ErrorKind::Postgres => "postgres error", + ErrorKind::Compute => { + "non-retryable compute connection error (or exhausted retry capacity)" + } + } + } + + pub fn to_metric_label(&self) -> &'static str { + match self { + ErrorKind::User => "user", + ErrorKind::ClientDisconnect => "clientdisconnect", + ErrorKind::RateLimit => "ratelimit", + ErrorKind::Service => "service", + ErrorKind::ControlPlane => "controlplane", + ErrorKind::Postgres => "postgres", + ErrorKind::Compute => "compute", } } } + +pub trait ReportableError: fmt::Display + Send + 'static { + fn get_error_kind(&self) -> ErrorKind; +} + +impl ReportableError for tokio::time::error::Elapsed { + fn get_error_kind(&self) -> ErrorKind { + ErrorKind::RateLimit + } +} diff --git a/proxy/src/metrics.rs b/proxy/src/metrics.rs index e2d96a9c27..ccf89f9b05 100644 --- a/proxy/src/metrics.rs +++ b/proxy/src/metrics.rs @@ -274,3 +274,22 @@ pub static CONNECTING_ENDPOINTS: Lazy> = Lazy::new(|| { ) .unwrap() }); + +pub static ERROR_BY_KIND: Lazy = Lazy::new(|| { + register_int_counter_vec!( + "proxy_errors_total", + "Number of errors by a given classification", + &["type"], + ) + .unwrap() +}); + +pub static ENDPOINT_ERRORS_BY_KIND: Lazy> = Lazy::new(|| { + register_hll_vec!( + 32, + "proxy_endpoints_affected_by_errors", + "Number of endpoints affected by errors of a given classification", + &["type"], + ) + .unwrap() +}); diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs index b3b221d3e2..50e22ec72a 100644 --- a/proxy/src/proxy.rs +++ b/proxy/src/proxy.rs @@ -13,9 +13,10 @@ use crate::{ compute, config::{ProxyConfig, TlsConfig}, context::RequestMonitoring, + error::ReportableError, metrics::{NUM_CLIENT_CONNECTION_GAUGE, NUM_CONNECTION_REQUESTS_GAUGE}, protocol2::WithClientIp, - proxy::{handshake::handshake, passthrough::proxy_pass}, + proxy::handshake::{handshake, HandshakeData}, rate_limiter::EndpointRateLimiter, stream::{PqStream, Stream}, EndpointCacheKey, @@ -28,14 +29,17 @@ use pq_proto::{BeMessage as Be, StartupMessageParams}; use regex::Regex; use smol_str::{format_smolstr, SmolStr}; use std::sync::Arc; +use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use tokio_util::sync::CancellationToken; use tracing::{error, info, info_span, Instrument}; -use self::connect_compute::{connect_to_compute, TcpMechanism}; +use self::{ + connect_compute::{connect_to_compute, TcpMechanism}, + passthrough::ProxyPassthrough, +}; const ERR_INSECURE_CONNECTION: &str = "connection is insecure (try using `sslmode=require`)"; -const ERR_PROTO_VIOLATION: &str = "protocol violation"; pub async fn run_until_cancelled( f: F, @@ -98,14 +102,14 @@ pub async fn task_main( bail!("missing required client IP"); } - let mut ctx = RequestMonitoring::new(session_id, peer_addr, "tcp", &config.region); - socket .inner .set_nodelay(true) .context("failed to set socket option")?; - handle_client( + let mut ctx = RequestMonitoring::new(session_id, peer_addr, "tcp", &config.region); + + let res = handle_client( config, &mut ctx, cancel_map, @@ -113,7 +117,26 @@ pub async fn task_main( ClientMode::Tcp, endpoint_rate_limiter, ) - .await + .await; + + match res { + Err(e) => { + // todo: log and push to ctx the error kind + ctx.set_error_kind(e.get_error_kind()); + ctx.log(); + Err(e.into()) + } + Ok(None) => { + ctx.set_success(); + ctx.log(); + Ok(()) + } + Ok(Some(p)) => { + ctx.set_success(); + ctx.log(); + p.proxy_pass().await + } + } } .unwrap_or_else(move |e| { // Acknowledge that the task has finished with an error. @@ -169,6 +192,37 @@ impl ClientMode { } } +#[derive(Debug, Error)] +// almost all errors should be reported to the user, but there's a few cases where we cannot +// 1. Cancellation: we are not allowed to tell the client any cancellation statuses for security reasons +// 2. Handshake: handshake reports errors if it can, otherwise if the handshake fails due to protocol violation, +// we cannot be sure the client even understands our error message +// 3. PrepareClient: The client disconnected, so we can't tell them anyway... +pub enum ClientRequestError { + #[error("{0}")] + Cancellation(#[from] cancellation::CancelError), + #[error("{0}")] + Handshake(#[from] handshake::HandshakeError), + #[error("{0}")] + HandshakeTimeout(#[from] tokio::time::error::Elapsed), + #[error("{0}")] + PrepareClient(#[from] std::io::Error), + #[error("{0}")] + ReportedError(#[from] crate::stream::ReportedError), +} + +impl ReportableError for ClientRequestError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + ClientRequestError::Cancellation(e) => e.get_error_kind(), + ClientRequestError::Handshake(e) => e.get_error_kind(), + ClientRequestError::HandshakeTimeout(_) => crate::error::ErrorKind::RateLimit, + ClientRequestError::ReportedError(e) => e.get_error_kind(), + ClientRequestError::PrepareClient(_) => crate::error::ErrorKind::ClientDisconnect, + } + } +} + pub async fn handle_client( config: &'static ProxyConfig, ctx: &mut RequestMonitoring, @@ -176,7 +230,7 @@ pub async fn handle_client( stream: S, mode: ClientMode, endpoint_rate_limiter: Arc, -) -> anyhow::Result<()> { +) -> Result>, ClientRequestError> { info!( protocol = ctx.protocol, "handling interactive connection from client" @@ -193,11 +247,16 @@ pub async fn handle_client( let tls = config.tls_config.as_ref(); let pause = ctx.latency_timer.pause(); - let do_handshake = handshake(stream, mode.handshake_tls(tls), &cancel_map); + let do_handshake = handshake(stream, mode.handshake_tls(tls)); let (mut stream, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake).await?? { - Some(x) => x, - None => return Ok(()), // it's a cancellation request + HandshakeData::Startup(stream, params) => (stream, params), + HandshakeData::Cancel(cancel_key_data) => { + return Ok(cancel_map + .cancel_session(cancel_key_data) + .await + .map(|()| None)?) + } }; drop(pause); @@ -222,7 +281,7 @@ pub async fn handle_client( if !endpoint_rate_limiter.check(ep) { return stream .throw_error(auth::AuthError::too_many_connections()) - .await; + .await?; } } @@ -242,7 +301,7 @@ pub async fn handle_client( let app = params.get("application_name"); let params_span = tracing::info_span!("", ?user, ?db, ?app); - return stream.throw_error(e).instrument(params_span).await; + return stream.throw_error(e).instrument(params_span).await?; } }; @@ -268,7 +327,13 @@ pub async fn handle_client( let (stream, read_buf) = stream.into_inner(); node.stream.write_all(&read_buf).await?; - proxy_pass(ctx, stream, node.stream, aux).await + Ok(Some(ProxyPassthrough { + client: stream, + compute: node, + aux, + req: _request_gauge, + conn: _client_gauge, + })) } /// Finish client connection initialization: confirm auth success, send params, etc. @@ -277,7 +342,7 @@ async fn prepare_client_connection( node: &compute::PostgresConnection, session: &cancellation::Session, stream: &mut PqStream, -) -> anyhow::Result<()> { +) -> Result<(), std::io::Error> { // Register compute's query cancellation token and produce a new, unique one. // The new token (cancel_key_data) will be sent to the client. let cancel_key_data = session.enable_query_cancellation(node.cancel_closure.clone()); diff --git a/proxy/src/proxy/handshake.rs b/proxy/src/proxy/handshake.rs index 1ad8da20d7..4665e07d23 100644 --- a/proxy/src/proxy/handshake.rs +++ b/proxy/src/proxy/handshake.rs @@ -1,15 +1,60 @@ -use anyhow::{bail, Context}; -use pq_proto::{BeMessage as Be, FeStartupPacket, StartupMessageParams}; +use pq_proto::{BeMessage as Be, CancelKeyData, FeStartupPacket, StartupMessageParams}; +use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite}; use tracing::info; use crate::{ - cancellation::CancelMap, config::TlsConfig, - proxy::{ERR_INSECURE_CONNECTION, ERR_PROTO_VIOLATION}, - stream::{PqStream, Stream}, + error::ReportableError, + proxy::ERR_INSECURE_CONNECTION, + stream::{PqStream, Stream, StreamUpgradeError}, }; +#[derive(Error, Debug)] +pub enum HandshakeError { + #[error("data is sent before server replied with EncryptionResponse")] + EarlyData, + + #[error("protocol violation")] + ProtocolViolation, + + #[error("missing certificate")] + MissingCertificate, + + #[error("{0}")] + StreamUpgradeError(#[from] StreamUpgradeError), + + #[error("{0}")] + Io(#[from] std::io::Error), + + #[error("{0}")] + ReportedError(#[from] crate::stream::ReportedError), +} + +impl ReportableError for HandshakeError { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + HandshakeError::EarlyData => crate::error::ErrorKind::User, + HandshakeError::ProtocolViolation => crate::error::ErrorKind::User, + // This error should not happen, but will if we have no default certificate and + // the client sends no SNI extension. + // If they provide SNI then we can be sure there is a certificate that matches. + HandshakeError::MissingCertificate => crate::error::ErrorKind::Service, + HandshakeError::StreamUpgradeError(upgrade) => match upgrade { + StreamUpgradeError::AlreadyTls => crate::error::ErrorKind::Service, + StreamUpgradeError::Io(_) => crate::error::ErrorKind::ClientDisconnect, + }, + HandshakeError::Io(_) => crate::error::ErrorKind::ClientDisconnect, + HandshakeError::ReportedError(e) => e.get_error_kind(), + } + } +} + +pub enum HandshakeData { + Startup(PqStream>, StartupMessageParams), + Cancel(CancelKeyData), +} + /// Establish a (most probably, secure) connection with the client. /// For better testing experience, `stream` can be any object satisfying the traits. /// It's easier to work with owned `stream` here as we need to upgrade it to TLS; @@ -18,8 +63,7 @@ use crate::{ pub async fn handshake( stream: S, mut tls: Option<&TlsConfig>, - cancel_map: &CancelMap, -) -> anyhow::Result>, StartupMessageParams)>> { +) -> Result, HandshakeError> { // Client may try upgrading to each protocol only once let (mut tried_ssl, mut tried_gss) = (false, false); @@ -49,14 +93,14 @@ pub async fn handshake( // pipelining in our node js driver. We should probably // support that by chaining read_buf with the stream. if !read_buf.is_empty() { - bail!("data is sent before server replied with EncryptionResponse"); + return Err(HandshakeError::EarlyData); } let tls_stream = raw.upgrade(tls.to_server_config()).await?; let (_, tls_server_end_point) = tls .cert_resolver .resolve(tls_stream.get_ref().1.server_name()) - .context("missing certificate")?; + .ok_or(HandshakeError::MissingCertificate)?; stream = PqStream::new(Stream::Tls { tls: Box::new(tls_stream), @@ -64,7 +108,7 @@ pub async fn handshake( }); } } - _ => bail!(ERR_PROTO_VIOLATION), + _ => return Err(HandshakeError::ProtocolViolation), }, GssEncRequest => match stream.get_ref() { Stream::Raw { .. } if !tried_gss => { @@ -73,23 +117,23 @@ pub async fn handshake( // Currently, we don't support GSSAPI stream.write_message(&Be::EncryptionResponse(false)).await?; } - _ => bail!(ERR_PROTO_VIOLATION), + _ => return Err(HandshakeError::ProtocolViolation), }, StartupMessage { params, .. } => { // Check that the config has been consumed during upgrade // OR we didn't provide it at all (for dev purposes). if tls.is_some() { - stream.throw_error_str(ERR_INSECURE_CONNECTION).await?; + return stream + .throw_error_str(ERR_INSECURE_CONNECTION, crate::error::ErrorKind::User) + .await?; } info!(session_type = "normal", "successful handshake"); - break Ok(Some((stream, params))); + break Ok(HandshakeData::Startup(stream, params)); } CancelRequest(cancel_key_data) => { - cancel_map.cancel_session(cancel_key_data).await?; - info!(session_type = "cancellation", "successful handshake"); - break Ok(None); + break Ok(HandshakeData::Cancel(cancel_key_data)); } } } diff --git a/proxy/src/proxy/passthrough.rs b/proxy/src/proxy/passthrough.rs index 53e0c3c8f3..b7018c6fb5 100644 --- a/proxy/src/proxy/passthrough.rs +++ b/proxy/src/proxy/passthrough.rs @@ -1,9 +1,11 @@ use crate::{ + compute::PostgresConnection, console::messages::MetricsAuxInfo, - context::RequestMonitoring, metrics::NUM_BYTES_PROXIED_COUNTER, + stream::Stream, usage_metrics::{Ids, USAGE_METRICS}, }; +use metrics::IntCounterPairGuard; use tokio::io::{AsyncRead, AsyncWrite}; use tracing::info; use utils::measured_stream::MeasuredStream; @@ -11,14 +13,10 @@ use utils::measured_stream::MeasuredStream; /// Forward bytes in both directions (client <-> compute). #[tracing::instrument(skip_all)] pub async fn proxy_pass( - ctx: &mut RequestMonitoring, client: impl AsyncRead + AsyncWrite + Unpin, compute: impl AsyncRead + AsyncWrite + Unpin, aux: MetricsAuxInfo, ) -> anyhow::Result<()> { - ctx.set_success(); - ctx.log(); - let usage = USAGE_METRICS.register(Ids { endpoint_id: aux.endpoint_id.clone(), branch_id: aux.branch_id.clone(), @@ -51,3 +49,18 @@ pub async fn proxy_pass( Ok(()) } + +pub struct ProxyPassthrough { + pub client: Stream, + pub compute: PostgresConnection, + pub aux: MetricsAuxInfo, + + pub req: IntCounterPairGuard, + pub conn: IntCounterPairGuard, +} + +impl ProxyPassthrough { + pub async fn proxy_pass(self) -> anyhow::Result<()> { + proxy_pass(self.client, self.compute.stream, self.aux).await + } +} diff --git a/proxy/src/proxy/tests.rs b/proxy/src/proxy/tests.rs index 656cabac75..3e961afb41 100644 --- a/proxy/src/proxy/tests.rs +++ b/proxy/src/proxy/tests.rs @@ -163,11 +163,11 @@ async fn dummy_proxy( tls: Option, auth: impl TestAuth + Send, ) -> anyhow::Result<()> { - let cancel_map = CancelMap::default(); let client = WithClientIp::new(client); - let (mut stream, _params) = handshake(client, tls.as_ref(), &cancel_map) - .await? - .context("handshake failed")?; + let mut stream = match handshake(client, tls.as_ref()).await? { + HandshakeData::Startup(stream, _) => stream, + HandshakeData::Cancel(_) => bail!("cancellation not supported"), + }; auth.authenticate(&mut stream).await?; diff --git a/proxy/src/proxy/tests/mitm.rs b/proxy/src/proxy/tests/mitm.rs index a0a84a1dc0..ed89e51754 100644 --- a/proxy/src/proxy/tests/mitm.rs +++ b/proxy/src/proxy/tests/mitm.rs @@ -35,12 +35,10 @@ async fn proxy_mitm( tokio::spawn(async move { // begin handshake with end_server let end_server = connect_tls(server2, client_config2.make_tls_connect().unwrap()).await; - // process handshake with end_client - let (end_client, startup) = - handshake(client1, Some(&server_config1), &CancelMap::default()) - .await - .unwrap() - .unwrap(); + let (end_client, startup) = match handshake(client1, Some(&server_config1)).await.unwrap() { + HandshakeData::Startup(stream, params) => (stream, params), + HandshakeData::Cancel(_) => panic!("cancellation not supported"), + }; let mut end_server = tokio_util::codec::Framed::new(end_server, PgFrame); let (end_client, buf) = end_client.framed.into_inner(); diff --git a/proxy/src/sasl.rs b/proxy/src/sasl.rs index da1cf21c6a..1cf8b53e11 100644 --- a/proxy/src/sasl.rs +++ b/proxy/src/sasl.rs @@ -10,7 +10,7 @@ mod channel_binding; mod messages; mod stream; -use crate::error::UserFacingError; +use crate::error::{ReportableError, UserFacingError}; use std::io; use thiserror::Error; @@ -48,6 +48,18 @@ impl UserFacingError for Error { } } +impl ReportableError for Error { + fn get_error_kind(&self) -> crate::error::ErrorKind { + match self { + Error::ChannelBindingFailed(_) => crate::error::ErrorKind::User, + Error::ChannelBindingBadMethod(_) => crate::error::ErrorKind::User, + Error::BadClientMessage(_) => crate::error::ErrorKind::User, + Error::MissingBinding => crate::error::ErrorKind::Service, + Error::Io(_) => crate::error::ErrorKind::ClientDisconnect, + } + } +} + /// A convenient result type for SASL exchange. pub type Result = std::result::Result; diff --git a/proxy/src/serverless.rs b/proxy/src/serverless.rs index 58aa925a6a..a20600b94a 100644 --- a/proxy/src/serverless.rs +++ b/proxy/src/serverless.rs @@ -109,10 +109,9 @@ pub async fn task_main( let make_svc = hyper::service::make_service_fn( |stream: &tokio_rustls::server::TlsStream>| { - let (io, tls) = stream.get_ref(); + let (io, _) = stream.get_ref(); let client_addr = io.client_addr(); let remote_addr = io.inner.remote_addr(); - let sni_name = tls.server_name().map(|s| s.to_string()); let backend = backend.clone(); let ws_connections = ws_connections.clone(); let endpoint_rate_limiter = endpoint_rate_limiter.clone(); @@ -125,7 +124,6 @@ pub async fn task_main( }; Ok(MetricService::new(hyper::service::service_fn( move |req: Request| { - let sni_name = sni_name.clone(); let backend = backend.clone(); let ws_connections = ws_connections.clone(); let endpoint_rate_limiter = endpoint_rate_limiter.clone(); @@ -141,7 +139,6 @@ pub async fn task_main( ws_connections, cancel_map, session_id, - sni_name, peer_addr.ip(), endpoint_rate_limiter, ) @@ -210,7 +207,6 @@ async fn request_handler( ws_connections: TaskTracker, cancel_map: Arc, session_id: uuid::Uuid, - sni_hostname: Option, peer_addr: IpAddr, endpoint_rate_limiter: Arc, ) -> Result, ApiError> { @@ -230,11 +226,11 @@ async fn request_handler( ws_connections.spawn( async move { - let mut ctx = RequestMonitoring::new(session_id, peer_addr, "ws", &config.region); + let ctx = RequestMonitoring::new(session_id, peer_addr, "ws", &config.region); if let Err(e) = websocket::serve_websocket( config, - &mut ctx, + ctx, websocket, cancel_map, host, @@ -251,9 +247,9 @@ async fn request_handler( // Return the response so the spawned future can continue. Ok(response) } else if request.uri().path() == "/sql" && request.method() == Method::POST { - let mut ctx = RequestMonitoring::new(session_id, peer_addr, "http", &config.region); + let ctx = RequestMonitoring::new(session_id, peer_addr, "http", &config.region); - sql_over_http::handle(config, &mut ctx, request, sni_hostname, backend).await + sql_over_http::handle(config, ctx, request, backend).await } else if request.uri().path() == "/sql" && request.method() == Method::OPTIONS { Response::builder() .header("Allow", "OPTIONS, POST") diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index 466a74f0ea..03257e9161 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -1,6 +1,5 @@ use std::{sync::Arc, time::Duration}; -use anyhow::Context; use async_trait::async_trait; use tracing::info; @@ -8,7 +7,10 @@ use crate::{ auth::{backend::ComputeCredentialKeys, check_peer_addr_is_in_list, AuthError}, compute, config::ProxyConfig, - console::CachedNodeInfo, + console::{ + errors::{GetAuthInfoError, WakeComputeError}, + CachedNodeInfo, + }, context::RequestMonitoring, proxy::connect_compute::ConnectMechanism, }; @@ -66,7 +68,7 @@ impl PoolingBackend { conn_info: ConnInfo, keys: ComputeCredentialKeys, force_new: bool, - ) -> anyhow::Result> { + ) -> Result, HttpConnError> { let maybe_client = if !force_new { info!("pool: looking for an existing connection"); self.pool.get(ctx, &conn_info).await? @@ -90,7 +92,7 @@ impl PoolingBackend { let mut node_info = backend .wake_compute(ctx) .await? - .context("missing cache entry from wake_compute")?; + .ok_or(HttpConnError::NoComputeInfo)?; match keys { #[cfg(any(test, feature = "testing"))] @@ -114,6 +116,23 @@ impl PoolingBackend { } } +#[derive(Debug, thiserror::Error)] +pub enum HttpConnError { + #[error("pooled connection closed at inconsistent state")] + ConnectionClosedAbruptly(#[from] tokio::sync::watch::error::SendError), + #[error("could not connection to compute")] + ConnectionError(#[from] tokio_postgres::Error), + + #[error("could not get auth info")] + GetAuthInfo(#[from] GetAuthInfoError), + #[error("user not authenticated")] + AuthError(#[from] AuthError), + #[error("wake_compute returned error")] + WakeCompute(#[from] WakeComputeError), + #[error("wake_compute returned nothing")] + NoComputeInfo, +} + struct TokioMechanism { pool: Arc>, conn_info: ConnInfo, @@ -124,7 +143,7 @@ struct TokioMechanism { impl ConnectMechanism for TokioMechanism { type Connection = Client; type ConnectError = tokio_postgres::Error; - type Error = anyhow::Error; + type Error = HttpConnError; async fn connect_once( &self, diff --git a/proxy/src/serverless/conn_pool.rs b/proxy/src/serverless/conn_pool.rs index a7b2c532d2..f92793096b 100644 --- a/proxy/src/serverless/conn_pool.rs +++ b/proxy/src/serverless/conn_pool.rs @@ -28,6 +28,8 @@ use crate::{ use tracing::{debug, error, warn, Span}; use tracing::{info, info_span, Instrument}; +use super::backend::HttpConnError; + pub const APP_NAME: SmolStr = SmolStr::new_inline("/sql_over_http"); #[derive(Debug, Clone)] @@ -358,7 +360,7 @@ impl GlobalConnPool { self: &Arc, ctx: &mut RequestMonitoring, conn_info: &ConnInfo, - ) -> anyhow::Result>> { + ) -> Result>, HttpConnError> { let mut client: Option> = None; let endpoint_pool = self.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key()); diff --git a/proxy/src/serverless/json.rs b/proxy/src/serverless/json.rs index a089d34040..c22c63e85b 100644 --- a/proxy/src/serverless/json.rs +++ b/proxy/src/serverless/json.rs @@ -60,6 +60,20 @@ fn json_array_to_pg_array(value: &Value) -> Option { } } +#[derive(Debug, thiserror::Error)] +pub enum JsonConversionError { + #[error("internal error compute returned invalid data: {0}")] + AsTextError(tokio_postgres::Error), + #[error("parse int error: {0}")] + ParseIntError(#[from] std::num::ParseIntError), + #[error("parse float error: {0}")] + ParseFloatError(#[from] std::num::ParseFloatError), + #[error("parse json error: {0}")] + ParseJsonError(#[from] serde_json::Error), + #[error("unbalanced array")] + UnbalancedArray, +} + // // Convert postgres row with text-encoded values to JSON object // @@ -68,7 +82,7 @@ pub fn pg_text_row_to_json( columns: &[Type], raw_output: bool, array_mode: bool, -) -> Result { +) -> Result { let iter = row .columns() .iter() @@ -76,7 +90,7 @@ pub fn pg_text_row_to_json( .enumerate() .map(|(i, (column, typ))| { let name = column.name(); - let pg_value = row.as_text(i)?; + let pg_value = row.as_text(i).map_err(JsonConversionError::AsTextError)?; let json_value = if raw_output { match pg_value { Some(v) => Value::String(v.to_string()), @@ -92,10 +106,10 @@ pub fn pg_text_row_to_json( // drop keys and aggregate into array let arr = iter .map(|r| r.map(|(_key, val)| val)) - .collect::, anyhow::Error>>()?; + .collect::, JsonConversionError>>()?; Ok(Value::Array(arr)) } else { - let obj = iter.collect::, anyhow::Error>>()?; + let obj = iter.collect::, JsonConversionError>>()?; Ok(Value::Object(obj)) } } @@ -103,7 +117,7 @@ pub fn pg_text_row_to_json( // // Convert postgres text-encoded value to JSON value // -fn pg_text_to_json(pg_value: Option<&str>, pg_type: &Type) -> Result { +fn pg_text_to_json(pg_value: Option<&str>, pg_type: &Type) -> Result { if let Some(val) = pg_value { if let Kind::Array(elem_type) = pg_type.kind() { return pg_array_parse(val, elem_type); @@ -142,7 +156,7 @@ fn pg_text_to_json(pg_value: Option<&str>, pg_type: &Type) -> Result Result { +fn pg_array_parse(pg_array: &str, elem_type: &Type) -> Result { _pg_array_parse(pg_array, elem_type, false).map(|(v, _)| v) } @@ -150,7 +164,7 @@ fn _pg_array_parse( pg_array: &str, elem_type: &Type, nested: bool, -) -> Result<(Value, usize), anyhow::Error> { +) -> Result<(Value, usize), JsonConversionError> { let mut pg_array_chr = pg_array.char_indices(); let mut level = 0; let mut quote = false; @@ -170,7 +184,7 @@ fn _pg_array_parse( entry: &mut String, entries: &mut Vec, elem_type: &Type, - ) -> Result<(), anyhow::Error> { + ) -> Result<(), JsonConversionError> { if !entry.is_empty() { // While in usual postgres response we get nulls as None and everything else // as Some(&str), in arrays we get NULL as unquoted 'NULL' string (while @@ -234,7 +248,7 @@ fn _pg_array_parse( } if level != 0 { - return Err(anyhow::anyhow!("unbalanced array")); + return Err(JsonConversionError::UnbalancedArray); } Ok((Value::Array(entries), 0)) diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 25e8813625..401022347e 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use anyhow::bail; -use anyhow::Context; use futures::pin_mut; use futures::StreamExt; use hyper::body::HttpBody; @@ -29,9 +28,11 @@ use utils::http::json::json_response; use crate::auth::backend::ComputeUserInfo; use crate::auth::endpoint_sni; +use crate::auth::ComputeUserInfoParseError; use crate::config::ProxyConfig; use crate::config::TlsConfig; use crate::context::RequestMonitoring; +use crate::error::ReportableError; use crate::metrics::HTTP_CONTENT_LENGTH; use crate::metrics::NUM_CONNECTION_REQUESTS_GAUGE; use crate::proxy::NeonOptions; @@ -41,7 +42,6 @@ use super::backend::PoolingBackend; use super::conn_pool::ConnInfo; use super::json::json_to_pg_text; use super::json::pg_text_row_to_json; -use super::SERVERLESS_DRIVER_SNI; #[derive(serde::Deserialize)] #[serde(rename_all = "camelCase")] @@ -86,67 +86,70 @@ where Ok(json_to_pg_text(json)) } +#[derive(Debug, thiserror::Error)] +pub enum ConnInfoError { + #[error("invalid header: {0}")] + InvalidHeader(&'static str), + #[error("invalid connection string: {0}")] + UrlParseError(#[from] url::ParseError), + #[error("incorrect scheme")] + IncorrectScheme, + #[error("missing database name")] + MissingDbName, + #[error("invalid database name")] + InvalidDbName, + #[error("missing username")] + MissingUsername, + #[error("missing password")] + MissingPassword, + #[error("missing hostname")] + MissingHostname, + #[error("invalid hostname: {0}")] + InvalidEndpoint(#[from] ComputeUserInfoParseError), + #[error("malformed endpoint")] + MalformedEndpoint, +} + fn get_conn_info( ctx: &mut RequestMonitoring, headers: &HeaderMap, - sni_hostname: Option, tls: &TlsConfig, -) -> Result { +) -> Result { let connection_string = headers .get("Neon-Connection-String") - .ok_or(anyhow::anyhow!("missing connection string"))? - .to_str()?; + .ok_or(ConnInfoError::InvalidHeader("Neon-Connection-String"))? + .to_str() + .map_err(|_| ConnInfoError::InvalidHeader("Neon-Connection-String"))?; let connection_url = Url::parse(connection_string)?; let protocol = connection_url.scheme(); if protocol != "postgres" && protocol != "postgresql" { - return Err(anyhow::anyhow!( - "connection string must start with postgres: or postgresql:" - )); + return Err(ConnInfoError::IncorrectScheme); } let mut url_path = connection_url .path_segments() - .ok_or(anyhow::anyhow!("missing database name"))?; + .ok_or(ConnInfoError::MissingDbName)?; - let dbname = url_path - .next() - .ok_or(anyhow::anyhow!("invalid database name"))?; + let dbname = url_path.next().ok_or(ConnInfoError::InvalidDbName)?; let username = RoleName::from(connection_url.username()); if username.is_empty() { - return Err(anyhow::anyhow!("missing username")); + return Err(ConnInfoError::MissingUsername); } ctx.set_user(username.clone()); let password = connection_url .password() - .ok_or(anyhow::anyhow!("no password"))?; - - // TLS certificate selector now based on SNI hostname, so if we are running here - // we are sure that SNI hostname is set to one of the configured domain names. - let sni_hostname = sni_hostname.ok_or(anyhow::anyhow!("no SNI hostname set"))?; + .ok_or(ConnInfoError::MissingPassword)?; let hostname = connection_url .host_str() - .ok_or(anyhow::anyhow!("no host"))?; + .ok_or(ConnInfoError::MissingHostname)?; - let host_header = headers - .get("host") - .and_then(|h| h.to_str().ok()) - .and_then(|h| h.split(':').next()); - - // sni_hostname has to be either the same as hostname or the one used in serverless driver. - if !check_matches(&sni_hostname, hostname)? { - return Err(anyhow::anyhow!("mismatched SNI hostname and hostname")); - } else if let Some(h) = host_header { - if h != sni_hostname { - return Err(anyhow::anyhow!("mismatched host header and hostname")); - } - } - - let endpoint = endpoint_sni(hostname, &tls.common_names)?.context("malformed endpoint")?; + let endpoint = + endpoint_sni(hostname, &tls.common_names)?.ok_or(ConnInfoError::MalformedEndpoint)?; ctx.set_endpoint_id(endpoint.clone()); let pairs = connection_url.query_pairs(); @@ -173,36 +176,27 @@ fn get_conn_info( }) } -fn check_matches(sni_hostname: &str, hostname: &str) -> Result { - if sni_hostname == hostname { - return Ok(true); - } - let (sni_hostname_first, sni_hostname_rest) = sni_hostname - .split_once('.') - .ok_or_else(|| anyhow::anyhow!("Unexpected sni format."))?; - let (_, hostname_rest) = hostname - .split_once('.') - .ok_or_else(|| anyhow::anyhow!("Unexpected hostname format."))?; - Ok(sni_hostname_rest == hostname_rest && sni_hostname_first == SERVERLESS_DRIVER_SNI) -} - // TODO: return different http error codes pub async fn handle( config: &'static ProxyConfig, - ctx: &mut RequestMonitoring, + mut ctx: RequestMonitoring, request: Request, - sni_hostname: Option, backend: Arc, ) -> Result, ApiError> { let result = tokio::time::timeout( config.http_config.request_timeout, - handle_inner(config, ctx, request, sni_hostname, backend), + handle_inner(config, &mut ctx, request, backend), ) .await; let mut response = match result { Ok(r) => match r { - Ok(r) => r, + Ok(r) => { + ctx.set_success(); + r + } Err(e) => { + // TODO: ctx.set_error_kind(e.get_error_type()); + let mut message = format!("{:?}", e); let db_error = e .downcast_ref::() @@ -278,7 +272,9 @@ pub async fn handle( )? } }, - Err(_) => { + Err(e) => { + ctx.set_error_kind(e.get_error_kind()); + let message = format!( "HTTP-Connection timed out, execution time exeeded {} seconds", config.http_config.request_timeout.as_secs() @@ -290,6 +286,7 @@ pub async fn handle( )? } }; + response.headers_mut().insert( "Access-Control-Allow-Origin", hyper::http::HeaderValue::from_static("*"), @@ -302,7 +299,6 @@ async fn handle_inner( config: &'static ProxyConfig, ctx: &mut RequestMonitoring, request: Request, - sni_hostname: Option, backend: Arc, ) -> anyhow::Result> { let _request_gauge = NUM_CONNECTION_REQUESTS_GAUGE @@ -318,12 +314,7 @@ async fn handle_inner( // let headers = request.headers(); // TLS config should be there. - let conn_info = get_conn_info( - ctx, - headers, - sni_hostname, - config.tls_config.as_ref().unwrap(), - )?; + let conn_info = get_conn_info(ctx, headers, config.tls_config.as_ref().unwrap())?; info!( user = conn_info.user_info.user.as_str(), project = conn_info.user_info.endpoint.as_str(), @@ -487,8 +478,6 @@ async fn handle_inner( } }; - ctx.set_success(); - ctx.log(); let metrics = client.metrics(); // how could this possibly fail diff --git a/proxy/src/serverless/websocket.rs b/proxy/src/serverless/websocket.rs index f68b35010a..062dd440b2 100644 --- a/proxy/src/serverless/websocket.rs +++ b/proxy/src/serverless/websocket.rs @@ -2,7 +2,7 @@ use crate::{ cancellation::CancelMap, config::ProxyConfig, context::RequestMonitoring, - error::io_error, + error::{io_error, ReportableError}, proxy::{handle_client, ClientMode}, rate_limiter::EndpointRateLimiter, }; @@ -131,23 +131,41 @@ impl AsyncBufRead for WebSocketRw { pub async fn serve_websocket( config: &'static ProxyConfig, - ctx: &mut RequestMonitoring, + mut ctx: RequestMonitoring, websocket: HyperWebsocket, cancel_map: Arc, hostname: Option, endpoint_rate_limiter: Arc, ) -> anyhow::Result<()> { let websocket = websocket.await?; - handle_client( + let res = handle_client( config, - ctx, + &mut ctx, cancel_map, WebSocketRw::new(websocket), ClientMode::Websockets { hostname }, endpoint_rate_limiter, ) - .await?; - Ok(()) + .await; + + match res { + Err(e) => { + // todo: log and push to ctx the error kind + ctx.set_error_kind(e.get_error_kind()); + ctx.log(); + Err(e.into()) + } + Ok(None) => { + ctx.set_success(); + ctx.log(); + Ok(()) + } + Ok(Some(p)) => { + ctx.set_success(); + ctx.log(); + p.proxy_pass().await + } + } } #[cfg(test)] diff --git a/proxy/src/stream.rs b/proxy/src/stream.rs index f48b3fe39f..0d639d2c07 100644 --- a/proxy/src/stream.rs +++ b/proxy/src/stream.rs @@ -1,6 +1,5 @@ use crate::config::TlsServerEndPoint; -use crate::error::UserFacingError; -use anyhow::bail; +use crate::error::{ErrorKind, ReportableError, UserFacingError}; use bytes::BytesMut; use pq_proto::framed::{ConnectionError, Framed}; @@ -73,6 +72,30 @@ impl PqStream { } } +#[derive(Debug)] +pub struct ReportedError { + source: anyhow::Error, + error_kind: ErrorKind, +} + +impl std::fmt::Display for ReportedError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.source.fmt(f) + } +} + +impl std::error::Error for ReportedError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.source.source() + } +} + +impl ReportableError for ReportedError { + fn get_error_kind(&self) -> ErrorKind { + self.error_kind + } +} + impl PqStream { /// Write the message into an internal buffer, but don't flush the underlying stream. pub fn write_message_noflush(&mut self, message: &BeMessage<'_>) -> io::Result<&mut Self> { @@ -98,24 +121,52 @@ impl PqStream { /// Write the error message using [`Self::write_message`], then re-throw it. /// Allowing string literals is safe under the assumption they might not contain any runtime info. /// This method exists due to `&str` not implementing `Into`. - pub async fn throw_error_str(&mut self, error: &'static str) -> anyhow::Result { - tracing::info!("forwarding error to user: {error}"); - self.write_message(&BeMessage::ErrorResponse(error, None)) - .await?; - bail!(error) + pub async fn throw_error_str( + &mut self, + msg: &'static str, + error_kind: ErrorKind, + ) -> Result { + tracing::info!( + kind = error_kind.to_metric_label(), + msg, + "forwarding error to user" + ); + + // already error case, ignore client IO error + let _: Result<_, std::io::Error> = self + .write_message(&BeMessage::ErrorResponse(msg, None)) + .await; + + Err(ReportedError { + source: anyhow::anyhow!(msg), + error_kind, + }) } /// Write the error message using [`Self::write_message`], then re-throw it. /// Trait [`UserFacingError`] acts as an allowlist for error types. - pub async fn throw_error(&mut self, error: E) -> anyhow::Result + pub async fn throw_error(&mut self, error: E) -> Result where E: UserFacingError + Into, { + let error_kind = error.get_error_kind(); let msg = error.to_string_client(); - tracing::info!("forwarding error to user: {msg}"); - self.write_message(&BeMessage::ErrorResponse(&msg, None)) - .await?; - bail!(error) + tracing::info!( + kind=error_kind.to_metric_label(), + error=%error, + msg, + "forwarding error to user" + ); + + // already error case, ignore client IO error + let _: Result<_, std::io::Error> = self + .write_message(&BeMessage::ErrorResponse(&msg, None)) + .await; + + Err(ReportedError { + source: anyhow::anyhow!(error), + error_kind, + }) } } From 1bb9abebf2cc380fa5ef0b876280afd2d120c257 Mon Sep 17 00:00:00 2001 From: Arseny Sher Date: Fri, 9 Feb 2024 16:41:43 +0300 Subject: [PATCH 61/72] Remove WAL segments from s3 in batches. Do list-delete operations in batches instead of doing full list first, to ensure deletion makes progress even if there are a lot of files to remove. To this end, add max_keys limit to remote storage list_files. --- libs/remote_storage/src/azure_blob.rs | 16 +++++++- libs/remote_storage/src/lib.rs | 38 +++++++++++++------ libs/remote_storage/src/local_fs.rs | 13 +++++-- libs/remote_storage/src/s3_bucket.rs | 21 +++++++++- libs/remote_storage/src/simulate_failures.rs | 7 +++- libs/remote_storage/tests/common/tests.rs | 15 ++++++-- libs/remote_storage/tests/test_real_s3.rs | 2 +- .../src/tenant/remote_timeline_client.rs | 2 +- .../tenant/remote_timeline_client/download.rs | 4 +- safekeeper/src/wal_backup.rs | 29 ++++++++++++-- 10 files changed, 119 insertions(+), 28 deletions(-) diff --git a/libs/remote_storage/src/azure_blob.rs b/libs/remote_storage/src/azure_blob.rs index c6d5224706..df6d45dde1 100644 --- a/libs/remote_storage/src/azure_blob.rs +++ b/libs/remote_storage/src/azure_blob.rs @@ -191,6 +191,7 @@ impl RemoteStorage for AzureBlobStorage { &self, prefix: Option<&RemotePath>, mode: ListingMode, + max_keys: Option, ) -> anyhow::Result { // get the passed prefix or if it is not set use prefix_in_bucket value let list_prefix = prefix @@ -223,6 +224,8 @@ impl RemoteStorage for AzureBlobStorage { let mut response = builder.into_stream(); let mut res = Listing::default(); + // NonZeroU32 doesn't support subtraction apparently + let mut max_keys = max_keys.map(|mk| mk.get()); while let Some(l) = response.next().await { let entry = l.map_err(to_download_error)?; let prefix_iter = entry @@ -235,7 +238,18 @@ impl RemoteStorage for AzureBlobStorage { .blobs .blobs() .map(|k| self.name_to_relative_path(&k.name)); - res.keys.extend(blob_iter); + + for key in blob_iter { + res.keys.push(key); + if let Some(mut mk) = max_keys { + assert!(mk > 0); + mk -= 1; + if mk == 0 { + return Ok(res); // limit reached + } + max_keys = Some(mk); + } + } } Ok(res) } diff --git a/libs/remote_storage/src/lib.rs b/libs/remote_storage/src/lib.rs index b6648931ac..5a0b74e406 100644 --- a/libs/remote_storage/src/lib.rs +++ b/libs/remote_storage/src/lib.rs @@ -16,7 +16,12 @@ mod simulate_failures; mod support; use std::{ - collections::HashMap, fmt::Debug, num::NonZeroUsize, pin::Pin, sync::Arc, time::SystemTime, + collections::HashMap, + fmt::Debug, + num::{NonZeroU32, NonZeroUsize}, + pin::Pin, + sync::Arc, + time::SystemTime, }; use anyhow::{bail, Context}; @@ -155,7 +160,7 @@ pub trait RemoteStorage: Send + Sync + 'static { prefix: Option<&RemotePath>, ) -> Result, DownloadError> { let result = self - .list(prefix, ListingMode::WithDelimiter) + .list(prefix, ListingMode::WithDelimiter, None) .await? .prefixes; Ok(result) @@ -171,11 +176,17 @@ pub trait RemoteStorage: Send + Sync + 'static { /// whereas, /// list_prefixes("foo/bar/") = ["cat", "dog"] /// See `test_real_s3.rs` for more details. + /// + /// max_keys limits max number of keys returned; None means unlimited. async fn list_files( &self, prefix: Option<&RemotePath>, + max_keys: Option, ) -> Result, DownloadError> { - let result = self.list(prefix, ListingMode::NoDelimiter).await?.keys; + let result = self + .list(prefix, ListingMode::NoDelimiter, max_keys) + .await? + .keys; Ok(result) } @@ -183,6 +194,7 @@ pub trait RemoteStorage: Send + Sync + 'static { &self, prefix: Option<&RemotePath>, _mode: ListingMode, + max_keys: Option, ) -> Result; /// Streams the local file contents into remote into the remote storage entry. @@ -341,27 +353,31 @@ impl GenericRemoteStorage> { &self, prefix: Option<&RemotePath>, mode: ListingMode, + max_keys: Option, ) -> anyhow::Result { match self { - Self::LocalFs(s) => s.list(prefix, mode).await, - Self::AwsS3(s) => s.list(prefix, mode).await, - Self::AzureBlob(s) => s.list(prefix, mode).await, - Self::Unreliable(s) => s.list(prefix, mode).await, + Self::LocalFs(s) => s.list(prefix, mode, max_keys).await, + Self::AwsS3(s) => s.list(prefix, mode, max_keys).await, + Self::AzureBlob(s) => s.list(prefix, mode, max_keys).await, + Self::Unreliable(s) => s.list(prefix, mode, max_keys).await, } } // A function for listing all the files in a "directory" // Example: // list_files("foo/bar") = ["foo/bar/a.txt", "foo/bar/b.txt"] + // + // max_keys limits max number of keys returned; None means unlimited. pub async fn list_files( &self, folder: Option<&RemotePath>, + max_keys: Option, ) -> Result, DownloadError> { match self { - Self::LocalFs(s) => s.list_files(folder).await, - Self::AwsS3(s) => s.list_files(folder).await, - Self::AzureBlob(s) => s.list_files(folder).await, - Self::Unreliable(s) => s.list_files(folder).await, + Self::LocalFs(s) => s.list_files(folder, max_keys).await, + Self::AwsS3(s) => s.list_files(folder, max_keys).await, + Self::AzureBlob(s) => s.list_files(folder, max_keys).await, + Self::Unreliable(s) => s.list_files(folder, max_keys).await, } } diff --git a/libs/remote_storage/src/local_fs.rs b/libs/remote_storage/src/local_fs.rs index 3ebea76181..f53ba9db07 100644 --- a/libs/remote_storage/src/local_fs.rs +++ b/libs/remote_storage/src/local_fs.rs @@ -4,7 +4,9 @@ //! This storage used in tests, but can also be used in cases when a certain persistent //! volume is mounted to the local FS. -use std::{borrow::Cow, future::Future, io::ErrorKind, pin::Pin, time::SystemTime}; +use std::{ + borrow::Cow, future::Future, io::ErrorKind, num::NonZeroU32, pin::Pin, time::SystemTime, +}; use anyhow::{bail, ensure, Context}; use bytes::Bytes; @@ -162,6 +164,7 @@ impl RemoteStorage for LocalFs { &self, prefix: Option<&RemotePath>, mode: ListingMode, + max_keys: Option, ) -> Result { let mut result = Listing::default(); @@ -178,6 +181,9 @@ impl RemoteStorage for LocalFs { !path.is_dir() }) .collect(); + if let Some(max_keys) = max_keys { + result.keys.truncate(max_keys.get() as usize); + } return Ok(result); } @@ -790,12 +796,12 @@ mod fs_tests { let child = upload_dummy_file(&storage, "grandparent/parent/child", None).await?; let uncle = upload_dummy_file(&storage, "grandparent/uncle", None).await?; - let listing = storage.list(None, ListingMode::NoDelimiter).await?; + let listing = storage.list(None, ListingMode::NoDelimiter, None).await?; assert!(listing.prefixes.is_empty()); assert_eq!(listing.keys, [uncle.clone(), child.clone()].to_vec()); // Delimiter: should only go one deep - let listing = storage.list(None, ListingMode::WithDelimiter).await?; + let listing = storage.list(None, ListingMode::WithDelimiter, None).await?; assert_eq!( listing.prefixes, @@ -808,6 +814,7 @@ mod fs_tests { .list( Some(&RemotePath::from_string("timelines/some_timeline/grandparent").unwrap()), ListingMode::WithDelimiter, + None, ) .await?; assert_eq!( diff --git a/libs/remote_storage/src/s3_bucket.rs b/libs/remote_storage/src/s3_bucket.rs index 2b33a6ffd1..dee5750cac 100644 --- a/libs/remote_storage/src/s3_bucket.rs +++ b/libs/remote_storage/src/s3_bucket.rs @@ -7,6 +7,7 @@ use std::{ borrow::Cow, collections::HashMap, + num::NonZeroU32, pin::Pin, sync::Arc, task::{Context, Poll}, @@ -408,8 +409,11 @@ impl RemoteStorage for S3Bucket { &self, prefix: Option<&RemotePath>, mode: ListingMode, + max_keys: Option, ) -> Result { let kind = RequestKind::List; + // s3 sdk wants i32 + let mut max_keys = max_keys.map(|mk| mk.get() as i32); let mut result = Listing::default(); // get the passed prefix or if it is not set use prefix_in_bucket value @@ -433,13 +437,20 @@ impl RemoteStorage for S3Bucket { let _guard = self.permit(kind).await; let started_at = start_measuring_requests(kind); + // min of two Options, returning Some if one is value and another is + // None (None is smaller than anything, so plain min doesn't work). + let request_max_keys = self + .max_keys_per_list_response + .into_iter() + .chain(max_keys.into_iter()) + .min(); let mut request = self .client .list_objects_v2() .bucket(self.bucket_name.clone()) .set_prefix(list_prefix.clone()) .set_continuation_token(continuation_token) - .set_max_keys(self.max_keys_per_list_response); + .set_max_keys(request_max_keys); if let ListingMode::WithDelimiter = mode { request = request.delimiter(REMOTE_STORAGE_PREFIX_SEPARATOR.to_string()); @@ -469,6 +480,14 @@ impl RemoteStorage for S3Bucket { let object_path = object.key().expect("response does not contain a key"); let remote_path = self.s3_object_to_relative_path(object_path); result.keys.push(remote_path); + if let Some(mut mk) = max_keys { + assert!(mk > 0); + mk -= 1; + if mk == 0 { + return Ok(result); // limit reached + } + max_keys = Some(mk); + } } result.prefixes.extend( diff --git a/libs/remote_storage/src/simulate_failures.rs b/libs/remote_storage/src/simulate_failures.rs index 14bdb5ed4d..3dfa16b64e 100644 --- a/libs/remote_storage/src/simulate_failures.rs +++ b/libs/remote_storage/src/simulate_failures.rs @@ -4,6 +4,7 @@ use bytes::Bytes; use futures::stream::Stream; use std::collections::HashMap; +use std::num::NonZeroU32; use std::sync::Mutex; use std::time::SystemTime; use std::{collections::hash_map::Entry, sync::Arc}; @@ -113,20 +114,22 @@ impl RemoteStorage for UnreliableWrapper { async fn list_files( &self, folder: Option<&RemotePath>, + max_keys: Option, ) -> Result, DownloadError> { self.attempt(RemoteOp::ListPrefixes(folder.cloned())) .map_err(DownloadError::Other)?; - self.inner.list_files(folder).await + self.inner.list_files(folder, max_keys).await } async fn list( &self, prefix: Option<&RemotePath>, mode: ListingMode, + max_keys: Option, ) -> Result { self.attempt(RemoteOp::ListPrefixes(prefix.cloned())) .map_err(DownloadError::Other)?; - self.inner.list(prefix, mode).await + self.inner.list(prefix, mode, max_keys).await } async fn upload( diff --git a/libs/remote_storage/tests/common/tests.rs b/libs/remote_storage/tests/common/tests.rs index abccc24c97..6d062f3898 100644 --- a/libs/remote_storage/tests/common/tests.rs +++ b/libs/remote_storage/tests/common/tests.rs @@ -1,8 +1,8 @@ use anyhow::Context; use camino::Utf8Path; use remote_storage::RemotePath; -use std::collections::HashSet; use std::sync::Arc; +use std::{collections::HashSet, num::NonZeroU32}; use test_context::test_context; use tracing::debug; @@ -103,7 +103,7 @@ async fn list_files_works(ctx: &mut MaybeEnabledStorageWithSimpleTestBlobs) -> a let base_prefix = RemotePath::new(Utf8Path::new("folder1")).context("common_prefix construction")?; let root_files = test_client - .list_files(None) + .list_files(None, None) .await .context("client list root files failure")? .into_iter() @@ -113,8 +113,17 @@ async fn list_files_works(ctx: &mut MaybeEnabledStorageWithSimpleTestBlobs) -> a ctx.remote_blobs.clone(), "remote storage list_files on root mismatches with the uploads." ); + + // Test that max_keys limit works. In total there are about 21 files (see + // upload_simple_remote_data call in test_real_s3.rs). + let limited_root_files = test_client + .list_files(None, Some(NonZeroU32::new(2).unwrap())) + .await + .context("client list root files failure")?; + assert_eq!(limited_root_files.len(), 2); + let nested_remote_files = test_client - .list_files(Some(&base_prefix)) + .list_files(Some(&base_prefix), None) .await .context("client list nested files failure")? .into_iter() diff --git a/libs/remote_storage/tests/test_real_s3.rs b/libs/remote_storage/tests/test_real_s3.rs index fc52dabc36..3dc8347c83 100644 --- a/libs/remote_storage/tests/test_real_s3.rs +++ b/libs/remote_storage/tests/test_real_s3.rs @@ -70,7 +70,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow: } async fn list_files(client: &Arc) -> anyhow::Result> { - Ok(retry(|| client.list_files(None)) + Ok(retry(|| client.list_files(None, None)) .await .context("list root files failure")? .into_iter() diff --git a/pageserver/src/tenant/remote_timeline_client.rs b/pageserver/src/tenant/remote_timeline_client.rs index 0c7dd68c3f..e17dea01a8 100644 --- a/pageserver/src/tenant/remote_timeline_client.rs +++ b/pageserver/src/tenant/remote_timeline_client.rs @@ -1151,7 +1151,7 @@ impl RemoteTimelineClient { let remaining = download_retry( || async { self.storage_impl - .list_files(Some(&timeline_storage_path)) + .list_files(Some(&timeline_storage_path), None) .await }, "list remaining files", diff --git a/pageserver/src/tenant/remote_timeline_client/download.rs b/pageserver/src/tenant/remote_timeline_client/download.rs index 33287fc8f4..e755cd08f3 100644 --- a/pageserver/src/tenant/remote_timeline_client/download.rs +++ b/pageserver/src/tenant/remote_timeline_client/download.rs @@ -220,7 +220,7 @@ pub async fn list_remote_timelines( || { download_cancellable( &cancel, - storage.list(Some(&remote_path), ListingMode::WithDelimiter), + storage.list(Some(&remote_path), ListingMode::WithDelimiter, None), ) }, &format!("list timelines for {tenant_shard_id}"), @@ -373,7 +373,7 @@ pub(super) async fn download_index_part( let index_prefix = remote_index_path(tenant_shard_id, timeline_id, Generation::none()); let indices = download_retry( - || async { storage.list_files(Some(&index_prefix)).await }, + || async { storage.list_files(Some(&index_prefix), None).await }, "list index_part files", cancel, ) diff --git a/safekeeper/src/wal_backup.rs b/safekeeper/src/wal_backup.rs index df99244770..dbdc742d26 100644 --- a/safekeeper/src/wal_backup.rs +++ b/safekeeper/src/wal_backup.rs @@ -10,6 +10,7 @@ use utils::id::NodeId; use std::cmp::min; use std::collections::{HashMap, HashSet}; +use std::num::NonZeroU32; use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -546,6 +547,10 @@ pub async fn delete_timeline(ttid: &TenantTimelineId) -> Result<()> { let ttid_path = Utf8Path::new(&ttid.tenant_id.to_string()).join(ttid.timeline_id.to_string()); let remote_path = RemotePath::new(&ttid_path)?; + // see DEFAULT_MAX_KEYS_PER_LIST_RESPONSE + // const Option unwrap is not stable, otherwise it would be const. + let batch_size: NonZeroU32 = NonZeroU32::new(1000).unwrap(); + // A backoff::retry is used here for two reasons: // - To provide a backoff rather than busy-polling the API on errors // - To absorb transient 429/503 conditions without hitting our error @@ -557,8 +562,26 @@ pub async fn delete_timeline(ttid: &TenantTimelineId) -> Result<()> { let token = CancellationToken::new(); // not really used backoff::retry( || async { - let files = storage.list_files(Some(&remote_path)).await?; - storage.delete_objects(&files).await + // Do list-delete in batch_size batches to make progress even if there a lot of files. + // Alternatively we could make list_files return iterator, but it is more complicated and + // I'm not sure deleting while iterating is expected in s3. + loop { + let files = storage + .list_files(Some(&remote_path), Some(batch_size)) + .await?; + if files.is_empty() { + return Ok(()); // done + } + // (at least) s3 results are sorted, so can log min/max: + // "List results are always returned in UTF-8 binary order." + info!( + "deleting batch of {} WAL segments [{}-{}]", + files.len(), + files.first().unwrap().object_name().unwrap_or(""), + files.last().unwrap().object_name().unwrap_or("") + ); + storage.delete_objects(&files).await?; + } }, |_| false, 3, @@ -594,7 +617,7 @@ pub async fn copy_s3_segments( let remote_path = RemotePath::new(&relative_dst_path)?; - let files = storage.list_files(Some(&remote_path)).await?; + let files = storage.list_files(Some(&remote_path), None).await?; let uploaded_segments = &files .iter() .filter_map(|file| file.object_name().map(ToOwned::to_owned)) From ca818c8bd76d815f0d41eb61fdb8fb9b826ffe54 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Fri, 9 Feb 2024 20:09:37 +0100 Subject: [PATCH 62/72] fix(test_ondemand_download_timetravel): occasionally fails with slightly higher physical size (#6687) --- test_runner/regress/test_ondemand_download.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test_runner/regress/test_ondemand_download.py b/test_runner/regress/test_ondemand_download.py index af2d7aae88..3a197875dd 100644 --- a/test_runner/regress/test_ondemand_download.py +++ b/test_runner/regress/test_ondemand_download.py @@ -197,6 +197,14 @@ def test_ondemand_download_timetravel(neon_env_builder: NeonEnvBuilder): ##### Stop the first pageserver instance, erase all its data env.endpoints.stop_all() + # Stop safekeepers and take another checkpoint. The endpoints might + # have written a few more bytes during shutdown. + for sk in env.safekeepers: + sk.stop() + + client.timeline_checkpoint(tenant_id, timeline_id) + current_lsn = Lsn(client.timeline_detail(tenant_id, timeline_id)["last_record_lsn"]) + # wait until pageserver has successfully uploaded all the data to remote storage wait_for_upload(client, tenant_id, timeline_id, current_lsn) From cbd3a32d4d4275338c851dd158e0cb950d64ee91 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Fri, 9 Feb 2024 19:22:23 +0000 Subject: [PATCH 63/72] proxy: decode username and password (#6700) ## Problem usernames and passwords can be URL 'percent' encoded in the connection string URL provided by serverless driver. ## Summary of changes Decode the parameters when getting conn info --- Cargo.lock | 2 ++ Cargo.toml | 1 + proxy/Cargo.toml | 4 +++- proxy/src/serverless/backend.rs | 2 +- proxy/src/serverless/conn_pool.rs | 7 ++++--- proxy/src/serverless/sql_over_http.rs | 10 ++++++++-- test_runner/fixtures/neon_fixtures.py | 6 +++--- test_runner/regress/test_proxy.py | 12 ++++++++++++ 8 files changed, 34 insertions(+), 10 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a2939e6c75..83afdaf66f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4125,6 +4125,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "smallvec", "smol_str", "socket2 0.5.5", "sync_wrapper", @@ -4143,6 +4144,7 @@ dependencies = [ "tracing-subscriber", "tracing-utils", "url", + "urlencoding", "utils", "uuid", "walkdir", diff --git a/Cargo.toml b/Cargo.toml index 6a2c3fa563..ebc3dfa7b1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -171,6 +171,7 @@ tracing-opentelemetry = "0.20.0" tracing-subscriber = { version = "0.3", default_features = false, features = ["smallvec", "fmt", "tracing-log", "std", "env-filter", "json"] } twox-hash = { version = "1.6.3", default-features = false } url = "2.2" +urlencoding = "2.1" uuid = { version = "1.6.1", features = ["v4", "v7", "serde"] } walkdir = "2.3.2" webpki-roots = "0.25" diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index 83cab381b3..0777d361d2 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -60,6 +60,8 @@ scopeguard.workspace = true serde.workspace = true serde_json.workspace = true sha2.workspace = true +smol_str.workspace = true +smallvec.workspace = true socket2.workspace = true sync_wrapper.workspace = true task-local-extensions.workspace = true @@ -76,6 +78,7 @@ tracing-subscriber.workspace = true tracing-utils.workspace = true tracing.workspace = true url.workspace = true +urlencoding.workspace = true utils.workspace = true uuid.workspace = true webpki-roots.workspace = true @@ -84,7 +87,6 @@ native-tls.workspace = true postgres-native-tls.workspace = true postgres-protocol.workspace = true redis.workspace = true -smol_str.workspace = true workspace_hack.workspace = true diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index 03257e9161..8285da68d7 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -48,7 +48,7 @@ impl PoolingBackend { } }; let auth_outcome = - crate::auth::validate_password_and_exchange(conn_info.password.as_bytes(), secret)?; + crate::auth::validate_password_and_exchange(&conn_info.password, secret)?; match auth_outcome { crate::sasl::Outcome::Success(key) => Ok(key), crate::sasl::Outcome::Failure(reason) => { diff --git a/proxy/src/serverless/conn_pool.rs b/proxy/src/serverless/conn_pool.rs index f92793096b..f4e5b145c5 100644 --- a/proxy/src/serverless/conn_pool.rs +++ b/proxy/src/serverless/conn_pool.rs @@ -3,6 +3,7 @@ use futures::{future::poll_fn, Future}; use metrics::IntCounterPairGuard; use parking_lot::RwLock; use rand::Rng; +use smallvec::SmallVec; use smol_str::SmolStr; use std::{collections::HashMap, pin::pin, sync::Arc, sync::Weak, time::Duration}; use std::{ @@ -36,7 +37,7 @@ pub const APP_NAME: SmolStr = SmolStr::new_inline("/sql_over_http"); pub struct ConnInfo { pub user_info: ComputeUserInfo, pub dbname: DbName, - pub password: SmolStr, + pub password: SmallVec<[u8; 16]>, } impl ConnInfo { @@ -731,7 +732,7 @@ mod tests { options: Default::default(), }, dbname: "dbname".into(), - password: "password".into(), + password: "password".as_bytes().into(), }; let ep_pool = Arc::downgrade(&pool.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key())); @@ -788,7 +789,7 @@ mod tests { options: Default::default(), }, dbname: "dbname".into(), - password: "password".into(), + password: "password".as_bytes().into(), }; let ep_pool = Arc::downgrade(&pool.get_or_create_endpoint_pool(&conn_info.endpoint_cache_key())); diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 401022347e..54424360c4 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -100,6 +100,8 @@ pub enum ConnInfoError { InvalidDbName, #[error("missing username")] MissingUsername, + #[error("invalid username: {0}")] + InvalidUsername(#[from] std::string::FromUtf8Error), #[error("missing password")] MissingPassword, #[error("missing hostname")] @@ -134,7 +136,7 @@ fn get_conn_info( let dbname = url_path.next().ok_or(ConnInfoError::InvalidDbName)?; - let username = RoleName::from(connection_url.username()); + let username = RoleName::from(urlencoding::decode(connection_url.username())?); if username.is_empty() { return Err(ConnInfoError::MissingUsername); } @@ -143,6 +145,7 @@ fn get_conn_info( let password = connection_url .password() .ok_or(ConnInfoError::MissingPassword)?; + let password = urlencoding::decode_binary(password.as_bytes()); let hostname = connection_url .host_str() @@ -172,7 +175,10 @@ fn get_conn_info( Ok(ConnInfo { user_info, dbname: dbname.into(), - password: password.into(), + password: match password { + std::borrow::Cow::Borrowed(b) => b.into(), + std::borrow::Cow::Owned(b) => b.into(), + }, }) } diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 9996853525..231eebff52 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -23,7 +23,7 @@ from itertools import chain, product from pathlib import Path from types import TracebackType from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union, cast -from urllib.parse import urlparse +from urllib.parse import quote, urlparse import asyncpg import backoff @@ -2822,8 +2822,8 @@ class NeonProxy(PgProtocol): def http_query(self, query, args, **kwargs): # TODO maybe use default values if not provided - user = kwargs["user"] - password = kwargs["password"] + user = quote(kwargs["user"]) + password = quote(kwargs["password"]) expected_code = kwargs.get("expected_code") connstr = f"postgresql://{user}:{password}@{self.domain}:{self.proxy_port}/postgres" diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py index 49a0450f0c..884643cef0 100644 --- a/test_runner/regress/test_proxy.py +++ b/test_runner/regress/test_proxy.py @@ -462,6 +462,18 @@ def test_sql_over_http_pool(static_proxy: NeonProxy): assert "password authentication failed for user" in res["message"] +def test_sql_over_http_urlencoding(static_proxy: NeonProxy): + static_proxy.safe_psql("create user \"http+auth$$\" with password '%+$^&*@!' superuser") + + static_proxy.http_query( + "select 1", + [], + user="http+auth$$", + password="%+$^&*@!", + expected_code=200, + ) + + # Beginning a transaction should not impact the next query, # which might come from a completely different client. def test_http_pool_begin(static_proxy: NeonProxy): From 1a4dd58b70ad1bf82c4daae520f4550612f91120 Mon Sep 17 00:00:00 2001 From: Sasha Krassovsky Date: Fri, 9 Feb 2024 11:22:53 -0900 Subject: [PATCH 64/72] Grant pg_monitor to neon_superuser (#6691) ## Problem The people want pg_monitor https://github.com/neondatabase/neon/issues/6682 ## Summary of changes Gives the people pg_monitor --- compute_tools/src/spec.rs | 1 + test_runner/regress/test_migrations.py | 4 ++-- test_runner/regress/test_neon_superuser.py | 18 ++++++++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs index 3df5f10e23..9c731f257c 100644 --- a/compute_tools/src/spec.rs +++ b/compute_tools/src/spec.rs @@ -776,6 +776,7 @@ BEGIN END IF; END $$;"#, + "GRANT pg_monitor TO neon_superuser WITH ADMIN OPTION", ]; let mut query = "CREATE SCHEMA IF NOT EXISTS neon_migration"; diff --git a/test_runner/regress/test_migrations.py b/test_runner/regress/test_migrations.py index 8954810451..7cc3024ec6 100644 --- a/test_runner/regress/test_migrations.py +++ b/test_runner/regress/test_migrations.py @@ -15,7 +15,7 @@ def test_migrations(neon_simple_env: NeonEnv): endpoint.wait_for_migrations() - num_migrations = 3 + num_migrations = 4 with endpoint.cursor() as cur: cur.execute("SELECT id FROM neon_migration.migration_id") @@ -24,7 +24,7 @@ def test_migrations(neon_simple_env: NeonEnv): with open(log_path, "r") as log_file: logs = log_file.read() - assert "INFO handle_migrations: Ran 3 migrations" in logs + assert f"INFO handle_migrations: Ran {num_migrations} migrations" in logs endpoint.stop() endpoint.start() diff --git a/test_runner/regress/test_neon_superuser.py b/test_runner/regress/test_neon_superuser.py index 34f1e64b34..ca8ada4ddb 100644 --- a/test_runner/regress/test_neon_superuser.py +++ b/test_runner/regress/test_neon_superuser.py @@ -76,3 +76,21 @@ def test_neon_superuser(neon_simple_env: NeonEnv, pg_version: PgVersion): assert [r[0] for r in res] == [10, 20, 30, 40] wait_until(10, 0.5, check_that_changes_propagated) + + # Test that pg_monitor is working for neon_superuser role + cur.execute("SELECT query from pg_stat_activity LIMIT 1") + assert cur.fetchall()[0][0] != "" + # Test that pg_monitor is not working for non neon_superuser role without grant + cur.execute("CREATE ROLE not_a_superuser LOGIN PASSWORD 'Password42!'") + cur.execute("GRANT not_a_superuser TO neon_superuser WITH ADMIN OPTION") + cur.execute("SET ROLE not_a_superuser") + cur.execute("SELECT query from pg_stat_activity LIMIT 1") + assert cur.fetchall()[0][0] == "" + cur.execute("RESET ROLE") + # Test that pg_monitor is working for non neon_superuser role with grant + cur.execute("GRANT pg_monitor TO not_a_superuser") + cur.execute("SET ROLE not_a_superuser") + cur.execute("SELECT query from pg_stat_activity LIMIT 1") + assert cur.fetchall()[0][0] != "" + cur.execute("RESET ROLE") + cur.execute("DROP ROLE not_a_superuser") From 5779c7908abaadb0c96a5087423e2082101924b9 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Fri, 9 Feb 2024 23:22:40 +0100 Subject: [PATCH 65/72] revert two recent `heavier_once_cell` changes (#6704) This PR reverts - https://github.com/neondatabase/neon/pull/6589 - https://github.com/neondatabase/neon/pull/6652 because there's a performance regression that's particularly visible at high layer counts. Most likely it's because the switch to RwLock inflates the ``` inner: heavier_once_cell::OnceCell, ``` size from 48 to 88 bytes, which, by itself is almost a doubling of the cache footprint, and probably the fact that it's now larger than a cache line also doesn't help. See this chat on the Neon discord for more context: https://discord.com/channels/1176467419317940276/1204714372295958548/1205541184634617906 I'm reverting 6652 as well because it might also have perf implications, and we're getting close to the next release. We should re-do its changes after the next release, though. cc @koivunej cc @ivaxer --- libs/utils/src/sync/heavier_once_cell.rs | 322 ++++--------------- pageserver/src/tenant/storage_layer/layer.rs | 24 +- pageserver/src/tenant/timeline.rs | 2 +- 3 files changed, 81 insertions(+), 267 deletions(-) diff --git a/libs/utils/src/sync/heavier_once_cell.rs b/libs/utils/src/sync/heavier_once_cell.rs index 81625b907e..0ccaf4e716 100644 --- a/libs/utils/src/sync/heavier_once_cell.rs +++ b/libs/utils/src/sync/heavier_once_cell.rs @@ -1,6 +1,6 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, - Arc, + Arc, Mutex, MutexGuard, }; use tokio::sync::Semaphore; @@ -12,7 +12,7 @@ use tokio::sync::Semaphore; /// /// [`OwnedSemaphorePermit`]: tokio::sync::OwnedSemaphorePermit pub struct OnceCell { - inner: tokio::sync::RwLock>, + inner: Mutex>, initializers: AtomicUsize, } @@ -50,7 +50,7 @@ impl OnceCell { let sem = Semaphore::new(1); sem.close(); Self { - inner: tokio::sync::RwLock::new(Inner { + inner: Mutex::new(Inner { init_semaphore: Arc::new(sem), value: Some(value), }), @@ -61,113 +61,56 @@ impl OnceCell { /// Returns a guard to an existing initialized value, or uniquely initializes the value before /// returning the guard. /// - /// Initializing might wait on any existing [`GuardMut::take_and_deinit`] deinitialization. + /// Initializing might wait on any existing [`Guard::take_and_deinit`] deinitialization. /// /// Initialization is panic-safe and cancellation-safe. - pub async fn get_mut_or_init(&self, factory: F) -> Result, E> + pub async fn get_or_init(&self, factory: F) -> Result, E> where F: FnOnce(InitPermit) -> Fut, Fut: std::future::Future>, { - loop { - let sem = { - let guard = self.inner.write().await; - if guard.value.is_some() { - return Ok(GuardMut(guard)); - } - guard.init_semaphore.clone() - }; - - { - let permit = { - // increment the count for the duration of queued - let _guard = CountWaitingInitializers::start(self); - sem.acquire().await - }; - - let Ok(permit) = permit else { - let guard = self.inner.write().await; - if !Arc::ptr_eq(&sem, &guard.init_semaphore) { - // there was a take_and_deinit in between - continue; - } - assert!( - guard.value.is_some(), - "semaphore got closed, must be initialized" - ); - return Ok(GuardMut(guard)); - }; - - permit.forget(); + let sem = { + let guard = self.inner.lock().unwrap(); + if guard.value.is_some() { + return Ok(Guard(guard)); } + guard.init_semaphore.clone() + }; - let permit = InitPermit(sem); - let (value, _permit) = factory(permit).await?; + let permit = { + // increment the count for the duration of queued + let _guard = CountWaitingInitializers::start(self); + sem.acquire_owned().await + }; - let guard = self.inner.write().await; + match permit { + Ok(permit) => { + let permit = InitPermit(permit); + let (value, _permit) = factory(permit).await?; - return Ok(Self::set0(value, guard)); + let guard = self.inner.lock().unwrap(); + + Ok(Self::set0(value, guard)) + } + Err(_closed) => { + let guard = self.inner.lock().unwrap(); + assert!( + guard.value.is_some(), + "semaphore got closed, must be initialized" + ); + return Ok(Guard(guard)); + } } } - /// Returns a guard to an existing initialized value, or uniquely initializes the value before - /// returning the guard. - /// - /// Initialization is panic-safe and cancellation-safe. - pub async fn get_or_init(&self, factory: F) -> Result, E> - where - F: FnOnce(InitPermit) -> Fut, - Fut: std::future::Future>, - { - loop { - let sem = { - let guard = self.inner.read().await; - if guard.value.is_some() { - return Ok(GuardRef(guard)); - } - guard.init_semaphore.clone() - }; - - { - let permit = { - // increment the count for the duration of queued - let _guard = CountWaitingInitializers::start(self); - sem.acquire().await - }; - - let Ok(permit) = permit else { - let guard = self.inner.read().await; - if !Arc::ptr_eq(&sem, &guard.init_semaphore) { - // there was a take_and_deinit in between - continue; - } - assert!( - guard.value.is_some(), - "semaphore got closed, must be initialized" - ); - return Ok(GuardRef(guard)); - }; - - permit.forget(); - } - - let permit = InitPermit(sem); - let (value, _permit) = factory(permit).await?; - - let guard = self.inner.write().await; - - return Ok(Self::set0(value, guard).downgrade()); - } - } - - /// Assuming a permit is held after previous call to [`GuardMut::take_and_deinit`], it can be used + /// Assuming a permit is held after previous call to [`Guard::take_and_deinit`], it can be used /// to complete initializing the inner value. /// /// # Panics /// /// If the inner has already been initialized. - pub async fn set(&self, value: T, _permit: InitPermit) -> GuardMut<'_, T> { - let guard = self.inner.write().await; + pub fn set(&self, value: T, _permit: InitPermit) -> Guard<'_, T> { + let guard = self.inner.lock().unwrap(); // cannot assert that this permit is for self.inner.semaphore, but we can assert it cannot // give more permits right now. @@ -179,31 +122,21 @@ impl OnceCell { Self::set0(value, guard) } - fn set0(value: T, mut guard: tokio::sync::RwLockWriteGuard<'_, Inner>) -> GuardMut<'_, T> { + fn set0(value: T, mut guard: std::sync::MutexGuard<'_, Inner>) -> Guard<'_, T> { if guard.value.is_some() { drop(guard); unreachable!("we won permit, must not be initialized"); } guard.value = Some(value); guard.init_semaphore.close(); - GuardMut(guard) + Guard(guard) } /// Returns a guard to an existing initialized value, if any. - pub async fn get_mut(&self) -> Option> { - let guard = self.inner.write().await; + pub fn get(&self) -> Option> { + let guard = self.inner.lock().unwrap(); if guard.value.is_some() { - Some(GuardMut(guard)) - } else { - None - } - } - - /// Returns a guard to an existing initialized value, if any. - pub async fn get(&self) -> Option> { - let guard = self.inner.read().await; - if guard.value.is_some() { - Some(GuardRef(guard)) + Some(Guard(guard)) } else { None } @@ -235,9 +168,9 @@ impl<'a, T> Drop for CountWaitingInitializers<'a, T> { /// Uninteresting guard object to allow short-lived access to inspect or clone the held, /// initialized value. #[derive(Debug)] -pub struct GuardMut<'a, T>(tokio::sync::RwLockWriteGuard<'a, Inner>); +pub struct Guard<'a, T>(MutexGuard<'a, Inner>); -impl std::ops::Deref for GuardMut<'_, T> { +impl std::ops::Deref for Guard<'_, T> { type Target = T; fn deref(&self) -> &Self::Target { @@ -248,7 +181,7 @@ impl std::ops::Deref for GuardMut<'_, T> { } } -impl std::ops::DerefMut for GuardMut<'_, T> { +impl std::ops::DerefMut for Guard<'_, T> { fn deref_mut(&mut self) -> &mut Self::Target { self.0 .value @@ -257,59 +190,34 @@ impl std::ops::DerefMut for GuardMut<'_, T> { } } -impl<'a, T> GuardMut<'a, T> { +impl<'a, T> Guard<'a, T> { /// Take the current value, and a new permit for it's deinitialization. /// /// The permit will be on a semaphore part of the new internal value, and any following /// [`OnceCell::get_or_init`] will wait on it to complete. pub fn take_and_deinit(&mut self) -> (T, InitPermit) { let mut swapped = Inner::default(); - let sem = swapped.init_semaphore.clone(); - sem.try_acquire().expect("we just created this").forget(); + let permit = swapped + .init_semaphore + .clone() + .try_acquire_owned() + .expect("we just created this"); std::mem::swap(&mut *self.0, &mut swapped); swapped .value - .map(|v| (v, InitPermit(sem))) - .expect("guard is not created unless value has been initialized") - } - - pub fn downgrade(self) -> GuardRef<'a, T> { - GuardRef(self.0.downgrade()) - } -} - -#[derive(Debug)] -pub struct GuardRef<'a, T>(tokio::sync::RwLockReadGuard<'a, Inner>); - -impl std::ops::Deref for GuardRef<'_, T> { - type Target = T; - - fn deref(&self) -> &Self::Target { - self.0 - .value - .as_ref() + .map(|v| (v, InitPermit(permit))) .expect("guard is not created unless value has been initialized") } } /// Type held by OnceCell (de)initializing task. -pub struct InitPermit(Arc); - -impl Drop for InitPermit { - fn drop(&mut self) { - debug_assert_eq!(self.0.available_permits(), 0); - self.0.add_permits(1); - } -} +pub struct InitPermit(tokio::sync::OwnedSemaphorePermit); #[cfg(test)] mod tests { - use futures::Future; - use super::*; use std::{ convert::Infallible, - pin::{pin, Pin}, sync::atomic::{AtomicUsize, Ordering}, time::Duration, }; @@ -340,7 +248,7 @@ mod tests { barrier.wait().await; let won = { let g = cell - .get_mut_or_init(|permit| { + .get_or_init(|permit| { counters.factory_got_to_run.fetch_add(1, Ordering::Relaxed); async { counters.future_polled.fetch_add(1, Ordering::Relaxed); @@ -387,11 +295,7 @@ mod tests { let cell = cell.clone(); let deinitialization_started = deinitialization_started.clone(); async move { - let (answer, _permit) = cell - .get_mut() - .await - .expect("initialized to value") - .take_and_deinit(); + let (answer, _permit) = cell.get().expect("initialized to value").take_and_deinit(); assert_eq!(answer, initial); deinitialization_started.wait().await; @@ -402,7 +306,7 @@ mod tests { deinitialization_started.wait().await; let started_at = tokio::time::Instant::now(); - cell.get_mut_or_init(|permit| async { Ok::<_, Infallible>((reinit, permit)) }) + cell.get_or_init(|permit| async { Ok::<_, Infallible>((reinit, permit)) }) .await .unwrap(); @@ -414,21 +318,21 @@ mod tests { jh.await.unwrap(); - assert_eq!(*cell.get_mut().await.unwrap(), reinit); + assert_eq!(*cell.get().unwrap(), reinit); } - #[tokio::test] - async fn reinit_with_deinit_permit() { + #[test] + fn reinit_with_deinit_permit() { let cell = Arc::new(OnceCell::new(42)); - let (mol, permit) = cell.get_mut().await.unwrap().take_and_deinit(); - cell.set(5, permit).await; - assert_eq!(*cell.get_mut().await.unwrap(), 5); + let (mol, permit) = cell.get().unwrap().take_and_deinit(); + cell.set(5, permit); + assert_eq!(*cell.get().unwrap(), 5); - let (five, permit) = cell.get_mut().await.unwrap().take_and_deinit(); + let (five, permit) = cell.get().unwrap().take_and_deinit(); assert_eq!(5, five); - cell.set(mol, permit).await; - assert_eq!(*cell.get_mut().await.unwrap(), 42); + cell.set(mol, permit); + assert_eq!(*cell.get().unwrap(), 42); } #[tokio::test] @@ -436,13 +340,13 @@ mod tests { let cell = OnceCell::default(); for _ in 0..10 { - cell.get_mut_or_init(|_permit| async { Err("whatever error") }) + cell.get_or_init(|_permit| async { Err("whatever error") }) .await .unwrap_err(); } let g = cell - .get_mut_or_init(|permit| async { Ok::<_, Infallible>(("finally success", permit)) }) + .get_or_init(|permit| async { Ok::<_, Infallible>(("finally success", permit)) }) .await .unwrap(); assert_eq!(*g, "finally success"); @@ -454,7 +358,7 @@ mod tests { let barrier = tokio::sync::Barrier::new(2); - let initializer = cell.get_mut_or_init(|permit| async { + let initializer = cell.get_or_init(|permit| async { barrier.wait().await; futures::future::pending::<()>().await; @@ -468,102 +372,12 @@ mod tests { // now initializer is dropped - assert!(cell.get_mut().await.is_none()); + assert!(cell.get().is_none()); let g = cell - .get_mut_or_init(|permit| async { Ok::<_, Infallible>(("now initialized", permit)) }) + .get_or_init(|permit| async { Ok::<_, Infallible>(("now initialized", permit)) }) .await .unwrap(); assert_eq!(*g, "now initialized"); } - - #[tokio::test(start_paused = true)] - async fn reproduce_init_take_deinit_race() { - init_take_deinit_scenario(|cell, factory| { - Box::pin(async { - cell.get_or_init(factory).await.unwrap(); - }) - }) - .await; - } - - #[tokio::test(start_paused = true)] - async fn reproduce_init_take_deinit_race_mut() { - init_take_deinit_scenario(|cell, factory| { - Box::pin(async { - cell.get_mut_or_init(factory).await.unwrap(); - }) - }) - .await; - } - - type BoxedInitFuture = Pin>>>; - type BoxedInitFunction = Box BoxedInitFuture>; - - /// Reproduce an assertion failure with both initialization methods. - /// - /// This has interesting generics to be generic between `get_or_init` and `get_mut_or_init`. - /// Alternative would be a macro_rules! but that is the last resort. - async fn init_take_deinit_scenario(init_way: F) - where - F: for<'a> Fn( - &'a OnceCell<&'static str>, - BoxedInitFunction<&'static str, Infallible>, - ) -> Pin + 'a>>, - { - let cell = OnceCell::default(); - - // acquire the init_semaphore only permit to drive initializing tasks in order to waiting - // on the same semaphore. - let permit = cell - .inner - .read() - .await - .init_semaphore - .clone() - .try_acquire_owned() - .unwrap(); - - let mut t1 = pin!(init_way( - &cell, - Box::new(|permit| Box::pin(async move { Ok(("t1", permit)) })), - )); - - let mut t2 = pin!(init_way( - &cell, - Box::new(|permit| Box::pin(async move { Ok(("t2", permit)) })), - )); - - // drive t2 first to the init_semaphore - tokio::select! { - _ = &mut t2 => unreachable!("it cannot get permit"), - _ = tokio::time::sleep(Duration::from_secs(3600 * 24 * 7 * 365)) => {} - } - - // followed by t1 in the init_semaphore - tokio::select! { - _ = &mut t1 => unreachable!("it cannot get permit"), - _ = tokio::time::sleep(Duration::from_secs(3600 * 24 * 7 * 365)) => {} - } - - // now let t2 proceed and initialize - drop(permit); - t2.await; - - let (s, permit) = { cell.get_mut().await.unwrap().take_and_deinit() }; - assert_eq!("t2", s); - - // now originally t1 would see the semaphore it has as closed. it cannot yet get a permit from - // the new one. - tokio::select! { - _ = &mut t1 => unreachable!("it cannot get permit"), - _ = tokio::time::sleep(Duration::from_secs(3600 * 24 * 7 * 365)) => {} - } - - // only now we get to initialize it - drop(permit); - t1.await; - - assert_eq!("t1", *cell.get().await.unwrap()); - } } diff --git a/pageserver/src/tenant/storage_layer/layer.rs b/pageserver/src/tenant/storage_layer/layer.rs index 52c0f8abdc..dd9de99477 100644 --- a/pageserver/src/tenant/storage_layer/layer.rs +++ b/pageserver/src/tenant/storage_layer/layer.rs @@ -300,8 +300,8 @@ impl Layer { }) } - pub(crate) async fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { - self.0.info(reset).await + pub(crate) fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { + self.0.info(reset) } pub(crate) fn access_stats(&self) -> &LayerAccessStats { @@ -612,10 +612,10 @@ impl LayerInner { let mut rx = self.status.subscribe(); let strong = { - match self.inner.get_mut().await { + match self.inner.get() { Some(mut either) => { self.wanted_evicted.store(true, Ordering::Relaxed); - ResidentOrWantedEvicted::downgrade(&mut either) + either.downgrade() } None => return Err(EvictionError::NotFound), } @@ -641,7 +641,7 @@ impl LayerInner { // use however late (compared to the initial expressing of wanted) as the // "outcome" now LAYER_IMPL_METRICS.inc_broadcast_lagged(); - match self.inner.get_mut().await { + match self.inner.get() { Some(_) => Err(EvictionError::Downloaded), None => Ok(()), } @@ -759,7 +759,7 @@ impl LayerInner { // use the already held initialization permit because it is impossible to hit the // below paths anymore essentially limiting the max loop iterations to 2. let (value, init_permit) = download(init_permit).await?; - let mut guard = self.inner.set(value, init_permit).await; + let mut guard = self.inner.set(value, init_permit); let (strong, _upgraded) = guard .get_and_upgrade() .expect("init creates strong reference, we held the init permit"); @@ -767,7 +767,7 @@ impl LayerInner { } let (weak, permit) = { - let mut locked = self.inner.get_mut_or_init(download).await?; + let mut locked = self.inner.get_or_init(download).await?; if let Some((strong, upgraded)) = locked.get_and_upgrade() { if upgraded { @@ -989,12 +989,12 @@ impl LayerInner { } } - async fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { + fn info(&self, reset: LayerAccessStatsReset) -> HistoricLayerInfo { let layer_file_name = self.desc.filename().file_name(); // this is not accurate: we could have the file locally but there was a cancellation // and now we are not in sync, or we are currently downloading it. - let remote = self.inner.get_mut().await.is_none(); + let remote = self.inner.get().is_none(); let access_stats = self.access_stats.as_api_model(reset); @@ -1053,7 +1053,7 @@ impl LayerInner { LAYER_IMPL_METRICS.inc_eviction_cancelled(EvictionCancelled::LayerGone); return; }; - match tokio::runtime::Handle::current().block_on(this.evict_blocking(version)) { + match this.evict_blocking(version) { Ok(()) => LAYER_IMPL_METRICS.inc_completed_evictions(), Err(reason) => LAYER_IMPL_METRICS.inc_eviction_cancelled(reason), } @@ -1061,7 +1061,7 @@ impl LayerInner { } } - async fn evict_blocking(&self, only_version: usize) -> Result<(), EvictionCancelled> { + fn evict_blocking(&self, only_version: usize) -> Result<(), EvictionCancelled> { // deleted or detached timeline, don't do anything. let Some(timeline) = self.timeline.upgrade() else { return Err(EvictionCancelled::TimelineGone); @@ -1070,7 +1070,7 @@ impl LayerInner { // to avoid starting a new download while we evict, keep holding on to the // permit. let _permit = { - let maybe_downloaded = self.inner.get_mut().await; + let maybe_downloaded = self.inner.get(); let (_weak, permit) = match maybe_downloaded { Some(mut guard) => { diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 735b8003b4..f96679ca69 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -1268,7 +1268,7 @@ impl Timeline { let mut historic_layers = Vec::new(); for historic_layer in layer_map.iter_historic_layers() { let historic_layer = guard.get_from_desc(&historic_layer); - historic_layers.push(historic_layer.info(reset).await); + historic_layers.push(historic_layer.info(reset)); } LayerMapInfo { From 0fd3cd27cb7ac66df5938bf219e9f12ce7b78c8a Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Fri, 9 Feb 2024 17:37:30 +0200 Subject: [PATCH 66/72] Tighten up the check for garbage after end-of-tar. Turn the warning into an error, if there is garbage after the end of imported tar file. However, it's normal for 'tar' to append extra empty blocks to the end, so tolerate those without warnings or errors. --- pageserver/src/page_service.rs | 17 ++++++++++++----- test_runner/regress/test_import.py | 10 +++------- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs index 6fc38a76d4..7b660b5eca 100644 --- a/pageserver/src/page_service.rs +++ b/pageserver/src/page_service.rs @@ -91,8 +91,8 @@ const ACTIVE_TENANT_TIMEOUT: Duration = Duration::from_millis(30000); /// `tokio_tar` already read the first such block. Read the second all-zeros block, /// and check that there is no more data after the EOF marker. /// -/// XXX: Currently, any trailing data after the EOF marker prints a warning. -/// Perhaps it should be a hard error? +/// 'tar' command can also write extra blocks of zeros, up to a record +/// size, controlled by the --record-size argument. Ignore them too. async fn read_tar_eof(mut reader: (impl AsyncRead + Unpin)) -> anyhow::Result<()> { use tokio::io::AsyncReadExt; let mut buf = [0u8; 512]; @@ -113,17 +113,24 @@ async fn read_tar_eof(mut reader: (impl AsyncRead + Unpin)) -> anyhow::Result<() anyhow::bail!("invalid tar EOF marker"); } - // Drain any data after the EOF marker + // Drain any extra zero-blocks after the EOF marker let mut trailing_bytes = 0; + let mut seen_nonzero_bytes = false; loop { let nbytes = reader.read(&mut buf).await?; trailing_bytes += nbytes; + if !buf.iter().all(|&x| x == 0) { + seen_nonzero_bytes = true; + } if nbytes == 0 { break; } } - if trailing_bytes > 0 { - warn!("ignored {trailing_bytes} unexpected bytes after the tar archive"); + if seen_nonzero_bytes { + anyhow::bail!("unexpected non-zero bytes after the tar archive"); + } + if trailing_bytes % 512 != 0 { + anyhow::bail!("unexpected number of zeros ({trailing_bytes}), not divisible by tar block size (512 bytes), after the tar archive"); } Ok(()) } diff --git a/test_runner/regress/test_import.py b/test_runner/regress/test_import.py index 3519cbbaab..7942f5cc9b 100644 --- a/test_runner/regress/test_import.py +++ b/test_runner/regress/test_import.py @@ -95,7 +95,6 @@ def test_import_from_vanilla(test_output_dir, pg_bin, vanilla_pg, neon_env_build ".*InternalServerError.*Tenant .* not found.*", ".*InternalServerError.*Timeline .* not found.*", ".*InternalServerError.*Cannot delete timeline which has child timelines.*", - ".*ignored .* unexpected bytes after the tar archive.*", ] ) @@ -142,12 +141,9 @@ def test_import_from_vanilla(test_output_dir, pg_bin, vanilla_pg, neon_env_build with pytest.raises(RuntimeError): import_tar(corrupt_base_tar, wal_tar) - # A tar with trailing garbage is currently accepted. It prints a warnings - # to the pageserver log, however. Check that. - import_tar(base_plus_garbage_tar, wal_tar) - assert env.pageserver.log_contains( - ".*WARN.*ignored .* unexpected bytes after the tar archive.*" - ) + # Importing a tar with trailing garbage fails + with pytest.raises(RuntimeError): + import_tar(base_plus_garbage_tar, wal_tar) client = env.pageserver.http_client() timeline_delete_wait_completed(client, tenant, timeline) From df5e2729a9ac3ddd80876e0d40e3ba55b95ebf0c Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Fri, 9 Feb 2024 17:37:34 +0200 Subject: [PATCH 67/72] Remove now unused allowlisted errors. I'm not sure when we stopped emitting these, but they don't seem to be needed anymore. --- test_runner/regress/test_import.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/test_runner/regress/test_import.py b/test_runner/regress/test_import.py index 7942f5cc9b..db385b3e73 100644 --- a/test_runner/regress/test_import.py +++ b/test_runner/regress/test_import.py @@ -98,15 +98,6 @@ def test_import_from_vanilla(test_output_dir, pg_bin, vanilla_pg, neon_env_build ] ) - env.pageserver.allowed_errors.extend( - [ - # FIXME: we should clean up pageserver to not print this - ".*exited with error: unexpected message type: CopyData.*", - # FIXME: Is this expected? - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*", - ] - ) - def import_tar(base, wal): env.neon_cli.raw_cli( [ From 12b39c9db95ec52353ab2bb3e21bc4a12306ce2b Mon Sep 17 00:00:00 2001 From: John Spray Date: Sat, 10 Feb 2024 11:56:52 +0000 Subject: [PATCH 68/72] control_plane: add debug APIs for force-dropping tenant/node (#6702) ## Problem When debugging/supporting this service, we sometimes need it to just forget about a tenant or node, e.g. because of an issue cleanly tearing them down. For example, if I create a tenant with a PlacementPolicy that can't be scheduled on the nodes we have, we would never be able to schedule it for a DELETE to work. ## Summary of changes - Add APIs for dropping nodes and tenants that do no teardown other than removing the entity from the DB and removing any references to it. --- control_plane/attachment_service/src/http.rs | 19 +++++++++ .../attachment_service/src/persistence.rs | 13 ++++++- .../attachment_service/src/service.rs | 39 +++++++++++++++++++ .../attachment_service/src/tenant_state.rs | 14 +++++++ test_runner/regress/test_sharding_service.py | 24 ++++++++++++ 5 files changed, 108 insertions(+), 1 deletion(-) diff --git a/control_plane/attachment_service/src/http.rs b/control_plane/attachment_service/src/http.rs index 8501e4980f..38785d3a98 100644 --- a/control_plane/attachment_service/src/http.rs +++ b/control_plane/attachment_service/src/http.rs @@ -280,6 +280,12 @@ async fn handle_node_list(req: Request) -> Result, ApiError json_response(StatusCode::OK, state.service.node_list().await?) } +async fn handle_node_drop(req: Request) -> Result, ApiError> { + let state = get_state(&req); + let node_id: NodeId = parse_request_param(&req, "node_id")?; + json_response(StatusCode::OK, state.service.node_drop(node_id).await?) +} + async fn handle_node_configure(mut req: Request) -> Result, ApiError> { let node_id: NodeId = parse_request_param(&req, "node_id")?; let config_req = json_request::(&mut req).await?; @@ -320,6 +326,13 @@ async fn handle_tenant_shard_migrate( ) } +async fn handle_tenant_drop(req: Request) -> Result, ApiError> { + let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?; + let state = get_state(&req); + + json_response(StatusCode::OK, state.service.tenant_drop(tenant_id).await?) +} + /// Status endpoint is just used for checking that our HTTP listener is up async fn handle_status(_req: Request) -> Result, ApiError> { json_response(StatusCode::OK, ()) @@ -402,6 +415,12 @@ pub fn make_router( request_span(r, handle_attach_hook) }) .post("/debug/v1/inspect", |r| request_span(r, handle_inspect)) + .post("/debug/v1/tenant/:tenant_id/drop", |r| { + request_span(r, handle_tenant_drop) + }) + .post("/debug/v1/node/:node_id/drop", |r| { + request_span(r, handle_node_drop) + }) .get("/control/v1/tenant/:tenant_id/locate", |r| { tenant_service_handler(r, handle_tenant_locate) }) diff --git a/control_plane/attachment_service/src/persistence.rs b/control_plane/attachment_service/src/persistence.rs index 623d625767..457dc43232 100644 --- a/control_plane/attachment_service/src/persistence.rs +++ b/control_plane/attachment_service/src/persistence.rs @@ -260,7 +260,6 @@ impl Persistence { /// Ordering: call this _after_ deleting the tenant on pageservers, but _before_ dropping state for /// the tenant from memory on this server. - #[allow(unused)] pub(crate) async fn delete_tenant(&self, del_tenant_id: TenantId) -> DatabaseResult<()> { use crate::schema::tenant_shards::dsl::*; self.with_conn(move |conn| -> DatabaseResult<()> { @@ -273,6 +272,18 @@ impl Persistence { .await } + pub(crate) async fn delete_node(&self, del_node_id: NodeId) -> DatabaseResult<()> { + use crate::schema::nodes::dsl::*; + self.with_conn(move |conn| -> DatabaseResult<()> { + diesel::delete(nodes) + .filter(node_id.eq(del_node_id.0 as i64)) + .execute(conn)?; + + Ok(()) + }) + .await + } + /// When a tenant invokes the /re-attach API, this function is responsible for doing an efficient /// batched increment of the generations of all tenants whose generation_pageserver is equal to /// the node that called /re-attach. diff --git a/control_plane/attachment_service/src/service.rs b/control_plane/attachment_service/src/service.rs index 0331087e0d..95efa8ecd7 100644 --- a/control_plane/attachment_service/src/service.rs +++ b/control_plane/attachment_service/src/service.rs @@ -1804,6 +1804,45 @@ impl Service { Ok(TenantShardMigrateResponse {}) } + /// This is for debug/support only: we simply drop all state for a tenant, without + /// detaching or deleting it on pageservers. + pub(crate) async fn tenant_drop(&self, tenant_id: TenantId) -> Result<(), ApiError> { + self.persistence.delete_tenant(tenant_id).await?; + + let mut locked = self.inner.write().unwrap(); + let mut shards = Vec::new(); + for (tenant_shard_id, _) in locked.tenants.range(TenantShardId::tenant_range(tenant_id)) { + shards.push(*tenant_shard_id); + } + + for shard in shards { + locked.tenants.remove(&shard); + } + + Ok(()) + } + + /// This is for debug/support only: we simply drop all state for a tenant, without + /// detaching or deleting it on pageservers. We do not try and re-schedule any + /// tenants that were on this node. + /// + /// TODO: proper node deletion API that unhooks things more gracefully + pub(crate) async fn node_drop(&self, node_id: NodeId) -> Result<(), ApiError> { + self.persistence.delete_node(node_id).await?; + + let mut locked = self.inner.write().unwrap(); + + for shard in locked.tenants.values_mut() { + shard.deref_node(node_id); + } + + let mut nodes = (*locked.nodes).clone(); + nodes.remove(&node_id); + locked.nodes = Arc::new(nodes); + + Ok(()) + } + pub(crate) async fn node_list(&self) -> Result, ApiError> { // It is convenient to avoid taking the big lock and converting Node to a serializable // structure, by fetching from storage instead of reading in-memory state. diff --git a/control_plane/attachment_service/src/tenant_state.rs b/control_plane/attachment_service/src/tenant_state.rs index c0ab076a55..1646ed9fcd 100644 --- a/control_plane/attachment_service/src/tenant_state.rs +++ b/control_plane/attachment_service/src/tenant_state.rs @@ -534,4 +534,18 @@ impl TenantState { seq: self.sequence, }) } + + // If we had any state at all referring to this node ID, drop it. Does not + // attempt to reschedule. + pub(crate) fn deref_node(&mut self, node_id: NodeId) { + if self.intent.attached == Some(node_id) { + self.intent.attached = None; + } + + self.intent.secondary.retain(|n| n != &node_id); + + self.observed.locations.remove(&node_id); + + debug_assert!(!self.intent.all_pageservers().contains(&node_id)); + } } diff --git a/test_runner/regress/test_sharding_service.py b/test_runner/regress/test_sharding_service.py index babb0d261c..248d992851 100644 --- a/test_runner/regress/test_sharding_service.py +++ b/test_runner/regress/test_sharding_service.py @@ -387,3 +387,27 @@ def test_sharding_service_compute_hook( assert notifications[1] == expect wait_until(10, 1, received_restart_notification) + + +def test_sharding_service_debug_apis(neon_env_builder: NeonEnvBuilder): + """ + Verify that occasional-use debug APIs work as expected. This is a lightweight test + that just hits the endpoints to check that they don't bitrot. + """ + + neon_env_builder.num_pageservers = 2 + env = neon_env_builder.init_start() + + tenant_id = TenantId.generate() + env.attachment_service.tenant_create(tenant_id, shard_count=2, shard_stripe_size=8192) + + # These APIs are intentionally not implemented as methods on NeonAttachmentService, as + # they're just for use in unanticipated circumstances. + env.attachment_service.request( + "POST", f"{env.attachment_service_api}/debug/v1/node/{env.pageservers[1].id}/drop" + ) + assert len(env.attachment_service.node_list()) == 1 + + env.attachment_service.request( + "POST", f"{env.attachment_service_api}/debug/v1/tenant/{tenant_id}/drop" + ) From da626fb1facd77b1159e55c5aaa39cc28ed3ed41 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Sat, 10 Feb 2024 10:48:11 +0200 Subject: [PATCH 69/72] tests: Remove "postgres is running on ... branch" messages It seems like useless chatter. The endpoint.start() itself prints a "Running command ... neon_local endpoint start" message too. --- test_runner/regress/test_ancestor_branch.py | 2 -- test_runner/regress/test_backpressure.py | 1 - test_runner/regress/test_branch_behind.py | 1 - test_runner/regress/test_clog_truncate.py | 2 -- test_runner/regress/test_config.py | 2 -- test_runner/regress/test_createdropdb.py | 2 -- test_runner/regress/test_createuser.py | 2 -- test_runner/regress/test_ddl_forwarding.py | 1 - test_runner/regress/test_fullbackup.py | 1 - test_runner/regress/test_gc_aggressive.py | 1 - test_runner/regress/test_layer_bloating.py | 1 - test_runner/regress/test_lfc_resize.py | 1 - test_runner/regress/test_logical_replication.py | 2 -- test_runner/regress/test_lsn_mapping.py | 2 -- test_runner/regress/test_multixact.py | 3 --- test_runner/regress/test_neon_extension.py | 3 --- test_runner/regress/test_old_request_lsn.py | 1 - test_runner/regress/test_parallel_copy.py | 2 -- test_runner/regress/test_pitr_gc.py | 1 - test_runner/regress/test_read_validation.py | 2 -- test_runner/regress/test_readonly_node.py | 1 - test_runner/regress/test_recovery.py | 1 - test_runner/regress/test_subxacts.py | 8 +------- test_runner/regress/test_timeline_size.py | 6 ------ test_runner/regress/test_twophase.py | 1 - test_runner/regress/test_vm_bits.py | 2 -- 26 files changed, 1 insertion(+), 51 deletions(-) diff --git a/test_runner/regress/test_ancestor_branch.py b/test_runner/regress/test_ancestor_branch.py index 0e390ba9e5..d16d2d6a24 100644 --- a/test_runner/regress/test_ancestor_branch.py +++ b/test_runner/regress/test_ancestor_branch.py @@ -45,7 +45,6 @@ def test_ancestor_branch(neon_env_builder: NeonEnvBuilder): # Create branch1. env.neon_cli.create_branch("branch1", "main", tenant_id=tenant, ancestor_start_lsn=lsn_100) endpoint_branch1 = env.endpoints.create_start("branch1", tenant_id=tenant) - log.info("postgres is running on 'branch1' branch") branch1_cur = endpoint_branch1.connect().cursor() branch1_timeline = TimelineId(query_scalar(branch1_cur, "SHOW neon.timeline_id")) @@ -68,7 +67,6 @@ def test_ancestor_branch(neon_env_builder: NeonEnvBuilder): # Create branch2. env.neon_cli.create_branch("branch2", "branch1", tenant_id=tenant, ancestor_start_lsn=lsn_200) endpoint_branch2 = env.endpoints.create_start("branch2", tenant_id=tenant) - log.info("postgres is running on 'branch2' branch") branch2_cur = endpoint_branch2.connect().cursor() branch2_timeline = TimelineId(query_scalar(branch2_cur, "SHOW neon.timeline_id")) diff --git a/test_runner/regress/test_backpressure.py b/test_runner/regress/test_backpressure.py index bc3faf9271..819912dd05 100644 --- a/test_runner/regress/test_backpressure.py +++ b/test_runner/regress/test_backpressure.py @@ -107,7 +107,6 @@ def test_backpressure_received_lsn_lag(neon_env_builder: NeonEnvBuilder): # which is needed for backpressure_lsns() to work endpoint.respec(skip_pg_catalog_updates=False) endpoint.start() - log.info("postgres is running on 'test_backpressure' branch") # setup check thread check_stop_event = threading.Event() diff --git a/test_runner/regress/test_branch_behind.py b/test_runner/regress/test_branch_behind.py index 9879254897..46c74a26b8 100644 --- a/test_runner/regress/test_branch_behind.py +++ b/test_runner/regress/test_branch_behind.py @@ -21,7 +21,6 @@ def test_branch_behind(neon_env_builder: NeonEnvBuilder): # Branch at the point where only 100 rows were inserted branch_behind_timeline_id = env.neon_cli.create_branch("test_branch_behind") endpoint_main = env.endpoints.create_start("test_branch_behind") - log.info("postgres is running on 'test_branch_behind' branch") main_cur = endpoint_main.connect().cursor() diff --git a/test_runner/regress/test_clog_truncate.py b/test_runner/regress/test_clog_truncate.py index f22eca02cc..26e6e336b9 100644 --- a/test_runner/regress/test_clog_truncate.py +++ b/test_runner/regress/test_clog_truncate.py @@ -25,7 +25,6 @@ def test_clog_truncate(neon_simple_env: NeonEnv): ] endpoint = env.endpoints.create_start("test_clog_truncate", config_lines=config) - log.info("postgres is running on test_clog_truncate branch") # Install extension containing function needed for test endpoint.safe_psql("CREATE EXTENSION neon_test_utils") @@ -62,7 +61,6 @@ def test_clog_truncate(neon_simple_env: NeonEnv): "test_clog_truncate_new", "test_clog_truncate", ancestor_start_lsn=lsn_after_truncation ) endpoint2 = env.endpoints.create_start("test_clog_truncate_new") - log.info("postgres is running on test_clog_truncate_new branch") # check that new node doesn't contain truncated segment pg_xact_0000_path_new = os.path.join(endpoint2.pg_xact_dir_path(), "0000") diff --git a/test_runner/regress/test_config.py b/test_runner/regress/test_config.py index 0ea5784b67..4bb7df1e6a 100644 --- a/test_runner/regress/test_config.py +++ b/test_runner/regress/test_config.py @@ -1,6 +1,5 @@ from contextlib import closing -from fixtures.log_helper import log from fixtures.neon_fixtures import NeonEnv @@ -13,7 +12,6 @@ def test_config(neon_simple_env: NeonEnv): # change config endpoint = env.endpoints.create_start("test_config", config_lines=["log_min_messages=debug1"]) - log.info("postgres is running on test_config branch") with closing(endpoint.connect()) as conn: with conn.cursor() as cur: diff --git a/test_runner/regress/test_createdropdb.py b/test_runner/regress/test_createdropdb.py index 500d19cf31..f741a9fc87 100644 --- a/test_runner/regress/test_createdropdb.py +++ b/test_runner/regress/test_createdropdb.py @@ -20,7 +20,6 @@ def test_createdb(neon_simple_env: NeonEnv, strategy: str): env.neon_cli.create_branch("test_createdb", "empty") endpoint = env.endpoints.create_start("test_createdb") - log.info("postgres is running on 'test_createdb' branch") with endpoint.cursor() as cur: # Cause a 'relmapper' change in the original branch @@ -65,7 +64,6 @@ def test_dropdb(neon_simple_env: NeonEnv, test_output_dir): env = neon_simple_env env.neon_cli.create_branch("test_dropdb", "empty") endpoint = env.endpoints.create_start("test_dropdb") - log.info("postgres is running on 'test_dropdb' branch") with endpoint.cursor() as cur: cur.execute("CREATE DATABASE foodb") diff --git a/test_runner/regress/test_createuser.py b/test_runner/regress/test_createuser.py index f1bc405287..17d9824f52 100644 --- a/test_runner/regress/test_createuser.py +++ b/test_runner/regress/test_createuser.py @@ -1,4 +1,3 @@ -from fixtures.log_helper import log from fixtures.neon_fixtures import NeonEnv from fixtures.utils import query_scalar @@ -10,7 +9,6 @@ def test_createuser(neon_simple_env: NeonEnv): env = neon_simple_env env.neon_cli.create_branch("test_createuser", "empty") endpoint = env.endpoints.create_start("test_createuser") - log.info("postgres is running on 'test_createuser' branch") with endpoint.cursor() as cur: # Cause a 'relmapper' change in the original branch diff --git a/test_runner/regress/test_ddl_forwarding.py b/test_runner/regress/test_ddl_forwarding.py index 7174487e68..50da673d87 100644 --- a/test_runner/regress/test_ddl_forwarding.py +++ b/test_runner/regress/test_ddl_forwarding.py @@ -296,7 +296,6 @@ def test_ddl_forwarding_invalid_db(neon_simple_env: NeonEnv): # Some non-existent url config_lines=["neon.console_url=http://localhost:9999/unknown/api/v0/roles_and_databases"], ) - log.info("postgres is running on 'test_ddl_forwarding_invalid_db' branch") with endpoint.cursor() as cur: cur.execute("SET neon.forward_ddl = false") diff --git a/test_runner/regress/test_fullbackup.py b/test_runner/regress/test_fullbackup.py index a456c06862..9a22084671 100644 --- a/test_runner/regress/test_fullbackup.py +++ b/test_runner/regress/test_fullbackup.py @@ -26,7 +26,6 @@ def test_fullbackup( env.neon_cli.create_branch("test_fullbackup") endpoint_main = env.endpoints.create_start("test_fullbackup") - log.info("postgres is running on 'test_fullbackup' branch") with endpoint_main.cursor() as cur: timeline = TimelineId(query_scalar(cur, "SHOW neon.timeline_id")) diff --git a/test_runner/regress/test_gc_aggressive.py b/test_runner/regress/test_gc_aggressive.py index ef68049ee7..c5070ee815 100644 --- a/test_runner/regress/test_gc_aggressive.py +++ b/test_runner/regress/test_gc_aggressive.py @@ -71,7 +71,6 @@ def test_gc_aggressive(neon_env_builder: NeonEnvBuilder): env = neon_env_builder.init_start() timeline = env.neon_cli.create_branch("test_gc_aggressive", "main") endpoint = env.endpoints.create_start("test_gc_aggressive") - log.info("postgres is running on test_gc_aggressive branch") with endpoint.cursor() as cur: # Create table, and insert the first 100 rows diff --git a/test_runner/regress/test_layer_bloating.py b/test_runner/regress/test_layer_bloating.py index 70b115ad61..bf5834b665 100644 --- a/test_runner/regress/test_layer_bloating.py +++ b/test_runner/regress/test_layer_bloating.py @@ -21,7 +21,6 @@ def test_layer_bloating(neon_simple_env: NeonEnv, vanilla_pg): "test_logical_replication", config_lines=["log_statement=all"] ) - log.info("postgres is running on 'test_logical_replication' branch") pg_conn = endpoint.connect() cur = pg_conn.cursor() diff --git a/test_runner/regress/test_lfc_resize.py b/test_runner/regress/test_lfc_resize.py index 5c68a63d06..2a3442448a 100644 --- a/test_runner/regress/test_lfc_resize.py +++ b/test_runner/regress/test_lfc_resize.py @@ -23,7 +23,6 @@ def test_lfc_resize(neon_simple_env: NeonEnv, pg_bin: PgBin): ) n_resize = 10 scale = 10 - log.info("postgres is running on 'test_lfc_resize' branch") def run_pgbench(connstr: str): log.info(f"Start a pgbench workload on pg {connstr}") diff --git a/test_runner/regress/test_logical_replication.py b/test_runner/regress/test_logical_replication.py index 059ddf79ec..eff0b124d3 100644 --- a/test_runner/regress/test_logical_replication.py +++ b/test_runner/regress/test_logical_replication.py @@ -26,7 +26,6 @@ def test_logical_replication(neon_simple_env: NeonEnv, vanilla_pg): "test_logical_replication", config_lines=["log_statement=all"] ) - log.info("postgres is running on 'test_logical_replication' branch") pg_conn = endpoint.connect() cur = pg_conn.cursor() @@ -315,7 +314,6 @@ def test_slots_and_branching(neon_simple_env: NeonEnv): # Create branch ws. env.neon_cli.create_branch("ws", "main", tenant_id=tenant) ws_branch = env.endpoints.create_start("ws", tenant_id=tenant) - log.info("postgres is running on 'ws' branch") # Check that we can create slot with the same name ws_cur = ws_branch.connect().cursor() diff --git a/test_runner/regress/test_lsn_mapping.py b/test_runner/regress/test_lsn_mapping.py index 50d7c74af0..5813231aab 100644 --- a/test_runner/regress/test_lsn_mapping.py +++ b/test_runner/regress/test_lsn_mapping.py @@ -28,7 +28,6 @@ def test_lsn_mapping(neon_env_builder: NeonEnvBuilder): timeline_id = env.neon_cli.create_branch("test_lsn_mapping", tenant_id=tenant_id) endpoint_main = env.endpoints.create_start("test_lsn_mapping", tenant_id=tenant_id) timeline_id = endpoint_main.safe_psql("show neon.timeline_id")[0][0] - log.info("postgres is running on 'main' branch") cur = endpoint_main.connect().cursor() @@ -114,7 +113,6 @@ def test_ts_of_lsn_api(neon_env_builder: NeonEnvBuilder): new_timeline_id = env.neon_cli.create_branch("test_ts_of_lsn_api") endpoint_main = env.endpoints.create_start("test_ts_of_lsn_api") - log.info("postgres is running on 'test_ts_of_lsn_api' branch") cur = endpoint_main.connect().cursor() # Create table, and insert rows, each in a separate transaction diff --git a/test_runner/regress/test_multixact.py b/test_runner/regress/test_multixact.py index 9db463dc4a..88f7a5db59 100644 --- a/test_runner/regress/test_multixact.py +++ b/test_runner/regress/test_multixact.py @@ -1,4 +1,3 @@ -from fixtures.log_helper import log from fixtures.neon_fixtures import NeonEnv, check_restored_datadir_content from fixtures.utils import query_scalar @@ -18,7 +17,6 @@ def test_multixact(neon_simple_env: NeonEnv, test_output_dir): env.neon_cli.create_branch("test_multixact", "empty") endpoint = env.endpoints.create_start("test_multixact") - log.info("postgres is running on 'test_multixact' branch") cur = endpoint.connect().cursor() cur.execute( """ @@ -78,7 +76,6 @@ def test_multixact(neon_simple_env: NeonEnv, test_output_dir): env.neon_cli.create_branch("test_multixact_new", "test_multixact", ancestor_start_lsn=lsn) endpoint_new = env.endpoints.create_start("test_multixact_new") - log.info("postgres is running on 'test_multixact_new' branch") next_multixact_id_new = endpoint_new.safe_psql( "SELECT next_multixact_id FROM pg_control_checkpoint()" )[0][0] diff --git a/test_runner/regress/test_neon_extension.py b/test_runner/regress/test_neon_extension.py index 998f84f968..62225e7b92 100644 --- a/test_runner/regress/test_neon_extension.py +++ b/test_runner/regress/test_neon_extension.py @@ -1,6 +1,5 @@ from contextlib import closing -from fixtures.log_helper import log from fixtures.neon_fixtures import NeonEnvBuilder @@ -14,8 +13,6 @@ def test_neon_extension(neon_env_builder: NeonEnvBuilder): endpoint_main.respec(skip_pg_catalog_updates=False) endpoint_main.start() - log.info("postgres is running on 'test_create_extension_neon' branch") - with closing(endpoint_main.connect()) as conn: with conn.cursor() as cur: cur.execute("SELECT extversion from pg_extension where extname='neon'") diff --git a/test_runner/regress/test_old_request_lsn.py b/test_runner/regress/test_old_request_lsn.py index 9b0bab5125..391305c58a 100644 --- a/test_runner/regress/test_old_request_lsn.py +++ b/test_runner/regress/test_old_request_lsn.py @@ -20,7 +20,6 @@ def test_old_request_lsn(neon_env_builder: NeonEnvBuilder): env = neon_env_builder.init_start() env.neon_cli.create_branch("test_old_request_lsn", "main") endpoint = env.endpoints.create_start("test_old_request_lsn") - log.info("postgres is running on test_old_request_lsn branch") pg_conn = endpoint.connect() cur = pg_conn.cursor() diff --git a/test_runner/regress/test_parallel_copy.py b/test_runner/regress/test_parallel_copy.py index 6f74d50b92..b33e387a66 100644 --- a/test_runner/regress/test_parallel_copy.py +++ b/test_runner/regress/test_parallel_copy.py @@ -1,7 +1,6 @@ import asyncio from io import BytesIO -from fixtures.log_helper import log from fixtures.neon_fixtures import Endpoint, NeonEnv @@ -44,7 +43,6 @@ def test_parallel_copy(neon_simple_env: NeonEnv, n_parallel=5): env = neon_simple_env env.neon_cli.create_branch("test_parallel_copy", "empty") endpoint = env.endpoints.create_start("test_parallel_copy") - log.info("postgres is running on 'test_parallel_copy' branch") # Create test table conn = endpoint.connect() diff --git a/test_runner/regress/test_pitr_gc.py b/test_runner/regress/test_pitr_gc.py index c2ea5b332a..539ef3eda7 100644 --- a/test_runner/regress/test_pitr_gc.py +++ b/test_runner/regress/test_pitr_gc.py @@ -16,7 +16,6 @@ def test_pitr_gc(neon_env_builder: NeonEnvBuilder): env = neon_env_builder.init_start() endpoint_main = env.endpoints.create_start("main") - log.info("postgres is running on 'main' branch") main_pg_conn = endpoint_main.connect() main_cur = main_pg_conn.cursor() diff --git a/test_runner/regress/test_read_validation.py b/test_runner/regress/test_read_validation.py index d695410efc..effb7e83f9 100644 --- a/test_runner/regress/test_read_validation.py +++ b/test_runner/regress/test_read_validation.py @@ -18,7 +18,6 @@ def test_read_validation(neon_simple_env: NeonEnv): env.neon_cli.create_branch("test_read_validation", "empty") endpoint = env.endpoints.create_start("test_read_validation") - log.info("postgres is running on 'test_read_validation' branch") with closing(endpoint.connect()) as con: with con.cursor() as c: @@ -145,7 +144,6 @@ def test_read_validation_neg(neon_simple_env: NeonEnv): env.pageserver.allowed_errors.append(".*invalid LSN\\(0\\) in request.*") endpoint = env.endpoints.create_start("test_read_validation_neg") - log.info("postgres is running on 'test_read_validation_neg' branch") with closing(endpoint.connect()) as con: with con.cursor() as c: diff --git a/test_runner/regress/test_readonly_node.py b/test_runner/regress/test_readonly_node.py index 2d641e36a7..b7c8f36107 100644 --- a/test_runner/regress/test_readonly_node.py +++ b/test_runner/regress/test_readonly_node.py @@ -16,7 +16,6 @@ def test_readonly_node(neon_simple_env: NeonEnv): env = neon_simple_env env.neon_cli.create_branch("test_readonly_node", "empty") endpoint_main = env.endpoints.create_start("test_readonly_node") - log.info("postgres is running on 'test_readonly_node' branch") env.pageserver.allowed_errors.append(".*basebackup .* failed: invalid basebackup lsn.*") diff --git a/test_runner/regress/test_recovery.py b/test_runner/regress/test_recovery.py index 9d7a4a8fd6..6aac1e1d84 100644 --- a/test_runner/regress/test_recovery.py +++ b/test_runner/regress/test_recovery.py @@ -19,7 +19,6 @@ def test_pageserver_recovery(neon_env_builder: NeonEnvBuilder): env.neon_cli.create_branch("test_pageserver_recovery", "main") endpoint = env.endpoints.create_start("test_pageserver_recovery") - log.info("postgres is running on 'test_pageserver_recovery' branch") with closing(endpoint.connect()) as conn: with conn.cursor() as cur: diff --git a/test_runner/regress/test_subxacts.py b/test_runner/regress/test_subxacts.py index eb96a8faa4..10cb00c780 100644 --- a/test_runner/regress/test_subxacts.py +++ b/test_runner/regress/test_subxacts.py @@ -1,4 +1,3 @@ -from fixtures.log_helper import log from fixtures.neon_fixtures import NeonEnv, check_restored_datadir_content @@ -13,15 +12,10 @@ def test_subxacts(neon_simple_env: NeonEnv, test_output_dir): env.neon_cli.create_branch("test_subxacts", "empty") endpoint = env.endpoints.create_start("test_subxacts") - log.info("postgres is running on 'test_subxacts' branch") pg_conn = endpoint.connect() cur = pg_conn.cursor() - cur.execute( - """ - CREATE TABLE t1(i int, j int); - """ - ) + cur.execute("CREATE TABLE t1(i int, j int);") cur.execute("select pg_switch_wal();") diff --git a/test_runner/regress/test_timeline_size.py b/test_runner/regress/test_timeline_size.py index cd7203bba6..a3f99948d3 100644 --- a/test_runner/regress/test_timeline_size.py +++ b/test_runner/regress/test_timeline_size.py @@ -43,7 +43,6 @@ def test_timeline_size(neon_simple_env: NeonEnv): client.timeline_wait_logical_size(env.initial_tenant, new_timeline_id) endpoint_main = env.endpoints.create_start("test_timeline_size") - log.info("postgres is running on 'test_timeline_size' branch") with closing(endpoint_main.connect()) as conn: with conn.cursor() as cur: @@ -79,7 +78,6 @@ def test_timeline_size_createdropdb(neon_simple_env: NeonEnv): ) endpoint_main = env.endpoints.create_start("test_timeline_size_createdropdb") - log.info("postgres is running on 'test_timeline_size_createdropdb' branch") with closing(endpoint_main.connect()) as conn: with conn.cursor() as cur: @@ -162,8 +160,6 @@ def test_timeline_size_quota_on_startup(neon_env_builder: NeonEnvBuilder): ) endpoint_main.start() - log.info("postgres is running on 'test_timeline_size_quota_on_startup' branch") - with closing(endpoint_main.connect()) as conn: with conn.cursor() as cur: cur.execute("CREATE TABLE foo (t text)") @@ -231,8 +227,6 @@ def test_timeline_size_quota(neon_env_builder: NeonEnvBuilder): endpoint_main.respec(skip_pg_catalog_updates=False) endpoint_main.start() - log.info("postgres is running on 'test_timeline_size_quota' branch") - with closing(endpoint_main.connect()) as conn: with conn.cursor() as cur: cur.execute("CREATE TABLE foo (t text)") diff --git a/test_runner/regress/test_twophase.py b/test_runner/regress/test_twophase.py index 305271c715..dd76689008 100644 --- a/test_runner/regress/test_twophase.py +++ b/test_runner/regress/test_twophase.py @@ -13,7 +13,6 @@ def test_twophase(neon_simple_env: NeonEnv): endpoint = env.endpoints.create_start( "test_twophase", config_lines=["max_prepared_transactions=5"] ) - log.info("postgres is running on 'test_twophase' branch") conn = endpoint.connect() cur = conn.cursor() diff --git a/test_runner/regress/test_vm_bits.py b/test_runner/regress/test_vm_bits.py index 06c30b8d81..1377bed6f6 100644 --- a/test_runner/regress/test_vm_bits.py +++ b/test_runner/regress/test_vm_bits.py @@ -14,7 +14,6 @@ def test_vm_bit_clear(neon_simple_env: NeonEnv): env.neon_cli.create_branch("test_vm_bit_clear", "empty") endpoint = env.endpoints.create_start("test_vm_bit_clear") - log.info("postgres is running on 'test_vm_bit_clear' branch") pg_conn = endpoint.connect() cur = pg_conn.cursor() @@ -93,7 +92,6 @@ def test_vm_bit_clear(neon_simple_env: NeonEnv): # server at the right point-in-time avoids that full-page image. endpoint_new = env.endpoints.create_start("test_vm_bit_clear_new") - log.info("postgres is running on 'test_vm_bit_clear_new' branch") pg_new_conn = endpoint_new.connect() cur_new = pg_new_conn.cursor() From 241dcbf70ce117a8b956fb990f13fee67029a197 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Sat, 10 Feb 2024 10:50:52 +0200 Subject: [PATCH 70/72] tests: Remove "Running in ..." log message from every CLI call It's always the same directory, the test's "repo" directory. --- test_runner/fixtures/neon_fixtures.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 231eebff52..31acb045ae 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1400,7 +1400,6 @@ class AbstractNeonCli(abc.ABC): args = [bin_neon] + arguments log.info('Running command "{}"'.format(" ".join(args))) - log.info(f'Running in "{self.env.repo_dir}"') env_vars = os.environ.copy() env_vars["NEON_REPO_DIR"] = str(self.env.repo_dir) From d77583c86ab3cf4d5b555d86a7b665c1457f97c8 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Sat, 10 Feb 2024 11:10:48 +0200 Subject: [PATCH 71/72] tests: Remove obsolete allowlist entries Commit 9a6c0be823 removed the code that printed these warnings: marking {} as locally complete, while it doesnt exist in remote index No timelines to attach received Remove those warnings from all the allowlists in tests. --- test_runner/regress/test_compatibility.py | 5 ----- test_runner/regress/test_import.py | 5 ----- test_runner/regress/test_remote_storage.py | 3 --- test_runner/regress/test_tenant_relocation.py | 2 -- test_runner/regress/test_tenants.py | 1 - .../regress/test_tenants_with_remote_storage.py | 16 ---------------- test_runner/regress/test_wal_acceptor.py | 10 ---------- 7 files changed, 42 deletions(-) diff --git a/test_runner/regress/test_compatibility.py b/test_runner/regress/test_compatibility.py index d5d70951be..826821e52b 100644 --- a/test_runner/regress/test_compatibility.py +++ b/test_runner/regress/test_compatibility.py @@ -112,11 +112,6 @@ def test_create_snapshot( env = neon_env_builder.init_start() endpoint = env.endpoints.create_start("main") - # FIXME: Is this expected? - env.pageserver.allowed_errors.append( - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*" - ) - pg_bin.run_capture(["pgbench", "--initialize", "--scale=10", endpoint.connstr()]) pg_bin.run_capture(["pgbench", "--time=60", "--progress=2", endpoint.connstr()]) pg_bin.run_capture( diff --git a/test_runner/regress/test_import.py b/test_runner/regress/test_import.py index db385b3e73..ec57860033 100644 --- a/test_runner/regress/test_import.py +++ b/test_runner/regress/test_import.py @@ -159,11 +159,6 @@ def test_import_from_pageserver_small( neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.LOCAL_FS) env = neon_env_builder.init_start() - # FIXME: Is this expected? - env.pageserver.allowed_errors.append( - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*" - ) - timeline = env.neon_cli.create_branch("test_import_from_pageserver_small") endpoint = env.endpoints.create_start("test_import_from_pageserver_small") diff --git a/test_runner/regress/test_remote_storage.py b/test_runner/regress/test_remote_storage.py index 98b2e856ec..32b4f54fbd 100644 --- a/test_runner/regress/test_remote_storage.py +++ b/test_runner/regress/test_remote_storage.py @@ -73,9 +73,6 @@ def test_remote_storage_backup_and_restore( env.pageserver.allowed_errors.extend( [ - # FIXME: Is this expected? - ".*marking .* as locally complete, while it doesnt exist in remote index.*", - ".*No timelines to attach received.*", ".*Failed to get local tenant state.*", # FIXME retry downloads without throwing errors ".*failed to load remote timeline.*", diff --git a/test_runner/regress/test_tenant_relocation.py b/test_runner/regress/test_tenant_relocation.py index 80b4fab1d3..f4eb6b092d 100644 --- a/test_runner/regress/test_tenant_relocation.py +++ b/test_runner/regress/test_tenant_relocation.py @@ -213,8 +213,6 @@ def test_tenant_relocation( env.pageservers[0].allowed_errors.extend( [ - # FIXME: Is this expected? - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*", # Needed for detach polling on the original pageserver f".*NotFound: tenant {tenant_id}.*", # We will dual-attach in this test, so stale generations are expected diff --git a/test_runner/regress/test_tenants.py b/test_runner/regress/test_tenants.py index ba391a69d8..bf317808ee 100644 --- a/test_runner/regress/test_tenants.py +++ b/test_runner/regress/test_tenants.py @@ -285,7 +285,6 @@ def test_pageserver_with_empty_tenants(neon_env_builder: NeonEnvBuilder): env.pageserver.allowed_errors.extend( [ - ".*marking .* as locally complete, while it doesnt exist in remote index.*", ".*load failed.*list timelines directory.*", ] ) diff --git a/test_runner/regress/test_tenants_with_remote_storage.py b/test_runner/regress/test_tenants_with_remote_storage.py index 6f05d7f7cb..1c693a0df5 100644 --- a/test_runner/regress/test_tenants_with_remote_storage.py +++ b/test_runner/regress/test_tenants_with_remote_storage.py @@ -61,11 +61,6 @@ async def all_tenants_workload(env: NeonEnv, tenants_endpoints): def test_tenants_many(neon_env_builder: NeonEnvBuilder): env = neon_env_builder.init_start() - # FIXME: Is this expected? - env.pageserver.allowed_errors.append( - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*" - ) - tenants_endpoints: List[Tuple[TenantId, Endpoint]] = [] for _ in range(1, 5): @@ -117,14 +112,6 @@ def test_tenants_attached_after_download(neon_env_builder: NeonEnvBuilder): ##### First start, insert secret data and upload it to the remote storage env = neon_env_builder.init_start() - env.pageserver.allowed_errors.extend( - [ - # FIXME: Are these expected? - ".*No timelines to attach received.*", - ".*marking .* as locally complete, while it doesnt exist in remote index.*", - ] - ) - pageserver_http = env.pageserver.http_client() endpoint = env.endpoints.create_start("main") @@ -223,9 +210,6 @@ def test_tenant_redownloads_truncated_file_on_startup( env.pageserver.allowed_errors.extend( [ ".*removing local file .* because .*", - # FIXME: Are these expected? - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*", - ".*No timelines to attach received.*", ] ) diff --git a/test_runner/regress/test_wal_acceptor.py b/test_runner/regress/test_wal_acceptor.py index dab446fcfd..3d7bba6153 100644 --- a/test_runner/regress/test_wal_acceptor.py +++ b/test_runner/regress/test_wal_acceptor.py @@ -280,11 +280,6 @@ def test_broker(neon_env_builder: NeonEnvBuilder): tenant_id = env.initial_tenant timeline_id = env.neon_cli.create_branch("test_broker", "main") - # FIXME: Is this expected? - env.pageserver.allowed_errors.append( - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*" - ) - endpoint = env.endpoints.create_start("test_broker") endpoint.safe_psql("CREATE TABLE t(key int primary key, value text)") @@ -342,11 +337,6 @@ def test_wal_removal(neon_env_builder: NeonEnvBuilder, auth_enabled: bool): neon_env_builder.auth_enabled = auth_enabled env = neon_env_builder.init_start() - # FIXME: Is this expected? - env.pageserver.allowed_errors.append( - ".*init_tenant_mgr: marking .* as locally complete, while it doesnt exist in remote index.*" - ) - tenant_id = env.initial_tenant timeline_id = env.neon_cli.create_branch("test_safekeepers_wal_removal") endpoint = env.endpoints.create_start("test_safekeepers_wal_removal") From e5daf366ac92a5398c09ea956ba03ac03848d3f8 Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Sat, 10 Feb 2024 11:25:47 +0200 Subject: [PATCH 72/72] tests: Remove unnecessary port config with VanillaPostgres class VanillaPostgres constructor prints the "port={port}" line to the config file, no need to do it in the callers. The TODO comment that it would be nice if VanillaPostgres could pick the port by itself is still valid though. --- test_runner/fixtures/neon_fixtures.py | 1 + test_runner/regress/test_fullbackup.py | 6 ------ test_runner/regress/test_timeline_size.py | 1 - 3 files changed, 1 insertion(+), 7 deletions(-) diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 31acb045ae..faa8effe10 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -2458,6 +2458,7 @@ def pg_bin(test_output_dir: Path, pg_distrib_dir: Path, pg_version: PgVersion) - return PgBin(test_output_dir, pg_distrib_dir, pg_version) +# TODO make port an optional argument class VanillaPostgres(PgProtocol): def __init__(self, pgdatadir: Path, pg_bin: PgBin, port: int, init: bool = True): super().__init__(host="localhost", port=port, dbname="postgres") diff --git a/test_runner/regress/test_fullbackup.py b/test_runner/regress/test_fullbackup.py index 9a22084671..d5f898492b 100644 --- a/test_runner/regress/test_fullbackup.py +++ b/test_runner/regress/test_fullbackup.py @@ -66,12 +66,6 @@ def test_fullbackup( # Restore from the backup and find the data we inserted port = port_distributor.get_port() with VanillaPostgres(restored_dir_path, pg_bin, port, init=False) as vanilla_pg: - # TODO make port an optional argument - vanilla_pg.configure( - [ - f"port={port}", - ] - ) vanilla_pg.start() num_rows_found = vanilla_pg.safe_psql("select count(*) from tbl;", user="cloud_admin")[0][0] assert num_rows == num_rows_found diff --git a/test_runner/regress/test_timeline_size.py b/test_runner/regress/test_timeline_size.py index a3f99948d3..0788c49c7b 100644 --- a/test_runner/regress/test_timeline_size.py +++ b/test_runner/regress/test_timeline_size.py @@ -579,7 +579,6 @@ def test_timeline_size_metrics( pg_bin = PgBin(test_output_dir, pg_distrib_dir, pg_version) port = port_distributor.get_port() with VanillaPostgres(pgdatadir, pg_bin, port) as vanilla_pg: - vanilla_pg.configure([f"port={port}"]) vanilla_pg.start() # Create database based on template0 because we can't connect to template0