diff --git a/.bazelrc b/.bazelrc index 0ea65008..fdae62c2 100644 --- a/.bazelrc +++ b/.bazelrc @@ -19,12 +19,12 @@ test:run_all_tests --test_verbose_timeout_warnings build:noexcept --copt=-fno-exceptions # Grant exceptions to some dependencies so they can use exceptions build:noexcept --per_file_copt=.*boost.*@-fexceptions -build:noexcept --per_file_copt=.*cc/aws/proxy.*@-fexceptions -build:noexcept --per_file_copt=.*cc/roma.*@-fexceptions +build:noexcept --per_file_copt=.*src/aws/proxy.*@-fexceptions +build:noexcept --per_file_copt=.*src/roma.*@-fexceptions build:noexcept --per_file_copt=.*oneTBB.*@-fexceptions build:noexcept --per_file_copt=.*com_github_nghttp2_nghttp2.*@-fexceptions -build:noexcept --per_file_copt=.*cc/core.*@-fexceptions -build:noexcept --per_file_copt=.*cc/cpio.*@-fexceptions +build:noexcept --per_file_copt=.*src/core.*@-fexceptions +build:noexcept --per_file_copt=.*src/cpio.*@-fexceptions test --test_output=errors # Disable ICU linking for googleurl. @@ -40,9 +40,17 @@ build:clang --host_cxxopt=-std=c++17 build:clang --client_env=BAZEL_CXXOPTS=-std=c++17 build:clang --per_file_copt=external/nitrokmscli_.*\.c@-Wno-int-conversion +build:cpp_nowarn --copt=-Werror +build:cpp_nowarn --per_file_copt=external/.*@-Wno-error + +build:clang-tidy --aspects @bazel_clang_tidy//clang_tidy:clang_tidy.bzl%clang_tidy_aspect +build:clang-tidy --output_groups=report +build:clang-tidy --@bazel_clang_tidy//:clang_tidy_config=//:clang_tidy_config + # Required to use protos in wasm_cc_binary/inline_wasm_cc_binary build:emscripten --per_file_copt=.*zlib.*@-Wno-deprecated-non-prototype build:emscripten --per_file_copt=.*utf8_range.*@-Wno-unused-function +build:emscripten --per_file_copt=.*protobuf.*@-Wno-deprecated-declarations # Address sanitizer, set action_env to segregate cache entries build:asan --action_env=PRIVACY_SANDBOX_SERVERS_ASAN=1 @@ -106,9 +114,14 @@ build:aws_platform --@google_privacysandbox_servers_common//:platform=aws build:gcp_platform --//:platform=gcp build:gcp_platform --@google_privacysandbox_servers_common//:platform=gcp +# --config prod_mode: builds the service in prod mode +build:prod_mode --//:mode=prod +build:prod_mode --@google_privacysandbox_servers_common//:build_flavor=prod + +# --config prod_mode: builds the service in prod mode +build:nonprod_mode --//:mode=nonprod +build:nonprod_mode --@google_privacysandbox_servers_common//:build_flavor=non_prod + try-import %workspace%/builders/.coverage.bazelrc coverage --test_tag_filters=-nocoverage coverage --test_size_filters=-enormous - -build:non_prod --@google_privacysandbox_servers_common//:build_flavor=non_prod -build:prod --@google_privacysandbox_servers_common//:build_flavor=prod diff --git a/.bazelversion b/.bazelversion index 91e4a9f2..f22d756d 100644 --- a/.bazelversion +++ b/.bazelversion @@ -1 +1 @@ -6.3.2 +6.5.0 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c2e6ace2..fe88324b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,7 +27,7 @@ exclude: (?x)^( fail_fast: true repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.4.0 + rev: v4.5.0 hooks: - id: end-of-file-fixer - id: fix-byte-order-marker @@ -43,6 +43,12 @@ repos: - id: check-executables-have-shebangs - id: detect-private-key +- repo: https://github.com/tcort/markdown-link-check + rev: v3.11.2 + hooks: + - id: markdown-link-check + args: [-c .precommit_configs/markdown-link-check.json] + - repo: https://github.com/jumanjihouse/pre-commit-hooks rev: 3.0.0 hooks: @@ -54,12 +60,12 @@ repos: exclude: ^(google_internal|builders/images)/.*$ - repo: https://github.com/bufbuild/buf - rev: v1.26.1 + rev: v1.29.0 hooks: - id: buf-format - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v16.0.6 + rev: v17.0.6 hooks: - id: clang-format types_or: @@ -100,7 +106,7 @@ repos: - terraform - repo: https://github.com/pre-commit/mirrors-prettier - rev: v3.0.3 + rev: v3.1.0 hooks: - id: prettier types_or: @@ -113,7 +119,7 @@ repos: )$ - repo: https://github.com/DavidAnson/markdownlint-cli2 - rev: v0.9.2 + rev: v0.12.1 hooks: - id: markdownlint-cli2 name: lint markdown @@ -148,7 +154,13 @@ repos: - --quiet - repo: https://github.com/psf/black - rev: 23.7.0 + rev: 24.2.0 hooks: - id: black name: black python formatter + +- repo: https://github.com/tekwizely/pre-commit-golang + rev: v1.0.0-rc.1 + hooks: + - id: go-fmt + name: go format diff --git a/.precommit_configs/markdown-link-check.json b/.precommit_configs/markdown-link-check.json new file mode 100644 index 00000000..98b5f35e --- /dev/null +++ b/.precommit_configs/markdown-link-check.json @@ -0,0 +1,27 @@ +{ + "ignorePatterns": [ + { + "pattern": "^tg/" + }, + { + "pattern": "^http://localhost" + }, + { + "pattern": "^https://demo.kv-server.your-domain.example/" + }, + { + "pattern": "^demo.kv-server.your-domain.example:8443" + } + ], + "replacementPatterns": [ + { + "pattern": "^/", + "replacement": "{{BASEURL}}/" + } + ], + "timeout": "20s", + "retryOn429": true, + "retryCount": 5, + "fallbackRetryDelay": "30s", + "aliveStatusCodes": [200, 206] +} diff --git a/BUILD.bazel b/BUILD.bazel index 51c4a0a9..6c8c9cb4 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +load("@bazel_skylib//lib:selects.bzl", "selects") load("@bazel_skylib//rules:common_settings.bzl", "string_flag") load("@io_bazel_rules_go//go:def.bzl", "nogo") @@ -107,6 +108,69 @@ config_setting( ], ) +string_flag( + name = "mode", + build_setting_default = "prod", + values = [ + "prod", + "nonprod", + ], +) + +config_setting( + name = "prod_mode", + flag_values = { + ":mode": "prod", + }, + visibility = [ + "//components:__subpackages__", + "//tools:__subpackages__", + ], +) + +config_setting( + name = "nonprod_mode", + flag_values = { + ":mode": "nonprod", + }, + visibility = [ + "//components:__subpackages__", + "//tools:__subpackages__", + ], +) + +selects.config_setting_group( + name = "aws_prod", + match_all = [ + "//:aws_platform", + "//:prod_mode", + ], +) + +selects.config_setting_group( + name = "aws_nonprod", + match_all = [ + "//:aws_platform", + "//:nonprod_mode", + ], +) + +selects.config_setting_group( + name = "gcp_prod", + match_all = [ + "//:gcp_platform", + "//:prod_mode", + ], +) + +selects.config_setting_group( + name = "gcp_nonprod", + match_all = [ + "//:gcp_platform", + "//:nonprod_mode", + ], +) + exports_files( [".bazelversion"], ) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e93e13d..6699b468 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,116 @@ All notable changes to this project will be documented in this file. See [commit-and-tag-version](https://github.com/absolute-version/commit-and-tag-version) for commit guidelines. +## 0.16.0 (2024-04-05) + + +### Features + +* Add cache hit or miss metrics +* Add coorindator specific terraform parameters +* Add data loading prefix allowlist parameter +* Add default PAS UDF +* Add E2E latency for GetKeyValues and GetKeyValueSet in sharded lookup +* Add file groups and file group reader logic +* Add go fmt to pre-commit +* Add key prefix support to blob storage client +* Add LogContext and ConsentedDebugConfiguration proto to v2 API and internal lookup API +* Add prod and nonprod build flag +* Add request context to wrap metrics context +* Add support for configuring directory allowlist +* Add wiring for prefix allowlist (actual impl in follow up cl) +* Allow overrides for coordinators endpoints in nonprod mode +* Allow to disable v1 key not found entry in response +* Create separate metrics context map for internal lookup server +* Deprecate metrics recorder for internal lookup +* Deprecate metrics recorder for internal server +* Deprecate metrics recorder for sharded lookup +* deprecate metrics recorder for V1 server and handler +* Deprecate metrics recorder from cache +* Enable simulation system send realtime udpates +* Enable TCMalloc for KV Server and benchmarks +* Explicitly enable core dumps +* Implement deletion cutoff max timestamp per directory +* Load data files and allow notifications from configured prefix +* Load prefix files on startup and handle prefix blob notifications +* Log common request metrics +* Migrate from glog to absl log +* Partition data loading metrics by delta file name +* Pass request context from hooks to downstream components +* Pass request context to udf hooks +* Read telemetry config from cloud parameter +* Revamp AWS metrics dashboard +* Revamp GCP metrics dashboard +* Set udf_min_log_level from parameter store. +* Support content type proto for v2 api +* Support content type proto for v2 api response +* Update cache interface and blob data location to pass prefix +* Update start_after to use a map from prefix to start_after +* Use file groups for loading snapshots +* Write logs to an Otel endpoint + + +### Bug Fixes + +* Actually load all files in a snapshot file group +* **AWS:** Filter out unavailable zones. +* Correct an error in kokoro_release. +* Correct format for image tag. +* Correct typo for internal dev's service_mesh_address. +* Correct typos in GCP deployment guide. +* Crash server if default UDF fails to load. +* Delete non-active certificate before creating a new one. +* Fix filtering logic for prefixed blobs +* Fix permissions for data-loading-blob-prefix-allowlist +* Make GCP nat optional. +* Parse delta filename from notification before validating it +* Remove glog dependency for record_utils +* Remove temp dir only if it's successfully created. +* Rename class to ThreadManager +* Set retain_initial_value_of_delta_metric flag for aws metrics exporter +* Update a outdated hyperlink. +* Update common repo to pick up the AWS metrics dimension fix +* Update GCP Terraform with ability to delete unhealthy instance. +* Update tf variables to use gorekore instead of kelvingorekore +* Use blob key instead of prefixed basename + + +### GCP: Fixes + +* **GCP:** Make sure server is connected to otel collector before reaching to ready state + + +### GCP: Features + +* **GCP:** Applying Terraform pulls docker image with new tag. +* **GCP:** Make service mesh address configurable. +* **GCP:** Make subnet ip cidr configurable. +* **GCP:** Make xlb/envoy optional. + + +### Documentation + +* Add ad retrieval explainer. +* Add docs for directory support +* Add PA and PAS folders +* Add PAS developer guide +* Add public docs for file groups +* Ads retreival explainer update. + + +### Dependencies + +* **deps:** Add clang-tidy bazel config +* **deps:** Add cpp_nowarn bazel config +* **deps:** Upgrade bazel to 6.5.0 +* **deps:** Upgrade build-system to 0.55.1 +* **deps:** Upgrade build-system to 0.55.2 +* **deps:** Upgrade build-system to 0.57.0 +* **deps:** Upgrade data-plane-shared repo +* **deps:** Upgrade data-plane-shared repo to 1684674 2024-02-09 +* **deps:** Upgrade data-plane-shared-libraries to 1fbac46 +* **deps:** Upgrade pre-commit hooks + ## 0.15.0 (2024-01-23) diff --git a/README.md b/README.md index 80d8e09f..d1632784 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ versions are released. The query API conforms to the [API explainer](https://github.com/WICG/turtledove/blob/main/FLEDGE_Key_Value_Server_API.md). At the moment, to load data, instead of calling the mutation API, you would place the data as files into a location that can be directly read by the server. See more details in the -[data loading guide](/docs/loading_data.md). +[data loading guide](/docs/data_loading/loading_data.md). Currently, this service can be deployed to 1 region of your choice with more regions to be added soon. Monitoring and alerts are currently unavailable. @@ -72,18 +72,18 @@ changes. - [FLEDGE K/V server API explainer](https://github.com/WICG/turtledove/blob/main/FLEDGE_Key_Value_Server_API.md) - [FLEDGE K/V server trust model](https://github.com/privacysandbox/fledge-docs/blob/main/key_value_service_trust_model.md) - [Local server quickstart guide](/docs/developing_the_server.md) -- [AWS server user deployment documentation](/docs/deploying_on_aws.md) -- [GCP server user deployment documentation](/docs/deploying_on_gcp.md) -- [Integrating the K/V server with FLEDGE](/docs/integrating_with_fledge.md) -- [FLEDGE K/V server sharding explainer](https://github.com/privacysandbox/fledge-docs/blob/main/key_value_sharding.md) +- [AWS server user deployment documentation](/docs/deployment/deploying_on_aws.md) +- [GCP server user deployment documentation](/docs/deployment/deploying_on_gcp.md) +- [Integrating the K/V server with FLEDGE](/docs/protected_audience/integrating_with_fledge.md) +- [FLEDGE K/V server sharding explainer](https://github.com/privacysandbox/protected-auction-services-docs/blob/main/key_value_service_sharding.md) - Operating documentation - - [Data loading API and operations](/docs/loading_data.md) + - [Data loading API and operations](/docs/data_loading/loading_data.md) - [Generating and loading UDF files](/docs/generating_udf_files.md) - Error handling explainer (_to be published_) - Developer guide - [Codebase structure](/docs/repo_layout.md) - - [Working with Terraform](/production/terraform/README.md) - - [Contributing to the codebase](/docs/CONTRIBUTING.md) + - [Working with Terraform](/docs/deployment/working_with_terraform.md) + - [Contributing to the codebase](/docs/contributing.md) - [Code of conduct](/docs/CODE_OF_CONDUCT.md) - [Change log](/CHANGELOG.md) diff --git a/WORKSPACE b/WORKSPACE index 4deae703..da073316 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -13,11 +13,11 @@ python_deps("//builders/bazel") http_archive( name = "google_privacysandbox_servers_common", - # commit f198e86 2024-01-16 - sha256 = "979d467165bef950ed69bc913e9ea743bfc068714fb43cf61e02b52b910a5561", - strip_prefix = "data-plane-shared-libraries-f198e86307028ad98c38a8ed1c72189eefa97334", + # commit b34fe82 2024-04-03 + sha256 = "2afc7017723efb9d34b6ed713be03dbdf9b45de8ba585d2ea314eb3a52903d0a", + strip_prefix = "data-plane-shared-libraries-b34fe821b982e06446df617edb7a6e3041c8b0db", urls = [ - "https://github.com/privacysandbox/data-plane-shared-libraries/archive/f198e86307028ad98c38a8ed1c72189eefa97334.zip", + "https://github.com/privacysandbox/data-plane-shared-libraries/archive/b34fe821b982e06446df617edb7a6e3041c8b0db.zip", ], ) @@ -51,6 +51,20 @@ load( cpp_repositories() +http_archive( + name = "io_bazel_rules_docker", + sha256 = "b1e80761a8a8243d03ebca8845e9cc1ba6c82ce7c5179ce2b295cd36f7e394bf", + urls = ["https://github.com/bazelbuild/rules_docker/releases/download/v0.25.0/rules_docker-v0.25.0.tar.gz"], +) + +load("@io_bazel_rules_docker//repositories:repositories.bzl", container_repositories = "repositories") + +container_repositories() + +load("@io_bazel_rules_docker//repositories:deps.bzl", io_bazel_rules_docker_deps = "deps") + +io_bazel_rules_docker_deps() + load("//third_party_deps:container_deps.bzl", "container_deps") container_deps() @@ -110,14 +124,13 @@ load("//third_party_deps:python_deps.bzl", "python_repositories") python_repositories() # Load the starlark macro, which will define your dependencies. -load("@word2vec//:requirements.bzl", "install_deps") +load("@latency_benchmark//:requirements.bzl", latency_benchmark_install_deps = "install_deps") +load("@word2vec//:requirements.bzl", word2vec_install_deps = "install_deps") # Call it to define repos for your requirements. -install_deps() +latency_benchmark_install_deps() -load("//third_party_deps:rules_closure_repositories.bzl", "rules_closure_repositories") - -rules_closure_repositories() +word2vec_install_deps() # Use nogo to run `go vet` with bazel load("@io_bazel_rules_go//go:deps.bzl", "go_register_toolchains", "go_rules_dependencies") @@ -125,3 +138,15 @@ load("@io_bazel_rules_go//go:deps.bzl", "go_register_toolchains", "go_rules_depe go_rules_dependencies() go_register_toolchains(nogo = "@//:kv_nogo") + +# setup container_structure_test +http_archive( + name = "container_structure_test", + sha256 = "2da13da4c4fec9d4627d4084b122be0f4d118bd02dfa52857ff118fde88e4faa", + strip_prefix = "container-structure-test-1.16.0", + urls = ["https://github.com/GoogleContainerTools/container-structure-test/archive/v1.16.0.zip"], +) + +load("@container_structure_test//:repositories.bzl", "container_structure_test_register_toolchain") + +container_structure_test_register_toolchain(name = "cst") diff --git a/builders/CHANGELOG.md b/builders/CHANGELOG.md index ecf0d25b..e323d8a3 100644 --- a/builders/CHANGELOG.md +++ b/builders/CHANGELOG.md @@ -2,6 +2,114 @@ All notable changes to this project will be documented in this file. See [commit-and-tag-version](https://github.com/absolute-version/commit-and-tag-version) for commit guidelines. +## 0.57.0 (2024-03-10) + + +### Features + +* Add a generic pylintrc +* Add clang-tidy to build-debian + +## 0.56.0 (2024-02-29) + + +### Features + +* Add pylint to presubmit + + +### Bug Fixes + +* Clean bazel build and mod caches +* Pin okigan/awscurl to v0.29 + +## 0.55.2 (2024-02-23) + + +### Bug Fixes + +* Add gmock to .clang-format + +## 0.55.1 (2024-02-22) + + +### Bug Fixes + +* Do not invoke normalize-bazel-symlinks for cbuild --cmd + +## 0.55.0 (2024-02-22) + + +### Bug Fixes + +* Normalize bazel symlinks to a resolved path +* Pass the correct path for normalize-bazel-symlinks + +## 0.54.0 (2024-02-09) + + +### Features + +* Set cbuild workdir to pwd relative to root workspace + +## 0.53.0 (2024-01-25) + + +### Features + +* Add support to collect-coverage tool for custom lcov report + + +### Bug Fixes + +* Improve --cmd-profiler support + +## 0.52.0 (2023-12-02) + + +### Features + +* add python3.9 dev to bazel-debian + +## 0.51.0 (2023-11-30) + + +### Bug Fixes + +* Clean go build cache at the end of image build script + + +### Dependencies + +* **deps:** Upgrade bazelisk to 1.19.0 + +## 0.50.0 (2023-11-06) + + +### Features + +* Add openssh-client to build-debian image + +## 0.49.1 (2023-10-30) + + +### Bug Fixes + +* Add tools/wrk2 wrapper script + +## 0.49.0 (2023-10-27) + + +### Features + +* Add wrk2 to test-tools image +* Extend normalize-bazel-symlink to normalize within containers + + +### Dependencies + +* **deps:** Update versions in test-tools image + ## 0.48.0 (2023-10-11) diff --git a/builders/etc/.clang-format b/builders/etc/.clang-format index 8c45f0ce..3232a3be 100644 --- a/builders/etc/.clang-format +++ b/builders/etc/.clang-format @@ -4,8 +4,8 @@ DerivePointerAlignment: false SortIncludes: true IncludeBlocks: Regroup IncludeCategories: - # gtest, this should be put first in tests. - - Regex: '^ from OS. - Regex: '^<[_A-Za-z0-9-]+\.h>' diff --git a/builders/etc/.clang-tidy b/builders/etc/.clang-tidy new file mode 100644 index 00000000..2e9829d0 --- /dev/null +++ b/builders/etc/.clang-tidy @@ -0,0 +1,8 @@ +Checks: > + -*, + bugprone-*, + -bugprone-narrowing-conversions, + performance-*, + +WarningsAsErrors: '*' +... diff --git a/builders/etc/.pylintrc b/builders/etc/.pylintrc new file mode 100644 index 00000000..967bd3ee --- /dev/null +++ b/builders/etc/.pylintrc @@ -0,0 +1,118 @@ +[MESSAGES CONTROL] + +# List of checkers and warnings to enable. +enable= + indexing-exception, + old-raise-syntax, + +# List of checkers and warnings to disable. +# TODO: Shrink this list to as small as possible. +disable= + attribute-defined-outside-init, + bad-option-value, + bare-except, + broad-except, + c-extension-no-member, + design, + file-ignored, + fixme, + global-statement, + import-error, + import-outside-toplevel, + locally-disabled, + misplaced-comparison-constant, + multiple-imports, + no-self-use, + relative-import, + similarities, + suppressed-message, + ungrouped-imports, + unsubscriptable-object, + useless-object-inheritance, # Remove once all bots are on Python 3. + useless-suppression, + wrong-import-order, + wrong-import-position, + # FIXME: To be removed. Leftovers from Python 3 migration. + consider-using-with, + raise-missing-from, + super-with-arguments, + use-a-generator, + consider-using-generator, + consider-using-f-string, # Too much legacy usages. + unspecified-encoding, # Too much legacy usage. + broad-exception-raised, + +[BASIC] + +# Regular expression which should only match the name +# of functions or classes which do not require a docstring. +no-docstring-rgx=(__.*__|main) + +# Min length in lines of a function that requires a docstring. +docstring-min-length=10 + +# Regular expression which should only match correct module names. The +# leading underscore is sanctioned for private modules by Google's style +# guide. +# +# There are exceptions to the basic rule (_?[a-z][a-z0-9_]*) to cover +# requirements of Python's module system and of the presubmit framework. +module-rgx=^(_?[a-z][a-z0-9_]*)|__init__|__main__|PRESUBMIT|PRESUBMIT_unittest$ + +# Regular expression which should only match correct module level names. +const-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$ + +# Regular expression which should only match correct class attribute. +class-attribute-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$ + +# Regular expression which should only match correct class names. +class-rgx=^_?[A-Z][a-zA-Z0-9]*$ + +# Regular expression which should only match correct function names. +# 'camel_case' and 'snake_case' group names are used for consistency of naming +# styles across functions and methods. +function-rgx=^(?:(?PsetUp|tearDown|setUpModule|tearDownModule)|(?P_?[A-Z][a-zA-Z0-9]*)|(?P_?[a-z][a-z0-9_]*))$ + +# Regular expression which should only match correct method names. +# 'camel_case' and 'snake_case' group names are used for consistency of naming +# styles across functions and methods. 'exempt' indicates a name which is +# consistent with all naming styles. +method-rgx=(?x) + ^(?:(?P_[a-z0-9_]+__|runTest|setUp|tearDown|setUpTestCase + |tearDownTestCase|setupSelf|tearDownClass|setUpClass + |(test|assert)_*[A-Z0-9][a-zA-Z0-9_]*|next) + |(?P_{0,2}[A-Z][a-zA-Z0-9_]*) + |(?P_{0,2}[a-z][a-z0-9_]*))$ + +# Regular expression which should only match correct instance attribute names. +attr-rgx=^_{0,2}[a-z][a-z0-9_]*$ + +# Regular expression which should only match correct argument names. +argument-rgx=^[a-z][a-z0-9_]*$ + +# Regular expression which should only match correct variable names. +variable-rgx=^[a-z][a-z0-9_]*$ + +# Regular expression which should only match correct list comprehension / +# generator expression variable names. +inlinevar-rgx=^[a-z][a-z0-9_]*$ + +# Good variable names which should always be accepted, separated by a comma. +good-names=main,_,maxDiff + +# Bad variable names which should always be refused, separated by a comma. +bad-names= + +# FIXME: Renable this. +# List of builtins function names that should not be used, separated by a comma. +#bad-builtin=input + +[FORMAT] + +# Maximum number of characters on a single line. +max-line-length=90 + +# Maximum number of lines in a module +max-module-lines=99999 + +[TYPECHECK] diff --git a/builders/images/build-amazonlinux2/install_apps b/builders/images/build-amazonlinux2/install_apps index e0f46ccc..a5e6ff35 100755 --- a/builders/images/build-amazonlinux2/install_apps +++ b/builders/images/build-amazonlinux2/install_apps @@ -86,6 +86,11 @@ function install_clang() { clang --version } +function cleanup() { + cd / + go clean -cache -modcache +} + if [[ ${VERBOSE} -eq 1 ]]; then printf "=== SHELL ENVIRONMENT ===\n" env @@ -99,3 +104,4 @@ install_golang "${BUILD_ARCH}" install_gcc install_packer install_python +cleanup diff --git a/builders/images/build-amazonlinux2023/install_apps b/builders/images/build-amazonlinux2023/install_apps index 722e89f6..51c64377 100755 --- a/builders/images/build-amazonlinux2023/install_apps +++ b/builders/images/build-amazonlinux2023/install_apps @@ -85,6 +85,11 @@ function install_clang() { clang --version } +function cleanup() { + cd / + go clean -cache -modcache +} + if [[ ${VERBOSE} -eq 1 ]]; then printf "=== SHELL ENVIRONMENT ===\n" env @@ -98,3 +103,4 @@ install_clang install_golang "${BUILD_ARCH}" install_gcc install_packer +cleanup diff --git a/builders/images/build-debian/install_apps b/builders/images/build-debian/install_apps index c7d8e4a6..fba47ca1 100755 --- a/builders/images/build-debian/install_apps +++ b/builders/images/build-debian/install_apps @@ -44,8 +44,8 @@ function apt_update() { } function install_python() { - DEBIAN_FRONTEND=noninteractive apt-get --quiet install -y --no-install-recommends \ - python3.9-venv="3.9.*" + apt-get --quiet install -y --no-install-recommends \ + python3.9-venv="3.9.*" python3.9-dev mkdir -p /opt/bin update-alternatives \ --force \ @@ -63,19 +63,21 @@ function install_python() { } function install_misc() { - DEBIAN_FRONTEND=noninteractive apt-get --quiet install -y --no-install-recommends \ + apt-get --quiet install -y --no-install-recommends \ apt-transport-https="2.0.*" \ bsdmainutils \ ca-certificates \ chrpath="0.16-*" \ libcurl4="7.68.*" \ curl="7.68.*" \ - file="1:5*" \ + file="1:5.*" \ gettext="0.19.*" \ git="1:2.25.*" \ gnupg="2.2.*" \ + google-perftools="2.*" \ locales="2.31-*" \ lsb-release="11.1.*" \ + openssh-client="1:8.2*" \ patch="2.7.*" \ rename="1.10-*" \ software-properties-common="0.99.*" \ @@ -121,11 +123,19 @@ function install_docker() { apt-get --quiet install -y --no-install-recommends docker-ce docker-ce-cli containerd.io } -function clean_debian() { +function install_clang_tidy() { + apt-get --quiet install -y clang-tidy + printf "clang-tidy version: %s\n" "$(clang-tidy --version)" + printf "clang-tidy config: %s\n" "$(clang-tidy -dump-config)" +} + +function cleanup() { apt-get --quiet autoremove -y apt-get autoclean apt-get clean rm -rf /var/lib/apt/lists + cd / + go clean -cache -modcache } if [[ ${VERBOSE} -eq 1 ]]; then @@ -133,10 +143,13 @@ if [[ ${VERBOSE} -eq 1 ]]; then env fi +declare -x -r DEBIAN_FRONTEND=noninteractive + apt_update install_misc install_clang +install_clang_tidy install_golang "${BUILD_ARCH}" install_docker "${BUILD_ARCH}" install_python # should run after other install_* -clean_debian +cleanup diff --git a/builders/images/install_golang_apps b/builders/images/install_golang_apps index a10d33e5..9d75eb94 100755 --- a/builders/images/install_golang_apps +++ b/builders/images/install_golang_apps @@ -31,7 +31,7 @@ while [[ $# -gt 0 ]]; do done function install_bazelisk() { - go install github.com/bazelbuild/bazelisk@v1.13.2 + go install github.com/bazelbuild/bazelisk@v1.19.0 BAZELISK="$(go env GOPATH)"/bin/bazelisk if [[ -n ${BAZEL_PATH} ]] && [[ -d ${BAZEL_PATH} ]]; then ln -s "${BAZELISK}" "${BAZEL_PATH}"/bazel @@ -43,4 +43,10 @@ function install_bazelisk() { rm -rf /bazel_root/* } +function cleanup() { + cd / + go clean -cache -modcache +} + install_bazelisk +cleanup diff --git a/builders/images/presubmit/install_apps b/builders/images/presubmit/install_apps index 64808cc1..c0d8d0d8 100755 --- a/builders/images/presubmit/install_apps +++ b/builders/images/presubmit/install_apps @@ -37,10 +37,7 @@ while [[ $# -gt 0 ]]; do set -o xtrace shift ;; - -h | --help) - usage 0 - break - ;; + -h | --help) usage 0 ;; *) usage 0 ;; esac done @@ -85,7 +82,9 @@ function install_docker() { function install_precommit() { /usr/bin/python3.9 -m venv "${PRE_COMMIT_VENV_DIR}" - "${PRE_COMMIT_VENV_DIR}"/bin/pip install pre-commit~=3.1 + "${PRE_COMMIT_VENV_DIR}"/bin/pip install \ + pre-commit~=3.1 \ + pylint~=3.1.0 "${PRE_COMMIT_TOOL}" --version # initialize pre-commit cache, which needs a git repo (a temporary will suffice) @@ -102,11 +101,13 @@ function install_precommit() { rm -rf "${GIT_REPO}" } -function clean_debian() { +function cleanup() { apt-get --quiet autoremove -y apt-get autoclean apt-get clean rm -rf /var/lib/apt/lists + cd / + go clean -cache -modcache } if [[ ${VERBOSE} -eq 1 ]]; then @@ -126,4 +127,4 @@ install_packages install_golang "${ARCH}" install_docker install_precommit -clean_debian +cleanup diff --git a/builders/images/test-tools/Dockerfile b/builders/images/test-tools/Dockerfile index 21bb383a..45fab3cf 100644 --- a/builders/images/test-tools/Dockerfile +++ b/builders/images/test-tools/Dockerfile @@ -12,31 +12,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -FROM alpine:3.16 as builder +FROM alpine:3.18 as slowhttptest_builder # hadolint ignore=DL3018 -RUN apk add --no-cache build-base git openssl-dev autoconf automake +RUN apk add --no-cache autoconf automake build-base git openssl-dev WORKDIR /build -ADD https://github.com/shekyan/slowhttptest/archive/refs/tags/v1.9.0.tar.gz /build -RUN tar xz --strip-components 1 -f v1.9.0.tar.gz && ls -l && ./configure && make +ADD https://github.com/shekyan/slowhttptest/archive/refs/tags/v1.9.0.tar.gz /build/src.tar.gz +RUN tar xz --strip-components 1 -f src.tar.gz && ./configure && make -FROM golang:1.19.4-alpine3.17 AS golang -ENV BUILD_ARCH="${TARGETARCH}" \ - GOBIN=/usr/local/go/bin +FROM alpine:3.18 as wrk_builder +ARG TARGETARCH +ENV BUILD_ARCH="${TARGETARCH}" +COPY build_wrk /build/ +WORKDIR /build +ADD https://github.com/giltene/wrk2/archive/44a94c17d8e6a0bac8559b53da76848e430cb7a7.tar.gz /build/src.tar.gz +RUN /build/build_wrk + +FROM golang:1.21-alpine3.18 AS golang +ENV GOBIN=/usr/local/go/bin COPY build_golang_apps /scripts/ RUN /scripts/build_golang_apps -FROM fullstorydev/grpcurl:v1.8.7 AS grpcurl -FROM alpine:3.17.2 +FROM fullstorydev/grpcurl:v1.8.9-alpine AS grpcurl +FROM alpine:3.18 COPY --from=golang /usr/local/go/bin/* /usr/local/bin/ COPY --from=grpcurl /bin/grpcurl /usr/local/bin/ - ARG TARGETARCH ENV BUILD_ARCH="${TARGETARCH}" \ PATH="${PATH}:/usr/local/go/bin" \ GOBIN=/usr/local/go/bin - COPY install_apps /scripts/ - RUN /scripts/install_apps -COPY --from=builder /build/src/slowhttptest /usr/bin/ +COPY --from=slowhttptest_builder /build/src/slowhttptest /usr/bin/ +COPY --from=wrk_builder /build/wrk /usr/bin/wrk2 diff --git a/builders/images/test-tools/build_wrk b/builders/images/test-tools/build_wrk new file mode 100755 index 00000000..d843dd06 --- /dev/null +++ b/builders/images/test-tools/build_wrk @@ -0,0 +1,27 @@ +#!/bin/busybox sh + +set -o errexit +set -o xtrace + +install_no_wrk() { + cat </build/wrk +#!/bin/busybox sh +PROGRAM_NAME="\$(basename \$0)" +printf "Error: %s not supported on Aarch64\n" "\${PROGRAM_NAME}" +printf "build_arch: %s\n" "${BUILD_ARCH}" >/dev/stderr +exit 1 +EOF + chmod 755 /build/wrk +} + +install_wrk() { + apk add --no-cache build-base git openssl-dev zlib-dev + tar xz --strip-components 1 -f src.tar.gz + make -j6 +} + +if [[ ${BUILD_ARCH} == arm64 ]]; then + install_no_wrk +else + install_wrk +fi diff --git a/builders/tests/data/hashes/build-amazonlinux2 b/builders/tests/data/hashes/build-amazonlinux2 index 770d4751..7aed7b0f 100644 --- a/builders/tests/data/hashes/build-amazonlinux2 +++ b/builders/tests/data/hashes/build-amazonlinux2 @@ -1 +1 @@ -1feb61b0cf40e2797fc6a3425c3e50ce928ca194f3338d9a7155e51a4917312a +3efa00f3a5dbe0a4708be523aa32aca91dcd56d403d3ff32e0202756b8321b3b diff --git a/builders/tests/data/hashes/build-amazonlinux2023 b/builders/tests/data/hashes/build-amazonlinux2023 index 036c5cd3..1bcc412e 100644 --- a/builders/tests/data/hashes/build-amazonlinux2023 +++ b/builders/tests/data/hashes/build-amazonlinux2023 @@ -1 +1 @@ -a1c165d019f546e1fcad06506da3b274094b61a9add5576bf3fed6e5a45fb8d4 +57396ff1c765f7b63905963cfe4498912f7f75b5cb9f7bc36bd6879af69872e7 diff --git a/builders/tests/data/hashes/build-debian b/builders/tests/data/hashes/build-debian index 1e7e375d..c89114f2 100644 --- a/builders/tests/data/hashes/build-debian +++ b/builders/tests/data/hashes/build-debian @@ -1 +1 @@ -a10495fe1e23472aea6770b0c9760e1cdfc011a4883ff5fb3d488774e6995ae6 +38cc8a23a6a56eb6567bef3685100cd3be1c0491dcc8b953993c42182da3fa40 diff --git a/builders/tests/data/hashes/presubmit b/builders/tests/data/hashes/presubmit index 8e2b4d7c..a35c6c86 100644 --- a/builders/tests/data/hashes/presubmit +++ b/builders/tests/data/hashes/presubmit @@ -1 +1 @@ -40b437991cd3ec239fc0d0eef162005dc62784def9c40548f3ea235d99c2e5a0 +d9dab1c798d51f79e68fd8eb3bb83312086808d789bbc09d0f2dbf708ef5f114 diff --git a/builders/tests/data/hashes/test-tools b/builders/tests/data/hashes/test-tools index 72a3fa38..fc5e0b5c 100644 --- a/builders/tests/data/hashes/test-tools +++ b/builders/tests/data/hashes/test-tools @@ -1 +1 @@ -e19aaa4e08668be8056a6064e27be159ee7e158b3a90d75b7c4e5483369ebe41 +dd1ec6137d4dd22fec555044cd85f484adfa6c7b686880ea5449cff936bad34e diff --git a/builders/tests/run-tests b/builders/tests/run-tests index f98389a4..037067d2 100755 --- a/builders/tests/run-tests +++ b/builders/tests/run-tests @@ -243,7 +243,9 @@ fi generate_hashes_cbuild # CLI tests - +if [[ ${VERBOSE} -eq 1 ]]; then + set -o xtrace +fi cli_tests_misc for img in ${IMAGE_LIST}; do cli_tests_test-tools "${img}" diff --git a/builders/tools/awscurl b/builders/tools/awscurl index 355d58e1..975dbb17 100755 --- a/builders/tools/awscurl +++ b/builders/tools/awscurl @@ -29,8 +29,21 @@ set -o errexit +TOOLS_DIR="$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" +readonly TOOLS_DIR + +function check_arch() { + local -r ARCH="$("${TOOLS_DIR}"/get-architecture)" + if [[ ${ARCH} == arm64 ]]; then + printf "architecture %s not supported\n" "${ARCH}" &>/dev/stderr + exit 0 + fi +} + +check_arch + # shellcheck disable=SC1090 -source "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")"/builder.sh +source "${TOOLS_DIR}"/builder.sh declare -a ENV_VARS builder::add_aws_env_vars ENV_VARS @@ -68,4 +81,4 @@ docker run \ --volume "${HOME}"/.aws/:/home/.aws/ \ --volume "${WORKSPACE}":/src/workspace \ --workdir /src/workspace/"${REL_PWD}" \ - ghcr.io/okigan/awscurl:latest "$@" + ghcr.io/okigan/awscurl:v0.29 "$@" diff --git a/builders/tools/cbuild b/builders/tools/cbuild index f0e4255f..c40a3aed 100755 --- a/builders/tools/cbuild +++ b/builders/tools/cbuild @@ -134,13 +134,6 @@ if ! [[ " ${IMAGE_LIST[*]} " =~ " ${IMAGE} " ]]; then usage 1 fi -if [[ ${WITH_CMD_PROFILER} -eq 1 ]]; then - if [[ ${IMAGE} != build-debian ]]; then - printf "error: --cmd-profiler is only compatible with build-debian\n" &>/dev/stderr - usage 1 - fi - CMD_PROFILER="LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libprofiler.so " -fi TOOLS_DIR="$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" readonly TOOLS_DIR # shellcheck disable=SC1090 @@ -160,6 +153,9 @@ if [[ ${VERBOSE} -eq 1 ]]; then printf "mounting workspace into container: %s\n" "${WORKSPACE_MOUNT}" &>/dev/stderr fi +TOOLS_RELDIR="$(realpath "${TOOLS_DIR}" --relative-to="${PWD}")" +readonly TOOLS_RELDIR + if [[ -n ${CBUILD_IMAGE} ]]; then IMAGE_TAGGED=${CBUILD_IMAGE} else @@ -167,12 +163,20 @@ else fi readonly IMAGE_TAGGED +PWD_WORKSPACE_REL_PATH="$(realpath --relative-base="${WORKSPACE}" "${PWD}")" +readonly PWD_WORKSPACE_REL_PATH +WORKDIR=/src/workspace +if [[ ${PWD_WORKSPACE_REL_PATH:0:1} != / ]]; then + WORKDIR="/src/workspace/${PWD_WORKSPACE_REL_PATH}" +fi +readonly WORKDIR + declare -a DOCKER_RUN_ARGS DOCKER_RUN_ARGS+=( "--rm" "--entrypoint=/bin/bash" "--volume=${WORKSPACE_MOUNT}:/src/workspace" - "--workdir=/src/workspace" + "--workdir=${WORKDIR}" "--network=${DOCKER_NETWORK}" "$(echo "${EXTRA_DOCKER_RUN_ARGS}" | envsubst)" ) @@ -181,6 +185,17 @@ if [[ ${DOCKER_SECCOMP_UNCONFINED} -eq 1 ]]; then DOCKER_RUN_ARGS+=("--security-opt=seccomp=unconfined") fi +if [[ ${WITH_CMD_PROFILER} -eq 1 ]]; then + if [[ ${IMAGE} != build-debian ]]; then + printf "error: --cmd-profiler is only compatible with build-debian\n" &>/dev/stderr + usage 1 + fi + DOCKER_RUN_ARGS+=( + "--env=CMD_PROFILER=LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libprofiler.so" + "--env=CPUPROFILE=benchmark.prof" + ) +fi + # inside the docker build images, /bazel_root is the bazel cache dir, per the system-wide bazelrc readonly BAZEL_ROOT=/bazel_root if [[ ${WITH_SHARED_CACHE} -eq 0 ]]; then @@ -222,10 +237,16 @@ if [[ -z ${CMD} ]]; then ${DOCKER_RUN_ARGS[@]} \ "${IMAGE_TAGGED}" \ --login +elif [[ ${WITH_CMD_PROFILER} -eq 1 ]]; then + # shellcheck disable=SC2068 + docker run \ + ${DOCKER_RUN_ARGS[@]} \ + "${IMAGE_TAGGED}" \ + --login -c "'${TOOLS_RELDIR}'/normalize-bazel-symlinks; env \${CMD_PROFILER} ${CMD}" else # shellcheck disable=SC2068 docker run \ ${DOCKER_RUN_ARGS[@]} \ "${IMAGE_TAGGED}" \ - --login -c "${CMD_PROFILER}${CMD}" + --login -c "${CMD}" fi diff --git a/builders/tools/collect-coverage b/builders/tools/collect-coverage index 2083fbdb..746185c2 100755 --- a/builders/tools/collect-coverage +++ b/builders/tools/collect-coverage @@ -14,12 +14,50 @@ # See the License for the specific language governing permissions and # limitations under the License. -# environment variables (all optional): +# environment variables: # WORKSPACE Set the path to the workspace (repo root) set -o pipefail set -o errexit +declare LCOV_REPORT="${WORKSPACE}/bazel-out/_coverage/_coverage_report.dat" +declare COVERAGE_FILENAME=coverage.zip + +function usage() { + local exitval=${1-1} + cat &>/dev/stderr << USAGE +usage: + $0 + --lcov_report path to lcov report relative to the WORKSPACE + --coverage_output_filename name of ZIP file that will contain artifacts from coverage collection +USAGE + # shellcheck disable=SC2086 + exit ${exitval} +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --lcov_report) + LCOV_REPORT="${WORKSPACE}/$2" + shift 2 || usage + ;; + --coverage_output_filename) + COVERAGE_FILENAME="$2" + shift 2 || usage + if [[ ${COVERAGE_FILENAME##*.} != zip ]]; then + printf "error: --coverage_output_filename must be a ZIP file\n" &>/dev/stderr + exit 1 + fi + ;; + -h | --help) + usage 0 + ;; + *) + usage + ;; + esac +done + trap _cleanup EXIT function _cleanup() { declare -r -i status=$? @@ -33,10 +71,9 @@ function _cleanup() { function generate_coverage_report() { local -r cov_dir="$(mktemp --tmpdir="${WORKSPACE}" --directory coverage-XXXX)" trap 'rm -rf "${cov_dir}"' RETURN EXIT - local -r cov_dat="${WORKSPACE}/bazel-out/_coverage/_coverage_report.dat" - cp "${cov_dat}" "${cov_dir}" + cp "${LCOV_REPORT}" "${cov_dir}"/_coverage_report.dat local -r dist_dir="${WORKSPACE}"/dist - cp "${cov_dat}" "${dist_dir}" + cp "${LCOV_REPORT}" "${dist_dir}"/_coverage_report.dat chmod -x {"${cov_dir}","${dist_dir}"}/_coverage_report.dat "${TOOLS_DIR}"/lcov --list dist/_coverage_report.dat >"${dist_dir}"/coverage_report.txt @@ -64,10 +101,5 @@ source "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")"/builder.sh TOOLS_DIR="$(builder::get_tools_dir)" readonly TOOLS_DIR -declare COVERAGE_FILENAME="$1" -if [[ ${COVERAGE_FILENAME##*.} != zip ]]; then - COVERAGE_FILENAME=coverage.zip -fi - generate_coverage_report "${TOOLS_DIR}"/normalize-dist diff --git a/builders/tools/normalize-bazel-symlinks b/builders/tools/normalize-bazel-symlinks index a66f5f34..8506ac96 100755 --- a/builders/tools/normalize-bazel-symlinks +++ b/builders/tools/normalize-bazel-symlinks @@ -12,11 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -if [[ -f /.dockerenv ]]; then - printf "Running inside Docker container? This script is only designed to be executed outside docker\n" &>/dev/stderr - exit 1 -fi - declare -r BAZEL_CACHE_DIR="${HOME}/.cache/bazel" function normalize_symlink() { @@ -26,11 +21,28 @@ function normalize_symlink() { return fi local -r link_path="$(readlink "${link_name}")" - local -r output_user_root="${link_path///bazel_root\/}" + local -r output_user_root="${link_path///bazel_root/}" + rm -f "${link_name}" + ln -s "$(realpath "${BAZEL_CACHE_DIR}/${output_user_root}")" "${link_name}" +} + +function normalize_symlink_docker() { + declare -r link_name="$1" + if readlink --canonicalize-existing "${link_name}" &>/dev/null ; then + printf "symlink %s resolves fully, skipping\n" "${link_name}" + return + fi + local -r link_path="$(readlink "${link_name}")" + local -r output_user_root="${link_path##*/.cache/bazel/}" rm -f "${link_name}" - ln -s "${BAZEL_CACHE_DIR}/${output_user_root}" "${link_name}" + ln -s "$(realpath "/bazel_root/${output_user_root}")" "${link_name}" } +declare _normalize_fn=normalize_symlink +if [[ -f /.dockerenv ]]; then + _normalize_fn=normalize_symlink_docker +fi + declare -a -r LINK_DIRS=( bazel-bin bazel-out @@ -38,5 +50,7 @@ declare -a -r LINK_DIRS=( bazel-workspace ) for link in "${LINK_DIRS[@]}"; do - normalize_symlink "${link}" + if [[ -L ${link} ]]; then + ${_normalize_fn} "${link}" + fi done diff --git a/builders/tools/pre-commit b/builders/tools/pre-commit index be8c5c7c..7bfc7058 100755 --- a/builders/tools/pre-commit +++ b/builders/tools/pre-commit @@ -96,14 +96,13 @@ function __init() { export CBUILD_IMAGE="${IMAGE_TAGGED}" local -r ARCH="$("${TOOLS_DIR}"/get-architecture)" + SKIP_HOOKS="${SKIP}" if [[ ${ARCH} == arm64 ]]; then if [[ -z ${SKIP} ]]; then SKIP_HOOKS="terraform-fmt" else - SKIP_HOOKS="${SKIP},terraform-fmt" + SKIP_HOOKS+=",terraform-fmt" fi - else - SKIP_HOOKS="${SKIP}" fi if [[ -n ${SKIP_HOOKS} ]]; then printf "Skipping pre-commit hooks: %s\n" "${SKIP_HOOKS}" diff --git a/builders/tools/wrk2 b/builders/tools/wrk2 new file mode 120000 index 00000000..d609abda --- /dev/null +++ b/builders/tools/wrk2 @@ -0,0 +1 @@ +test-tool \ No newline at end of file diff --git a/builders/version.txt b/builders/version.txt index fdae41d2..78756de3 100644 --- a/builders/version.txt +++ b/builders/version.txt @@ -1 +1 @@ -0.48.0 \ No newline at end of file +0.57.0 \ No newline at end of file diff --git a/components/cloud_config/BUILD.bazel b/components/cloud_config/BUILD.bazel index 04ce12eb..8cd15ad6 100644 --- a/components/cloud_config/BUILD.bazel +++ b/components/cloud_config/BUILD.bazel @@ -30,10 +30,10 @@ cc_library( ], deps = [ "//public:constants", - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:marshalling", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -49,14 +49,14 @@ cc_library( "parameter_client.h", ], deps = [ - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//scp/cc/public/core/interface:errors", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/interface/parameter_client", + "@google_privacysandbox_servers_common//src/public/core/interface:errors", + "@google_privacysandbox_servers_common//src/public/cpio/interface/parameter_client", ], ) @@ -82,13 +82,13 @@ cc_test( ], "//:gcp_platform": [ "@com_google_googletest//:gtest", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/interface:cpio", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/mock/parameter_client:parameter_client_mock", + "@google_privacysandbox_servers_common//src/public/cpio/interface:cpio", + "@google_privacysandbox_servers_common//src/public/cpio/mock/parameter_client:parameter_client_mock", ], "//:local_platform": [ "//components/data/common:mocks", "//components/util:sleepfor_mock", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_github_grpc_grpc//:grpc++", ], }) + [ @@ -114,13 +114,13 @@ cc_library( "@aws_sdk_cpp//:ssm", ], "//:gcp_platform": [ - "@google_privacysandbox_servers_common//scp/cc/public/core/interface:errors", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/interface:cpio", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/interface/parameter_client", + "@google_privacysandbox_servers_common//src/public/core/interface:errors", + "@google_privacysandbox_servers_common//src/public/cpio/interface:cpio", + "@google_privacysandbox_servers_common//src/public/cpio/interface/parameter_client", ], "//conditions:default": [], }) + [ - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", ], @@ -147,21 +147,22 @@ cc_library( "//components/errors:gcp_error_util", "@com_github_googleapis_google_cloud_cpp//:compute_instances", "@com_google_absl//absl/flags:flag", - "@com_google_absl//absl/log:check", "@com_google_absl//absl/synchronization", - "@google_privacysandbox_servers_common//scp/cc/public/core/interface:execution_result", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/interface/instance_client", + "@google_privacysandbox_servers_common//src/public/core/interface:execution_result", + "@google_privacysandbox_servers_common//src/public/cpio/interface/instance_client", ], "//:local_instance": [ "@com_google_absl//absl/flags:flag", ], "//conditions:default": [], }) + [ - "@com_github_google_glog//:glog", + "//components/errors:error_tag", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) diff --git a/components/cloud_config/instance_client_aws.cc b/components/cloud_config/instance_client_aws.cc index 77d70b73..0191e891 100644 --- a/components/cloud_config/instance_client_aws.cc +++ b/components/cloud_config/instance_client_aws.cc @@ -16,6 +16,8 @@ #include #include +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" @@ -37,12 +39,17 @@ #include "aws/ec2/model/DescribeTagsResponse.h" #include "aws/ec2/model/Filter.h" #include "components/cloud_config/instance_client.h" +#include "components/errors/error_tag.h" #include "components/errors/error_util_aws.h" -#include "glog/logging.h" namespace kv_server { namespace { +enum class ErrorTag : int { + kGetAwsHttpResourceError = 1, + kAutoScalingSizeError = 2 +}; + using Aws::AutoScaling::Model::DescribeAutoScalingGroupsRequest; using Aws::AutoScaling::Model::Instance; using Aws::AutoScaling::Model::LifecycleState; @@ -93,8 +100,10 @@ absl::StatusOr GetAwsHttpResource( if (result.GetResponseCode() == Aws::Http::HttpResponseCode::OK) { return Aws::Utils::StringUtils::Trim(result.GetPayload().c_str()); } - return absl::Status(HttpResponseCodeToStatusCode(result.GetResponseCode()), - "Failed to get AWS Http resource."); + return StatusWithErrorTag( + absl::Status(HttpResponseCodeToStatusCode(result.GetResponseCode()), + "Failed to get AWS Http resource."), + __FILE__, ErrorTag::kGetAwsHttpResourceError); } absl::StatusOr GetImdsToken( @@ -142,7 +151,8 @@ absl::StatusOr GetAutoScalingGroupName( "Could not get auto scaling instances for instance ", instance_id, ". Retrieved ", outcome.GetResult().GetAutoScalingInstances().size(), " auto scaling groups."); - return absl::NotFoundError(error_msg); + return StatusWithErrorTag(absl::NotFoundError(error_msg), __FILE__, + ErrorTag::kAutoScalingSizeError); } return outcome.GetResult() .GetAutoScalingInstances()[0] diff --git a/components/cloud_config/instance_client_gcp.cc b/components/cloud_config/instance_client_gcp.cc index cd622517..4839204e 100644 --- a/components/cloud_config/instance_client_gcp.cc +++ b/components/cloud_config/instance_client_gcp.cc @@ -19,6 +19,7 @@ #include "absl/flags/flag.h" #include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" @@ -26,11 +27,10 @@ #include "absl/synchronization/notification.h" #include "components/cloud_config/instance_client.h" #include "components/errors/error_util_gcp.h" -#include "glog/logging.h" #include "google/cloud/compute/instances/v1/instances_client.h" -#include "scp/cc/public/core/interface/execution_result.h" -#include "scp/cc/public/cpio/interface/instance_client/instance_client_interface.h" -#include "src/cpp/util/status_macro/status_macros.h" +#include "src/public/core/interface/execution_result.h" +#include "src/public/cpio/interface/instance_client/instance_client_interface.h" +#include "src/util/status_macro/status_macros.h" ABSL_FLAG(std::string, shard_num, "0", "Shard number."); @@ -275,7 +275,7 @@ class GcpInstanceClient : public InstanceClient { if (result.Successful()) { resource_name = std::string{response.instance_resource_name()}; } else { - LOG(ERROR) << "Faild to get instance resource name: " + LOG(ERROR) << "Failed to get instance resource name: " << GetErrorMessage(result.status_code); } diff --git a/components/cloud_config/instance_client_local.cc b/components/cloud_config/instance_client_local.cc index 510565f8..32a30f19 100644 --- a/components/cloud_config/instance_client_local.cc +++ b/components/cloud_config/instance_client_local.cc @@ -17,10 +17,10 @@ #include #include "absl/flags/flag.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "components/cloud_config/instance_client.h" -#include "glog/logging.h" ABSL_FLAG(std::string, environment, "local", "Environment name."); ABSL_FLAG(std::string, shard_num, "0", "Shard number."); diff --git a/components/cloud_config/parameter_client_aws.cc b/components/cloud_config/parameter_client_aws.cc index 2a5655ac..bb13143a 100644 --- a/components/cloud_config/parameter_client_aws.cc +++ b/components/cloud_config/parameter_client_aws.cc @@ -20,6 +20,7 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/numbers.h" @@ -31,7 +32,6 @@ #include "aws/ssm/model/GetParameterResult.h" #include "components/cloud_config/parameter_client.h" #include "components/errors/error_util_aws.h" -#include "glog/logging.h" namespace kv_server { namespace { diff --git a/components/cloud_config/parameter_client_gcp.cc b/components/cloud_config/parameter_client_gcp.cc index ba9762f1..a5aee5e4 100644 --- a/components/cloud_config/parameter_client_gcp.cc +++ b/components/cloud_config/parameter_client_gcp.cc @@ -20,16 +20,17 @@ #include "absl/container/flat_hash_map.h" #include "absl/flags/flag.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "absl/synchronization/blocking_counter.h" #include "components/cloud_config/parameter_client.h" -#include "glog/logging.h" -#include "scp/cc/public/core/interface/errors.h" -#include "scp/cc/public/core/interface/execution_result.h" -#include "scp/cc/public/cpio/interface/parameter_client/parameter_client_interface.h" +#include "src/public/core/interface/errors.h" +#include "src/public/core/interface/execution_result.h" +#include "src/public/cpio/interface/parameter_client/parameter_client_interface.h" namespace kv_server { namespace { @@ -106,6 +107,8 @@ class GcpParameterClient : public ParameterClient { << " with error: " << status; return status; } + LOG(INFO) << "Got parameter: " << parameter_name + << " with value: " << param_value; return param_value; } diff --git a/components/cloud_config/parameter_client_gcp_test.cc b/components/cloud_config/parameter_client_gcp_test.cc index 7e076998..1f81c585 100644 --- a/components/cloud_config/parameter_client_gcp_test.cc +++ b/components/cloud_config/parameter_client_gcp_test.cc @@ -22,9 +22,9 @@ #include "components/cloud_config/parameter_client.h" #include "gtest/gtest.h" -#include "scp/cc/public/cpio/interface/error_codes.h" -#include "scp/cc/public/cpio/interface/parameter_client/parameter_client_interface.h" -#include "scp/cc/public/cpio/mock/parameter_client/mock_parameter_client.h" +#include "src/public/cpio/interface/error_codes.h" +#include "src/public/cpio/interface/parameter_client/parameter_client_interface.h" +#include "src/public/cpio/mock/parameter_client/mock_parameter_client.h" namespace kv_server { namespace { diff --git a/components/cloud_config/parameter_client_local.cc b/components/cloud_config/parameter_client_local.cc index 7fb9882f..ea8bf5f4 100644 --- a/components/cloud_config/parameter_client_local.cc +++ b/components/cloud_config/parameter_client_local.cc @@ -51,11 +51,22 @@ ABSL_FLAG(bool, route_v1_to_v2, false, ABSL_FLAG(std::string, data_loading_file_format, std::string(kv_server::kFileFormats[static_cast( kv_server::FileFormat::kRiegeli)]), - "File format of the input data files."); + "File format of the input data files. See /public/constants.h for " + "possible values."); ABSL_FLAG(std::int32_t, logging_verbosity_level, 0, "Loggging verbosity level."); ABSL_FLAG(absl::Duration, udf_timeout, absl::Seconds(5), "Timeout for one UDF invocation"); +ABSL_FLAG(int32_t, udf_min_log_level, 0, + "Minimum logging level for UDFs. Info=0, Warn=1, Error=2. Default is " + "0(info)."); +ABSL_FLAG(bool, enable_otel_logger, false, "Whether to enable otel logger."); +ABSL_FLAG(std::string, telemetry_config, "mode: EXPERIMENT", + "Telemetry configuration for exporting raw or noised metrics"); +ABSL_FLAG(std::string, data_loading_prefix_allowlist, "", + "Allowlist for blob prefixes."); +ABSL_FLAG(bool, add_missing_keys_v1, false, + "Whether to add missing keys for v1."); namespace kv_server { namespace { @@ -79,6 +90,11 @@ class LocalParameterClient : public ParameterClient { absl::GetFlag(FLAGS_realtime_directory)}); string_flag_values_.insert({"kv-server-local-data-loading-file-format", absl::GetFlag(FLAGS_data_loading_file_format)}); + string_flag_values_.insert({"kv-server-local-telemetry-config", + absl::GetFlag(FLAGS_telemetry_config)}); + string_flag_values_.insert( + {"kv-server-local-data-loading-blob-prefix-allowlist", + absl::GetFlag(FLAGS_data_loading_prefix_allowlist)}); // Insert more string flag values here. int32_t_flag_values_.insert( @@ -111,13 +127,19 @@ class LocalParameterClient : public ParameterClient { int32_t_flag_values_.insert( {"kv-server-local-udf-timeout-millis", absl::ToInt64Milliseconds(absl::GetFlag(FLAGS_udf_timeout))}); + int32_t_flag_values_.insert({"kv-server-local-udf-min-log-level", + absl::GetFlag(FLAGS_udf_min_log_level)}); // Insert more int32 flag values here. bool_flag_values_.insert({"kv-server-local-route-v1-to-v2", absl::GetFlag(FLAGS_route_v1_to_v2)}); + bool_flag_values_.insert({"kv-server-local-add-missing-keys-v1", + absl::GetFlag(FLAGS_add_missing_keys_v1)}); bool_flag_values_.insert({"kv-server-local-use-real-coordinators", false}); bool_flag_values_.insert( {"kv-server-local-use-external-metrics-collector-endpoint", false}); bool_flag_values_.insert({"kv-server-local-use-sharding-key-regex", false}); + bool_flag_values_.insert({"kv-server-local-enable-otel-logger", + absl::GetFlag(FLAGS_enable_otel_logger)}); // Insert more bool flag values here. } diff --git a/components/cloud_config/parameter_client_local_test.cc b/components/cloud_config/parameter_client_local_test.cc index 839e8a4f..62393754 100644 --- a/components/cloud_config/parameter_client_local_test.cc +++ b/components/cloud_config/parameter_client_local_test.cc @@ -104,12 +104,24 @@ TEST(ParameterClientLocal, ExpectedFlagDefaultsArePresent) { ASSERT_TRUE(statusor.ok()); EXPECT_EQ(5000, *statusor); } + { + const auto statusor = + client->GetInt32Parameter("kv-server-local-udf-min-log-level"); + ASSERT_TRUE(statusor.ok()); + EXPECT_EQ(0, *statusor); + } { const auto statusor = client->GetBoolParameter("kv-server-local-route-v1-to-v2"); ASSERT_TRUE(statusor.ok()); EXPECT_EQ(false, *statusor); } + { + const auto statusor = + client->GetBoolParameter("kv-server-local-add-missing-keys-v1"); + ASSERT_TRUE(statusor.ok()); + EXPECT_EQ(false, *statusor); + } { const auto statusor = client->GetBoolParameter("kv-server-local-use-real-coordinators"); @@ -122,6 +134,18 @@ TEST(ParameterClientLocal, ExpectedFlagDefaultsArePresent) { ASSERT_TRUE(statusor.ok()); EXPECT_EQ(false, *statusor); } + { + const auto statusor = + client->GetBoolParameter("kv-server-local-enable-otel-logger"); + ASSERT_TRUE(statusor.ok()); + EXPECT_EQ(false, *statusor); + } + { + const auto statusor = + client->GetParameter("kv-server-local-telemetry-config"); + ASSERT_TRUE(statusor.ok()); + EXPECT_EQ("mode: EXPERIMENT", *statusor); + } } } // namespace diff --git a/components/data/blob_storage/BUILD.bazel b/components/data/blob_storage/BUILD.bazel index 3cb4deeb..3597529b 100644 --- a/components/data/blob_storage/BUILD.bazel +++ b/components/data/blob_storage/BUILD.bazel @@ -25,12 +25,12 @@ cc_library( hdrs = ["seeking_input_streambuf.h"], deps = [ "//components/telemetry:server_definition", - "@com_github_google_glog//:glog", "@com_google_absl//absl/base", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", ], ) @@ -43,8 +43,31 @@ cc_test( ":seeking_input_streambuf", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", + ], +) + +cc_library( + name = "blob_prefix_allowlist", + srcs = ["blob_prefix_allowlist.cc"], + hdrs = ["blob_prefix_allowlist.h"], + deps = [ + "//public/data_loading:filename_utils", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", + ], +) + +cc_test( + name = "blob_prefix_allowlist_test", + size = "small", + srcs = ["blob_prefix_allowlist_test.cc"], + deps = [ + ":blob_prefix_allowlist", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", ], ) @@ -72,8 +95,9 @@ cc_library( ], "//conditions:default": [], }) + [ + ":blob_prefix_allowlist", ":seeking_input_streambuf", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -109,7 +133,6 @@ cc_test( "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/flags:flag", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -160,10 +183,9 @@ cc_test( ], }) + [ ":blob_storage_change_notifier", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -176,6 +198,7 @@ cc_library( "delta_file_notifier.h", ], deps = [ + ":blob_prefix_allowlist", ":blob_storage_change_notifier", ":blob_storage_client", "//components/data/common:thread_manager", @@ -183,13 +206,14 @@ cc_library( "//components/util:sleepfor", "//public:constants", "//public/data_loading:filename_utils", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/functional:bind_front", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/synchronization", - "@google_privacysandbox_servers_common//src/cpp/util:duration", + "@google_privacysandbox_servers_common//src/util:duration", ], ) diff --git a/components/data/blob_storage/blob_prefix_allowlist.cc b/components/data/blob_storage/blob_prefix_allowlist.cc new file mode 100644 index 00000000..1a0e0b23 --- /dev/null +++ b/components/data/blob_storage/blob_prefix_allowlist.cc @@ -0,0 +1,68 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "components/data/blob_storage/blob_prefix_allowlist.h" + +#include +#include + +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" + +namespace kv_server { +namespace { +constexpr std::string_view kBlobNameDelimiter = "/"; +constexpr std::string_view kPrefixListDelimiter = ","; +} // namespace + +BlobPrefixAllowlist::BlobName ParseBlobName(std::string_view blob_name) { + if (blob_name.empty()) { + return BlobPrefixAllowlist::BlobName{}; + } + std::string blob_name_copy(blob_name); + std::reverse(blob_name_copy.begin(), blob_name_copy.end()); + std::vector name_parts = absl::StrSplit( + blob_name_copy, + absl::MaxSplits(/*delimiter=*/kBlobNameDelimiter, /*limit=*/1)); + for (auto& name_part : name_parts) { + std::reverse(name_part.begin(), name_part.end()); + } + auto prefix = name_parts.size() == 1 ? "" : std::move(name_parts.back()); + return BlobPrefixAllowlist::BlobName{.prefix = std::move(prefix), + .key = std::move(name_parts.front())}; +} + +BlobPrefixAllowlist::BlobPrefixAllowlist(std::string_view allowed_prefixes) { + std::vector prefixes = + absl::StrSplit(allowed_prefixes, kPrefixListDelimiter); + // We always allow reading blobs at the bucket level. + allowed_prefixes_.insert(std::string(kDefaultBlobPrefix)); + allowed_prefixes_.insert(prefixes.begin(), prefixes.end()); +} + +bool BlobPrefixAllowlist::Contains(std::string_view prefix) const { + return allowed_prefixes_.contains(prefix); +} + +bool BlobPrefixAllowlist::ContainsBlobPrefix(std::string_view blob_name) const { + return Contains(ParseBlobName(blob_name).prefix); +} + +const absl::flat_hash_set& BlobPrefixAllowlist::Prefixes() const { + return allowed_prefixes_; +} + +} // namespace kv_server diff --git a/components/data/blob_storage/blob_prefix_allowlist.h b/components/data/blob_storage/blob_prefix_allowlist.h new file mode 100644 index 00000000..52ccad78 --- /dev/null +++ b/components/data/blob_storage/blob_prefix_allowlist.h @@ -0,0 +1,66 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef COMPONENTS_DATA_BLOB_STORAGE_BLOB_PREFIX_ALLOWLIST_H_ +#define COMPONENTS_DATA_BLOB_STORAGE_BLOB_PREFIX_ALLOWLIST_H_ + +#include +#include + +#include "absl/container/flat_hash_set.h" + +namespace kv_server { + +// Default prefix for blobs at the bucket level. +constexpr std::string_view kDefaultBlobPrefix = ""; + +// List of blob prefixes that are allowlisted for data loading. Blobs with +// prefix not included in this are ignored. +class BlobPrefixAllowlist { + public: + struct BlobName { + std::string prefix; + std::string key; + }; + + BlobPrefixAllowlist() : BlobPrefixAllowlist("") {} + explicit BlobPrefixAllowlist(std::string_view allowed_prefixes); + + // Returns true if `prefix` is allowlisted, meaning that blobs with this + // prefix are eligible for data loading. + [[nodiscard]] bool Contains(std::string_view prefix) const; + // Returns true if `blob_name`'s is allowlisted, meaning that the blob is + // eligible for data loading. + [[nodiscard]] bool ContainsBlobPrefix(std::string_view blob_name) const; + // Returns the set of prefixes contained in this list. + [[nodiscard]] const absl::flat_hash_set& Prefixes() const; + + private: + absl::flat_hash_set allowed_prefixes_; +}; + +// Parses a `blob_name` into it's corresponding name parts. +// +// For example: +// (1) blob_name="prefix1/DELTA_1705430864435450" => {.prefix="prefix1", +// .key="DELTA_1705430864435450"} +// (2) blob_name="DELTA_1705430864435450" => {.prefix="", +// .key="DELTA_1705430864435450"} +BlobPrefixAllowlist::BlobName ParseBlobName(std::string_view blob_name); + +} // namespace kv_server + +#endif // COMPONENTS_DATA_BLOB_STORAGE_BLOB_PREFIX_ALLOWLIST_H_ diff --git a/components/data/blob_storage/blob_prefix_allowlist_test.cc b/components/data/blob_storage/blob_prefix_allowlist_test.cc new file mode 100644 index 00000000..8ebcdea2 --- /dev/null +++ b/components/data/blob_storage/blob_prefix_allowlist_test.cc @@ -0,0 +1,62 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "components/data/blob_storage/blob_prefix_allowlist.h" + +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +namespace kv_server { +namespace { + +using testing::UnorderedElementsAre; + +TEST(BlobPrefixAllowlistTest, ValidateParsingBlobName) { + auto result = ParseBlobName("DELTA_1705430864435450"); + EXPECT_EQ(result.key, "DELTA_1705430864435450"); + EXPECT_EQ(result.prefix, ""); + + result = ParseBlobName("prefix/DELTA_1705430864435450"); + EXPECT_EQ(result.key, "DELTA_1705430864435450"); + EXPECT_EQ(result.prefix, "prefix"); + + result = ParseBlobName("prefix1/prefix2/DELTA_1705430864435450"); + EXPECT_EQ(result.key, "DELTA_1705430864435450"); + EXPECT_EQ(result.prefix, "prefix1/prefix2"); +} + +TEST(BlobPrefixAllowlistTest, ValidateContainsPrefix) { + auto allowlist = BlobPrefixAllowlist("prefix1,prefix1/prefix2"); + EXPECT_TRUE(allowlist.Contains("")); + EXPECT_TRUE(allowlist.ContainsBlobPrefix("DELTA_1705430864435450")); + EXPECT_TRUE(allowlist.Contains("prefix1")); + EXPECT_TRUE(allowlist.ContainsBlobPrefix("prefix1/DELTA_1705430864435450")); + EXPECT_TRUE(allowlist.Contains("prefix1/prefix2")); + EXPECT_TRUE( + allowlist.ContainsBlobPrefix("prefix1/prefix2/DELTA_1705430864435450")); + EXPECT_FALSE(allowlist.Contains("non-existant")); + EXPECT_FALSE( + allowlist.ContainsBlobPrefix("non-existant/DELTA_1705430864435450")); +} + +TEST(BlobPrefixAllowlistTest, ValidateGettingAllPrefixes) { + auto allowlist = BlobPrefixAllowlist("prefix1,prefix1/prefix2"); + EXPECT_THAT(allowlist.Prefixes(), + UnorderedElementsAre("", "prefix1", "prefix1/prefix2")); +} + +} // namespace +} // namespace kv_server diff --git a/components/data/blob_storage/blob_storage_change_notifier_s3.cc b/components/data/blob_storage/blob_storage_change_notifier_s3.cc index 4e02db89..0f24f007 100644 --- a/components/data/blob_storage/blob_storage_change_notifier_s3.cc +++ b/components/data/blob_storage/blob_storage_change_notifier_s3.cc @@ -13,13 +13,13 @@ // limitations under the License. #include "absl/functional/bind_front.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "aws/core/utils/json/JsonSerializer.h" #include "components/data/blob_storage/blob_storage_change_notifier.h" #include "components/data/common/change_notifier.h" #include "components/telemetry/server_definition.h" -#include "glog/logging.h" namespace kv_server { namespace { @@ -47,10 +47,7 @@ class S3BlobStorageChangeNotifier : public BlobStorageChangeNotifier { if (!parsedMessage.ok()) { LOG(ERROR) << "Failed to parse JSON. Error: " << parsedMessage.status() << " Message:" << message; - LogIfError(KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kAwsJsonParseError), 1}})); + LogServerErrorMetric(kAwsJsonParseError); continue; } parsed_notifications.push_back(std::move(*parsedMessage)); diff --git a/components/data/blob_storage/blob_storage_change_notifier_s3_test.cc b/components/data/blob_storage/blob_storage_change_notifier_s3_test.cc index 472c7e04..0709d015 100644 --- a/components/data/blob_storage/blob_storage_change_notifier_s3_test.cc +++ b/components/data/blob_storage/blob_storage_change_notifier_s3_test.cc @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "aws/sqs/SQSClient.h" #include "aws/sqs/model/ReceiveMessageRequest.h" @@ -21,7 +22,6 @@ #include "components/data/common/msg_svc.h" #include "components/telemetry/server_definition.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" #include "gmock/gmock.h" #include "gtest/gtest.h" diff --git a/components/data/blob_storage/blob_storage_client.h b/components/data/blob_storage/blob_storage_client.h index 796fba03..68cb3f27 100644 --- a/components/data/blob_storage/blob_storage_client.h +++ b/components/data/blob_storage/blob_storage_client.h @@ -26,6 +26,7 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" namespace kv_server { @@ -43,10 +44,12 @@ class BlobStorageClient { public: struct DataLocation { std::string bucket; + std::string prefix; std::string key; bool operator==(const DataLocation& other) const { - return key == other.key && bucket == other.bucket; + return prefix == other.prefix && key == other.key && + bucket == other.bucket; } }; struct ListOptions { @@ -56,7 +59,7 @@ class BlobStorageClient { // Options for the underyling storage client. struct ClientOptions { - ClientOptions() {} + ClientOptions() = default; int64_t max_connections = std::thread::hardware_concurrency(); int64_t max_range_bytes = 8 * 1024 * 1024; // 8MB }; @@ -81,7 +84,9 @@ class BlobStorageClient { inline std::ostream& operator<<( std::ostream& os, const BlobStorageClient::DataLocation& location) { - os << location.bucket << "/" << location.key; + location.prefix.empty() + ? os << location.bucket << "/" << location.key + : os << location.bucket << "/" << location.prefix << "/" << location.key; return os; } diff --git a/components/data/blob_storage/blob_storage_client_gcp.cc b/components/data/blob_storage/blob_storage_client_gcp.cc index 9e17c4eb..5be23c51 100644 --- a/components/data/blob_storage/blob_storage_client_gcp.cc +++ b/components/data/blob_storage/blob_storage_client_gcp.cc @@ -24,18 +24,25 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" +#include "components/data/blob_storage/blob_prefix_allowlist.h" #include "components/data/blob_storage/blob_storage_client.h" #include "components/data/blob_storage/seeking_input_streambuf.h" #include "components/errors/error_util_gcp.h" -#include "glog/logging.h" #include "google/cloud/storage/client.h" namespace kv_server { namespace { +std::string AppendPrefix(const std::string& value, const std::string& prefix) { + return prefix.empty() ? value : absl::StrCat(prefix, "/", value); +} + class GcpBlobInputStreamBuf : public SeekingInputStreambuf { public: GcpBlobInputStreamBuf(google::cloud::storage::Client& client, @@ -50,8 +57,8 @@ class GcpBlobInputStreamBuf : public SeekingInputStreambuf { protected: absl::StatusOr SizeImpl() override { - auto object_metadata = - client_.GetObjectMetadata(location_.bucket, location_.key); + auto object_metadata = client_.GetObjectMetadata( + location_.bucket, AppendPrefix(location_.key, location_.prefix)); if (!object_metadata) { return GoogleErrorStatusToAbslStatus(object_metadata.status()); } @@ -61,7 +68,7 @@ class GcpBlobInputStreamBuf : public SeekingInputStreambuf { absl::StatusOr ReadChunk(int64_t offset, int64_t chunk_size, char* dest_buffer) override { auto stream = client_.ReadObject( - location_.bucket, location_.key, + location_.bucket, AppendPrefix(location_.key, location_.prefix), google::cloud::storage::ReadRange(offset, offset + chunk_size)); if (!stream.status().ok()) { return GoogleErrorStatusToAbslStatus(stream.status()); @@ -84,7 +91,8 @@ class GcpBlobReader : public BlobReader { : BlobReader(), streambuf_(client, location, GetOptions([this, location](absl::Status status) { - LOG(ERROR) << "Blob " << location.key + LOG(ERROR) << "Blob " + << AppendPrefix(location.key, location.prefix) << " failed stream with: " << status; is_.setstate(std::ios_base::badbit); })), @@ -117,7 +125,8 @@ std::unique_ptr GcpBlobStorageClient::GetBlobReader( absl::Status GcpBlobStorageClient::PutBlob(BlobReader& blob_reader, DataLocation location) { - auto blob_ostream = client_->WriteObject(location.bucket, location.key); + auto blob_ostream = client_->WriteObject( + location.bucket, AppendPrefix(location.key, location.prefix)); if (!blob_ostream) { return GoogleErrorStatusToAbslStatus(blob_ostream.last_status()); } @@ -129,16 +138,19 @@ absl::Status GcpBlobStorageClient::PutBlob(BlobReader& blob_reader, } absl::Status GcpBlobStorageClient::DeleteBlob(DataLocation location) { - google::cloud::Status status = - client_->DeleteObject(location.bucket, location.key); + google::cloud::Status status = client_->DeleteObject( + location.bucket, AppendPrefix(location.key, location.prefix)); return status.ok() ? absl::OkStatus() : GoogleErrorStatusToAbslStatus(status); } absl::StatusOr> GcpBlobStorageClient::ListBlobs( DataLocation location, ListOptions options) { - auto list_object_reader = client_->ListObjects( - location.bucket, google::cloud::storage::Prefix(options.prefix), - google::cloud::storage::StartOffset(options.start_after)); + auto list_object_reader = + client_->ListObjects(location.bucket, + google::cloud::storage::Prefix( + AppendPrefix(options.prefix, location.prefix)), + google::cloud::storage::StartOffset(AppendPrefix( + options.start_after, location.prefix))); std::vector keys; if (list_object_reader.begin() == list_object_reader.end()) { return keys; @@ -149,13 +161,13 @@ absl::StatusOr> GcpBlobStorageClient::ListBlobs( << std::move(object_metadata).status().message(); continue; } - // Manually exclude the starting name as the StartOffset option is - // inclusive. - if (object_metadata->name() == options.start_after) { + // inclusive and also drop blobs with different prefix. + auto blob = ParseBlobName(object_metadata->name()); + if (blob.key == options.start_after || blob.prefix != location.prefix) { continue; } - keys.push_back(object_metadata->name()); + keys.push_back(std::move(blob.key)); } std::sort(keys.begin(), keys.end()); return keys; diff --git a/components/data/blob_storage/blob_storage_client_gcp_test.cc b/components/data/blob_storage/blob_storage_client_gcp_test.cc index 510fe0d2..e08c952e 100644 --- a/components/data/blob_storage/blob_storage_client_gcp_test.cc +++ b/components/data/blob_storage/blob_storage_client_gcp_test.cc @@ -39,12 +39,13 @@ #include "google/cloud/storage/testing/canonical_errors.h" #include "google/cloud/storage/testing/mock_client.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { using ::google::cloud::storage::testing::canonical_errors::TransientError; +using testing::AllOf; +using testing::Property; class GcpBlobStorageClientTest : public ::testing::Test { protected: @@ -191,5 +192,30 @@ TEST_F(GcpBlobStorageClientTest, ListBlobWithNoNewObject) { EXPECT_TRUE(my_blobs->begin() == my_blobs->end()); } +TEST_F(GcpBlobStorageClientTest, DeleteBlobWithPrefixSucceeds) { + namespace gcs = google::cloud::storage; + std::shared_ptr mock = + std::make_shared(); + std::unique_ptr mock_client = std::make_unique( + gcs::internal::ClientImplDetails::CreateWithoutDecorations(mock)); + EXPECT_CALL(*mock, + DeleteObject(AllOf( + Property(&gcs::internal::DeleteObjectRequest::bucket_name, + "test_bucket"), + Property(&gcs::internal::DeleteObjectRequest::object_name, + "test_prefix/test_object")))) + .WillOnce(testing::Return( + google::cloud::make_status_or(gcs::internal::EmptyResponse{}))); + + std::unique_ptr client = + std::make_unique(std::move(mock_client)); + BlobStorageClient::DataLocation location{ + .bucket = "test_bucket", + .prefix = "test_prefix", + .key = "test_object", + }; + EXPECT_TRUE(client->DeleteBlob(location).ok()); +} + } // namespace } // namespace kv_server diff --git a/components/data/blob_storage/blob_storage_client_local.cc b/components/data/blob_storage/blob_storage_client_local.cc index d55d0829..8f823a38 100644 --- a/components/data/blob_storage/blob_storage_client_local.cc +++ b/components/data/blob_storage/blob_storage_client_local.cc @@ -22,11 +22,11 @@ #include #include +#include "absl/log/log.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "components/data/blob_storage/blob_storage_client.h" #include "components/data/blob_storage/seeking_input_streambuf.h" -#include "glog/logging.h" namespace kv_server { namespace { @@ -87,9 +87,13 @@ absl::Status FileBlobStorageClient::DeleteBlob(DataLocation location) { } absl::StatusOr> FileBlobStorageClient::ListBlobs( DataLocation location, ListOptions options) { + auto directory_name = + location.prefix.empty() + ? location.bucket + : absl::StrCat(location.bucket, "/", location.prefix); { std::error_code error_code; - std::filesystem::directory_entry directory{location.bucket, error_code}; + std::filesystem::directory_entry directory{directory_name, error_code}; if (error_code) { return absl::InternalError(absl::StrCat("Error getting directory entry: ", error_code.message())); @@ -98,7 +102,7 @@ absl::StatusOr> FileBlobStorageClient::ListBlobs( std::error_code error_code; std::vector blob_names; for (const auto& dir_entry : - std::filesystem::directory_iterator(location.bucket, error_code)) { + std::filesystem::directory_iterator(directory_name, error_code)) { if (dir_entry.is_directory()) { continue; } @@ -120,7 +124,10 @@ absl::StatusOr> FileBlobStorageClient::ListBlobs( std::filesystem::path FileBlobStorageClient::GetFullPath( const DataLocation& location) { - return std::filesystem::path(location.bucket) / location.key; + return location.prefix.empty() + ? std::filesystem::path(location.bucket) / location.key + : std::filesystem::path(location.bucket) / location.prefix / + location.key; } namespace { diff --git a/components/data/blob_storage/blob_storage_client_local_test.cc b/components/data/blob_storage/blob_storage_client_local_test.cc index fb00ac1b..e8b714d2 100644 --- a/components/data/blob_storage/blob_storage_client_local_test.cc +++ b/components/data/blob_storage/blob_storage_client_local_test.cc @@ -33,6 +33,11 @@ namespace kv_server { namespace { +void CreateSubDir(std::string_view subdir_name) { + std::filesystem::create_directory( + std::filesystem::path(::testing::TempDir()) / subdir_name); +} + void CreateFileInTmpDir(const std::string& filename) { const std::filesystem::path path = std::filesystem::path(::testing::TempDir()) / filename; @@ -121,6 +126,42 @@ TEST(LocalBlobStorageClientTest, PutBlob) { client->PutBlob(*from_blob_reader, to).code()); } +TEST(LocalBlobStorageClientTest, DeleteBlobWithPrefix) { + std::unique_ptr client = + std::make_unique(); + CreateSubDir("prefix"); + BlobStorageClient::DataLocation location{ + .bucket = ::testing::TempDir(), + .prefix = "prefix", + .key = "object", + }; + CreateFileInTmpDir("prefix/object"); + auto status = client->DeleteBlob(location); + EXPECT_TRUE(status.ok()) << status; + location.key = "non-existent-object"; + status = client->DeleteBlob(location); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInternal) << status; +} + +TEST(LocalBlobStorageClientTest, ListSubDirectoryWithFiles) { + std::unique_ptr client = + std::make_unique(); + CreateSubDir("prefix"); + CreateFileInTmpDir("prefix/object1"); + CreateFileInTmpDir("prefix/object2"); + CreateFileInTmpDir("prefix/object3"); + BlobStorageClient::DataLocation location{ + .bucket = ::testing::TempDir(), + .prefix = "prefix", + }; + kv_server::BlobStorageClient::ListOptions options; + auto status_or = client->ListBlobs(location, options); + ASSERT_TRUE(status_or.ok()) << status_or.status(); + EXPECT_EQ(*status_or, + std::vector({"object1", "object2", "object3"})); +} + // TODO(237669491): Add tests here } // namespace diff --git a/components/data/blob_storage/blob_storage_client_s3.cc b/components/data/blob_storage/blob_storage_client_s3.cc index e3556e5a..984e593c 100644 --- a/components/data/blob_storage/blob_storage_client_s3.cc +++ b/components/data/blob_storage/blob_storage_client_s3.cc @@ -19,7 +19,10 @@ #include #include +#include "absl/log/log.h" +#include "absl/strings/match.h" #include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" #include "aws/core/Aws.h" #include "aws/core/utils/threading/Executor.h" #include "aws/s3/S3Client.h" @@ -32,14 +35,18 @@ #include "aws/s3/model/PutObjectRequest.h" #include "aws/transfer/TransferHandle.h" #include "aws/transfer/TransferManager.h" +#include "components/data/blob_storage/blob_prefix_allowlist.h" #include "components/data/blob_storage/blob_storage_client.h" #include "components/data/blob_storage/seeking_input_streambuf.h" #include "components/errors/error_util_aws.h" -#include "glog/logging.h" namespace kv_server { namespace { +std::string AppendPrefix(const std::string& value, const std::string& prefix) { + return prefix.empty() ? value : absl::StrCat(prefix, "/", value); +} + // Sequentially load byte range data with a fixed amount of memory usage. class S3BlobInputStreamBuf : public SeekingInputStreambuf { public: @@ -57,7 +64,7 @@ class S3BlobInputStreamBuf : public SeekingInputStreambuf { absl::StatusOr SizeImpl() override { Aws::S3::Model::HeadObjectRequest request; request.SetBucket(location_.bucket); - request.SetKey(location_.key); + request.SetKey(AppendPrefix(location_.key, location_.prefix)); auto outcome = client_.HeadObject(request); if (!outcome.IsSuccess()) { return AwsErrorToStatus(outcome.GetError()); @@ -69,7 +76,7 @@ class S3BlobInputStreamBuf : public SeekingInputStreambuf { char* dest_buffer) override { Aws::S3::Model::GetObjectRequest request; request.SetBucket(location_.bucket); - request.SetKey(location_.key); + request.SetKey(AppendPrefix(location_.key, location_.prefix)); request.SetRange(GetRange(offset, chunk_size)); auto outcome = client_.GetObject(request); if (!outcome.IsSuccess()) { @@ -160,7 +167,7 @@ absl::Status S3BlobStorageClient::PutBlob(BlobReader& reader, // The owner of the stream is the caller. auto handle = transfer_manager_->UploadFile( std::shared_ptr(iostream.get(), [](std::iostream*) {}), - location.bucket, location.key, "", {}); + location.bucket, AppendPrefix(location.key, location.prefix), "", {}); handle->WaitUntilFinished(); const bool success = handle->GetStatus() == Aws::Transfer::TransferStatus::COMPLETED; @@ -170,7 +177,7 @@ absl::Status S3BlobStorageClient::PutBlob(BlobReader& reader, absl::Status S3BlobStorageClient::DeleteBlob(DataLocation location) { Aws::S3::Model::DeleteObjectRequest request; request.SetBucket(std::move(location.bucket)); - request.SetKey(std::move(location.key)); + request.SetKey(AppendPrefix(location.key, location.prefix)); const auto outcome = client_->DeleteObject(request); return outcome.IsSuccess() ? absl::OkStatus() : AwsErrorToStatus(outcome.GetError()); @@ -181,10 +188,12 @@ absl::StatusOr> S3BlobStorageClient::ListBlobs( Aws::S3::Model::ListObjectsV2Request request; request.SetBucket(std::move(location.bucket)); if (!options.prefix.empty()) { - request.SetPrefix(std::move(options.prefix)); + request.SetPrefix( + AppendPrefix(/*value=*/options.prefix, /*prefix=*/location.prefix)); } if (!options.start_after.empty()) { - request.SetStartAfter(std::move(options.start_after)); + request.SetStartAfter(AppendPrefix(/*value=*/options.start_after, + /*prefix=*/location.prefix)); } bool done = false; std::vector keys; @@ -196,7 +205,10 @@ absl::StatusOr> S3BlobStorageClient::ListBlobs( const Aws::Vector objects = outcome.GetResult().GetContents(); for (const Aws::S3::Model::Object& object : objects) { - keys.push_back(object.GetKey()); + if (auto blob = ParseBlobName(object.GetKey()); + blob.prefix == location.prefix) { + keys.emplace_back(std::move(blob.key)); + } } done = !outcome.GetResult().GetIsTruncated(); if (!done) { diff --git a/components/data/blob_storage/blob_storage_client_s3_test.cc b/components/data/blob_storage/blob_storage_client_s3_test.cc index 955382a5..7d20884e 100644 --- a/components/data/blob_storage/blob_storage_client_s3_test.cc +++ b/components/data/blob_storage/blob_storage_client_s3_test.cc @@ -31,12 +31,15 @@ #include "components/data/blob_storage/blob_storage_client.h" #include "components/telemetry/server_definition.h" #include "components/util/platform_initializer.h" +#include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { +using testing::AllOf; +using testing::Property; + constexpr int64_t kMaxRangeBytes = 1024 * 1024 * 8; class MockS3Client : public ::Aws::S3::S3Client { @@ -158,5 +161,64 @@ TEST_F(BlobStorageClientS3Test, ListBlobsFails) { client->ListBlobs(location, list_options).status().code()); } +TEST_F(BlobStorageClientS3Test, DeleteBlobWithPrefixSucceeds) { + auto mock_s3_client = std::make_shared(); + Aws::S3::Model::DeleteObjectResult result; // An empty result means success + EXPECT_CALL( + *mock_s3_client, + DeleteObject(::testing::AllOf( + testing::Property(&Aws::S3::Model::DeleteObjectRequest::GetBucket, + "bucket"), + testing::Property(&Aws::S3::Model::DeleteObjectRequest::GetKey, + "prefix/object")))) + .WillOnce(::testing::Return(result)); + std::unique_ptr client = + std::make_unique(mock_s3_client, kMaxRangeBytes); + BlobStorageClient::DataLocation location{ + .bucket = "bucket", + .prefix = "prefix", + .key = "object", + }; + EXPECT_TRUE(client->DeleteBlob(location).ok()); +} + +TEST_F(BlobStorageClientS3Test, ListBlobsWithPrefixSucceeds) { + auto mock_s3_client = std::make_shared(); + { + Aws::S3::Model::ListObjectsV2Result + result; // An empty result means success. + Aws::S3::Model::Object object_to_return; + object_to_return.SetKey("directory1/DELTA_1699834075511696"); + Aws::Vector objects_to_return = {object_to_return}; + result.SetContents(objects_to_return); + EXPECT_CALL( + *mock_s3_client, + ListObjectsV2( + AllOf(Property(&Aws::S3::Model::ListObjectsV2Request::GetBucket, + "bucket"), + Property(&Aws::S3::Model::ListObjectsV2Request::GetPrefix, + "directory1/DELTA"), + Property(&Aws::S3::Model::ListObjectsV2Request::GetStartAfter, + "directory1/DELTA_1699834075511695")))) + .WillOnce(::testing::Return(result)); + } + + std::unique_ptr client = + std::make_unique(mock_s3_client, kMaxRangeBytes); + BlobStorageClient::DataLocation location{ + .bucket = "bucket", + .prefix = "directory1", + }; + BlobStorageClient::ListOptions list_options{ + .prefix = "DELTA", + .start_after = "DELTA_1699834075511695", + }; + absl::StatusOr> response = + client->ListBlobs(location, list_options); + ASSERT_TRUE(response.ok()); + EXPECT_THAT(*response, + testing::UnorderedElementsAreArray({"DELTA_1699834075511696"})); +} + } // namespace } // namespace kv_server diff --git a/components/data/blob_storage/delta_file_notifier.cc b/components/data/blob_storage/delta_file_notifier.cc index ac67cdb4..cdef6df9 100644 --- a/components/data/blob_storage/delta_file_notifier.cc +++ b/components/data/blob_storage/delta_file_notifier.cc @@ -19,14 +19,15 @@ #include #include +#include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "components/data/common/thread_manager.h" #include "components/errors/retry.h" -#include "glog/logging.h" #include "public/constants.h" #include "public/data_loading/filename_utils.h" -#include "src/cpp/util/duration.h" +#include "src/util/duration.h" namespace kv_server { namespace { @@ -39,24 +40,27 @@ class DeltaFileNotifierImpl : public DeltaFileNotifier { explicit DeltaFileNotifierImpl(BlobStorageClient& client, const absl::Duration poll_frequency, std::unique_ptr sleep_for, - SteadyClock& clock) - : thread_manager_(TheadManager::Create("Delta file notifier")), + SteadyClock& clock, + BlobPrefixAllowlist blob_prefix_allowlist) + : thread_manager_(ThreadManager::Create("Delta file notifier")), client_(client), poll_frequency_(poll_frequency), sleep_for_(std::move(sleep_for)), - clock_(clock) {} + clock_(clock), + blob_prefix_allowlist_(std::move(blob_prefix_allowlist)) {} absl::Status Start( BlobStorageChangeNotifier& change_notifier, - BlobStorageClient::DataLocation location, std::string start_after, - std::function callback) override { - return thread_manager_->Start([this, location = std::move(location), - start_after = std::move(start_after), - callback = std::move(callback), - &change_notifier]() mutable { - Watch(change_notifier, std::move(location), std::move(start_after), - std::move(callback)); - }); + BlobStorageClient::DataLocation location, + absl::flat_hash_map&& prefix_start_after_map, + std::function callback) override { + return thread_manager_->Start( + [this, location = std::move(location), + prefix_start_after_map = std::move(prefix_start_after_map), + callback = std::move(callback), &change_notifier]() mutable { + Watch(change_notifier, std::move(location), + std::move(prefix_start_after_map), std::move(callback)); + }); } absl::Status Stop() override { @@ -71,128 +75,179 @@ class DeltaFileNotifierImpl : public DeltaFileNotifier { // Returns max DeltaFile in alphabetical order from notification // Empty string if wait_duration exceeded. absl::StatusOr WaitForNotification( - BlobStorageChangeNotifier& change_notifier, - absl::Duration wait_duration) { - absl::StatusOr> changes = - change_notifier.GetNotifications( - wait_duration, [this]() { return thread_manager_->ShouldStop(); }); - if (!changes.ok()) { - return changes.status(); - } - std::string_view max_change = ""; - for (const auto& change : *changes) { - if (change > max_change && IsDeltaFilename(change)) { - max_change = change; - } - } - return std::string(max_change); - } - + BlobStorageChangeNotifier& change_notifier, absl::Duration wait_duration); // Returns true if we need to check for new blobs. Returns the error on // failure. absl::StatusOr ShouldListBlobs( BlobStorageChangeNotifier& change_notifier, ExpiringFlag& expiring_flag, - std::string_view last_key) { - if (!expiring_flag.Get()) { - VLOG(5) << "Backup poll"; - return true; + const absl::flat_hash_map& + prefix_start_after_map); + void Watch( + BlobStorageChangeNotifier& change_notifier, + BlobStorageClient::DataLocation location, + absl::flat_hash_map&& prefix_start_after_map, + std::function callback); + + std::unique_ptr thread_manager_; + BlobStorageClient& client_; + const absl::Duration poll_frequency_; + std::unique_ptr sleep_for_; + SteadyClock& clock_; + BlobPrefixAllowlist blob_prefix_allowlist_; +}; + +absl::StatusOr DeltaFileNotifierImpl::WaitForNotification( + BlobStorageChangeNotifier& change_notifier, absl::Duration wait_duration) { + absl::StatusOr> changes = + change_notifier.GetNotifications( + wait_duration, [this]() { return thread_manager_->ShouldStop(); }); + if (!changes.ok()) { + return changes.status(); + } + std::string_view max_change = ""; + for (const auto& change : *changes) { + if (auto blob = ParseBlobName(change); + change > max_change && IsDeltaFilename(blob.key)) { + max_change = change; } - absl::StatusOr notification_key = - WaitForNotification(change_notifier, expiring_flag.GetTimeRemaining()); - if (notification_key.ok() && *notification_key < last_key) { - // Ignore notifications for keys we've already seen. - return false; + } + return std::string(max_change); +} + +// Returns true if we need to check for new blobs. Returns the error on +// failure. +absl::StatusOr DeltaFileNotifierImpl::ShouldListBlobs( + BlobStorageChangeNotifier& change_notifier, ExpiringFlag& expiring_flag, + const absl::flat_hash_map& + prefix_start_after_map) { + if (!expiring_flag.Get()) { + VLOG(5) << "Backup poll"; + return true; + } + absl::StatusOr notification_key = + WaitForNotification(change_notifier, expiring_flag.GetTimeRemaining()); + // Don't poll on error. A backup poll will trigger if necessary. + if (absl::IsDeadlineExceeded(notification_key.status())) { + // Deadline exceeded while waiting, trigger backup poll + VLOG(5) << "Backup poll"; + return true; + } + if (!notification_key.ok()) { + return notification_key.status(); + } + auto notification_blob = ParseBlobName(*notification_key); + if (auto iter = prefix_start_after_map.find(notification_blob.prefix); + iter != prefix_start_after_map.end() && + notification_blob.key < iter->second) { + // Ignore notifications for keys we've already seen. + return false; + } + // Return True if there is a Delta file notification + // False is returned on DeadlineExceeded. + return blob_prefix_allowlist_.Contains(notification_blob.prefix) && + IsDeltaFilename(notification_blob.key); +} + +absl::flat_hash_map> ListPrefixDeltaFiles( + BlobStorageClient::DataLocation location, + const BlobPrefixAllowlist& prefix_allowlist, + const absl::flat_hash_map& prefix_start_after_map, + BlobStorageClient& blob_client) { + absl::flat_hash_map> prefix_blobs_map; + for (const auto& blob_prefix : prefix_allowlist.Prefixes()) { + location.prefix = blob_prefix; + auto iter = prefix_start_after_map.find(blob_prefix); + absl::StatusOr> result = blob_client.ListBlobs( + location, + {.prefix = std::string(FilePrefix()), + .start_after = + (iter == prefix_start_after_map.end()) ? "" : iter->second}); + if (!result.ok()) { + LOG(ERROR) << "Failed to list " << location << ": " << result.status(); + continue; } - if (absl::IsDeadlineExceeded(notification_key.status())) { - // Deadline exceeded while waiting, trigger backup poll - VLOG(5) << "Backup poll"; - return true; + if (result->empty()) { + continue; } - // Don't poll on error. A backup poll will trigger if necessary. - if (!notification_key.ok()) { - return notification_key.status(); + auto& prefix_blobs = prefix_blobs_map[blob_prefix]; + prefix_blobs.reserve(result->size()); + for (const auto& blob : *result) { + prefix_blobs.push_back(blob); } - // Return True if there is a Delta file notification - // False is returned on DeadlineExceeded. - return IsDeltaFilename(*notification_key); } + return prefix_blobs_map; +} - void Watch(BlobStorageChangeNotifier& change_notifier, - BlobStorageClient::DataLocation location, std::string start_after, - std::function callback) { - LOG(INFO) << "Started to watch " << location; - std::string last_key = std::move(start_after); - // Flag starts expired, and forces an initial poll. - ExpiringFlag expiring_flag(clock_); - uint32_t sequential_failures = 0; - while (!thread_manager_->ShouldStop()) { - const absl::StatusOr should_list_blobs = - ShouldListBlobs(change_notifier, expiring_flag, last_key); - if (!should_list_blobs.ok()) { - ++sequential_failures; - const absl::Duration backoff_time = - std::min(expiring_flag.GetTimeRemaining(), - ExponentialBackoffForRetry(sequential_failures)); - LOG(ERROR) << "Failed to get delta file notifications: " - << should_list_blobs.status() << ". Waiting for " - << backoff_time; - if (!sleep_for_->Duration(backoff_time)) { - LOG(ERROR) << "Failed to sleep for " << backoff_time - << ". SleepFor invalid."; - } - continue; - } - sequential_failures = 0; - if (!*should_list_blobs) { - continue; - } - absl::StatusOr> result = client_.ListBlobs( - location, {.prefix = std::string(FilePrefix()), - .start_after = last_key}); - if (!result.ok()) { - LOG(ERROR) << "Failed to list " << location << ": " << result.status(); - continue; +void DeltaFileNotifierImpl::Watch( + BlobStorageChangeNotifier& change_notifier, + BlobStorageClient::DataLocation location, + absl::flat_hash_map&& prefix_start_after_map, + std::function callback) { + LOG(INFO) << "Started to watch " << location; + // Flag starts expired, and forces an initial poll. + ExpiringFlag expiring_flag(clock_); + uint32_t sequential_failures = 0; + while (!thread_manager_->ShouldStop()) { + const absl::StatusOr should_list_blobs = + ShouldListBlobs(change_notifier, expiring_flag, prefix_start_after_map); + if (!should_list_blobs.ok()) { + ++sequential_failures; + const absl::Duration backoff_time = + std::min(expiring_flag.GetTimeRemaining(), + ExponentialBackoffForRetry(sequential_failures)); + LOG(ERROR) << "Failed to get delta file notifications: " + << should_list_blobs.status() << ". Waiting for " + << backoff_time; + if (!sleep_for_->Duration(backoff_time)) { + LOG(ERROR) << "Failed to sleep for " << backoff_time + << ". SleepFor invalid."; } - // Set expiring flag before callback for unit test simplicity. - // Fake clock is moved forward in callback so flag must be set beforehand. - expiring_flag.Set(poll_frequency_); - int delta_file_count = 0; - for (const std::string& key : *result) { - if (!IsDeltaFilename(key)) { + continue; + } + sequential_failures = 0; + if (!*should_list_blobs) { + continue; + } + // Set expiring flag before callback for unit test simplicity. + // Fake clock is moved forward in callback so flag must be set beforehand. + expiring_flag.Set(poll_frequency_); + int delta_file_count = 0; + auto prefix_blobs_map = ListPrefixDeltaFiles( + location, blob_prefix_allowlist_, prefix_start_after_map, client_); + for (const auto& [prefix, prefix_blobs] : prefix_blobs_map) { + for (const auto& blob : prefix_blobs) { + if (!IsDeltaFilename(blob)) { continue; } - callback(key); - last_key = key; + callback(prefix.empty() ? blob : absl::StrCat(prefix, "/", blob)); + prefix_start_after_map[prefix] = blob; delta_file_count++; } - if (!delta_file_count) { - VLOG(2) << "No new file found"; - } + } + if (delta_file_count == 0) { + VLOG(2) << "No new file found"; } } - - std::unique_ptr thread_manager_; - BlobStorageClient& client_; - const absl::Duration poll_frequency_; - std::unique_ptr sleep_for_; - SteadyClock& clock_; -}; +} } // namespace std::unique_ptr DeltaFileNotifier::Create( - BlobStorageClient& client, const absl::Duration poll_frequency) { - return std::make_unique(client, poll_frequency, - std::make_unique(), - SteadyClock::RealClock()); + BlobStorageClient& client, const absl::Duration poll_frequency, + BlobPrefixAllowlist blob_prefix_allowlist) { + return std::make_unique( + client, poll_frequency, std::make_unique(), + SteadyClock::RealClock(), std::move(blob_prefix_allowlist)); } // For test only std::unique_ptr DeltaFileNotifier::Create( BlobStorageClient& client, const absl::Duration poll_frequency, - std::unique_ptr sleep_for, SteadyClock& clock) { - return std::make_unique(client, poll_frequency, - std::move(sleep_for), clock); + std::unique_ptr sleep_for, SteadyClock& clock, + BlobPrefixAllowlist blob_prefix_allowlist) { + return std::make_unique( + client, poll_frequency, std::move(sleep_for), clock, + std::move(blob_prefix_allowlist)); } } // namespace kv_server diff --git a/components/data/blob_storage/delta_file_notifier.h b/components/data/blob_storage/delta_file_notifier.h index c18a647f..6df74bce 100644 --- a/components/data/blob_storage/delta_file_notifier.h +++ b/components/data/blob_storage/delta_file_notifier.h @@ -20,12 +20,14 @@ #include #include +#include "absl/container/flat_hash_map.h" +#include "components/data/blob_storage/blob_prefix_allowlist.h" #include "components/data/blob_storage/blob_storage_change_notifier.h" #include "components/data/blob_storage/blob_storage_client.h" #include "components/data/common/thread_manager.h" #include "components/errors/retry.h" #include "components/util/sleepfor.h" -#include "src/cpp/util/duration.h" +#include "src/util/duration.h" namespace kv_server { @@ -46,8 +48,9 @@ class DeltaFileNotifier { // the constructor. virtual absl::Status Start( BlobStorageChangeNotifier& change_notifier, - BlobStorageClient::DataLocation location, std::string start_after, - std::function callback) = 0; + BlobStorageClient::DataLocation location, + absl::flat_hash_map&& prefix_start_after_map, + std::function callback) = 0; // Blocks until `IsRunning` is False. virtual absl::Status Stop() = 0; @@ -58,13 +61,15 @@ class DeltaFileNotifier { static std::unique_ptr Create( BlobStorageClient& client, - const absl::Duration poll_frequency = absl::Minutes(5)); + const absl::Duration poll_frequency = absl::Minutes(5), + BlobPrefixAllowlist blob_prefix_allowlist = BlobPrefixAllowlist("")); // Used for test static std::unique_ptr Create( BlobStorageClient& client, const absl::Duration poll_frequency, std::unique_ptr sleep_for, - privacy_sandbox::server_common::SteadyClock& clock); + privacy_sandbox::server_common::SteadyClock& clock, + BlobPrefixAllowlist blob_prefix_allowlist = BlobPrefixAllowlist("")); }; } // namespace kv_server diff --git a/components/data/blob_storage/delta_file_notifier_test.cc b/components/data/blob_storage/delta_file_notifier_test.cc index d9e359b5..b6e28224 100644 --- a/components/data/blob_storage/delta_file_notifier_test.cc +++ b/components/data/blob_storage/delta_file_notifier_test.cc @@ -26,6 +26,7 @@ #include "public/data_loading/filename_utils.h" using testing::_; +using testing::AllOf; using testing::Field; using testing::Return; @@ -34,14 +35,17 @@ namespace { using privacy_sandbox::server_common::SimulatedSteadyClock; +constexpr std::string_view kBlobPrefix1 = "prefix1"; + class DeltaFileNotifierTest : public ::testing::Test { protected: void SetUp() override { std::unique_ptr mock_sleep_for = std::make_unique(); sleep_for_ = mock_sleep_for.get(); - notifier_ = DeltaFileNotifier::Create( - client_, poll_frequency_, std::move(mock_sleep_for), sim_clock_); + notifier_ = DeltaFileNotifier::Create(client_, poll_frequency_, + std::move(mock_sleep_for), sim_clock_, + BlobPrefixAllowlist(kBlobPrefix1)); } MockBlobStorageClient client_; @@ -59,20 +63,21 @@ TEST_F(DeltaFileNotifierTest, NotRunning) { TEST_F(DeltaFileNotifierTest, StartFailure) { BlobStorageClient::DataLocation location = {.bucket = "testbucket"}; - absl::Status status = - notifier_->Start(change_notifier_, {.bucket = "testbucket"}, initial_key_, - [](const std::string&) {}); + absl::Status status = notifier_->Start( + change_notifier_, {.bucket = "testbucket"}, + {std::make_pair("", initial_key_)}, [](const std::string&) {}); ASSERT_TRUE(status.ok()); status = notifier_->Start(change_notifier_, {.bucket = "testbucket"}, - initial_key_, [](const std::string&) {}); + {std::make_pair("", initial_key_)}, + [](const std::string&) {}); ASSERT_FALSE(status.ok()); } TEST_F(DeltaFileNotifierTest, StartsAndStops) { BlobStorageClient::DataLocation location = {.bucket = "testbucket"}; - absl::Status status = - notifier_->Start(change_notifier_, {.bucket = "testbucket"}, initial_key_, - [](const std::string&) {}); + absl::Status status = notifier_->Start( + change_notifier_, {.bucket = "testbucket"}, + {std::make_pair("", initial_key_)}, [](const std::string&) {}); ASSERT_TRUE(status.ok()); EXPECT_TRUE(notifier_->IsRunning()); status = notifier_->Stop(); @@ -99,6 +104,13 @@ TEST_F(DeltaFileNotifierTest, NotifiesWithNewFiles) { Field(&BlobStorageClient::ListOptions::start_after, ToDeltaFileName(3).value()))) .WillOnce(Return(std::vector({ToDeltaFileName(4).value()}))); + EXPECT_CALL( + client_, + ListBlobs( + AllOf(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::DataLocation::prefix, kBlobPrefix1)), + Field(&BlobStorageClient::ListOptions::start_after, ""))) + .WillRepeatedly(Return(std::vector())); absl::Notification finished; testing::MockFunction callback; @@ -112,9 +124,9 @@ TEST_F(DeltaFileNotifierTest, NotifiesWithNewFiles) { finished.Notify(); }); - absl::Status status = - notifier_->Start(change_notifier_, {.bucket = "testbucket"}, initial_key_, - callback.AsStdFunction()); + absl::Status status = notifier_->Start( + change_notifier_, {.bucket = "testbucket"}, + {std::make_pair("", initial_key_)}, callback.AsStdFunction()); ASSERT_TRUE(status.ok()); EXPECT_TRUE(notifier_->IsRunning()); finished.WaitForNotification(); @@ -147,6 +159,13 @@ TEST_F(DeltaFileNotifierTest, NotifiesWithInvalidFilesIngored) { Field(&BlobStorageClient::ListOptions::start_after, ToDeltaFileName(3).value()))) .WillOnce(Return(std::vector({ToDeltaFileName(4).value()}))); + EXPECT_CALL( + client_, + ListBlobs( + AllOf(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::DataLocation::prefix, kBlobPrefix1)), + Field(&BlobStorageClient::ListOptions::start_after, ""))) + .WillRepeatedly(Return(std::vector())); EXPECT_CALL( client_, ListBlobs(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), @@ -169,9 +188,9 @@ TEST_F(DeltaFileNotifierTest, NotifiesWithInvalidFilesIngored) { finished.Notify(); }); - absl::Status status = - notifier_->Start(change_notifier_, {.bucket = "testbucket"}, initial_key_, - callback.AsStdFunction()); + absl::Status status = notifier_->Start( + change_notifier_, {.bucket = "testbucket"}, + {std::make_pair("", initial_key_)}, callback.AsStdFunction()); ASSERT_TRUE(status.ok()); EXPECT_TRUE(notifier_->IsRunning()); finished.WaitForNotification(); @@ -194,7 +213,13 @@ TEST_F(DeltaFileNotifierTest, GetChangesFailure) { ToDeltaFileName(1).value()))) .WillOnce(Return(std::vector({}))) .WillOnce(Return(std::vector({ToDeltaFileName(1).value()}))); - + EXPECT_CALL( + client_, + ListBlobs( + AllOf(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::DataLocation::prefix, kBlobPrefix1)), + Field(&BlobStorageClient::ListOptions::start_after, ""))) + .WillRepeatedly(Return(std::vector())); absl::Notification finished; testing::MockFunction callback; EXPECT_CALL(callback, Call).Times(1).WillOnce([&](const std::string& key) { @@ -208,9 +233,9 @@ TEST_F(DeltaFileNotifierTest, GetChangesFailure) { .Times(1) .WillOnce(Return(true)); - absl::Status status = - notifier_->Start(change_notifier_, {.bucket = "testbucket"}, initial_key_, - callback.AsStdFunction()); + absl::Status status = notifier_->Start( + change_notifier_, {.bucket = "testbucket"}, + {std::make_pair("", initial_key_)}, callback.AsStdFunction()); ASSERT_TRUE(status.ok()); EXPECT_TRUE(notifier_->IsRunning()); finished.WaitForNotification(); @@ -241,7 +266,13 @@ TEST_F(DeltaFileNotifierTest, BackupPoll) { Field(&BlobStorageClient::ListOptions::start_after, ToDeltaFileName(3).value()))) .WillOnce(Return(std::vector({ToDeltaFileName(4).value()}))); - + EXPECT_CALL( + client_, + ListBlobs( + AllOf(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::DataLocation::prefix, kBlobPrefix1)), + Field(&BlobStorageClient::ListOptions::start_after, ""))) + .WillRepeatedly(Return(std::vector())); absl::Notification finished; testing::MockFunction callback; EXPECT_CALL(callback, Call) @@ -261,8 +292,70 @@ TEST_F(DeltaFileNotifierTest, BackupPoll) { finished.Notify(); }); + absl::Status status = notifier_->Start( + change_notifier_, {.bucket = "testbucket"}, + {std::make_pair("", initial_key_)}, callback.AsStdFunction()); + ASSERT_TRUE(status.ok()); + EXPECT_TRUE(notifier_->IsRunning()); + finished.WaitForNotification(); + status = notifier_->Stop(); + ASSERT_TRUE(status.ok()); + EXPECT_FALSE(notifier_->IsRunning()); +} + +TEST_F(DeltaFileNotifierTest, NotifiesWithNewPrefixedFiles) { + BlobStorageClient::DataLocation location = {.bucket = "testbucket"}; + EXPECT_CALL(change_notifier_, GetNotifications(_, _)) + .WillOnce(Return(std::vector({ToDeltaFileName(3).value()}))) + .WillOnce(Return(std::vector({ToDeltaFileName(4).value()}))) + .WillRepeatedly(Return(std::vector())); + EXPECT_CALL( + client_, + ListBlobs(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::ListOptions::start_after, + ToDeltaFileName(1).value()))) + .WillOnce(Return(std::vector({}))) + .WillOnce(Return(std::vector({ToDeltaFileName(3).value()}))); + EXPECT_CALL( + client_, + ListBlobs(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::ListOptions::start_after, + ToDeltaFileName(3).value()))) + .WillOnce(Return(std::vector({ToDeltaFileName(4).value()}))); + EXPECT_CALL( + client_, + ListBlobs( + AllOf(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::DataLocation::prefix, kBlobPrefix1)), + Field(&BlobStorageClient::ListOptions::start_after, + ToDeltaFileName(10).value()))) + .WillOnce( + Return(std::vector({ToDeltaFileName(11).value()}))); + EXPECT_CALL( + client_, + ListBlobs( + AllOf(Field(&BlobStorageClient::DataLocation::bucket, "testbucket"), + Field(&BlobStorageClient::DataLocation::prefix, kBlobPrefix1)), + Field(&BlobStorageClient::ListOptions::start_after, + ToDeltaFileName(11).value()))) + .WillRepeatedly(Return(std::vector())); + + absl::Notification finished; + testing::MockFunction callback; + EXPECT_CALL(callback, Call(ToDeltaFileName(3).value())); + EXPECT_CALL(callback, Call(absl::StrCat(kBlobPrefix1, "/", + ToDeltaFileName(11).value()))); + EXPECT_CALL(callback, Call(ToDeltaFileName(4).value())).WillOnce([&]() { + finished.Notify(); + }); + absl::Status status = - notifier_->Start(change_notifier_, {.bucket = "testbucket"}, initial_key_, + notifier_->Start(change_notifier_, {.bucket = "testbucket"}, + { + std::make_pair("", initial_key_), + std::make_pair(std::string(kBlobPrefix1), + ToDeltaFileName(10).value()), + }, callback.AsStdFunction()); ASSERT_TRUE(status.ok()); EXPECT_TRUE(notifier_->IsRunning()); diff --git a/components/data/blob_storage/seeking_input_streambuf.cc b/components/data/blob_storage/seeking_input_streambuf.cc index 531ad8cf..fdb3bba1 100644 --- a/components/data/blob_storage/seeking_input_streambuf.cc +++ b/components/data/blob_storage/seeking_input_streambuf.cc @@ -21,11 +21,11 @@ #include #include "absl/base/optimization.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "components/telemetry/server_definition.h" -#include "glog/logging.h" namespace kv_server { namespace { @@ -57,7 +57,9 @@ std::streampos SeekingInputStreambuf::seekpos(std::streampos pos, std::streampos SeekingInputStreambuf::seekoff(std::streamoff off, std::ios_base::seekdir dir, std::ios_base::openmode which) { - auto start_time = absl::Now(); + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); const auto size = Size(); if (ABSL_PREDICT_FALSE(!size.ok())) { MaybeReportError(size.status()); @@ -105,17 +107,14 @@ std::streampos SeekingInputStreambuf::seekoff(std::streamoff off, buffer_.data() + (new_position - BufferStartPosition()), buffer_.data() + buffer_.length()); } - auto duration = absl::Now() - start_time; - LogIfError(KVServerContextMap() - ->SafeMetric() - .LogHistogram( - absl::ToDoubleMicroseconds(duration))); - MaybeVerboseLogLatency(kSeekoffEventName, duration); + MaybeVerboseLogLatency(kSeekoffEventName, latency_recorder.GetLatency()); return std::streampos(std::streamoff(new_position)); } std::streambuf::int_type SeekingInputStreambuf::underflow() { - auto start_time = absl::Now(); + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); const auto size = Size(); if (ABSL_PREDICT_FALSE(!size.ok())) { MaybeReportError(size.status()); @@ -149,12 +148,7 @@ std::streambuf::int_type SeekingInputStreambuf::underflow() { buffer_.resize(total_bytes_read); } setg(buffer_.data(), buffer_.data(), buffer_.data() + buffer_.length()); - auto duration = absl::Now() - start_time; - LogIfError(KVServerContextMap() - ->SafeMetric() - .LogHistogram( - absl::ToDoubleMicroseconds(duration))); - MaybeVerboseLogLatency(kUnderflowEventName, duration); + MaybeVerboseLogLatency(kUnderflowEventName, latency_recorder.GetLatency()); return traits_type::to_int_type(buffer_[0]); } @@ -175,7 +169,9 @@ int64_t SeekingInputStreambuf::BufferCursorPosition() { } absl::StatusOr SeekingInputStreambuf::Size() { - auto start_time = absl::Now(); + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); if (ABSL_PREDICT_TRUE(src_cached_size_ >= 0)) { return src_cached_size_; } @@ -184,12 +180,7 @@ absl::StatusOr SeekingInputStreambuf::Size() { return size.status(); } src_cached_size_ = *size; - auto duration = absl::Now() - start_time; - LogIfError(KVServerContextMap() - ->SafeMetric() - .LogHistogram( - absl::ToDoubleMicroseconds(duration))); - MaybeVerboseLogLatency(kSizeEventName, duration); + MaybeVerboseLogLatency(kSizeEventName, latency_recorder.GetLatency()); return *size; } diff --git a/components/data/blob_storage/seeking_input_streambuf.h b/components/data/blob_storage/seeking_input_streambuf.h index b7bc0235..c76a0c5f 100644 --- a/components/data/blob_storage/seeking_input_streambuf.h +++ b/components/data/blob_storage/seeking_input_streambuf.h @@ -20,7 +20,7 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" #ifndef COMPONENTS_DATA_BLOB_STORAGE_SEEKING_INPUT_STREAMBUF_H_ #define COMPONENTS_DATA_BLOB_STORAGE_SEEKING_INPUT_STREAMBUF_H_ diff --git a/components/data/blob_storage/seeking_input_streambuf_test.cc b/components/data/blob_storage/seeking_input_streambuf_test.cc index 242d8795..845f6995 100644 --- a/components/data/blob_storage/seeking_input_streambuf_test.cc +++ b/components/data/blob_storage/seeking_input_streambuf_test.cc @@ -27,7 +27,6 @@ #include "absl/strings/str_format.h" #include "components/telemetry/server_definition.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/data/common/BUILD.bazel b/components/data/common/BUILD.bazel index feb343cc..05830260 100644 --- a/components/data/common/BUILD.bazel +++ b/components/data/common/BUILD.bazel @@ -49,7 +49,7 @@ cc_library( "//conditions:default": [], }) + [ ":msg_svc_util", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/random", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", @@ -88,11 +88,11 @@ cc_library( ":message_service", "//components/telemetry:server_definition", "//components/util:sleepfor", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", ], ) @@ -123,10 +123,9 @@ cc_test( ], }) + [ ":change_notifier", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -142,13 +141,13 @@ cc_library( "//components/errors:retry", "//public:constants", "//public/data_loading:filename_utils", - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/functional:bind_front", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/synchronization", - "@google_privacysandbox_servers_common//src/cpp/util:duration", + "@google_privacysandbox_servers_common//src/util:duration", ], ) @@ -175,6 +174,7 @@ cc_library( "//components/data/blob_storage:delta_file_notifier", "//components/data/realtime:realtime_notifier", "//components/data/realtime:realtime_thread_pool_manager", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_googletest//:gtest", ], ) diff --git a/components/data/common/change_notifier_aws.cc b/components/data/common/change_notifier_aws.cc index d8f0fa0e..9e1ee74c 100644 --- a/components/data/common/change_notifier_aws.cc +++ b/components/data/common/change_notifier_aws.cc @@ -18,6 +18,7 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" @@ -42,8 +43,7 @@ #include "components/data/common/msg_svc.h" #include "components/errors/error_util_aws.h" #include "components/telemetry/server_definition.h" -#include "glog/logging.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" namespace kv_server { namespace { @@ -78,11 +78,7 @@ class AwsChangeNotifier : public ChangeNotifier { absl::Status status = queue_manager_->SetupQueue(); if (!status.ok()) { LOG(ERROR) << "Could not set up queue for topic " << sns_arn_; - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kAwsChangeNotifierQueueSetupFailure), 1}})); + LogServerErrorMetric(kAwsChangeNotifierQueueSetupFailure); return status; } } @@ -130,10 +126,7 @@ class AwsChangeNotifier : public ChangeNotifier { } else { LOG(ERROR) << "Failed to TagQueue with " << kLastUpdatedTag << ": " << tag << " " << status; - LogIfError(KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kAwsChangeNotifierTagFailure), 1}})); + LogServerErrorMetric(kAwsChangeNotifierTagFailure); } } } @@ -155,13 +148,7 @@ class AwsChangeNotifier : public ChangeNotifier { if (!outcome.IsSuccess()) { LOG(ERROR) << "Failed to receive message from SQS: " << outcome.GetError().GetMessage(); - - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kAwsChangeNotifierMessagesReceivingFailure), - 1}})); + LogServerErrorMetric(kAwsChangeNotifierMessagesReceivingFailure); if (!outcome.GetError().ShouldRetry()) { // Handle case where recreating Queue will resolve the issue. // Example: Queue accidentally deleted. @@ -186,11 +173,7 @@ class AwsChangeNotifier : public ChangeNotifier { } DeleteMessages(GetSqsUrl(), messages); if (keys.empty()) { - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kAwsChangeNotifierMessagesDataLoss), 1}})); + LogServerErrorMetric(kAwsChangeNotifierMessagesDataLoss); return absl::DataLossError("All messages invalid."); } return keys; @@ -216,12 +199,7 @@ class AwsChangeNotifier : public ChangeNotifier { if (!outcome.IsSuccess()) { LOG(ERROR) << "Failed to delete message from SQS: " << outcome.GetError().GetMessage(); - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kAwsChangeNotifierMessagesDeletionFailure), - 1}})); + LogServerErrorMetric(kAwsChangeNotifierMessagesDeletionFailure); } } diff --git a/components/data/common/change_notifier_aws_test.cc b/components/data/common/change_notifier_aws_test.cc index 5970ccd6..2db5e4f6 100644 --- a/components/data/common/change_notifier_aws_test.cc +++ b/components/data/common/change_notifier_aws_test.cc @@ -16,6 +16,7 @@ #include +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "aws/sqs/SQSClient.h" #include "aws/sqs/model/DeleteMessageBatchRequest.h" @@ -24,7 +25,6 @@ #include "components/data/common/msg_svc.h" #include "components/telemetry/server_definition.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" #include "gmock/gmock.h" #include "gtest/gtest.h" diff --git a/components/data/common/change_notifier_gcp.cc b/components/data/common/change_notifier_gcp.cc index 2c9718ba..b9f4c19c 100644 --- a/components/data/common/change_notifier_gcp.cc +++ b/components/data/common/change_notifier_gcp.cc @@ -19,12 +19,12 @@ #include #include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/time/clock.h" #include "components/data/common/change_notifier.h" #include "components/util/sleepfor.h" -#include "glog/logging.h" namespace kv_server { namespace { diff --git a/components/data/common/change_notifier_gcp_test.cc b/components/data/common/change_notifier_gcp_test.cc index 4ee34e84..94c58ad7 100644 --- a/components/data/common/change_notifier_gcp_test.cc +++ b/components/data/common/change_notifier_gcp_test.cc @@ -20,11 +20,10 @@ #include #include +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "components/data/common/change_notifier.h" -#include "glog/logging.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/data/common/change_notifier_local.cc b/components/data/common/change_notifier_local.cc index 0015c858..94d8a09d 100644 --- a/components/data/common/change_notifier_local.cc +++ b/components/data/common/change_notifier_local.cc @@ -16,13 +16,13 @@ #include #include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/time/clock.h" #include "components/data/common/change_notifier.h" #include "components/util/sleepfor.h" -#include "glog/logging.h" namespace kv_server { namespace { diff --git a/components/data/common/change_notifier_local_test.cc b/components/data/common/change_notifier_local_test.cc index 94fc27eb..073fe472 100644 --- a/components/data/common/change_notifier_local_test.cc +++ b/components/data/common/change_notifier_local_test.cc @@ -20,11 +20,10 @@ #include #include +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "components/data/common/change_notifier.h" -#include "glog/logging.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/data/common/mocks.h b/components/data/common/mocks.h index aa0ea669..707b0a43 100644 --- a/components/data/common/mocks.h +++ b/components/data/common/mocks.h @@ -19,6 +19,7 @@ #include #include +#include "absl/container/flat_hash_map.h" #include "components/data/blob_storage/blob_storage_client.h" #include "components/data/blob_storage/delta_file_notifier.h" #include "components/data/common/change_notifier.h" @@ -50,12 +51,10 @@ class MockBlobStorageChangeNotifier : public BlobStorageChangeNotifier { class MockDeltaFileNotifier : public DeltaFileNotifier { public: MockDeltaFileNotifier() : DeltaFileNotifier() {} - MOCK_METHOD(absl::Status, Start, - (BlobStorageChangeNotifier & change_notifier, - BlobStorageClient::DataLocation location, - std::string start_after, - std::function callback), + (BlobStorageChangeNotifier&, BlobStorageClient::DataLocation, + (absl::flat_hash_map&&), + std::function), (override)); MOCK_METHOD(absl::Status, Stop, (), (override)); MOCK_METHOD(bool, IsRunning, (), (const, override)); diff --git a/components/data/common/msg_svc_gcp.cc b/components/data/common/msg_svc_gcp.cc index 99ba9b4b..892c298e 100644 --- a/components/data/common/msg_svc_gcp.cc +++ b/components/data/common/msg_svc_gcp.cc @@ -14,13 +14,13 @@ #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "components/data/common/msg_svc.h" #include "components/data/common/msg_svc_util.h" #include "components/errors/error_util_gcp.h" -#include "glog/logging.h" #include "google/cloud/pubsub/message.h" #include "google/cloud/pubsub/subscriber.h" #include "google/cloud/pubsub/subscription_admin_client.h" diff --git a/components/data/common/thread_manager.cc b/components/data/common/thread_manager.cc index 1cfd6e42..59c6c6b0 100644 --- a/components/data/common/thread_manager.cc +++ b/components/data/common/thread_manager.cc @@ -21,21 +21,21 @@ #include #include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "components/errors/retry.h" -#include "glog/logging.h" #include "public/constants.h" #include "public/data_loading/filename_utils.h" -#include "src/cpp/util/duration.h" +#include "src/util/duration.h" namespace kv_server { namespace { -class TheadManagerImpl : public TheadManager { +class ThreadManagerImpl : public ThreadManager { public: - explicit TheadManagerImpl(std::string thread_name) + explicit ThreadManagerImpl(std::string thread_name) : thread_name_(std::move(thread_name)) {} - ~TheadManagerImpl() { + ~ThreadManagerImpl() { if (!IsRunning()) return; VLOG(8) << thread_name_ << " In destructor. Attempting to stop the thread."; if (const auto s = Stop(); !s.ok()) { @@ -78,8 +78,8 @@ class TheadManagerImpl : public TheadManager { } // namespace -std::unique_ptr TheadManager::Create(std::string thread_name) { - return std::make_unique(std::move(thread_name)); +std::unique_ptr ThreadManager::Create(std::string thread_name) { + return std::make_unique(std::move(thread_name)); } } // namespace kv_server diff --git a/components/data/common/thread_manager.h b/components/data/common/thread_manager.h index abd9f24b..123e9d43 100644 --- a/components/data/common/thread_manager.h +++ b/components/data/common/thread_manager.h @@ -24,11 +24,11 @@ namespace kv_server { -class TheadManager { +class ThreadManager { public: - virtual ~TheadManager() = default; + virtual ~ThreadManager() = default; - // Checks if the TheadManager is already running. + // Checks if the ThreadManager is already running. // If not, starts a thread on which `watch` is executed. // Start and Stop should be called on the same thread as // the constructor. @@ -43,7 +43,7 @@ class TheadManager { virtual bool ShouldStop() = 0; - static std::unique_ptr Create(std::string thread_name); + static std::unique_ptr Create(std::string thread_name); }; } // namespace kv_server diff --git a/components/data/file_group/BUILD.bazel b/components/data/file_group/BUILD.bazel new file mode 100644 index 00000000..28d96438 --- /dev/null +++ b/components/data/file_group/BUILD.bazel @@ -0,0 +1,73 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") + +package(default_visibility = [ + "//components:__subpackages__", +]) + +cc_library( + name = "file_group", + srcs = ["file_group.cc"], + hdrs = ["file_group.h"], + deps = [ + "//public:base_types_cc_proto", + "//public/data_loading:filename_utils", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", + ], +) + +cc_test( + name = "file_group_test", + size = "small", + srcs = ["file_group_test.cc"], + deps = [ + ":file_group", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "file_group_search_utils", + srcs = ["file_group_search_utils.cc"], + hdrs = ["file_group_search_utils.h"], + deps = [ + ":file_group", + "//components/data/blob_storage:blob_storage_client", + "//public:base_types_cc_proto", + "//public:constants", + "//public/data_loading:filename_utils", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", + ], +) + +cc_test( + name = "file_group_search_utils_test", + size = "small", + srcs = ["file_group_search_utils_test.cc"], + deps = [ + ":file_group", + ":file_group_search_utils", + "//components/data/blob_storage:blob_storage_client", + "//components/data/common:mocks", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/components/data/file_group/file_group.cc b/components/data/file_group/file_group.cc new file mode 100644 index 00000000..93954366 --- /dev/null +++ b/components/data/file_group/file_group.cc @@ -0,0 +1,141 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "components/data/file_group/file_group.h" + +#include "absl/status/statusor.h" +#include "absl/strings/numbers.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" +#include "public/constants.h" +#include "public/data_loading/filename_utils.h" +#include "src/util/status_macro/status_macros.h" + +namespace kv_server { +namespace { + +constexpr int32_t kFileTypeStrIndex = 0; +constexpr int32_t kLogicalCommitTimeStrIndex = 1; +constexpr int32_t kFileGroupFileIndexStrIndex = 2; +constexpr int32_t kFileGroupSizeStrIndex = 4; +constexpr int32_t kFileGroupFileNameNumComponents = 5; +constexpr int32_t kStandardDeltaOrSnapshotFileGroupSize = 1; + +absl::StatusOr ToInt64(absl::string_view digits) { + if (int64_t result; absl::SimpleAtoi(digits, &result)) { + return result; + } + return absl::InvalidArgumentError( + absl::StrCat("Input ", digits, " is not numeric.")); +} + +} // namespace + +absl::StatusOr FileGroup::ParseMetadataFromFilename( + std::string_view filename) { + Metadata metadata; + std::vector name_parts = + absl::StrSplit(filename, kFileComponentDelimiter); + FileType::Enum file_type; + if (!FileType_Enum_Parse(name_parts[kFileTypeStrIndex], &file_type)) { + return absl::InvalidArgumentError( + absl::StrCat("File type: ", name_parts[kFileTypeStrIndex])); + } + PS_ASSIGN_OR_RETURN(auto logical_commit_time, + ToInt64(name_parts[kLogicalCommitTimeStrIndex])); + metadata.file_type = file_type; + metadata.logical_commit_time = logical_commit_time; + metadata.basename = + absl::StrCat(name_parts[kFileTypeStrIndex], kFileComponentDelimiter, + name_parts[kLogicalCommitTimeStrIndex]); + if (name_parts.size() == kFileGroupFileNameNumComponents) { + PS_ASSIGN_OR_RETURN(auto file_index, + ToInt64(name_parts[kFileGroupFileIndexStrIndex])); + PS_ASSIGN_OR_RETURN(auto file_group_size, + ToInt64(name_parts[kFileGroupSizeStrIndex])); + if (file_index >= file_group_size) { + return absl::InvalidArgumentError( + absl::StrCat(filename, " is invalid. index: ", file_index, + " must be < ", file_group_size)); + } + metadata.group_size = file_group_size; + } else { + metadata.group_size = kStandardDeltaOrSnapshotFileGroupSize; + } + return metadata; +} + +absl::Status FileGroup::ValidateIncomingMetadata( + const Metadata& existing_metadata, const Metadata& incoming_metadata) { + if (existing_metadata.group_size != incoming_metadata.group_size) { + return absl::InvalidArgumentError( + absl::StrCat("Wrong group size: ", incoming_metadata.group_size, + ", group size is: ", existing_metadata.group_size)); + } + if (existing_metadata.logical_commit_time != + incoming_metadata.logical_commit_time) { + return absl::InvalidArgumentError(absl::StrCat( + "Wrong logical commit time: ", incoming_metadata.logical_commit_time, + ", group commit time is: ", existing_metadata.logical_commit_time)); + } + if (existing_metadata.file_type != incoming_metadata.file_type) { + return absl::InvalidArgumentError(absl::StrCat( + "Wrong file type: ", FileType_Enum_Name(incoming_metadata.file_type), + ", group file type is: ", + FileType_Enum_Name(existing_metadata.file_type))); + } + return absl::OkStatus(); +} + +int32_t FileGroup::Size() const { return cached_metadata_.group_size; } + +FileType::Enum FileGroup::Type() const { return cached_metadata_.file_type; } + +std::string_view FileGroup::Basename() const { + return cached_metadata_.basename; +} + +const absl::flat_hash_set& FileGroup::Filenames() const { + return files_set_; +} + +FileGroup::FileStatus FileGroup::GetStatus() const { + return !files_set_.empty() && Size() == files_set_.size() + ? FileStatus::kComplete + : FileStatus::kPending; +} + +absl::Status FileGroup::AddFile(std::string_view filename) { + if (GetStatus() == FileStatus::kComplete || files_set_.contains(filename)) { + return absl::OkStatus(); + } + if (!IsDeltaFilename(filename) && !IsSnapshotFilename(filename) && + !IsFileGroupFileName(filename)) { + return absl::InvalidArgumentError(absl::StrCat( + "File name: ", filename, " is not supported for file groups.")); + } + PS_ASSIGN_OR_RETURN(auto incoming_metadata, + ParseMetadataFromFilename(filename)); + if (cached_metadata_.file_type == FileType::FILE_TYPE_UNSPECIFIED) { + cached_metadata_ = incoming_metadata; + } + PS_RETURN_IF_ERROR( + ValidateIncomingMetadata(cached_metadata_, incoming_metadata)); + files_set_.emplace(filename); + return absl::OkStatus(); +} + +} // namespace kv_server diff --git a/components/data/file_group/file_group.h b/components/data/file_group/file_group.h new file mode 100644 index 00000000..ab6e3ee2 --- /dev/null +++ b/components/data/file_group/file_group.h @@ -0,0 +1,88 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef COMPONENTS_DATA_FILE_GROUP_FILE_GROUP_H_ +#define COMPONENTS_DATA_FILE_GROUP_FILE_GROUP_H_ + +#include +#include +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" +#include "public/base_types.pb.h" + +namespace kv_server { + +// A "file group" is a group of related snapshot/delta files that share a common +// prefix and are treated as a single file for reading purposes by the KV +// server. +// +// Note that this class is not thread safe. +class FileGroup { + public: + enum class FileStatus : int8_t { + // Means some files are not yet added to the group. + kPending = 0, + // Means all files are added to the group. + kComplete, + }; + + // Returns the total number of files in the group (including files that are + // not yet added to the group but are supposed to be). + [[nodiscard]] int32_t Size() const; + + // Returns type of files contained by this group. Currently this can only be + // either `FileType::DELTA` or `FileType::SNAPSHOT`. + [[nodiscard]] FileType::Enum Type() const; + + [[nodiscard]] std::string_view Basename() const; + + // Returns a list of filenames that have been added to the group so far. + [[nodiscard]] const absl::flat_hash_set& Filenames() const; + + // Attempts to add `filename` to this file group. + // Returns: + // - absl::OKStatus: when filename is added successfully or already exists in + // the group. + // - absl::InvalidArgumentError: when `filename` does not belong to this file + // group or is not supported for file groups. + absl::Status AddFile(std::string_view filename); + + [[nodiscard]] FileStatus GetStatus() const; + + private: + struct Metadata { + std::string basename; + int32_t group_size = 0; + int64_t logical_commit_time = -1; + FileType::Enum file_type = FileType::FILE_TYPE_UNSPECIFIED; + }; + + absl::StatusOr ParseMetadataFromFilename(std::string_view filename); + absl::Status ValidateIncomingMetadata(const Metadata& existing_metadata, + const Metadata& incoming_metadata); + + // Cache metadata so we don't have to repeatedly parse one of the existing + // files to get the group's `file_type`, `logical_commit_time` and`group_size` + Metadata cached_metadata_; + absl::flat_hash_set files_set_; +}; + +} // namespace kv_server + +#endif // COMPONENTS_DATA_FILE_GROUP_FILE_GROUP_H_ diff --git a/components/data/file_group/file_group_search_utils.cc b/components/data/file_group/file_group_search_utils.cc new file mode 100644 index 00000000..9a0d2def --- /dev/null +++ b/components/data/file_group/file_group_search_utils.cc @@ -0,0 +1,100 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "components/data/file_group/file_group_search_utils.h" + +#include +#include + +#include "absl/strings/match.h" +#include "public/data_loading/filename_utils.h" +#include "src/util/status_macro/status_macros.h" + +namespace kv_server { +namespace { + +BlobStorageClient::ListOptions CreateBlobListOptions( + const FileGroupFilter& filter) { + BlobStorageClient::ListOptions options; + if (!filter.start_after_basename.empty()) { + options.start_after = filter.start_after_basename; + } + if (filter.file_type.has_value()) { + options.prefix = FileType::Enum_Name(*filter.file_type); + } + return options; +} + +// Assumes that filenames are lexicographically ordered. +absl::StatusOr> CreateFileGroupsFromBlobNames( + const std::vector& filenames) { + std::vector file_groups; + if (filenames.empty()) { + return file_groups; + } + FileGroup current_group; + PS_RETURN_IF_ERROR(current_group.AddFile(filenames[0])); + for (int32_t file_index = 1; file_index < filenames.size(); file_index++) { + const auto& filename = filenames[file_index]; + if (!absl::StartsWith(filename, current_group.Basename())) { + file_groups.push_back(current_group); + current_group = FileGroup(); + } + PS_RETURN_IF_ERROR(current_group.AddFile(filename)); + } + file_groups.push_back(current_group); + return file_groups; +} + +} // namespace + +absl::StatusOr> FindMostRecentFileGroup( + const BlobStorageClient::DataLocation& location, + const FileGroupFilter& filter, BlobStorageClient& blob_client) { + PS_ASSIGN_OR_RETURN(auto file_groups, + FindFileGroups(location, filter, blob_client)); + if (!file_groups.empty()) { + return file_groups.back(); + } + return std::nullopt; +} + +absl::StatusOr> FindFileGroups( + const BlobStorageClient::DataLocation& location, + const FileGroupFilter& filter, BlobStorageClient& blob_client) { + PS_ASSIGN_OR_RETURN( + auto blob_names, + blob_client.ListBlobs(location, CreateBlobListOptions(filter))); + std::vector supported_blob_names; + std::copy_if( + blob_names.begin(), blob_names.end(), + std::back_inserter(supported_blob_names), [](std::string_view blob_name) { + return IsDeltaFilename(blob_name) || IsSnapshotFilename(blob_name) || + IsFileGroupFileName(blob_name); + }); + PS_ASSIGN_OR_RETURN(auto file_groups, + CreateFileGroupsFromBlobNames(supported_blob_names)); + if (!filter.status.has_value()) { + return file_groups; + } + std::vector result; + std::copy_if(file_groups.begin(), file_groups.end(), + std::back_inserter(result), + [&filter](const FileGroup& file_group) { + return *filter.status == file_group.GetStatus(); + }); + return result; +} + +} // namespace kv_server diff --git a/components/data/file_group/file_group_search_utils.h b/components/data/file_group/file_group_search_utils.h new file mode 100644 index 00000000..9fedbd99 --- /dev/null +++ b/components/data/file_group/file_group_search_utils.h @@ -0,0 +1,66 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef COMPONENTS_DATA_FILE_GROUP_FILE_GROUP_SEARCH_UTILS_H_ +#define COMPONENTS_DATA_FILE_GROUP_FILE_GROUP_SEARCH_UTILS_H_ + +#include +#include +#include + +#include "absl/status/statusor.h" +#include "components/data/blob_storage/blob_storage_client.h" +#include "components/data/file_group/file_group.h" +#include "public/base_types.pb.h" + +namespace kv_server { + +struct FileGroupFilter { + // Only return file groups with a basename more recent than this one. + // Return all file groups when `start_after_basename.empty()`. + std::string start_after_basename; + + // Only return file groups of this type. + // Return all file groups when `!type.has_value()`. + std::optional file_type; + + // Only consider file groups with this upload status. + // Return all file groups when `!upload_status.has_value()`. + std::optional status; +}; + +// Finds the most recent file group in blob storage that matches the `filter` +// criteria. +// +// Returns a descriptive error on failure. +absl::StatusOr> FindMostRecentFileGroup( + const BlobStorageClient::DataLocation& location, + const FileGroupFilter& filter, BlobStorageClient& blob_client); + +// Finds all file groups in blob storage that meet the `filter` criteria. +// Returned file groups are lexicographically ordered by basename. +// +// Returns (when status is ok): +// - file group: when file group matching `filer` is found. +// - std::nullopt: when file group matching `filter` is not found. +// Returns a descriptive error on failure. +absl::StatusOr> FindFileGroups( + const BlobStorageClient::DataLocation& location, + const FileGroupFilter& filter, BlobStorageClient& blob_client); + +} // namespace kv_server + +#endif // COMPONENTS_DATA_FILE_GROUP_FILE_GROUP_SEARCH_UTILS_H_ diff --git a/components/data/file_group/file_group_search_utils_test.cc b/components/data/file_group/file_group_search_utils_test.cc new file mode 100644 index 00000000..2294c7f9 --- /dev/null +++ b/components/data/file_group/file_group_search_utils_test.cc @@ -0,0 +1,152 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "components/data/file_group/file_group_search_utils.h" + +#include "components/data/blob_storage/blob_storage_client.h" +#include "components/data/common/mocks.h" +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +namespace kv_server { +namespace { + +using testing::_; +using testing::AllOf; +using testing::Field; +using testing::Return; +using testing::UnorderedElementsAre; + +TEST(FileGroupSearchUtilsTest, ValidateFindingFileGroupsBlobClientError) { + BlobStorageClient::DataLocation location{.bucket = "bucket"}; + MockBlobStorageClient blob_client; + EXPECT_CALL(blob_client, ListBlobs(_, _)) + .Times(2) + .WillRepeatedly( + Return(absl::InvalidArgumentError("bucket does not exist."))); + auto file_groups = FindFileGroups(location, FileGroupFilter{}, blob_client); + EXPECT_FALSE(file_groups.ok()) << file_groups.status(); + EXPECT_EQ(file_groups.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_EQ(file_groups.status().message(), "bucket does not exist."); + auto file_group = + FindMostRecentFileGroup(location, FileGroupFilter{}, blob_client); + EXPECT_FALSE(file_group.ok()) << file_group.status(); + EXPECT_EQ(file_group.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_EQ(file_group.status().message(), "bucket does not exist."); +} + +TEST(FileGroupSearchUtilsTest, + ValidateFindingFileGroupsBlobClientEmptyResponse) { + FileGroupFilter filter{.file_type = FileType::SNAPSHOT}; + BlobStorageClient::DataLocation location{.bucket = "bucket"}; + MockBlobStorageClient blob_client; + EXPECT_CALL( + blob_client, + ListBlobs(Field(&BlobStorageClient::DataLocation::bucket, "bucket"), + Field(&BlobStorageClient::ListOptions::prefix, "SNAPSHOT"))) + .Times(2) + .WillRepeatedly(Return(std::vector{})); + auto file_group = FindMostRecentFileGroup(location, filter, blob_client); + EXPECT_TRUE(file_group.ok()) << file_group.status(); + EXPECT_FALSE(file_group->has_value()); // No file group found. + auto file_groups = FindFileGroups(location, filter, blob_client); + EXPECT_TRUE(file_groups.ok()) << file_groups.status(); + EXPECT_TRUE(file_groups->empty()); +} + +TEST(FileGroupSearchUtilsTest, ValidateFindingSnapshotFileGroups) { + MockBlobStorageClient blob_client; + EXPECT_CALL( + blob_client, + ListBlobs(Field(&BlobStorageClient::DataLocation::bucket, "bucket"), + Field(&BlobStorageClient::ListOptions::prefix, "SNAPSHOT"))) + .WillOnce(Return(std::vector{ + "SNAPSHOT_1705430864435450_00000_OF_000003", + "SNAPSHOT_1705430864435450_00001_OF_000003", + "SNAPSHOT_1705430864435450_00002_OF_000003", + "SNAPSHOT_1705430864435451_00000_OF_000002"})); + auto file_groups = FindFileGroups( + BlobStorageClient::DataLocation{.bucket = "bucket"}, + FileGroupFilter{.file_type = FileType::SNAPSHOT}, blob_client); + EXPECT_TRUE(file_groups.ok()) << file_groups.status(); + ASSERT_EQ(file_groups->size(), 2); + // Verify first file group + EXPECT_EQ((*file_groups)[0].Size(), 3); + EXPECT_EQ((*file_groups)[0].Basename(), "SNAPSHOT_1705430864435450"); + EXPECT_EQ((*file_groups)[0].GetStatus(), FileGroup::FileStatus::kComplete); + EXPECT_EQ((*file_groups)[0].Type(), FileType::SNAPSHOT); + EXPECT_THAT( + (*file_groups)[0].Filenames(), + UnorderedElementsAre("SNAPSHOT_1705430864435450_00000_OF_000003", + "SNAPSHOT_1705430864435450_00001_OF_000003", + "SNAPSHOT_1705430864435450_00002_OF_000003")); + // Verify second file group + EXPECT_EQ((*file_groups)[1].Size(), 2); + EXPECT_EQ((*file_groups)[1].Basename(), "SNAPSHOT_1705430864435451"); + EXPECT_EQ((*file_groups)[1].GetStatus(), FileGroup::FileStatus::kPending); + EXPECT_EQ((*file_groups)[1].Type(), FileType::SNAPSHOT); + EXPECT_THAT( + (*file_groups)[1].Filenames(), + UnorderedElementsAre("SNAPSHOT_1705430864435451_00000_OF_000002")); +} + +TEST(FileGroupSearchUtilsTest, ValidateFindingMostRecentSnapshotFileGroups) { + MockBlobStorageClient blob_client; + EXPECT_CALL( + blob_client, + ListBlobs( + Field(&BlobStorageClient::DataLocation::bucket, "bucket"), + AllOf(Field(&BlobStorageClient::ListOptions::prefix, "SNAPSHOT"), + Field(&BlobStorageClient::ListOptions::start_after, + "start_after")))) + .Times(2) + .WillRepeatedly(Return(std::vector{ + "SNAPSHOT_1705430864435450_00000_OF_000003", + "SNAPSHOT_1705430864435450_00001_OF_000003", + "SNAPSHOT_1705430864435450_00002_OF_000003", + "SNAPSHOT_1705430864435451_00000_OF_000002"})); + auto file_group = FindMostRecentFileGroup( + BlobStorageClient::DataLocation{.bucket = "bucket"}, + FileGroupFilter{.start_after_basename = "start_after", + .file_type = FileType::SNAPSHOT, + .status = FileGroup::FileStatus::kComplete}, + blob_client); + ASSERT_TRUE(file_group.ok()) << file_group.status(); + EXPECT_EQ((*file_group)->Size(), 3); + EXPECT_EQ((*file_group)->Basename(), "SNAPSHOT_1705430864435450"); + EXPECT_EQ((*file_group)->GetStatus(), FileGroup::FileStatus::kComplete); + EXPECT_EQ((*file_group)->Type(), FileType::SNAPSHOT); + EXPECT_THAT( + (*file_group)->Filenames(), + UnorderedElementsAre("SNAPSHOT_1705430864435450_00000_OF_000003", + "SNAPSHOT_1705430864435450_00001_OF_000003", + "SNAPSHOT_1705430864435450_00002_OF_000003")); + file_group = FindMostRecentFileGroup( + BlobStorageClient::DataLocation{.bucket = "bucket"}, + FileGroupFilter{.start_after_basename = "start_after", + .file_type = FileType::SNAPSHOT, + .status = FileGroup::FileStatus::kPending}, + blob_client); + ASSERT_TRUE(file_group.ok()) << file_group.status(); + EXPECT_EQ((*file_group)->Size(), 2); + EXPECT_EQ((*file_group)->Basename(), "SNAPSHOT_1705430864435451"); + EXPECT_EQ((*file_group)->GetStatus(), FileGroup::FileStatus::kPending); + EXPECT_EQ((*file_group)->Type(), FileType::SNAPSHOT); + EXPECT_THAT( + (*file_group)->Filenames(), + UnorderedElementsAre("SNAPSHOT_1705430864435451_00000_OF_000002")); +} + +} // namespace +} // namespace kv_server diff --git a/components/data/file_group/file_group_test.cc b/components/data/file_group/file_group_test.cc new file mode 100644 index 00000000..2d37566b --- /dev/null +++ b/components/data/file_group/file_group_test.cc @@ -0,0 +1,122 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "components/data/file_group/file_group.h" + +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +namespace kv_server { +namespace { + +TEST(FileGroupTest, ValidateDefaultFileGroup) { + FileGroup file_group; + EXPECT_EQ(file_group.Type(), FileType::FILE_TYPE_UNSPECIFIED); + EXPECT_EQ(file_group.Size(), 0); + EXPECT_EQ(file_group.GetStatus(), FileGroup::FileStatus::kPending); + EXPECT_EQ(file_group.Basename(), ""); + EXPECT_TRUE(file_group.Filenames().empty()); +} + +TEST(FileGroupTest, ValidateAddingFilesUntilGroupIsComplete) { + auto basename = "SNAPSHOT_1705430864435450"; + auto filename0 = "SNAPSHOT_1705430864435450_00000_OF_000002"; + FileGroup file_group; + auto status = file_group.AddFile(filename0); + EXPECT_TRUE(status.ok()) << status; + EXPECT_EQ(file_group.Basename(), basename); + EXPECT_EQ(file_group.Size(), 2); + EXPECT_EQ(file_group.Type(), FileType::SNAPSHOT); + EXPECT_EQ(file_group.GetStatus(), FileGroup::FileStatus::kPending); + EXPECT_THAT(file_group.Filenames(), testing::UnorderedElementsAre(filename0)); + // Add the remaining file + auto filename1 = "SNAPSHOT_1705430864435450_00001_OF_000002"; + status = file_group.AddFile(filename1); + EXPECT_TRUE(status.ok()) << status; + EXPECT_EQ(file_group.Basename(), basename); + EXPECT_EQ(file_group.Size(), 2); + EXPECT_EQ(file_group.Type(), FileType::SNAPSHOT); + EXPECT_EQ(file_group.GetStatus(), FileGroup::FileStatus::kComplete); + EXPECT_THAT(file_group.Filenames(), + testing::UnorderedElementsAre(filename0, filename1)); +} + +TEST(FileGroupTest, ValidateAddingSingleDeltaFile) { + auto delta_file = "DELTA_1705430864435450"; + FileGroup file_group; + auto status = file_group.AddFile(delta_file); + EXPECT_TRUE(status.ok()) << status; + EXPECT_EQ(file_group.Basename(), delta_file); + EXPECT_EQ(file_group.Size(), 1); + EXPECT_EQ(file_group.Type(), FileType::DELTA); + EXPECT_EQ(file_group.GetStatus(), FileGroup::FileStatus::kComplete); + EXPECT_THAT(file_group.Filenames(), + testing::UnorderedElementsAre(delta_file)); +} + +TEST(FileGroupTest, ValidateAddingSingleSnapshotFile) { + auto snapshot_file = "SNAPSHOT_1705430864435450"; + FileGroup file_group; + auto status = file_group.AddFile(snapshot_file); + EXPECT_TRUE(status.ok()) << status; + EXPECT_EQ(file_group.Basename(), snapshot_file); + EXPECT_EQ(file_group.Size(), 1); + EXPECT_EQ(file_group.Type(), FileType::SNAPSHOT); + EXPECT_EQ(file_group.GetStatus(), FileGroup::FileStatus::kComplete); + EXPECT_THAT(file_group.Filenames(), + testing::UnorderedElementsAre(snapshot_file)); +} + +TEST(FileGroupTest, ValidateAddingInvalidFilenames) { + FileGroup file_group; + auto status = file_group.AddFile(""); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + status = file_group.AddFile("UNKNOWN_1705430864435450_00000_OF_000010"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + status = file_group.AddFile("SNAPSHOT_1705430864435450_00000"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + status = file_group.AddFile("SNAPSHOT_1705430864435450_00000_OF"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + status = file_group.AddFile("DELTA_1705430864435450_OF_000010"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + status = file_group.AddFile("SNAPSHOT_1705430864435450_00010_OF_000010"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); +} + +TEST(FileGroupTest, ValidateAddingFilesWithMismatchingMetadata) { + FileGroup file_group; + auto status = file_group.AddFile("SNAPSHOT_1705430864435450_00000_OF_000010"); + EXPECT_TRUE(status.ok()) << status; + // check logical commit time mismatch + status = file_group.AddFile("SNAPSHOT_1705430000000000_00000_OF_000010"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + // check file type mismatch + status = file_group.AddFile("DELTA_1705430864435450_00001_OF_000010"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + // check group size mismatch + status = file_group.AddFile("SNAPSHOT_1705430864435450_00001_OF_000020"); + EXPECT_FALSE(status.ok()) << status; + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); +} + +} // namespace +} // namespace kv_server diff --git a/components/data/realtime/BUILD.bazel b/components/data/realtime/BUILD.bazel index 5c8f5e34..f8dc66c1 100644 --- a/components/data/realtime/BUILD.bazel +++ b/components/data/realtime/BUILD.bazel @@ -40,7 +40,7 @@ cc_library( "//components/data/common:change_notifier", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//src/cpp/telemetry", + "@google_privacysandbox_servers_common//src/telemetry", ], ) @@ -64,10 +64,9 @@ cc_test( ":delta_file_record_change_notifier", "//components/data/common:change_notifier", "//components/data/common:mocks", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -95,11 +94,11 @@ cc_library( "//components/errors:retry", "//components/util:sleepfor", "//public:constants", - "@google_privacysandbox_servers_common//src/cpp/util:duration", + "@google_privacysandbox_servers_common//src/util:duration", ], }) + [ "//components/data/common:thread_manager", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", ], @@ -130,7 +129,6 @@ cc_test( "//public/data_loading:filename_utils", "@com_github_grpc_grpc//:grpc++", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -151,8 +149,6 @@ cc_library( ":realtime_notifier", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -180,6 +176,5 @@ cc_test( "//components/util:sleepfor_mock", "@com_github_grpc_grpc//:grpc++", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) diff --git a/components/data/realtime/delta_file_record_change_notifier.h b/components/data/realtime/delta_file_record_change_notifier.h index 6fd66c6f..d48cda20 100644 --- a/components/data/realtime/delta_file_record_change_notifier.h +++ b/components/data/realtime/delta_file_record_change_notifier.h @@ -26,7 +26,7 @@ #include "absl/status/statusor.h" #include "absl/time/time.h" #include "components/data/common/change_notifier.h" -#include "src/cpp/telemetry/telemetry.h" +#include "src/telemetry/telemetry.h" namespace kv_server { diff --git a/components/data/realtime/delta_file_record_change_notifier_aws.cc b/components/data/realtime/delta_file_record_change_notifier_aws.cc index 9bc72b14..5d0cb768 100644 --- a/components/data/realtime/delta_file_record_change_notifier_aws.cc +++ b/components/data/realtime/delta_file_record_change_notifier_aws.cc @@ -14,6 +14,7 @@ #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/escaping.h" @@ -21,8 +22,7 @@ #include "components/data/common/change_notifier.h" #include "components/data/realtime/delta_file_record_change_notifier.h" #include "components/telemetry/server_definition.h" -#include "glog/logging.h" -#include "src/cpp/telemetry/telemetry.h" +#include "src/telemetry/telemetry.h" namespace kv_server { namespace { @@ -65,12 +65,7 @@ class AwsDeltaFileRecordChangeNotifier : public DeltaFileRecordChangeNotifier { if (!parsedMessage.ok()) { LOG(ERROR) << "Failed to parse JSON: " << message << ", error: " << parsedMessage.status(); - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kDeltaFileRecordChangeNotifierParsingFailure), - 1}})); + LogServerErrorMetric(kDeltaFileRecordChangeNotifierParsingFailure); continue; } nc.realtime_messages.push_back(RealtimeMessage{ diff --git a/components/data/realtime/delta_file_record_change_notifier_aws_test.cc b/components/data/realtime/delta_file_record_change_notifier_aws_test.cc index 131e0444..e2f91460 100644 --- a/components/data/realtime/delta_file_record_change_notifier_aws_test.cc +++ b/components/data/realtime/delta_file_record_change_notifier_aws_test.cc @@ -25,7 +25,6 @@ #include "components/data/common/mocks.h" #include "components/data/realtime/delta_file_record_change_notifier.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" using testing::_; diff --git a/components/data/realtime/delta_file_record_change_notifier_local.cc b/components/data/realtime/delta_file_record_change_notifier_local.cc index a7ad5906..f444cdf3 100644 --- a/components/data/realtime/delta_file_record_change_notifier_local.cc +++ b/components/data/realtime/delta_file_record_change_notifier_local.cc @@ -14,11 +14,11 @@ #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "components/data/common/change_notifier.h" #include "components/data/realtime/delta_file_record_change_notifier.h" -#include "glog/logging.h" namespace kv_server { namespace { diff --git a/components/data/realtime/delta_file_record_change_notifier_local_test.cc b/components/data/realtime/delta_file_record_change_notifier_local_test.cc index 6a33768c..ef3589e3 100644 --- a/components/data/realtime/delta_file_record_change_notifier_local_test.cc +++ b/components/data/realtime/delta_file_record_change_notifier_local_test.cc @@ -23,7 +23,6 @@ #include "components/data/common/mocks.h" #include "components/data/realtime/delta_file_record_change_notifier.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/data/realtime/realtime_notifier_aws.cc b/components/data/realtime/realtime_notifier_aws.cc index 19b5f917..56f4cc6d 100644 --- a/components/data/realtime/realtime_notifier_aws.cc +++ b/components/data/realtime/realtime_notifier_aws.cc @@ -18,15 +18,15 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "components/data/common/thread_manager.h" #include "components/data/realtime/delta_file_record_change_notifier.h" #include "components/data/realtime/realtime_notifier.h" #include "components/errors/retry.h" -#include "glog/logging.h" #include "public/constants.h" -#include "src/cpp/telemetry/telemetry.h" -#include "src/cpp/util/duration.h" +#include "src/telemetry/telemetry.h" +#include "src/util/duration.h" namespace kv_server { namespace { @@ -36,7 +36,7 @@ class RealtimeNotifierImpl : public RealtimeNotifier { explicit RealtimeNotifierImpl( std::unique_ptr sleep_for, std::unique_ptr change_notifier) - : thread_manager_(TheadManager::Create("Realtime notifier")), + : thread_manager_(ThreadManager::Create("Realtime notifier")), sleep_for_(std::move(sleep_for)), change_notifier_(std::move(change_notifier)) {} @@ -82,48 +82,50 @@ class RealtimeNotifierImpl : public RealtimeNotifier { ExponentialBackoffForRetry(sequential_failures); LOG(ERROR) << "Failed to get realtime notifications: " << updates.status() << ". Waiting for " << backoff_time; - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kRealtimeGetNotificationsFailure), 1}})); + LogServerErrorMetric(kRealtimeGetNotificationsFailure); if (!sleep_for_->Duration(backoff_time)) { LOG(ERROR) << "Failed to sleep for " << backoff_time << ". SleepFor invalid."; - LogIfError(KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - {{std::string(kRealtimeSleepFailure), 1}})); + LogServerErrorMetric(kRealtimeSleepFailure); } continue; } sequential_failures = 0; for (const auto& realtime_message : updates->realtime_messages) { - auto count = callback(realtime_message.parsed_notification); - if (count.ok()) { - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - static_cast(count->total_updated_records + - count->total_deleted_records))); + if (auto count = callback(realtime_message.parsed_notification); + !count.ok()) { + LOG(ERROR) << "Data loading callback failed: " << count.status(); + LogServerErrorMetric(kRealtimeMessageApplicationFailure); + } + auto e2e_cloud_provided_latency = absl::ToDoubleMicroseconds( + absl::Now() - realtime_message.notifications_sns_inserted); + // we're getting this value based on two different clocks. Opentelemetry + // does not allow negative values for histograms. However, not logging + // this will affect the pvalues, so the next best thing is set it to 0. + if (e2e_cloud_provided_latency < 0) { + e2e_cloud_provided_latency = 0; } LogIfError( KVServerContextMap() ->SafeMetric() .LogHistogram( - absl::ToDoubleMicroseconds( - absl::Now() - - (realtime_message.notifications_sns_inserted)))); + e2e_cloud_provided_latency)); if (realtime_message.notifications_inserted) { - auto e2eDuration = - absl::Now() - (realtime_message.notifications_inserted).value(); + // we're getting this value based on two different clocks. + // Opentelemetry does not allow negative values for histograms. + // However, not logging this will affect the pvalues, so the next best + // thing is set it to 0. + auto e2e_latency = absl::ToDoubleMicroseconds( + absl::Now() - realtime_message.notifications_inserted.value()); + if (e2e_latency < 0) { + e2e_latency = 0; + } LogIfError(KVServerContextMap() ->SafeMetric() .LogHistogram( - absl::ToDoubleMicroseconds(e2eDuration))); + e2e_latency)); } } LogIfError(KVServerContextMap() @@ -138,7 +140,7 @@ class RealtimeNotifierImpl : public RealtimeNotifier { } } - std::unique_ptr thread_manager_; + std::unique_ptr thread_manager_; std::unique_ptr sleep_for_; std::unique_ptr change_notifier_; }; diff --git a/components/data/realtime/realtime_notifier_aws_test.cc b/components/data/realtime/realtime_notifier_aws_test.cc index 54a12c5e..238acdc1 100644 --- a/components/data/realtime/realtime_notifier_aws_test.cc +++ b/components/data/realtime/realtime_notifier_aws_test.cc @@ -24,7 +24,6 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" #include "public/data_loading/filename_utils.h" -#include "src/cpp/telemetry/mocks.h" using testing::_; using testing::Field; diff --git a/components/data/realtime/realtime_notifier_gcp.cc b/components/data/realtime/realtime_notifier_gcp.cc index d76c2ea8..073ed9a4 100644 --- a/components/data/realtime/realtime_notifier_gcp.cc +++ b/components/data/realtime/realtime_notifier_gcp.cc @@ -14,15 +14,15 @@ #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/escaping.h" #include "components/data/common/msg_svc.h" #include "components/data/common/thread_manager.h" #include "components/data/realtime/realtime_notifier.h" -#include "glog/logging.h" #include "google/cloud/pubsub/message.h" #include "google/cloud/pubsub/subscriber.h" -#include "src/cpp/telemetry/telemetry.h" +#include "src/telemetry/telemetry.h" namespace kv_server { namespace { @@ -37,7 +37,7 @@ class RealtimeNotifierGcp : public RealtimeNotifier { public: explicit RealtimeNotifierGcp(std::unique_ptr gcp_subscriber, std::unique_ptr sleep_for) - : thread_manager_(TheadManager::Create("Realtime notifier")), + : thread_manager_(ThreadManager::Create("Realtime notifier")), sleep_for_(std::move(sleep_for)), gcp_subscriber_(std::move(gcp_subscriber)) {} @@ -113,24 +113,14 @@ class RealtimeNotifierGcp : public RealtimeNotifier { auto start = absl::Now(); std::string string_decoded; if (!absl::Base64Unescape(m.data(), &string_decoded)) { - LogIfError( - KVServerContextMap()->SafeMetric().LogUpDownCounter( - {{std::string(kRealtimeDecodeMessageFailure), 1}})); + LogServerErrorMetric(kRealtimeDecodeMessageFailure); LOG(ERROR) << "The body of the message is not a base64 encoded string."; std::move(h).ack(); return; } - auto count = callback(string_decoded); - if (count.ok()) { - LogIfError(KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - static_cast(count->total_updated_records + - count->total_deleted_records))); - } else { - LogIfError( - KVServerContextMap()->SafeMetric().LogUpDownCounter( - {{std::string(kRealtimeMessageApplicationFailure), 1}})); + if (auto count = callback(string_decoded); !count.ok()) { + LOG(ERROR) << "Data loading callback failed: " << count.status(); + LogServerErrorMetric(kRealtimeMessageApplicationFailure); } RecordGcpSuppliedE2ELatency(m); RecordProducerSuppliedE2ELatency(m); @@ -156,7 +146,7 @@ class RealtimeNotifierGcp : public RealtimeNotifier { LOG(INFO) << "Realtime updater stopped watching."; } - std::unique_ptr thread_manager_; + std::unique_ptr thread_manager_; mutable absl::Mutex mutex_; future session_ ABSL_GUARDED_BY(mutex_); std::unique_ptr sleep_for_; diff --git a/components/data/realtime/realtime_notifier_gcp_test.cc b/components/data/realtime/realtime_notifier_gcp_test.cc index 31d731c1..fa6e5603 100644 --- a/components/data/realtime/realtime_notifier_gcp_test.cc +++ b/components/data/realtime/realtime_notifier_gcp_test.cc @@ -29,7 +29,6 @@ #include "google/cloud/pubsub/subscriber.h" #include "gtest/gtest.h" #include "public/data_loading/filename_utils.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/data/realtime/realtime_thread_pool_manager.h b/components/data/realtime/realtime_thread_pool_manager.h index 8c1540e8..07371c0e 100644 --- a/components/data/realtime/realtime_thread_pool_manager.h +++ b/components/data/realtime/realtime_thread_pool_manager.h @@ -25,7 +25,7 @@ #include "absl/status/statusor.h" #include "components/data/common/notifier_metadata.h" #include "components/data/realtime/realtime_notifier.h" -#include "src/cpp/telemetry/telemetry.h" +#include "src/telemetry/telemetry.h" namespace kv_server { diff --git a/components/data/realtime/realtime_thread_pool_manager_aws_test.cc b/components/data/realtime/realtime_thread_pool_manager_aws_test.cc index 4fa49b1f..3ee12dc2 100644 --- a/components/data/realtime/realtime_thread_pool_manager_aws_test.cc +++ b/components/data/realtime/realtime_thread_pool_manager_aws_test.cc @@ -24,7 +24,6 @@ #include "components/util/sleepfor_mock.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/data/realtime/realtime_thread_pool_manager_gcp_test.cc b/components/data/realtime/realtime_thread_pool_manager_gcp_test.cc index 584083b9..b56e4282 100644 --- a/components/data/realtime/realtime_thread_pool_manager_gcp_test.cc +++ b/components/data/realtime/realtime_thread_pool_manager_gcp_test.cc @@ -25,7 +25,6 @@ #include "google/cloud/pubsub/mocks/mock_subscriber_connection.h" #include "google/cloud/pubsub/subscriber.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/data_server/cache/BUILD.bazel b/components/data_server/cache/BUILD.bazel index 06112b56..d6eaecc1 100644 --- a/components/data_server/cache/BUILD.bazel +++ b/components/data_server/cache/BUILD.bazel @@ -40,6 +40,7 @@ cc_library( ], deps = [ ":get_key_value_set_result_impl", + "//components/util:request_context", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", ], @@ -57,12 +58,11 @@ cc_library( ":cache", ":get_key_value_set_result_impl", "//public:base_types_cc_proto", - "@com_github_google_glog//:glog", "@com_google_absl//absl/base", "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", "@com_google_absl//absl/synchronization", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -80,8 +80,7 @@ cc_test( "@com_google_absl//absl/container:flat_hash_map", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", ], ) diff --git a/components/data_server/cache/cache.h b/components/data_server/cache/cache.h index 1136e1f0..4341cff7 100644 --- a/components/data_server/cache/cache.h +++ b/components/data_server/cache/cache.h @@ -27,6 +27,7 @@ #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "components/data_server/cache/get_key_value_set_result.h" +#include "components/util/request_context.h" namespace kv_server { @@ -38,35 +39,42 @@ class Cache { // Looks up and returns key-value pairs for the given keys. virtual absl::flat_hash_map GetKeyValuePairs( + const RequestContext& request_context, const absl::flat_hash_set& key_list) const = 0; // Looks up and returns key-value set result for the given key set. virtual std::unique_ptr GetKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const = 0; - // Inserts or updates the key with the new value. + // Inserts or updates the key with the new value for a given prefix virtual void UpdateKeyValue(std::string_view key, std::string_view value, - int64_t logical_commit_time) = 0; + int64_t logical_commit_time, + std::string_view prefix = "") = 0; - // Inserts or updates values in the set for a given key, if a value exists, - // updates its timestamp to the latest logical commit time. + // Inserts or updates values in the set for a given key and prefix, if a value + // exists, updates its timestamp to the latest logical commit time. virtual void UpdateKeyValueSet(std::string_view key, absl::Span value_set, - int64_t logical_commit_time) = 0; + int64_t logical_commit_time, + std::string_view prefix = "") = 0; - // Deletes a particular (key, value) pair. - virtual void DeleteKey(std::string_view key, int64_t logical_commit_time) = 0; + // Deletes a particular (key, value) pair for a given prefix. + virtual void DeleteKey(std::string_view key, int64_t logical_commit_time, + std::string_view prefix = "") = 0; - // Deletes values in the set for a given key. The deletion, this object - // still exist and is marked "deleted", in case there are - // late-arriving updates to this value. + // Deletes values in the set for a given key and prefix. The deletion, this + // object still exist and is marked "deleted", in case there are late-arriving + // updates to this value. virtual void DeleteValuesInSet(std::string_view key, absl::Span value_set, - int64_t logical_commit_time) = 0; + int64_t logical_commit_time, + std::string_view prefix = "") = 0; // Removes the values that were deleted before the specified - // logical_commit_time. - virtual void RemoveDeletedKeys(int64_t logical_commit_time) = 0; + // logical_commit_time for a given prefix. + virtual void RemoveDeletedKeys(int64_t logical_commit_time, + std::string_view prefix = "") = 0; }; } // namespace kv_server diff --git a/components/data_server/cache/key_value_cache.cc b/components/data_server/cache/key_value_cache.cc index 3d4866ff..f984de36 100644 --- a/components/data_server/cache/key_value_cache.cc +++ b/components/data_server/cache/key_value_cache.cc @@ -19,32 +19,20 @@ #include #include "absl/container/flat_hash_map.h" +#include "absl/log/log.h" #include "absl/memory/memory.h" #include "absl/synchronization/mutex.h" #include "components/data_server/cache/cache.h" #include "components/data_server/cache/get_key_value_set_result.h" -#include "glog/logging.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { -using privacy_sandbox::server_common::MetricsRecorder; -using privacy_sandbox::server_common::ScopeLatencyRecorder; - -constexpr char kGetKeyValuePairsEvent[] = "GetKeyValuePairs"; -constexpr char kGetKeyValueSetEvent[] = "GetKeyValueSet"; -constexpr char kUpdateKeyValueEvent[] = "UpdateKeyValue"; -constexpr char kUpdateKeyValueSetEvent[] = "UpdateKeyValueSet"; -constexpr char kDeleteKeyEvent[] = "DeleteKey"; -constexpr char kDeleteValuesInSetEvent[] = "DeleteValuesInSet"; -constexpr char kRemoveDeletedKeysEvent[] = "RemoveDeletedKeys"; -constexpr char kCleanUpKeyValueMapEvent[] = "CleanUpKeyValueMap"; -constexpr char kCleanUpKeyValueSetMapEvent[] = "CleanUpKeyValueSetMap"; - absl::flat_hash_map KeyValueCache::GetKeyValuePairs( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const { - ScopeLatencyRecorder latency_recorder(kGetKeyValuePairsEvent, - metrics_recorder_); + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetInternalLookupMetricsContext()); absl::flat_hash_map kv_pairs; absl::ReaderMutexLock lock(&mutex_); for (std::string_view key : key_set) { @@ -57,16 +45,24 @@ absl::flat_hash_map KeyValueCache::GetKeyValuePairs( kv_pairs.insert_or_assign(key, *(key_iter->second.value)); } } + if (kv_pairs.empty()) { + LogCacheAccessMetrics(request_context, kKeyValueCacheMiss); + } else { + LogCacheAccessMetrics(request_context, kKeyValueCacheHit); + } return kv_pairs; } std::unique_ptr KeyValueCache::GetKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const { - ScopeLatencyRecorder latency_recorder(kGetKeyValueSetEvent, - metrics_recorder_); + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetInternalLookupMetricsContext()); // lock the cache map absl::ReaderMutexLock lock(&set_map_mutex_); auto result = GetKeyValueSetResult::Create(); + bool cache_hit = false; for (const auto& key : key_set) { VLOG(8) << "Getting key: " << key; const auto key_itr = key_to_value_set_map_.find(key); @@ -81,25 +77,35 @@ std::unique_ptr KeyValueCache::GetKeyValueSet( } // Add key value set to the result result->AddKeyValueSet(key, std::move(value_set), std::move(set_lock)); + cache_hit = true; } } + if (cache_hit) { + LogCacheAccessMetrics(request_context, kKeyValueSetCacheHit); + } else { + LogCacheAccessMetrics(request_context, kKeyValueSetCacheMiss); + } return result; } // Replaces the current key-value entry with the new key-value entry. void KeyValueCache::UpdateKeyValue(std::string_view key, std::string_view value, - int64_t logical_commit_time) { - ScopeLatencyRecorder latency_recorder(kUpdateKeyValueEvent, - metrics_recorder_); + int64_t logical_commit_time, + std::string_view prefix) { + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); VLOG(9) << "Received update for [" << key << "] at " << logical_commit_time << ". value will be set to: " << value; absl::MutexLock lock(&mutex_); - if (logical_commit_time <= max_cleanup_logical_commit_time_) { + auto max_cleanup_logical_commit_time = + max_cleanup_logical_commit_time_map_[prefix]; + + if (logical_commit_time <= max_cleanup_logical_commit_time) { VLOG(1) << "Skipping the update as its logical_commit_time: " << logical_commit_time << " is not newer than the current cutoff time:" - << max_cleanup_logical_commit_time_; + << max_cleanup_logical_commit_time; return; } @@ -119,10 +125,15 @@ void KeyValueCache::UpdateKeyValue(std::string_view key, std::string_view value, key_iter->second.last_logical_commit_time < logical_commit_time && key_iter->second.value == nullptr) { // should always have this, but checking just in case - auto dl_key_iter = - deleted_nodes_.find(key_iter->second.last_logical_commit_time); - if (dl_key_iter != deleted_nodes_.end() && dl_key_iter->second == key) { - deleted_nodes_.erase(dl_key_iter); + + if (auto prefix_deleted_nodes_iter = deleted_nodes_map_.find(prefix); + prefix_deleted_nodes_iter != deleted_nodes_map_.end()) { + auto dl_key_iter = prefix_deleted_nodes_iter->second.find( + key_iter->second.last_logical_commit_time); + if (dl_key_iter != prefix_deleted_nodes_iter->second.end() && + dl_key_iter->second == key) { + prefix_deleted_nodes_iter->second.erase(dl_key_iter); + } } } @@ -132,9 +143,10 @@ void KeyValueCache::UpdateKeyValue(std::string_view key, std::string_view value, void KeyValueCache::UpdateKeyValueSet( std::string_view key, absl::Span input_value_set, - int64_t logical_commit_time) { - ScopeLatencyRecorder latency_recorder(kUpdateKeyValueSetEvent, - metrics_recorder_); + int64_t logical_commit_time, std::string_view prefix) { + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); VLOG(9) << "Received update for [" << key << "] at " << logical_commit_time; std::unique_ptr key_lock; absl::flat_hash_map* existing_value_set; @@ -142,11 +154,14 @@ void KeyValueCache::UpdateKeyValueSet( { absl::MutexLock lock_map(&set_map_mutex_); - if (logical_commit_time <= max_cleanup_logical_commit_time_for_set_cache_) { + auto max_cleanup_logical_commit_time = + max_cleanup_logical_commit_time_map_for_set_cache_[prefix]; + + if (logical_commit_time <= max_cleanup_logical_commit_time) { VLOG(1) << "Skipping the update as its logical_commit_time: " << logical_commit_time << " is older than the current cutoff time:" - << max_cleanup_logical_commit_time_for_set_cache_; + << max_cleanup_logical_commit_time; return; } else if (input_value_set.empty()) { VLOG(1) << "Skipping the update as it has no value in the set."; @@ -191,11 +206,14 @@ void KeyValueCache::UpdateKeyValueSet( // end locking key } -void KeyValueCache::DeleteKey(std::string_view key, - int64_t logical_commit_time) { - ScopeLatencyRecorder latency_recorder(kDeleteKeyEvent, metrics_recorder_); +void KeyValueCache::DeleteKey(std::string_view key, int64_t logical_commit_time, + std::string_view prefix) { + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); absl::MutexLock lock(&mutex_); - if (logical_commit_time <= max_cleanup_logical_commit_time_) { + auto max_cleanup_logical_commit_time = + max_cleanup_logical_commit_time_map_[prefix]; + if (logical_commit_time <= max_cleanup_logical_commit_time) { return; } const auto key_iter = map_.find(key); @@ -208,23 +226,25 @@ void KeyValueCache::DeleteKey(std::string_view key, map_.insert_or_assign( key, {.value = nullptr, .last_logical_commit_time = logical_commit_time}); - - auto result = deleted_nodes_.emplace(logical_commit_time, key); + auto result = deleted_nodes_map_[prefix].emplace(logical_commit_time, key); } } void KeyValueCache::DeleteValuesInSet(std::string_view key, absl::Span value_set, - int64_t logical_commit_time) { - ScopeLatencyRecorder latency_recorder(kDeleteValuesInSetEvent, - metrics_recorder_); + int64_t logical_commit_time, + std::string_view prefix) { + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); std::unique_ptr key_lock; absl::flat_hash_map* existing_value_set; // The max cleanup time needs to be locked before doing this comparison { absl::MutexLock lock_map(&set_map_mutex_); - - if (logical_commit_time <= max_cleanup_logical_commit_time_for_set_cache_ || + auto max_cleanup_logical_commit_time = + max_cleanup_logical_commit_time_map_for_set_cache_[prefix]; + if (logical_commit_time <= max_cleanup_logical_commit_time || value_set.empty()) { return; } @@ -243,7 +263,7 @@ void KeyValueCache::DeleteValuesInSet(std::string_view key, key_to_value_set_map_.emplace(key, std::move(mutex_value_map_pair)); // Add to deleted set nodes for (const std::string_view value : value_set) { - deleted_set_nodes_[logical_commit_time][key].emplace(value); + deleted_set_nodes_map_[prefix][logical_commit_time][key].emplace(value); } return; } @@ -273,25 +293,36 @@ void KeyValueCache::DeleteValuesInSet(std::string_view key, key_lock.reset(); absl::MutexLock lock_map(&set_map_mutex_); for (const std::string_view value : values_to_delete) { - deleted_set_nodes_[logical_commit_time][key].emplace(value); + deleted_set_nodes_map_[prefix][logical_commit_time][key].emplace(value); } } } -void KeyValueCache::RemoveDeletedKeys(int64_t logical_commit_time) { - ScopeLatencyRecorder latency_recorder(kRemoveDeletedKeysEvent, - metrics_recorder_); - CleanUpKeyValueMap(logical_commit_time); - CleanUpKeyValueSetMap(logical_commit_time); +void KeyValueCache::RemoveDeletedKeys(int64_t logical_commit_time, + std::string_view prefix) { + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); + CleanUpKeyValueMap(logical_commit_time, prefix); + CleanUpKeyValueSetMap(logical_commit_time, prefix); } -void KeyValueCache::CleanUpKeyValueMap(int64_t logical_commit_time) { - ScopeLatencyRecorder latency_recorder(kCleanUpKeyValueMapEvent, - metrics_recorder_); +void KeyValueCache::CleanUpKeyValueMap(int64_t logical_commit_time, + std::string_view prefix) { + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); absl::MutexLock lock(&mutex_); - auto it = deleted_nodes_.begin(); + if (max_cleanup_logical_commit_time_map_[prefix] < logical_commit_time) { + max_cleanup_logical_commit_time_map_[prefix] = logical_commit_time; + } + auto deleted_nodes_per_prefix = deleted_nodes_map_.find(prefix); + if (deleted_nodes_per_prefix == deleted_nodes_map_.end()) { + return; + } + auto it = deleted_nodes_per_prefix->second.begin(); - while (it != deleted_nodes_.end()) { + while (it != deleted_nodes_per_prefix->second.end()) { if (it->first > logical_commit_time) { break; } @@ -305,17 +336,30 @@ void KeyValueCache::CleanUpKeyValueMap(int64_t logical_commit_time) { ++it; } - deleted_nodes_.erase(deleted_nodes_.begin(), it); - max_cleanup_logical_commit_time_ = - std::max(max_cleanup_logical_commit_time_, logical_commit_time); + deleted_nodes_per_prefix->second.erase( + deleted_nodes_per_prefix->second.begin(), it); + if (deleted_nodes_per_prefix->second.empty()) { + deleted_nodes_map_.erase(prefix); + } } -void KeyValueCache::CleanUpKeyValueSetMap(int64_t logical_commit_time) { - ScopeLatencyRecorder latency_recorder(kCleanUpKeyValueSetMapEvent, - metrics_recorder_); +void KeyValueCache::CleanUpKeyValueSetMap(int64_t logical_commit_time, + std::string_view prefix) { + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); absl::MutexLock lock_set_map(&set_map_mutex_); - auto delete_itr = deleted_set_nodes_.begin(); - while (delete_itr != deleted_set_nodes_.end()) { + if (max_cleanup_logical_commit_time_map_for_set_cache_[prefix] < + logical_commit_time) { + max_cleanup_logical_commit_time_map_for_set_cache_[prefix] = + logical_commit_time; + } + auto deleted_nodes_per_prefix = deleted_set_nodes_map_.find(prefix); + if (deleted_nodes_per_prefix == deleted_set_nodes_map_.end()) { + return; + } + auto delete_itr = deleted_nodes_per_prefix->second.begin(); + while (delete_itr != deleted_nodes_per_prefix->second.end()) { if (delete_itr->first > logical_commit_time) { break; } @@ -341,13 +385,22 @@ void KeyValueCache::CleanUpKeyValueSetMap(int64_t logical_commit_time) { } ++delete_itr; } - deleted_set_nodes_.erase(deleted_set_nodes_.begin(), delete_itr); - max_cleanup_logical_commit_time_for_set_cache_ = std::max( - max_cleanup_logical_commit_time_for_set_cache_, logical_commit_time); + deleted_nodes_per_prefix->second.erase( + deleted_nodes_per_prefix->second.begin(), delete_itr); + if (deleted_nodes_per_prefix->second.empty()) { + deleted_set_nodes_map_.erase(prefix); + } +} + +void KeyValueCache::LogCacheAccessMetrics( + const RequestContext& request_context, + std::string_view cache_access_event) const { + LogIfError( + request_context.GetInternalLookupMetricsContext() + .AccumulateMetric(1, cache_access_event)); } -std::unique_ptr KeyValueCache::Create( - MetricsRecorder& metrics_recorder) { - return absl::WrapUnique(new KeyValueCache(metrics_recorder)); +std::unique_ptr KeyValueCache::Create() { + return absl::WrapUnique(new KeyValueCache()); } } // namespace kv_server diff --git a/components/data_server/cache/key_value_cache.h b/components/data_server/cache/key_value_cache.h index dd327427..aeae908c 100644 --- a/components/data_server/cache/key_value_cache.h +++ b/components/data_server/cache/key_value_cache.h @@ -32,53 +32,54 @@ #include "components/data_server/cache/cache.h" #include "components/data_server/cache/get_key_value_set_result.h" #include "public/base_types.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { // In-memory datastore. // One cache object is only for keys in one namespace. class KeyValueCache : public Cache { public: - KeyValueCache( - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) - : metrics_recorder_(metrics_recorder) {} - // Looks up and returns key-value pairs for the given keys. absl::flat_hash_map GetKeyValuePairs( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const override; // Looks up and returns key-value set result for the given key set. std::unique_ptr GetKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const override; - // Inserts or updates the key with the new value. + // Inserts or updates the key with the new value for a given prefix void UpdateKeyValue(std::string_view key, std::string_view value, - int64_t logical_commit_time) override; + int64_t logical_commit_time, + std::string_view prefix = "") override; - // Inserts or updates values in the set for a given key, if a value exists, - // updates its timestamp to the latest logical commit time. + // Inserts or updates values in the set for a given key and prefix, if a value + // exists, updates its timestamp to the latest logical commit time. void UpdateKeyValueSet(std::string_view key, absl::Span input_value_set, - int64_t logical_commit_time) override; + int64_t logical_commit_time, + std::string_view prefix = "") override; - // Deletes a particular (key, value) pair. - void DeleteKey(std::string_view key, int64_t logical_commit_time) override; + // Deletes a particular (key, value) pair for a given prefix. + void DeleteKey(std::string_view key, int64_t logical_commit_time, + std::string_view prefix = "") override; - // Deletes values in the set for a given key. The deletion, this object - // still exist and is marked "deleted", in case there are - // late-arriving updates to this value. + // Deletes values in the set for a given key and prefix. The deletion, this + // object still exist and is marked "deleted", in case there are late-arriving + // updates to this value. void DeleteValuesInSet(std::string_view key, absl::Span value_set, - int64_t logical_commit_time) override; + int64_t logical_commit_time, + std::string_view prefix = "") override; // Removes the values that were deleted before the specified - // logical_commit_time. + // logical_commit_time for a given prefix. // TODO: b/267182790 -- Cache cleanup should be done periodically from a // background thread - void RemoveDeletedKeys(int64_t logical_commit_time) override; + void RemoveDeletedKeys(int64_t logical_commit_time, + std::string_view prefix = "") override; - static std::unique_ptr Create( - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder); + static std::unique_ptr Create(); private: struct CacheValue { @@ -113,18 +114,25 @@ class KeyValueCache : public Cache { // Sorted mapping from the logical timestamp to a key, for nodes that were // deleted We keep this to do proper and efficient clean up in map_. - std::multimap deleted_nodes_ ABSL_GUARDED_BY(mutex_); + // The key in the inner map is the prefix and the value is the keys of key + // value pairs deleted from that prefix + + absl::flat_hash_map> + deleted_nodes_map_ ABSL_GUARDED_BY(mutex_); - // The maximum value that was passed to RemoveDeletedKeys. - int64_t max_cleanup_logical_commit_time_ ABSL_GUARDED_BY(mutex_) = 0; + // The key is the prefix and the value is the + // maximum timestamp that was passed to RemoveDeletedKeys. + absl::flat_hash_map max_cleanup_logical_commit_time_map_ + ABSL_GUARDED_BY(mutex_); - // The maximum value of logical commit time that is used to do update/delete - // for key-value set map. + // The key is the prefix and the value is the maximum + // logical commit time that is used to do update/delete for key-value set map. // TODO(b/284474892) Need to evaluate if we really need to make this variable // guarded b mutex, if not, we may want to remove it and use one // max_cleanup_logical_commit_time in update/deletion for both maps - int64_t max_cleanup_logical_commit_time_for_set_cache_ - ABSL_GUARDED_BY(set_map_mutex_) = 0; + absl::flat_hash_map + max_cleanup_logical_commit_time_map_for_set_cache_ + ABSL_GUARDED_BY(set_map_mutex_); // Mapping from a key to its value map. The key in the inner map is the // value string, and value is the ValueMeta. The inner map allows value @@ -136,23 +144,30 @@ class KeyValueCache : public Cache { std::unique_ptr>>> key_to_value_set_map_ ABSL_GUARDED_BY(set_map_mutex_); - // Sorted mapping from logical timestamp to key-value_set map to keep track of + // The key of outer map is the prefix, and value is the sorted mapping + // from logical timestamp to key-value_set map to keep track of // deleted key-values to handle out of order update case. In the inner map, // the key string is the key for the values, and the string // in the flat_hash_set is the value - absl::btree_map>> - deleted_set_nodes_ ABSL_GUARDED_BY(set_map_mutex_); - - // Removes deleted keys from key-value map - void CleanUpKeyValueMap(int64_t logical_commit_time); - - // Removes deleted key-values from key-value_set map - void CleanUpKeyValueSetMap(int64_t logical_commit_time); + absl::flat_hash_map< + std::string, + absl::btree_map< + int64_t, + absl::flat_hash_map>>> + deleted_set_nodes_map_ ABSL_GUARDED_BY(set_map_mutex_); + + // Removes deleted keys from key-value map for a given prefix + void CleanUpKeyValueMap(int64_t logical_commit_time, std::string_view prefix); + + // Removes deleted key-values from key-value_set map for a given prefix + void CleanUpKeyValueSetMap(int64_t logical_commit_time, + std::string_view prefix); + // Logs cache access metrics for cache hit or miss counts. The cache access + // event name is defined in server_definition.h file + void LogCacheAccessMetrics(const RequestContext& request_context, + std::string_view cache_access_event) const; friend class KeyValueCacheTestPeer; - - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; }; } // namespace kv_server diff --git a/components/data_server/cache/key_value_cache_test.cc b/components/data_server/cache/key_value_cache_test.cc index 2bc09b39..12ee3d6f 100644 --- a/components/data_server/cache/key_value_cache_test.cc +++ b/components/data_server/cache/key_value_cache_test.cc @@ -30,8 +30,7 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" #include "public/base_types.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" namespace kv_server { @@ -39,9 +38,12 @@ class KeyValueCacheTestPeer { public: KeyValueCacheTestPeer() = delete; static std::multimap ReadDeletedNodes( - const KeyValueCache& c) { + const KeyValueCache& c, std::string_view prefix = "") { absl::MutexLock lock(&c.mutex_); - return c.deleted_nodes_; + auto map_itr = c.deleted_nodes_map_.find(prefix); + return map_itr == c.deleted_nodes_map_.end() + ? std::multimap() + : c.deleted_nodes_map_.find(prefix)->second; } static absl::flat_hash_map& ReadNodes(KeyValueCache& c) { @@ -49,18 +51,24 @@ class KeyValueCacheTestPeer { return c.map_; } - static int GetDeletedSetNodesMapSize(const KeyValueCache& c) { + static int GetDeletedSetNodesMapSize(const KeyValueCache& c, + std::string prefix = "") { absl::MutexLock lock(&c.set_map_mutex_); - return c.deleted_set_nodes_.size(); + auto map_itr = c.deleted_set_nodes_map_.find(prefix); + return map_itr == c.deleted_set_nodes_map_.end() ? 0 + : map_itr->second.size(); } static absl::flat_hash_set ReadDeletedSetNodesForTimestamp( - const KeyValueCache& c, int64_t logical_commit_time, - std::string_view key) { + const KeyValueCache& c, int64_t logical_commit_time, std::string_view key, + std::string_view prefix = "") { absl::MutexLock lock(&c.set_map_mutex_); - return c.deleted_set_nodes_.find(logical_commit_time) - ->second.find(key) - ->second; + auto map_itr = c.deleted_set_nodes_map_.find(prefix); + return map_itr == c.deleted_set_nodes_map_.end() + ? absl::flat_hash_set() + : map_itr->second.find(logical_commit_time) + ->second.find(key) + ->second; } static int GetCacheKeyValueSetMapSize(KeyValueCache& c) { @@ -82,7 +90,7 @@ class KeyValueCacheTestPeer { } static void CallCacheCleanup(KeyValueCache& c, int64_t logical_commit_time) { - c.CleanUpKeyValueMap(logical_commit_time); + c.RemoveDeletedKeys(logical_commit_time); } }; @@ -91,24 +99,33 @@ namespace { using privacy_sandbox::server_common::TelemetryProvider; using testing::UnorderedElementsAre; -TEST(CacheTest, RetrievesMatchingEntry) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +class CacheTest : public ::testing::Test { + protected: + CacheTest() { + InitMetricsContextMap(); + scope_metrics_context_ = std::make_unique(); + request_context_ = + std::make_unique(*scope_metrics_context_); + } + RequestContext& GetRequestContext() { return *request_context_; } + std::unique_ptr scope_metrics_context_; + std::unique_ptr request_context_; +}; + +TEST_F(CacheTest, RetrievesMatchingEntry) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 1); absl::flat_hash_set keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(keys); + cache->GetKeyValuePairs(GetRequestContext(), keys); absl::flat_hash_set wrong_keys = {"wrong_key"}; - EXPECT_FALSE(cache->GetKeyValuePairs(keys).empty()); - EXPECT_TRUE(cache->GetKeyValuePairs(wrong_keys).empty()); + EXPECT_FALSE(cache->GetKeyValuePairs(GetRequestContext(), keys).empty()); + EXPECT_TRUE(cache->GetKeyValuePairs(GetRequestContext(), wrong_keys).empty()); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); } -TEST(CacheTest, GetWithMultipleKeysReturnsMatchingValues) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, GetWithMultipleKeysReturnsMatchingValues) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("key1", "value1", 1); cache->UpdateKeyValue("key2", "value2", 2); cache->UpdateKeyValue("key3", "value3", 3); @@ -116,116 +133,101 @@ TEST(CacheTest, GetWithMultipleKeysReturnsMatchingValues) { absl::flat_hash_set full_keys = {"key1", "key2"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(full_keys); + cache->GetKeyValuePairs(GetRequestContext(), full_keys); EXPECT_EQ(kv_pairs.size(), 2); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("key1", "value1"), KVPairEq("key2", "value2"))); } -TEST(CacheTest, GetAfterUpdateReturnsNewValue) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, GetAfterUpdateReturnsNewValue) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 1); absl::flat_hash_set keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(keys); + cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); cache->UpdateKeyValue("my_key", "my_new_value", 2); - kv_pairs = cache->GetKeyValuePairs(keys); + kv_pairs = cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_EQ(kv_pairs.size(), 1); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_new_value"))); } -TEST(CacheTest, GetAfterUpdateDifferentKeyReturnsSameValue) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, GetAfterUpdateDifferentKeyReturnsSameValue) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 1); cache->UpdateKeyValue("new_key", "new_value", 2); absl::flat_hash_set keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(keys); + cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); } -TEST(CacheTest, GetForEmptyCacheReturnsEmptyList) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, GetForEmptyCacheReturnsEmptyList) { + std::unique_ptr cache = KeyValueCache::Create(); absl::flat_hash_set keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(keys); + cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_EQ(kv_pairs.size(), 0); } -TEST(CacheTest, GetForCacheReturnsValueSet) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, GetForCacheReturnsValueSet) { + std::unique_ptr cache = KeyValueCache::Create(); std::vector values = {"v1", "v2"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_THAT(value_set, UnorderedElementsAre("v1", "v2")); } -TEST(CacheTest, GetForCacheMissingKeyReturnsEmptySet) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, GetForCacheMissingKeyReturnsEmptySet) { + std::unique_ptr cache = KeyValueCache::Create(); std::vector values = {"v1", "v2"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); auto get_key_value_set_result = - cache->GetKeyValueSet({"missing_key", "my_key"}); + cache->GetKeyValueSet(GetRequestContext(), {"missing_key", "my_key"}); EXPECT_EQ(get_key_value_set_result->GetValueSet("missing_key").size(), 0); EXPECT_THAT(get_key_value_set_result->GetValueSet("my_key"), UnorderedElementsAre("v1", "v2")); } -TEST(DeleteKeyTest, RemovesKeyEntry) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyTestRemovesKeyEntry) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 1); cache->DeleteKey("my_key", 2); absl::flat_hash_set full_keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(full_keys); + cache->GetKeyValuePairs(GetRequestContext(), full_keys); EXPECT_EQ(kv_pairs.size(), 0); } -TEST(DeleteKeyValueSetTest, WrongkeyDoesNotRemoveEntry) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyValueSetWrongkeyDoesNotRemoveEntry) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 1); cache->DeleteKey("wrong_key", 1); absl::flat_hash_set keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(keys); + cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); } -TEST(DeleteKeyValueSetTest, RemovesValueEntry) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyValueSetRemovesValueEntry) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1", "v2", "v3"}; std::vector values_to_delete = {"v1", "v2"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("my_key", absl::Span(values_to_delete), 2); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_THAT(value_set, UnorderedElementsAre("v3")); auto value_meta_v3 = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v3"); @@ -243,18 +245,15 @@ TEST(DeleteKeyValueSetTest, RemovesValueEntry) { EXPECT_EQ(value_meta_v2_deleted.is_deleted, true); } -TEST(DeleteKeyValueSetTest, WrongKeyDoesNotRemoveKeyValueEntry) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyValueSetWrongKeyDoesNotRemoveKeyValueEntry) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1", "v2", "v3"}; std::vector values_to_delete = {"v1"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("wrong_key", absl::Span(values_to_delete), 2); std::unique_ptr result = - cache->GetKeyValueSet({"my_key", "wrong_key"}); + cache->GetKeyValueSet(GetRequestContext(), {"my_key", "wrong_key"}); EXPECT_THAT(result->GetValueSet("my_key"), UnorderedElementsAre("v1", "v2", "v3")); EXPECT_EQ(result->GetValueSet("wrong_key").size(), 0); @@ -276,18 +275,16 @@ TEST(DeleteKeyValueSetTest, WrongKeyDoesNotRemoveKeyValueEntry) { EXPECT_EQ(value_meta_v1_deleted_for_wrong_key.is_deleted, true); } -TEST(DeleteKeyValueSetTest, WrongValueDoesNotRemoveEntry) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyValueSetWrongValueDoesNotRemoveEntry) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1", "v2", "v3"}; std::vector values_to_delete = {"v4"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("my_key", absl::Span(values_to_delete), 2); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_THAT(value_set, UnorderedElementsAre("v1", "v2", "v3")); auto value_meta_v1 = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v1"); @@ -304,85 +301,73 @@ TEST(DeleteKeyValueSetTest, WrongValueDoesNotRemoveEntry) { EXPECT_EQ(value_set_in_cache_size, 4); } -TEST(CacheTest, OutOfOrderUpdateAfterUpdateWorks) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, OutOfOrderUpdateAfterUpdateWorks) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 2); absl::flat_hash_set keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(keys); + cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); cache->UpdateKeyValue("my_key", "my_new_value", 1); - kv_pairs = cache->GetKeyValuePairs(keys); + kv_pairs = cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_EQ(kv_pairs.size(), 1); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); } -TEST(DeleteKeyTest, OutOfOrderDeleteAfterUpdateWorks) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyOutOfOrderDeleteAfterUpdateWorks) { + std::unique_ptr cache = KeyValueCache::Create(); cache->DeleteKey("my_key", 2); cache->UpdateKeyValue("my_key", "my_value", 1); absl::flat_hash_set full_keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(full_keys); + cache->GetKeyValuePairs(GetRequestContext(), full_keys); EXPECT_EQ(kv_pairs.size(), 0); } -TEST(DeleteKeyTest, OutOfOrderUpdateAfterDeleteWorks) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyOutOfOrderUpdateAfterDeleteWorks) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 2); cache->DeleteKey("my_key", 1); absl::flat_hash_set full_keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(full_keys); + cache->GetKeyValuePairs(GetRequestContext(), full_keys); EXPECT_EQ(kv_pairs.size(), 1); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); } -TEST(DeleteKeyTest, InOrderUpdateAfterDeleteWorks) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyInOrderUpdateAfterDeleteWorks) { + std::unique_ptr cache = KeyValueCache::Create(); cache->DeleteKey("my_key", 1); cache->UpdateKeyValue("my_key", "my_value", 2); absl::flat_hash_set full_keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(full_keys); + cache->GetKeyValuePairs(GetRequestContext(), full_keys); EXPECT_EQ(kv_pairs.size(), 1); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key", "my_value"))); } -TEST(DeleteKeyTest, InOrderDeleteAfterUpdateWorks) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); +TEST_F(CacheTest, DeleteKeyInOrderDeleteAfterUpdateWorks) { + std::unique_ptr cache = KeyValueCache::Create(); cache->UpdateKeyValue("my_key", "my_value", 1); cache->DeleteKey("my_key", 2); absl::flat_hash_set full_keys = {"my_key"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(full_keys); + cache->GetKeyValuePairs(GetRequestContext(), full_keys); EXPECT_EQ(kv_pairs.size(), 0); } -TEST(UpdateKeyValueSetTest, UpdateAfterUpdateWithSameValue) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, UpdateSetTestUpdateAfterUpdateWithSameValue) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->UpdateKeyValueSet("my_key", absl::Span(values), 2); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_THAT(value_set, UnorderedElementsAre("v1")); auto value_meta = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v1"); @@ -390,11 +375,8 @@ TEST(UpdateKeyValueSetTest, UpdateAfterUpdateWithSameValue) { EXPECT_EQ(value_meta.is_deleted, false); } -TEST(UpdateKeyValueSetTest, UpdateAfterUpdateWithDifferentValue) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, UpdateSetTestUpdateAfterUpdateWithDifferentValue) { + std::unique_ptr cache = std::make_unique(); std::vector first_value = {"v1"}; std::vector second_value = {"v2"}; cache->UpdateKeyValueSet("my_key", absl::Span(first_value), @@ -402,7 +384,8 @@ TEST(UpdateKeyValueSetTest, UpdateAfterUpdateWithDifferentValue) { cache->UpdateKeyValueSet("my_key", absl::Span(second_value), 2); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_THAT(value_set, UnorderedElementsAre("v1", "v2")); auto value_meta_v1 = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v1"); @@ -414,16 +397,14 @@ TEST(UpdateKeyValueSetTest, UpdateAfterUpdateWithDifferentValue) { EXPECT_EQ(value_meta_v2.is_deleted, false); } -TEST(InOrderUpdateKeyValueSetTest, InsertAfterDeleteExpectInsert) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, InOrderUpdateSetInsertAfterDeleteExpectInsert) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1"}; cache->DeleteValuesInSet("my_key", absl::Span(values), 1); cache->UpdateKeyValueSet("my_key", absl::Span(values), 2); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_THAT(value_set, UnorderedElementsAre("v1")); auto value_meta = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v1"); @@ -431,16 +412,14 @@ TEST(InOrderUpdateKeyValueSetTest, InsertAfterDeleteExpectInsert) { EXPECT_EQ(value_meta.is_deleted, false); } -TEST(InOrderUpdateKeyValueSetTest, DeleteAfterInsert) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, InOrderUpdateSetDeleteAfterInsert) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("my_key", absl::Span(values), 2); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_EQ(value_set.size(), 0); auto value_meta_v1 = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v1"); @@ -448,16 +427,14 @@ TEST(InOrderUpdateKeyValueSetTest, DeleteAfterInsert) { EXPECT_EQ(value_meta_v1.is_deleted, true); } -TEST(OutOfOrderUpdateKeyValueSetTest, InsertAfterDeleteExpectNoInsert) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, OutOfOrderUpdateSetInsertAfterDeleteExpectNoInsert) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1"}; cache->DeleteValuesInSet("my_key", absl::Span(values), 2); cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_EQ(value_set.size(), 0); auto value_meta = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v1"); @@ -465,16 +442,14 @@ TEST(OutOfOrderUpdateKeyValueSetTest, InsertAfterDeleteExpectNoInsert) { EXPECT_EQ(value_meta.is_deleted, true); } -TEST(OutOfOrderUpdateKeyValueSetTest, DeleteAfterInsertExpectNoDelete) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, OutOfOrderUpdateSetDeleteAfterInsertExpectNoDelete) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 2); cache->DeleteValuesInSet("my_key", absl::Span(values), 1); absl::flat_hash_set value_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_THAT(value_set, UnorderedElementsAre("v1")); auto value_meta_v1 = KeyValueCacheTestPeer::GetSetValueMeta(*cache, "my_key", "v1"); @@ -482,22 +457,16 @@ TEST(OutOfOrderUpdateKeyValueSetTest, DeleteAfterInsertExpectNoDelete) { EXPECT_EQ(value_meta_v1.is_deleted, false); } -TEST(CleanUpTimestamps, InsertAKeyDoesntUpdateDeletedNodes) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsInsertAKeyDoesntUpdateDeletedNodes) { + std::unique_ptr cache = std::make_unique(); cache->UpdateKeyValue("my_key", "my_value", 1); auto deleted_nodes = KeyValueCacheTestPeer::ReadDeletedNodes(*cache); EXPECT_EQ(deleted_nodes.size(), 0); } -TEST(CleanUpTimestamps, RemoveDeletedKeysRemovesOldRecords) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsRemoveDeletedKeysRemovesOldRecords) { + std::unique_ptr cache = std::make_unique(); cache->UpdateKeyValue("my_key", "my_value", 1); cache->DeleteKey("my_key", 2); @@ -510,11 +479,8 @@ TEST(CleanUpTimestamps, RemoveDeletedKeysRemovesOldRecords) { EXPECT_EQ(nodes.size(), 0); } -TEST(CleanUpTimestamps, RemoveDeletedKeysDoesntAffectNewRecords) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsRemoveDeletedKeysDoesntAffectNewRecords) { + std::unique_ptr cache = std::make_unique(); cache->UpdateKeyValue("my_key", "my_value", 5); cache->DeleteKey("my_key", 6); @@ -527,12 +493,9 @@ TEST(CleanUpTimestamps, RemoveDeletedKeysDoesntAffectNewRecords) { EXPECT_EQ(range.first->second, "my_key"); } -TEST(CleanUpTimestamps, - RemoveDeletedKeysRemovesOldRecordsDoesntAffectNewRecords) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, + CleanupRemoveDeletedKeysRemovesOldRecordsDoesntAffectNewRecords) { + std::unique_ptr cache = std::make_unique(); cache->UpdateKeyValue("my_key1", "my_value", 1); cache->UpdateKeyValue("my_key2", "my_value", 2); cache->UpdateKeyValue("my_key3", "my_value", 3); @@ -559,17 +522,14 @@ TEST(CleanUpTimestamps, "my_key1", "my_key2", "my_key3", "my_key4", "my_key5", }; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(full_keys); + cache->GetKeyValuePairs(GetRequestContext(), full_keys); EXPECT_EQ(kv_pairs.size(), 2); EXPECT_THAT(kv_pairs, UnorderedElementsAre(KVPairEq("my_key4", "my_value"), KVPairEq("my_key5", "my_value"))); } -TEST(CleanUpTimestamps, CantInsertOldRecordsAfterCleanup) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsCantInsertOldRecordsAfterCleanup) { + std::unique_ptr cache = std::make_unique(); cache->UpdateKeyValue("my_key1", "my_value", 10); cache->DeleteKey("my_key1", 12); cache->RemoveDeletedKeys(13); @@ -582,15 +542,12 @@ TEST(CleanUpTimestamps, CantInsertOldRecordsAfterCleanup) { absl::flat_hash_set keys = {"my_key1"}; absl::flat_hash_map kv_pairs = - cache->GetKeyValuePairs(keys); + cache->GetKeyValuePairs(GetRequestContext(), keys); EXPECT_EQ(kv_pairs.size(), 0); } -TEST(CleanUpTimestampsForSetCache, InsertKeyValueSetDoesntUpdateDeletedNodes) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsInsertKeyValueSetDoesntUpdateDeletedNodes) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"my_value"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); int deleted_nodes_map_size = @@ -598,11 +555,8 @@ TEST(CleanUpTimestampsForSetCache, InsertKeyValueSetDoesntUpdateDeletedNodes) { EXPECT_EQ(deleted_nodes_map_size, 0); } -TEST(CleanUpTimestampsForSetCache, DeleteKeyValueSetExpectUpdateDeletedNodes) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsDeleteKeyValueSetExpectUpdateDeletedNodes) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"my_value"}; cache->DeleteValuesInSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("another_key", absl::Span(values), @@ -620,11 +574,8 @@ TEST(CleanUpTimestampsForSetCache, DeleteKeyValueSetExpectUpdateDeletedNodes) { 1); } -TEST(CleanUpTimestampsForSetCache, RemoveDeletedKeyValuesRemovesOldRecords) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsRemoveDeletedKeyValuesRemovesOldRecords) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"my_value"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("my_key", absl::Span(values), 2); @@ -639,12 +590,9 @@ TEST(CleanUpTimestampsForSetCache, RemoveDeletedKeyValuesRemovesOldRecords) { EXPECT_EQ(KeyValueCacheTestPeer::GetCacheKeyValueSetMapSize(*cache), 0); } -TEST(CleanUpTimestampsForSetCache, - RemoveDeletedKeyValuesDoesntAffectNewRecords) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, + CleanupTimestampsRemoveDeletedKeyValuesDoesntAffectNewRecords) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"my_value"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 5); cache->DeleteValuesInSet("my_key", absl::Span(values), 6); @@ -660,12 +608,10 @@ TEST(CleanUpTimestampsForSetCache, 1); } -TEST(CleanUpTimestampsForSetCache, - RemoveDeletedKeysRemovesOldRecordsDoesntAffectNewRecords) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F( + CacheTest, + CleanupSetCacheRemoveDeletedKeysRemovesOldRecordsDoesntAffectNewRecords) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"v1", "v2"}; std::vector values_to_delete = {"v1"}; cache->UpdateKeyValueSet("my_key1", absl::Span(values), 1); @@ -689,8 +635,8 @@ TEST(CleanUpTimestampsForSetCache, "my_key2") .size(), 1); - auto get_value_set_result = - cache->GetKeyValueSet({"my_key1", "my_key4", "my_key3"}); + auto get_value_set_result = cache->GetKeyValueSet( + GetRequestContext(), {"my_key1", "my_key4", "my_key3"}); EXPECT_THAT(get_value_set_result->GetValueSet("my_key4"), UnorderedElementsAre("v1", "v2")); EXPECT_THAT(get_value_set_result->GetValueSet("my_key3"), @@ -699,11 +645,8 @@ TEST(CleanUpTimestampsForSetCache, UnorderedElementsAre("v2")); } -TEST(CleanUpTimestampsForSetCache, CantInsertOldRecordsAfterCleanup) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsSetCacheCantInsertOldRecordsAfterCleanup) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"my_value"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("my_key", absl::Span(values), 2); @@ -717,15 +660,13 @@ TEST(CleanUpTimestampsForSetCache, CantInsertOldRecordsAfterCleanup) { cache->UpdateKeyValueSet("my_key", absl::Span(values), 2); absl::flat_hash_set kv_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_EQ(kv_set.size(), 0); } -TEST(CleanUpTimestampsForSetCache, CantAddOldDeletedRecordsAfterCleanup) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = - std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, CleanupTimestampsCantAddOldDeletedRecordsAfterCleanup) { + std::unique_ptr cache = std::make_unique(); std::vector values = {"my_value"}; cache->UpdateKeyValueSet("my_key", absl::Span(values), 1); cache->DeleteValuesInSet("my_key", absl::Span(values), 2); @@ -755,14 +696,13 @@ TEST(CleanUpTimestampsForSetCache, CantAddOldDeletedRecordsAfterCleanup) { EXPECT_EQ(value_meta.last_logical_commit_time, 4); absl::flat_hash_set kv_set = - cache->GetKeyValueSet({"my_key"})->GetValueSet("my_key"); + cache->GetKeyValueSet(GetRequestContext(), {"my_key"}) + ->GetValueSet("my_key"); EXPECT_EQ(kv_set.size(), 0); } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndGet) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentGetAndGet) { + auto cache = std::make_unique(); absl::flat_hash_set keys_lookup_request = {"key1", "key2"}; std::vector values_for_key1 = {"v1"}; std::vector values_for_key2 = {"v2"}; @@ -771,9 +711,10 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndGet) { cache->UpdateKeyValueSet("key2", absl::Span(values_for_key2), 1); absl::Notification start; - auto lookup_fn = [&cache, &keys_lookup_request, &start]() { + auto request_context = GetRequestContext(); + auto lookup_fn = [&cache, &keys_lookup_request, &start, &request_context]() { start.WaitForNotification(); - auto result = cache->GetKeyValueSet(keys_lookup_request); + auto result = cache->GetKeyValueSet(request_context, keys_lookup_request); EXPECT_THAT(result->GetValueSet("key1"), UnorderedElementsAre("v1")); EXPECT_THAT(result->GetValueSet("key2"), UnorderedElementsAre("v2")); }; @@ -788,19 +729,19 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndGet) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndUpdateExpectNoUpdate) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentGetAndUpdateExpectNoUpdate) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1"}; std::vector existing_values = {"v1"}; cache->UpdateKeyValueSet("key1", absl::Span(existing_values), 3); absl::Notification start; - auto lookup_fn = [&cache, &keys, &start]() { + auto request_context = GetRequestContext(); + auto lookup_fn = [&cache, &keys, &start, &request_context]() { start.WaitForNotification(); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); }; std::vector new_values = {"v1"}; auto update_fn = [&cache, &new_values, &start]() { @@ -820,19 +761,19 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndUpdateExpectNoUpdate) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndUpdateExpectUpdate) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentGetAndUpdateExpectUpdate) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1", "key2"}; std::vector existing_values = {"v1"}; cache->UpdateKeyValueSet("key1", absl::Span(existing_values), 1); absl::Notification start; - auto lookup_fn = [&cache, &keys, &start]() { + auto request_context = GetRequestContext(); + auto lookup_fn = [&cache, &keys, &start, &request_context]() { start.WaitForNotification(); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); }; std::vector new_values_for_key2 = {"v2"}; auto update_fn = [&cache, &new_values_for_key2, &start]() { @@ -853,19 +794,19 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndUpdateExpectUpdate) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndDeleteExpectNoDelete) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentGetAndDeleteExpectNoDelete) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1"}; std::vector existing_values = {"v1"}; cache->UpdateKeyValueSet("key1", absl::Span(existing_values), 3); absl::Notification start; - auto lookup_fn = [&cache, &keys, &start]() { + auto request_context = GetRequestContext(); + auto lookup_fn = [&cache, &keys, &start, &request_context]() { start.WaitForNotification(); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); }; std::vector delete_values = {"v1"}; auto delete_fn = [&cache, &delete_values, &start]() { @@ -886,10 +827,8 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndDeleteExpectNoDelete) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndCleanUp) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentGetAndCleanUp) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1", "key2"}; std::vector existing_values = {"v1"}; cache->UpdateKeyValueSet("key1", @@ -899,11 +838,16 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndCleanUp) { cache->DeleteValuesInSet("key2", absl::Span(existing_values), 2); absl::Notification start; - auto lookup_fn = [&cache, &keys, &start]() { + auto request_context = GetRequestContext(); + auto lookup_fn = [&cache, &keys, &start, &request_context]() { start.WaitForNotification(); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); - EXPECT_EQ(cache->GetKeyValueSet(keys)->GetValueSet("key2").size(), 0); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); + EXPECT_EQ(cache->GetKeyValueSet(request_context, keys) + ->GetValueSet("key2") + .size(), + 0); }; auto cleanup_fn = [&cache, &start]() { // clean up old records @@ -922,29 +866,32 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetAndCleanUp) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentUpdateAndUpdateExpectUpdateBoth) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentUpdateAndUpdateExpectUpdateBoth) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1", "key2"}; std::vector values_for_key1 = {"v1"}; absl::Notification start; - auto update_key1 = [&cache, &keys, &values_for_key1, &start]() { + auto request_context = GetRequestContext(); + auto update_key1 = [&cache, &keys, &values_for_key1, &start, + &request_context]() { start.WaitForNotification(); // expect new value is inserted for key1 cache->UpdateKeyValueSet("key1", absl::Span(values_for_key1), 1); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); }; std::vector values_for_key2 = {"v2"}; - auto update_key2 = [&cache, &keys, &values_for_key2, &start]() { + auto update_key2 = [&cache, &keys, &values_for_key2, &start, + &request_context]() { // expect new value is inserted for key2 start.WaitForNotification(); cache->UpdateKeyValueSet("key2", absl::Span(values_for_key2), 2); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key2"), - UnorderedElementsAre("v2")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key2"), + UnorderedElementsAre("v2")); }; std::vector threads; for (int i = 0; i < std::min(20, (int)std::thread::hardware_concurrency()); @@ -958,20 +905,21 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentUpdateAndUpdateExpectUpdateBoth) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentUpdateAndDelete) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentUpdateAndDelete) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1", "key2"}; std::vector values_for_key1 = {"v1"}; absl::Notification start; - auto update_key1 = [&cache, &keys, &values_for_key1, &start]() { + auto request_context = GetRequestContext(); + auto update_key1 = [&cache, &keys, &values_for_key1, &start, + &request_context]() { start.WaitForNotification(); // expect new value is inserted for key1 cache->UpdateKeyValueSet("key1", absl::Span(values_for_key1), 1); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); }; // Update existing value for key2 std::vector existing_values_for_key2 = {"v1", "v2"}; @@ -979,13 +927,15 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentUpdateAndDelete) { "key2", absl::Span(existing_values_for_key2), 1); std::vector values_to_delete_for_key2 = {"v1"}; - auto delete_key2 = [&cache, &keys, &values_to_delete_for_key2, &start]() { + auto delete_key2 = [&cache, &keys, &values_to_delete_for_key2, &start, + &request_context]() { start.WaitForNotification(); // expect value is deleted for key2 cache->DeleteValuesInSet( "key2", absl::Span(values_to_delete_for_key2), 2); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key2"), - UnorderedElementsAre("v2")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key2"), + UnorderedElementsAre("v2")); }; std::vector threads; @@ -1000,23 +950,24 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentUpdateAndDelete) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentUpdateAndCleanUp) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentUpdateAndCleanUp) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1"}; std::vector values_for_key1 = {"v1"}; absl::Notification start; - auto update_fn = [&cache, &keys, &values_for_key1, &start]() { + auto request_context = GetRequestContext(); + auto update_fn = [&cache, &keys, &values_for_key1, &start, + &request_context]() { start.WaitForNotification(); cache->UpdateKeyValueSet("key1", - absl::Span(values_for_key1), 1); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); + absl::Span(values_for_key1), 2); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); }; auto cleanup_fn = [&cache, &start]() { start.WaitForNotification(); - KeyValueCacheTestPeer::CallCacheCleanup(*cache, 2); + KeyValueCacheTestPeer::CallCacheCleanup(*cache, 1); }; std::vector threads; @@ -1031,25 +982,28 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentUpdateAndCleanUp) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentDeleteAndCleanUp) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentDeleteAndCleanUp) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1"}; std::vector values_for_key1 = {"v1"}; cache->UpdateKeyValueSet("key1", absl::Span(values_for_key1), 1); absl::Notification start; - auto delete_fn = [&cache, &keys, &values_for_key1, &start]() { + auto request_context = GetRequestContext(); + auto delete_fn = [&cache, &keys, &values_for_key1, &start, + &request_context]() { start.WaitForNotification(); // expect new value is deleted for key1 cache->DeleteValuesInSet("key1", absl::Span(values_for_key1), 2); - EXPECT_EQ(cache->GetKeyValueSet(keys)->GetValueSet("key1").size(), 0); + EXPECT_EQ(cache->GetKeyValueSet(request_context, keys) + ->GetValueSet("key1") + .size(), + 0); }; auto cleanup_fn = [&cache, &start]() { start.WaitForNotification(); - KeyValueCacheTestPeer::CallCacheCleanup(*cache, 2); + KeyValueCacheTestPeer::CallCacheCleanup(*cache, 1); }; std::vector threads; for (int i = 0; i < std::min(20, (int)std::thread::hardware_concurrency()); @@ -1063,10 +1017,8 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentDeleteAndCleanUp) { } } -TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetUpdateDeleteCleanUp) { - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - auto cache = std::make_unique(*noop_metrics_recorder); +TEST_F(CacheTest, ConcurrentGetUpdateDeleteCleanUp) { + auto cache = std::make_unique(); absl::flat_hash_set keys = {"key1", "key2"}; std::vector existing_values_for_key1 = {"v1"}; std::vector existing_values_for_key2 = {"v1"}; @@ -1090,13 +1042,14 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetUpdateDeleteCleanUp) { }; auto cleanup = [&cache, &start]() { start.WaitForNotification(); - KeyValueCacheTestPeer::CallCacheCleanup(*cache, 2); + KeyValueCacheTestPeer::CallCacheCleanup(*cache, 1); }; - - auto lookup_for_key1 = [&cache, &keys, &start]() { + auto request_context = GetRequestContext(); + auto lookup_for_key1 = [&cache, &keys, &start, &request_context]() { start.WaitForNotification(); - EXPECT_THAT(cache->GetKeyValueSet(keys)->GetValueSet("key1"), - UnorderedElementsAre("v1")); + EXPECT_THAT( + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key1"), + UnorderedElementsAre("v1")); }; std::vector threads; @@ -1112,8 +1065,216 @@ TEST(ConcurrentSetMemoryAccessTest, ConcurrentGetUpdateDeleteCleanUp) { thread.join(); } auto look_up_result_for_key2 = - cache->GetKeyValueSet(keys)->GetValueSet("key2"); + cache->GetKeyValueSet(request_context, keys)->GetValueSet("key2"); EXPECT_THAT(look_up_result_for_key2, UnorderedElementsAre("v2")); } + +TEST_F(CacheTest, MultiplePrefixKeyValueUpdates) { + std::unique_ptr cache = KeyValueCache::Create(); + // Call remove deleted keys for prefix1 to update the max delete cutoff + // timestamp + cache->RemoveDeletedKeys(1, "prefix1"); + cache->UpdateKeyValue("prefix1-key", "value1", 2, "prefix1"); + cache->UpdateKeyValue("prefix2-key", "value2", 1, "prefix2"); + absl::flat_hash_map kv_pairs = + cache->GetKeyValuePairs(GetRequestContext(), + {"prefix1-key", "prefix2-key"}); + EXPECT_EQ(kv_pairs.size(), 2); + EXPECT_EQ(kv_pairs["prefix1-key"], "value1"); + EXPECT_EQ(kv_pairs["prefix2-key"], "value2"); +} + +TEST_F(CacheTest, MultiplePrefixKeyValueNoUpdateForAnother) { + std::unique_ptr cache = KeyValueCache::Create(); + // Call remove deleted keys for prefix1 to update the max delete cutoff + // timestamp + cache->RemoveDeletedKeys(2, "prefix1"); + // Expect no update for prefix1 + cache->UpdateKeyValue("prefix1-key", "value1", 1, "prefix1"); + cache->UpdateKeyValue("prefix2-key", "value2", 1, "prefix2"); + absl::flat_hash_map kv_pairs = + cache->GetKeyValuePairs(GetRequestContext(), + {"prefix1-key", "prefix2-key"}); + EXPECT_EQ(kv_pairs.size(), 1); + EXPECT_EQ(kv_pairs["prefix2-key"], "value2"); +} + +TEST_F(CacheTest, MultiplePrefixKeyValueNoDeleteForAnother) { + std::unique_ptr cache = KeyValueCache::Create(); + // Call remove deleted keys for prefix1 to update the max delete cutoff + // timestamp + cache->RemoveDeletedKeys(2, "prefix1"); + cache->UpdateKeyValue("prefix1-key", "value1", 3, "prefix1"); + // Expect no deletion + cache->DeleteKey("prefix1-key", 1, "prefix1"); + cache->UpdateKeyValue("prefix2-key", "value2", 1, "prefix2"); + absl::flat_hash_map kv_pairs = + cache->GetKeyValuePairs(GetRequestContext(), + {"prefix1-key", "prefix2-key"}); + EXPECT_EQ(kv_pairs.size(), 2); + EXPECT_EQ(kv_pairs["prefix1-key"], "value1"); + EXPECT_EQ(kv_pairs["prefix2-key"], "value2"); +} + +TEST_F(CacheTest, MultiplePrefixKeyValueDeletesAndUpdates) { + std::unique_ptr cache = std::make_unique(); + cache->DeleteKey("prefix1-key", 2, "prefix1"); + cache->UpdateKeyValue("prefix1-key", "value1", 1, "prefix1"); + cache->UpdateKeyValue("prefix2-key", "value2", 1, "prefix2"); + absl::flat_hash_map kv_pairs = + cache->GetKeyValuePairs(GetRequestContext(), + {"prefix1-key", "prefix2-key"}); + EXPECT_EQ(kv_pairs.size(), 1); + EXPECT_EQ(kv_pairs["prefix2-key"], "value2"); + auto deleted_nodes = + KeyValueCacheTestPeer::ReadDeletedNodes(*cache, "prefix1"); + EXPECT_EQ(deleted_nodes.size(), 1); + deleted_nodes = KeyValueCacheTestPeer::ReadDeletedNodes(*cache, "prefix2"); + EXPECT_EQ(deleted_nodes.size(), 0); +} + +TEST_F(CacheTest, MultiplePrefixKeyValueUpdatesAndDeletes) { + std::unique_ptr cache = std::make_unique(); + cache->UpdateKeyValue("prefix1-key", "value1", 2, "prefix1"); + // Expects no deletes + cache->DeleteKey("prefix1-key", 1, "prefix1"); + cache->UpdateKeyValue("prefix2-key", "value2", 1, "prefix2"); + absl::flat_hash_map kv_pairs = + cache->GetKeyValuePairs(GetRequestContext(), + {"prefix1-key", "prefix2-key"}); + EXPECT_EQ(kv_pairs.size(), 2); + EXPECT_EQ(kv_pairs["prefix1-key"], "value1"); + EXPECT_EQ(kv_pairs["prefix2-key"], "value2"); + auto deleted_nodes = KeyValueCacheTestPeer::ReadDeletedNodes(*cache); + EXPECT_EQ(deleted_nodes.size(), 0); +} + +TEST_F(CacheTest, MultiplePrefixKeyValueSetUpdates) { + std::unique_ptr cache = std::make_unique(); + std::vector values1 = {"v1", "v2"}; + std::vector values2 = {"v3", "v4"}; + // Call remove deleted keys for prefix1 to update the max delete cutoff + // timestamp + cache->RemoveDeletedKeys(1, "prefix1"); + cache->UpdateKeyValueSet("prefix1-key", absl::Span(values1), + 2, "prefix1"); + cache->UpdateKeyValueSet("prefix2-key", absl::Span(values2), + 1, "prefix2"); + + auto get_value_set_result = cache->GetKeyValueSet( + GetRequestContext(), {"prefix1-key", "prefix2-key"}); + EXPECT_THAT(get_value_set_result->GetValueSet("prefix1-key"), + UnorderedElementsAre("v1", "v2")); + EXPECT_THAT(get_value_set_result->GetValueSet("prefix2-key"), + UnorderedElementsAre("v3", "v4")); +} + +TEST_F(CacheTest, MultipleKeyValueSetNoUpdateForAnother) { + std::unique_ptr cache = std::make_unique(); + std::vector values1 = {"v1", "v2"}; + std::vector values2 = {"v3", "v4"}; + // Call remove deleted keys for prefix1 to update the max delete cutoff + // timestamp + cache->RemoveDeletedKeys(2, "prefix1"); + cache->UpdateKeyValueSet("prefix1-key", absl::Span(values1), + 1, "prefix1"); + cache->UpdateKeyValueSet("prefix2-key", absl::Span(values2), + 1, "prefix2"); + auto get_value_set_result = cache->GetKeyValueSet( + GetRequestContext(), {"prefix1-key", "prefix2-key"}); + EXPECT_EQ(get_value_set_result->GetValueSet("prefix1-key").size(), 0); + EXPECT_THAT(get_value_set_result->GetValueSet("prefix2-key"), + UnorderedElementsAre("v3", "v4")); +} + +TEST_F(CacheTest, MultiplePrefixKeyValueSetDeletesAndUpdates) { + std::unique_ptr cache = std::make_unique(); + std::vector values1 = {"v1", "v2"}; + std::vector values_to_delete = {"v1"}; + std::vector values2 = {"v3", "v4"}; + cache->DeleteValuesInSet("prefix1-key", + absl::Span(values_to_delete), 2, + "prefix1"); + cache->UpdateKeyValueSet("prefix1-key", absl::Span(values1), + 1, "prefix1"); + cache->UpdateKeyValueSet("prefix2-key", absl::Span(values2), + 1, "prefix2"); + auto get_value_set_result = cache->GetKeyValueSet( + GetRequestContext(), {"prefix1-key", "prefix2-key"}); + EXPECT_THAT(get_value_set_result->GetValueSet("prefix1-key"), + UnorderedElementsAre("v2")); + EXPECT_THAT(get_value_set_result->GetValueSet("prefix2-key"), + UnorderedElementsAre("v3", "v4")); + EXPECT_EQ(KeyValueCacheTestPeer::GetDeletedSetNodesMapSize(*cache, "prefix1"), + 1); + EXPECT_EQ(KeyValueCacheTestPeer::GetDeletedSetNodesMapSize(*cache, "prefix2"), + 0); +} + +TEST_F(CacheTest, MultiplePrefixKeyValueSetUpdatesAndDeletes) { + std::unique_ptr cache = std::make_unique(); + std::vector values1 = {"v1", "v2"}; + std::vector values_to_delete = {"v1"}; + std::vector values2 = {"v3", "v4"}; + + cache->UpdateKeyValueSet("prefix1-key", absl::Span(values1), + 2, "prefix1"); + cache->UpdateKeyValueSet("prefix2-key", absl::Span(values2), + 1, "prefix2"); + // Expect no deletes + cache->DeleteValuesInSet("prefix1-key", + absl::Span(values_to_delete), 1, + "prefix1"); + auto get_value_set_result = cache->GetKeyValueSet( + GetRequestContext(), {"prefix1-key", "prefix2-key"}); + EXPECT_THAT(get_value_set_result->GetValueSet("prefix1-key"), + UnorderedElementsAre("v1", "v2")); + EXPECT_THAT(get_value_set_result->GetValueSet("prefix2-key"), + UnorderedElementsAre("v3", "v4")); + EXPECT_EQ(KeyValueCacheTestPeer::GetDeletedSetNodesMapSize(*cache, "prefix1"), + 0); + EXPECT_EQ(KeyValueCacheTestPeer::GetDeletedSetNodesMapSize(*cache, "prefix2"), + 0); +} + +TEST_F(CacheTest, MultiplePrefixTimestampKeyValueCleanUps) { + std::unique_ptr cache = std::make_unique(); + cache->UpdateKeyValue("prefix1-key", "value", 2, "prefix1"); + cache->DeleteKey("prefix1-key", 3, "prefix1"); + cache->UpdateKeyValue("prefix2-key", "value", 2, "prefix2"); + cache->DeleteKey("prefix2-key", 5, "prefix2"); + auto deleted_nodes_for_prefix1 = + KeyValueCacheTestPeer::ReadDeletedNodes(*cache, "prefix1"); + EXPECT_EQ(deleted_nodes_for_prefix1.size(), 1); + cache->RemoveDeletedKeys(4, "prefix1"); + cache->RemoveDeletedKeys(4, "prefix2"); + deleted_nodes_for_prefix1 = + KeyValueCacheTestPeer::ReadDeletedNodes(*cache, "prefix1"); + EXPECT_EQ(deleted_nodes_for_prefix1.size(), 0); + auto deleted_nodes_for_prefix2 = + KeyValueCacheTestPeer::ReadDeletedNodes(*cache, "prefix2"); + EXPECT_EQ(deleted_nodes_for_prefix2.size(), 1); +} +TEST_F(CacheTest, MultiplePrefixTimestampKeyValueSetCleanUps) { + std::unique_ptr cache = std::make_unique(); + std::vector values = {"v1", "v2"}; + std::vector values_to_delete = {"v1"}; + cache->UpdateKeyValueSet("prefix1-key", absl::Span(values), + 2, "prefix1"); + cache->UpdateKeyValueSet("prefix2-key", absl::Span(values), + 2, "prefix2"); + cache->DeleteValuesInSet("prefix1-key", + absl::Span(values_to_delete), 3, + "prefix1"); + cache->DeleteValuesInSet("prefix2-key", + absl::Span(values_to_delete), 5, + "prefix2"); + cache->RemoveDeletedKeys(4, "prefix1"); + cache->RemoveDeletedKeys(4, "prefix2"); + EXPECT_EQ(KeyValueCacheTestPeer::GetDeletedSetNodesMapSize(*cache, "prefix1"), + 0); + EXPECT_EQ(KeyValueCacheTestPeer::GetDeletedSetNodesMapSize(*cache, "prefix2"), + 1); +} } // namespace } // namespace kv_server diff --git a/components/data_server/cache/mocks.h b/components/data_server/cache/mocks.h index aec1d60a..134c71fc 100644 --- a/components/data_server/cache/mocks.h +++ b/components/data_server/cache/mocks.h @@ -33,24 +33,30 @@ MATCHER_P2(KVPairEq, key, value, "") { class MockCache : public Cache { public: MOCK_METHOD((absl::flat_hash_map), GetKeyValuePairs, - (const absl::flat_hash_set&), + (const RequestContext& request_context, + const absl::flat_hash_set&), (const, override)); MOCK_METHOD((std::unique_ptr), GetKeyValueSet, - (const absl::flat_hash_set&), + (const RequestContext& request_context, + const absl::flat_hash_set&), (const, override)); MOCK_METHOD(void, UpdateKeyValue, - (std::string_view key, std::string_view value, int64_t ts), + (std::string_view key, std::string_view value, int64_t ts, + std::string_view prefix), (override)); MOCK_METHOD(void, UpdateKeyValueSet, (std::string_view key, absl::Span value_set, - int64_t logical_commit_time), + int64_t logical_commit_time, std::string_view prefix), (override)); MOCK_METHOD(void, DeleteValuesInSet, (std::string_view key, absl::Span value_set, - int64_t logical_commit_time), + int64_t logical_commit_time, std::string_view prefix), + (override)); + MOCK_METHOD(void, DeleteKey, + (std::string_view key, int64_t ts, std::string_view prefix), + (override)); + MOCK_METHOD(void, RemoveDeletedKeys, (int64_t ts, std::string_view prefix), (override)); - MOCK_METHOD(void, DeleteKey, (std::string_view key, int64_t ts), (override)); - MOCK_METHOD(void, RemoveDeletedKeys, (int64_t ts), (override)); }; class MockGetKeyValueSetResult : public GetKeyValueSetResult { diff --git a/components/data_server/cache/noop_key_value_cache.h b/components/data_server/cache/noop_key_value_cache.h index b26b3fda..5fc2afd9 100644 --- a/components/data_server/cache/noop_key_value_cache.h +++ b/components/data_server/cache/noop_key_value_cache.h @@ -26,23 +26,30 @@ namespace kv_server { class NoOpKeyValueCache : public Cache { public: absl::flat_hash_map GetKeyValuePairs( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const override { return {}; }; std::unique_ptr GetKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const override { return std::make_unique(); } void UpdateKeyValue(std::string_view key, std::string_view value, - int64_t logical_commit_time) override {} + int64_t logical_commit_time, + std::string_view prefix) override {} void UpdateKeyValueSet(std::string_view key, absl::Span value_set, - int64_t logical_commit_time) override {} - void DeleteKey(std::string_view key, int64_t logical_commit_time) override {} + int64_t logical_commit_time, + std::string_view prefix) override {} + void DeleteKey(std::string_view key, int64_t logical_commit_time, + std::string_view prefix) override {} void DeleteValuesInSet(std::string_view key, absl::Span value_set, - int64_t logical_commit_time) override {} - void RemoveDeletedKeys(int64_t logical_commit_time) override {} + int64_t logical_commit_time, + std::string_view prefix) override {} + void RemoveDeletedKeys(int64_t logical_commit_time, + std::string_view prefix) override {} static std::unique_ptr Create() { return std::make_unique(); } diff --git a/components/data_server/data_loading/BUILD.bazel b/components/data_server/data_loading/BUILD.bazel index e33a3ed4..a15a8102 100644 --- a/components/data_server/data_loading/BUILD.bazel +++ b/components/data_server/data_loading/BUILD.bazel @@ -27,9 +27,11 @@ cc_library( "data_orchestrator.h", ], deps = [ + "//components/data/blob_storage:blob_prefix_allowlist", "//components/data/blob_storage:blob_storage_change_notifier", "//components/data/blob_storage:blob_storage_client", "//components/data/blob_storage:delta_file_notifier", + "//components/data/file_group:file_group_search_utils", "//components/data/realtime:realtime_notifier", "//components/data/realtime:realtime_thread_pool_manager", "//components/data_server/cache", @@ -42,13 +44,14 @@ cc_library( "//public/data_loading/readers:riegeli_stream_io", "//public/data_loading/readers:stream_record_reader_factory", "//public/sharding:key_sharder", - "@com_github_google_glog//:glog", "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/functional:bind_front", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", - "@google_privacysandbox_servers_common//src/cpp/telemetry:tracing", + "@google_privacysandbox_servers_common//src/telemetry:tracing", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) @@ -68,9 +71,8 @@ cc_test( "//public/data_loading:records_utils", "//public/test_util:mocks", "//public/test_util:proto_matcher", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) diff --git a/components/data_server/data_loading/data_orchestrator.cc b/components/data_server/data_loading/data_orchestrator.cc index dca1ffcd..f787b309 100644 --- a/components/data_server/data_loading/data_orchestrator.cc +++ b/components/data_server/data_loading/data_orchestrator.cc @@ -19,19 +19,27 @@ #include #include +#include "absl/container/flat_hash_map.h" #include "absl/functional/bind_front.h" +#include "absl/log/log.h" #include "absl/strings/str_cat.h" +#include "components/data/file_group/file_group_search_utils.h" #include "components/errors/retry.h" -#include "glog/logging.h" #include "public/constants.h" #include "public/data_loading/data_loading_generated.h" #include "public/data_loading/filename_utils.h" #include "public/data_loading/records_utils.h" #include "public/sharding/sharding_function.h" -#include "src/cpp/telemetry/tracing.h" +#include "src/telemetry/tracing.h" +#include "src/util/status_macro/status_macros.h" namespace kv_server { namespace { +// TODO(b/321716836): use the default prefix to apply cache updates for realtime +// for now. This needs to be removed after we are done with directory support +// for file updates. +constexpr std::string_view kDefaultPrefixForRealTimeUpdates = ""; +constexpr std::string_view kDefaultDataSourceForRealtimeUpdates = "realtime"; using privacy_sandbox::server_common::TraceWithStatusOr; @@ -46,36 +54,41 @@ class BlobRecordStream : public RecordStream { std::unique_ptr blob_reader_; }; -void LogDataLoadingMetrics(const DataLoadingStats& data_loading_stats) { - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - static_cast(data_loading_stats.total_updated_records))); - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - static_cast(data_loading_stats.total_deleted_records))); - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogUpDownCounter( - static_cast(data_loading_stats.total_dropped_records))); +void LogDataLoadingMetrics(std::string_view source, + const DataLoadingStats& data_loading_stats) { + LogIfError(KVServerContextMap() + ->SafeMetric() + .LogUpDownCounter( + {{std::string(source), + static_cast( + data_loading_stats.total_updated_records)}})); + LogIfError(KVServerContextMap() + ->SafeMetric() + .LogUpDownCounter( + {{std::string(source), + static_cast( + data_loading_stats.total_deleted_records)}})); + LogIfError(KVServerContextMap() + ->SafeMetric() + .LogUpDownCounter( + {{std::string(source), + static_cast( + data_loading_stats.total_dropped_records)}})); } -absl::Status ApplyUpdateMutation(const KeyValueMutationRecord& record, +absl::Status ApplyUpdateMutation(std::string_view prefix, + const KeyValueMutationRecord& record, Cache& cache) { if (record.value_type() == Value::StringValue) { cache.UpdateKeyValue(record.key()->string_view(), GetRecordValue(record), - record.logical_commit_time()); + record.logical_commit_time(), prefix); return absl::OkStatus(); } if (record.value_type() == Value::StringSet) { auto values = GetRecordValue>(record); cache.UpdateKeyValueSet(record.key()->string_view(), absl::MakeSpan(values), - record.logical_commit_time()); + record.logical_commit_time(), prefix); return absl::OkStatus(); } return absl::InvalidArgumentError( @@ -83,16 +96,18 @@ absl::Status ApplyUpdateMutation(const KeyValueMutationRecord& record, " has unsupported value type: ", record.value_type())); } -absl::Status ApplyDeleteMutation(const KeyValueMutationRecord& record, +absl::Status ApplyDeleteMutation(std::string_view prefix, + const KeyValueMutationRecord& record, Cache& cache) { if (record.value_type() == Value::StringValue) { - cache.DeleteKey(record.key()->string_view(), record.logical_commit_time()); + cache.DeleteKey(record.key()->string_view(), record.logical_commit_time(), + prefix); return absl::OkStatus(); } if (record.value_type() == Value::StringSet) { auto values = GetRecordValue>(record); cache.DeleteValuesInSet(record.key()->string_view(), absl::MakeSpan(values), - record.logical_commit_time()); + record.logical_commit_time(), prefix); return absl::OkStatus(); } return absl::InvalidArgumentError( @@ -123,11 +138,12 @@ bool ShouldProcessRecord(const KeyValueMutationRecord& record, } absl::Status ApplyKeyValueMutationToCache( - const KeyValueMutationRecord& record, Cache& cache, int64_t& max_timestamp, - DataLoadingStats& data_loading_stats) { + std::string_view prefix, const KeyValueMutationRecord& record, Cache& cache, + int64_t& max_timestamp, DataLoadingStats& data_loading_stats) { switch (record.mutation_type()) { case KeyValueMutationType::Update: { - if (auto status = ApplyUpdateMutation(record, cache); !status.ok()) { + if (auto status = ApplyUpdateMutation(prefix, record, cache); + !status.ok()) { return status; } max_timestamp = std::max(max_timestamp, record.logical_commit_time()); @@ -135,7 +151,8 @@ absl::Status ApplyKeyValueMutationToCache( break; } case KeyValueMutationType::Delete: { - if (auto status = ApplyDeleteMutation(record, cache); !status.ok()) { + if (auto status = ApplyDeleteMutation(prefix, record, cache); + !status.ok()) { return status; } max_timestamp = std::max(max_timestamp, record.logical_commit_time()); @@ -151,12 +168,13 @@ absl::Status ApplyKeyValueMutationToCache( } absl::StatusOr LoadCacheWithData( + std::string_view data_source, std::string_view prefix, StreamRecordReader& record_reader, Cache& cache, int64_t& max_timestamp, const int32_t server_shard_num, const int32_t num_shards, UdfClient& udf_client, const KeySharder& key_sharder) { DataLoadingStats data_loading_stats; const auto process_data_record_fn = - [&cache, &max_timestamp, &data_loading_stats, server_shard_num, + [prefix, &cache, &max_timestamp, &data_loading_stats, server_shard_num, num_shards, &udf_client, &key_sharder](const DataRecord& data_record) { if (data_record.record_type() == Record::KeyValueMutationRecord) { const auto* record = data_record.record_as_KeyValueMutationRecord(); @@ -166,8 +184,8 @@ absl::StatusOr LoadCacheWithData( // this will get us in a loop return absl::OkStatus(); } - return ApplyKeyValueMutationToCache(*record, cache, max_timestamp, - data_loading_stats); + return ApplyKeyValueMutationToCache( + prefix, *record, cache, max_timestamp, data_loading_stats); } else if (data_record.record_type() == Record::UserDefinedFunctionsConfig) { const auto* udf_config = @@ -180,18 +198,16 @@ absl::StatusOr LoadCacheWithData( .logical_commit_time = udf_config->logical_commit_time(), .version = udf_config->version()}); } - LOG(ERROR) << "Received unsupported record "; - return absl::InvalidArgumentError("Record type not supported."); + return absl::InvalidArgumentError("Received unsupported record."); }; - - auto status = record_reader.ReadStreamRecords( + // TODO(b/314302953): ReadStreamRecords will skip over individual records that + // have errors. We should pass the file name to the function so that it will + // appear in error logs. + PS_RETURN_IF_ERROR(record_reader.ReadStreamRecords( [&process_data_record_fn](std::string_view raw) { return DeserializeDataRecord(raw, process_data_record_fn); - }); - if (!status.ok()) { - return status; - } - LogDataLoadingMetrics(data_loading_stats); + })); + LogDataLoadingMetrics(data_source, data_loading_stats); return data_loading_stats; } @@ -208,14 +224,12 @@ absl::StatusOr LoadCacheWithDataFromFile( return std::make_unique( options.blob_client.GetBlobReader(location)); }); - auto metadata = record_reader->GetKVFileMetadata(); - if (!metadata.ok()) { - return metadata.status(); - } - if (metadata->has_sharding_metadata() && - metadata->sharding_metadata().shard_num() != options.shard_num) { + PS_ASSIGN_OR_RETURN(auto metadata, record_reader->GetKVFileMetadata(), + _ << "Blob " << location); + if (metadata.has_sharding_metadata() && + metadata.sharding_metadata().shard_num() != options.shard_num) { LOG(INFO) << "Blob " << location << " belongs to shard num " - << metadata->sharding_metadata().shard_num() + << metadata.sharding_metadata().shard_num() << " but server shard num is " << options.shard_num << " Skipping it."; return DataLoadingStats{ @@ -224,14 +238,20 @@ absl::StatusOr LoadCacheWithDataFromFile( .total_dropped_records = 0, }; } - auto status = LoadCacheWithData(*record_reader, cache, max_timestamp, - options.shard_num, options.num_shards, - options.udf_client, options.key_sharder); - if (status.ok()) { - cache.RemoveDeletedKeys(max_timestamp); - } - return status; + std::string file_name = + location.prefix.empty() + ? location.key + : absl::StrCat(location.prefix, "/", location.key); + PS_ASSIGN_OR_RETURN( + auto data_loading_stats, + LoadCacheWithData(file_name, location.prefix, *record_reader, cache, + max_timestamp, options.shard_num, options.num_shards, + options.udf_client, options.key_sharder), + _ << "Blob: " << location); + cache.RemoveDeletedKeys(max_timestamp, location.prefix); + return data_loading_stats; } + absl::StatusOr TraceLoadCacheWithDataFromFile( BlobStorageClient::DataLocation location, const DataOrchestrator::Options& options) { @@ -241,6 +261,7 @@ absl::StatusOr TraceLoadCacheWithDataFromFile( }, "LoadCacheWithDataFromFile", {{"bucket", std::move(location.bucket)}, + {"prefix", std::move(location.prefix)}, {"key", std::move(location.key)}}); } @@ -248,9 +269,11 @@ class DataOrchestratorImpl : public DataOrchestrator { public: // `last_basename` is the last file seen during init. The cache is up to // date until this file. - DataOrchestratorImpl(Options options, std::string last_basename) + DataOrchestratorImpl( + Options options, + absl::flat_hash_map prefix_last_basenames) : options_(std::move(options)), - last_basename_of_init_(std::move(last_basename)) {} + prefix_last_basenames_(std::move(prefix_last_basenames)) {} ~DataOrchestratorImpl() override { if (!data_loader_thread_) return; @@ -270,38 +293,43 @@ class DataOrchestratorImpl : public DataOrchestrator { LOG(INFO) << "Stopped loading new data"; } - static absl::StatusOr Init(Options& options) { - auto ending_delta_file = LoadSnapshotFiles(options); - if (!ending_delta_file.ok()) { - return ending_delta_file.status(); - } - auto maybe_filenames = options.blob_client.ListBlobs( - {.bucket = options.data_bucket}, - {.prefix = std::string(FilePrefix()), - .start_after = *ending_delta_file}); - if (!maybe_filenames.ok()) { - return maybe_filenames.status(); + static absl::StatusOr> Init( + Options& options) { + auto ending_delta_files = LoadSnapshotFiles(options); + if (!ending_delta_files.ok()) { + return ending_delta_files.status(); } - LOG(INFO) << "Initializing cache with " << maybe_filenames->size() - << " delta files from " << options.data_bucket; - - std::string last_basename = std::move(*ending_delta_file); - for (auto&& basename : std::move(*maybe_filenames)) { - if (!IsDeltaFilename(basename)) { - LOG(WARNING) << "Saw a file " << basename - << " not in delta file format. Skipping it."; - continue; + for (const auto& prefix : options.blob_prefix_allowlist.Prefixes()) { + auto location = BlobStorageClient::DataLocation{ + .bucket = options.data_bucket, .prefix = prefix}; + auto iter = ending_delta_files->find(prefix); + auto maybe_filenames = options.blob_client.ListBlobs( + location, + {.prefix = std::string(FilePrefix()), + .start_after = + iter != ending_delta_files->end() ? iter->second : ""}); + if (!maybe_filenames.ok()) { + return maybe_filenames.status(); } - last_basename = basename; - if (const auto s = TraceLoadCacheWithDataFromFile( - {.bucket = options.data_bucket, .key = std::move(basename)}, - options); - !s.ok()) { - return s.status(); + LOG(INFO) << "Initializing cache with " << maybe_filenames->size() + << " delta files from " << location; + for (auto&& basename : std::move(*maybe_filenames)) { + auto blob = BlobStorageClient::DataLocation{ + .bucket = options.data_bucket, .prefix = prefix, .key = basename}; + if (!IsDeltaFilename(blob.key)) { + LOG(WARNING) << "Saw a file " << blob + << " not in delta file format. Skipping it."; + continue; + } + (*ending_delta_files)[prefix] = blob.key; + if (const auto s = TraceLoadCacheWithDataFromFile(blob, options); + !s.ok()) { + return s.status(); + } + LOG(INFO) << "Done loading " << blob; } - LOG(INFO) << "Done loading " << last_basename; } - return last_basename; + return ending_delta_files; } absl::Status Start() override { @@ -309,9 +337,10 @@ class DataOrchestratorImpl : public DataOrchestrator { return absl::OkStatus(); } LOG(INFO) << "Transitioning to state ContinuouslyLoadNewData"; + auto prefix_last_basenames = prefix_last_basenames_; absl::Status status = options_.delta_notifier.Start( options_.change_notifier, {.bucket = options_.data_bucket}, - last_basename_of_init_, + std::move(prefix_last_basenames), absl::bind_front(&DataOrchestratorImpl::EnqueueNewFilesToProcess, this)); if (!status.ok()) { @@ -324,8 +353,10 @@ class DataOrchestratorImpl : public DataOrchestrator { [this, &cache = options_.cache, &delta_stream_reader_factory = options_.delta_stream_reader_factory]( const std::string& message_body) { - return LoadCacheWithHighPriorityUpdates(delta_stream_reader_factory, - message_body, cache); + return LoadCacheWithHighPriorityUpdates( + kDefaultDataSourceForRealtimeUpdates, + kDefaultPrefixForRealTimeUpdates, delta_stream_reader_factory, + message_body, cache); }); } @@ -354,16 +385,25 @@ class DataOrchestratorImpl : public DataOrchestrator { unprocessed_basenames_.pop_back(); } LOG(INFO) << "Loading " << basename; - if (!IsDeltaFilename(basename)) { + auto blob = ParseBlobName(basename); + if (!IsDeltaFilename(blob.key)) { LOG(WARNING) << "Received file with invalid name: " << basename; continue; } + if (!options_.blob_prefix_allowlist.Contains(blob.prefix)) { + LOG(WARNING) << "Received file with prefix not allowlisted: " + << basename; + continue; + } RetryUntilOk( - [this, &basename] { + [this, &basename, &blob] { // TODO: distinguish status. Some can be retried while others // are fatal. return TraceLoadCacheWithDataFromFile( - {.bucket = options_.data_bucket, .key = basename}, options_); + {.bucket = options_.data_bucket, + .prefix = blob.prefix, + .key = blob.key}, + options_); }, "LoadNewFile", LogStatusSafeMetricsFn()); } @@ -379,67 +419,68 @@ class DataOrchestratorImpl : public DataOrchestrator { // Loads snapshot files if there are any. // Returns the latest delta file to be included in a snapshot. - static absl::StatusOr LoadSnapshotFiles(const Options& options) { - absl::StatusOr> snapshots = - options.blob_client.ListBlobs( - {.bucket = options.data_bucket}, - {.prefix = FilePrefix().data()}); - if (!snapshots.ok()) { - return snapshots.status(); - } - LOG(INFO) << "Initializing cache with snapshot file(s) from: " - << options.data_bucket; - std::string ending_delta_file; - for (int64_t s = snapshots->size() - 1; s >= 0; s--) { - std::string_view snapshot = snapshots->at(s); - if (!IsSnapshotFilename(snapshot)) { - LOG(WARNING) << "Saw a file " << snapshot - << " not in snapshot file format. Skipping it."; - continue; - } - BlobStorageClient::DataLocation location{.bucket = options.data_bucket, - .key = snapshot.data()}; - auto record_reader = - options.delta_stream_reader_factory.CreateConcurrentReader( - /*stream_factory=*/[&location, &options]() { - return std::make_unique( - options.blob_client.GetBlobReader(location)); - }); - auto metadata = record_reader->GetKVFileMetadata(); - if (!metadata.ok()) { - return metadata.status(); - } - if (metadata->has_sharding_metadata() && - metadata->sharding_metadata().shard_num() != options.shard_num) { - LOG(INFO) << "Snapshot " << location << " belongs to shard num " - << metadata->sharding_metadata().shard_num() - << " but server shard num is " << options.shard_num - << ". Skipping it."; + static absl::StatusOr> + LoadSnapshotFiles(const Options& options) { + absl::flat_hash_map ending_delta_files; + for (const auto& prefix : options.blob_prefix_allowlist.Prefixes()) { + auto location = BlobStorageClient::DataLocation{ + .bucket = options.data_bucket, .prefix = prefix}; + LOG(INFO) << "Initializing cache with snapshot file(s) from: " + << location; + PS_ASSIGN_OR_RETURN( + auto snapshot_group, + FindMostRecentFileGroup( + location, + FileGroupFilter{.file_type = FileType::SNAPSHOT, + .status = FileGroup::FileStatus::kComplete}, + options.blob_client)); + if (!snapshot_group.has_value()) { + LOG(INFO) << "No snapshot files found in: " << location; continue; } - LOG(INFO) << "Loading snapshot file: " << location; - if (auto status = TraceLoadCacheWithDataFromFile(location, options); - !status.ok()) { - return status.status(); - } - if (metadata->snapshot().ending_delta_file() > ending_delta_file) { - ending_delta_file = std::move(metadata->snapshot().ending_delta_file()); + for (const auto& snapshot : snapshot_group->Filenames()) { + auto snapshot_blob = BlobStorageClient::DataLocation{ + .bucket = options.data_bucket, .prefix = prefix, .key = snapshot}; + auto record_reader = + options.delta_stream_reader_factory.CreateConcurrentReader( + /*stream_factory=*/[&snapshot_blob, &options]() { + return std::make_unique( + options.blob_client.GetBlobReader(snapshot_blob)); + }); + PS_ASSIGN_OR_RETURN(auto metadata, record_reader->GetKVFileMetadata()); + if (metadata.has_sharding_metadata() && + metadata.sharding_metadata().shard_num() != options.shard_num) { + LOG(INFO) << "Snapshot " << snapshot_blob << " belongs to shard num " + << metadata.sharding_metadata().shard_num() + << " but server shard num is " << options.shard_num + << ". Skipping it."; + continue; + } + LOG(INFO) << "Loading snapshot file: " << snapshot_blob; + PS_ASSIGN_OR_RETURN( + auto stats, TraceLoadCacheWithDataFromFile(snapshot_blob, options)); + if (auto iter = ending_delta_files.find(prefix); + iter == ending_delta_files.end() || + metadata.snapshot().ending_delta_file() > iter->second) { + ending_delta_files[prefix] = metadata.snapshot().ending_delta_file(); + } + LOG(INFO) << "Done loading snapshot file: " << snapshot_blob; } - LOG(INFO) << "Done loading snapshot file: " << location; - break; } - return ending_delta_file; + return ending_delta_files; } absl::StatusOr LoadCacheWithHighPriorityUpdates( + std::string_view data_source, std::string_view prefix, StreamRecordReaderFactory& delta_stream_reader_factory, const std::string& record_string, Cache& cache) { std::istringstream is(record_string); int64_t max_timestamp = 0; auto record_reader = delta_stream_reader_factory.CreateReader(is); - return LoadCacheWithData(*record_reader, cache, max_timestamp, - options_.shard_num, options_.num_shards, - options_.udf_client, options_.key_sharder); + return LoadCacheWithData(data_source, prefix, *record_reader, cache, + max_timestamp, options_.shard_num, + options_.num_shards, options_.udf_client, + options_.key_sharder); } const Options options_; @@ -448,19 +489,19 @@ class DataOrchestratorImpl : public DataOrchestrator { std::unique_ptr data_loader_thread_; bool stop_ ABSL_GUARDED_BY(mu_) = false; // last basename of file in initialization. - const std::string last_basename_of_init_; + absl::flat_hash_map prefix_last_basenames_; }; } // namespace absl::StatusOr> DataOrchestrator::TryCreate( Options options) { - const auto maybe_last_basename = DataOrchestratorImpl::Init(options); - if (!maybe_last_basename.ok()) { - return maybe_last_basename.status(); + const auto prefix_last_basenames = DataOrchestratorImpl::Init(options); + if (!prefix_last_basenames.ok()) { + return prefix_last_basenames.status(); } auto orchestrator = std::make_unique( - std::move(options), std::move(maybe_last_basename.value())); + std::move(options), std::move(prefix_last_basenames.value())); return orchestrator; } } // namespace kv_server diff --git a/components/data_server/data_loading/data_orchestrator.h b/components/data_server/data_loading/data_orchestrator.h index 769f4917..d9a17c51 100644 --- a/components/data_server/data_loading/data_orchestrator.h +++ b/components/data_server/data_loading/data_orchestrator.h @@ -22,6 +22,7 @@ #include "absl/functional/any_invocable.h" #include "absl/status/status.h" +#include "components/data/blob_storage/blob_prefix_allowlist.h" #include "components/data/blob_storage/blob_storage_change_notifier.h" #include "components/data/blob_storage/blob_storage_client.h" #include "components/data/blob_storage/delta_file_notifier.h" @@ -56,6 +57,7 @@ class DataOrchestrator { const int32_t shard_num = 0; const int32_t num_shards = 1; const KeySharder key_sharder; + BlobPrefixAllowlist blob_prefix_allowlist; }; // Creates initial state. Scans the bucket and initializes the cache with data diff --git a/components/data_server/data_loading/data_orchestrator_test.cc b/components/data_server/data_loading/data_orchestrator_test.cc index a1ce263c..913e6324 100644 --- a/components/data_server/data_loading/data_orchestrator_test.cc +++ b/components/data_server/data_loading/data_orchestrator_test.cc @@ -18,6 +18,7 @@ #include #include +#include "absl/log/log.h" #include "absl/synchronization/notification.h" #include "components/data/common/mocks.h" #include "components/data/realtime/realtime_notifier.h" @@ -25,7 +26,6 @@ #include "components/data_server/cache/mocks.h" #include "components/udf/code_config.h" #include "components/udf/mocks.h" -#include "glog/logging.h" #include "gmock/gmock.h" #include "google/protobuf/text_format.h" #include "gtest/gtest.h" @@ -36,8 +36,8 @@ #include "public/sharding/sharding_function.h" #include "public/test_util/mocks.h" #include "public/test_util/proto_matcher.h" -#include "src/cpp/telemetry/mocks.h" +using kv_server::BlobPrefixAllowlist; using kv_server::BlobStorageChangeNotifier; using kv_server::BlobStorageClient; using kv_server::CodeConfig; @@ -70,8 +70,10 @@ using testing::_; using testing::AllOf; using testing::ByMove; using testing::Field; +using testing::Pair; using testing::Return; using testing::ReturnRef; +using testing::UnorderedElementsAre; namespace { // using google::protobuf::TextFormat; @@ -103,8 +105,9 @@ class DataOrchestratorTest : public ::testing::Test { .udf_client = udf_client_, .delta_stream_reader_factory = delta_stream_reader_factory_, .realtime_thread_pool_manager = realtime_thread_pool_manager_, - .key_sharder = kv_server::KeySharder( - kv_server::ShardingFunction{/*seed=*/""})}) {} + .key_sharder = + kv_server::KeySharder(kv_server::ShardingFunction{/*seed=*/""}), + .blob_prefix_allowlist = kv_server::BlobPrefixAllowlist("")}) {} MockBlobStorageClient blob_client_; MockDeltaFileNotifier notifier_; @@ -290,7 +293,9 @@ TEST_F(DataOrchestratorTest, InitCache_SkipsInvalidKVMutation) { ASSERT_TRUE(maybe_orchestrator.ok()); const std::string last_basename = ToDeltaFileName(1).value(); - EXPECT_CALL(notifier_, Start(_, GetTestLocation(), last_basename, _)) + EXPECT_CALL(notifier_, + Start(_, GetTestLocation(), + UnorderedElementsAre(Pair("", last_basename)), _)) .WillOnce(Return(absl::UnknownError(""))); EXPECT_FALSE((*maybe_orchestrator)->Start().ok()); } @@ -352,15 +357,17 @@ TEST_F(DataOrchestratorTest, InitCacheSuccess) { .WillOnce(Return(ByMove(std::move(update_reader)))) .WillOnce(Return(ByMove(std::move(delete_reader)))); - EXPECT_CALL(cache_, UpdateKeyValue("bar", "bar value", 3)).Times(1); - EXPECT_CALL(cache_, DeleteKey("bar", 3)).Times(1); - EXPECT_CALL(cache_, RemoveDeletedKeys(3)).Times(2); + EXPECT_CALL(cache_, UpdateKeyValue("bar", "bar value", 3, _)).Times(1); + EXPECT_CALL(cache_, DeleteKey("bar", 3, _)).Times(1); + EXPECT_CALL(cache_, RemoveDeletedKeys(3, _)).Times(2); auto maybe_orchestrator = DataOrchestrator::TryCreate(options_); ASSERT_TRUE(maybe_orchestrator.ok()); const std::string last_basename = ToDeltaFileName(2).value(); - EXPECT_CALL(notifier_, Start(_, GetTestLocation(), last_basename, _)) + EXPECT_CALL(notifier_, + Start(_, GetTestLocation(), + UnorderedElementsAre(Pair("", last_basename)), _)) .WillOnce(Return(absl::UnknownError(""))); EXPECT_FALSE((*maybe_orchestrator)->Start().ok()); } @@ -410,7 +417,9 @@ TEST_F(DataOrchestratorTest, UpdateUdfCodeSuccess) { ASSERT_TRUE(maybe_orchestrator.ok()); const std::string last_basename = ToDeltaFileName(1).value(); - EXPECT_CALL(notifier_, Start(_, GetTestLocation(), last_basename, _)) + EXPECT_CALL(notifier_, + Start(_, GetTestLocation(), + UnorderedElementsAre(Pair("", last_basename)), _)) .WillOnce(Return(absl::UnknownError(""))); EXPECT_FALSE((*maybe_orchestrator)->Start().ok()); } @@ -460,7 +469,9 @@ TEST_F(DataOrchestratorTest, UpdateUdfCodeFails_OrchestratorContinues) { ASSERT_TRUE(maybe_orchestrator.ok()); const std::string last_basename = ToDeltaFileName(1).value(); - EXPECT_CALL(notifier_, Start(_, GetTestLocation(), last_basename, _)) + EXPECT_CALL(notifier_, + Start(_, GetTestLocation(), + UnorderedElementsAre(Pair("", last_basename)), _)) .WillOnce(Return(absl::UnknownError(""))); EXPECT_FALSE((*maybe_orchestrator)->Start().ok()); } @@ -473,10 +484,11 @@ TEST_F(DataOrchestratorTest, StartLoading) { auto orchestrator = std::move(maybe_orchestrator.value()); const std::string last_basename = ""; - EXPECT_CALL(notifier_, Start(_, GetTestLocation(), last_basename, _)) - .WillOnce([](BlobStorageChangeNotifier& change_notifier, - BlobStorageClient::DataLocation location, - std::string start_after, + EXPECT_CALL(notifier_, + Start(_, GetTestLocation(), + absl::flat_hash_map(), _)) + .WillOnce([](BlobStorageChangeNotifier&, BlobStorageClient::DataLocation, + absl::flat_hash_map, std::function callback) { callback(ToDeltaFileName(6).value()); callback(ToDeltaFileName(7).value()); @@ -528,9 +540,9 @@ TEST_F(DataOrchestratorTest, StartLoading) { .WillOnce(Return(ByMove(std::move(update_reader)))) .WillOnce(Return(ByMove(std::move(delete_reader)))); - EXPECT_CALL(cache_, UpdateKeyValue("bar", "bar value", 3)).Times(1); - EXPECT_CALL(cache_, DeleteKey("bar", 3)).Times(1); - EXPECT_CALL(cache_, RemoveDeletedKeys(3)).Times(2); + EXPECT_CALL(cache_, UpdateKeyValue("bar", "bar value", 3, _)).Times(1); + EXPECT_CALL(cache_, DeleteKey("bar", 3, _)).Times(1); + EXPECT_CALL(cache_, RemoveDeletedKeys(3, _)).Times(2); EXPECT_TRUE(orchestrator->Start().ok()); LOG(INFO) << "Created ContinuouslyLoadNewData"; @@ -605,9 +617,9 @@ TEST_F(DataOrchestratorTest, InitCacheShardedSuccessSkipRecord) { .WillOnce(Return(ByMove(std::move(update_reader)))) .WillOnce(Return(ByMove(std::move(delete_reader)))); - EXPECT_CALL(strict_cache, RemoveDeletedKeys(0)).Times(1); - EXPECT_CALL(strict_cache, DeleteKey("shard2", 3)).Times(1); - EXPECT_CALL(strict_cache, RemoveDeletedKeys(3)).Times(1); + EXPECT_CALL(strict_cache, RemoveDeletedKeys(0, _)).Times(1); + EXPECT_CALL(strict_cache, DeleteKey("shard2", 3, _)).Times(1); + EXPECT_CALL(strict_cache, RemoveDeletedKeys(3, _)).Times(1); auto sharded_options = DataOrchestrator::Options{ .data_bucket = GetTestLocation().bucket, @@ -621,7 +633,8 @@ TEST_F(DataOrchestratorTest, InitCacheShardedSuccessSkipRecord) { .shard_num = 1, .num_shards = 2, .key_sharder = - kv_server::KeySharder(kv_server::ShardingFunction{/*seed=*/""})}; + kv_server::KeySharder(kv_server::ShardingFunction{/*seed=*/""}), + .blob_prefix_allowlist = BlobPrefixAllowlist("")}; auto maybe_orchestrator = DataOrchestrator::TryCreate(sharded_options); ASSERT_TRUE(maybe_orchestrator.ok()); @@ -659,4 +672,24 @@ TEST_F(DataOrchestratorTest, InitCacheSkipsSnapshotFilesForOtherShards) { EXPECT_TRUE(DataOrchestrator::TryCreate(options_).ok()); } +TEST_F(DataOrchestratorTest, VerifyLoadingDataFromPrefixes) { + for (auto file_type : + {FilePrefix(), FilePrefix()}) { + for (auto prefix : {"", "prefix1", "prefix2"}) { + EXPECT_CALL( + blob_client_, + ListBlobs( + BlobStorageClient::DataLocation{.bucket = "testbucket", + .prefix = prefix}, + AllOf(Field(&BlobStorageClient::ListOptions::start_after, ""), + Field(&BlobStorageClient::ListOptions::prefix, file_type)))) + .WillOnce(Return(std::vector({}))); + } + } + auto options = options_; + options.blob_prefix_allowlist = BlobPrefixAllowlist("prefix1,prefix2"); + auto maybe_orchestrator = DataOrchestrator::TryCreate(options); + ASSERT_TRUE(maybe_orchestrator.ok()); +} + } // namespace diff --git a/components/data_server/request_handler/BUILD.bazel b/components/data_server/request_handler/BUILD.bazel index 82926f77..292869da 100644 --- a/components/data_server/request_handler/BUILD.bazel +++ b/components/data_server/request_handler/BUILD.bazel @@ -32,15 +32,14 @@ cc_library( deps = [ ":get_values_adapter", "//components/data_server/cache", + "//components/util:request_context", "//public:base_types_cc_proto", "//public:constants", "//public/query:get_values_cc_grpc", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -60,8 +59,6 @@ cc_test( "//public/test_util:proto_matcher", "@com_github_grpc_grpc//:grpc++", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -77,20 +74,21 @@ cc_library( ":compression", ":ohttp_server_encryptor", "//components/data_server/cache", + "//components/telemetry:server_definition", "//components/udf:udf_client", + "//components/util:request_context", "//public:api_schema_cc_proto", "//public:base_types_cc_proto", "//public/query/v2:get_values_v2_cc_grpc", - "@com_github_google_glog//:glog", "@com_github_google_quiche//quiche:binary_http_unstable_api", "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/telemetry", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) @@ -109,8 +107,8 @@ cc_library( deps = [ "@brotli//:brotlidec", "@brotli//:brotlienc", - "@com_github_google_glog//:glog", "@com_github_google_quiche//quiche:quiche_unstable_api", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", ], ) @@ -129,7 +127,7 @@ cc_test( srcs = ["compression_brotli_test.cc"], deps = [ ":compression", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_googletest//:gtest_main", ], ) @@ -140,7 +138,6 @@ cc_test( srcs = [ "get_values_v2_handler_test.cc", ], - env = {"GLOG_v": "9"}, linkstatic = True, deps = [ ":get_values_v2_handler", @@ -150,14 +147,12 @@ cc_test( "//components/udf:udf_client", "//public/query/v2:get_values_v2_cc_grpc", "//public/test_util:proto_matcher", - "@com_github_google_glog//:glog", "@com_github_google_quiche//quiche:binary_http_unstable_api", "@com_github_google_quiche//quiche:oblivious_http_unstable_api", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", "@nlohmann_json//:lib", ], ) @@ -178,10 +173,10 @@ cc_library( "//public/applications/pa:response_utils", "//public/query:get_values_cc_grpc", "//public/query/v2:get_values_v2_cc_grpc", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) @@ -201,9 +196,7 @@ cc_test( "//public/test_util:proto_matcher", "@com_github_grpc_grpc//:grpc++", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", "@nlohmann_json//:lib", ], ) @@ -264,7 +257,7 @@ cc_library( "@com_github_google_quiche//quiche:oblivious_http_unstable_api", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:key_fetcher_manager", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:key_fetcher_manager", ], ) @@ -281,7 +274,7 @@ cc_library( "@com_github_google_quiche//quiche:oblivious_http_unstable_api", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:key_fetcher_manager", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:key_fetcher_manager", ], ) @@ -295,6 +288,6 @@ cc_test( ":ohttp_client_encryptor", ":ohttp_server_encryptor", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", ], ) diff --git a/components/data_server/request_handler/compression.cc b/components/data_server/request_handler/compression.cc index 3b3924cc..7defc60e 100644 --- a/components/data_server/request_handler/compression.cc +++ b/components/data_server/request_handler/compression.cc @@ -13,9 +13,9 @@ // limitations under the License. #include "components/data_server/request_handler/compression.h" +#include "absl/log/log.h" #include "components/data_server/request_handler/compression_brotli.h" #include "components/data_server/request_handler/uncompressed.h" -#include "glog/logging.h" #include "quiche/common/quiche_data_writer.h" namespace kv_server { diff --git a/components/data_server/request_handler/compression_brotli.cc b/components/data_server/request_handler/compression_brotli.cc index 31afe4ae..ba902ca4 100644 --- a/components/data_server/request_handler/compression_brotli.cc +++ b/components/data_server/request_handler/compression_brotli.cc @@ -18,10 +18,10 @@ #include #include +#include "absl/log/log.h" #include "absl/strings/str_join.h" #include "brotli/decode.h" #include "brotli/encode.h" -#include "glog/logging.h" #include "quiche/common/quiche_data_writer.h" namespace kv_server { diff --git a/components/data_server/request_handler/compression_brotli_test.cc b/components/data_server/request_handler/compression_brotli_test.cc index 1f474aae..3c5daf5b 100644 --- a/components/data_server/request_handler/compression_brotli_test.cc +++ b/components/data_server/request_handler/compression_brotli_test.cc @@ -17,8 +17,8 @@ #include #include +#include "absl/log/log.h" #include "components/data_server/request_handler/uncompressed.h" -#include "glog/logging.h" #include "gmock/gmock.h" #include "gtest/gtest.h" diff --git a/components/data_server/request_handler/get_values_adapter.cc b/components/data_server/request_handler/get_values_adapter.cc index 2703297a..0c86b5c9 100644 --- a/components/data_server/request_handler/get_values_adapter.cc +++ b/components/data_server/request_handler/get_values_adapter.cc @@ -21,13 +21,13 @@ #include #include +#include "absl/log/log.h" #include "components/data_server/request_handler/v2_response_data.pb.h" -#include "glog/logging.h" #include "google/protobuf/util/json_util.h" #include "public/api_schema.pb.h" #include "public/applications/pa/api_overlay.pb.h" #include "public/applications/pa/response_utils.h" -#include "src/cpp/util/status_macro/status_macros.h" +#include "src/util/status_macro/status_macros.h" namespace kv_server { namespace { diff --git a/components/data_server/request_handler/get_values_adapter_test.cc b/components/data_server/request_handler/get_values_adapter_test.cc index aba170e7..15d77b47 100644 --- a/components/data_server/request_handler/get_values_adapter_test.cc +++ b/components/data_server/request_handler/get_values_adapter_test.cc @@ -26,16 +26,13 @@ #include "public/applications/pa/api_overlay.pb.h" #include "public/applications/pa/response_utils.h" #include "public/test_util/proto_matcher.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" namespace kv_server { namespace { using google::protobuf::TextFormat; -using privacy_sandbox::server_common::MockMetricsRecorder; using testing::_; using testing::Return; @@ -54,8 +51,9 @@ class GetValuesAdapterTest : public ::testing::Test { protected: void SetUp() override { v2_handler_ = std::make_unique( - mock_udf_client_, mock_metrics_recorder_, fake_key_fetcher_manager_); + mock_udf_client_, fake_key_fetcher_manager_); get_values_adapter_ = GetValuesAdapter::Create(std::move(v2_handler_)); + InitMetricsContextMap(); } privacy_sandbox::server_common::FakeKeyFetcherManager @@ -63,7 +61,6 @@ class GetValuesAdapterTest : public ::testing::Test { std::unique_ptr get_values_adapter_; std::unique_ptr v2_handler_; MockUdfClient mock_udf_client_; - MockMetricsRecorder mock_metrics_recorder_; }; TEST_F(GetValuesAdapterTest, EmptyRequestReturnsEmptyResponse) { @@ -71,8 +68,9 @@ TEST_F(GetValuesAdapterTest, EmptyRequestReturnsEmptyResponse) { TextFormat::ParseFromString(kEmptyMetadata, &udf_metadata); nlohmann::json output = nlohmann::json::parse(R"({"keyGroupOutputs": {}})"); - EXPECT_CALL(mock_udf_client_, - ExecuteCode(EqualsProto(udf_metadata), testing::IsEmpty())) + EXPECT_CALL( + mock_udf_client_, + ExecuteCode(testing::_, EqualsProto(udf_metadata), testing::IsEmpty())) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -133,7 +131,7 @@ data { )", &key_group_outputs); EXPECT_CALL(mock_udf_client_, - ExecuteCode(EqualsProto(udf_metadata), + ExecuteCode(testing::_, EqualsProto(udf_metadata), testing::ElementsAre(EqualsProto(arg)))) .WillOnce(Return( application_pa::KeyGroupOutputsToJson(key_group_outputs).value())); @@ -236,7 +234,7 @@ data { &key_group_outputs); EXPECT_CALL( mock_udf_client_, - ExecuteCode(EqualsProto(udf_metadata), + ExecuteCode(testing::_, EqualsProto(udf_metadata), testing::ElementsAre(EqualsProto(arg1), EqualsProto(arg2)))) .WillOnce(Return( application_pa::KeyGroupOutputsToJson(key_group_outputs).value())); @@ -265,7 +263,7 @@ TEST_F(GetValuesAdapterTest, V2ResponseIsNullReturnsError) { nlohmann::json output = R"({ "keyGroupOutpus": [] })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -286,7 +284,7 @@ TEST_F(GetValuesAdapterTest, KeyGroupOutputWithEmptyKVsReturnsOk) { }], "udfOutputApiVersion": 1 })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -307,7 +305,7 @@ TEST_F(GetValuesAdapterTest, KeyGroupOutputWithInvalidNamespaceTagIsIgnored) { }], "udfOutputApiVersion": 1 })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -328,7 +326,7 @@ TEST_F(GetValuesAdapterTest, KeyGroupOutputWithNoCustomTagIsIgnored) { }], "udfOutputApiVersion": 1 })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -349,7 +347,7 @@ TEST_F(GetValuesAdapterTest, KeyGroupOutputWithNoNamespaceTagIsIgnored) { }], "udfOutputApiVersion": 1 })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -375,7 +373,7 @@ TEST_F(GetValuesAdapterTest, }], "udfOutputApiVersion": 1 })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -409,7 +407,7 @@ TEST_F(GetValuesAdapterTest, KeyGroupOutputHasDifferentValueTypesReturnsOk) { }], "udfOutputApiVersion": 1 })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; @@ -489,7 +487,7 @@ TEST_F(GetValuesAdapterTest, ValueWithStatusSuccess) { }], "udfOutputApiVersion": 1 })"_json; - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _)) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, _)) .WillOnce(Return(output.dump())); v1::GetValuesRequest v1_request; diff --git a/components/data_server/request_handler/get_values_handler.cc b/components/data_server/request_handler/get_values_handler.cc index 319a5fed..f8cce490 100644 --- a/components/data_server/request_handler/get_values_handler.cc +++ b/components/data_server/request_handler/get_values_handler.cc @@ -19,20 +19,16 @@ #include #include +#include "absl/log/log.h" #include "absl/strings/str_replace.h" #include "absl/strings/str_split.h" #include "components/data_server/request_handler/get_values_adapter.h" -#include "glog/logging.h" #include "grpcpp/grpcpp.h" #include "public/constants.h" #include "public/query/get_values.grpc.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry.h" #include "src/google/protobuf/message.h" #include "src/google/protobuf/struct.pb.h" - -constexpr char* kCacheKeyHit = "CacheKeyHit"; -constexpr char* kCacheKeyMiss = "CacheKeyMiss"; +#include "src/telemetry/telemetry.h" namespace kv_server { namespace { @@ -40,8 +36,6 @@ using google::protobuf::RepeatedPtrField; using google::protobuf::Struct; using google::protobuf::Value; using grpc::StatusCode; -using privacy_sandbox::server_common::GetTracer; -using privacy_sandbox::server_common::MetricsRecorder; using v1::GetValuesRequest; using v1::GetValuesResponse; using v1::KeyValueService; @@ -58,25 +52,25 @@ absl::flat_hash_set GetKeys( return key_list; } -void ProcessKeys(const RepeatedPtrField& keys, const Cache& cache, - MetricsRecorder& metrics_recorder, - google::protobuf::Map& - result_struct) { +void ProcessKeys( + const RequestContext& request_context, + const RepeatedPtrField& keys, const Cache& cache, + google::protobuf::Map& result_struct, + bool add_missing_keys_v1) { if (keys.empty()) return; auto actual_keys = GetKeys(keys); - auto kv_pairs = cache.GetKeyValuePairs(actual_keys); - - if (kv_pairs.empty()) - metrics_recorder.IncrementEventCounter(kCacheKeyMiss); - else - metrics_recorder.IncrementEventCounter(kCacheKeyHit); + auto kv_pairs = cache.GetKeyValuePairs(request_context, actual_keys); + // TODO(b/326118416): Record cache hit and miss metrics for (const auto& key : actual_keys) { v1::V1SingleLookupResult result; const auto key_iter = kv_pairs.find(key); if (key_iter == kv_pairs.end()) { - auto status = result.mutable_status(); - status->set_code(static_cast(absl::StatusCode::kNotFound)); - status->set_message("Key not found"); + if (add_missing_keys_v1) { + auto status = result.mutable_status(); + status->set_code(static_cast(absl::StatusCode::kNotFound)); + status->set_message("Key not found"); + result_struct[key] = std::move(result); + } } else { Value value_proto; absl::Status status = google::protobuf::util::JsonStringToMessage( @@ -90,40 +84,41 @@ void ProcessKeys(const RepeatedPtrField& keys, const Cache& cache, value.set_string_value(std::move(key_iter->second)); *result.mutable_value() = std::move(value); } + result_struct[key] = std::move(result); } - result_struct[key] = std::move(result); } } } // namespace -grpc::Status GetValuesHandler::GetValues(const GetValuesRequest& request, +grpc::Status GetValuesHandler::GetValues(const RequestContext& request_context, + const GetValuesRequest& request, GetValuesResponse* response) const { if (use_v2_) { VLOG(5) << "Using V2 adapter for " << request.DebugString(); return adapter_.CallV2Handler(request, *response); } - if (!request.kv_internal().empty()) { VLOG(5) << "Processing kv_internal for " << request.DebugString(); - ProcessKeys(request.kv_internal(), cache_, metrics_recorder_, - *response->mutable_kv_internal()); + ProcessKeys(request_context, request.kv_internal(), cache_, + *response->mutable_kv_internal(), add_missing_keys_v1_); } if (!request.keys().empty()) { VLOG(5) << "Processing keys for " << request.DebugString(); - ProcessKeys(request.keys(), cache_, metrics_recorder_, - *response->mutable_keys()); + ProcessKeys(request_context, request.keys(), cache_, + *response->mutable_keys(), add_missing_keys_v1_); } if (!request.render_urls().empty()) { VLOG(5) << "Processing render_urls for " << request.DebugString(); - ProcessKeys(request.render_urls(), cache_, metrics_recorder_, - *response->mutable_render_urls()); + ProcessKeys(request_context, request.render_urls(), cache_, + *response->mutable_render_urls(), add_missing_keys_v1_); } if (!request.ad_component_render_urls().empty()) { VLOG(5) << "Processing ad_component_render_urls for " << request.DebugString(); - ProcessKeys(request.ad_component_render_urls(), cache_, metrics_recorder_, - *response->mutable_ad_component_render_urls()); + ProcessKeys(request_context, request.ad_component_render_urls(), cache_, + *response->mutable_ad_component_render_urls(), + add_missing_keys_v1_); } return grpc::Status::OK; } diff --git a/components/data_server/request_handler/get_values_handler.h b/components/data_server/request_handler/get_values_handler.h index 880eed06..ef2be6cf 100644 --- a/components/data_server/request_handler/get_values_handler.h +++ b/components/data_server/request_handler/get_values_handler.h @@ -24,7 +24,6 @@ #include "components/data_server/request_handler/get_values_adapter.h" #include "grpcpp/grpcpp.h" #include "public/query/get_values.grpc.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" #include "src/google/protobuf/struct.pb.h" namespace kv_server { @@ -33,26 +32,24 @@ namespace kv_server { // See the Service proto definition for details. class GetValuesHandler { public: - explicit GetValuesHandler( - const Cache& cache, const GetValuesAdapter& adapter, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder, - bool use_v2) + explicit GetValuesHandler(const Cache& cache, const GetValuesAdapter& adapter, + bool use_v2, bool add_missing_keys_v1 = true) : cache_(std::move(cache)), adapter_(std::move(adapter)), - metrics_recorder_(metrics_recorder), - use_v2_(use_v2) {} + use_v2_(use_v2), + add_missing_keys_v1_(add_missing_keys_v1) {} // TODO: Implement hostname, ad/render url lookups. - grpc::Status GetValues(const v1::GetValuesRequest& request, + grpc::Status GetValues(const RequestContext& request_context, + const v1::GetValuesRequest& request, v1::GetValuesResponse* response) const; private: const Cache& cache_; const GetValuesAdapter& adapter_; - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; - // If true, routes requests through V2 (UDF). Otherwise, calls cache. const bool use_v2_; + const bool add_missing_keys_v1_; }; } // namespace kv_server diff --git a/components/data_server/request_handler/get_values_handler_test.cc b/components/data_server/request_handler/get_values_handler_test.cc index 3197839b..76f0d40a 100644 --- a/components/data_server/request_handler/get_values_handler_test.cc +++ b/components/data_server/request_handler/get_values_handler_test.cc @@ -28,15 +28,12 @@ #include "grpcpp/grpcpp.h" #include "gtest/gtest.h" #include "public/test_util/proto_matcher.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { using google::protobuf::TextFormat; using grpc::StatusCode; -using privacy_sandbox::server_common::MockMetricsRecorder; using testing::_; using testing::DoAll; using testing::Return; @@ -48,13 +45,21 @@ using v1::GetValuesResponse; class GetValuesHandlerTest : public ::testing::Test { protected: + GetValuesHandlerTest() { + InitMetricsContextMap(); + scope_metrics_context_ = std::make_unique(); + request_context_ = + std::make_unique(*scope_metrics_context_); + } MockCache mock_cache_; - MockMetricsRecorder mock_metrics_recorder_; MockGetValuesAdapter mock_get_values_adapter_; + RequestContext& GetRequestContext() { return *request_context_; } + std::unique_ptr scope_metrics_context_; + std::unique_ptr request_context_; }; TEST_F(GetValuesHandlerTest, ReturnsExistingKeyTwice) { - EXPECT_CALL(mock_cache_, GetKeyValuePairs(UnorderedElementsAre("my_key"))) + EXPECT_CALL(mock_cache_, GetKeyValuePairs(_, UnorderedElementsAre("my_key"))) .Times(2) .WillRepeatedly(Return(absl::flat_hash_map{ {"my_key", "my_value"}})); @@ -62,9 +67,9 @@ TEST_F(GetValuesHandlerTest, ReturnsExistingKeyTwice) { request.add_keys("my_key"); GetValuesResponse response; GetValuesHandler handler(mock_cache_, mock_get_values_adapter_, - mock_metrics_recorder_, /*use_v2=*/false); - const auto result = handler.GetValues(request, &response); + const auto result = + handler.GetValues(GetRequestContext(), request, &response); ASSERT_TRUE(result.ok()) << "code: " << result.error_code() << ", msg: " << result.error_message(); @@ -77,13 +82,13 @@ TEST_F(GetValuesHandlerTest, ReturnsExistingKeyTwice) { &expected); EXPECT_THAT(response, EqualsProto(expected)); - ASSERT_TRUE(handler.GetValues(request, &response).ok()); + ASSERT_TRUE(handler.GetValues(GetRequestContext(), request, &response).ok()); EXPECT_THAT(response, EqualsProto(expected)); } TEST_F(GetValuesHandlerTest, RepeatedKeys) { EXPECT_CALL(mock_cache_, - GetKeyValuePairs(UnorderedElementsAre("key1", "key2", "key3"))) + GetKeyValuePairs(_, UnorderedElementsAre("key1", "key2", "key3"))) .Times(1) .WillRepeatedly(Return( absl::flat_hash_map{{"key1", "value1"}})); @@ -91,9 +96,8 @@ TEST_F(GetValuesHandlerTest, RepeatedKeys) { request.add_keys("key1,key2,key3"); GetValuesResponse response; GetValuesHandler handler(mock_cache_, mock_get_values_adapter_, - mock_metrics_recorder_, /*use_v2=*/false); - ASSERT_TRUE(handler.GetValues(request, &response).ok()); + ASSERT_TRUE(handler.GetValues(GetRequestContext(), request, &response).ok()); GetValuesResponse expected; TextFormat::ParseFromString( @@ -115,9 +119,34 @@ TEST_F(GetValuesHandlerTest, RepeatedKeys) { EXPECT_THAT(response, EqualsProto(expected)); } +TEST_F(GetValuesHandlerTest, RepeatedKeysSkipEmpty) { + EXPECT_CALL(mock_cache_, + GetKeyValuePairs(_, UnorderedElementsAre("key1", "key2", "key3"))) + .Times(1) + .WillRepeatedly(Return( + absl::flat_hash_map{{"key1", "value1"}})); + GetValuesRequest request; + request.add_keys("key1,key2,key3"); + GetValuesResponse response; + GetValuesHandler handler(mock_cache_, mock_get_values_adapter_, + /*use_v2=*/false, /*add_missing_keys_v1=*/false); + ASSERT_TRUE(handler.GetValues(GetRequestContext(), request, &response).ok()); + + GetValuesResponse expected; + TextFormat::ParseFromString( + R"pb( + keys { + key: "key1" + value { value { string_value: "value1" } } + } + )pb", + &expected); + EXPECT_THAT(response, EqualsProto(expected)); +} + TEST_F(GetValuesHandlerTest, ReturnsMultipleExistingKeysSameNamespace) { EXPECT_CALL(mock_cache_, - GetKeyValuePairs(UnorderedElementsAre("key1", "key2"))) + GetKeyValuePairs(_, UnorderedElementsAre("key1", "key2"))) .Times(1) .WillOnce(Return(absl::flat_hash_map{ {"key1", "value1"}, {"key2", "value2"}})); @@ -126,9 +155,8 @@ TEST_F(GetValuesHandlerTest, ReturnsMultipleExistingKeysSameNamespace) { request.add_keys("key2"); GetValuesResponse response; GetValuesHandler handler(mock_cache_, mock_get_values_adapter_, - mock_metrics_recorder_, /*use_v2=*/false); - ASSERT_TRUE(handler.GetValues(request, &response).ok()); + ASSERT_TRUE(handler.GetValues(GetRequestContext(), request, &response).ok()); GetValuesResponse expected; TextFormat::ParseFromString(R"pb( @@ -145,11 +173,11 @@ TEST_F(GetValuesHandlerTest, ReturnsMultipleExistingKeysSameNamespace) { } TEST_F(GetValuesHandlerTest, ReturnsMultipleExistingKeysDifferentNamespace) { - EXPECT_CALL(mock_cache_, GetKeyValuePairs(UnorderedElementsAre("key1"))) + EXPECT_CALL(mock_cache_, GetKeyValuePairs(_, UnorderedElementsAre("key1"))) .Times(1) .WillOnce(Return( absl::flat_hash_map{{"key1", "value1"}})); - EXPECT_CALL(mock_cache_, GetKeyValuePairs(UnorderedElementsAre("key2"))) + EXPECT_CALL(mock_cache_, GetKeyValuePairs(_, UnorderedElementsAre("key2"))) .Times(1) .WillOnce(Return( absl::flat_hash_map{{"key2", "value2"}})); @@ -158,9 +186,8 @@ TEST_F(GetValuesHandlerTest, ReturnsMultipleExistingKeysDifferentNamespace) { request.add_ad_component_render_urls("key2"); GetValuesResponse response; GetValuesHandler handler(mock_cache_, mock_get_values_adapter_, - mock_metrics_recorder_, /*use_v2=*/false); - ASSERT_TRUE(handler.GetValues(request, &response).ok()); + ASSERT_TRUE(handler.GetValues(GetRequestContext(), request, &response).ok()); GetValuesResponse expected; TextFormat::ParseFromString(R"pb(render_urls { @@ -252,7 +279,7 @@ TEST_F(GetValuesHandlerTest, TestResponseOnDifferentValueFormats) { })json"; EXPECT_CALL(mock_cache_, - GetKeyValuePairs(UnorderedElementsAre("key1", "key2", "key3"))) + GetKeyValuePairs(_, UnorderedElementsAre("key1", "key2", "key3"))) .Times(1) .WillOnce(Return(absl::flat_hash_map{ {"key1", value1}, {"key2", value2}, {"key3", value3}})); @@ -263,9 +290,8 @@ TEST_F(GetValuesHandlerTest, TestResponseOnDifferentValueFormats) { request.add_keys("key3"); GetValuesResponse response; GetValuesHandler handler(mock_cache_, mock_get_values_adapter_, - mock_metrics_recorder_, /*use_v2=*/false); - ASSERT_TRUE(handler.GetValues(request, &response).ok()); + ASSERT_TRUE(handler.GetValues(GetRequestContext(), request, &response).ok()); GetValuesResponse expected_from_pb; TextFormat::ParseFromString(response_pb_string, &expected_from_pb); EXPECT_THAT(response, EqualsProto(expected_from_pb)); @@ -292,9 +318,8 @@ TEST_F(GetValuesHandlerTest, CallsV2Adapter) { request.add_keys("key1"); GetValuesResponse response; GetValuesHandler handler(mock_cache_, mock_get_values_adapter_, - mock_metrics_recorder_, /*use_v2=*/true); - ASSERT_TRUE(handler.GetValues(request, &response).ok()); + ASSERT_TRUE(handler.GetValues(GetRequestContext(), request, &response).ok()); EXPECT_THAT(response, EqualsProto(adapter_response)); } diff --git a/components/data_server/request_handler/get_values_v2_handler.cc b/components/data_server/request_handler/get_values_v2_handler.cc index 3b33aab2..06d3b5c7 100644 --- a/components/data_server/request_handler/get_values_v2_handler.cc +++ b/components/data_server/request_handler/get_values_v2_handler.cc @@ -21,10 +21,11 @@ #include #include "absl/algorithm/container.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/strings/ascii.h" #include "components/data_server/request_handler/ohttp_server_encryptor.h" -#include "glog/logging.h" +#include "components/telemetry/server_definition.h" #include "google/protobuf/util/json_util.h" #include "grpcpp/grpcpp.h" #include "public/base_types.pb.h" @@ -33,11 +34,8 @@ #include "quiche/binary_http/binary_http_message.h" #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h" #include "quiche/oblivious_http/oblivious_http_gateway.h" -#include "src/cpp/telemetry/telemetry.h" -#include "src/cpp/util/status_macro/status_macros.h" - -constexpr char* kCacheKeyV2Hit = "CacheKeyHit"; -constexpr char* kCacheKeyV2Miss = "CacheKeyMiss"; +#include "src/telemetry/telemetry.h" +#include "src/util/status_macro/status_macros.h" namespace kv_server { namespace { @@ -75,16 +73,35 @@ grpc::Status GetValuesV2Handler::GetValuesHttp( GetValuesHttp(request.raw_body().data(), *response->mutable_data())); } -absl::Status GetValuesV2Handler::GetValuesHttp( - std::string_view request, std::string& json_response) const { +absl::Status GetValuesV2Handler::GetValuesHttp(std::string_view request, + std::string& response, + ContentType content_type) const { v2::GetValuesRequest request_proto; - PS_RETURN_IF_ERROR( - google::protobuf::util::JsonStringToMessage(request, &request_proto)); + if (content_type == ContentType::kJson) { + PS_RETURN_IF_ERROR( + google::protobuf::util::JsonStringToMessage(request, &request_proto)); + } else { // proto + if (!request_proto.ParseFromString(request)) { + auto error_message = + "Cannot parse request as a valid serilized proto object."; + VLOG(4) << error_message; + return absl::InvalidArgumentError(error_message); + } + } VLOG(9) << "Converted the http request to proto: " << request_proto.DebugString(); v2::GetValuesResponse response_proto; PS_RETURN_IF_ERROR(GetValues(request_proto, &response_proto)); - return MessageToJsonString(response_proto, &json_response); + if (content_type == ContentType::kJson) { + return MessageToJsonString(response_proto, &response); + } + // content_type == proto + if (!response_proto.SerializeToString(&response)) { + auto error_message = "Cannot serialize the response as a proto."; + VLOG(4) << error_message; + return absl::InvalidArgumentError(error_message); + } + return absl::OkStatus(); } grpc::Status GetValuesV2Handler::BinaryHttpGetValues( @@ -94,22 +111,38 @@ grpc::Status GetValuesV2Handler::BinaryHttpGetValues( *response->mutable_data())); } +GetValuesV2Handler::ContentType GetValuesV2Handler::GetContentType( + const quiche::BinaryHttpRequest& deserialized_req) const { + for (const auto& header : deserialized_req.GetHeaderFields()) { + if (absl::AsciiStrToLower(header.name) == kContentTypeHeader && + absl::AsciiStrToLower(header.value) == + kContentEncodingProtoHeaderValue) { + return ContentType::kProto; + } + } + return ContentType::kJson; +} + absl::StatusOr GetValuesV2Handler::BuildSuccessfulGetValuesBhttpResponse( std::string_view bhttp_request_body) const { VLOG(9) << "Handling the binary http layer"; - PS_ASSIGN_OR_RETURN(quiche::BinaryHttpRequest maybe_deserialized_req, + PS_ASSIGN_OR_RETURN(quiche::BinaryHttpRequest deserialized_req, quiche::BinaryHttpRequest::Create(bhttp_request_body), _ << "Failed to deserialize binary http request"); - VLOG(3) << "BinaryHttpGetValues request: " - << maybe_deserialized_req.DebugString(); - - std::string json_response; + VLOG(3) << "BinaryHttpGetValues request: " << deserialized_req.DebugString(); + std::string response; + auto content_type = GetContentType(deserialized_req); PS_RETURN_IF_ERROR( - GetValuesHttp(maybe_deserialized_req.body(), json_response)); - + GetValuesHttp(deserialized_req.body(), response, content_type)); quiche::BinaryHttpResponse bhttp_response(200); - bhttp_response.set_body(std::move(json_response)); + if (content_type == ContentType::kProto) { + bhttp_response.AddHeaderField({ + .name = std::string(kContentTypeHeader), + .value = std::string(kContentEncodingProtoHeaderValue), + }); + } + bhttp_response.set_body(std::move(response)); return bhttp_response; } @@ -159,6 +192,7 @@ grpc::Status GetValuesV2Handler::ObliviousGetValues( } void GetValuesV2Handler::ProcessOnePartition( + RequestContext request_context, const google::protobuf::Struct& req_metadata, const v2::RequestPartition& req_partition, v2::ResponsePartition& resp_partition) const { @@ -166,7 +200,8 @@ void GetValuesV2Handler::ProcessOnePartition( UDFExecutionMetadata udf_metadata; *udf_metadata.mutable_request_metadata() = req_metadata; const auto maybe_output_string = udf_client_.ExecuteCode( - std::move(udf_metadata), req_partition.arguments()); + std::move(request_context), std::move(udf_metadata), + req_partition.arguments()); if (!maybe_output_string.ok()) { resp_partition.mutable_status()->set_code( static_cast(maybe_output_string.status().code())); @@ -181,8 +216,11 @@ void GetValuesV2Handler::ProcessOnePartition( grpc::Status GetValuesV2Handler::GetValues( const v2::GetValuesRequest& request, v2::GetValuesResponse* response) const { + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); if (request.partitions().size() == 1) { - ProcessOnePartition(request.metadata(), request.partitions(0), + ProcessOnePartition(std::move(request_context), request.metadata(), + request.partitions(0), *response->mutable_single_partition()); return grpc::Status::OK; } diff --git a/components/data_server/request_handler/get_values_v2_handler.h b/components/data_server/request_handler/get_values_v2_handler.h index 32960e29..52f1c0d6 100644 --- a/components/data_server/request_handler/get_values_v2_handler.h +++ b/components/data_server/request_handler/get_values_v2_handler.h @@ -26,15 +26,26 @@ #include "absl/strings/escaping.h" #include "components/data_server/cache/cache.h" #include "components/data_server/request_handler/compression.h" +#include "components/telemetry/server_definition.h" #include "components/udf/udf_client.h" +#include "components/util/request_context.h" #include "grpcpp/grpcpp.h" #include "public/query/v2/get_values_v2.grpc.pb.h" #include "quiche/binary_http/binary_http_message.h" -#include "src/cpp/encryption/key_fetcher/src/key_fetcher_manager.h" -#include "src/cpp/telemetry/metrics_recorder.h" +#include "src/encryption/key_fetcher/key_fetcher_manager.h" namespace kv_server { +// Content Type Header Name. Can be set for bhttp request to proto or json +// values below. +inline constexpr std::string_view kContentTypeHeader = "content-type"; +// Protobuf Content Type Header Value. +inline constexpr std::string_view kContentEncodingProtoHeaderValue = + "application/protobuf"; +// Json Content Type Header Value. +inline constexpr std::string_view kContentEncodingJsonHeaderValue = + "application/json"; + // Handles the request family of *GetValues. // See the Service proto definition for details. class GetValuesV2Handler { @@ -42,14 +53,12 @@ class GetValuesV2Handler { // Accepts a functor to create compression blob builder for testing purposes. explicit GetValuesV2Handler( const UdfClient& udf_client, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder, privacy_sandbox::server_common::KeyFetcherManagerInterface& key_fetcher_manager, std::function create_compression_group_concatenator = &CompressionGroupConcatenator::Create) : udf_client_(udf_client), - metrics_recorder_(metrics_recorder), create_compression_group_concatenator_( std::move(create_compression_group_concatenator)), key_fetcher_manager_(key_fetcher_manager) {} @@ -81,8 +90,16 @@ class GetValuesV2Handler { google::api::HttpBody* response) const; private: - absl::Status GetValuesHttp(std::string_view request, - std::string& json_response) const; + enum class ContentType { + kJson = 0, + kProto, + }; + ContentType GetContentType( + const quiche::BinaryHttpRequest& deserialized_req) const; + + absl::Status GetValuesHttp( + std::string_view request, std::string& json_response, + ContentType content_type = ContentType::kJson) const; // On success, returns a BinaryHttpResponse with a successful response. The // reason that this is a separate function is so that the error status @@ -99,14 +116,14 @@ class GetValuesV2Handler { std::string& response) const; // Invokes UDF to process one partition. - void ProcessOnePartition(const google::protobuf::Struct& req_metadata, + void ProcessOnePartition(RequestContext request_context, + const google::protobuf::Struct& req_metadata, const v2::RequestPartition& req_partition, v2::ResponsePartition& resp_partition) const; const UdfClient& udf_client_; std::function create_compression_group_concatenator_; - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; privacy_sandbox::server_common::KeyFetcherManagerInterface& key_fetcher_manager_; }; diff --git a/components/data_server/request_handler/get_values_v2_handler_test.cc b/components/data_server/request_handler/get_values_v2_handler_test.cc index f351bb7c..9c7376f7 100644 --- a/components/data_server/request_handler/get_values_v2_handler_test.cc +++ b/components/data_server/request_handler/get_values_v2_handler_test.cc @@ -19,10 +19,10 @@ #include #include +#include "absl/log/log.h" #include "components/data_server/cache/cache.h" #include "components/data_server/cache/mocks.h" #include "components/udf/mocks.h" -#include "glog/logging.h" #include "gmock/gmock.h" #include "google/protobuf/text_format.h" #include "grpcpp/grpcpp.h" @@ -33,16 +33,13 @@ #include "quiche/binary_http/binary_http_message.h" #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h" #include "quiche/oblivious_http/oblivious_http_client.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" namespace kv_server { namespace { using google::protobuf::TextFormat; using grpc::StatusCode; -using privacy_sandbox::server_common::MockMetricsRecorder; using testing::_; using testing::Return; using testing::ReturnRef; @@ -57,13 +54,25 @@ enum class ProtocolType { kObliviousHttp, }; +struct TestingParameters { + ProtocolType protocol_type; + const std::string_view content_type; +}; + class GetValuesHandlerTest : public ::testing::Test, - public ::testing::WithParamInterface { + public ::testing::WithParamInterface { protected: + void SetUp() override { InitMetricsContextMap(); } template bool IsUsing() { - return GetParam() == protocol_type; + auto param = GetParam(); + return param.protocol_type == protocol_type; + } + + bool IsProtobufContent() { + auto param = GetParam(); + return param.content_type == kContentEncodingProtoHeaderValue; } class PlainRequest { @@ -85,8 +94,15 @@ class GetValuesHandlerTest class BHTTPRequest { public: - explicit BHTTPRequest(PlainRequest plain_request) { + explicit BHTTPRequest(PlainRequest plain_request, + bool is_protobuf_content) { quiche::BinaryHttpRequest req_bhttp_layer({}); + if (is_protobuf_content) { + req_bhttp_layer.AddHeaderField({ + .name = std::string(kContentTypeHeader), + .value = std::string(kContentEncodingProtoHeaderValue), + }); + } req_bhttp_layer.set_body(plain_request.RequestBody()); auto maybe_serialized = req_bhttp_layer.Serialize(); EXPECT_TRUE(maybe_serialized.ok()); @@ -119,16 +135,32 @@ class GetValuesHandlerTest return maybe_res_bhttp_layer->status_code(); } - std::string Unwrap() const { + std::string Unwrap(bool is_protobuf_content) const { const absl::StatusOr maybe_res_bhttp_layer = quiche::BinaryHttpResponse::Create(response_.data()); EXPECT_TRUE(maybe_res_bhttp_layer.ok()) << "quiche::BinaryHttpResponse::Create failed: " << maybe_res_bhttp_layer.status(); + if (maybe_res_bhttp_layer->status_code() == 200 & is_protobuf_content) { + EXPECT_TRUE(HasHeader(*maybe_res_bhttp_layer, kContentTypeHeader, + kContentEncodingProtoHeaderValue)); + } return std::string(maybe_res_bhttp_layer->body()); } private: + bool HasHeader(const quiche::BinaryHttpResponse& response, + const std::string_view header_key, + const std::string_view header_value) const { + for (const auto& header : response.GetHeaderFields()) { + if (absl::AsciiStrToLower(header.name) == header_key && + absl::AsciiStrToLower(header.value) == header_value) { + return true; + } + } + return false; + } + google::api::HttpBody response_; }; @@ -176,7 +208,7 @@ class GetValuesHandlerTest // ../encryption/key_fetcher/src/fake_key_fetcher_manager.h uint8_t key_id = 64; auto maybe_config = quiche::ObliviousHttpHeaderKeyConfig::Create( - key_id, 0x0020, 0x0001, 0x0001); + key_id, kKEMParameter, kKDFParameter, kAEADParameter); EXPECT_TRUE(maybe_config.ok()); auto client = @@ -212,7 +244,7 @@ class GetValuesHandlerTest return handler->GetValuesHttp(plain_request.Build(), response); } - BHTTPRequest bhttp_request(std::move(plain_request)); + BHTTPRequest bhttp_request(std::move(plain_request), IsProtobufContent()); BHTTPResponse bresponse; if (IsUsing()) { @@ -237,20 +269,38 @@ class GetValuesHandlerTest *bhttp_response_code = bresponse.ResponseCode(); } - response->set_data(bresponse.Unwrap()); + response->set_data(bresponse.Unwrap(IsProtobufContent())); return grpc::Status::OK; } MockUdfClient mock_udf_client_; - MockMetricsRecorder mock_metrics_recorder_; privacy_sandbox::server_common::FakeKeyFetcherManager fake_key_fetcher_manager_; }; -INSTANTIATE_TEST_SUITE_P(GetValuesHandlerTest, GetValuesHandlerTest, - testing::Values(ProtocolType::kPlain, - ProtocolType::kBinaryHttp, - ProtocolType::kObliviousHttp)); +INSTANTIATE_TEST_SUITE_P( + GetValuesHandlerTest, GetValuesHandlerTest, + testing::Values( + TestingParameters{ + .protocol_type = ProtocolType::kPlain, + .content_type = kContentEncodingJsonHeaderValue, + }, + TestingParameters{ + .protocol_type = ProtocolType::kBinaryHttp, + .content_type = kContentEncodingJsonHeaderValue, + }, + TestingParameters{ + .protocol_type = ProtocolType::kObliviousHttp, + .content_type = kContentEncodingJsonHeaderValue, + }, + TestingParameters{ + .protocol_type = ProtocolType::kBinaryHttp, + .content_type = kContentEncodingProtoHeaderValue, + }, + TestingParameters{ + .protocol_type = ProtocolType::kObliviousHttp, + .content_type = kContentEncodingProtoHeaderValue, + })); TEST_P(GetValuesHandlerTest, Success) { UDFExecutionMetadata udf_metadata; @@ -326,11 +376,11 @@ data { )"); EXPECT_CALL( mock_udf_client_, - ExecuteCode(EqualsProto(udf_metadata), + ExecuteCode(_, EqualsProto(udf_metadata), testing::ElementsAre(EqualsProto(arg1), EqualsProto(arg2)))) .WillOnce(Return(output.dump())); - const std::string core_request_body = R"( + std::string core_request_body = R"( { "metadata": { "hostname": "example.com" @@ -365,9 +415,16 @@ data { )"; google::api::HttpBody response; - GetValuesV2Handler handler(mock_udf_client_, mock_metrics_recorder_, - fake_key_fetcher_manager_); + GetValuesV2Handler handler(mock_udf_client_, fake_key_fetcher_manager_); int16_t bhttp_response_code = 0; + if (IsProtobufContent()) { + v2::GetValuesRequest request_proto; + ASSERT_TRUE(google::protobuf::util::JsonStringToMessage(core_request_body, + &request_proto) + .ok()); + ASSERT_TRUE(request_proto.SerializeToString(&core_request_body)); + } + const auto result = GetValuesBasedOnProtocol(core_request_body, &response, &bhttp_response_code, &handler); ASSERT_EQ(bhttp_response_code, 200); @@ -377,24 +434,34 @@ data { v2::GetValuesResponse actual_response, expected_response; expected_response.mutable_single_partition()->set_string_output( output.dump()); - - ASSERT_TRUE(google::protobuf::util::JsonStringToMessage(response.data(), - &actual_response) - .ok()); + if (IsProtobufContent()) { + ASSERT_TRUE(actual_response.ParseFromString(response.data())); + } else { + ASSERT_TRUE(google::protobuf::util::JsonStringToMessage(response.data(), + &actual_response) + .ok()); + } EXPECT_THAT(actual_response, EqualsProto(expected_response)); } TEST_P(GetValuesHandlerTest, NoPartition) { - const std::string core_request_body = R"( + std::string core_request_body = R"( { "metadata": { "hostname": "example.com" } })"; google::api::HttpBody response; - GetValuesV2Handler handler(mock_udf_client_, mock_metrics_recorder_, - fake_key_fetcher_manager_); + GetValuesV2Handler handler(mock_udf_client_, fake_key_fetcher_manager_); int16_t bhttp_response_code = 0; + + if (IsProtobufContent()) { + v2::GetValuesRequest request_proto; + ASSERT_TRUE(google::protobuf::util::JsonStringToMessage(core_request_body, + &request_proto) + .ok()); + ASSERT_TRUE(request_proto.SerializeToString(&core_request_body)); + } const auto result = GetValuesBasedOnProtocol(core_request_body, &response, &bhttp_response_code, &handler); if (IsUsing()) { @@ -407,10 +474,10 @@ TEST_P(GetValuesHandlerTest, NoPartition) { } TEST_P(GetValuesHandlerTest, UdfFailureForOnePartition) { - EXPECT_CALL(mock_udf_client_, ExecuteCode(_, testing::IsEmpty())) + EXPECT_CALL(mock_udf_client_, ExecuteCode(_, _, testing::IsEmpty())) .WillOnce(Return(absl::InternalError("UDF execution error"))); - const std::string core_request_body = R"( + std::string core_request_body = R"( { "partitions": [ { @@ -421,9 +488,17 @@ TEST_P(GetValuesHandlerTest, UdfFailureForOnePartition) { )"; google::api::HttpBody response; - GetValuesV2Handler handler(mock_udf_client_, mock_metrics_recorder_, - fake_key_fetcher_manager_); + GetValuesV2Handler handler(mock_udf_client_, fake_key_fetcher_manager_); int16_t bhttp_response_code = 0; + + if (IsProtobufContent()) { + v2::GetValuesRequest request_proto; + ASSERT_TRUE(google::protobuf::util::JsonStringToMessage(core_request_body, + &request_proto) + .ok()); + ASSERT_TRUE(request_proto.SerializeToString(&core_request_body)); + } + const auto result = GetValuesBasedOnProtocol(core_request_body, &response, &bhttp_response_code, &handler); ASSERT_EQ(bhttp_response_code, 200); @@ -436,9 +511,13 @@ TEST_P(GetValuesHandlerTest, UdfFailureForOnePartition) { resp_status->set_code(13); resp_status->set_message("UDF execution error"); - ASSERT_TRUE(google::protobuf::util::JsonStringToMessage(response.data(), - &actual_response) - .ok()); + if (IsProtobufContent()) { + ASSERT_TRUE(actual_response.ParseFromString(response.data())); + } else { + ASSERT_TRUE(google::protobuf::util::JsonStringToMessage(response.data(), + &actual_response) + .ok()); + } EXPECT_THAT(actual_response, EqualsProto(expected_response)); } @@ -450,11 +529,11 @@ TEST_F(GetValuesHandlerTest, PureGRPCTest) { arguments { data { string_value: "ECHO" } } })pb", &req); - GetValuesV2Handler handler(mock_udf_client_, mock_metrics_recorder_, - fake_key_fetcher_manager_); + GetValuesV2Handler handler(mock_udf_client_, fake_key_fetcher_manager_); EXPECT_CALL(mock_udf_client_, - ExecuteCode(testing::_, testing::ElementsAre(EqualsProto( - req.partitions(0).arguments(0))))) + ExecuteCode(_, _, + testing::ElementsAre( + EqualsProto(req.partitions(0).arguments(0))))) .WillOnce(Return("ECHO")); v2::GetValuesResponse resp; const auto result = handler.GetValues(req, &resp); @@ -475,11 +554,11 @@ TEST_F(GetValuesHandlerTest, PureGRPCTestFailure) { arguments { data { string_value: "ECHO" } } })pb", &req); - GetValuesV2Handler handler(mock_udf_client_, mock_metrics_recorder_, - fake_key_fetcher_manager_); + GetValuesV2Handler handler(mock_udf_client_, fake_key_fetcher_manager_); EXPECT_CALL(mock_udf_client_, - ExecuteCode(testing::_, testing::ElementsAre(EqualsProto( - req.partitions(0).arguments(0))))) + ExecuteCode(_, _, + testing::ElementsAre( + EqualsProto(req.partitions(0).arguments(0))))) .WillOnce(Return(absl::InternalError("UDF execution error"))); v2::GetValuesResponse resp; const auto result = handler.GetValues(req, &resp); diff --git a/components/data_server/request_handler/ohttp_client_encryptor.cc b/components/data_server/request_handler/ohttp_client_encryptor.cc index cb42717d..81f33d0c 100644 --- a/components/data_server/request_handler/ohttp_client_encryptor.cc +++ b/components/data_server/request_handler/ohttp_client_encryptor.cc @@ -16,8 +16,8 @@ #include +#include "absl/log/log.h" #include "absl/strings/escaping.h" -#include "glog/logging.h" #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h" namespace kv_server { @@ -77,8 +77,8 @@ absl::StatusOr OhttpClientEncryptor::EncryptRequest( return serialized_encrypted_req; } -absl::StatusOr -OhttpClientEncryptor::DecryptResponse(std::string encrypted_payload) { +absl::StatusOr OhttpClientEncryptor::DecryptResponse( + std::string encrypted_payload) { if (!http_client_.has_value() || !http_request_context_.has_value()) { return absl::InternalError( "Emtpy `http_client_` or `http_request_context_`. You should call " @@ -89,6 +89,6 @@ OhttpClientEncryptor::DecryptResponse(std::string encrypted_payload) { if (!decrypted_response.ok()) { return decrypted_response.status(); } - return *decrypted_response; + return std::move(*decrypted_response).ConsumePlaintextData(); } } // namespace kv_server diff --git a/components/data_server/request_handler/ohttp_client_encryptor.h b/components/data_server/request_handler/ohttp_client_encryptor.h index 773fea32..1e023516 100644 --- a/components/data_server/request_handler/ohttp_client_encryptor.h +++ b/components/data_server/request_handler/ohttp_client_encryptor.h @@ -22,7 +22,7 @@ #include "absl/strings/escaping.h" #include "public/constants.h" #include "quiche/oblivious_http/oblivious_http_client.h" -#include "src/cpp/encryption/key_fetcher/src/key_fetcher_manager.h" +#include "src/encryption/key_fetcher/key_fetcher_manager.h" namespace kv_server { @@ -44,12 +44,7 @@ class OhttpClientEncryptor { absl::StatusOr EncryptRequest(std::string payload); // Decrypts incoming reponse. Since OHTTP is stateful, this method should be // called after EncryptRequest. - // In order to avoid an extra copy, leaking the `ObliviousHttpResponse`. - // Note that we have a CL for the underlying library that might allow us to - // not do leak this object and not do the copy. If/when that's merged, we - // should refactor this back to returning a string. - absl::StatusOr DecryptResponse( - std::string encrypted_payload); + absl::StatusOr DecryptResponse(std::string encrypted_payload); private: ::privacy_sandbox::server_common::CloudPlatform cloud_platform_ = diff --git a/components/data_server/request_handler/ohttp_encryptor_test.cc b/components/data_server/request_handler/ohttp_encryptor_test.cc index ce54f4f4..73cdd15f 100644 --- a/components/data_server/request_handler/ohttp_encryptor_test.cc +++ b/components/data_server/request_handler/ohttp_encryptor_test.cc @@ -17,8 +17,8 @@ #include "components/data_server/request_handler/ohttp_client_encryptor.h" #include "components/data_server/request_handler/ohttp_server_encryptor.h" #include "gtest/gtest.h" -#include "src/cpp/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" namespace kv_server { namespace { @@ -43,7 +43,7 @@ TEST(OhttpEncryptorTest, FullCircleSuccess) { auto response_decrypted_status = client_encryptor.DecryptResponse(*response_encrypted_status); ASSERT_TRUE(response_decrypted_status.ok()); - EXPECT_EQ(kTestResponse, response_decrypted_status->GetPlaintextData()); + EXPECT_EQ(kTestResponse, *response_decrypted_status); } TEST(OhttpEncryptorTest, ServerDecryptRequestFails) { diff --git a/components/data_server/request_handler/ohttp_server_encryptor.cc b/components/data_server/request_handler/ohttp_server_encryptor.cc index 689022f5..8280e845 100644 --- a/components/data_server/request_handler/ohttp_server_encryptor.cc +++ b/components/data_server/request_handler/ohttp_server_encryptor.cc @@ -16,7 +16,7 @@ #include -#include "glog/logging.h" +#include "absl/log/log.h" #include "quiche/oblivious_http/common/oblivious_http_header_key_config.h" namespace kv_server { diff --git a/components/data_server/request_handler/ohttp_server_encryptor.h b/components/data_server/request_handler/ohttp_server_encryptor.h index b35595c6..f888284f 100644 --- a/components/data_server/request_handler/ohttp_server_encryptor.h +++ b/components/data_server/request_handler/ohttp_server_encryptor.h @@ -22,7 +22,7 @@ #include "absl/strings/escaping.h" #include "public/constants.h" #include "quiche/oblivious_http/oblivious_http_gateway.h" -#include "src/cpp/encryption/key_fetcher/src/key_fetcher_manager.h" +#include "src/encryption/key_fetcher/key_fetcher_manager.h" namespace kv_server { diff --git a/components/data_server/request_handler/uncompressed.cc b/components/data_server/request_handler/uncompressed.cc index 7a190f00..cbbe32d5 100644 --- a/components/data_server/request_handler/uncompressed.cc +++ b/components/data_server/request_handler/uncompressed.cc @@ -15,7 +15,7 @@ #include -#include "glog/logging.h" +#include "absl/log/log.h" #include "quiche/common/quiche_data_writer.h" namespace kv_server { diff --git a/components/data_server/server/BUILD.bazel b/components/data_server/server/BUILD.bazel index 32b3ea06..c1cb1072 100644 --- a/components/data_server/server/BUILD.bazel +++ b/components/data_server/server/BUILD.bazel @@ -28,7 +28,6 @@ cc_library( "//components/data_server/request_handler:get_values_handler", "//public/query:get_values_cc_grpc", "@com_github_grpc_grpc//:grpc++", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -72,7 +71,6 @@ cc_test( "@com_google_absl//absl/flags:parse", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -84,8 +82,8 @@ cc_library( "//components/cloud_config:instance_client", "//components/data_server/server:parameter_fetcher", "//components/errors:retry", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", ], ) @@ -102,7 +100,6 @@ cc_test( "//components/cloud_config:parameter_client", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -119,7 +116,6 @@ cc_library( "//components/data_server/request_handler:get_values_v2_handler", "//public/query/v2:get_values_v2_cc_grpc", "@com_github_grpc_grpc//:grpc++", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -153,6 +149,7 @@ cc_library( "//components/internal_server:sharded_lookup", "//components/sharding:cluster_mappings_manager", "//components/telemetry:kv_telemetry", + "//components/telemetry:open_telemetry_sink", "//components/telemetry:server_definition", "//components/udf:udf_client", "//components/udf:udf_config_builder", @@ -168,17 +165,16 @@ cc_library( "//public/query:get_values_cc_grpc", "//public/sharding:key_sharder", "//public/udf:constants", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", "@com_github_grpc_grpc//:grpc++_reflection", # for grpc_cli "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", "@com_google_absl//absl/functional:bind_front", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:init", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry", + "@google_privacysandbox_servers_common//src/telemetry:init", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", ], ) @@ -198,13 +194,13 @@ cc_test( "@com_google_absl//absl/flags:parse", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) cc_binary( name = "server", srcs = ["main.cc"], + malloc = "@com_google_tcmalloc//tcmalloc", visibility = [ "//google_internal/production/packaging:__subpackages__", "//production/packaging:__subpackages__", @@ -214,12 +210,15 @@ cc_binary( ":server_lib", "//components/sharding:shard_manager", "//components/util:version_linkstamp", - "@com_github_google_glog//:glog", "@com_google_absl//absl/debugging:failure_signal_handler", "@com_google_absl//absl/debugging:symbolize", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/strings", + "@google_privacysandbox_servers_common//src/util:rlimit_core_config", ], ) @@ -258,33 +257,59 @@ cc_library( cc_library( name = "key_fetcher_factory", srcs = select({ - "//:aws_platform": [ + "//:aws_nonprod": [ + "key_fetcher_factory_cloud.cc", + "nonprod_key_fetcher_factory_aws.cc", + "nonprod_key_fetcher_factory_cloud.cc", + ], + "//:aws_prod": [ "key_fetcher_factory_aws.cc", "key_fetcher_factory_cloud.cc", ], - "//:gcp_platform": [ + "//:gcp_nonprod": [ + "key_fetcher_factory_cloud.cc", + "key_fetcher_utils_gcp.cc", + "nonprod_key_fetcher_factory_cloud.cc", + "nonprod_key_fetcher_factory_gcp.cc", + ], + "//:gcp_prod": [ "key_fetcher_factory_cloud.cc", "key_fetcher_factory_gcp.cc", + "key_fetcher_utils_gcp.cc", ], "//:local_platform": [ "key_fetcher_factory_local.cc", ], }), - hdrs = [ + hdrs = select({ + "//:gcp_platform": [ + "key_fetcher_utils_gcp.h", + ], + "//:nonprod_mode": [ + "nonprod_key_fetcher_factory_cloud.h", + ], + "//conditions:default": [], + }) + [ "key_fetcher_factory.h", ], - deps = [ - ":parameter_fetcher", - "@com_google_absl//absl/flags:flag", - "@com_google_absl//absl/flags:parse", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/time", - "@google_privacysandbox_servers_common//src/cpp/concurrent:executor", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:private_key_fetcher", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:public_key_fetcher", - ], + deps = + select({ + "//:gcp_platform": [ + ":key_fetcher_utils_gcp", + ], + "//conditions:default": [], + }) + [ + ":parameter_fetcher", + "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/time", + "@google_privacysandbox_servers_common//src/concurrent:executor", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:key_fetcher_manager", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:private_key_fetcher", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:public_key_fetcher", + ], ) cc_library( @@ -306,7 +331,16 @@ cc_library( "//components/udf/hooks:get_values_hook", "//components/udf/hooks:run_query_hook", "//public/sharding:key_sharder", - "@com_github_google_glog//:glog", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", + "@com_google_absl//absl/log", + ], +) + +cc_library( + name = "key_fetcher_utils_gcp", + srcs = ["key_fetcher_utils_gcp.cc"], + hdrs = ["key_fetcher_utils_gcp.h"], + deps = [ + ":parameter_fetcher", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:key_fetcher_manager", ], ) diff --git a/components/data_server/server/key_fetcher_factory.h b/components/data_server/server/key_fetcher_factory.h index be8ae41b..e5f61f73 100644 --- a/components/data_server/server/key_fetcher_factory.h +++ b/components/data_server/server/key_fetcher_factory.h @@ -15,12 +15,14 @@ */ #include +#include +#include #ifndef COMPONENTS_DATA_SERVER_SERVER_KEY_FETCHER_FACTORY_H_ #define COMPONENTS_DATA_SERVER_SERVER_KEY_FETCHER_FACTORY_H_ #include "components/data_server/server/parameter_fetcher.h" -#include "src/cpp/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" namespace kv_server { // Constructs KeyFetcherManager. @@ -56,6 +58,8 @@ class CloudKeyFetcherFactory : public KeyFetcherFactory { virtual google::scp::cpio::PrivateKeyVendingEndpoint GetSecondaryKeyFetchingEndpoint( const ParameterFetcher& parameter_fetcher) const; + virtual std::vector GetPublicKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const; virtual ::privacy_sandbox::server_common::CloudPlatform GetCloudPlatform() const; }; diff --git a/components/data_server/server/key_fetcher_factory_cloud.cc b/components/data_server/server/key_fetcher_factory_cloud.cc index d655904c..810b17ed 100644 --- a/components/data_server/server/key_fetcher_factory_cloud.cc +++ b/components/data_server/server/key_fetcher_factory_cloud.cc @@ -20,22 +20,14 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" #include "absl/flags/usage.h" +#include "absl/log/log.h" #include "components/data_server/server/key_fetcher_factory.h" -#include "glog/logging.h" +#include "src/concurrent/event_engine_executor.h" #include "src/core/lib/event_engine/default_event_engine.h" -#include "src/cpp/concurrent/event_engine_executor.h" -#include "src/cpp/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" ABSL_FLAG(std::string, public_key_endpoint, "", "Public key endpoint."); -ABSL_FLAG(std::string, primary_coordinator_private_key_endpoint, "", - "Primary coordinator private key endpoint."); -ABSL_FLAG(std::string, secondary_coordinator_private_key_endpoint, "", - "Secondary coordinator private key endpoint."); -ABSL_FLAG(std::string, primary_coordinator_region, "", - "Primary coordinator region."); -ABSL_FLAG(std::string, secondary_coordinator_region, "", - "Secondary coordinator region."); namespace kv_server { using ::google::scp::cpio::PrivateKeyVendingEndpoint; @@ -56,6 +48,17 @@ constexpr std::string_view kPrimaryCoordinatorAccountIdentityParameterSuffix = "primary-coordinator-account-identity"; constexpr std::string_view kSecondaryCoordinatorAccountIdentityParameterSuffix = "secondary-coordinator-account-identity"; +constexpr std::string_view + kPrimaryCoordinatorPrivateKeyEndpointParameterSuffix = + "primary-coordinator-private-key-endpoint"; +constexpr std::string_view kPrimaryCoordinatorRegionParameterSuffix = + "primary-coordinator-region"; +constexpr std::string_view + kSecondaryCoordinatoPrivateKeyEndpointParameterSuffix = + "secondary-coordinator-private-key-endpoint"; +constexpr std::string_view kSecondaryCoordinatorRegionParameterSuffix = + "secondary-coordinator-region"; + // Setting these to match // ..fledge/servers/bidding-auction-server/+/main:services/common/constants/common_service_flags.cc constexpr absl::Duration kPrivateKeyCacheTtl = absl::Hours(24 * 45); // 45 days @@ -63,17 +66,20 @@ constexpr absl::Duration kKeyRefreshFlowRunFrequency = absl::Hours(3); PrivateKeyVendingEndpoint GetKeyFetchingEndpoint( const ParameterFetcher& parameter_fetcher, - std::string_view account_identity_prefix, std::string_view service_endpoint, - absl::string_view region) { + std::string_view account_identity_prefix, + std::string_view private_key_endpoint_prefix, + absl::string_view region_prefix) { PrivateKeyVendingEndpoint endpoint; endpoint.account_identity = parameter_fetcher.GetParameter(account_identity_prefix); LOG(INFO) << "Retrieved " << account_identity_prefix << " parameter: " << endpoint.account_identity; - endpoint.private_key_vending_service_endpoint = service_endpoint; - LOG(INFO) << "Service endpoint: " << service_endpoint; - endpoint.service_region = region; - LOG(INFO) << "Region: " << region; + endpoint.private_key_vending_service_endpoint = + parameter_fetcher.GetParameter(private_key_endpoint_prefix); + LOG(INFO) << "Service endpoint: " + << endpoint.private_key_vending_service_endpoint; + endpoint.service_region = parameter_fetcher.GetParameter(region_prefix); + LOG(INFO) << "Region: " << endpoint.service_region; return endpoint; } } // namespace @@ -87,10 +93,8 @@ CloudKeyFetcherFactory::CreateKeyFetcherManager( "and private keys"; return std::make_unique(); } - auto publicKeyEndpointParameter = absl::GetFlag(FLAGS_public_key_endpoint); - LOG(INFO) << "Retrieved public_key_endpoint parameter: " - << publicKeyEndpointParameter; - std::vector endpoints = {publicKeyEndpointParameter}; + std::vector endpoints = + GetPublicKeyFetchingEndpoint(parameter_fetcher); std::unique_ptr public_key_fetcher = PublicKeyFetcherFactory::Create({{GetCloudPlatform(), endpoints}}); auto primary = GetPrimaryKeyFetchingEndpoint(parameter_fetcher); @@ -109,12 +113,21 @@ CloudKeyFetcherFactory::CreateKeyFetcherManager( return manager; } +std::vector CloudKeyFetcherFactory::GetPublicKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const { + auto publicKeyEndpointParameter = absl::GetFlag(FLAGS_public_key_endpoint); + LOG(INFO) << "Retrieved public_key_endpoint parameter: " + << publicKeyEndpointParameter; + std::vector endpoints = {publicKeyEndpointParameter}; + return endpoints; +} + PrivateKeyVendingEndpoint CloudKeyFetcherFactory::GetPrimaryKeyFetchingEndpoint( const ParameterFetcher& parameter_fetcher) const { return GetKeyFetchingEndpoint( parameter_fetcher, kPrimaryCoordinatorAccountIdentityParameterSuffix, - absl::GetFlag(FLAGS_primary_coordinator_private_key_endpoint), - absl::GetFlag(FLAGS_primary_coordinator_region)); + kPrimaryCoordinatorPrivateKeyEndpointParameterSuffix, + kPrimaryCoordinatorRegionParameterSuffix); } PrivateKeyVendingEndpoint @@ -122,8 +135,8 @@ CloudKeyFetcherFactory::GetSecondaryKeyFetchingEndpoint( const ParameterFetcher& parameter_fetcher) const { return GetKeyFetchingEndpoint( parameter_fetcher, kSecondaryCoordinatorAccountIdentityParameterSuffix, - absl::GetFlag(FLAGS_secondary_coordinator_private_key_endpoint), - absl::GetFlag(FLAGS_secondary_coordinator_region)); + kPrimaryCoordinatorAccountIdentityParameterSuffix, + kSecondaryCoordinatorRegionParameterSuffix); } CloudPlatform CloudKeyFetcherFactory::GetCloudPlatform() const { diff --git a/components/data_server/server/key_fetcher_factory_gcp.cc b/components/data_server/server/key_fetcher_factory_gcp.cc index 9e6a62df..d4fe2ada 100644 --- a/components/data_server/server/key_fetcher_factory_gcp.cc +++ b/components/data_server/server/key_fetcher_factory_gcp.cc @@ -12,45 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "absl/log/log.h" #include "components/data_server/server/key_fetcher_factory.h" -#include "glog/logging.h" +#include "components/data_server/server/key_fetcher_utils_gcp.h" namespace kv_server { namespace { using ::google::scp::cpio::PrivateKeyVendingEndpoint; using ::privacy_sandbox::server_common::CloudPlatform; -constexpr std::string_view kPrimaryKeyServiceCloudFunctionUrlSuffix = - "primary-key-service-cloud-function-url"; -constexpr std::string_view kPrimaryWorkloadIdentityPoolProviderSuffix = - "primary-workload-identity-pool-provider"; -constexpr std::string_view kSecondaryKeyServiceCloudFunctionUrlSuffix = - "secondary-key-service-cloud-function-url"; -constexpr std::string_view kSecondaryWorkloadIdentityPoolProviderSuffix = - "secondary-workload-identity-pool-provider"; - -void SetGcpSpecificParameters(PrivateKeyVendingEndpoint& endpoint, - const ParameterFetcher& parameter_fetcher, - const std::string_view cloudfunction_prefix, - const std::string_view wip_provider) { - endpoint.gcp_private_key_vending_service_cloudfunction_url = - parameter_fetcher.GetParameter(cloudfunction_prefix); - LOG(INFO) << "Retrieved " << cloudfunction_prefix << " parameter: " - << endpoint.gcp_private_key_vending_service_cloudfunction_url; - endpoint.gcp_wip_provider = parameter_fetcher.GetParameter(wip_provider); - LOG(INFO) << "Retrieved " << wip_provider - << " parameter: " << endpoint.gcp_wip_provider; -} - class KeyFetcherFactoryGcp : public CloudKeyFetcherFactory { PrivateKeyVendingEndpoint GetPrimaryKeyFetchingEndpoint( const ParameterFetcher& parameter_fetcher) const override { PrivateKeyVendingEndpoint endpoint = CloudKeyFetcherFactory::GetPrimaryKeyFetchingEndpoint( parameter_fetcher); - SetGcpSpecificParameters(endpoint, parameter_fetcher, - kPrimaryKeyServiceCloudFunctionUrlSuffix, - kPrimaryWorkloadIdentityPoolProviderSuffix); + UpdatePrimaryGcpEndpoint(endpoint, parameter_fetcher); return endpoint; } @@ -59,9 +36,7 @@ class KeyFetcherFactoryGcp : public CloudKeyFetcherFactory { PrivateKeyVendingEndpoint endpoint = CloudKeyFetcherFactory::GetSecondaryKeyFetchingEndpoint( parameter_fetcher); - SetGcpSpecificParameters(endpoint, parameter_fetcher, - kSecondaryKeyServiceCloudFunctionUrlSuffix, - kSecondaryWorkloadIdentityPoolProviderSuffix); + UpdateSecondaryGcpEndpoint(endpoint, parameter_fetcher); return endpoint; } diff --git a/components/data_server/server/key_fetcher_factory_local.cc b/components/data_server/server/key_fetcher_factory_local.cc index 5ae07dc3..9c1545b0 100644 --- a/components/data_server/server/key_fetcher_factory_local.cc +++ b/components/data_server/server/key_fetcher_factory_local.cc @@ -17,7 +17,7 @@ #include "absl/flags/flag.h" #include "components/data_server/server/key_fetcher_factory.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" ABSL_FLAG(std::string, public_key_endpoint, "", "Public key endpoint."); ABSL_FLAG(std::string, primary_coordinator_private_key_endpoint, "", diff --git a/components/data_server/server/key_fetcher_utils_gcp.cc b/components/data_server/server/key_fetcher_utils_gcp.cc new file mode 100644 index 00000000..8fedbfa8 --- /dev/null +++ b/components/data_server/server/key_fetcher_utils_gcp.cc @@ -0,0 +1,50 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "components/data_server/server/key_fetcher_utils_gcp.h" + +#include "absl/log/log.h" + +namespace kv_server { + +void SetGcpSpecificParameters(PrivateKeyVendingEndpoint& endpoint, + const ParameterFetcher& parameter_fetcher, + const std::string_view cloudfunction_prefix, + const std::string_view wip_provider) { + endpoint.gcp_private_key_vending_service_cloudfunction_url = + parameter_fetcher.GetParameter(cloudfunction_prefix); + LOG(INFO) << "Retrieved " << cloudfunction_prefix << " parameter: " + << endpoint.gcp_private_key_vending_service_cloudfunction_url; + endpoint.gcp_wip_provider = parameter_fetcher.GetParameter(wip_provider); + LOG(INFO) << "Retrieved " << wip_provider + << " parameter: " << endpoint.gcp_wip_provider; +} + +void UpdatePrimaryGcpEndpoint(PrivateKeyVendingEndpoint& endpoint, + const ParameterFetcher& parameter_fetcher) { + SetGcpSpecificParameters(endpoint, parameter_fetcher, + kPrimaryKeyServiceCloudFunctionUrlSuffix, + kPrimaryWorkloadIdentityPoolProviderSuffix); +} + +void UpdateSecondaryGcpEndpoint(PrivateKeyVendingEndpoint& endpoint, + const ParameterFetcher& parameter_fetcher) { + SetGcpSpecificParameters(endpoint, parameter_fetcher, + kSecondaryKeyServiceCloudFunctionUrlSuffix, + kSecondaryWorkloadIdentityPoolProviderSuffix); +} + +} // namespace kv_server diff --git a/components/data_server/server/key_fetcher_utils_gcp.h b/components/data_server/server/key_fetcher_utils_gcp.h new file mode 100644 index 00000000..1b508aae --- /dev/null +++ b/components/data_server/server/key_fetcher_utils_gcp.h @@ -0,0 +1,44 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef COMPONENTS_DATA_SERVER_KEY_FETCHER_UTILS_GCP_H_ +#define COMPONENTS_DATA_SERVER_KEY_FETCHER_UTILS_GCP_H_ + +#include "components/data_server/server/parameter_fetcher.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" + +namespace kv_server { + +using ::google::scp::cpio::PrivateKeyVendingEndpoint; + +constexpr std::string_view kPrimaryKeyServiceCloudFunctionUrlSuffix = + "primary-key-service-cloud-function-url"; +constexpr std::string_view kPrimaryWorkloadIdentityPoolProviderSuffix = + "primary-workload-identity-pool-provider"; +constexpr std::string_view kSecondaryKeyServiceCloudFunctionUrlSuffix = + "secondary-key-service-cloud-function-url"; +constexpr std::string_view kSecondaryWorkloadIdentityPoolProviderSuffix = + "secondary-workload-identity-pool-provider"; + +void UpdatePrimaryGcpEndpoint(PrivateKeyVendingEndpoint& endpoint, + const ParameterFetcher& parameter_fetcher); + +void UpdateSecondaryGcpEndpoint(PrivateKeyVendingEndpoint& endpoint, + const ParameterFetcher& parameter_fetcher); + +} // namespace kv_server + +#endif // COMPONENTS_DATA_SERVER_KEY_FETCHER_UTILS_GCP_H_ diff --git a/components/data_server/server/key_value_service_impl.cc b/components/data_server/server/key_value_service_impl.cc index f29b4963..35fc0e08 100644 --- a/components/data_server/server/key_value_service_impl.cc +++ b/components/data_server/server/key_value_service_impl.cc @@ -20,18 +20,12 @@ #include "components/data_server/request_handler/get_values_handler.h" #include "public/query/get_values.grpc.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry.h" - -constexpr char* kGetValuesSuccess = "GetValuesSuccess"; namespace kv_server { using google::protobuf::Struct; using google::protobuf::Value; using grpc::CallbackServerContext; -using privacy_sandbox::server_common::MetricsRecorder; -using privacy_sandbox::server_common::ScopeLatencyRecorder; using v1::GetValuesRequest; using v1::GetValuesResponse; using v1::KeyValueService; @@ -39,25 +33,13 @@ using v1::KeyValueService; grpc::ServerUnaryReactor* KeyValueServiceImpl::GetValues( CallbackServerContext* context, const GetValuesRequest* request, GetValuesResponse* response) { - ScopeLatencyRecorder latency_recorder(std::string(kGetValuesV1Latency), - metrics_recorder_); - - grpc::Status status = handler_.GetValues(*request, response); - - if (status.ok()) { - metrics_recorder_.IncrementEventStatus(kGetValuesSuccess, absl::OkStatus()); - } else { - // TODO: use implicit conversion when it becomes available externally - // https://g3doc.corp.google.com/net/grpc/g3doc/grpc_prod/cpp/status_mapping.md?cl=head - absl::StatusCode absl_status_code = - static_cast(status.error_code()); - absl::Status absl_status = - absl::Status(absl_status_code, status.error_message()); - metrics_recorder_.IncrementEventStatus(kGetValuesSuccess, absl_status); - } - + auto request_received_time = absl::Now(); + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); + grpc::Status status = handler_.GetValues(request_context, *request, response); auto* reactor = context->DefaultReactor(); reactor->Finish(status); + LogRequestCommonSafeMetrics(request, response, status, request_received_time); return reactor; } diff --git a/components/data_server/server/key_value_service_impl.h b/components/data_server/server/key_value_service_impl.h index de9d0bca..a57818d5 100644 --- a/components/data_server/server/key_value_service_impl.h +++ b/components/data_server/server/key_value_service_impl.h @@ -24,23 +24,15 @@ #include "components/data_server/request_handler/get_values_handler.h" #include "grpcpp/grpcpp.h" #include "public/query/get_values.grpc.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { -constexpr char* kGetValuesV1Latency = "GetValuesV1Latency"; - // Implements Key-Value service. class KeyValueServiceImpl final : public kv_server::v1::KeyValueService::CallbackService { public: - explicit KeyValueServiceImpl( - GetValuesHandler handler, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) - : handler_(std::move(handler)), metrics_recorder_(metrics_recorder) { - metrics_recorder_.RegisterHistogram( - kGetValuesV1Latency, "GetValues V1 service latency", "nanosecond"); - } + explicit KeyValueServiceImpl(GetValuesHandler handler) + : handler_(std::move(handler)) {} grpc::ServerUnaryReactor* GetValues( grpc::CallbackServerContext* context, @@ -49,7 +41,6 @@ class KeyValueServiceImpl final private: GetValuesHandler handler_; - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; }; } // namespace kv_server diff --git a/components/data_server/server/key_value_service_v2_impl.cc b/components/data_server/server/key_value_service_v2_impl.cc index 35934e31..36932962 100644 --- a/components/data_server/server/key_value_service_v2_impl.cc +++ b/components/data_server/server/key_value_service_v2_impl.cc @@ -17,8 +17,7 @@ #include #include "public/query/v2/get_values_v2.grpc.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry.h" +#include "src/telemetry/telemetry.h" namespace kv_server { namespace { @@ -37,10 +36,11 @@ grpc::ServerUnaryReactor* HandleRequest( CallbackServerContext* context, const RequestT* request, ResponseT* response, const GetValuesV2Handler& handler, HandlerFunctionT handler_function) { + auto request_received_time = absl::Now(); grpc::Status status = (handler.*handler_function)(*request, response); - auto* reactor = context->DefaultReactor(); reactor->Finish(status); + LogRequestCommonSafeMetrics(request, response, status, request_received_time); return reactor; } diff --git a/components/data_server/server/key_value_service_v2_impl.h b/components/data_server/server/key_value_service_v2_impl.h index 3f20ba53..c95d7b3b 100644 --- a/components/data_server/server/key_value_service_v2_impl.h +++ b/components/data_server/server/key_value_service_v2_impl.h @@ -24,7 +24,6 @@ #include "components/data_server/request_handler/get_values_v2_handler.h" #include "grpcpp/grpcpp.h" #include "public/query/v2/get_values_v2.grpc.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { @@ -32,10 +31,8 @@ namespace kv_server { class KeyValueServiceV2Impl final : public v2::KeyValueService::CallbackService { public: - explicit KeyValueServiceV2Impl( - GetValuesV2Handler handler, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) - : handler_(std::move(handler)), metrics_recorder_(metrics_recorder) {} + explicit KeyValueServiceV2Impl(GetValuesV2Handler handler) + : handler_(std::move(handler)) {} grpc::ServerUnaryReactor* GetValuesHttp( grpc::CallbackServerContext* context, @@ -58,7 +55,6 @@ class KeyValueServiceV2Impl final private: const GetValuesV2Handler handler_; - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; }; } // namespace kv_server diff --git a/components/data_server/server/lifecycle_heartbeat.cc b/components/data_server/server/lifecycle_heartbeat.cc index 52998d19..593faf80 100644 --- a/components/data_server/server/lifecycle_heartbeat.cc +++ b/components/data_server/server/lifecycle_heartbeat.cc @@ -17,8 +17,8 @@ #include #include +#include "absl/log/log.h" #include "components/errors/retry.h" -#include "glog/logging.h" namespace kv_server { diff --git a/components/data_server/server/lifecycle_heartbeat_test.cc b/components/data_server/server/lifecycle_heartbeat_test.cc index 2e4e3a76..936a1f7b 100644 --- a/components/data_server/server/lifecycle_heartbeat_test.cc +++ b/components/data_server/server/lifecycle_heartbeat_test.cc @@ -21,7 +21,6 @@ #include "components/data_server/server/mocks.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { diff --git a/components/data_server/server/main.cc b/components/data_server/server/main.cc index e260cc74..15caa8aa 100644 --- a/components/data_server/server/main.cc +++ b/components/data_server/server/main.cc @@ -17,10 +17,13 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" #include "absl/flags/usage.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "absl/strings/str_cat.h" #include "components/data_server/server/server.h" #include "components/util/build_info.h" -#include "glog/logging.h" +#include "src/util/rlimit_core_config.h" ABSL_FLAG(bool, buildinfo, false, "Print build info."); @@ -35,12 +38,15 @@ int main(int argc, char** argv) { // 3. Production versions of the K/V Server run inside Trusted Execution // Environments, which restrict where STDOUT and STDERR are visible to. absl::InitializeSymbolizer(argv[0]); + privacysandbox::server_common::SetRLimits({ + .enable_core_dumps = true, + }); { absl::FailureSignalHandlerOptions options; absl::InstallFailureSignalHandler(options); } - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); absl::SetProgramUsageMessage(absl::StrCat( "FLEDGE Key/Value Server. Sample usage:\n", argv[0], " --port=50051")); absl::ParseCommandLine(argc, argv); diff --git a/components/data_server/server/mocks.h b/components/data_server/server/mocks.h index 4ba36d19..8734fce2 100644 --- a/components/data_server/server/mocks.h +++ b/components/data_server/server/mocks.h @@ -25,7 +25,6 @@ #include "components/data_server/server/parameter_fetcher.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { diff --git a/components/data_server/server/nonprod_key_fetcher_factory_aws.cc b/components/data_server/server/nonprod_key_fetcher_factory_aws.cc new file mode 100644 index 00000000..d50eb289 --- /dev/null +++ b/components/data_server/server/nonprod_key_fetcher_factory_aws.cc @@ -0,0 +1,21 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "components/data_server/server/nonprod_key_fetcher_factory_cloud.h" + +namespace kv_server { +std::unique_ptr KeyFetcherFactory::Create() { + return std::make_unique(); +} +} // namespace kv_server diff --git a/components/data_server/server/nonprod_key_fetcher_factory_cloud.cc b/components/data_server/server/nonprod_key_fetcher_factory_cloud.cc new file mode 100644 index 00000000..a877ec57 --- /dev/null +++ b/components/data_server/server/nonprod_key_fetcher_factory_cloud.cc @@ -0,0 +1,48 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "components/data_server/server/nonprod_key_fetcher_factory_cloud.h" + +#include +#include +#include +#include + +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/flags/usage.h" +#include "absl/log/log.h" +#include "components/data_server/server/key_fetcher_factory.h" +#include "src/concurrent/event_engine_executor.h" +#include "src/core/lib/event_engine/default_event_engine.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" + +namespace kv_server { +using ::google::scp::cpio::PrivateKeyVendingEndpoint; + +constexpr std::string_view kPublicKeyEndpointParameterSuffix = + "public-key-endpoint"; + +std::vector +NonprodCloudKeyFetcherFactory::GetPublicKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const { + auto publicKeyEndpointParameter = + parameter_fetcher.GetParameter(kPublicKeyEndpointParameterSuffix); + LOG(INFO) << "Retrieved public_key_endpoint parameter: " + << publicKeyEndpointParameter; + std::vector endpoints = {publicKeyEndpointParameter}; + return endpoints; +} +} // namespace kv_server diff --git a/components/data_server/server/nonprod_key_fetcher_factory_cloud.h b/components/data_server/server/nonprod_key_fetcher_factory_cloud.h new file mode 100644 index 00000000..8f9bb4a5 --- /dev/null +++ b/components/data_server/server/nonprod_key_fetcher_factory_cloud.h @@ -0,0 +1,46 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +#ifndef COMPONENTS_DATA_SERVER_SERVER_NONPROD_KEY_FETCHER_FACTORY_CLOUD_H_ +#define COMPONENTS_DATA_SERVER_SERVER_NONPROD_KEY_FETCHER_FACTORY_CLOUD_H_ + +#include "components/data_server/server/key_fetcher_factory.h" +#include "components/data_server/server/parameter_fetcher.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" + +namespace kv_server { + +// Constructs KeyFetcherManager. This factory allows to override the public key +// endpoint. This is a security risk for produciton build. Which is why this +// implementation is only allowed in `nonprod` build mode. +class NonprodCloudKeyFetcherFactory : public CloudKeyFetcherFactory { + protected: + google::scp::cpio::PrivateKeyVendingEndpoint GetPrimaryKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const override; + google::scp::cpio::PrivateKeyVendingEndpoint GetSecondaryKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const override; + std::vector GetPublicKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const override; + ::privacy_sandbox::server_common::CloudPlatform GetCloudPlatform() + const override; +}; + +} // namespace kv_server +#endif // COMPONENTS_DATA_SERVER_SERVER_NONPROD_KEY_FETCHER_FACTORY_CLOUD_H_ diff --git a/components/data_server/server/nonprod_key_fetcher_factory_gcp.cc b/components/data_server/server/nonprod_key_fetcher_factory_gcp.cc new file mode 100644 index 00000000..47880038 --- /dev/null +++ b/components/data_server/server/nonprod_key_fetcher_factory_gcp.cc @@ -0,0 +1,53 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "absl/log/log.h" +#include "components/data_server/server/key_fetcher_factory.h" +#include "components/data_server/server/key_fetcher_utils_gcp.h" +#include "components/data_server/server/nonprod_key_fetcher_factory_cloud.h" + +namespace kv_server { +namespace { +using ::google::scp::cpio::PrivateKeyVendingEndpoint; +using ::privacy_sandbox::server_common::CloudPlatform; + +class KeyFetcherFactoryGcpNonProd : public NonprodCloudKeyFetcherFactory { + PrivateKeyVendingEndpoint GetPrimaryKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const override { + PrivateKeyVendingEndpoint endpoint = + NonprodCloudKeyFetcherFactory::GetPrimaryKeyFetchingEndpoint( + parameter_fetcher); + UpdatePrimaryGcpEndpoint(endpoint, parameter_fetcher); + return endpoint; + } + + PrivateKeyVendingEndpoint GetSecondaryKeyFetchingEndpoint( + const ParameterFetcher& parameter_fetcher) const override { + PrivateKeyVendingEndpoint endpoint = + NonprodCloudKeyFetcherFactory::GetSecondaryKeyFetchingEndpoint( + parameter_fetcher); + UpdateSecondaryGcpEndpoint(endpoint, parameter_fetcher); + return endpoint; + } + + CloudPlatform GetCloudPlatform() const override { + return CloudPlatform::kGcp; + } +}; +} // namespace + +std::unique_ptr KeyFetcherFactory::Create() { + return std::make_unique(); +} +} // namespace kv_server diff --git a/components/data_server/server/parameter_fetcher_aws.cc b/components/data_server/server/parameter_fetcher_aws.cc index b1cb3d9c..6d4f0e55 100644 --- a/components/data_server/server/parameter_fetcher_aws.cc +++ b/components/data_server/server/parameter_fetcher_aws.cc @@ -14,8 +14,8 @@ #include +#include "absl/log/log.h" #include "components/data_server/server/parameter_fetcher.h" -#include "glog/logging.h" namespace kv_server { diff --git a/components/data_server/server/parameter_fetcher_gcp.cc b/components/data_server/server/parameter_fetcher_gcp.cc index e3cd546b..5bcc9fbf 100644 --- a/components/data_server/server/parameter_fetcher_gcp.cc +++ b/components/data_server/server/parameter_fetcher_gcp.cc @@ -17,9 +17,9 @@ #include +#include "absl/log/log.h" #include "absl/strings/str_format.h" #include "components/data_server/server/parameter_fetcher.h" -#include "glog/logging.h" namespace kv_server { constexpr std::string_view kEnvironment = "environment"; diff --git a/components/data_server/server/parameter_fetcher_gcp_test.cc b/components/data_server/server/parameter_fetcher_gcp_test.cc index 824c3371..62418fa3 100644 --- a/components/data_server/server/parameter_fetcher_gcp_test.cc +++ b/components/data_server/server/parameter_fetcher_gcp_test.cc @@ -19,7 +19,6 @@ #include "components/data_server/server/parameter_fetcher.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { diff --git a/components/data_server/server/parameter_fetcher_local.cc b/components/data_server/server/parameter_fetcher_local.cc index ef47dc44..b22e1e1c 100644 --- a/components/data_server/server/parameter_fetcher_local.cc +++ b/components/data_server/server/parameter_fetcher_local.cc @@ -14,8 +14,8 @@ #include +#include "absl/log/log.h" #include "components/data_server/server/parameter_fetcher.h" -#include "glog/logging.h" namespace kv_server { diff --git a/components/data_server/server/parameter_fetcher_local_test.cc b/components/data_server/server/parameter_fetcher_local_test.cc index 1783162b..8a844bcf 100644 --- a/components/data_server/server/parameter_fetcher_local_test.cc +++ b/components/data_server/server/parameter_fetcher_local_test.cc @@ -21,7 +21,6 @@ #include "components/data_server/server/parameter_fetcher.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { diff --git a/components/data_server/server/server.cc b/components/data_server/server/server.cc index 53b7ea0b..2b88b059 100644 --- a/components/data_server/server/server.cc +++ b/components/data_server/server/server.cc @@ -20,8 +20,11 @@ #include "absl/flags/parse.h" #include "absl/flags/usage.h" #include "absl/functional/bind_front.h" +#include "absl/log/log.h" +#include "absl/log/log_sink_registry.h" #include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "components/data/blob_storage/blob_prefix_allowlist.h" #include "components/data_server/request_handler/get_values_adapter.h" #include "components/data_server/request_handler/get_values_handler.h" #include "components/data_server/request_handler/get_values_v2_handler.h" @@ -40,7 +43,7 @@ #include "components/udf/hooks/run_query_hook.h" #include "components/udf/udf_config_builder.h" #include "components/util/build_info.h" -#include "glog/logging.h" +#include "google/protobuf/text_format.h" #include "grpcpp/ext/proto_server_reflection_plugin.h" #include "grpcpp/health_check_service_interface.h" #include "public/constants.h" @@ -48,10 +51,10 @@ #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" #include "public/data_loading/readers/stream_record_reader_factory.h" #include "public/udf/constants.h" -#include "src/cpp/telemetry/init.h" -#include "src/cpp/telemetry/telemetry.h" -#include "src/cpp/telemetry/telemetry_provider.h" #include "src/google/protobuf/struct.pb.h" +#include "src/telemetry/init.h" +#include "src/telemetry/telemetry.h" +#include "src/telemetry/telemetry_provider.h" ABSL_FLAG(uint16_t, port, 50051, "Port the server is listening on. Defaults to 50051."); @@ -91,14 +94,22 @@ constexpr absl::string_view kLoggingVerbosityLevelParameterSuffix = "logging-verbosity-level"; constexpr absl::string_view kUdfTimeoutMillisParameterSuffix = "udf-timeout-millis"; +constexpr absl::string_view kUdfMinLogLevelParameterSuffix = + "udf-min-log-level"; constexpr absl::string_view kUseShardingKeyRegexParameterSuffix = "use-sharding-key-regex"; constexpr absl::string_view kShardingKeyRegexParameterSuffix = "sharding-key-regex"; constexpr absl::string_view kRouteV1ToV2Suffix = "route-v1-to-v2"; +constexpr absl::string_view kAddMissingKeysV1Suffix = "add-missing-keys-v1"; constexpr absl::string_view kAutoscalerHealthcheck = "autoscaler-healthcheck"; constexpr absl::string_view kLoadbalancerHealthcheck = "loadbalancer-healthcheck"; +constexpr absl::string_view kEnableOtelLoggerParameterSuffix = + "enable-otel-logger"; +constexpr std::string_view kDataLoadingBlobPrefixAllowlistSuffix = + "data-loading-blob-prefix-allowlist"; +constexpr std::string_view kTelemetryConfigSuffix = "telemetry-config"; opentelemetry::sdk::metrics::PeriodicExportingMetricReaderOptions GetMetricsOptions(const ParameterClient& parameter_client, @@ -124,19 +135,18 @@ GetMetricsOptions(const ParameterClient& parameter_client, return metrics_options; } -absl::Status CheckMetricsCollectorEndPointConnection( +void CheckMetricsCollectorEndPointConnection( std::string_view collector_endpoint) { auto channel = grpc::CreateChannel(std::string(collector_endpoint), grpc::InsecureChannelCredentials()); - // TODO(b/300137699): make the connection timeout a parameter - if (!channel->WaitForConnected( - gpr_time_add(gpr_now(GPR_CLOCK_REALTIME), - gpr_time_from_seconds(120, GPR_TIMESPAN)))) { - return absl::DeadlineExceededError( - "Timeout waiting for metrics collector connection"); - } - LOG(INFO) << "Metrics collector is connected"; - return absl::OkStatus(); + RetryUntilOk( + [channel]() { + if (channel->GetState(true) != GRPC_CHANNEL_READY) { + return absl::UnavailableError("metrics collector is not connected"); + } + return absl::OkStatus(); + }, + "Checking connection to metrics collector", LogMetricsNoOpCallback()); } absl::optional GetMetricsCollectorEndPoint( @@ -159,29 +169,41 @@ absl::optional GetMetricsCollectorEndPoint( privacy_sandbox::server_common::telemetry::TelemetryConfig GetServerTelemetryConfig(const ParameterClient& parameter_client, const std::string& environment) { - // TODO(b/304306398): Read telemetry config from parameter + ParameterFetcher parameter_fetcher(environment, parameter_client); + auto config_string = parameter_fetcher.GetParameter(kTelemetryConfigSuffix); privacy_sandbox::server_common::telemetry::TelemetryConfig config; - config.set_mode( - privacy_sandbox::server_common::telemetry::TelemetryConfig::EXPERIMENT); + if (!google::protobuf::TextFormat::ParseFromString(config_string, &config)) { + LOG(ERROR) << "Invalid proto format for telemetry config " << config_string + << ", fall back to prod config mode"; + config.set_mode( + privacy_sandbox::server_common::telemetry::TelemetryConfig::PROD); + } return config; } +BlobPrefixAllowlist GetBlobPrefixAllowlist( + const ParameterFetcher& parameter_fetcher) { + const auto prefix_allowlist = parameter_fetcher.GetParameter( + kDataLoadingBlobPrefixAllowlistSuffix, /*default_value=*/""); + LOG(INFO) << "Retrieved " << kDataLoadingBlobPrefixAllowlistSuffix + << " parameter: " << prefix_allowlist; + return BlobPrefixAllowlist(prefix_allowlist); +} + } // namespace Server::Server() - : metrics_recorder_( - TelemetryProvider::GetInstance().CreateMetricsRecorder()), - string_get_values_hook_( + : string_get_values_hook_( GetValuesHook::Create(GetValuesHook::OutputType::kString)), binary_get_values_hook_( GetValuesHook::Create(GetValuesHook::OutputType::kBinary)), run_query_hook_(RunQueryHook::Create()) {} -// Because the cache relies on metrics_recorder_, this function needs to be +// Because the cache relies on telemetry, this function needs to be // called right after telemetry has been initialized but before anything that // requires the cache has been initialized. void Server::InitializeKeyValueCache() { - cache_ = KeyValueCache::Create(*metrics_recorder_); + cache_ = KeyValueCache::Create(); cache_->UpdateKeyValue( "hi", "Hello, world! If you are seeing this, it means you can " @@ -189,6 +211,24 @@ void Server::InitializeKeyValueCache() { /*logical_commit_time = */ 1); } +void Server::InitOtelLogger( + ::opentelemetry::sdk::resource::Resource server_info, + absl::optional collector_endpoint, + const ParameterFetcher& parameter_fetcher) { + const bool enable_otel_logger = + parameter_fetcher.GetBoolParameter(kEnableOtelLoggerParameterSuffix); + LOG(INFO) << "Retrieved " << kEnableOtelLoggerParameterSuffix + << " parameter: " << enable_otel_logger; + if (!enable_otel_logger) { + return; + } + log_provider_ = privacy_sandbox::server_common::ConfigurePrivateLogger( + server_info, collector_endpoint); + open_telemetry_sink_ = std::make_unique( + log_provider_->GetLogger(kServiceName.data())); + absl::AddLogSink(open_telemetry_sink_.get()); +} + void Server::InitializeTelemetry(const ParameterClient& parameter_client, InstanceClient& instance_client) { std::string instance_id = RetryUntilOk( @@ -200,11 +240,7 @@ void Server::InitializeTelemetry(const ParameterClient& parameter_client, auto metrics_collector_endpoint = GetMetricsCollectorEndPoint(parameter_client, environment_); if (metrics_collector_endpoint.has_value()) { - if (const absl::Status status = CheckMetricsCollectorEndPointConnection( - metrics_collector_endpoint.value()); - !status.ok()) { - LOG(ERROR) << "Error in connecting metrics collector: " << status; - } + CheckMetricsCollectorEndPointConnection(metrics_collector_endpoint.value()); } LOG(INFO) << "Done retrieving metrics collector endpoint"; BuildDependentConfig telemetry_config( @@ -216,16 +252,26 @@ void Server::InitializeTelemetry(const ParameterClient& parameter_client, environment_), metrics_options, metrics_collector_endpoint)); AddSystemMetric(context_map); + + auto* internal_lookup_context_map = InternalLookupServerContextMap( + telemetry_config, + ConfigurePrivateMetrics( + CreateKVAttributes(instance_id, std::to_string(shard_num_), + environment_), + metrics_options, metrics_collector_endpoint)); + // TODO(b/300137699): Deprecate ConfigureMetrics once all metrics are migrated // to new telemetry API ConfigureMetrics( CreateKVAttributes(instance_id, std::to_string(shard_num_), environment_), metrics_options, metrics_collector_endpoint); - ConfigureTracer(CreateKVAttributes(std::move(instance_id), - std::to_string(shard_num_), environment_), - metrics_collector_endpoint); - - metrics_recorder_ = TelemetryProvider::GetInstance().CreateMetricsRecorder(); + ConfigureTracer( + CreateKVAttributes(instance_id, std::to_string(shard_num_), environment_), + metrics_collector_endpoint); + ParameterFetcher parameter_fetcher(environment_, parameter_client); + InitOtelLogger(CreateKVAttributes(std::move(instance_id), + std::to_string(shard_num_), environment_), + metrics_collector_endpoint, parameter_fetcher); LOG(INFO) << "Done init telemetry"; } @@ -247,13 +293,13 @@ absl::Status Server::CreateDefaultInstancesIfNecessaryAndGetEnvironment( parameter_fetcher.GetInt32Parameter(kUdfNumWorkersParameterSuffix); int32_t udf_timeout_ms = parameter_fetcher.GetInt32Parameter(kUdfTimeoutMillisParameterSuffix); + int32_t udf_min_log_level = + parameter_fetcher.GetInt32Parameter(kUdfMinLogLevelParameterSuffix); // updating verbosity level flag as early as we can, as it affects all logging // downstream. - // see - // https://github.com/google/glog/blob/931323df212c46e3a01b743d761c6ab8dc9f0d09/README.rst#setting-flags - FLAGS_v = parameter_fetcher.GetInt32Parameter( - kLoggingVerbosityLevelParameterSuffix); + absl::SetGlobalVLogLevel(parameter_fetcher.GetInt32Parameter( + kLoggingVerbosityLevelParameterSuffix)); if (udf_client != nullptr) { udf_client_ = std::move(udf_client); return absl::OkStatus(); @@ -267,10 +313,10 @@ absl::Status Server::CreateDefaultInstancesIfNecessaryAndGetEnvironment( .RegisterStringGetValuesHook(*string_get_values_hook_) .RegisterBinaryGetValuesHook(*binary_get_values_hook_) .RegisterRunQueryHook(*run_query_hook_) - .RegisterLoggingHook() + .RegisterLoggingFunction() .SetNumberOfWorkers(number_of_workers) .Config()), - absl::Milliseconds(udf_timeout_ms)); + absl::Milliseconds(udf_timeout_ms), udf_min_log_level); if (udf_client_or_status.ok()) { udf_client_ = std::move(*udf_client_or_status); } @@ -346,7 +392,11 @@ absl::Status Server::InitOnceInstancesAreCreated() { return status; } - SetDefaultUdfCodeObject(); + if (absl::Status status = SetDefaultUdfCodeObject(); !status.ok()) { + return absl::InternalError( + "Error setting default UDF. Please contact Google to fix the default " + "UDF or retry starting the server."); + } num_shards_ = parameter_fetcher.GetInt32Parameter(kNumShardsParameterSuffix); LOG(INFO) << "Retrieved " << kNumShardsParameterSuffix @@ -368,12 +418,11 @@ absl::Status Server::InitOnceInstancesAreCreated() { SetQueueManager(metadata, message_service_blob_.get()); grpc_server_ = CreateAndStartGrpcServer(); - local_lookup_ = CreateLocalLookup(*cache_, *metrics_recorder_); + local_lookup_ = CreateLocalLookup(*cache_); auto key_sharder = GetKeySharder(parameter_fetcher); auto server_initializer = GetServerInitializer( - num_shards_, *metrics_recorder_, *key_fetcher_manager_, *local_lookup_, - environment_, shard_num_, *instance_client_, *cache_, parameter_fetcher, - key_sharder); + num_shards_, *key_fetcher_manager_, *local_lookup_, environment_, + shard_num_, *instance_client_, *cache_, parameter_fetcher, key_sharder); remote_lookup_ = server_initializer->CreateAndStartRemoteLookupServer(); { auto status_or_notifier = @@ -564,6 +613,7 @@ std::unique_ptr Server::CreateDataOrchestrator( .shard_num = shard_num_, .num_shards = num_shards_, .key_sharder = std::move(key_sharder), + .blob_prefix_allowlist = GetBlobPrefixAllowlist(parameter_fetcher), }); }, "CreateDataOrchestrator", metrics_callback); @@ -571,18 +621,19 @@ std::unique_ptr Server::CreateDataOrchestrator( void Server::CreateGrpcServices(const ParameterFetcher& parameter_fetcher) { const bool use_v2 = parameter_fetcher.GetBoolParameter(kRouteV1ToV2Suffix); + const bool add_missing_keys_v1 = + parameter_fetcher.GetBoolParameter(kAddMissingKeysV1Suffix); LOG(INFO) << "Retrieved " << kRouteV1ToV2Suffix << " parameter: " << use_v2; get_values_adapter_ = GetValuesAdapter::Create(std::make_unique( - *udf_client_, *metrics_recorder_, *key_fetcher_manager_)); - GetValuesHandler handler(*cache_, *get_values_adapter_, *metrics_recorder_, - use_v2); - grpc_services_.push_back(std::make_unique( - std::move(handler), *metrics_recorder_)); - GetValuesV2Handler v2handler(*udf_client_, *metrics_recorder_, - *key_fetcher_manager_); - grpc_services_.push_back(std::make_unique( - std::move(v2handler), *metrics_recorder_)); + *udf_client_, *key_fetcher_manager_)); + GetValuesHandler handler(*cache_, *get_values_adapter_, use_v2, + add_missing_keys_v1); + grpc_services_.push_back( + std::make_unique(std::move(handler))); + GetValuesV2Handler v2handler(*udf_client_, *key_fetcher_manager_); + grpc_services_.push_back( + std::make_unique(std::move(v2handler))); } std::unique_ptr Server::CreateAndStartGrpcServer() { @@ -595,6 +646,13 @@ std::unique_ptr Server::CreateAndStartGrpcServer() { builder.AddListeningPort(server_address, grpc::InsecureServerCredentials()); // Register "service" as the instance through which we'll communicate with // clients. In this case it corresponds to a *synchronous* service. + + // Increase metadata size, this includes, for example the HTTP URL. Default is + // 8KB. + builder.AddChannelArgument(GRPC_ARG_ABSOLUTE_MAX_METADATA_SIZE, + 32 * 1024); // Set to 32KB + builder.AddChannelArgument(GRPC_ARG_MAX_METADATA_SIZE, + 32 * 1024); // Set to 32KB for (auto& service : grpc_services_) { builder.RegisterService(service.get()); } @@ -608,15 +666,13 @@ std::unique_ptr Server::CreateAndStartGrpcServer() { return std::move(server); } -void Server::SetDefaultUdfCodeObject() { +absl::Status Server::SetDefaultUdfCodeObject() { const absl::Status status = udf_client_->SetCodeObject( CodeConfig{.js = kDefaultUdfCodeSnippet, .udf_handler_name = kDefaultUdfHandlerName, .logical_commit_time = kDefaultLogicalCommitTime, .version = kDefaultVersion}); - if (!status.ok()) { - LOG(ERROR) << "Error setting code object: " << status; - } + return status; } std::unique_ptr Server::CreateDeltaFileNotifier( @@ -627,7 +683,8 @@ std::unique_ptr Server::CreateDeltaFileNotifier( << " parameter: " << backup_poll_frequency_secs; return DeltaFileNotifier::Create(*blob_client_, - absl::Seconds(backup_poll_frequency_secs)); + absl::Seconds(backup_poll_frequency_secs), + GetBlobPrefixAllowlist(parameter_fetcher)); } } // namespace kv_server diff --git a/components/data_server/server/server.h b/components/data_server/server/server.h index 47c4ed40..494a2dd4 100644 --- a/components/data_server/server/server.h +++ b/components/data_server/server/server.h @@ -37,6 +37,7 @@ #include "components/internal_server/lookup.h" #include "components/sharding/cluster_mappings_manager.h" #include "components/sharding/shard_manager.h" +#include "components/telemetry/open_telemetry_sink.h" #include "components/udf/hooks/get_values_hook.h" #include "components/udf/hooks/run_query_hook.h" #include "components/udf/udf_client.h" @@ -45,8 +46,7 @@ #include "public/base_types.pb.h" #include "public/query/get_values.grpc.pb.h" #include "public/sharding/key_sharder.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry.h" +#include "src/telemetry/telemetry.h" namespace kv_server { @@ -98,11 +98,14 @@ class Server { absl::Status InitializeUdfHooks(); std::unique_ptr CreateAndStartRemoteLookupServer(); - void SetDefaultUdfCodeObject(); + absl::Status SetDefaultUdfCodeObject(); void InitializeTelemetry(const ParameterClient& parameter_client, InstanceClient& instance_client); absl::Status CreateShardManager(); + void InitOtelLogger(::opentelemetry::sdk::resource::Resource server_info, + absl::optional collector_endpoint, + const ParameterFetcher& parameter_fetcher); // This must be first, otherwise the AWS SDK will crash when it's called: PlatformInitializer platform_initializer_; @@ -110,8 +113,6 @@ class Server { std::unique_ptr parameter_client_; std::unique_ptr instance_client_; std::string environment_; - std::unique_ptr - metrics_recorder_; std::vector> grpc_services_; std::unique_ptr grpc_server_; std::unique_ptr cache_; @@ -153,6 +154,8 @@ class Server { std::unique_ptr key_fetcher_manager_; + std::unique_ptr log_provider_; + std::unique_ptr open_telemetry_sink_; }; } // namespace kv_server diff --git a/components/data_server/server/server_initializer.cc b/components/data_server/server/server_initializer.cc index f3550dc6..4499017f 100644 --- a/components/data_server/server/server_initializer.cc +++ b/components/data_server/server/server_initializer.cc @@ -18,13 +18,12 @@ #include +#include "absl/log/log.h" #include "components/internal_server/constants.h" #include "components/internal_server/local_lookup.h" #include "components/internal_server/lookup_server_impl.h" -#include "components/internal_server/remote_lookup_client.h" #include "components/internal_server/sharded_lookup.h" -#include "glog/logging.h" -#include "src/cpp/encryption/key_fetcher/src/key_fetcher_manager.h" +#include "src/encryption/key_fetcher/key_fetcher_manager.h" namespace kv_server { namespace { @@ -45,9 +44,7 @@ absl::Status InitializeUdfHooksInternal( class NonshardedServerInitializer : public ServerInitializer { public: - explicit NonshardedServerInitializer(MetricsRecorder& metrics_recorder, - Cache& cache) - : metrics_recorder_(metrics_recorder), cache_(cache) {} + explicit NonshardedServerInitializer(Cache& cache) : cache_(cache) {} RemoteLookup CreateAndStartRemoteLookupServer() override { RemoteLookup remote_lookup; @@ -59,9 +56,8 @@ class NonshardedServerInitializer : public ServerInitializer { GetValuesHook& binary_get_values_hook, RunQueryHook& run_query_hook) override { ShardManagerState shard_manager_state; - auto lookup_supplier = [&cache = cache_, - &metrics_recorder = metrics_recorder_]() { - return CreateLocalLookup(cache, metrics_recorder); + auto lookup_supplier = [&cache = cache_]() { + return CreateLocalLookup(cache); }; InitializeUdfHooksInternal(std::move(lookup_supplier), string_get_values_hook, binary_get_values_hook, @@ -70,20 +66,17 @@ class NonshardedServerInitializer : public ServerInitializer { } private: - MetricsRecorder& metrics_recorder_; Cache& cache_; }; class ShardedServerInitializer : public ServerInitializer { public: explicit ShardedServerInitializer( - MetricsRecorder& metrics_recorder, KeyFetcherManagerInterface& key_fetcher_manager, Lookup& local_lookup, std::string environment, int32_t num_shards, int32_t current_shard_num, InstanceClient& instance_client, ParameterFetcher& parameter_fetcher, KeySharder key_sharder) - : metrics_recorder_(metrics_recorder), - key_fetcher_manager_(key_fetcher_manager), + : key_fetcher_manager_(key_fetcher_manager), local_lookup_(local_lookup), environment_(environment), num_shards_(num_shards), @@ -95,7 +88,7 @@ class ShardedServerInitializer : public ServerInitializer { RemoteLookup CreateAndStartRemoteLookupServer() override { RemoteLookup remote_lookup; remote_lookup.remote_lookup_service = std::make_unique( - local_lookup_, key_fetcher_manager_, metrics_recorder_); + local_lookup_, key_fetcher_manager_); grpc::ServerBuilder remote_lookup_server_builder; auto remoteLookupServerAddress = absl::StrCat(kLocalIp, ":", kRemoteLookupServerPort); @@ -122,10 +115,9 @@ class ShardedServerInitializer : public ServerInitializer { num_shards = num_shards_, current_shard_num = current_shard_num_, &shard_manager = *maybe_shard_state->shard_manager, - &metrics_recorder = metrics_recorder_, &key_sharder = key_sharder_]() { return CreateShardedLookup(local_lookup, num_shards, current_shard_num, - shard_manager, metrics_recorder, key_sharder); + shard_manager, key_sharder); }; InitializeUdfHooksInternal(std::move(lookup_supplier), string_get_values_hook, binary_get_values_hook, @@ -143,8 +135,8 @@ class ShardedServerInitializer : public ServerInitializer { shard_manager_state.shard_manager = TraceRetryUntilOk( [&cluster_mappings_manager = *shard_manager_state.cluster_mappings_manager, - &num_shards = num_shards_, &key_fetcher_manager = key_fetcher_manager_, - &metrics_recorder = metrics_recorder_] { + &num_shards = num_shards_, + &key_fetcher_manager = key_fetcher_manager_] { // It might be that the cluster mappings that are passed don't pass // validation. E.g. a particular cluster might not have any // replicas @@ -153,7 +145,7 @@ class ShardedServerInitializer : public ServerInitializer { // at that point in time might have new replicas spun up. return ShardManager::Create( num_shards, key_fetcher_manager, - cluster_mappings_manager.GetClusterMappings(), metrics_recorder); + cluster_mappings_manager.GetClusterMappings()); }, "GetShardManager", LogStatusSafeMetricsFn()); auto start_status = shard_manager_state.cluster_mappings_manager->Start( @@ -163,8 +155,6 @@ class ShardedServerInitializer : public ServerInitializer { } return std::move(shard_manager_state); } - - MetricsRecorder& metrics_recorder_; KeyFetcherManagerInterface& key_fetcher_manager_; Lookup& local_lookup_; std::string environment_; @@ -178,20 +168,18 @@ class ShardedServerInitializer : public ServerInitializer { } // namespace std::unique_ptr GetServerInitializer( - int64_t num_shards, MetricsRecorder& metrics_recorder, - KeyFetcherManagerInterface& key_fetcher_manager, Lookup& local_lookup, - std::string environment, int32_t current_shard_num, + int64_t num_shards, KeyFetcherManagerInterface& key_fetcher_manager, + Lookup& local_lookup, std::string environment, int32_t current_shard_num, InstanceClient& instance_client, Cache& cache, ParameterFetcher& parameter_fetcher, KeySharder key_sharder) { CHECK_GT(num_shards, 0) << "num_shards must be greater than 0"; if (num_shards == 1) { - return std::make_unique(metrics_recorder, - cache); + return std::make_unique(cache); } return std::make_unique( - metrics_recorder, key_fetcher_manager, local_lookup, environment, - num_shards, current_shard_num, instance_client, parameter_fetcher, + key_fetcher_manager, local_lookup, environment, num_shards, + current_shard_num, instance_client, parameter_fetcher, std::move(key_sharder)); } } // namespace kv_server diff --git a/components/data_server/server/server_initializer.h b/components/data_server/server/server_initializer.h index cae89c93..c490c377 100644 --- a/components/data_server/server/server_initializer.h +++ b/components/data_server/server/server_initializer.h @@ -29,8 +29,7 @@ #include "components/udf/hooks/run_query_hook.h" #include "grpcpp/grpcpp.h" #include "public/sharding/key_sharder.h" -#include "src/cpp/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" -#include "src/cpp/telemetry/metrics_recorder.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" namespace kv_server { @@ -59,7 +58,7 @@ class ServerInitializer { }; std::unique_ptr GetServerInitializer( - int64_t num_shards, MetricsRecorder& metrics_recorder, + int64_t num_shards, privacy_sandbox::server_common::KeyFetcherManagerInterface& key_fetcher_manager, Lookup& local_lookup, std::string environment, int32_t current_shard_num, diff --git a/components/data_server/server/server_local_test.cc b/components/data_server/server/server_local_test.cc index 35d00469..d6ae6cd8 100644 --- a/components/data_server/server/server_local_test.cc +++ b/components/data_server/server/server_local_test.cc @@ -53,6 +53,12 @@ void RegisterRequiredTelemetryExpectations(MockParameterClient& client) { EXPECT_CALL(client, GetInt32Parameter( "kv-server-environment-backup-poll-frequency-secs")) .WillOnce(::testing::Return(123)); + EXPECT_CALL(client, + GetBoolParameter("kv-server-environment-enable-otel-logger")) + .WillOnce(::testing::Return(false)); + EXPECT_CALL(client, GetParameter("kv-server-environment-telemetry-config", + testing::Eq(std::nullopt))) + .WillOnce(::testing::Return("mode: EXPERIMENT")); } void InitializeMetrics() { @@ -109,9 +115,15 @@ TEST(ServerLocalTest, InitFailsWithNoDeltaDirectory) { EXPECT_CALL(*parameter_client, GetInt32Parameter("kv-server-environment-udf-timeout-millis")) .WillOnce(::testing::Return(5000)); + EXPECT_CALL(*parameter_client, + GetInt32Parameter("kv-server-environment-udf-min-log-level")) + .WillOnce(::testing::Return(0)); EXPECT_CALL(*parameter_client, GetBoolParameter("kv-server-environment-route-v1-to-v2")) .WillOnce(::testing::Return(false)); + EXPECT_CALL(*parameter_client, + GetBoolParameter("kv-server-environment-add-missing-keys-v1")) + .WillOnce(::testing::Return(false)); EXPECT_CALL( *parameter_client, GetInt32Parameter("kv-server-environment-logging-verbosity-level")) @@ -119,6 +131,11 @@ TEST(ServerLocalTest, InitFailsWithNoDeltaDirectory) { EXPECT_CALL(*parameter_client, GetBoolParameter("kv-server-environment-use-sharding-key-regex")) .WillOnce(::testing::Return(false)); + EXPECT_CALL( + *parameter_client, + GetParameter("kv-server-environment-data-loading-blob-prefix-allowlist", + ::testing::Eq(""))) + .WillOnce(::testing::Return("")); kv_server::Server server; absl::Status status = server.Init(std::move(parameter_client), std::move(instance_client), @@ -167,9 +184,15 @@ TEST(ServerLocalTest, InitPassesWithDeltaDirectoryAndRealtimeDirectory) { EXPECT_CALL(*parameter_client, GetInt32Parameter("kv-server-environment-udf-timeout-millis")) .WillOnce(::testing::Return(5000)); + EXPECT_CALL(*parameter_client, + GetInt32Parameter("kv-server-environment-udf-min-log-level")) + .WillOnce(::testing::Return(0)); EXPECT_CALL(*parameter_client, GetBoolParameter("kv-server-environment-route-v1-to-v2")) .WillOnce(::testing::Return(false)); + EXPECT_CALL(*parameter_client, + GetBoolParameter("kv-server-environment-add-missing-keys-v1")) + .WillOnce(::testing::Return(false)); EXPECT_CALL( *parameter_client, GetInt32Parameter("kv-server-environment-logging-verbosity-level")) @@ -179,7 +202,12 @@ TEST(ServerLocalTest, InitPassesWithDeltaDirectoryAndRealtimeDirectory) { .WillOnce(::testing::Return(false)); EXPECT_CALL(*mock_udf_client, SetCodeObject(_)) .WillOnce(testing::Return(absl::OkStatus())); - + EXPECT_CALL( + *parameter_client, + GetParameter("kv-server-environment-data-loading-blob-prefix-allowlist", + ::testing::Eq(""))) + .Times(2) + .WillRepeatedly(::testing::Return("")); kv_server::Server server; absl::Status status = server.Init(std::move(parameter_client), std::move(instance_client), @@ -231,15 +259,26 @@ TEST(ServerLocalTest, GracefulServerShutdown) { EXPECT_CALL(*parameter_client, GetInt32Parameter("kv-server-environment-udf-timeout-millis")) .WillOnce(::testing::Return(5000)); + EXPECT_CALL(*parameter_client, + GetInt32Parameter("kv-server-environment-udf-min-log-level")) + .WillOnce(::testing::Return(0)); EXPECT_CALL(*parameter_client, GetBoolParameter("kv-server-environment-route-v1-to-v2")) .WillOnce(::testing::Return(false)); + EXPECT_CALL(*parameter_client, + GetBoolParameter("kv-server-environment-add-missing-keys-v1")) + .WillOnce(::testing::Return(false)); EXPECT_CALL(*parameter_client, GetBoolParameter("kv-server-environment-use-sharding-key-regex")) .WillOnce(::testing::Return(false)); EXPECT_CALL(*mock_udf_client, SetCodeObject(_)) .WillOnce(testing::Return(absl::OkStatus())); - + EXPECT_CALL( + *parameter_client, + GetParameter("kv-server-environment-data-loading-blob-prefix-allowlist", + ::testing::Eq(""))) + .Times(2) + .WillRepeatedly(::testing::Return("")); kv_server::Server server; absl::Status status = server.Init(std::move(parameter_client), std::move(instance_client), @@ -291,9 +330,15 @@ TEST(ServerLocalTest, ForceServerShutdown) { EXPECT_CALL(*parameter_client, GetInt32Parameter("kv-server-environment-udf-timeout-millis")) .WillOnce(::testing::Return(5000)); + EXPECT_CALL(*parameter_client, + GetInt32Parameter("kv-server-environment-udf-min-log-level")) + .WillOnce(::testing::Return(0)); EXPECT_CALL(*parameter_client, GetBoolParameter("kv-server-environment-route-v1-to-v2")) .WillOnce(::testing::Return(false)); + EXPECT_CALL(*parameter_client, + GetBoolParameter("kv-server-environment-add-missing-keys-v1")) + .WillOnce(::testing::Return(false)); EXPECT_CALL( *parameter_client, GetInt32Parameter("kv-server-environment-logging-verbosity-level")) @@ -301,10 +346,14 @@ TEST(ServerLocalTest, ForceServerShutdown) { EXPECT_CALL(*parameter_client, GetBoolParameter("kv-server-environment-use-sharding-key-regex")) .WillOnce(::testing::Return(false)); - EXPECT_CALL(*mock_udf_client, SetCodeObject(_)) .WillOnce(testing::Return(absl::OkStatus())); - + EXPECT_CALL( + *parameter_client, + GetParameter("kv-server-environment-data-loading-blob-prefix-allowlist", + ::testing::Eq(""))) + .Times(2) + .WillRepeatedly(::testing::Return("")); kv_server::Server server; absl::Status status = server.Init(std::move(parameter_client), std::move(instance_client), diff --git a/components/errors/BUILD.bazel b/components/errors/BUILD.bazel index de95448e..29aaccd8 100644 --- a/components/errors/BUILD.bazel +++ b/components/errors/BUILD.bazel @@ -17,6 +17,7 @@ load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") package(default_visibility = [ "//components:__subpackages__", "//production/packaging:__subpackages__", + "//public/data_loading:__subpackages__", ]) cc_library( @@ -86,10 +87,10 @@ cc_library( deps = [ "//components/telemetry:server_definition", "//components/util:sleepfor", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/time", - "@google_privacysandbox_servers_common//src/cpp/telemetry:tracing", + "@google_privacysandbox_servers_common//src/telemetry:tracing", ], ) @@ -103,6 +104,17 @@ cc_test( ":retry", "//components/util:sleepfor_mock", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + ], +) + +cc_library( + name = "error_tag", + srcs = [ + "error_tag.h", + ], + deps = [ + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:cord", ], ) diff --git a/components/errors/error_tag.h b/components/errors/error_tag.h new file mode 100644 index 00000000..0cca4a2e --- /dev/null +++ b/components/errors/error_tag.h @@ -0,0 +1,45 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef COMPONENTS_ERRORS_ERROR_TAG_H_ +#define COMPONENTS_ERRORS_ERROR_TAG_H_ + +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/strings/cord.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" + +namespace kv_server { + +// Sets the payload for an absl::Status +// The payload key is the file_name extracted from file_path +// The payload value is the error_tag_enum +template +inline absl::Status StatusWithErrorTag(absl::Status status, + std::string_view file_path, + T error_tag_enum) { + std::vector file_segments = absl::StrSplit(file_path, "/"); + status.SetPayload(file_segments.back(), + absl::Cord(absl::StrCat(error_tag_enum))); + return status; +} + +} // namespace kv_server +#endif // COMPONENTS_ERRORS_ERROR_TAG_H_ diff --git a/components/errors/retry.h b/components/errors/retry.h index 6879deb1..8f033af0 100644 --- a/components/errors/retry.h +++ b/components/errors/retry.h @@ -19,12 +19,12 @@ #include #include +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/time/time.h" #include "components/telemetry/server_definition.h" #include "components/util/sleepfor.h" -#include "glog/logging.h" -#include "src/cpp/telemetry/tracing.h" +#include "src/telemetry/tracing.h" namespace kv_server { @@ -91,7 +91,7 @@ class RetryableWithMax { // Retries functors that return an absl::StatusOr until they are `ok`. // The value of type T is returned by this function. -// `metrics_recorder` is optional. +// `metrics_callback` is optional. template typename std::invoke_result_t>::value_type RetryUntilOk( Func&& f, std::string task_name, @@ -105,7 +105,7 @@ typename std::invoke_result_t>::value_type RetryUntilOk( } // Same as above `RetryUntilOk`, wrapped in an `opentelemetry::trace::Span`. // Each individual retry of `func` is also traced. -// `metrics_recorder` is optional. +// `metrics_callback` is optional. template typename std::invoke_result_t>::value_type TraceRetryUntilOk( @@ -126,7 +126,7 @@ TraceRetryUntilOk( } // Retries functors that return an absl::Status until they are `ok`. -// `metrics_recorder` is optional. +// `metrics_callback` is optional. inline void RetryUntilOk( std::function func, std::string task_name, const absl::AnyInvocable& @@ -140,7 +140,7 @@ inline void RetryUntilOk( // Starts and `opentelemetry::trace::Span` and Calls `RetryUntilOk`. // Each individual retry of `func` is also traced. -// `metrics_recorder` is optional. +// `metrics_callback` is optional. void TraceRetryUntilOk(std::function func, std::string task_name, const absl::AnyInvocable func, // Retries functors that return an absl::StatusOr until they are `ok` or // max_attempts is reached. Retry starts at 1. -// `metrics_recorder` is optional. +// `metrics_callback` is optional. template typename std::invoke_result_t> RetryWithMax( Func&& f, std::string task_name, int max_attempts, diff --git a/components/errors/retry_test.cc b/components/errors/retry_test.cc index b9d2f363..d5349bdb 100644 --- a/components/errors/retry_test.cc +++ b/components/errors/retry_test.cc @@ -23,7 +23,6 @@ #include "components/util/sleepfor_mock.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/components/internal_server/BUILD.bazel b/components/internal_server/BUILD.bazel index 1e44b772..ae82b089 100644 --- a/components/internal_server/BUILD.bazel +++ b/components/internal_server/BUILD.bazel @@ -39,8 +39,7 @@ cc_library( "//components/query:scanner", "@com_github_grpc_grpc//:grpc++", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", + "@google_privacysandbox_servers_common//src/telemetry", ], ) @@ -60,8 +59,7 @@ cc_test( "//public/test_util:proto_matcher", "@com_github_grpc_grpc//:grpc++", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", ], ) @@ -93,6 +91,7 @@ proto_library( srcs = ["lookup.proto"], deps = [ "@com_google_googleapis//google/rpc:status_proto", + "@google_privacysandbox_servers_common//src/logger:logger_proto", ], ) @@ -126,6 +125,7 @@ cc_library( hdrs = ["lookup.h"], deps = [ ":internal_lookup_cc_proto", + "//components/util:request_context", "@com_google_absl//absl/status:statusor", ], ) @@ -141,9 +141,8 @@ cc_library( "//components/data_server/cache", "//components/query:driver", "//components/query:scanner", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -161,13 +160,11 @@ cc_library( "//components/query:scanner", "//components/sharding:shard_manager", "//public/sharding:key_sharder", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:statusor", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -185,8 +182,7 @@ cc_test( "//components/sharding:mocks", "//public/test_util:proto_matcher", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", ], ) @@ -203,13 +199,12 @@ cc_library( ":internal_lookup_cc_grpc", ":string_padder", "//components/data_server/request_handler:ohttp_client_encryptor", - "@com_github_google_glog//:glog", + "//components/util:request_context", "@com_github_grpc_grpc//:grpc++", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -222,8 +217,8 @@ cc_library( "string_padder.h", ], deps = [ - "@com_github_google_glog//:glog", "@com_github_google_quiche//quiche:quiche_unstable_api", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -256,8 +251,7 @@ cc_test( "//components/data_server/cache:mocks", "//public/test_util:proto_matcher", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", ], ) @@ -273,6 +267,5 @@ cc_test( "//public/test_util:proto_matcher", "@com_google_googletest//:gtest_main", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) diff --git a/components/internal_server/local_lookup.cc b/components/internal_server/local_lookup.cc index 872791b7..45522e49 100644 --- a/components/internal_server/local_lookup.cc +++ b/components/internal_server/local_lookup.cc @@ -20,51 +20,49 @@ #include #include +#include "absl/log/log.h" #include "components/data_server/cache/cache.h" #include "components/internal_server/lookup.h" #include "components/internal_server/lookup.pb.h" #include "components/query/driver.h" #include "components/query/scanner.h" -#include "glog/logging.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { namespace { -using privacy_sandbox::server_common::MetricsRecorder; -using privacy_sandbox::server_common::ScopeLatencyRecorder; - -constexpr char kKeySetNotFound[] = "KeysetNotFound"; -constexpr char kLocalRunQuery[] = "LocalRunQuery"; - class LocalLookup : public Lookup { public: - explicit LocalLookup(const Cache& cache, MetricsRecorder& metrics_recorder) - : cache_(cache), metrics_recorder_(metrics_recorder) {} + explicit LocalLookup(const Cache& cache) : cache_(cache) {} absl::StatusOr GetKeyValues( + const RequestContext& request_context, const absl::flat_hash_set& keys) const override { - return ProcessKeys(keys); + return ProcessKeys(request_context, keys); } absl::StatusOr GetKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const override { - return ProcessKeysetKeys(key_set); + return ProcessKeysetKeys(request_context, key_set); } absl::StatusOr RunQuery( - std::string query) const override { - return ProcessQuery(query); + const RequestContext& request_context, std::string query) const override { + return ProcessQuery(request_context, query); } private: InternalLookupResponse ProcessKeys( + const RequestContext& request_context, const absl::flat_hash_set& keys) const { + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetInternalLookupMetricsContext()); InternalLookupResponse response; if (keys.empty()) { return response; } - auto kv_pairs = cache_.GetKeyValuePairs(keys); + auto kv_pairs = cache_.GetKeyValuePairs(request_context, keys); for (const auto& key : keys) { SingleLookupResult result; @@ -72,7 +70,7 @@ class LocalLookup : public Lookup { if (key_iter == kv_pairs.end()) { auto status = result.mutable_status(); status->set_code(static_cast(absl::StatusCode::kNotFound)); - status->set_message("Key not found"); + status->set_message(absl::StrCat("Key not found: ", key)); } else { result.set_value(std::move(key_iter->second)); } @@ -82,20 +80,23 @@ class LocalLookup : public Lookup { } absl::StatusOr ProcessKeysetKeys( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const { + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetInternalLookupMetricsContext()); InternalLookupResponse response; if (key_set.empty()) { return response; } - auto key_value_set_result = cache_.GetKeyValueSet(key_set); + auto key_value_set_result = cache_.GetKeyValueSet(request_context, key_set); for (const auto& key : key_set) { SingleLookupResult result; const auto value_set = key_value_set_result->GetValueSet(key); if (value_set.empty()) { auto status = result.mutable_status(); status->set_code(static_cast(absl::StatusCode::kNotFound)); - status->set_message("Key not found"); - metrics_recorder_.IncrementEventCounter(kKeySetNotFound); + status->set_message(absl::StrCat("Key not found: ", key)); } else { auto keyset_values = result.mutable_keyset_values(); keyset_values->mutable_values()->Add(value_set.begin(), @@ -107,9 +108,10 @@ class LocalLookup : public Lookup { } absl::StatusOr ProcessQuery( - std::string query) const { - ScopeLatencyRecorder latency_recorder(std::string(kLocalRunQuery), - metrics_recorder_); + const RequestContext& request_context, std::string query) const { + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetInternalLookupMetricsContext()); if (query.empty()) return absl::OkStatus(); std::unique_ptr get_key_value_set_result; kv_server::Driver driver([&get_key_value_set_result](std::string_view key) { @@ -121,30 +123,32 @@ class LocalLookup : public Lookup { kv_server::Parser parse(driver, scanner); int parse_result = parse(); if (parse_result) { + LogInternalLookupRequestErrorMetric( + request_context.GetInternalLookupMetricsContext(), + kLocalRunQueryParsingFailure); return absl::InvalidArgumentError("Parsing failure."); } get_key_value_set_result = - cache_.GetKeyValueSet(driver.GetRootNode()->Keys()); + cache_.GetKeyValueSet(request_context, driver.GetRootNode()->Keys()); auto result = driver.GetResult(); if (!result.ok()) { + LogInternalLookupRequestErrorMetric( + request_context.GetInternalLookupMetricsContext(), + kLocalRunQueryFailure); return result.status(); } InternalRunQueryResponse response; response.mutable_elements()->Assign(result->begin(), result->end()); return response; } - const Cache& cache_; - MetricsRecorder& metrics_recorder_; }; } // namespace -std::unique_ptr CreateLocalLookup( - const Cache& cache, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) { - return std::make_unique(cache, metrics_recorder); +std::unique_ptr CreateLocalLookup(const Cache& cache) { + return std::make_unique(cache); } } // namespace kv_server diff --git a/components/internal_server/local_lookup.h b/components/internal_server/local_lookup.h index d2027f9a..2bc5138c 100644 --- a/components/internal_server/local_lookup.h +++ b/components/internal_server/local_lookup.h @@ -21,13 +21,10 @@ #include "components/data_server/cache/cache.h" #include "components/internal_server/lookup.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { -std::unique_ptr CreateLocalLookup( - const Cache& cache, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder); +std::unique_ptr CreateLocalLookup(const Cache& cache); } // namespace kv_server diff --git a/components/internal_server/local_lookup_test.cc b/components/internal_server/local_lookup_test.cc index ad78aad3..bddb646b 100644 --- a/components/internal_server/local_lookup_test.cc +++ b/components/internal_server/local_lookup_test.cc @@ -24,30 +24,37 @@ #include "google/protobuf/text_format.h" #include "gtest/gtest.h" #include "public/test_util/proto_matcher.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { using google::protobuf::TextFormat; -using privacy_sandbox::server_common::MockMetricsRecorder; using testing::_; using testing::Return; using testing::ReturnRef; class LocalLookupTest : public ::testing::Test { protected: + LocalLookupTest() { + InitMetricsContextMap(); + scope_metrics_context_ = std::make_unique(); + request_context_ = + std::make_unique(*scope_metrics_context_); + } + RequestContext& GetRequestContext() { return *request_context_; } + std::unique_ptr scope_metrics_context_; + std::unique_ptr request_context_; MockCache mock_cache_; - MockMetricsRecorder mock_metrics_recorder_; }; TEST_F(LocalLookupTest, GetKeyValues_KeysFound_Success) { - EXPECT_CALL(mock_cache_, GetKeyValuePairs(_)) + EXPECT_CALL(mock_cache_, GetKeyValuePairs(_, _)) .WillOnce(Return(absl::flat_hash_map{ {"key1", "value1"}, {"key2", "value2"}})); - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->GetKeyValues({"key1", "key2"}); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = + local_lookup->GetKeyValues(GetRequestContext(), {"key1", "key2"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -65,12 +72,13 @@ TEST_F(LocalLookupTest, GetKeyValues_KeysFound_Success) { } TEST_F(LocalLookupTest, GetKeyValues_DuplicateKeys_Success) { - EXPECT_CALL(mock_cache_, GetKeyValuePairs(_)) + EXPECT_CALL(mock_cache_, GetKeyValuePairs(_, _)) .WillOnce(Return(absl::flat_hash_map{ {"key1", "value1"}, {"key2", "value2"}})); - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->GetKeyValues({"key1", "key1"}); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = + local_lookup->GetKeyValues(GetRequestContext(), {"key1", "key1"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -84,12 +92,13 @@ TEST_F(LocalLookupTest, GetKeyValues_DuplicateKeys_Success) { } TEST_F(LocalLookupTest, GetKeyValues_KeyMissing_ReturnsStatusForKey) { - EXPECT_CALL(mock_cache_, GetKeyValuePairs(_)) + EXPECT_CALL(mock_cache_, GetKeyValuePairs(_, _)) .WillOnce(Return( absl::flat_hash_map{{"key1", "value1"}})); - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->GetKeyValues({"key1", "key2"}); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = + local_lookup->GetKeyValues(GetRequestContext(), {"key1", "key2"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -100,7 +109,7 @@ TEST_F(LocalLookupTest, GetKeyValues_KeyMissing_ReturnsStatusForKey) { } kv_pairs { key: "key2" - value { status { code: 5 message: "Key not found" } } + value { status { code: 5 message: "Key not found: key2" } } } )pb", &expected); @@ -108,8 +117,8 @@ TEST_F(LocalLookupTest, GetKeyValues_KeyMissing_ReturnsStatusForKey) { } TEST_F(LocalLookupTest, GetKeyValues_EmptyRequest_ReturnsEmptyResponse) { - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->GetKeyValues({}); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = local_lookup->GetKeyValues(GetRequestContext(), {}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -122,11 +131,11 @@ TEST_F(LocalLookupTest, GetKeyValueSets_KeysFound_Success) { EXPECT_CALL(*mock_get_key_value_set_result, GetValueSet("key1")) .WillOnce( Return(absl::flat_hash_set{"value1", "value2"})); - EXPECT_CALL(mock_cache_, GetKeyValueSet(_)) + EXPECT_CALL(mock_cache_, GetKeyValueSet(_, _)) .WillOnce(Return(std::move(mock_get_key_value_set_result))); - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->GetKeyValueSet({"key1"}); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = local_lookup->GetKeyValueSet(GetRequestContext(), {"key1"}); EXPECT_TRUE(response.ok()); std::vector expected_resulting_set = {"value1", "value2"}; @@ -139,18 +148,18 @@ TEST_F(LocalLookupTest, GetKeyValueSets_SetEmpty_Success) { std::make_unique(); EXPECT_CALL(*mock_get_key_value_set_result, GetValueSet("key1")) .WillOnce(Return(absl::flat_hash_set{})); - EXPECT_CALL(mock_cache_, GetKeyValueSet(_)) + EXPECT_CALL(mock_cache_, GetKeyValueSet(_, _)) .WillOnce(Return(std::move(mock_get_key_value_set_result))); - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->GetKeyValueSet({"key1"}); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = local_lookup->GetKeyValueSet(GetRequestContext(), {"key1"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; TextFormat::ParseFromString( R"pb(kv_pairs { key: "key1" - value { status { code: 5 message: "Key not found" } } + value { status { code: 5 message: "Key not found: key1" } } } )pb", &expected); @@ -158,8 +167,8 @@ TEST_F(LocalLookupTest, GetKeyValueSets_SetEmpty_Success) { } TEST_F(LocalLookupTest, GetKeyValueSet_EmptyRequest_ReturnsEmptyResponse) { - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->GetKeyValueSet({}); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = local_lookup->GetKeyValueSet(GetRequestContext(), {}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -174,12 +183,13 @@ TEST_F(LocalLookupTest, RunQuery_Success) { EXPECT_CALL(*mock_get_key_value_set_result, GetValueSet("someset")) .WillOnce( Return(absl::flat_hash_set{"value1", "value2"})); - EXPECT_CALL(mock_cache_, - GetKeyValueSet(absl::flat_hash_set{"someset"})) + EXPECT_CALL( + mock_cache_, + GetKeyValueSet(_, absl::flat_hash_set{"someset"})) .WillOnce(Return(std::move(mock_get_key_value_set_result))); - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->RunQuery(query); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = local_lookup->RunQuery(GetRequestContext(), query); EXPECT_TRUE(response.ok()); InternalRunQueryResponse expected; @@ -190,8 +200,8 @@ TEST_F(LocalLookupTest, RunQuery_Success) { TEST_F(LocalLookupTest, RunQuery_ParsingError_Error) { std::string query = "someset|("; - auto local_lookup = CreateLocalLookup(mock_cache_, mock_metrics_recorder_); - auto response = local_lookup->RunQuery(query); + auto local_lookup = CreateLocalLookup(mock_cache_); + auto response = local_lookup->RunQuery(GetRequestContext(), query); EXPECT_FALSE(response.ok()); EXPECT_EQ(response.status().code(), absl::StatusCode::kInvalidArgument); } diff --git a/components/internal_server/lookup.h b/components/internal_server/lookup.h index dde3917e..dd3356e7 100644 --- a/components/internal_server/lookup.h +++ b/components/internal_server/lookup.h @@ -24,6 +24,7 @@ #include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" #include "components/internal_server/lookup.pb.h" +#include "components/util/request_context.h" namespace kv_server { @@ -33,13 +34,15 @@ class Lookup { virtual ~Lookup() = default; virtual absl::StatusOr GetKeyValues( + const RequestContext& request_context, const absl::flat_hash_set& keys) const = 0; virtual absl::StatusOr GetKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const = 0; virtual absl::StatusOr RunQuery( - std::string query) const = 0; + const RequestContext& request_context, std::string query) const = 0; }; } // namespace kv_server diff --git a/components/internal_server/lookup.proto b/components/internal_server/lookup.proto index bc73389c..11d5c0ee 100644 --- a/components/internal_server/lookup.proto +++ b/components/internal_server/lookup.proto @@ -17,6 +17,7 @@ syntax = "proto3"; package kv_server; import "google/rpc/status.proto"; +import "src/logger/logger.proto"; // Internal Lookup Service API. service InternalLookupService { @@ -39,6 +40,10 @@ message InternalLookupRequest { // False means values are looked up. // True means value sets are looked up. bool lookup_sets = 2; + // Context useful for logging and tracing requests + privacy_sandbox.server_common.LogContext log_context = 3; + // Consented debugging configuration + privacy_sandbox.server_common.ConsentedDebugConfiguration consented_debug_config = 4; } // Encrypted and padded lookup request for internal datastore. @@ -92,6 +97,10 @@ message KeysetValues { message InternalRunQueryRequest { // Query to run. optional string query = 1; + // Context useful for logging and tracing requests + privacy_sandbox.server_common.LogContext log_context = 2; + // Consented debugging configuration + privacy_sandbox.server_common.ConsentedDebugConfiguration consented_debug_config = 3; } // Run Query response. diff --git a/components/internal_server/lookup_server_impl.cc b/components/internal_server/lookup_server_impl.cc index ba5bbaff..ed9b5715 100644 --- a/components/internal_server/lookup_server_impl.cc +++ b/components/internal_server/lookup_server_impl.cc @@ -21,50 +21,45 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "components/data_server/request_handler/ohttp_server_encryptor.h" #include "components/internal_server/lookup.grpc.pb.h" #include "components/internal_server/lookup.h" #include "components/internal_server/string_padder.h" -#include "glog/logging.h" #include "google/protobuf/message.h" #include "grpcpp/grpcpp.h" -#include "src/cpp/telemetry/telemetry.h" namespace kv_server { using google::protobuf::RepeatedPtrField; -using privacy_sandbox::server_common::ScopeLatencyRecorder; using grpc::StatusCode; -constexpr char kKeySetNotFound[] = "KeysetNotFound"; -constexpr char kDecryptionError[] = "DecryptionError"; -constexpr char kUnpaddingError[] = "UnpaddingError"; -constexpr char kEncryptionError[] = "EncryptionError"; -constexpr char kDeserializationError[] = "DeserializationError"; -constexpr char kRunQueryError[] = "RunQueryError"; -constexpr char kSecureLookup[] = "SecureLookup"; grpc::Status LookupServiceImpl::ToInternalGrpcStatus( - const absl::Status& status, const char* eventName) const { - metrics_recorder_.IncrementEventCounter(eventName); + const RequestContext& request_context, const absl::Status& status, + std::string_view error_code) const { + LogInternalLookupRequestErrorMetric( + request_context.GetInternalLookupMetricsContext(), error_code); return grpc::Status(StatusCode::INTERNAL, absl::StrCat(status.code(), " : ", status.message())); } -void LookupServiceImpl::ProcessKeys(const RepeatedPtrField& keys, +void LookupServiceImpl::ProcessKeys(const RequestContext& request_context, + const RepeatedPtrField& keys, InternalLookupResponse& response) const { if (keys.empty()) return; absl::flat_hash_set key_set; for (const auto& key : keys) { key_set.insert(key); } - auto lookup_result = lookup_.GetKeyValues(key_set); + auto lookup_result = lookup_.GetKeyValues(request_context, key_set); if (lookup_result.ok()) { response = *std::move(lookup_result); } } void LookupServiceImpl::ProcessKeysetKeys( + const RequestContext& request_context, const RepeatedPtrField& keys, InternalLookupResponse& response) const { if (keys.empty()) return; @@ -72,7 +67,7 @@ void LookupServiceImpl::ProcessKeysetKeys( for (const auto& key : keys) { key_list.insert(key); } - auto key_value_set_result = lookup_.GetKeyValueSet(key_list); + auto key_value_set_result = lookup_.GetKeyValueSet(request_context, key_list); if (key_value_set_result.ok()) { response = *std::move(key_value_set_result); } @@ -81,11 +76,13 @@ void LookupServiceImpl::ProcessKeysetKeys( grpc::Status LookupServiceImpl::InternalLookup( grpc::ServerContext* context, const InternalLookupRequest* request, InternalLookupResponse* response) { + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); if (context->IsCancelled()) { return grpc::Status(grpc::StatusCode::CANCELLED, "Deadline exceeded or client cancelled, abandoning."); } - ProcessKeys(request->keys(), *response); + ProcessKeys(request_context, request->keys(), *response); return grpc::Status::OK; } @@ -93,8 +90,13 @@ grpc::Status LookupServiceImpl::SecureLookup( grpc::ServerContext* context, const SecureLookupRequest* secure_lookup_request, SecureLookupResponse* secure_response) { - ScopeLatencyRecorder latency_recorder(std::string(kSecureLookup), - metrics_recorder_); + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); + LogIfError(request_context.GetInternalLookupMetricsContext() + .AccumulateMetric(1)); + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetInternalLookupMetricsContext()); if (context->IsCancelled()) { return grpc::Status(grpc::StatusCode::CANCELLED, "Deadline exceeded or client cancelled, abandoning."); @@ -105,17 +107,18 @@ grpc::Status LookupServiceImpl::SecureLookup( auto padded_serialized_request_maybe = encryptor.DecryptRequest(secure_lookup_request->ohttp_request()); if (!padded_serialized_request_maybe.ok()) { - return ToInternalGrpcStatus(padded_serialized_request_maybe.status(), - kDecryptionError); + return ToInternalGrpcStatus(request_context, + padded_serialized_request_maybe.status(), + kRequestDecryptionFailure); } VLOG(9) << "SecureLookup decrypted"; auto serialized_request_maybe = kv_server::Unpad(*padded_serialized_request_maybe); if (!serialized_request_maybe.ok()) { - metrics_recorder_.IncrementEventCounter(kDeserializationError); - return ToInternalGrpcStatus(serialized_request_maybe.status(), - kUnpaddingError); + return ToInternalGrpcStatus(request_context, + serialized_request_maybe.status(), + kRequestUnpaddingError); } VLOG(9) << "SecureLookup unpadded"; @@ -125,7 +128,8 @@ grpc::Status LookupServiceImpl::SecureLookup( "Failed parsing incoming request"); } - auto payload_to_encrypt = GetPayload(request.lookup_sets(), request.keys()); + auto payload_to_encrypt = + GetPayload(request_context, request.lookup_sets(), request.keys()); if (payload_to_encrypt.empty()) { // we cannot encrypt an empty payload. Note, that soon we will add logic // to pad responses, so this branch will never be hit. @@ -134,20 +138,22 @@ grpc::Status LookupServiceImpl::SecureLookup( auto encrypted_response_payload = encryptor.EncryptResponse(payload_to_encrypt); if (!encrypted_response_payload.ok()) { - return ToInternalGrpcStatus(encrypted_response_payload.status(), - kEncryptionError); + return ToInternalGrpcStatus(request_context, + encrypted_response_payload.status(), + kResponseEncryptionFailure); } secure_response->set_ohttp_response(*encrypted_response_payload); return grpc::Status::OK; } std::string LookupServiceImpl::GetPayload( - const bool lookup_sets, const RepeatedPtrField& keys) const { + const RequestContext& request_context, const bool lookup_sets, + const RepeatedPtrField& keys) const { InternalLookupResponse response; if (lookup_sets) { - ProcessKeysetKeys(keys, response); + ProcessKeysetKeys(request_context, keys, response); } else { - ProcessKeys(keys, response); + ProcessKeys(request_context, keys, response); } return response.SerializeAsString(); } @@ -155,13 +161,17 @@ std::string LookupServiceImpl::GetPayload( grpc::Status LookupServiceImpl::InternalRunQuery( grpc::ServerContext* context, const InternalRunQueryRequest* request, InternalRunQueryResponse* response) { + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); if (context->IsCancelled()) { return grpc::Status(grpc::StatusCode::CANCELLED, "Deadline exceeded or client cancelled, abandoning."); } - const auto process_result = lookup_.RunQuery(request->query()); + const auto process_result = + lookup_.RunQuery(request_context, request->query()); if (!process_result.ok()) { - return ToInternalGrpcStatus(process_result.status(), kRunQueryError); + return ToInternalGrpcStatus(request_context, process_result.status(), + kInternalRunQueryRequestFailure); } *response = *std::move(process_result); return grpc::Status::OK; diff --git a/components/internal_server/lookup_server_impl.h b/components/internal_server/lookup_server_impl.h index 4046a9e6..83fb9dfc 100644 --- a/components/internal_server/lookup_server_impl.h +++ b/components/internal_server/lookup_server_impl.h @@ -21,24 +21,20 @@ #include "components/internal_server/lookup.grpc.pb.h" #include "components/internal_server/lookup.h" +#include "components/util/request_context.h" #include "grpcpp/grpcpp.h" -#include "src/cpp/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" +#include "src/telemetry/telemetry.h" namespace kv_server { // Implements the internal lookup service for the data store. class LookupServiceImpl final : public kv_server::InternalLookupService::Service { public: - LookupServiceImpl( - const Lookup& lookup, - privacy_sandbox::server_common::KeyFetcherManagerInterface& - key_fetcher_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) - : lookup_(lookup), - key_fetcher_manager_(key_fetcher_manager), - metrics_recorder_(metrics_recorder) {} + LookupServiceImpl(const Lookup& lookup, + privacy_sandbox::server_common::KeyFetcherManagerInterface& + key_fetcher_manager) + : lookup_(lookup), key_fetcher_manager_(key_fetcher_manager) {} ~LookupServiceImpl() override = default; @@ -58,19 +54,21 @@ class LookupServiceImpl final private: std::string GetPayload( - const bool lookup_sets, + const RequestContext& request_context, const bool lookup_sets, const google::protobuf::RepeatedPtrField& keys) const; - void ProcessKeys(const google::protobuf::RepeatedPtrField& keys, + void ProcessKeys(const RequestContext& request_context, + const google::protobuf::RepeatedPtrField& keys, InternalLookupResponse& response) const; void ProcessKeysetKeys( + const RequestContext& request_context, const google::protobuf::RepeatedPtrField& keys, InternalLookupResponse& response) const; - grpc::Status ToInternalGrpcStatus(const absl::Status& status, - const char* eventName) const; + grpc::Status ToInternalGrpcStatus(const RequestContext& request_context, + const absl::Status& status, + std::string_view error_code) const; const Lookup& lookup_; privacy_sandbox::server_common::KeyFetcherManagerInterface& key_fetcher_manager_; - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; }; } // namespace kv_server diff --git a/components/internal_server/lookup_server_impl_test.cc b/components/internal_server/lookup_server_impl_test.cc index 47df74b1..be2118fd 100644 --- a/components/internal_server/lookup_server_impl_test.cc +++ b/components/internal_server/lookup_server_impl_test.cc @@ -29,14 +29,12 @@ #include "grpcpp/grpcpp.h" #include "gtest/gtest.h" #include "public/test_util/proto_matcher.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" namespace kv_server { namespace { using google::protobuf::TextFormat; -using privacy_sandbox::server_common::MockMetricsRecorder; using testing::_; using testing::Return; using testing::ReturnRef; @@ -45,15 +43,15 @@ class LookupServiceImplTest : public ::testing::Test { protected: LookupServiceImplTest() { lookup_service_ = std::make_unique( - mock_lookup_, fake_key_fetcher_manager_, mock_metrics_recorder_); + mock_lookup_, fake_key_fetcher_manager_); grpc::ServerBuilder builder; builder.RegisterService(lookup_service_.get()); server_ = (builder.BuildAndStart()); stub_ = InternalLookupService::NewStub( server_->InProcessChannel(grpc::ChannelArguments())); + InitMetricsContextMap(); } - ~LookupServiceImplTest() { server_->Shutdown(); server_->Wait(); @@ -64,7 +62,6 @@ class LookupServiceImplTest : public ::testing::Test { std::unique_ptr lookup_service_; std::unique_ptr server_; std::unique_ptr stub_; - MockMetricsRecorder mock_metrics_recorder_; }; TEST_F(LookupServiceImplTest, InternalLookup_Success) { @@ -82,7 +79,7 @@ TEST_F(LookupServiceImplTest, InternalLookup_Success) { } )pb", &expected); - EXPECT_CALL(mock_lookup_, GetKeyValues(_)).WillOnce(Return(expected)); + EXPECT_CALL(mock_lookup_, GetKeyValues(_, _)).WillOnce(Return(expected)); InternalLookupResponse response; grpc::ClientContext context; @@ -96,7 +93,7 @@ TEST_F(LookupServiceImplTest, InternalLookupRequest request; request.add_keys("key1"); request.add_keys("key2"); - EXPECT_CALL(mock_lookup_, GetKeyValues(_)) + EXPECT_CALL(mock_lookup_, GetKeyValues(_, _)) .WillOnce(Return(absl::UnknownError("Some error"))); InternalLookupResponse response; @@ -114,7 +111,7 @@ TEST_F(LookupServiceImplTest, InternalRunQuery_Success) { InternalRunQueryResponse expected; expected.add_elements("value1"); expected.add_elements("value2"); - EXPECT_CALL(mock_lookup_, RunQuery(_)).WillOnce(Return(expected)); + EXPECT_CALL(mock_lookup_, RunQuery(_, _)).WillOnce(Return(expected)); InternalRunQueryResponse response; grpc::ClientContext context; grpc::Status status = stub_->InternalRunQuery(&context, request, &response); @@ -126,7 +123,7 @@ TEST_F(LookupServiceImplTest, InternalRunQuery_Success) { TEST_F(LookupServiceImplTest, InternalRunQuery_LookupError_Failure) { InternalRunQueryRequest request; request.set_query("fail|||||now"); - EXPECT_CALL(mock_lookup_, RunQuery(_)) + EXPECT_CALL(mock_lookup_, RunQuery(_, _)) .WillOnce(Return(absl::UnknownError("Some error"))); InternalRunQueryResponse response; grpc::ClientContext context; diff --git a/components/internal_server/mocks.h b/components/internal_server/mocks.h index 058273d8..1d049b15 100644 --- a/components/internal_server/mocks.h +++ b/components/internal_server/mocks.h @@ -33,7 +33,8 @@ class MockRemoteLookupClient : public RemoteLookupClient { public: MockRemoteLookupClient() : RemoteLookupClient() {} MOCK_METHOD(absl::StatusOr, GetValues, - (std::string_view serialized_message, int32_t padding_length), + (const RequestContext& request_context, + std::string_view serialized_message, int32_t padding_length), (const, override)); MOCK_METHOD(std::string_view, GetIpAddress, (), (const, override)); }; @@ -41,13 +42,15 @@ class MockRemoteLookupClient : public RemoteLookupClient { class MockLookup : public Lookup { public: MOCK_METHOD(absl::StatusOr, GetKeyValues, - (const absl::flat_hash_set&), + (const RequestContext&, + const absl::flat_hash_set&), (const, override)); MOCK_METHOD(absl::StatusOr, GetKeyValueSet, - (const absl::flat_hash_set&), + (const RequestContext&, + const absl::flat_hash_set&), (const, override)); MOCK_METHOD(absl::StatusOr, RunQuery, - (std::string query), (const, override)); + (const RequestContext&, std::string query), (const, override)); }; } // namespace kv_server diff --git a/components/internal_server/remote_lookup_client.h b/components/internal_server/remote_lookup_client.h index d505c612..ece86fa8 100644 --- a/components/internal_server/remote_lookup_client.h +++ b/components/internal_server/remote_lookup_client.h @@ -23,9 +23,8 @@ #include "absl/status/statusor.h" #include "components/internal_server/lookup.grpc.pb.h" -#include "src/cpp/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry.h" +#include "components/util/request_context.h" +#include "src/encryption/key_fetcher/interface/key_fetcher_manager_interface.h" namespace kv_server { @@ -38,18 +37,17 @@ class RemoteLookupClient { // figure out the correct padding length across multiple requests. That helps // with preventing double serialization. virtual absl::StatusOr GetValues( + const RequestContext& request_context, std::string_view serialized_message, int32_t padding_length) const = 0; virtual std::string_view GetIpAddress() const = 0; static std::unique_ptr Create( std::string ip_address, privacy_sandbox::server_common::KeyFetcherManagerInterface& - key_fetcher_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder); + key_fetcher_manager); static std::unique_ptr Create( std::unique_ptr stub, privacy_sandbox::server_common::KeyFetcherManagerInterface& - key_fetcher_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder); + key_fetcher_manager); }; } // namespace kv_server diff --git a/components/internal_server/remote_lookup_client_impl.cc b/components/internal_server/remote_lookup_client_impl.cc index 95a21fcd..8640f39d 100644 --- a/components/internal_server/remote_lookup_client_impl.cc +++ b/components/internal_server/remote_lookup_client_impl.cc @@ -14,6 +14,7 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" @@ -22,19 +23,11 @@ #include "components/internal_server/lookup.grpc.pb.h" #include "components/internal_server/remote_lookup_client.h" #include "components/internal_server/string_padder.h" -#include "glog/logging.h" #include "grpcpp/grpcpp.h" namespace kv_server { namespace { -using privacy_sandbox::server_common::ScopeLatencyRecorder; - -constexpr char kEncryptionFailure[] = "EncryptionFailure"; -constexpr char kSecureLookupFailure[] = "SecureLookupFailure"; -constexpr char kDecryptionFailure[] = "DecryptionFailure"; -constexpr char kRemoteLookupGetValues[] = "RemoteLookupGetValues"; - class RemoteLookupClientImpl : public RemoteLookupClient { public: RemoteLookupClientImpl(const RemoteLookupClientImpl&) = delete; @@ -43,34 +36,32 @@ class RemoteLookupClientImpl : public RemoteLookupClient { explicit RemoteLookupClientImpl( std::string ip_address, privacy_sandbox::server_common::KeyFetcherManagerInterface& - key_fetcher_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) + key_fetcher_manager) : ip_address_( absl::StrFormat("%s:%s", ip_address, kRemoteLookupServerPort)), stub_(InternalLookupService::NewStub(grpc::CreateChannel( ip_address_, grpc::InsecureChannelCredentials()))), - key_fetcher_manager_(key_fetcher_manager), - metrics_recorder_(metrics_recorder) {} + key_fetcher_manager_(key_fetcher_manager) {} explicit RemoteLookupClientImpl( std::unique_ptr stub, privacy_sandbox::server_common::KeyFetcherManagerInterface& - key_fetcher_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) - : stub_(std::move(stub)), - key_fetcher_manager_(key_fetcher_manager), - metrics_recorder_(metrics_recorder) {} + key_fetcher_manager) + : stub_(std::move(stub)), key_fetcher_manager_(key_fetcher_manager) {} absl::StatusOr GetValues( + const RequestContext& request_context, std::string_view serialized_message, int32_t padding_length) const override { - ScopeLatencyRecorder latency_recorder(std::string(kRemoteLookupGetValues), - metrics_recorder_); + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetUdfRequestMetricsContext()); OhttpClientEncryptor encryptor(key_fetcher_manager_); auto encrypted_padded_serialized_request_maybe = encryptor.EncryptRequest(Pad(serialized_message, padding_length)); if (!encrypted_padded_serialized_request_maybe.ok()) { - metrics_recorder_.IncrementEventCounter(kEncryptionFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kRemoteRequestEncryptionFailure); return encrypted_padded_serialized_request_maybe.status(); } SecureLookupRequest secure_lookup_request; @@ -81,7 +72,8 @@ class RemoteLookupClientImpl : public RemoteLookupClient { grpc::Status status = stub_->SecureLookup(&context, secure_lookup_request, &secure_response); if (!status.ok()) { - metrics_recorder_.IncrementEventCounter(kSecureLookupFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kRemoteSecureLookupFailure); LOG(ERROR) << status.error_code() << ": " << status.error_message(); return absl::Status((absl::StatusCode)status.error_code(), status.error_message()); @@ -95,12 +87,11 @@ class RemoteLookupClientImpl : public RemoteLookupClient { auto decrypted_response_maybe = encryptor.DecryptResponse(std::move(secure_response.ohttp_response())); if (!decrypted_response_maybe.ok()) { - metrics_recorder_.IncrementEventCounter(kDecryptionFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kResponseEncryptionFailure); return decrypted_response_maybe.status(); } - - if (!response.ParseFromString( - decrypted_response_maybe->GetPlaintextData())) { + if (!response.ParseFromString(*decrypted_response_maybe)) { return absl::InvalidArgumentError("Failed parsing the response."); } return response; @@ -113,7 +104,6 @@ class RemoteLookupClientImpl : public RemoteLookupClient { std::unique_ptr stub_; privacy_sandbox::server_common::KeyFetcherManagerInterface& key_fetcher_manager_; - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; }; } // namespace @@ -121,18 +111,16 @@ class RemoteLookupClientImpl : public RemoteLookupClient { std::unique_ptr RemoteLookupClient::Create( std::string ip_address, privacy_sandbox::server_common::KeyFetcherManagerInterface& - key_fetcher_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) { - return std::make_unique( - std::move(ip_address), key_fetcher_manager, metrics_recorder); + key_fetcher_manager) { + return std::make_unique(std::move(ip_address), + key_fetcher_manager); } std::unique_ptr RemoteLookupClient::Create( std::unique_ptr stub, privacy_sandbox::server_common::KeyFetcherManagerInterface& - key_fetcher_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) { - return std::make_unique( - std::move(stub), key_fetcher_manager, metrics_recorder); + key_fetcher_manager) { + return std::make_unique(std::move(stub), + key_fetcher_manager); } } // namespace kv_server diff --git a/components/internal_server/remote_lookup_client_impl_test.cc b/components/internal_server/remote_lookup_client_impl_test.cc index e9093453..56f9b7fe 100644 --- a/components/internal_server/remote_lookup_client_impl_test.cc +++ b/components/internal_server/remote_lookup_client_impl_test.cc @@ -21,14 +21,12 @@ #include "grpcpp/grpcpp.h" #include "gtest/gtest.h" #include "public/test_util/proto_matcher.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" namespace kv_server { namespace { using google::protobuf::TextFormat; -using privacy_sandbox::server_common::MockMetricsRecorder; using testing::_; using testing::Return; @@ -36,27 +34,33 @@ class RemoteLookupClientImplTest : public ::testing::Test { protected: RemoteLookupClientImplTest() { lookup_service_ = std::make_unique( - mock_lookup_, fake_key_fetcher_manager_, mock_metrics_recorder_); + mock_lookup_, fake_key_fetcher_manager_); grpc::ServerBuilder builder; builder.RegisterService(lookup_service_.get()); server_ = (builder.BuildAndStart()); remote_lookup_client_ = RemoteLookupClient::Create( InternalLookupService::NewStub( server_->InProcessChannel(grpc::ChannelArguments())), - fake_key_fetcher_manager_, mock_metrics_recorder_); + fake_key_fetcher_manager_); + InitMetricsContextMap(); + scope_metrics_context_ = std::make_unique(); + request_context_ = + std::make_unique(*scope_metrics_context_); } ~RemoteLookupClientImplTest() { server_->Shutdown(); server_->Wait(); } + RequestContext& GetRequestContext() { return *request_context_; } MockLookup mock_lookup_; - MockMetricsRecorder mock_metrics_recorder_; privacy_sandbox::server_common::FakeKeyFetcherManager fake_key_fetcher_manager_; std::unique_ptr lookup_service_; std::unique_ptr server_; std::unique_ptr remote_lookup_client_; + std::unique_ptr scope_metrics_context_; + std::unique_ptr request_context_; }; TEST_F(RemoteLookupClientImplTest, EncryptedPaddedSuccessfulCall) { @@ -77,10 +81,10 @@ TEST_F(RemoteLookupClientImplTest, EncryptedPaddedSuccessfulCall) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_lookup_, GetKeyValues(_)) + EXPECT_CALL(mock_lookup_, GetKeyValues(_, _)) .WillOnce(Return(local_lookup_response)); - auto response_status = - remote_lookup_client_->GetValues(serialized_message, padding_length); + auto response_status = remote_lookup_client_->GetValues( + GetRequestContext(), serialized_message, padding_length); EXPECT_TRUE(response_status.ok()); InternalLookupResponse response = *response_status; InternalLookupResponse expected; @@ -104,8 +108,8 @@ TEST_F(RemoteLookupClientImplTest, EncryptedPaddedEmptySuccessfulCall) { request.set_lookup_sets(false); std::string serialized_message = request.SerializeAsString(); int32_t padding_length = 10; - auto response_status = - remote_lookup_client_->GetValues(serialized_message, padding_length); + auto response_status = remote_lookup_client_->GetValues( + GetRequestContext(), serialized_message, padding_length); EXPECT_TRUE(response_status.ok()); InternalLookupResponse response = *response_status; InternalLookupResponse expected; @@ -129,11 +133,11 @@ TEST_F(RemoteLookupClientImplTest, EncryptedPaddedSuccessfulKeysettLookup) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_lookup_, GetKeyValueSet(_)) + EXPECT_CALL(mock_lookup_, GetKeyValueSet(_, _)) .WillOnce(Return(local_lookup_response)); - auto response_status = - remote_lookup_client_->GetValues(serialized_message, padding_length); + auto response_status = remote_lookup_client_->GetValues( + GetRequestContext(), serialized_message, padding_length); EXPECT_TRUE(response_status.ok()); InternalLookupResponse response = *response_status; @@ -159,8 +163,8 @@ TEST_F(RemoteLookupClientImplTest, request.set_lookup_sets(true); std::string serialized_message = request.SerializeAsString(); int32_t padding_length = 10; - auto response_status = - remote_lookup_client_->GetValues(serialized_message, padding_length); + auto response_status = remote_lookup_client_->GetValues( + GetRequestContext(), serialized_message, padding_length); EXPECT_TRUE(response_status.ok()); InternalLookupResponse response = *response_status; diff --git a/components/internal_server/sharded_lookup.cc b/components/internal_server/sharded_lookup.cc index c2c7b00d..b102cdca 100644 --- a/components/internal_server/sharded_lookup.cc +++ b/components/internal_server/sharded_lookup.cc @@ -22,41 +22,20 @@ #include #include "absl/log/check.h" +#include "absl/log/log.h" #include "components/internal_server/lookup.h" #include "components/internal_server/lookup.pb.h" #include "components/internal_server/remote_lookup_client.h" #include "components/query/driver.h" #include "components/query/scanner.h" #include "components/sharding/shard_manager.h" -#include "glog/logging.h" +#include "components/util/request_context.h" #include "pir/hashing/sha256_hash_family.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { namespace { using google::protobuf::RepeatedPtrField; -using privacy_sandbox::server_common::MetricsRecorder; -using privacy_sandbox::server_common::ScopeLatencyRecorder; - -constexpr char kShardedLookupGrpcFailure[] = "ShardedLookupGrpcFailure"; -constexpr char kInternalRunQuery[] = "InternalRunQuery"; -constexpr char kInternalRunQueryQueryFailure[] = "InternalRunQueryQueryFailure"; -constexpr char kInternalRunQueryKeysetRetrievalFailure[] = - "InternalRunQueryKeysetRetrievalFailure"; -constexpr char kInternalRunQueryParsingFailure[] = - "InternalRunQueryParsingFailure"; -constexpr char kInternalRunQueryMissingKeyset[] = - "InternalRunQueryMissingKeyset"; -constexpr char kInternalRunQueryEmtpyQuery[] = "InternalRunQueryEmtpyQuery"; -constexpr char kKeySetNotFound[] = "KeysetNotFound"; -constexpr char kShardedLookupServerKeyCollisionOnCollection[] = - "ShardedLookupServerKeyCollisionOnCollection"; -constexpr char kLookupClientMissing[] = "LookupClientMissing"; -constexpr char kShardedLookupServerRequestFailed[] = - "ShardedLookupServerRequestFailed"; -constexpr char kLookupFuturesCreationFailure[] = "LookupFuturesCreationFailure"; -constexpr char kShardedLookupFailure[] = "ShardedLookupFailure"; void UpdateResponse( const std::vector& key_list, @@ -90,16 +69,14 @@ void SetRequestFailed(const std::vector& key_list, class ShardedLookup : public Lookup { public: - explicit ShardedLookup( - const Lookup& local_lookup, const int32_t num_shards, - const int32_t current_shard_num, const ShardManager& shard_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder, - KeySharder key_sharder) + explicit ShardedLookup(const Lookup& local_lookup, const int32_t num_shards, + const int32_t current_shard_num, + const ShardManager& shard_manager, + KeySharder key_sharder) : local_lookup_(local_lookup), num_shards_(num_shards), current_shard_num_(current_shard_num), shard_manager_(shard_manager), - metrics_recorder_(metrics_recorder), key_sharder_(std::move(key_sharder)) { CHECK_GT(num_shards, 1) << "num_shards for ShardedLookup must be > 1"; } @@ -113,21 +90,30 @@ class ShardedLookup : public Lookup { // return an empty response and `Internal` error as the status for the gRPC // status code. absl::StatusOr GetKeyValues( + const RequestContext& request_context, const absl::flat_hash_set& keys) const override { - return ProcessShardedKeys(keys); + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetUdfRequestMetricsContext()); + return ProcessShardedKeys(request_context, keys); } absl::StatusOr GetKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& keys) const override { + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetUdfRequestMetricsContext()); InternalLookupResponse response; if (keys.empty()) { return response; } absl::flat_hash_map> key_sets; - auto get_key_value_set_result_maybe = GetShardedKeyValueSet(keys); + auto get_key_value_set_result_maybe = + GetShardedKeyValueSet(request_context, keys); if (!get_key_value_set_result_maybe.ok()) { - metrics_recorder_.IncrementEventCounter( - kInternalRunQueryKeysetRetrievalFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedGetKeyValueSetKeySetRetrievalFailure); return get_key_value_set_result_maybe.status(); } key_sets = *std::move(get_key_value_set_result_maybe); @@ -138,7 +124,8 @@ class ShardedLookup : public Lookup { if (key_iter == key_sets.end()) { auto status = result.mutable_status(); status->set_code(static_cast(absl::StatusCode::kNotFound)); - metrics_recorder_.IncrementEventCounter(kKeySetNotFound); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedGetKeyValueSetKeySetNotFound); } else { auto keyset_values = result.mutable_keyset_values(); keyset_values->mutable_values()->Add(key_iter->second.begin(), @@ -150,23 +137,25 @@ class ShardedLookup : public Lookup { } absl::StatusOr RunQuery( - std::string query) const override { - ScopeLatencyRecorder latency_recorder(std::string(kInternalRunQuery), - metrics_recorder_); + const RequestContext& request_context, std::string query) const override { + ScopeLatencyMetricsRecorder + latency_recorder(request_context.GetUdfRequestMetricsContext()); InternalRunQueryResponse response; if (query.empty()) { - metrics_recorder_.IncrementEventCounter(kInternalRunQueryEmtpyQuery); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedRunQueryEmptyQuery); return response; } absl::flat_hash_map> keysets; - auto& metrics_recorder = metrics_recorder_; - kv_server::Driver driver([&keysets, - &metrics_recorder](std::string_view key) { + kv_server::Driver driver([&keysets, this, + &request_context](std::string_view key) { const auto key_iter = keysets.find(key); if (key_iter == keysets.end()) { VLOG(8) << "Driver can't find " << key << "key_set. Returning empty."; - metrics_recorder.IncrementEventCounter(kInternalRunQueryMissingKeyset); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedRunQueryMissingKeySet); absl::flat_hash_set set; return set; } else { @@ -180,20 +169,22 @@ class ShardedLookup : public Lookup { kv_server::Parser parse(driver, scanner); int parse_result = parse(); if (parse_result) { - metrics_recorder_.IncrementEventCounter(kInternalRunQueryParsingFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedRunQueryParsingFailure); return absl::InvalidArgumentError("Parsing failure."); } auto get_key_value_set_result_maybe = - GetShardedKeyValueSet(driver.GetRootNode()->Keys()); + GetShardedKeyValueSet(request_context, driver.GetRootNode()->Keys()); if (!get_key_value_set_result_maybe.ok()) { - metrics_recorder_.IncrementEventCounter( - kInternalRunQueryKeysetRetrievalFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedRunQueryKeySetRetrievalFailure); return get_key_value_set_result_maybe.status(); } keysets = std::move(*get_key_value_set_result_maybe); auto result = driver.GetResult(); if (!result.ok()) { - metrics_recorder_.IncrementEventCounter(kInternalRunQueryQueryFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedRunQueryFailure); return result.status(); } VLOG(8) << "Driver results for query " << query; @@ -267,13 +258,18 @@ class ShardedLookup : public Lookup { absl::StatusOr< std::vector>>> - GetLookupFutures(const std::vector& shard_lookup_inputs, + GetLookupFutures(const RequestContext& request_context, + const std::vector& shard_lookup_inputs, std::function( const std::vector& key_list)> get_local_future) const { std::vector>> responses; for (int shard_num = 0; shard_num < num_shards_; shard_num++) { auto& shard_lookup_input = shard_lookup_inputs[shard_num]; + LogIfError(request_context.GetUdfRequestMetricsContext() + .AccumulateMetric( + (int)shard_lookup_input.keys.size(), + std::to_string(shard_num))); if (shard_num == current_shard_num_) { // Eventually this whole branch will go away. responses.push_back(std::async(std::launch::async, get_local_future, @@ -281,13 +277,17 @@ class ShardedLookup : public Lookup { } else { const auto client = shard_manager_.Get(shard_num); if (client == nullptr) { - metrics_recorder_.IncrementEventCounter(kLookupClientMissing); + LogUdfRequestErrorMetric( + request_context.GetUdfRequestMetricsContext(), + kLookupClientMissing); return absl::InternalError("Internal lookup client is unavailable."); } responses.push_back(std::async( std::launch::async, - [client](std::string_view serialized_request, int32_t padding) { - return client->GetValues(serialized_request, padding); + [client, &request_context](std::string_view serialized_request, + int32_t padding) { + return client->GetValues(request_context, serialized_request, + padding); }, shard_lookup_input.serialized_request, shard_lookup_input.padding)); } @@ -296,14 +296,16 @@ class ShardedLookup : public Lookup { } absl::StatusOr GetLocalValues( + const RequestContext& request_context, const std::vector& key_list) const { InternalLookupResponse response; absl::flat_hash_set keys(key_list.begin(), key_list.end()); - return local_lookup_.GetKeyValues(keys); + return local_lookup_.GetKeyValues(request_context, keys); } absl::StatusOr GetLocalKeyValuesSet( + const RequestContext& request_context, const std::vector& key_list) const { if (key_list.empty()) { InternalLookupResponse response; @@ -317,10 +319,11 @@ class ShardedLookup : public Lookup { // a sepration between UDF and Data servers. absl::flat_hash_set key_list_set(key_list.begin(), key_list.end()); - return local_lookup_.GetKeyValueSet(key_list_set); + return local_lookup_.GetKeyValueSet(request_context, key_list_set); } absl::StatusOr ProcessShardedKeys( + const RequestContext& request_context, const absl::flat_hash_set& keys) const { InternalLookupResponse response; if (keys.empty()) { @@ -328,9 +331,10 @@ class ShardedLookup : public Lookup { } const auto shard_lookup_inputs = ShardKeys(keys, false); auto responses = - GetLookupFutures(shard_lookup_inputs, - [this](const std::vector& key_list) { - return GetLocalValues(key_list); + GetLookupFutures(request_context, shard_lookup_inputs, + [this, &request_context]( + const std::vector& key_list) { + return GetLocalValues(request_context, key_list); }); if (!responses.ok()) { return responses.status(); @@ -341,8 +345,8 @@ class ShardedLookup : public Lookup { auto result = (*responses)[shard_num].get(); if (!result.ok()) { // mark all keys as internal failure - metrics_recorder_.IncrementEventCounter( - kShardedLookupServerRequestFailed); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedKeyValueRequestFailure); SetRequestFailed(shard_lookup_input.keys, response); continue; } @@ -353,6 +357,7 @@ class ShardedLookup : public Lookup { } void CollectKeySets( + const RequestContext& request_context, absl::flat_hash_map>& key_sets, InternalLookupResponse& keysets_lookup_response) const { @@ -371,8 +376,9 @@ class ShardedLookup : public Lookup { auto [_, inserted] = key_sets.insert_or_assign(key, std::move(value_set)); if (!inserted) { - metrics_recorder_.IncrementEventCounter( - kShardedLookupServerKeyCollisionOnCollection); + LogUdfRequestErrorMetric( + request_context.GetUdfRequestMetricsContext(), + kShardedKeyCollisionOnKeySetCollection); LOG(ERROR) << "Key collision, when collecting results from shards: " << key; } @@ -384,15 +390,18 @@ class ShardedLookup : public Lookup { absl::StatusOr< absl::flat_hash_map>> GetShardedKeyValueSet( + const RequestContext& request_context, const absl::flat_hash_set& key_set) const { const auto shard_lookup_inputs = ShardKeys(key_set, true); - auto responses = - GetLookupFutures(shard_lookup_inputs, - [this](const std::vector& key_list) { - return GetLocalKeyValuesSet(key_list); - }); + auto responses = GetLookupFutures( + request_context, shard_lookup_inputs, + [this, + &request_context](const std::vector& key_list) { + return GetLocalKeyValuesSet(request_context, key_list); + }); if (!responses.ok()) { - metrics_recorder_.IncrementEventCounter(kLookupFuturesCreationFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kLookupClientMissing); return responses.status(); } // process responses @@ -401,10 +410,11 @@ class ShardedLookup : public Lookup { auto& shard_lookup_input = shard_lookup_inputs[shard_num]; auto result = (*responses)[shard_num].get(); if (!result.ok()) { - metrics_recorder_.IncrementEventCounter(kShardedLookupFailure); + LogUdfRequestErrorMetric(request_context.GetUdfRequestMetricsContext(), + kShardedKeyValueSetRequestFailure); return result.status(); } - CollectKeySets(key_sets, *result); + CollectKeySets(request_context, key_sets, *result); } return key_sets; } @@ -414,20 +424,19 @@ class ShardedLookup : public Lookup { const int32_t current_shard_num_; const std::string hashing_seed_; const ShardManager& shard_manager_; - MetricsRecorder& metrics_recorder_; KeySharder key_sharder_; }; } // namespace -std::unique_ptr CreateShardedLookup( - const Lookup& local_lookup, const int32_t num_shards, - const int32_t current_shard_num, const ShardManager& shard_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder, - KeySharder key_sharder) { - return std::make_unique( - local_lookup, num_shards, current_shard_num, shard_manager, - metrics_recorder, std::move(key_sharder)); +std::unique_ptr CreateShardedLookup(const Lookup& local_lookup, + const int32_t num_shards, + const int32_t current_shard_num, + const ShardManager& shard_manager, + KeySharder key_sharder) { + return std::make_unique(local_lookup, num_shards, + current_shard_num, shard_manager, + std::move(key_sharder)); } } // namespace kv_server diff --git a/components/internal_server/sharded_lookup.h b/components/internal_server/sharded_lookup.h index 207748b4..2619af0a 100644 --- a/components/internal_server/sharded_lookup.h +++ b/components/internal_server/sharded_lookup.h @@ -23,15 +23,14 @@ #include "components/internal_server/lookup.h" #include "components/sharding/shard_manager.h" #include "public/sharding/key_sharder.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { -std::unique_ptr CreateShardedLookup( - const Lookup& local_lookup, const int32_t num_shards, - const int32_t current_shard_num, const ShardManager& shard_manager, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder, - KeySharder key_sharder); +std::unique_ptr CreateShardedLookup(const Lookup& local_lookup, + const int32_t num_shards, + const int32_t current_shard_num, + const ShardManager& shard_manager, + KeySharder key_sharder); } // namespace kv_server diff --git a/components/internal_server/sharded_lookup_test.cc b/components/internal_server/sharded_lookup_test.cc index 7099c1e0..a1ddda53 100644 --- a/components/internal_server/sharded_lookup_test.cc +++ b/components/internal_server/sharded_lookup_test.cc @@ -26,24 +26,30 @@ #include "google/protobuf/text_format.h" #include "gtest/gtest.h" #include "public/test_util/proto_matcher.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { using google::protobuf::TextFormat; -using privacy_sandbox::server_common::MockMetricsRecorder; using testing::_; using testing::Return; using testing::ReturnRef; class ShardedLookupTest : public ::testing::Test { protected: + ShardedLookupTest() { + InitMetricsContextMap(); + scope_metrics_context_ = std::make_unique(); + request_context_ = + std::make_unique(*scope_metrics_context_); + } + RequestContext& GetRequestContext() { return *request_context_; } + std::unique_ptr scope_metrics_context_; + std::unique_ptr request_context_; int32_t num_shards_ = 2; int32_t shard_num_ = 0; MockLookup mock_local_lookup_; - MockMetricsRecorder mock_metrics_recorder_; KeySharder key_sharder_ = KeySharder(ShardingFunction{/*seed=*/""}); }; @@ -55,7 +61,7 @@ TEST_F(ShardedLookupTest, GetKeyValues_Success) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValues(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValues(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -77,7 +83,7 @@ TEST_F(ShardedLookupTest, GetKeyValues_Success) { key_list_remote.end()); const std::string serialized_request = request.SerializeAsString(); EXPECT_CALL(*mock_remote_lookup_client_1, - GetValues(serialized_request, 0)) + GetValues(_, serialized_request, 0)) .WillOnce([&]() { InternalLookupResponse resp; SingleLookupResult result; @@ -88,10 +94,11 @@ TEST_F(ShardedLookupTest, GetKeyValues_Success) { return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValues({"key1", "key4"}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = + sharded_lookup->GetKeyValues(GetRequestContext(), {"key1", "key4"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -117,7 +124,7 @@ TEST_F(ShardedLookupTest, GetKeyValues_KeyMissing_ReturnsStatus) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValues(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValues(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -139,8 +146,9 @@ TEST_F(ShardedLookupTest, GetKeyValues_KeyMissing_ReturnsStatus) { key_list_remote.end()); const std::string serialized_request = request.SerializeAsString(); - EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, 0)) - .WillOnce([=](const std::string_view serialized_message, + EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, _, 0)) + .WillOnce([=](const RequestContext& request_context, + const std::string_view serialized_message, const int32_t padding_length) { InternalLookupRequest request; EXPECT_TRUE(request.ParseFromString(serialized_message)); @@ -161,10 +169,11 @@ TEST_F(ShardedLookupTest, GetKeyValues_KeyMissing_ReturnsStatus) { return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValues({"key1", "key4", "key5"}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->GetKeyValues(GetRequestContext(), + {"key1", "key4", "key5"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -196,10 +205,10 @@ TEST_F(ShardedLookupTest, GetKeyValues_EmptyRequest_ReturnsEmptyResponse) { std::make_unique(), [](const std::string& ip) { return std::make_unique(); }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValues({}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->GetKeyValues(GetRequestContext(), {}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -215,7 +224,7 @@ TEST_F(ShardedLookupTest, GetKeyValues_FailedDownstreamRequest) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValues(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValues(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -236,16 +245,17 @@ TEST_F(ShardedLookupTest, GetKeyValues_FailedDownstreamRequest) { key_list_remote.end()); const std::string serialized_request = request.SerializeAsString(); EXPECT_CALL(*mock_remote_lookup_client_1, - GetValues(serialized_request, 0)) + GetValues(_, serialized_request, 0)) .WillOnce([]() { return absl::DeadlineExceededError("too long"); }); return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValues({"key1", "key4"}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = + sharded_lookup->GetKeyValues(GetRequestContext(), {"key1", "key4"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -294,8 +304,9 @@ TEST_F(ShardedLookupTest, GetKeyValues_ReturnsKeysFromCachePadding) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValues(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValues(_, _)) .WillOnce([&key_list, &local_lookup_response]( + const RequestContext& request_context, absl::flat_hash_set key_list_input) { EXPECT_THAT(key_list, testing::UnorderedElementsAreArray(key_list_input)); @@ -319,9 +330,9 @@ TEST_F(ShardedLookupTest, GetKeyValues_ReturnsKeysFromCachePadding) { request.mutable_keys()->Assign(key_list_remote.begin(), key_list_remote.end()); const std::string serialized_request = request.SerializeAsString(); - EXPECT_CALL(*mock_remote_lookup_client_1, - GetValues(testing::_, testing::_)) + EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, _, _)) .WillOnce([total_length, key_list_remote]( + const RequestContext& request_context, const std::string_view serialized_message, const int32_t padding_length) { EXPECT_EQ(total_length, @@ -356,8 +367,9 @@ TEST_F(ShardedLookupTest, GetKeyValues_ReturnsKeysFromCachePadding) { key_list_remote.end()); const std::string serialized_request = request.SerializeAsString(); EXPECT_CALL(*mock_remote_lookup_client_1, - GetValues(serialized_request, testing::_)) - .WillOnce([&](const std::string_view serialized_message, + GetValues(_, serialized_request, _)) + .WillOnce([&](const RequestContext& request_context, + const std::string_view serialized_message, const int32_t padding_length) { InternalLookupResponse resp; return resp; @@ -374,8 +386,9 @@ TEST_F(ShardedLookupTest, GetKeyValues_ReturnsKeysFromCachePadding) { request.mutable_keys()->Assign(key_list_remote.begin(), key_list_remote.end()); const std::string serialized_request = request.SerializeAsString(); - EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, testing::_)) - .WillOnce([=](const std::string_view serialized_message, + EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, _, _)) + .WillOnce([=](const RequestContext& request_context, + const std::string_view serialized_message, const int32_t padding_length) { InternalLookupRequest request; EXPECT_TRUE(request.ParseFromString(serialized_message)); @@ -402,10 +415,10 @@ TEST_F(ShardedLookupTest, GetKeyValues_ReturnsKeysFromCachePadding) { return std::make_unique(); }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValues(keys); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->GetKeyValues(GetRequestContext(), keys); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -455,7 +468,7 @@ TEST_F(ShardedLookupTest, GetKeyValueSets_KeysFound_Success) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -478,7 +491,7 @@ TEST_F(ShardedLookupTest, GetKeyValueSets_KeysFound_Success) { key_list_remote.end()); request.set_lookup_sets(true); const std::string serialized_request = request.SerializeAsString(); - EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, 0)) + EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, _, 0)) .WillOnce([&]() { InternalLookupResponse resp; TextFormat::ParseFromString( @@ -494,10 +507,11 @@ TEST_F(ShardedLookupTest, GetKeyValueSets_KeysFound_Success) { return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValueSet({"key1", "key4"}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = + sharded_lookup->GetKeyValueSet(GetRequestContext(), {"key1", "key4"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -524,7 +538,7 @@ TEST_F(ShardedLookupTest, GetKeyValueSets_KeysMissing_ReturnsStatus) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -547,8 +561,9 @@ TEST_F(ShardedLookupTest, GetKeyValueSets_KeysMissing_ReturnsStatus) { key_list_remote.end()); request.set_lookup_sets(true); const std::string serialized_request = request.SerializeAsString(); - EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, 0)) - .WillOnce([=](const std::string_view serialized_message, + EXPECT_CALL(*mock_remote_lookup_client_1, GetValues(_, _, 0)) + .WillOnce([=](const RequestContext& request_context, + const std::string_view serialized_message, const int32_t padding_length) { InternalLookupRequest request; EXPECT_TRUE(request.ParseFromString(serialized_message)); @@ -569,10 +584,11 @@ TEST_F(ShardedLookupTest, GetKeyValueSets_KeysMissing_ReturnsStatus) { return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValueSet({"key1", "key4", "key5"}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->GetKeyValueSet(GetRequestContext(), + {"key1", "key4", "key5"}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -604,10 +620,10 @@ TEST_F(ShardedLookupTest, GetKeyValueSet_EmptyRequest_ReturnsEmptyResponse) { std::make_unique(), [](const std::string& ip) { return std::make_unique(); }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValueSet({}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->GetKeyValueSet(GetRequestContext(), {}); EXPECT_TRUE(response.ok()); InternalLookupResponse expected; @@ -623,7 +639,7 @@ TEST_F(ShardedLookupTest, GetKeyValueSet_FailedDownstreamRequest) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -645,16 +661,17 @@ TEST_F(ShardedLookupTest, GetKeyValueSet_FailedDownstreamRequest) { request.set_lookup_sets(true); const std::string serialized_request = request.SerializeAsString(); EXPECT_CALL(*mock_remote_lookup_client_1, - GetValues(serialized_request, 0)) + GetValues(_, serialized_request, 0)) .WillOnce([]() { return absl::DeadlineExceededError("too long"); }); return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->GetKeyValueSet({"key1", "key4"}); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = + sharded_lookup->GetKeyValueSet(GetRequestContext(), {"key1", "key4"}); EXPECT_FALSE(response.ok()); EXPECT_EQ(response.status().code(), absl::StatusCode::kDeadlineExceeded); } @@ -668,7 +685,7 @@ TEST_F(ShardedLookupTest, RunQuery_Success) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -691,7 +708,7 @@ TEST_F(ShardedLookupTest, RunQuery_Success) { request.set_lookup_sets(true); const std::string serialized_request = request.SerializeAsString(); EXPECT_CALL(*mock_remote_lookup_client_1, - GetValues(serialized_request, 0)) + GetValues(_, serialized_request, 0)) .WillOnce([&]() { InternalLookupResponse resp; TextFormat::ParseFromString( @@ -707,10 +724,10 @@ TEST_F(ShardedLookupTest, RunQuery_Success) { return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->RunQuery("key1|key4"); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->RunQuery(GetRequestContext(), "key1|key4"); EXPECT_TRUE(response.ok()); EXPECT_THAT(response.value().elements(), @@ -726,7 +743,7 @@ TEST_F(ShardedLookupTest, RunQuery_MissingKeySet_IgnoresMissingSet_Success) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -749,7 +766,7 @@ TEST_F(ShardedLookupTest, RunQuery_MissingKeySet_IgnoresMissingSet_Success) { request.set_lookup_sets(true); const std::string serialized_request = request.SerializeAsString(); EXPECT_CALL(*mock_remote_lookup_client_1, - GetValues(serialized_request, 0)) + GetValues(_, serialized_request, 0)) .WillOnce([&]() { InternalLookupResponse resp; TextFormat::ParseFromString( @@ -765,10 +782,10 @@ TEST_F(ShardedLookupTest, RunQuery_MissingKeySet_IgnoresMissingSet_Success) { return mock_remote_lookup_client_1; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->RunQuery("key1|key4"); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->RunQuery(GetRequestContext(), "key1|key4"); EXPECT_TRUE(response.ok()); EXPECT_THAT(response.value().elements(), @@ -784,7 +801,7 @@ TEST_F(ShardedLookupTest, RunQuery_ShardedLookupFails_Error) { } )pb", &local_lookup_response); - EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_)) + EXPECT_CALL(mock_local_lookup_, GetKeyValueSet(_, _)) .WillOnce(Return(local_lookup_response)); std::vector> cluster_mappings; @@ -796,10 +813,10 @@ TEST_F(ShardedLookupTest, RunQuery_ShardedLookupFails_Error) { std::make_unique(), [](const std::string& ip) { return nullptr; }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->RunQuery("key1|key4"); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->RunQuery(GetRequestContext(), "key1|key4"); EXPECT_FALSE(response.ok()); EXPECT_THAT(response.status().code(), absl::StatusCode::kInternal); @@ -816,10 +833,10 @@ TEST_F(ShardedLookupTest, RunQuery_ParseError_ReturnStatus) { return std::make_unique(); }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->RunQuery("key1|"); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->RunQuery(GetRequestContext(), "key1|"); EXPECT_FALSE(response.ok()); EXPECT_EQ(response.status().code(), absl::StatusCode::kInvalidArgument); @@ -836,10 +853,10 @@ TEST_F(ShardedLookupTest, RunQuery_EmptyRequest_EmptyResponse) { return std::make_unique(); }); - auto sharded_lookup = CreateShardedLookup( - mock_local_lookup_, num_shards_, shard_num_, *(*shard_manager), - mock_metrics_recorder_, key_sharder_); - auto response = sharded_lookup->RunQuery(""); + auto sharded_lookup = + CreateShardedLookup(mock_local_lookup_, num_shards_, shard_num_, + *(*shard_manager), key_sharder_); + auto response = sharded_lookup->RunQuery(GetRequestContext(), ""); EXPECT_TRUE(response.ok()); EXPECT_TRUE(response.value().elements().empty()); } diff --git a/components/internal_server/string_padder.cc b/components/internal_server/string_padder.cc index 663b2a4d..88ddcf56 100644 --- a/components/internal_server/string_padder.cc +++ b/components/internal_server/string_padder.cc @@ -15,7 +15,7 @@ #include -#include "glog/logging.h" +#include "absl/log/log.h" #include "quiche/common/quiche_data_reader.h" #include "quiche/common/quiche_data_writer.h" diff --git a/components/sharding/BUILD.bazel b/components/sharding/BUILD.bazel index 901c2024..940030e0 100644 --- a/components/sharding/BUILD.bazel +++ b/components/sharding/BUILD.bazel @@ -33,8 +33,6 @@ cc_library( "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", ], ) @@ -50,8 +48,7 @@ cc_test( ":shard_manager", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", ], ) @@ -99,7 +96,6 @@ cc_test( "//components/data_server/server:mocks", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/encryption/key_fetcher/src:fake_key_fetcher_manager", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/encryption/key_fetcher:fake_key_fetcher_manager", ], ) diff --git a/components/sharding/cluster_mappings_manager.cc b/components/sharding/cluster_mappings_manager.cc index 101732c8..826422fc 100644 --- a/components/sharding/cluster_mappings_manager.cc +++ b/components/sharding/cluster_mappings_manager.cc @@ -26,7 +26,7 @@ ClusterMappingsManager::ClusterMappingsManager( : environment_{std::move(environment)}, num_shards_{num_shards}, instance_client_{instance_client}, - thread_manager_(TheadManager::Create("Cluster mappings updater")), + thread_manager_(ThreadManager::Create("Cluster mappings updater")), sleep_for_(std::move(sleep_for)), update_interval_millis_(update_interval_millis) { CHECK_GT(num_shards, 1) << "num_shards for ShardedLookup must be > 1"; diff --git a/components/sharding/cluster_mappings_manager.h b/components/sharding/cluster_mappings_manager.h index cb98e0ee..0c6b8ebc 100644 --- a/components/sharding/cluster_mappings_manager.h +++ b/components/sharding/cluster_mappings_manager.h @@ -68,7 +68,7 @@ class ClusterMappingsManager { std::string environment_; int32_t num_shards_; InstanceClient& instance_client_; - std::unique_ptr thread_manager_; + std::unique_ptr thread_manager_; std::unique_ptr sleep_for_; int32_t update_interval_millis_; }; diff --git a/components/sharding/cluster_mappings_manager_aws_test.cc b/components/sharding/cluster_mappings_manager_aws_test.cc index 7fe5cd2d..738e609f 100644 --- a/components/sharding/cluster_mappings_manager_aws_test.cc +++ b/components/sharding/cluster_mappings_manager_aws_test.cc @@ -23,8 +23,7 @@ #include "components/sharding/mocks.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" namespace kv_server { namespace { @@ -171,7 +170,6 @@ TEST_F(ClusterMappingsAwsTest, RetrieveMappingsWithRetrySuccessfully) { TEST_F(ClusterMappingsAwsTest, UpdateMappings) { std::string environment = "testenv"; int32_t num_shards = 2; - privacy_sandbox::server_common::MockMetricsRecorder mock_metrics_recorder; privacy_sandbox::server_common::FakeKeyFetcherManager fake_key_fetcher_manager; auto instance_client = std::make_unique(); @@ -179,9 +177,8 @@ TEST_F(ClusterMappingsAwsTest, UpdateMappings) { for (int i = 0; i < num_shards; i++) { cluster_mappings.push_back({"some_ip"}); } - auto shard_manager_status = - ShardManager::Create(num_shards, fake_key_fetcher_manager, - cluster_mappings, mock_metrics_recorder); + auto shard_manager_status = ShardManager::Create( + num_shards, fake_key_fetcher_manager, cluster_mappings); ASSERT_TRUE(shard_manager_status.ok()); auto shard_manager = std::move(*shard_manager_status); absl::Notification finished; @@ -205,7 +202,7 @@ TEST_F(ClusterMappingsAwsTest, UpdateMappings) { std::vector instances{ii1}; return instances; }) - .WillOnce([&](DescribeInstanceGroupInput& input) { + .WillRepeatedly([&](DescribeInstanceGroupInput& input) { auto aws_describe_instance_group_input = std::get_if(&input); absl::flat_hash_set instance_group_names_expected = { @@ -240,7 +237,7 @@ TEST_F(ClusterMappingsAwsTest, UpdateMappings) { std::vector instances{ii1}; return instances; }) - .WillOnce( + .WillRepeatedly( [&](const absl::flat_hash_set& instance_group_names) { absl::flat_hash_set instance_group_names_expected = { "id20"}; diff --git a/components/sharding/cluster_mappings_manager_gcp_test.cc b/components/sharding/cluster_mappings_manager_gcp_test.cc index 5700b5cd..c081af1a 100644 --- a/components/sharding/cluster_mappings_manager_gcp_test.cc +++ b/components/sharding/cluster_mappings_manager_gcp_test.cc @@ -23,8 +23,7 @@ #include "components/sharding/mocks.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" namespace kv_server { namespace { @@ -133,7 +132,6 @@ TEST_F(ClusterMappingsGcpTest, UpdateMappings) { std::string environment = "testenv"; std::string project_id = "some-project-id"; int32_t num_shards = 2; - privacy_sandbox::server_common::MockMetricsRecorder mock_metrics_recorder; privacy_sandbox::server_common::FakeKeyFetcherManager fake_key_fetcher_manager; auto instance_client = std::make_unique(); @@ -141,9 +139,8 @@ TEST_F(ClusterMappingsGcpTest, UpdateMappings) { for (int i = 0; i < num_shards; i++) { cluster_mappings.push_back({"some_ip"}); } - auto shard_manager_status = - ShardManager::Create(num_shards, fake_key_fetcher_manager, - cluster_mappings, mock_metrics_recorder); + auto shard_manager_status = ShardManager::Create( + num_shards, fake_key_fetcher_manager, cluster_mappings); ASSERT_TRUE(shard_manager_status.ok()); auto shard_manager = std::move(*shard_manager_status); absl::Notification finished; @@ -160,7 +157,7 @@ TEST_F(ClusterMappingsGcpTest, UpdateMappings) { std::vector instances{ii1}; return instances; }) - .WillOnce([&](DescribeInstanceGroupInput& input) { + .WillRepeatedly([&](DescribeInstanceGroupInput& input) { auto gcp_describe_instance_group_input = std::get_if(&input); EXPECT_EQ(gcp_describe_instance_group_input->project_id, project_id); diff --git a/components/sharding/shard_manager.cc b/components/sharding/shard_manager.cc index d76470ff..3a7e567f 100644 --- a/components/sharding/shard_manager.cc +++ b/components/sharding/shard_manager.cc @@ -140,17 +140,15 @@ absl::StatusOr> ShardManager::Create( int32_t num_shards, privacy_sandbox::server_common::KeyFetcherManagerInterface& key_fetcher_manager, - const std::vector>& cluster_mappings, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) { + const std::vector>& cluster_mappings) { auto validationStatus = ValidateMapping(num_shards, cluster_mappings); if (!validationStatus.ok()) { return validationStatus; } auto shard_manager = std::make_unique( cluster_mappings.size(), - [&key_fetcher_manager, &metrics_recorder](const std::string& ip) { - return RemoteLookupClient::Create(ip, key_fetcher_manager, - metrics_recorder); + [&key_fetcher_manager](const std::string& ip) { + return RemoteLookupClient::Create(ip, key_fetcher_manager); }, std::make_unique()); shard_manager->InsertBatch(std::move(cluster_mappings)); diff --git a/components/sharding/shard_manager.h b/components/sharding/shard_manager.h index cdf76608..ea4ff8ca 100644 --- a/components/sharding/shard_manager.h +++ b/components/sharding/shard_manager.h @@ -25,8 +25,6 @@ #include "absl/container/flat_hash_set.h" #include "components/internal_server/remote_lookup_client.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry.h" namespace kv_server { // This class is useful for testing ShardManager @@ -56,8 +54,7 @@ class ShardManager { int32_t num_shards, privacy_sandbox::server_common::KeyFetcherManagerInterface& key_fetcher_manager, - const std::vector>& cluster_mappings, - privacy_sandbox::server_common::MetricsRecorder& metrics_recorder); + const std::vector>& cluster_mappings); static absl::StatusOr> Create( int32_t num_shards, const std::vector>& cluster_mappings, diff --git a/components/sharding/shard_manager_test.cc b/components/sharding/shard_manager_test.cc index 971c5133..86cd0c8c 100644 --- a/components/sharding/shard_manager_test.cc +++ b/components/sharding/shard_manager_test.cc @@ -22,26 +22,22 @@ #include "components/sharding/mocks.h" #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/encryption/key_fetcher/src/fake_key_fetcher_manager.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/encryption/key_fetcher/fake_key_fetcher_manager.h" namespace kv_server { namespace { using privacy_sandbox::server_common::FakeKeyFetcherManager; -using privacy_sandbox::server_common::MockMetricsRecorder; class ShardManagerTest : public ::testing::Test { protected: FakeKeyFetcherManager fake_key_fetcher_manager_; - MockMetricsRecorder mock_metrics_recorder_; }; TEST_F(ShardManagerTest, CreationNotInitialized) { std::vector> cluster_mappings; - auto shard_manager = - ShardManager::Create(4, fake_key_fetcher_manager_, - std::move(cluster_mappings), mock_metrics_recorder_); + auto shard_manager = ShardManager::Create(4, fake_key_fetcher_manager_, + std::move(cluster_mappings)); ASSERT_FALSE(shard_manager.ok()); } @@ -51,9 +47,8 @@ TEST_F(ShardManagerTest, CreationInitialized) { for (int i = 0; i < num_shards; i++) { cluster_mappings.push_back({"some_ip"}); } - auto shard_manager = - ShardManager::Create(num_shards, fake_key_fetcher_manager_, - std::move(cluster_mappings), mock_metrics_recorder_); + auto shard_manager = ShardManager::Create( + num_shards, fake_key_fetcher_manager_, std::move(cluster_mappings)); ASSERT_TRUE(shard_manager.ok()); } @@ -63,9 +58,8 @@ TEST_F(ShardManagerTest, CreationNotInitializedMissingClusters) { for (int i = 0; i < 2; i++) { cluster_mappings.push_back({"some_ip"}); } - auto shard_manager = - ShardManager::Create(num_shards, fake_key_fetcher_manager_, - std::move(cluster_mappings), mock_metrics_recorder_); + auto shard_manager = ShardManager::Create( + num_shards, fake_key_fetcher_manager_, std::move(cluster_mappings)); ASSERT_FALSE(shard_manager.ok()); } @@ -76,9 +70,8 @@ TEST_F(ShardManagerTest, CreationNotInitializedMissingReplicas) { cluster_mappings.push_back({"some_ip"}); } cluster_mappings.push_back({}); - auto shard_manager = - ShardManager::Create(num_shards, fake_key_fetcher_manager_, - std::move(cluster_mappings), mock_metrics_recorder_); + auto shard_manager = ShardManager::Create( + num_shards, fake_key_fetcher_manager_, std::move(cluster_mappings)); ASSERT_FALSE(shard_manager.ok()); } @@ -88,9 +81,8 @@ TEST_F(ShardManagerTest, InsertRetrieveSuccess) { for (int i = 0; i < num_shards; i++) { cluster_mappings.push_back({"some_ip"}); } - auto shard_manager = - ShardManager::Create(num_shards, fake_key_fetcher_manager_, - std::move(cluster_mappings), mock_metrics_recorder_); + auto shard_manager = ShardManager::Create( + num_shards, fake_key_fetcher_manager_, std::move(cluster_mappings)); ASSERT_TRUE(shard_manager.ok()); EXPECT_EQ(absl::StrCat("some_ip:", kRemoteLookupServerPort), (*shard_manager)->Get(0)->GetIpAddress()); @@ -102,9 +94,8 @@ TEST_F(ShardManagerTest, InsertMissingReplicasRetrieveSuccess) { for (int i = 0; i < num_shards; i++) { cluster_mappings.push_back({"some_ip"}); } - auto shard_manager = - ShardManager::Create(num_shards, fake_key_fetcher_manager_, - std::move(cluster_mappings), mock_metrics_recorder_); + auto shard_manager = ShardManager::Create( + num_shards, fake_key_fetcher_manager_, std::move(cluster_mappings)); std::vector> cluster_mappings_2; for (int i = 0; i < 3; i++) { cluster_mappings_2.push_back({"some_ip"}); @@ -117,7 +108,6 @@ TEST_F(ShardManagerTest, InsertMissingReplicasRetrieveSuccess) { TEST_F(ShardManagerTest, InsertRetrieveTwoVersions) { auto random_generator = std::make_unique(); - MockMetricsRecorder mock_metrics_recorder_; EXPECT_CALL(*random_generator, Get(testing::_)) .WillOnce([]() { return 0; }) .WillOnce([]() { return 1; }); @@ -129,11 +119,8 @@ TEST_F(ShardManagerTest, InsertRetrieveTwoVersions) { cluster_mappings.push_back({"some_ip_3"}); } auto& fake_key_fetcher_manager = fake_key_fetcher_manager_; - auto& mock_metrics_recorder = mock_metrics_recorder_; - auto client_factory = [&fake_key_fetcher_manager, - &mock_metrics_recorder](const std::string& ip) { - return RemoteLookupClient::Create(ip, fake_key_fetcher_manager, - mock_metrics_recorder); + auto client_factory = [&fake_key_fetcher_manager](const std::string& ip) { + return RemoteLookupClient::Create(ip, fake_key_fetcher_manager); }; auto shard_manager = ShardManager::Create(4, std::move(cluster_mappings), diff --git a/components/telemetry/BUILD.bazel b/components/telemetry/BUILD.bazel index cfece9c8..67289032 100644 --- a/components/telemetry/BUILD.bazel +++ b/components/telemetry/BUILD.bazel @@ -51,7 +51,23 @@ cc_library( ], deps = [ ":error_code", - "@google_privacysandbox_servers_common//src/cpp/metric:context_map", - "@google_privacysandbox_servers_common//src/cpp/util:read_system", + "@google_privacysandbox_servers_common//src/core/common/uuid", + "@google_privacysandbox_servers_common//src/metric:context_map", + "@google_privacysandbox_servers_common//src/util:duration", + "@google_privacysandbox_servers_common//src/util:read_system", + ], +) + +cc_library( + name = "open_telemetry_sink", + srcs = [ + "open_telemetry_sink.cc", + ], + hdrs = [ + "open_telemetry_sink.h", + ], + deps = [ + "@com_google_absl//absl/log", + "@google_privacysandbox_servers_common//src/telemetry", ], ) diff --git a/components/telemetry/error_code.h b/components/telemetry/error_code.h index 55f1070c..4bf282ac 100644 --- a/components/telemetry/error_code.h +++ b/components/telemetry/error_code.h @@ -56,11 +56,92 @@ inline constexpr absl::string_view kRealtimeSleepFailure = inline constexpr absl::string_view kRealtimeDecodeMessageFailure = "RealtimeDecodeMessageFailure"; -// TODO(b/304311724): Revisit the error codes and provide utililities to make -// it easier to log error metrics +// Failure in decrypting request in the internal secure lookup +inline constexpr absl::string_view kRequestDecryptionFailure = + "RequestDecryptionFailure"; +// Error in unpadding request in the internal secure lookup +inline constexpr absl::string_view kRequestUnpaddingError = + "RequestUnpaddingError"; +// Failure in encrypting response in internal secure lookup +inline constexpr absl::string_view kResponseEncryptionFailure = + "ResponseEncryptionFailure"; +// Internal run query request failure +inline constexpr std::string_view kInternalRunQueryRequestFailure = + "InternalRunQueryRequestFailure"; +// Failure in executing run query in local lookup +inline constexpr std::string_view kLocalRunQueryFailure = + "LocalRunQueryFailure"; +// Missing keyset in the run query in local lookup +inline constexpr std::string_view kLocalRunQueryMissingKeySet = + "LocalRunQueryMissingKeySet"; +// Query parsing failure in the run query in local lookup +inline constexpr std::string_view kLocalRunQueryParsingFailure = + "LocalRunQueryParsingFailure"; + +// Lookup client missing in the sharded lookup +inline constexpr std::string_view kLookupClientMissing = "LookupClientMissing"; +// Failure in creating lookup futures +inline constexpr std::string_view kLookupFuturesCreationFailure = + "LookupFuturesCreationFailure"; +inline constexpr std::string_view kRemoteRequestEncryptionFailure = + "RemoteRequestEncryptionFailure"; +inline constexpr std::string_view kRemoteResponseDecryptionFailure = + "RemoteResponseDecryptionFailure"; +inline constexpr std::string_view kRemoteSecureLookupFailure = + "RemoteSecureLookupFailure"; +// Sharded GetKeyValues request failure +inline constexpr std::string_view kShardedKeyValueRequestFailure = + "ShardedKeyValueRequestFailure"; +// Sharded GetKeyValueSet request failure +inline constexpr std::string_view kShardedKeyValueSetRequestFailure = + "ShardedKeyValueSetRequestFailure"; +// Key collisions in collecting results from sharded GetKeyValueSet requests +inline constexpr std::string_view kShardedKeyCollisionOnKeySetCollection = + "ShardedKeyCollisionOnKeySetCollection"; +// Empty query encountered in the sharded lookup +inline constexpr std::string_view kShardedRunQueryEmptyQuery = + "ShardedRunQueryEmptyQuery"; +// Failure in running query in sharded lookup +inline constexpr std::string_view kShardedRunQueryFailure = + "ShardedRunQueryFailure"; +// Key set not found error in the GetValueKeySet in sharded lookup +inline constexpr std::string_view kShardedGetKeyValueSetKeySetNotFound = + "ShardedGetKeyValueSetKeySetNotFound"; +// Key set retrieval failure in the GetValueKeySet in sharded lookup +inline constexpr std::string_view kShardedGetKeyValueSetKeySetRetrievalFailure = + "ShardedGetKeyValueSetKeySetRetrievalFailure"; +// Failure in key set retrieval in the run query in sharded lookup +inline constexpr std::string_view kShardedRunQueryKeySetRetrievalFailure = + "ShardedRunQueryKeySetRetrievalFailure"; +// Missing keyset in the run query in sharded lookup +inline constexpr std::string_view kShardedRunQueryMissingKeySet = + "ShardedRunQueryMissingKeySet"; +// Query parsing failure in the run query in sharded lookup +inline constexpr std::string_view kShardedRunQueryParsingFailure = + "ShardedRunQueryParsingFailure"; // Strings must be sorted, this is required by the API of partitioned metrics -inline constexpr absl::string_view kChangeNotifierErrorCode[] = { +inline constexpr absl::string_view kKVUdfRequestErrorCode[] = { + kLookupClientMissing, + kLookupFuturesCreationFailure, + kRemoteRequestEncryptionFailure, + kRemoteResponseDecryptionFailure, + kRemoteSecureLookupFailure, + kShardedGetKeyValueSetKeySetNotFound, + kShardedGetKeyValueSetKeySetRetrievalFailure, + kShardedKeyCollisionOnKeySetCollection, + kShardedKeyValueRequestFailure, + kShardedKeyValueSetRequestFailure, + kShardedRunQueryEmptyQuery, + kShardedRunQueryFailure, + kShardedRunQueryKeySetRetrievalFailure, + kShardedRunQueryMissingKeySet, + kShardedRunQueryParsingFailure, +}; + +// Non request related server error +// Strings must be sorted, this is required by the API of partitioned metrics +inline constexpr std::string_view kKVServerErrorCode[] = { kAwsChangeNotifierMessagesDataLoss, kAwsChangeNotifierMessagesDeletionFailure, kAwsChangeNotifierMessagesReceivingFailure, @@ -68,16 +149,20 @@ inline constexpr absl::string_view kChangeNotifierErrorCode[] = { kAwsChangeNotifierTagFailure, kAwsJsonParseError, kDeltaFileRecordChangeNotifierParsingFailure, -}; - -// Strings must be sorted, this is required by the API of partitioned metrics -inline constexpr absl::string_view kRealtimeErrorCode[] = { kRealtimeDecodeMessageFailure, kRealtimeGetNotificationsFailure, kRealtimeMessageApplicationFailure, kRealtimeSleepFailure, }; +// Strings must be sorted, this is required by the API of partitioned metrics +inline constexpr absl::string_view kInternalLookupRequestErrorCode[] = { + kInternalRunQueryRequestFailure, kLocalRunQueryFailure, + kLocalRunQueryMissingKeySet, kLocalRunQueryParsingFailure, + kRequestDecryptionFailure, kRequestUnpaddingError, + kResponseEncryptionFailure, +}; + } // namespace kv_server #endif // COMPONENTS_TELEMETRY_ERROR_CODE_H_ diff --git a/components/telemetry/init_aws.cc b/components/telemetry/init_aws.cc index 050406eb..35020066 100644 --- a/components/telemetry/init_aws.cc +++ b/components/telemetry/init_aws.cc @@ -16,8 +16,8 @@ #include "opentelemetry/exporters/otlp/otlp_grpc_exporter_factory.h" #include "opentelemetry/exporters/otlp/otlp_grpc_metric_exporter_factory.h" #include "opentelemetry/sdk/trace/random_id_generator_factory.h" -#include "src/cpp/telemetry/init.h" -#include "src/cpp/telemetry/trace_generator_aws.h" +#include "src/telemetry/init.h" +#include "src/telemetry/trace_generator_aws.h" namespace kv_server { std::unique_ptr CreateSpanExporter() { diff --git a/components/telemetry/init_local_ostream.cc b/components/telemetry/init_local_ostream.cc index 8c5ec034..4a0b8a4f 100644 --- a/components/telemetry/init_local_ostream.cc +++ b/components/telemetry/init_local_ostream.cc @@ -16,7 +16,7 @@ #include "opentelemetry/exporters/ostream/span_exporter_factory.h" #include "opentelemetry/sdk/metrics/export/periodic_exporting_metric_reader.h" #include "opentelemetry/sdk/trace/random_id_generator_factory.h" -#include "src/cpp/telemetry/init.h" +#include "src/telemetry/init.h" namespace kv_server { diff --git a/components/telemetry/init_local_otlp.cc b/components/telemetry/init_local_otlp.cc index 8a870ff5..d0754012 100644 --- a/components/telemetry/init_local_otlp.cc +++ b/components/telemetry/init_local_otlp.cc @@ -15,14 +15,14 @@ #include "opentelemetry/exporters/otlp/otlp_grpc_exporter_factory.h" #include "opentelemetry/exporters/otlp/otlp_grpc_metric_exporter_factory.h" #include "opentelemetry/sdk/trace/random_id_generator_factory.h" -#include "src/cpp/telemetry/init.h" +#include "src/telemetry/init.h" // To use Jaeger first run a local instance of the collector // https://www.jaegertracing.io/docs/1.42/getting-started/ // Then build run server with flags for local and otlp. Ex: // `bazel run //components/data_server/server:server --//:instance=local // --//:platform=aws -// --@google_privacysandbox_servers_common//src/cpp/telemetry:local_otel_export=otlp +// --@google_privacysandbox_servers_common//src/telemetry:local_otel_export=otlp // -- // --environment="test"` namespace kv_server { diff --git a/components/telemetry/local_otlp_config/README.md b/components/telemetry/local_otlp_config/README.md index e51ce8d4..6a1cfbd0 100644 --- a/components/telemetry/local_otlp_config/README.md +++ b/components/telemetry/local_otlp_config/README.md @@ -2,8 +2,8 @@ The server can be run locally with 2 differt compile time flags. -- default: `--@google_privacysandbox_servers_common//src/cpp/telemetry:local_otel_export=ostream` -- alternative: `--@google_privacysandbox_servers_common//src/cpp/telemetry:local_otel_export=otlp` +- default: `--@google_privacysandbox_servers_common//src/telemetry:local_otel_export=ostream` +- alternative: `--@google_privacysandbox_servers_common//src/telemetry:local_otel_export=otlp` ## OTLP diff --git a/components/telemetry/open_telemetry_sink.cc b/components/telemetry/open_telemetry_sink.cc new file mode 100644 index 00000000..a7a8b1dc --- /dev/null +++ b/components/telemetry/open_telemetry_sink.cc @@ -0,0 +1,28 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "components/telemetry/open_telemetry_sink.h" + +#include + +namespace kv_server { +OpenTelemetrySink::OpenTelemetrySink( + opentelemetry::nostd::shared_ptr logger) + : logger_(std::move(logger)) {} +void OpenTelemetrySink::Send(const absl::LogEntry& entry) { + logger_->EmitLogRecord(entry.text_message_with_prefix_and_newline_c_str()); +} +// opentelemetry::logs::Logger doesn't have flush +void OpenTelemetrySink::Flush() {} +} // namespace kv_server diff --git a/components/telemetry/open_telemetry_sink.h b/components/telemetry/open_telemetry_sink.h new file mode 100644 index 00000000..0334f950 --- /dev/null +++ b/components/telemetry/open_telemetry_sink.h @@ -0,0 +1,38 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef COMPONENTS_TELEMETRY_OPEN_TELEMETRY_SINK_H_ +#define COMPONENTS_TELEMETRY_OPEN_TELEMETRY_SINK_H_ + +#include + +#include "absl/log/log.h" +#include "src/telemetry/telemetry.h" + +namespace kv_server { + +class OpenTelemetrySink : public absl::LogSink { + public: + explicit OpenTelemetrySink( + opentelemetry::nostd::shared_ptr logger); + void Send(const absl::LogEntry& entry) override; + void Flush() override; + opentelemetry::nostd::shared_ptr logger_; +}; + +} // namespace kv_server + +#endif // COMPONENTS_TELEMETRY_OPEN_TELEMETRY_SINK_H_ diff --git a/components/telemetry/server_definition.h b/components/telemetry/server_definition.h index ba8a4e8a..339dd851 100644 --- a/components/telemetry/server_definition.h +++ b/components/telemetry/server_definition.h @@ -21,12 +21,18 @@ #include #include +#include "absl/time/time.h" #include "components/telemetry/error_code.h" -#include "src/cpp/metric/context_map.h" -#include "src/cpp/util/read_system.h" +#include "src/core/common/uuid/uuid.h" +#include "src/metric/context_map.h" +#include "src/util/duration.h" +#include "src/util/read_system.h" namespace kv_server { +constexpr std::string_view kKVServerServiceName = "KVServer"; +constexpr std::string_view kInternalLookupServiceName = "InternalLookupServer"; + // When server is running in debug mode, all unsafe metrics will be logged // safely without DP noise applied. Therefore for now it is okay to set DP // upper and lower bounds for all unsafe metrics to a default value. But @@ -37,6 +43,11 @@ namespace kv_server { constexpr int kCounterDPLowerBound = 1; constexpr int kCounterDPUpperBound = 10; +constexpr int kErrorCounterDPLowerBound = 0; +constexpr int kErrorCounterDPUpperBound = 1; +constexpr int kErrorMaxPartitionsContributed = 1; +constexpr double kErrorMinNoiseToOutput = 0.99; + constexpr int kMicroSecondsLowerBound = 1; constexpr int kMicroSecondsUpperBound = 2'000'000'000; @@ -69,6 +80,15 @@ inline constexpr absl::string_view kAbslStatusStrings[] = { "UNIMPLEMENTED", "UNKNOWN"}; +inline constexpr std::string_view kKeyValueCacheHit = "KeyValueCacheHit"; +inline constexpr std::string_view kKeyValueCacheMiss = "KeyValueCacheMiss"; +inline constexpr std::string_view kKeyValueSetCacheHit = "KeyValueSetCacheHit"; +inline constexpr std::string_view kKeyValueSetCacheMiss = + "KeyValueSetCacheMiss"; +inline constexpr std::string_view kCacheAccessEvents[] = { + kKeyValueCacheHit, kKeyValueCacheMiss, kKeyValueSetCacheHit, + kKeyValueSetCacheMiss}; + inline constexpr privacy_sandbox::server_common::metrics::PrivacyBudget privacy_total_budget{/*epsilon*/ 5}; @@ -76,156 +96,100 @@ inline constexpr privacy_sandbox::server_common::metrics::PrivacyBudget // and should be logged unsafe with DP(differential privacy) noises. inline constexpr privacy_sandbox::server_common::metrics::Definition< int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kInternalRunQueryKeySetRetrievalFailure( - "InternalRunQueryKeySetRetrievalFailure", - "Number of key set internal retrieval failures during internal" - "run query processing", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kKeysNotFoundInKeySetsInShardedLookup( - "KeysNotFoundInKeySetsInShardedLookup", - "Number of keys not found in the result key set in the sharded lookup", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kKeysNotFoundInKeySetsInLocalLookup( - "KeysNotFoundInKeySetsInLocalLookup", - "Number of keys not found in the result key set in the local lookup", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kInternalRunQueryEmtpyQuery("InternalRunQueryEmtpyQuery", - "Number of empty queries encountered during " - "internal run query processing", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kInternalRunQueryMissingKeySet( - "InternalRunQueryMissingKeySet", - "Number of missing keys not found in the key set " - "during internal run query processing", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kInternalRunQueryParsingFailure( - "InternalRunQueryParsingFailure", - "Number of failures in parsing query during " - "internal run query processing", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kLookupClientMissing( - "LookupClientMissing", - "Number of missing internal lookup clients encountered during " - "sharded lookup", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kShardedLookupServerRequestFailed( - "ShardedLookupServerRequestFailed", - "Number of failed server requests in the sharded lookup", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kShardedLookupServerKeyCollisionOnCollection( - "ShardedLookupServerKeyCollisionOnCollection", - "Number of key collisions when collecting results from shards", - kCounterDPUpperBound, kCounterDPLowerBound); + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kKVUdfRequestError("KVUdfRequestError", + "Errors in processing KV server V2 request", + "error_code", kErrorMaxPartitionsContributed, + kKVUdfRequestErrorCode, kErrorCounterDPUpperBound, + kErrorCounterDPLowerBound, kErrorMinNoiseToOutput); +// Metric definitions for request level metrics that are privacy impacting +// and should be logged unsafe with DP(differential privacy) noises. inline constexpr privacy_sandbox::server_common::metrics::Definition< int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kLookupFuturesCreationFailure( - "LookupFuturesCreationFailure", - "Number of failures in creating lookup futures in the sharded lookup", + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kShardedLookupKeyCountByShard( + "ShardedLookupKeyCountByShard", "Keys count by shard number", + "key_shard_num", 1 /*max_partitions_contributed*/, + privacy_sandbox::server_common::metrics::kEmptyPublicPartition, kCounterDPUpperBound, kCounterDPLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kShardedLookupFailure("ShardedLookupFailure", - "Number of lookup failures in the sharded lookup", - kCounterDPUpperBound, kCounterDPLowerBound); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kRemoteClientEncryptionFailure( - "RemoteClientEncryptionFailure", - "Number of request encryption failures in the remote lookup client", - kCounterDPUpperBound, kCounterDPLowerBound); + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kShardedLookupGetKeyValuesLatencyInMicros( + "ShardedLookupGetKeyValuesLatencyInMicros", + "Latency in executing GetKeyValues in the sharded lookup", + kLatencyInMicroSecondsBoundaries, kMicroSecondsUpperBound, + kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kRemoteClientSecureLookupFailure( - "RemoteClientSecureLookupFailure", - "Number of secure lookup failures in the remote lookup client", - kCounterDPUpperBound, kCounterDPLowerBound); + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kShardedLookupGetKeyValueSetLatencyInMicros( + "ShardedLookupGetKeyValueSetLatencyInMicros", + "Latency in executing GetKeyValueSet in the sharded lookup", + kLatencyInMicroSecondsBoundaries, kMicroSecondsUpperBound, + kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kRemoteClientDecryptionFailure( - "RemoteClientDecryptionFailure", - "Number of response decryption failures in the remote lookup client", - kCounterDPUpperBound, kCounterDPLowerBound); + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kShardedLookupRunQueryLatencyInMicros( + "ShardedLookupRunQueryLatencyInMicros", + "Latency in executing RunQuery in the sharded lookup", + kLatencyInMicroSecondsBoundaries, kMicroSecondsUpperBound, + kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kInternalClientDecryptionFailure( - "InternalClientEncryptionFailure", - "Number of request decryption failures in the internal lookup client", - kCounterDPUpperBound, kCounterDPLowerBound); + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kRemoteLookupGetValuesLatencyInMicros( + "RemoteLookupGetValuesLatencyInMicros", + "Latency in get values in the remote lookup", + kLatencyInMicroSecondsBoundaries, kMicroSecondsUpperBound, + kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kInternalClientUnpaddingRequestError( - "InternalClientUnpaddingRequestError", - "Number of unpadding errors in the request deserialization in the " - "internal lookup client", - kCounterDPUpperBound, kCounterDPLowerBound); + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kInternalRunQueryLatencyInMicros("InternalRunQueryLatencyInMicros", + "Latency in internal run query call", + kLatencyInMicroSecondsBoundaries, + kMicroSecondsUpperBound, + kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, privacy_sandbox::server_common::metrics::Instrument::kHistogram> - kShardedLookupRunQueryLatencyInMicros( - "ShardedLookupRunQueryLatencyInMicros", - "Latency in executing run query in the sharded lookup", + kInternalGetKeyValuesLatencyInMicros( + "InternalGetKeyValuesLatencyInMicros", + "Latency in internal get key values call", kLatencyInMicroSecondsBoundaries, kMicroSecondsUpperBound, kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, privacy_sandbox::server_common::metrics::Instrument::kHistogram> - kRemoteLookupGetValuesLatencyInMicros( - "RemoteLookupGetValuesLatencyInMicros", - "Latency in get values in the remote lookup", + kInternalGetKeyValueSetLatencyInMicros( + "InternalGetKeyValueSetLatencyInMicros", + "Latency in internal get key value set call", kLatencyInMicroSecondsBoundaries, kMicroSecondsUpperBound, kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kInternalLookupRequestError("InternalLookupRequestError", + "Errors in processing internal lookup request", + "error_code", kErrorMaxPartitionsContributed, + kInternalLookupRequestErrorCode, + kErrorCounterDPUpperBound, + kErrorCounterDPLowerBound, + kErrorMinNoiseToOutput); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, privacy_sandbox::server_common::metrics::Instrument::kHistogram> kInternalSecureLookupLatencyInMicros("InternalSecureLookupLatencyInMicros", "Latency in internal secure lookup", @@ -234,7 +198,7 @@ inline constexpr privacy_sandbox::server_common::metrics::Definition< kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, privacy_sandbox::server_common::metrics::Instrument::kHistogram> kGetValuePairsLatencyInMicros("GetValuePairsLatencyInMicros", "Latency in executing GetValuePairs in cache", @@ -243,7 +207,7 @@ inline constexpr privacy_sandbox::server_common::metrics::Definition< kMicroSecondsLowerBound); inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + double, privacy_sandbox::server_common::metrics::Privacy::kImpacting, privacy_sandbox::server_common::metrics::Instrument::kHistogram> kGetKeyValueSetLatencyInMicros( "GetKeyValueSetLatencyInMicros", @@ -251,7 +215,24 @@ inline constexpr privacy_sandbox::server_common::metrics::Definition< kLatencyInMicroSecondsBoundaries, kMicroSecondsUpperBound, kMicroSecondsLowerBound); +inline constexpr privacy_sandbox::server_common::metrics::Definition< + int, privacy_sandbox::server_common::metrics::Privacy::kImpacting, + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kCacheAccessEventCount("CacheAccessEventCount", + "Count of cache hit or miss events by request", + "cache_access", 4 /*max_partitions_contributed*/, + kCacheAccessEvents, kCounterDPUpperBound, + kCounterDPLowerBound); + // Metric definitions for safe metrics that are not privacy impacting +inline constexpr privacy_sandbox::server_common::metrics::Definition< + int, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kRequestFailedCountByStatus( + "request.failed_count_by_status", + "Total number of requests that resulted in failure partitioned by " + "Error Code", + "error_code", kAbslStatusStrings); inline constexpr privacy_sandbox::server_common::metrics::Definition< int, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> @@ -306,12 +287,6 @@ inline constexpr privacy_sandbox::server_common::metrics::Definition< "Describe instances status", "status", kAbslStatusStrings); -inline constexpr privacy_sandbox::server_common::metrics::Definition< - double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kRealtimeTotalRowsUpdated("RealtimeTotalRowsUpdated", - "Realtime total rows updated count"); - inline constexpr privacy_sandbox::server_common::metrics::Definition< double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, privacy_sandbox::server_common::metrics::Instrument::kHistogram> @@ -349,15 +324,8 @@ inline constexpr privacy_sandbox::server_common::metrics::Definition< inline constexpr privacy_sandbox::server_common::metrics::Definition< int, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> - kChangeNotifierErrors("ChangeNotifierErrors", - "Errors in the change notifier", "error_code", - kChangeNotifierErrorCode); - -inline constexpr privacy_sandbox::server_common::metrics::Definition< - int, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, - privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> - kRealtimeErrors("RealtimeErrors", "Errors in realtime data loading", - "error_code", kRealtimeErrorCode); + kKVServerError("KVServerError", "Non request related server errors", + "error_code", kKVServerErrorCode); inline constexpr privacy_sandbox::server_common::metrics::Definition< double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, @@ -384,21 +352,33 @@ inline constexpr privacy_sandbox::server_common::metrics::Definition< inline constexpr privacy_sandbox::server_common::metrics::Definition< double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kTotalRowsDroppedInDataLoading("TotalRowsDroppedInDataLoading", - "Total rows dropped during data loading"); + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kTotalRowsDroppedInDataLoading( + "TotalRowsDroppedInDataLoading", + "Total rows dropped during data loading from data source," + "data source can be a data file or realtime", + "data_source", + privacy_sandbox::server_common::metrics::kEmptyPublicPartition); inline constexpr privacy_sandbox::server_common::metrics::Definition< double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kTotalRowsUpdatedInDataLoading("TotalRowsUpdatedInDataLoading", - "Total rows updated during data loading"); + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kTotalRowsUpdatedInDataLoading( + "TotalRowsUpdatedInDataLoading", + "Total rows updated during data loading from data source," + "data source can be a data file or realtime ", + "data_source", + privacy_sandbox::server_common::metrics::kEmptyPublicPartition); inline constexpr privacy_sandbox::server_common::metrics::Definition< double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, - privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> - kTotalRowsDeletedInDataLoading("TotalRowsDeletedInDataLoading", - "Total rows deleted during data loading"); + privacy_sandbox::server_common::metrics::Instrument::kPartitionedCounter> + kTotalRowsDeletedInDataLoading( + "TotalRowsDeletedInDataLoading", + "Total rows deleted during data loading from data source," + "data source can be a data file or realtime", + "data_source", + privacy_sandbox::server_common::metrics::kEmptyPublicPartition); inline constexpr privacy_sandbox::server_common::metrics::Definition< double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, @@ -424,58 +404,148 @@ inline constexpr privacy_sandbox::server_common::metrics::Definition< "Latency in ConcurrentStreamRecordReader reading byte range", kLatencyInMicroSecondsBoundaries); +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kUpdateKeyValueLatency("UpdateKeyValueLatency", + "Latency in key value update", + kLatencyInMicroSecondsBoundaries); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kUpdateKeyValueSetLatency("UpdateKeyValueSetLatency", + "Latency in key value set update", + kLatencyInMicroSecondsBoundaries); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kDeleteKeyLatency("DeleteKeyLatency", "Latency in deleting key", + kLatencyInMicroSecondsBoundaries); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kDeleteValuesInSetLatency("DeleteValuesInSetLatency", + "Latency in deleting values in set", + kLatencyInMicroSecondsBoundaries); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kRemoveDeletedKeyLatency( + "RemoveDeletedKeyLatency", + "Latency in removing deleted keys in the clean up process", + kLatencyInMicroSecondsBoundaries); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kCleanUpKeyValueMapLatency("CleanUpKeyValueMapLatency", + "Latency in cleaning up key value map", + kLatencyInMicroSecondsBoundaries); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + double, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kHistogram> + kCleanUpKeyValueSetMapLatency("CleanUpKeyValueSetMapLatency", + "Latency in cleaning up key value set map", + kLatencyInMicroSecondsBoundaries); + +inline constexpr privacy_sandbox::server_common::metrics::Definition< + int, privacy_sandbox::server_common::metrics::Privacy::kNonImpacting, + privacy_sandbox::server_common::metrics::Instrument::kUpDownCounter> + kSecureLookupRequestCount( + "SecureLookupRequestCount", + "Number of secure lookup requests received from remote server"); + +// KV server metrics list contains contains non request related safe metrics +// and request metrics collected before stage of internal lookups inline constexpr const privacy_sandbox::server_common::metrics::DefinitionName* kKVServerMetricList[] = { // Unsafe metrics - &kInternalRunQueryKeySetRetrievalFailure, - &kKeysNotFoundInKeySetsInShardedLookup, - &kKeysNotFoundInKeySetsInLocalLookup, &kInternalRunQueryEmtpyQuery, - &kInternalRunQueryMissingKeySet, &kInternalRunQueryParsingFailure, - &kLookupClientMissing, &kShardedLookupServerRequestFailed, - &kShardedLookupServerKeyCollisionOnCollection, - &kLookupFuturesCreationFailure, &kShardedLookupFailure, - &kRemoteClientEncryptionFailure, &kRemoteClientSecureLookupFailure, - &kRemoteClientDecryptionFailure, &kInternalClientDecryptionFailure, - &kInternalClientUnpaddingRequestError, + &kKVUdfRequestError, &kShardedLookupKeyCountByShard, + &kShardedLookupGetKeyValuesLatencyInMicros, + &kShardedLookupGetKeyValueSetLatencyInMicros, &kShardedLookupRunQueryLatencyInMicros, &kRemoteLookupGetValuesLatencyInMicros, - &kInternalSecureLookupLatencyInMicros, &kGetValuePairsLatencyInMicros, - &kGetKeyValueSetLatencyInMicros, // Safe metrics + &kKVServerError, + &privacy_sandbox::server_common::metrics::kTotalRequestCount, &privacy_sandbox::server_common::metrics::kServerTotalTimeMs, - &kGetParameterStatus, &kCompleteLifecycleStatus, - &kCreateDataOrchestratorStatus, &kStartDataOrchestratorStatus, - &kLoadNewFilesStatus, &kGetShardManagerStatus, - &kDescribeInstanceGroupInstancesStatus, &kDescribeInstancesStatus, - &kRealtimeTotalRowsUpdated, + &privacy_sandbox::server_common::metrics::kRequestByte, + &privacy_sandbox::server_common::metrics::kResponseByte, + &kRequestFailedCountByStatus, &kGetParameterStatus, + &kCompleteLifecycleStatus, &kCreateDataOrchestratorStatus, + &kStartDataOrchestratorStatus, &kLoadNewFilesStatus, + &kGetShardManagerStatus, &kDescribeInstanceGroupInstancesStatus, + &kDescribeInstancesStatus, &kReceivedLowLatencyNotificationsE2ECloudProvided, &kReceivedLowLatencyNotificationsE2E, &kReceivedLowLatencyNotifications, - &kChangeNotifierErrors, &kRealtimeErrors, &kAwsSqsReceiveMessageLatency, - &kSeekingInputStreambufSeekoffLatency, + &kAwsSqsReceiveMessageLatency, &kSeekingInputStreambufSeekoffLatency, &kSeekingInputStreambufSizeLatency, &kSeekingInputStreambufUnderflowLatency, &kTotalRowsDroppedInDataLoading, &kTotalRowsUpdatedInDataLoading, &kTotalRowsDeletedInDataLoading, &kConcurrentStreamRecordReaderReadShardRecordsLatency, &kConcurrentStreamRecordReaderReadStreamRecordsLatency, - &kConcurrentStreamRecordReaderReadByteRangeLatency}; + &kConcurrentStreamRecordReaderReadByteRangeLatency, + &kUpdateKeyValueLatency, &kUpdateKeyValueSetLatency, &kDeleteKeyLatency, + &kDeleteValuesInSetLatency, &kRemoveDeletedKeyLatency, + &kCleanUpKeyValueMapLatency, &kCleanUpKeyValueSetMapLatency}; + +// Internal lookup service metrics list contains metrics collected in the +// internal lookup server. This separation from KV metrics list allows all +// lookup requests (local and request from remote KV server) to contribute to +// the same set of metrics, so that the noise of unsafe metric won't be skewed +// for particular batch of requests, e.g server request that requires only +// remote lookups +inline constexpr const privacy_sandbox::server_common::metrics::DefinitionName* + kInternalLookupServiceMetricsList[] = { + // Safe metrics + &kSecureLookupRequestCount, + // Unsafe metrics + &kInternalLookupRequestError, &kInternalRunQueryLatencyInMicros, + &kInternalGetKeyValuesLatencyInMicros, + &kInternalGetKeyValueSetLatencyInMicros, + &kInternalSecureLookupLatencyInMicros, &kGetValuePairsLatencyInMicros, + &kGetKeyValueSetLatencyInMicros, &kCacheAccessEventCount}; inline constexpr absl::Span< const privacy_sandbox::server_common::metrics::DefinitionName* const> kKVServerMetricSpan = kKVServerMetricList; +inline constexpr absl::Span< + const privacy_sandbox::server_common::metrics::DefinitionName* const> + kInternalLookupServiceMetricsSpan = kInternalLookupServiceMetricsList; + inline auto* KVServerContextMap( std::optional< privacy_sandbox::server_common::telemetry::BuildDependentConfig> config = std::nullopt, std::unique_ptr provider = nullptr, - absl::string_view service = "", absl::string_view version = "") { + absl::string_view service = kKVServerServiceName, + absl::string_view version = "") { return privacy_sandbox::server_common::metrics::GetContextMap< const std::string, kKVServerMetricSpan>(std::move(config), std::move(provider), service, version, privacy_total_budget); } +inline auto* InternalLookupServerContextMap( + std::optional< + privacy_sandbox::server_common::telemetry::BuildDependentConfig> + config = std::nullopt, + std::unique_ptr provider = nullptr, + absl::string_view service = kInternalLookupServiceName, + absl::string_view version = "") { + return privacy_sandbox::server_common::metrics::GetContextMap< + const std::string, kInternalLookupServiceMetricsSpan>( + std::move(config), std::move(provider), service, version, + privacy_total_budget); +} + template inline void AddSystemMetric(T* context_map) { context_map->AddObserverable( @@ -519,8 +589,154 @@ inline void InitMetricsContextMap() { kv_server::KVServerContextMap( privacy_sandbox::server_common::telemetry::BuildDependentConfig( config_proto)); + kv_server::InternalLookupServerContextMap( + privacy_sandbox::server_common::telemetry::BuildDependentConfig( + config_proto)); +} + +using UdfRequestMetricsContext = + privacy_sandbox::server_common::metrics::ServerContext; +using InternalLookupMetricsContext = + privacy_sandbox::server_common::metrics::ServerContext< + kInternalLookupServiceMetricsSpan>; +using ServerSafeMetricsContext = + privacy_sandbox::server_common::metrics::ServerSafeContext< + kKVServerMetricSpan>; + +inline void LogUdfRequestErrorMetric(UdfRequestMetricsContext& metrics_context, + std::string_view error_code) { + LogIfError( + metrics_context.AccumulateMetric(1, error_code)); +} + +inline void LogInternalLookupRequestErrorMetric( + InternalLookupMetricsContext& metrics_context, + std::string_view error_code) { + LogIfError(metrics_context.AccumulateMetric( + 1, error_code)); } +// Logs non-request related error metrics +inline void LogServerErrorMetric(std::string_view error_code) { + LogIfError( + KVServerContextMap()->SafeMetric().LogUpDownCounter( + {{std::string(error_code), 1}})); +} + +// Logs common safe request metrics +template +inline void LogRequestCommonSafeMetrics( + const RequestT* request, const ResponseT* response, + const grpc::Status& grpc_request_status, + const absl::Time& request_received_time) { + LogIfError( + KVServerContextMap() + ->SafeMetric() + .LogUpDownCounter< + privacy_sandbox::server_common::metrics::kTotalRequestCount>(1)); + if (auto request_status = + privacy_sandbox::server_common::ToAbslStatus(grpc_request_status); + !request_status.ok()) { + LogIfError(KVServerContextMap() + ->SafeMetric() + .LogUpDownCounter( + {{absl::StatusCodeToString(request_status.code()), 1}})); + } + LogIfError(KVServerContextMap() + ->SafeMetric() + .template LogHistogram< + privacy_sandbox::server_common::metrics::kRequestByte>( + (int)request->ByteSizeLong())); + LogIfError(KVServerContextMap() + ->SafeMetric() + .template LogHistogram< + privacy_sandbox::server_common::metrics::kResponseByte>( + (int)response->ByteSizeLong())); + int duration_ms = + (absl::Now() - request_received_time) / absl::Milliseconds(1); + LogIfError( + KVServerContextMap() + ->SafeMetric() + .LogHistogram< + privacy_sandbox::server_common::metrics::kServerTotalTimeMs>( + duration_ms)); +} + +// ScopeMetricsContext provides metrics context ties to the request and +// should have the same lifetime of the request. +// The purpose of this class is to avoid explicit creating and deleting metrics +// context from context map. The metrics context associated with the request +// will be destroyed after ScopeMetricsContext goes out of scope. +class ScopeMetricsContext { + public: + explicit ScopeMetricsContext( + std::string request_id = google::scp::core::common::ToString( + google::scp::core::common::Uuid::GenerateUuid())) + : request_id_(std::move(request_id)) { + // Create a metrics context in the context map and + // associated it with request id + KVServerContextMap()->Get(&request_id_); + CHECK_OK([this]() { + // Remove the metrics context for request_id to transfer the ownership + // of metrics context to the ScopeMetricsContext. This is to ensure that + // metrics context has the same lifetime with RequestContext and be + // destroyed when ScopeMetricsContext goes out of scope. + PS_ASSIGN_OR_RETURN(udf_request_metrics_context_, + KVServerContextMap()->Remove(&request_id_)); + return absl::OkStatus(); + }()) << "Udf request metrics context is not initialized"; + InternalLookupServerContextMap()->Get(&request_id_); + CHECK_OK([this]() { + // Remove the metrics context for request_id to transfer the ownership + // of metrics context to the ScopeMetricsContext. This is to ensure that + // metrics context has the same lifetime with RequestContext and be + // destroyed when ScopeMetricsContext goes out of scope. + PS_ASSIGN_OR_RETURN( + internal_lookup_metrics_context_, + InternalLookupServerContextMap()->Remove(&request_id_)); + return absl::OkStatus(); + }()) << "Internal lookup metrics context is not initialized"; + } + UdfRequestMetricsContext& GetUdfRequestMetricsContext() const { + return *udf_request_metrics_context_; + } + InternalLookupMetricsContext& GetInternalLookupMetricsContext() const { + return *internal_lookup_metrics_context_; + } + + private: + const std::string request_id_; + // Metrics context has the same lifetime of server request context + std::unique_ptr udf_request_metrics_context_; + std::unique_ptr + internal_lookup_metrics_context_; +}; + +// Measures the latency of a block of code. The latency is recorded in +// microseconds as histogram metrics when the object of this class goes +// out of scope. The metric can be either safe or unsafe metric. +template +class ScopeLatencyMetricsRecorder { + public: + explicit ScopeLatencyMetricsRecorder( + ContextT& metrics_context, + std::unique_ptr stopwatch = + std::make_unique()) + : metrics_context_(metrics_context) { + stopwatch_ = std::move(stopwatch); + } + ~ScopeLatencyMetricsRecorder() { + LogIfError(metrics_context_.template LogHistogram( + absl::ToDoubleMicroseconds(stopwatch_->GetElapsedTime()))); + } + // Returns the latency so far + absl::Duration GetLatency() { return stopwatch_->GetElapsedTime(); } + + private: + ContextT& metrics_context_; + std::unique_ptr stopwatch_; +}; + } // namespace kv_server #endif // COMPONENTS_TELEMETRY_SERVER_DEFINITION_H_ diff --git a/components/tools/BUILD.bazel b/components/tools/BUILD.bazel index f18e4c3c..91a7483f 100644 --- a/components/tools/BUILD.bazel +++ b/components/tools/BUILD.bazel @@ -24,7 +24,7 @@ pkg_tar( name = "data_loading_analyzer_binaries", srcs = [ ":data_loading_analyzer", - "@google_privacysandbox_servers_common//scp/cc/aws/proxy/src:proxify_layer", + "@google_privacysandbox_servers_common//src/aws/proxy:libnsm_and_proxify_tar", ], package_dir = "/opt/privacysandbox/bin", ) @@ -74,12 +74,10 @@ cc_binary( "//public/data_loading/readers:riegeli_stream_record_reader_factory", "//public/sharding:key_sharder", "@com_github_google_flatbuffers//:flatbuffers", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", ], ) @@ -111,7 +109,7 @@ cc_library( cc_binary( name = "blob_storage_util", srcs = [ - "blob_storage_util_aws.cc", + "blob_storage_util.cc", ], deps = [ ":blob_storage_commands", @@ -167,10 +165,10 @@ cc_binary( "//components/errors:aws_error_util", "@aws_sdk_cpp//:core", "@aws_sdk_cpp//:ec2", - "@com_github_google_glog//:glog", "@com_google_absl//absl/cleanup", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", ], @@ -189,9 +187,11 @@ cc_binary( deps = [ ":concurrent_publishing_engine", "//components/util:platform_initializer", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/strings", ], ) @@ -202,6 +202,9 @@ cc_library( "//:gcp_platform": [ "publisher_service_gcp.cc", ], + "//:local_platform": [ + "publisher_service_local.cc", + ], "//conditions:default": [ "publisher_service_aws.cc", ], @@ -220,15 +223,17 @@ cc_library( "//components/util:platform_initializer", "@com_github_googleapis_google_cloud_cpp//:pubsub", ], + "//:local_platform": [ + ], "//conditions:default": [ "//components/errors:aws_error_util", "@aws_sdk_cpp//:sns", ], }) + [ "//components/data/common:message_service", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -275,6 +280,8 @@ cc_binary( "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", "@com_google_absl//absl/flags:usage", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/strings", ], ) @@ -288,8 +295,8 @@ cc_library( visibility = ["//tools:__subpackages__"], deps = [ ":publisher_service", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/util:duration", + "@google_privacysandbox_servers_common//src/util:duration", ], ) diff --git a/components/tools/benchmarks/BUILD.bazel b/components/tools/benchmarks/BUILD.bazel index a0c5b33a..ccc611e4 100644 --- a/components/tools/benchmarks/BUILD.bazel +++ b/components/tools/benchmarks/BUILD.bazel @@ -25,7 +25,7 @@ cc_library( deps = [ "//public/data_loading:records_utils", "//public/data_loading/writers:delta_record_stream_writer", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -47,6 +47,7 @@ cc_test( cc_binary( name = "data_loading_benchmark", srcs = ["data_loading_benchmark.cc"], + malloc = "@com_google_tcmalloc//tcmalloc", deps = [ ":benchmark_util", "//components/data/blob_storage:blob_storage_client", @@ -57,36 +58,38 @@ cc_binary( "//public/data_loading:data_loading_fbs", "//public/data_loading:records_utils", "//public/data_loading/readers:riegeli_stream_io", - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", "@com_google_benchmark//:benchmark", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", ], ) cc_binary( name = "cache_benchmark", srcs = ["cache_benchmark.cc"], + malloc = "@com_google_tcmalloc//tcmalloc", deps = [ ":benchmark_util", "//components/data_server/cache", "//components/data_server/cache:key_value_cache", "//components/data_server/cache:noop_key_value_cache", - "@com_github_google_glog//:glog", + "//components/util:request_context", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", "@com_google_benchmark//:benchmark", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", ], ) diff --git a/components/tools/benchmarks/cache_benchmark.cc b/components/tools/benchmarks/cache_benchmark.cc index 662ddc73..3aa3133f 100644 --- a/components/tools/benchmarks/cache_benchmark.cc +++ b/components/tools/benchmarks/cache_benchmark.cc @@ -23,6 +23,9 @@ #include "absl/container/flat_hash_map.h" #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" @@ -32,9 +35,6 @@ #include "components/data_server/cache/key_value_cache.h" #include "components/data_server/cache/noop_key_value_cache.h" #include "components/tools/benchmarks/benchmark_util.h" -#include "glog/logging.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry_provider.h" ABSL_FLAG(std::vector, record_size, std::vector({"1"}), @@ -69,8 +69,6 @@ namespace { using kv_server::benchmark::AsyncTask; using kv_server::benchmark::GenerateRandomString; using kv_server::benchmark::ParseInt64List; -using privacy_sandbox::server_common::MetricsRecorder; -using privacy_sandbox::server_common::TelemetryProvider; // Format variables used to generate benchmark names. // @@ -104,8 +102,8 @@ Cache* GetNoOpCache() { return cache; } -Cache* GetLockBasedCache(MetricsRecorder& metrics_recorder) { - static auto* const cache = KeyValueCache::Create(metrics_recorder).release(); +Cache* GetLockBasedCache() { + static auto* const cache = KeyValueCache::Create().release(); return cache; } @@ -173,8 +171,11 @@ void BM_GetKeyValuePairs(::benchmark::State& state, BenchmarkArgs args) { } auto keys = GetKeys(args.query_size); auto keys_view = ToContainerView>(keys); + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); for (auto _ : state) { - ::benchmark::DoNotOptimize(args.cache->GetKeyValuePairs(keys_view)); + ::benchmark::DoNotOptimize( + args.cache->GetKeyValuePairs(request_context, keys_view)); } state.counters[std::string(kReadsPerSec)] = ::benchmark::Counter(state.iterations(), ::benchmark::Counter::kIsRate); @@ -199,8 +200,11 @@ void BM_GetKeyValueSet(::benchmark::State& state, BenchmarkArgs args) { } auto keys = GetKeys(args.query_size); auto keys_view = ToContainerView>(keys); + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); for (auto _ : state) { - ::benchmark::DoNotOptimize(args.cache->GetKeyValueSet(keys_view)); + ::benchmark::DoNotOptimize( + args.cache->GetKeyValueSet(request_context, keys_view)); } state.counters[std::string(kReadsPerSec)] = ::benchmark::Counter(state.iterations(), ::benchmark::Counter::kIsRate); @@ -209,14 +213,16 @@ void BM_GetKeyValueSet(::benchmark::State& state, BenchmarkArgs args) { void BM_UpdateKeyValue(::benchmark::State& state, BenchmarkArgs args) { uint seed = args.concurrent_tasks; std::vector reader_tasks; + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); if (state.thread_index() == 0 && args.concurrent_tasks) { auto num_readers = args.concurrent_tasks; reader_tasks.reserve(num_readers); while (num_readers-- > 0) { - reader_tasks.emplace_back([args, &seed]() { + reader_tasks.emplace_back([args, &seed, &request_context]() { auto key = std::to_string(rand_r(&seed) % args.keyspace_size); args.cache->GetKeyValuePairs( - absl::flat_hash_set({key})); + request_context, absl::flat_hash_set({key})); }); } } @@ -232,13 +238,15 @@ void BM_UpdateKeyValue(::benchmark::State& state, BenchmarkArgs args) { void BM_UpdateKeyValueSet(::benchmark::State& state, BenchmarkArgs args) { uint seed = args.concurrent_tasks; std::vector reader_tasks; + auto scope_metrics_context = std::make_unique(); + RequestContext request_context(*scope_metrics_context); if (state.thread_index() == 0 && args.concurrent_tasks) { auto num_readers = args.concurrent_tasks; reader_tasks.reserve(num_readers); while (num_readers-- > 0) { - reader_tasks.emplace_back([args, &seed]() { + reader_tasks.emplace_back([args, &seed, &request_context]() { auto key = std::to_string(rand_r(&seed) % args.keyspace_size); - args.cache->GetKeyValueSet({key}); + args.cache->GetKeyValueSet(request_context, {key}); }); } } @@ -267,7 +275,7 @@ void RegisterBenchmark( } } -void RegisterReadBenchmarks(MetricsRecorder& metrics_recorder) { +void RegisterReadBenchmarks() { auto query_sizes = ParseInt64List(absl::GetFlag(FLAGS_query_size)); auto set_query_sizes = ParseInt64List(absl::GetFlag(FLAGS_set_query_size)); auto record_sizes = ParseInt64List(absl::GetFlag(FLAGS_record_size)); @@ -286,7 +294,7 @@ void RegisterReadBenchmarks(MetricsRecorder& metrics_recorder) { absl::StrFormat(kNoOpCacheGetKeyValuePairsFmt, query_size, record_size, num_writers), args, BM_GetKeyValuePairs); - args.cache = GetLockBasedCache(metrics_recorder); + args.cache = GetLockBasedCache(); ::kv_server::RegisterBenchmark( absl::StrFormat(kLockBasedCacheGetKeyValuePairsFmt, query_size, record_size, num_writers), @@ -298,7 +306,7 @@ void RegisterReadBenchmarks(MetricsRecorder& metrics_recorder) { absl::StrFormat(kNoOpCacheGetKeyValueSetFmt, query_size, set_query_size, record_size, num_writers), args, BM_GetKeyValueSet); - args.cache = GetLockBasedCache(metrics_recorder); + args.cache = GetLockBasedCache(); ::kv_server::RegisterBenchmark( absl::StrFormat(kLockBasedCacheGetKeyValueSetFmt, query_size, set_query_size, record_size, num_writers), @@ -309,7 +317,7 @@ void RegisterReadBenchmarks(MetricsRecorder& metrics_recorder) { } } -void RegisterWriteBenchmarks(MetricsRecorder& metrics_recorder) { +void RegisterWriteBenchmarks() { auto keyspace_sizes = ParseInt64List(absl::GetFlag(FLAGS_keyspace_size)); auto record_sizes = ParseInt64List(absl::GetFlag(FLAGS_record_size)); auto set_query_sizes = ParseInt64List(absl::GetFlag(FLAGS_set_query_size)); @@ -328,7 +336,7 @@ void RegisterWriteBenchmarks(MetricsRecorder& metrics_recorder) { absl::StrFormat(kNoOpCacheUpdateKeyValueFmt, keyspace_size, record_size, num_readers), args, BM_UpdateKeyValue); - args.cache = GetLockBasedCache(metrics_recorder); + args.cache = GetLockBasedCache(); ::kv_server::RegisterBenchmark( absl::StrFormat(kLockBasedCacheUpdateKeyValueFmt, keyspace_size, record_size, num_readers), @@ -340,7 +348,7 @@ void RegisterWriteBenchmarks(MetricsRecorder& metrics_recorder) { absl::StrFormat(kNoOpCacheUpdateKeyValueSetFmt, keyspace_size, set_query_size, record_size, num_readers), args, BM_UpdateKeyValueSet); - args.cache = GetLockBasedCache(metrics_recorder); + args.cache = GetLockBasedCache(); ::kv_server::RegisterBenchmark( absl::StrFormat(kLockBasedCacheUpdateKeyValueSetFmt, keyspace_size, set_query_size, record_size, @@ -357,19 +365,18 @@ void RegisterWriteBenchmarks(MetricsRecorder& metrics_recorder) { // Microbenchmarks for Cache impelementations. Sample run: // -// GLOG_logtostderr=1 bazel run -c opt \ +// bazel run -c opt \ // //components/tools/benchmarks:cache_benchmark \ // --//:instance=local \ // --//:platform=local -- \ -// --benchmark_counters_tabular=true +// --benchmark_counters_tabular=true --stderrthreshold=0 int main(int argc, char** argv) { - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); ::benchmark::Initialize(&argc, argv); absl::ParseCommandLine(argc, argv); - auto noop_metrics_recorder = - ::kv_server::TelemetryProvider::GetInstance().CreateMetricsRecorder(); - ::kv_server::RegisterReadBenchmarks(*noop_metrics_recorder); - ::kv_server::RegisterWriteBenchmarks(*noop_metrics_recorder); + kv_server::InitMetricsContextMap(); + ::kv_server::RegisterReadBenchmarks(); + ::kv_server::RegisterWriteBenchmarks(); ::benchmark::RunSpecifiedBenchmarks(); ::benchmark::Shutdown(); return 0; diff --git a/components/tools/benchmarks/data_loading_benchmark.cc b/components/tools/benchmarks/data_loading_benchmark.cc index 1f71ebca..cf95b0ea 100644 --- a/components/tools/benchmarks/data_loading_benchmark.cc +++ b/components/tools/benchmarks/data_loading_benchmark.cc @@ -23,6 +23,9 @@ #include "absl/container/flat_hash_map.h" #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/time/time.h" @@ -33,11 +36,9 @@ #include "components/data_server/cache/noop_key_value_cache.h" #include "components/tools/benchmarks/benchmark_util.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" #include "public/data_loading/data_loading_generated.h" #include "public/data_loading/readers/riegeli_stream_io.h" #include "public/data_loading/records_utils.h" -#include "src/cpp/telemetry/telemetry_provider.h" ABSL_FLAG(std::string, data_directory, "", "Data directory or bucket to store benchmark input data files in."); @@ -83,8 +84,6 @@ using kv_server::RecordStream; using kv_server::Value; using kv_server::benchmark::ParseInt64List; using kv_server::benchmark::WriteRecords; -using privacy_sandbox::server_common::MetricsRecorder; -using privacy_sandbox::server_common::TelemetryProvider; constexpr std::string_view kNoOpCacheNameFormat = "BM_DataLoading_NoOpCache/tds:%d/conns:%d/buf:%d"; @@ -169,11 +168,7 @@ void RegisterBenchmarks() { RegisterBenchmark(absl::StrFormat(kNoOpCacheNameFormat, num_threads, num_connections, byte_range_mb), args); - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - args.create_cache_fn = [&noop_metrics_recorder]() { - return KeyValueCache::Create(*noop_metrics_recorder); - }; + args.create_cache_fn = []() { return KeyValueCache::Create(); }; RegisterBenchmark(absl::StrFormat(kMutexCacheNameFormat, num_threads, num_connections, byte_range_mb), args); @@ -283,7 +278,7 @@ void BM_LoadDataIntoCache(benchmark::State& state, BenchmarkArgs args) { // Sample usage: // -// GLOG_logtostderr=1 bazel run \ +// bazel run \ // components/tools/benchmarks:data_loading_benchmark \ // --//:instance=local --//:platform=local -- \ // --benchmark_time_unit=ms \ @@ -295,10 +290,10 @@ void BM_LoadDataIntoCache(benchmark::State& state, BenchmarkArgs args) { // --record_size=1000 \ // --args_client_max_range_mb=8 \ // --args_client_max_connections=64 \ -// --args_reader_worker_threads=16,32,64 +// --args_reader_worker_threads=16,32,64 --stderrthreshold=0 int main(int argc, char** argv) { ::kv_server::PlatformInitializer platform_initializer; - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); ::benchmark::Initialize(&argc, argv); absl::ParseCommandLine(argc, argv); if (absl::GetFlag(FLAGS_data_directory).empty()) { diff --git a/components/tools/blob_storage_change_watcher_aws.cc b/components/tools/blob_storage_change_watcher_aws.cc index 559f88a4..48477cce 100644 --- a/components/tools/blob_storage_change_watcher_aws.cc +++ b/components/tools/blob_storage_change_watcher_aws.cc @@ -20,7 +20,7 @@ #include "components/data/blob_storage/blob_storage_change_notifier.h" #include "components/telemetry/server_definition.h" #include "components/util/platform_initializer.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" ABSL_FLAG(std::string, sns_arn, "", "sns_arn"); diff --git a/components/tools/blob_storage_commands.cc b/components/tools/blob_storage_commands.cc index b660a507..c555eb0c 100644 --- a/components/tools/blob_storage_commands.cc +++ b/components/tools/blob_storage_commands.cc @@ -24,7 +24,7 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" #include "components/data/blob_storage/blob_storage_client.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" namespace kv_server { namespace blob_storage_commands { @@ -55,12 +55,16 @@ class StdinBlobReader : public BlobReader { using kv_server::BlobStorageClient; -bool CatObjects(std::string bucket_or_directory, absl::Span keys) { +bool CatObjects(std::string bucket_or_directory, std::string prefix, + absl::Span keys) { std::unique_ptr blob_storage_client_factory = BlobStorageClientFactory::Create(); std::unique_ptr client = blob_storage_client_factory->CreateBlobStorageClient(); - BlobStorageClient::DataLocation location = {std::move(bucket_or_directory)}; + BlobStorageClient::DataLocation location = { + .bucket = std::move(bucket_or_directory), + .prefix = std::move(prefix), + }; for (const auto& key : keys) { location.key = key; auto reader = client->GetBlobReader(location); @@ -69,12 +73,16 @@ bool CatObjects(std::string bucket_or_directory, absl::Span keys) { return true; } -bool DeleteObjects(std::string bucket_or_directory, absl::Span keys) { +bool DeleteObjects(std::string bucket_or_directory, std::string prefix, + absl::Span keys) { std::unique_ptr blob_storage_client_factory = BlobStorageClientFactory::Create(); std::unique_ptr client = blob_storage_client_factory->CreateBlobStorageClient(); - BlobStorageClient::DataLocation location = {std::move(bucket_or_directory)}; + BlobStorageClient::DataLocation location = { + .bucket = std::move(bucket_or_directory), + .prefix = std::move(prefix), + }; for (const auto& key : keys) { location.key = key; const absl::Status status = client->DeleteBlob(location); @@ -86,13 +94,15 @@ bool DeleteObjects(std::string bucket_or_directory, absl::Span keys) { return true; } -bool ListObjects(std::string bucket_or_directory) { +bool ListObjects(std::string bucket_or_directory, std::string prefix) { std::unique_ptr blob_storage_client_factory = BlobStorageClientFactory::Create(); std::unique_ptr client = blob_storage_client_factory->CreateBlobStorageClient(); const BlobStorageClient::DataLocation location = { - std::move(bucket_or_directory)}; + .bucket = std::move(bucket_or_directory), + .prefix = std::move(prefix), + }; const absl::StatusOr> keys = client->ListBlobs(location, {}); if (!keys.ok()) { diff --git a/components/tools/blob_storage_commands.h b/components/tools/blob_storage_commands.h index d3339cff..e61e48dc 100644 --- a/components/tools/blob_storage_commands.h +++ b/components/tools/blob_storage_commands.h @@ -28,13 +28,15 @@ namespace kv_server { namespace blob_storage_commands { // Print the contents of blobs to stdout; returns success. -bool CatObjects(std::string bucket_or_directory, absl::Span keys); +bool CatObjects(std::string bucket_or_directory, std::string prefix, + absl::Span keys); // Delete the blobs; returns success. -bool DeleteObjects(std::string bucket_or_directory, absl::Span keys); +bool DeleteObjects(std::string bucket_or_directory, std::string prefix, + absl::Span keys); // Print to stdout a list of blobs that are in this location; returns success. -bool ListObjects(std::string bucket_or_directory); +bool ListObjects(std::string bucket_or_directory, std::string prefix); // Create a new BlobReader that will read from either a file or stdin if the // source is "-". diff --git a/components/tools/blob_storage_util_aws.cc b/components/tools/blob_storage_util.cc similarity index 91% rename from components/tools/blob_storage_util_aws.cc rename to components/tools/blob_storage_util.cc index 1db257e1..4ee70581 100644 --- a/components/tools/blob_storage_util_aws.cc +++ b/components/tools/blob_storage_util.cc @@ -21,9 +21,10 @@ #include "components/data/blob_storage/blob_storage_client.h" #include "components/tools/blob_storage_commands.h" #include "components/util/platform_initializer.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" ABSL_FLAG(std::string, bucket, "", "cloud storage bucket name"); +ABSL_FLAG(std::string, prefix, "", "object prefix name"); using kv_server::BlobReader; using kv_server::BlobStorageClient; @@ -102,13 +103,15 @@ int main(int argc, char** argv) { std::cerr << "Must specify bucket" << std::endl; return -1; } + std::string prefix = absl::GetFlag(FLAGS_prefix); absl::string_view operation = commands[1]; if (operation == "ls") { if (commands.size() != 2) { std::cerr << "ls does not take any extra arguments." << std::endl; return -1; } - return kv_server::blob_storage_commands::ListObjects(std::move(bucket)) + return kv_server::blob_storage_commands::ListObjects(std::move(bucket), + std::move(prefix)) ? 0 : -1; } @@ -118,7 +121,8 @@ int main(int argc, char** argv) { return -1; } return kv_server::blob_storage_commands::DeleteObjects( - std::move(bucket), absl::MakeSpan(commands).subspan(2)) + std::move(bucket), std::move(prefix), + absl::MakeSpan(commands).subspan(2)) ? 0 : -1; } @@ -128,7 +132,8 @@ int main(int argc, char** argv) { return -1; } return kv_server::blob_storage_commands::CatObjects( - std::move(bucket), absl::MakeSpan(commands).subspan(2)) + std::move(bucket), std::move(prefix), + absl::MakeSpan(commands).subspan(2)) ? 0 : -1; } diff --git a/components/tools/concurrent_publishing_engine.cc b/components/tools/concurrent_publishing_engine.cc index 07591b29..3994e8f1 100644 --- a/components/tools/concurrent_publishing_engine.cc +++ b/components/tools/concurrent_publishing_engine.cc @@ -37,14 +37,24 @@ void ConcurrentPublishingEngine::Start() { } void ConcurrentPublishingEngine::Stop() { + { + absl::MutexLock l(&mutex_); + stop_ = true; + } for (auto& publisher_thread : publishers_) { publisher_thread->join(); } } -std::optional ConcurrentPublishingEngine::Pop() { - absl::MutexLock lock(&mutex_); - if (updates_queue_.empty()) { +bool ConcurrentPublishingEngine::HasNewMessageToProcess() const { + return !updates_queue_.empty() || stop_; +} + +std::optional ConcurrentPublishingEngine::Pop( + absl::Condition& has_new_event) { + absl::MutexLock lock(&mutex_, has_new_event); + if (stop_) { + LOG(INFO) << "Thread for new file processing stopped"; return std::nullopt; } auto file_encoded = updates_queue_.front(); @@ -52,20 +62,31 @@ std::optional ConcurrentPublishingEngine::Pop() { return file_encoded; } +bool ConcurrentPublishingEngine::ShouldStop() { + absl::MutexLock l(&mutex_); + return stop_; +} + void ConcurrentPublishingEngine::ConsumeAndPublish(int thread_idx) { const auto start = clock_.Now(); int delta_file_index = 1; auto batch_end = clock_.Now() + absl::Seconds(1); - auto message = Pop(); auto maybe_msg_service = PublisherService::Create(notifier_metadata_); if (!maybe_msg_service.ok()) { LOG(ERROR) << "Failed creating a publisher service"; return; } auto msg_service = std::move(*maybe_msg_service); - while (message.has_value()) { - LOG(INFO) << ": Inserting to the SNS: " << delta_file_index - << " Thread idx " << thread_idx; + absl::Condition has_new_event( + this, &ConcurrentPublishingEngine::HasNewMessageToProcess); + + while (!ShouldStop()) { + auto message = Pop(has_new_event); + if (!message.has_value()) { + return; + } + VLOG(9) << ": Inserting to the SNS: " << delta_file_index << " Thread idx " + << thread_idx; auto status = msg_service->Publish(message->message, message->shard_num); if (!status.ok()) { LOG(ERROR) << status; @@ -79,7 +100,6 @@ void ConcurrentPublishingEngine::ConsumeAndPublish(int thread_idx) { } batch_end += absl::Seconds(1); } - message = Pop(); } int64_t elapsed_seconds = absl::ToInt64Seconds(clock_.Now() - start); diff --git a/components/tools/concurrent_publishing_engine.h b/components/tools/concurrent_publishing_engine.h index 317971ba..0785f7c5 100644 --- a/components/tools/concurrent_publishing_engine.h +++ b/components/tools/concurrent_publishing_engine.h @@ -21,9 +21,9 @@ #include #include +#include "absl/log/log.h" #include "components/tools/publisher_service.h" -#include "glog/logging.h" -#include "src/cpp/util/duration.h" +#include "src/util/duration.h" namespace kv_server { @@ -56,13 +56,16 @@ class ConcurrentPublishingEngine { delete; private: - std::optional Pop(); + std::optional Pop(absl::Condition& has_new_event); + bool ShouldStop(); void ConsumeAndPublish(int thread_idx); const int insertion_num_threads_; const NotifierMetadata notifier_metadata_; const int files_insertion_rate_; absl::Mutex& mutex_; + bool stop_ ABSL_GUARDED_BY(mutex_) = false; + bool HasNewMessageToProcess() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); std::queue& updates_queue_ ABSL_GUARDED_BY(mutex_); std::vector> publishers_; privacy_sandbox::server_common::SteadyClock& clock_ = diff --git a/components/tools/data_loading_analyzer.cc b/components/tools/data_loading_analyzer.cc index 334f8942..53ff75a0 100644 --- a/components/tools/data_loading_analyzer.cc +++ b/components/tools/data_loading_analyzer.cc @@ -19,6 +19,7 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/log.h" #include "absl/strings/str_cat.h" #include "components/data/blob_storage/blob_storage_client.h" #include "components/data/blob_storage/delta_file_notifier.h" @@ -27,13 +28,11 @@ #include "components/data_server/data_loading/data_orchestrator.h" #include "components/udf/noop_udf_client.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" #include "public/base_types.pb.h" #include "public/data_loading/data_loading_generated.h" #include "public/data_loading/readers/riegeli_stream_io.h" #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" #include "public/sharding/key_sharder.h" -#include "src/cpp/telemetry/telemetry_provider.h" ABSL_FLAG(std::vector, operations, std::vector({"PASS_THROUGH", "READ_ONLY", "CACHE"}), @@ -44,10 +43,6 @@ ABSL_FLAG(std::string, bucket, "performance-test-data-bucket", namespace kv_server { namespace { -using privacy_sandbox::server_common::GetTracer; -using privacy_sandbox::server_common::MetricsRecorder; -using privacy_sandbox::server_common::TelemetryProvider; - class NoopBlobStorageChangeNotifier : public BlobStorageChangeNotifier { public: absl::StatusOr> GetNotifications( @@ -146,9 +141,7 @@ std::vector OperationsFromFlag() { absl::Status InitOnce(Operation operation) { std::unique_ptr noop_udf_client = NewNoopUdfClient(); InitMetricsContextMap(); - auto noop_metrics_recorder = - TelemetryProvider::GetInstance().CreateMetricsRecorder(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); + std::unique_ptr cache = KeyValueCache::Create(); std::unique_ptr blob_storage_client_factory = BlobStorageClientFactory::Create(); diff --git a/components/tools/delta_file_record_change_watcher.cc b/components/tools/delta_file_record_change_watcher.cc index 0e209af9..de01a605 100644 --- a/components/tools/delta_file_record_change_watcher.cc +++ b/components/tools/delta_file_record_change_watcher.cc @@ -26,7 +26,7 @@ #include "public/data_loading/filename_utils.h" #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" #include "public/data_loading/records_utils.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" ABSL_FLAG(std::string, sns_arn, "", "sns_arn"); diff --git a/components/tools/delta_file_watcher_aws.cc b/components/tools/delta_file_watcher_aws.cc index bf4cea44..9a951a00 100644 --- a/components/tools/delta_file_watcher_aws.cc +++ b/components/tools/delta_file_watcher_aws.cc @@ -21,7 +21,7 @@ #include "components/data/blob_storage/delta_file_notifier.h" #include "components/data/common/thread_manager.h" #include "components/util/platform_initializer.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" ABSL_FLAG(std::string, bucket, "", "cloud storage bucket name"); ABSL_FLAG(std::string, sns_arn, "", "sns_arn"); @@ -74,7 +74,8 @@ int main(int argc, char** argv) { } const absl::Status status = notifier->Start( - **status_or_change_notifier, {.bucket = std::move(bucket)}, "", + **status_or_change_notifier, {.bucket = std::move(bucket)}, + /*prefix_start_after_map=*/{std::make_pair("", "")}, [](const std::string& key) { std::cout << key << std::endl; }); if (!status.ok()) { std::cerr << "Failed to start notifier: " << status << std::endl; diff --git a/components/tools/get_region_aws.cc b/components/tools/get_region_aws.cc index bee8805b..561bdd15 100644 --- a/components/tools/get_region_aws.cc +++ b/components/tools/get_region_aws.cc @@ -20,6 +20,7 @@ #include "absl/cleanup/cleanup.h" #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" @@ -30,7 +31,6 @@ #include "aws/core/internal/AWSHttpResourceClient.h" #include "aws/core/utils/Outcome.h" #include "components/errors/error_util_aws.h" -#include "glog/logging.h" ABSL_FLAG(std::string, output_file, "", "output_file"); diff --git a/components/tools/get_region_local.cc b/components/tools/get_region_local.cc index 7c86ffc8..176a4bd3 100644 --- a/components/tools/get_region_local.cc +++ b/components/tools/get_region_local.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "glog/logging.h" +#include "absl/log/log.h" int main(int argc, char** argv) { LOG(FATAL) << "The get_region tool is not available for --platform==local"; diff --git a/components/tools/publisher_service_local.cc b/components/tools/publisher_service_local.cc new file mode 100644 index 00000000..1783fde1 --- /dev/null +++ b/components/tools/publisher_service_local.cc @@ -0,0 +1,48 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "components/tools/publisher_service.h" + +namespace kv_server { +namespace { + +class LocalPublisherService : public PublisherService { + public: + LocalPublisherService() {} + + absl::Status Publish(const std::string& body, std::optional shard_num) { + return absl::UnimplementedError( + "PublisherService is not implemented for local"); + } + + absl::StatusOr BuildNotifierMetadataAndSetQueue() { + return absl::UnimplementedError( + "PublisherService is not implemented for local"); + } +}; + +} // namespace + +absl::StatusOr> PublisherService::Create( + NotifierMetadata notifier_metadata) { + return std::make_unique(); +} + +absl::StatusOr PublisherService::GetNotifierMetadata() { + return absl::UnimplementedError( + "PublisherService is not implemented for local"); +} +} // namespace kv_server diff --git a/components/tools/realtime_notifier.cc b/components/tools/realtime_notifier.cc index a46e76ee..65667625 100644 --- a/components/tools/realtime_notifier.cc +++ b/components/tools/realtime_notifier.cc @@ -19,6 +19,8 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" #include "absl/flags/usage.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" #include "absl/strings/str_join.h" #include "components/data/common/msg_svc.h" #include "components/tools/publisher_service.h" @@ -27,7 +29,7 @@ #include "public/data_loading/filename_utils.h" #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" #include "public/data_loading/records_utils.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" ABSL_FLAG(std::string, local_directory, "", "Local directory"); @@ -122,7 +124,7 @@ absl::Status Run() { int main(int argc, char* argv[]) { const std::vector commands = absl::ParseCommandLine(argc, argv); - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); const absl::Status status = kv_server::Run(); if (!status.ok()) { LOG(FATAL) << "Failed to run: " << status; diff --git a/components/tools/realtime_updates_publisher.cc b/components/tools/realtime_updates_publisher.cc index 45f8c2d3..65854d17 100644 --- a/components/tools/realtime_updates_publisher.cc +++ b/components/tools/realtime_updates_publisher.cc @@ -19,12 +19,14 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "absl/strings/substitute.h" #include "components/data/common/msg_svc.h" #include "components/tools/concurrent_publishing_engine.h" #include "components/tools/publisher_service.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" ABSL_FLAG(std::string, deltas_folder_path, "", "Path to the folder with delta files"); @@ -93,7 +95,7 @@ absl::Status Run() { // `tools/serving_data_generator/generate_load_test_data`. int main(int argc, char** argv) { const std::vector commands = absl::ParseCommandLine(argc, argv); - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); const absl::Status status = kv_server::Run(); if (!status.ok()) { LOG(FATAL) << "Failed to run: " << status; diff --git a/components/tools/sharding_correctness_validator/BUILD.bazel b/components/tools/sharding_correctness_validator/BUILD.bazel index 6a3d0091..160abb3d 100644 --- a/components/tools/sharding_correctness_validator/BUILD.bazel +++ b/components/tools/sharding_correctness_validator/BUILD.bazel @@ -20,9 +20,11 @@ cc_binary( deps = [ "//public/applications/pa:response_utils", "//public/query/cpp:grpc_client", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/random", "@com_google_absl//absl/strings", ], diff --git a/components/tools/sharding_correctness_validator/validator.cc b/components/tools/sharding_correctness_validator/validator.cc index e0755045..5299bf79 100644 --- a/components/tools/sharding_correctness_validator/validator.cc +++ b/components/tools/sharding_correctness_validator/validator.cc @@ -16,10 +16,12 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "absl/random/random.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" -#include "glog/logging.h" #include "public/applications/pa/response_utils.h" #include "public/query/cpp/grpc_client.h" @@ -30,6 +32,7 @@ ABSL_FLAG(int, number_of_requests_to_make, 1, "Number of requests to make"); ABSL_FLAG(int, value_size, 10000, "Specify the size of value for the key"); ABSL_FLAG(int, batch_size, 10, "Batch size"); ABSL_FLAG(std::string, key_prefix, "foo", "Key prefix"); +ABSL_FLAG(bool, use_tls, false, "Whether to use TLS for grpc calls."); namespace kv_server { namespace { @@ -121,11 +124,18 @@ void Validate() { const int qps = absl::GetFlag(FLAGS_qps); const int number_of_requests_to_make = absl::GetFlag(FLAGS_number_of_requests_to_make); + const bool use_tls = absl::GetFlag(FLAGS_use_tls); int requests_made_this_second = 0; int total_requests_made = 0; auto batch_end = absl::Now() + absl::Seconds(1); - std::unique_ptr stub = - GrpcClient::CreateStub(kv_endpoint, grpc::InsecureChannelCredentials()); + std::unique_ptr stub; + if (use_tls) { + stub = GrpcClient::CreateStub( + kv_endpoint, grpc::SslCredentials(grpc::SslCredentialsOptions())); + } else { + stub = + GrpcClient::CreateStub(kv_endpoint, grpc::InsecureChannelCredentials()); + } GrpcClient client(*stub); while (total_requests_made < number_of_requests_to_make) { auto random_index = Get(inclusive_upper_bound / batch_size); @@ -168,7 +178,7 @@ void Validate() { int main(int argc, char** argv) { const std::vector commands = absl::ParseCommandLine(argc, argv); - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); kv_server::Validate(); if (kv_server::total_failures > 0 || kv_server::total_mismatches > 0) { diff --git a/components/udf/BUILD.bazel b/components/udf/BUILD.bazel index 630b048a..5f7e7525 100644 --- a/components/udf/BUILD.bazel +++ b/components/udf/BUILD.bazel @@ -48,8 +48,8 @@ cc_library( "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//scp/cc/roma/interface:roma_interface_lib", - "@google_privacysandbox_servers_common//scp/cc/roma/roma_service:roma_service_lib", + "@google_privacysandbox_servers_common//src/roma/interface", + "@google_privacysandbox_servers_common//src/roma/roma_service", ], ) @@ -83,8 +83,8 @@ cc_library( "//components/udf/hooks:get_values_hook", "//components/udf/hooks:logging_hook", "//components/udf/hooks:run_query_hook", - "@google_privacysandbox_servers_common//scp/cc/roma/interface:roma_interface_lib", - "@google_privacysandbox_servers_common//scp/cc/roma/roma_service:roma_service_lib", + "@google_privacysandbox_servers_common//src/roma/interface", + "@google_privacysandbox_servers_common//src/roma/roma_service", ], ) @@ -105,12 +105,15 @@ cc_test( "//components/internal_server:mocks", "//components/udf/hooks:get_values_hook", "//components/udf/hooks:run_query_hook", + "//public/query/v2:get_values_v2_cc_proto", "//public/test_util:proto_matcher", + "//public/udf:constants", + "@com_google_absl//absl/log:scoped_mock_log", "@com_google_absl//absl/status", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//scp/cc/roma/interface:roma_interface_lib", - "@google_privacysandbox_servers_common//scp/cc/roma/roma_service:roma_service_lib", + "@google_privacysandbox_servers_common//src/roma/interface", + "@google_privacysandbox_servers_common//src/roma/roma_service", ], ) @@ -123,6 +126,6 @@ cc_library( ":udf_client", "@com_google_absl//absl/status", "@com_google_googletest//:gtest", - "@google_privacysandbox_servers_common//scp/cc/roma/interface:roma_interface_lib", + "@google_privacysandbox_servers_common//src/roma/interface", ], ) diff --git a/components/udf/hooks/BUILD.bazel b/components/udf/hooks/BUILD.bazel index b18d74b8..c3a37094 100644 --- a/components/udf/hooks/BUILD.bazel +++ b/components/udf/hooks/BUILD.bazel @@ -31,15 +31,14 @@ cc_library( "//components/internal_server:internal_lookup_cc_proto", "//components/internal_server:local_lookup", "//components/internal_server:lookup", + "//components/util:request_context", "//public/udf:binary_get_values_cc_proto", - "@com_github_google_glog//:glog", "@com_google_absl//absl/functional:any_invocable", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//scp/cc/roma/interface:roma_interface_lib", - "@google_privacysandbox_servers_common//src/cpp/telemetry", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", + "@google_privacysandbox_servers_common//src/roma/interface", "@nlohmann_json//:lib", ], ) @@ -55,12 +54,13 @@ cc_library( deps = [ "//components/internal_server:internal_lookup_cc_proto", "//components/internal_server:lookup", - "@com_github_google_glog//:glog", + "//components/util:request_context", "@com_google_absl//absl/functional:any_invocable", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//scp/cc/roma/interface:roma_function_binding_io_cc_proto", - "@google_privacysandbox_servers_common//scp/cc/roma/interface:roma_interface_lib", + "@google_privacysandbox_servers_common//src/roma/interface", + "@google_privacysandbox_servers_common//src/roma/interface:function_binding_io_cc_proto", "@nlohmann_json//:lib", ], ) @@ -71,7 +71,8 @@ cc_library( "logging_hook.h", ], deps = [ - "@com_github_google_glog//:glog", + "//components/util:request_context", + "@com_google_absl//absl/log", ], ) @@ -88,8 +89,6 @@ cc_test( "@com_google_absl//absl/status", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", ], ) diff --git a/components/udf/hooks/get_values_hook.cc b/components/udf/hooks/get_values_hook.cc index 8d356071..e35fa682 100644 --- a/components/udf/hooks/get_values_hook.cc +++ b/components/udf/hooks/get_values_hook.cc @@ -21,16 +21,16 @@ #include #include "absl/functional/any_invocable.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/strings/str_join.h" #include "components/data_server/cache/cache.h" #include "components/internal_server/local_lookup.h" #include "components/internal_server/lookup.pb.h" -#include "glog/logging.h" +#include "components/telemetry/server_definition.h" #include "google/protobuf/util/json_util.h" #include "nlohmann/json.hpp" #include "public/udf/binary_get_values.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" namespace kv_server { namespace { @@ -38,7 +38,6 @@ namespace { using google::protobuf::json::MessageToJsonString; using google::scp::roma::FunctionBindingPayload; using google::scp::roma::proto::FunctionBindingIoProto; -using privacy_sandbox::server_common::MetricsRecorder; constexpr char kOkStatusMessage[] = "ok"; @@ -129,7 +128,7 @@ class GetValuesHookImpl : public GetValuesHook { } } - void operator()(FunctionBindingPayload<>& payload) { + void operator()(FunctionBindingPayload& payload) { VLOG(9) << "Called getValues hook"; if (lookup_ == nullptr) { SetStatus(absl::StatusCode::kInternal, @@ -154,7 +153,7 @@ class GetValuesHookImpl : public GetValuesHook { VLOG(9) << "Calling internal lookup client"; absl::StatusOr response_or_status = - lookup_->GetKeyValues(keys); + lookup_->GetKeyValues(payload.metadata, keys); if (!response_or_status.ok()) { SetStatus(response_or_status.status().code(), response_or_status.status().message(), payload.io_proto); diff --git a/components/udf/hooks/get_values_hook.h b/components/udf/hooks/get_values_hook.h index dba7a7e7..24d576ac 100644 --- a/components/udf/hooks/get_values_hook.h +++ b/components/udf/hooks/get_values_hook.h @@ -25,13 +25,11 @@ #include "absl/functional/any_invocable.h" #include "components/data_server/cache/cache.h" #include "components/internal_server/lookup.h" -#include "roma/config/src/function_binding_object_v2.h" -#include "src/cpp/telemetry/metrics_recorder.h" +#include "components/util/request_context.h" +#include "src/roma/config/function_binding_object_v2.h" namespace kv_server { -using privacy_sandbox::server_common::MetricsRecorder; - // Functor that acts as a wrapper for the internal lookup client call. class GetValuesHook { public: @@ -48,7 +46,7 @@ class GetValuesHook { // This is registered with v8 and is exposed to the UDF. Internally, it calls // the internal lookup client. virtual void operator()( - google::scp::roma::FunctionBindingPayload<>& payload) = 0; + google::scp::roma::FunctionBindingPayload& payload) = 0; static std::unique_ptr Create(OutputType output_type); }; diff --git a/components/udf/hooks/get_values_hook_test.cc b/components/udf/hooks/get_values_hook_test.cc index b74f8d12..3d6f2b03 100644 --- a/components/udf/hooks/get_values_hook_test.cc +++ b/components/udf/hooks/get_values_hook_test.cc @@ -38,7 +38,12 @@ using google::scp::roma::proto::FunctionBindingIoProto; using testing::_; using testing::Return; -TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesValue) { +class GetValuesHookTest : public ::testing::Test { + protected: + void SetUp() override { InitMetricsContextMap(); } +}; + +TEST_F(GetValuesHookTest, StringOutput_SuccessfullyProcessesValue) { absl::flat_hash_set keys = {"key1", "key2"}; InternalLookupResponse lookup_response; TextFormat::ParseFromString(R"pb(kv_pairs { @@ -51,7 +56,7 @@ TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesValue) { })pb", &lookup_response); auto mock_lookup = std::make_unique(); - EXPECT_CALL(*mock_lookup, GetKeyValues(keys)) + EXPECT_CALL(*mock_lookup, GetKeyValues(_, keys)) .WillOnce(Return(lookup_response)); FunctionBindingIoProto io; @@ -60,7 +65,9 @@ TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesValue) { auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kString); get_values_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + FunctionBindingPayload payload{ + io, RequestContext(metrics_context)}; (*get_values_hook)(payload); nlohmann::json result_json = @@ -77,7 +84,7 @@ TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesValue) { EXPECT_EQ(result_json["status"]["message"], "ok"); } -TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesResultsWithStatus) { +TEST_F(GetValuesHookTest, StringOutput_SuccessfullyProcessesResultsWithStatus) { absl::flat_hash_set keys = {"key1"}; InternalLookupResponse lookup_response; TextFormat::ParseFromString( @@ -88,7 +95,7 @@ TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesResultsWithStatus) { &lookup_response); auto mock_lookup = std::make_unique(); - EXPECT_CALL(*mock_lookup, GetKeyValues(keys)) + EXPECT_CALL(*mock_lookup, GetKeyValues(_, keys)) .WillOnce(Return(lookup_response)); FunctionBindingIoProto io; @@ -97,7 +104,9 @@ TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesResultsWithStatus) { auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kString); get_values_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + FunctionBindingPayload payload{ + io, RequestContext(metrics_context)}; (*get_values_hook)(payload); nlohmann::json expected = @@ -105,10 +114,10 @@ TEST(GetValuesHookTest, StringOutput_SuccessfullyProcessesResultsWithStatus) { EXPECT_EQ(io.output_string(), expected.dump()); } -TEST(GetValuesHookTest, StringOutput_LookupReturnsError) { +TEST_F(GetValuesHookTest, StringOutput_LookupReturnsError) { absl::flat_hash_set keys = {"key1"}; auto mock_lookup = std::make_unique(); - EXPECT_CALL(*mock_lookup, GetKeyValues(keys)) + EXPECT_CALL(*mock_lookup, GetKeyValues(_, keys)) .WillOnce(Return(absl::UnknownError("Some error"))); FunctionBindingIoProto io; @@ -117,14 +126,16 @@ TEST(GetValuesHookTest, StringOutput_LookupReturnsError) { auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kString); get_values_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + FunctionBindingPayload payload{ + io, RequestContext(metrics_context)}; (*get_values_hook)(payload); nlohmann::json expected = R"({"code":2,"message":"Some error"})"_json; EXPECT_EQ(io.output_string(), expected.dump()); } -TEST(GetValuesHookTest, StringOutput_InputIsNotListOfStrings) { +TEST_F(GetValuesHookTest, StringOutput_InputIsNotListOfStrings) { absl::flat_hash_set keys = {"key1"}; auto mock_lookup = std::make_unique(); @@ -133,7 +144,9 @@ TEST(GetValuesHookTest, StringOutput_InputIsNotListOfStrings) { auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kString); get_values_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + FunctionBindingPayload payload{ + io, RequestContext(metrics_context)}; (*get_values_hook)(payload); nlohmann::json expected = @@ -141,7 +154,7 @@ TEST(GetValuesHookTest, StringOutput_InputIsNotListOfStrings) { EXPECT_EQ(io.output_string(), expected.dump()); } -TEST(GetValuesHookTest, BinaryOutput_SuccessfullyProcessesValue) { +TEST_F(GetValuesHookTest, BinaryOutput_SuccessfullyProcessesValue) { absl::flat_hash_set keys = {"key1", "key2"}; InternalLookupResponse lookup_response; TextFormat::ParseFromString( @@ -155,7 +168,7 @@ TEST(GetValuesHookTest, BinaryOutput_SuccessfullyProcessesValue) { })pb", &lookup_response); auto mock_lookup = std::make_unique(); - EXPECT_CALL(*mock_lookup, GetKeyValues(keys)) + EXPECT_CALL(*mock_lookup, GetKeyValues(_, keys)) .WillOnce(Return(lookup_response)); FunctionBindingIoProto io; @@ -164,7 +177,9 @@ TEST(GetValuesHookTest, BinaryOutput_SuccessfullyProcessesValue) { auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kBinary); get_values_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + FunctionBindingPayload payload{ + io, RequestContext(metrics_context)}; (*get_values_hook)(payload); EXPECT_TRUE(io.has_output_bytes()); @@ -187,10 +202,10 @@ TEST(GetValuesHookTest, BinaryOutput_SuccessfullyProcessesValue) { EXPECT_THAT(response.kv_pairs().at("key2"), EqualsProto(value_with_status)); } -TEST(GetValuesHookTest, BinaryOutput_LookupReturnsError) { +TEST_F(GetValuesHookTest, BinaryOutput_LookupReturnsError) { absl::flat_hash_set keys = {"key1"}; auto mock_lookup = std::make_unique(); - EXPECT_CALL(*mock_lookup, GetKeyValues(keys)) + EXPECT_CALL(*mock_lookup, GetKeyValues(_, keys)) .WillOnce(Return(absl::UnknownError("Some error"))); FunctionBindingIoProto io; @@ -199,7 +214,9 @@ TEST(GetValuesHookTest, BinaryOutput_LookupReturnsError) { auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kBinary); get_values_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + FunctionBindingPayload payload{ + io, RequestContext(metrics_context)}; (*get_values_hook)(payload); EXPECT_TRUE(io.has_output_bytes()); diff --git a/components/udf/hooks/logging_hook.h b/components/udf/hooks/logging_hook.h index 296c358e..3e541600 100644 --- a/components/udf/hooks/logging_hook.h +++ b/components/udf/hooks/logging_hook.h @@ -20,17 +20,16 @@ #include #include -#include "glog/logging.h" +#include "absl/log/log.h" +#include "components/util/request_context.h" namespace kv_server { -// UDF hook for logging a string. -// TODO(b/285331079): Disable for production builds. -inline void LogMessage(google::scp::roma::FunctionBindingPayload<>& payload) { - if (payload.io_proto.has_input_string()) { - LOG(INFO) << payload.io_proto.input_string(); - } - payload.io_proto.set_output_string(""); +// Logging function to register with Roma. +inline void LoggingFunction(absl::LogSeverity severity, + const RequestContext& context, + std::string_view msg) { + LOG(LEVEL(severity)) << msg; } } // namespace kv_server diff --git a/components/udf/hooks/run_query_hook.cc b/components/udf/hooks/run_query_hook.cc index f878ced7..c2aa113f 100644 --- a/components/udf/hooks/run_query_hook.cc +++ b/components/udf/hooks/run_query_hook.cc @@ -21,9 +21,9 @@ #include #include "absl/functional/any_invocable.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "components/internal_server/lookup.h" -#include "glog/logging.h" #include "nlohmann/json.hpp" namespace kv_server { @@ -39,7 +39,7 @@ class RunQueryHookImpl : public RunQueryHook { } } - void operator()(FunctionBindingPayload<>& payload) { + void operator()(FunctionBindingPayload& payload) { if (lookup_ == nullptr) { nlohmann::json status; status["code"] = absl::StatusCode::kInternal; @@ -62,7 +62,7 @@ class RunQueryHookImpl : public RunQueryHook { VLOG(9) << "Calling internal run query client"; absl::StatusOr response_or_status = - lookup_->RunQuery(payload.io_proto.input_string()); + lookup_->RunQuery(payload.metadata, payload.io_proto.input_string()); if (!response_or_status.ok()) { LOG(ERROR) << "Internal run query returned error: " diff --git a/components/udf/hooks/run_query_hook.h b/components/udf/hooks/run_query_hook.h index 1ab871c8..35560c3e 100644 --- a/components/udf/hooks/run_query_hook.h +++ b/components/udf/hooks/run_query_hook.h @@ -24,7 +24,8 @@ #include "absl/functional/any_invocable.h" #include "components/internal_server/lookup.h" -#include "roma/config/src/function_binding_object_v2.h" +#include "components/util/request_context.h" +#include "src/roma/config/function_binding_object_v2.h" namespace kv_server { @@ -42,7 +43,7 @@ class RunQueryHook { // This is registered with v8 and is exposed to the UDF. Internally, it calls // the internal query client. virtual void operator()( - google::scp::roma::FunctionBindingPayload<>& payload) = 0; + google::scp::roma::FunctionBindingPayload& payload) = 0; static std::unique_ptr Create(); }; diff --git a/components/udf/hooks/run_query_hook_test.cc b/components/udf/hooks/run_query_hook_test.cc index d1a1d54d..985513de 100644 --- a/components/udf/hooks/run_query_hook_test.cc +++ b/components/udf/hooks/run_query_hook_test.cc @@ -36,41 +36,50 @@ using testing::_; using testing::Return; using testing::UnorderedElementsAreArray; -TEST(RunQueryHookTest, SuccessfullyProcessesValue) { +class RunQueryHookTest : public ::testing::Test { + protected: + void SetUp() override { InitMetricsContextMap(); } +}; + +TEST_F(RunQueryHookTest, SuccessfullyProcessesValue) { std::string query = "Q"; InternalRunQueryResponse run_query_response; TextFormat::ParseFromString(R"pb(elements: "a" elements: "b")pb", &run_query_response); auto mock_lookup = std::make_unique(); - EXPECT_CALL(*mock_lookup, RunQuery(query)) + EXPECT_CALL(*mock_lookup, RunQuery(_, query)) .WillOnce(Return(run_query_response)); FunctionBindingIoProto io; TextFormat::ParseFromString(R"pb(input_string: "Q")pb", &io); auto run_query_hook = RunQueryHook::Create(); run_query_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + RequestContext request_context(metrics_context); + FunctionBindingPayload payload{io, request_context}; (*run_query_hook)(payload); EXPECT_THAT(io.output_list_of_string().data(), UnorderedElementsAreArray({"a", "b"})); } -TEST(GetValuesHookTest, RunQueryClientReturnsError) { +TEST_F(RunQueryHookTest, RunQueryClientReturnsError) { std::string query = "Q"; auto mock_lookup = std::make_unique(); - EXPECT_CALL(*mock_lookup, RunQuery(query)) + EXPECT_CALL(*mock_lookup, RunQuery(_, query)) .WillOnce(Return(absl::UnknownError("Some error"))); FunctionBindingIoProto io; TextFormat::ParseFromString(R"pb(input_string: "Q")pb", &io); auto run_query_hook = RunQueryHook::Create(); run_query_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + RequestContext request_context(metrics_context); + FunctionBindingPayload payload{io, request_context}; (*run_query_hook)(payload); EXPECT_TRUE(io.output_list_of_string().data().empty()); } -TEST(GetValuesHookTest, InputIsNotString) { +TEST_F(RunQueryHookTest, InputIsNotString) { auto mock_lookup = std::make_unique(); FunctionBindingIoProto io; @@ -78,7 +87,9 @@ TEST(GetValuesHookTest, InputIsNotString) { &io); auto run_query_hook = RunQueryHook::Create(); run_query_hook->FinishInit(std::move(mock_lookup)); - FunctionBindingPayload<> payload{io, {}}; + ScopeMetricsContext metrics_context; + RequestContext request_context(metrics_context); + FunctionBindingPayload payload{io, request_context}; (*run_query_hook)(payload); EXPECT_THAT( diff --git a/components/udf/mocks.h b/components/udf/mocks.h index cba13db7..8e8bbbc5 100644 --- a/components/udf/mocks.h +++ b/components/udf/mocks.h @@ -25,16 +25,16 @@ #include "components/udf/code_config.h" #include "components/udf/udf_client.h" #include "gmock/gmock.h" -#include "roma/interface/roma.h" +#include "src/roma/interface/roma.h" namespace kv_server { class MockUdfClient : public UdfClient { public: MOCK_METHOD((absl::StatusOr), ExecuteCode, - (std::vector), (const, override)); + (RequestContext, std::vector), (const, override)); MOCK_METHOD((absl::StatusOr), ExecuteCode, - (UDFExecutionMetadata&&, + (RequestContext, UDFExecutionMetadata&&, const google::protobuf::RepeatedPtrField&), (const, override)); MOCK_METHOD((absl::Status), Stop, (), (override)); diff --git a/components/udf/noop_udf_client.cc b/components/udf/noop_udf_client.cc index ceacba4f..e257a7e7 100644 --- a/components/udf/noop_udf_client.cc +++ b/components/udf/noop_udf_client.cc @@ -24,18 +24,19 @@ #include "absl/status/statusor.h" #include "components/udf/code_config.h" #include "components/udf/udf_client.h" -#include "roma/config/src/config.h" +#include "src/roma/config/config.h" namespace kv_server { namespace { class NoopUdfClientImpl : public UdfClient { public: - absl::StatusOr ExecuteCode(std::vector keys) const { + absl::StatusOr ExecuteCode(RequestContext request_context, + std::vector keys) const { return ""; } absl::StatusOr ExecuteCode( - UDFExecutionMetadata&&, + RequestContext request_context, UDFExecutionMetadata&&, const google::protobuf::RepeatedPtrField& arguments) const { return ""; } diff --git a/components/udf/udf_client.cc b/components/udf/udf_client.cc index d8ecb289..8569c7f5 100644 --- a/components/udf/udf_client.cc +++ b/components/udf/udf_client.cc @@ -22,15 +22,15 @@ #include #include "absl/flags/flag.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/synchronization/notification.h" #include "absl/time/time.h" -#include "glog/logging.h" #include "google/protobuf/util/json_util.h" -#include "roma/config/src/config.h" -#include "roma/interface/roma.h" -#include "roma/roma_service/roma_service.h" +#include "src/roma/config/config.h" +#include "src/roma/interface/roma.h" +#include "src/roma/roma_service/roma_service.h" namespace kv_server { @@ -54,13 +54,16 @@ constexpr int kUdfInterfaceVersion = 1; class UdfClientImpl : public UdfClient { public: - explicit UdfClientImpl(Config<>&& config = Config(), - absl::Duration udf_timeout = absl::Seconds(5)) - : udf_timeout_(udf_timeout), roma_service_(std::move(config)) {} + explicit UdfClientImpl( + Config&& config = Config(), + absl::Duration udf_timeout = absl::Seconds(5), int udf_min_log_level = 0) + : udf_timeout_(udf_timeout), + roma_service_(std::move(config)), + udf_min_log_level_(udf_min_log_level) {} // Converts the arguments into plain JSON strings to pass to Roma. absl::StatusOr ExecuteCode( - UDFExecutionMetadata&& execution_metadata, + RequestContext request_context, UDFExecutionMetadata&& execution_metadata, const google::protobuf::RepeatedPtrField& arguments) const { execution_metadata.set_udf_interface_version(kUdfInterfaceVersion); std::vector string_args; @@ -88,27 +91,29 @@ class UdfClientImpl : public UdfClient { } string_args.push_back(json_arg); } - return ExecuteCode(std::move(string_args)); + return ExecuteCode(std::move(request_context), std::move(string_args)); } - absl::StatusOr ExecuteCode(std::vector keys) const { + absl::StatusOr ExecuteCode( + RequestContext request_context, std::vector input) const { std::shared_ptr response_status = std::make_shared(); std::shared_ptr result = std::make_shared(); std::shared_ptr notification = std::make_shared(); - InvocationStrRequest<> invocation_request = - BuildInvocationRequest(std::move(keys)); - VLOG(9) << "Executing UDF"; + auto invocation_request = + BuildInvocationRequest(std::move(request_context), std::move(input)); + VLOG(9) << "Executing UDF with input arg(s): " + << absl::StrJoin(invocation_request.input, ","); const auto status = roma_service_.Execute( - std::make_unique>(invocation_request), + std::make_unique>( + std::move(invocation_request)), [notification, response_status, - result](std::unique_ptr> response) { - if (response->ok()) { - auto& code_response = **response; - *result = std::move(code_response.resp); + result](absl::StatusOr response) { + if (response.ok()) { + *result = std::move(response->resp); } else { - response_status->Update(std::move(response->status())); + response_status->Update(std::move(response.status())); } notification->Notify(); }); @@ -149,10 +154,9 @@ class UdfClientImpl : public UdfClient { code_config.version); absl::Status load_status = roma_service_.LoadCodeObj( std::make_unique(code_object), - [notification, response_status]( - std::unique_ptr> resp) { - if (!resp->ok()) { - response_status->Update(std::move(resp->status())); + [notification, response_status](absl::StatusOr resp) { + if (!resp.ok()) { + response_status->Update(std::move(resp.status())); } notification->Notify(); }); @@ -166,12 +170,14 @@ class UdfClientImpl : public UdfClient { return absl::InternalError("Timed out setting UDF code object."); } if (!response_status->ok()) { - LOG(ERROR) << "Error setting UDF Code object: " << *response_status; + LOG(ERROR) << "Error compiling UDF code object. " << *response_status; return *response_status; } handler_name_ = std::move(code_config.udf_handler_name); logical_commit_time_ = code_config.logical_commit_time; version_ = code_config.version; + VLOG(5) << "Successfully set UDF code object with handler_name " + << handler_name_; return absl::OkStatus(); } @@ -184,13 +190,16 @@ class UdfClientImpl : public UdfClient { } private: - InvocationStrRequest<> BuildInvocationRequest( - std::vector keys) const { + InvocationStrRequest BuildInvocationRequest( + RequestContext request_context, std::vector input) const { return {.id = kInvocationRequestId, .version_string = absl::StrCat("v", version_), .handler_name = handler_name_, - .tags = {{kTimeoutDurationTag, FormatDuration(udf_timeout_)}}, - .input = std::move(keys)}; + .tags = {{std::string(kTimeoutDurationTag), + FormatDuration(udf_timeout_)}}, + .input = std::move(input), + .metadata = std::move(request_context), + .min_log_level = absl::LogSeverity(udf_min_log_level_)}; } CodeObject BuildCodeObject(std::string js, std::string wasm, @@ -205,6 +214,7 @@ class UdfClientImpl : public UdfClient { int64_t logical_commit_time_ = -1; int64_t version_ = 1; const absl::Duration udf_timeout_; + int udf_min_log_level_; // Per b/299667930, RomaService has been extended to support metadata storage // as a side effect of RomaService::Execute(), making it no longer const. // However, UDFClient::ExecuteCode() remains logically const, so RomaService @@ -212,15 +222,16 @@ class UdfClientImpl : public UdfClient { // concerns about mutable or go/totw/174, RomaService is thread-safe, so // losing the thread-safety of usage within a const function is a lesser // concern. - mutable RomaService<> roma_service_; + mutable RomaService roma_service_; }; } // namespace absl::StatusOr> UdfClient::Create( - Config<>&& config, absl::Duration udf_timeout) { - auto udf_client = - std::make_unique(std::move(config), udf_timeout); + Config&& config, absl::Duration udf_timeout, + int udf_min_log_level) { + auto udf_client = std::make_unique( + std::move(config), udf_timeout, udf_min_log_level); const auto init_status = udf_client->Init(); if (!init_status.ok()) { return init_status; diff --git a/components/udf/udf_client.h b/components/udf/udf_client.h index 99d274fb..abb9e8bd 100644 --- a/components/udf/udf_client.h +++ b/components/udf/udf_client.h @@ -23,11 +23,13 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "components/telemetry/server_definition.h" #include "components/udf/code_config.h" +#include "components/util/request_context.h" #include "google/protobuf/message.h" #include "public/api_schema.pb.h" -#include "roma/config/src/config.h" -#include "roma/interface/roma.h" +#include "src/roma/config/config.h" +#include "src/roma/interface/roma.h" namespace kv_server { @@ -41,12 +43,12 @@ class UdfClient { // UDF signature. ABSL_DEPRECATED("Use ExecuteCode(metadata, arguments) instead") virtual absl::StatusOr ExecuteCode( - std::vector keys) const = 0; + RequestContext request_context, std::vector keys) const = 0; // Executes the UDF. Code object must be set before making // this call. virtual absl::StatusOr ExecuteCode( - UDFExecutionMetadata&& execution_metadata, + RequestContext request_context, UDFExecutionMetadata&& execution_metadata, const google::protobuf::RepeatedPtrField& arguments) const = 0; @@ -60,8 +62,9 @@ class UdfClient { // Creates a UDF executor. This calls Roma::Init, which forks. static absl::StatusOr> Create( - google::scp::roma::Config<>&& config = google::scp::roma::Config(), - absl::Duration udf_timeout = absl::Seconds(5)); + google::scp::roma::Config&& config = + google::scp::roma::Config(), + absl::Duration udf_timeout = absl::Seconds(5), int udf_min_log_level = 0); }; } // namespace kv_server diff --git a/components/udf/udf_client_test.cc b/components/udf/udf_client_test.cc index 1b193ea0..755d5a00 100644 --- a/components/udf/udf_client_test.cc +++ b/components/udf/udf_client_test.cc @@ -20,6 +20,7 @@ #include #include +#include "absl/log/scoped_mock_log.h" #include "absl/status/statusor.h" #include "components/internal_server/mocks.h" #include "components/udf/code_config.h" @@ -30,14 +31,15 @@ #include "gmock/gmock.h" #include "google/protobuf/text_format.h" #include "gtest/gtest.h" -#include "roma/config/src/config.h" -#include "roma/interface/roma.h" +#include "public/query/v2/get_values_v2.pb.h" +#include "public/udf/constants.h" +#include "src/roma/config/config.h" +#include "src/roma/interface/roma.h" using google::protobuf::TextFormat; using google::scp::roma::Config; using google::scp::roma::FunctionBindingObjectV2; using google::scp::roma::FunctionBindingPayload; -using google::scp::roma::WasmDataType; using testing::_; using testing::Return; @@ -45,19 +47,24 @@ namespace kv_server { namespace { absl::StatusOr> CreateUdfClient() { - Config config; + Config config; config.number_of_workers = 1; return UdfClient::Create(std::move(config)); } -TEST(UdfClientTest, UdfClient_Create_Success) { +class UdfClientTest : public ::testing::Test { + protected: + void SetUp() override { InitMetricsContextMap(); } +}; + +TEST_F(UdfClientTest, UdfClient_Create_Success) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); absl::Status stop = udf_client.value()->Stop(); EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsCallSucceeds) { +TEST_F(UdfClientTest, JsCallSucceeds) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -68,8 +75,9 @@ TEST(UdfClientTest, JsCallSucceeds) { .version = 1, }); EXPECT_TRUE(code_obj_status.ok()); - - absl::StatusOr result = udf_client.value()->ExecuteCode({}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = + udf_client.value()->ExecuteCode(RequestContext(metrics_context), {}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Hello world!")"); @@ -77,7 +85,7 @@ TEST(UdfClientTest, JsCallSucceeds) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, RepeatedJsCallsSucceed) { +TEST_F(UdfClientTest, RepeatedJsCallsSucceed) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -88,12 +96,14 @@ TEST(UdfClientTest, RepeatedJsCallsSucceed) { .version = 1, }); EXPECT_TRUE(code_obj_status.ok()); - - absl::StatusOr result1 = udf_client.value()->ExecuteCode({}); + ScopeMetricsContext metrics_context; + absl::StatusOr result1 = + udf_client.value()->ExecuteCode(RequestContext(metrics_context), {}); EXPECT_TRUE(result1.ok()); EXPECT_EQ(*result1, R"("Hello world!")"); - absl::StatusOr result2 = udf_client.value()->ExecuteCode({}); + absl::StatusOr result2 = + udf_client.value()->ExecuteCode(RequestContext(metrics_context), {}); EXPECT_TRUE(result2.ok()); EXPECT_EQ(*result2, R"("Hello world!")"); @@ -101,7 +111,7 @@ TEST(UdfClientTest, RepeatedJsCallsSucceed) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsEchoCallSucceeds) { +TEST_F(UdfClientTest, JsEchoCallSucceeds) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -112,9 +122,9 @@ TEST(UdfClientTest, JsEchoCallSucceeds) { .version = 1, }); EXPECT_TRUE(code_obj_status.ok()); - - absl::StatusOr result = - udf_client.value()->ExecuteCode({R"("ECHO")"}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {R"("ECHO")"}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Hello world! \"ECHO\"")"); @@ -122,7 +132,7 @@ TEST(UdfClientTest, JsEchoCallSucceeds) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string) { +TEST_F(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -141,8 +151,9 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string) { arg.mutable_data()->set_string_value("ECHO"); return arg; }()); - absl::StatusOr result = - udf_client.value()->ExecuteCode({}, args); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {}, args); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Hello world! \"ECHO\"")"); @@ -150,7 +161,7 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged) { +TEST_F(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -170,8 +181,9 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged) { arg.mutable_data()->set_string_value("ECHO"); return arg; }()); - absl::StatusOr result = - udf_client.value()->ExecuteCode({}, args); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {}, args); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Hello world! {\"tags\":[\"tag1\"],\"data\":\"ECHO\"}")"); @@ -179,7 +191,7 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged) { absl::Status stop = udf_client.value()->Stop(); EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged_list) { +TEST_F(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged_list) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -201,8 +213,9 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged_list) { list_value->add_values()->set_string_value("key2"); return arg; }()); - absl::StatusOr result = - udf_client.value()->ExecuteCode({}, args); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {}, args); EXPECT_TRUE(result.ok()); EXPECT_EQ( *result, @@ -212,7 +225,7 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_string_tagged_list) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_struct) { +TEST_F(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_struct) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -232,8 +245,9 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_struct) { .set_string_value("value"); return arg; }()); - absl::StatusOr result = - udf_client.value()->ExecuteCode({}, args); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {}, args); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Hello world! {\"key\":\"value\"}")"); @@ -241,20 +255,20 @@ TEST(UdfClientTest, JsEchoCallSucceeds_SimpleUDFArg_struct) { EXPECT_TRUE(stop.ok()); } -static void udfCbEcho(FunctionBindingPayload<>& payload) { +static void udfCbEcho(FunctionBindingPayload& payload) { payload.io_proto.set_output_string("Echo: " + payload.io_proto.input_string()); } -TEST(UdfClientTest, JsEchoHookCallSucceeds) { - auto function_object = std::make_unique>(); +TEST_F(UdfClientTest, JsEchoHookCallSucceeds) { + auto function_object = + std::make_unique>(); function_object->function_name = "echo"; function_object->function = udfCbEcho; - Config config; + Config config; config.number_of_workers = 1; config.RegisterFunctionBinding(std::move(function_object)); - absl::StatusOr> udf_client = UdfClient::Create(std::move(config)); EXPECT_TRUE(udf_client.ok()); @@ -266,9 +280,9 @@ TEST(UdfClientTest, JsEchoHookCallSucceeds) { .version = 1, }); EXPECT_TRUE(code_obj_status.ok()); - - absl::StatusOr result = - udf_client.value()->ExecuteCode({R"("I'm a key")"}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {R"("I'm a key")"}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Hello world! Echo: I'm a key")"); @@ -276,7 +290,7 @@ TEST(UdfClientTest, JsEchoHookCallSucceeds) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsStringInWithGetValuesHookSucceeds) { +TEST_F(UdfClientTest, JsStringInWithGetValuesHookSucceeds) { auto mock_lookup = std::make_unique(); InternalLookupResponse response; @@ -285,7 +299,7 @@ TEST(UdfClientTest, JsStringInWithGetValuesHookSucceeds) { value { value: "value1" } })pb", &response); - ON_CALL(*mock_lookup, GetKeyValues(_)).WillByDefault(Return(response)); + ON_CALL(*mock_lookup, GetKeyValues(_, _)).WillByDefault(Return(response)); auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kString); @@ -315,9 +329,9 @@ TEST(UdfClientTest, JsStringInWithGetValuesHookSucceeds) { .version = 1, }); EXPECT_TRUE(code_obj_status.ok()); - - absl::StatusOr result = - udf_client.value()->ExecuteCode({R"("key1")"}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {R"("key1")"}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Key: key1, Value: value1")"); @@ -325,7 +339,7 @@ TEST(UdfClientTest, JsStringInWithGetValuesHookSucceeds) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsJSONObjectInWithGetValuesHookSucceeds) { +TEST_F(UdfClientTest, JsJSONObjectInWithGetValuesHookSucceeds) { auto mock_lookup = std::make_unique(); InternalLookupResponse response; @@ -334,7 +348,7 @@ TEST(UdfClientTest, JsJSONObjectInWithGetValuesHookSucceeds) { value { value: "value1" } })pb", &response); - ON_CALL(*mock_lookup, GetKeyValues(_)).WillByDefault(Return(response)); + ON_CALL(*mock_lookup, GetKeyValues(_, _)).WillByDefault(Return(response)); auto get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kString); @@ -366,9 +380,9 @@ TEST(UdfClientTest, JsJSONObjectInWithGetValuesHookSucceeds) { .version = 1, }); EXPECT_TRUE(code_obj_status.ok()); - - absl::StatusOr result = - udf_client.value()->ExecuteCode({R"({"keys":["key1"]})"}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {R"({"keys":["key1"]})"}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("Key: key1, Value: value1")"); @@ -376,12 +390,12 @@ TEST(UdfClientTest, JsJSONObjectInWithGetValuesHookSucceeds) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsJSONObjectInWithRunQueryHookSucceeds) { +TEST_F(UdfClientTest, JsJSONObjectInWithRunQueryHookSucceeds) { auto mock_lookup = std::make_unique(); InternalRunQueryResponse response; TextFormat::ParseFromString(R"pb(elements: "a")pb", &response); - ON_CALL(*mock_lookup, RunQuery(_)).WillByDefault(Return(response)); + ON_CALL(*mock_lookup, RunQuery(_, _)).WillByDefault(Return(response)); auto run_query_hook = RunQueryHook::Create(); run_query_hook->FinishInit(std::move(mock_lookup)); @@ -405,9 +419,9 @@ TEST(UdfClientTest, JsJSONObjectInWithRunQueryHookSucceeds) { .version = 1, }); EXPECT_TRUE(code_obj_status.ok()); - - absl::StatusOr result = - udf_client.value()->ExecuteCode({R"({"keys":["key1"]})"}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {R"({"keys":["key1"]})"}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"(["a"])"); @@ -415,19 +429,21 @@ TEST(UdfClientTest, JsJSONObjectInWithRunQueryHookSucceeds) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, JsCallsLogMessageTwiceSucceeds) { +TEST_F(UdfClientTest, JsCallsLoggingFunctionSucceeds) { UdfConfigBuilder config_builder; absl::StatusOr> udf_client = - UdfClient::Create(std::move( - config_builder.RegisterLoggingHook().SetNumberOfWorkers(1).Config())); + UdfClient::Create(std::move(config_builder.RegisterLoggingFunction() + .SetNumberOfWorkers(1) + .Config())); EXPECT_TRUE(udf_client.ok()); absl::Status code_obj_status = udf_client.value()->SetCodeObject(CodeConfig{ .js = R"( function hello(input) { - const a = logMessage("first message"); - const b = logMessage("second message"); - return a + b; + const a = console.error("Error message"); + const b = console.warn("Warning message"); + const c = console.log("Info message"); + return ""; } )", .udf_handler_name = "hello", @@ -436,16 +452,26 @@ TEST(UdfClientTest, JsCallsLogMessageTwiceSucceeds) { }); EXPECT_TRUE(code_obj_status.ok()); - absl::StatusOr result = - udf_client.value()->ExecuteCode({R"({"keys":["key1"]})"}); + absl::ScopedMockLog log; + EXPECT_CALL(log, Log(absl::LogSeverity::kError, testing::_, "Error message")); + EXPECT_CALL(log, + Log(absl::LogSeverity::kWarning, testing::_, "Warning message")); + EXPECT_CALL(log, Log(absl::LogSeverity::kInfo, testing::_, "Info message")); + log.StartCapturingLogs(); + + ScopeMetricsContext metrics_context; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), {R"({"keys":["key1"]})"}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("")"); + log.StopCapturingLogs(); + absl::Status stop = udf_client.value()->Stop(); EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, UpdatesCodeObjectTwice) { +TEST_F(UdfClientTest, UpdatesCodeObjectTwice) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -464,8 +490,9 @@ TEST(UdfClientTest, UpdatesCodeObjectTwice) { .version = 2, }); EXPECT_TRUE(status.ok()); - - absl::StatusOr result = udf_client.value()->ExecuteCode({}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = + udf_client.value()->ExecuteCode(RequestContext(metrics_context), {}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("2")"); @@ -473,7 +500,7 @@ TEST(UdfClientTest, UpdatesCodeObjectTwice) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, IgnoresCodeObjectWithSameCommitTime) { +TEST_F(UdfClientTest, IgnoresCodeObjectWithSameCommitTime) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -492,8 +519,9 @@ TEST(UdfClientTest, IgnoresCodeObjectWithSameCommitTime) { .version = 1, }); EXPECT_TRUE(status.ok()); - - absl::StatusOr result = udf_client.value()->ExecuteCode({}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = + udf_client.value()->ExecuteCode(RequestContext(metrics_context), {}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("1")"); @@ -501,7 +529,7 @@ TEST(UdfClientTest, IgnoresCodeObjectWithSameCommitTime) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, IgnoresCodeObjectWithSmallerCommitTime) { +TEST_F(UdfClientTest, IgnoresCodeObjectWithSmallerCommitTime) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); @@ -520,8 +548,9 @@ TEST(UdfClientTest, IgnoresCodeObjectWithSmallerCommitTime) { .version = 1, }); EXPECT_TRUE(status.ok()); - - absl::StatusOr result = udf_client.value()->ExecuteCode({}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = + udf_client.value()->ExecuteCode(RequestContext(metrics_context), {}); EXPECT_TRUE(result.ok()); EXPECT_EQ(*result, R"("1")"); @@ -529,11 +558,12 @@ TEST(UdfClientTest, IgnoresCodeObjectWithSmallerCommitTime) { EXPECT_TRUE(stop.ok()); } -TEST(UdfClientTest, CodeObjectNotSetError) { +TEST_F(UdfClientTest, CodeObjectNotSetError) { auto udf_client = CreateUdfClient(); EXPECT_TRUE(udf_client.ok()); - - absl::StatusOr result = udf_client.value()->ExecuteCode({}); + ScopeMetricsContext metrics_context; + absl::StatusOr result = + udf_client.value()->ExecuteCode(RequestContext(metrics_context), {}); EXPECT_FALSE(result.ok()); EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); @@ -541,6 +571,200 @@ TEST(UdfClientTest, CodeObjectNotSetError) { EXPECT_TRUE(stop.ok()); } +TEST_F(UdfClientTest, MetadataPassedSuccesfully) { + UdfConfigBuilder config_builder; + absl::StatusOr> udf_client = UdfClient::Create( + std::move(config_builder.SetNumberOfWorkers(1).Config())); + EXPECT_TRUE(udf_client.ok()); + absl::Status code_obj_status = udf_client.value()->SetCodeObject(CodeConfig{ + .js = R"( + function hello(metadata) { + if(metadata.requestMetadata && + metadata.requestMetadata.is_pas) + { + return "true"; + } + return "false"; + } + )", + .udf_handler_name = "hello", + .logical_commit_time = 1, + .version = 1, + }); + EXPECT_TRUE(code_obj_status.ok()); + v2::GetValuesRequest req; + (*(req.mutable_metadata()->mutable_fields()))["is_pas"].set_string_value( + "true"); + UDFExecutionMetadata udf_metadata; + *udf_metadata.mutable_request_metadata() = *req.mutable_metadata(); + ScopeMetricsContext metrics_context; + google::protobuf::RepeatedPtrField args; + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), std::move(udf_metadata), args); + EXPECT_TRUE(result.ok()); + EXPECT_EQ(*result, R"("true")"); + + UDFExecutionMetadata udf_metadata_non_pas; + result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), std::move(udf_metadata_non_pas), args); + EXPECT_TRUE(result.ok()); + EXPECT_EQ(*result, R"("false")"); + absl::Status stop = udf_client.value()->Stop(); + EXPECT_TRUE(stop.ok()); +} + +TEST_F(UdfClientTest, DefaultUdfPASucceeds) { + auto mock_lookup = std::make_unique(); + InternalLookupResponse response; + TextFormat::ParseFromString(R"pb(kv_pairs { + key: "key1" + value { value: "value1" } + })pb", + &response); + ON_CALL(*mock_lookup, GetKeyValues(_, _)).WillByDefault(Return(response)); + + auto get_values_hook = + GetValuesHook::Create(GetValuesHook::OutputType::kString); + get_values_hook->FinishInit(std::move(mock_lookup)); + UdfConfigBuilder config_builder; + absl::StatusOr> udf_client = UdfClient::Create( + std::move(config_builder.RegisterStringGetValuesHook(*get_values_hook) + .SetNumberOfWorkers(1) + .Config())); + EXPECT_TRUE(udf_client.ok()); + absl::Status code_obj_status = udf_client.value()->SetCodeObject(CodeConfig{ + .js = kDefaultUdfCodeSnippet, + .udf_handler_name = kDefaultUdfHandlerName, + .logical_commit_time = kDefaultLogicalCommitTime, + .version = kDefaultVersion, + }); + EXPECT_TRUE(code_obj_status.ok()); + ScopeMetricsContext metrics_context; + UDFExecutionMetadata udf_metadata; + google::protobuf::RepeatedPtrField args; + args.Add([] { + UDFArgument arg; + TextFormat::ParseFromString(R"( + data { + list_value { + values { + string_value: "key1" + } + } + })", + &arg); + return arg; + }()); + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), std::move(udf_metadata), args); + EXPECT_TRUE(result.ok()); + EXPECT_EQ( + *result, + R"({"keyGroupOutputs":[{"keyValues":{"key1":{"value":"value1"}}}],"udfOutputApiVersion":1})"); + absl::Status stop = udf_client.value()->Stop(); + EXPECT_TRUE(stop.ok()); +} + +TEST_F(UdfClientTest, DefaultUdfPasKeyLookupFails) { + auto mock_lookup = std::make_unique(); + absl::Status status = absl::InvalidArgumentError("Error!"); + ON_CALL(*mock_lookup, GetKeyValues(_, _)).WillByDefault(Return(status)); + auto get_values_hook = + GetValuesHook::Create(GetValuesHook::OutputType::kString); + get_values_hook->FinishInit(std::move(mock_lookup)); + UdfConfigBuilder config_builder; + absl::StatusOr> udf_client = UdfClient::Create( + std::move(config_builder.RegisterStringGetValuesHook(*get_values_hook) + .SetNumberOfWorkers(1) + .Config())); + EXPECT_TRUE(udf_client.ok()); + absl::Status code_obj_status = udf_client.value()->SetCodeObject(CodeConfig{ + .js = kDefaultUdfCodeSnippet, + .udf_handler_name = kDefaultUdfHandlerName, + .logical_commit_time = kDefaultLogicalCommitTime, + .version = kDefaultVersion, + }); + EXPECT_TRUE(code_obj_status.ok()); + ScopeMetricsContext metrics_context; + v2::GetValuesRequest req; + (*(req.mutable_metadata()->mutable_fields()))["is_pas"].set_string_value( + "true"); + UDFExecutionMetadata udf_metadata; + *udf_metadata.mutable_request_metadata() = *req.mutable_metadata(); + google::protobuf::RepeatedPtrField args; + args.Add([] { + UDFArgument arg; + TextFormat::ParseFromString(R"( + data { + list_value { + values { + string_value: "key1" + } + } + })", + &arg); + return arg; + }()); + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), std::move(udf_metadata), args); + EXPECT_FALSE(result.ok()); + absl::Status stop = udf_client.value()->Stop(); + EXPECT_TRUE(stop.ok()); +} + +TEST_F(UdfClientTest, DefaultUdfPasSucceeds) { + auto mock_lookup = std::make_unique(); + InternalLookupResponse response; + TextFormat::ParseFromString(R"pb(kv_pairs { + key: "key1" + value { value: "value1" } + })pb", + &response); + ON_CALL(*mock_lookup, GetKeyValues(_, _)).WillByDefault(Return(response)); + auto get_values_hook = + GetValuesHook::Create(GetValuesHook::OutputType::kString); + get_values_hook->FinishInit(std::move(mock_lookup)); + UdfConfigBuilder config_builder; + absl::StatusOr> udf_client = UdfClient::Create( + std::move(config_builder.RegisterStringGetValuesHook(*get_values_hook) + .SetNumberOfWorkers(1) + .Config())); + EXPECT_TRUE(udf_client.ok()); + absl::Status code_obj_status = udf_client.value()->SetCodeObject(CodeConfig{ + .js = kDefaultUdfCodeSnippet, + .udf_handler_name = kDefaultUdfHandlerName, + .logical_commit_time = kDefaultLogicalCommitTime, + .version = kDefaultVersion, + }); + EXPECT_TRUE(code_obj_status.ok()); + ScopeMetricsContext metrics_context; + v2::GetValuesRequest req; + (*(req.mutable_metadata()->mutable_fields()))["is_pas"].set_string_value( + "true"); + UDFExecutionMetadata udf_metadata; + *udf_metadata.mutable_request_metadata() = *req.mutable_metadata(); + google::protobuf::RepeatedPtrField args; + args.Add([] { + UDFArgument arg; + TextFormat::ParseFromString(R"( + data { + list_value { + values { + string_value: "key1" + } + } + })", + &arg); + return arg; + }()); + absl::StatusOr result = udf_client.value()->ExecuteCode( + RequestContext(metrics_context), std::move(udf_metadata), args); + EXPECT_TRUE(result.ok()); + EXPECT_EQ(*result, R"({"key1":{"value":"value1"}})"); + absl::Status stop = udf_client.value()->Stop(); + EXPECT_TRUE(stop.ok()); +} + } // namespace } // namespace kv_server diff --git a/components/udf/udf_config_builder.cc b/components/udf/udf_config_builder.cc index 42178634..e6bc3fbe 100644 --- a/components/udf/udf_config_builder.cc +++ b/components/udf/udf_config_builder.cc @@ -24,9 +24,9 @@ #include "components/udf/hooks/get_values_hook.h" #include "components/udf/hooks/logging_hook.h" #include "components/udf/hooks/run_query_hook.h" -#include "roma/config/src/config.h" -#include "roma/config/src/function_binding_object_v2.h" -#include "roma/interface/roma.h" +#include "src/roma/config/config.h" +#include "src/roma/config/function_binding_object_v2.h" +#include "src/roma/interface/roma.h" namespace kv_server { namespace { @@ -38,15 +38,17 @@ using google::scp::roma::FunctionBindingPayload; constexpr char kStringGetValuesHookJsName[] = "getValues"; constexpr char kBinaryGetValuesHookJsName[] = "getValuesBinary"; constexpr char kRunQueryHookJsName[] = "runQuery"; -constexpr char kLoggingHookJsName[] = "logMessage"; -std::unique_ptr> GetValuesFunctionObject( - GetValuesHook& get_values_hook, std::string handler_name) { +std::unique_ptr> +GetValuesFunctionObject(GetValuesHook& get_values_hook, + std::string handler_name) { auto get_values_function_object = - std::make_unique>(); + std::make_unique>(); get_values_function_object->function_name = std::move(handler_name); get_values_function_object->function = - [&get_values_hook](FunctionBindingPayload<>& in) { get_values_hook(in); }; + [&get_values_hook](FunctionBindingPayload& in) { + get_values_hook(in); + }; return get_values_function_object; } @@ -69,19 +71,18 @@ UdfConfigBuilder& UdfConfigBuilder::RegisterBinaryGetValuesHook( UdfConfigBuilder& UdfConfigBuilder::RegisterRunQueryHook( RunQueryHook& run_query_hook) { auto run_query_function_object = - std::make_unique>(); + std::make_unique>(); run_query_function_object->function_name = kRunQueryHookJsName; run_query_function_object->function = - [&run_query_hook](FunctionBindingPayload<>& in) { run_query_hook(in); }; + [&run_query_hook](FunctionBindingPayload& in) { + run_query_hook(in); + }; config_.RegisterFunctionBinding(std::move(run_query_function_object)); return *this; } -UdfConfigBuilder& UdfConfigBuilder::RegisterLoggingHook() { - auto logging_function_object = std::make_unique>(); - logging_function_object->function_name = kLoggingHookJsName; - logging_function_object->function = LogMessage; - config_.RegisterFunctionBinding(std::move(logging_function_object)); +UdfConfigBuilder& UdfConfigBuilder::RegisterLoggingFunction() { + config_.SetLoggingFunction(LoggingFunction); return *this; } @@ -91,6 +92,8 @@ UdfConfigBuilder& UdfConfigBuilder::SetNumberOfWorkers( return *this; } -google::scp::roma::Config<>& UdfConfigBuilder::Config() { return config_; } +google::scp::roma::Config& UdfConfigBuilder::Config() { + return config_; +} } // namespace kv_server diff --git a/components/udf/udf_config_builder.h b/components/udf/udf_config_builder.h index 55f4ed69..9ad033fb 100644 --- a/components/udf/udf_config_builder.h +++ b/components/udf/udf_config_builder.h @@ -17,7 +17,7 @@ #include "components/udf/hooks/get_values_hook.h" #include "components/udf/hooks/run_query_hook.h" -#include "roma/config/src/config.h" +#include "src/roma/config/config.h" namespace kv_server { @@ -29,13 +29,13 @@ class UdfConfigBuilder { UdfConfigBuilder& RegisterRunQueryHook(RunQueryHook& run_query_hook); - UdfConfigBuilder& RegisterLoggingHook(); + UdfConfigBuilder& RegisterLoggingFunction(); UdfConfigBuilder& SetNumberOfWorkers(int number_of_workers); - google::scp::roma::Config<>& Config(); + google::scp::roma::Config& Config(); private: - google::scp::roma::Config<> config_; + google::scp::roma::Config config_; }; } // namespace kv_server diff --git a/components/util/BUILD.bazel b/components/util/BUILD.bazel index d098a0ea..a017a0dd 100644 --- a/components/util/BUILD.bazel +++ b/components/util/BUILD.bazel @@ -47,7 +47,7 @@ selects.config_setting_group( name = "local_otel_otlp", match_all = [ "//:local_instance", - "@google_privacysandbox_servers_common//src/cpp/telemetry:local_otel_export_otlp", + "@google_privacysandbox_servers_common//src/telemetry:local_otel_export_otlp", ], ) @@ -55,7 +55,7 @@ selects.config_setting_group( name = "local_otel_ostream", match_all = [ "//:local_instance", - "@google_privacysandbox_servers_common//src/cpp/telemetry:local_otel_export_ostream", + "@google_privacysandbox_servers_common//src/telemetry:local_otel_export_ostream", ], ) @@ -102,8 +102,8 @@ cc_library( local_defines = local_defines, visibility = ["//visibility:private"], deps = [ - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", ], ) @@ -136,15 +136,15 @@ cc_library( "//:aws_platform": [ "//components/errors:aws_error_util", "@aws_sdk_cpp//:core", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/interface:cpio", + "@google_privacysandbox_servers_common//src/public/cpio/interface:cpio", ], "//:gcp_platform": [ - "@google_privacysandbox_servers_common//scp/cc/public/core/interface:errors", - "@google_privacysandbox_servers_common//scp/cc/public/cpio/interface:cpio", + "@google_privacysandbox_servers_common//src/public/core/interface:errors", + "@google_privacysandbox_servers_common//src/public/cpio/interface:cpio", ], "//conditions:default": [], }) + [ - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", ], @@ -171,7 +171,7 @@ cc_library( "sleepfor.h", ], deps = [ - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/time", ], @@ -198,3 +198,16 @@ cc_library( "@com_google_googletest//:gtest", ], ) + +cc_library( + name = "request_context", + srcs = [ + "request_context.cc", + ], + hdrs = [ + "request_context.h", + ], + deps = [ + "//components/telemetry:server_definition", + ], +) diff --git a/components/util/build_info.cc b/components/util/build_info.cc index d78a2e1b..202377dc 100644 --- a/components/util/build_info.cc +++ b/components/util/build_info.cc @@ -18,7 +18,7 @@ #include -#include "glog/logging.h" +#include "absl/log/log.h" namespace kv_server { diff --git a/components/util/platform_initializer_aws.cc b/components/util/platform_initializer_aws.cc index a9519f63..1f2fc025 100644 --- a/components/util/platform_initializer_aws.cc +++ b/components/util/platform_initializer_aws.cc @@ -12,10 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "absl/log/log.h" #include "aws/core/Aws.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" -#include "public/cpio/interface/cpio.h" +#include "src/public/cpio/interface/cpio.h" namespace kv_server { using google::scp::cpio::Cpio; diff --git a/components/util/platform_initializer_gcp.cc b/components/util/platform_initializer_gcp.cc index 0bb071ae..e15be624 100644 --- a/components/util/platform_initializer_gcp.cc +++ b/components/util/platform_initializer_gcp.cc @@ -14,11 +14,12 @@ #include "absl/flags/declare.h" #include "absl/flags/flag.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" -#include "public/cpio/interface/cpio.h" -#include "scp/cc/public/core/interface/errors.h" -#include "scp/cc/public/core/interface/execution_result.h" +#include "src/public/core/interface/errors.h" +#include "src/public/core/interface/execution_result.h" +#include "src/public/cpio/interface/cpio.h" // This flag is added to allow for a local instance to use GCP as the cloud // platform. Ideally, this would be fetched from the parameter_client, but the diff --git a/components/util/request_context.cc b/components/util/request_context.cc new file mode 100644 index 00000000..aec75902 --- /dev/null +++ b/components/util/request_context.cc @@ -0,0 +1,31 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "components/util/request_context.h" + +#include + +#include "components/telemetry/server_definition.h" + +namespace kv_server { + +UdfRequestMetricsContext& RequestContext::GetUdfRequestMetricsContext() const { + return udf_request_metrics_context_; +} +InternalLookupMetricsContext& RequestContext::GetInternalLookupMetricsContext() + const { + return internal_lookup_metrics_context_; +} + +} // namespace kv_server diff --git a/components/util/request_context.h b/components/util/request_context.h new file mode 100644 index 00000000..b58036c2 --- /dev/null +++ b/components/util/request_context.h @@ -0,0 +1,52 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef COMPONENTS_UTIL_REQUEST_CONTEXT_H_ +#define COMPONENTS_UTIL_REQUEST_CONTEXT_H_ + +#include +#include +#include + +#include "components/telemetry/server_definition.h" + +namespace kv_server { + +// RequestContext holds the reference of udf request metrics context and +// internal lookup request context that ties to a single +// request. The request_id can be either passed from upper stream or assigned +// from uuid generated when RequestContext is constructed. + +class RequestContext { + public: + explicit RequestContext(const ScopeMetricsContext& metrics_context) + : udf_request_metrics_context_( + metrics_context.GetUdfRequestMetricsContext()), + internal_lookup_metrics_context_( + metrics_context.GetInternalLookupMetricsContext()) {} + UdfRequestMetricsContext& GetUdfRequestMetricsContext() const; + InternalLookupMetricsContext& GetInternalLookupMetricsContext() const; + + ~RequestContext() = default; + + private: + UdfRequestMetricsContext& udf_request_metrics_context_; + InternalLookupMetricsContext& internal_lookup_metrics_context_; +}; + +} // namespace kv_server + +#endif // COMPONENTS_UTIL_REQUEST_CONTEXT_H_ diff --git a/components/util/sleepfor.cc b/components/util/sleepfor.cc index bdf84ebf..2687c32c 100644 --- a/components/util/sleepfor.cc +++ b/components/util/sleepfor.cc @@ -14,7 +14,7 @@ #include "components/util/sleepfor.h" -#include "glog/logging.h" +#include "absl/log/log.h" namespace kv_server { diff --git a/docs/AWS_Terraform_vars.md b/docs/AWS_Terraform_vars.md index 144c0e83..cdb80321 100644 --- a/docs/AWS_Terraform_vars.md +++ b/docs/AWS_Terraform_vars.md @@ -1,5 +1,9 @@ # AWS Key Value Server Terraform vars documentation +- **add_missing_keys_v1** + + Add missing keys v1. + - **autoscaling_desired_capacity** Number of Amazon EC2 instances that should be running in the autoscaling group @@ -24,6 +28,10 @@ If you want to import an existing public certificate into ACM, follow these steps to [import the certificate](https://docs.aws.amazon.com/acm/latest/userguide/import-certificate.html). +- **data_loading_blob_prefix_allowlist** + + A comma separated list of prefixes (i.e., directories) where data is loaded from. + - **data_loading_file_format** Data file format for blob storage and realtime updates. See /public/constants.h for possible @@ -107,6 +115,14 @@ Primary coordinator account identity. +- **primary_coordinator_private_key_endpoint** + + Primary coordinator private key endpoint. + +- **primary_coordinator_region** + + Primary coordinator region. + - **prometheus_service_region** Specifies which region to find Prometheus service and use. Not all regions have Prometheus @@ -122,6 +138,10 @@ from that region is created before this terraform file is applied. That can be done by running the Key Value service terraform file in that region. +- **public_key_endpoint** + + Public key endpoint. Can only be overriden in non-prod mode. + - **realtime_updater_num_threads** The number of threads to process real time updates. @@ -149,8 +169,7 @@ - **run_server_outside_tee** - Whether to run the server outside the TEE. Not suitable for production. This runs the server in - a Docker container and can be useful for debugging. + Whether to run the server outside the TEE. - **s3_delta_file_bucket_name** @@ -170,6 +189,14 @@ Secondary coordinator account identity. +- **secondary_coordinator_private_key_endpoint** + + Secondary coordinator private key endpoint. + +- **secondary_coordinator_region** + + Secondary coordinator region. + - **server_port** Set the port of the EC2 parent instance (that hosts the Nitro Enclave instance). @@ -191,6 +218,17 @@ Source ips allowed to send ssh traffic to the ssh instance. +- **telemetry_config** + + Telemetry configuration to control whether metrics are raw or noised. Options are: mode: + PROD(noised metrics), mode: EXPERIMENT(raw metrics), mode: COMPARE(both raw and noised metrics), + mode: OFF(no metrics) + +- **udf_min_log_level** + + Minimum log level for UDFs. Info = 0, Warn = 1, Error = 2. The UDF will only attempt to log for + min_log_level and above. Default is 0 (info). + - **udf_num_workers** Total number of workers for UDF execution diff --git a/docs/GCP_Terraform_vars.md b/docs/GCP_Terraform_vars.md index 84f1b568..3ad47098 100644 --- a/docs/GCP_Terraform_vars.md +++ b/docs/GCP_Terraform_vars.md @@ -1,5 +1,9 @@ # GCP Key Value Server Terraform vars documentation +- **add_missing_keys_v1** + + Add missing keys v1. + - **backup_poll_frequency_secs** Backup poll frequency for delta file notifier in seconds. @@ -32,10 +36,19 @@ Directory to watch for files. +- **data_loading_blob_prefix_allowlist** + + A comma separated list of prefixes (i.e., directories) where data is loaded from. + - **data_loading_num_threads** Number of parallel threads for reading and loading data files. +- **enable_external_traffic** + + Whether to serve external traffic. If disabled, only internal traffic via service mesh will be + served. + - **environment** Assigned environment name to group related resources. Also servers as gcp image tag. @@ -43,7 +56,7 @@ - **envoy_port** External load balancer will send traffic to this port. Envoy will forward traffic to - kv_service_port. Must match envoy.yaml. + kv_service_port. Must match envoy.yaml. Ignored if `enable_external_traffic` is false. - **existing_service_mesh** @@ -102,6 +115,14 @@ Account identity for the primary coordinator. +- **primary_coordinator_private_key_endpoint** + + Primary coordinator private key endpoint. + +- **primary_coordinator_region** + + Primary coordinator region. + - **primary_key_service_cloud_function_url** Primary workload identity pool provider. @@ -114,6 +135,10 @@ GCP project id. +- **public_key_endpoint** + + Public key endpoint. Can only be overriden in non-prod mode. + - **realtime_updater_num_threads** Amount of realtime updates threads locally. @@ -122,6 +147,15 @@ Regions to deploy to. +- **regions_cidr_blocks** + + A set of CIDR ranges for all specified regions. The number of blocks here should correspond to + the number of regions. + +- **regions_use_existing_nat** + + Regions that use existing nat. No new nats will be created for regions specified here. + - **route_v1_to_v2** Whether to route V1 requests through V2. @@ -130,6 +164,14 @@ Account identity for the secondary coordinator. +- **secondary_coordinator_private_key_endpoint** + + Secondary coordinator private key endpoint. + +- **secondary_coordinator_region** + + Secondary coordinator region. + - **secondary_key_service_cloud_function_url** Secondary key service cloud function url. @@ -140,24 +182,35 @@ - **server_dns_zone** - Dns zone for Kv-serer. + Dns zone for Kv-serer. Ignored if `enable_external_traffic` is false. - **server_domain_ssl_certificate_id** - Ssl certificate id of the Kv-server domain. + Ssl certificate id of the Kv-server domain. Ignored if `enable_external_traffic` is false. - **server_url** - Kv-serer URL. Example: kv-server-environment.example.com + Kv-serer URL. Example: kv-server-environment.example.com. Ignored if `enable_external_traffic` + is false. - **service_account_email** Email of the service account that be used by all instances. +- **service_mesh_address** + + Service mesh address of the KV server. + - **tee_impersonate_service_accounts** Tee can impersonate these service accounts. Necessary for coordinators. +- **telemetry_config** + + Telemetry configuration to control whether metrics are raw or noised. Options are: mode: + PROD(noised metrics), mode: EXPERIMENT(raw metrics), mode: COMPARE(both raw and noised metrics), + mode: OFF(no metrics) + - **udf_num_workers** Number of workers for UDF execution. diff --git a/docs/ad_retrieval_overview.md b/docs/ad_retrieval_overview.md new file mode 100644 index 00000000..2cf44adf --- /dev/null +++ b/docs/ad_retrieval_overview.md @@ -0,0 +1,4 @@ +The content of this file has been moved to +[docs/protected_app_signals/ad_retrieval_overview.md](/docs/protected_app_signals/ad_retrieval_overview.md). + +You can also find other latest documentation related to Protected App Signals in that directory. diff --git a/docs/assets/ad_retrieval_filter_funnel.png b/docs/assets/ad_retrieval_filter_funnel.png new file mode 100644 index 00000000..bf78273a Binary files /dev/null and b/docs/assets/ad_retrieval_filter_funnel.png differ diff --git a/docs/assets/ad_retrieval_udf.png b/docs/assets/ad_retrieval_udf.png new file mode 100644 index 00000000..52618a87 Binary files /dev/null and b/docs/assets/ad_retrieval_udf.png differ diff --git a/docs/assets/ad_retrieval_use_case_overview.png b/docs/assets/ad_retrieval_use_case_overview.png new file mode 100644 index 00000000..ae60a605 Binary files /dev/null and b/docs/assets/ad_retrieval_use_case_overview.png differ diff --git a/docs/assets/ad_retrieval_walkthrough.png b/docs/assets/ad_retrieval_walkthrough.png new file mode 100644 index 00000000..b016d433 Binary files /dev/null and b/docs/assets/ad_retrieval_walkthrough.png differ diff --git a/docs/data_loading/data_format_specification.md b/docs/data_loading/data_format_specification.md new file mode 100644 index 00000000..6d091f7a --- /dev/null +++ b/docs/data_loading/data_format_specification.md @@ -0,0 +1,78 @@ +# Event/Record + +The KV server data is represented by events or records. These are used interchangeably in the +codebase. In the context of data loading, they are also called mutations. + +A data mutation record consists of a key, a value and some metadata. When loaded into the server, +they can be queried by the query API. + +The record format is [Flatbuffers](https://flatbuffers.dev/). The schema is defined +[here](/public/data_loading/data_loading.fbs). + +- Each mutation event is associated with a `logical_commit_timestamp`, larger timestamp indicates + a more recent record. +- There are two types of mutation events: + 1. UPDATE which inserts/modifies a key/value record + 2. DELETE which deletes an existing key/value record. + +# Data files + +Files are used as containers of records. The server reads the files to load the records. + +There are two types data files consumed by the server: + +1. delta files: incremental updates +1. snapshot files: full rewrite + +In both cases, newer key/value pairs supersede existing key/value pairs. + +## Delta files + +Delta filename must conform to the regular expression `DELTA_\d{16}`. See +[constants.h](/public/constants.h) for the most up-to-date format. More recent delta files are +lexicographically greater than older delta files. Delta files have the following properties: + +- Consists of key/value mutation events (updates/deletes) for a fixed time window. +- `logical_commit_timestamp` of the records have no relation with their file's name. It is + acceptable to also use timestamps in file names for ordering purposes for your convenience but + the system makes no assumption on the relation between the record timestamps and the file names. +- There are no enforced size limits for delta files, but smaller files are faster to read. +- Server instances continually watch for newer delta files and update their in-memory caches. + +## Snapshot files + +Snapshot filename must conform to the regular expression `SNAPSHOT_\d{16}`. See +[constants.h](/public/constants.h). for the most up-to-date format. More recent snapshot files are +lexicographically greater than older snapshot files. Snpashot files have the following properties: + +- Uses the same file format as delta files and are only read at server startup time. +- Generated from: + - compacting a set of delta files by merging multiple mutation events for the same key such + that the resulting snapshot consists of only UPDATE mutation events. + - compacting a base snapshot file together with a set of delta files that are not in the base + snapshot file. +- Contains the entire set of key/value records since the beginning of time. +- There are no enforced size limits for snapshot files. + +See [File groups](file_groups.md#file-groups) on how to improve snapshot generation scalability and +throughput. + +## File format + +### Avro + +The system supports [Avro](https://avro.apache.org/) as the file format. It should be an Avro file +containing one or more Flatbuffers records. + +The Avro schema is the primitive string schema. + +Each row is a serialized Flatbuffers record. + +- [C++ example](/public/data_loading/readers/avro_stream_io_test.cc) +- [Java example](https://github.com/privacysandbox/protected-auction-key-value-service/issues/39): + This is not maintained by the dev team and may be out of date. + +#### Riegeli + +The system also supports [Riegeli](https://github.com/google/riegeli). Similarly each file contains +one or more serialized Flatbuffers records. diff --git a/docs/data_loading_capabilities.md b/docs/data_loading/data_loading_capabilities.md similarity index 100% rename from docs/data_loading_capabilities.md rename to docs/data_loading/data_loading_capabilities.md diff --git a/docs/data_loading/file_groups.md b/docs/data_loading/file_groups.md new file mode 100644 index 00000000..e6118b87 --- /dev/null +++ b/docs/data_loading/file_groups.md @@ -0,0 +1,44 @@ +# File groups + +## What are file groups? + +A file group is a group of split files treated logically as one file for data loading purposes by +the KV server. Each split can be uploaded to the blob storage bucket as soon as it's completely +generated which improves scalability and throughput for data generation pipelines. + +> Currently, file groups are only supported for [SNAPSHOT](loading_data.md#snapshot-files) files. + +## Naming scheme + +FILE GROUP NAME FORMAT: +``_``_``_OF_`` + +VARIABLE DESCRIPTION: + +- `FILE_TYPE`: type of the file. Currently, this can only be SNAPSHOT. +- `LOGICAL_TIMESTAMP`: is a 16 digit monotonic clock that only moves forward. a larger value means + a more recent file. +- `PART_FILE_INDEX`: a 5 digit number that represents the index of the file in the file group. + each index can only be used with a single part file name, loading part files sharing an index is + not guaranteed to be correct. valid range is `[0..NUM_PART_FILES-1]`. +- `NUM_PART_FILES`: a 6 digit number that represents the total number of part files in a file + group. valid range is `[1..100,000]`. + +VALID EXAMPLES: + +- `SNAPSHOT_1705430864435450_00000_OF_000010` is the first part file in a snapshot file group with + 10 part files. +- `SNAPSHOT_1705430864435450_00009_OF_000010` is the last part file in a snapshot file group with + 10 part files. + +There is a util function to generate split file names conforming to this naming scheme from +components [ToFileGroupFileName(...)](/public/data_loading/filename_utils.h#L67) + +## Server startup and snapshot file groups + +During server startup, the server looks for the most recent, complete snapshot file group and loads +that first. Then, it continues with loading most recent delta files not included in the snapshot +file group. A file group is considered complete is all of it's part files are in the storage at load +time. For example, the following set of part files a complete snapshot file group of size 2: + +`complete_group = ["SNAPSHOT_1705430864435450_00000_OF_000002", "SNAPSHOT_1705430864435450_00001_OF_000002"]` diff --git a/docs/loading_data.md b/docs/data_loading/loading_data.md similarity index 67% rename from docs/loading_data.md rename to docs/data_loading/loading_data.md index 4dff2d83..deebff18 100644 --- a/docs/loading_data.md +++ b/docs/data_loading/loading_data.md @@ -1,66 +1,31 @@ -> FLEDGE has been renamed to Protected Audience API. To learn more about the name change, see the -> [blog post](https://privacysandbox.com/intl/en_us/news/protected-audience-api-our-new-name-for-fledge) +# Loading data into the Key/Value server -# Load data into the FLEDGE Key/Value server +There are two ways to populate data in the server. -The FLEDGE Key/Value server is used to send real-time signals to the buyers and the sellers during a -FLEDGE auction. +- The standard path is by uploading files to a cloud file storage service. The standard upload is + the authoritative, high bandwidth and persistent source of truth. +- The other way is via a low latency path. To apply such an update, you should send an update to a + dedicated broadcast topic. -There are two ways to populate data in the server. The standard path is by uploading files to a -cloud file storage service. The standard upload is the authoritative, high bandwidth and persistent -source of truth. - -The other way is via a low latency path. To apply such an update, you should send an update to a -dedicated broadcast topic. - -This doc explains the expected file format, and processes to perform the common data loading -operations. Please note the following: +The data format specification is defined [here](/docs/data_loading/data_format_specification.md). +This doc talks about how to load the data. - We provide a [C++ library reference implementation](#using-the-c-reference-library-to-read-and-write-data-files) and a [CLI tool](#using-the-cli-tool-to-generate-delta-and-snapshot-files) that can be used to generate (or write) and read data files. - The reference library and CLI tool are ready to use as-is, or you can write your own libraries. -- The data generation part is a general process that applies to all cloud providers, but the - uploading instructions are for AWS only. - -# Data files - -There are two types data files consumed by the server, (1) delta files and (2) snapshot files. In -both cases, newer key/value pairs supersede existing key/value pairs. - -## Delta files - -Delta filename must conform to the regular expression `DELTA_\d{16}`. See -[constants.h](../public/constants.h) for the most up-to-date format. More recent delta files are -lexicographically greater than older delta files. Delta files have the following properties: - -- Consists of key/value mutation events (updates/deletes) for a fixed time window. The events are - in the format of Flatbuffers ([Schema](/public/data_loading/data_loading.fbs)). -- Each mutation event is associated with a `logical_commit_timestamp`, larger timestamp indicates - a more recent record. -- `logical_commit_timestamp` of the records have no relation with their file's name. It is - acceptable to also use timestamps in file names for ordering purposes for your convenience but - the system makes no assumption on the relation between the record timestamps and the file names. -- There are two types of mutation events: (1) UPDATE which introduces/modifies a key/value record, - and (2) DELETE which deletes an existing key/value record. -- There are no enforced size limits for delta files, but smaller files are faster to read. -- Server instances continually watch for newer delta files and update their in-memory caches. - -## Snapshot files - -Snapshot filename must conform to the regular expression `SNAPSHOT_\d{16}`. See -[constants.h](../public/constants.h). for the most up-to-date format. More recent snapshot files are -lexicographically greater than older snapshot files. SNpashot files have the following properties: - -- Uses the same file format as delta files and are only read at server startup time. -- Generated from: - - compacting a set of delta files by merging multiple mutation events for the same key such - that the resulting snapshot consists of only UPDATE mutation events. - - compacting a base snapshot file together with a set of delta files that are not in the base - snapshot file. -- Contains the entire set of key/value records since the beginning of time. -- There are no enforced size limits for snapshot files. +- The data generation part is a general process that applies to all cloud providers. + +# Before you start: choose your file format + +Currently the files can be in one of multiple formats. When deploying the system, set the data +format parameter to instruct the system to read data files as the specified format. + +- For AWS: Set the [Terraform var](/docs/AWS_Terraform_vars.md) data_loading_file_format. +- For Local: Set flag `--data_loading_file_format` + ([defined here](/components/cloud_config/parameter_client_local.cc)). +- For GCP: To be supported. # Experimenting with sample data @@ -77,9 +42,7 @@ Confirm that the sample data file `DELTA_\d{16}` has been generated. # Using the CLI tool to generate delta and snapshot files -The data CLI is located under: `//tools/data_cli`. First build the cli using the following command -(Note that to build the cli to use `generate_snapshot` command with data in AWS S3, use -`--//:platform=aws`.): +The data CLI is located under: `//tools/data_cli`. First build the cli using the following command: ```sh -$ builders/tools/bazel-debian run //production/packaging/tools:copy_to_dist --//:instance=local --//:platform=local @@ -178,10 +141,8 @@ And to generate a snapshot from a set of delta files, run the following command values with your own values): ```sh --$ export GLOG_logtostderr=1; -export DATA_DIR=; +-$ export DATA_DIR=; docker run -it --rm \ - --env GLOG_logtostderr \ --volume=/tmp:/tmp \ --volume=$DATA_DIR:$DATA_DIR \ --user $(id -u ${USER}):$(id -g ${USER}) \ @@ -193,35 +154,32 @@ docker run -it --rm \ --starting_file=DELTA_0000000000000001 \ --ending_delta_file=DELTA_0000000000000010 \ --snapshot_file=SNAPSHOT_0000000000000001 + --stderrthreshold=0 ``` The output snapshot file will be written to `$DATA_DIR`. # Using the C++ reference library to read and write data files -Data files are written using the [Riegeli](https://github.com/google/riegeli) format and data -records are stored as [Flatbuffers](https://google.github.io/flatbuffers/). The record schema is -here: [Flatbuffer record schema](public/data_loading/data_loading.fbs). - The C++ reference library implementation can be found under: -[C++ data file readers](../public/data_loading/readers) and -[C++ data file writers](../public/data_loading/writers). To write snapshot files, you can use -[Snapshot writer](../public/data_loading/writers/snapshot_stream_writer.h) and to write delta files, -you can use [Delta writer](../public/data_loading/writers/delta_record_stream_writer.h). Both files -can be read using the -[data file reader](../public/data_loading/readers/delta_record_stream_reader.h). The source and -destination of the provided readers and writers are required to be `std::iostream` objects. +[C++ data file readers](/public/data_loading/readers) and +[C++ data file writers](/public/data_loading/writers). To write snapshot files, you can use +[Snapshot writer](/public/data_loading/writers/snapshot_stream_writer.h) and to write delta files, +you can use [Delta writer](/public/data_loading/writers/delta_record_stream_writer.h). Both files +can be read using the [data file reader](/public/data_loading/readers/delta_record_stream_reader.h). +The source and destination of the provided readers and writers are required to be `std::iostream` +objects. # Writing your own C++ data libraries Feel free to use the C++ reference library provided above as examples if you want to write your own data library. Keep the following things in mind: -- Make sure the output files adhere to the [delta](#delta-files) and [snapshot](#snapshot-files) - file properties listed above. +- Make sure the output files adhere to the + [specification](/docs/data_loading/data_format_specification.md). - Snapshot files must be written with metadata specifying the starting and ending filenames of records included in the snapshot. See - [SnpashotMetadata proto](../public/data_loading/riegeli_metadata.proto). + [SnpashotMetadata proto](/public/data_loading/riegeli_metadata.proto). # Upload data files to AWS @@ -242,7 +200,7 @@ You can use the AWS CLI to upload the sample data to S3, or you can also use the Confirm that the file is present in the S3 bucket: -![the delta file listed in the S3 console](assets/s3_delta_file.png) +![the delta file listed in the S3 console](/docs/assets/s3_delta_file.png) ## Upload data files to GCP @@ -250,7 +208,7 @@ Similar to AWS, the server in GCP watches a Google Cloud Storage (GCS) bucket co Terraform config. New files in the bucket will be automatically uploaded to the server. You can uploade files to your GCS bucket through Google Cloud Console. -![files listed in the Google Cloud Console](assets/gcp_gcs_bucket.png) +![files listed in the Google Cloud Console](/docs/assets/gcp_gcs_bucket.png) Alternatively, you can use [gsutil tool](https://cloud.google.com/storage/docs/gsutil) to upload files to GCS. For example: @@ -260,11 +218,44 @@ export GCS_BUCKET=your-gcs-bucket-id gsutil cp DELTA_* gs://${GCS_BUCKET} ``` -## Integrating file uploading with your data source for AWS +## Organizing data files using prefixes + +### Intended use case + +By default, the server automatically monitors and loads data files uploaded to the S3 or GCS bucket. +However, since the server expects the file names to monotonically increase and will not read files +older than the previous file it has read, it can be challenging if there are more than one data +ingestion pipelines creating delta files independently, because this would require them to +coordinate the file names to make sure the server reads all the files. And the coordination adds +extra complexity and latency. + +This feature allows multiple data ingestion pipelines to operate completely independently. + +### Allow listing prefixes + +Prefixes need to be allow listed before the server can continuously monitor and load new files under +them. The allowlist is a comma separated list of prefixes and is controlled via a server flag +`data_loading_blob_prefix_allowlist`: + +- For AWS, see docs here: [AWS vars](/docs/AWS_Terraform_vars.md) +- For GCP, see docs here: [GCP vars](/docs/GCP_Terraform_vars.md) +- For local, set the flag: `--data_loading_blob_prefix_allowlist` + [defined here](/components/cloud_config/parameter_client_local.cc). + +For example, to add `prefix1` and `prefix2` to the allow list, set +`data_loading_blob_prefix_allowlist` to `prefix1,prefix2`. With this setup the server will continue +to load data files at the main bucket level, and will also monitor and load files that start with +`prefix1` or `prefix2`, e.g., `prefix1/DELTA_001` and `prefix2/DELTA_001`. + +### Important things to note -AWS provides libraries to communicate with S3, such as the -[C++ SDK](https://aws.amazon.com/sdk-for-cpp/). As soon as a file is uploaded to a watched bucket it -will be read into the service, assuming that it has a higher logical commit timestamp. +- Records from files with different prefixes are merged in the internal cache so two records with + the same key (from different prefixes) will conflict with each other. These collisions should be + managed when writing records into data files, e.g., use keys `prefix1:foo0` and `prefix2:foo0` + for writing the records to data files and at query time. +- The server keeps track of the most recent loaded file separately for each prefix. +- The server does garbage collection of deleted records using a separate max cutoff timestamp for + each prefix. # Realtime updates @@ -273,7 +264,7 @@ delta file to a dedicated broadcast topic. ## AWS -![Realtime design](assets/realtime_design.png) +![Realtime design](/docs/assets/realtime_design.png) In the case of AWS it is a Simple Notification Service (SNS) topic. That topic is created in terraform @@ -293,9 +284,9 @@ The setup is similar to AWS above. The differences are in terminology: - Sqs->Subscription - EC2->VM -In the case of AWS it is a PubSub topic. That topic is created in terraform -[here](../production/terraform/gcp/services/realtime/main.tf) Delta files contain multiple rows, -which allows you to batch multiple updates together. There is a +In the case of GCP it is a PubSub topic. That topic is created in terraform +[here](/production/terraform/gcp/services/realtime/main.tf) Delta files contain multiple rows, which +allows you to batch multiple updates together. There is a [limit](https://cloud.google.com/pubsub/quotas#resource_limits) of 10MB for the message size. Each data server is subscribed to the topic through @@ -310,10 +301,10 @@ data loading path. If it is not, then that update can be lost, for example, duri The standard upload is the authoritative and persistent source of truth, and the low latency update allows to speed up the update latency. -![Realtime sequence](assets/realtime_sequence.png) +![Realtime sequence](/docs/assets/realtime_sequence.png) As per the diagram below, first you should -[write](<(#using-the-c-reference-library-to-read-and-write-data-files)>) the updates to a delta file +[write](#using-the-c-reference-library-to-read-and-write-data-files) the updates to a delta file that will be uploaded via a standard path later. The purpose of this step is to guarantee that this record won't be missed later. @@ -366,5 +357,5 @@ gcloud pubsub topics publish "$topic_arn" --message "$file" ### Cpp -Check out this sample [tool](../components/tools/realtime_updates_publisher.cc) on how to insert low +Check out this sample [tool](/components/tools/realtime_updates_publisher.cc) on how to insert low latency updates. diff --git a/docs/realtime_updates_capabilities.md b/docs/data_loading/realtime_updates_capabilities.md similarity index 95% rename from docs/realtime_updates_capabilities.md rename to docs/data_loading/realtime_updates_capabilities.md index bbf46ce3..85834d3a 100644 --- a/docs/realtime_updates_capabilities.md +++ b/docs/data_loading/realtime_updates_capabilities.md @@ -4,9 +4,9 @@ A parameter ([AWS](https://github.com/privacysandbox/fledge-key-value-service/blob/7f3710b1f1c944d7879718a334afd5cb8f80f3d9/production/terraform/aws/environments/kv_server.tf#L51), -[GCP](../docs/GCP_Terraform_vars.md#L96)) sets the size of the thread pool that reads off a queue. -The bigger that number is, the smaller the batch size can be. It is preferred to use a larger batch -size where possible. +[GCP](/docs/GCP_Terraform_vars.md#L96)) sets the size of the thread pool that reads off a queue. The +bigger that number is, the smaller the batch size can be. It is preferred to use a larger batch size +where possible. ### AWS @@ -46,7 +46,7 @@ While similar logic applies, the GCP SDK has superior performance due to To get to a higher QPS we can have multiple threads reading off a queue. This is a parameter ([AWS](https://github.com/privacysandbox/fledge-key-value-service/blob/7f3710b1f1c944d7879718a334afd5cb8f80f3d9/production/terraform/aws/environments/kv_server.tf#L51), -[GCP](../docs/GCP_Terraform_vars.md#L96)) that our solution exposes. It can be increased to match +[GCP](/docs/GCP_Terraform_vars.md#L96)) that our solution exposes. It can be increased to match specific QPS requirements and underlying hardware - based on the number of cores. ## Batching @@ -219,5 +219,6 @@ histogram_quantile(0.5,rate(Latency_bucket{event="ReceivedLowLatencyNotification You can query the prometheus the same way it's done for AWS. Note that KV server doesn't expose `AwsSqsReceiveMessageLatency`, and `AWS` in the metric name should be substituted with `GCP`. -You can also use the UI [dashboard](../production/terraform/gcp/realtime_pubsub_dashboard.json). -Make sure to replace PROJECT_ID and ENVIRONMENT with your values. +You can also use the UI +[dashboard](/production/terraform/gcp/dashboards/realtime_pubsub_dashboard.json). Make sure to +replace PROJECT_ID and ENVIRONMENT with your values. diff --git a/docs/deployment/deploying_locally.md b/docs/deployment/deploying_locally.md index f36276be..486090e9 100644 --- a/docs/deployment/deploying_locally.md +++ b/docs/deployment/deploying_locally.md @@ -5,13 +5,13 @@ This article is for adtech engineers who want to test the Key/Value server locally. Deploying production servers in this way is not recommended, please see the -[AWS deployment guide](deploying_on_aws.md) instead. +[AWS deployment guide](deploying_on_aws.md) or [GCP deployment guide](deploying_on_gcp.md) instead. To learn more about FLEDGE and the Key/Value server, take a look at the following documents: - [FLEDGE Key/Value server explainer](https://github.com/WICG/turtledove/blob/main/FLEDGE_Key_Value_Server_API.md) - [FLEDGE Key/Value server trust model](https://github.com/privacysandbox/fledge-docs/blob/main/key_value_service_trust_model.md) -- [FLEDGE explainer](https://developer.chrome.com/en/docs/privacy-sandbox/fledge/) +- [FLEDGE explainer](https://developers.google.com/privacy-sandbox/relevance/protected-audience) - [FLEDGE API developer guide](https://developer.chrome.com/blog/fledge-api/) > The instructions written in this document are for running a test Key/Value server that does @@ -51,7 +51,7 @@ From the Key/Value server repo folder, execute the following command: We provide a default UDF implementation that is loaded into the server at startup. -To use your own UDF, refer to the [UDF Delta file documentation](./generating_udf_files.md) to +To use your own UDF, refer to the [UDF Delta file documentation](/docs/generating_udf_files.md) to generate a UDF delta file. Include the delta file in your local delta directory (see below). @@ -73,10 +73,9 @@ their contents on startup and continue to watch them while it is running. ## Start the server ```sh -GLOG_alsologtostderr=1 GLOG_v=4 \ ./bazel-bin/components/data_server/server/server \ --delta_directory=/tmp/deltas \ - --realtime_directory=/tmp/realtime + --realtime_directory=/tmp/realtime --v=4 --stderrthreshold=0 ``` The server will start up and begin listening for new delta and realtime files in the directories diff --git a/docs/deployment/deploying_on_aws.md b/docs/deployment/deploying_on_aws.md index c9fcacca..83d0e38b 100644 --- a/docs/deployment/deploying_on_aws.md +++ b/docs/deployment/deploying_on_aws.md @@ -118,7 +118,7 @@ Take a note of the AMI ID from the output as it will be used for Terraform later We provide a default UDF implementation that is loaded into the server at startup. -To use your own UDF, refer to the [UDF Delta file documentation](./generating_udf_files.md) to +To use your own UDF, refer to the [UDF Delta file documentation](/docs/generating_udf_files.md) to generate a UDF delta file. Upload this UDF delta file to the S3 bucket that will be used for delta files before attempting to @@ -245,8 +245,8 @@ scope for this documentation. # Loading data into the server -Refer to the [FLEDGE Key/Value data loading guide documentation](./loading_data.md) for loading data -to be queried into the server. +Refer to the [FLEDGE Key/Value data loading guide documentation](/docs/data_loading/loading_data.md) +for loading data to be queried into the server. # Common operations @@ -297,7 +297,7 @@ grpcurl --protoset dist/query_api_descriptor_set.pb -d '{"raw_body": {"data": "' ## SSH into EC2 -![how a single SSH instance is used to log into multiple server instances](assets/ssh_instance.png) +![how a single SSH instance is used to log into multiple server instances](../assets/ssh_instance.png) ### Step 1: SSH into the SSH EC2 instance @@ -309,7 +309,7 @@ proceeding. We will need either the instance id (if connecting using EC2 instanc the public IP dns (if connecting using own key and SSH client) of the SSH instance and both can be retrieved from the EC2 dashboard. -![where to find the instance id or public dns for the EC2 instance](assets/ec2_instance_id_and_public_dns.png) +![where to find the instance id or public dns for the EC2 instance](../assets/ec2_instance_id_and_public_dns.png) Confirm that you can SSH into your SSH EC2 instance by following the instructions on [Connect using EC2 Instance Connect](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-connect-methods.html). @@ -477,8 +477,8 @@ Alternatively, you can SSH into an existing server instance and start the Docker 1. Run the docker container ```sh - docker run -d --init --rm --env GLOG_v=5 --network host --security-opt=seccomp=unconfined \ - --entrypoint=/init_server_basic bazel/production/packaging/aws/data_server:server_docker_image -- --port 50051 + docker run -d --init --rm --network host --security-opt=seccomp=unconfined \ + --entrypoint=/init_server_basic bazel/production/packaging/aws/data_server:server_docker_image -- --port 50051 --v=5 ``` ## Viewing Telemetry @@ -537,4 +537,4 @@ The resources are allocated by specifying the per-TEE values in the terraform va ## How is private communication configured? -See this [doc](private_communication_aws.md) for more details. +See this [doc](/docs/private_communication_aws.md) for more details. diff --git a/docs/deployment/deploying_on_gcp.md b/docs/deployment/deploying_on_gcp.md index b5f62035..e258f70c 100644 --- a/docs/deployment/deploying_on_gcp.md +++ b/docs/deployment/deploying_on_gcp.md @@ -149,8 +149,8 @@ deploy to, and update the `[[REGION]].tfvars.json` with Terraform variables for The description of each variable is described in [GCP Terraform Vars doc](/docs/GCP_Terraform_vars.md). -Note that variable `tls-key` and `tls-cert` are not in `[[REGION]].tfvars.json`. Please supply these -with a `secrets.auto.tfvar` file under `production/terraform/gcp/environments/`. +Note that variable `tls_key` and `tls_cert` are not in `[[REGION]].tfvars.json`. Please supply these +with a `secrets.auto.tfvars` file under `production/terraform/gcp/environments/`. Update the `[[REGION]].backend.conf`: @@ -158,6 +158,31 @@ Update the `[[REGION]].backend.conf`: [Set up GCS bucket for Terraform states](#set-up-gcs-bucket-for-terraform-states) step. - `prefix` - Set a path/to/directory to contain the Terraform state. +## Bidding and Auction services integration within the same VPC + +If you're integrating with Bidding and Auction services (B&A), you are likely going to be reusing +the same VPC (virtual private cloud) and Service Mesh (internal LB). + +Hence, you need to set these two parameters to true: `use_existing_service_mesh`, +`use_existing_vpc`. + +You also need to set these parameters to proper values: `existing_service_mesh`, `existing_vpc_id`. +Example value: + +`existing_service_mesh`: `projects/your-project/locations/global/meshes/your-mesh` +`existing_vpc_id`: `projects/your-project/global/networks/your-vpc` + +Other things to keep in mind + +- CIDR range needs to be different for each server deployment (B&A, each kv). It is specified with + `regions_cidr_blocks` +- `enable_external_traffic` can be set to false. In this case, several other terraform vars can be + [ignored](../GCP_Terraform_vars.md) +- `regions_use_existing_nat` -- note that the NAT for each region can only be set up once under + the same VPC. If a NAT has already been set up in an existing server deployment (by specifying + an empty set of the `regions_use_existing_nat`), other sever deployments in the same region(s) + need to specify the region(s) in the `regions_use_existing_nat` var. + ## Apply Terraform From the Key/Value server repo folder, run: @@ -183,13 +208,13 @@ builders/tools/terraform -chdir=production/terraform/gcp/environments apply --va At the end, to destroy all the GCP resources: ```sh -builders/tools/terraform -chdir=production/terraform/gcp/environments destroy --var-file=--var-file=${ENVIRONMENT}/${REGION}.tfvars.json +builders/tools/terraform -chdir=production/terraform/gcp/environments destroy --var-file=${ENVIRONMENT}/${REGION}.tfvars.json ``` # Loading data into the server -Refer to the [FLEDGE Key/Value data loading guide documentation](./loading_data.md) for loading data -to be queried into the server. +Refer to the [FLEDGE Key/Value data loading guide documentation](/docs/data_loading/loading_data.md) +for loading data to be queried into the server. # Common operations @@ -259,7 +284,7 @@ BODY='{ "metadata": { "hostname": "example.com" }, "partitions": [{ "id": 0, "co ### Option 2: Via service mesh In short, service mesh is an internal load balancer. So, it is only available internally within a -VPC. The Kv-server is set up with a service mesh as a backend service and we can query it internally +VPC. The KV Server is set up with a service mesh as a backend service and we can query it internally using proxyless gRPC. Normally, you would send proxyless gRPC queries from other internal servers (e.g., a bidding server). For this demo, however, we are setting up a temporary server within the same VPC and service mesh as the Kv-server for demonstration purposes. @@ -272,7 +297,7 @@ export ENVIRONMENT=your_environment ``` Then, use the following command to set up a client VM within the same VPC network as your already -deployed Kv-server. Note that you still need to manually replace `[[your_environment]]` at the end +deployed KV Server. Note that you still need to manually replace `[[your_environment]]` at the end of the script. ```sh diff --git a/production/terraform/README.md b/docs/deployment/working_with_terraform.md similarity index 71% rename from production/terraform/README.md rename to docs/deployment/working_with_terraform.md index 16d29562..32b0a8fb 100644 --- a/production/terraform/README.md +++ b/docs/deployment/working_with_terraform.md @@ -1,8 +1,9 @@ ## Directory Structure -### AWS +### AWS/GCP -The root directory for Terraform templates is `${project_root}/production/terraform/aws`. +The root directory for Terraform templates is `${project_root}/production/terraform/${platform}`, +where `${platform}` is eigher `aws` or `gcp`. Within the root directory, the templates are organized as follows: @@ -16,4 +17,5 @@ Within the root directory, the templates are organized as follows: ### Usage Terraform operations should be performed on a particular environment. See the -[environment guide](/production/terraform/aws/environments/README.md) for actual usage. +[AWS environment guide](/production/terraform/aws/environments/README.md) and +[GCP environment guide](/production/terraform/gcp/environments/README.md) for actual usage. diff --git a/docs/developing_the_server.md b/docs/developing_the_server.md index f7d66502..ad636737 100644 --- a/docs/developing_the_server.md +++ b/docs/developing_the_server.md @@ -3,7 +3,7 @@ # FLEDGE K/V Server developer guide -## Data Server +## Develop and run the server for AWS platform in your local machine The data server provides the read API for the KV service. @@ -11,7 +11,7 @@ The data server provides the read API for the KV service. > Attention: The server can run locally (in or outside of Docker) while specifying `aws` as platform, in which case it will > contact AWS based on the local AWS credentials. However, this requires the AWS environment to be -> set up first following the [AWS deployment guide](/docs/deploying_on_aws.md). You might need to +> set up first following the [AWS deployment guide](/docs/deployment/deploying_on_aws.md). You might need to > set up the following parameters in the AWS System Manager: > > | Parameter Name | Value | @@ -76,7 +76,7 @@ The data server provides the read API for the KV service. 1. Run the container. Port 50051 can be used to query the server directly through gRPC. --environment must be specified. The server will still read data from S3 and the server uses environment to find the S3 bucket. The environment is configured as part of the - [AWS deployment process](/docs/deploying_on_aws.md). + [AWS deployment process](/docs/deployment/deploying_on_aws.md). Set region. The region should be where your environment is deployed: @@ -137,12 +137,65 @@ grpc_cli call localhost:50051 kv_server.v1.KeyValueService.GetValues \ curl http://localhost:51052/v1/getvalues?kv_internal=hi ``` +## Develop and run the server for GCP platform in your local machine + +The server can run locally while specifying `gcp` as platform. However, certain GCP resources (such +as parameters, GCS data bucket) are still required and please follow +[GCP deployment guide](/docs/deployment/deploying_on_gcp.md) to set up the GCP environment first. + +### Run the server locally inside a docker container + +#### Build the image + +From the kv-server repo folder, execute the following command + +```sh +builders/tools/bazel-debian run //production/packaging/gcp/data_server:copy_to_dist --config=local_instance --config=gcp_platform +``` + +#### Load the image into docker + +```sh +docker load -i dist/server_docker_image.tar +``` + +#### Start the server + +```sh +docker run --init -v "$HOME/.config/gcloud/application_default_credentials.json":/root/.config/gcloud/application_default_credentials.json:ro --network host --add-host=host.docker.internal:host-gateway --privileged --rm bazel/production/packaging/gcp/data_server:server_docker_image --gcp_project_id=${GCP_PROJECT_ID} --environment=${GCP_ENVIRONMENT} +``` + +where `${GCP_PROJECT_ID}` is your GCP project_id and `${GCP_ENVIRONMENT}` is the environment name +for your GCP resources. + +### Interact with the server + +- If the parameter `enable_external_traffic` (Terraform variable) is set to true, we can query the + server via the envoy port: + +```sh +./grpcurl -insecure -d '{"kv_internal":"hi"}' localhost:51052 kv_server.v1.KeyValueService.GetValues +``` + +- Alternatively, if `enable_external_traffic` is false, we can directly query the server port: + +```sh +./grpcurl -plaintext -d '{"kv_internal":"hi"}' localhost:50051 kv_server.v1.KeyValueService.GetValues +``` + +Note that you may need to set the path to your `grpcurl` tool, or install `grpcurl` if you haven't +done so already + +```sh +curl -L https://github.com/fullstorydev/grpcurl/releases/download/v1.8.1/grpcurl_1.8.1_linux_x86_64.tar.gz | tar -xz +``` + ## Develop and run the server inside AWS enclave The KV service instance should be set up by following the deployment guide -([AWS](/docs/deploying_on_aws.md)). For faster iteration, enclave image of the server is also -produced under `dist/`. Once the system has been started, iterating on changes to the server itself -only requires restarting the enclave image: +([AWS](/docs/deployment/deploying_on_aws.md)). For faster iteration, enclave image of the server is +also produced under `dist/`. Once the system has been started, iterating on changes to the server +itself only requires restarting the enclave image: 1. Copy the new enclave EIF to an AWS EC2 instance that supports nitro enclave. Note: The system has a SSH instance that a developer can access. From there the user can access actual server EC2 diff --git a/docs/generating_udf_files.md b/docs/generating_udf_files.md index b242098f..e714f6f9 100644 --- a/docs/generating_udf_files.md +++ b/docs/generating_udf_files.md @@ -47,10 +47,11 @@ Tools to generate UDF delta files and test them are in the `tools/udf` directory 1. Build the executables: ```sh - -$ builders/tools/bazel-debian run //production/packaging/tools:copy_to_dist_udf + -$ builders/tools/bazel-debian build -c opt //tools/udf/udf_generator:udf_delta_file_generator ``` -2. Generate a UDF delta file using the `dist/debian/udf_delta_file_generator` executable. +2. Generate a UDF delta file using the `bazel-bin/tools/udf/udf_generator/udf_delta_file_generator` + executable. Flags: @@ -63,13 +64,13 @@ Tools to generate UDF delta files and test them are in the `tools/udf` directory Example: ```sh - -$ dist/debian/udf_delta_file_generator --output_dir="$PWD" --udf_file_path="path/to/my/udf/udf.js" + -$ bazel-bin/tools/udf/udf_generator/udf_delta_file_generator --output_dir="$PWD" --udf_file_path="path/to/my/udf/udf.js" ``` ### Option 2. Generating your own delta file You can use other options to generate delta files, e.g. using the -[`data_cli` tool](./loading_data.md). +[`data_cli` tool](/docs/data_loading/loading_data.md). The delta file must have a `DataRecord` with a `UserDefinedFunctionsConfig` as its record. diff --git a/docs/inline_wasm_udfs.md b/docs/inline_wasm_udfs.md index 6fde4d48..6d9a35e4 100644 --- a/docs/inline_wasm_udfs.md +++ b/docs/inline_wasm_udfs.md @@ -187,7 +187,7 @@ To test the UDF delta file, use the provided UDF tools. ```sh UDF_DELTA=path/to/udf/delta - GLOG_v=10 dist/debian/udf_delta_file_tester --input_arguments="$TEST_KEY" --kv_delta_file_path="$KV_DELTA" --udf_delta_file_path="$UDF_DELTA" + dist/debian/udf_delta_file_tester --input_arguments="$TEST_KEY" --kv_delta_file_path="$KV_DELTA" --udf_delta_file_path="$UDF_DELTA" --v=10 ``` See the [generating UDF files doc](./generating_udf_files.md#3-test-the-udf-delta-file) for more diff --git a/docs/profiling_the_server.md b/docs/profiling_the_server.md index 25140ec6..797f349e 100644 --- a/docs/profiling_the_server.md +++ b/docs/profiling_the_server.md @@ -53,7 +53,6 @@ docker run \ --env CPUPROFILE=/data/profiles/server.cpu.prof \ --env CPUPROFILESIGNAL=12 \ --env LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libprofiler.so \ - --env GLOG_logtostder=1 \ --volume=/data:/data \ --network=host \ --security-opt=seccomp=unconfined \ @@ -61,7 +60,8 @@ docker run \ --cpus=4 \ --entrypoint=/server \ bazel/production/packaging/local/data_server:server_profiling_docker_image \ - --port 50051 -delta_directory=/data --realtime_directory=/data/realtime + --port 50051 -delta_directory=/data --realtime_directory=/data/realtime \ + --stderrthreshold=0 ``` **STEP 2:** Run the following command to profile the server for 10 seconds and generate a CPU @@ -106,7 +106,6 @@ docker run \ --env HEAPPROFILE=/data/profiles/server.heap.hprof \ --env CPUPROFILESIGNAL=12 \ --env LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so \ - --env GLOG_logtostder=1 \ --volume=/data:/data \ --network=host \ --security-opt=seccomp=unconfined \ @@ -114,7 +113,8 @@ docker run \ --cpus=4 \ --entrypoint=/server \ bazel/production/packaging/local/data_server:server_profiling_docker_image \ - --port 50051 -delta_directory=/data --realtime_directory=/data/realtime + --port 50051 -delta_directory=/data --realtime_directory=/data/realtime \ + --stderrthreshold=0 ``` The command above may generate a number of heap profiles during startup depending on how much data @@ -185,7 +185,6 @@ List of pre-defined events (to be used in -e): ```bash docker run \ -it --init --rm --name=server-profiling-container \ - --env GLOG_logtostder=1 \ --volume=/data:/data \ --network=host \ --security-opt=seccomp=unconfined \ @@ -193,7 +192,8 @@ docker run \ --cpus=4 \ --entrypoint=/server \ bazel/production/packaging/local/data_server:server_profiling_docker_image \ - --port 50051 -delta_directory=/data --realtime_directory=/data/realtime + --port 50051 -delta_directory=/data --realtime_directory=/data/realtime \ + --stderrthreshold=0 ``` **STEP 2:** Run the following command to grap the PID of the server process running inside the diff --git a/docs/protected_app_signals/ad_retrieval_overview.md b/docs/protected_app_signals/ad_retrieval_overview.md new file mode 100644 index 00000000..ddfef42f --- /dev/null +++ b/docs/protected_app_signals/ad_retrieval_overview.md @@ -0,0 +1,407 @@ +## Background + +This document provides a detailed overview of the Ad Retrieval server, which is a server-side +component of the Protected App Signals (PAS) API. The Ad Retrieval server must run within a trusted +execution environment (TEE). + +The PAS API flow at a high level is: + +1. Ad tech companies would curate signals from a variety of sources, including first-party data, + contextual data, and third-party data. +1. The signals would be stored securely on the device. +1. Ad tech companies would have access to the signals during a Protected Auction to serve relevant + ads. +1. **Ad tech custom logic deployed in Trusted execution environments can access the signals to do + real-time ad retrieval and bidding** +1. Buyers submit ads with bids to sellers for final scoring to choose a winning ad to render. + +The Ad retrieval service is a proposed solution for performing ad and creative targeting, as listed +in the step #4. In general, ad retrieval typically includes ads matching, filtering, scoring, +ranking and top-k selection. It is developed based on the Trusted Execution Environment (TEE) on +selected public cloud platforms. + +Today, the ad retrieval service is implemented as part of the Privacy Sandbox's TEE key value +service. The privacy characteristics conform to +[the KV service trust model](https://github.com/privacysandbox/fledge-docs/blob/main/key_value_service_trust_model.md). + +## Use case overview + +Today ad techs curate data from multiple sources and use this data to choose relevant ads for the +user. Ad requests are typically accompanied by an AdID and publisher provided data (ex OpenRTB) +which is used to determine most relevant ads for the user based on a user profile keyed on the AdiD. + +Ad candidates are filtered down to a few relevant ones for an ad request. This is the retrieval +phase and the focus of this document. Bids are computed for the selected set of ads. Top bids are +sent to the seller who would score and pick the winning ad. The winning ad is then rendered by the +client. + +![alt_text](../assets/ad_retrieval_use_case_overview.png 'use case overview') + +#### Figure 1: High level layout of the retrieval flow + +An ad tech has up to N active ad campaigns at any given time. To choose relevant ads for each ad +retrieval request an ad tech needs to filter the ads to the top K number of ads. To do this ad techs +typically follow a multi stage filtering process that looks something like the following: + +![alt_text](../assets/ad_retrieval_filter_funnel.png 'image_tooltip') + +### Coarse-grained selection + +Ad techs filter a large number of ads (e.g., 10,000s) to a few (e.g., 1,000s) using device and +publisher provided signals, such as geo, ad slot type size, language, etc. + +### Filtering + +This filtering will further reduce the number of ads that are eligible to be shown. For example, ad +tech uses real-time signals, such as the amount of budget left for campaigns and other information, +status of the campaign (active / inactive etc.) to further reduce the number of ads from thousands +to hundreds. + +### Lightweight scoring and Top-K selection + +From the remaining ad candidates, ad techs can reduce the number of results returned by scoring, +sorting, and truncating the list to the top K results. The lightweight scores can be computed using +inputs such as the embedding vectors sent in via the request and per-ad embedding vectors stored in +the Ads dataset. The set of candidates fetched at this stage can be further narrowed down during a +bidding phase using more powerful inference capabilities. The bidding phase is outside the scope of +the retrieval service. + +Note: This is just a general overview of the ad tech selection process. The specific steps and +processes involved may vary depending on the individual ad tech implementation. + +## Ad retrieval walkthrough + +![alt_text](../assets/ad_retrieval_walkthrough.png 'walkthrough') + + + +#### Figure 2: Typical setup in one region. Components in and surrounded by purple are part of the PAS system. Components in yellow are developed by the ad tech operator, proprietary to the ad tech and may differ vastly between one ad tech and another. + + + +### How to deploy the system and load data to it + +This section describes how data is loaded into the server. + +#### Deployment + +The ad tech builds the service system by downloading the source code from the +[Github repository](https://github.com/privacysandbox/fledge-key-value-service) and following the +documentation in the repository. + +The ad tech deploys the system to a supported public cloud of their choice. At time of publication, +the system will be available on GCP and AWS. + +The service runs in the same network as the Bidding & Auction services. The Bidding & Auction +services' configuration is updated to send Ad retrieval requests to the Ad Retrieval service. + +#### Data loading + +The ad retrieval service consumes data generated from ad techs' own custom data generation systems. +The retrieval service development team anticipates that various forms of data generation systems +will be used, such as periodic batch files, constant streaming files or pub/sub: + +- For data that require low latency propagation, such as budgets, the ad tech bundles the data + into a specific format of file and pushes the data file into the Cloud's pub/sub system, which + then pushes the data to the servers. +- For data that does not require low latency propagation, the ad tech writes them into the same + format of files, uploads them to the Cloud's blob storage system such as AWS S3, and the service + will pick up the data. + +For privacy reasons, the service loads data into its RAM and serves requests from there. + +For details, see +"[Loading data into the key value server](https://team.git.corp.google.com/kiwi-air-force-eng-team/kv-server/+/refs/heads/main/docs/loading_data.md)". + +#### Data model + +The dataset in the server RAM is organized as simple key value pairs. The key is a string and value +can either be a string or a set of strings. + +##### Design pattern: Ad hierarchy + +The data model being simple KV pairs does not mean all data must be flat. The ad tech may use a +hierarchical model such as Campaign -> Ad Group -> Creative. The ad tech can use these as key prefix +to represent "tables" such as `"campaign_123_adgroup"`. + +At request processing time, it is up to the ad tech to perform "joins". The ad tech can first look +up certain campaigns, then look up ad groups for those campaigns. This allows ad techs to store +campaign-level information only once in the campaign entry. + +##### Design pattern: use "set" to represent a group of values, such as ads or ad groups that share a certain feature + +There are 2 categories of values: singular values and sets. + +- Singular values are the classic building blocks of key value lookups. The server can load rows + of data where each row is a key value pair and the value is a singular value. It is the most + versatile way to model the data. +- Sets are supported to optimize common operations used in ad matching, which includes union, + intersection and difference. During data loading, each row can specify which type the row uses. + In this case the key value pair is a key and a set of values. + +A typical use case of sets is to have the key be a feature (or the lack of such feature) and its set +the group of ads/campaigns/etc that satisfies this feature. During request processing, the request +may be looking for ads that have M features and do not have N features. The user-defined retrieval +logic (described below) would perform an intersection and difference of the ads of the features. +Using sets and the accompanying RunQuery API (described below) would save the user the effort of +implementing this. + +##### Design pattern: versioning + +Cross-key transactions are not supported. A mutation to one key value pair is independent of other +updates. If a collection of ads metadata requires query integrity, such that the queried ads +metadata of certain campaign or customer must be from the same data snapshot (This only applies to +when the upstream data generation is periodic batch based), one way to ensure it is to append a +per-snapshot version to each key, and have a specific key/value pair for the version. During +lookups, the version key is first queried so the version can be used to construct the actual ad +metadata query. + +The ad tech defines the order of mutations by adding a logical timestamp to each mutation. The +service accepts/evicts records based on the timestamp but there is no guarantee that the service +receives the mutations in chronological order. + +### Processing requests with ad tech-specific User Defined Functions + +This phase explains the components 2 ("input signals") and 3 ("retrieval logic") in +[Figure 1](#figure-1-high-level-layout-of-the-retrieval-flow) + +Ad techs cannot change the server code inside TEE. But the server provides a flexible way to execute +custom logic through "User defined functions" ("UDF"). Today, UDF supports JavaScript or WebAssembly +(WASM) which supports many languages. + +A UDF is not a single function. Similar to running a JavaScript on a web page load where it can call +other dependency libraries during the execution, a UDF entrypoint is invoked by the server on each +retrieval request and it can call other functions possibly defined in other files and make use of +all the language constructs like polymorphism, templating, etc. A UDF refers to the collection of +all the code that will run inside the server. + +The ad tech writes the UDF code in their preferred way, converts it into a file with specific format +recognizable by the service, and stores it in a Cloud blob storage location. The service monitors +the location and picks up the UDF as it appears. + +There is one UDF invocation per request (one request per auction). + +#### UDF API + +In the initial version, the server uses a JavaScript wrapper layer even if WASM is used. So for +`byte[]` inputs, the input will be a base64 encoded string. + +```javascript +string HandleRequest( + requestMetadata, + protectedSignals, + deviceMetadata, + contextualSignals, + contextualAdIds, +) +``` + +- requestMetadata: JSON. Per-request server metadata to the UDF. Empty for now. +- protectedSignals: arbitrary string, originating from the device, and passed from the bidding + service. The Protected App Signals can be decoded in the UDF to produce embeddings for top K ads + ranking, and other information useful for the retrieval. Note: at this stage the Protected App + Signals would be decoded and unencrypted. +- deviceMetadata: JSON object containing device metadata forwarded by the Seller's Ad Service. See + the + [B&A documentation](https://github.com/privacysandbox/fledge-docs/blob/main/bidding_auction_services_api.md#metadata-forwarded-by-sellers-ad-service) + for further details. + - `X-Accept-Language`: language used on the device. + - `X-User-Agent`: User Agent used on the device. + - `X-BnA-Client-IP`: Device IP address. + - Example: + +```javascript +{ + "X-Accept-Language": "en-US", + "X-User-Agent": "ExampleAgent", + "X-BnA-Client-IP": "1.1.1.1" +} +``` + +- [contextualSignals:](https://github.com/privacysandbox/bidding-auction-servers/blob/b222e359f09de60f0994090f7a57aa796e927345/api/bidding_auction_servers.proto#L945) + arbitrary string originated from the contextual bidding server operated by the same DSP. The UDF + is expected to be able to decode the string and use it. Contextual Signals may contain any + information such as ML model version information for the protected embedding passed in via + Protected App Signals. +- contextualAdIds: JSON object containing an optional list of ad ids. +- Output: string. This will be sent back to the bidding service and passed into the `generateBid` + function for bid generation. + +![alt_text](../assets/ad_retrieval_udf.png 'ad_retrieval_udf') + +#### Figure 3: Zoomed-in view of request path. + +#### API available to UDF + +While the server invokes the UDF to process requests, the UDF has access to a few APIs provided by +the server to assist certain operations. The usage will be explained concretely in the example +below. + +- `getValues([key_strings])`: Given a list of keys, performs lookups in the loaded dataset and + returns a list of values corresponding to the keys. +- `runQuery(query_string)`: UDF can construct a query to perform set operations, such as union, + intersection and difference. The query uses keys to represent the sets. The keys are defined as + the sets are loaded into the dataset. See the exact grammar + [here](https://github.com/privacysandbox/fledge-key-value-service/blob/main/components/query/parser.yy). + +For more information, see +[the UDF spec](https://github.com/privacysandbox/fledge-docs/blob/main/key_value_service_user_defined_functions.md). + +#### Example + +The server is loaded with a large amount of ad candidates and their associated metadata. In this +case we limit the metadata to an embedding and a budget. The embedding is a vector of floats and the +budget is an integer USD cents. + +We also have loaded the server with a mapping for each Protected App Signals and Device Metadata to +every ad that is associated with it (ex. `Games->[ad_group1, ad_group2, ad_group3,...]`) + +The Protected App Signals contains a list of the of apps that have store signals on the device by +the adtech. As a reminder signals are only availbe to teh adtech whom stored the signals.. + +Ad tech writes the UDF for example in C++. They can use all the C++ primitives such as classes, +polymorphism, templates, etc. They use a compiler to compile the code into WASM. The entrypoint of +the code has access to the input. + +The UDF goes through multiple stages of logic. + +#### Stage 1: Ad matching + +The ad tech wants to find some ads related to the types of apps installed on a device, in this +example "games" and "news". This can be done by first looking up the mapping, finding the list of +ads associated with games and news then performing an intersection of the 2 lists. + +The service provides a few APIs to UDFs to access loaded data. The above operation can be done by +using the "RunQuery" API. The ad tech constructs a query to find the intersection of ad groups +associated with certain app types: `"games & news"` and calls `RunQuery("games & news")` which +returns `["ad group1", "ad group3"]`. + +##### Design pattern: Campaign liveness + +The ad tech wants to exclude certain ad groups from bidding if budget disallows. They model this +with a key/value pair where the key is "`disabled_adgroups`" and the value is a list of ad groups +`["ad group123", "ad group789", ...]`. + +The ad tech updates this key with the low latency pub/sub mechanism whenever an ad group encounters +a budget issue. During ad matching, the ad tech performs a difference in its query to exclude these +ad groups, e.g., `RunQuery("(games & news) - disabled_adgroups")`. + +#### Stage 2: Filtering + +A large set of ad candidates is produced by the matching stage. These candidates are then fed into a +filtering stage, which can contain any number of filters. One example of a filter is negative +filtering, which is used to prevent an app-install ad from being shown for an app that is already +installed on the device. + +The device can store the information of "apps already installed on the device" as the protected app +signals sent to the ad retrieval server. The specific design of the information can take various +forms, such as a collection of hashes of the package names of the apps. + +If the filter computes that the corresponding app of an ad candidate exists on the device, the ad +candidate is dropped. + +#### User constraints + +If the dataset is large, the ad tech would need to enable sharding to store it in multiple clusters +of machines. RunQuery and GetValues APIs would perform remote calls which are slower than local +executions when the dataset is small. + +### Scoring and Top-K selection (Last UDF stage) + +This operation corresponds to the previously mentioned lightweight scoring and Top-K selection, +which happens after the ad matching and filtering. + +#### Design pattern: Dot product selection + +Suppose the ad tech wants to perform dot-product based scoring. + +After filtering, the UDF has a relatively small amount of ads and their metadata. + +The UDF logic uses GetValues API to query the ad embeddings of these ads, previously loaded into the +system through earlier phases. + +In the request input there are the protected user embeddings generated by the Bidding Service using +Protected App Signals. The UDF could perform a dot product between user and contextual embedding and +an ad embedding that would come from the retrieval data to rank the results and select top K ads. + +The UDF may perform another GetValues call on the K ads to look up additional metadata. The UDF then +aggregates them and constructs the final response according to the UDF return value format and +returns it. The service takes the return value and returns it to the calling Bidding service, which +performs further bid generation logic. + +#### Design pattern: versioning of embeddings + +The ad tech should specify which version of model/embedding should be used so in the case they are +using +[model factorization](https://developer.android.com/design-for-safety/privacy-sandbox/protected-app-signals#model-factorization) +the ad embeddings and user embeddings' versions match. The version can be passed in as part of the +per_buyer_signals in the request. In the dataset, the embeddings' keys have a version. The UDF +constructs the keys by conforming to the naming convention defined by the ad tech. For example the +version in the request may be `"5"` and the UDF can build a key: `"ad_123_embedding_v5"`. + +#### Design pattern: linear regression ML selection + +Instead of a dot product, the ad tech may link some ML libraries with their UDF and perform ML-based +selection. While in general this follows the same flow as the dot-product, there are multiple +constraints for this approach. + +#### User constraints: + +The retrieval server does not provide additional ML support on the server framework level. I.e., all +ML operations can only be performed within the UDF space. If WASM has constraints supporting e.g., +Tensorflow, it will be constraints of the retrieval service. There is also no accelerator. + +Same as data and UDF, the ML model must be looked up before using it. The lookup may be a remote +call. And the model needs to be small to not contend with other RAM usage, which is prominent in the +retrieval server. + +Data loading has propagation delay. It may take minutes for every replica of the retrieval server to +have key/value pairs of a certain version. When setting the version in the request, it is better to +not set the latest version right away lest the lookups fail. This applies in general to all lookups. +However, the server will provide visibility to the state of the data loading. + +### Ad retrieval output + +The UDF should return a string on success. The string is returned to the bidding server which then +passes it to the `generateBid()` function. Although the string can just be a simple string, most +likely the string should be a serialized object whose schema is defined by each ad tech on their +own. There is no constraint on the schema as long as the ad tech's `generateBid()` logic can +recognize and use the string. + +The reason that the string likely should be a serialized object is we expect that the use case would +require the ad retrieval server output to contain a dictionary of ad candidates and their +information. The ad candidates can be keyed by certain unique ids and the information may contain +the ad embeddings, and various pieces to construct the final rendering URL, such as the available +sizes of the ad. The `generateBid()` function may then generate the final render URL +`https://example.com/ads/4/123?w=1200&h=628` using contextual information. + +Example schema: + +```javascript +{ // Ad candidates and their information + "123": { + "campaign": 4, + "sizes": [[1200, 628], [1200, 1500]], + "embeddings": "ZXhhbXBsZSBlbWJlZGRpbmc=", + "rank": 1 + }, + "456": { + "campaign": 6, + "sizes": [[1200, 628], [1200, 1500]], + "embeddings": "YW5vdGhlciBleGFtcGxlIGVtYmVkZGluZw==" + "rank": 2 + } +} +``` + +## Concrete example and specification + +The service codebase has an end-to-end example in the context of information retrieval. + +Documentation: + +- High level: +- Lower level: + +Getting started & Example: + diff --git a/docs/protected_app_signals/examples/BUILD.bazel b/docs/protected_app_signals/examples/BUILD.bazel new file mode 100644 index 00000000..7ff15a52 --- /dev/null +++ b/docs/protected_app_signals/examples/BUILD.bazel @@ -0,0 +1,63 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +load("@bazel_skylib//rules:run_binary.bzl", "run_binary") + +package(default_visibility = [ + "//production/packaging/tools:__subpackages__", + "//tools:__subpackages__", +]) + +run_binary( + name = "generate_delta", + srcs = [ + ":ad_retrieval.csv", + ], + outs = [ + "DELTA_0000000000000001", + ], + args = [ + "format_data", + "--input_file", + "$(location :ad_retrieval.csv)", + "--input_format", + "CSV", + "--output_file", + "$(location DELTA_0000000000000001)", + "--output_format", + "DELTA", + ], + tags = ["manual"], + tool = "//tools/data_cli", +) + +run_binary( + name = "ad_retrieval_udf", + srcs = [ + ":ad_retrieval_udf.js", + ], + outs = [ + "DELTA_0000000000000002", + ], + args = [ + "--udf_file_path", + "$(location :ad_retrieval_udf.js)", + "--output_path", + "$(location DELTA_0000000000000002)", + "--logical_commit_time", + "1700000000", + ], + tags = ["manual"], + tool = "//tools/udf/udf_generator:udf_delta_file_generator", +) diff --git a/docs/protected_app_signals/examples/ad_retrieval.csv b/docs/protected_app_signals/examples/ad_retrieval.csv new file mode 100644 index 00000000..1d774d4a --- /dev/null +++ b/docs/protected_app_signals/examples/ad_retrieval.csv @@ -0,0 +1,2 @@ +key,logical_commit_time,mutation_type,value,value_type +ad_id1,1680815895468056,UPDATE,ad_id1_value,string diff --git a/docs/protected_app_signals/examples/ad_retrieval_udf.js b/docs/protected_app_signals/examples/ad_retrieval_udf.js new file mode 100644 index 00000000..389f5ecd --- /dev/null +++ b/docs/protected_app_signals/examples/ad_retrieval_udf.js @@ -0,0 +1,30 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +function HandleRequest(requestMetadata, protectedSignals, deviceMetadata, contextualSignals, contextualAdIds) { + let protectedSignalsKeys = []; + const parsedProtectedSignals = JSON.parse(protectedSignals); + for (const [key, value] of Object.entries(parsedProtectedSignals)) { + protectedSignalsKeys.push(key); + } + const kv_result = JSON.parse(getValues(protectedSignalsKeys)); + if (kv_result.hasOwnProperty('kvPairs')) { + return kv_result.kvPairs; + } + const error_message = 'Error executing handle PAS:' + JSON.stringify(kv_result); + console.error(error_message); + throw new Error(error_message); +} diff --git a/docs/protected_app_signals/onboarding_dev_guide.md b/docs/protected_app_signals/onboarding_dev_guide.md new file mode 100644 index 00000000..a0b3fe2b --- /dev/null +++ b/docs/protected_app_signals/onboarding_dev_guide.md @@ -0,0 +1,210 @@ +# Protected App Signals, Ad Retrieval developer guide + +This Ad Retrieval guide is a natural extension of this PAS +[guide](https://developer.android.com/design-for-safety/privacy-sandbox/guides/protected-audience/protected-app-signals#ad-retrieval) +which assumes a BYOS set up, but it could use a TEE Ad Retrieval server instead. + +This guide provides sample retrieval and lookup usecases. To support these usecases, details on how +to set up a test environment in the same VPC as the B&A server and reuse the same mesh, upload a +sample UDF function and upload a sample Delta file are given. + +For an overview of the Protected App Signals API, read the design +[proposal.](https://developer.android.com/design-for-safety/privacy-sandbox/protected-app-signals) + +For an overview of the Key Value/ Ad Retrieval Protected App Signals, read the the following +[doc.](/docs/protected_app_signals/ad_retrieval_overview.md) + +## Deploying an Ad Retrieval server + +[GCP.](/docs/deployment/deploying_on_gcp.md) Please follow the `B&A integration within the same VPC` +section there. + +AWS support is coming later. + +## Example Setup + +This script is an extension of the example scenario listed +[here.](https://developer.android.com/design-for-safety/privacy-sandbox/guides/protected-audience/protected-app-signals#example-setup) + +Consider the following scenario: using the Protected App Signals API, an ad tech stores relevant +signals based on user app usage. In our example, signals are stored that represent in-app purchases +from several apps. During an auction, the encrypted signals are collected and passed into a +Protected Auction running in B&A. The buyer's UDFs running in B&A use the signals to fetch ad +candidates and compute a bid. + +The fetching happens by calling the Ad Retrieval server. + +### Loading data + +For privacy reasons, the server loads all necessary data in the form of files into its RAM and +serves all requests with the in-RAM dataset. It loads the existing data at startup and updates its +in-RAM dataset as new files appear. + +The files are called "Delta files", which are similar to database journal files. Each delta file has +many rows of records and each record can be an update or delete. + +The delta file type is [Riegeli](https://github.com/google/riegeli) and the record format is +[Flatbuffers](https://flatbuffers.dev/). + +Now let's add some data. There is some example data in +[examples/ad_retrieval.csv](./examples/ad_retrieval.csv). The +[build definition](./examples/BUILD.bazel) has predefined the command to use `data_cli` to generate +the data. + +```sh +./builders/tools/bazel-debian build //docs/protected_app_signals/examples:generate_delta +``` + +##### GCP + +Upload the delta file to the bucket + +```sh +export GCS_BUCKET=your-gcs-bucket-id +gsutil cp bazel-bin/docs/protected_app_signals/examples/DELTA_0000000000000001 gs://${GCS_BUCKET} +``` + +More [details](../data_loading/loading_data.md#upload-data-files-to-gcp) + +### Retrieval Path + +Note that this is a simplistic example created for an illustrative purpose. The retrieval case can +get more complicated. + +See this [example](/getting_started/examples/sample_word2vec/). The sample demonstrates how you can +query for a set of words, taking advantage of the native set query support, and sort them based on +scoring criteria defined by word similarities, using embeddings. + +### UDF + +This [script](examples/ad_retrieval_udf.js) looks up and returns values for the keys specified in +`protectedSignals`. + +```javascript +function HandleRequest(requestMetadata, protectedSignals, deviceMetadata, contextualSignals) { + let protectedSignalsKeys = []; + const parsedProtectedSignals = JSON.parse(protectedSignals); + for (const [key, value] of Object.entries(parsedProtectedSignals)) { + protectedSignalsKeys.push(key); + } + const kv_result = JSON.parse(getValues(protectedSignalsKeys)); + if (kv_result.hasOwnProperty('kvPairs')) { + return kv_result.kvPairs; + } + const error_message = 'Error executing handle PAS:' + JSON.stringify(kv_result); + console.error(error_message); + throw new Error(error_message); +} +``` + +#### Loading the UDF + +UDFs are also ingested through Delta files. Build the delta file: + +```sh +./builders/tools/bazel-debian build //docs/protected_app_signals/examples:ad_retrieval_udf +``` + +##### GCP + +Upload the delta file to the bucket + +```sh +export GCS_BUCKET=your-gcs-bucket-id +gsutil cp bazel-bin/docs/protected_app_signals/examples/DELTA_0000000000000002 gs://${GCS_BUCKET} +``` + +More [details](../data_loading/loading_data.md#upload-data-files-to-gcp) + +#### Input + +```proto +partitions { + arguments { + data { + string_value: "{\"ad_id1\":1}" + } + } + arguments { + data { + struct_value { + fields { + key: "X-Accept-Language" + value { + string_value: "en-US,en;q=0.9" + } + } + fields { + key: "X-BnA-Client-IP" + value { + string_value: "104.133.126.32" + } + } + fields { + key: "X-User-Agent" + value { + string_value: "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36" + } + } + } + } + } + arguments { + data { + string_value: "{\"h3\": \"1\"}" + } + } + arguments { + } +} +``` + +#### Output + +```proto +single_partition { + string_output: "{\"ad_id1\":{\"value\":\"ad_id1_value\"}}" +} +``` + +### Lookup Path + +The list of IDs is provided by the buyer in the contextual path. The server will lookup the data +associated with the IDs. + +No need to load a UDF here, since the default UDF for PAS usecase should handle that. + +#### Input + +```proto +metadata { + fields { + key: "is_pas" + value { + string_value: "true" + } + } +} +partitions { + arguments { + data { + list_value { + values { + string_value: "ad_id1" + } + values { + string_value: "ad_id2" + } + } + } + } +} +``` + +#### Output + +```proto +single_partition { + string_output: "{\"ad_id1\":{\"value\":\"ad_id1_value\"},\"ad_id2\":{\"status\":{\"code\":5,\"message\":\"Key not found: ad_id2\"}}}" +} +``` diff --git a/docs/integrating_with_fledge.md b/docs/protected_audience/integrating_with_fledge.md similarity index 97% rename from docs/integrating_with_fledge.md rename to docs/protected_audience/integrating_with_fledge.md index f4e51b3f..ad603585 100644 --- a/docs/integrating_with_fledge.md +++ b/docs/protected_audience/integrating_with_fledge.md @@ -51,7 +51,7 @@ navigator.joinAdInterestGroup(interestGroup, 7 * kSecsPerDay); The ad interest group properties can be inspected in DevTools. Open the **Application** tab and select **Interest Group** from the sidebar. -![inspecting bidding signals in devtools](assets/devtools_bidding_signals.png) +![inspecting bidding signals in devtools](../assets/devtools_bidding_signals.png) When the auction is executed at a later time, the browser will use the keys defined when the user was added to an interest group to query the Key/Value server. The trusted bidding signal values will @@ -177,7 +177,7 @@ scoreAd(adMetadata, bid, adSelectionConfig, sellerSignals, trustedScoringSignals # Bidding and Auction Services with Key-Value Integration -![architecture of FLEDGE services](assets/fledge_services_architecture.png) +![architecture of FLEDGE services](../assets/fledge_services_architecture.png) Bidding & Auction services is a way to allow FLEDGE computation to take place on cloud servers in a trusted execution environment, rather than running locally on a user's device. Bidding & Auction diff --git a/docs/working_with_terraform.md b/docs/working_with_terraform.md deleted file mode 120000 index 61432b3f..00000000 --- a/docs/working_with_terraform.md +++ /dev/null @@ -1 +0,0 @@ -../production/terraform/README.md \ No newline at end of file diff --git a/getting_started/examples/sample_word2vec/README.md b/getting_started/examples/sample_word2vec/README.md index 875810ca..3b1c56df 100644 --- a/getting_started/examples/sample_word2vec/README.md +++ b/getting_started/examples/sample_word2vec/README.md @@ -75,10 +75,9 @@ Build the local server: Run the local server: ```sh -GLOG_alsologtostderr=1 \ ./bazel-bin/components/data_server/server/server \ --delta_directory=/tmp/deltas \ - --realtime_directory=/tmp/realtime + --realtime_directory=/tmp/realtime --stderrthreshold=0 ``` ## Send a query diff --git a/getting_started/quick_start.md b/getting_started/quick_start.md index ab89db41..befbbafc 100644 --- a/getting_started/quick_start.md +++ b/getting_started/quick_start.md @@ -137,8 +137,8 @@ The delta file type is [Riegeli](https://github.com/google/riegeli) and the reco [Flatbuffers](https://flatbuffers.dev/). Now let's add some data. There is some example data in -[tools/udf/udf_tester/example_data.csv](/tools/udf/udf_tester/example_data.csv). The -[build definition](/getting_started/examples/canonical_examples/BUILD.bazel) has predefined the +[/getting_started/examples/canonical_examples/example_data.csv](/getting_started/examples/canonical_examples/example_data.csv). +The [build definition](/getting_started/examples/canonical_examples/BUILD.bazel) has predefined the command to use `data_cli` to generate the data. ```sh @@ -173,7 +173,7 @@ curl http://localhost:51052/v1/getvalues?keys=example_key } ``` -See [here](/docs/loading_data.md) for more information about data loading. +See [here](/docs/data_loading/loading_data.md) for more information about data loading. ## Use `User Defined Functions (UDF)` to process requests @@ -205,7 +205,7 @@ cp bazel-bin/getting_started/examples/canonical_examples/DELTA_0000000000000002 (Similar to the data file, the UDF file can also be generated with building the tool specified in the build target and running it with your own command line flags. See -[details](docs/generating_udf_files.md).) +[details](/docs/generating_udf_files.md).) And query: @@ -263,7 +263,7 @@ function getKeyGroupOutputs(hostname, udf_arguments) { } function HandleRequest(executionMetadata, ...udf_arguments) { - logMessage(JSON.stringify(executionMetadata)); + console.log(JSON.stringify(executionMetadata)); const keyGroupOutputs = getKeyGroupOutputs(executionMetadata.requestMetadata.hostname, udf_arguments); return {keyGroupOutputs, udfOutputApiVersion: 1}; } @@ -315,8 +315,8 @@ At this point we have looked at all the basic components. See the following spec advanced topics and features: - [Writing WebAssembly User defined functions:](/docs/inline_wasm_udfs.md) -- [Deploying on AWS](/docs/deploying_on_aws.md) -- [Deploying on GCP](/docs/deploying_on_gcp.md) -- [Sharding](/docs/sharding.md) -- [Working with Terraform](/docs/working_with_terraform.md) +- [Deploying on AWS](/docs/deployment/deploying_on_aws.md) +- [Deploying on GCP](/docs/deployment/deploying_on_gcp.md) +- [Sharding](https://github.com/privacysandbox/protected-auction-services-docs/blob/main/key_value_service_sharding.md) +- [Working with Terraform](/docs/deployment/working_with_terraform.md) - [UDF binary data API](/docs/udf_read_apis_with_binary_data.md) diff --git a/getting_started/quick_start_assets/docker-compose.yaml b/getting_started/quick_start_assets/docker-compose.yaml index bf8bbd4f..8b886131 100644 --- a/getting_started/quick_start_assets/docker-compose.yaml +++ b/getting_started/quick_start_assets/docker-compose.yaml @@ -19,11 +19,9 @@ services: context: ../.. dockerfile: getting_started/quick_start_assets/Dockerfile.server network_mode: host - environment: - - GLOG_alsologtostderr=1 - - GLOG_v=9 security_opt: - seccomp:unconfined + command: --v=9 --stderrthreshold=0 volumes: - ../../dist/deltas:/tmp/deltas diff --git a/infrastructure/testing/BUILD.bazel b/infrastructure/testing/BUILD.bazel index b7bc24d0..76ea12e2 100644 --- a/infrastructure/testing/BUILD.bazel +++ b/infrastructure/testing/BUILD.bazel @@ -22,12 +22,12 @@ cc_binary( deps = [ ":protocol_testing_helper_server_cc_grpc", "//public:constants", - "@com_github_google_glog//:glog", "@com_github_google_quiche//quiche:binary_http_unstable_api", "@com_github_google_quiche//quiche:oblivious_http_unstable_api", "@com_github_grpc_grpc//:grpc++", "@com_github_grpc_grpc//:grpc++_reflection", # for grpc_cli "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", ], ) diff --git a/infrastructure/testing/protocol_testing_helper_server.cc b/infrastructure/testing/protocol_testing_helper_server.cc index 4eac2f2f..ec136142 100644 --- a/infrastructure/testing/protocol_testing_helper_server.cc +++ b/infrastructure/testing/protocol_testing_helper_server.cc @@ -16,10 +16,10 @@ #include #include "absl/flags/flag.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/strings/escaping.h" #include "absl/strings/str_cat.h" -#include "glog/logging.h" #include "grpcpp/ext/proto_server_reflection_plugin.h" #include "grpcpp/grpcpp.h" #include "infrastructure/testing/protocol_testing_helper_server.grpc.pb.h" diff --git a/production/packaging/aws/build_and_test b/production/packaging/aws/build_and_test index ff75fefe..6be00c9b 100755 --- a/production/packaging/aws/build_and_test +++ b/production/packaging/aws/build_and_test @@ -33,6 +33,8 @@ function _print_runtime() { exit ${STATUS} } +declare MODE=prod + function usage() { local exitval=${1-1} cat &>/dev/stderr < Mode can be prod or nonprod. Default: ${MODE} environment variables (all optional): WORKSPACE Set the path to the workspace (repo root) @@ -71,6 +74,11 @@ while [[ $# -gt 0 ]]; do BUILD_AND_TEST_ARGS+=("--with-tests") shift ;; + --mode) + MODE="$2" + shift + shift + ;; --verbose) BUILD_AND_TEST_ARGS+=("--verbose") set -o xtrace @@ -115,7 +123,7 @@ if ! [[ -r ${WORKSPACE}/production/packaging/build_and_test_all_in_docker && -x printf "build_and_test script not found at location: %s/production/packaging/build_and_test_all_in_docker\n" "${WORKSPACE}" &>/dev/stderr fail "build_and_test not found" fi -if ! "${WORKSPACE}"/production/packaging/build_and_test_all_in_docker "${BUILD_AND_TEST_ARGS[@]}" --instance aws; then +if ! "${WORKSPACE}"/production/packaging/build_and_test_all_in_docker "${BUILD_AND_TEST_ARGS[@]}" --instance aws --mode "${MODE}"; then fail "Failed to run build_and_test_all_in_docker" fi diff --git a/production/packaging/aws/data_server/BUILD.bazel b/production/packaging/aws/data_server/BUILD.bazel index 5c1d5dc7..b60c3563 100644 --- a/production/packaging/aws/data_server/BUILD.bazel +++ b/production/packaging/aws/data_server/BUILD.bazel @@ -12,16 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -load( - "@io_bazel_rules_docker//container:container.bzl", - "container_image", - "container_layer", -) -load("@io_bazel_rules_docker//contrib:test.bzl", "container_test") +load("@container_structure_test//:defs.bzl", "container_structure_test") +load("@rules_oci//oci:defs.bzl", "oci_image", "oci_tarball") load( "@rules_pkg//pkg:mappings.bzl", "pkg_attributes", "pkg_files", + "pkg_mklink", ) load("@rules_pkg//pkg:tar.bzl", "pkg_tar") load("@rules_pkg//pkg:zip.bzl", "pkg_zip") @@ -40,7 +37,7 @@ pkg_files( pkg_files( name = "kmstool_enclave_executables", srcs = [ - "@google_privacysandbox_servers_common//scp/cc/cpio/client_providers/kms_client_provider/src/aws:kms_cli", + "@google_privacysandbox_servers_common//src/cpio/client_providers/kms_client_provider/aws:kms_cli", ], attributes = pkg_attributes(mode = "0555"), prefix = "/cpio/bin", @@ -49,16 +46,25 @@ pkg_files( pkg_files( name = "kmstool_enclave_libs", srcs = [ - "@google_privacysandbox_servers_common//scp/cc/cpio/client_providers/kms_client_provider/src/aws:libnsm_so", + "@google_privacysandbox_servers_common//src/cpio/client_providers/kms_client_provider/aws:libnsm_so", ], attributes = pkg_attributes(mode = "0444"), prefix = "/cpio/lib", ) +# Create a symlink between where kmstool_enclave_cli expects shell to be +# (/bin/sh) and where it actually is on our image (/busybox/sh). +pkg_mklink( + name = "busybox_sh_symlink", + link_name = "/bin/sh", + target = "/busybox/sh", +) + server_binaries = [ ":kmstool_enclave_executables", ":kmstool_enclave_libs", ":server_executables", + ":busybox_sh_symlink", ] pkg_zip( @@ -75,43 +81,18 @@ pkg_tar( pkg_tar( name = "libnsm-tar", srcs = [ - "@google_privacysandbox_servers_common//scp/cc/cpio/client_providers/kms_client_provider/src/aws:libnsm_so", + "@google_privacysandbox_servers_common//src/cpio/client_providers/kms_client_provider/aws:libnsm_so", ], mode = "0444", package_dir = "/cpio/lib", visibility = ["//visibility:public"], ) -container_layer( - name = "server_binary_layer", - directory = "/", - tars = [ - ":libnsm-tar", - ":server_binaries_tar", - ], -) - -# Create a symlink between where kmstool_enclave_cli expects shell to be -# (/bin/sh) and where it actually is on our image (/busybox/sh). -container_layer( - name = "kmstool_enclave_cli_layer", - symlinks = { - "/bin/sh": "/busybox/sh", - }, - tars = [ - ":libnsm-tar", - ], -) - # This image target is meant for testing running the server in an enclave using. # # See project README.md on how to run the image. -container_image( +oci_image( name = "server_docker_image", - architecture = select({ - "@platforms//cpu:arm64": "arm64", - "@platforms//cpu:x86_64": "amd64", - }), base = select({ "@platforms//cpu:arm64": "@runtime-debian-debug-nonroot-arm64//image", "@platforms//cpu:x86_64": "@runtime-debian-debug-nonroot-amd64//image", @@ -123,32 +104,34 @@ container_image( "--", # Note: these ports must correspond with those specified in envoy.yaml. "--port=50051", - # These affect PCR0, so changing these would result in the loss of - # access to private keys for decryption. - "--public_key_endpoint='https://d3gf5400xe31j1.cloudfront.net/v1alpha/publicKeys'", - "--primary_coordinator_private_key_endpoint='https://uun5qzrqvj.execute-api.us-east-1.amazonaws.com/stage/v1alpha/encryptionKeys'", - "--secondary_coordinator_private_key_endpoint='https://ddqkl8ay59.execute-api.us-east-1.amazonaws.com/stage/v1alpha/encryptionKeys'", - "--primary_coordinator_region='us-east-1'", - "--secondary_coordinator_region='us-east-1'", + # These affect PCR0, so changing these would result in the loss of ability to communicate with + # the downstream components + "--public_key_endpoint=https://publickeyservice.pa.aws.privacysandboxservices.com/.well-known/protected-auction/v1/public-keys", + "--stderrthreshold=0", ], - env = {"GLOG_logtostderr": "1"}, - layers = [ - "@google_privacysandbox_servers_common//scp/cc/aws/proxy/src:proxify_layer", + tars = [ + "@google_privacysandbox_servers_common//src/aws/proxy:libnsm_and_proxify_tar", "//production/packaging/aws/resolv:resolv_config_layer", - ":server_binary_layer", - ":kmstool_enclave_cli_layer", + ":libnsm-tar", + ":server_binaries_tar", ], ) -container_test( +oci_tarball( + name = "server_docker_tarball", + image = ":server_docker_image", + repo_tags = ["bazel/production/packaging/aws/data_server:server_docker_image"], +) + +container_structure_test( name = "structure_test", size = "medium", configs = ["test/structure.yaml"], driver = "tar", - image = ":server_docker_image", + image = ":server_docker_tarball", ) -container_test( +container_structure_test( name = "commands_test", size = "small", configs = ["test/commands.yaml"], @@ -166,14 +149,15 @@ genrule( name = "copy_to_dist", srcs = [ ":server_artifacts", - ":server_docker_image.tar", + ":server_docker_tarball", "//public/query:query_api_descriptor_set", ], outs = ["copy_to_dist.bin"], cmd_bash = """cat << EOF > '$@' mkdir -p dist/debian cp $(execpath :server_artifacts) dist/debian -cp $(execpath :server_docker_image.tar) $(execpath //public/query:query_api_descriptor_set) dist +cp $(execpath :server_docker_tarball) dist/server_docker_image.tar +cp $(execpath //public/query:query_api_descriptor_set) dist # retain previous server_docker_image.tar location as a symlink ln -rsf dist/server_docker_image.tar dist/debian/server_docker_image.tar builders/tools/normalize-dist diff --git a/production/packaging/aws/data_server/ami/BUILD.bazel b/production/packaging/aws/data_server/ami/BUILD.bazel index b64d93b4..3e294b2f 100644 --- a/production/packaging/aws/data_server/ami/BUILD.bazel +++ b/production/packaging/aws/data_server/ami/BUILD.bazel @@ -21,7 +21,7 @@ pkg_zip( "//components/aws:sqs_lambda.tar", "//production/packaging/aws/otel_collector:aws-otel-collector.rpm", "//production/packaging/aws/otel_collector:aws_otel_collector_cfg", - "@google_privacysandbox_servers_common//scp/cc/aws/proxy/src:proxy", + "@google_privacysandbox_servers_common//src/aws/proxy", ], ) diff --git a/production/packaging/aws/data_server/bin/BUILD.bazel b/production/packaging/aws/data_server/bin/BUILD.bazel index 4397d2c3..f105dfb8 100644 --- a/production/packaging/aws/data_server/bin/BUILD.bazel +++ b/production/packaging/aws/data_server/bin/BUILD.bazel @@ -7,6 +7,7 @@ package(default_visibility = [ cc_binary( name = "init_server_basic", srcs = ["init_server_main.cc"], + malloc = "@com_google_tcmalloc//tcmalloc", deps = [ "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", diff --git a/production/packaging/aws/data_server/nitro-pcr0/amd64.json b/production/packaging/aws/data_server/nitro-pcr0/amd64.json index ad2c912a..615805cc 100644 --- a/production/packaging/aws/data_server/nitro-pcr0/amd64.json +++ b/production/packaging/aws/data_server/nitro-pcr0/amd64.json @@ -1 +1 @@ -{"PCR0":"25c773b4d9cc6f2316dfd9565b2895367360db465d5519f240a5ffa3059952d3ffb72d75e2cf508bdb08652e696aaafb"} +{"PCR0":"76ece8efa6f8c5723edd64b3a0464ec72411651be7642c951c73bcbc42863a8f1c376970726109c0e718b2f0d326e39d"} diff --git a/production/packaging/aws/otel_collector/otel_collector_config.yaml b/production/packaging/aws/otel_collector/otel_collector_config.yaml index 887e5136..9c8aca32 100644 --- a/production/packaging/aws/otel_collector/otel_collector_config.yaml +++ b/production/packaging/aws/otel_collector/otel_collector_config.yaml @@ -37,14 +37,20 @@ processors: send_batch_size: 50 batch/metrics: timeout: 60s + batch/logs: + timeout: 60s exporters: + awscloudwatchlogs: + log_group_name: "kv-server-log-group" + log_stream_name: "kv-server-log-stream" awsxray: index_all_attributes: true awsemf: namespace: 'KV-Server' resource_to_telemetry_conversion: enabled: true + retain_initial_value_of_delta_metric: true prometheusremotewrite: endpoint: "https://aps-workspaces.$REGION.amazonaws.com/workspaces/$WORKSPACE_ID/api/v1/remote_write" auth: @@ -65,5 +71,8 @@ service: receivers: [otlp] processors: [batch/metrics] exporters: [prometheusremotewrite] - + logs: + receivers: [otlp] + processors: [batch/logs] + exporters: [awscloudwatchlogs] extensions: [health_check,sigv4auth] diff --git a/production/packaging/build_and_test_all_in_docker b/production/packaging/build_and_test_all_in_docker index 4ebca8fd..4cca3b85 100755 --- a/production/packaging/build_and_test_all_in_docker +++ b/production/packaging/build_and_test_all_in_docker @@ -45,6 +45,7 @@ declare -i SKIP_PRECOMMIT=0 declare -i RUN_TESTS=0 declare INSTANCE=local declare PLATFORM=aws +declare MODE=prod function usage() { local -r -i exitval=${1-1} @@ -53,6 +54,7 @@ usage: ${BASH_SOURCE[0]} --instance Instance can be local or aws. Default: ${INSTANCE} --platform Platform can be aws or local. Default: ${PLATFORM} + --mode Mode can be prod or nonprod. Default: ${MODE} --precommit-only Run precommit then exit --no-precommit Skip precommit checks --with-tests Also runs tests before building @@ -103,6 +105,11 @@ while [[ $# -gt 0 ]]; do shift shift ;; + --mode) + MODE="$2" + shift + shift + ;; --verbose) VERBOSE=1 shift @@ -112,7 +119,7 @@ while [[ $# -gt 0 ]]; do esac done -BAZEL_EXTRA_ARGS="${BAZEL_EXTRA_ARGS} --config=${INSTANCE}_instance --config=${PLATFORM}_platform" +BAZEL_EXTRA_ARGS="${BAZEL_EXTRA_ARGS} --config=${INSTANCE}_instance --config=${PLATFORM}_platform --config=${MODE}_mode" if [[ ${VERBOSE} -eq 1 ]]; then set -o xtrace @@ -135,7 +142,7 @@ set -o errexit trap _collect_logs EXIT function _collect_logs() { declare -r -i status\$? - declare -r filename=${INSTANCE}-${PLATFORM}-logs.zip + declare -r filename=${INSTANCE}-${PLATFORM}-${MODE}-logs.zip printf 'Collecting bazel logs [%s]... [status: %d]\n' \${filename@Q} \${status} &>/dev/stderr bazel ${BAZEL_STARTUP_ARGS} run ${BAZEL_EXTRA_ARGS} //:collect-logs \${filename} exit \${status} diff --git a/production/packaging/gcp/build_and_test b/production/packaging/gcp/build_and_test index 845b8378..36d7fae9 100755 --- a/production/packaging/gcp/build_and_test +++ b/production/packaging/gcp/build_and_test @@ -33,6 +33,8 @@ function _print_runtime() { exit ${STATUS} } +declare MODE=prod + function usage() { local exitval=${1-1} cat &>/dev/stderr < Mode can be prod or nonprod. Default: ${MODE} environment variables (all optional): WORKSPACE Set the path to the workspace (repo root) @@ -64,6 +67,11 @@ while [[ $# -gt 0 ]]; do BUILD_AND_TEST_ARGS+=("--with-tests") shift ;; + --mode) + MODE="$2" + shift + shift + ;; --verbose) BUILD_AND_TEST_ARGS+=("--verbose") set -o xtrace @@ -102,6 +110,6 @@ fi source "${BUILDER}" || fail "Failed to source builder.sh" printf "==== Running build_and_test_all_in_docker =====\n" -if ! "${WORKSPACE}"/production/packaging/build_and_test_all_in_docker "${BUILD_AND_TEST_ARGS[@]}" --instance gcp --platform gcp; then +if ! "${WORKSPACE}"/production/packaging/build_and_test_all_in_docker "${BUILD_AND_TEST_ARGS[@]}" --instance gcp --platform gcp --mode "${MODE}"; then fail "Failed to run build_and_test_all_in_docker" fi diff --git a/production/packaging/gcp/data_server/BUILD.bazel b/production/packaging/gcp/data_server/BUILD.bazel index c5cbc5f8..6e9e90ff 100644 --- a/production/packaging/gcp/data_server/BUILD.bazel +++ b/production/packaging/gcp/data_server/BUILD.bazel @@ -128,17 +128,12 @@ container_image( }), entrypoint = [ "/init_server_basic", - # These affect PCR0, so changing these would result in the loss of - # access to private keys for decryption. - "--public_key_endpoint='https://publickeyservice-test1.bas-kms.xyz/v1alpha/publicKeys'", - "--primary_coordinator_private_key_endpoint='https://privatekeyservice-test1.bas-kms.xyz/v1alpha/encryptionKeys'", - "--secondary_coordinator_private_key_endpoint='https://privatekeyservice-test2.bas-kms.xyz/v1alpha/encryptionKeys'", - "--primary_coordinator_region='us-central1'", - "--secondary_coordinator_region='us-central1'", + # These affect PCR0, so changing these would result in the loss of ability to communicate with + # the downstream components + "--public_key_endpoint=https://publickeyservice-a.pa-3.gcp.privacysandboxservices.com/.well-known/protected-auction/v1/public-keys", + "--stderrthreshold=0", ], env = { - "GLOG_logtostderr": "1", - "GLOG_stderrthreshold": "0", "GRPC_DNS_RESOLVER": "native", }, labels = {"tee.launch_policy.log_redirect": "always"}, diff --git a/production/packaging/gcp/data_server/bin/BUILD.bazel b/production/packaging/gcp/data_server/bin/BUILD.bazel index 0e4ca491..c8adb46f 100644 --- a/production/packaging/gcp/data_server/bin/BUILD.bazel +++ b/production/packaging/gcp/data_server/bin/BUILD.bazel @@ -11,6 +11,7 @@ cc_binary( "//:local_instance": ["-DINSTANCE_LOCAL=1"], "//conditions:default": [], }), + malloc = "@com_google_tcmalloc//tcmalloc", deps = [ "//components/cloud_config:instance_client", "//components/cloud_config:parameter_client", diff --git a/production/packaging/gcp/data_server/bin/init_server_main.cc b/production/packaging/gcp/data_server/bin/init_server_main.cc index 519a5ac1..34df2cb1 100644 --- a/production/packaging/gcp/data_server/bin/init_server_main.cc +++ b/production/packaging/gcp/data_server/bin/init_server_main.cc @@ -42,7 +42,7 @@ ABSL_FLAG(std::string, environment, "NOT_SPECIFIED", "Environment name."); // The environment flag is defined for local instance (in instance_client_local) // but not for GCP instance, hence the difference here. -void PrepareTlsKeyCertForEnvoy() { +bool PrepareTlsKeyCertForEnvoy() { // Initializes GCP platform and its parameter client. kv_server::PlatformInitializer platform_initializer; std::unique_ptr parameter_client = @@ -58,19 +58,50 @@ void PrepareTlsKeyCertForEnvoy() { "GetEnvironment", kv_server::LogMetricsNoOpCallback()); } + kv_server::ParameterFetcher parameter_fetcher(environment, *parameter_client); + if (bool enable_external_traffic = + parameter_fetcher.GetBoolParameter("enable-external-traffic"); + !enable_external_traffic) { + return false; // Envoy is not needed if we don't serve external traffic. + } + // Prepares TLS cert and key for the connection between XLB and envoy. Per // GCP's protocol, the certificate here is not verified and it's ok to use a // self-signed cert. // (https://cloud.google.com/load-balancing/docs/ssl-certificates/encryption-to-the-backends#secure-protocol-considerations) - kv_server::ParameterFetcher parameter_fetcher(environment, *parameter_client); std::string tls_key_str = parameter_fetcher.GetParameter("tls-key"); std::string tls_cert_str = parameter_fetcher.GetParameter("tls-cert"); + if (tls_key_str == "NOT_PROVIDED" || tls_cert_str == "NOT_PROVIDED") { + LOG(ERROR) << "TLS key/cert are not provided!"; + exit(1); + } std::ofstream key_file("/etc/envoy/key.pem"); std::ofstream cert_file("/etc/envoy/cert.pem"); key_file << tls_key_str; cert_file << tls_cert_str; key_file.close(); cert_file.close(); + return true; +} + +void StartKvServer(int argc, char* argv[]) { + std::vector server_exec_args = {"/server"}; + for (int i = 1; i < argc; ++i) { + server_exec_args.push_back(argv[i]); + } + server_exec_args.push_back(nullptr); + LOG(INFO) << "Starting KV-Server"; + execv(server_exec_args[0], &server_exec_args[0]); + LOG(ERROR) << "Server failure:" << std::strerror(errno); +} + +void StartEnvoy() { + LOG(INFO) << "Starting Envoy"; + std::vector envoy_exec_args = {"/usr/local/bin/envoy", "--config-path", + "/etc/envoy/envoy.yaml", "-l", "warn"}; + envoy_exec_args.push_back(nullptr); + execv(envoy_exec_args[0], &envoy_exec_args[0]); + LOG(ERROR) << "Envoy failure:" << std::strerror(errno); } int main(int argc, char* argv[]) { @@ -80,30 +111,20 @@ int main(int argc, char* argv[]) { // errors. The unrecognized flags will later be used by "/server". absl::ParseAbseilFlagsOnly(argc, argv, positional_args, unrecognized_flags); - PrepareTlsKeyCertForEnvoy(); - - // Starts Envoy and server in separate processes - if (const pid_t pid = fork(); pid == 1) { - LOG(ERROR) << "Fork failure!"; - return errno; - } else if (pid == 0) { - LOG(INFO) << "Starting Envoy"; - std::vector envoy_exec_args = { - "/usr/local/bin/envoy", "--config-path", "/etc/envoy/envoy.yaml", "-l", - "warn"}; - envoy_exec_args.push_back(nullptr); - execv(envoy_exec_args[0], &envoy_exec_args[0]); - LOG(ERROR) << "Envoy failure:" << std::strerror(errno); - } else { - sleep(5); - std::vector server_exec_args = {"/server"}; - for (int i = 1; i < argc; ++i) { - server_exec_args.push_back(argv[i]); + if (PrepareTlsKeyCertForEnvoy()) { + // Starts Envoy and server in separate processes + if (const pid_t pid = fork(); pid == 1) { + LOG(ERROR) << "Fork failure!"; + return errno; + } else if (pid == 0) { + StartEnvoy(); + } else { + sleep(5); + StartKvServer(argc, argv); } - server_exec_args.push_back(nullptr); - LOG(INFO) << "Starting KV-Server"; - execv(server_exec_args[0], &server_exec_args[0]); - LOG(ERROR) << "Server failure:" << std::strerror(errno); + } else { + // Only starts server if envoy is not needed. + StartKvServer(argc, argv); } return errno; } diff --git a/production/packaging/local/data_server/BUILD.bazel b/production/packaging/local/data_server/BUILD.bazel index 79bb3edf..c2d02cf3 100644 --- a/production/packaging/local/data_server/BUILD.bazel +++ b/production/packaging/local/data_server/BUILD.bazel @@ -55,8 +55,8 @@ container_image( "--port=50051", "--delta_directory=/data", "--realtime_directory=/data/realtime", + "--stderrthreshold=0", ], - env = {"GLOG_logtostderr": "1"}, layers = [ ":profiling_tools_layer", "//production/packaging/gcp/data_server:server_binary_layer", @@ -85,8 +85,8 @@ container_image( "--port=50051", "--delta_directory=/data", "--realtime_directory=/data/realtime", + "--stderrthreshold=0", ], - env = {"GLOG_logtostderr": "1"}, layers = [ "//production/packaging/gcp/data_server:server_binary_layer", ], diff --git a/production/packaging/tools/request_simulation/bin/start_request_simulation_system b/production/packaging/tools/request_simulation/bin/start_request_simulation_system index e11dbeb6..ecd35254 100644 --- a/production/packaging/tools/request_simulation/bin/start_request_simulation_system +++ b/production/packaging/tools/request_simulation/bin/start_request_simulation_system @@ -15,4 +15,4 @@ set -o errexit -GLOG_logtostderr=1 /request_simulation/bin/request_simulation_system_main ${EXTRA_FLAGS} +/request_simulation/bin/request_simulation_system_main ${EXTRA_FLAGS} --stderrthreshold=0 diff --git a/production/terraform/aws/environments/demo/us-east-1.tfvars.json b/production/terraform/aws/environments/demo/us-east-1.tfvars.json index 9ba17556..aa0cb895 100644 --- a/production/terraform/aws/environments/demo/us-east-1.tfvars.json +++ b/production/terraform/aws/environments/demo/us-east-1.tfvars.json @@ -1,9 +1,11 @@ { + "add_missing_keys_v1": true, "autoscaling_desired_capacity": 4, "autoscaling_max_size": 6, "autoscaling_min_size": 4, "backup_poll_frequency_secs": 300, "certificate_arn": "cert-arn", + "data_loading_blob_prefix_allowlist": ",", "data_loading_file_format": "riegeli", "data_loading_num_threads": 16, "enclave_cpu_count": 2, @@ -13,6 +15,7 @@ "healthcheck_healthy_threshold": 3, "healthcheck_interval_sec": 30, "healthcheck_unhealthy_threshold": 3, + "http_api_paths": ["/v1/*", "/v2/*", "/healthcheck"], "instance_ami_id": "ami-0000000", "instance_type": "m5.xlarge", "logging_verbosity_level": 0, @@ -21,7 +24,10 @@ "metrics_export_timeout_millis": 500, "num_shards": 1, "primary_coordinator_account_identity": "", + "primary_coordinator_private_key_endpoint": "https://privatekeyservice-a.pa-1.aws.privacysandboxservices.com/v1alpha", + "primary_coordinator_region": "us-east-1", "prometheus_service_region": "us-east-1", + "public_key_endpoint": "https://publickeyservice.staging-pa-1.aws.privacysandboxservices.com/v1alpha/publicKeys", "realtime_updater_num_threads": 4, "region": "us-east-1", "root_domain": "demo-server.com", @@ -32,11 +38,15 @@ "s3client_max_connections": 64, "s3client_max_range_bytes": 8388608, "secondary_coordinator_account_identity": "", + "secondary_coordinator_private_key_endpoint": "https://privatekeyservice-b.pa-2.aws.privacysandboxservices.com/v1alpha", + "secondary_coordinator_region": "us-east-1", "server_port": 51052, "sqs_cleanup_image_uri": "123456789.dkr.ecr.us-east-1.amazonaws.com/sqs_lambda:latest", "sqs_cleanup_schedule": "rate(6 hours)", "sqs_queue_timeout_secs": 86400, "ssh_source_cidr_blocks": ["0.0.0.0/0"], + "telemetry_config": "mode: PROD", + "udf_min_log_level": 0, "udf_num_workers": 2, "use_external_metrics_collector_endpoint": false, "use_real_coordinators": false, diff --git a/production/terraform/aws/environments/demo/us-west-1.tfvars.json b/production/terraform/aws/environments/demo/us-west-1.tfvars.json index 1b5870af..ba64ac5c 100644 --- a/production/terraform/aws/environments/demo/us-west-1.tfvars.json +++ b/production/terraform/aws/environments/demo/us-west-1.tfvars.json @@ -1,9 +1,11 @@ { + "add_missing_keys_v1": true, "autoscaling_desired_capacity": 4, "autoscaling_max_size": 6, "autoscaling_min_size": 4, "backup_poll_frequency_secs": 300, "certificate_arn": "cert-arn", + "data_loading_blob_prefix_allowlist": ",", "data_loading_file_format": "riegeli", "data_loading_num_threads": 16, "enclave_cpu_count": 2, @@ -20,7 +22,10 @@ "metrics_export_timeout_millis": 500, "num_shards": 1, "primary_coordinator_account_identity": "", + "primary_coordinator_private_key_endpoint": "https://privatekeyservice-a.pa-1.aws.privacysandboxservices.com/v1alpha", + "primary_coordinator_region": "us-east-1", "prometheus_service_region": "us-east-1", + "public_key_endpoint": "https://publickeyservice.staging-pa-1.aws.privacysandboxservices.com/v1alpha/publicKeys", "realtime_updater_num_threads": 4, "region": "us-west-1", "root_domain": "demo-server.com", @@ -30,6 +35,8 @@ "s3client_max_connections": 64, "s3client_max_range_bytes": 8388608, "secondary_coordinator_account_identity": "", + "secondary_coordinator_private_key_endpoint": "https://privatekeyservice-b.pa-2.aws.privacysandboxservices.com/v1alpha", + "secondary_coordinator_region": "us-east-1", "server_port": 51052, "sqs_cleanup_image_uri": "123456789.dkr.ecr.us-east-1.amazonaws.com/sqs_lambda:latest", "sqs_cleanup_schedule": "rate(6 hours)", diff --git a/production/terraform/aws/environments/kv_server.tf b/production/terraform/aws/environments/kv_server.tf index 8931fa89..2cc19747 100644 --- a/production/terraform/aws/environments/kv_server.tf +++ b/production/terraform/aws/environments/kv_server.tf @@ -38,6 +38,7 @@ module "kv_server" { enclave_memory_mib = var.enclave_memory_mib enclave_enable_debug_mode = var.enclave_enable_debug_mode run_server_outside_tee = var.run_server_outside_tee + add_missing_keys_v1 = var.add_missing_keys_v1 # Variables related to autoscaling and load balancing. autoscaling_desired_capacity = var.autoscaling_desired_capacity @@ -66,33 +67,42 @@ module "kv_server" { metrics_collector_endpoint = var.metrics_collector_endpoint metrics_export_interval_millis = var.metrics_export_interval_millis metrics_export_timeout_millis = var.metrics_export_timeout_millis + telemetry_config = var.telemetry_config # Variables related to prometheus service prometheus_service_region = var.prometheus_service_region prometheus_workspace_id = var.prometheus_workspace_id # Variables related to data loading. - data_loading_num_threads = var.data_loading_num_threads - s3client_max_connections = var.s3client_max_connections - s3client_max_range_bytes = var.s3client_max_range_bytes - data_loading_file_format = var.data_loading_file_format + data_loading_num_threads = var.data_loading_num_threads + s3client_max_connections = var.s3client_max_connections + s3client_max_range_bytes = var.s3client_max_range_bytes + data_loading_file_format = var.data_loading_file_format + data_loading_blob_prefix_allowlist = var.data_loading_blob_prefix_allowlist # Variables related to sharding. num_shards = var.num_shards use_sharding_key_regex = var.use_sharding_key_regex sharding_key_regex = var.sharding_key_regex - # Variables related to UDF exeuction. + # Variables related to UDF execution. udf_num_workers = var.udf_num_workers udf_timeout_millis = var.udf_timeout_millis + udf_min_log_level = var.udf_min_log_level # Variables related to coordinators - use_real_coordinators = var.use_real_coordinators - primary_coordinator_account_identity = var.primary_coordinator_account_identity - secondary_coordinator_account_identity = var.secondary_coordinator_account_identity + use_real_coordinators = var.use_real_coordinators + primary_coordinator_account_identity = var.primary_coordinator_account_identity + secondary_coordinator_account_identity = var.secondary_coordinator_account_identity + primary_coordinator_private_key_endpoint = var.primary_coordinator_private_key_endpoint + secondary_coordinator_private_key_endpoint = var.secondary_coordinator_private_key_endpoint + primary_coordinator_region = var.primary_coordinator_region + secondary_coordinator_region = var.secondary_coordinator_region + public_key_endpoint = var.public_key_endpoint # Variables related to logging logging_verbosity_level = var.logging_verbosity_level + enable_otel_logger = var.enable_otel_logger } output "kv_server_url" { diff --git a/production/terraform/aws/environments/kv_server_variables.tf b/production/terraform/aws/environments/kv_server_variables.tf index 4e4951cb..48a68fc9 100644 --- a/production/terraform/aws/environments/kv_server_variables.tf +++ b/production/terraform/aws/environments/kv_server_variables.tf @@ -156,6 +156,12 @@ variable "metrics_export_timeout_millis" { type = number } +variable "telemetry_config" { + description = "Telemetry configuration to control whether metrics are raw or noised. Options are: mode: PROD(noised metrics), mode: EXPERIMENT(raw metrics), mode: COMPARE(both raw and noised metrics), mode: OFF(no metrics)" + default = "mode: PROD" + type = string +} + variable "realtime_updater_num_threads" { description = "Number of realtime threads." type = number @@ -202,6 +208,10 @@ variable "route_v1_requests_to_v2" { type = bool } +variable "add_missing_keys_v1" { + description = "Add missing keys v1." + type = bool +} variable "use_real_coordinators" { description = "Use real coordinators." @@ -262,3 +272,46 @@ variable "udf_timeout_millis" { default = 5000 type = number } + +variable "udf_min_log_level" { + description = "Minimum log level for UDFs. Info = 0, Warn = 1, Error = 2. The UDF will only attempt to log for min_log_level and above. Default is 0(info)." + default = 0 + type = number +} + +variable "enable_otel_logger" { + description = "Whether to enable otel logger." + type = bool + default = true +} + +variable "data_loading_blob_prefix_allowlist" { + description = "A comma separated list of prefixes (i.e., directories) where data is loaded from." + default = "," + type = string +} + +variable "primary_coordinator_private_key_endpoint" { + description = "Primary coordinator private key endpoint." + type = string +} + +variable "primary_coordinator_region" { + description = "Primary coordinator region." + type = string +} + +variable "secondary_coordinator_private_key_endpoint" { + description = "Secondary coordinator private key endpoint." + type = string +} + +variable "secondary_coordinator_region" { + description = "Secondary coordinator region." + type = string +} + +variable "public_key_endpoint" { + description = "Public key endpoint. Can only be overriden in non-prod mode." + type = string +} diff --git a/production/terraform/aws/modules/kv_server/main.tf b/production/terraform/aws/modules/kv_server/main.tf index bf09cf69..63043ecb 100644 --- a/production/terraform/aws/modules/kv_server/main.tf +++ b/production/terraform/aws/modules/kv_server/main.tf @@ -141,32 +141,44 @@ module "ssh" { } module "parameter" { - source = "../../services/parameter" - service = local.service - environment = var.environment - s3_bucket_parameter_value = module.data_storage.s3_data_bucket_id - bucket_update_sns_arn_parameter_value = module.data_storage.sns_data_updates_topic_arn - realtime_sns_arn_parameter_value = module.data_storage.sns_realtime_topic_arn - backup_poll_frequency_secs_parameter_value = var.backup_poll_frequency_secs - use_external_metrics_collector_endpoint = var.use_external_metrics_collector_endpoint - metrics_collector_endpoint = var.metrics_collector_endpoint - metrics_export_interval_millis_parameter_value = var.metrics_export_interval_millis - metrics_export_timeout_millis_parameter_value = var.metrics_export_timeout_millis - realtime_updater_num_threads_parameter_value = var.realtime_updater_num_threads - data_loading_num_threads_parameter_value = var.data_loading_num_threads - s3client_max_connections_parameter_value = var.s3client_max_connections - s3client_max_range_bytes_parameter_value = var.s3client_max_range_bytes - num_shards_parameter_value = var.num_shards - udf_num_workers_parameter_value = var.udf_num_workers - udf_timeout_millis_parameter_value = var.udf_timeout_millis - route_v1_requests_to_v2_parameter_value = var.route_v1_requests_to_v2 - use_real_coordinators_parameter_value = var.use_real_coordinators - primary_coordinator_account_identity_parameter_value = var.primary_coordinator_account_identity - secondary_coordinator_account_identity_parameter_value = var.secondary_coordinator_account_identity - data_loading_file_format_parameter_value = var.data_loading_file_format - logging_verbosity_level_parameter_value = var.logging_verbosity_level - use_sharding_key_regex_parameter_value = var.use_sharding_key_regex - sharding_key_regex_parameter_value = var.sharding_key_regex + source = "../../services/parameter" + service = local.service + environment = var.environment + s3_bucket_parameter_value = module.data_storage.s3_data_bucket_id + bucket_update_sns_arn_parameter_value = module.data_storage.sns_data_updates_topic_arn + realtime_sns_arn_parameter_value = module.data_storage.sns_realtime_topic_arn + backup_poll_frequency_secs_parameter_value = var.backup_poll_frequency_secs + use_external_metrics_collector_endpoint = var.use_external_metrics_collector_endpoint + metrics_collector_endpoint = var.metrics_collector_endpoint + metrics_export_interval_millis_parameter_value = var.metrics_export_interval_millis + metrics_export_timeout_millis_parameter_value = var.metrics_export_timeout_millis + telemetry_config = var.telemetry_config + realtime_updater_num_threads_parameter_value = var.realtime_updater_num_threads + data_loading_num_threads_parameter_value = var.data_loading_num_threads + s3client_max_connections_parameter_value = var.s3client_max_connections + s3client_max_range_bytes_parameter_value = var.s3client_max_range_bytes + num_shards_parameter_value = var.num_shards + udf_num_workers_parameter_value = var.udf_num_workers + udf_timeout_millis_parameter_value = var.udf_timeout_millis + udf_min_log_level_parameter_value = var.udf_min_log_level + route_v1_requests_to_v2_parameter_value = var.route_v1_requests_to_v2 + add_missing_keys_v1_parameter_value = var.add_missing_keys_v1 + use_real_coordinators_parameter_value = var.use_real_coordinators + primary_coordinator_account_identity_parameter_value = var.primary_coordinator_account_identity + secondary_coordinator_account_identity_parameter_value = var.secondary_coordinator_account_identity + primary_coordinator_private_key_endpoint_parameter_value = var.primary_coordinator_private_key_endpoint + secondary_coordinator_private_key_endpoint_parameter_value = var.secondary_coordinator_private_key_endpoint + primary_coordinator_region_parameter_value = var.primary_coordinator_region + secondary_coordinator_region_parameter_value = var.secondary_coordinator_region + public_key_endpoint_parameter_value = var.public_key_endpoint + + + data_loading_file_format_parameter_value = var.data_loading_file_format + logging_verbosity_level_parameter_value = var.logging_verbosity_level + use_sharding_key_regex_parameter_value = var.use_sharding_key_regex + sharding_key_regex_parameter_value = var.sharding_key_regex + enable_otel_logger_parameter_value = var.enable_otel_logger + data_loading_blob_prefix_allowlist = var.data_loading_blob_prefix_allowlist } module "security_group_rules" { @@ -203,6 +215,7 @@ module "iam_role_policies" { module.parameter.use_external_metrics_collector_endpoint_arn, module.parameter.metrics_export_interval_millis_parameter_arn, module.parameter.metrics_export_timeout_millis_parameter_arn, + module.parameter.telemetry_config_parameter_arn, module.parameter.realtime_updater_num_threads_parameter_arn, module.parameter.data_loading_num_threads_parameter_arn, module.parameter.s3client_max_connections_parameter_arn, @@ -210,15 +223,24 @@ module "iam_role_policies" { module.parameter.num_shards_parameter_arn, module.parameter.udf_num_workers_parameter_arn, module.parameter.route_v1_requests_to_v2_parameter_arn, + module.parameter.add_missing_keys_v1_parameter_arn, module.parameter.data_loading_file_format_parameter_arn, module.parameter.logging_verbosity_level_parameter_arn, module.parameter.use_real_coordinators_parameter_arn, module.parameter.use_sharding_key_regex_parameter_arn, - module.parameter.udf_timeout_millis_parameter_arn] + module.parameter.udf_timeout_millis_parameter_arn, + module.parameter.udf_min_log_level_parameter_arn, + module.parameter.enable_otel_logger_parameter_arn, + module.parameter.data_loading_blob_prefix_allowlist_parameter_arn] coordinator_parameter_arns = ( var.use_real_coordinators ? [ module.parameter.primary_coordinator_account_identity_parameter_arn, - module.parameter.secondary_coordinator_account_identity_parameter_arn + module.parameter.secondary_coordinator_account_identity_parameter_arn, + module.parameter.primary_coordinator_private_key_endpoint_parameter_arn, + module.parameter.secondary_coordinator_private_key_endpoint_parameter_arn, + module.parameter.primary_coordinator_region_parameter_arn, + module.parameter.secondary_coordinator_region_parameter_arn, + module.parameter.public_key_endpoint_parameter_arn ] : [] ) metrics_collector_endpoint_arns = ( @@ -245,4 +267,5 @@ module "iam_group_policies" { module "dashboards" { source = "../../services/dashboard" environment = var.environment + region = var.region } diff --git a/production/terraform/aws/modules/kv_server/variables.tf b/production/terraform/aws/modules/kv_server/variables.tf index 4ea1b94d..39364af9 100644 --- a/production/terraform/aws/modules/kv_server/variables.tf +++ b/production/terraform/aws/modules/kv_server/variables.tf @@ -78,7 +78,7 @@ variable "autoscaling_min_size" { variable "autoscaling_wait_for_capacity_timeout" { type = string - default = "10m" + default = "20m" } variable "sqs_cleanup_image_uri" { @@ -157,6 +157,12 @@ variable "metrics_export_timeout_millis" { type = number } +variable "telemetry_config" { + description = "Telemetry configuration to control whether metrics are raw or noised. Options are: mode: PROD(noised metrics), mode: EXPERIMENT(raw metrics), mode: COMPARE(both raw and noised metrics), mode: OFF(no metrics)" + default = "mode: PROD" + type = string +} + variable "realtime_updater_num_threads" { description = "The number of threads for the realtime updater." type = number @@ -203,6 +209,11 @@ variable "route_v1_requests_to_v2" { type = bool } +variable "add_missing_keys_v1" { + description = "Add missing keys v1." + type = bool +} + variable "use_real_coordinators" { description = "Will use real coordinators. `enclave_enable_debug_mode` should be set to `false` if the attestation check is enabled for coordinators. Attestation check is enabled on all production instances, and might be disabled for testing purposes only on staging/dev environments." type = bool @@ -258,3 +269,44 @@ variable "udf_timeout_millis" { default = 5000 type = number } + +variable "udf_min_log_level" { + description = "Minimum log level for UDFs. Info = 0, Warn = 1, Error = 2. The UDF will only attempt to log for min_log_level and above. Default is 0(info)." + type = number +} + +variable "enable_otel_logger" { + description = "Whether to enable otel logger." + type = bool + default = true +} + +variable "data_loading_blob_prefix_allowlist" { + description = "A comma separated list of prefixes (i.e., directories) where data is loaded from." + type = string +} + +variable "primary_coordinator_private_key_endpoint" { + description = "Primary coordinator private key endpoint." + type = string +} + +variable "secondary_coordinator_private_key_endpoint" { + description = "Secondary coordinator private key endpoint." + type = string +} + +variable "primary_coordinator_region" { + description = "Primary coordinator region." + type = string +} + +variable "secondary_coordinator_region" { + description = "Secondary coordinator region." + type = string +} + +variable "public_key_endpoint" { + description = "Public key endpoint. Can only be overriden in non-prod mode." + type = string +} diff --git a/production/terraform/aws/services/dashboard/main.tf b/production/terraform/aws/services/dashboard/main.tf index 3858c265..76dcf871 100644 --- a/production/terraform/aws/services/dashboard/main.tf +++ b/production/terraform/aws/services/dashboard/main.tf @@ -23,223 +23,509 @@ resource "aws_cloudwatch_dashboard" "environment_dashboard" { { "widgets": [ { - "height": 9, - "width": 9, + "height": 10, + "width": 12, "y": 0, - "x": 9, + "x": 0, "type": "metric", "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"request.count\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", "view": "timeSeries", "stacked": false, - "region": "us-east-1", + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "request.count [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 0, + "x": 12, + "type": "metric", + "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,OTelLib,deployment.environment,event,host.arch,service.instance.id,service.name,service.version,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} CacheKey', 'Average', 300)", "id": "e1", "period": 300, "label": "$${PROP('Dim.event')} $${PROP('Dim.service.instance.id')}" } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"SecureLookupRequestCount\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], - "title": "Cache hits", - "period": 300, + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, "yAxis": { "left": { - "showUnits": false, - "min": 0 + "min": 0, + "showUnits": false } - } + }, + "title": "Secure lookup request count [MEAN]" } }, { - "height": 9, - "width": 9, - "y": 0, + "height": 10, + "width": 12, + "y": 10, "x": 0, "type": "metric", "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"request.failed_count_by_status\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.error_code')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", "view": "timeSeries", "stacked": false, - "region": "us-east-1", + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "request.failed_count_by_status [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 10, + "x": 12, + "type": "metric", + "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,OTelLib,deployment.environment,event,host.arch,service.instance.id,service.name,service.version,status,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} GetValuesSuccess', 'Average', 300)", "id": "e1", "period": 300, "label": "$${PROP('Dim.service.instance.id')}" } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"request.duration_ms\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], - "title": "GetValuesSuccess", - "period": 300, + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, "yAxis": { "left": { - "showUnits": false, - "min": 0 + "min": 0, + "showUnits": false } - } + }, + "title": "request.duration_ms [MEAN]" } }, { - "height": 9, - "width": 9, - "y": 9, + "height": 10, + "width": 12, + "y": 20, "x": 0, "type": "metric", "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,OTelLib,deployment.environment,event,host.arch,service.instance.id,service.name,service.version,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} GetValuesV1Latency', 'Average', 1)", "id": "e1", "label": "$${PROP('Dim.service.instance.id')}" } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"request.size_bytes\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], + "region": "${var.region}", "view": "timeSeries", "stacked": false, - "region": "us-east-1", - "stat": "Average", - "period": 300, - "title": "GetValuesV1Latency (nanoseconds)", "yAxis": { "left": { - "label": "", - "showUnits": false, - "min": 0 + "min": 0, + "showUnits": false } - } + }, + "title": "request.size_bytes [MEAN]" } }, { + "height": 10, + "width": 12, + "y": 20, + "x": 12, "type": "metric", - "x": 9, - "y": 9, - "width": 9, - "height": 9, "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,OTelLib,deployment.environment,event,host.arch,service.instance.id,service.name,service.version,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} ConcurrentStreamRecordReader', 'Average', 300)", "id": "e1", "period": 300 } ], - [ { "expression": "SEARCH('{KV-Server,OTelLib,deployment.environment,event,host.arch,service.instance.id,service.name,service.version,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} AwsSqsReceiveMessageLatency', 'Average', 300)", "id": "e2", "period": 300 } ], - [ { "expression": "SEARCH('{KV-Server,OTelLib,deployment.environment,event,host.arch,service.instance.id,service.name,service.version,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} SeekingInputStreambuf', 'Average', 300)", "id": "e3", "period": 300 } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"response.size_bytes\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], + "region": "${var.region}", "view": "timeSeries", "stacked": false, - "region": "us-east-1", - "stat": "Average", - "period": 300, - "title": "Read latency", "yAxis": { "left": { - "showUnits": false, - "min": 0 + "min": 0, + "showUnits": false + } + }, + "title": "response.size_bytes [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 30, + "x": 0, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"KVUdfRequestError\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.error_code')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false } - } + }, + "title": "Request Errors [MEAN]" } }, { + "height": 10, + "width": 12, + "y": 30, + "x": 12, "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"InternalLookupRequestError\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.error_code')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Internal Request Errors [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 40, "x": 0, - "y": 18, - "width": 9, - "height": 9, + "type": "metric", "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,Noise,deployment.environment,host.arch,label,service.instance.id,service.name,service.version,shard_number,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} total cores', 'Average', 60)", "id": "e1", "period": 60 } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"KVServerError\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.error_code')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], + "region": "${var.region}", "view": "timeSeries", "stacked": false, - "region": "us-east-1", - "stat": "Average", - "period": 60, - "title": "system.cpu.total_cores[MEAN]", "yAxis": { "left": { - "showUnits": false, - "min": 0 + "min": 0, + "showUnits": false } - } + }, + "title": "Server Non-request Errors [MEAN]" } }, { + "height": 10, + "width": 12, + "y": 40, + "x": 12, "type": "metric", - "x": 9, - "y": 18, - "width": 9, - "height": 9, "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,Noise,deployment.environment,host.arch,label,service.instance.id,service.name,service.version,shard_number,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} main process utilization', 'Average', 60)", "id": "e1", "period": 60 } ], - [ { "expression": "SEARCH('{KV-Server,Noise,deployment.environment,host.arch,label,service.instance.id,service.name,service.version,shard_number,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} total utilization', 'Average', 60)", "id": "e2", "period": 60 } ], - [ { "expression": "SEARCH('{KV-Server,Noise,deployment.environment,host.arch,label,service.instance.id,service.name,service.version,shard_number,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} total load', 'Average', 60)", "id": "e3", "period": 60 } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=(\"ShardedLookupGetKeyValuesLatencyInMicros\" OR \"ShardedLookupGetKeyValueSetLatencyInMicros\" OR \"ShardedLookupRunQueryLatencyInMicros\") Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('MetricName')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], + "region": "${var.region}", "view": "timeSeries", "stacked": false, - "region": "us-east-1", - "stat": "Average", - "period": 60, - "title": "system.cpu.percent[MEAN]", "yAxis": { "left": { - "showUnits": false, - "min": 0 + "min": 0, + "showUnits": false } - } + }, + "title": "Sharded Lookup Latency Microseconds [MEAN]" } }, { + "height": 10, + "width": 12, + "y": 50, + "x": 0, "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"ShardedLookupKeyCountByShard\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.key_shard_num')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Sharded Lookup Key Count By Shard [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 50, + "x": 12, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=(\"InternalGetKeyValuesLatencyInMicros\" OR \"InternalGetKeyValueSetLatencyInMicros\" OR \"InternalRunQueryLatencyInMicros\") Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('MetricName')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Internal Lookup Latency Microseconds [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 60, "x": 0, - "y": 27, - "width": 9, - "height": 9, + "type": "metric", "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,Noise,deployment.environment,host.arch,label,service.instance.id,service.name,service.version,shard_number,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} system.memory.usage main process', 'Average', 60)", "id": "e1", "period": 60, "label": "$${PROP('Dim.service.instance.id')}" } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=(\"GetValuePairsLatencyInMicros\" OR \"GetKeyValueSetLatencyInMicros\") Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('MetricName')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], + "region": "${var.region}", "view": "timeSeries", "stacked": false, - "region": "us-east-1", - "stat": "Average", - "period": 60, - "title": "system.memory.usage_kb for main process[MEAN]", + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Cache Query Latency Microseconds [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 60, + "x": 12, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"CacheAccessEventCount\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.cache_access')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Cache Access Event Count [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 70, + "x": 0, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} status Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.status')} $${PROP('MetricName')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Server Retryable Operation Status Count [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 70, + "x": 12, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} data_source data_source=(NOT \"realtime\") Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.data_source')} $${PROP('MetricName')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "File Update Stats [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 80, + "x": 0, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=(\"ConcurrentStreamRecordReaderReadShardRecordsLatency\" OR \"ConcurrentStreamRecordReaderReadStreamRecordsLatency\" OR \"ConcurrentStreamRecordReaderReadByteRangeLatency\") Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('MetricName')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Data Reader Latency Microseconds[MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 80, + "x": 12, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=(\"UpdateKeyValueLatency\" OR \"UpdateKeyValueSetLatency\" OR \"DeleteKeyLatency\" OR \"DeleteValuesInSetLatency\" OR \"RemoveDeletedKeyLatency\") Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('MetricName')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Cache Update Latency Microseconds[MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 90, + "x": 0, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} data_source data_source=\"realtime\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Realtime Update Stats [MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 90, + "x": 12, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=${var.environment} MetricName=\"ReceivedLowLatencyNotifications\" Noise=(\"Raw\" OR \"Noised\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "region": "${var.region}", + "view": "timeSeries", + "stacked": false, + "yAxis": { + "left": { + "min": 0, + "showUnits": false + } + }, + "title": "Realtime Message Processing Latency Microseconds[MEAN]" + } + }, + { + "height": 10, + "width": 12, + "y": 100, + "x": 0, + "type": "metric", + "properties": { + "metrics": [ + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=\"${var.environment}\" Noise=(\"Raw\" OR \"Noised\") MetricName=\"system.cpu.percent\" label=(\"total utilization\" OR \"main process utilization\" OR \"total load\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.label')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] + ], + "view": "timeSeries", + "stacked": false, + "region": "${var.region}", "yAxis": { "left": { "showUnits": false, "min": 0 } - } + }, + "title": "system.cpu.percent [MEAN]" } }, { + "height": 10, + "width": 12, + "y": 100, + "x": 12, "type": "metric", - "x": 9, - "y": 27, - "width": 9, - "height": 9, "properties": { "metrics": [ - [ { "expression": "SEARCH('{KV-Server,Noise,deployment.environment,host.arch,label,service.instance.id,service.name,service.version,shard_number,telemetry.sdk.language,telemetry.sdk.name,telemetry.sdk.version} ${var.environment} system.memory.usage MemAvailable', 'Average', 60)", "id": "e1", "period": 60, "label": "$${PROP('Dim.service.instance.id')}" } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=\"${var.environment}\" Noise=(\"Raw\" OR \"Noised\") MetricName=\"system.memory.usage_kb\" label=(\"MemTotal:\" OR \"main process\")', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.label')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], "view": "timeSeries", "stacked": false, - "region": "us-east-1", - "stat": "Average", - "period": 60, - "title": "system.memory.usage_kb for MemAvailable[MEAN]", + "region": "${var.region}", "yAxis": { "left": { "showUnits": false, "min": 0 } - } + }, + "title": "system.memory.usage_kb [MEAN]" } }, { - "type": "metric", + "height": 10, + "width": 12, + "y": 110, "x": 0, - "y": 36, - "width": 9, - "height": 9, + "type": "metric", "properties": { "metrics": [ - [ { "expression": "SELECT COUNT(ChangeNotifierErrors) FROM SCHEMA(\"KV-Server\", Noise,\"deployment.environment\",error_code,\"host.arch\",\"service.instance.id\",\"service.name\",\"service.version\",shard_number,\"telemetry.sdk.language\",\"telemetry.sdk.name\",\"telemetry.sdk.version\") WHERE \"deployment.environment\" = '${var.environment}' GROUP BY error_code, \"service.instance.id\"", "label": "error rate", "id": "m1", "stat": "Average", "visible": false } ], - [ { "expression": "DIFF(m1)", "label": "error count", "id": "e1" } ] + [ { "expression": "REMOVE_EMPTY(SEARCH('service.name=\"kv-server\" deployment.environment=\"${var.environment}\" Noise=(\"Raw\" OR \"Noised\") MetricName=\"system.cpu.percent\" label=\"total cpu cores\"', 'Average', 60))", "id": "e1", "label": "$${PROP('Dim.Noise')} $${PROP('Dim.label')} $${PROP('Dim.service.instance.id')} $${PROP('Dim.shard_number')}" } ] ], "view": "timeSeries", "stacked": false, - "region": "us-east-1", - "title": "Change notifier error count", - "period": 60, - "stat": "Average", + "region": "${var.region}", "yAxis": { "left": { - "min": 0, - "showUnits": false + "showUnits": false, + "min": 0 } - } + }, + "title": "system.cpu.total_cores [MEAN]" } } ] diff --git a/production/terraform/aws/services/dashboard/variables.tf b/production/terraform/aws/services/dashboard/variables.tf index ec4e6179..7c46199b 100644 --- a/production/terraform/aws/services/dashboard/variables.tf +++ b/production/terraform/aws/services/dashboard/variables.tf @@ -19,3 +19,8 @@ variable "environment" { description = "Assigned environment name to group related resources." type = string } + +variable "region" { + description = "AWS region to deploy to." + type = string +} diff --git a/production/terraform/aws/services/iam_role_policies/main.tf b/production/terraform/aws/services/iam_role_policies/main.tf index cf1c5dd2..3d1d95e1 100644 --- a/production/terraform/aws/services/iam_role_policies/main.tf +++ b/production/terraform/aws/services/iam_role_policies/main.tf @@ -147,7 +147,7 @@ data "aws_iam_policy_document" "sqs_cleanup_lambda_policy_doc" { "sns:List*", "sns:Unsubscribe" ] - resources = [var.sns_data_updates_topic_arn, var.sns_realtime_topic_arn] + resources = ["*"] } statement { sid = "AllowLambdaToManageSQSQueues" diff --git a/production/terraform/aws/services/parameter/main.tf b/production/terraform/aws/services/parameter/main.tf index 58348672..7d6d6450 100644 --- a/production/terraform/aws/services/parameter/main.tf +++ b/production/terraform/aws/services/parameter/main.tf @@ -79,6 +79,13 @@ resource "aws_ssm_parameter" "metrics_export_timeout_millis_parameter" { overwrite = true } +resource "aws_ssm_parameter" "telemetry_config_parameter" { + name = "${var.service}-${var.environment}-telemetry-config" + type = "String" + value = var.telemetry_config + overwrite = true +} + resource "aws_ssm_parameter" "realtime_updater_num_threads_parameter" { name = "${var.service}-${var.environment}-realtime-updater-num-threads" type = "String" @@ -128,6 +135,13 @@ resource "aws_ssm_parameter" "route_v1_requests_to_v2_parameter" { overwrite = true } +resource "aws_ssm_parameter" "add_missing_keys_v1_parameter" { + name = "${var.service}-${var.environment}-add-missing-keys-v1" + type = "String" + value = var.add_missing_keys_v1_parameter_value + overwrite = true +} + resource "aws_ssm_parameter" "use_real_coordinators_parameter" { name = "${var.service}-${var.environment}-use-real-coordinators" type = "String" @@ -151,6 +165,46 @@ resource "aws_ssm_parameter" "secondary_coordinator_account_identity_parameter" overwrite = true } +resource "aws_ssm_parameter" "primary_coordinator_private_key_endpoint_parameter" { + count = (var.use_real_coordinators_parameter_value) ? 1 : 0 + name = "${var.service}-${var.environment}-primary-coordinator-private-key-endpoint" + type = "String" + value = var.primary_coordinator_private_key_endpoint_parameter_value + overwrite = true +} + +resource "aws_ssm_parameter" "secondary_coordinator_private_key_endpoint_parameter" { + count = (var.use_real_coordinators_parameter_value) ? 1 : 0 + name = "${var.service}-${var.environment}-primary-coordinator-region" + type = "String" + value = var.secondary_coordinator_private_key_endpoint_parameter_value + overwrite = true +} + +resource "aws_ssm_parameter" "primary_coordinator_region_parameter" { + count = (var.use_real_coordinators_parameter_value) ? 1 : 0 + name = "${var.service}-${var.environment}-secondary-coordinator-private-key-endpoint" + type = "String" + value = var.primary_coordinator_region_parameter_value + overwrite = true +} + +resource "aws_ssm_parameter" "secondary_coordinator_region_parameter" { + count = (var.use_real_coordinators_parameter_value) ? 1 : 0 + name = "${var.service}-${var.environment}-secondary-coordinator-region" + type = "String" + value = var.secondary_coordinator_region_parameter_value + overwrite = true +} + +resource "aws_ssm_parameter" "public_key_endpoint_parameter" { + count = (var.use_real_coordinators_parameter_value) ? 1 : 0 + name = "${var.service}-${var.environment}-public-key-endpoint" + type = "String" + value = var.public_key_endpoint_parameter_value + overwrite = true +} + resource "aws_ssm_parameter" "data_loading_file_format_parameter" { name = "${var.service}-${var.environment}-data-loading-file-format" type = "String" @@ -186,3 +240,24 @@ resource "aws_ssm_parameter" "udf_timeout_millis_parameter" { value = var.udf_timeout_millis_parameter_value overwrite = true } + +resource "aws_ssm_parameter" "udf_min_log_level_parameter" { + name = "${var.service}-${var.environment}-udf-min-log-level" + type = "String" + value = var.udf_min_log_level_parameter_value + overwrite = true +} + +resource "aws_ssm_parameter" "enable_otel_logger_parameter" { + name = "${var.service}-${var.environment}-enable-otel-logger" + type = "String" + value = var.enable_otel_logger_parameter_value + overwrite = true +} + +resource "aws_ssm_parameter" "data_loading_blob_prefix_allowlist" { + name = "${var.service}-${var.environment}-data-loading-blob-prefix-allowlist" + type = "String" + value = var.data_loading_blob_prefix_allowlist + overwrite = true +} diff --git a/production/terraform/aws/services/parameter/outputs.tf b/production/terraform/aws/services/parameter/outputs.tf index 970408a4..32e9aacf 100644 --- a/production/terraform/aws/services/parameter/outputs.tf +++ b/production/terraform/aws/services/parameter/outputs.tf @@ -54,6 +54,10 @@ output "metrics_export_timeout_millis_parameter_arn" { value = aws_ssm_parameter.metrics_export_timeout_millis_parameter.arn } +output "telemetry_config_parameter_arn" { + value = aws_ssm_parameter.telemetry_config_parameter.arn +} + output "realtime_updater_num_threads_parameter_arn" { value = aws_ssm_parameter.realtime_updater_num_threads_parameter.arn } @@ -82,6 +86,10 @@ output "route_v1_requests_to_v2_parameter_arn" { value = aws_ssm_parameter.route_v1_requests_to_v2_parameter.arn } +output "add_missing_keys_v1_parameter_arn" { + value = aws_ssm_parameter.add_missing_keys_v1_parameter.arn +} + output "use_real_coordinators_parameter_arn" { value = aws_ssm_parameter.use_real_coordinators_parameter.arn } @@ -94,6 +102,26 @@ output "secondary_coordinator_account_identity_parameter_arn" { value = (var.use_real_coordinators_parameter_value) ? aws_ssm_parameter.secondary_coordinator_account_identity_parameter[0].arn : "" } +output "primary_coordinator_private_key_endpoint_parameter_arn" { + value = (var.use_real_coordinators_parameter_value) ? aws_ssm_parameter.primary_coordinator_private_key_endpoint_parameter[0].arn : "" +} + +output "secondary_coordinator_private_key_endpoint_parameter_arn" { + value = (var.use_real_coordinators_parameter_value) ? aws_ssm_parameter.secondary_coordinator_private_key_endpoint_parameter[0].arn : "" +} + +output "primary_coordinator_region_parameter_arn" { + value = (var.use_real_coordinators_parameter_value) ? aws_ssm_parameter.primary_coordinator_region_parameter[0].arn : "" +} + +output "secondary_coordinator_region_parameter_arn" { + value = (var.use_real_coordinators_parameter_value) ? aws_ssm_parameter.secondary_coordinator_region_parameter[0].arn : "" +} + +output "public_key_endpoint_parameter_arn" { + value = (var.use_real_coordinators_parameter_value) ? aws_ssm_parameter.public_key_endpoint_parameter[0].arn : "" +} + output "data_loading_file_format_parameter_arn" { value = aws_ssm_parameter.data_loading_file_format_parameter.arn } @@ -113,3 +141,15 @@ output "sharding_key_regex_parameter_arn" { output "udf_timeout_millis_parameter_arn" { value = aws_ssm_parameter.udf_timeout_millis_parameter.arn } + +output "udf_min_log_level_parameter_arn" { + value = aws_ssm_parameter.udf_min_log_level_parameter.arn +} + +output "enable_otel_logger_parameter_arn" { + value = aws_ssm_parameter.enable_otel_logger_parameter.arn +} + +output "data_loading_blob_prefix_allowlist_parameter_arn" { + value = aws_ssm_parameter.data_loading_blob_prefix_allowlist.arn +} diff --git a/production/terraform/aws/services/parameter/variables.tf b/production/terraform/aws/services/parameter/variables.tf index 91d36660..32228f3b 100644 --- a/production/terraform/aws/services/parameter/variables.tf +++ b/production/terraform/aws/services/parameter/variables.tf @@ -59,6 +59,11 @@ variable "metrics_export_timeout_millis_parameter_value" { type = number } +variable "telemetry_config" { + description = "Telemetry config for exporting raw or noised metrics" + type = string +} + variable "realtime_updater_num_threads_parameter_value" { description = "Amount of realtime notifier threads." type = number @@ -94,6 +99,11 @@ variable "route_v1_requests_to_v2_parameter_value" { type = bool } +variable "add_missing_keys_v1_parameter_value" { + description = "Add missing keys v1." + type = bool +} + variable "use_real_coordinators_parameter_value" { description = "Number of parallel threads for reading and loading data files." type = bool @@ -138,3 +148,43 @@ variable "udf_timeout_millis_parameter_value" { description = "UDF execution timeout in milliseconds." type = number } + +variable "udf_min_log_level_parameter_value" { + description = "Minimum log level for UDFs. Info = 0, Warn = 1, Error = 2. The UDF will only attempt to log for min_log_level and above. Default is 0(info)." + type = number +} + +variable "enable_otel_logger_parameter_value" { + description = "Whether to enable otel logger." + type = bool +} + +variable "data_loading_blob_prefix_allowlist" { + description = "A comma separated list of prefixes (i.e., directories) where data is loaded from." + type = string +} + +variable "primary_coordinator_private_key_endpoint_parameter_value" { + description = "Primary coordinator private key endpoint." + type = string +} + +variable "secondary_coordinator_private_key_endpoint_parameter_value" { + description = "Secondary coordinator private key endpoint." + type = string +} + +variable "primary_coordinator_region_parameter_value" { + description = "Primary coordinator region." + type = string +} + +variable "secondary_coordinator_region_parameter_value" { + description = "Secondary coordinator region." + type = string +} + +variable "public_key_endpoint_parameter_value" { + description = "Public key endpoint. Can only be overriden in non-prod mode." + type = string +} diff --git a/production/terraform/aws/services/security_group_rules/main.tf b/production/terraform/aws/services/security_group_rules/main.tf index a528b9d2..efe01002 100644 --- a/production/terraform/aws/services/security_group_rules/main.tf +++ b/production/terraform/aws/services/security_group_rules/main.tf @@ -157,3 +157,12 @@ resource "aws_security_group_rule" "allow_ec2_to_ec2_endpoint_ingress" { type = "ingress" source_security_group_id = var.instances_security_group_id } + +resource "aws_security_group_rule" "allow_ec2_secure_tcp_egress" { + from_port = 443 + protocol = "TCP" + security_group_id = var.instances_security_group_id + to_port = 443 + type = "egress" + cidr_blocks = ["0.0.0.0/0"] +} diff --git a/production/terraform/gcp/environments/README.md b/production/terraform/gcp/environments/README.md new file mode 100644 index 00000000..e2f209bc --- /dev/null +++ b/production/terraform/gcp/environments/README.md @@ -0,0 +1,23 @@ +# Demo Environment + +An existing demo environment exists for the KV server and can be used as a reference for creating +your own environment (e.g. "dev" and "prod"). + +## Synopsis + +```bash +nano demo/us-east1.backend.conf # Customize backend parameters +nano demo/us-east1.tfvars.json # Customize input variables to suit your demo environment. +terraform init --backend-config=demo/us-east1.backend.conf --var-file=demo/us-east1.tfvars.json --reconfigure +terraform plan --var-file=demo/us-east1.tfvars.json +terraform apply --var-file=demo/us-east1.tfvars.json +``` + +## Configuration Files + +The files which should be modified for your purposes for each environment that you create are: + +- [us-east1.tfvars.json](demo/us-east1.tfvars.json) - an example configuration file for the KV + server. +- [us-east1.backend.conf](demo/us-east1.backend.conf) - contains terraform state bucket location - + should be edited to point to a state bucket you control. diff --git a/production/terraform/gcp/environments/demo/us-east1.tfvars.json b/production/terraform/gcp/environments/demo/us-east1.tfvars.json index b0ec981a..b0cf074d 100644 --- a/production/terraform/gcp/environments/demo/us-east1.tfvars.json +++ b/production/terraform/gcp/environments/demo/us-east1.tfvars.json @@ -1,4 +1,5 @@ { + "add_missing_keys_v1": true, "backup_poll_frequency_secs": 5, "collector_dns_zone": "your-dns-zone-name", "collector_domain_name": "your-domain-name", @@ -7,7 +8,9 @@ "collector_service_port": 4317, "cpu_utilization_percent": 0.9, "data_bucket_id": "your-delta-file-bucket", + "data_loading_blob_prefix_allowlist": ",", "data_loading_num_threads": 16, + "enable_external_traffic": true, "environment": "demo", "envoy_port": 51052, "existing_service_mesh": "", @@ -24,20 +27,29 @@ "min_replicas_per_service_region": 1, "num_shards": 1, "primary_coordinator_account_identity": "EMPTY_STRING", + "primary_coordinator_private_key_endpoint": "https://privatekeyservice-a.pa-3.gcp.privacysandboxservices.com/v1alpha/encryptionKeys", + "primary_coordinator_region": "us-central1", "primary_key_service_cloud_function_url": "EMPTY_STRING", "primary_workload_identity_pool_provider": "EMPTY_STRING", "project_id": "your-project-id", + "public_key_endpoint": "https://publickeyservice.stg-pa.gcp.pstest.dev/.well-known/protected-auction/v1/public-keys", "realtime_updater_num_threads": 1, "regions": ["us-east1"], + "regions_cidr_blocks": ["10.0.3.0/24"], + "regions_use_existing_nat": [], "route_v1_to_v2": false, "secondary_coordinator_account_identity": "EMPTY_STRING", + "secondary_coordinator_private_key_endpoint": "https://privatekeyservice.pa-4.gcp.privacysandboxservices.com/v1alpha/encryptionKeys", + "secondary_coordinator_region": "us-central1", "secondary_key_service_cloud_function_url": "EMPTY_STRING", "secondary_workload_identity_pool_provider": "EMPTY_STRING", "server_dns_zone": "your-server-dns-zone-name", "server_domain_ssl_certificate_id": "your-server-ssl-certificate-id", "server_url": "your-kv-server-url", "service_account_email": "your-service-account-email", + "service_mesh_address": "xds:///kv-service-host", "tee_impersonate_service_accounts": "", + "telemetry_config": "mode: EXPERIMENT", "udf_num_workers": 2, "use_confidential_space_debug_image": false, "use_existing_service_mesh": false, diff --git a/production/terraform/gcp/environments/kv_server.tf b/production/terraform/gcp/environments/kv_server.tf index 6c337b4e..332507de 100644 --- a/production/terraform/gcp/environments/kv_server.tf +++ b/production/terraform/gcp/environments/kv_server.tf @@ -37,6 +37,8 @@ module "kv_server" { service = local.kv_service service_account_email = var.service_account_email regions = var.regions + regions_cidr_blocks = var.regions_cidr_blocks + regions_use_existing_nat = var.regions_use_existing_nat gcp_image_tag = var.gcp_image_tag gcp_image_repo = var.gcp_image_repo kv_service_port = var.kv_service_port @@ -63,34 +65,47 @@ module "kv_server" { num_shards = var.num_shards use_existing_service_mesh = var.use_existing_service_mesh existing_service_mesh = var.existing_service_mesh + service_mesh_address = var.service_mesh_address + enable_external_traffic = var.enable_external_traffic parameters = { - data-bucket-id = var.data_bucket_id - launch-hook = "${local.kv_service}-${var.environment}-launch-hook" - use-external-metrics-collector-endpoint = var.use_external_metrics_collector_endpoint - metrics-collector-endpoint = "${var.environment}-${var.collector_service_name}.${var.collector_domain_name}:${var.collector_service_port}" - metrics-export-interval-millis = var.metrics_export_interval_millis - metrics-export-timeout-millis = var.metrics_export_timeout_millis - backup-poll-frequency-secs = var.backup_poll_frequency_secs - realtime-updater-num-threads = var.realtime_updater_num_threads - data-loading-num-threads = var.data_loading_num_threads - num-shards = var.num_shards - udf-num-workers = var.udf_num_workers - udf-timeout-millis = var.udf_timeout_millis - route-v1-to-v2 = var.route_v1_to_v2 - use-real-coordinators = var.use_real_coordinators - environment = var.environment - project-id = var.project_id - primary-key-service-cloud-function-url = var.primary_key_service_cloud_function_url - primary-workload-identity-pool-provider = var.primary_workload_identity_pool_provider - secondary-key-service-cloud-function-url = var.secondary_key_service_cloud_function_url - secondary-workload-identity-pool-provider = var.secondary_workload_identity_pool_provider - primary-coordinator-account-identity = var.primary_coordinator_account_identity - secondary-coordinator-account-identity = var.secondary_coordinator_account_identity - logging-verbosity-level = var.logging_verbosity_level - use-sharding-key-regex = var.use_sharding_key_regex - sharding-key-regex = var.sharding_key_regex - tls-key = var.tls_key - tls-cert = var.tls_cert + data-bucket-id = var.data_bucket_id + launch-hook = "${local.kv_service}-${var.environment}-launch-hook" + use-external-metrics-collector-endpoint = var.use_external_metrics_collector_endpoint + metrics-collector-endpoint = "${var.environment}-${var.collector_service_name}.${var.collector_domain_name}:${var.collector_service_port}" + metrics-export-interval-millis = var.metrics_export_interval_millis + metrics-export-timeout-millis = var.metrics_export_timeout_millis + backup-poll-frequency-secs = var.backup_poll_frequency_secs + realtime-updater-num-threads = var.realtime_updater_num_threads + data-loading-num-threads = var.data_loading_num_threads + num-shards = var.num_shards + udf-num-workers = var.udf_num_workers + udf-timeout-millis = var.udf_timeout_millis + udf-min-log-level = var.udf_min_log_level + route-v1-to-v2 = var.route_v1_to_v2 + add-missing-keys-v1 = var.add_missing_keys_v1 + use-real-coordinators = var.use_real_coordinators + environment = var.environment + project-id = var.project_id + primary-key-service-cloud-function-url = var.primary_key_service_cloud_function_url + primary-workload-identity-pool-provider = var.primary_workload_identity_pool_provider + secondary-key-service-cloud-function-url = var.secondary_key_service_cloud_function_url + secondary-workload-identity-pool-provider = var.secondary_workload_identity_pool_provider + primary-coordinator-account-identity = var.primary_coordinator_account_identity + secondary-coordinator-account-identity = var.secondary_coordinator_account_identity + primary-coordinator-private-key-endpoint = var.primary_coordinator_private_key_endpoint + primary-coordinator-region = var.primary_coordinator_region + secondary-coordinator-private-key-endpoint = var.secondary_coordinator_private_key_endpoint + secondary-coordinator-region = var.secondary_coordinator_region + public-key-endpoint = var.public_key_endpoint + logging-verbosity-level = var.logging_verbosity_level + use-sharding-key-regex = var.use_sharding_key_regex + sharding-key-regex = var.sharding_key_regex + tls-key = var.tls_key + tls-cert = var.tls_cert + enable-otel-logger = var.enable_otel_logger + enable-external-traffic = var.enable_external_traffic + telemetry-config = var.telemetry_config + data-loading-blob-prefix-allowlist = var.data_loading_blob_prefix_allowlist } } diff --git a/production/terraform/gcp/environments/kv_server_variables.tf b/production/terraform/gcp/environments/kv_server_variables.tf index 0cc6f400..8c0a8ba0 100644 --- a/production/terraform/gcp/environments/kv_server_variables.tf +++ b/production/terraform/gcp/environments/kv_server_variables.tf @@ -34,6 +34,16 @@ variable "regions" { type = set(string) } +variable "regions_cidr_blocks" { + description = "A set of CIDR ranges for all specified regions. The number of blocks here should correspond to the number of regions." + type = set(string) +} + +variable "regions_use_existing_nat" { + description = "Regions that use existing nat. No new nats will be created for regions specified here." + type = set(string) +} + variable "gcp_image_tag" { description = "Tag of the gcp docker image uploaded to the artifact registry." type = string @@ -50,32 +60,34 @@ variable "kv_service_port" { } variable "envoy_port" { - description = "External load balancer will send traffic to this port. Envoy will forward traffic to kv_service_port. Must match envoy.yaml." + description = "External load balancer will send traffic to this port. Envoy will forward traffic to kv_service_port. Must match envoy.yaml. Ignored if `enable_external_traffic` is false." type = number } variable "server_url" { - description = "Kv-serer URL. Example: kv-server-environment.example.com" + description = "Kv-serer URL. Example: kv-server-environment.example.com. Ignored if `enable_external_traffic` is false." type = string } variable "server_dns_zone" { - description = "Google Cloud Dns zone for Kv-serer." + description = "Google Cloud Dns zone for Kv-serer. Ignored if `enable_external_traffic` is false." type = string } variable "server_domain_ssl_certificate_id" { - description = "Ssl certificate id of the Kv-server URL." + description = "Ssl certificate id of the Kv-server URL. Ignored if `enable_external_traffic` is false." type = string } variable "tls_key" { - description = "TLS key. Please specify this variable in a tfvars file (e.g., secrets.auto.tfvars) under the `environments` directory." + description = "TLS key. Please specify this variable in a tfvars file (e.g., secrets.auto.tfvars) under the `environments` directory. Ignored if `enable_external_traffic` is false." + default = "NOT_PROVIDED" type = string } variable "tls_cert" { - description = "TLS cert. Please specify this variable in a tfvars file (e.g., secrets.auto.tfvars) under the `environments` directory." + description = "TLS cert. Please specify this variable in a tfvars file (e.g., secrets.auto.tfvars) under the `environments` directory. Ignored if `enable_external_traffic` is false." + default = "NOT_PROVIDED" type = string } @@ -162,11 +174,22 @@ variable "udf_timeout_millis" { description = "UDF execution timeout in milliseconds." } +variable "udf_min_log_level" { + type = number + default = 0 + description = "Minimum log level for UDFs. Info = 0, Warn = 1, Error = 2. The UDF will only attempt to log for min_log_level and above. Default is 0(info)." +} + variable "route_v1_to_v2" { type = bool description = "Whether to route V1 requests through V2." } +variable "add_missing_keys_v1" { + type = bool + description = "Add missing keys for v1." +} + variable "use_real_coordinators" { type = bool description = "Use real coordinators." @@ -274,3 +297,58 @@ variable "sharding_key_regex" { default = "EMPTY_STRING" type = string } + +variable "service_mesh_address" { + description = "Service mesh address of the KV server." + default = "xds:///kv-service-host" + type = string +} + +variable "enable_otel_logger" { + description = "Whether to use otel logger." + default = true + type = bool +} + +variable "enable_external_traffic" { + description = "Whether to serve external traffic. If disabled, only internal traffic via service mesh will be served." + default = true + type = bool +} + +variable "telemetry_config" { + description = "Telemetry configuration to control whether metrics are raw or noised. Options are: mode: PROD(noised metrics), mode: EXPERIMENT(raw metrics), mode: COMPARE(both raw and noised metrics), mode: OFF(no metrics)" + default = "mode: PROD" + type = string +} + +variable "data_loading_blob_prefix_allowlist" { + description = "A comma separated list of prefixes (i.e., directories) where data is loaded from." + default = "," + type = string +} + +variable "primary_coordinator_private_key_endpoint" { + description = "Primary coordinator private key endpoint." + type = string +} + +variable "secondary_coordinator_private_key_endpoint" { + description = "Secondary coordinator private key endpoint." + type = string +} + +variable "primary_coordinator_region" { + description = "Primary coordinator region." + type = string +} + +variable "secondary_coordinator_region" { + description = "Secondary coordinator region." + type = string +} + +variable "public_key_endpoint" { + description = "Public key endpoint. Can only be overriden in non-prod mode." + type = string +} diff --git a/production/terraform/gcp/environments/terraform.tf b/production/terraform/gcp/environments/terraform.tf index d36e1551..9d99b89e 100644 --- a/production/terraform/gcp/environments/terraform.tf +++ b/production/terraform/gcp/environments/terraform.tf @@ -25,7 +25,7 @@ terraform { google-beta = { source = "hashicorp/google-beta" - version = "5.0.0" + version = "5.12.0" } } } diff --git a/production/terraform/gcp/modules/kv_server/main.tf b/production/terraform/gcp/modules/kv_server/main.tf index b79130ed..a43531d9 100644 --- a/production/terraform/gcp/modules/kv_server/main.tf +++ b/production/terraform/gcp/modules/kv_server/main.tf @@ -14,18 +14,17 @@ * limitations under the License. */ -locals { - kv_server_address = "xds:///kv-service-host" -} - module "networking" { - source = "../../services/networking" - service = var.service - environment = var.environment - regions = var.regions - collector_service_name = var.collector_service_name - use_existing_vpc = var.use_existing_vpc - existing_vpc_id = var.existing_vpc_id + source = "../../services/networking" + service = var.service + environment = var.environment + regions = var.regions + regions_cidr_blocks = var.regions_cidr_blocks + regions_use_existing_nat = var.regions_use_existing_nat + collector_service_name = var.collector_service_name + use_existing_vpc = var.use_existing_vpc + existing_vpc_id = var.existing_vpc_id + enable_external_traffic = var.enable_external_traffic } module "security" { @@ -59,6 +58,7 @@ module "autoscaling" { parameters = var.parameters tee_impersonate_service_accounts = var.tee_impersonate_service_accounts shard_num = count.index + enable_external_traffic = var.enable_external_traffic } module "metrics_collector_autoscaling" { @@ -91,16 +91,19 @@ module "service_mesh" { service = var.service environment = var.environment service_port = var.kv_service_port - kv_server_address = local.kv_server_address + kv_server_address = var.service_mesh_address project_id = var.project_id instance_groups = flatten(module.autoscaling[*].kv_server_instance_groups) collector_forwarding_rule = module.metrics_collector.collector_forwarding_rule collector_tcp_proxy = module.metrics_collector.collector_tcp_proxy use_existing_service_mesh = var.use_existing_service_mesh existing_service_mesh = var.existing_service_mesh + enable_external_traffic = var.enable_external_traffic } module "external_load_balancing" { + count = var.enable_external_traffic ? 1 : 0 + source = "../../services/external_load_balancing" service = var.service environment = var.environment diff --git a/production/terraform/gcp/modules/kv_server/variables.tf b/production/terraform/gcp/modules/kv_server/variables.tf index 2746b1d5..bdcd0fc7 100644 --- a/production/terraform/gcp/modules/kv_server/variables.tf +++ b/production/terraform/gcp/modules/kv_server/variables.tf @@ -44,6 +44,16 @@ variable "regions" { type = set(string) } +variable "regions_cidr_blocks" { + description = "A set of CIDR ranges for all specified regions. The number of blocks here should correspond to the number of regions." + type = set(string) +} + +variable "regions_use_existing_nat" { + description = "Regions that use existing nat. No new nats will be created for regions specified here." + type = set(string) +} + variable "gcp_image_repo" { description = "A URL to a docker image repo containing the key-value service" type = string @@ -175,3 +185,15 @@ variable "server_domain_ssl_certificate_id" { description = "Ssl certificate id of the Kv-server domain." type = string } + +variable "service_mesh_address" { + description = "Service mesh address of the KV server." + default = "xds:///kv-service-host" + type = string +} + +variable "enable_external_traffic" { + description = "Whether to serve external traffic. If disabled, only internal traffic via service mesh will be served." + default = true + type = bool +} diff --git a/production/terraform/gcp/services/autoscaling/main.tf b/production/terraform/gcp/services/autoscaling/main.tf index 4d039298..3b10e1bc 100644 --- a/production/terraform/gcp/services/autoscaling/main.tf +++ b/production/terraform/gcp/services/autoscaling/main.tf @@ -24,7 +24,7 @@ resource "google_compute_instance_template" "kv_server" { for_each = var.subnets region = each.value.region - name = "${var.service}-${var.environment}-${var.shard_num}-instance-lt" + name_prefix = "${var.service}-${var.environment}-${var.shard_num}-instance-lt-" machine_type = var.machine_type tags = ["allow-hc", "allow-ssh", "allow-backend-ingress", "allow-all-egress"] @@ -81,7 +81,6 @@ resource "google_compute_instance_template" "kv_server" { lifecycle { create_before_destroy = true - ignore_changes = [name] replace_triggered_by = [null_resource.kv_parameters] } @@ -90,6 +89,8 @@ resource "google_compute_instance_template" "kv_server" { } resource "google_compute_region_instance_group_manager" "kv_server" { + provider = google-beta + for_each = google_compute_instance_template.kv_server name = "${var.service}-${var.environment}-${each.value.region}-${var.shard_num}-mig" region = each.value.region @@ -103,9 +104,12 @@ resource "google_compute_region_instance_group_manager" "kv_server" { port = var.service_port } - named_port { - name = "envoy" - port = var.envoy_port + dynamic "named_port" { + for_each = var.enable_external_traffic ? toset([1]) : toset([]) + content { + name = "envoy" + port = var.envoy_port + } } base_instance_name = "${var.service}-${var.environment}" diff --git a/production/terraform/gcp/services/autoscaling/variables.tf b/production/terraform/gcp/services/autoscaling/variables.tf index d070aa23..8b25de98 100644 --- a/production/terraform/gcp/services/autoscaling/variables.tf +++ b/production/terraform/gcp/services/autoscaling/variables.tf @@ -110,3 +110,9 @@ variable "shard_num" { description = "Shard number." type = string } + +variable "enable_external_traffic" { + description = "Whether to serve external traffic. If disabled, only internal traffic via service mesh will be served." + default = true + type = bool +} diff --git a/production/terraform/gcp/services/dashboards/main.tf b/production/terraform/gcp/services/dashboards/main.tf index f5cd84d1..88f7f0bf 100644 --- a/production/terraform/gcp/services/dashboards/main.tf +++ b/production/terraform/gcp/services/dashboards/main.tf @@ -20,9 +20,1273 @@ resource "google_monitoring_dashboard" "environment_dashboard" { "columns": 48, "tiles": [ { - "height": 19, + "height": 20, "widget": { - "title": "system.cpu.total_cores [MEAN]", + "title": "request.count [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/request.count\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 0 + }, + { + "height": 20, + "widget": { + "title": "Secure lookup request count [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/SecureLookupRequestCount\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 0 + }, + { + "height": 20, + "widget": { + "title": "request.failed_count_by_status [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/request.failed_count_by_status\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 20 + }, + { + "height": 20, + "widget": { + "title": "request.duration_ms [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/request.duration_ms\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 20 + }, + { + "height": 20, + "widget": { + "title": "request.size_bytes [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/request.size_bytes\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 40 + }, + { + "height": 20, + "widget": { + "title": "response.size_bytes [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/response.size_bytes\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 40 + }, + { + "height": 20, + "widget": { + "title": "Request Errors [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/KVUdfRequestError\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"error_code\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 60 + }, + { + "height": 20, + "widget": { + "title": "Internal Request Errors [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/InternalLookupRequestError\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"error_code\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 60 + }, + { + "height": 20, + "widget": { + "title": "Server Non-Request Errors [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/KVServerError\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"error_code\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 80 + }, + { + "height": 20, + "widget": { + "title": "Sharded Lookup Key Count By Shard [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/ShardedLookupKeyCountByShard\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 80 + }, + { + "height": 20, + "widget": { + "title": "Sharded Lookup GetKeyValues Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/ShardedLookupGetKeyValuesLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 100 + }, + { + "height": 20, + "widget": { + "title": "Sharded Lookup GetKeyValueSet Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/ShardedLookupGetKeyValueSetLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 100 + }, + { + "height": 20, + "widget": { + "title": "Sharded Lookup RunQuery Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/ShardedLookupRunQueryLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 120 + }, + { + "height": 20, + "widget": { + "title": "Internal GetKeyValues Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/InternalGetKeyValuesLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 120 + }, + { + "height": 20, + "widget": { + "title": "Internal GetKeyValueSet Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/InternalGetKeyValueSetLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 140 + }, + { + "height": 20, + "widget": { + "title": "Internal RunQuery Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/InternalRunQueryLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 140 + }, + { + "height": 20, + "widget": { + "title": "Cache GetKeyValuePairs Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/GetValuePairsLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 160 + }, + { + "height": 20, + "widget": { + "title": "Cache GetKeyValueSet Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/GetKeyValueSetLatencyInMicros\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 160 + }, + { + "height": 20, + "widget": { + "title": "Cache Access Event Count [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/CacheAccessEventCount\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"cache_access\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 180 + }, + { + "height": 20, + "widget": { + "title": "Server Retryable Operation Status Count [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/GetParameterStatus\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"status\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/CompleteLifecycleStatus\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"status\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/CreateDataOrchestratorStatus\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"status\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/StartDataOrchestratorStatus\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"status\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/LoadNewFilesStatus\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"status\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/GetShardManagerStatus\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"status\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/DescribeInstanceGroupInstancesStatus\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"status\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 180 + }, + { + "height": 20, + "widget": { + "title": "File Update Stats [MEAN]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/TotalRowsUpdatedInDataLoading\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\" metric.label.\"data_source\"!=\"realtime\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"data_source\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/TotalRowsDeletedInDataLoading\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\" metric.label.\"data_source\"!=\"realtime\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"data_source\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/TotalRowsDroppedInDataLoading\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\" metric.label.\"data_source\"!=\"realtime\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"data_source\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 200 + }, + { + "height": 20, + "widget": { + "title": "Data Reader Latency [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/ConcurrentStreamRecordReaderReadShardRecordsLatency\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/ConcurrentStreamRecordReaderReadStreamRecordsLatency\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 200 + }, + { + "height": 20, + "widget": { + "title": "Cache UpdateKeyValue Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/UpdateKeyValueLatency\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 220 + }, + { + "height": 20, + "widget": { + "title": "Cache UpdateKeyValueSet Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/UpdateKeyValueSetLatency\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 220 + }, + { + "height": 20, + "widget": { + "title": "Cache DeleteKey Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/DeleteKeyLatency\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 240 + }, + { + "height": 20, + "widget": { + "title": "Cache DeleteValuesInSet Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/DeleteValuesInSetLatency\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 240 + }, + { + "height": 20, + "widget": { + "title": "Cache RemoveDeletedKey Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/RemoveDeletedKeyLatency\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 0, + "yPos": 260 + }, + { + "height": 20, + "widget": { + "title": "Realtime Update Stats [MEAN]", "xyChart": { "chartOptions": {}, "dataSets": [ @@ -33,18 +1297,105 @@ resource "google_monitoring_dashboard" "environment_dashboard" { "timeSeriesQuery": { "timeSeriesFilter": { "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/TotalRowsUpdatedInDataLoading\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\" metric.label.\"data_source\"=\"realtime\"", + "secondaryAggregation": { "alignmentPeriod": "60s", "crossSeriesReducer": "REDUCE_MEAN", "groupByFields": [ - "metric.label.\"service_name\"", - "metric.label.\"deployment_environment\"", + "metric.label.\"Noise\"", "metric.label.\"shard_number\"", - "resource.label.\"task_id\"", - "metric.label.\"service_version\"" + "metric.label.\"service_instance_id\"" ], "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" }, - "filter": "metric.type=\"workload.googleapis.com/system.cpu.percent\" resource.type=\"generic_task\" metric.label.\"label\"=\"total cpu cores\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + "filter": "metric.type=\"workload.googleapis.com/TotalRowsDeletedInDataLoading\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\" metric.label.\"data_source\"=\"realtime\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "perSeriesAligner": "ALIGN_RATE" + }, + "filter": "metric.type=\"workload.googleapis.com/TotalRowsDroppedInDataLoading\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\" metric.label.\"data_source\"=\"realtime\"", + "secondaryAggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + } + } + } + } + ], + "yAxis": { + "scale": "LINEAR" + } + } + }, + "width": 24, + "xPos": 24, + "yPos": 260 + }, + { + "height": 20, + "widget": { + "title": "Realtime Message Processing Latency Microseconds [95TH PERCENTILE]", + "xyChart": { + "chartOptions": {}, + "dataSets": [ + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_PERCENTILE_95", + "groupByFields": [ + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_DELTA" + }, + "filter": "metric.type=\"workload.googleapis.com/ReceivedLowLatencyNotifications\" resource.type=\"generic_task\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" } } } @@ -56,10 +1407,10 @@ resource "google_monitoring_dashboard" "environment_dashboard" { }, "width": 24, "xPos": 0, - "yPos": 0 + "yPos": 280 }, { - "height": 19, + "height": 20, "widget": { "title": "system.cpu.percent [MEAN]", "xyChart": { @@ -75,12 +1426,10 @@ resource "google_monitoring_dashboard" "environment_dashboard" { "alignmentPeriod": "60s", "crossSeriesReducer": "REDUCE_MEAN", "groupByFields": [ - "metric.label.\"service_name\"", - "metric.label.\"deployment_environment\"", "metric.label.\"label\"", + "metric.label.\"Noise\"", "metric.label.\"shard_number\"", - "resource.label.\"task_id\"", - "metric.label.\"service_version\"" + "metric.label.\"service_instance_id\"" ], "perSeriesAligner": "ALIGN_MEAN" }, @@ -96,12 +1445,12 @@ resource "google_monitoring_dashboard" "environment_dashboard" { }, "width": 24, "xPos": 24, - "yPos": 0 + "yPos": 280 }, { - "height": 19, + "height": 20, "widget": { - "title": "system.memory.usage_kb for main process [MEAN]", + "title": "system.memory.usage_kb [MEAN]", "xyChart": { "chartOptions": {}, "dataSets": [ @@ -115,17 +1464,37 @@ resource "google_monitoring_dashboard" "environment_dashboard" { "alignmentPeriod": "60s", "crossSeriesReducer": "REDUCE_MEAN", "groupByFields": [ - "metric.label.\"service_name\"", - "metric.label.\"deployment_environment\"", + "metric.label.\"label\"", + "metric.label.\"Noise\"", "metric.label.\"shard_number\"", - "resource.label.\"task_id\"", - "metric.label.\"service_version\"" + "metric.label.\"service_instance_id\"" ], "perSeriesAligner": "ALIGN_MEAN" }, "filter": "metric.type=\"workload.googleapis.com/system.memory.usage_kb\" resource.type=\"generic_task\" metric.label.\"label\"=\"main process\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" } } + }, + { + "minAlignmentPeriod": "60s", + "plotType": "LINE", + "targetAxis": "Y1", + "timeSeriesQuery": { + "timeSeriesFilter": { + "aggregation": { + "alignmentPeriod": "60s", + "crossSeriesReducer": "REDUCE_MEAN", + "groupByFields": [ + "metric.label.\"label\"", + "metric.label.\"Noise\"", + "metric.label.\"shard_number\"", + "metric.label.\"service_instance_id\"" + ], + "perSeriesAligner": "ALIGN_MEAN" + }, + "filter": "metric.type=\"workload.googleapis.com/system.memory.usage_kb\" resource.type=\"generic_task\" metric.label.\"label\"=\"MemTotal:\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + } + } } ], "yAxis": { @@ -135,12 +1504,12 @@ resource "google_monitoring_dashboard" "environment_dashboard" { }, "width": 24, "xPos": 0, - "yPos": 19 + "yPos": 300 }, { - "height": 19, + "height": 20, "widget": { - "title": "system.memory.usage_kb for MemAvailable: [MEAN]", + "title": "system.cpu.total_cores [MEAN]", "xyChart": { "chartOptions": {}, "dataSets": [ @@ -154,16 +1523,13 @@ resource "google_monitoring_dashboard" "environment_dashboard" { "alignmentPeriod": "60s", "crossSeriesReducer": "REDUCE_MEAN", "groupByFields": [ - "metric.label.\"service_name\"", - "metric.label.\"deployment_environment\"", - "metric.label.\"label\"", + "metric.label.\"Noise\"", "metric.label.\"shard_number\"", - "resource.label.\"task_id\"", - "metric.label.\"service_version\"" + "metric.label.\"service_instance_id\"" ], "perSeriesAligner": "ALIGN_MEAN" }, - "filter": "metric.type=\"workload.googleapis.com/system.memory.usage_kb\" resource.type=\"generic_task\" metric.label.\"label\"=\"MemAvailable:\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" + "filter": "metric.type=\"workload.googleapis.com/system.cpu.percent\" resource.type=\"generic_task\" metric.label.\"label\"=\"total cpu cores\" metric.label.\"deployment_environment\"=\"${var.environment}\" metric.label.\"service_name\"=\"kv-server\"" } } } @@ -175,7 +1541,7 @@ resource "google_monitoring_dashboard" "environment_dashboard" { }, "width": 24, "xPos": 24, - "yPos": 19 + "yPos": 300 } ] } diff --git a/production/terraform/gcp/services/networking/main.tf b/production/terraform/gcp/services/networking/main.tf index bd749005..87c8fb07 100644 --- a/production/terraform/gcp/services/networking/main.tf +++ b/production/terraform/gcp/services/networking/main.tf @@ -27,7 +27,7 @@ resource "google_compute_subnetwork" "kv_server" { network = var.use_existing_vpc ? var.existing_vpc_id : google_compute_network.kv_server[0].id purpose = "PRIVATE" region = each.value - ip_cidr_range = "10.${each.key}.3.0/24" + ip_cidr_range = tolist(var.regions_cidr_blocks)[each.key] } resource "google_compute_router" "kv_server" { @@ -39,7 +39,10 @@ resource "google_compute_router" "kv_server" { } resource "google_compute_router_nat" "kv_server" { - for_each = google_compute_router.kv_server + for_each = { + for key, value in google_compute_router.kv_server : key => value + if !contains(var.regions_use_existing_nat, value.region) + } name = "${var.service}-${var.environment}-${each.value.region}-nat" router = each.value.name @@ -59,6 +62,7 @@ resource "google_compute_global_address" "collector" { } resource "google_compute_global_address" "kv_server" { + count = var.enable_external_traffic ? 1 : 0 name = "${var.service}-${var.environment}-xlb-ip" ip_version = "IPV4" } diff --git a/production/terraform/gcp/services/networking/outputs.tf b/production/terraform/gcp/services/networking/outputs.tf index feb20b94..3bb29175 100644 --- a/production/terraform/gcp/services/networking/outputs.tf +++ b/production/terraform/gcp/services/networking/outputs.tf @@ -28,5 +28,5 @@ output "collector_ip_address" { } output "server_ip_address" { - value = google_compute_global_address.kv_server.address + value = one(google_compute_global_address.kv_server[*].address) } diff --git a/production/terraform/gcp/services/networking/variables.tf b/production/terraform/gcp/services/networking/variables.tf index 95d04e5d..00bf42ec 100644 --- a/production/terraform/gcp/services/networking/variables.tf +++ b/production/terraform/gcp/services/networking/variables.tf @@ -29,6 +29,16 @@ variable "regions" { type = set(string) } +variable "regions_cidr_blocks" { + description = "A set of CIDR ranges for all specified regions. The number of blocks here should correspond to the number of regions." + type = set(string) +} + +variable "regions_use_existing_nat" { + description = "Regions that use existing nat. No new nats will be created for regions specified here." + type = set(string) +} + variable "collector_service_name" { description = "Assigned name of metrics collection service" type = string @@ -43,3 +53,9 @@ variable "existing_vpc_id" { description = "Existing vpc id. This would only be used if use_existing_vpc is true." type = string } + +variable "enable_external_traffic" { + description = "Whether to serve external traffic. If disabled, only internal traffic via service mesh will be served." + default = true + type = bool +} diff --git a/production/terraform/gcp/services/service_mesh/variables.tf b/production/terraform/gcp/services/service_mesh/variables.tf index 96acce7c..8d35f935 100644 --- a/production/terraform/gcp/services/service_mesh/variables.tf +++ b/production/terraform/gcp/services/service_mesh/variables.tf @@ -63,3 +63,9 @@ variable "existing_service_mesh" { description = "Existing service mesh. This would only be used if use_existing_service_mesh is true." type = string } + +variable "enable_external_traffic" { + description = "Whether to serve external traffic. If disabled, only internal traffic via service mesh will be served." + default = true + type = bool +} diff --git a/public/applications/pa/BUILD.bazel b/public/applications/pa/BUILD.bazel index 19c877cc..ff1bdd0b 100644 --- a/public/applications/pa/BUILD.bazel +++ b/public/applications/pa/BUILD.bazel @@ -37,7 +37,7 @@ cc_library( deps = [ "api_overlay_cc_proto", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) diff --git a/public/applications/pa/response_utils.cc b/public/applications/pa/response_utils.cc index c4e5ed82..bbc44965 100644 --- a/public/applications/pa/response_utils.cc +++ b/public/applications/pa/response_utils.cc @@ -15,7 +15,7 @@ #include "public/applications/pa/response_utils.h" #include "google/protobuf/util/json_util.h" -#include "src/cpp/util/status_macro/status_macros.h" +#include "src/util/status_macro/status_macros.h" namespace kv_server::application_pa { diff --git a/public/applications/pas/retrieval_request_builder.cc b/public/applications/pas/retrieval_request_builder.cc index 612a390f..45f2d420 100644 --- a/public/applications/pas/retrieval_request_builder.cc +++ b/public/applications/pas/retrieval_request_builder.cc @@ -16,14 +16,20 @@ namespace kv_server::application_pas { +v2::GetValuesRequest GetRequest() { + static const std::string* kClient = new std::string("Retrieval.20231018"); + v2::GetValuesRequest req; + req.set_client_version(*kClient); + (*(req.mutable_metadata()->mutable_fields()))["is_pas"].set_string_value( + "true"); + return req; +} + v2::GetValuesRequest BuildRetrievalRequest( std::string protected_signals, absl::flat_hash_map device_metadata, std::string contextual_signals, std::vector optional_ad_ids) { - static const std::string* kClient = new std::string("Retrieval.20231018"); - - v2::GetValuesRequest req; - req.set_client_version(*kClient); + v2::GetValuesRequest req = GetRequest(); v2::RequestPartition* partition = req.add_partitions(); { auto* protected_signals_arg = partition->add_arguments(); @@ -56,4 +62,17 @@ v2::GetValuesRequest BuildRetrievalRequest( return req; } +v2::GetValuesRequest BuildLookupRequest(std::vector ad_ids) { + v2::GetValuesRequest req = GetRequest(); + v2::RequestPartition* partition = req.add_partitions(); + auto* ad_id_arg = partition->add_arguments(); + for (auto&& item : std::move(ad_ids)) { + ad_id_arg->mutable_data() + ->mutable_list_value() + ->add_values() + ->set_string_value(std::move(item)); + } + return req; +} + } // namespace kv_server::application_pas diff --git a/public/applications/pas/retrieval_request_builder.h b/public/applications/pas/retrieval_request_builder.h index 7dd64298..99e2f7ce 100644 --- a/public/applications/pas/retrieval_request_builder.h +++ b/public/applications/pas/retrieval_request_builder.h @@ -35,6 +35,9 @@ v2::GetValuesRequest BuildRetrievalRequest( std::string contextual_signals, std::vector optional_ad_ids = {}); +// Builds a GetValuesRequest. Stores the input arguments into the request. +v2::GetValuesRequest BuildLookupRequest(std::vector ad_ids); + } // namespace kv_server::application_pas #endif // PUBLIC_APPLICATIONS_PAS_RETRIEVAL_REQUEST_BUILDER_H_ diff --git a/public/constants.cc b/public/constants.cc index f2105691..6fbeac95 100644 --- a/public/constants.cc +++ b/public/constants.cc @@ -44,4 +44,11 @@ const std::regex& LogicalShardingConfigFileFormatRegex() { return *regex; } +const std::regex& FileGroupFilenameFormatRegex() { + static const std::regex* const regex = new std::regex(absl::StrFormat( + R"((DELTA|SNAPSHOT)_\d{%d}_\d{%d}_OF_\d{%d})", kLogicalTimeDigits, + kFileGroupFileIndexDigits, kFileGroupSizeDigits)); + return *regex; +} + } // namespace kv_server diff --git a/public/constants.h b/public/constants.h index 10356049..ea2a346a 100644 --- a/public/constants.h +++ b/public/constants.h @@ -38,6 +38,12 @@ constexpr std::string_view kFileComponentDelimiter = "_"; // Number of digits in logical time in file basename. constexpr int kLogicalTimeDigits = 16; +// Number of digits used to represent the index of a file in a file group. +constexpr int kFileGroupFileIndexDigits = 5; + +// Number of digits used to represent the size of a file group. +constexpr int kFileGroupSizeDigits = 6; + // "DELTA_\d{16}" // The first component represents the file type. // @@ -65,10 +71,26 @@ constexpr char kQueryArgDelimiter = ','; // indicates a more recent snapshot. const std::regex& SnapshotFileFormatRegex(); +// Returns a compiled file group file name regex defined as follows: +// +// Compiled regex = "(DELTA|SNAPSHOT)_\d{16}_\d{5}_OF_\d{6}". Regex parts +// correspond to the following parts: +// NAME REGEX: +// ___OF_ +// WHERE: +// (1) FILE_TYPE - type of file. Valid values: [DELTA, SNAPSHOT]. +// (2) LOGICAL_TIMESTAMP - strictly increasing 16 digit number that represents +// logical time, a larger value means a more recent file. +// (3) PART_FILE_INDEX - zero-based 5 digit index of file in the file group. +// Valid range: [0..NUM_PART_FILES). +// (4) NUM_PART_FILES = a 6 digit number representing the total number of +// files in a file group. Valid range: [1..100,000] +const std::regex& FileGroupFilenameFormatRegex(); + // X25519 public key used to test/debug/demo the ObliviousGetValues query API. // ObliviousGetValues requests encrypted with this key can be processed by the // server. -// For cross code base consistency, matches the test key we use in the commong +// For cross code base consistency, matches the test key we use in the common // repo, see ../encryption/key_fetcher/src/fake_key_fetcher_manager.h constexpr std::string_view kTestPublicKey = "f3b7b2f1764f5c077effecad2afd86154596e63f7375ea522761b881e6c3c323"; @@ -76,12 +98,12 @@ constexpr std::string_view kTestPublicKey = // Parameters used to configure Oblivious HTTP according to // https://github.com/WICG/turtledove/blob/main/FLEDGE_Key_Value_Server_API.md#encryption // -// KEM: DHKEM(X25519, HKDF-SHA256) 0x0020 +// KEM: DHKEM(X25519, HKDF-SHA256) const uint16_t kKEMParameter = 0x0020; -// KDF: HKDF-SHA256 0x0001 +// KDF: HKDF-SHA256 const uint16_t kKDFParameter = 0x0001; -// AEAD: AES-128-GCM 0X0001 -const uint16_t kAEADParameter = 0x0001; +// AEAD: AES-256-GCM +const uint16_t kAEADParameter = 0x0002; constexpr std::string_view kServiceName = "kv-server"; diff --git a/public/data_loading/BUILD.bazel b/public/data_loading/BUILD.bazel index 8d0caf97..3286748e 100644 --- a/public/data_loading/BUILD.bazel +++ b/public/data_loading/BUILD.bazel @@ -56,7 +56,7 @@ cc_library( deps = [ ":data_loading_fbs", ":record_utils", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", ], @@ -68,7 +68,7 @@ cc_library( hdrs = ["record_utils.h"], deps = [ ":data_loading_fbs", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -112,8 +112,9 @@ cc_library( srcs = ["filename_utils.cc"], hdrs = ["filename_utils.h"], deps = [ + "//public:base_types_cc_proto", "//public:constants", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/public/data_loading/aggregation/BUILD.bazel b/public/data_loading/aggregation/BUILD.bazel index 673b48da..48af8024 100644 --- a/public/data_loading/aggregation/BUILD.bazel +++ b/public/data_loading/aggregation/BUILD.bazel @@ -25,8 +25,8 @@ cc_library( "//public/data_loading:records_utils", "//public/data_loading/writers:delta_record_stream_writer", "//public/data_loading/writers:delta_record_writer", - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -51,6 +51,7 @@ cc_test( cc_binary( name = "record_aggregator_benchmarks", srcs = ["record_aggregator_benchmarks.cc"], + malloc = "@com_google_tcmalloc//tcmalloc", deps = [ ":record_aggregator", "//public/data_loading:records_utils", diff --git a/public/data_loading/aggregation/record_aggregator.cc b/public/data_loading/aggregation/record_aggregator.cc index 893cf36b..4a2acee3 100644 --- a/public/data_loading/aggregation/record_aggregator.cc +++ b/public/data_loading/aggregation/record_aggregator.cc @@ -21,11 +21,11 @@ #include #include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" #include "absl/memory/memory.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" -#include "glog/logging.h" #include "public/data_loading/data_loading_generated.h" #include "public/data_loading/writers/delta_record_stream_writer.h" #include "public/data_loading/writers/delta_record_writer.h" diff --git a/public/data_loading/csv/BUILD.bazel b/public/data_loading/csv/BUILD.bazel index b36e5a7b..396907f0 100644 --- a/public/data_loading/csv/BUILD.bazel +++ b/public/data_loading/csv/BUILD.bazel @@ -67,7 +67,7 @@ cc_library( "@com_google_riegeli//riegeli/bytes:istream_reader", "@com_google_riegeli//riegeli/csv:csv_reader", "@com_google_riegeli//riegeli/csv:csv_record", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) @@ -75,11 +75,10 @@ cc_test( name = "csv_delta_record_stream_reader_test", size = "small", srcs = ["csv_delta_record_stream_reader_test.cc"], - env = {"GLOG_v": "10"}, deps = [ ":csv_delta_record_stream_reader", ":csv_delta_record_stream_writer", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_googletest//:gtest_main", ], ) diff --git a/public/data_loading/csv/csv_delta_record_stream_reader.cc b/public/data_loading/csv/csv_delta_record_stream_reader.cc index 6662c912..433c34ff 100644 --- a/public/data_loading/csv/csv_delta_record_stream_reader.cc +++ b/public/data_loading/csv/csv_delta_record_stream_reader.cc @@ -18,11 +18,11 @@ #include +#include "absl/log/log.h" #include "absl/strings/ascii.h" #include "absl/strings/escaping.h" #include "absl/strings/match.h" #include "absl/strings/str_split.h" -#include "glog/logging.h" #include "public/data_loading/record_utils.h" namespace kv_server { diff --git a/public/data_loading/csv/csv_delta_record_stream_reader.h b/public/data_loading/csv/csv_delta_record_stream_reader.h index 1dc17028..31ca6d1b 100644 --- a/public/data_loading/csv/csv_delta_record_stream_reader.h +++ b/public/data_loading/csv/csv_delta_record_stream_reader.h @@ -21,7 +21,7 @@ #include #include -#include "glog/logging.h" +#include "absl/log/log.h" #include "public/data_loading/csv/constants.h" #include "public/data_loading/readers/delta_record_reader.h" #include "public/data_loading/record_utils.h" diff --git a/public/data_loading/csv/csv_delta_record_stream_reader_test.cc b/public/data_loading/csv/csv_delta_record_stream_reader_test.cc index 0cab98cb..5ec95cb0 100644 --- a/public/data_loading/csv/csv_delta_record_stream_reader_test.cc +++ b/public/data_loading/csv/csv_delta_record_stream_reader_test.cc @@ -18,7 +18,7 @@ #include -#include "glog/logging.h" +#include "absl/log/log.h" #include "gmock/gmock.h" #include "gtest/gtest.h" #include "public/data_loading/csv/csv_delta_record_stream_writer.h" diff --git a/public/data_loading/csv/csv_delta_record_stream_writer.cc b/public/data_loading/csv/csv_delta_record_stream_writer.cc index f884a389..1fc1d970 100644 --- a/public/data_loading/csv/csv_delta_record_stream_writer.cc +++ b/public/data_loading/csv/csv_delta_record_stream_writer.cc @@ -16,10 +16,10 @@ #include "public/data_loading/csv/csv_delta_record_stream_writer.h" +#include "absl/log/log.h" #include "absl/strings/escaping.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" -#include "glog/logging.h" #include "public/data_loading/data_loading_generated.h" namespace kv_server { diff --git a/public/data_loading/filename_utils.cc b/public/data_loading/filename_utils.cc index c1b1fff0..291a0a16 100644 --- a/public/data_loading/filename_utils.cc +++ b/public/data_loading/filename_utils.cc @@ -17,9 +17,9 @@ #include #include // NOLINT +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/str_format.h" -#include "glog/logging.h" #include "public/constants.h" namespace kv_server { @@ -84,4 +84,31 @@ absl::StatusOr ToLogicalShardingConfigFilename( return result; } +bool IsFileGroupFileName(std::string_view filename) { + return std::regex_match(filename.begin(), filename.end(), + FileGroupFilenameFormatRegex()); +} + +absl::StatusOr ToFileGroupFileName(FileType::Enum file_type, + uint64_t logical_commit_time, + uint64_t file_index, + uint64_t file_group_size) { + if (file_type != FileType::DELTA && file_type != FileType::SNAPSHOT) { + return absl::InvalidArgumentError( + absl::StrCat("File groups are not supported for file type: ", + FileType_Enum_Name(file_type))); + } + if (file_index >= file_group_size) { + return absl::InvalidArgumentError( + absl::StrCat("file index: ", file_index, + " must be less than file group size: ", file_group_size)); + } + return absl::StrFormat( + "%s%s%0*d%s%0*d%sOF%s%0*d", FileType::Enum_Name(file_type), + kFileComponentDelimiter, kLogicalTimeDigits, logical_commit_time, + kFileComponentDelimiter, kFileGroupFileIndexDigits, file_index, + kFileComponentDelimiter, kFileComponentDelimiter, kFileGroupSizeDigits, + file_group_size); +} + } // namespace kv_server diff --git a/public/data_loading/filename_utils.h b/public/data_loading/filename_utils.h index a331f7ba..3405ce63 100644 --- a/public/data_loading/filename_utils.h +++ b/public/data_loading/filename_utils.h @@ -21,6 +21,7 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "public/base_types.pb.h" namespace kv_server { @@ -55,6 +56,19 @@ bool IsLogicalShardingConfigFilename(std::string_view basename); absl::StatusOr ToLogicalShardingConfigFilename( uint64_t logical_commit_time); +// Returns true if `filename` is a valid file group filename. +// +// Valid file group filenames conform to the regex return by +// `FileGroupFilenameFormatRegex()` in constants.h +bool IsFileGroupFileName(std::string_view filename); + +// Attempts to construct a valid file group file name from the inputs. +// Returns a descriptive `absl::InvalidArgumentError` on error. +absl::StatusOr ToFileGroupFileName(FileType::Enum file_type, + uint64_t logical_commit_time, + uint64_t file_index, + uint64_t file_group_size); + } // namespace kv_server #endif // PUBLIC_DATA_LOADING_FILENAME_UTILS_H_ diff --git a/public/data_loading/filename_utils_test.cc b/public/data_loading/filename_utils_test.cc index d9e78ea7..e8882ecc 100644 --- a/public/data_loading/filename_utils_test.cc +++ b/public/data_loading/filename_utils_test.cc @@ -95,5 +95,41 @@ TEST(SnapshotFilename, ToLogicalShardingConfigFilename) { ("LOGICAL_SHARDING_CONFIG_1234512345123451")); } +TEST(FileGroupFilename, IsFileGroupFileName) { + EXPECT_FALSE(IsFileGroupFileName("")); + EXPECT_FALSE(IsFileGroupFileName("DELTA")); + EXPECT_FALSE(IsFileGroupFileName("SNAPSHOT")); + EXPECT_FALSE(IsFileGroupFileName("DELTA_1234512345123451")); + EXPECT_FALSE(IsFileGroupFileName("SNAPSHOT_1234512345123451")); + EXPECT_FALSE(IsFileGroupFileName("SNAPSHOT_1234512345123451_00000")); + EXPECT_TRUE(IsFileGroupFileName("DELTA_1234512345123451_00000_OF_000100")); + EXPECT_TRUE(IsFileGroupFileName("SNAPSHOT_1234512345123451_00000_OF_000100")); +} + +TEST(FileGroupFilename, ToFileGroupFileNameInvalidInputs) { + auto filename = + ToFileGroupFileName(FileType::DELTA, /*logical_commit_time=*/10, + /*file_index=*/11, /*file_group_size=*/10); + EXPECT_FALSE(filename.ok()) << filename.status(); + EXPECT_THAT(filename.status().code(), absl::StatusCode::kInvalidArgument); + filename = ToFileGroupFileName(FileType::FILE_TYPE_UNSPECIFIED, + /*logical_commit_time=*/10, + /*file_index=*/1, /*file_group_size=*/10); + EXPECT_FALSE(filename.ok()) << filename.status(); + EXPECT_THAT(filename.status().code(), absl::StatusCode::kInvalidArgument); +} + +TEST(FileGroupFilename, ToFileGroupFileNameValidInputs) { + auto filename = + ToFileGroupFileName(FileType::DELTA, /*logical_commit_time=*/10, + /*file_index=*/1, /*file_group_size=*/10); + ASSERT_TRUE(filename.ok()) << filename.status(); + EXPECT_THAT(*filename, "DELTA_0000000000000010_00001_OF_000010"); + filename = ToFileGroupFileName(FileType::SNAPSHOT, /*logical_commit_time=*/10, + /*file_index=*/0, /*file_group_size=*/10); + ASSERT_TRUE(filename.ok()) << filename.status(); + EXPECT_THAT(*filename, "SNAPSHOT_0000000000000010_00000_OF_000010"); +} + } // namespace } // namespace kv_server diff --git a/public/data_loading/readers/BUILD.bazel b/public/data_loading/readers/BUILD.bazel index 2a680fdc..6cb118e1 100644 --- a/public/data_loading/readers/BUILD.bazel +++ b/public/data_loading/readers/BUILD.bazel @@ -33,15 +33,16 @@ cc_library( ":stream_record_reader", "//components/telemetry:server_definition", "//public/data_loading:riegeli_metadata_cc_proto", - "@com_github_google_glog//:glog", "@com_google_absl//absl/base", "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_riegeli//riegeli/bytes:istream_reader", "@com_google_riegeli//riegeli/records:record_reader", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", ], ) @@ -60,13 +61,14 @@ cc_library( "//components/telemetry:server_definition", "//public/data_loading:riegeli_metadata_cc_proto", "@avro//:avrocpp", - "@com_github_google_glog//:glog", "@com_google_absl//absl/base", "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", ], ) @@ -75,7 +77,7 @@ cc_library( hdrs = ["stream_record_reader_factory.h"], deps = [ ":stream_record_reader", - "@google_privacysandbox_servers_common//src/cpp/telemetry:telemetry_provider", + "@google_privacysandbox_servers_common//src/telemetry:telemetry_provider", ], ) @@ -124,7 +126,6 @@ cc_test( "@com_google_riegeli//riegeli/bytes:ostream_writer", "@com_google_riegeli//riegeli/bytes:string_writer", "@com_google_riegeli//riegeli/records:record_writer", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -145,7 +146,6 @@ cc_test( "//public/test_util:mocks", "//public/test_util:proto_matcher", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -174,7 +174,7 @@ cc_library( ":delta_record_reader", "//public/data_loading:records_utils", "@avro//:avrocpp", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", ], diff --git a/public/data_loading/readers/avro_stream_io.cc b/public/data_loading/readers/avro_stream_io.cc index cc5a7426..cc648df0 100644 --- a/public/data_loading/readers/avro_stream_io.cc +++ b/public/data_loading/readers/avro_stream_io.cc @@ -14,6 +14,7 @@ #include "public/data_loading/readers/avro_stream_io.h" +#include "absl/log/check.h" #include "third_party/avro/api/DataFile.hh" #include "third_party/avro/api/Schema.hh" #include "third_party/avro/api/Stream.hh" @@ -25,15 +26,19 @@ AvroStreamReader::AvroStreamReader(std::istream& data_input) absl::Status AvroStreamReader::ReadStreamRecords( const std::function& callback) { - avro::InputStreamPtr input_stream = avro::istreamInputStream(data_input_); - avro::DataFileReader reader(std::move(input_stream)); + try { + avro::InputStreamPtr input_stream = avro::istreamInputStream(data_input_); + avro::DataFileReader reader(std::move(input_stream)); - std::string record; - absl::Status overall_status; - while (reader.read(record)) { - overall_status.Update(callback(record)); + std::string record; + absl::Status overall_status; + while (reader.read(record)) { + overall_status.Update(callback(record)); + } + return overall_status; + } catch (const std::exception& e) { + return absl::InternalError(e.what()); } - return overall_status; } AvroConcurrentStreamRecordReader::AvroConcurrentStreamRecordReader( @@ -95,7 +100,10 @@ AvroConcurrentStreamRecordReader::BuildByteRanges() { // stream are read. absl::Status AvroConcurrentStreamRecordReader::ReadStreamRecords( const std::function& callback) { - auto start_time = absl::Now(); + ScopeLatencyMetricsRecorder< + ServerSafeMetricsContext, + kConcurrentStreamRecordReaderReadStreamRecordsLatency> + latency_recorder(KVServerContextMap()->SafeMetric()); auto byte_ranges = BuildByteRanges(); if (!byte_ranges.ok() || byte_ranges->empty()) { return byte_ranges.status(); @@ -122,14 +130,9 @@ absl::Status AvroConcurrentStreamRecordReader::ReadStreamRecords( } total_records_read += curr_byte_range_result->num_records_read; } - auto duration = absl::Now() - start_time; VLOG(2) << "Done reading " << total_records_read << " records in " - << absl::ToDoubleMilliseconds(duration) << " ms."; - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogHistogram( - absl::ToDoubleMicroseconds(duration))); + << absl::ToDoubleMilliseconds(latency_recorder.GetLatency()) + << " ms."; return absl::OkStatus(); } @@ -153,7 +156,9 @@ AvroConcurrentStreamRecordReader::ReadByteRange( VLOG(2) << "Reading byte_range: " << "[" << byte_range.begin_offset << "," << byte_range.end_offset << "]"; - auto start_time = absl::Now(); + ScopeLatencyMetricsRecorder + latency_recorder(KVServerContextMap()->SafeMetric()); auto record_stream = stream_factory_(); VLOG(9) << "creating input stream"; avro::InputStreamPtr input_stream = @@ -182,15 +187,10 @@ AvroConcurrentStreamRecordReader::ReadByteRange( << overall_status; return overall_status; } - auto duration = absl::Now() - start_time; VLOG(2) << "Done reading " << num_records_read << " records in byte_range: [" << byte_range.begin_offset << "," << byte_range.end_offset << "] in " - << absl::ToDoubleMilliseconds(duration) << " ms."; - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogHistogram( - absl::ToDoubleMicroseconds(duration))); + << absl::ToDoubleMilliseconds(latency_recorder.GetLatency()) + << " ms."; ByteRangeResult result; result.num_records_read = num_records_read; return result; diff --git a/public/data_loading/readers/avro_stream_io.h b/public/data_loading/readers/avro_stream_io.h index 38ae8c30..de7b8c35 100644 --- a/public/data_loading/readers/avro_stream_io.h +++ b/public/data_loading/readers/avro_stream_io.h @@ -27,14 +27,14 @@ #include "absl/base/optimization.h" #include "absl/cleanup/cleanup.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "components/telemetry/server_definition.h" -#include "glog/logging.h" #include "public/data_loading/readers/stream_record_reader.h" #include "public/data_loading/riegeli_metadata.pb.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" namespace kv_server { diff --git a/public/data_loading/readers/avro_stream_io_test.cc b/public/data_loading/readers/avro_stream_io_test.cc index 181abfdb..59e4e5b0 100644 --- a/public/data_loading/readers/avro_stream_io_test.cc +++ b/public/data_loading/readers/avro_stream_io_test.cc @@ -19,7 +19,6 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" #include "third_party/avro/api/DataFile.hh" #include "third_party/avro/api/Schema.hh" #include "third_party/avro/api/ValidSchema.hh" @@ -30,6 +29,11 @@ namespace { constexpr std::string_view kTestRecord = "testrecord"; constexpr int64_t kIterations = 1024 * 1024 * 9; +void WriteInvalidFile(const std::vector& records, + std::ostream& dest_stream) { + dest_stream << "invalid"; +} + void WriteAvroToFile(const std::vector& records, std::ostream& dest_stream) { avro::OutputStreamPtr avro_output_stream = @@ -122,6 +126,44 @@ TEST(AvroStreamIO, SequentialReading) { record_reader.ReadStreamRecords(record_callback.AsStdFunction()); EXPECT_TRUE(status.ok()) << status; } +TEST(AvroStreamIO, ConcurrentReadingInvalidFile) { + kv_server::InitMetricsContextMap(); + constexpr std::string_view kFileName = "ConcurrentReading.invalid"; + const std::filesystem::path path = + std::filesystem::path(::testing::TempDir()) / kFileName; + std::ofstream output_stream(path); + WriteInvalidFile({kTestRecord}, output_stream); + output_stream.close(); + + AvroConcurrentStreamRecordReader::Options options; + AvroConcurrentStreamRecordReader record_reader( + [&path] { return std::make_unique(path); }, options); + + testing::MockFunction record_callback; + EXPECT_CALL(record_callback, Call).Times(0); + auto status = + record_reader.ReadStreamRecords(record_callback.AsStdFunction()); + EXPECT_FALSE(status.ok()); +} + +TEST(AvroStreamIO, SequentialReadingInvalidFile) { + kv_server::InitMetricsContextMap(); + constexpr std::string_view kFileName = "SequentialReading.invalid"; + const std::filesystem::path path = + std::filesystem::path(::testing::TempDir()) / kFileName; + std::ofstream output_stream(path); + WriteInvalidFile({kTestRecord}, output_stream); + output_stream.close(); + + std::ifstream is(path); + AvroStreamReader record_reader(is); + + testing::MockFunction record_callback; + EXPECT_CALL(record_callback, Call).Times(0); + auto status = + record_reader.ReadStreamRecords(record_callback.AsStdFunction()); + EXPECT_FALSE(status.ok()); +} } // namespace } // namespace kv_server diff --git a/public/data_loading/readers/riegeli_stream_io.h b/public/data_loading/readers/riegeli_stream_io.h index d863f4e1..b15ea78f 100644 --- a/public/data_loading/readers/riegeli_stream_io.h +++ b/public/data_loading/readers/riegeli_stream_io.h @@ -27,16 +27,17 @@ #include "absl/base/optimization.h" #include "absl/cleanup/cleanup.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "components/telemetry/server_definition.h" -#include "glog/logging.h" #include "public/data_loading/readers/stream_record_reader.h" #include "public/data_loading/riegeli_metadata.pb.h" #include "riegeli/bytes/istream_reader.h" #include "riegeli/records/record_reader.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" namespace kv_server { @@ -57,7 +58,8 @@ class RiegeliStreamReader : public StreamRecordReader { riegeli::RecordsMetadata metadata; if (!reader_.ReadMetadata(metadata)) { if (reader_.ok()) { - return absl::UnavailableError("Metadata not found"); + return absl::UnavailableError( + "Metadata not found. Please ensure metadata is set properly."); } return reader_.status(); } @@ -74,7 +76,9 @@ class RiegeliStreamReader : public StreamRecordReader { RecordT record; absl::Status overall_status; while (reader_.ReadRecord(record)) { - overall_status.Update(callback(record)); + const auto callback_status = callback(record); + LOG_IF(WARNING, !callback_status.ok()); + overall_status.Update(callback_status); } if (!overall_status.ok()) { LOG(ERROR) << overall_status; @@ -247,7 +251,10 @@ ConcurrentStreamRecordReader::BuildShards() { template absl::Status ConcurrentStreamRecordReader::ReadStreamRecords( const std::function& callback) { - auto start_time = absl::Now(); + ScopeLatencyMetricsRecorder< + ServerSafeMetricsContext, + kConcurrentStreamRecordReaderReadStreamRecordsLatency> + latency_recorder(KVServerContextMap()->SafeMetric()); auto shards = BuildShards(); if (!shards.ok() || shards->empty()) { return shards.status(); @@ -285,14 +292,9 @@ absl::Status ConcurrentStreamRecordReader::ReadStreamRecords( total_records_read += curr_shard_result->num_records_read; prev_shard_result = curr_shard_result; } - auto duration = absl::Now() - start_time; VLOG(2) << "Done reading " << total_records_read << " records in " - << absl::ToDoubleMilliseconds(duration) << " ms."; - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogHistogram( - absl::ToDoubleMicroseconds(duration))); + << absl::ToDoubleMilliseconds(latency_recorder.GetLatency()) + << " ms."; return absl::OkStatus(); } @@ -303,7 +305,10 @@ ConcurrentStreamRecordReader::ReadShardRecords( const std::function& record_callback) { VLOG(2) << "Reading shard: " << "[" << shard.start_pos << "," << shard.end_pos << "]"; - auto start_time = absl::Now(); + ScopeLatencyMetricsRecorder< + ServerSafeMetricsContext, + kConcurrentStreamRecordReaderReadShardRecordsLatency> + latency_recorder(KVServerContextMap()->SafeMetric()); auto record_stream = stream_factory_(); riegeli::RecordReader> record_reader( riegeli::IStreamReader(&record_stream->Stream()), @@ -334,15 +339,10 @@ ConcurrentStreamRecordReader::ReadShardRecords( } shard_result.next_shard_first_record_pos = next_record_pos; shard_result.num_records_read = num_records_read; - auto duration = absl::Now() - start_time; VLOG(2) << "Done reading " << num_records_read << " records in shard: [" << shard.start_pos << "," << shard.end_pos << "] in " - << absl::ToDoubleMilliseconds(duration) << " ms."; - LogIfError( - KVServerContextMap() - ->SafeMetric() - .LogHistogram( - absl::ToDoubleMicroseconds(duration))); + << absl::ToDoubleMilliseconds(latency_recorder.GetLatency()) + << " ms."; return shard_result; } diff --git a/public/data_loading/readers/riegeli_stream_io_test.cc b/public/data_loading/readers/riegeli_stream_io_test.cc index e231f034..7bb4998a 100644 --- a/public/data_loading/readers/riegeli_stream_io_test.cc +++ b/public/data_loading/readers/riegeli_stream_io_test.cc @@ -20,6 +20,7 @@ #include #include +#include "absl/log/check.h" #include "absl/status/status.h" #include "gmock/gmock.h" #include "google/protobuf/text_format.h" @@ -30,7 +31,6 @@ #include "riegeli/bytes/ostream_writer.h" #include "riegeli/bytes/string_writer.h" #include "riegeli/records/record_writer.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { @@ -268,6 +268,7 @@ class NonSeekingStringBlobStream : public RecordStream { }; TEST(ConcurrentStreamRecordReaderTest, FailsToReadNonSeekingStream) { + kv_server::InitMetricsContextMap(); std::string content; auto writer = riegeli::RecordWriter(riegeli::StringWriter(&content), riegeli::RecordWriterBase::Options()); diff --git a/public/data_loading/readers/stream_record_reader_factory.h b/public/data_loading/readers/stream_record_reader_factory.h index 5821666f..3910d749 100644 --- a/public/data_loading/readers/stream_record_reader_factory.h +++ b/public/data_loading/readers/stream_record_reader_factory.h @@ -20,7 +20,7 @@ #include #include "public/data_loading/readers/stream_record_reader.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" namespace kv_server { diff --git a/public/data_loading/record_utils.cc b/public/data_loading/record_utils.cc index 8f30922f..75d1ee4c 100644 --- a/public/data_loading/record_utils.cc +++ b/public/data_loading/record_utils.cc @@ -14,9 +14,9 @@ #include "public/data_loading/record_utils.h" +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" -#include "glog/logging.h" namespace kv_server { namespace { diff --git a/public/data_loading/records_utils.cc b/public/data_loading/records_utils.cc index eeb14eb6..73c42b99 100644 --- a/public/data_loading/records_utils.cc +++ b/public/data_loading/records_utils.cc @@ -16,8 +16,8 @@ #include +#include "absl/log/log.h" #include "absl/status/statusor.h" -#include "glog/logging.h" namespace kv_server { namespace { diff --git a/public/data_loading/writers/BUILD.bazel b/public/data_loading/writers/BUILD.bazel index 1a1e0f8f..3c3189bb 100644 --- a/public/data_loading/writers/BUILD.bazel +++ b/public/data_loading/writers/BUILD.bazel @@ -33,7 +33,7 @@ cc_library( deps = [ ":delta_record_writer", "//public/data_loading:records_utils", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_riegeli//riegeli/bytes:ostream_writer", @@ -48,7 +48,7 @@ cc_library( ":delta_record_writer", "//public/data_loading:records_utils", "@avro//:avrocpp", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", ], @@ -64,7 +64,6 @@ cc_test( "//public/data_loading/readers:riegeli_stream_record_reader_factory", "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) @@ -102,7 +101,7 @@ cc_library( deps = [ "//public/data_loading:records_utils", "//public/sharding:sharding_function", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -127,7 +126,7 @@ cc_library( hdrs = ["delta_record_limiting_file_writer.h"], deps = [ ":delta_record_writer", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_riegeli//riegeli/bytes:fd_writer", diff --git a/public/data_loading/writers/avro_delta_record_stream_writer.h b/public/data_loading/writers/avro_delta_record_stream_writer.h index 7e0b663e..c816162f 100644 --- a/public/data_loading/writers/avro_delta_record_stream_writer.h +++ b/public/data_loading/writers/avro_delta_record_stream_writer.h @@ -22,9 +22,9 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "glog/logging.h" #include "public/data_loading/records_utils.h" #include "public/data_loading/writers/delta_record_writer.h" #include "third_party/avro/api/DataFile.hh" diff --git a/public/data_loading/writers/delta_record_limiting_file_writer.cc b/public/data_loading/writers/delta_record_limiting_file_writer.cc index 415ca5fc..0db8e406 100644 --- a/public/data_loading/writers/delta_record_limiting_file_writer.cc +++ b/public/data_loading/writers/delta_record_limiting_file_writer.cc @@ -14,7 +14,7 @@ #include "public/data_loading/writers/delta_record_limiting_file_writer.h" -#include "glog/logging.h" +#include "absl/log/log.h" namespace kv_server { diff --git a/public/data_loading/writers/delta_record_stream_writer.h b/public/data_loading/writers/delta_record_stream_writer.h index df08766e..4567f559 100644 --- a/public/data_loading/writers/delta_record_stream_writer.h +++ b/public/data_loading/writers/delta_record_stream_writer.h @@ -21,9 +21,9 @@ #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "glog/logging.h" #include "public/data_loading/records_utils.h" #include "public/data_loading/writers/delta_record_writer.h" #include "riegeli/bytes/ostream_writer.h" diff --git a/public/data_loading/writers/delta_record_stream_writer_test.cc b/public/data_loading/writers/delta_record_stream_writer_test.cc index 771f37c1..08c71eaf 100644 --- a/public/data_loading/writers/delta_record_stream_writer_test.cc +++ b/public/data_loading/writers/delta_record_stream_writer_test.cc @@ -24,7 +24,6 @@ #include "public/data_loading/readers/riegeli_stream_io.h" #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" #include "public/data_loading/records_utils.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/public/query/cpp/BUILD.bazel b/public/query/cpp/BUILD.bazel index 0ca5e2d7..d0298d21 100644 --- a/public/query/cpp/BUILD.bazel +++ b/public/query/cpp/BUILD.bazel @@ -23,7 +23,7 @@ cc_library( deps = [ "//public/query/v2:get_values_v2_cc_proto", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) diff --git a/public/query/cpp/client_utils.cc b/public/query/cpp/client_utils.cc index 267a22a8..84eb0168 100644 --- a/public/query/cpp/client_utils.cc +++ b/public/query/cpp/client_utils.cc @@ -15,7 +15,7 @@ #include "public/query/cpp/client_utils.h" #include "google/protobuf/util/json_util.h" -#include "src/cpp/util/status_macro/status_macros.h" +#include "src/util/status_macro/status_macros.h" namespace kv_server { diff --git a/public/query/v2/BUILD.bazel b/public/query/v2/BUILD.bazel index 0f11e3e9..ce6623c7 100644 --- a/public/query/v2/BUILD.bazel +++ b/public/query/v2/BUILD.bazel @@ -37,6 +37,7 @@ proto_library( "@com_google_googleapis//google/api:httpbody_proto", "@com_google_googleapis//google/rpc:status_proto", "@com_google_protobuf//:struct_proto", + "@google_privacysandbox_servers_common//src/logger:logger_proto", ], ) diff --git a/public/query/v2/get_values_v2.proto b/public/query/v2/get_values_v2.proto index ca13819e..f161711d 100644 --- a/public/query/v2/get_values_v2.proto +++ b/public/query/v2/get_values_v2.proto @@ -21,6 +21,7 @@ import "google/api/httpbody.proto"; import "google/protobuf/struct.proto"; import "google/rpc/status.proto"; import "public/api_schema.proto"; +import "src/logger/logger.proto"; // Key Value Service API V2. // Spec: @@ -107,6 +108,10 @@ message GetValuesRequest { // Metadata that is useful for all partitions in a request. google.protobuf.Struct metadata = 2; repeated RequestPartition partitions = 3; + // Context useful for logging and tracing requests + privacy_sandbox.server_common.LogContext log_context = 4; + // Consented debugging configuration + privacy_sandbox.server_common.ConsentedDebugConfiguration consented_debug_config = 5; } message ResponsePartition { diff --git a/public/test_util/BUILD.bazel b/public/test_util/BUILD.bazel index 2444921e..f5ee9f3f 100644 --- a/public/test_util/BUILD.bazel +++ b/public/test_util/BUILD.bazel @@ -34,6 +34,5 @@ cc_library( "//public/data_loading/readers:riegeli_stream_io", "//public/data_loading/readers:stream_record_reader", "@com_google_googletest//:gtest", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) diff --git a/public/test_util/mocks.h b/public/test_util/mocks.h index 1a9e192d..936a409c 100644 --- a/public/test_util/mocks.h +++ b/public/test_util/mocks.h @@ -23,7 +23,6 @@ #include "public/data_loading/readers/riegeli_stream_io.h" #include "public/data_loading/readers/stream_record_reader.h" #include "public/data_loading/riegeli_metadata.pb.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { diff --git a/public/udf/BUILD.bazel b/public/udf/BUILD.bazel index 70a056e5..51ed56ca 100644 --- a/public/udf/BUILD.bazel +++ b/public/udf/BUILD.bazel @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -load("@io_bazel_rules_closure//closure:defs.bzl", "closure_js_proto_library") +load("@google_privacysandbox_servers_common//src/roma/tools/api_plugin:roma_api.bzl", "declare_roma_api", "js_proto_library") load("@rules_cc//cc:defs.bzl", "cc_library", "cc_proto_library") load("@rules_proto//proto:defs.bzl", "proto_library") @@ -30,12 +30,16 @@ cc_proto_library( deps = [":binary_get_values_proto"], ) -# Library that can be added as a dep from closure_js_libary or *_js_binary rules -closure_js_proto_library( +kv_api = declare_roma_api( + cc_protos = [":binary_get_values_cc_proto"], + proto_basename = "binary_get_values", + protos = [":binary_get_values_proto"], +) + +# Library that can be added as a dep from closure_js_library or *_js_binary rules +js_proto_library( name = "binary_get_values_js_proto", - srcs = ["binary_get_values.proto"], - import_style = "IMPORT_COMMONJS", - protocbin = "@com_google_protobuf_for_closure//:protoc", + roma_api = kv_api, ) cc_library( diff --git a/public/udf/constants.h b/public/udf/constants.h index 8cf6e6dc..8a7215dc 100644 --- a/public/udf/constants.h +++ b/public/udf/constants.h @@ -50,10 +50,35 @@ function getKeyGroupOutputs(udf_arguments) { return keyGroupOutputs; } +function handlePas(udf_arguments) { + if (udf_arguments.length != 1) { + const error_message = + 'For PAS default UDF exactly one argument should be provided, but was provided ' + udf_arguments.length; + console.error(error_message); + throw new Error(error_message); + } + const kv_result = JSON.parse(getValues(udf_arguments[0])); + if (kv_result.hasOwnProperty("kvPairs")) { + return kv_result.kvPairs; + } + const error_message = "Error executing handle PAS:" + + JSON.stringify(kv_result); + console.error(error_message); + throw new Error(error_message); +} -function HandleRequest(executionMetadata, ...udf_arguments) { +function handlePA(udf_arguments) { const keyGroupOutputs = getKeyGroupOutputs(udf_arguments); - return {keyGroupOutputs, udfOutputApiVersion: 1}; + return { keyGroupOutputs, udfOutputApiVersion: 1 }; +} + +function HandleRequest(executionMetadata, ...udf_arguments) { + if(executionMetadata.requestMetadata && + executionMetadata.requestMetadata.is_pas) { + console.log('Executing PAS branch'); + return handlePas(udf_arguments); + } + return handlePA(udf_arguments); } )"; diff --git a/third_party_deps/latency_benchmark_requirements.txt b/third_party_deps/latency_benchmark_requirements.txt new file mode 100644 index 00000000..aa1d793b --- /dev/null +++ b/third_party_deps/latency_benchmark_requirements.txt @@ -0,0 +1,6 @@ +numpy==1.25.1 +python-dateutil==2.8.2 +pytz==2020.1 +pandas==2.2.0 +tzdata==2022.7 +six diff --git a/third_party_deps/python_deps.bzl b/third_party_deps/python_deps.bzl index 6573d202..130255d3 100644 --- a/third_party_deps/python_deps.bzl +++ b/third_party_deps/python_deps.bzl @@ -19,3 +19,8 @@ def python_repositories(): name = "word2vec", requirements_lock = "//third_party_deps:word2vec_requirements.txt", ) + + pip_parse( + name = "latency_benchmark", + requirements_lock = "//third_party_deps:latency_benchmark_requirements.txt", + ) diff --git a/third_party_deps/rules_closure_repositories.bzl b/third_party_deps/rules_closure_repositories.bzl deleted file mode 100644 index 8c26eab1..00000000 --- a/third_party_deps/rules_closure_repositories.bzl +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright 2023 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") - -def rules_closure_repositories(): - # This is currently needed for closure_js_proto_library, - # since the server is using protobuf 3.23*., which no longer has javascript support. - # rules_closure uses this version of protoc. - # To make `closure_js_proto_library` work, we need to pass - # closure_js_proto_library( - # ... - # protocbin = "@com_google_protobuf_for_closure//:protoc" - # ) - http_archive( - name = "com_google_protobuf_for_closure", - strip_prefix = "protobuf-3.19.1", - sha256 = "87407cd28e7a9c95d9f61a098a53cf031109d451a7763e7dd1253abf8b4df422", - urls = [ - "https://github.com/protocolbuffers/protobuf/archive/v3.19.1.tar.gz", - ], - ) diff --git a/tools/bidding_auction_data_generator/BUILD.bazel b/tools/bidding_auction_data_generator/BUILD.bazel index 806916fe..41a29b1a 100644 --- a/tools/bidding_auction_data_generator/BUILD.bazel +++ b/tools/bidding_auction_data_generator/BUILD.bazel @@ -63,8 +63,8 @@ cc_library( hdrs = ["value_fetch_util.h"], deps = [ "//public/query:get_values_cc_grpc", - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_protobuf//:protobuf", "@curl", @@ -76,7 +76,7 @@ cc_library( srcs = ["http_url_fetch_client.cc"], hdrs = ["http_url_fetch_client.h"], deps = [ - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@curl", @@ -91,9 +91,9 @@ cc_library( ":http_url_fetch_client", ":value_fetch_util", "//public/query:get_values_cc_grpc", - "@com_github_google_glog//:glog", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_protobuf//:protobuf", @@ -124,7 +124,6 @@ cc_test( "//public/data_loading/readers:riegeli_stream_record_reader_factory", "//public/data_loading/writers:delta_record_stream_writer", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", ], ) diff --git a/tools/bidding_auction_data_generator/bidding_auction_data_cli.cc b/tools/bidding_auction_data_generator/bidding_auction_data_cli.cc index 420f9eba..eb582a3f 100644 --- a/tools/bidding_auction_data_generator/bidding_auction_data_cli.cc +++ b/tools/bidding_auction_data_generator/bidding_auction_data_cli.cc @@ -20,7 +20,7 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" #include "absl/flags/usage.h" -#include "glog/logging.h" +#include "absl/log/log.h" #include "public/data_loading/filename_utils.h" #include "custom_audience_data_parser.h" diff --git a/tools/bidding_auction_data_generator/delta_key_value_writer_test.cc b/tools/bidding_auction_data_generator/delta_key_value_writer_test.cc index e38ba313..0183218d 100644 --- a/tools/bidding_auction_data_generator/delta_key_value_writer_test.cc +++ b/tools/bidding_auction_data_generator/delta_key_value_writer_test.cc @@ -21,7 +21,6 @@ #include "gtest/gtest.h" #include "public/data_loading/readers/delta_record_stream_reader.h" #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" -#include "src/cpp/telemetry/mocks.h" namespace kv_server { namespace { diff --git a/tools/bidding_auction_data_generator/http_url_fetch_client.cc b/tools/bidding_auction_data_generator/http_url_fetch_client.cc index 6e88aa05..1c09e6a4 100644 --- a/tools/bidding_auction_data_generator/http_url_fetch_client.cc +++ b/tools/bidding_auction_data_generator/http_url_fetch_client.cc @@ -19,9 +19,9 @@ #include #include +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "curl/multi.h" -#include "glog/logging.h" namespace kv_server { namespace { diff --git a/tools/bidding_auction_data_generator/http_value_retriever.cc b/tools/bidding_auction_data_generator/http_value_retriever.cc index 3c4c3c2f..6607d62b 100644 --- a/tools/bidding_auction_data_generator/http_value_retriever.cc +++ b/tools/bidding_auction_data_generator/http_value_retriever.cc @@ -24,9 +24,9 @@ #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "glog/logging.h" #include "google/protobuf/util/json_util.h" #include "public/query/get_values.pb.h" #include "tools/bidding_auction_data_generator/value_fetch_util.h" diff --git a/tools/data_cli/BUILD.bazel b/tools/data_cli/BUILD.bazel index 2b68fcb2..8c5bdd40 100644 --- a/tools/data_cli/BUILD.bazel +++ b/tools/data_cli/BUILD.bazel @@ -15,6 +15,7 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = [ + "//docs/protected_app_signals:__subpackages__", "//getting_started:__subpackages__", "//production/packaging/tools:__subpackages__", "//testing:__subpackages__", @@ -33,5 +34,7 @@ cc_binary( "//tools/data_cli/commands:generate_snapshot_command", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", ], ) diff --git a/tools/data_cli/commands/BUILD.bazel b/tools/data_cli/commands/BUILD.bazel index 09d1490a..774bb3f0 100644 --- a/tools/data_cli/commands/BUILD.bazel +++ b/tools/data_cli/commands/BUILD.bazel @@ -45,12 +45,12 @@ cc_library( "//public/data_loading/writers:delta_record_stream_writer", "//public/data_loading/writers:delta_record_writer", "//public/sharding:sharding_function", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) diff --git a/tools/data_cli/commands/format_data_command.cc b/tools/data_cli/commands/format_data_command.cc index 47b05de9..6ea65fe4 100644 --- a/tools/data_cli/commands/format_data_command.cc +++ b/tools/data_cli/commands/format_data_command.cc @@ -31,7 +31,7 @@ #include "public/data_loading/writers/avro_delta_record_stream_writer.h" #include "public/data_loading/writers/delta_record_stream_writer.h" #include "public/sharding/sharding_function.h" -#include "src/cpp/util/status_macro/status_macros.h" +#include "src/util/status_macro/status_macros.h" namespace kv_server { namespace { diff --git a/tools/data_cli/commands/generate_snapshot_command.cc b/tools/data_cli/commands/generate_snapshot_command.cc index ae4a4bfa..7072a9a4 100644 --- a/tools/data_cli/commands/generate_snapshot_command.cc +++ b/tools/data_cli/commands/generate_snapshot_command.cc @@ -36,7 +36,7 @@ #include "public/data_loading/readers/delta_record_stream_reader.h" #include "public/data_loading/riegeli_metadata.pb.h" #include "public/sharding/sharding_function.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/telemetry_provider.h" namespace kv_server { namespace { diff --git a/tools/data_cli/data_cli.cc b/tools/data_cli/data_cli.cc index f6c79fcc..f7ec1c91 100644 --- a/tools/data_cli/data_cli.cc +++ b/tools/data_cli/data_cli.cc @@ -19,8 +19,10 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" #include "absl/flags/usage.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "components/util/platform_initializer.h" -#include "glog/logging.h" #include "tools/data_cli/commands/command.h" #include "tools/data_cli/commands/format_data_command.h" #include "tools/data_cli/commands/generate_snapshot_command.h" @@ -134,18 +136,19 @@ bool IsSupportedCommand(std::string_view command) { // Sample run using bazel: // -// GLOG_logtostderr=1 GLOG_v=3 bazel run \ +// bazel run \ // //tools/data_cli:data_cli \ // --//:instance=local --//:platform=local -- \ // format_data \ // --input_file=/data/DELTA_1689344645643610 \ // --input_format=DELTA \ // --output_format=CSV \ -// --output_file=/data/DELTA_1689344645643610.csv +// --output_file=/data/DELTA_1689344645643610.csv \ +// --v=3 --stderrthreshold=0 int main(int argc, char** argv) { kv_server::PlatformInitializer initializer; - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); absl::SetProgramUsageMessage(kUsageMessage); const std::vector commands = absl::ParseCommandLine(argc, argv); if (commands.size() < 2) { diff --git a/tools/latency_benchmarking/BUILD.bazel b/tools/latency_benchmarking/BUILD.bazel new file mode 100644 index 00000000..f3025200 --- /dev/null +++ b/tools/latency_benchmarking/BUILD.bazel @@ -0,0 +1,23 @@ +load("@rules_python//python:defs.bzl", "py_binary") + +py_binary( + name = "generate_requests", + srcs = ["generate_requests.py"], + tags = ["manual"], +) + +py_binary( + name = "create_csv_summary", + srcs = ["create_csv_summary.py"], + deps = [ + "@latency_benchmark_pandas//:pkg", + ], +) + +py_binary( + name = "merge_csvs", + srcs = ["merge_csvs.py"], + deps = [ + "@latency_benchmark_pandas//:pkg", + ], +) diff --git a/tools/latency_benchmarking/README.md b/tools/latency_benchmarking/README.md new file mode 100644 index 00000000..6372288d --- /dev/null +++ b/tools/latency_benchmarking/README.md @@ -0,0 +1,405 @@ +# Latency Benchmarking Tool + +## Requirements + +- Install ghz + +The `run_benchmarks` script uses [ghz](https://ghz.sh/docs/intro). + +Follow instructions to [install ghz](https://ghz.sh/docs/install) and make sure it is installed +correctly: + +```sh +ghz -v +``` + +- Follow deployment guides + +To use the tool you will need to have a complete deployment setup. + +Please read through the deployment guide for the relevant cloud provider +([AWS](/docs/deployment/deploying_on_aws.md) or [GCP](/docs/deployment/deploying_on_gcp.md)). + +At the minimum, you will need to go through the steps up until the `terraform init` command. +Ideally, follow the entire guide and make sure your deployment setup works. + +## Tools + +Both the `run_benchmarks` and `deploy_and_benchmark` scripts are meant as a convenience to run +benchmarks against the server with different parameters/configurations. Instead of running multiple +iterations of benchmarks manually, the scripts take certain input parameters, run benchmarks, and +summarize the result. + +> Note: The script does not inspect the server's response. Should the server time out or not +> respond, the script currently does not attempt to rerun the benchmarks, it will simply continue. +> To find out which benchmarks need to be rerun, inspect the summary csvs. + +### run_benchmarks + +This script assumes that a server has already been deployed. + +It creates one or multiple requests, runs ghz for each request, and outputs the summary in a CSV +file. + +The request body is designed to work with the server's [default UDF](/public/udf/constants.h). To +ensure the script works, any custom UDF will need to be able to handle the same inputs (outputs are +not affected). + +```js +{"metadata": {...}, "partitions": [{"id": 0, "compressionGroupId": 0, "arguments": [{"tags": ["custom", "keys"], "data": }]}]} +``` + +The tool takes a directory of snapshot files or delta files with key value mutations of type +`Update`. + +For each snapshot file in `snapshot-dir`, the script reads its keys and may generate multiple +requests, one for each element in `number-of-lookup-keys-list`. The number of lookup keys element +indicates how many keys from the snapshot file should be included in a request. + +For example, given a `snapshot-dir` with 2 snapshot files and a `number-of-lookup-keys-list` with 3 +elements, the script will generate 6 requests. + +Each benchmark iteration will about 40 seconds to execute. In total, the script would take about 4-5 +minutes to complete the benchmark + 1-2 minutes for pre- and post processing. + +#### Reference + +Usage: + +```sh +./tools/latency_benchmarking/run_benchmarks +``` + +Flags: + +- `--server-address` + + Required. gRPC host and port. + + Example: `--server-address my-server:8443` + +- `--snapshot-dir` + + Required if `--snapshot-csv-dir` is not provided. Full path to a directory of snapshot files. + These snapshot files are converted to CSVs using the + [`data cli`](/docs/data_loading/loading_data.md). The keys from the snapshot file are used to + create requests. + +- `--snapshot-csv-dir` + + Required if `--snapshot-dir` is not provided. Full path to a directory of only snapshot CSV + files (i.e. converted using the [`data cli`](/docs/data_loading/loading_data.md)). This avoids + doing the conversion in the tool, which saves some time for multiple runs on the same snapshot + files. + +- `--number-of-lookup-keys-list` (Optional) + + A list of number of keys (in quotes, as a string) to include in a request. The tool will iterate + through snapshot csvs and the `--number-of-lookup-keys-list` to construct a request for each + combination. Default is `"1 5 10"`. + + Example: `--number-of-lookup-keys-list "1 10"` + +- `--benchmark-duration` (Optional) + + How long each benchmark iteration should run for. Default is `"5s"`. + +- `--ghz-tags` (Optional) + + Tags to include for all ghz runs. The csv summary will include this tag for all iterations. + + Example: `-- ghz-tags '{"my_tag":"my_value"}'` + +- `--request-metadata-json` (Optional) + + The request metadata json object to use for all requests. Default is empty `{}`. + + Example: `--request-metadata-json '{"metadata_key":"metadata_value"}'` + +- `--filter-snapshot-by-sets` (Optional) + + Whether to filter snapshot csvs using `value_type=string_set` only. This will create requests + that only include keys of sets. Default is false. + +#### Example + +Start from the workspace root. + +```sh +SNAPSHOT_DIR=/path/to/snapshot/dir +NUMBER_OF_LOOKUP_KEYS_LIST="1 10 100" +SERVER_ADDRESS="demo.kv-server.your-domain.example:8443" +./tools/latency_benchmarking/run_benchmarks \ +--server-address ${SERVER_ADDRESS} \ +--snapshot-dir ${SNAPSHOT_DIR} \ +--number-of-lookup-keys-list "${NUMBER_OF_LOOKUP_KEYS_LIST}" +``` + +The summary can be found in +`dist/tools/latency_benchmarking/output//summary.csv`. + +### deploy_and_benchmark + +This script will deploy a terraform configuration and call the `run_benchmarks` script. + +The tool will _not_ upload key value delta/snapshot files. Please ensure those are already loaded +into the blob storage bucket that the server will be reading from. + +It includes almost all of the `run_benchmarks` flags in addition to more deployment parameters to +iterate through. + +- If provided with a terraform overrides file with sets of terraform variable overrides, the + script will deploy once per set of variables and run benchmarks. +- If given a directory with UDF deltas, it will upload each UDF to the data bucket, and run + benchmarks against it. +- If given a json lines file with request metadata, it will iterate through each request metadata + and run benchmarks against it. + +> The number of iterations can increase dramatically with each added parameter. For example, given 3 +> sets of terraform overrides, 3 UDF deltas, and 3 request metadata jsons, `run_benchmarks` will +> execute 27 times. If in addition to that, the `number-of-lookup-keys` has 3 elements and +> `snapshot-dir` has 3 files, the script will be running 243 iterations of ghz, once for each +> combination of parameters. + +#### Reference + +Usage: + +```sh +./tools/latency_benchmarking/deploy_and_benchmark +``` + +Flags: + +- `--cloud-provider` + + Required. Cloud provider. Options: "aws" or "gcp" + +- `--tf-var-file` + + Required. Full path to `tfvars.json` file. + +- `--tf-backend-config` + + Required. Full path to `backend.conf` file. + +- `--server-url` + + Required. URL to deployed server. + + Example: "" + +- `--snapshot-dir` + + Required. Full path to a directory of snapshot files. The keys from the snapshot file are used + to create requests. + +- `--csv-output` (Optional) + + Full path to csv output. Contains a summary of all ghz runs. + +- `--tf-overrides` (Optional) + + File with terraform variable overrides. Each line in the file counts as a set of variable + overrides. If provided, the tool will deploy once per line. The expected format for each line is + a comma-separated list of terraform variable overrides, e.g. + `instance_type=c5.4xlarge,enclave_cpu_count=12` + + Example: `/tools/latency_benchmarking/example/aws_tf_overrides.txt` + +- `--udf-delta-dir` (Optional) + + Full path to directory of udf delta files to use for benchmarking. The tool will upload each + file in the directory to the given `data-bucket` and run benchmarks after. + + Note that all udf deltas are uploaded once per deployment and removed by the end of each + deployment. Since the server does not restart in between udf uploads, it will only pick up udf + deltas if the logical_commit_times and versions are set correctly. + + See the [udf delta documentation](/docs/generating_udf_files.md) for more info. + +- `--data-bucket` (Optional) + + The data bucket to upload the udf files too. This should be the same data bucket that the server + is reading delta files from. + +- `--request-metadata-json-file` (Optional) + + Path to JSON lines file to iterate through to generate requests. + + Example: `/tools/latency_benchmarking/example/request_metadata.jsonl` + +- `--number-of-lookup-keys-list` (Optional) + + A list of number of keys (in quotes, as a string) to include in a request. The tool will iterate + through snapshot csvs and the `--number-of-lookup-keys-list` to construct a request for each + combination. Default is `"1 10 50 100"`. + + Example: `--number-of-lookup-keys-list "1 10"` + +- `--minimum-server-wait-secs` (Optional) + + The amount of time that the tool should wait after the `terraform apply` command before checking + whether the server is healthy. + + If this is too low, the tool may still be pinging the previous deployment's server, since it may + take a while to tear down the instance. + + Default is `90`. + +- `--extra-server-wait-timeout` (Optional) + + The amount of time that the tool should continue checking server health on top of the + `minimum-server-wait-secs`. + + The tool will continue pinging the server until the timeout is reached or the server is healthy. + + If the delta files + + Default is `5m`. + + Examples: `5m`, `600s` + +- `--benchmark-duration` (Optional) + + How long each benchmark iteration should run for. Default is `"5s"`. + +- `--ghz-tags` (Optional) + + Tags to include for all ghz runs. The csv summary will include this tag for all iterations. + + Example: `-- ghz-tags '{"my_tag":"my_value"}'` + +- `--cleanup-deployment` (Optional) + + Whether to call terraform destroy once the tool exits. Default false. + +- `--filter-snapshot-by-sets` (Optional) + + Whether to filter snapshot csvs using `value_type=string_set` only. This will create requests + that only include keys of sets. Default is false. + +#### Example + +Start from the workspace root. + +1. Generate SNAPSHOT/DELTA files to upload to data storage: + + ```sh + ./tools/serving_data_generator/generate_test_riegeli_data + GENERATED_DELTA=/path/to/delta/dir/GENERATED_DELTA_FILE + ``` + + For AWS: + + ```sh + aws s3 cp $GENERATED_DELTA s3://bucket_name + ``` + + For GCP: + + ```sh + gcloud storage cp $GENERATED_DELTA gs://bucket_name + ``` + +1. Set your cloud provider + + AWS: + + ```sh + CLOUD_PROVIDER="aws" + ``` + + GCP: + + ```sh + CLOUD_PROVIDER="gcp" + ``` + +1. Provide a directory of SNAPSHOT (or DELTA) files with keys that should be sent in the request to + the server. This SNAPSHOT (or DELTA) file should only have `UPDATE` mutations. The tool will + iterate through each SNAPSHOT file, select keys from that file, and run benchmarks with those + keys. + + For this example, we'll use the generated DELTA files from step 1 + + ```sh + SNAPSHOT_DIR=/path/to/delta/dir/ + ``` + +1. Set up your terraform config files + + ```sh + TF_VAR_FILE=/path/to/my.tfvars.json + TF_BACKEND_CONFIG=/path/to/my.backend.conf + ``` + +1. Set the server url + + ```sh + SERVER_URL="https://demo.kv-server.your-domain.example/" + ``` + +1. (optional) Provide a file with sets of terraform variables to be overriden. For each set of + terraform variables, the script will `terraform apply` once with the given variables and run + benchmarks against the deployed server. The variable override file should have the following + format: + + - Each line should be in the form + + ```txt + variable_name1=variable_valueA,variable_name2=variable_valueB + ``` + + - Each line is considered a set of variables to be overriden in one `terraform apply` command + + - For an example, see `/latency_benchmarking/example/aws_tf_overrides.txt`. + + ```sh + TF_OVERRIDES=/path/to/tf_variable_overrides.txt + ``` + +1. (optional) Write UDFs If given a directory of UDF delta files, the tool iterates through each + one, uploads it to the given data bucket, runs benchmarks, then removes the UDF delta from the + data bucket. + + ```sh + UDF_DELTA_DIR=/path/to/udf_deltas/ + DATA_BUCKET=s3://bucket_name + ``` + +1. Run the script and wait for the result + + ```sh + ./tools/latency_benchmarking/deploy_and_benchmark \ + --cloud-provider ${CLOUD_PROVIDER} \ + --server-url ${SERVER_URL} \ + --snapshot-dir ${SNAPSHOT_DIR} \ + --tf-var-file ${TF_VAR_FILE} \ + --tf-backend-config ${TF_BACKEND_CONFIG} \ + --tf-overrides ${TF_OVERRIDES} \ + --csv-output ${PWD}/my_summary.csv + ``` + + The result will be in `my_summary.csv`. + + ```sh + ls my_summary.csv + ``` + +## Appendix + +### Things of note when deploying + +- Make sure the DELTA/SNAPSHOT files follow the expected logic for time stamps and naming, + otherwise the data may not load properly. See the + [data format specification](/docs/data_loading/data_format_specification.md) and + [udf delta doc](/docs/generating_udf_files.md) for more info. + +- Make sure the chosen instance sizes have enough memory to hold your data. + +- On AWS, some larger instances have multiple NUMA clusters. However, an enclave can only run on + one NUMA cluster. You may want to consider [sharding](/docs/sharding/sharding.md). + +- Depending on the number of iterations the script runs, it may take a long time to finish. Each + benchmark (ghz) iteration takes 40-60s + a few minutes of overhead per deployment. diff --git a/tools/latency_benchmarking/create_csv_summary.py b/tools/latency_benchmarking/create_csv_summary.py new file mode 100644 index 00000000..b5522bc0 --- /dev/null +++ b/tools/latency_benchmarking/create_csv_summary.py @@ -0,0 +1,78 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import argparse +import json +import pandas as pd + +""" +Iterates through ghz_output.json files and collects relevant data in a CSV file. +""" + + +def _ExtractGhzInfo(ghz_output_df: pd.DataFrame) -> pd.DataFrame: + keys = ["date", "count", "total", "average", "fastest", "slowest", "rps"] + prefixes = ("tags", "statusCodeDistribution") + return ghz_output_df.loc[ + :, + ghz_output_df.columns.str.startswith(prefixes) + | ghz_output_df.columns.isin(keys), + ] + + +def JsonToDataFrame(ghz_result_dir: str) -> pd.DataFrame: + """Reads a ghz_output.json files and outputs a dataframe with relevant columns. + + The dataframe will contain: + - metadata: date, user tags added to the ghz command + - overall stats: # requests, total time spent, rps, fastest, slowest + - status code distributions + + Example: + date count total average fastest slowest rps statusCodeDistribution.OK statusCodeDistribution.Unavailable + 0 2024-01-30T21:27:06Z 13130 5000416406 37281592 6313913 100341707 2625.781322 13021 109 + 0 2024-01-30T21:27:06Z 13130 5000416406 37281592 6313913 100341707 2625.781322 13021 109 + 0 2024-01-30T21:27:06Z 13130 5000416406 37281592 6313913 100341707 2625.781322 13021 109 + """ + output_dfs = [] + for root, dirs, files in os.walk(ghz_result_dir): + for name in files: + if name == "ghz_output.json": + fp = os.path.join(root, name) + try: + with open(fp, "r") as f: + json_data = json.loads(f.read()) + ghz_output_df = pd.json_normalize(json_data) + output_dfs.append(_ExtractGhzInfo(ghz_output_df)) + except FileNotFoundError: + print(f"File not found: {fp}") + return pd.concat(output_dfs) + + +def Main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--ghz-result-dir", + dest="ghz_result_dir", + type=str, + help="Directory containing subdirectories with ghz.json output files.", + ) + args = parser.parse_args() + output_df = JsonToDataFrame(args.ghz_result_dir) + output_df.to_csv(os.path.join(args.ghz_result_dir, "summary.csv"), index=False) + + +if __name__ == "__main__": + Main() diff --git a/tools/latency_benchmarking/deploy_and_benchmark b/tools/latency_benchmarking/deploy_and_benchmark new file mode 100755 index 00000000..730eacf8 --- /dev/null +++ b/tools/latency_benchmarking/deploy_and_benchmark @@ -0,0 +1,393 @@ +#!/bin/bash +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -o pipefail +set -o errexit + +# shellcheck disable=SC1090 +source ./builders/tools/builder.sh + +START=$(date +%s) +WORKSPACE="$(git rev-parse --show-toplevel)" + +# Script input params +EXTRA_SERVER_WAIT_TIMEOUT="5m" +MINIMUM_SERVER_WAIT_SECS=90 +NUMBER_OF_LOOKUP_KEYS_LIST="1 10 50 100" +GHZ_TAGS="{}" +FILTER_BY_SETS=0 + +# Additional GHZ tags per deployment/udf +DEPLOYMENT_GHZ_TAGS="{}" +UDF_GHZ_TAGS="{}" + +# Input to run_benchmark script +declare RUN_BENCHMARK_ARGS +BENCHMARK_DURATION="5s" + +# Deployment vars +declare SERVER_ADDRESS +declare SERVER_ENDPOINT + +# CSV I/O +declare -a BENCHMARK_CSVS +CSV_OUTPUT="${WORKSPACE}/dist/tools/latency_benchmarking/output/output.csv" +declare -r DOCKER_OUTPUT_CSV="/tmp/latency_benchmarking/output/deploy_and_benchmark/output.csv" +declare -r SNAPSHOT_CSV_DIR="${WORKSPACE}/dist/tools/latency_benchmarking/deploy_and_benchmark/snapshot_csvs/${START}" +declare -r DOCKER_SNAPSHOT_CSV_DIR="/tmp/latency_benchmarking/deploy_and_benchmark/snapshot_csvs/" +declare -r DOCKER_SNAPSHOT_DIR="/tmp/latency_benchmarking/deploy_and_benchmark/snapshots/" + +# run_benchmarks writes output to this directory +CSV_SUMMARY_INPUT_DIR="${WORKSPACE}/dist/tools/latency_benchmarking/output" +DOCKER_INPUT_DIR="/tmp/latency_benchmarking/output" +readonly DOCKER_INPUT_DIR + +DESTROY_INSTANCES=0 +TF_DEPLOY_SUCCESS=0 + +trap _destroy EXIT +function _destroy() { + if [[ ${DESTROY_INSTANCES} == 1 && ${TF_DEPLOY_SUCCESS} == 1 ]]; then + printf "Running terraform destroy\n" + builders/tools/terraform \ + -chdir="${WORKSPACE}"/production/terraform/"${CLOUD_PROVIDER}"/environments \ + destroy --var-file="${TF_VAR_FILE}" --auto-approve >/dev/null + fi +} + +trap _trap ERR +function _trap() { + local -r -i STATUS=$? + FAILED_COMMAND="${BASH_COMMAND}" + printf "Failed command: %s\n" "${FAILED_COMMAND}" + exit ${STATUS} +} + +function usage() { + local -r -i exitval=${1-1} + cat &>/dev/stderr < + [--cloud-provider] (Required) Cloud provider. Options: "aws" or "gcp" + [--tf-var-file] (Required) Full path to tfvars.json file. + [--tf-backend-config] (Required) Full path to tf backend.conf file. + [--server-url] (Required) URL of deployed server. + [--snapshot-dir] (Required) Full path to a directory of snapshot files. + [--csv-output] (Optional) Path to output file for summary of benchmarks. + [--tf-overrides] (Optional) Path to file with terraform variable overrides. + [--udf-delta-dir] (Optional) Full path to directory of udf delta files. + [--data-bucket] (Optional) Data bucket to upload the udf files too. + [--request-metadata-json-file] (Optional) Path to JSON lines file with request metadata. + [--number-of-lookup-keys-list] (Optional) List of number of keys to include in a request. + [--minimum-server-wait-secs] (Optional) Amount of time to wait before checking server health. + [--extra-server-wait-timeout] (Optional) Timeout for waiting for server to become healthy. + [--benchmark-duration] (Optional) Duration of each benchmark. Default "5s". + [--ghz-tags] (Optional) Tags to include in the ghz run. + [--cleanup-deployment] (Optional) Whether to call terraform destroy once the tool exits. + [--filter-snapshot-by-sets] (Optional) Whether to filter snapshots by sets when creating requests. +USAGE + # shellcheck disable=SC2086 + exit ${exitval} +} + +function convert_snapshots_to_csvs() { + mkdir -p "${SNAPSHOT_CSV_DIR}" + # Iterate through snapshot files and convert them to CSV + for SNAPSHOT_FILE in "${SNAPSHOT_DIR}"/*; do + SNAPSHOT_FILENAME=$(basename "${SNAPSHOT_FILE}") + EXTRA_DOCKER_RUN_ARGS+=" --volume ${SNAPSHOT_CSV_DIR}:${DOCKER_SNAPSHOT_CSV_DIR} --volume ${SNAPSHOT_DIR}:${DOCKER_SNAPSHOT_DIR} " \ + builders/tools/bazel-debian run //tools/data_cli:data_cli format_data \ + -- \ + --input_file "${DOCKER_SNAPSHOT_DIR}/${SNAPSHOT_FILENAME}" \ + --input_format DELTA \ + --output_file "${DOCKER_SNAPSHOT_CSV_DIR}/${SNAPSHOT_FILENAME}.csv" \ + --output_format CSV + done +} + +function set_benchmark_args() { + local -a RUN_BENCHMARK_GHZ_TAGS + RUN_BENCHMARK_GHZ_TAGS=$(jq -s -c 'add' <(echo "${GHZ_TAGS}") \ + <(echo "${DEPLOYMENT_GHZ_TAGS}") \ + <(echo "${UDF_GHZ_TAGS}")) + RUN_BENCHMARK_ARGS=( + --number-of-lookup-keys-list "${NUMBER_OF_LOOKUP_KEYS_LIST[@]}" + --server-address "${SERVER_ADDRESS}" + --ghz-tags "${RUN_BENCHMARK_GHZ_TAGS}" + --snapshot-csv-dir "${SNAPSHOT_CSV_DIR}" + --benchmark-duration "${BENCHMARK_DURATION}" + ) + if [[ -v "${REQUEST_METADATA_JSON}" ]]; then + RUN_BENCHMARK_ARGS+=( + --request-metadata-json "${REQUEST_METADATA_JSON}" + ) + fi + if [[ "${FILTER_BY_SETS}" = 1 ]]; then + RUN_BENCHMARK_ARGS+=( + --filter-snapshot-by-sets + ) + fi +} + +function run_benchmarks() { + if [[ -z "${REQUEST_METADATA_JSON_FILE}" ]]; then + set_benchmark_args + printf "BENCHMARK ARGS: %s\n" "${RUN_BENCHMARK_ARGS[*]}" + + local BENCHMARK_OUTPUT + BENCHMARK_OUTPUT=$(./tools/latency_benchmarking/run_benchmarks "${RUN_BENCHMARK_ARGS[@]}") + BENCHMARK_CSVS+=( + "$(echo "${BENCHMARK_OUTPUT}" | tail -n 1 2>&1 | tee /dev/tty)" + ) + else + while IFS= read -r REQUEST_METADATA_JSON; do + set_benchmark_args + printf "BENCHMARK ARGS: %s\n" "${RUN_BENCHMARK_ARGS[*]}" + + local BENCHMARK_OUTPUT + BENCHMARK_OUTPUT=$(./tools/latency_benchmarking/run_benchmarks "${RUN_BENCHMARK_ARGS[@]}") + BENCHMARK_CSVS+=( + "$(echo "${BENCHMARK_OUTPUT}" | tail -n 1 2>&1 | tee /dev/tty)" + ) + done < "${REQUEST_METADATA_JSON_FILE}" + fi +} + +function set_server_address() { + # Build HTTP server endpoint from tf output + local SERVER_HOSTNAME + SERVER_ENDPOINT="${SERVER_URL}/v1/getvalues?keys=hi" + # Build gRPC server address from tf output + SERVER_HOSTNAME=$([[ "${SERVER_URL}" =~ https://(.*) ]] && echo "${BASH_REMATCH[1]}") + SERVER_ADDRESS="${SERVER_HOSTNAME}:8443" +} + +function upload_file_to_bucket() { + if [[ "${CLOUD_PROVIDER}" == "aws" ]]; then + EXTRA_DOCKER_RUN_ARGS+=" --volume ${1}:/tmp/deltas/${1}" \ + builders/tools/aws-cli s3 cp "/tmp/deltas/${1}" "${2}" + elif [[ "${CLOUD_PROVIDER}" == "gcp" ]]; then + gcloud storage cp "${1}" "${2}" + else + echo "Cloud provider not supported" + exit 1 + fi +} + +function remove_file_from_bucket() { + if [[ "${CLOUD_PROVIDER}" == "aws" ]]; then + builders/tools/aws-cli s3 rm "${1}" + elif [[ "${CLOUD_PROVIDER}" == "gcp" ]]; then + gcloud storage rm "${1}" + else + echo "Cloud provider not supported" + exit 1 + fi +} + +function deploy_and_benchmark() { + printf "Running terraform init\n" + builders/tools/terraform \ + -chdir=production/terraform/"${CLOUD_PROVIDER}"/environments \ + init --backend-config="${TF_BACKEND_CONFIG}" \ + --var-file="${TF_VAR_FILE}" --reconfigure -input=false \ + >/dev/null + + printf "Running terraform apply with var file: %s\n" "${TF_VAR_FILE}" + printf "and var overrides: %s\n" "${VAR_OVERRIDES[*]}" + builders/tools/terraform \ + -chdir=production/terraform/"${CLOUD_PROVIDER}"/environments \ + apply --var-file="${TF_VAR_FILE}" "${VAR_OVERRIDES[@]}" \ + -auto-approve + TF_DEPLOY_SUCCESS=1 + printf "Done applying terraform, waiting for server to be ready\n" + + set_server_address + + # Wait for potential instance teardown before periodically checking if server is ready + sleep "${MINIMUM_SERVER_WAIT_SECS}" + + # Wait for server to be up + timeout --foreground "${EXTRA_SERVER_WAIT_TIMEOUT}" bash -c \ + "until curl --output /dev/null --silent --fail ${SERVER_ENDPOINT};do sleep 15; done" + printf "Server ready\n" + + # Iterate through each given UDF + # Upload delta file, run benchmarks, then remove it + if [[ -d "${UDF_DELTA_DIR}" ]]; then + for FILE in "${UDF_DELTA_DIR}"/*; do + local FILENAME + FILENAME=$(basename "${FILE}") + printf "Uploading UDF file %s to cloud storage %s\n" "${FILE}" "${DATA_BUCKET}" + upload_file_to_bucket "${FILE}" "${DATA_BUCKET}/${FILENAME}" + # Allow some time for server to pick up new UDF + sleep 60 + UDF_GHZ_TAGS="{ \"udf_delta_file\": \"${FILENAME}\" }" + run_benchmarks + printf "Removing file from cloud storage %s\n" "${DATA_BUCKET}/${FILENAME}" + remove_file_from_bucket "${DATA_BUCKET}/${FILENAME}" + done + else + run_benchmarks + fi +} + +function merge_benchmark_csvs() { + # Map csv summary files from run_benchmarks to docker volume paths + # so that we can access them from within builders/tools/bazel-debian + # i.e. csv_output_dir/summary.csv -> docker_mounted_dir/summary.csv + declare -a DOCKER_BENCHMARK_CSVS + for BENCHMARK_CSV in "${BENCHMARK_CSVS[@]}"; do + DOCKER_BENCHMARK_CSVS+=( + "${DOCKER_INPUT_DIR}${BENCHMARK_CSV#"${CSV_SUMMARY_INPUT_DIR}"}" + ) + done + touch "${CSV_OUTPUT}" + # Run merge_csvs python script + EXTRA_DOCKER_RUN_ARGS+=" --volume ${CSV_SUMMARY_INPUT_DIR}:${DOCKER_INPUT_DIR} --volume ${CSV_OUTPUT}:${DOCKER_OUTPUT_CSV} " \ + builders/tools/bazel-debian run //tools/latency_benchmarking:merge_csvs \ + -- \ + --csv-inputs "${DOCKER_BENCHMARK_CSVS[@]}" \ + --csv-output "${DOCKER_OUTPUT_CSV}" + + printf "Results in: %s\n" "${CSV_OUTPUT}" +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --tf-var-file) + TF_VAR_FILE="$2" + shift 2 || usage + ;; + --tf-backend-config) + TF_BACKEND_CONFIG="$2" + shift 2 || usage + ;; + --tf-overrides) + TF_OVERRIDES="$2" + shift 2 + ;; + --csv-output) + CSV_OUTPUT="$2" + shift 2 + ;; + --snapshot-dir) + SNAPSHOT_DIR="$2" + shift 2 + ;; + --udf-delta-dir) + UDF_DELTA_DIR="$2" + shift 2 + ;; + --data-bucket) + DATA_BUCKET="$2" + shift 2 + ;; + --extra-server-wait-timeout) + EXTRA_SERVER_WAIT_TIMEOUT="$2" + shift 2 + ;; + --number-of-lookup-keys-list) + NUMBER_OF_LOOKUP_KEYS_LIST="$2" + shift 2 + ;; + --benchmark-duration) + BENCHMARK_DURATION="$2" + shift 2 + ;; + --ghz-tags) + GHZ_TAGS="$2" + shift 2 + ;; + --minimum-server-wait-secs) + MINIMUM_SERVER_WAIT_SECS="$2" + shift 2 + ;; + --cloud-provider) + CLOUD_PROVIDER="$2" + shift 2 + ;; + --server-url) + SERVER_URL="$2" + shift 2 + ;; + --cleanup-deployment) + DESTROY_INSTANCES=1 + shift + ;; + --request-metadata-json-file) + REQUEST_METADATA_JSON_FILE="$2" + shift 2 + ;; + --filter-snapshot-by-sets) + FILTER_BY_SETS=1 + shift + ;; + -h | --help) usage 0 ;; + *) usage ;; + esac +done + +if [[ -z "${CLOUD_PROVIDER}" || ! ("${CLOUD_PROVIDER}" == "aws" || "${CLOUD_PROVIDER}" == "gcp") ]]; then + printf "cloud-provider must equal \"aws\" or \"gcp\"" + exit 1 +fi + +if [[ -z "${SERVER_URL}" ]]; then + printf "server-url must be provided" + exit 1 +fi + +# Check for SNAPSHOT_DIR. If not available, exit. +if [[ -z "${SNAPSHOT_DIR}" || ! -d "${SNAPSHOT_DIR}" ]]; then + printf "snapshot-dir not found:%s\n" "${SNAPSHOT_DIR}" + exit 1 +fi + +# Check that files, if provided, are readable +if [[ -v "${TF_OVERRIDES}" && ! -r "${TF_OVERRIDES}" ]]; then + printf "tf-overrides file not readable: %s\n" "${TF_OVERRIDES}" + exit 1 +fi + +if [[ -v "${REQUEST_METADATA_JSON_FILE}" && ! -r "${REQUEST_METADATA_JSON_FILE}" ]]; then + printf "request-metadata-json-file file not readable: %s\n" "${REQUEST_METADATA_JSON_FILE}" + exit 1 +fi + +convert_snapshots_to_csvs + +# No terraform variable overrides, deploy and benchmark without overrides +if [[ -z "${TF_OVERRIDES}" ]]; then + deploy_and_benchmark +else + # Terraform override file found: + # Each row defines a set of overrides for terraform variables. + # Pass the overrides to deploy_and_benchmark function and set + # them as tags in the ghz call. + while IFS=',' read -ra VARS; do + declare -a VAR_OVERRIDES=() + DEPLOYMENT_GHZ_TAGS="{}" + for VAR in "${VARS[@]}"; do + VAR_OVERRIDES+=(-var "${VAR}") + OVERRIDE_VAR_GHZ_TAG=$(echo "${VAR}" | jq -s -R 'split("\n") | .[0] | split("=") | {(.[0]): .[1]}') + DEPLOYMENT_GHZ_TAGS=$(echo "${DEPLOYMENT_GHZ_TAGS} ${OVERRIDE_VAR_GHZ_TAG}" | jq -s -c 'add') + done + deploy_and_benchmark + done <"${TF_OVERRIDES}" +fi + +# Benchmarks done, merge CSVs +merge_benchmark_csvs diff --git a/tools/latency_benchmarking/example/aws_tf_overrides.txt b/tools/latency_benchmarking/example/aws_tf_overrides.txt new file mode 100644 index 00000000..c37884e6 --- /dev/null +++ b/tools/latency_benchmarking/example/aws_tf_overrides.txt @@ -0,0 +1,6 @@ +instance_type=c5.4xlarge,enclave_cpu_count=12,enclave_memory_mib=24576,num_shards=1,udf_num_workers=12,instance_ami_id=ami-0c3ea8d5ff3ef70d2 +instance_type=c5.12xlarge,enclave_cpu_count=36,enclave_memory_mib=73728,num_shards=1,udf_num_workers=36,instance_ami_id=ami-0c3ea8d5ff3ef70d2 +instance_type=m5.4xlarge,enclave_cpu_count=12,enclave_memory_mib=49152,num_shards=1,udf_num_workers=12,instance_ami_id=ami-0c3ea8d5ff3ef70d2 +instance_type=m5.12xlarge,enclave_cpu_count=36,enclave_memory_mib=147456,num_shards=1,udf_num_workers=36,instance_ami_id=ami-0c3ea8d5ff3ef70d2 +instance_type=c5.4xlarge,enclave_cpu_count=12,enclave_memory_mib=24576,num_shards=2,udf_num_workers=12,instance_ami_id=ami-0c3ea8d5ff3ef70d2 +instance_type=m5.4xlarge,enclave_cpu_count=12,enclave_memory_mib=49152,num_shards=2,udf_num_workers=12,instance_ami_id=ami-0c3ea8d5ff3ef70d2 diff --git a/tools/latency_benchmarking/example/gcp_tf_overrides.txt b/tools/latency_benchmarking/example/gcp_tf_overrides.txt new file mode 100644 index 00000000..19c25231 --- /dev/null +++ b/tools/latency_benchmarking/example/gcp_tf_overrides.txt @@ -0,0 +1,2 @@ +machine_type=n2d-standard-4,num_shards=1,udf_num_workers=4 +machine_type=n2d-standard-8,num_shards=1,udf_num_workers=8 diff --git a/tools/latency_benchmarking/example/kv_data/BUILD.bazel b/tools/latency_benchmarking/example/kv_data/BUILD.bazel new file mode 100644 index 00000000..91bf4777 --- /dev/null +++ b/tools/latency_benchmarking/example/kv_data/BUILD.bazel @@ -0,0 +1,52 @@ +load("@bazel_skylib//rules:run_binary.bzl", "run_binary") + +package(default_visibility = [ + "//tools:__subpackages__", +]) + +num_kv_records = 100000 + +num_kv_set_records = 25000 + +num_values_in_set = 5 + +total_records = num_kv_records + num_kv_set_records + +# Rule to generate example KV delta that includes +# num_kv_records key value records and +# num_kv_set_records set records each of size num_values_in_set +[ + run_binary( + name = "generate_{}_delta".format(size), + outs = [ + "DELTA_100000000000000{}".format(num), + ], + args = [ + "--output_dir", + "/src/workspace", + "--output_filename", + "$(location DELTA_100000000000000{})".format(num), + "--key", + "value{}_".format(size), + "--num_records", + "{}".format(num_kv_records), + "--value_size", + "{}".format(size), + "--timestamp", + "{}".format(total_records * num), + "--generate_set_record", + "--set_value_key", + "setvalue{}_".format(size), + "--num_set_records", + "{}".format(num_kv_set_records), + "--num_values_in_set", + "{}".format(num_values_in_set), + ], + tags = ["manual"], + tool = "//tools/serving_data_generator:test_serving_data_generator", + ) + for num, size in enumerate( + ("10", "20", "50", "100", "500", "1000", "5000"), + 1, + ) +] diff --git a/tools/latency_benchmarking/example/request_metadata.jsonl b/tools/latency_benchmarking/example/request_metadata.jsonl new file mode 100644 index 00000000..ca9fcb1b --- /dev/null +++ b/tools/latency_benchmarking/example/request_metadata.jsonl @@ -0,0 +1,4 @@ +{"useGetValuesBinary": false} +{"useGetValuesBinary": false, "lookup_batch_size": 50} +{"useGetValuesBinary": true} +{"useGetValuesBinary": true, "lookup_batch_size": 50} diff --git a/tools/latency_benchmarking/example/request_metadata_run_query.jsonl b/tools/latency_benchmarking/example/request_metadata_run_query.jsonl new file mode 100644 index 00000000..b8966d1d --- /dev/null +++ b/tools/latency_benchmarking/example/request_metadata_run_query.jsonl @@ -0,0 +1,2 @@ +{"runQuery": true, "useGetValuesBinary": false} +{"runQuery": true, "useGetValuesBinary": true} diff --git a/tools/latency_benchmarking/example/udf_code/BUILD.bazel b/tools/latency_benchmarking/example/udf_code/BUILD.bazel new file mode 100644 index 00000000..1745c297 --- /dev/null +++ b/tools/latency_benchmarking/example/udf_code/BUILD.bazel @@ -0,0 +1,44 @@ +load("@io_bazel_rules_closure//closure:defs.bzl", "closure_js_library") +load("//tools/udf/closure_js:closure_to_delta.bzl", "closure_to_delta") +load("//tools/udf/inline_wasm:wasm.bzl", "cc_inline_wasm_udf_delta") + +closure_js_library( + name = "benchmark_udf_js_lib", + srcs = [ + "benchmark_udf.js", # User-defined functions with handler/entrypoint + "externs.js", # Register APIs called from UDF as functions without definitions for closure. + ], + convention = "NONE", + suppress = [ + "missingSourcesWarnings", + ], + deps = [ + "//public/udf:binary_get_values_js_proto", + ], +) + +# Generates a UDF delta file using the given closure_js_lib target +# builders/tools/bazel-debian run \ +# //tools/latency_benchmarking/example/udf_code:benchmark_udf_js_delta +closure_to_delta( + name = "benchmark_udf_js_delta", + closure_js_library_target = ":benchmark_udf_js_lib", + logical_commit_time = "200000001", + output_file_name = "DELTA_2000000000000001", +) + +# builders/tools/bazel-debian run --config=emscripten \ +# //tools/latency_benchmarking/example/udf_code:benchmark_cpp_wasm_udf_delta +cc_inline_wasm_udf_delta( + name = "benchmark_cpp_wasm_udf_delta", + srcs = ["benchmark_cpp_wasm_udf.cc"], + custom_udf_js = "benchmark_cpp_wasm_udf.js", + logical_commit_time = "200000002", + output_file_name = "DELTA_2000000000000002", + tags = ["manual"], + deps = [ + "//public/udf:binary_get_values_cc_proto", + "@com_google_absl//absl/status:statusor", + "@nlohmann_json//:lib", + ], +) diff --git a/tools/latency_benchmarking/example/udf_code/benchmark_cpp_wasm_udf.cc b/tools/latency_benchmarking/example/udf_code/benchmark_cpp_wasm_udf.cc new file mode 100644 index 00000000..b2e278bc --- /dev/null +++ b/tools/latency_benchmarking/example/udf_code/benchmark_cpp_wasm_udf.cc @@ -0,0 +1,290 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#include "absl/status/statusor.h" +#include "emscripten/bind.h" +#include "nlohmann/json.hpp" +#include "public/udf/binary_get_values.pb.h" + +namespace { + +constexpr int kTopK = 4; +constexpr int kLookupNKeysFromRunQuery = 100; + +// Calls getValues for a list of keys and processes the response. +absl::StatusOr> GetKvPairs( + const emscripten::val& get_values_cb, + const std::vector& lookup_data) { + std::vector kv_pairs_json; + for (const auto& keys : lookup_data) { + // `get_values_cb` takes an emscripten::val of type array and returns + // an emscripten::val of type string. + // We need to cast the emscripten::val string to a C++ string. + // However, getValues actually returns a serialized JSON object, + // so we need to parse the string to use it. + const std::string get_values_result = get_values_cb(keys).as(); + const nlohmann::json get_values_json = + nlohmann::json::parse(get_values_result, nullptr, + /*allow_exceptions=*/false, + /*ignore_comments=*/true); + if (get_values_json.is_discarded() || + !get_values_json.contains("kvPairs")) { + return absl::InvalidArgumentError("No kvPairs returned"); + } + kv_pairs_json.push_back(get_values_json); + } + return kv_pairs_json; +} + +absl::StatusOr GetKvPairsAsEmVal( + const emscripten::val& get_values_cb, + const std::vector& lookup_data) { + emscripten::val kv_pairs = emscripten::val::object(); + kv_pairs.set("udfApi", "getValues"); + + const auto kv_pairs_json_or_status = GetKvPairs(get_values_cb, lookup_data); + if (!kv_pairs_json_or_status.ok()) { + return kv_pairs_json_or_status.status(); + } + for (const auto& kv_pair_json : *kv_pairs_json_or_status) { + for (auto& [k, v] : kv_pair_json["kvPairs"].items()) { + if (v.contains("value")) { + emscripten::val value = emscripten::val::object(); + value.set("value", emscripten::val(v["value"].get())); + kv_pairs.set(k, value); + } + } + } + return kv_pairs; +} + +absl::StatusOr> GetKvPairsAsMap( + const emscripten::val& get_values_cb, + const std::vector& lookup_data) { + std::map kv_map; + kv_map["udfApi"] = "getValues"; + + const auto kv_pairs_json_or_status = GetKvPairs(get_values_cb, lookup_data); + if (!kv_pairs_json_or_status.ok()) { + return kv_pairs_json_or_status.status(); + } + for (const auto& kv_pair_json : *kv_pairs_json_or_status) { + for (auto& [k, v] : kv_pair_json["kvPairs"].items()) { + if (v.contains("value")) { + kv_map[k] = v["value"]; + } + } + } + return kv_map; +} + +// Calls getValuesBinary for a list of keys and processes the response. +absl::StatusOr> +GetKvPairsBinary(const emscripten::val& get_values_binary_cb, + const std::vector& lookup_data) { + std::vector responses; + for (const auto& keys : lookup_data) { + // `get_values_cb` takes an emscripten::val of type array and returns + // an emscripten::val of type Uint8Array that contains a serialized + // BinaryGetValuesResponse. + std::vector get_values_binary_result_array = + emscripten::convertJSArrayToNumberVector( + get_values_binary_cb(keys)); + kv_server::BinaryGetValuesResponse response; + if (!response.ParseFromArray(get_values_binary_result_array.data(), + get_values_binary_result_array.size())) { + return absl::InternalError( + absl::StrCat("Could not parse binary response to proto")); + } + responses.push_back(response); + } + return responses; +} + +absl::StatusOr GetKvPairsBinaryAsEmVal( + const emscripten::val& get_values_binary_cb, + const std::vector& lookup_data) { + emscripten::val kv_pairs = emscripten::val::object(); + kv_pairs.set("udfApi", "getValuesBinary"); + const auto responses_or_status = + GetKvPairsBinary(get_values_binary_cb, lookup_data); + if (!responses_or_status.ok()) { + return responses_or_status.status(); + } + for (const auto& response : *responses_or_status) { + for (auto& [k, v] : response.kv_pairs()) { + kv_pairs.set(k, v.data()); + } + } + return kv_pairs; +} + +absl::StatusOr> GetKvPairsBinaryAsMap( + const emscripten::val& get_values_binary_cb, + const std::vector& lookup_data) { + std::map kv_map; + kv_map["udfApi"] = "getValuesBinary"; + const auto responses_or_status = + GetKvPairsBinary(get_values_binary_cb, lookup_data); + if (!responses_or_status.ok()) { + return responses_or_status.status(); + } + for (const auto& response : *responses_or_status) { + for (auto& [k, v] : response.kv_pairs()) { + kv_map[k] = v.data(); + } + } + return kv_map; +} + +std::vector MaybeSplitDataByBatchSize( + const emscripten::val& request_metadata, const emscripten::val& data) { + if (!request_metadata.hasOwnProperty("lookup_batch_size")) { + return std::vector({data}); + } + const int batch_size = request_metadata["lookup_batch_size"].as(); + const int data_length = data["length"].as(); + if (batch_size >= data_length) { + return std::vector({data}); + } + std::vector batches; + for (int i = 0; i < data_length; i += batch_size) { + batches.emplace_back( + data.call("slice", i, i + batch_size)); + } + return batches; +} + +// I/O processing, similar to +// tools/latency_benchmarking/example/udf_code/benchmark_udf.js +emscripten::val GetKeyGroupOutputs(const emscripten::val& get_values_cb, + const emscripten::val& get_values_binary_cb, + const emscripten::val& request_metadata, + const emscripten::val& udf_arguments) { + emscripten::val key_group_outputs = emscripten::val::array(); + // Convert a JS array to a std::vector so we can iterate through it. + const std::vector key_groups = + emscripten::vecFromJSArray(udf_arguments); + for (const auto& key_group : key_groups) { + emscripten::val key_group_output = emscripten::val::object(); + const emscripten::val data = + key_group.hasOwnProperty("tags") ? key_group["data"] : key_group; + + std::vector lookup_data = + MaybeSplitDataByBatchSize(request_metadata, data); + absl::StatusOr kv_pairs; + kv_pairs = request_metadata.hasOwnProperty("useGetValuesBinary") && + request_metadata["useGetValuesBinary"].as() + ? GetKvPairsBinaryAsEmVal(get_values_binary_cb, lookup_data) + : GetKvPairsAsEmVal(get_values_cb, lookup_data); + if (kv_pairs.ok()) { + key_group_output.set("keyValues", *kv_pairs); + key_group_outputs.call("push", key_group_output); + } + } + return key_group_outputs; +} + +emscripten::val HandleGetValuesFlow(const emscripten::val& get_values_cb, + const emscripten::val& get_values_binary_cb, + const emscripten::val& request_metadata, + const emscripten::val& udf_arguments) { + emscripten::val result = emscripten::val::object(); + emscripten::val key_group_outputs = GetKeyGroupOutputs( + get_values_cb, get_values_binary_cb, request_metadata, udf_arguments); + result.set("keyGroupOutputs", key_group_outputs); + result.set("udfOutputApiVersion", emscripten::val(1)); + return result; +} + +// The run query flow performs the following steps: +// 1. Compute the set union of given arguments using `runQuery` API +// 2. Call `getValues`/`getValuesBinary` with first +// `lookup_n_keys_from_runquery` keys +// 3. Sort returned KVs +// 4. Return top 5 KV pairs +emscripten::val HandleRunQueryFlow(const emscripten::val& get_values_cb, + const emscripten::val& get_values_binary_cb, + const emscripten::val& run_query_cb, + const emscripten::val& request_metadata, + const emscripten::val& udf_arguments) { + emscripten::val result = emscripten::val::object(); + result.set("udfOutputApiVersion", emscripten::val(1)); + if (udf_arguments["length"].as() <= 0) { + return result; + } + // Union of all sets in udf_arguments + emscripten::val set_keys = udf_arguments[0].hasOwnProperty("data") + ? udf_arguments[0]["data"] + : udf_arguments[0]; + emscripten::val query = + set_keys.call("join", emscripten::val("|")); + emscripten::val keys = run_query_cb(query); + int n = kLookupNKeysFromRunQuery; + if (request_metadata.hasOwnProperty("lookup_n_keys_from_runquery")) { + n = request_metadata["lookup_n_keys_from_runquery"].as(); + } + emscripten::val lookup_keys = keys.call("slice", 0, n); + std::vector lookup_data = + MaybeSplitDataByBatchSize(request_metadata, lookup_keys); + absl::StatusOr> kv_map = + request_metadata.hasOwnProperty("useGetValuesBinary") && + request_metadata["useGetValuesBinary"].as() + ? GetKvPairsBinaryAsMap(get_values_binary_cb, lookup_data) + : GetKvPairsAsMap(get_values_cb, lookup_data); + if (!kv_map.ok()) { + return result; + } + int i = 0; + emscripten::val key_group_outputs = emscripten::val::array(); + emscripten::val key_group_output = emscripten::val::object(); + emscripten::val kv_pairs = emscripten::val::object(); + // select only kTopK + for (const auto& [k, v] : *kv_map) { + emscripten::val value = emscripten::val::object(); + value.set("value", v); + kv_pairs.set(k, value); + if (++i > kTopK) { + break; + } + } + key_group_output.set("keyValues", kv_pairs); + key_group_outputs.call("push", key_group_output); + result.set("keyGroupOutputs", key_group_outputs); + return result; +} + +} // namespace + +emscripten::val HandleRequestCc(const emscripten::val& get_values_cb, + const emscripten::val& get_values_binary_cb, + const emscripten::val& run_query_cb, + const emscripten::val& request_metadata, + const emscripten::val& udf_arguments) { + if (request_metadata.hasOwnProperty("runQuery") && + request_metadata["runQuery"].as()) { + return HandleRunQueryFlow(get_values_cb, get_values_binary_cb, run_query_cb, + request_metadata, udf_arguments); + } + return HandleGetValuesFlow(get_values_cb, get_values_binary_cb, + request_metadata, udf_arguments); +} + +EMSCRIPTEN_BINDINGS(HandleRequestExample) { + emscripten::function("handleRequestCc", &HandleRequestCc); +} diff --git a/tools/latency_benchmarking/example/udf_code/benchmark_cpp_wasm_udf.js b/tools/latency_benchmarking/example/udf_code/benchmark_cpp_wasm_udf.js new file mode 100644 index 00000000..4efae9ff --- /dev/null +++ b/tools/latency_benchmarking/example/udf_code/benchmark_cpp_wasm_udf.js @@ -0,0 +1,27 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +async function HandleRequest(executionMetadata, ...udf_arguments) { + const module = await getModule(); + const result = module.handleRequestCc( + getValues, + getValuesBinary, + runQuery, + executionMetadata.requestMetadata, + udf_arguments + ); + return result; +} diff --git a/tools/latency_benchmarking/example/udf_code/benchmark_udf.js b/tools/latency_benchmarking/example/udf_code/benchmark_udf.js new file mode 100644 index 00000000..688f32a3 --- /dev/null +++ b/tools/latency_benchmarking/example/udf_code/benchmark_udf.js @@ -0,0 +1,236 @@ +/** + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * Import the generated JS code based on public/udf/binary_get_values.proto + */ +goog.require('proto.kv_server.BinaryGetValuesResponse'); + +/** + * @param {!Object} value + * @param {!Object} keyGroupOutputs + * @return {Object} + * @suppress {reportUnknownTypes} + */ +function addValueToOutput(value, keyGroupOutputs) { + let keyGroupOutput = value; + keyGroupOutputs.push(keyGroupOutput); + return null; +} + +/** + * @param {!Object} getValuesBinaryProto + * @param {!Object} keyValuesOutput + * @return {Object} + * @suppress {reportUnknownTypes} + */ +function processGetValuesBinary(getValuesBinaryProto, keyValuesOutput) { + var kvMap = getValuesBinaryProto.getKvPairsMap(); + const keys = kvMap.keys(); + for (const key of keys) { + const val = kvMap.get(key); + const hasData = val.hasData(); + if (hasData) { + const valUint8Arr = val.getData_asU8(); + const valString = String.fromCharCode.apply(null, valUint8Arr); + keyValuesOutput[key] = { value: valString }; + } + } + return keyValuesOutput; +} + +/** + * @param {!Array>} lookupData + * @return {Object} + */ +function callGetValuesBinary(lookupData) { + let keyValuesOutput = {}; + for (const data of lookupData) { + const getValuesUnparsed = getValuesBinary(data); + const getValuesBinaryProto = proto.kv_server.BinaryGetValuesResponse.deserializeBinary(getValuesUnparsed); + processGetValuesBinary(getValuesBinaryProto, keyValuesOutput); + } + keyValuesOutput['getValuesApi'] = 'getValuesBinary'; + return keyValuesOutput; +} + +/** + * @param {!Array>} lookupData + * @return {Object} + */ +function callGetValues(lookupData) { + const keyValuesOutput = {}; + for (const data of lookupData) { + const getValuesResult = JSON.parse(getValues(data)); + // getValuesResult returns "kvPairs" when successful and "code" on failure. + // Ignore failures and only add successful getValuesResult lookups to + // output. + if (getValuesResult.hasOwnProperty('kvPairs')) { + const kvPairs = getValuesResult['kvPairs']; + for (const key in kvPairs) { + if (kvPairs[key].hasOwnProperty('value')) { + keyValuesOutput[key] = { value: kvPairs[key]['value'] }; + } + } + } + } + keyValuesOutput['getValuesApi'] = 'getValues'; + return keyValuesOutput; +} + +/** + * @param {!Array} data + * @param {number} batchSize + * @return {!Array>} + */ + +function splitDataByBatchSize(data, batchSize) { + const batches = []; + for (let i = 0; i < data.length; i += batchSize) { + batches.push(data.slice(i, i + batchSize)); + } + return batches; +} + +/** + * @param {Object} requestMetadata + * @param {!Array} udf_arguments + * @return {Object} + */ +function handleGetValuesRequest(requestMetadata, udf_arguments) { + let keyGroupOutputs = []; + for (let argument of udf_arguments) { + let keyGroupOutput = {}; + let data = argument.hasOwnProperty('tags') ? argument['data'] : argument; + let lookupData = [data]; + if (requestMetadata.hasOwnProperty('lookup_batch_size')) { + batchSize = parseInt(requestMetadata['lookup_batch_size'], 10); + lookupData = splitDataByBatchSize(data, batchSize); + } + keyGroupOutput['keyValues'] = requestMetadataKeyIsTrue(requestMetadata, 'useGetValuesBinary') + ? callGetValuesBinary(lookupData) + : callGetValues(lookupData); + keyGroupOutputs.push(keyGroupOutput); + } + return keyGroupOutputs; +} + +function requestMetadataKeyIsTrue(requestMetadata, key) { + return requestMetadata.hasOwnProperty(key) && requestMetadata[key] == true; +} + +/** + * Handles the runQuery flow: + * 1. compute set union on arguments using `runQuery` + * 2. getValues/getValuesBinary on first `lookup_n_keys_from_runquery` keys + * 3. sort returned KVs by key length + * 4. return top 5 KVs + * + * @param {!Object} requestMetadata + * @param {!Array} udf_arguments + * @returns {Object} + */ +function handleRunQueryFlow(requestMetadata, udf_arguments) { + var result = {}; + result['udfOutputApiVersion'] = 1; + if (!udf_arguments.length) { + return result; + } + + // Union all the sets in the udf_arguments + const setKeys = udf_arguments[0].hasOwnProperty('data') ? udf_arguments[0]['data'] : udf_arguments[0]; + const keys = runQuery(setKeys.join('|')); + const n = parseInt(requestMetadata['lookup_n_keys_from_runquery'] || '100', 10); + + const keyValuesOutput = handleGetValuesRequest(requestMetadata, [keys.slice(0, n)]); + if (!keyValuesOutput.length || !keyValuesOutput[0].hasOwnProperty('keyValues')) { + return result; + } + let top5keyValuesArray = Object.entries(keyValuesOutput[0]['keyValues']) + .sort((a, b) => a[0].length - b[0].length) + .slice(0, 5); + result['keyGroupOutputs'] = Object.fromEntries(top5keyValuesArray); + return result; +} + +/** + * Handles the getValues flow (without runQuery&sorting). + * + * @param {!Object} requestMetadata + * @param {!Array} udf_arguments + * @returns {Object} + */ +function handleGetValuesFlow(requestMetadata, udf_arguments) { + const keyGroupOutputs = handleGetValuesRequest(requestMetadata, udf_arguments); + var result = {}; + result['keyGroupOutputs'] = keyGroupOutputs; + result['udfOutputApiVersion'] = 1; + return result; +} + +/** + * Entry point for code snippet execution. + * + * This is a sample UDF js that calls different UDF APIs depending + * on executionMetadata.requestMetadata fields. + * + * If `executionMetadata.requestMetadata.useGetValuesBinary` is set to true, + * the UDF will use the `getValuesBinary` API to retrieve key-values. + * Otherwise, the UDF will use the `getValues` API to retrieve key-values. + * + * If `executionMetadata.requestMetadata.runQuery` is set to true, + * the UDF will do the following steps: + * 1. Compute the set union of all elements in the first `udf_argument` + * 2. Call `getValues`/`getValuesBinary` on the returned keys from step 1. + * 3. Return the top 5 key values (sorted by key length) from step 2. + * Otherwise, the UDF will just act as a pass-through UDF that calls + * `getValues`/`getValuesBinary` on the `udf_arguments` and return the + * retrieved key-value pairs. + * + * If `requestMetadata` has a `lookup_batch_size`, it will batch + * getValues/getValuesBinary calls by the provided `lookup_batch_size`. + * For example, if the input has 800 keys and `lookup_batch_size` is 300, + * the UDF will make 3 getValues/getValuesBinary calls: + * 2 call with 300 keys + * 1 call with 200 keys + * + * The Closure Compiler will see the @export JSDoc annotation below and + * generate the following synthetic code in the global scope: + * + * goog.exportSymbol('HandleRequest', HandleRequest); + * + * That makes the minified version of this function still available in the + * global namespace. + * + * You can also use @export for property names on classes, to prevent them + * from being minified. + * + * @export + * @param {Object} executionMetadata + * @param {...!Object} udf_arguments + * @return {Object} + * @suppress {checkTypes} + */ +function HandleRequest(executionMetadata, ...udf_arguments) { + const requestMetadata = executionMetadata.hasOwnProperty('requestMetadata') + ? executionMetadata['requestMetadata'] + : {}; + + if (requestMetadataKeyIsTrue(requestMetadata, 'runQuery')) { + return handleRunQueryFlow(requestMetadata, udf_arguments); + } + return handleGetValuesFlow(requestMetadata, udf_arguments); +} diff --git a/tools/latency_benchmarking/example/udf_code/externs.js b/tools/latency_benchmarking/example/udf_code/externs.js new file mode 100644 index 00000000..3b982ca0 --- /dev/null +++ b/tools/latency_benchmarking/example/udf_code/externs.js @@ -0,0 +1,34 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * These functions are provided by the K/V server as a part of the + * UDF API and need to be marked as @externs for the Closure Compiler. + * + * @externs + */ +function getValues(keyList) {} // Note the empty body. +function runQuery(query) {} // Note the empty body. + +/** + * Returns a serialized BinaryGetValuesResponse for the keyList. + * + * This function is provided by the K/V server as a part of the + * UDF API and needs to be marked as @externs for the Closure Compiler. + * + * @externs + */ +function getValuesBinary(keyList) {} // Note the empty body. diff --git a/tools/latency_benchmarking/generate_requests.py b/tools/latency_benchmarking/generate_requests.py new file mode 100644 index 00000000..e55b5733 --- /dev/null +++ b/tools/latency_benchmarking/generate_requests.py @@ -0,0 +1,160 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import csv +import argparse +import json +import base64 + +from pathlib import Path +from typing import Any, Iterator + +""" +Generates JSON requests to be sent to the KV server V2 API. + +For each N in a number-of-keys-list, it will + * select N keys from the given snapshot.csv file + * generate a request body with the list of keys + * save request under //n=/request.json +""" + + +def _BuildRequest(data: list[str], metadata: dict[str, str]) -> dict[str, Any]: + """Build the HTTP body that contains the base64 encoded request body as data.""" + arguments = [] + argument = {"tags": ["custom", "keys"], "data": data} + arguments.append(argument) + + body = { + "metadata": metadata, + "partitions": [{"id": 0, "compressionGroupId": 0, "arguments": arguments}], + } + body_base64_string = base64.b64encode(json.dumps(body).encode()).decode() + http_body = {"raw_body": {"data": body_base64_string}} + return json.dumps(http_body) + + +def WriteRequests( + keys: list[str], + number_of_keys_list: list[int], + output_dir: str, + metadata: dict[str, str], +) -> None: + """Writes the requests to JSON files. + + Args: + keys: List of all keys. + number_of_keys_list: List of number of keys to use in request. One request will generated per item. + output_dir: Base output dir to write to. + metadata: Metadata to include in the request body, as per V2 API. + """ + for n in number_of_keys_list: + if n > len(keys): + print( + f"Warning: List of provided lookup keys ({len(keys)}) is smaller than number of keys to be included in request ({n}). Skipping...\n" + ) + continue + + request = _BuildRequest(keys[:n], metadata) + # Write to an output file at /n=/request.json + output_dir_n = os.path.join(output_dir, f"{n=}") + Path(output_dir_n).mkdir(parents=True, exist_ok=True) + with open(os.path.join(output_dir_n, "request.json"), "w") as f: + f.write(request) + + +def _ReadCsv(snapshot_csv_file: str) -> Iterator[str]: + with open(snapshot_csv_file, "r") as csvfile: + reader = csv.DictReader(csvfile) + for row in reader: + yield row + + +def ReadKeys( + snapshot_csv_file: str, max_number_of_keys: int, filter_by_sets: bool +) -> list[str]: + """Read keys from CSV file. Only include update and string type mutations. + + Args: + snapshot_csv_file: Path to snapshot CSV file. + max_number_of_keys: Maximum number of keys to read. + filter_by_sets: Whether to only use rows with value_type "string_set" + + Returns: + List of unique set of keys. + """ + keys = set() + for row in _ReadCsv(snapshot_csv_file): + if filter_by_sets and row["value_type"].lower() != "string_set": + continue + if not filter_by_sets and row["value_type"].lower() != "string": + continue + if row["mutation_type"].lower() == "update": + keys.add(row["key"]) + if len(keys) >= max_number_of_keys: + break + return list(keys) + + +def Main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--number-of-keys-list", + dest="number_of_keys_list", + type=int, + nargs="+", + help="List of number of keys to include in ghz request", + ) + parser.add_argument( + "--output-dir", + dest="output_dir", + type=str, + default="/tmp/latency_benchmarking", + help="Output directory for benchmarks", + ) + parser.add_argument( + "--snapshot-csv-dir", + dest="snapshot_csv_dir", + default="snapshot.csv", + help="Directory with snapshot CSVs with KVMutation update entries.", + ) + parser.add_argument( + "--metadata", + dest="metadata", + type=str, + default="{}", + help="Request metadata as a JSON object.", + ) + parser.add_argument( + "--filter-by-sets", + dest="filter_by_sets", + action="store_true", + help="Whether to only use keys of sets from the input to build the requests", + ) + args = parser.parse_args() + metadata = json.loads(args.metadata) + if not isinstance(metadata, dict): + raise ValueError("metadata is not a JSON object") + for filename in os.listdir(args.snapshot_csv_dir): + snapshot_csv_file = os.path.join(args.snapshot_csv_dir, filename) + keys = ReadKeys( + snapshot_csv_file, max(args.number_of_keys_list), args.filter_by_sets + ) + output_dir_for_snapshot = os.path.join(args.output_dir, filename) + WriteRequests(keys, args.number_of_keys_list, output_dir_for_snapshot, metadata) + + +if __name__ == "__main__": + Main() diff --git a/tools/latency_benchmarking/merge_csvs.py b/tools/latency_benchmarking/merge_csvs.py new file mode 100644 index 00000000..fcccbf19 --- /dev/null +++ b/tools/latency_benchmarking/merge_csvs.py @@ -0,0 +1,41 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import pandas as pd + +""" +Merges list of CSV files into one +""" + + +def Main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--csv-inputs", + dest="csv_inputs", + type=str, + nargs="+", + help="List of csv input files.", + ) + parser.add_argument( + "--csv-output", dest="csv_output", type=str, help="Output CSV file to write to." + ) + args = parser.parse_args() + df = pd.concat([pd.read_csv(file) for file in args.csv_inputs], ignore_index=True) + df.to_csv(args.csv_output, index=False) + + +if __name__ == "__main__": + Main() diff --git a/tools/latency_benchmarking/run_benchmarks b/tools/latency_benchmarking/run_benchmarks new file mode 100755 index 00000000..c67d11fb --- /dev/null +++ b/tools/latency_benchmarking/run_benchmarks @@ -0,0 +1,197 @@ +#!/bin/bash +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -o pipefail +set -o errexit + +START=$(date +%s) +readonly START + +WORKSPACE="$(git rev-parse --show-toplevel)" + +BASE_OUTPUT_DIR="${WORKSPACE}/dist/tools/latency_benchmarking/output/${START}" + +NUMBER_OF_LOOKUP_KEYS="1 5 10" +REQUEST_METADATA="{}" +FILTER_BY_SETS=0 +FILTER_BY_SETS_JSON='{"filter_snapshot_by_sets": "false"}' +BENCHMARK_DURATION="5s" + +DOCKER_OUTPUT_DIR="/tmp/latency_benchmarking/output" +readonly DOCKER_OUTPUT_DIR + +DOCKER_SNAPSHOT_DIR="/tmp/latency_benchmarking/snapshots" +DOCKER_SNAPSHOT_CSV_DIR="/tmp/latency_benchmarking/snapshot_csvs" + +function usage() { + local -r -i exitval=${1-1} + cat &>/dev/stderr < + [--server-address] (Required) gRPC host and port. + [--snapshot-dir] or (Required) Full path to either snapshot-dir or snapshot-csv-dir + [--snapshot-csv-dir] is required. + [--number-of-lookup-keys-list] (Optional) List of number of keys to include in a request. + [--benchmark-duration] (Optional) Duration of each benchmark. Default "5s". + [--ghz-tags] (Optional) Tags to include in the ghz run. + [--request-metadata-json] (Optional) Request metadata to set for all requests + [--filter-snapshot-by-sets] (Optional) Whether to filter snapshots by sets when creating requests. +USAGE + # shellcheck disable=SC2086 + exit ${exitval} +} + +function generate_requests() { + # Create output dirs before calling bazel-debian + # If the dir is created in the docker container, we lose write permission + for SNAPSHOT_CSV in "${SNAPSHOT_CSV_DIR}"/*; do + SNAPSHOT_CSV_FILENAME=$(basename "${SNAPSHOT_CSV}") + for N in "${NUMBER_OF_LOOKUP_KEYS_LIST[@]}"; do + mkdir -p "${BASE_OUTPUT_DIR}/${SNAPSHOT_CSV_FILENAME}/n=${N}" + done + done + + local -a GENERATE_REQUESTS_ARGS + GENERATE_REQUESTS_ARGS+=(--output-dir "${DOCKER_OUTPUT_DIR}") + GENERATE_REQUESTS_ARGS+=(--number-of-keys-list "${NUMBER_OF_LOOKUP_KEYS_LIST[@]}") + GENERATE_REQUESTS_ARGS+=(--metadata "${REQUEST_METADATA}") + GENERATE_REQUESTS_ARGS+=(--snapshot-csv-dir "${DOCKER_SNAPSHOT_CSV_DIR}") + if [[ "${FILTER_BY_SETS}" = 1 ]]; then + GENERATE_REQUESTS_ARGS+=(--filter-by-sets) + FILTER_BY_SETS_JSON='{"filter_snapshot_by_sets": "true"}' + fi + + # Mount the output dir to docker and write requests to output dir for each item in + # `NUMBER_OF_LOOKUP_KEYS_LIST`. + # This will write a json request for each NUMBER_OF_LOOKUP_KEY=N to + # dist/tools/latency_benchmarking/output/${START}/${SNAPSHOT_FILENAME}/n=${N}/request.json + EXTRA_DOCKER_RUN_ARGS+=" --volume ${BASE_OUTPUT_DIR}:${DOCKER_OUTPUT_DIR} --volume ${SNAPSHOT_CSV_DIR}:${DOCKER_SNAPSHOT_CSV_DIR} " \ + builders/tools/bazel-debian run //tools/latency_benchmarking:generate_requests \ + -- "${GENERATE_REQUESTS_ARGS[@]}" +} + +function run_ghz_for_requests() { + # Iterate through the generated request.json files and call `ghz` to benchmark server at ${SERVER_ADDRESS} + for N in "${NUMBER_OF_LOOKUP_KEYS_LIST[@]}"; do + DIR="${OUTPUT_DIR}/n=${N}" + REQUEST_JSON="${DIR}"/request.json + if [[ ! -f "${REQUEST_JSON}" ]]; then + continue + fi + + printf "Running ghz for number of keys %s\n" "${N}" + BASE_GHZ_TAGS='{"number_of_lookup_keys": "'"${N}"'", "keys_from_file": "'"${SNAPSHOT_CSV_FILENAME}"'"}' + REQUEST_METADATA_TAGS=$(echo "${REQUEST_METADATA}" | jq 'with_entries(.key |= "request_metadata."+.) | with_entries( .value |= @json)') + TAGS=$(echo "${GHZ_TAGS} ${BASE_GHZ_TAGS} ${REQUEST_METADATA_TAGS} ${FILTER_BY_SETS_JSON}" | jq -s -c 'add') + GHZ_OUTPUT_JSON_FILE="${DIR}/ghz_output.json" + ghz --protoset "${WORKSPACE}/dist/query_api_descriptor_set.pb" \ + -D "${REQUEST_JSON}" \ + --duration="${BENCHMARK_DURATION}" \ + --skipFirst=100 \ + --concurrency=100 \ + --format=pretty \ + --tags "${TAGS}" \ + --output "${GHZ_OUTPUT_JSON_FILE}" \ + --call kv_server.v2.KeyValueService/GetValuesHttp \ + "${SERVER_ADDRESS}" + # In case the server hangs from previous requests, wait before sending next batch + sleep 30 + done +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --server-address) + SERVER_ADDRESS="$2" + shift 2 || usage + ;; + --number-of-lookup-keys-list) + NUMBER_OF_LOOKUP_KEYS="$2" + shift 2 + ;; + --benchmark-duration) + BENCHMARK_DURATION="$2" + shift 2 + ;; + --ghz-tags) + GHZ_TAGS="$2" + shift 2 + ;; + --snapshot-dir) + SNAPSHOT_DIR="$2" + shift 2 + ;; + --snapshot-csv-dir) + SNAPSHOT_CSV_DIR="$2" + shift 2 + ;; + --request-metadata-json) + REQUEST_METADATA="$2" + shift 2 + ;; + --filter-snapshot-by-sets) + FILTER_BY_SETS=1 + shift + ;; + -h | --help) usage 0 ;; + *) usage ;; + esac +done + +# Check for SNAPSHOT_DIR or SNAPSHOT_CSV_DIR. If not available, exit. +if [[ (-z "${SNAPSHOT_DIR}" || ! -d "${SNAPSHOT_DIR}") && (-z "${SNAPSHOT_CSV_DIR}" || ! -d "${SNAPSHOT_CSV_DIR}") ]]; then + printf "snapshot-dir not found:%s\n" "${SNAPSHOT_DIR}" + exit 1 +fi + +IFS=' ' read -ra NUMBER_OF_LOOKUP_KEYS_LIST <<<"${NUMBER_OF_LOOKUP_KEYS}" +if [[ -d "${SNAPSHOT_CSV_DIR}" ]]; then + generate_requests +else + # No snapshot csv given, iterate through snapshot dir to create csvs + SNAPSHOT_CSV_DIR="${BASE_OUTPUT_DIR}/snapshot_csvs/" + # Iterate through snapshot files and convert them to CSV + mkdir -p "${SNAPSHOT_CSV_DIR}" + for SNAPSHOT_FILE in "${SNAPSHOT_DIR}"/*; do + SNAPSHOT_FILENAME=$(basename "${SNAPSHOT_FILE}") + EXTRA_DOCKER_RUN_ARGS+=" --volume ${SNAPSHOT_CSV_DIR}:${DOCKER_SNAPSHOT_CSV_DIR} --volume ${SNAPSHOT_DIR}:${DOCKER_SNAPSHOT_DIR} " \ + builders/tools/bazel-debian run //tools/data_cli:data_cli format_data \ + -- \ + --input_file "${DOCKER_SNAPSHOT_DIR}/${SNAPSHOT_FILENAME}" \ + --input_format DELTA \ + --output_file "${DOCKER_SNAPSHOT_CSV_DIR}/${SNAPSHOT_FILENAME}.csv" \ + --output_format CSV + done + generate_requests +fi + +# Iterate through each snapshot file and +# 1. create requests under dist/tools/latency_benchmarking/output/${START}/${SNAPSHOT_FILENAME} +# 2. for each request, run benchmarks with ghz +for SNAPSHOT_CSV_FILE in "${SNAPSHOT_CSV_DIR}"/*; do + SNAPSHOT_CSV_FILENAME=$(basename "${SNAPSHOT_CSV_FILE}") + OUTPUT_DIR="${BASE_OUTPUT_DIR}/${SNAPSHOT_CSV_FILENAME}" + run_ghz_for_requests +done + +# Go through all the ghz results in dist/tools/latency_benchmarking/output/${START} +# and collect the summary in a csv +EXTRA_DOCKER_RUN_ARGS+=" --volume ${BASE_OUTPUT_DIR}:${DOCKER_OUTPUT_DIR} " \ + builders/tools/bazel-debian run //tools/latency_benchmarking:create_csv_summary \ + -- \ + --ghz-result-dir ${DOCKER_OUTPUT_DIR} + +echo "Result in:" +echo "${BASE_OUTPUT_DIR}/summary.csv" diff --git a/tools/request_simulation/BUILD.bazel b/tools/request_simulation/BUILD.bazel index 1a8f01c7..f09a4173 100644 --- a/tools/request_simulation/BUILD.bazel +++ b/tools/request_simulation/BUILD.bazel @@ -47,7 +47,7 @@ cc_library( "//components/util:sleepfor", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//src/cpp/util:duration", + "@google_privacysandbox_servers_common//src/util:duration", ], ) @@ -80,9 +80,9 @@ cc_library( deps = [ "//components/data/common:thread_manager", "//components/util:sleepfor", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//test/core/util:grpc_test_util_base", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", + "@com_google_absl//absl/log", + "@google_privacysandbox_servers_common//src/telemetry:metrics_recorder", ], ) @@ -95,9 +95,9 @@ cc_library( ":metrics_collector", ":rate_limiter", "//components/data/common:thread_manager", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/functional:any_invocable", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", ], ) @@ -118,10 +118,10 @@ cc_library( "//public/data_loading/readers:riegeli_stream_io", "//public/data_loading/readers:stream_record_reader_factory", "@com_github_google_flatbuffers//:flatbuffers", - "@com_github_google_glog//:glog", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", - "@google_privacysandbox_servers_common//src/cpp/telemetry:metrics_recorder", - "@google_privacysandbox_servers_common//src/cpp/telemetry:tracing", + "@google_privacysandbox_servers_common//src/telemetry:metrics_recorder", + "@google_privacysandbox_servers_common//src/telemetry:tracing", ], ) @@ -135,8 +135,8 @@ cc_library( hdrs = ["request_simulation_parameter_fetcher.h"], deps = [ "//components/data/common:message_service", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/log", ], ) @@ -147,8 +147,10 @@ cc_library( deps = [ ":client_worker", ":delta_based_request_generator", + ":detla_based_realtime_updates_publisher", ":grpc_client", ":metrics_collector", + ":realtime_message_batcher", ":request_generation_util", ":request_simulation_parameter_fetcher", ":synthetic_request_generator", @@ -163,9 +165,9 @@ cc_library( "//public/data_loading/readers:riegeli_stream_record_reader_factory", "//public/query:get_values_cc_grpc", "//tools/request_simulation/request:raw_request_cc_proto", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", "@com_google_absl//absl/functional:any_invocable", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", ], ) @@ -189,11 +191,13 @@ cc_binary( ], deps = [ ":request_simulation_system", - "@com_github_google_glog//:glog", "@com_google_absl//absl/debugging:failure_signal_handler", "@com_google_absl//absl/debugging:symbolize", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/strings", ], ) @@ -267,7 +271,7 @@ cc_test( "//public/testing:fake_key_value_service_impl", "//tools/request_simulation/request:raw_request_cc_proto", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/telemetry:mocks", ], ) @@ -279,7 +283,7 @@ cc_test( ":metrics_collector", "//components/util:sleepfor_mock", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/telemetry:mocks", ], ) @@ -302,7 +306,7 @@ cc_test( "//tools/request_simulation:mocks", "//tools/request_simulation/request:raw_request_cc_proto", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/telemetry:mocks", ], ) @@ -315,7 +319,7 @@ cc_test( "//components/data/common:mocks", "//public/test_util:mocks", "@com_google_googletest//:gtest_main", - "@google_privacysandbox_servers_common//src/cpp/telemetry:mocks", + "@google_privacysandbox_servers_common//src/telemetry:mocks", ], ) @@ -327,7 +331,8 @@ cc_library( "//components/tools:concurrent_publishing_engine", "//public/data_loading/writers:delta_record_limiting_file_writer", "//public/sharding:key_sharder", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@com_google_absl//absl/log:check", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) @@ -357,8 +362,8 @@ cc_library( "//components/data/realtime:realtime_notifier", "//public/data_loading:records_utils", "//public/data_loading/readers:riegeli_stream_record_reader_factory", - "@com_github_google_glog//:glog", "@com_google_absl//absl/functional:any_invocable", + "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", ], ) diff --git a/tools/request_simulation/client_worker.h b/tools/request_simulation/client_worker.h index cb679b48..d4c82ae4 100644 --- a/tools/request_simulation/client_worker.h +++ b/tools/request_simulation/client_worker.h @@ -22,8 +22,8 @@ #include #include "absl/functional/any_invocable.h" +#include "absl/log/log.h" #include "components/data/common/thread_manager.h" -#include "glog/logging.h" #include "grpcpp/grpcpp.h" #include "tools/request_simulation/grpc_client.h" #include "tools/request_simulation/message_queue.h" @@ -67,7 +67,7 @@ class ClientWorker { metrics_collector_(metrics_collector), request_converter_(std::move(request_converter)), thread_manager_( - TheadManager::Create(absl::StrCat("Client worker ", id))) { + ThreadManager::Create(absl::StrCat("Client worker ", id))) { grpc_client_ = std::make_unique>( channel, request_timeout, is_client_channel); } @@ -93,7 +93,7 @@ class ClientWorker { std::unique_ptr> grpc_client_; absl::AnyInvocable request_converter_; // Thread manager to start or stop the request sending thread. - std::unique_ptr thread_manager_; + std::unique_ptr thread_manager_; }; template diff --git a/tools/request_simulation/client_worker_test.cc b/tools/request_simulation/client_worker_test.cc index a787a225..a9c45ea1 100644 --- a/tools/request_simulation/client_worker_test.cc +++ b/tools/request_simulation/client_worker_test.cc @@ -25,8 +25,8 @@ #include "grpcpp/grpcpp.h" #include "gtest/gtest.h" #include "public/testing/fake_key_value_service_impl.h" -#include "src/cpp/telemetry/mocks.h" -#include "src/cpp/util/duration.h" +#include "src/telemetry/mocks.h" +#include "src/util/duration.h" #include "tools/request_simulation/mocks.h" #include "tools/request_simulation/request/raw_request.pb.h" #include "tools/request_simulation/request_generation_util.h" diff --git a/tools/request_simulation/delta_based_request_generator.cc b/tools/request_simulation/delta_based_request_generator.cc index eef90c2f..37ceaaad 100644 --- a/tools/request_simulation/delta_based_request_generator.cc +++ b/tools/request_simulation/delta_based_request_generator.cc @@ -20,7 +20,7 @@ #include "public/data_loading/data_loading_generated.h" #include "public/data_loading/filename_utils.h" #include "public/data_loading/records_utils.h" -#include "src/cpp/telemetry/tracing.h" +#include "src/telemetry/tracing.h" using privacy_sandbox::server_common::MetricsRecorder; using privacy_sandbox::server_common::TraceWithStatusOr; @@ -42,7 +42,8 @@ class BlobRecordStream : public RecordStream { absl::Status DeltaBasedRequestGenerator::Start() { LOG(INFO) << "Start monitor and load delta files"; absl::Status status = options_.delta_notifier.Start( - options_.change_notifier, {.bucket = options_.data_bucket}, "", + options_.change_notifier, {.bucket = options_.data_bucket}, + {std::make_pair("", "")}, absl::bind_front(&DeltaBasedRequestGenerator::EnqueueNewFilesToProcess, this)); if (!status.ok()) { diff --git a/tools/request_simulation/delta_based_request_generator.h b/tools/request_simulation/delta_based_request_generator.h index f87ff4da..ada28ca2 100644 --- a/tools/request_simulation/delta_based_request_generator.h +++ b/tools/request_simulation/delta_based_request_generator.h @@ -30,7 +30,7 @@ #include "components/data/realtime/realtime_notifier.h" #include "public/data_loading/readers/riegeli_stream_io.h" #include "public/data_loading/readers/stream_record_reader_factory.h" -#include "src/cpp/telemetry/metrics_recorder.h" +#include "src/telemetry/metrics_recorder.h" #include "tools/request_simulation/message_queue.h" #include "tools/request_simulation/request_generation_util.h" @@ -57,7 +57,7 @@ class DeltaBasedRequestGenerator { privacy_sandbox::server_common::MetricsRecorder& metrics_recorder) : options_(std::move(options)), data_load_thread_manager_( - TheadManager::Create("Delta file loading thread")), + ThreadManager::Create("Delta file loading thread")), request_generation_fn_(std::move(request_generation_fn)), metrics_recorder_(metrics_recorder) {} ~DeltaBasedRequestGenerator() = default; @@ -92,7 +92,7 @@ class DeltaBasedRequestGenerator { absl::Mutex mu_; std::deque unprocessed_basenames_ ABSL_GUARDED_BY(mu_); bool stop_ ABSL_GUARDED_BY(mu_) = false; - std::unique_ptr data_load_thread_manager_; + std::unique_ptr data_load_thread_manager_; // Callback function to generate KV request from a given key absl::AnyInvocable request_generation_fn_; privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; diff --git a/tools/request_simulation/delta_based_request_generator_test.cc b/tools/request_simulation/delta_based_request_generator_test.cc index ee758088..ecc92de7 100644 --- a/tools/request_simulation/delta_based_request_generator_test.cc +++ b/tools/request_simulation/delta_based_request_generator_test.cc @@ -26,7 +26,7 @@ #include "public/data_loading/filename_utils.h" #include "public/data_loading/records_utils.h" #include "public/test_util/mocks.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/telemetry/mocks.h" #include "tools/request_simulation/request_generation_util.h" using kv_server::BlobStorageChangeNotifier; @@ -58,8 +58,10 @@ using testing::_; using testing::AllOf; using testing::ByMove; using testing::Field; +using testing::Pair; using testing::Return; using testing::ReturnRef; +using testing::UnorderedElementsAre; namespace { @@ -103,10 +105,11 @@ TEST_F(GenerateRequestsFromDeltaFilesTest, LoadingDataFromDeltaFiles) { DeltaBasedRequestGenerator request_generator( std::move(options_), std::move(GetRequestGenFn()), metrics_recorder_); const std::string last_basename = ""; - EXPECT_CALL(notifier_, Start(_, GetTestLocation(), last_basename, _)) - .WillOnce([](BlobStorageChangeNotifier& change_notifier, - BlobStorageClient::DataLocation location, - std::string start_after, + EXPECT_CALL(notifier_, + Start(_, GetTestLocation(), + UnorderedElementsAre(Pair("", last_basename)), _)) + .WillOnce([](BlobStorageChangeNotifier&, BlobStorageClient::DataLocation, + absl::flat_hash_map, std::function callback) { callback(ToDeltaFileName(1).value()); LOG(INFO) << "Notified 1 file"; diff --git a/tools/request_simulation/detla_based_realtime_updates_publisher.cc b/tools/request_simulation/detla_based_realtime_updates_publisher.cc index fad28107..0abfae97 100644 --- a/tools/request_simulation/detla_based_realtime_updates_publisher.cc +++ b/tools/request_simulation/detla_based_realtime_updates_publisher.cc @@ -16,6 +16,7 @@ #include "absl/functional/bind_front.h" #include "public/data_loading/filename_utils.h" +#include "src/util/status_macro/status_macros.h" using privacy_sandbox::server_common::TraceWithStatusOr; @@ -33,16 +34,20 @@ class BlobRecordStream : public RecordStream { }; } // namespace +// b/321746753 -- this class can potentially benefit from refactoring absl::Status DeltaBasedRealtimeUpdatesPublisher::Start() { LOG(INFO) << "Start monitor delta files and publish them as realtime updates"; - absl::Status status = options_.delta_notifier.Start( - options_.change_notifier, {.bucket = options_.data_bucket}, "", + PS_RETURN_IF_ERROR(options_.delta_notifier.Start( + options_.change_notifier, {.bucket = options_.data_bucket}, + {std::make_pair("", "")}, absl::bind_front( - &DeltaBasedRealtimeUpdatesPublisher::EnqueueNewFilesToProcess, this)); - if (!status.ok()) { - return status; - } - return data_load_thread_manager_->Start([this]() { ProcessNewFiles(); }); + &DeltaBasedRealtimeUpdatesPublisher::EnqueueNewFilesToProcess, + this))); + PS_RETURN_IF_ERROR( + data_load_thread_manager_->Start([this]() { ProcessNewFiles(); })); + + return rt_publisher_thread_manager_->Start( + [this]() { concurrent_publishing_engine_->Start(); }); } absl::Status DeltaBasedRealtimeUpdatesPublisher::Stop() { @@ -57,6 +62,15 @@ absl::Status DeltaBasedRealtimeUpdatesPublisher::Stop() { LOG(ERROR) << "Failed to stop notify: " << status; } } + + if (rt_publisher_thread_manager_->IsRunning()) { + concurrent_publishing_engine_->Stop(); + if (const auto status = rt_publisher_thread_manager_->Stop(); + !status.ok()) { + LOG(ERROR) << "Failed to stop realtime publisher thread manager: " + << status; + } + } LOG(INFO) << "Delta notifier stopped"; LOG(INFO) << "Stopping publishing realtime updates from"; return data_load_thread_manager_->Stop(); @@ -112,10 +126,6 @@ DeltaBasedRealtimeUpdatesPublisher::CreateRealtimeMessagesAndAddToQueue( return std::make_unique( blob_client.GetBlobReader(location)); }); - auto metadata = record_reader->GetKVFileMetadata(); - if (!metadata.ok()) { - return metadata.status(); - } DataLoadingStats data_loading_stats; const auto process_data_record_fn = [this, &data_loading_stats](const DataRecord& data_record) { @@ -136,13 +146,10 @@ DeltaBasedRealtimeUpdatesPublisher::CreateRealtimeMessagesAndAddToQueue( } } }; - auto status = record_reader->ReadStreamRecords( + PS_RETURN_IF_ERROR(record_reader->ReadStreamRecords( [&process_data_record_fn](std::string_view raw) { return DeserializeDataRecord(raw, process_data_record_fn); - }); - if (!status.ok()) { - return status; - } + })); return data_loading_stats; } bool DeltaBasedRealtimeUpdatesPublisher::HasNewEventToProcess() const { diff --git a/tools/request_simulation/detla_based_realtime_updates_publisher.h b/tools/request_simulation/detla_based_realtime_updates_publisher.h index 1812a82b..8059a27c 100644 --- a/tools/request_simulation/detla_based_realtime_updates_publisher.h +++ b/tools/request_simulation/detla_based_realtime_updates_publisher.h @@ -51,11 +51,15 @@ class DeltaBasedRealtimeUpdatesPublisher { }; DeltaBasedRealtimeUpdatesPublisher( std::unique_ptr realtime_message_batcher, + std::unique_ptr concurrent_publishing_engine, Options options) : realtime_message_batcher_(std::move(realtime_message_batcher)), + concurrent_publishing_engine_(std::move(concurrent_publishing_engine)), options_(std::move(options)), data_load_thread_manager_( - TheadManager::Create("Realtime sharded publisher thread")) {} + ThreadManager::Create("Realtime sharded publisher thread")), + rt_publisher_thread_manager_( + ThreadManager::Create("Publisher thread")) {} ~DeltaBasedRealtimeUpdatesPublisher() = default; // DeltaBasedRealtimeUpdatesPublisher is neither copyable nor movable. @@ -75,7 +79,8 @@ class DeltaBasedRealtimeUpdatesPublisher { private: std::unique_ptr realtime_message_batcher_; - // Checks if there is a new event such as new file or stop event to process + std::unique_ptr concurrent_publishing_engine_; + // Checks if there is new event such as new file or stop event to process bool HasNewEventToProcess() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_); absl::Mutex mu_; Options options_; @@ -90,7 +95,8 @@ class DeltaBasedRealtimeUpdatesPublisher { BlobStorageClient::DataLocation location); std::deque unprocessed_basenames_ ABSL_GUARDED_BY(mu_); bool stop_ ABSL_GUARDED_BY(mu_) = false; - std::unique_ptr data_load_thread_manager_; + std::unique_ptr data_load_thread_manager_; + std::unique_ptr rt_publisher_thread_manager_; }; } // namespace kv_server diff --git a/tools/request_simulation/grpc_client.h b/tools/request_simulation/grpc_client.h index 50f49544..a29762ff 100644 --- a/tools/request_simulation/grpc_client.h +++ b/tools/request_simulation/grpc_client.h @@ -135,7 +135,7 @@ class GrpcClient { std::shared_ptr response) { if (is_client_channel_ && grpc_channel_->GetState(true) != GRPC_CHANNEL_READY) { - return absl::FailedPreconditionError("GRPC channel is disconnected"); + return absl::UnavailableError("GRPC channel is disconnected"); } std::shared_ptr notification = std::make_shared(); diff --git a/tools/request_simulation/main.cc b/tools/request_simulation/main.cc index 8e0aa5fd..0690a362 100644 --- a/tools/request_simulation/main.cc +++ b/tools/request_simulation/main.cc @@ -17,10 +17,12 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" #include "absl/flags/usage.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "absl/strings/str_cat.h" -#include "glog/logging.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry_provider.h" +#include "src/telemetry/metrics_recorder.h" +#include "src/telemetry/telemetry_provider.h" #include "tools/request_simulation/grpc_client.h" #include "tools/request_simulation/request_simulation_system.h" @@ -34,7 +36,7 @@ int main(int argc, char** argv) { absl::FailureSignalHandlerOptions options; absl::InstallFailureSignalHandler(options); } - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); absl::SetProgramUsageMessage(absl::StrCat( "Key Value Server Request Simulation System. Sample usage:\n", argv[0])); kv_server::RequestSimulationSystem::InitializeTelemetry(); diff --git a/tools/request_simulation/metrics_collector.cc b/tools/request_simulation/metrics_collector.cc index efe882d5..fdfae0ae 100644 --- a/tools/request_simulation/metrics_collector.cc +++ b/tools/request_simulation/metrics_collector.cc @@ -16,7 +16,7 @@ #include -#include "glog/logging.h" +#include "absl/log/log.h" ABSL_FLAG(absl::Duration, metrics_report_interval, absl::Minutes(1), "The interval for reporting metrics"); @@ -41,7 +41,7 @@ MetricsCollector::MetricsCollector( requests_with_error_response_per_interval_(0), report_interval_(std::move(absl::GetFlag(FLAGS_metrics_report_interval))), report_thread_manager_( - TheadManager::Create("Metrics periodic report thread")), + ThreadManager::Create("Metrics periodic report thread")), metrics_recorder_(metrics_recorder), sleep_for_(std::move(sleep_for)) { histogram_per_interval_ = grpc_histogram_create(kDefaultHistogramResolution, diff --git a/tools/request_simulation/metrics_collector.h b/tools/request_simulation/metrics_collector.h index 77becc19..49721a15 100644 --- a/tools/request_simulation/metrics_collector.h +++ b/tools/request_simulation/metrics_collector.h @@ -22,7 +22,7 @@ #include "absl/flags/flag.h" #include "components/data/common/thread_manager.h" #include "components/util/sleepfor.h" -#include "src/cpp/telemetry/metrics_recorder.h" +#include "src/telemetry/metrics_recorder.h" #include "test/core/util/histogram.h" namespace kv_server { @@ -77,7 +77,7 @@ class MetricsCollector { mutable std::atomic requests_with_ok_response_per_interval_; mutable std::atomic requests_with_error_response_per_interval_; absl::Duration report_interval_; - std::unique_ptr report_thread_manager_; + std::unique_ptr report_thread_manager_; privacy_sandbox::server_common::MetricsRecorder& metrics_recorder_; std::unique_ptr sleep_for_; grpc_histogram* histogram_per_interval_ ABSL_GUARDED_BY(mutex_); diff --git a/tools/request_simulation/metrics_collector_test.cc b/tools/request_simulation/metrics_collector_test.cc index 0b88b22a..014e6708 100644 --- a/tools/request_simulation/metrics_collector_test.cc +++ b/tools/request_simulation/metrics_collector_test.cc @@ -21,7 +21,7 @@ #include "components/util/sleepfor_mock.h" #include "gtest/gtest.h" -#include "src/cpp/telemetry/mocks.h" +#include "src/telemetry/mocks.h" namespace kv_server { diff --git a/tools/request_simulation/rate_limiter.h b/tools/request_simulation/rate_limiter.h index 67420580..9d0fcfb3 100644 --- a/tools/request_simulation/rate_limiter.h +++ b/tools/request_simulation/rate_limiter.h @@ -25,7 +25,7 @@ #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" #include "components/util/sleepfor.h" -#include "src/cpp/util/duration.h" +#include "src/util/duration.h" namespace kv_server { // A simple permit-based rate limiter. The permits are refilled at given rate diff --git a/tools/request_simulation/realtime_message_batcher.cc b/tools/request_simulation/realtime_message_batcher.cc index 61ea61c0..af20231d 100644 --- a/tools/request_simulation/realtime_message_batcher.cc +++ b/tools/request_simulation/realtime_message_batcher.cc @@ -19,7 +19,9 @@ #include #include -#include "src/cpp/util/status_macro/status_macros.h" +#include "absl/log/check.h" +#include "absl/strings/substitute.h" +#include "src/util/status_macro/status_macros.h" namespace kv_server { namespace { @@ -89,8 +91,9 @@ RealtimeMessage RealtimeMessageBatcher::GetMessage(int shard_num) const { shard_num_opt = shard_num; } RealtimeMessage rm{ - .message = std::string(std::istreambuf_iterator(file_stream), - std::istreambuf_iterator()), + .message = absl::Base64Escape( + std::string(std::istreambuf_iterator(file_stream), + std::istreambuf_iterator())), .shard_num = shard_num_opt, }; return rm; diff --git a/tools/request_simulation/realtime_message_batcher_test.cc b/tools/request_simulation/realtime_message_batcher_test.cc index c553610a..c5a486de 100644 --- a/tools/request_simulation/realtime_message_batcher_test.cc +++ b/tools/request_simulation/realtime_message_batcher_test.cc @@ -18,6 +18,7 @@ #include "absl/strings/match.h" #include "absl/strings/str_cat.h" +#include "absl/strings/substitute.h" #include "gtest/gtest.h" #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" @@ -27,7 +28,9 @@ namespace { absl::StatusOr> Convert( RealtimeMessage rt) { std::vector rows; - std::istringstream is(rt.message); + std::string string_decoded; + absl::Base64Unescape(rt.message, &string_decoded); + std::istringstream is(string_decoded); auto delta_stream_reader_factory = std::make_unique(); auto record_reader = delta_stream_reader_factory->CreateReader(is); diff --git a/tools/request_simulation/request_generation_util.cc b/tools/request_simulation/request_generation_util.cc index 16ec1388..8cd8f91a 100644 --- a/tools/request_simulation/request_generation_util.cc +++ b/tools/request_simulation/request_generation_util.cc @@ -29,8 +29,7 @@ constexpr std::string_view kKVV2KeyValueDSPRequestBodyFormat = R"json( std::vector GenerateRandomKeys(int number_of_keys, int key_size) { std::vector result; for (int i = 0; i < number_of_keys; ++i) { - result.push_back( - std::string(key_size, 'A' + (std::rand() % number_of_keys))); + result.push_back(std::string(key_size, 'A' + (std::rand() % 26))); } return result; } diff --git a/tools/request_simulation/request_simulation_parameter_fetcher.h b/tools/request_simulation/request_simulation_parameter_fetcher.h index bf7b118e..94475144 100644 --- a/tools/request_simulation/request_simulation_parameter_fetcher.h +++ b/tools/request_simulation/request_simulation_parameter_fetcher.h @@ -26,6 +26,7 @@ class RequestSimulationParameterFetcher { RequestSimulationParameterFetcher() = default; virtual ~RequestSimulationParameterFetcher() = default; virtual NotifierMetadata GetBlobStorageNotifierMetadata() const; + virtual NotifierMetadata GetRealtimeNotifierMetadata() const; }; } // namespace kv_server diff --git a/tools/request_simulation/request_simulation_parameter_fetcher_aws.cc b/tools/request_simulation/request_simulation_parameter_fetcher_aws.cc index 86f98e2f..16d71d5a 100644 --- a/tools/request_simulation/request_simulation_parameter_fetcher_aws.cc +++ b/tools/request_simulation/request_simulation_parameter_fetcher_aws.cc @@ -15,13 +15,17 @@ #include #include "absl/flags/flag.h" -#include "glog/logging.h" +#include "absl/log/log.h" #include "tools/request_simulation/request_simulation_parameter_fetcher.h" ABSL_FLAG(std::string, s3_bucket_sns_arn, "", "The Amazon Resource Name(ARN) for the SNS topic" "configured for the S3 bucket"); +ABSL_FLAG(std::string, realtime_sns_arn, "", + "The Amazon Resource Name(ARN) for the SNS topic" + "configured for the realtime updates"); + namespace kv_server { NotifierMetadata @@ -31,4 +35,11 @@ RequestSimulationParameterFetcher::GetBlobStorageNotifierMetadata() const { return AwsNotifierMetadata{"BlobNotifier_", std::move(bucket_sns_arn)}; } +NotifierMetadata +RequestSimulationParameterFetcher::GetRealtimeNotifierMetadata() const { + std::string realtime_sns_arn = absl::GetFlag(FLAGS_realtime_sns_arn); + LOG(INFO) << "The sns arn for s3 bucket is " << realtime_sns_arn; + return AwsNotifierMetadata{"QueueNotifier_", std::move(realtime_sns_arn)}; +} + } // namespace kv_server diff --git a/tools/request_simulation/request_simulation_parameter_fetcher_gcp.cc b/tools/request_simulation/request_simulation_parameter_fetcher_gcp.cc index 61623f2c..6581fcf2 100644 --- a/tools/request_simulation/request_simulation_parameter_fetcher_gcp.cc +++ b/tools/request_simulation/request_simulation_parameter_fetcher_gcp.cc @@ -18,10 +18,12 @@ #include #include "absl/flags/flag.h" -#include "glog/logging.h" +#include "absl/log/log.h" #include "tools/request_simulation/request_simulation_parameter_fetcher.h" ABSL_FLAG(std::string, delta_dir, "", "The local directory for delta files"); +ABSL_FLAG(std::string, realtime_delta_dir, "", + "The local directory for realtime delta files"); namespace kv_server { @@ -32,4 +34,11 @@ RequestSimulationParameterFetcher::GetBlobStorageNotifierMetadata() const { return LocalNotifierMetadata{.local_directory = std::move(directory)}; } +NotifierMetadata +RequestSimulationParameterFetcher::GetRealtimeNotifierMetadata() const { + std::string directory = absl::GetFlag(FLAGS_realtime_delta_dir); + LOG(INFO) << "The local realtim delta file directory is " << directory; + return LocalNotifierMetadata{.local_directory = std::move(directory)}; +} + } // namespace kv_server diff --git a/tools/request_simulation/request_simulation_parameter_fetcher_local.cc b/tools/request_simulation/request_simulation_parameter_fetcher_local.cc index 9fd87521..cbc6dcdb 100644 --- a/tools/request_simulation/request_simulation_parameter_fetcher_local.cc +++ b/tools/request_simulation/request_simulation_parameter_fetcher_local.cc @@ -15,10 +15,12 @@ #include #include "absl/flags/flag.h" -#include "glog/logging.h" +#include "absl/log/log.h" #include "tools/request_simulation/request_simulation_parameter_fetcher.h" ABSL_FLAG(std::string, delta_dir, "", "The local directory for delta files"); +ABSL_FLAG(std::string, realtime_delta_dir, "", + "The local directory for realtime delta files"); namespace kv_server { @@ -29,4 +31,11 @@ RequestSimulationParameterFetcher::GetBlobStorageNotifierMetadata() const { return LocalNotifierMetadata{.local_directory = std::move(directory)}; } +NotifierMetadata +RequestSimulationParameterFetcher::GetRealtimeNotifierMetadata() const { + std::string directory = absl::GetFlag(FLAGS_realtime_delta_dir); + LOG(INFO) << "The local realtim delta file directory is " << directory; + return LocalNotifierMetadata{.local_directory = std::move(directory)}; +} + } // namespace kv_server diff --git a/tools/request_simulation/request_simulation_system.cc b/tools/request_simulation/request_simulation_system.cc index 34a66a96..83a85095 100644 --- a/tools/request_simulation/request_simulation_system.cc +++ b/tools/request_simulation/request_simulation_system.cc @@ -15,16 +15,20 @@ #include "tools/request_simulation/request_simulation_system.h" #include +#include #include #include #include -#include "glog/logging.h" +#include "absl/log/log.h" +#include "components/tools/concurrent_publishing_engine.h" #include "grpcpp/grpcpp.h" #include "opentelemetry/sdk/resource/resource.h" #include "opentelemetry/sdk/resource/semantic_conventions.h" #include "public/data_loading/readers/riegeli_stream_record_reader_factory.h" #include "public/query/get_values.grpc.pb.h" +#include "src/util/status_macro/status_macros.h" +#include "tools/request_simulation/realtime_message_batcher.h" #include "tools/request_simulation/request/raw_request.pb.h" #include "tools/request_simulation/request_generation_util.h" #include "tools/request_simulation/request_simulation_parameter_fetcher.h" @@ -40,7 +44,7 @@ ABSL_FLAG(int, concurrency, 10, "Number of concurrent requests sent to the server," "this number will be limited by the maximum concurrent threads" "supported by state of the machine"); -ABSL_FLAG(absl::Duration, request_timeout, absl::Seconds(5), +ABSL_FLAG(absl::Duration, request_timeout, absl::Seconds(300), "The timeout duration for getting response for the request"); ABSL_FLAG(int64_t, synthetic_requests_fill_qps, 1000, "The per second rate of synthetic requests generated by the " @@ -81,6 +85,13 @@ ABSL_FLAG( ABSL_FLAG(int32_t, data_loading_num_threads, 1, "Number of parallel threads for reading and loading data files."); ABSL_FLAG(std::string, delta_file_bucket, "", "The name of delta file bucket"); +ABSL_FLAG(int32_t, num_shards, 1, "Number of shards on the kv-server"); +ABSL_FLAG(int32_t, realtime_message_size_kb, 10, + "Realtime message size threshold in kb"); +ABSL_FLAG(int32_t, realtime_publisher_insertion_num_threads, 1, + "Number of threads used to write to pubsub in parallel."); +ABSL_FLAG(int32_t, realtime_publisher_files_insertion_rate, 15, + "Number of messages sent per insertion thread to pubsub per second"); namespace kv_server { @@ -180,6 +191,7 @@ absl::Status RequestSimulationSystem::Init( } blob_storage_client_ = CreateBlobClient(); delta_file_notifier_ = CreateDeltaFileNotifier(); + realtime_delta_file_notifier_ = CreateDeltaFileNotifier(); delta_stream_reader_factory_ = CreateStreamRecordReaderFactory(); auto blob_notifier_meta_data = @@ -192,7 +204,7 @@ absl::Status RequestSimulationSystem::Init( SetQueueManager(blob_notifier_meta_data, message_service_blob_.get()); { auto status_or_notifier = - BlobStorageChangeNotifier::Create(std::move(blob_notifier_meta_data)); + BlobStorageChangeNotifier::Create(blob_notifier_meta_data); if (!status_or_notifier.ok()) { // The ChangeNotifier is required to read delta files, if it's not // available that's a critical error and so return immediately. @@ -202,7 +214,9 @@ absl::Status RequestSimulationSystem::Init( } blob_change_notifier_ = std::move(*status_or_notifier); } - + PS_ASSIGN_OR_RETURN( + realtime_blob_change_notifier_, + BlobStorageChangeNotifier::Create(std::move(blob_notifier_meta_data))); delta_based_request_generator_ = std::make_unique( DeltaBasedRequestGenerator::Options{ .data_bucket = absl::GetFlag(FLAGS_delta_file_bucket), @@ -212,6 +226,34 @@ absl::Status RequestSimulationSystem::Init( .change_notifier = *blob_change_notifier_, .delta_stream_reader_factory = *delta_stream_reader_factory_}, CreateRequestFromKeyFn(), metrics_recorder_); + PS_ASSIGN_OR_RETURN(realtime_message_batcher_, + RealtimeMessageBatcher::Create( + realtime_messages_, realtime_messages_mutex_, + absl::GetFlag(FLAGS_num_shards), + absl::GetFlag(FLAGS_realtime_message_size_kb))); + auto notifier_metadata = parameter_fetcher_->GetRealtimeNotifierMetadata(); + const int realtime_publisher_insertion_num_threads = + absl::GetFlag(FLAGS_realtime_publisher_insertion_num_threads); + const int realtime_publisher_files_insertion_rate = + absl::GetFlag(FLAGS_realtime_publisher_files_insertion_rate); + concurrent_publishing_engine_ = std::make_unique( + realtime_publisher_insertion_num_threads, std::move(notifier_metadata), + realtime_publisher_files_insertion_rate, realtime_messages_mutex_, + realtime_messages_); + delta_based_realtime_updates_publisher_ = + std::make_unique( + std::move(realtime_message_batcher_), + std::move(concurrent_publishing_engine_), + DeltaBasedRealtimeUpdatesPublisher::Options{ + .data_bucket = absl::GetFlag(FLAGS_delta_file_bucket), + .blob_client = *blob_storage_client_, + .delta_notifier = *realtime_delta_file_notifier_, + .change_notifier = *realtime_blob_change_notifier_, + .delta_stream_reader_factory = *delta_stream_reader_factory_, + .realtime_messages = realtime_messages_, + .realtime_messages_mutex = realtime_messages_mutex_, + }); + LOG(INFO) << "Request simulation system is initialized," "target server address is " << server_address_ << " and server method is " << server_method_; @@ -233,6 +275,17 @@ absl::Status RequestSimulationSystem::InitializeGrpcClientWorkers() { absl::GetFlag(FLAGS_server_auth_mode)); bool is_client_channel = absl::GetFlag(FLAGS_is_client_channel); + if (is_client_channel) { + RetryUntilOk( + [channel]() { + if (channel->GetState(true) != GRPC_CHANNEL_READY) { + return absl::UnavailableError("GRPC channel is disconnected"); + } + return absl::OkStatus(); + }, + "check grpc connection in start up", LogMetricsNoOpCallback()); + } + auto request_timeout = absl::GetFlag(FLAGS_request_timeout); for (int i = 0; i < num_of_workers; ++i) { auto request_converter = [](const std::string& request_body) { RawRequest request; @@ -241,7 +294,7 @@ absl::Status RequestSimulationSystem::InitializeGrpcClientWorkers() { }; auto worker = std::make_unique>( - i, channel, server_method_, absl::Seconds(1), request_converter, + i, channel, server_method_, request_timeout, request_converter, *message_queue_, *grpc_request_rate_limiter_, *metrics_collector_, is_client_channel); grpc_client_workers_.push_back(std::move(worker)); @@ -254,6 +307,8 @@ absl::Status RequestSimulationSystem::Start() { if (auto status = delta_based_request_generator_->Start(); !status.ok()) { return status; } + LOG(INFO) << "Starting delta based realtime updates publisher"; + PS_RETURN_IF_ERROR(delta_based_realtime_updates_publisher_->Start()); if (synthetic_requests_fill_qps_ > 0) { LOG(INFO) << "Starting synthetic request generator"; if (auto status = synthetic_request_generator_->Start(); !status.ok()) { @@ -276,6 +331,8 @@ absl::Status RequestSimulationSystem::Start() { return absl::OkStatus(); } absl::Status RequestSimulationSystem::Stop() { + LOG(INFO) << "Stopping delta based realtime updates publisher"; + PS_RETURN_IF_ERROR(delta_based_realtime_updates_publisher_->Stop()); LOG(INFO) << "Stopping delta based request generator"; if (auto status = delta_based_request_generator_->Stop(); !status.ok()) { return status; diff --git a/tools/request_simulation/request_simulation_system.h b/tools/request_simulation/request_simulation_system.h index 43d7c1f7..0d7579d6 100644 --- a/tools/request_simulation/request_simulation_system.h +++ b/tools/request_simulation/request_simulation_system.h @@ -18,6 +18,7 @@ #define TOOLS_REQUEST_SIMULATION_REQUEST_SIMULATION_SYSTEM_H_ #include +#include #include #include #include @@ -33,10 +34,11 @@ #include "grpcpp/grpcpp.h" #include "public/data_loading/readers/riegeli_stream_io.h" #include "public/query/get_values.grpc.pb.h" -#include "src/cpp/telemetry/metrics_recorder.h" +#include "src/telemetry/metrics_recorder.h" #include "test/core/util/histogram.h" #include "tools/request_simulation/client_worker.h" #include "tools/request_simulation/delta_based_request_generator.h" +#include "tools/request_simulation/detla_based_realtime_updates_publisher.h" #include "tools/request_simulation/message_queue.h" #include "tools/request_simulation/rate_limiter.h" #include "tools/request_simulation/request/raw_request.pb.h" @@ -73,6 +75,8 @@ namespace kv_server { // 4. Client workers that send requests to the target server. // The number of client workers is determined by the given concurrency // parameter. +// 5. A delta based request generator that reads keys from delta file and +// publishes realtime updates to the specified SNS/pubsub endpoint. // // Once the system successfully starts, the system will continuously generates // requests and send requests to the target server. @@ -147,16 +151,25 @@ class RequestSimulationSystem { std::unique_ptr blob_storage_client_; std::unique_ptr message_service_blob_; std::unique_ptr blob_change_notifier_; + std::unique_ptr realtime_blob_change_notifier_; std::unique_ptr delta_file_notifier_; + std::unique_ptr realtime_delta_file_notifier_; std::unique_ptr delta_stream_reader_factory_; std::unique_ptr message_queue_; std::unique_ptr synthetic_request_generator_rate_limiter_; std::unique_ptr grpc_request_rate_limiter_; std::unique_ptr synthetic_request_generator_; std::unique_ptr delta_based_request_generator_; + std::unique_ptr + delta_based_realtime_updates_publisher_; std::vector>> grpc_client_workers_; std::unique_ptr parameter_fetcher_; + std::queue realtime_messages_; + absl::Mutex realtime_messages_mutex_; + std::unique_ptr realtime_message_batcher_; + std::unique_ptr concurrent_publishing_engine_; + bool is_running; friend class RequestSimulationSystemTestPeer; }; diff --git a/tools/request_simulation/request_simulation_system_local_test.cc b/tools/request_simulation/request_simulation_system_local_test.cc index 62843058..30b3eaf7 100644 --- a/tools/request_simulation/request_simulation_system_local_test.cc +++ b/tools/request_simulation/request_simulation_system_local_test.cc @@ -23,8 +23,8 @@ #include "grpcpp/grpcpp.h" #include "gtest/gtest.h" #include "public/testing/fake_key_value_service_impl.h" -#include "src/cpp/telemetry/mocks.h" -#include "src/cpp/util/duration.h" +#include "src/telemetry/mocks.h" +#include "src/util/duration.h" #include "tools/request_simulation/mocks.h" #include "tools/request_simulation/request_simulation_system.h" @@ -54,6 +54,7 @@ class MockRequestSimulationParameterFetcher : public RequestSimulationParameterFetcher { public: MOCK_METHOD(NotifierMetadata, GetBlobStorageNotifierMetadata, (), (const)); + MOCK_METHOD(NotifierMetadata, GetRealtimeNotifierMetadata, (), (const)); }; namespace { @@ -164,6 +165,10 @@ TEST_F(SimulationSystemTest, TestSimulationSystemRunning) { GetBlobStorageNotifierMetadata()) .WillRepeatedly(Return( LocalNotifierMetadata{.local_directory = ::testing::TempDir()})); + EXPECT_CALL(*mock_request_simulation_parameter_fetcher_, + GetRealtimeNotifierMetadata()) + .WillRepeatedly(Return( + LocalNotifierMetadata{.local_directory = ::testing::TempDir()})); RequestSimulationSystem system( metrics_recorder_, sim_clock_, channel_creation_fn, std::move(mock_request_simulation_parameter_fetcher_)); diff --git a/tools/request_simulation/synthetic_request_generator.h b/tools/request_simulation/synthetic_request_generator.h index 97aae750..e4b9a2e0 100644 --- a/tools/request_simulation/synthetic_request_generator.h +++ b/tools/request_simulation/synthetic_request_generator.h @@ -44,7 +44,7 @@ class SyntheticRequestGenerator { MessageQueue& message_queue, RateLimiter& rate_limiter, std::unique_ptr sleep_for, int64_t requests_fill_qps, absl::AnyInvocable request_body_generation_fn) - : thread_manager_(TheadManager::Create("Synthetic request generator")), + : thread_manager_(ThreadManager::Create("Synthetic request generator")), message_queue_(message_queue), rate_limiter_(rate_limiter), sleep_for_(std::move(sleep_for)), @@ -65,7 +65,7 @@ class SyntheticRequestGenerator { private: // The actual function that generates requests void GenerateRequests(); - std::unique_ptr thread_manager_; + std::unique_ptr thread_manager_; kv_server::MessageQueue& message_queue_; RateLimiter& rate_limiter_; std::unique_ptr sleep_for_; diff --git a/tools/server_diagnostic/helloworld_server/BUILD.bazel b/tools/server_diagnostic/helloworld_server/BUILD.bazel index 81af4c4b..a78f154d 100644 --- a/tools/server_diagnostic/helloworld_server/BUILD.bazel +++ b/tools/server_diagnostic/helloworld_server/BUILD.bazel @@ -23,11 +23,11 @@ cc_binary( deps = [ "//public/query:get_values_cc_grpc", "//public/query/v2:get_values_v2_cc_grpc", - "@com_github_google_glog//:glog", "@com_github_grpc_grpc//:grpc++", "@com_github_grpc_grpc//:grpc++_reflection", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", ], ) diff --git a/tools/server_diagnostic/helloworld_server/helloworld.cc b/tools/server_diagnostic/helloworld_server/helloworld.cc index 9474576a..4a110cae 100644 --- a/tools/server_diagnostic/helloworld_server/helloworld.cc +++ b/tools/server_diagnostic/helloworld_server/helloworld.cc @@ -13,7 +13,7 @@ // limitations under the License. #include "absl/flags/flag.h" -#include "glog/logging.h" +#include "absl/log/log.h" #include "grpcpp/ext/proto_server_reflection_plugin.h" #include "grpcpp/grpcpp.h" #include "public/query/get_values.grpc.pb.h" diff --git a/tools/serving_data_generator/BUILD.bazel b/tools/serving_data_generator/BUILD.bazel index 245f8e75..99f7e36d 100644 --- a/tools/serving_data_generator/BUILD.bazel +++ b/tools/serving_data_generator/BUILD.bazel @@ -14,6 +14,10 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") +package(default_visibility = [ + "//tools:__subpackages__", +]) + cc_binary( name = "test_serving_data_generator", srcs = ["test_serving_data_generator.cc"], @@ -23,9 +27,11 @@ cc_binary( "//public/data_loading:records_utils", "//public/data_loading:riegeli_metadata_cc_proto", "//public/sharding:sharding_function", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/log:initialize", "@com_google_absl//absl/strings", "@com_google_riegeli//riegeli/bytes:ostream_writer", "@com_google_riegeli//riegeli/records:record_writer", diff --git a/tools/serving_data_generator/test_serving_data_generator.cc b/tools/serving_data_generator/test_serving_data_generator.cc index eb5aa9fb..aebae233 100644 --- a/tools/serving_data_generator/test_serving_data_generator.cc +++ b/tools/serving_data_generator/test_serving_data_generator.cc @@ -18,8 +18,10 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/flags.h" +#include "absl/log/initialize.h" +#include "absl/log/log.h" #include "absl/strings/substitute.h" -#include "glog/logging.h" #include "google/protobuf/text_format.h" #include "public/data_loading/data_loading_generated.h" #include "public/data_loading/filename_utils.h" @@ -29,18 +31,22 @@ #include "riegeli/bytes/ostream_writer.h" #include "riegeli/records/record_writer.h" -ABSL_FLAG(std::string, key, "foo", "Specify the key for lookups"); +ABSL_FLAG(std::string, key, "foo", "Specify the key prefix for lookups"); ABSL_FLAG(int, value_size, 10, "Specify the size of value for the key"); ABSL_FLAG(std::string, output_dir, "", "Output file directory"); +ABSL_FLAG(std::string, output_filename, "", "Output file name"); ABSL_FLAG(int, num_records, 5, "Number of records to generate"); ABSL_FLAG(int, num_shards, 1, "Number of shards"); ABSL_FLAG(int, shard_number, 0, "Shard number"); ABSL_FLAG(int64_t, timestamp, absl::ToUnixMicros(absl::Now()), - "Record timestamp"); + "Record timestamp. Increases by 1 for each record."); ABSL_FLAG(bool, generate_set_record, false, "Whether to generate set record or not"); +ABSL_FLAG(std::string, set_value_key, "bar", + "Specify the set value key prefix for lookups"); ABSL_FLAG(int, num_values_in_set, 10, "Number of values in the set to generate"); +ABSL_FLAG(int, num_set_records, 5, "Number of records to generate"); using kv_server::DataRecordStruct; using kv_server::KeyValueMutationRecordStruct; @@ -51,16 +57,27 @@ using kv_server::ToDeltaFileName; using kv_server::ToFlatBufferBuilder; using kv_server::ToStringView; -void WriteKeyValueRecords(std::string_view key, int value_size, - riegeli::RecordWriterBase& writer) { - const int repetition = absl::GetFlag(FLAGS_num_records); - int64_t timestamp = absl::GetFlag(FLAGS_timestamp); +void WriteKeyValueRecord(std::string_view key, std::string_view value, + int64_t logical_commit_time, + riegeli::RecordWriterBase& writer) { + auto kv_record = KeyValueMutationRecordStruct{ + KeyValueMutationType::Update, logical_commit_time, key, value}; + writer.WriteRecord(ToStringView( + ToFlatBufferBuilder(DataRecordStruct{.record = std::move(kv_record)}))); +} + +std::vector WriteKeyValueRecords( + std::string_view key, int value_size, int64_t timestamp, + riegeli::RecordWriterBase& writer) { + const int num_records = absl::GetFlag(FLAGS_num_records); const int64_t num_shards = absl::GetFlag(FLAGS_num_shards); const int64_t current_shard_number = absl::GetFlag(FLAGS_shard_number); + std::vector keys; std::string query(" "); - for (int i = 0; i < repetition; ++i) { + for (int i = 0; i < num_records; ++i) { const std::string value(value_size, 'A' + (i % 50)); const std::string actual_key = absl::StrCat(key, i); + keys.emplace_back(actual_key); if (num_shards > 1) { kv_server::ShardingFunction sharding_func(""); auto shard_number = @@ -69,39 +86,39 @@ void WriteKeyValueRecords(std::string_view key, int value_size, continue; } } - auto kv_record = KeyValueMutationRecordStruct{ - KeyValueMutationType::Update, timestamp++, actual_key, value}; - writer.WriteRecord(ToStringView( - ToFlatBufferBuilder(DataRecordStruct{.record = std::move(kv_record)}))); + WriteKeyValueRecord(actual_key, value, timestamp++, writer); absl::StrAppend(&query, "\"", actual_key, "\"", ", "); } LOG(INFO) << "Print keys to query " << query; LOG(INFO) << "write done"; + return keys; } -void WriteKeyValueSetRecords(std::string_view key, int value_size, +void WriteKeyValueSetRecords(const std::vector& keys, + std::string_view set_value_key_prefix, + int64_t timestamp, riegeli::RecordWriterBase& writer) { - const int repetition = absl::GetFlag(FLAGS_num_records); - int64_t timestamp = absl::GetFlag(FLAGS_timestamp); + const int num_set_records = absl::GetFlag(FLAGS_num_set_records); const int num_values_in_set = absl::GetFlag(FLAGS_num_values_in_set); + const int keys_max_index = keys.size() - 1; std::string query(" "); - for (int i = 0; i < repetition; ++i) { + for (int i = 0; i < num_set_records; ++i) { std::vector set_copy; for (int j = 0; j < num_values_in_set; ++j) { - const std::string value(value_size, 'A' + (j % 50)); - set_copy.emplace_back( - absl::StrCat(value, std::to_string(std::rand() % num_values_in_set))); + // Add a random element from keys + set_copy.emplace_back(keys[std::rand() % keys_max_index]); } std::vector set; for (const auto& v : set_copy) { set.emplace_back(v); } - absl::StrAppend(&query, absl::StrCat(key, i), " | "); + std::string set_value_key = absl::StrCat(set_value_key_prefix, i); + absl::StrAppend(&query, set_value_key, " | "); KeyValueMutationRecordStruct record; record.value = set; record.mutation_type = KeyValueMutationType::Update; record.logical_commit_time = timestamp++; - record.key = absl::StrCat(key, i); + record.key = set_value_key; writer.WriteRecord(ToStringView( ToFlatBufferBuilder(DataRecordStruct{.record = std::move(record)}))); } @@ -123,12 +140,15 @@ KVFileMetadata GetKVFileMetadata() { int main(int argc, char** argv) { const std::vector commands = absl::ParseCommandLine(argc, argv); - google::InitGoogleLogging(argv[0]); + absl::InitializeLog(); const std::string output_dir = absl::GetFlag(FLAGS_output_dir); + std::string output_filename = absl::GetFlag(FLAGS_output_filename); auto write_records = [](std::ostream* os) { const std::string key = absl::GetFlag(FLAGS_key); const int value_size = absl::GetFlag(FLAGS_value_size); + const std::string set_value_key_prefix = absl::GetFlag(FLAGS_set_value_key); + int64_t timestamp = absl::GetFlag(FLAGS_timestamp); auto os_writer = riegeli::OStreamWriter(os); riegeli::RecordWriterBase::Options options; @@ -138,34 +158,39 @@ int main(int argc, char** argv) { GetKVFileMetadata(); options.set_metadata(std::move(metadata)); auto record_writer = riegeli::RecordWriter(std::move(os_writer), options); + const auto keys = + WriteKeyValueRecords(key, value_size, timestamp, record_writer); if (absl::GetFlag(FLAGS_generate_set_record)) { - WriteKeyValueSetRecords(key, value_size, record_writer); - } else { - WriteKeyValueRecords(key, value_size, record_writer); + timestamp += keys.size(); + WriteKeyValueSetRecords(keys, set_value_key_prefix, timestamp, + record_writer); } - record_writer.Close(); }; if (output_dir == "-") { LOG(INFO) << "Writing records to console"; - write_records(&std::cout); - } else { + return 0; + } + + if (output_filename.empty()) { absl::Time now = absl::Now(); if (const auto maybe_name = ToDeltaFileName(absl::ToUnixMicros(now)); !maybe_name.ok()) { LOG(ERROR) << "Unable to construct file name: " << maybe_name.status(); return -1; } else { - const std::string outfile = - absl::StrCat(output_dir, "/", maybe_name.value()); - LOG(INFO) << "Writing records to " << outfile; - - std::ofstream ofs(outfile); - write_records(&ofs); - ofs.close(); + output_filename = *maybe_name; } } + + const std::string outfile = absl::StrCat(output_dir, "/", output_filename); + LOG(INFO) << "Writing records to " << outfile; + + std::ofstream ofs(outfile); + write_records(&ofs); + ofs.close(); + return 0; } diff --git a/tools/udf/closure_js/closure_to_delta.bzl b/tools/udf/closure_js/closure_to_delta.bzl index 81f49975..5ebc1094 100644 --- a/tools/udf/closure_js/closure_to_delta.bzl +++ b/tools/udf/closure_js/closure_to_delta.bzl @@ -20,6 +20,7 @@ def closure_to_delta( closure_js_library_target, custom_udf_js_handler = "HandleRequest", output_file_name = "DELTA_0000000000000009", + logical_commit_time = None, # Not passing a logical_commit_time will default to now. udf_tool = "//tools/udf/udf_generator:udf_delta_file_generator", tags = ["manual"]): """Generate a JS UDF delta file from a given closure_js_library target and put it under dist/ @@ -30,6 +31,7 @@ def closure_to_delta( closure_js_library_target = ":my_js_lib", custom_udf_js_handler = "MyHandlerName", output_file_name = "DELTA_0000000000000009", + logical_commit_time="123123123", ) Args: @@ -39,6 +41,7 @@ def closure_to_delta( output_file_name: Name of UDF delta file output udf_tool: BUILD target for the udf_delta_file_generator. Defaults to `//tools/udf/udf_generator:udf_delta_file_generator` + logical_commit_time: Logical commit timestamp for UDF config. Optional, defaults to now. tags: tags to propagate to rules """ closure_js_binary( @@ -50,6 +53,8 @@ def closure_to_delta( tags = tags, ) + logical_commit_time_args = [] if logical_commit_time == None else ["--logical_commit_time", logical_commit_time] + run_binary( name = "{}_udf_delta".format(name), srcs = [ @@ -65,7 +70,7 @@ def closure_to_delta( "$(location {})".format(output_file_name), "--udf_handler_name", custom_udf_js_handler, - ], + ] + logical_commit_time_args, tool = udf_tool, visibility = ["//visibility:private"], tags = tags, @@ -78,8 +83,8 @@ def closure_to_delta( ], outs = ["{}_copy_to_dist.bin".format(name)], cmd_bash = """cat << EOF > '$@' -mkdir -p dist -cp $(location {}_udf_delta) dist +mkdir -p dist/deltas +cp $(location {}_udf_delta) dist/deltas builders/tools/normalize-dist EOF""".format(name), executable = True, diff --git a/tools/udf/closure_js/examples/get_values_binary/udf.js b/tools/udf/closure_js/examples/get_values_binary/udf.js index 8980f573..63c63ba3 100644 --- a/tools/udf/closure_js/examples/get_values_binary/udf.js +++ b/tools/udf/closure_js/examples/get_values_binary/udf.js @@ -56,6 +56,33 @@ function getKeyGroupOutputs(udf_arguments) { return keyGroupOutputs; } +/** + * @param {!Array} udf_arguments + * @suppress {reportUnknownTypes} + * @return {Object} + */ +function handlePas(udf_arguments) { + if (udf_arguments.length != 1) { + const error_message = + 'For PAS default UDF exactly one argument should be provided, but was provided ' + udf_arguments.length; + console.error(error_message); + throw new Error(error_message); + } + var serializedGetValuesBinary = /** @type {!Array} */ (getValuesBinary(udf_arguments[0])); + var getValuesBinaryProto = proto.kv_server.BinaryGetValuesResponse.deserializeBinary(serializedGetValuesBinary); + return getValuesBinaryProto.getKvPairsMap(); +} + +/** + * @param {!Array} udf_arguments + * @suppress {reportUnknownTypes} + * @return {Object} + */ +function handlePa(udf_arguments) { + const keyGroupOutputs = getKeyGroupOutputs(udf_arguments); + return { keyGroupOutputs, udfOutputApiVersion: 1 }; +} + /** * Entry point for code snippet execution. * @@ -71,11 +98,15 @@ function getKeyGroupOutputs(udf_arguments) { * from being minified. * * @export + * @suppress {reportUnknownTypes} * @param {!Object} executionMetadata * @param {...?} udf_arguments * @return {Object} */ function HandleRequest(executionMetadata, ...udf_arguments) { - const keyGroupOutputs = getKeyGroupOutputs(udf_arguments); - return { keyGroupOutputs: keyGroupOutputs, udfOutputApiVersion: 1 }; + if (executionMetadata.requestMetadata && executionMetadata.requestMetadata.is_pas) { + console.log('Executing PAS branch'); + return handlePas(udf_arguments); + } + return handlePa(udf_arguments); } diff --git a/tools/udf/inline_wasm/examples/get_values_binary_proto/my_udf.js b/tools/udf/inline_wasm/examples/get_values_binary_proto/my_udf.js index fbb998e6..285d4bbd 100644 --- a/tools/udf/inline_wasm/examples/get_values_binary_proto/my_udf.js +++ b/tools/udf/inline_wasm/examples/get_values_binary_proto/my_udf.js @@ -16,11 +16,11 @@ async function HandleRequest(executionMetadata, ...udf_arguments) { const module = await getModule(); - logMessage('Done loading WASM Module'); + console.log('Done loading WASM Module'); // Pass in the getValuesBinary function for the C++ code to call. // getValuesBinary returns a Uint8Array, which emscripten converts to std::string const result = module.handleRequestCc(getValuesBinary, udf_arguments); - logMessage('handleRequestCc result: ' + JSON.stringify(result)); + console.log('handleRequestCc result: ' + JSON.stringify(result)); return result; } diff --git a/tools/udf/inline_wasm/examples/hello_world/BUILD.bazel b/tools/udf/inline_wasm/examples/hello_world/BUILD.bazel index 60acdc5c..e1936042 100644 --- a/tools/udf/inline_wasm/examples/hello_world/BUILD.bazel +++ b/tools/udf/inline_wasm/examples/hello_world/BUILD.bazel @@ -36,7 +36,7 @@ inline_wasm_udf_delta( wasm_binary = ":hello.wasm", ) -# builders/tools/bazel-debian run \ +# builders/tools/bazel-debian run --config=emscripten \ # //tools/udf/inline_wasm/examples/hello_world:hello_delta_cc cc_inline_wasm_udf_delta( name = "hello_delta_cc", diff --git a/tools/udf/inline_wasm/examples/hello_world/my_udf.js b/tools/udf/inline_wasm/examples/hello_world/my_udf.js index 81a9a3c7..e14ef923 100644 --- a/tools/udf/inline_wasm/examples/hello_world/my_udf.js +++ b/tools/udf/inline_wasm/examples/hello_world/my_udf.js @@ -41,9 +41,9 @@ function getKeyGroupOutputs(udf_arguments, module) { } async function HandleRequest(executionMetadata, ...udf_arguments) { - logMessage('Handling request'); + console.log('Handling request'); const module = await getModule(); - logMessage('Done loading WASM Module'); + console.log('Done loading WASM Module'); const keyGroupOutputs = getKeyGroupOutputs(udf_arguments, module); return { keyGroupOutputs, udfOutputApiVersion: 1 }; } diff --git a/tools/udf/inline_wasm/examples/js_call/my_udf.js b/tools/udf/inline_wasm/examples/js_call/my_udf.js index 038a5503..a156d031 100644 --- a/tools/udf/inline_wasm/examples/js_call/my_udf.js +++ b/tools/udf/inline_wasm/examples/js_call/my_udf.js @@ -15,12 +15,12 @@ */ async function HandleRequest(executionMetadata, ...input) { - logMessage('Handling request'); + console.log('Handling request'); const module = await getModule(); - logMessage('Done loading WASM Module'); + console.log('Done loading WASM Module'); // Pass in the getValues function for the C++ code to call. const result = module.handleRequestCc(getValues, input); - logMessage('handleRequestCc result: ' + JSON.stringify(result)); + console.log('handleRequestCc result: ' + JSON.stringify(result)); return result; } diff --git a/tools/udf/inline_wasm/examples/protobuf/my_udf.js b/tools/udf/inline_wasm/examples/protobuf/my_udf.js index 89d97825..57a5145a 100644 --- a/tools/udf/inline_wasm/examples/protobuf/my_udf.js +++ b/tools/udf/inline_wasm/examples/protobuf/my_udf.js @@ -42,9 +42,9 @@ function getKeyGroupOutputs(udf_arguments, module) { } async function HandleRequest(executionMetadata, ...udf_arguments) { - logMessage('Handling request'); + console.log('Handling request'); const module = await getModule(); - logMessage('Done loading WASM Module'); + console.log('Done loading WASM Module'); const keyGroupOutputs = getKeyGroupOutputs(udf_arguments, module); return { keyGroupOutputs, udfOutputApiVersion: 1 }; } diff --git a/tools/udf/inline_wasm/wasm.bzl b/tools/udf/inline_wasm/wasm.bzl index 1deb4f3f..07d49b24 100644 --- a/tools/udf/inline_wasm/wasm.bzl +++ b/tools/udf/inline_wasm/wasm.bzl @@ -22,7 +22,7 @@ def inline_wasm_udf_delta( custom_udf_js, custom_udf_js_handler = "HandleRequest", output_file_name = "DELTA_0000000000000005", - logical_commit_time = "123123123", + logical_commit_time = None, udf_tool = "//tools/udf/udf_generator:udf_delta_file_generator", tags = ["manual"]): """Generate a JS + inline WASM UDF delta file and put it under dist/ directory @@ -52,8 +52,7 @@ def inline_wasm_udf_delta( output_file_name: Name of UDF delta file output. Recommended to follow DELTA file naming convention. Defaults to `DELTA_0000000000000005` - logical_commit_time: Logical commit timestamp for UDF config. - Defaults to `123123123`. + logical_commit_time: Logical commit timestamp for UDF config. Optional, defaults to now. udf_tool: build target for the udf_delta_file_generator. Defaults to `//tools/udf/udf_generator:udf_delta_file_generator` tags: tags to propagate to rules @@ -79,6 +78,8 @@ def inline_wasm_udf_delta( tags = tags, ) + logical_commit_time_args = [] if logical_commit_time == None else ["--logical_commit_time", logical_commit_time] + run_binary( name = "{}_udf_delta".format(name), srcs = [ @@ -92,11 +93,9 @@ def inline_wasm_udf_delta( "$(location {}_generated)".format(name), "--output_path", "$(location {})".format(output_file_name), - "--logical_commit_time", - logical_commit_time, "--udf_handler_name", custom_udf_js_handler, - ], + ] + logical_commit_time_args, tool = udf_tool, visibility = ["//visibility:private"], tags = tags, @@ -110,11 +109,12 @@ def inline_wasm_udf_delta( ], outs = ["{}_copy_to_dist.bin".format(name)], cmd_bash = """cat << EOF > '$@' -mkdir -p dist/debian -cp $(location {}_udf_delta) dist -cp $(location {}_generated) dist +mkdir -p dist/deltas +mkdir -p dist/udfs +cp $(location {name}_udf_delta) dist/deltas +cp $(location {name}_generated) dist/udfs builders/tools/normalize-dist -EOF""".format(name, name), +EOF""".format(name = name), executable = True, local = True, message = "Copying {} dist directory".format(output_file_name), @@ -127,7 +127,7 @@ def cc_inline_wasm_udf_delta( custom_udf_js, custom_udf_js_handler = "HandleRequest", output_file_name = "DELTA_0000000000000005", - logical_commit_time = "123123123", + logical_commit_time = None, udf_tool = "//tools/udf/udf_generator:udf_delta_file_generator", deps = [], tags = ["manual"], @@ -167,7 +167,6 @@ def cc_inline_wasm_udf_delta( Recommended to follow DELTA file naming convention. Defaults to `DELTA_0000000000000005` logical_commit_time: Logical commit timestamp for UDF config. - Defaults to `123123123`. udf_tool: build target for the udf_delta_file_generator. Defaults to `//tools/udf/udf_generator:udf_delta_file_generator` tags: tags to propagate to rules diff --git a/tools/udf/sample_udf/udf.js b/tools/udf/sample_udf/udf.js index d840e1d2..d7cbc7ef 100644 --- a/tools/udf/sample_udf/udf.js +++ b/tools/udf/sample_udf/udf.js @@ -38,7 +38,31 @@ function getKeyGroupOutputs(udf_arguments) { return keyGroupOutputs; } -function HandleRequest(executionMetadata, ...udf_arguments) { +function handlePas(udf_arguments) { + if (udf_arguments.length != 1) { + let error_message = + 'For PAS default UDF exactly one argument should be provided, but was provided ' + udf_arguments.length; + console.error(error_message); + return error_message; + } + const kv_result = JSON.parse(getValues(udf_arguments[0])); + if (kv_result.hasOwnProperty('kvPairs')) { + return kv_result.kvPairs; + } + let error_message_lookup = 'Failed looking up values'; + console.error(error_message_lookup); + return error_message_lookup; +} + +function handlePA(udf_arguments) { const keyGroupOutputs = getKeyGroupOutputs(udf_arguments); return { keyGroupOutputs, udfOutputApiVersion: 1 }; } + +function HandleRequest(executionMetadata, ...udf_arguments) { + if (executionMetadata.requestMetadata && executionMetadata.requestMetadata.is_pas) { + console.log('Executing PAS branch'); + return handlePas(udf_arguments); + } + return handlePA(udf_arguments); +} diff --git a/tools/udf/udf_generator/BUILD.bazel b/tools/udf/udf_generator/BUILD.bazel index 17b0296a..78e56bf3 100644 --- a/tools/udf/udf_generator/BUILD.bazel +++ b/tools/udf/udf_generator/BUILD.bazel @@ -15,6 +15,7 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = [ + "//docs/protected_app_signals:__subpackages__", "//getting_started:__subpackages__", "//production/packaging/tools:__subpackages__", "//testing:__subpackages__", @@ -40,9 +41,9 @@ cc_binary( "//public/data_loading/writers:delta_record_stream_writer", "//public/data_loading/writers:delta_record_writer", "//public/udf:constants", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", "@com_google_riegeli//riegeli/bytes:ostream_writer", "@com_google_riegeli//riegeli/records:record_writer", diff --git a/tools/udf/udf_generator/udf_delta_file_generator.cc b/tools/udf/udf_generator/udf_delta_file_generator.cc index c37fed04..b4e0e559 100644 --- a/tools/udf/udf_generator/udf_delta_file_generator.cc +++ b/tools/udf/udf_generator/udf_delta_file_generator.cc @@ -20,8 +20,8 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/log.h" #include "absl/strings/substitute.h" -#include "glog/logging.h" #include "google/protobuf/text_format.h" #include "public/constants.h" #include "public/data_loading/data_loading_generated.h" @@ -40,8 +40,8 @@ ABSL_FLAG(std::string, output_path, "", "Output path. If specified, output_dir is ignored. If '-', output is " "written to " "console."); -ABSL_FLAG(int64_t, logical_commit_time, 123123123, - "Record logical_commit_time. Default is 123123123."); +ABSL_FLAG(int64_t, logical_commit_time, absl::ToUnixMicros(absl::Now()), + "Record logical_commit_time. Default is current timestamp."); ABSL_FLAG(int64_t, code_snippet_version, 2, "UDF version. Default is 2."); ABSL_FLAG(std::string, data_loading_file_format, std::string(kv_server::kFileFormats[static_cast( diff --git a/tools/udf/udf_tester/BUILD.bazel b/tools/udf/udf_tester/BUILD.bazel index 98d1b458..de2873b8 100644 --- a/tools/udf/udf_tester/BUILD.bazel +++ b/tools/udf/udf_tester/BUILD.bazel @@ -33,11 +33,11 @@ cc_binary( "//public/data_loading/readers:delta_record_stream_reader", "//public/query/v2:get_values_v2_cc_proto", "//public/udf:constants", - "@com_github_google_glog//:glog", "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", "@com_google_protobuf//:protobuf", - "@google_privacysandbox_servers_common//src/cpp/util/status_macro:status_macros", + "@google_privacysandbox_servers_common//src/util/status_macro:status_macros", ], ) diff --git a/tools/udf/udf_tester/README.md b/tools/udf/udf_tester/README.md index cb868663..3496d3ea 100644 --- a/tools/udf/udf_tester/README.md +++ b/tools/udf/udf_tester/README.md @@ -4,7 +4,8 @@ This binary directly invokes UDF which can access data input. It requires two delta files: -- Delta file with key-value pairs to be stored in memory ([docs](/docs/loading_data.md)) +- Delta file with key-value pairs to be stored in memory + ([docs](/docs/data_loading/loading_data.md)) - Delta file with the UDF configuration ([docs](/docs/generating_udf_files.md)). ## Flags: diff --git a/tools/udf/udf_tester/udf_delta_file_tester.cc b/tools/udf/udf_tester/udf_delta_file_tester.cc index ea498e95..d7e178fa 100644 --- a/tools/udf/udf_tester/udf_delta_file_tester.cc +++ b/tools/udf/udf_tester/udf_delta_file_tester.cc @@ -16,6 +16,7 @@ #include "absl/flags/flag.h" #include "absl/flags/parse.h" +#include "absl/log/log.h" #include "absl/strings/substitute.h" #include "components/data_server/cache/cache.h" #include "components/data_server/cache/key_value_cache.h" @@ -23,15 +24,13 @@ #include "components/udf/hooks/get_values_hook.h" #include "components/udf/udf_client.h" #include "components/udf/udf_config_builder.h" -#include "glog/logging.h" #include "google/protobuf/util/json_util.h" #include "public/data_loading/data_loading_generated.h" #include "public/data_loading/readers/delta_record_stream_reader.h" #include "public/query/v2/get_values_v2.pb.h" #include "public/udf/constants.h" -#include "src/cpp/telemetry/metrics_recorder.h" -#include "src/cpp/telemetry/telemetry_provider.h" -#include "src/cpp/util/status_macro/status_macros.h" +#include "src/telemetry/telemetry_provider.h" +#include "src/util/status_macro/status_macros.h" ABSL_FLAG(std::string, kv_delta_file_path, "", "Path to delta file with KV pairs."); @@ -138,9 +137,9 @@ void ShutdownUdf(UdfClient& udf_client) { absl::Status TestUdf(const std::string& kv_delta_file_path, const std::string& udf_delta_file_path, const std::string& input_arguments) { + InitMetricsContextMap(); LOG(INFO) << "Loading cache from delta file: " << kv_delta_file_path; - auto noop_metrics_recorder = MetricsRecorder::CreateNoop(); - std::unique_ptr cache = KeyValueCache::Create(*noop_metrics_recorder); + std::unique_ptr cache = KeyValueCache::Create(); PS_RETURN_IF_ERROR(LoadCacheFromFile(kv_delta_file_path, *cache)) << "Error loading cache from file"; @@ -154,20 +153,18 @@ absl::Status TestUdf(const std::string& kv_delta_file_path, UdfConfigBuilder config_builder; auto string_get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kString); - string_get_values_hook->FinishInit( - CreateLocalLookup(*cache, *noop_metrics_recorder)); + string_get_values_hook->FinishInit(CreateLocalLookup(*cache)); auto binary_get_values_hook = GetValuesHook::Create(GetValuesHook::OutputType::kBinary); - binary_get_values_hook->FinishInit( - CreateLocalLookup(*cache, *noop_metrics_recorder)); + binary_get_values_hook->FinishInit(CreateLocalLookup(*cache)); auto run_query_hook = RunQueryHook::Create(); - run_query_hook->FinishInit(CreateLocalLookup(*cache, *noop_metrics_recorder)); + run_query_hook->FinishInit(CreateLocalLookup(*cache)); absl::StatusOr> udf_client = UdfClient::Create(std::move( config_builder.RegisterStringGetValuesHook(*string_get_values_hook) .RegisterBinaryGetValuesHook(*binary_get_values_hook) .RegisterRunQueryHook(*run_query_hook) - .RegisterLoggingHook() + .RegisterLoggingFunction() .SetNumberOfWorkers(1) .Config())); PS_RETURN_IF_ERROR(udf_client.status()) @@ -188,8 +185,9 @@ absl::Status TestUdf(const std::string& kv_delta_file_path, JsonStringToMessage(req_partition_json, &req_partition); LOG(INFO) << "Calling UDF for partition: " << req_partition.DebugString(); - auto udf_result = - udf_client.value()->ExecuteCode({}, req_partition.arguments()); + auto metrics_context = std::make_unique(); + auto udf_result = udf_client.value()->ExecuteCode( + RequestContext(*metrics_context), {}, req_partition.arguments()); if (!udf_result.ok()) { LOG(ERROR) << "UDF execution failed: " << udf_result.status(); ShutdownUdf(*udf_client.value()); diff --git a/version.txt b/version.txt index 7092c7c4..d183d4ac 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -0.15.0 \ No newline at end of file +0.16.0 \ No newline at end of file