snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)
Some checks failed
codecov / upload (push) Failing after 13m59s
CodeQL / Analyze (push) Failing after 1s
govulncheck / govulncheck (.) (push) Failing after 1s
govulncheck / govulncheck (./cmd/protoc-gen-go-grpc) (push) Failing after 1s
govulncheck / govulncheck (./gcp/observability) (push) Failing after 1s
govulncheck / govulncheck (./security/advancedtls) (push) Failing after 1s
govulncheck / govulncheck (./stats/opencensus) (push) Failing after 1s
Testing / vet-proto (push) Failing after 1s
Testing / tests (arm, latest) (push) Waiting to run
Testing / extras (latest) (push) Failing after 1s
Testing / static checks (latest-1) (push) Failing after 1s
Testing / tests (-race, latest) (push) Failing after 1s
Testing / tests (i386, latest) (push) Failing after 1s
Testing / tests (latest) (push) Failing after 1s
Testing / tests (latest-1) (push) Failing after 1s
Lock Threads / lock (push) Failing after 1s
Stale bot / stale (push) Failing after 1s

This commit is contained in:
upstream-mirror
2026-09-15 20:40:15 +05:00
commit e241b33f26
1365 changed files with 339570 additions and 0 deletions

13
.gemini/config.yaml Normal file
View File

@@ -0,0 +1,13 @@
have_fun: false
memory_config:
disabled: false
code_review:
disable: false
comment_severity_threshold: MEDIUM
max_review_comments: -1
pull_request_opened:
help: false
summary: false
code_review: false
include_drafts: false
ignore_patterns: []

16
.gemini/styleguide.md Normal file
View File

@@ -0,0 +1,16 @@
# Code Review Guidelines
* Primary Style: Strictly adhere to the Google Go Style Guide
(https://google.github.io/styleguide/go/). Evaluate all pull requests against
these conventions.
* gRPC and Protobuf Imports: When importing generated gRPC service and Protobuf
message code from the same package, you **MUST** use separate named imports to
distinguish them. Alias the message code import with a `pb` suffix and the
service code import with a `grpc` suffix.
```go
import (
testgrpc "google.golang.org/grpc/interop/grpc_testing"
testpb "google.golang.org/grpc/interop/grpc_testing"
)
```

27
.github/ISSUE_TEMPLATE/bug.md vendored Normal file
View File

@@ -0,0 +1,27 @@
---
name: Bug Report
about: Report a non-security bug. For suspected security vulnerabilities or crashes, please use "Report a Security Vulnerability", below.
labels: 'Type: Bug'
---
NOTE: if you are reporting is a potential security vulnerability or a crash,
please follow our CVE process at
https://github.com/grpc/proposal/blob/master/P4-grpc-cve-process.md instead of
filing an issue here.
Please see the FAQ in our main README.md, then answer the questions below
before submitting your issue.
### What version of gRPC are you using?
### What version of Go are you using (`go version`)?
### What operating system (Linux, Windows, …) and version?
### What did you do?
If possible, provide a recipe for reproducing the error.
### What did you expect to see?
### What did you see instead?

16
.github/ISSUE_TEMPLATE/feature.md vendored Normal file
View File

@@ -0,0 +1,16 @@
---
name: Feature Request
about: Suggest an idea for gRPC-Go
labels: 'Type: Feature'
---
Please see the FAQ in our main README.md before submitting your issue.
### Use case(s) - what problem will this feature solve?
### Proposed Solution
### Alternatives Considered
### Additional Context

8
.github/ISSUE_TEMPLATE/question.md vendored Normal file
View File

@@ -0,0 +1,8 @@
---
name: Question
about: Ask a question about gRPC-Go
labels: 'Type: Question'
---
Please see the FAQ in our main README.md before submitting your issue.

25
.github/codecov.yml vendored Normal file
View File

@@ -0,0 +1,25 @@
coverage:
status:
project:
default:
informational: true
patch:
default:
informational: true
ignore:
# All 'pb.go's.
- "**/*.pb.go"
# Tests and test related files.
- "**/test"
- "**/testdata"
- "**/testutils"
- "benchmark"
- "interop"
# Other submodules.
- "cmd"
- "examples"
- "gcp"
- "security"
- "stats/opencensus"
comment:
layout: "header, diff, files"

4
.github/pull_request_template.md vendored Normal file
View File

@@ -0,0 +1,4 @@
Thank you for your PR. Please read and follow
https://github.com/grpc/grpc-go/blob/master/CONTRIBUTING.md, especially the
"Guidelines for Pull Requests" section, and then delete this text before
entering your PR description.

37
.github/workflows/codeql-analysis.yml vendored Normal file
View File

@@ -0,0 +1,37 @@
name: "CodeQL"
on:
push:
branches: [ master ]
schedule:
- cron: '24 20 * * 3'
permissions:
contents: read
jobs:
analyze:
name: Analyze
runs-on: ubuntu-latest
timeout-minutes: 30
permissions:
security-events: write
pull-requests: read
actions: read
strategy:
fail-fast: false
steps:
- name: Checkout repository
uses: actions/checkout@v4
# Initializes the CodeQL tools for scanning.
- name: Initialize CodeQL
uses: github/codeql-action/init@v2
with:
languages: go
- name: Perform CodeQL Analysis
uses: github/codeql-action/analyze@v2

26
.github/workflows/coverage.yml vendored Normal file
View File

@@ -0,0 +1,26 @@
name: codecov
on: [push, pull_request]
permissions:
contents: read
jobs:
upload:
runs-on: ubuntu-latest
steps:
- name: Install checkout
uses: actions/checkout@v4
- name: Install checkout
uses: actions/setup-go@v5
with:
go-version: "stable"
- name: Run coverage
run: go test -coverprofile=coverage.out -coverpkg=./... ./...
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
fail_ci_if_error: true

60
.github/workflows/deps.yml vendored Normal file
View File

@@ -0,0 +1,60 @@
name: Dependency Changes
# Trigger on PRs.
on:
pull_request:
permissions:
contents: read
jobs:
# Compare dependencies before and after this PR.
dependencies:
runs-on: ubuntu-latest
timeout-minutes: 10
strategy:
fail-fast: true
steps:
- name: Checkout repo
uses: actions/checkout@v4
with:
fetch-depth: 0
- name: Setup Go
uses: actions/setup-go@v5
with:
go-version: stable
cache-dependency-path: "**/*go.sum"
# Run the commands to generate dependencies before and after and compare.
- name: Compare dependencies
run: |
set -eu
TEMP_DIR="$(mktemp -d)"
# GITHUB_BASE_REF is set when the job is triggered by a PR.
TARGET_REF="${GITHUB_BASE_REF:-master}"
mkdir "${TEMP_DIR}/after"
scripts/gen-deps.sh "${TEMP_DIR}/after"
git checkout "origin/${TARGET_REF}"
mkdir "${TEMP_DIR}/before"
scripts/gen-deps.sh "${TEMP_DIR}/before"
echo -e " \nComparing dependencies..."
cd "${TEMP_DIR}"
# Run grep in a sub-shell since bash does not support ! in the middle of a pipe.
if diff -u0 -r "before" "after" | bash -c '! grep -v "@@"'; then
echo "No changes detected."
exit 0
fi
# Print packages in `after` but not `before`.
for x in $(ls -1 after | grep -vF "$(ls -1 before)"); do
echo -e " \nDependencies of new package $x:"
cat "after/$x"
done
echo -e " \nChanges detected; exiting with error."
exit 1

27
.github/workflows/govulncheck.yaml vendored Normal file
View File

@@ -0,0 +1,27 @@
name: govulncheck
on:
push:
pull_request:
schedule: # daily at 10:22 UTC
- cron: '22 10 * * *'
workflow_dispatch:
permissions:
contents: read
jobs:
govulncheck:
runs-on: ubuntu-latest
timeout-minutes: 10
strategy:
fail-fast: false
matrix:
work-dir: ['.',
'./gcp/observability',
'./cmd/protoc-gen-go-grpc',
'./security/advancedtls',
'./stats/opencensus']
steps:
- name: Run govulncheck
uses: golang/govulncheck-action@v1
with:
check-latest: true
work-dir: ${{ matrix.work-dir }}

22
.github/workflows/lock.yml vendored Normal file
View File

@@ -0,0 +1,22 @@
name: 'Lock Threads'
on:
workflow_dispatch:
schedule:
- cron: '22 1 * * *'
permissions:
contents: read
jobs:
lock:
runs-on: ubuntu-latest
permissions:
issues: write
pull-requests: write
steps:
- uses: dessant/lock-threads@v5
with:
github-token: ${{ github.token }}
issue-inactive-days: 180
pr-inactive-days: 180

55
.github/workflows/pr-validation.yml vendored Normal file
View File

@@ -0,0 +1,55 @@
name: PR Validation
on:
pull_request:
types: [opened, edited, synchronize, labeled, unlabeled, milestoned, demilestoned]
permissions:
contents: read
jobs:
validate:
name: Validate PR
runs-on: ubuntu-latest
steps:
- name: Validate Description
uses: actions/github-script@v6
with:
script: |
const body = context.payload.pull_request.body;
const requiredRegex = new RegExp('^RELEASE NOTES:\\s*([Nn][Oo][Nn][Ee]|[Nn]/[Aa]|\\n(\\*|-)\\s*.+)$', 'm');
if (!requiredRegex.test(body)) {
core.setFailed(`
The PR description must include a RELEASE NOTES section.
It should be in one of the following formats:
- "RELEASE NOTES: none" (case-insensitive)
- "RELEASE NOTES: N/A" (case-insensitive)
- A bulleted list under "RELEASE NOTES:", for example:
RELEASE NOTES:
* my_package: Fix bug causing crash...
`);
}
- name: Validate Label
uses: actions/github-script@v6
with:
script: |
const labels = context.payload.pull_request.labels.map(label => label.name);
const requiredRegex = new RegExp('^Type:');
const hasRequiredLabel = labels.some(label => requiredRegex.test(label));
if (!hasRequiredLabel) {
core.setFailed("This PR must have a label starting with 'Type:'.");
}
- name: Validate Milestone
uses: actions/github-script@v6
with:
script: |
const milestone = context.payload.pull_request.milestone;
if (!milestone) {
core.setFailed("This PR must be associated with a milestone.");
} else {
const requiredRegex = new RegExp('Release$');
if (!requiredRegex.test(milestone.title)) {
core.setFailed("The milestone for this PR must end with 'Release'.");
}
}

63
.github/workflows/release.yml vendored Normal file
View File

@@ -0,0 +1,63 @@
name: Release
on:
release:
types: [published]
permissions:
contents: read
jobs:
release:
permissions:
contents: write # to upload release asset (actions/upload-release-asset)
name: Release cmd/protoc-gen-go-grpc
runs-on: ubuntu-latest
if: startsWith(github.event.release.tag_name, 'cmd/protoc-gen-go-grpc/')
strategy:
matrix:
goos: [linux, darwin, windows]
goarch: [386, amd64, arm64]
exclude:
- goos: darwin
goarch: 386
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Set up Go
uses: actions/setup-go@v5
- name: Download dependencies
run: |
cd cmd/protoc-gen-go-grpc
go mod download
- name: Prepare build directory
run: |
mkdir -p build/
cp README.md build/
cp LICENSE build/
- name: Build
env:
GOOS: ${{ matrix.goos }}
GOARCH: ${{ matrix.goarch }}
run: |
cd cmd/protoc-gen-go-grpc
go build -trimpath -o $GITHUB_WORKSPACE/build
- name: Create package
id: package
run: |
PACKAGE_NAME=protoc-gen-go-grpc.${GITHUB_REF#refs/tags/cmd/protoc-gen-go-grpc/}.${{ matrix.goos }}.${{ matrix.goarch }}.tar.gz
tar -czvf $PACKAGE_NAME -C build .
echo "name=${PACKAGE_NAME}" >> $GITHUB_OUTPUT
- name: Upload asset
run: |
gh release upload ${{ github.event.release.tag_name }} ./${{ steps.package.outputs.name }}
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

33
.github/workflows/stale.yml vendored Normal file
View File

@@ -0,0 +1,33 @@
name: Stale bot
on:
workflow_dispatch:
schedule:
- cron: "44 */2 * * *"
permissions:
contents: read
jobs:
stale:
runs-on: ubuntu-latest
permissions:
issues: write
pull-requests: write
steps:
- uses: actions/stale@v8
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
days-before-stale: 6
days-before-close: 7
only-labels: 'Status: Requires Reporter Clarification'
stale-issue-label: 'stale'
stale-pr-label: 'stale'
operations-per-run: 999
stale-issue-message: >
This issue is labeled as requiring an update from the reporter, and no update has been received
after 6 days. If no update is provided in the next 7 days, this issue will be automatically closed.
stale-pr-message: >
This PR is labeled as requiring an update from the reporter, and no update has been received
after 6 days. If no update is provided in the next 7 days, this issue will be automatically closed.

135
.github/workflows/testing.yml vendored Normal file
View File

@@ -0,0 +1,135 @@
name: Testing
# Trigger on pushes, PRs (excluding documentation changes), and nightly.
on:
push:
pull_request:
schedule:
- cron: 0 0 * * * # daily at 00:00
permissions:
contents: read
# Always force the use of Go modules
env:
GO111MODULE: on
jobs:
# Check generated protos match their source repos (optional for PRs).
vet-proto:
runs-on: ubuntu-latest
timeout-minutes: 20
steps:
- name: Checkout repo
uses: actions/checkout@v4
# Setup the environment.
- name: Setup Go
uses: actions/setup-go@v5
with:
go-version: '1.26'
cache-dependency-path: "**/go.sum"
# Run the vet-proto checks.
- name: vet-proto
run: ./scripts/vet-proto.sh -install && ./scripts/vet-proto.sh
# Run the main gRPC-Go tests.
tests:
name: ${{ matrix.display_name }}
# Use the matrix variable to set the runner, with 'ubuntu-latest' as the
# default.
runs-on: ${{ matrix.runner || 'ubuntu-latest' }}
timeout-minutes: 20
strategy:
fail-fast: false
matrix:
include:
- type: vet
goversion: '1.25'
display_name: 'static checks (latest-1)'
- type: extras
goversion: '1.26'
display_name: 'extras (latest)'
- type: tests
goversion: '1.26'
display_name: 'tests (latest)'
- type: tests
goversion: '1.26'
testflags: -race
display_name: 'tests (-race, latest)'
- type: tests
goversion: '1.26'
goarch: 386
display_name: 'tests (i386, latest)'
- type: tests
goversion: '1.26'
goarch: arm64
runner: ubuntu-24.04-arm
display_name: 'tests (arm, latest)'
- type: tests
goversion: '1.25'
display_name: 'tests (latest-1)'
steps:
# Setup the environment.
- name: Setup GOARCH
if: matrix.goarch != ''
run: echo "GOARCH=${{ matrix.goarch }}" >> $GITHUB_ENV
- name: Setup GRPC environment
if: matrix.grpcenv != ''
run: echo "${{ matrix.grpcenv }}" >> $GITHUB_ENV
- name: Checkout repo
uses: actions/checkout@v4
- name: Setup Go
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.goversion }}
cache-dependency-path: "**/*go.sum"
# Only run vet for 'vet' runs.
- name: Run vet.sh
if: matrix.type == 'vet'
run: ./scripts/vet.sh -install && ./scripts/vet.sh
# Main tests run for everything except when testing "extras"
# (where we run a reduced set of tests).
- name: Run tests
if: matrix.type == 'tests'
run: |
go version
go test ${{ matrix.testflags }} -cpu 1,4 -timeout 7m ./...
cd "${GITHUB_WORKSPACE}"
for MOD_FILE in $(find . -name 'go.mod' | grep -Ev '^\./go\.mod'); do
pushd "$(dirname ${MOD_FILE})"
go test ${{ matrix.testflags }} -cpu 1,4 -timeout 2m ./...
popd
done
# Non-core gRPC tests (examples, interop, etc)
- name: Run extras tests
if: matrix.type == 'extras'
run: |
export TERM=${TERM:-xterm}
go version
echo -e "\n-- Running Examples --"
examples/examples_test.sh
echo -e "\n-- Running AdvancedTLS Examples --"
security/advancedtls/examples/examples_test.sh
echo -e "\n-- Running Interop Test --"
interop/interop_test.sh
echo -e "\n-- Running xDS E2E Test --"
internal/xds/test/e2e/run.sh
echo -e "\n-- Running protoc-gen-go-grpc test --"
./scripts/vet-proto.sh -install
cmd/protoc-gen-go-grpc/protoc-gen-go-grpc_test.sh

1
AUTHORS Normal file
View File

@@ -0,0 +1 @@
Google Inc.

3
CODE-OF-CONDUCT.md Normal file
View File

@@ -0,0 +1,3 @@
## Community Code of Conduct
gRPC follows the [CNCF Code of Conduct](https://github.com/cncf/foundation/blob/master/code-of-conduct.md).

182
CONTRIBUTING.md Normal file
View File

@@ -0,0 +1,182 @@
# How to contribute
We welcome your patches and contributions to gRPC! Please read the gRPC
organization's [governance
rules](https://github.com/grpc/grpc-community/blob/master/governance.md) before
proceeding.
If you are new to GitHub, please start by reading [Pull Request howto](https://help.github.com/articles/about-pull-requests/)
## Legal requirements
In order to protect both you and ourselves, you will need to sign the
[Contributor License
Agreement](https://identity.linuxfoundation.org/projects/cncf). When you create
your first PR, a link will be added as a comment that contains the steps needed
to complete this process.
## Getting Started
A great way to start is by searching through our open issues. [Unassigned issues
labeled as "help
wanted"](https://github.com/grpc/grpc-go/issues?q=sort%3Aupdated-desc%20is%3Aissue%20is%3Aopen%20label%3A%22Status%3A%20Help%20Wanted%22%20no%3Aassignee)
are especially nice for first-time contributors, as they should be well-defined
problems that already have agreed-upon solutions.
## Generative AI Policy
AI tools have the ability to produce more code than is possible for the gRPC
team to read, understand, review, and accept into the repository. For this
reason, we request that all contributions adhere to the following rules:
1. **No AI-Generated Interactions:** All communication in the repo must be
authored by a human. _Exception: AIs may be used for directed writing
assistance or translation._ Absolutely no automated agents are allowed to
directly publish to GitHub.
2. **Author Ownership and Accountability:** Code contributions are expected to
be fully owned and understood by the human contributor. If the code was
produced by generative AI, the author is expected to have reviewed and
understood it in its entirety before submitting it for review. This includes
all content: production code, tests, examples, tools, etc.
In addition to the above requirements, any AI-assisted contributions must also
comply with the [Linux Foundation Generative AI
Policy](https://www.linuxfoundation.org/legal/generative-ai). This includes
confirming that all contributions are legally allowed to be contributed to the
gRPC project under the applicable license terms.
## Code Style
We follow [Google's published Go style
guide](https://google.github.io/styleguide/go/). Note that there are three
primary documents that make up this style guide; please follow them as closely
as possible. If a reviewer recommends something that contradicts those
guidelines, there may be valid reasons to do so, but it should be rare.
## Guidelines for Pull Requests
Please read the following carefully to ensure your contributions can be merged
smoothly and quickly.
### PR Contents
- Create **small PRs** that are narrowly focused on **addressing a single
concern**. We often receive PRs that attempt to fix several things at the same
time, and if one part of the PR has a problem, that will hold up the entire
PR.
- If your change does not address an **open issue** with an **agreed
resolution**, consider opening an issue and discussing it first. If you are
suggesting a behavioral or API change, consider starting with a [gRFC
proposal](https://github.com/grpc/proposal). Many new features that are not
bug fixes will require cross-language agreement.
- If you want to fix **formatting or style**, consider whether your changes are
an obvious improvement or might be considered a personal preference. If a
style change is based on preference, it likely will not be accepted. If it
corrects widely agreed-upon anti-patterns, then please do create a PR and
explain the benefits of the change.
- For correcting **misspellings**, please be aware that we use some terms that
are sometimes flagged by spell checkers. As an example, "if an only if" is
often written as "iff". Please do not make spelling correction changes unless
you are certain they are misspellings.
- **All tests need to be passing** before your change can be merged. We
recommend you run tests locally before creating your PR to catch breakages
early on:
- `./scripts/vet.sh` to catch vet errors.
- `go test -cpu 1,4 -timeout 7m ./...` to run the tests.
- `go test -race -cpu 1,4 -timeout 7m ./...` to run tests in race mode.
Note that we have a multi-module repo, so `go test` commands may need to be
run from the root of each module in order to cause all tests to run.
*Alternatively*, you may find it easier to push your changes to your fork on
GitHub, which will trigger a GitHub Actions run that you can use to verify
everything is passing.
- Note that there are two GitHub actions checks that need not be green:
1. We test the freshness of the generated proto code we maintain via the
`vet-proto` check. If the source proto files are updated, but our repo is
not updated, an optional checker will fail. This will be fixed by our team
in a separate PR and will not prevent the merge of your PR.
2. We run a checker that will fail if there is any change in dependencies of
an exported package via the `dependencies` check. If new dependencies are
added that are not appropriate, we may not accept your PR (see below).
- If you are adding a **new file**, make sure it has the **copyright message**
template at the top as a comment. You can copy the message from an existing
file and update the year.
- The grpc package should only depend on standard Go packages and a small number
of exceptions. **If your contribution introduces new dependencies**, you will
need a discussion with gRPC-Go maintainers.
### PR Descriptions
- **PR titles** should start with the name of the component being addressed, or
the type of change. Examples: transport, client, server, round_robin, xds,
cleanup, deps.
- Read and follow the **guidelines for PR titles and descriptions** here:
https://google.github.io/eng-practices/review/developer/cl-descriptions.html
*particularly* the sections "First Line" and "Body is Informative".
Note: your PR description will be used as the git commit message in a
squash-and-merge if your PR is approved. We may make changes to this as
necessary.
- **Does this PR relate to an open issue?** On the first line, please use the
tag `Fixes #<issue>` to ensure the issue is closed when the PR is merged. Or
use `Updates #<issue>` if the PR is related to an open issue, but does not fix
it. Consider filing an issue if one does not already exist.
- PR descriptions *must* conclude with **release notes** as follows:
```
RELEASE NOTES:
* <component>: <summary>
```
This need not match the PR title.
The summary must:
* be something that gRPC users will understand.
* clearly explain the feature being added, the issue being fixed, or the
behavior being changed, etc. If fixing a bug, be clear about how the bug
can be triggered by an end-user.
* begin with a capital letter and use complete sentences.
* be as short as possible to describe the change being made.
If a PR is *not* end-user visible -- e.g. a cleanup, testing change, or
GitHub-related, use `RELEASE NOTES: n/a`.
### PR Process
- Please **self-review** your code changes before sending your PR. This will
prevent simple, obvious errors from causing delays.
- Maintain a **clean commit history** and use **meaningful commit messages**.
PRs with messy commit histories are difficult to review and won't be merged.
Before sending your PR, ensure your changes are based on top of the latest
`upstream/master` commits, and avoid rebasing in the middle of a code review.
You should **never use `git push -f`** unless absolutely necessary during a
review, as it can interfere with GitHub's tracking of comments.
- Unless your PR is trivial, you should **expect reviewer comments** that you
will need to address before merging. We'll label the PR as `Status: Requires
Reporter Clarification` if we expect you to respond to these comments in a
timely manner. If the PR remains inactive for 6 days, it will be marked as
`stale`, and we will automatically close it after 7 days if we don't hear back
from you. Please feel free to ping issues or bugs if you do not get a response
within a week.

View File

@@ -0,0 +1,183 @@
## Anti-Patterns of Client creation
### How to properly create a `ClientConn`: `grpc.NewClient`
[`grpc.NewClient`](https://pkg.go.dev/google.golang.org/grpc#NewClient) is the
function in the gRPC library that creates a virtual connection from a client
application to a gRPC server. It takes a target URI (which represents the name
of a logical backend service and resolves to one or more physical addresses) and
a list of options, and returns a
[`ClientConn`](https://pkg.go.dev/google.golang.org/grpc#ClientConn) object that
represents the virtual connection to the server. The `ClientConn` contains one
or more actual connections to real servers and attempts to maintain these
connections by automatically reconnecting to them when they break. `NewClient`
was introduced in gRPC-Go v1.63.
### The wrong way: `grpc.Dial`
[`grpc.Dial`](https://pkg.go.dev/google.golang.org/grpc#Dial) is a deprecated
function that also creates the same virtual connection pool as `grpc.NewClient`.
However, unlike `grpc.NewClient`, it immediately starts connecting and supports
a few additional `DialOption`s that control this initial connection attempt.
These are: `WithBlock`, `WithTimeout`, `WithReturnConnectionError`, and
`FailOnNonTempDialError`.
That `grpc.Dial` creates connections immediately is not a problem in and of
itself, but this behavior differs from how gRPC works in all other languages,
and it can be convenient to have a constructor that does not perform I/O. It
can also be confusing to users, as most people expect a function called `Dial`
to create _a_ connection which may need to be recreated if it is lost.
`grpc.Dial` uses "passthrough" as the default name resolver for backward
compatibility while `grpc.NewClient` uses "dns" as its default name resolver.
This subtle difference is important to legacy systems that also specified a
custom dialer and expected it to receive the target string directly.
For these reasons, using `grpc.Dial` is discouraged. Even though it is marked
as deprecated, we will continue to support it until a v2 is released (and no
plans for a v2 exist at the time this was written).
### Especially bad: using deprecated `DialOptions`
`FailOnNonTempDialError`, `WithBlock`, and `WithReturnConnectionError` are three
`DialOption`s that are only supported by `Dial` because they only affect the
behavior of `Dial` itself. `WithBlock` causes `Dial` to wait until the
`ClientConn` reports its `State` as `connectivity.Connected`. The other two deal
with returning connection errors before the timeout (`WithTimeout` or on the
context when using `DialContext`).
The reason these options can be a problem is that connections with a
`ClientConn` are dynamic -- they may come and go over time. If your client
successfully connects, the server could go down 1 second later, and your RPCs
will fail. "Knowing you are connected" does not tell you much in this regard.
Additionally, _all_ RPCs created on an "idle" or a "connecting" `ClientConn`
will wait until their deadline or until a connection is established before
failing. This means that you don't need to check that a `ClientConn` is "ready"
before starting your RPCs. By default, RPCs will fail if the `ClientConn`
enters the "transient failure" state, but setting `WaitForReady(true)` on a
call will cause it to queue even in the "transient failure" state, and it will
only ever fail due to a deadline, a server response, or a connection loss after
the RPC was sent to a server.
Some users of `Dial` use it as a way to validate the configuration of their
system. If you wish to maintain this behavior but migrate to `NewClient`, you
can call `GetState`, then `Connect` if the state is `Idle` and
`WaitForStateChange` until the channel is connected. However, if this fails,
it does not mean that your configuration was bad - it could also mean the
service is not reachable by the client due to connectivity reasons.
## Best practices for error handling in gRPC
Instead of relying on failures at dial time, we strongly encourage developers to
rely on errors from RPCs. When a client makes an RPC, it can receive an error
response from the server. These errors can provide valuable information about
what went wrong, including information about network issues, server-side errors,
and incorrect usage of the gRPC API.
By handling errors from RPCs correctly, developers can write more reliable and
robust gRPC applications. Here are some best practices for error handling in
gRPC:
- Always check for error responses from RPCs and handle them appropriately.
- Use the `status` field of the error response to determine the type of error
that occurred.
- When retrying failed RPCs, consider using the built-in retry mechanism
provided by gRPC-Go, if available, instead of manually implementing retries.
Refer to the [gRPC-Go retry example
documentation](https://github.com/grpc/grpc-go/blob/master/examples/features/retry/README.md)
for more information. Note that this is not a substitute for client-side
retries as errors that occur after an RPC starts on a server cannot be
retried through gRPC's built-in mechanism.
- If making an outgoing RPC from a server handler, be sure to translate the
status code before returning the error from your method handler. For example,
if the error is an `INVALID_ARGUMENT` status code, that probably means
your service has a bug (otherwise it shouldn't have triggered this error), in
which case `INTERNAL` is more appropriate to return back to your users.
### Example: Handling errors from an RPC
The following code snippet demonstrates how to handle errors from an RPC in
gRPC:
```go
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
res, err := client.MyRPC(ctx, &MyRequest{})
if err != nil {
// Handle the error appropriately,
// log it & return an error to the caller, etc.
log.Printf("Error calling MyRPC: %v", err)
return nil, err
}
// Use the response as appropriate
log.Printf("MyRPC response: %v", res)
```
To determine the type of error that occurred, you can use the status field of
the error response:
```go
resp, err := client.MakeRPC(context.TODO(), request)
if err != nil {
if status, ok := status.FromError(err); ok {
// Handle the error based on its status code
if status.Code() == codes.NotFound {
log.Println("Requested resource not found")
} else {
log.Printf("RPC error: %v", status.Message())
}
} else {
// Handle non-RPC errors
log.Printf("Non-RPC error: %v", err)
}
return
}
// Use the response as needed
log.Printf("Response received: %v", resp)
```
### Example: Using a backoff strategy
When retrying failed RPCs, use a backoff strategy to avoid overwhelming the
server or exacerbating network issues:
```go
var res *MyResponse
var err error
retryableStatusCodes := map[codes.Code]bool{
codes.Unavailable: true, // etc
}
// Retry the RPC a maximum number of times.
for i := 0; i < maxRetries; i++ {
// Make the RPC.
res, err = client.MyRPC(context.TODO(), &MyRequest{})
// Check if the RPC was successful.
if !retryableStatusCodes[status.Code(err)] {
// The RPC was successful or errored in a non-retryable way;
// do not retry.
break
}
// The RPC is retryable; wait for a backoff period before retrying.
backoff := time.Duration(i+1) * time.Second
log.Printf("Error calling MyRPC: %v; retrying in %v", err, backoff)
time.Sleep(backoff)
}
// Check if the RPC was successful after all retries.
if err != nil {
// All retries failed, so handle the error appropriately
log.Printf("Error calling MyRPC: %v", err)
return nil, err
}
// Use the response as appropriate.
log.Printf("MyRPC response: %v", res)
```

View File

@@ -0,0 +1,75 @@
# Benchmark
gRPC-Go comes with a set of benchmarking utilities to measure performance.
These utilities can be found in the `benchmark` directory within the project's
root directory.
The main utility, aptly named `benchmain`, supports a host of configurable
parameters to simulate various environments and workloads. For example, if your
server's workload is primarily streaming RPCs with large messages with
compression turned on, invoking `benchmain` in the following way may closely
simulate your application:
```bash
$ go run google.golang.org/grpc/benchmark/benchmain/main.go \
-workloads=streaming \
-reqSizeBytes=1024 \
-respSizeBytes=1024 \
-compression=gzip
```
Pass the `-h` flag to the `benchmain` utility to see other flags and workloads
that are supported.
## Varying Payload Sizes (Weighted Random Distribution)
The `benchmain` utility supports two flags, `-reqPayloadCurveFiles` and
`-respPayloadCurveFiles`, that can be used to specify histograms representing
a weighted random distribution of request and response payload sizes,
respectively. This is useful to simulate workloads with arbitrary payload
sizes.
The options take a comma-separated list of file paths as value. Each file must
be a valid CSV file with three columns in each row. Each row represents a range
of payload sizes (first two columns) and the weight associated with that range
(third column). For example, consider the below file:
```csv
1,32,12.5
128,256,12.5
1024,2048,25.0
```
Assume that `benchmain` is invoked like so:
```bash
$ go run google.golang.org/grpc/benchmark/benchmain/main.go \
-workloads=unary \
-reqPayloadCurveFiles=/path/to/csv \
-respPayloadCurveFiles=/path/to/csv
```
This tells the `benchmain` utility to generate unary RPC requests with a 25%
probability of payload sizes in the ranges 1-32 bytes, 25% probability in the
128-256 bytes range, and 50% probability in the 1024-2048 bytes range. RPC
requests outside these ranges will not be generated.
You may specify multiple CSV files delimited by a comma. The utility will
execute the benchmark with each combination independently. That is, the
following command will execute four benchmarks:
```bash
$ go run google.golang.org/grpc/benchmark/benchmain/main.go \
-workloads=unary \
-reqPayloadCurveFiles=/path/to/csv1,/path/to/csv2 \
-respPayloadCurveFiles=/path/to/csv3,/path/to/csv4
```
You may also combine `PayloadCurveFiles` with `SizeBytes` options. For example:
```
$ go run google.golang.org/grpc/benchmark/benchmain/main.go \
-workloads=unary \
-reqPayloadCurveFiles=/path/to/csv \
-respSizeBytes=1
```

View File

@@ -0,0 +1,80 @@
# Compression
The preferred method for configuring message compression on both clients and
servers is to use
[`encoding.RegisterCompressor`](https://godoc.org/google.golang.org/grpc/encoding#RegisterCompressor)
to register an implementation of a compression algorithm. See
`grpc/encoding/gzip/gzip.go` for an example of how to implement one.
Once a compressor has been registered on the client-side, RPCs may be sent using
it via the
[`UseCompressor`](https://godoc.org/google.golang.org/grpc#UseCompressor)
`CallOption`. Remember that `CallOption`s may be turned into defaults for all
calls from a `ClientConn` by using the
[`WithDefaultCallOptions`](https://godoc.org/google.golang.org/grpc#WithDefaultCallOptions)
`DialOption`. If `UseCompressor` is used and the corresponding compressor has
not been installed, an `Internal` error will be returned to the application
before the RPC is sent.
Server-side, registered compressors will be used automatically to decode request
messages and encode the responses. Servers currently always respond using the
same compression method specified by the client. If the corresponding
compressor has not been registered, an `Unimplemented` status will be returned
to the client.
## Deprecated API
There is a deprecated API for setting compression as well. It is not
recommended for use. However, if you were previously using it, the following
section may be helpful in understanding how it works in combination with the new
API.
### Client-Side
There are two legacy functions and one new function to configure compression:
```go
func WithCompressor(grpc.Compressor) DialOption {}
func WithDecompressor(grpc.Decompressor) DialOption {}
func UseCompressor(name) CallOption {}
```
For outgoing requests, the following rules are applied in order:
1. If `UseCompressor` is used, messages will be compressed using the compressor
named.
* If the compressor named is not registered, an Internal error is returned
back to the client before sending the RPC.
* If UseCompressor("identity"), no compressor will be used, but "identity"
will be sent in the header to the server.
1. If `WithCompressor` is used, messages will be compressed using that
compressor implementation.
1. Otherwise, outbound messages will be uncompressed.
For incoming responses, the following rules are applied in order:
1. If `WithDecompressor` is used and it matches the message's encoding, it will
be used.
1. If a registered compressor matches the response's encoding, it will be used.
1. Otherwise, the stream will be closed and an `Unimplemented` status error will
be returned to the application.
### Server-Side
There are two legacy functions to configure compression:
```go
func RPCCompressor(grpc.Compressor) ServerOption {}
func RPCDecompressor(grpc.Decompressor) ServerOption {}
```
For incoming requests, the following rules are applied in order:
1. If `RPCDecompressor` is used and that decompressor matches the request's
encoding: it will be used.
1. If a registered compressor matches the request's encoding, it will be used.
1. Otherwise, an `Unimplemented` status will be returned to the client.
For outgoing responses, the following rules are applied in order:
1. If `RPCCompressor` is used, that compressor will be used to compress all
response messages.
1. If compression was used for the incoming request and a registered compressor
supports it, that same compression method will be used for the outgoing
response.
1. Otherwise, no compression will be used for the outgoing response.

View File

@@ -0,0 +1,39 @@
# Concurrency
In general, gRPC-go provides a concurrency-friendly API. What follows are some
guidelines.
## Clients
A [ClientConn][client-conn] can safely be accessed concurrently. Using
[helloworld][helloworld] as an example, one could share the `ClientConn` across
multiple goroutines to create multiple `GreeterClient` types. In this case,
RPCs would be sent in parallel. `GreeterClient`, generated from the proto
definitions and wrapping `ClientConn`, is also concurrency safe, and may be
directly shared in the same way. Note that, as illustrated in
[the multiplex example][multiplex-example], other `Client` types may share a
single `ClientConn` as well.
## Streams
When using streams, one must take care to avoid calling either `SendMsg` or
`RecvMsg` multiple times against the same [Stream][stream] from different
goroutines. In other words, it's safe to have a goroutine calling `SendMsg` and
another goroutine calling `RecvMsg` on the same stream at the same time. But it
is not safe to call `SendMsg` on the same stream in different goroutines, or to
call `RecvMsg` on the same stream in different goroutines.
## Servers
Each RPC handler attached to a registered server will be invoked in its own
goroutine. For example, [SayHello][say-hello] will be invoked in its own
goroutine. The same is true for service handlers for streaming RPCs, as seen
in the route guide example [here][route-guide-stream]. Similar to clients,
multiple services can be registered to the same server.
[helloworld]: https://github.com/grpc/grpc-go/blob/master/examples/helloworld/greeter_client/main.go#L43
[client-conn]: https://godoc.org/google.golang.org/grpc#ClientConn
[stream]: https://godoc.org/google.golang.org/grpc#Stream
[say-hello]: https://github.com/grpc/grpc-go/blob/master/examples/helloworld/greeter_server/main.go#L41
[route-guide-stream]: https://github.com/grpc/grpc-go/blob/master/examples/route_guide/server/server.go#L126
[multiplex-example]: https://github.com/grpc/grpc-go/tree/master/examples/features/multiplex

146
Documentation/encoding.md Normal file
View File

@@ -0,0 +1,146 @@
# Encoding
The gRPC API for sending and receiving is based upon *messages*. However,
messages cannot be transmitted directly over a network; they must first be
converted into *bytes*. This document describes how gRPC-Go converts messages
into bytes and vice-versa for the purposes of network transmission.
## Codecs (Serialization and Deserialization)
A `Codec` contains code to serialize a message into a byte slice (`Marshal`) and
deserialize a byte slice back into a message (`Unmarshal`). `Codec`s are
registered by name into a global registry maintained in the `encoding` package.
### Implementing a `Codec`
A typical `Codec` will be implemented in its own package with an `init` function
that registers itself, and is imported anonymously. For example:
```go
package proto
import "google.golang.org/grpc/encoding"
func init() {
encoding.RegisterCodec(protoCodec{})
}
// ... implementation of protoCodec ...
```
For an example, gRPC's implementation of the `proto` codec can be found in
[`encoding/proto`](https://godoc.org/google.golang.org/grpc/encoding/proto).
### Using a `Codec`
By default, gRPC registers and uses the "proto" codec, so it is not necessary to
do this in your own code to send and receive proto messages. To use another
`Codec` from a client or server:
```go
package myclient
import _ "path/to/another/codec"
```
`Codec`s, by definition, must be symmetric, so the same desired `Codec` should
be registered in both client and server binaries.
On the client-side, to specify a `Codec` to use for message transmission, the
`CallOption` `CallContentSubtype` should be used as follows:
```go
response, err := myclient.MyCall(ctx, request, grpc.CallContentSubtype("mycodec"))
```
As a reminder, all `CallOption`s may be converted into `DialOption`s that become
the default for all RPCs sent through a client using `grpc.WithDefaultCallOptions`:
```go
myclient := grpc.NewClient(target, grpc.WithDefaultCallOptions(grpc.CallContentSubtype("mycodec")))
```
When specified in either of these ways, messages will be encoded using this
codec and sent along with headers indicating the codec (`content-type` set to
`application/grpc+<codec name>`).
On the server-side, using a `Codec` is as simple as registering it into the
global registry (i.e. `import`ing it). If a message is encoded with the content
sub-type supported by a registered `Codec`, it will be used automatically for
decoding the request and encoding the response. Otherwise, for
backward-compatibility reasons, gRPC will attempt to use the "proto" codec. In
an upcoming change (tracked in [this
issue](https://github.com/grpc/grpc-go/issues/1824)), such requests will be
rejected with status code `Unimplemented` instead.
## Compressors (Compression and Decompression)
Sometimes, the resulting serialization of a message is not space-efficient, and
it may be beneficial to compress this byte stream before transmitting it over
the network. To facilitate this operation, gRPC supports a mechanism for
performing compression and decompression.
A `Compressor` contains code to compress and decompress by wrapping `io.Writer`s
and `io.Reader`s, respectively. (The form of `Compress` and `Decompress` were
chosen to most closely match Go's standard package
[implementations](https://golang.org/pkg/compress/) of compressors). Like
`Codec`s, `Compressor`s are registered by name into a global registry maintained
in the `encoding` package.
### Implementing a `Compressor`
A typical `Compressor` will be implemented in its own package with an `init`
function that registers itself, and is imported anonymously. For example:
```go
package gzip
import "google.golang.org/grpc/encoding"
func init() {
encoding.RegisterCompressor(compressor{})
}
// ... implementation of compressor ...
```
An implementation of a `gzip` compressor can be found in
[`encoding/gzip`](https://godoc.org/google.golang.org/grpc/encoding/gzip).
### Using a `Compressor`
By default, gRPC does not register or use any compressors. To use a
`Compressor` from a client or server:
```go
package myclient
import _ "google.golang.org/grpc/encoding/gzip"
```
`Compressor`s, by definition, must be symmetric, so the same desired
`Compressor` should be registered in both client and server binaries.
On the client-side, to specify a `Compressor` to use for message transmission,
the `CallOption` `UseCompressor` should be used as follows:
```go
response, err := myclient.MyCall(ctx, request, grpc.UseCompressor("gzip"))
```
As a reminder, all `CallOption`s may be converted into `DialOption`s that become
the default for all RPCs sent through a client using `grpc.WithDefaultCallOptions`:
```go
myclient := grpc.NewClient(target, grpc.WithDefaultCallOptions(grpc.UseCompressor("gzip")))
```
When specified in either of these ways, messages will be compressed using this
compressor and sent along with headers indicating the compressor
(`content-coding` set to `<compressor name>`).
On the server-side, using a `Compressor` is as simple as registering it into the
global registry (i.e. `import`ing it). If a message is compressed with the
content coding supported by a registered `Compressor`, it will be used
automatically for decompressing the request and compressing the response.
Otherwise, the request will be rejected with status code `Unimplemented`.

View File

@@ -0,0 +1,83 @@
# Authentication
As outlined in the
[gRPC authentication guide](https://grpc.io/docs/guides/auth.html), there are a
number of different mechanisms for asserting identity between a client and
server. We'll present some code-samples here demonstrating how to provide TLS
support encryption and identity assertions as well as passing OAuth2 tokens to
services that support it.
## Enabling TLS on a gRPC client
```go
conn, err := grpc.NewClient(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")))
```
## Enabling TLS on a gRPC server
```go
creds, err := credentials.NewServerTLSFromFile(certFile, keyFile)
if err != nil {
log.Fatalf("Failed to generate credentials %v", err)
}
lis, err := net.Listen("tcp", ":0")
server := grpc.NewServer(grpc.Creds(creds))
...
server.Serve(lis)
```
## OAuth2
For an example of how to configure client and server to use OAuth2 tokens, see
[here](https://github.com/grpc/grpc-go/tree/master/examples/features/authentication).
### Validating a token on the server
Clients may use
[metadata.MD](https://godoc.org/google.golang.org/grpc/metadata#MD) to store
tokens and other authentication-related data. To gain access to the
`metadata.MD` object, a server may use
[metadata.FromIncomingContext](https://godoc.org/google.golang.org/grpc/metadata#FromIncomingContext).
With a reference to `metadata.MD` on the server, one needs to simply look up the
`authorization` key. Note, all keys stored within `metadata.MD` are normalized
to lowercase. See [here](https://godoc.org/google.golang.org/grpc/metadata#New).
It is possible to configure token validation for all RPCs using an interceptor.
A server may configure either a
[grpc.UnaryInterceptor](https://godoc.org/google.golang.org/grpc#UnaryInterceptor)
or a
[grpc.StreamInterceptor](https://godoc.org/google.golang.org/grpc#StreamInterceptor).
### Adding a token to all outgoing client RPCs
To send an OAuth2 token with each RPC, a client may configure the
`grpc.DialOption`
[grpc.WithPerRPCCredentials](https://godoc.org/google.golang.org/grpc#WithPerRPCCredentials).
Alternatively, a client may also use the `grpc.CallOption`
[grpc.PerRPCCredentials](https://godoc.org/google.golang.org/grpc#PerRPCCredentials)
on each invocation of an RPC.
To create a `credentials.PerRPCCredentials`, use
[oauth.TokenSource](https://godoc.org/google.golang.org/grpc/credentials/oauth#TokenSource).
Note, the OAuth2 implementation of `grpc.PerRPCCredentials` requires a client to
use
[grpc.WithTransportCredentials](https://godoc.org/google.golang.org/grpc#WithTransportCredentials)
to prevent any insecure transmission of tokens.
## Authenticating with Google
### Google Compute Engine (GCE)
```go
conn, err := grpc.NewClient(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), grpc.WithPerRPCCredentials(oauth.NewComputeEngine()))
```
### JWT
```go
jwtCreds, err := oauth.NewServiceAccountFromFile(*serviceAccountKeyFile, *oauthScope)
if err != nil {
log.Fatalf("Failed to create JWT credentials: %v", err)
}
conn, err := grpc.NewClient(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), grpc.WithPerRPCCredentials(jwtCreds))
```

View File

@@ -0,0 +1,244 @@
# Metadata
gRPC supports sending metadata between client and server.
This doc shows how to send and receive metadata in gRPC-go.
## Background
Four kinds of service method:
- [Unary RPC](https://grpc.io/docs/guides/concepts.html#unary-rpc)
- [Server streaming RPC](https://grpc.io/docs/guides/concepts.html#server-streaming-rpc)
- [Client streaming RPC](https://grpc.io/docs/guides/concepts.html#client-streaming-rpc)
- [Bidirectional streaming RPC](https://grpc.io/docs/guides/concepts.html#bidirectional-streaming-rpc)
And concept of [metadata].
## Constructing metadata
A metadata can be created using package [metadata].
The type MD is actually a map from string to a list of strings:
```go
type MD map[string][]string
```
Metadata can be read like a normal map.
Note that the value type of this map is `[]string`,
so that users can attach multiple values using a single key.
### Creating a new metadata
A metadata can be created from a `map[string]string` using function `New`:
```go
md := metadata.New(map[string]string{"key1": "val1", "key2": "val2"})
```
Another way is to use `Pairs`.
Values with the same key will be merged into a list:
```go
md := metadata.Pairs(
"key1", "val1",
"key1", "val1-2", // "key1" will have map value []string{"val1", "val1-2"}
"key2", "val2",
)
```
__Note:__ all the keys will be automatically converted to lowercase,
so "key1" and "kEy1" will be the same key and their values will be merged into the same list.
This happens for both `New` and `Pairs`.
### Storing binary data in metadata
In metadata, keys are always strings. But values can be strings or binary data.
To store binary data value in metadata, simply add "-bin" suffix to the key.
The values with "-bin" suffixed keys will be encoded when creating the metadata:
```go
md := metadata.Pairs(
"key", "string value",
"key-bin", string([]byte{96, 102}), // this binary data will be encoded (base64) before sending
// and will be decoded after being transferred.
)
```
## Sending and receiving metadata - client side
Client side metadata sending and receiving examples are available
[here](../examples/features/metadata/client/main.go).
### Sending metadata
There are two ways to send metadata to the server. The recommended way is to append kv pairs to the context using
`AppendToOutgoingContext`. This can be used with or without existing metadata on the context. When there is no prior
metadata, metadata is added; when metadata already exists on the context, kv pairs are merged in.
```go
// create a new context with some metadata
ctx := metadata.AppendToOutgoingContext(ctx, "k1", "v1", "k1", "v2", "k2", "v3")
// later, add some more metadata to the context (e.g. in an interceptor)
ctx := metadata.AppendToOutgoingContext(ctx, "k3", "v4")
// make unary RPC
response, err := client.SomeRPC(ctx, someRequest)
// or make streaming RPC
stream, err := client.SomeStreamingRPC(ctx)
```
Alternatively, metadata may be attached to the context using `NewOutgoingContext`. However, this
replaces any existing metadata in the context, so care must be taken to preserve the existing
metadata if desired. This is slower than using `AppendToOutgoingContext`. An example of this
is below:
```go
// create a new context with some metadata
md := metadata.Pairs("k1", "v1", "k1", "v2", "k2", "v3")
ctx := metadata.NewOutgoingContext(context.Background(), md)
// later, add some more metadata to the context (e.g. in an interceptor)
send, _ := metadata.FromOutgoingContext(ctx)
newMD := metadata.Pairs("k3", "v3")
ctx = metadata.NewOutgoingContext(ctx, metadata.Join(send, newMD))
// make unary RPC
response, err := client.SomeRPC(ctx, someRequest)
// or make streaming RPC
stream, err := client.SomeStreamingRPC(ctx)
```
### Receiving metadata
Metadata that a client can receive includes header and trailer.
#### Unary call
Header and trailer sent along with a unary call can be retrieved using function
[Header] and [Trailer] in [CallOption]:
```go
var header, trailer metadata.MD // variable to store header and trailer
r, err := client.SomeRPC(
ctx,
someRequest,
grpc.Header(&header), // will retrieve header
grpc.Trailer(&trailer), // will retrieve trailer
)
// do something with header and trailer
```
#### Streaming call
For streaming calls including:
- Server streaming RPC
- Client streaming RPC
- Bidirectional streaming RPC
Header and trailer can be retrieved from the returned stream using function
`Header` and `Trailer` in interface [ClientStream]:
```go
stream, err := client.SomeStreamingRPC(ctx)
// retrieve header
header, err := stream.Header()
// retrieve trailer
trailer := stream.Trailer()
```
## Sending and receiving metadata - server side
Server side metadata sending and receiving examples are available
[here](../examples/features/metadata/server/main.go).
### Receiving metadata
To read metadata sent by the client, the server needs to retrieve it from RPC
context using [FromIncomingContext].
If it is a unary call, the RPC handler's context can be used.
For streaming calls, the server needs to get context from the stream.
#### Unary call
```go
func (s *server) SomeRPC(ctx context.Context, in *pb.someRequest) (*pb.someResponse, error) {
md, ok := metadata.FromIncomingContext(ctx)
// do something with metadata
}
```
#### Streaming call
```go
func (s *server) SomeStreamingRPC(stream pb.Service_SomeStreamingRPCServer) error {
md, ok := metadata.FromIncomingContext(stream.Context()) // get context from stream
// do something with metadata
}
```
### Sending metadata
#### Unary call
To send header and trailer to client in unary call, the server can call
[SetHeader] and [SetTrailer] functions in module [grpc].
These two functions take a context as the first parameter.
It should be the RPC handler's context or one derived from it:
```go
func (s *server) SomeRPC(ctx context.Context, in *pb.someRequest) (*pb.someResponse, error) {
// create and set header
header := metadata.Pairs("header-key", "val")
grpc.SetHeader(ctx, header)
// create and set trailer
trailer := metadata.Pairs("trailer-key", "val")
grpc.SetTrailer(ctx, trailer)
}
```
#### Streaming call
For streaming calls, header and trailer can be sent using function
[SetHeader] and [SetTrailer] in interface [ServerStream]:
```go
func (s *server) SomeStreamingRPC(stream pb.Service_SomeStreamingRPCServer) error {
// create and set header
header := metadata.Pairs("header-key", "val")
stream.SetHeader(header)
// create and set trailer
trailer := metadata.Pairs("trailer-key", "val")
stream.SetTrailer(trailer)
}
```
**Important**
Do not use
[FromOutgoingContext] on the server to write metadata to be sent to the client.
[FromOutgoingContext] is for client-side use only.
## Updating metadata from a server interceptor
An example for updating metadata from a server interceptor is
available [here](../examples/features/metadata_interceptor/server/main.go).
[FromIncomingContext]: <https://pkg.go.dev/google.golang.org/grpc/metadata#FromIncomingContext>
[SetHeader]: <https://godoc.org/google.golang.org/grpc#SetHeader>
[SetTrailer]: https://godoc.org/google.golang.org/grpc#SetTrailer
[FromOutgoingContext]: https://pkg.go.dev/google.golang.org/grpc/metadata#FromOutgoingContext
[ServerStream]: https://godoc.org/google.golang.org/grpc#ServerStream
[grpc]: https://godoc.org/google.golang.org/grpc
[ClientStream]: https://godoc.org/google.golang.org/grpc#ClientStream
[Header]: https://godoc.org/google.golang.org/grpc#Header
[Trailer]: https://godoc.org/google.golang.org/grpc#Trailer
[CallOption]: https://godoc.org/google.golang.org/grpc#CallOption
[metadata]: https://godoc.org/google.golang.org/grpc/metadata

View File

@@ -0,0 +1,57 @@
# Keepalive
gRPC sends http2 pings on the transport to detect if the connection is down. If
the ping is not acknowledged by the other side within a certain period, the
connection will be closed. Note that pings are only necessary when there's no
activity on the connection.
For how to configure keepalive, see
https://godoc.org/google.golang.org/grpc/keepalive for the options.
## Why do I need this?
Keepalive can be useful to detect TCP level connection failures. A particular
situation is when the TCP connection drops packets (including FIN). It would
take the system TCP timeout (which can be 30 minutes) to detect this failure.
Keepalive would allow gRPC to detect this failure much sooner.
Another usage is (as the name suggests) to keep the connection alive. For
example in cases where the L4 proxies are configured to kill "idle" connections.
Sending pings would make the connections not "idle".
## What should I set?
It should be sufficient for most users to set [client
parameters](https://godoc.org/google.golang.org/grpc/keepalive) as a [dial
option](https://godoc.org/google.golang.org/grpc#WithKeepaliveParams).
## What will happen?
(The behavior described here is specific for gRPC-go, it might be slightly
different in other languages.)
When there's no activity on a connection (note that an ongoing stream results in
__no activity__ when there's no message being sent), after `Time`, a ping will
be sent by the client and the server will send a ping ack when it gets the ping.
Client will wait for `Timeout`, and check if there's any activity on the
connection during this period (a ping ack is an activity).
## What about server side?
Server has similar `Time` and `Timeout` settings as client. Server can also
configure connection max-age. See [server
parameters](https://godoc.org/google.golang.org/grpc/keepalive#ServerParameters)
for details.
### Enforcement policy
[Enforcement
policy](https://godoc.org/google.golang.org/grpc/keepalive#EnforcementPolicy) is
a special setting on server side to protect server from malicious or misbehaving
clients.
Server sends GOAWAY with ENHANCE_YOUR_CALM and close the connection when bad
behaviors are detected:
- Client sends too frequent pings
- Client sends pings when there's no stream and this is disallowed by server
config

View File

@@ -0,0 +1,49 @@
# Log Levels
This document describes the different log levels supported by the grpc-go
library, and under what conditions they should be used.
### Info
Info messages are for informational purposes and may aid in the debugging of
applications or the gRPC library.
Examples:
- The name resolver received an update.
- The balancer updated its picker.
- Significant gRPC state is changing.
At verbosity of 0 (the default), any single info message should not be output
more than once every 5 minutes under normal operation.
### Warning
Warning messages indicate problems that are non-fatal for the application, but
could lead to unexpected behavior or subsequent errors.
Examples:
- Resolver could not resolve target name.
- Error received while connecting to a server.
- Lost or corrupt connection with remote endpoint.
### Error
Error messages represent errors in the usage of gRPC that cannot be returned to
the application as errors, or internal gRPC-Go errors that are recoverable.
Internal errors are detected during gRPC tests and will result in test failures.
Examples:
- Invalid arguments passed to a function that cannot return an error.
- An internal error that cannot be returned or would be inappropriate to return
to the user.
### Fatal
Fatal errors are severe internal errors that are unrecoverable. These lead
directly to panics, and are avoided as much as possible.
Example:
- Internal invariant was violated.
- User attempted an action that cannot return an error gracefully, but would
lead to an invalid state if performed.

25
Documentation/proxy.md Normal file
View File

@@ -0,0 +1,25 @@
# Proxy
HTTP CONNECT proxies are supported by default in gRPC. The proxy address can be
specified by the environment variables `HTTPS_PROXY` and `NO_PROXY`. (Note that
these environment variables are case insensitive.)
**NOTE**: Using CONNECT proxies via https is not supported. gRPC performs a
plaintext CONNECT handshake to establish a tunnel and does not support the
additional encryption required to secure the initial connection to the proxy
itself.
When using TLS, the gRPC traffic is encrypted end-to-end between the client and
the destination server. Even when using a CONNECT proxy without https, the
security is not compromised, as the proxy only sees the destination address and
cannot intercept the encrypted gRPC data.
## Custom proxy
Currently, proxy support is implemented in the default dialer. It does one more
handshake (a CONNECT handshake in the case of HTTP CONNECT proxy) on the
connection before giving it to gRPC.
If the default proxy doesn't work for you, replace the default dialer with your
custom proxy dialer. This can be done using
[`WithContextDialer`](https://pkg.go.dev/google.golang.org/grpc#WithContextDialer).

View File

@@ -0,0 +1,68 @@
# RPC Errors
All service method handlers should return `nil` or errors from the
`status.Status` type. Clients have direct access to the errors.
Upon encountering an error, a gRPC server method handler should create a
`status.Status`. In typical usage, one would use [status.New][new-status]
passing in an appropriate [codes.Code][code] as well as a description of the
error to produce a `status.Status`. Calling [status.Err][status-err] converts
the `status.Status` type into an `error`. As a convenience method, there is also
[status.Error][status-error] which obviates the conversion step. Compare:
```
st := status.New(codes.NotFound, "some description")
err := st.Err()
// vs.
err := status.Error(codes.NotFound, "some description")
```
## Adding additional details to errors
In some cases, it may be necessary to add details for a particular error on the
server side. The [status.WithDetails][with-details] method exists for this
purpose. Clients may then read those details by first converting the plain
`error` type back to a [status.Status][status] and then using
[status.Details][details].
## Example
The [example][example] demonstrates the API discussed above and shows how to add
information about rate limits to the error message using `status.Status`.
To run the example, first start the server:
```
$ go run examples/rpc_errors/server/main.go
```
In a separate session, run the client:
```
$ go run examples/rpc_errors/client/main.go
```
On the first run of the client, all is well:
```
2018/03/12 19:39:33 Greeting: Hello world
```
Upon running the client a second time, the client exceeds the rate limit and
receives an error with details:
```
2018/03/19 16:42:01 Quota failure: violations:<subject:"name:world" description:"Limit one greeting per person" >
exit status 1
```
[status]: https://godoc.org/google.golang.org/grpc/status#Status
[new-status]: https://godoc.org/google.golang.org/grpc/status#New
[code]: https://godoc.org/google.golang.org/grpc/codes#Code
[with-details]: https://godoc.org/google.golang.org/grpc/internal/status#Status.WithDetails
[details]: https://godoc.org/google.golang.org/grpc/internal/status#Status.Details
[status-err]: https://godoc.org/google.golang.org/grpc/internal/status#Status.Err
[status-error]: https://godoc.org/google.golang.org/grpc/status#Error
[example]: https://github.com/grpc/grpc-go/tree/master/examples/features/error_details

View File

@@ -0,0 +1,171 @@
# gRPC Server Reflection Tutorial
gRPC Server Reflection provides information about publicly-accessible gRPC
services on a server, and assists clients at runtime to construct RPC requests
and responses without precompiled service information. It is used by
[gRPCurl](https://github.com/fullstorydev/grpcurl), which can be used to
introspect server protos and send/receive test RPCs.
## Enable Server Reflection
gRPC-go Server Reflection is implemented in package
[reflection](https://github.com/grpc/grpc-go/tree/master/reflection). To enable
server reflection, you need to import this package and register reflection
service on your gRPC server.
For example, to enable server reflection in `example/helloworld`, we need to
make the following changes:
```diff
--- a/examples/helloworld/greeter_server/main.go
+++ b/examples/helloworld/greeter_server/main.go
@@ -40,6 +40,7 @@ import (
"google.golang.org/grpc"
pb "google.golang.org/grpc/examples/helloworld/helloworld"
+ "google.golang.org/grpc/reflection"
)
const (
@@ -61,6 +62,8 @@ func main() {
}
s := grpc.NewServer()
pb.RegisterGreeterService(s, &pb.GreeterService{SayHello: sayHello})
+ // Register reflection service on gRPC server.
+ reflection.Register(s)
if err := s.Serve(lis); err != nil {
log.Fatalf("failed to serve: %v", err)
}
```
An example server with reflection registered can be found at
`examples/features/reflection/server`.
## gRPCurl
After enabling Server Reflection in a server application, you can use gRPCurl
to check its services. gRPCurl is built with Go and has packages available.
Instructions on how to install and use gRPCurl can be found at
[gRPCurl Installation](https://github.com/fullstorydev/grpcurl#installation).
## Use gRPCurl to check services
First, start the helloworld server in grpc-go directory:
```sh
$ cd <grpc-go-directory>/examples
$ go run features/reflection/server/main.go
```
output:
```sh
server listening at [::]:50051
```
After installing gRPCurl, open a new terminal and run the commands from the new
terminal.
**NOTE:** gRPCurl expects a TLS-encrypted connection by default. For all of
the commands below, use the `-plaintext` flag to use an unencrypted connection.
### List services and methods
The `list` command lists services exposed at a given port:
- List all the services exposed at a given port
```sh
$ grpcurl -plaintext localhost:50051 list
```
output:
```sh
grpc.examples.echo.Echo
grpc.reflection.v1alpha.ServerReflection
helloworld.Greeter
```
- List all the methods of a service
The `list` command lists methods given the full service name (in the format of
\<package\>.\<service\>).
```sh
$ grpcurl -plaintext localhost:50051 list helloworld.Greeter
```
output:
```sh
helloworld.Greeter.SayHello
```
### Describe services and methods
- Describe all services
The `describe` command inspects a service given its full name (in the format
of \<package\>.\<service\>).
```sh
$ grpcurl -plaintext localhost:50051 describe helloworld.Greeter
```
output:
```sh
helloworld.Greeter is a service:
service Greeter {
rpc SayHello ( .helloworld.HelloRequest ) returns ( .helloworld.HelloReply );
}
```
- Describe all methods of a service
The `describe` command inspects a method given its full name (in the format of
\<package\>.\<service\>.\<method\>).
```sh
$ grpcurl -plaintext localhost:50051 describe helloworld.Greeter.SayHello
```
output:
```sh
helloworld.Greeter.SayHello is a method:
rpc SayHello ( .helloworld.HelloRequest ) returns ( .helloworld.HelloReply );
```
### Inspect message types
We can use the `describe` command to inspect request/response types given the
full name of the type (in the format of \<package\>.\<type\>).
- Get information about the request type
```sh
$ grpcurl -plaintext localhost:50051 describe helloworld.HelloRequest
```
output:
```sh
helloworld.HelloRequest is a message:
message HelloRequest {
string name = 1;
}
```
### Call a remote method
We can send RPCs to a server and get responses using the full method name (in
the format of \<package\>.\<service\>.\<method\>). The `-d <string>` flag
represents the request data and the `-format text` flag indicates that the
request data is in text format.
- Call a unary method
```sh
$ grpcurl -plaintext -format text -d 'name: "gRPCurl"' \
localhost:50051 helloworld.Greeter.SayHello
```
output:
```sh
message: "Hello gRPCurl"
```

View File

@@ -0,0 +1,34 @@
# Versioning and Releases
Note: This document references terminology defined at http://semver.org.
## Release Frequency
Regular MINOR releases of gRPC-Go are performed every six weeks. Patch releases
to the previous two MINOR releases may be performed on demand or if serious
security problems are discovered.
## Versioning Policy
The gRPC-Go versioning policy follows the Semantic Versioning 2.0.0
specification, with the following exceptions:
- A MINOR version will not _necessarily_ add new functionality.
- MINOR releases will not break backward compatibility, except in the following
circumstances:
- An API was marked as EXPERIMENTAL upon its introduction.
- An API was marked as DEPRECATED in the initial MAJOR release.
- An API is inherently flawed and cannot provide correct or secure behavior.
In these cases, APIs MAY be changed or removed without a MAJOR release.
Otherwise, backward compatibility will be preserved by MINOR releases.
For an API marked as DEPRECATED, an alternative will be available (if
appropriate) for at least three months prior to its removal.
## Release History
Please see our release history on GitHub:
https://github.com/grpc/grpc-go/releases

1
GOVERNANCE.md Normal file
View File

@@ -0,0 +1 @@
This repository is governed by the gRPC organization's [governance rules](https://github.com/grpc/grpc-community/blob/master/governance.md).

202
LICENSE Normal file
View File

@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
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.

36
MAINTAINERS.md Normal file
View File

@@ -0,0 +1,36 @@
This page lists all active maintainers of this repository. If you were a
maintainer and would like to add your name to the Emeritus list, please send us a
PR.
See [GOVERNANCE.md](https://github.com/grpc/grpc-community/blob/master/governance.md)
for governance guidelines and how to become a maintainer.
See [CONTRIBUTING.md](https://github.com/grpc/grpc-community/blob/master/CONTRIBUTING.md)
for general contribution guidelines.
## Maintainers (in alphabetical order)
- [arjan-bal](https://github.com/arjan-bal), Google LLC
- [arvindbr8](https://github.com/arvindbr8), Google LLC
- [atollena](https://github.com/atollena), Datadog, Inc.
- [dfawley](https://github.com/dfawley), Google LLC
- [easwars](https://github.com/easwars), Google LLC
- [gtcooke94](https://github.com/gtcooke94), Google LLC
## Emeritus Maintainers (in alphabetical order)
- [adelez](https://github.com/adelez)
- [aranjans](https://github.com/aranjans)
- [canguler](https://github.com/canguler)
- [cesarghali](https://github.com/cesarghali)
- [erm-g](https://github.com/erm-g)
- [iamqizhao](https://github.com/iamqizhao)
- [jeanbza](https://github.com/jeanbza)
- [jtattermusch](https://github.com/jtattermusch)
- [lyuxuan](https://github.com/lyuxuan)
- [makmukhi](https://github.com/makmukhi)
- [matt-kwong](https://github.com/matt-kwong)
- [menghanl](https://github.com/menghanl)
- [nicolasnoble](https://github.com/nicolasnoble)
- [purnesh42h](https://github.com/purnesh42h)
- [srini100](https://github.com/srini100)
- [yongni](https://github.com/yongni)
- [zasweq](https://github.com/zasweq)

49
Makefile Normal file
View File

@@ -0,0 +1,49 @@
all: vet test testrace
build:
go build google.golang.org/grpc/...
clean:
go clean -i google.golang.org/grpc/...
deps:
GO111MODULE=on go get -d -v google.golang.org/grpc/...
proto:
@ if ! which protoc > /dev/null; then \
echo "error: protoc not installed" >&2; \
exit 1; \
fi
go generate google.golang.org/grpc/...
test:
go test -cpu 1,4 -timeout 7m google.golang.org/grpc/...
testsubmodule:
cd security/advancedtls && go test -cpu 1,4 -timeout 7m google.golang.org/grpc/security/advancedtls/...
cd security/authorization && go test -cpu 1,4 -timeout 7m google.golang.org/grpc/security/authorization/...
testrace:
go test -race -cpu 1,4 -timeout 7m google.golang.org/grpc/...
testdeps:
GO111MODULE=on go get -d -v -t google.golang.org/grpc/...
vet: vetdeps
./scripts/vet.sh
vetdeps:
./scripts/vet.sh -install
.PHONY: \
all \
build \
clean \
deps \
proto \
test \
testsubmodule \
testrace \
testdeps \
vet \
vetdeps

13
NOTICE.txt Normal file
View File

@@ -0,0 +1,13 @@
Copyright 2014 gRPC authors.
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.

108
README.md Normal file
View File

@@ -0,0 +1,108 @@
# gRPC-Go
[![GoDoc](https://pkg.go.dev/badge/google.golang.org/grpc)][API]
[![GoReportCard](https://goreportcard.com/badge/grpc/grpc-go)](https://goreportcard.com/report/github.com/grpc/grpc-go)
[![codecov](https://codecov.io/gh/grpc/grpc-go/graph/badge.svg)](https://codecov.io/gh/grpc/grpc-go)
The [Go][] implementation of [gRPC][]: A high performance, open source, general
RPC framework that puts mobile and HTTP/2 first. For more information see the
[Go gRPC docs][], or jump directly into the [quick start][].
## Prerequisites
- **[Go][]**: any one of the **two latest major** [releases][go-releases].
## Installation
Simply add the following import to your code, and then `go [build|run|test]`
will automatically fetch the necessary dependencies:
```go
import "google.golang.org/grpc"
```
> **Note:** If you are trying to access `grpc-go` from **China**, see the
> [FAQ](#FAQ) below.
## Learn more
- [Go gRPC docs][], which include a [quick start][] and [API
reference][API] among other resources
- [Low-level technical docs](Documentation) from this repository
- [Performance benchmark][]
- [Examples](examples)
- [Contribution guidelines](CONTRIBUTING.md)
## FAQ
### I/O Timeout Errors
The `golang.org` domain may be blocked from some countries. `go get` usually
produces an error like the following when this happens:
```console
$ go get -u google.golang.org/grpc
package google.golang.org/grpc: unrecognized import path "google.golang.org/grpc" (https fetch: Get https://google.golang.org/grpc?go-get=1: dial tcp 216.239.37.1:443: i/o timeout)
```
To build Go code, there are several options:
- Set up a VPN and access google.golang.org through that.
- With Go module support: it is possible to use the `replace` feature of `go
mod` to create aliases for golang.org packages. In your project's directory:
```sh
go mod edit -replace=google.golang.org/grpc=github.com/grpc/grpc-go@latest
go mod tidy
go mod vendor
go build -mod=vendor
```
Again, this will need to be done for all transitive dependencies hosted on
golang.org as well. For details, refer to [golang/go issue
#28652](https://github.com/golang/go/issues/28652).
### Compiling error, undefined: grpc.SupportPackageIsVersion
Please update to the latest version of gRPC-Go using
`go get google.golang.org/grpc`.
### How to turn on logging
The default logger is controlled by environment variables. Turn everything on
like this:
```console
$ export GRPC_GO_LOG_VERBOSITY_LEVEL=99
$ export GRPC_GO_LOG_SEVERITY_LEVEL=info
```
### The RPC failed with error `"code = Unavailable desc = transport is closing"`
This error means the connection the RPC is using was closed, and there are many
possible reasons, including:
1. mis-configured transport credentials, connection failed on handshaking
1. bytes disrupted, possibly by a proxy in between
1. server shutdown
1. Keepalive parameters caused connection shutdown, for example if you have
configured your server to terminate connections regularly to [trigger DNS
lookups](https://github.com/grpc/grpc-go/issues/3170#issuecomment-552517779).
If this is the case, you may want to increase your
[MaxConnectionAgeGrace](https://pkg.go.dev/google.golang.org/grpc/keepalive?tab=doc#ServerParameters),
to allow longer RPC calls to finish.
It can be tricky to debug this because the error happens on the client side but
the root cause of the connection being closed is on the server side. Turn on
logging on __both client and server__, and see if there are any transport
errors.
[API]: https://pkg.go.dev/google.golang.org/grpc
[Go]: https://golang.org
[Go module]: https://github.com/golang/go/wiki/Modules
[gRPC]: https://grpc.io
[Go gRPC docs]: https://grpc.io/docs/languages/go
[Performance benchmark]: https://performance-dot-grpc-testing.appspot.com/explore?dashboard=5180705743044608
[quick start]: https://grpc.io/docs/languages/go/quickstart
[go-releases]: https://golang.org/doc/devel/release.html

3
SECURITY.md Normal file
View File

@@ -0,0 +1,3 @@
# Security Policy
For information on gRPC Security Policy and reporting potential security issues, please see [gRPC CVE Process](https://github.com/grpc/proposal/blob/master/P4-grpc-cve-process.md).

58
admin/admin.go Normal file
View File

@@ -0,0 +1,58 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
// Package admin provides a convenient method for registering a collection of
// administration services to a gRPC server. The services registered are:
//
// - Channelz: https://github.com/grpc/proposal/blob/master/A14-channelz.md
//
// - CSDS: https://github.com/grpc/proposal/blob/master/A40-csds-support.md
//
// # Experimental
//
// Notice: All APIs in this package are experimental and may be removed in a
// later release.
package admin
import (
"google.golang.org/grpc"
channelzservice "google.golang.org/grpc/channelz/service"
internaladmin "google.golang.org/grpc/internal/admin"
)
func init() {
// Add a list of default services to admin here. Optional services, like
// CSDS, will be added by other packages.
internaladmin.AddService(func(registrar grpc.ServiceRegistrar) (func(), error) {
channelzservice.RegisterChannelzServiceToServer(registrar)
return nil, nil
})
}
// Register registers the set of admin services to the given server.
//
// The returned cleanup function should be called to clean up the resources
// allocated for the service handlers after the server is stopped.
//
// Note that if `s` is not a *grpc.Server or a *xds.GRPCServer, CSDS will not be
// registered because CSDS generated code is old and doesn't support interface
// `grpc.ServiceRegistrar`.
// https://github.com/envoyproxy/go-control-plane/issues/403
func Register(s grpc.ServiceRegistrar) (cleanup func(), _ error) {
return internaladmin.Register(s)
}

34
admin/admin_test.go Normal file
View File

@@ -0,0 +1,34 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package admin_test
import (
"testing"
"google.golang.org/grpc/admin/test"
"google.golang.org/grpc/codes"
)
func TestRegisterNoCSDS(t *testing.T) {
test.RunRegisterTests(t, test.ExpectedStatusCodes{
ChannelzCode: codes.OK,
// CSDS is not registered because xDS isn't imported.
CSDSCode: codes.Unimplemented,
})
}

38
admin/test/admin_test.go Normal file
View File

@@ -0,0 +1,38 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
// This file has the same content as admin_test.go, difference is that this is
// in another package, and it imports "xds", so we can test that csds is
// registered when xds is imported.
package test_test
import (
"testing"
"google.golang.org/grpc/admin/test"
"google.golang.org/grpc/codes"
_ "google.golang.org/grpc/xds"
)
func TestRegisterWithCSDS(t *testing.T) {
test.RunRegisterTests(t, test.ExpectedStatusCodes{
ChannelzCode: codes.OK,
CSDSCode: codes.OK,
})
}

104
admin/test/utils.go Normal file
View File

@@ -0,0 +1,104 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
// Package test contains test only functions for package admin. It's used by
// admin/admin_test.go and admin/test/admin_test.go.
package test
import (
"context"
"net"
"testing"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/admin"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/status"
v3statusgrpc "github.com/envoyproxy/go-control-plane/envoy/service/status/v3"
v3statuspb "github.com/envoyproxy/go-control-plane/envoy/service/status/v3"
channelzgrpc "google.golang.org/grpc/channelz/grpc_channelz_v1"
channelzpb "google.golang.org/grpc/channelz/grpc_channelz_v1"
)
const (
defaultTestTimeout = 10 * time.Second
)
// ExpectedStatusCodes contains the expected status code for each RPC (can be
// OK).
type ExpectedStatusCodes struct {
ChannelzCode codes.Code
CSDSCode codes.Code
}
// RunRegisterTests makes a client, runs the RPCs, and compares the status
// codes.
func RunRegisterTests(t *testing.T, ec ExpectedStatusCodes) {
lis, err := net.Listen("tcp", "localhost:0")
if err != nil {
t.Fatalf("cannot create listener: %v", err)
}
server := grpc.NewServer()
defer server.Stop()
cleanup, err := admin.Register(server)
if err != nil {
t.Fatalf("failed to register admin: %v", err)
}
defer cleanup()
go func() {
server.Serve(lis)
}()
conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient(%q) = %v", lis.Addr().String(), err)
}
t.Run("channelz", func(t *testing.T) {
if err := RunChannelz(conn); status.Code(err) != ec.ChannelzCode {
t.Fatalf("%s RPC failed with error %v, want code %v", "channelz", err, ec.ChannelzCode)
}
})
t.Run("csds", func(t *testing.T) {
if err := RunCSDS(conn); status.Code(err) != ec.CSDSCode {
t.Fatalf("%s RPC failed with error %v, want code %v", "CSDS", err, ec.CSDSCode)
}
})
}
// RunChannelz makes a channelz RPC.
func RunChannelz(conn *grpc.ClientConn) error {
c := channelzgrpc.NewChannelzClient(conn)
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
_, err := c.GetTopChannels(ctx, &channelzpb.GetTopChannelsRequest{}, grpc.WaitForReady(true))
return err
}
// RunCSDS makes a CSDS RPC.
func RunCSDS(conn *grpc.ClientConn) error {
c := v3statusgrpc.NewClientStatusDiscoveryServiceClient(conn)
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
_, err := c.FetchClientStatus(ctx, &v3statuspb.ClientStatusRequest{}, grpc.WaitForReady(true))
return err
}

174
attributes/attributes.go Normal file
View File

@@ -0,0 +1,174 @@
/*
*
* Copyright 2019 gRPC authors.
*
* 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.
*
*/
// Package attributes defines a generic key/value store used in various gRPC
// components.
//
// # Experimental
//
// Notice: This package is EXPERIMENTAL and may be changed or removed in a
// later release.
package attributes
import (
"fmt"
"iter"
"maps"
"strings"
)
// Attributes is an immutable struct for storing and retrieving generic
// key/value pairs. Keys must be hashable, and users should define their own
// types for keys. Values should not be modified after they are added to an
// Attributes or if they were received from one. If values implement 'Equal(o
// any) bool', it will be called by (*Attributes).Equal to determine whether
// two values with the same key should be considered equal.
type Attributes struct {
parent *Attributes
key, value any
}
// New returns a new Attributes containing the key/value pair.
func New(key, value any) *Attributes {
return &Attributes{
key: key,
value: value,
}
}
// WithValue returns a new Attributes containing the previous keys and values
// and the new key/value pair. If the same key appears multiple times, the
// last value overwrites all previous values for that key. value should not be
// modified later.
//
// Note that Attributes do not support deletion. Avoid using untyped nil values.
// Since the Value method returns an untyped nil when a key is absent, it is
// impossible to distinguish between a missing key and a key explicitly set to
// an untyped nil. If you need to represent a value being unset, consider
// storing a specific sentinel type or a wrapper struct with a boolean field
// indicating presence.
func (a *Attributes) WithValue(key, value any) *Attributes {
return &Attributes{
parent: a,
key: key,
value: value,
}
}
// Value returns the value associated with these attributes for key, or nil if
// no value is associated with key. The returned value should not be modified.
func (a *Attributes) Value(key any) any {
for cur := a; cur != nil; cur = cur.parent {
if cur.key == key {
return cur.value
}
}
return nil
}
// Equal returns whether a and o are equivalent. If 'Equal(o any) bool' is
// implemented for a value in the attributes, it is called to determine if the
// value matches the one stored in the other attributes. If Equal is not
// implemented, standard equality is used to determine if the two values are
// equal. Note that some types (e.g. maps) aren't comparable by default, so
// they must be wrapped in a struct, or in an alias type, with Equal defined.
func (a *Attributes) Equal(o *Attributes) bool {
if a == nil && o == nil {
return true
}
if a == nil || o == nil {
return false
}
if a == o {
return true
}
m := maps.Collect(o.all())
lenA := 0
for k, v := range a.all() {
lenA++
ov, ok := m[k]
if !ok {
// o missing element of a
return false
}
if eq, ok := v.(interface{ Equal(o any) bool }); ok {
if !eq.Equal(ov) {
return false
}
} else if v != ov {
// Fallback to a standard equality check if Value is unimplemented.
return false
}
}
return lenA == len(m)
}
// String prints the attribute map. If any key or values throughout the map
// implement fmt.Stringer, it calls that method and appends.
func (a *Attributes) String() string {
var sb strings.Builder
sb.WriteString("{")
first := true
for k, v := range a.all() {
if !first {
sb.WriteString(", ")
}
fmt.Fprintf(&sb, "%q: %q ", str(k), str(v))
first = false
}
sb.WriteString("}")
return sb.String()
}
func str(x any) (s string) {
if v, ok := x.(fmt.Stringer); ok {
return fmt.Sprint(v)
} else if v, ok := x.(string); ok {
return v
}
return fmt.Sprintf("<%p>", x)
}
// MarshalJSON helps implement the json.Marshaler interface, thereby rendering
// the Attributes correctly when printing (via pretty.JSON) structs containing
// Attributes as fields.
//
// Is it impossible to unmarshal attributes from a JSON representation and this
// method is meant only for debugging purposes.
func (a *Attributes) MarshalJSON() ([]byte, error) {
return []byte(a.String()), nil
}
// all returns an iterator that yields all key-value pairs in the Attributes
// chain. If a key appears multiple times, only the most recently added value
// is yielded.
func (a *Attributes) all() iter.Seq2[any, any] {
return func(yield func(any, any) bool) {
seen := map[any]bool{}
for cur := a; cur != nil; cur = cur.parent {
if seen[cur.key] {
continue
}
if !yield(cur.key, cur.value) {
return
}
seen[cur.key] = true
}
}
}

View File

@@ -0,0 +1,189 @@
/*
*
* Copyright 2019 gRPC authors.
*
* 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.
*
*/
package attributes_test
import (
"fmt"
"testing"
"google.golang.org/grpc/attributes"
)
type stringVal struct {
s string
}
func (s stringVal) Equal(o any) bool {
os, ok := o.(stringVal)
return ok && s.s == os.s
}
type stringerVal struct {
s string
}
func (s stringerVal) String() string {
return s.s
}
func ExampleAttributes() {
type keyOne struct{}
type keyTwo struct{}
a := attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"})
fmt.Println("Key one:", a.Value(keyOne{}))
fmt.Println("Key two:", a.Value(keyTwo{}))
// Output:
// Key one: 1
// Key two: {two}
}
func ExampleAttributes_WithValue() {
type keyOne struct{}
type keyTwo struct{}
a := attributes.New(keyOne{}, 1)
a = a.WithValue(keyTwo{}, stringVal{s: "two"})
fmt.Println("Key one:", a.Value(keyOne{}))
fmt.Println("Key two:", a.Value(keyTwo{}))
// Output:
// Key one: 1
// Key two: {two}
}
func ExampleAttributes_String() {
type key struct{}
var typedNil *stringerVal
a1 := attributes.New(key{}, typedNil) // typed nil implements [fmt.Stringer]
a2 := attributes.New(key{}, (*stringerVal)(nil)) // typed nil implements [fmt.Stringer]
a3 := attributes.New(key{}, (*stringVal)(nil)) // typed nil not implements [fmt.Stringer]
a4 := attributes.New(key{}, nil) // untyped nil
a5 := attributes.New(key{}, 1)
a6 := attributes.New(key{}, stringerVal{s: "two"})
a7 := attributes.New(key{}, stringVal{s: "two"})
a8 := attributes.New(1, true)
fmt.Println("a1:", a1.String())
fmt.Println("a2:", a2.String())
fmt.Println("a3:", a3.String())
fmt.Println("a4:", a4.String())
fmt.Println("a5:", a5.String())
fmt.Println("a6:", a6.String())
fmt.Println("a7:", a7.String())
fmt.Println("a8:", a8.String())
// Output:
// a1: {"<%!p(attributes_test.key={})>": "<nil>" }
// a2: {"<%!p(attributes_test.key={})>": "<nil>" }
// a3: {"<%!p(attributes_test.key={})>": "<0x0>" }
// a4: {"<%!p(attributes_test.key={})>": "<%!p(<nil>)>" }
// a5: {"<%!p(attributes_test.key={})>": "<%!p(int=1)>" }
// a6: {"<%!p(attributes_test.key={})>": "two" }
// a7: {"<%!p(attributes_test.key={})>": "<%!p(attributes_test.stringVal={two})>" }
// a8: {"<%!p(int=1)>": "<%!p(bool=true)>" }
}
// Test that two attributes with different content are not Equal.
func TestEqual(t *testing.T) {
type keyOne struct{}
type keyTwo struct{}
tests := []struct {
name string
a *attributes.Attributes
b *attributes.Attributes
want bool
}{
{
name: "different_first_value",
a: attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"}),
b: attributes.New(keyOne{}, 2).WithValue(keyTwo{}, stringVal{s: "two"}),
want: false,
},
{
name: "different_second_value",
a: attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "one"}),
b: attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"}),
want: false,
},
{
name: "same",
a: attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"}),
b: attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"}),
want: true,
},
{
name: "subset",
a: attributes.New(keyOne{}, 1),
b: attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"}),
want: false,
},
{
name: "superset",
a: attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"}),
b: attributes.New(keyTwo{}, stringVal{s: "two"}),
want: false,
},
{
name: "a_nil",
a: nil,
b: attributes.New(keyOne{}, 1),
want: false,
},
{
name: "b_nil",
a: attributes.New(keyOne{}, 1),
b: nil,
want: false,
},
{
name: "both_nil",
a: nil,
b: nil,
want: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := tt.a.Equal(tt.b); got != tt.want {
t.Errorf("%+v.Equal(%+v) = %v; want %v", tt.a, tt.b, got, tt.want)
}
// The Equal function should be symmetric, i.e. a.Equals(b) ==
// b.Equals(a).
if got := tt.b.Equal(tt.a); got != tt.want {
t.Errorf("%+v.Equal(%+v) = %v; want %v", tt.b, tt.a, got, tt.want)
}
})
}
}
func BenchmarkWithValue(b *testing.B) {
keys := make([]any, 10)
for i := range 10 {
keys[i] = i
}
b.ReportAllocs()
for b.Loop() {
// 50 endpoints
for range 50 {
a := attributes.New(keys[0], keys[0])
// 10 attributes each.
for j := 1; j < 10; j++ {
a = a.WithValue(keys[j], keys[j])
}
}
}
}

127
authz/audit/audit_logger.go Normal file
View File

@@ -0,0 +1,127 @@
/*
*
* Copyright 2023 gRPC authors.
*
* 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.
*
*/
// Package audit contains interfaces for audit logging during authorization.
package audit
import (
"encoding/json"
"sync"
)
// loggerBuilderRegistry holds a map of audit logger builders and a mutex
// to facilitate thread-safe reading/writing operations.
type loggerBuilderRegistry struct {
mu sync.Mutex
builders map[string]LoggerBuilder
}
var (
registry = loggerBuilderRegistry{
builders: make(map[string]LoggerBuilder),
}
)
// RegisterLoggerBuilder registers the builder in a global map
// using b.Name() as the key.
//
// This should only be called during initialization time (i.e. in an init()
// function). If multiple builders are registered with the same name,
// the one registered last will take effect.
func RegisterLoggerBuilder(b LoggerBuilder) {
registry.mu.Lock()
defer registry.mu.Unlock()
registry.builders[b.Name()] = b
}
// GetLoggerBuilder returns a builder with the given name.
// It returns nil if the builder is not found in the registry.
func GetLoggerBuilder(name string) LoggerBuilder {
registry.mu.Lock()
defer registry.mu.Unlock()
return registry.builders[name]
}
// Event contains information passed to the audit logger as part of an
// audit logging event.
type Event struct {
// FullMethodName is the full method name of the audited RPC, in the format
// of "/pkg.Service/Method". For example, "/helloworld.Greeter/SayHello".
FullMethodName string
// Principal is the identity of the caller. Currently it will only be
// available in certificate-based TLS authentication.
Principal string
// PolicyName is the authorization policy name or the xDS RBAC filter name.
PolicyName string
// MatchedRule is the matched rule or policy name in the xDS RBAC filter.
// It will be empty if there is no match.
MatchedRule string
// Authorized indicates whether the audited RPC is authorized or not.
Authorized bool
}
// LoggerConfig represents an opaque data structure holding an audit
// logger configuration. Concrete types representing configuration of specific
// audit loggers must embed this interface to implement it.
type LoggerConfig interface {
loggerConfig()
}
// Logger is the interface to be implemented by audit loggers.
//
// An audit logger is a logger instance that can be configured via the
// authorization policy API or xDS HTTP RBAC filters. When the authorization
// decision meets the condition for audit, all the configured audit loggers'
// Log() method will be invoked to log that event.
//
// Please refer to
// https://github.com/grpc/proposal/blob/master/A59-audit-logging.md for more
// details about audit logging.
type Logger interface {
// Log performs audit logging for the provided audit event.
//
// This method is invoked in the RPC path and therefore implementations
// must not block.
Log(*Event)
}
// LoggerBuilder is the interface to be implemented by audit logger
// builders that are used at runtime to configure and instantiate audit loggers.
//
// Users who want to implement their own audit logging logic should
// implement this interface, along with the Logger interface, and register
// it by calling RegisterLoggerBuilder() at init time.
//
// Please refer to
// https://github.com/grpc/proposal/blob/master/A59-audit-logging.md for more
// details about audit logging.
type LoggerBuilder interface {
// ParseLoggerConfig parses the given JSON bytes into a structured
// logger config this builder can use to build an audit logger.
ParseLoggerConfig(config json.RawMessage) (LoggerConfig, error)
// Build builds an audit logger with the given logger config.
// This will only be called with valid configs returned from
// ParseLoggerConfig() and any runtime issues such as failing to
// create a file should be handled by the logger implementation instead of
// failing the logger instantiation. So implementers need to make sure it
// can return a logger without error at this stage.
Build(LoggerConfig) Logger
// Name returns the name of logger built by this builder.
// This is used to register and pick the builder.
Name() string
}

View File

@@ -0,0 +1,370 @@
/*
*
* Copyright 2023 gRPC authors.
*
* 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.
*
*/
package audit_test
import (
"context"
"crypto/tls"
"crypto/x509"
"encoding/json"
"io"
"os"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc"
"google.golang.org/grpc/authz"
"google.golang.org/grpc/authz/audit"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/internal/stubserver"
testgrpc "google.golang.org/grpc/interop/grpc_testing"
testpb "google.golang.org/grpc/interop/grpc_testing"
"google.golang.org/grpc/status"
"google.golang.org/grpc/testdata"
_ "google.golang.org/grpc/authz/audit/stdout"
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
type statAuditLogger struct {
authzDecisionStat map[bool]int // Map to hold the counts of authorization decisions
lastEvent *audit.Event // Field to store last received event
}
func (s *statAuditLogger) Log(event *audit.Event) {
s.authzDecisionStat[event.Authorized]++
*s.lastEvent = *event
}
type loggerBuilder struct {
authzDecisionStat map[bool]int
lastEvent *audit.Event
}
func (loggerBuilder) Name() string {
return "stat_logger"
}
func (lb *loggerBuilder) Build(audit.LoggerConfig) audit.Logger {
return &statAuditLogger{
authzDecisionStat: lb.authzDecisionStat,
lastEvent: lb.lastEvent,
}
}
func (*loggerBuilder) ParseLoggerConfig(json.RawMessage) (audit.LoggerConfig, error) {
return nil, nil
}
// TestAuditLogger examines audit logging invocations using four different
// authorization policies. It covers scenarios including a disabled audit,
// auditing both 'allow' and 'deny' outcomes, and separately auditing 'allow'
// and 'deny' outcomes. Additionally, it checks if SPIFFE ID from a certificate
// is propagated correctly.
func (s) TestAuditLogger(t *testing.T) {
// Each test data entry contains an authz policy for a grpc server,
// how many 'allow' and 'deny' outcomes we expect (each test case makes 2
// unary calls and one client-streaming call), and a structure to check if
// the audit.Event fields are properly populated. Additionally, we specify
// directly which authz outcome we expect from each type of call.
tests := []struct {
name string
authzPolicy string
wantAuthzOutcomes map[bool]int
eventContent *audit.Event
wantUnaryCallCode codes.Code
wantStreamingCallCode codes.Code
}{
{
name: "No audit",
authzPolicy: `{
"name": "authz",
"allow_rules": [
{
"name": "allow_UnaryCall",
"request": {
"paths": [
"/grpc.testing.TestService/UnaryCall"
]
}
}
],
"audit_logging_options": {
"audit_condition": "NONE",
"audit_loggers": [
{
"name": "stat_logger",
"config": {},
"is_optional": false
}
]
}
}`,
wantAuthzOutcomes: map[bool]int{true: 0, false: 0},
wantUnaryCallCode: codes.OK,
wantStreamingCallCode: codes.PermissionDenied,
},
{
name: "Allow All Deny Streaming - Audit All",
authzPolicy: `{
"name": "authz",
"allow_rules": [
{
"name": "allow_all",
"request": {
"paths": [
"*"
]
}
}
],
"deny_rules": [
{
"name": "deny_all",
"request": {
"paths": [
"/grpc.testing.TestService/StreamingInputCall"
]
}
}
],
"audit_logging_options": {
"audit_condition": "ON_DENY_AND_ALLOW",
"audit_loggers": [
{
"name": "stat_logger",
"config": {},
"is_optional": false
},
{
"name": "stdout_logger",
"is_optional": false
}
]
}
}`,
wantAuthzOutcomes: map[bool]int{true: 2, false: 1},
eventContent: &audit.Event{
FullMethodName: "/grpc.testing.TestService/StreamingInputCall",
Principal: "spiffe://foo.bar.com/client/workload/1",
PolicyName: "authz",
MatchedRule: "authz_deny_all",
Authorized: false,
},
wantUnaryCallCode: codes.OK,
wantStreamingCallCode: codes.PermissionDenied,
},
{
name: "Allow Unary - Audit Allow",
authzPolicy: `{
"name": "authz",
"allow_rules": [
{
"name": "allow_UnaryCall",
"request": {
"paths": [
"/grpc.testing.TestService/UnaryCall"
]
}
}
],
"audit_logging_options": {
"audit_condition": "ON_ALLOW",
"audit_loggers": [
{
"name": "stat_logger",
"config": {},
"is_optional": false
}
]
}
}`,
wantAuthzOutcomes: map[bool]int{true: 2, false: 0},
wantUnaryCallCode: codes.OK,
wantStreamingCallCode: codes.PermissionDenied,
},
{
name: "Allow Typo - Audit Deny",
authzPolicy: `{
"name": "authz",
"allow_rules": [
{
"name": "allow_UnaryCall",
"request": {
"paths": [
"/grpc.testing.TestService/UnaryCall_Z"
]
}
}
],
"audit_logging_options": {
"audit_condition": "ON_DENY",
"audit_loggers": [
{
"name": "stat_logger",
"config": {},
"is_optional": false
}
]
}
}`,
wantAuthzOutcomes: map[bool]int{true: 0, false: 3},
wantUnaryCallCode: codes.PermissionDenied,
wantStreamingCallCode: codes.PermissionDenied,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
// Construct the credentials for the tests and the stub server
serverCreds := loadServerCreds(t)
clientCreds := loadClientCreds(t)
ss := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
_, err := stream.Recv()
if err != io.EOF {
return err
}
return nil
},
}
// Setup test statAuditLogger, gRPC test server with authzPolicy, unary
// and stream interceptors.
lb := &loggerBuilder{
authzDecisionStat: map[bool]int{true: 0, false: 0},
lastEvent: &audit.Event{},
}
audit.RegisterLoggerBuilder(lb)
i, _ := authz.NewStatic(test.authzPolicy)
s := grpc.NewServer(grpc.Creds(serverCreds), grpc.ChainUnaryInterceptor(i.UnaryInterceptor), grpc.ChainStreamInterceptor(i.StreamInterceptor))
defer s.Stop()
ss.S = s
stubserver.StartTestService(t, ss)
// Setup gRPC test client with certificates containing a SPIFFE Id.
cc, err := grpc.NewClient(ss.Address, grpc.WithTransportCredentials(clientCreds))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", ss.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if _, err := client.UnaryCall(ctx, &testpb.SimpleRequest{}); status.Code(err) != test.wantUnaryCallCode {
t.Errorf("Unexpected UnaryCall fail: got %v want %v", err, test.wantUnaryCallCode)
}
if _, err := client.UnaryCall(ctx, &testpb.SimpleRequest{}); status.Code(err) != test.wantUnaryCallCode {
t.Errorf("Unexpected UnaryCall fail: got %v want %v", err, test.wantUnaryCallCode)
}
stream, err := client.StreamingInputCall(ctx)
if err != nil {
t.Fatalf("StreamingInputCall failed: %v", err)
}
req := &testpb.StreamingInputCallRequest{
Payload: &testpb.Payload{
Body: []byte("hi"),
},
}
if err := stream.Send(req); err != nil && err != io.EOF {
t.Fatalf("stream.Send failed: %v", err)
}
if _, err := stream.CloseAndRecv(); status.Code(err) != test.wantStreamingCallCode {
t.Errorf("Unexpected stream.CloseAndRecv fail: got %v want %v", err, test.wantStreamingCallCode)
}
// Compare expected number of allows/denies with content of the internal
// map of statAuditLogger.
if diff := cmp.Diff(lb.authzDecisionStat, test.wantAuthzOutcomes); diff != "" {
t.Errorf("Authorization decisions do not match\ndiff (-got +want):\n%s", diff)
}
// Compare last event received by statAuditLogger with expected event.
if test.eventContent != nil {
if diff := cmp.Diff(lb.lastEvent, test.eventContent); diff != "" {
t.Errorf("Unexpected message\ndiff (-got +want):\n%s", diff)
}
}
})
}
}
// loadServerCreds constructs TLS containing server certs and CA
func loadServerCreds(t *testing.T) credentials.TransportCredentials {
t.Helper()
cert := loadKeys(t, "x509/server1_cert.pem", "x509/server1_key.pem")
certPool := loadCACerts(t, "x509/client_ca_cert.pem")
return credentials.NewTLS(&tls.Config{
ClientAuth: tls.RequireAndVerifyClientCert,
Certificates: []tls.Certificate{cert},
ClientCAs: certPool,
})
}
// loadClientCreds constructs TLS containing client certs and CA
func loadClientCreds(t *testing.T) credentials.TransportCredentials {
t.Helper()
cert := loadKeys(t, "x509/client_with_spiffe_cert.pem", "x509/client_with_spiffe_key.pem")
roots := loadCACerts(t, "x509/server_ca_cert.pem")
return credentials.NewTLS(&tls.Config{
Certificates: []tls.Certificate{cert},
RootCAs: roots,
ServerName: "x.test.example.com",
})
}
// loadKeys loads X509 key pair from the provided file paths.
// It is used for loading both client and server certificates for the test
func loadKeys(t *testing.T, certPath, key string) tls.Certificate {
t.Helper()
cert, err := tls.LoadX509KeyPair(testdata.Path(certPath), testdata.Path(key))
if err != nil {
t.Fatalf("tls.LoadX509KeyPair(%q, %q) failed: %v", certPath, key, err)
}
return cert
}
// loadCACerts loads CA certificates and constructs x509.CertPool
// It is used for loading both client and server CAs for the test
func loadCACerts(t *testing.T, certPath string) *x509.CertPool {
t.Helper()
ca, err := os.ReadFile(testdata.Path(certPath))
if err != nil {
t.Fatalf("os.ReadFile(%q) failed: %v", certPath, err)
}
roots := x509.NewCertPool()
if !roots.AppendCertsFromPEM(ca) {
t.Fatal("Failed to append certificates")
}
return roots
}

View File

@@ -0,0 +1,110 @@
/*
*
* Copyright 2023 gRPC authors.
*
* 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.
*
*/
// Package stdout defines an stdout audit logger.
package stdout
import (
"encoding/json"
"log"
"os"
"time"
"google.golang.org/grpc/authz/audit"
"google.golang.org/grpc/grpclog"
)
var grpcLogger = grpclog.Component("authz-audit")
// Name is the string to identify this logger type in the registry
const Name = "stdout_logger"
func init() {
audit.RegisterLoggerBuilder(&loggerBuilder{
goLogger: log.New(os.Stdout, "", 0),
})
}
type event struct {
FullMethodName string `json:"rpc_method"`
Principal string `json:"principal"`
PolicyName string `json:"policy_name"`
MatchedRule string `json:"matched_rule"`
Authorized bool `json:"authorized"`
Timestamp string `json:"timestamp"` // Time when the audit event is logged via Log method
}
// logger implements the audit.logger interface by logging to standard output.
type logger struct {
goLogger *log.Logger
}
// Log marshals the audit.Event to json and prints it to standard output.
func (l *logger) Log(event *audit.Event) {
jsonContainer := map[string]any{
"grpc_audit_log": convertEvent(event),
}
jsonBytes, err := json.Marshal(jsonContainer)
if err != nil {
grpcLogger.Errorf("failed to marshal AuditEvent data to JSON: %v", err)
return
}
l.goLogger.Println(string(jsonBytes))
}
// loggerConfig represents the configuration for the stdout logger.
// It is currently empty and implements the audit.Logger interface by embedding it.
type loggerConfig struct {
audit.LoggerConfig
}
type loggerBuilder struct {
goLogger *log.Logger
}
func (loggerBuilder) Name() string {
return Name
}
// Build returns a new instance of the stdout logger.
// Passed in configuration is ignored as the stdout logger does not
// expect any configuration to be provided.
func (lb *loggerBuilder) Build(audit.LoggerConfig) audit.Logger {
return &logger{
goLogger: lb.goLogger,
}
}
// ParseLoggerConfig is a no-op since the stdout logger does not accept any configuration.
func (*loggerBuilder) ParseLoggerConfig(config json.RawMessage) (audit.LoggerConfig, error) {
if len(config) != 0 && string(config) != "{}" {
grpcLogger.Warningf("Stdout logger doesn't support custom configs. Ignoring:\n%s", string(config))
}
return &loggerConfig{}, nil
}
func convertEvent(auditEvent *audit.Event) *event {
return &event{
FullMethodName: auditEvent.FullMethodName,
Principal: auditEvent.Principal,
PolicyName: auditEvent.PolicyName,
MatchedRule: auditEvent.MatchedRule,
Authorized: auditEvent.Authorized,
Timestamp: time.Now().Format(time.RFC3339Nano),
}
}

View File

@@ -0,0 +1,140 @@
/*
*
* Copyright 2023 gRPC authors.
*
* 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.
*
*/
package stdout
import (
"bytes"
"encoding/json"
"log"
"os"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc/authz/audit"
"google.golang.org/grpc/internal/grpctest"
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
func (s) TestStdoutLogger_Log(t *testing.T) {
tests := map[string]struct {
event *audit.Event
wantMessage string
wantErr string
}{
"few fields": {
event: &audit.Event{PolicyName: "test policy", Principal: "test principal"},
wantMessage: `{"fullMethodName":"","principal":"test principal","policyName":"test policy","matchedRule":"","authorized":false`,
},
"all fields": {
event: &audit.Event{
FullMethodName: "/helloworld.Greeter/SayHello",
Principal: "spiffe://example.org/ns/default/sa/default/backend",
PolicyName: "example-policy",
MatchedRule: "dev-access",
Authorized: true,
},
wantMessage: `{"fullMethodName":"/helloworld.Greeter/SayHello",` +
`"principal":"spiffe://example.org/ns/default/sa/default/backend","policyName":"example-policy",` +
`"matchedRule":"dev-access","authorized":true`,
},
}
for name, test := range tests {
t.Run(name, func(t *testing.T) {
before := time.Now().Unix()
var buf bytes.Buffer
builder := &loggerBuilder{goLogger: log.New(&buf, "", 0)}
auditLogger := builder.Build(nil)
auditLogger.Log(test.event)
var container map[string]any
if err := json.Unmarshal(buf.Bytes(), &container); err != nil {
t.Fatalf("Failed to unmarshal audit log event: %v", err)
}
innerEvent := extractEvent(container["grpc_audit_log"].(map[string]any))
if innerEvent.Timestamp == "" {
t.Fatalf("Resulted event has no timestamp: %v", innerEvent)
}
after := time.Now().Unix()
innerEventUnixTime, err := time.Parse(time.RFC3339Nano, innerEvent.Timestamp)
if err != nil {
t.Fatalf("Failed to convert event timestamp into Unix time format: %v", err)
}
if before > innerEventUnixTime.Unix() || after < innerEventUnixTime.Unix() {
t.Errorf("The audit event timestamp is outside of the test interval: test start %v, event timestamp %v, test end %v", before, innerEventUnixTime.Unix(), after)
}
if diff := cmp.Diff(trimEvent(innerEvent), test.event); diff != "" {
t.Fatalf("Unexpected message\ndiff (-got +want):\n%s", diff)
}
})
}
}
func (s) TestStdoutLoggerBuilder_NilConfig(t *testing.T) {
builder := &loggerBuilder{
goLogger: log.New(os.Stdout, "", log.LstdFlags),
}
config, err := builder.ParseLoggerConfig(nil)
if err != nil {
t.Fatalf("Failed to parse stdout logger configuration: %v", err)
}
if l := builder.Build(config); l == nil {
t.Fatal("Failed to build stdout audit logger")
}
}
func (s) TestStdoutLoggerBuilder_Registration(t *testing.T) {
if audit.GetLoggerBuilder("stdout_logger") == nil {
t.Fatal("stdout logger is not registered")
}
}
// extractEvent extracts an stdout.event from a map
// unmarshalled from a logged json message.
func extractEvent(container map[string]any) event {
return event{
FullMethodName: container["rpc_method"].(string),
Principal: container["principal"].(string),
PolicyName: container["policy_name"].(string),
MatchedRule: container["matched_rule"].(string),
Authorized: container["authorized"].(bool),
Timestamp: container["timestamp"].(string),
}
}
// trimEvent converts a logged stdout.event into an audit.Event
// by removing Timestamp field. It is used for comparing events during testing.
func trimEvent(testEvent event) *audit.Event {
return &audit.Event{
FullMethodName: testEvent.FullMethodName,
Principal: testEvent.Principal,
PolicyName: testEvent.PolicyName,
MatchedRule: testEvent.MatchedRule,
Authorized: testEvent.Authorized,
}
}

View File

@@ -0,0 +1,708 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package authz_test
import (
"context"
"crypto/tls"
"crypto/x509"
"io"
"os"
"testing"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/authz"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/internal/stubserver"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"google.golang.org/grpc/testdata"
testgrpc "google.golang.org/grpc/interop/grpc_testing"
testpb "google.golang.org/grpc/interop/grpc_testing"
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
var authzTests = map[string]struct {
authzPolicy string
md metadata.MD
wantStatus *status.Status
}{
"DeniesRPCMatchInDenyNoMatchInAllow": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_StreamingOutputCall",
"request": {
"paths":
[
"/grpc.testing.TestService/StreamingOutputCall"
]
}
}
],
"deny_rules":
[
{
"name": "deny_TestServiceCalls",
"request": {
"paths":
[
"/grpc.testing.TestService/*"
],
"headers":
[
{
"key": "key-abc",
"values":
[
"val-abc",
"val-def"
]
}
]
}
}
]
}`,
md: metadata.Pairs("key-abc", "val-abc"),
wantStatus: status.New(codes.PermissionDenied, "authz: unauthorized RPC request rejected"),
},
"DeniesRPCMatchInDenyAndAllow": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_all",
"request": {
"paths":
[
"*"
]
}
}
],
"deny_rules":
[
{
"name": "deny_all",
"request": {
"paths":
[
"*"
]
}
}
]
}`,
wantStatus: status.New(codes.PermissionDenied, "authz: unauthorized RPC request rejected"),
},
"AllowsRPCNoMatchInDenyMatchInAllow": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_all"
}
],
"deny_rules":
[
{
"name": "deny_TestServiceCalls",
"request": {
"paths":
[
"/grpc.testing.TestService/UnaryCall",
"/grpc.testing.TestService/StreamingInputCall"
],
"headers":
[
{
"key": "key-abc",
"values":
[
"val-abc",
"val-def"
]
}
]
}
}
]
}`,
md: metadata.Pairs("key-xyz", "val-xyz"),
wantStatus: status.New(codes.OK, ""),
},
"DeniesRPCNoMatchInDenyAndAllow": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_some_user",
"source": {
"principals":
[
"some_user"
]
}
}
],
"deny_rules":
[
{
"name": "deny_StreamingOutputCall",
"request": {
"paths":
[
"/grpc.testing.TestService/StreamingOutputCall"
]
}
}
]
}`,
wantStatus: status.New(codes.PermissionDenied, "authz: unauthorized RPC request rejected"),
},
"AllowsRPCEmptyDenyMatchInAllow": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_UnaryCall",
"request":
{
"paths":
[
"/grpc.testing.TestService/UnaryCall"
]
}
},
{
"name": "allow_StreamingInputCall",
"request":
{
"paths":
[
"/grpc.testing.TestService/StreamingInputCall"
]
}
}
]
}`,
wantStatus: status.New(codes.OK, ""),
},
"DeniesRPCEmptyDenyNoMatchInAllow": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_StreamingOutputCall",
"request":
{
"paths":
[
"/grpc.testing.TestService/StreamingOutputCall"
]
}
}
]
}`,
wantStatus: status.New(codes.PermissionDenied, "authz: unauthorized RPC request rejected"),
},
"DeniesRPCRequestWithPrincipalsFieldOnUnauthenticatedConnection": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_authenticated",
"source": {
"principals": ["*", ""]
}
}
]
}`,
wantStatus: status.New(codes.PermissionDenied, "authz: unauthorized RPC request rejected"),
},
"DeniesRPCRequestNoMatchInAllowFailsPresenceMatch": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_TestServiceCalls",
"request": {
"paths":
[
"/grpc.testing.TestService/*"
],
"headers":
[
{
"key": "key-abc",
"values":
[
"*"
]
}
]
}
}
]
}`,
md: metadata.Pairs("key-abc", ""),
wantStatus: status.New(codes.PermissionDenied, "authz: unauthorized RPC request rejected"),
},
}
func (s) TestStaticPolicyEnd2End(t *testing.T) {
for name, test := range authzTests {
t.Run(name, func(t *testing.T) {
// Start a gRPC server with gRPC authz unary and stream server interceptors.
i, _ := authz.NewStatic(test.authzPolicy)
stub := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
StreamingInputCallF: func(stream testgrpc.TestService_StreamingInputCallServer) error {
for {
_, err := stream.Recv()
if err == io.EOF {
return stream.SendAndClose(&testpb.StreamingInputCallResponse{})
}
if err != nil {
return err
}
}
},
S: grpc.NewServer(grpc.ChainUnaryInterceptor(i.UnaryInterceptor), grpc.ChainStreamInterceptor(i.StreamInterceptor)),
}
stubserver.StartTestService(t, stub)
defer stub.Stop()
// Establish a connection to the server.
cc, err := grpc.NewClient(stub.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", stub.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ctx = metadata.NewOutgoingContext(ctx, test.md)
// Verifying authorization decision for Unary RPC.
_, err = client.UnaryCall(ctx, &testpb.SimpleRequest{})
if got := status.Convert(err); got.Code() != test.wantStatus.Code() || got.Message() != test.wantStatus.Message() {
t.Fatalf("[UnaryCall] error want:{%v} got:{%v}", test.wantStatus.Err(), got.Err())
}
// Verifying authorization decision for Streaming RPC.
stream, err := client.StreamingInputCall(ctx)
if err != nil {
t.Fatalf("failed StreamingInputCall err: %v", err)
}
req := &testpb.StreamingInputCallRequest{
Payload: &testpb.Payload{
Body: []byte("hi"),
},
}
if err := stream.Send(req); err != nil && err != io.EOF {
t.Fatalf("failed stream.Send err: %v", err)
}
_, err = stream.CloseAndRecv()
if got := status.Convert(err); got.Code() != test.wantStatus.Code() || got.Message() != test.wantStatus.Message() {
t.Fatalf("[StreamingCall] error want:{%v} got:{%v}", test.wantStatus.Err(), got.Err())
}
})
}
}
func (s) TestAllowsRPCRequestWithPrincipalsFieldOnTLSAuthenticatedConnection(t *testing.T) {
authzPolicy := `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_authenticated",
"source": {
"principals": ["*", ""]
}
}
]
}`
// Start a gRPC server with gRPC authz unary server interceptor.
i, _ := authz.NewStatic(authzPolicy)
creds, err := credentials.NewServerTLSFromFile(testdata.Path("x509/server1_cert.pem"), testdata.Path("x509/server1_key.pem"))
if err != nil {
t.Fatalf("failed to generate credentials: %v", err)
}
stub := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
S: grpc.NewServer(grpc.Creds(creds), grpc.ChainUnaryInterceptor(i.UnaryInterceptor)),
}
stubserver.StartTestService(t, stub)
defer stub.S.Stop()
// Establish a connection to the server.
creds, err = credentials.NewClientTLSFromFile(testdata.Path("x509/server_ca_cert.pem"), "x.test.example.com")
if err != nil {
t.Fatalf("failed to load credentials: %v", err)
}
cc, err := grpc.NewClient(stub.Address, grpc.WithTransportCredentials(creds))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", stub.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// Verifying authorization decision.
if _, err = client.UnaryCall(ctx, &testpb.SimpleRequest{}); err != nil {
t.Fatalf("client.UnaryCall(_, _) = %v; want nil", err)
}
}
func (s) TestAllowsRPCRequestWithPrincipalsFieldOnMTLSAuthenticatedConnection(t *testing.T) {
authzPolicy := `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_authenticated",
"source": {
"principals": ["*", ""]
}
}
]
}`
// Start a gRPC server with gRPC authz unary server interceptor.
i, _ := authz.NewStatic(authzPolicy)
cert, err := tls.LoadX509KeyPair(testdata.Path("x509/server1_cert.pem"), testdata.Path("x509/server1_key.pem"))
if err != nil {
t.Fatalf("tls.LoadX509KeyPair(x509/server1_cert.pem, x509/server1_key.pem) failed: %v", err)
}
ca, err := os.ReadFile(testdata.Path("x509/client_ca_cert.pem"))
if err != nil {
t.Fatalf("os.ReadFile(x509/client_ca_cert.pem) failed: %v", err)
}
certPool := x509.NewCertPool()
if !certPool.AppendCertsFromPEM(ca) {
t.Fatal("failed to append certificates")
}
creds := credentials.NewTLS(&tls.Config{
ClientAuth: tls.RequireAndVerifyClientCert,
Certificates: []tls.Certificate{cert},
ClientCAs: certPool,
})
stub := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
S: grpc.NewServer(grpc.Creds(creds), grpc.ChainUnaryInterceptor(i.UnaryInterceptor)),
}
stubserver.StartTestService(t, stub)
defer stub.Stop()
// Establish a connection to the server.
cert, err = tls.LoadX509KeyPair(testdata.Path("x509/client1_cert.pem"), testdata.Path("x509/client1_key.pem"))
if err != nil {
t.Fatalf("tls.LoadX509KeyPair(x509/client1_cert.pem, x509/client1_key.pem) failed: %v", err)
}
ca, err = os.ReadFile(testdata.Path("x509/server_ca_cert.pem"))
if err != nil {
t.Fatalf("os.ReadFile(x509/server_ca_cert.pem) failed: %v", err)
}
roots := x509.NewCertPool()
if !roots.AppendCertsFromPEM(ca) {
t.Fatal("failed to append certificates")
}
creds = credentials.NewTLS(&tls.Config{
Certificates: []tls.Certificate{cert},
RootCAs: roots,
ServerName: "x.test.example.com",
})
cc, err := grpc.NewClient(stub.Address, grpc.WithTransportCredentials(creds))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", stub.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// Verifying authorization decision.
if _, err = client.UnaryCall(ctx, &testpb.SimpleRequest{}); err != nil {
t.Fatalf("client.UnaryCall(_, _) = %v; want nil", err)
}
}
func (s) TestFileWatcherEnd2End(t *testing.T) {
for name, test := range authzTests {
t.Run(name, func(t *testing.T) {
file := createTmpPolicyFile(t, name, []byte(test.authzPolicy))
i, _ := authz.NewFileWatcher(file, 1*time.Second)
defer i.Close()
stub := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
StreamingInputCallF: func(stream testgrpc.TestService_StreamingInputCallServer) error {
for {
_, err := stream.Recv()
if err == io.EOF {
return stream.SendAndClose(&testpb.StreamingInputCallResponse{})
}
if err != nil {
return err
}
}
},
// Start a gRPC server with gRPC authz unary and stream server interceptors.
S: grpc.NewServer(grpc.ChainUnaryInterceptor(i.UnaryInterceptor), grpc.ChainStreamInterceptor(i.StreamInterceptor)),
}
stubserver.StartTestService(t, stub)
defer stub.Stop()
// Establish a connection to the server.
cc, err := grpc.NewClient(stub.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", stub.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ctx = metadata.NewOutgoingContext(ctx, test.md)
// Verifying authorization decision for Unary RPC.
_, err = client.UnaryCall(ctx, &testpb.SimpleRequest{})
if got := status.Convert(err); got.Code() != test.wantStatus.Code() || got.Message() != test.wantStatus.Message() {
t.Fatalf("[UnaryCall] error want:{%v} got:{%v}", test.wantStatus.Err(), got.Err())
}
// Verifying authorization decision for Streaming RPC.
stream, err := client.StreamingInputCall(ctx)
if err != nil {
t.Fatalf("failed StreamingInputCall : %v", err)
}
req := &testpb.StreamingInputCallRequest{
Payload: &testpb.Payload{
Body: []byte("hi"),
},
}
if err := stream.Send(req); err != nil && err != io.EOF {
t.Fatalf("failed stream.Send : %v", err)
}
_, err = stream.CloseAndRecv()
if got := status.Convert(err); got.Code() != test.wantStatus.Code() || got.Message() != test.wantStatus.Message() {
t.Fatalf("[StreamingCall] error want:{%v} got:{%v}", test.wantStatus.Err(), got.Err())
}
})
}
}
func retryUntil(ctx context.Context, tsc testgrpc.TestServiceClient, want *status.Status) (lastErr error) {
for ctx.Err() == nil {
_, lastErr = tsc.UnaryCall(ctx, &testpb.SimpleRequest{})
if s := status.Convert(lastErr); s.Code() == want.Code() && s.Message() == want.Message() {
return nil
}
time.Sleep(20 * time.Millisecond)
}
return lastErr
}
func (s) TestFileWatcher_ValidPolicyRefresh(t *testing.T) {
valid1 := authzTests["DeniesRPCMatchInDenyAndAllow"]
file := createTmpPolicyFile(t, "valid_policy_refresh", []byte(valid1.authzPolicy))
i, _ := authz.NewFileWatcher(file, 100*time.Millisecond)
defer i.Close()
stub := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
// Start a gRPC server with gRPC authz unary server interceptor.
S: grpc.NewServer(grpc.ChainUnaryInterceptor(i.UnaryInterceptor)),
}
stubserver.StartTestService(t, stub)
defer stub.Stop()
// Establish a connection to the server.
cc, err := grpc.NewClient(stub.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", stub.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// Verifying authorization decision.
_, err = client.UnaryCall(ctx, &testpb.SimpleRequest{})
if got := status.Convert(err); got.Code() != valid1.wantStatus.Code() || got.Message() != valid1.wantStatus.Message() {
t.Fatalf("client.UnaryCall(_, _) = %v; want = %v", got.Err(), valid1.wantStatus.Err())
}
// Rewrite the file with a different valid authorization policy.
valid2 := authzTests["AllowsRPCEmptyDenyMatchInAllow"]
if err := os.WriteFile(file, []byte(valid2.authzPolicy), os.ModePerm); err != nil {
t.Fatalf("os.WriteFile(%q) failed: %v", file, err)
}
// Verifying authorization decision.
if got := retryUntil(ctx, client, valid2.wantStatus); got != nil {
t.Fatalf("client.UnaryCall(_, _) = %v; want = %v", got, valid2.wantStatus.Err())
}
}
func (s) TestFileWatcher_InvalidPolicySkipReload(t *testing.T) {
valid := authzTests["DeniesRPCMatchInDenyAndAllow"]
file := createTmpPolicyFile(t, "invalid_policy_skip_reload", []byte(valid.authzPolicy))
i, _ := authz.NewFileWatcher(file, 20*time.Millisecond)
defer i.Close()
stub := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
// Start a gRPC server with gRPC authz unary server interceptors.
S: grpc.NewServer(grpc.ChainUnaryInterceptor(i.UnaryInterceptor)),
}
stubserver.StartTestService(t, stub)
defer stub.Stop()
// Establish a connection to the server.
cc, err := grpc.NewClient(stub.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", stub.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// Verifying authorization decision.
_, err = client.UnaryCall(ctx, &testpb.SimpleRequest{})
if got := status.Convert(err); got.Code() != valid.wantStatus.Code() || got.Message() != valid.wantStatus.Message() {
t.Fatalf("client.UnaryCall(_, _) = %v; want = %v", got.Err(), valid.wantStatus.Err())
}
// Skips the invalid policy update, and continues to use the valid policy.
if err := os.WriteFile(file, []byte("{}"), os.ModePerm); err != nil {
t.Fatalf("os.WriteFile(%q) failed: %v", file, err)
}
// Wait 40 ms for background go routine to read updated files.
time.Sleep(40 * time.Millisecond)
// Verifying authorization decision.
_, err = client.UnaryCall(ctx, &testpb.SimpleRequest{})
if got := status.Convert(err); got.Code() != valid.wantStatus.Code() || got.Message() != valid.wantStatus.Message() {
t.Fatalf("client.UnaryCall(_, _) = %v; want = %v", got.Err(), valid.wantStatus.Err())
}
}
func (s) TestFileWatcher_RecoversFromReloadFailure(t *testing.T) {
valid1 := authzTests["DeniesRPCMatchInDenyAndAllow"]
file := createTmpPolicyFile(t, "recovers_from_reload_failure", []byte(valid1.authzPolicy))
i, _ := authz.NewFileWatcher(file, 100*time.Millisecond)
defer i.Close()
stub := &stubserver.StubServer{
UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
return &testpb.SimpleResponse{}, nil
},
S: grpc.NewServer(grpc.ChainUnaryInterceptor(i.UnaryInterceptor)),
}
stubserver.StartTestService(t, stub)
defer stub.Stop()
// Establish a connection to the server.
cc, err := grpc.NewClient(stub.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient(%v) failed: %v", stub.Address, err)
}
defer cc.Close()
client := testgrpc.NewTestServiceClient(cc)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// Verifying authorization decision.
_, err = client.UnaryCall(ctx, &testpb.SimpleRequest{})
if got := status.Convert(err); got.Code() != valid1.wantStatus.Code() || got.Message() != valid1.wantStatus.Message() {
t.Fatalf("client.UnaryCall(_, _) = %v; want = %v", got.Err(), valid1.wantStatus.Err())
}
// Skips the invalid policy update, and continues to use the valid policy.
if err := os.WriteFile(file, []byte("{}"), os.ModePerm); err != nil {
t.Fatalf("os.WriteFile(%q) failed: %v", file, err)
}
// Wait 120 ms for background go routine to read updated files.
time.Sleep(120 * time.Millisecond)
// Verifying authorization decision.
_, err = client.UnaryCall(ctx, &testpb.SimpleRequest{})
if got := status.Convert(err); got.Code() != valid1.wantStatus.Code() || got.Message() != valid1.wantStatus.Message() {
t.Fatalf("client.UnaryCall(_, _) = %v; want = %v", got.Err(), valid1.wantStatus.Err())
}
// Rewrite the file with a different valid authorization policy.
valid2 := authzTests["AllowsRPCEmptyDenyMatchInAllow"]
if err := os.WriteFile(file, []byte(valid2.authzPolicy), os.ModePerm); err != nil {
t.Fatalf("os.WriteFile(%q) failed: %v", file, err)
}
// Verifying authorization decision.
if got := retryUntil(ctx, client, valid2.wantStatus); got != nil {
t.Fatalf("client.UnaryCall(_, _) = %v; want = %v", got, valid2.wantStatus.Err())
}
}

View File

@@ -0,0 +1,211 @@
/*
* Copyright 2021 gRPC authors.
*
* 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.
*/
package authz
import (
"bytes"
"context"
"fmt"
"os"
"sync/atomic"
"time"
"unsafe"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/internal/xds/rbac"
"google.golang.org/grpc/status"
)
var logger = grpclog.Component("authz")
// StaticInterceptor contains engines used to make authorization decisions. It
// either contains two engines deny engine followed by an allow engine or only
// one allow engine.
type StaticInterceptor struct {
engines rbac.ChainEngine
}
// NewStatic returns a new StaticInterceptor from a static authorization policy
// JSON string.
func NewStatic(authzPolicy string) (*StaticInterceptor, error) {
rbacs, policyName, err := translatePolicy(authzPolicy)
if err != nil {
return nil, err
}
chainEngine, err := rbac.NewChainEngine(rbacs, policyName)
if err != nil {
return nil, err
}
return &StaticInterceptor{*chainEngine}, nil
}
// UnaryInterceptor intercepts incoming Unary RPC requests.
// Only authorized requests are allowed to pass. Otherwise, an unauthorized
// error is returned to the client.
func (i *StaticInterceptor) UnaryInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
err := i.engines.IsAuthorized(ctx)
if err != nil {
if status.Code(err) == codes.PermissionDenied {
if logger.V(2) {
logger.Infof("unauthorized RPC request rejected: %v", err)
}
return nil, status.Errorf(codes.PermissionDenied, "authz: unauthorized RPC request rejected")
}
return nil, err
}
return handler(ctx, req)
}
// StreamInterceptor intercepts incoming Stream RPC requests.
// Only authorized requests are allowed to pass. Otherwise, an unauthorized
// error is returned to the client.
func (i *StaticInterceptor) StreamInterceptor(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
err := i.engines.IsAuthorized(ss.Context())
if err != nil {
if status.Code(err) == codes.PermissionDenied {
if logger.V(2) {
logger.Infof("unauthorized RPC request rejected: %v", err)
}
return status.Errorf(codes.PermissionDenied, "authz: unauthorized RPC request rejected")
}
return err
}
return handler(srv, ss)
}
// FileWatcherInterceptor contains details used to make authorization decisions
// by watching a file path that contains authorization policy in JSON format.
type FileWatcherInterceptor struct {
options FileWatcherOptions
internalInterceptor unsafe.Pointer // *StaticInterceptor
policyContents []byte
cancel context.CancelFunc
}
// FileWatcherOptions contains configuration options for the
// FileWatcherInterceptor.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
type FileWatcherOptions struct {
// PolicyFile contains a JSON string of the authorization policy.
PolicyFile string
// RefreshDuration is the delay between policy refreshes.
RefreshDuration time.Duration
// OnPolicyUpdate is a callback to be invoked when a policy is
// loaded/updated. The loaded policy string is passed as an argument.
//
// The callback is executed synchronously, so should complete quickly or
// risk blocking future updates.
OnPolicyUpdate func(string)
}
// NewFileWatcher returns a new FileWatcherInterceptor from a policy file
// that contains JSON string of authorization policy and a refresh duration to
// specify the amount of time between policy refreshes.
func NewFileWatcher(file string, duration time.Duration) (*FileWatcherInterceptor, error) {
return NewFileWatcherWithOptions(FileWatcherOptions{PolicyFile: file, RefreshDuration: duration, OnPolicyUpdate: nil})
}
// NewFileWatcherWithOptions returns a new FileWatcherInterceptor from a set of
// options.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
func NewFileWatcherWithOptions(options FileWatcherOptions) (*FileWatcherInterceptor, error) {
if options.PolicyFile == "" {
return nil, fmt.Errorf("authz: authorization policy file path is empty")
}
if options.RefreshDuration <= time.Duration(0) {
return nil, fmt.Errorf("authz: requires refresh interval(%v) greater than 0s", options.RefreshDuration)
}
i := &FileWatcherInterceptor{options: options}
if err := i.updateInternalInterceptor(); err != nil {
return nil, err
}
ctx, cancel := context.WithCancel(context.Background())
i.cancel = cancel
// Create a background go routine for policy refresh.
go i.run(ctx)
return i, nil
}
func (i *FileWatcherInterceptor) run(ctx context.Context) {
ticker := time.NewTicker(i.options.RefreshDuration)
for {
if err := i.updateInternalInterceptor(); err != nil {
logger.Warningf("authorization policy reload status err: %v", err)
}
select {
case <-ctx.Done():
ticker.Stop()
return
case <-ticker.C:
}
}
}
// updateInternalInterceptor checks if the policy file that is watching has changed,
// and if so, updates the internalInterceptor with the policy. Unlike the
// constructor, if there is an error in reading the file or parsing the policy, the
// previous internalInterceptors will not be replaced.
func (i *FileWatcherInterceptor) updateInternalInterceptor() error {
policyContents, err := os.ReadFile(i.options.PolicyFile)
if err != nil {
return fmt.Errorf("policyFile(%s) read failed: %v", i.options.PolicyFile, err)
}
if bytes.Equal(i.policyContents, policyContents) {
return nil
}
i.policyContents = policyContents
policyContentsString := string(policyContents)
interceptor, err := NewStatic(policyContentsString)
if err != nil {
return err
}
atomic.StorePointer(&i.internalInterceptor, unsafe.Pointer(interceptor))
logger.Infof("authorization policy reload status: successfully loaded new policy %v", policyContentsString)
if i.options.OnPolicyUpdate != nil {
i.options.OnPolicyUpdate(policyContentsString)
}
return nil
}
// Close cleans up resources allocated by the interceptor.
func (i *FileWatcherInterceptor) Close() {
i.cancel()
}
// UnaryInterceptor intercepts incoming Unary RPC requests.
// Only authorized requests are allowed to pass. Otherwise, an unauthorized
// error is returned to the client.
func (i *FileWatcherInterceptor) UnaryInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
return ((*StaticInterceptor)(atomic.LoadPointer(&i.internalInterceptor))).UnaryInterceptor(ctx, req, info, handler)
}
// StreamInterceptor intercepts incoming Stream RPC requests.
// Only authorized requests are allowed to pass. Otherwise, an unauthorized
// error is returned to the client.
func (i *FileWatcherInterceptor) StreamInterceptor(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
return ((*StaticInterceptor)(atomic.LoadPointer(&i.internalInterceptor))).StreamInterceptor(srv, ss, info, handler)
}

View File

@@ -0,0 +1,174 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package authz_test
import (
"context"
"fmt"
"os"
"path"
"testing"
"time"
"google.golang.org/grpc/authz"
)
const defaultTestTimeout = 10 * time.Second
func createTmpPolicyFile(t *testing.T, dirSuffix string, policy []byte) string {
t.Helper()
// Create a temp directory. Passing an empty string for the first argument
// uses the system temp directory.
dir, err := os.MkdirTemp("", dirSuffix)
if err != nil {
t.Fatalf("os.MkdirTemp() failed: %v", err)
}
t.Logf("Using tmpdir: %s", dir)
// Write policy into file.
filename := path.Join(dir, "policy.json")
if err := os.WriteFile(filename, policy, os.ModePerm); err != nil {
t.Fatalf("os.WriteFile(%q) failed: %v", filename, err)
}
t.Logf("Wrote policy %s to file at %s", string(policy), filename)
return filename
}
func (s) TestNewStatic(t *testing.T) {
tests := map[string]struct {
authzPolicy string
wantErr error
}{
"InvalidPolicyFailsToCreateInterceptor": {
authzPolicy: `{}`,
wantErr: fmt.Errorf(`"name" is not present`),
},
"ValidPolicyCreatesInterceptor": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_all"
}
]
}`,
},
}
for name, test := range tests {
t.Run(name, func(t *testing.T) {
if _, err := authz.NewStatic(test.authzPolicy); fmt.Sprint(err) != fmt.Sprint(test.wantErr) {
t.Fatalf("NewStatic(%v) returned err: %v, want err: %v", test.authzPolicy, err, test.wantErr)
}
})
}
}
func (s) TestNewFileWatcher(t *testing.T) {
tests := map[string]struct {
authzPolicy string
refreshDuration time.Duration
wantErr error
}{
"InvalidRefreshDurationFailsToCreateInterceptor": {
refreshDuration: time.Duration(0),
wantErr: fmt.Errorf("authz: requires refresh interval(0s) greater than 0s"),
},
"InvalidPolicyFailsToCreateInterceptor": {
authzPolicy: `{}`,
refreshDuration: time.Duration(1),
wantErr: fmt.Errorf(`"name" is not present`),
},
"ValidPolicyCreatesInterceptor": {
authzPolicy: `{
"name": "authz",
"allow_rules":
[
{
"name": "allow_all"
}
]
}`,
refreshDuration: time.Duration(1),
},
}
for name, test := range tests {
t.Run(name, func(t *testing.T) {
file := createTmpPolicyFile(t, name, []byte(test.authzPolicy))
i, err := authz.NewFileWatcher(file, test.refreshDuration)
if fmt.Sprint(err) != fmt.Sprint(test.wantErr) {
t.Fatalf("NewFileWatcher(%v) returned err: %v, want err: %v", test.authzPolicy, err, test.wantErr)
}
if i != nil {
i.Close()
}
})
}
}
func (s) TestOnPolicyUpdate(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
updates := make(chan string, 1)
onPolicyUpdate := func(s string) {
updates <- s
}
content := `{"name": "foo1", "allow_rules":[{"name":"bar"}]}`
file := createTmpPolicyFile(t, "onpolicyupdate", []byte(content))
opts := authz.FileWatcherOptions{PolicyFile: file, RefreshDuration: 50 * time.Millisecond, OnPolicyUpdate: onPolicyUpdate}
i, err := authz.NewFileWatcherWithOptions(opts)
if err != nil {
t.Fatalf("NewFileWatcherWithOptions(%v) returned err: %v", opts, err)
}
defer i.Close()
select {
case <-ctx.Done():
t.Fatalf("Timeout waiting for policy update")
case update := <-updates:
if update != content {
t.Fatalf("Unexpected contents on first load of policy file: got=%v, want=%v", update, content)
}
}
// Tweak the file, expect an update.
content = `{"name": "foo2", "allow_rules":[{"name":"bar"}]}`
if err := os.WriteFile(file, []byte(content), os.ModePerm); err != nil {
t.Fatalf("os.WriteFile(%q) failed: %v", file, err)
}
select {
case <-ctx.Done():
t.Fatalf("Timeout waiting for policy update")
case update := <-updates:
if update != content {
t.Fatalf("Unexpected contents after policy file changed: got=%v, want=%v", update, content)
}
}
sCtx, sCancel := context.WithTimeout(ctx, 100*time.Millisecond)
defer sCancel()
select {
case <-updates:
t.Fatal("OnPolicyUpdate callback invoked more times than expected")
case <-sCtx.Done():
}
}

398
authz/rbac_translator.go Normal file
View File

@@ -0,0 +1,398 @@
/*
* Copyright 2021 gRPC authors.
*
* 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.
*/
// Package authz exposes methods to manage authorization within gRPC.
//
// # Experimental
//
// Notice: This package is EXPERIMENTAL and may be changed or removed
// in a later release.
package authz
import (
"bytes"
"encoding/json"
"fmt"
"strings"
v1xdsudpatypepb "github.com/cncf/xds/go/udpa/type/v1"
v3corepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
v3rbacpb "github.com/envoyproxy/go-control-plane/envoy/config/rbac/v3"
v3routepb "github.com/envoyproxy/go-control-plane/envoy/config/route/v3"
v3matcherpb "github.com/envoyproxy/go-control-plane/envoy/type/matcher/v3"
"google.golang.org/protobuf/types/known/anypb"
"google.golang.org/protobuf/types/known/structpb"
)
// This is used when converting a custom config from raw JSON to a TypedStruct
// The TypeURL of the TypeStruct will be "grpc.authz.audit_logging/<name>"
const typeURLPrefix = "grpc.authz.audit_logging/"
type header struct {
Key string
Values []string
}
type peer struct {
Principals []string
}
type request struct {
Paths []string
Headers []header
}
type rule struct {
Name string
Source peer
Request request
}
type auditLogger struct {
Name string `json:"name"`
Config *structpb.Struct `json:"config"`
IsOptional bool `json:"is_optional"`
}
type auditLoggingOptions struct {
AuditCondition string `json:"audit_condition"`
AuditLoggers []*auditLogger `json:"audit_loggers"`
}
// Represents the SDK authorization policy provided by user.
type authorizationPolicy struct {
Name string
DenyRules []rule `json:"deny_rules"`
AllowRules []rule `json:"allow_rules"`
AuditLoggingOptions auditLoggingOptions `json:"audit_logging_options"`
}
func principalOr(principals []*v3rbacpb.Principal) *v3rbacpb.Principal {
return &v3rbacpb.Principal{
Identifier: &v3rbacpb.Principal_OrIds{
OrIds: &v3rbacpb.Principal_Set{
Ids: principals,
},
},
}
}
func permissionOr(permission []*v3rbacpb.Permission) *v3rbacpb.Permission {
return &v3rbacpb.Permission{
Rule: &v3rbacpb.Permission_OrRules{
OrRules: &v3rbacpb.Permission_Set{
Rules: permission,
},
},
}
}
func permissionAnd(permission []*v3rbacpb.Permission) *v3rbacpb.Permission {
return &v3rbacpb.Permission{
Rule: &v3rbacpb.Permission_AndRules{
AndRules: &v3rbacpb.Permission_Set{
Rules: permission,
},
},
}
}
func getStringMatcher(value string) *v3matcherpb.StringMatcher {
switch {
case value == "*":
return &v3matcherpb.StringMatcher{
MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{
SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
}
case strings.HasSuffix(value, "*"):
prefix := strings.TrimSuffix(value, "*")
return &v3matcherpb.StringMatcher{
MatchPattern: &v3matcherpb.StringMatcher_Prefix{Prefix: prefix},
}
case strings.HasPrefix(value, "*"):
suffix := strings.TrimPrefix(value, "*")
return &v3matcherpb.StringMatcher{
MatchPattern: &v3matcherpb.StringMatcher_Suffix{Suffix: suffix},
}
default:
return &v3matcherpb.StringMatcher{
MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: value},
}
}
}
func getHeaderMatcher(key, value string) *v3routepb.HeaderMatcher {
switch {
case value == "*":
return &v3routepb.HeaderMatcher{
Name: key,
HeaderMatchSpecifier: &v3routepb.HeaderMatcher_SafeRegexMatch{
SafeRegexMatch: &v3matcherpb.RegexMatcher{Regex: ".+"}},
}
case strings.HasSuffix(value, "*"):
prefix := strings.TrimSuffix(value, "*")
return &v3routepb.HeaderMatcher{
Name: key,
HeaderMatchSpecifier: &v3routepb.HeaderMatcher_PrefixMatch{PrefixMatch: prefix},
}
case strings.HasPrefix(value, "*"):
suffix := strings.TrimPrefix(value, "*")
return &v3routepb.HeaderMatcher{
Name: key,
HeaderMatchSpecifier: &v3routepb.HeaderMatcher_SuffixMatch{SuffixMatch: suffix},
}
default:
return &v3routepb.HeaderMatcher{
Name: key,
HeaderMatchSpecifier: &v3routepb.HeaderMatcher_ExactMatch{ExactMatch: value},
}
}
}
func parsePrincipalNames(principalNames []string) []*v3rbacpb.Principal {
ps := make([]*v3rbacpb.Principal, 0, len(principalNames))
for _, principalName := range principalNames {
newPrincipalName := &v3rbacpb.Principal{
Identifier: &v3rbacpb.Principal_Authenticated_{
Authenticated: &v3rbacpb.Principal_Authenticated{
PrincipalName: getStringMatcher(principalName),
},
}}
ps = append(ps, newPrincipalName)
}
return ps
}
func parsePeer(source peer) *v3rbacpb.Principal {
if len(source.Principals) == 0 {
return &v3rbacpb.Principal{
Identifier: &v3rbacpb.Principal_Any{
Any: true,
},
}
}
return principalOr(parsePrincipalNames(source.Principals))
}
func parsePaths(paths []string) []*v3rbacpb.Permission {
ps := make([]*v3rbacpb.Permission, 0, len(paths))
for _, path := range paths {
newPath := &v3rbacpb.Permission{
Rule: &v3rbacpb.Permission_UrlPath{
UrlPath: &v3matcherpb.PathMatcher{
Rule: &v3matcherpb.PathMatcher_Path{Path: getStringMatcher(path)}}}}
ps = append(ps, newPath)
}
return ps
}
func parseHeaderValues(key string, values []string) []*v3rbacpb.Permission {
vs := make([]*v3rbacpb.Permission, 0, len(values))
for _, value := range values {
newHeader := &v3rbacpb.Permission{
Rule: &v3rbacpb.Permission_Header{
Header: getHeaderMatcher(key, value)}}
vs = append(vs, newHeader)
}
return vs
}
var unsupportedHeaders = map[string]bool{
"host": true,
"connection": true,
"keep-alive": true,
"proxy-authenticate": true,
"proxy-authorization": true,
"te": true,
"trailer": true,
"transfer-encoding": true,
"upgrade": true,
}
func unsupportedHeader(key string) bool {
return key[0] == ':' || strings.HasPrefix(key, "grpc-") || unsupportedHeaders[key]
}
func parseHeaders(headers []header) ([]*v3rbacpb.Permission, error) {
hs := make([]*v3rbacpb.Permission, 0, len(headers))
for i, header := range headers {
if header.Key == "" {
return nil, fmt.Errorf(`"headers" %d: "key" is not present`, i)
}
header.Key = strings.ToLower(header.Key)
if unsupportedHeader(header.Key) {
return nil, fmt.Errorf(`"headers" %d: unsupported "key" %s`, i, header.Key)
}
if len(header.Values) == 0 {
return nil, fmt.Errorf(`"headers" %d: "values" is not present`, i)
}
values := parseHeaderValues(header.Key, header.Values)
hs = append(hs, permissionOr(values))
}
return hs, nil
}
func parseRequest(request request) (*v3rbacpb.Permission, error) {
var and []*v3rbacpb.Permission
if len(request.Paths) > 0 {
and = append(and, permissionOr(parsePaths(request.Paths)))
}
if len(request.Headers) > 0 {
headers, err := parseHeaders(request.Headers)
if err != nil {
return nil, err
}
and = append(and, permissionAnd(headers))
}
if len(and) > 0 {
return permissionAnd(and), nil
}
return &v3rbacpb.Permission{
Rule: &v3rbacpb.Permission_Any{
Any: true,
},
}, nil
}
func parseRules(rules []rule, prefixName string) (map[string]*v3rbacpb.Policy, error) {
policies := make(map[string]*v3rbacpb.Policy)
for i, rule := range rules {
if rule.Name == "" {
return policies, fmt.Errorf(`%d: "name" is not present`, i)
}
permission, err := parseRequest(rule.Request)
if err != nil {
return nil, fmt.Errorf("%d: %v", i, err)
}
policyName := prefixName + "_" + rule.Name
policies[policyName] = &v3rbacpb.Policy{
Principals: []*v3rbacpb.Principal{parsePeer(rule.Source)},
Permissions: []*v3rbacpb.Permission{permission},
}
}
return policies, nil
}
// Parse auditLoggingOptions to the associated RBAC protos. The single
// auditLoggingOptions results in two different parsed protos, one for the allow
// policy and one for the deny policy
func (options *auditLoggingOptions) toProtos() (allow *v3rbacpb.RBAC_AuditLoggingOptions, deny *v3rbacpb.RBAC_AuditLoggingOptions, err error) {
allow = &v3rbacpb.RBAC_AuditLoggingOptions{}
deny = &v3rbacpb.RBAC_AuditLoggingOptions{}
if options.AuditCondition != "" {
rbacCondition, ok := v3rbacpb.RBAC_AuditLoggingOptions_AuditCondition_value[options.AuditCondition]
if !ok {
return nil, nil, fmt.Errorf("failed to parse AuditCondition %v. Allowed values {NONE, ON_DENY, ON_ALLOW, ON_DENY_AND_ALLOW}", options.AuditCondition)
}
allow.AuditCondition = v3rbacpb.RBAC_AuditLoggingOptions_AuditCondition(rbacCondition)
deny.AuditCondition = toDenyCondition(v3rbacpb.RBAC_AuditLoggingOptions_AuditCondition(rbacCondition))
}
for i, config := range options.AuditLoggers {
if config.Name == "" {
return nil, nil, fmt.Errorf("missing required field: name in audit_logging_options.audit_loggers[%v]", i)
}
if config.Config == nil {
config.Config = &structpb.Struct{}
}
typedStruct := &v1xdsudpatypepb.TypedStruct{
TypeUrl: typeURLPrefix + config.Name,
Value: config.Config,
}
customConfig, err := anypb.New(typedStruct)
if err != nil {
return nil, nil, fmt.Errorf("error parsing custom audit logger config: %v", err)
}
logger := &v3corepb.TypedExtensionConfig{Name: config.Name, TypedConfig: customConfig}
rbacConfig := v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
IsOptional: config.IsOptional,
AuditLogger: logger,
}
allow.LoggerConfigs = append(allow.LoggerConfigs, &rbacConfig)
deny.LoggerConfigs = append(deny.LoggerConfigs, &rbacConfig)
}
return allow, deny, nil
}
// Maps the AuditCondition coming from AuditLoggingOptions to the proper
// condition for the deny policy RBAC proto
func toDenyCondition(condition v3rbacpb.RBAC_AuditLoggingOptions_AuditCondition) v3rbacpb.RBAC_AuditLoggingOptions_AuditCondition {
// Mapping the overall policy AuditCondition to what it must be for the Deny and Allow RBAC
// See gRPC A59 for details - https://github.com/grpc/proposal/pull/346/files
// |Authorization Policy |DENY RBAC |ALLOW RBAC |
// |----------------------|-------------------|---------------------|
// |NONE |NONE |NONE |
// |ON_DENY |ON_DENY |ON_DENY |
// |ON_ALLOW |NONE |ON_ALLOW |
// |ON_DENY_AND_ALLOW |ON_DENY |ON_DENY_AND_ALLOW |
switch condition {
case v3rbacpb.RBAC_AuditLoggingOptions_NONE:
return v3rbacpb.RBAC_AuditLoggingOptions_NONE
case v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY:
return v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY
case v3rbacpb.RBAC_AuditLoggingOptions_ON_ALLOW:
return v3rbacpb.RBAC_AuditLoggingOptions_NONE
case v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY_AND_ALLOW:
return v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY
default:
return v3rbacpb.RBAC_AuditLoggingOptions_NONE
}
}
// translatePolicy translates SDK authorization policy in JSON format to two
// Envoy RBAC polices (deny followed by allow policy) or only one Envoy RBAC
// allow policy. Also returns the overall policy name. If the input policy
// cannot be parsed or is invalid, an error will be returned.
func translatePolicy(policyStr string) ([]*v3rbacpb.RBAC, string, error) {
policy := &authorizationPolicy{}
d := json.NewDecoder(bytes.NewReader([]byte(policyStr)))
d.DisallowUnknownFields()
if err := d.Decode(policy); err != nil {
return nil, "", fmt.Errorf("failed to unmarshal policy: %v", err)
}
if policy.Name == "" {
return nil, "", fmt.Errorf(`"name" is not present`)
}
if len(policy.AllowRules) == 0 {
return nil, "", fmt.Errorf(`"allow_rules" is not present`)
}
allowLogger, denyLogger, err := policy.AuditLoggingOptions.toProtos()
if err != nil {
return nil, "", err
}
rbacs := make([]*v3rbacpb.RBAC, 0, 2)
if len(policy.DenyRules) > 0 {
denyPolicies, err := parseRules(policy.DenyRules, policy.Name)
if err != nil {
return nil, "", fmt.Errorf(`"deny_rules" %v`, err)
}
denyRBAC := &v3rbacpb.RBAC{
Action: v3rbacpb.RBAC_DENY,
Policies: denyPolicies,
AuditLoggingOptions: denyLogger,
}
rbacs = append(rbacs, denyRBAC)
}
allowPolicies, err := parseRules(policy.AllowRules, policy.Name)
if err != nil {
return nil, "", fmt.Errorf(`"allow_rules" %v`, err)
}
allowRBAC := &v3rbacpb.RBAC{Action: v3rbacpb.RBAC_ALLOW, Policies: allowPolicies, AuditLoggingOptions: allowLogger}
return append(rbacs, allowRBAC), policy.Name, nil
}

File diff suppressed because it is too large Load Diff

61
backoff.go Normal file
View File

@@ -0,0 +1,61 @@
/*
*
* Copyright 2017 gRPC authors.
*
* 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.
*
*/
// See internal/backoff package for the backoff implementation. This file is
// kept for the exported types and API backward compatibility.
package grpc
import (
"time"
"google.golang.org/grpc/backoff"
)
// DefaultBackoffConfig uses values specified for backoff in
// https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
//
// Deprecated: use ConnectParams instead. Will be supported throughout 1.x.
var DefaultBackoffConfig = BackoffConfig{
MaxDelay: 120 * time.Second,
}
// BackoffConfig defines the parameters for the default gRPC backoff strategy.
//
// Deprecated: use ConnectParams instead. Will be supported throughout 1.x.
type BackoffConfig struct {
// MaxDelay is the upper bound of backoff delay.
MaxDelay time.Duration
}
// ConnectParams defines the parameters for connecting and retrying. Users are
// encouraged to use this instead of the BackoffConfig type defined above. See
// here for more details:
// https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
//
// # Experimental
//
// Notice: This type is EXPERIMENTAL and may be changed or removed in a
// later release.
type ConnectParams struct {
// Backoff specifies the configuration options for connection backoff.
Backoff backoff.Config
// MinConnectTimeout is the minimum amount of time we are willing to give a
// connection to complete.
MinConnectTimeout time.Duration
}

52
backoff/backoff.go Normal file
View File

@@ -0,0 +1,52 @@
/*
*
* Copyright 2019 gRPC authors.
*
* 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.
*
*/
// Package backoff provides configuration options for backoff.
//
// More details can be found at:
// https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
//
// All APIs in this package are experimental.
package backoff
import "time"
// Config defines the configuration options for backoff.
type Config struct {
// BaseDelay is the amount of time to backoff after the first failure.
BaseDelay time.Duration
// Multiplier is the factor with which to multiply backoffs after a
// failed retry. Should ideally be greater than 1.
Multiplier float64
// Jitter is the factor with which backoffs are randomized.
Jitter float64
// MaxDelay is the upper bound of backoff delay.
MaxDelay time.Duration
}
// DefaultConfig is a backoff configuration with the default values specified
// at https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
//
// This should be useful for callers who want to configure backoff with
// non-default values only for a subset of the options.
var DefaultConfig = Config{
BaseDelay: 1.0 * time.Second,
Multiplier: 1.6,
Jitter: 0.2,
MaxDelay: 120 * time.Second,
}

View File

@@ -0,0 +1,112 @@
/*
*
* Copyright 2026 gRPC authors.
*
* 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.
*
*/
// Package autosharding implements the autosharding load balancing policy.
package autosharding
import (
"encoding/json"
"fmt"
"time"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/endpointsharding"
iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/serviceconfig"
)
// Name is the name of the autosharding balancer.
const Name = "autosharding_experimental"
func init() {
balancer.Register(bb{})
}
// lbConfig is the balancer config for the autosharding balancer.
type lbConfig struct {
serviceconfig.LoadBalancingConfig `json:"-"`
ChannelFactoryKey string `json:"channelFactoryKey,omitempty"`
AutoShardingTarget string `json:"autoshardingTarget,omitempty"`
KeyHeaderName string `json:"keyHeaderName,omitempty"`
EnableFallback bool `json:"enableFallback,omitempty"`
InitialAssignmentTimeout iserviceconfig.Duration `json:"initialAssignmentTimeout,omitempty"`
}
type bb struct{}
func (bb) Name() string {
return Name
}
func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
lbConfig := &lbConfig{InitialAssignmentTimeout: iserviceconfig.Duration(60 * time.Second)}
if err := json.Unmarshal(s, lbConfig); err != nil {
return nil, fmt.Errorf("autosharding: unable to unmarshal LBConfig: %v", err)
}
if lbConfig.ChannelFactoryKey == "" {
return nil, fmt.Errorf("autosharding: channelFactoryKey field is required")
}
if lbConfig.AutoShardingTarget == "" {
return nil, fmt.Errorf("autosharding: autoshardingTarget field is required")
}
if lbConfig.KeyHeaderName == "" {
return nil, fmt.Errorf("autosharding: keyHeaderName field is required")
}
return lbConfig, nil
}
func (bb) Build(balancer.ClientConn, balancer.BuildOptions) balancer.Balancer {
return &autoshardingBalancer{}
}
// slice represents a key range and its assigned endpoints.
//
//lint:ignore U1000 Struct fields planned for future implementation
type slice struct {
startKey []byte // Inclusive start key of the key-range
endKey []byte // Exclusive, nil for sentinel/infinity
endpoints []int // Indices into assignment.endpointNames
}
// assignment represents a complete snapshot of sharding assignments, and is
// expected to cover the entire key range.
type assignment struct {
slices []slice // Sorted by startKey
endpointNames []string // Complete list of endpoint names
generation int64
}
// endpointState represents the state associated with an endpoint in the LB policy.
//
//lint:ignore U1000 Struct fields planned for future implementation
type endpointState struct {
index int // Index of the endpoint within the NR update
endpoint resolver.Endpoint // The actual endpoint returned by the NR
childState endpointsharding.ChildState // State as reported by the child policy
}
// endpointMap maps from endpoint hostname to endpoint state.
type endpointMap struct {
m map[string]*endpointState
}
type autoshardingBalancer struct {
balancer.Balancer
}

View File

@@ -0,0 +1,144 @@
/*
*
* Copyright 2026 gRPC authors.
*
* 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.
*
*/
package autosharding
import (
"encoding/json"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc/internal/grpctest"
iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
"google.golang.org/grpc/serviceconfig"
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
func (s) TestParseConfig_Success(t *testing.T) {
parser := bb{}
tests := []struct {
name string
input string
wantCfg serviceconfig.LoadBalancingConfig
}{
{
name: "all-fields",
input: `{
"channelFactoryKey": "factory-key",
"autoshardingTarget": "target",
"keyHeaderName": "header",
"enableFallback": true,
"initialAssignmentTimeout": "30s"
}`,
wantCfg: &lbConfig{
ChannelFactoryKey: "factory-key",
AutoShardingTarget: "target",
KeyHeaderName: "header",
EnableFallback: true,
InitialAssignmentTimeout: iserviceconfig.Duration(30 * time.Second),
},
},
{
name: "default-timeout",
input: `{
"channelFactoryKey": "factory-key",
"autoshardingTarget": "target",
"keyHeaderName": "header"
}`,
wantCfg: &lbConfig{
ChannelFactoryKey: "factory-key",
AutoShardingTarget: "target",
KeyHeaderName: "header",
InitialAssignmentTimeout: iserviceconfig.Duration(60 * time.Second),
},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
gotCfg, err := parser.ParseConfig(json.RawMessage(test.input))
if err != nil {
t.Fatalf("ParseConfig() error = %v, want nil", err)
}
if diff := cmp.Diff(test.wantCfg, gotCfg); diff != "" {
t.Errorf("ParseConfig() config diff (-want +got):\n%s", diff)
}
})
}
}
func (s) TestParseConfig_Failure(t *testing.T) {
parser := bb{}
tests := []struct {
name string
input string
}{
{
name: "invalid-json",
input: "{{invalidjson{{",
},
{
name: "invalid-duration",
input: `{
"channelFactoryKey": "factory-key",
"autoshardingTarget": "target",
"keyHeaderName": "header",
"initialAssignmentTimeout": "invalid"
}`,
},
{
name: "missing-channel-factory-key",
input: `{
"autoshardingTarget": "target",
"keyHeaderName": "header"
}`,
},
{
name: "missing-autosharding-target",
input: `{
"channelFactoryKey": "factory-key",
"keyHeaderName": "header"
}`,
},
{
name: "missing-key-header-name",
input: `{
"channelFactoryKey": "factory-key",
"autoshardingTarget": "target"
}`,
},
{
name: "empty-config",
input: `{}`,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if _, err := parser.ParseConfig(json.RawMessage(test.input)); err == nil {
t.Fatalf("ParseConfig() succeeded, want error")
}
})
}
}

View File

@@ -0,0 +1,120 @@
/*
*
* Copyright 2026 gRPC authors.
*
* 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.
*
*/
package autosharding
import (
"bytes"
"slices"
)
// sliceMapEntry represents an entry for a key-range in the sliceMap.
type sliceMapEntry struct {
startKey []byte // Inclusive start key of the key-range
endpoints []int // Indices into list[PickerEndpoint] in the Picker
}
// sliceMap is an immutable data structure used by the Picker to map request
// keys to their assigned key-range. Because it is immutable, it allows the
// Picker to access it without any explicit synchronization with the LB policy.
//
// The sliceMap is meant to be used by the Picker in conjunction with a list of
// pickerEndpoints such that the list can be swapped out, as long as the number
// and order of endpoints don't change.
type sliceMap struct {
slices []sliceMapEntry // Sorted by startKey
fallbackPool []int // Indices into list[PickerEndpoint] for all endpoints
generation int64 // Snapshot generation number
}
// lookup returns the index of the slice covering the given key.
//
// Because assignments are pre-validated to have no gaps and cover the full key
// range, and since slices is sorted by startKey, lookup boils down to a binary
// search, looking for the insertion point.
//
// If the key matches the startKey of a slice, that slice index is returned.
// Otherwise, the index of the slice immediately preceding the insertion point
// is returned (i.e., the slice covering the range [startKey, nextStartKey)).
//
// Returns -1 when the sliceMap is empty, which is the case before the first
// assignment is received.
func (sm *sliceMap) lookup(key []byte) int {
if len(sm.slices) == 0 {
return -1
}
idx, found := slices.BinarySearchFunc(sm.slices, key, func(e sliceMapEntry, k []byte) int {
return bytes.Compare(e.startKey, k)
})
if found {
return idx
}
// Key falls in range [slices[idx - 1].startKey, slices[idx].startKey).
return idx - 1
}
// buildSliceMap is used to generate a new sliceMap from the EndpointMap and
// Assignment when either of them change.
func buildSliceMap(endpointMap *endpointMap, assignment *assignment) *sliceMap {
sm := &sliceMap{}
// Populate fallbackPool with values [0, 1, 2, ... N-1] where N is the
// number of endpoints returned by the name resolver. The only reason to
// have this slice is to allow the Picker to share code that picks a random
// endpoint from a slice of indices, either from the fallbackPool or from a
// sliceMapEntry.
sm.fallbackPool = make([]int, len(endpointMap.m))
for i := range sm.fallbackPool {
sm.fallbackPool[i] = i
}
// If no assignment has been received yet (startup case), return early with
// empty slices.
if assignment == nil {
return sm
}
sm.generation = assignment.generation
// Build sliceMapEntry for each Slice in the assignment.
sm.slices = make([]sliceMapEntry, 0, len(assignment.slices))
for _, s := range assignment.slices {
entry := sliceMapEntry{
startKey: s.startKey,
endpoints: make([]int, 0, len(s.endpoints)),
}
// Populate the endpoints for this slice by looking up the endpoint
// names in the EndpointMap. If an endpoint name is not found in the
// EndpointMap, it is skipped. This ensures that the sliceMap only
// contains valid endpoints that are currently available.
for _, idx := range s.endpoints {
hostname := assignment.endpointNames[idx]
if es, ok := endpointMap.m[hostname]; ok {
entry.endpoints = append(entry.endpoints, es.index)
}
}
sm.slices = append(sm.slices, entry)
}
return sm
}

View File

@@ -0,0 +1,216 @@
/*
*
* Copyright 2026 gRPC authors.
*
* 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.
*
*/
package autosharding
import (
"bytes"
"fmt"
"slices"
"testing"
"github.com/google/go-cmp/cmp"
)
func (s) TestSliceMap_Lookup(t *testing.T) {
sm := &sliceMap{
slices: []sliceMapEntry{
{startKey: []byte("a"), endpoints: []int{1}},
{startKey: []byte("d"), endpoints: []int{2}},
{startKey: []byte("f"), endpoints: []int{3}},
},
}
tests := []struct {
name string
key []byte
want int
}{
{
name: "exact-match-first",
key: []byte("a"),
want: 0,
},
{
name: "exact-match-middle",
key: []byte("d"),
want: 1,
},
{
name: "exact-match-last",
key: []byte("f"),
want: 2,
},
{
name: "between-first-and-middle",
key: []byte("c"),
want: 0,
},
{
name: "between-middle-and-last",
key: []byte("e"),
want: 1,
},
{
name: "after-last",
key: []byte("g"),
want: 2,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
if got := sm.lookup(tc.key); got != tc.want {
t.Errorf("lookup(%q) = %d, want %d", tc.key, got, tc.want)
}
})
}
}
func (s) TestSliceMap_Lookup_Empty(t *testing.T) {
sm := &sliceMap{}
if got := sm.lookup([]byte("a")); got != -1 {
t.Errorf("lookup() on empty map = %d, want -1", got)
}
}
func (s) TestBuildSliceMap(t *testing.T) {
// Setup endpointMap with unsorted order in map to verify deterministic
// sorting of fallbackPool.
epMap := &endpointMap{
m: map[string]*endpointState{
"hostC": {index: 2},
"hostA": {index: 0},
"hostB": {index: 1},
},
}
tests := []struct {
name string
assignment *assignment
want *sliceMap
}{
{
name: "nil-assignment-startup",
assignment: nil,
want: &sliceMap{fallbackPool: []int{0, 1, 2}},
},
{
name: "valid-assignment",
assignment: &assignment{
endpointNames: []string{"hostA", "hostB", "hostC"},
slices: []slice{
{startKey: []byte("a"), endpoints: []int{0, 1}}, // hostA, hostB -> indices 0, 1
{startKey: []byte("m"), endpoints: []int{1, 2}}, // hostB, hostC -> indices 1, 2
},
generation: 42,
},
want: &sliceMap{
slices: []sliceMapEntry{
{startKey: []byte("a"), endpoints: []int{0, 1}},
{startKey: []byte("m"), endpoints: []int{1, 2}},
},
fallbackPool: []int{0, 1, 2},
generation: 42,
},
},
{
name: "assignment-with-unknown-host",
assignment: &assignment{
endpointNames: []string{"hostA", "hostUnknown", "hostC"},
slices: []slice{
{startKey: []byte("a"), endpoints: []int{0, 1}}, // hostUnknown is skipped
},
generation: 43,
},
want: &sliceMap{
slices: []sliceMapEntry{
{startKey: []byte("a"), endpoints: []int{0}}, // Only hostA (index 0)
},
fallbackPool: []int{0, 1, 2},
generation: 43,
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
got := buildSliceMap(epMap, tc.assignment)
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Errorf("buildSliceMap() diff (-want +got):\n%s", diff)
}
})
}
}
func (e sliceMapEntry) Equal(b sliceMapEntry) bool {
return bytes.Equal(e.startKey, b.startKey) && slices.Equal(e.endpoints, b.endpoints)
}
func (sm *sliceMap) Equal(b *sliceMap) bool {
if sm == nil || b == nil {
return sm == b
}
fallbackPoolEqual := slices.Equal(sm.fallbackPool, b.fallbackPool)
slicesEqual := slices.EqualFunc(sm.slices, b.slices, func(x, y sliceMapEntry) bool { return x.Equal(y) })
return sm.generation == b.generation && fallbackPoolEqual && slicesEqual
}
func BenchmarkSliceMap_Lookup(b *testing.B) {
for _, numSlices := range []int{1, 10, 100, 1000, 10000} {
for _, keySize := range []int{16, 32, 64, 128, 256, 512} {
b.Run(fmt.Sprintf("slices_%d_keySize_%d", numSlices, keySize), func(b *testing.B) {
sm := &sliceMap{
slices: make([]sliceMapEntry, numSlices),
}
for i := 0; i < numSlices; i++ {
// Generate lexicographically sorted keys to serve as slice
// boundaries. We multiply by 1000 to create ranges (gaps) between
// successive slices (e.g., Slice 0 covers [0, 1000), Slice 1 covers
// [1000, 2000)). This allows us to test lookups that fall
// *mid-slice*, rather than just exact matches.
//
// We use zero-padding on the left via "%0*d" (where * is the keySize
// width). This fills the key to its full length (e.g., 512 bytes)
// with leading zeros. Scanning long identical prefixes forces
// bytes.Compare to traverse the full length, simulating a
// conservative, worst-case latency scenario.
key := fmt.Appendf(nil, "%0*d", keySize, i*1000)
sm.slices[i] = sliceMapEntry{startKey: key, endpoints: []int{i}}
}
// Pre-generate a pool of lookup keys. This helps avoid
// measuring string formatting overhead inside the timer loop.
lookupKeys := make([][]byte, 10000)
for i := 0; i < 10000; i++ {
// Distribute the 10,000 lookup keys proportionally across the entire
// synthetic keyspace [0, numSlices * 1000). This ensures a good mix
// of boundary hits and interior hits.
val := i * (numSlices * 1000) / 10000
lookupKeys[i] = fmt.Appendf(nil, "%0*d", keySize, val)
}
var i int
for b.Loop() {
sm.lookup(lookupKeys[i%10000])
i++
}
})
}
}
}

394
balancer/balancer.go Normal file
View File

@@ -0,0 +1,394 @@
/*
*
* Copyright 2017 gRPC authors.
*
* 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.
*
*/
// Package balancer defines APIs for load balancing in gRPC.
// All APIs in this package are experimental.
package balancer
import (
"context"
"encoding/json"
"errors"
"net"
"strings"
"google.golang.org/grpc/channelz"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/credentials"
estats "google.golang.org/grpc/experimental/stats"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/internal/envconfig"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/serviceconfig"
)
var (
// m is a map from name to balancer builder.
m = make(map[string]Builder)
logger = grpclog.Component("balancer")
)
// Register registers the balancer builder to the balancer map. b.Name
// will be used as the name registered with this builder. If the Builder
// implements ConfigParser, ParseConfig will be called when new service
// configs are received by the resolver, and the result will be provided to the
// Balancer in UpdateClientConnState.
//
// NOTE: this function must only be called during initialization time (i.e. in
// an init() function), and is not thread-safe. If multiple Balancers are
// registered with the same name, the one registered last will take effect.
func Register(b Builder) {
name := b.Name()
if !envconfig.CaseSensitiveBalancerRegistries {
name = strings.ToLower(name)
if name != b.Name() {
logger.Warningf("Balancer registered with name %q. grpc-go has switched to case sensitive balancer registries. GRPC_GO_EXPERIMENTAL_CASE_SENSITIVE_BALANCER_REGISTRIES env variable will be removed in release v1.82.0", b.Name())
}
}
m[name] = b
}
// unregisterForTesting deletes the balancer with the given name from the
// balancer map.
//
// This function is not thread-safe.
func unregisterForTesting(name string) {
delete(m, name)
}
func init() {
internal.BalancerUnregister = unregisterForTesting
}
// Get returns the resolver builder registered with the given name.
// Note that the compare is done in a case-sensitive fashion.
// If no builder is register with the name, nil will be returned.
func Get(name string) Builder {
if !envconfig.CaseSensitiveBalancerRegistries {
lowerName := strings.ToLower(name)
if lowerName != name {
logger.Warningf("Balancer retrieved for name %q. grpc-go has switched to case sensitive balancer registries. GRPC_GO_EXPERIMENTAL_CASE_SENSITIVE_BALANCER_REGISTRIES env variable will be removed in release v1.82.0", name)
}
name = lowerName
}
if b, ok := m[name]; ok {
return b
}
return nil
}
// NewSubConnOptions contains options to create new SubConn.
type NewSubConnOptions struct {
// CredsBundle is the credentials bundle that will be used in the created
// SubConn. If it's nil, the original creds from grpc DialOptions will be
// used.
//
// Deprecated: Use the Attributes field in resolver.Address to pass
// arbitrary data to the credential handshaker.
CredsBundle credentials.Bundle
// HealthCheckEnabled indicates whether health check service should be
// enabled on this SubConn
HealthCheckEnabled bool
// StateListener is called when the state of the subconn changes. If nil,
// Balancer.UpdateSubConnState will be called instead. Will never be
// invoked until after Connect() is called on the SubConn created with
// these options.
StateListener func(SubConnState)
}
// State contains the balancer's state relevant to the gRPC ClientConn.
type State struct {
// State contains the connectivity state of the balancer, which is used to
// determine the state of the ClientConn.
ConnectivityState connectivity.State
// Picker is used to choose connections (SubConns) for RPCs.
Picker Picker
}
// ClientConn represents a gRPC ClientConn.
//
// This interface is to be implemented by gRPC. Users should not need a
// brand new implementation of this interface. For the situations like
// testing, the new implementation should embed this interface. This allows
// gRPC to add new methods to this interface.
//
// NOTICE: This interface is intended to be implemented by gRPC, or intercepted
// by custom load balancing polices. Users should not need their own complete
// implementation of this interface -- they should always delegate to a
// ClientConn passed to Builder.Build() by embedding it in their
// implementations. An embedded ClientConn must never be nil, or runtime panics
// will occur.
type ClientConn interface {
// NewSubConn is called by balancer to create a new SubConn.
// It doesn't block and wait for the connections to be established.
// Behaviors of the SubConn can be controlled by options.
//
// Deprecated: please be aware that in a future version, SubConns will only
// support one address per SubConn.
NewSubConn([]resolver.Address, NewSubConnOptions) (SubConn, error)
// RemoveSubConn removes the SubConn from ClientConn.
// The SubConn will be shutdown.
//
// Deprecated: use SubConn.Shutdown instead.
RemoveSubConn(SubConn)
// UpdateAddresses updates the addresses used in the passed in SubConn.
// gRPC checks if the currently connected address is still in the new list.
// If so, the connection will be kept. Else, the connection will be
// gracefully closed, and a new connection will be created.
//
// This may trigger a state transition for the SubConn.
//
// Deprecated: this method will be removed. Create new SubConns for new
// addresses instead.
UpdateAddresses(SubConn, []resolver.Address)
// UpdateState notifies gRPC that the balancer's internal state has
// changed.
//
// gRPC will update the connectivity state of the ClientConn, and will call
// Pick on the new Picker to pick new SubConns.
UpdateState(State)
// ResolveNow is called by balancer to notify gRPC to do a name resolving.
ResolveNow(resolver.ResolveNowOptions)
// Target returns the dial target for this ClientConn.
//
// Deprecated: Use the Target field in the BuildOptions instead.
Target() string
// MetricsRecorder provides the metrics recorder that balancers can use to
// record metrics. Balancer implementations which do not register metrics on
// metrics registry and record on them can ignore this method. The returned
// MetricsRecorder is guaranteed to never be nil.
MetricsRecorder() estats.MetricsRecorder
// EnforceClientConnEmbedding is included to force implementers to embed
// another implementation of this interface, allowing gRPC to add methods
// without breaking users.
internal.EnforceClientConnEmbedding
}
// BuildOptions contains additional information for Build.
type BuildOptions struct {
// DialCreds is the transport credentials to use when communicating with a
// remote load balancer server. Balancer implementations which do not
// communicate with a remote load balancer server can ignore this field.
DialCreds credentials.TransportCredentials
// CredsBundle is the credentials bundle to use when communicating with a
// remote load balancer server. Balancer implementations which do not
// communicate with a remote load balancer server can ignore this field.
CredsBundle credentials.Bundle
// Dialer is the custom dialer to use when communicating with a remote load
// balancer server. Balancer implementations which do not communicate with a
// remote load balancer server can ignore this field.
Dialer func(context.Context, string) (net.Conn, error)
// Authority is the server name to use as part of the authentication
// handshake when communicating with a remote load balancer server. Balancer
// implementations which do not communicate with a remote load balancer
// server can ignore this field.
Authority string
// ChannelzParent is the parent ClientConn's channelz channel.
ChannelzParent channelz.Identifier
// CustomUserAgent is the custom user agent set on the parent ClientConn.
// The balancer should set the same custom user agent if it creates a
// ClientConn.
CustomUserAgent string
// Target contains the parsed address info of the dial target. It is the
// same resolver.Target as passed to the resolver. See the documentation for
// the resolver.Target type for details about what it contains.
Target resolver.Target
}
// Builder creates a balancer.
type Builder interface {
// Build creates a new balancer with the ClientConn.
Build(cc ClientConn, opts BuildOptions) Balancer
// Name returns the name of balancers built by this builder.
// It will be used to pick balancers (for example in service config).
Name() string
}
// ConfigParser parses load balancer configs.
type ConfigParser interface {
// ParseConfig parses the JSON load balancer config provided into an
// internal form or returns an error if the config is invalid. For future
// compatibility reasons, unknown fields in the config should be ignored.
ParseConfig(LoadBalancingConfigJSON json.RawMessage) (serviceconfig.LoadBalancingConfig, error)
}
// PickInfo contains additional information for the Pick operation.
type PickInfo struct {
// FullMethodName is the method name that NewClientStream() is called
// with. The canonical format is /service/Method.
FullMethodName string
// Ctx is the RPC's context, and may contain relevant RPC-level information
// like the outgoing header metadata.
Ctx context.Context
}
// DoneInfo contains additional information for done.
type DoneInfo struct {
// Err is the rpc error the RPC finished with. It could be nil.
Err error
// Trailer contains the metadata from the RPC's trailer, if present.
Trailer metadata.MD
// BytesSent indicates if any bytes have been sent to the server.
BytesSent bool
// BytesReceived indicates if any byte has been received from the server.
BytesReceived bool
// ServerLoad is the load received from server. It's usually sent as part of
// trailing metadata.
//
// The only supported type now is *orca_v3.LoadReport.
ServerLoad any
}
var (
// ErrNoSubConnAvailable indicates no SubConn is available for pick().
// gRPC will block the RPC until a new picker is available via UpdateState().
ErrNoSubConnAvailable = errors.New("no SubConn is available")
// ErrTransientFailure indicates all SubConns are in TransientFailure.
// WaitForReady RPCs will block, non-WaitForReady RPCs will fail.
//
// Deprecated: return an appropriate error based on the last resolution or
// connection attempt instead. The behavior is the same for any non-gRPC
// status error.
ErrTransientFailure = errors.New("all SubConns are in TransientFailure")
)
// PickResult contains information related to a connection chosen for an RPC.
type PickResult struct {
// SubConn is the connection to use for this pick, if its state is Ready.
// If the state is not Ready, gRPC will block the RPC until a new Picker is
// provided by the balancer (using ClientConn.UpdateState). The SubConn
// must be one returned by ClientConn.NewSubConn.
SubConn SubConn
// Done is called when the RPC is completed. If the SubConn is not ready,
// this will be called with a nil parameter. If the SubConn is not a valid
// type, Done may not be called. May be nil if the balancer does not wish
// to be notified when the RPC completes.
Done func(DoneInfo)
// Metadata provides a way for LB policies to inject arbitrary per-call
// metadata. Any metadata returned here will be merged with existing
// metadata added by the client application.
//
// LB policies with child policies are responsible for propagating metadata
// injected by their children to the ClientConn, as part of Pick().
Metadata metadata.MD
}
// TransientFailureError returns e. It exists for backward compatibility and
// will be deleted soon.
//
// Deprecated: no longer necessary, picker errors are treated this way by
// default.
func TransientFailureError(e error) error { return e }
// Picker is used by gRPC to pick a SubConn to send an RPC.
// Balancer is expected to generate a new picker from its snapshot every time its
// internal state has changed.
//
// The pickers used by gRPC can be updated by ClientConn.UpdateState().
type Picker interface {
// Pick returns the connection to use for this RPC and related information.
//
// Pick should not block. If the balancer needs to do I/O or any blocking
// or time-consuming work to service this call, it should return
// ErrNoSubConnAvailable, and the Pick call will be repeated by gRPC when
// the Picker is updated (using ClientConn.UpdateState).
//
// If an error is returned:
//
// - If the error is ErrNoSubConnAvailable, gRPC will block until a new
// Picker is provided by the balancer (using ClientConn.UpdateState).
//
// - If the error is a status error (implemented by the grpc/status
// package), gRPC will terminate the RPC with the code and message
// provided.
//
// - For all other errors, wait for ready RPCs will wait, but non-wait for
// ready RPCs will be terminated with this error's Error() string and
// status code Unavailable.
Pick(info PickInfo) (PickResult, error)
}
// Balancer takes input from gRPC, manages SubConns, and collects and aggregates
// the connectivity states.
//
// It also generates and updates the Picker used by gRPC to pick SubConns for RPCs.
//
// UpdateClientConnState, ResolverError, UpdateSubConnState, and Close are
// guaranteed to be called synchronously from the same goroutine. There's no
// guarantee on picker.Pick, it may be called anytime.
type Balancer interface {
// UpdateClientConnState is called by gRPC when the state of the ClientConn
// changes. If the error returned is ErrBadResolverState, the ClientConn
// will begin calling ResolveNow on the active name resolver with
// exponential backoff until a subsequent call to UpdateClientConnState
// returns a nil error. Any other errors are currently ignored.
UpdateClientConnState(ClientConnState) error
// ResolverError is called by gRPC when the name resolver reports an error.
ResolverError(error)
// UpdateSubConnState is called by gRPC when the state of a SubConn
// changes.
//
// Deprecated: Use NewSubConnOptions.StateListener when creating the
// SubConn instead.
UpdateSubConnState(SubConn, SubConnState)
// Close closes the balancer. The balancer is not currently required to
// call SubConn.Shutdown for its existing SubConns; however, this will be
// required in a future release, so it is recommended.
Close()
// ExitIdle instructs the LB policy to reconnect to backends / exit the
// IDLE state, if appropriate and possible. Note that SubConns that enter
// the IDLE state will not reconnect until SubConn.Connect is called.
ExitIdle()
}
// ExitIdler is an optional interface for balancers to implement. If
// implemented, ExitIdle will be called when ClientConn.Connect is called, if
// the ClientConn is idle. If unimplemented, ClientConn.Connect will cause
// all SubConns to connect.
//
// Deprecated: All balancers must implement this interface. This interface will
// be removed in a future release.
type ExitIdler interface {
// ExitIdle instructs the LB policy to reconnect to backends / exit the
// IDLE state, if appropriate and possible. Note that SubConns that enter
// the IDLE state will not reconnect until SubConn.Connect is called.
ExitIdle()
}
// ClientConnState describes the state of a ClientConn relevant to the
// balancer.
type ClientConnState struct {
ResolverState resolver.State
// The parsed load balancing configuration returned by the builder's
// ParseConfig method, if implemented.
BalancerConfig serviceconfig.LoadBalancingConfig
}
// ErrBadResolverState may be returned by UpdateClientConnState to indicate a
// problem with the provided name resolver data.
var ErrBadResolverState = errors.New("bad resolver state")

260
balancer/base/balancer.go Normal file
View File

@@ -0,0 +1,260 @@
/*
*
* Copyright 2017 gRPC authors.
*
* 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.
*
*/
package base
import (
"errors"
"fmt"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/resolver"
)
var logger = grpclog.Component("balancer")
type baseBuilder struct {
name string
pickerBuilder PickerBuilder
config Config
}
func (bb *baseBuilder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
bal := &baseBalancer{
cc: cc,
pickerBuilder: bb.pickerBuilder,
subConns: resolver.NewAddressMapV2[balancer.SubConn](),
scStates: make(map[balancer.SubConn]connectivity.State),
csEvltr: &balancer.ConnectivityStateEvaluator{},
config: bb.config,
state: connectivity.Connecting,
}
// Initialize picker to a picker that always returns
// ErrNoSubConnAvailable, because when state of a SubConn changes, we
// may call UpdateState with this picker.
bal.picker = NewErrPicker(balancer.ErrNoSubConnAvailable)
return bal
}
func (bb *baseBuilder) Name() string {
return bb.name
}
type baseBalancer struct {
cc balancer.ClientConn
pickerBuilder PickerBuilder
csEvltr *balancer.ConnectivityStateEvaluator
state connectivity.State
subConns *resolver.AddressMapV2[balancer.SubConn]
scStates map[balancer.SubConn]connectivity.State
picker balancer.Picker
config Config
resolverErr error // the last error reported by the resolver; cleared on successful resolution
connErr error // the last connection error; cleared upon leaving TransientFailure
}
func (b *baseBalancer) ResolverError(err error) {
b.resolverErr = err
if b.subConns.Len() == 0 {
b.state = connectivity.TransientFailure
}
if b.state != connectivity.TransientFailure {
// The picker will not change since the balancer does not currently
// report an error.
return
}
b.regeneratePicker()
b.cc.UpdateState(balancer.State{
ConnectivityState: b.state,
Picker: b.picker,
})
}
func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
// TODO: handle s.ResolverState.ServiceConfig?
if logger.V(2) {
logger.Info("base.baseBalancer: got new ClientConn state: ", s)
}
// Successful resolution; clear resolver error and ensure we return nil.
b.resolverErr = nil
// addrsSet is the set converted from addrs, it's used for quick lookup of an address.
addrsSet := resolver.NewAddressMapV2[any]()
for _, a := range s.ResolverState.Addresses {
addrsSet.Set(a, nil)
if _, ok := b.subConns.Get(a); !ok {
// a is a new address (not existing in b.subConns).
var sc balancer.SubConn
opts := balancer.NewSubConnOptions{
HealthCheckEnabled: b.config.HealthCheck,
StateListener: func(scs balancer.SubConnState) { b.updateSubConnState(sc, scs) },
}
sc, err := b.cc.NewSubConn([]resolver.Address{a}, opts)
if err != nil {
logger.Warningf("base.baseBalancer: failed to create new SubConn: %v", err)
continue
}
b.subConns.Set(a, sc)
b.scStates[sc] = connectivity.Idle
b.csEvltr.RecordTransition(connectivity.Shutdown, connectivity.Idle)
sc.Connect()
}
}
for a, sc := range b.subConns.All() {
// a was removed by resolver.
if _, ok := addrsSet.Get(a); !ok {
sc.Shutdown()
b.subConns.Delete(a)
// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
// The entry will be deleted in updateSubConnState.
}
}
// If resolver state contains no addresses, return an error so ClientConn
// will trigger re-resolve. Also records this as a resolver error, so when
// the overall state turns transient failure, the error message will have
// the zero address information.
if len(s.ResolverState.Addresses) == 0 {
b.ResolverError(errors.New("produced zero addresses"))
return balancer.ErrBadResolverState
}
b.regeneratePicker()
b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
return nil
}
// mergeErrors builds an error from the last connection error and the last
// resolver error. Must only be called if b.state is TransientFailure.
func (b *baseBalancer) mergeErrors() error {
// connErr must always be non-nil unless there are no SubConns, in which
// case resolverErr must be non-nil.
if b.connErr == nil {
return fmt.Errorf("last resolver error: %v", b.resolverErr)
}
if b.resolverErr == nil {
return fmt.Errorf("last connection error: %v", b.connErr)
}
return fmt.Errorf("last connection error: %v; last resolver error: %v", b.connErr, b.resolverErr)
}
// regeneratePicker takes a snapshot of the balancer, and generates a picker
// from it. The picker is
// - errPicker if the balancer is in TransientFailure,
// - built by the pickerBuilder with all READY SubConns otherwise.
func (b *baseBalancer) regeneratePicker() {
if b.state == connectivity.TransientFailure {
b.picker = NewErrPicker(b.mergeErrors())
return
}
readySCs := make(map[balancer.SubConn]SubConnInfo)
// Filter out all ready SCs from full subConn map.
for addr, sc := range b.subConns.All() {
if st, ok := b.scStates[sc]; ok && st == connectivity.Ready {
readySCs[sc] = SubConnInfo{Address: addr}
}
}
b.picker = b.pickerBuilder.Build(PickerBuildInfo{ReadySCs: readySCs})
}
// UpdateSubConnState is a nop because a StateListener is always set in NewSubConn.
func (b *baseBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
logger.Errorf("base.baseBalancer: UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
}
func (b *baseBalancer) updateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
s := state.ConnectivityState
if logger.V(2) {
logger.Infof("base.baseBalancer: handle SubConn state change: %p, %v", sc, s)
}
oldS, ok := b.scStates[sc]
if !ok {
if logger.V(2) {
logger.Infof("base.baseBalancer: got state changes for an unknown SubConn: %p, %v", sc, s)
}
return
}
if oldS == connectivity.TransientFailure &&
(s == connectivity.Connecting || s == connectivity.Idle) {
// Once a subconn enters TRANSIENT_FAILURE, ignore subsequent IDLE or
// CONNECTING transitions to prevent the aggregated state from being
// always CONNECTING when many backends exist but are all down.
if s == connectivity.Idle {
sc.Connect()
}
return
}
b.scStates[sc] = s
switch s {
case connectivity.Idle:
sc.Connect()
case connectivity.Shutdown:
// When an address was removed by resolver, b called Shutdown but kept
// the sc's state in scStates. Remove state for this sc here.
delete(b.scStates, sc)
case connectivity.TransientFailure:
// Save error to be reported via picker.
b.connErr = state.ConnectionError
}
b.state = b.csEvltr.RecordTransition(oldS, s)
// Regenerate picker when one of the following happens:
// - this sc entered or left ready
// - the aggregated state of balancer is TransientFailure
// (may need to update error message)
if (s == connectivity.Ready) != (oldS == connectivity.Ready) ||
b.state == connectivity.TransientFailure {
b.regeneratePicker()
}
b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
}
// Close is a nop because base balancer doesn't have internal state to clean up,
// and it doesn't need to call Shutdown for the SubConns.
func (b *baseBalancer) Close() {
}
// ExitIdle is a nop because the base balancer attempts to stay connected to
// all SubConns at all times.
func (b *baseBalancer) ExitIdle() {
}
// NewErrPicker returns a Picker that always returns err on Pick().
func NewErrPicker(err error) balancer.Picker {
return &errPicker{err: err}
}
// NewErrPickerV2 is temporarily defined for backward compatibility reasons.
//
// Deprecated: use NewErrPicker instead.
var NewErrPickerV2 = NewErrPicker
type errPicker struct {
err error // Pick() always returns this err.
}
func (p *errPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
return balancer.PickResult{}, p.err
}

View File

@@ -0,0 +1,122 @@
/*
*
* Copyright 2020 gRPC authors.
*
* 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.
*
*/
package base
import (
"context"
"testing"
"time"
"google.golang.org/grpc/attributes"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/resolver"
)
type testClientConn struct {
balancer.ClientConn
newSubConn func([]resolver.Address, balancer.NewSubConnOptions) (balancer.SubConn, error)
}
func (c *testClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
return c.newSubConn(addrs, opts)
}
func (c *testClientConn) UpdateState(balancer.State) {}
type testSubConn struct {
balancer.SubConn
updateState func(balancer.SubConnState)
}
func (sc *testSubConn) UpdateAddresses([]resolver.Address) {}
func (sc *testSubConn) Connect() {}
func (sc *testSubConn) Shutdown() {}
func (sc *testSubConn) GetOrBuildProducer(balancer.ProducerBuilder) (balancer.Producer, func()) {
return nil, nil
}
// RegisterHealthListener is a no-op.
func (*testSubConn) RegisterHealthListener(func(balancer.SubConnState)) {}
// testPickBuilder creates balancer.Picker for test.
type testPickBuilder struct {
validate func(info PickerBuildInfo)
}
func (p *testPickBuilder) Build(info PickerBuildInfo) balancer.Picker {
p.validate(info)
return nil
}
func TestBaseBalancerReserveAttributes(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
validated := make(chan struct{}, 1)
v := func(info PickerBuildInfo) {
defer func() { validated <- struct{}{} }()
for _, sc := range info.ReadySCs {
if sc.Address.Addr == "1.1.1.1" {
if sc.Address.Attributes == nil {
t.Errorf("in picker.validate, got address %+v with nil attributes, want not nil", sc.Address)
}
foo, ok := sc.Address.Attributes.Value("foo").(string)
if !ok || foo != "2233niang" {
t.Errorf("in picker.validate, got address[1.1.1.1] with invalid attributes value %v, want 2233niang", sc.Address.Attributes.Value("foo"))
}
} else if sc.Address.Addr == "2.2.2.2" {
if sc.Address.Attributes != nil {
t.Error("in b.subConns, got address[2.2.2.2] with not nil attributes, want nil")
}
}
}
}
pickBuilder := &testPickBuilder{validate: v}
b := (&baseBuilder{pickerBuilder: pickBuilder}).Build(&testClientConn{
newSubConn: func(_ []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
return &testSubConn{updateState: opts.StateListener}, nil
},
}, balancer.BuildOptions{}).(*baseBalancer)
b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{
Addresses: []resolver.Address{
{Addr: "1.1.1.1", Attributes: attributes.New("foo", "2233niang")},
{Addr: "2.2.2.2", Attributes: nil},
},
},
})
select {
case <-validated:
case <-ctx.Done():
t.Fatalf("timed out waiting for UpdateClientConnState to call picker.Build")
}
for sc := range b.scStates {
sc.(*testSubConn).updateState(balancer.SubConnState{ConnectivityState: connectivity.Ready, ConnectionError: nil})
select {
case <-validated:
case <-ctx.Done():
t.Fatalf("timed out waiting for UpdateClientConnState to call picker.Build")
}
}
}

71
balancer/base/base.go Normal file
View File

@@ -0,0 +1,71 @@
/*
*
* Copyright 2017 gRPC authors.
*
* 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.
*
*/
// Package base defines a balancer base that can be used to build balancers with
// different picking algorithms.
//
// The base balancer creates a new SubConn for each resolved address. The
// provided picker will only be notified about READY SubConns.
//
// This package is the base of round_robin balancer, its purpose is to be used
// to build round_robin like balancers with complex picking algorithms.
// Balancers with more complicated logic should try to implement a balancer
// builder from scratch.
//
// All APIs in this package are experimental.
package base
import (
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/resolver"
)
// PickerBuilder creates balancer.Picker.
type PickerBuilder interface {
// Build returns a picker that will be used by gRPC to pick a SubConn.
Build(info PickerBuildInfo) balancer.Picker
}
// PickerBuildInfo contains information needed by the picker builder to
// construct a picker.
type PickerBuildInfo struct {
// ReadySCs is a map from all ready SubConns to the Addresses used to
// create them.
ReadySCs map[balancer.SubConn]SubConnInfo
}
// SubConnInfo contains information about a SubConn created by the base
// balancer.
type SubConnInfo struct {
Address resolver.Address // the address used to create this SubConn
}
// Config contains the config info about the base balancer builder.
type Config struct {
// HealthCheck indicates whether health checking should be enabled for this specific balancer.
HealthCheck bool
}
// NewBalancerBuilder returns a base balancer builder configured by the provided config.
func NewBalancerBuilder(name string, pb PickerBuilder, config Config) balancer.Builder {
return &baseBuilder{
name: name,
pickerBuilder: pb,
config: config,
}
}

View File

@@ -0,0 +1,74 @@
/*
*
* Copyright 2022 gRPC authors.
*
* 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.
*
*/
package balancer
import "google.golang.org/grpc/connectivity"
// ConnectivityStateEvaluator takes the connectivity states of multiple SubConns
// and returns one aggregated connectivity state.
//
// It's not thread safe.
type ConnectivityStateEvaluator struct {
numReady uint64 // Number of addrConns in ready state.
numConnecting uint64 // Number of addrConns in connecting state.
numTransientFailure uint64 // Number of addrConns in transient failure state.
numIdle uint64 // Number of addrConns in idle state.
}
// RecordTransition records state change happening in subConn and based on that
// it evaluates what aggregated state should be.
//
// - If at least one SubConn in Ready, the aggregated state is Ready;
// - Else if at least one SubConn in Connecting, the aggregated state is Connecting;
// - Else if at least one SubConn is Idle, the aggregated state is Idle;
// - Else if at least one SubConn is TransientFailure (or there are no SubConns), the aggregated state is Transient Failure.
//
// Shutdown is not considered.
func (cse *ConnectivityStateEvaluator) RecordTransition(oldState, newState connectivity.State) connectivity.State {
// Update counters.
for idx, state := range []connectivity.State{oldState, newState} {
updateVal := 2*uint64(idx) - 1 // -1 for oldState and +1 for new.
switch state {
case connectivity.Ready:
cse.numReady += updateVal
case connectivity.Connecting:
cse.numConnecting += updateVal
case connectivity.TransientFailure:
cse.numTransientFailure += updateVal
case connectivity.Idle:
cse.numIdle += updateVal
}
}
return cse.CurrentState()
}
// CurrentState returns the current aggregate conn state by evaluating the counters
func (cse *ConnectivityStateEvaluator) CurrentState() connectivity.State {
// Evaluate.
if cse.numReady > 0 {
return connectivity.Ready
}
if cse.numConnecting > 0 {
return connectivity.Connecting
}
if cse.numIdle > 0 {
return connectivity.Idle
}
return connectivity.TransientFailure
}

View File

@@ -0,0 +1,245 @@
/*
*
* Copyright 2022 gRPC authors.
*
* 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.
*
*/
package balancer
import (
"testing"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/internal/grpctest"
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
// TestRecordTransition_FirstStateChange tests the first call to
// RecordTransition where the `oldState` is usually set to `Shutdown` (a state
// that the ConnectivityStateEvaluator is set to ignore).
func (s) TestRecordTransition_FirstStateChange(t *testing.T) {
tests := []struct {
newState connectivity.State
wantState connectivity.State
}{
{
newState: connectivity.Idle,
wantState: connectivity.Idle,
},
{
newState: connectivity.Connecting,
wantState: connectivity.Connecting,
},
{
newState: connectivity.Ready,
wantState: connectivity.Ready,
},
{
newState: connectivity.TransientFailure,
wantState: connectivity.TransientFailure,
},
{
newState: connectivity.Shutdown,
wantState: connectivity.TransientFailure,
},
}
for _, test := range tests {
cse := &ConnectivityStateEvaluator{}
if gotState := cse.RecordTransition(connectivity.Shutdown, test.newState); gotState != test.wantState {
t.Fatalf("RecordTransition(%v, %v) = %v, want %v", connectivity.Shutdown, test.newState, gotState, test.wantState)
}
}
}
// TestRecordTransition_SameState tests the scenario where state transitions to
// the same state are recorded multiple times.
func (s) TestRecordTransition_SameState(t *testing.T) {
tests := []struct {
newState connectivity.State
wantState connectivity.State
}{
{
newState: connectivity.Idle,
wantState: connectivity.Idle,
},
{
newState: connectivity.Connecting,
wantState: connectivity.Connecting,
},
{
newState: connectivity.Ready,
wantState: connectivity.Ready,
},
{
newState: connectivity.TransientFailure,
wantState: connectivity.TransientFailure,
},
{
newState: connectivity.Shutdown,
wantState: connectivity.TransientFailure,
},
}
const numStateChanges = 5
for _, test := range tests {
cse := &ConnectivityStateEvaluator{}
var prevState, gotState connectivity.State
prevState = connectivity.Shutdown
for i := 0; i < numStateChanges; i++ {
gotState = cse.RecordTransition(prevState, test.newState)
prevState = test.newState
}
if gotState != test.wantState {
t.Fatalf("RecordTransition() = %v, want %v", gotState, test.wantState)
}
}
}
// TestRecordTransition_SingleSubConn_DifferentStates tests some common
// connectivity state change scenarios, on a single subConn.
func (s) TestRecordTransition_SingleSubConn_DifferentStates(t *testing.T) {
tests := []struct {
name string
states []connectivity.State
wantState connectivity.State
}{
{
name: "regular transition to ready",
states: []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready},
wantState: connectivity.Ready,
},
{
name: "regular transition to transient failure",
states: []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
wantState: connectivity.TransientFailure,
},
{
name: "regular transition to ready",
states: []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.Idle},
wantState: connectivity.Idle,
},
{
name: "transition from ready to transient failure",
states: []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure},
wantState: connectivity.TransientFailure,
},
{
name: "transition from transient failure back to ready",
states: []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure, connectivity.Ready},
wantState: connectivity.Ready,
},
{
// This state transition is usually suppressed at the LB policy level, by
// not calling RecordTransition.
name: "transition from transient failure back to idle",
states: []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure, connectivity.Idle},
wantState: connectivity.Idle,
},
{
// This state transition is usually suppressed at the LB policy level, by
// not calling RecordTransition.
name: "transition from transient failure back to connecting",
states: []connectivity.State{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure, connectivity.Connecting},
wantState: connectivity.Connecting,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
cse := &ConnectivityStateEvaluator{}
var prevState, gotState connectivity.State
prevState = connectivity.Shutdown
for _, newState := range test.states {
gotState = cse.RecordTransition(prevState, newState)
prevState = newState
}
if gotState != test.wantState {
t.Fatalf("RecordTransition() = %v, want %v", gotState, test.wantState)
}
})
}
}
// TestRecordTransition_MultipleSubConns_DifferentStates tests state transitions
// among multiple subConns, and verifies that the connectivity state aggregation
// algorithm produces the expected aggregate connectivity state.
func (s) TestRecordTransition_MultipleSubConns_DifferentStates(t *testing.T) {
tests := []struct {
name string
// Each entry in this slice corresponds to the state changes happening on an
// individual subConn.
subConnStates [][]connectivity.State
wantState connectivity.State
}{
{
name: "atleast one ready",
subConnStates: [][]connectivity.State{
{connectivity.Idle, connectivity.Connecting, connectivity.Ready},
{connectivity.Idle},
{connectivity.Idle, connectivity.Connecting},
{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
},
wantState: connectivity.Ready,
},
{
name: "atleast one connecting",
subConnStates: [][]connectivity.State{
{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.Connecting},
{connectivity.Idle},
{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
},
wantState: connectivity.Connecting,
},
{
name: "atleast one idle",
subConnStates: [][]connectivity.State{
{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.Idle},
{connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure},
},
wantState: connectivity.Idle,
},
{
name: "atleast one transient failure",
subConnStates: [][]connectivity.State{
{connectivity.Idle, connectivity.Connecting, connectivity.Ready, connectivity.TransientFailure},
{connectivity.TransientFailure},
},
wantState: connectivity.TransientFailure,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
cse := &ConnectivityStateEvaluator{}
var prevState, gotState connectivity.State
for _, scStates := range test.subConnStates {
prevState = connectivity.Shutdown
for _, newState := range scStates {
gotState = cse.RecordTransition(prevState, newState)
prevState = newState
}
}
if gotState != test.wantState {
t.Fatalf("RecordTransition() = %v, want %v", gotState, test.wantState)
}
})
}
}

View File

@@ -0,0 +1,385 @@
/*
*
* Copyright 2024 gRPC authors.
*
* 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.
*
*/
// Package endpointsharding implements a load balancing policy that manages
// homogeneous child policies each owning a single endpoint.
//
// # Experimental
//
// Notice: This package is EXPERIMENTAL and may be changed or removed in a
// later release.
package endpointsharding
import (
"errors"
rand "math/rand/v2"
"sync"
"sync/atomic"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/base"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/resolver"
)
var randIntN = rand.IntN
// ChildState is the state of a child balancer.
type ChildState struct {
Endpoint resolver.Endpoint // Endpoint of the child balancer.
State balancer.State // State of the child balancer.
ExitIdle func() // Function to exit the child balancer from IDLE state.
}
// Options configure the behaviour of the endpointsharding balancer.
type Options struct {
// DisableAutoReconnect allows the balancer to keep child balancer in the
// IDLE state until they are explicitly triggered to exit using the
// ChildState obtained from the endpointsharding picker. When set to false,
// the endpointsharding balancer will automatically call ExitIdle on child
// connections that report IDLE.
DisableAutoReconnect bool
}
// ChildBuilderFunc creates a new balancer with the ClientConn. It has the same
// type as the balancer.Builder.Build method.
type ChildBuilderFunc func(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer
// NewBalancer returns a load balancing policy that manages homogeneous child
// policies each owning a single endpoint. The endpointsharding balancer
// forwards the LoadBalancingConfig in ClientConn state updates to its children.
func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions, childBuilder ChildBuilderFunc, esOpts Options) balancer.Balancer {
return &endpointSharding{
cc: cc,
bOpts: opts,
esOpts: esOpts,
childBuilder: childBuilder,
endpoints: resolver.NewEndpointMap[*endpointState](),
}
}
// endpointSharding is a balancer that wraps child balancers. It creates a child
// balancer with child config for every unique Endpoint received. It updates the
// child states on any update from parent or child.
type endpointSharding struct {
cc balancer.ClientConn
bOpts balancer.BuildOptions
esOpts Options
childBuilder ChildBuilderFunc
// mu guards access to the below fields and guarantees mutual exclusion
// between top-down methods (like UpdateClientConnState, ResolverError etc,
// which are already serialized) and bottom-up methods (like UpdateState)
// that can be called concurrently.
//
// Lock ordering & deadlock prevention:
// Child callbacks (like UpdateState) acquire this mutex to push state updates
// to the parent. This establishes a strict lock ordering:
// [endpointState.childMu] -> [endpointSharding.mu]
//
// To prevent deadlocks, we must never invert this order. Therefore, we must
// never call any methods on a child balancer while holding this mutex. If we
// did, and that child balancer invoked UpdateState synchronously, it would
// attempt to re-acquire this mutex, causing a deadlock.
//
// Concurrency:
// Top-down operations need to iterate over all children. To ensure a
// single, clean aggregated update at the end of such operations, we inhibit
// intermediate updates from children. We grab this mutex briefly to set
// `inhibitChildUpdates = true` and immediately release it. This allows us
// to perform the potentially slow, top-down child updates without holding
// any parent locks. Once finished, we grab the mutex again to unset the
// flag and push the final aggregated state. An update from a child during
// this time will *only* update the child state, and will not access the
// endpoints map or push an aggregated state to the parent.
mu sync.Mutex
endpoints *resolver.EndpointMap[*endpointState]
inhibitChildUpdates bool
}
// rotateEndpoints returns a slice of all the input endpoints rotated a random
// amount.
func rotateEndpoints(es []resolver.Endpoint) []resolver.Endpoint {
n := len(es)
if n == 0 {
return es
}
r := randIntN(n)
// Make a copy to avoid mutating data beyond the end of es.
ret := make([]resolver.Endpoint, n)
copy(ret, es[r:])
copy(ret[n-r:], es[:r])
return ret
}
// UpdateClientConnState creates a child for new endpoints and deletes children
// for endpoints that are no longer present. It also updates all the children,
// and sends a single synchronous update of the childrens' aggregated state at
// the end of the UpdateClientConnState operation.
//
// Returns the first error found from a child, but fully processes the update.
func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState) error {
es.inhibitUpdatesFromChildren()
// Update/create child balancers for each endpoint in the update. Note that we
// don't hold the mutex here, but this is fine because inhibitChildUpdates is
// true, and therefore UpdateState will not access es.endpoints.
var retErr error
newEndpoints := resolver.NewEndpointMap[*endpointState]()
for _, endpoint := range rotateEndpoints(state.ResolverState.Endpoints) {
if _, ok := newEndpoints.Get(endpoint); ok {
// Skip duplicate endpoints.
continue
}
epState, ok := es.endpoints.Get(endpoint)
if ok {
// Endpoint child already exists, update the stored endpoint.
epState.endpoint = endpoint
} else {
// Endpoint child does not exist, create a new one.
epState = &endpointState{
ClientConn: es.cc,
parent: es,
endpoint: endpoint,
disableAutoReconnect: es.esOpts.DisableAutoReconnect,
}
epState.childLB = es.childBuilder(epState, es.bOpts)
}
// Update the endpoint state for the endpoint.
newEndpoints.Set(endpoint, epState)
if err := epState.updateClientConnState(balancer.ClientConnState{
BalancerConfig: state.BalancerConfig,
ResolverState: resolver.State{
Endpoints: []resolver.Endpoint{endpoint},
Attributes: state.ResolverState.Attributes,
},
}); err != nil && retErr == nil {
// Keep the first error found from any child.
retErr = err
}
}
// Delete old children that are no longer present.
for e, child := range es.endpoints.All() {
if _, ok := newEndpoints.Get(e); !ok {
child.close()
}
}
if newEndpoints.Len() == 0 {
retErr = balancer.ErrBadResolverState
}
es.mu.Lock()
es.endpoints = newEndpoints
es.inhibitChildUpdates = false
es.updateStateLocked()
es.mu.Unlock()
return retErr
}
// ResolverError forwards the resolver error to all of the endpointSharding's
// children and sends a single synchronous update of the childStates at the end
// of the ResolverError operation.
func (es *endpointSharding) ResolverError(err error) {
es.inhibitUpdatesFromChildren()
for _, child := range es.endpoints.All() {
child.resolverError(err)
}
es.allowUpdatesFromChildren()
}
func (es *endpointSharding) UpdateSubConnState(balancer.SubConn, balancer.SubConnState) {
// UpdateSubConnState is deprecated.
}
func (es *endpointSharding) Close() {
es.inhibitUpdatesFromChildren()
for _, child := range es.endpoints.All() {
child.close()
}
}
func (es *endpointSharding) ExitIdle() {
es.inhibitUpdatesFromChildren()
for _, child := range es.endpoints.All() {
child.exitIdle()
}
es.allowUpdatesFromChildren()
}
func (es *endpointSharding) inhibitUpdatesFromChildren() {
es.mu.Lock()
es.inhibitChildUpdates = true
es.mu.Unlock()
}
func (es *endpointSharding) allowUpdatesFromChildren() {
es.mu.Lock()
es.inhibitChildUpdates = false
es.updateStateLocked()
es.mu.Unlock()
}
// updateStateLocked updates this component's state. It sends the aggregated
// state, and a picker with round robin behavior with all the child states
// present if needed. This method must only be called when inhibitChildUpdates
// is false.
//
// Caller must hold es.mu.
func (es *endpointSharding) updateStateLocked() {
var readyPickers, connectingPickers, idlePickers, transientFailurePickers []balancer.Picker
childStates := make([]ChildState, 0, es.endpoints.Len())
for _, epState := range es.endpoints.All() {
childState := ChildState{
Endpoint: epState.endpoint,
State: epState.state,
ExitIdle: func() { go epState.exitIdle() },
}
childStates = append(childStates, childState)
childPicker := childState.State.Picker
switch childState.State.ConnectivityState {
case connectivity.Ready:
readyPickers = append(readyPickers, childPicker)
case connectivity.Connecting:
connectingPickers = append(connectingPickers, childPicker)
case connectivity.Idle:
idlePickers = append(idlePickers, childPicker)
case connectivity.TransientFailure:
transientFailurePickers = append(transientFailurePickers, childPicker)
// connectivity.Shutdown shouldn't appear.
}
}
// Construct the round robin picker based off the aggregated state. Whatever
// the aggregated state, use the pickers present that are currently in that
// state only.
var aggState connectivity.State
var pickers []balancer.Picker
if len(readyPickers) >= 1 {
aggState = connectivity.Ready
pickers = readyPickers
} else if len(connectingPickers) >= 1 {
aggState = connectivity.Connecting
pickers = connectingPickers
} else if len(idlePickers) >= 1 {
aggState = connectivity.Idle
pickers = idlePickers
} else if len(transientFailurePickers) >= 1 {
aggState = connectivity.TransientFailure
pickers = transientFailurePickers
} else {
aggState = connectivity.TransientFailure
pickers = []balancer.Picker{base.NewErrPicker(errors.New("no children to pick from"))}
} // No children (resolver error before valid update).
es.cc.UpdateState(balancer.State{
ConnectivityState: aggState,
Picker: &pickerWithChildStates{
pickers: pickers,
childStates: childStates,
next: uint32(randIntN(len(pickers))),
},
})
}
// pickerWithChildStates delegates to the pickers it holds in a round robin
// fashion. It also contains the childStates of all the endpointSharding's
// children.
type pickerWithChildStates struct {
pickers []balancer.Picker
childStates []ChildState
next uint32
}
func (p *pickerWithChildStates) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
nextIndex := atomic.AddUint32(&p.next, 1)
picker := p.pickers[nextIndex%uint32(len(p.pickers))]
return picker.Pick(info)
}
// ChildStatesFromPicker returns the state of all the children managed by the
// endpoint sharding balancer that created this picker.
func ChildStatesFromPicker(picker balancer.Picker) []ChildState {
p, ok := picker.(*pickerWithChildStates)
if !ok {
return nil
}
return p.childStates
}
// endpointState is the internal state maintained for each endpoint.
type endpointState struct {
balancer.ClientConn // Embedded to intercept UpdateState
parent *endpointSharding // Parent endpointsharding balancer.
endpoint resolver.Endpoint // Endpoint of the child balancer.
state balancer.State // State of the child balancer.
disableAutoReconnect bool // Whether to disable auto reconnect for this child.
childMu sync.Mutex // Guarantees mutual exclusion for Balancer API calls to the child balancer.
childLB balancer.Balancer // Child balancer.
closed bool // Tracks closure of the child balancer to ensure ExitIdle is not called after Close().
}
func (es *endpointState) UpdateState(state balancer.State) {
es.parent.mu.Lock()
es.state = state
if !es.parent.inhibitChildUpdates {
es.parent.updateStateLocked()
}
es.parent.mu.Unlock()
if state.ConnectivityState == connectivity.Idle && !es.disableAutoReconnect {
go es.exitIdle()
}
}
func (es *endpointState) updateClientConnState(state balancer.ClientConnState) error {
es.childMu.Lock()
err := es.childLB.UpdateClientConnState(state)
es.childMu.Unlock()
return err
}
func (es *endpointState) resolverError(err error) {
es.childMu.Lock()
es.childLB.ResolverError(err)
es.childMu.Unlock()
}
func (es *endpointState) close() {
es.childMu.Lock()
if !es.closed {
es.closed = true
es.childLB.Close()
}
es.childMu.Unlock()
}
func (es *endpointState) exitIdle() {
es.childMu.Lock()
if !es.closed {
es.childLB.ExitIdle()
}
es.childMu.Unlock()
}

View File

@@ -0,0 +1,353 @@
/*
*
* Copyright 2024 gRPC authors.
*
* 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.
*
*/
package endpointsharding_test
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"testing"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/backoff"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/endpointsharding"
"google.golang.org/grpc/balancer/pickfirst"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/internal/balancer/stub"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/internal/stubserver"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/internal/testutils/roundrobin"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/resolver/manual"
"google.golang.org/grpc/serviceconfig"
"google.golang.org/grpc/status"
testgrpc "google.golang.org/grpc/interop/grpc_testing"
testpb "google.golang.org/grpc/interop/grpc_testing"
)
var (
defaultTestTimeout = time.Second * 10
defaultTestShortTimeout = time.Millisecond * 10
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
var logger = grpclog.Component("endpoint-sharding-test")
func init() {
balancer.Register(fakePetioleBuilder{})
}
const fakePetioleName = "fake_petiole"
type fakePetioleBuilder struct{}
func (fakePetioleBuilder) Name() string {
return fakePetioleName
}
func (fakePetioleBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
fp := &fakePetiole{
ClientConn: cc,
bOpts: opts,
}
fp.Balancer = endpointsharding.NewBalancer(fp, opts, balancer.Get(pickfirst.Name).Build, endpointsharding.Options{})
return fp
}
func (fakePetioleBuilder) ParseConfig(json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
return nil, nil
}
// fakePetiole is a load balancer that wraps the endpointShardingBalancer, and
// forwards ClientConnUpdates with a child config of graceful switch that wraps
// pick first. It also intercepts UpdateState to make sure it can access the
// child state maintained by EndpointSharding.
type fakePetiole struct {
balancer.Balancer
balancer.ClientConn
bOpts balancer.BuildOptions
}
func (fp *fakePetiole) UpdateClientConnState(state balancer.ClientConnState) error {
if el := state.ResolverState.Endpoints; len(el) != 2 {
return fmt.Errorf("UpdateClientConnState wants two endpoints, got: %v", el)
}
return fp.Balancer.UpdateClientConnState(state)
}
func (fp *fakePetiole) UpdateState(state balancer.State) {
childStates := endpointsharding.ChildStatesFromPicker(state.Picker)
// Both child states should be present in the child picker. States and
// picker change over the lifecycle of test, but there should always be two.
if len(childStates) != 2 {
logger.Fatal(fmt.Errorf("length of child states received: %v, want 2", len(childStates)))
}
fp.ClientConn.UpdateState(state)
}
// TestEndpointShardingBasic tests the basic functionality of the endpoint
// sharding balancer. It specifies a petiole policy that is essentially a
// wrapper around the endpoint sharder. Two backends are started, with each
// backend's address specified in an endpoint. The petiole does not have a
// special picker, so it should fallback to the default behavior, which is to
// round_robin amongst the endpoint children that are in the aggregated state.
// It also verifies the petiole has access to the raw child state in case it
// wants to implement a custom picker. The test sends a resolver error to the
// endpointsharding balancer and verifies an error picker from the children
// is used while making an RPC.
func (s) TestEndpointShardingBasic(t *testing.T) {
backend1 := stubserver.StartTestService(t, nil)
defer backend1.Stop()
backend2 := stubserver.StartTestService(t, nil)
defer backend2.Stop()
mr := manual.NewBuilderWithScheme("e2e-test")
defer mr.Close()
json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, fakePetioleName)
sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
mr.InitialState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend1.Address}}},
{Addresses: []resolver.Address{{Addr: backend2.Address}}},
},
ServiceConfig: sc,
})
dOpts := []grpc.DialOption{
grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()),
// Use a large backoff delay to avoid the error picker being updated
// too quickly.
grpc.WithConnectParams(grpc.ConnectParams{
Backoff: backoff.Config{
BaseDelay: 2 * defaultTestTimeout,
Multiplier: float64(0),
Jitter: float64(0),
MaxDelay: 2 * defaultTestTimeout,
},
}),
}
cc, err := grpc.NewClient(mr.Scheme()+":///", dOpts...)
if err != nil {
t.Fatalf("Failed to create new client: %v", err)
}
defer cc.Close()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
client := testgrpc.NewTestServiceClient(cc)
// Assert a round robin distribution between the two spun up backends. This
// requires a poll and eventual consistency as both endpoint children do not
// start in state READY.
if err = roundrobin.CheckRoundRobinRPCs(ctx, client, []resolver.Address{{Addr: backend1.Address}, {Addr: backend2.Address}}); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Stopping both the backends should make the channel enter
// TransientFailure.
backend1.Stop()
backend2.Stop()
testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
// When the resolver reports an error, the picker should get updated to
// return the resolver error.
mr.CC().ReportError(errors.New("test error"))
testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
for ; ctx.Err() == nil; <-time.After(time.Millisecond) {
_, err := client.EmptyCall(ctx, &testpb.Empty{})
if err == nil {
t.Fatalf("EmptyCall succeeded when expected to fail with %q", "test error")
}
if strings.Contains(err.Error(), "test error") {
break
}
}
if ctx.Err() != nil {
t.Fatalf("Context timed out waiting for picker with resolver error.")
}
}
// Tests that endpointsharding doesn't automatically re-connect IDLE children.
// The test creates an endpoint with two servers and another with a single
// server. The active service in endpoint 1 is closed to make the child
// pickfirst enter IDLE state. The test verifies that the child pickfirst
// doesn't connect to the second address in the endpoint.
func (s) TestEndpointShardingReconnectDisabled(t *testing.T) {
backend1 := stubserver.StartTestService(t, nil)
defer backend1.Stop()
backend2 := stubserver.StartTestService(t, nil)
defer backend2.Stop()
backend3 := stubserver.StartTestService(t, nil)
defer backend3.Stop()
mr := manual.NewBuilderWithScheme("e2e-test")
defer mr.Close()
name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
bf := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
epOpts := endpointsharding.Options{DisableAutoReconnect: true}
bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build, epOpts)
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
stub.Register(name, bf)
json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, name)
sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
mr.InitialState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend1.Address}, {Addr: backend2.Address}}},
{Addresses: []resolver.Address{{Addr: backend3.Address}}},
},
ServiceConfig: sc,
})
cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("Failed to create new client: %v", err)
}
defer cc.Close()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
client := testgrpc.NewTestServiceClient(cc)
// Assert a round robin distribution between the two spun up backends. This
// requires a poll and eventual consistency as both endpoint children do not
// start in state READY.
if err = roundrobin.CheckRoundRobinRPCs(ctx, client, []resolver.Address{{Addr: backend1.Address}, {Addr: backend3.Address}}); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// On closing the first server, the first child balancer should enter
// IDLE. Since endpointsharding is configured not to auto-reconnect, it will
// remain IDLE and will not try to connect to the second backend in the same
// endpoint.
backend1.Stop()
// CheckRoundRobinRPCs waits for all the backends to become reachable, we
// call it to ensure the picker no longer sends RPCs to closed backend.
if err = roundrobin.CheckRoundRobinRPCs(ctx, client, []resolver.Address{{Addr: backend3.Address}}); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Verify requests go only to backend3 for a short time.
shortCtx, cancel := context.WithTimeout(ctx, defaultTestShortTimeout)
defer cancel()
for ; shortCtx.Err() == nil; <-time.After(time.Millisecond) {
var peer peer.Peer
if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
if status.Code(err) != codes.DeadlineExceeded {
t.Fatalf("EmptyCall() returned unexpected error %v", err)
}
break
}
if got, want := peer.Addr.String(), backend3.Address; got != want {
t.Fatalf("EmptyCall() went to unexpected backend: got %q, want %q", got, want)
}
}
}
// Tests that endpointsharding doesn't automatically re-connect IDLE children
// until cc.Connect() is called. The test creates an endpoint with a single
// address. The client is connected and the active server is closed to make the
// child pickfirst enter IDLE state. The test verifies that the child pickfirst
// doesn't re-connect automatically. The test calls cc.Connect() and verified
// that the balancer connects causing the channel to enter TransientFailure.
func (s) TestEndpointShardingExitIdle(t *testing.T) {
backend := stubserver.StartTestService(t, nil)
defer backend.Stop()
mr := manual.NewBuilderWithScheme("e2e-test")
defer mr.Close()
name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
bf := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
epOpts := endpointsharding.Options{DisableAutoReconnect: true}
bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build, epOpts)
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
ExitIdle: func(bd *stub.BalancerData) {
bd.ChildBalancer.ExitIdle()
},
}
stub.Register(name, bf)
json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, name)
sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
mr.InitialState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend.Address}}},
},
ServiceConfig: sc,
})
cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("Failed to create new client: %v", err)
}
defer cc.Close()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
client := testgrpc.NewTestServiceClient(cc)
if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
}
// On closing the first server, the first child balancer should enter
// IDLE. Since endpointsharding is configured not to auto-reconnect, it will
// remain IDLE and will not try to re-connect
backend.Stop()
testutils.AwaitState(ctx, t, cc, connectivity.Idle)
shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
defer shortCancel()
testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
// The balancer should try to re-connect and fail.
cc.Connect()
testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
}

View File

@@ -0,0 +1,83 @@
/*
*
* Copyright 2025 gRPC authors.
*
* 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.
*
*/
package endpointsharding
import (
"fmt"
"testing"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/resolver"
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
func (s) TestRotateEndpoints(t *testing.T) {
ep := func(addr string) resolver.Endpoint {
return resolver.Endpoint{Addresses: []resolver.Address{{Addr: addr}}}
}
endpoints := []resolver.Endpoint{ep("1"), ep("2"), ep("3"), ep("4"), ep("5")}
testCases := []struct {
rval int
want []resolver.Endpoint
}{
{
rval: 0,
want: []resolver.Endpoint{ep("1"), ep("2"), ep("3"), ep("4"), ep("5")},
},
{
rval: 1,
want: []resolver.Endpoint{ep("2"), ep("3"), ep("4"), ep("5"), ep("1")},
},
{
rval: 2,
want: []resolver.Endpoint{ep("3"), ep("4"), ep("5"), ep("1"), ep("2")},
},
{
rval: 3,
want: []resolver.Endpoint{ep("4"), ep("5"), ep("1"), ep("2"), ep("3")},
},
{
rval: 4,
want: []resolver.Endpoint{ep("5"), ep("1"), ep("2"), ep("3"), ep("4")},
},
}
defer func(r func(int) int) {
randIntN = r
}(randIntN)
for _, tc := range testCases {
t.Run(fmt.Sprint(tc.rval), func(t *testing.T) {
randIntN = func(int) int {
return tc.rval
}
got := rotateEndpoints(endpoints)
if fmt.Sprint(got) != fmt.Sprint(tc.want) {
t.Fatalf("rand=%v; rotateEndpoints(%v) = %v; want %v", tc.rval, endpoints, got, tc.want)
}
})
}
}

View File

@@ -0,0 +1,763 @@
// Copyright 2015 The gRPC Authors
//
// 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.
// This file defines the GRPCLB LoadBalancing protocol.
//
// The canonical version of this proto can be found at
// https://github.com/grpc/grpc-proto/blob/master/grpc/lb/v1/load_balancer.proto
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.12
// protoc v5.27.1
// source: grpc/lb/v1/load_balancer.proto
package grpc_lb_v1
import (
protoreflect "google.golang.org/protobuf/reflect/protoreflect"
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
durationpb "google.golang.org/protobuf/types/known/durationpb"
timestamppb "google.golang.org/protobuf/types/known/timestamppb"
reflect "reflect"
sync "sync"
unsafe "unsafe"
)
const (
// Verify that this generated code is sufficiently up-to-date.
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
// Verify that runtime/protoimpl is sufficiently up-to-date.
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
)
type LoadBalanceRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
// Types that are valid to be assigned to LoadBalanceRequestType:
//
// *LoadBalanceRequest_InitialRequest
// *LoadBalanceRequest_ClientStats
LoadBalanceRequestType isLoadBalanceRequest_LoadBalanceRequestType `protobuf_oneof:"load_balance_request_type"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *LoadBalanceRequest) Reset() {
*x = LoadBalanceRequest{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *LoadBalanceRequest) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*LoadBalanceRequest) ProtoMessage() {}
func (x *LoadBalanceRequest) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[0]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use LoadBalanceRequest.ProtoReflect.Descriptor instead.
func (*LoadBalanceRequest) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{0}
}
func (x *LoadBalanceRequest) GetLoadBalanceRequestType() isLoadBalanceRequest_LoadBalanceRequestType {
if x != nil {
return x.LoadBalanceRequestType
}
return nil
}
func (x *LoadBalanceRequest) GetInitialRequest() *InitialLoadBalanceRequest {
if x != nil {
if x, ok := x.LoadBalanceRequestType.(*LoadBalanceRequest_InitialRequest); ok {
return x.InitialRequest
}
}
return nil
}
func (x *LoadBalanceRequest) GetClientStats() *ClientStats {
if x != nil {
if x, ok := x.LoadBalanceRequestType.(*LoadBalanceRequest_ClientStats); ok {
return x.ClientStats
}
}
return nil
}
type isLoadBalanceRequest_LoadBalanceRequestType interface {
isLoadBalanceRequest_LoadBalanceRequestType()
}
type LoadBalanceRequest_InitialRequest struct {
// This message should be sent on the first request to the load balancer.
InitialRequest *InitialLoadBalanceRequest `protobuf:"bytes,1,opt,name=initial_request,json=initialRequest,proto3,oneof"`
}
type LoadBalanceRequest_ClientStats struct {
// The client stats should be periodically reported to the load balancer
// based on the duration defined in the InitialLoadBalanceResponse.
ClientStats *ClientStats `protobuf:"bytes,2,opt,name=client_stats,json=clientStats,proto3,oneof"`
}
func (*LoadBalanceRequest_InitialRequest) isLoadBalanceRequest_LoadBalanceRequestType() {}
func (*LoadBalanceRequest_ClientStats) isLoadBalanceRequest_LoadBalanceRequestType() {}
type InitialLoadBalanceRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
// The name of the load balanced service (e.g., service.googleapis.com). Its
// length should be less than 256 bytes.
// The name might include a port number. How to handle the port number is up
// to the balancer.
Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *InitialLoadBalanceRequest) Reset() {
*x = InitialLoadBalanceRequest{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *InitialLoadBalanceRequest) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*InitialLoadBalanceRequest) ProtoMessage() {}
func (x *InitialLoadBalanceRequest) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[1]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use InitialLoadBalanceRequest.ProtoReflect.Descriptor instead.
func (*InitialLoadBalanceRequest) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{1}
}
func (x *InitialLoadBalanceRequest) GetName() string {
if x != nil {
return x.Name
}
return ""
}
// Contains the number of calls finished for a particular load balance token.
type ClientStatsPerToken struct {
state protoimpl.MessageState `protogen:"open.v1"`
// See Server.load_balance_token.
LoadBalanceToken string `protobuf:"bytes,1,opt,name=load_balance_token,json=loadBalanceToken,proto3" json:"load_balance_token,omitempty"`
// The total number of RPCs that finished associated with the token.
NumCalls int64 `protobuf:"varint,2,opt,name=num_calls,json=numCalls,proto3" json:"num_calls,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *ClientStatsPerToken) Reset() {
*x = ClientStatsPerToken{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *ClientStatsPerToken) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*ClientStatsPerToken) ProtoMessage() {}
func (x *ClientStatsPerToken) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[2]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use ClientStatsPerToken.ProtoReflect.Descriptor instead.
func (*ClientStatsPerToken) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{2}
}
func (x *ClientStatsPerToken) GetLoadBalanceToken() string {
if x != nil {
return x.LoadBalanceToken
}
return ""
}
func (x *ClientStatsPerToken) GetNumCalls() int64 {
if x != nil {
return x.NumCalls
}
return 0
}
// Contains client level statistics that are useful to load balancing. Each
// count except the timestamp should be reset to zero after reporting the stats.
type ClientStats struct {
state protoimpl.MessageState `protogen:"open.v1"`
// The timestamp of generating the report.
Timestamp *timestamppb.Timestamp `protobuf:"bytes,1,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
// The total number of RPCs that started.
NumCallsStarted int64 `protobuf:"varint,2,opt,name=num_calls_started,json=numCallsStarted,proto3" json:"num_calls_started,omitempty"`
// The total number of RPCs that finished.
NumCallsFinished int64 `protobuf:"varint,3,opt,name=num_calls_finished,json=numCallsFinished,proto3" json:"num_calls_finished,omitempty"`
// The total number of RPCs that failed to reach a server except dropped RPCs.
NumCallsFinishedWithClientFailedToSend int64 `protobuf:"varint,6,opt,name=num_calls_finished_with_client_failed_to_send,json=numCallsFinishedWithClientFailedToSend,proto3" json:"num_calls_finished_with_client_failed_to_send,omitempty"`
// The total number of RPCs that finished and are known to have been received
// by a server.
NumCallsFinishedKnownReceived int64 `protobuf:"varint,7,opt,name=num_calls_finished_known_received,json=numCallsFinishedKnownReceived,proto3" json:"num_calls_finished_known_received,omitempty"`
// The list of dropped calls.
CallsFinishedWithDrop []*ClientStatsPerToken `protobuf:"bytes,8,rep,name=calls_finished_with_drop,json=callsFinishedWithDrop,proto3" json:"calls_finished_with_drop,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *ClientStats) Reset() {
*x = ClientStats{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[3]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *ClientStats) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*ClientStats) ProtoMessage() {}
func (x *ClientStats) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[3]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use ClientStats.ProtoReflect.Descriptor instead.
func (*ClientStats) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{3}
}
func (x *ClientStats) GetTimestamp() *timestamppb.Timestamp {
if x != nil {
return x.Timestamp
}
return nil
}
func (x *ClientStats) GetNumCallsStarted() int64 {
if x != nil {
return x.NumCallsStarted
}
return 0
}
func (x *ClientStats) GetNumCallsFinished() int64 {
if x != nil {
return x.NumCallsFinished
}
return 0
}
func (x *ClientStats) GetNumCallsFinishedWithClientFailedToSend() int64 {
if x != nil {
return x.NumCallsFinishedWithClientFailedToSend
}
return 0
}
func (x *ClientStats) GetNumCallsFinishedKnownReceived() int64 {
if x != nil {
return x.NumCallsFinishedKnownReceived
}
return 0
}
func (x *ClientStats) GetCallsFinishedWithDrop() []*ClientStatsPerToken {
if x != nil {
return x.CallsFinishedWithDrop
}
return nil
}
type LoadBalanceResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
// Types that are valid to be assigned to LoadBalanceResponseType:
//
// *LoadBalanceResponse_InitialResponse
// *LoadBalanceResponse_ServerList
// *LoadBalanceResponse_FallbackResponse
LoadBalanceResponseType isLoadBalanceResponse_LoadBalanceResponseType `protobuf_oneof:"load_balance_response_type"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *LoadBalanceResponse) Reset() {
*x = LoadBalanceResponse{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[4]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *LoadBalanceResponse) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*LoadBalanceResponse) ProtoMessage() {}
func (x *LoadBalanceResponse) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[4]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use LoadBalanceResponse.ProtoReflect.Descriptor instead.
func (*LoadBalanceResponse) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{4}
}
func (x *LoadBalanceResponse) GetLoadBalanceResponseType() isLoadBalanceResponse_LoadBalanceResponseType {
if x != nil {
return x.LoadBalanceResponseType
}
return nil
}
func (x *LoadBalanceResponse) GetInitialResponse() *InitialLoadBalanceResponse {
if x != nil {
if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_InitialResponse); ok {
return x.InitialResponse
}
}
return nil
}
func (x *LoadBalanceResponse) GetServerList() *ServerList {
if x != nil {
if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_ServerList); ok {
return x.ServerList
}
}
return nil
}
func (x *LoadBalanceResponse) GetFallbackResponse() *FallbackResponse {
if x != nil {
if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_FallbackResponse); ok {
return x.FallbackResponse
}
}
return nil
}
type isLoadBalanceResponse_LoadBalanceResponseType interface {
isLoadBalanceResponse_LoadBalanceResponseType()
}
type LoadBalanceResponse_InitialResponse struct {
// This message should be sent on the first response to the client.
InitialResponse *InitialLoadBalanceResponse `protobuf:"bytes,1,opt,name=initial_response,json=initialResponse,proto3,oneof"`
}
type LoadBalanceResponse_ServerList struct {
// Contains the list of servers selected by the load balancer. The client
// should send requests to these servers in the specified order.
ServerList *ServerList `protobuf:"bytes,2,opt,name=server_list,json=serverList,proto3,oneof"`
}
type LoadBalanceResponse_FallbackResponse struct {
// If this field is set, then the client should eagerly enter fallback
// mode (even if there are existing, healthy connections to backends).
FallbackResponse *FallbackResponse `protobuf:"bytes,3,opt,name=fallback_response,json=fallbackResponse,proto3,oneof"`
}
func (*LoadBalanceResponse_InitialResponse) isLoadBalanceResponse_LoadBalanceResponseType() {}
func (*LoadBalanceResponse_ServerList) isLoadBalanceResponse_LoadBalanceResponseType() {}
func (*LoadBalanceResponse_FallbackResponse) isLoadBalanceResponse_LoadBalanceResponseType() {}
type FallbackResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *FallbackResponse) Reset() {
*x = FallbackResponse{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[5]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *FallbackResponse) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*FallbackResponse) ProtoMessage() {}
func (x *FallbackResponse) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[5]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use FallbackResponse.ProtoReflect.Descriptor instead.
func (*FallbackResponse) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{5}
}
type InitialLoadBalanceResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
// This interval defines how often the client should send the client stats
// to the load balancer. Stats should only be reported when the duration is
// positive.
ClientStatsReportInterval *durationpb.Duration `protobuf:"bytes,2,opt,name=client_stats_report_interval,json=clientStatsReportInterval,proto3" json:"client_stats_report_interval,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *InitialLoadBalanceResponse) Reset() {
*x = InitialLoadBalanceResponse{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[6]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *InitialLoadBalanceResponse) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*InitialLoadBalanceResponse) ProtoMessage() {}
func (x *InitialLoadBalanceResponse) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[6]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use InitialLoadBalanceResponse.ProtoReflect.Descriptor instead.
func (*InitialLoadBalanceResponse) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{6}
}
func (x *InitialLoadBalanceResponse) GetClientStatsReportInterval() *durationpb.Duration {
if x != nil {
return x.ClientStatsReportInterval
}
return nil
}
type ServerList struct {
state protoimpl.MessageState `protogen:"open.v1"`
// Contains a list of servers selected by the load balancer. The list will
// be updated when server resolutions change or as needed to balance load
// across more servers. The client should consume the server list in order
// unless instructed otherwise via the client_config.
Servers []*Server `protobuf:"bytes,1,rep,name=servers,proto3" json:"servers,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *ServerList) Reset() {
*x = ServerList{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[7]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *ServerList) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*ServerList) ProtoMessage() {}
func (x *ServerList) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[7]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use ServerList.ProtoReflect.Descriptor instead.
func (*ServerList) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{7}
}
func (x *ServerList) GetServers() []*Server {
if x != nil {
return x.Servers
}
return nil
}
// Contains server information. When the drop field is not true, use the other
// fields.
type Server struct {
state protoimpl.MessageState `protogen:"open.v1"`
// A resolved address for the server, serialized in network-byte-order. It may
// either be an IPv4 or IPv6 address.
IpAddress []byte `protobuf:"bytes,1,opt,name=ip_address,json=ipAddress,proto3" json:"ip_address,omitempty"`
// A resolved port number for the server.
Port int32 `protobuf:"varint,2,opt,name=port,proto3" json:"port,omitempty"`
// An opaque but printable token for load reporting. The client must include
// the token of the picked server into the initial metadata when it starts a
// call to that server. The token is used by the server to verify the request
// and to allow the server to report load to the gRPC LB system. The token is
// also used in client stats for reporting dropped calls.
//
// Its length can be variable but must be less than 50 bytes.
LoadBalanceToken string `protobuf:"bytes,3,opt,name=load_balance_token,json=loadBalanceToken,proto3" json:"load_balance_token,omitempty"`
// Indicates whether this particular request should be dropped by the client.
// If the request is dropped, there will be a corresponding entry in
// ClientStats.calls_finished_with_drop.
Drop bool `protobuf:"varint,4,opt,name=drop,proto3" json:"drop,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *Server) Reset() {
*x = Server{}
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[8]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *Server) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*Server) ProtoMessage() {}
func (x *Server) ProtoReflect() protoreflect.Message {
mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[8]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use Server.ProtoReflect.Descriptor instead.
func (*Server) Descriptor() ([]byte, []int) {
return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{8}
}
func (x *Server) GetIpAddress() []byte {
if x != nil {
return x.IpAddress
}
return nil
}
func (x *Server) GetPort() int32 {
if x != nil {
return x.Port
}
return 0
}
func (x *Server) GetLoadBalanceToken() string {
if x != nil {
return x.LoadBalanceToken
}
return ""
}
func (x *Server) GetDrop() bool {
if x != nil {
return x.Drop
}
return false
}
var File_grpc_lb_v1_load_balancer_proto protoreflect.FileDescriptor
const file_grpc_lb_v1_load_balancer_proto_rawDesc = "" +
"\n" +
"\x1egrpc/lb/v1/load_balancer.proto\x12\n" +
"grpc.lb.v1\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\xc1\x01\n" +
"\x12LoadBalanceRequest\x12P\n" +
"\x0finitial_request\x18\x01 \x01(\v2%.grpc.lb.v1.InitialLoadBalanceRequestH\x00R\x0einitialRequest\x12<\n" +
"\fclient_stats\x18\x02 \x01(\v2\x17.grpc.lb.v1.ClientStatsH\x00R\vclientStatsB\x1b\n" +
"\x19load_balance_request_type\"/\n" +
"\x19InitialLoadBalanceRequest\x12\x12\n" +
"\x04name\x18\x01 \x01(\tR\x04name\"`\n" +
"\x13ClientStatsPerToken\x12,\n" +
"\x12load_balance_token\x18\x01 \x01(\tR\x10loadBalanceToken\x12\x1b\n" +
"\tnum_calls\x18\x02 \x01(\x03R\bnumCalls\"\xb0\x03\n" +
"\vClientStats\x128\n" +
"\ttimestamp\x18\x01 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12*\n" +
"\x11num_calls_started\x18\x02 \x01(\x03R\x0fnumCallsStarted\x12,\n" +
"\x12num_calls_finished\x18\x03 \x01(\x03R\x10numCallsFinished\x12]\n" +
"-num_calls_finished_with_client_failed_to_send\x18\x06 \x01(\x03R&numCallsFinishedWithClientFailedToSend\x12H\n" +
"!num_calls_finished_known_received\x18\a \x01(\x03R\x1dnumCallsFinishedKnownReceived\x12X\n" +
"\x18calls_finished_with_drop\x18\b \x03(\v2\x1f.grpc.lb.v1.ClientStatsPerTokenR\x15callsFinishedWithDropJ\x04\b\x04\x10\x05J\x04\b\x05\x10\x06\"\x90\x02\n" +
"\x13LoadBalanceResponse\x12S\n" +
"\x10initial_response\x18\x01 \x01(\v2&.grpc.lb.v1.InitialLoadBalanceResponseH\x00R\x0finitialResponse\x129\n" +
"\vserver_list\x18\x02 \x01(\v2\x16.grpc.lb.v1.ServerListH\x00R\n" +
"serverList\x12K\n" +
"\x11fallback_response\x18\x03 \x01(\v2\x1c.grpc.lb.v1.FallbackResponseH\x00R\x10fallbackResponseB\x1c\n" +
"\x1aload_balance_response_type\"\x12\n" +
"\x10FallbackResponse\"~\n" +
"\x1aInitialLoadBalanceResponse\x12Z\n" +
"\x1cclient_stats_report_interval\x18\x02 \x01(\v2\x19.google.protobuf.DurationR\x19clientStatsReportIntervalJ\x04\b\x01\x10\x02\"@\n" +
"\n" +
"ServerList\x12,\n" +
"\aservers\x18\x01 \x03(\v2\x12.grpc.lb.v1.ServerR\aserversJ\x04\b\x03\x10\x04\"\x83\x01\n" +
"\x06Server\x12\x1d\n" +
"\n" +
"ip_address\x18\x01 \x01(\fR\tipAddress\x12\x12\n" +
"\x04port\x18\x02 \x01(\x05R\x04port\x12,\n" +
"\x12load_balance_token\x18\x03 \x01(\tR\x10loadBalanceToken\x12\x12\n" +
"\x04drop\x18\x04 \x01(\bR\x04dropJ\x04\b\x05\x10\x062b\n" +
"\fLoadBalancer\x12R\n" +
"\vBalanceLoad\x12\x1e.grpc.lb.v1.LoadBalanceRequest\x1a\x1f.grpc.lb.v1.LoadBalanceResponse(\x010\x01BW\n" +
"\rio.grpc.lb.v1B\x11LoadBalancerProtoP\x01Z1google.golang.org/grpc/balancer/grpclb/grpc_lb_v1b\x06proto3"
var (
file_grpc_lb_v1_load_balancer_proto_rawDescOnce sync.Once
file_grpc_lb_v1_load_balancer_proto_rawDescData []byte
)
func file_grpc_lb_v1_load_balancer_proto_rawDescGZIP() []byte {
file_grpc_lb_v1_load_balancer_proto_rawDescOnce.Do(func() {
file_grpc_lb_v1_load_balancer_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_grpc_lb_v1_load_balancer_proto_rawDesc), len(file_grpc_lb_v1_load_balancer_proto_rawDesc)))
})
return file_grpc_lb_v1_load_balancer_proto_rawDescData
}
var file_grpc_lb_v1_load_balancer_proto_msgTypes = make([]protoimpl.MessageInfo, 9)
var file_grpc_lb_v1_load_balancer_proto_goTypes = []any{
(*LoadBalanceRequest)(nil), // 0: grpc.lb.v1.LoadBalanceRequest
(*InitialLoadBalanceRequest)(nil), // 1: grpc.lb.v1.InitialLoadBalanceRequest
(*ClientStatsPerToken)(nil), // 2: grpc.lb.v1.ClientStatsPerToken
(*ClientStats)(nil), // 3: grpc.lb.v1.ClientStats
(*LoadBalanceResponse)(nil), // 4: grpc.lb.v1.LoadBalanceResponse
(*FallbackResponse)(nil), // 5: grpc.lb.v1.FallbackResponse
(*InitialLoadBalanceResponse)(nil), // 6: grpc.lb.v1.InitialLoadBalanceResponse
(*ServerList)(nil), // 7: grpc.lb.v1.ServerList
(*Server)(nil), // 8: grpc.lb.v1.Server
(*timestamppb.Timestamp)(nil), // 9: google.protobuf.Timestamp
(*durationpb.Duration)(nil), // 10: google.protobuf.Duration
}
var file_grpc_lb_v1_load_balancer_proto_depIdxs = []int32{
1, // 0: grpc.lb.v1.LoadBalanceRequest.initial_request:type_name -> grpc.lb.v1.InitialLoadBalanceRequest
3, // 1: grpc.lb.v1.LoadBalanceRequest.client_stats:type_name -> grpc.lb.v1.ClientStats
9, // 2: grpc.lb.v1.ClientStats.timestamp:type_name -> google.protobuf.Timestamp
2, // 3: grpc.lb.v1.ClientStats.calls_finished_with_drop:type_name -> grpc.lb.v1.ClientStatsPerToken
6, // 4: grpc.lb.v1.LoadBalanceResponse.initial_response:type_name -> grpc.lb.v1.InitialLoadBalanceResponse
7, // 5: grpc.lb.v1.LoadBalanceResponse.server_list:type_name -> grpc.lb.v1.ServerList
5, // 6: grpc.lb.v1.LoadBalanceResponse.fallback_response:type_name -> grpc.lb.v1.FallbackResponse
10, // 7: grpc.lb.v1.InitialLoadBalanceResponse.client_stats_report_interval:type_name -> google.protobuf.Duration
8, // 8: grpc.lb.v1.ServerList.servers:type_name -> grpc.lb.v1.Server
0, // 9: grpc.lb.v1.LoadBalancer.BalanceLoad:input_type -> grpc.lb.v1.LoadBalanceRequest
4, // 10: grpc.lb.v1.LoadBalancer.BalanceLoad:output_type -> grpc.lb.v1.LoadBalanceResponse
10, // [10:11] is the sub-list for method output_type
9, // [9:10] is the sub-list for method input_type
9, // [9:9] is the sub-list for extension type_name
9, // [9:9] is the sub-list for extension extendee
0, // [0:9] is the sub-list for field type_name
}
func init() { file_grpc_lb_v1_load_balancer_proto_init() }
func file_grpc_lb_v1_load_balancer_proto_init() {
if File_grpc_lb_v1_load_balancer_proto != nil {
return
}
file_grpc_lb_v1_load_balancer_proto_msgTypes[0].OneofWrappers = []any{
(*LoadBalanceRequest_InitialRequest)(nil),
(*LoadBalanceRequest_ClientStats)(nil),
}
file_grpc_lb_v1_load_balancer_proto_msgTypes[4].OneofWrappers = []any{
(*LoadBalanceResponse_InitialResponse)(nil),
(*LoadBalanceResponse_ServerList)(nil),
(*LoadBalanceResponse_FallbackResponse)(nil),
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_grpc_lb_v1_load_balancer_proto_rawDesc), len(file_grpc_lb_v1_load_balancer_proto_rawDesc)),
NumEnums: 0,
NumMessages: 9,
NumExtensions: 0,
NumServices: 1,
},
GoTypes: file_grpc_lb_v1_load_balancer_proto_goTypes,
DependencyIndexes: file_grpc_lb_v1_load_balancer_proto_depIdxs,
MessageInfos: file_grpc_lb_v1_load_balancer_proto_msgTypes,
}.Build()
File_grpc_lb_v1_load_balancer_proto = out.File
file_grpc_lb_v1_load_balancer_proto_goTypes = nil
file_grpc_lb_v1_load_balancer_proto_depIdxs = nil
}

View File

@@ -0,0 +1,134 @@
// Copyright 2015 The gRPC Authors
//
// 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.
// This file defines the GRPCLB LoadBalancing protocol.
//
// The canonical version of this proto can be found at
// https://github.com/grpc/grpc-proto/blob/master/grpc/lb/v1/load_balancer.proto
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.6.2
// - protoc v5.27.1
// source: grpc/lb/v1/load_balancer.proto
package grpc_lb_v1
import (
context "context"
grpc "google.golang.org/grpc"
codes "google.golang.org/grpc/codes"
status "google.golang.org/grpc/status"
)
// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
// Requires gRPC-Go v1.64.0 or later.
const _ = grpc.SupportPackageIsVersion9
const (
LoadBalancer_BalanceLoad_FullMethodName = "/grpc.lb.v1.LoadBalancer/BalanceLoad"
)
// LoadBalancerClient is the client API for LoadBalancer service.
//
// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
type LoadBalancerClient interface {
// Bidirectional rpc to get a list of servers.
BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[LoadBalanceRequest, LoadBalanceResponse], error)
}
type loadBalancerClient struct {
cc grpc.ClientConnInterface
}
func NewLoadBalancerClient(cc grpc.ClientConnInterface) LoadBalancerClient {
return &loadBalancerClient{cc}
}
func (c *loadBalancerClient) BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[LoadBalanceRequest, LoadBalanceResponse], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &LoadBalancer_ServiceDesc.Streams[0], LoadBalancer_BalanceLoad_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[LoadBalanceRequest, LoadBalanceResponse]{ClientStream: stream}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type LoadBalancer_BalanceLoadClient = grpc.BidiStreamingClient[LoadBalanceRequest, LoadBalanceResponse]
// LoadBalancerServer is the server API for LoadBalancer service.
// All implementations should embed UnimplementedLoadBalancerServer
// for forward compatibility.
type LoadBalancerServer interface {
// Bidirectional rpc to get a list of servers.
BalanceLoad(grpc.BidiStreamingServer[LoadBalanceRequest, LoadBalanceResponse]) error
}
// UnimplementedLoadBalancerServer should be embedded to have
// forward compatible implementations.
//
// NOTE: this should be embedded by value instead of pointer to avoid a nil
// pointer dereference when methods are called.
type UnimplementedLoadBalancerServer struct{}
func (UnimplementedLoadBalancerServer) BalanceLoad(grpc.BidiStreamingServer[LoadBalanceRequest, LoadBalanceResponse]) error {
return status.Error(codes.Unimplemented, "method BalanceLoad not implemented")
}
func (UnimplementedLoadBalancerServer) testEmbeddedByValue() {}
// UnsafeLoadBalancerServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to LoadBalancerServer will
// result in compilation errors.
type UnsafeLoadBalancerServer interface {
mustEmbedUnimplementedLoadBalancerServer()
}
func RegisterLoadBalancerServer(s grpc.ServiceRegistrar, srv LoadBalancerServer) {
// If the following call panics, it indicates UnimplementedLoadBalancerServer was
// embedded by pointer and is nil. This will cause panics if an
// unimplemented method is ever invoked, so we test this at initialization
// time to prevent it from happening at runtime later due to I/O.
if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
t.testEmbeddedByValue()
}
s.RegisterService(&LoadBalancer_ServiceDesc, srv)
}
func _LoadBalancer_BalanceLoad_Handler(srv interface{}, stream grpc.ServerStream) error {
return srv.(LoadBalancerServer).BalanceLoad(&grpc.GenericServerStream[LoadBalanceRequest, LoadBalanceResponse]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type LoadBalancer_BalanceLoadServer = grpc.BidiStreamingServer[LoadBalanceRequest, LoadBalanceResponse]
// LoadBalancer_ServiceDesc is the grpc.ServiceDesc for LoadBalancer service.
// It's only intended for direct use with grpc.RegisterService,
// and not to be introspected or modified (even as a copy)
var LoadBalancer_ServiceDesc = grpc.ServiceDesc{
ServiceName: "grpc.lb.v1.LoadBalancer",
HandlerType: (*LoadBalancerServer)(nil),
Methods: []grpc.MethodDesc{},
Streams: []grpc.StreamDesc{
{
StreamName: "BalanceLoad",
Handler: _LoadBalancer_BalanceLoad_Handler,
ServerStreams: true,
ClientStreams: true,
},
},
Metadata: "grpc/lb/v1/load_balancer.proto",
}

535
balancer/grpclb/grpclb.go Normal file
View File

@@ -0,0 +1,535 @@
/*
*
* Copyright 2016 gRPC authors.
*
* 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.
*
*/
// Package grpclb defines a grpclb balancer.
//
// To install grpclb balancer, import this package as:
//
// import _ "google.golang.org/grpc/balancer/grpclb"
package grpclb
import (
"context"
"errors"
"fmt"
"sync"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/base"
grpclbstate "google.golang.org/grpc/balancer/grpclb/state"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/internal/backoff"
internalgrpclog "google.golang.org/grpc/internal/grpclog"
"google.golang.org/grpc/internal/pretty"
"google.golang.org/grpc/internal/resolver/dns"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/resolver/manual"
"google.golang.org/protobuf/types/known/durationpb"
lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
)
const (
lbTokenKey = "lb-token"
defaultFallbackTimeout = 10 * time.Second
grpclbName = "grpclb"
)
var errServerTerminatedConnection = errors.New("grpclb: failed to recv server list: server terminated connection")
var logger = grpclog.Component("grpclb")
func convertDuration(d *durationpb.Duration) time.Duration {
if d == nil {
return 0
}
return time.Duration(d.Seconds)*time.Second + time.Duration(d.Nanos)*time.Nanosecond
}
// Client API for LoadBalancer service.
// Mostly copied from generated pb.go file.
// To avoid circular dependency.
type loadBalancerClient struct {
cc *grpc.ClientConn
}
func (c *loadBalancerClient) BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (*balanceLoadClientStream, error) {
desc := &grpc.StreamDesc{
StreamName: "BalanceLoad",
ServerStreams: true,
ClientStreams: true,
}
stream, err := c.cc.NewStream(ctx, desc, "/grpc.lb.v1.LoadBalancer/BalanceLoad", opts...)
if err != nil {
return nil, err
}
x := &balanceLoadClientStream{stream}
return x, nil
}
type balanceLoadClientStream struct {
grpc.ClientStream
}
func (x *balanceLoadClientStream) Send(m *lbpb.LoadBalanceRequest) error {
return x.ClientStream.SendMsg(m)
}
func (x *balanceLoadClientStream) Recv() (*lbpb.LoadBalanceResponse, error) {
m := new(lbpb.LoadBalanceResponse)
if err := x.ClientStream.RecvMsg(m); err != nil {
return nil, err
}
return m, nil
}
func init() {
balancer.Register(newLBBuilder())
dns.EnableSRVLookups = true
}
// newLBBuilder creates a builder for grpclb.
func newLBBuilder() balancer.Builder {
return newLBBuilderWithFallbackTimeout(defaultFallbackTimeout)
}
// newLBBuilderWithFallbackTimeout creates a grpclb builder with the given
// fallbackTimeout. If no response is received from the remote balancer within
// fallbackTimeout, the backend addresses from the resolved address list will be
// used.
//
// Only call this function when a non-default fallback timeout is needed.
func newLBBuilderWithFallbackTimeout(fallbackTimeout time.Duration) balancer.Builder {
return &lbBuilder{
fallbackTimeout: fallbackTimeout,
}
}
type lbBuilder struct {
fallbackTimeout time.Duration
}
func (b *lbBuilder) Name() string {
return grpclbName
}
func (b *lbBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) balancer.Balancer {
// This generates a manual resolver builder with a fixed scheme. This
// scheme will be used to dial to remote LB, so we can send filtered
// address updates to remote LB ClientConn using this manual resolver.
mr := manual.NewBuilderWithScheme("grpclb-internal")
// ResolveNow() on this manual resolver is forwarded to the parent
// ClientConn, so when grpclb client loses contact with the remote balancer,
// the parent ClientConn's resolver will re-resolve.
mr.ResolveNowCallback = cc.ResolveNow
lb := &lbBalancer{
cc: newLBCacheClientConn(cc),
dialTarget: opt.Target.Endpoint(),
target: opt.Target.Endpoint(),
opt: opt,
fallbackTimeout: b.fallbackTimeout,
doneCh: make(chan struct{}),
manualResolver: mr,
subConns: make(map[resolver.Address]balancer.SubConn),
scStates: make(map[balancer.SubConn]connectivity.State),
picker: base.NewErrPicker(balancer.ErrNoSubConnAvailable),
clientStats: newRPCStats(),
backoff: backoff.DefaultExponential, // TODO: make backoff configurable.
}
lb.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[grpclb %p] ", lb))
var err error
if opt.CredsBundle != nil {
lb.grpclbClientConnCreds, err = opt.CredsBundle.NewWithMode(internal.CredsBundleModeBalancer)
if err != nil {
lb.logger.Warningf("Failed to create credentials used for connecting to grpclb: %v", err)
}
lb.grpclbBackendCreds, err = opt.CredsBundle.NewWithMode(internal.CredsBundleModeBackendFromBalancer)
if err != nil {
lb.logger.Warningf("Failed to create credentials used for connecting to backends returned by grpclb: %v", err)
}
}
return lb
}
type lbBalancer struct {
cc *lbCacheClientConn
dialTarget string // user's dial target
target string // same as dialTarget unless overridden in service config
opt balancer.BuildOptions
logger *internalgrpclog.PrefixLogger
usePickFirst bool
// grpclbClientConnCreds is the creds bundle to be used to connect to grpclb
// servers. If it's nil, use the TransportCredentials from BuildOptions
// instead.
grpclbClientConnCreds credentials.Bundle
// grpclbBackendCreds is the creds bundle to be used for addresses that are
// returned by grpclb server. If it's nil, don't set anything when creating
// SubConns.
grpclbBackendCreds credentials.Bundle
fallbackTimeout time.Duration
doneCh chan struct{}
// manualResolver is used in the remote LB ClientConn inside grpclb. When
// resolved address updates are received by grpclb, filtered updates will be
// sent to remote LB ClientConn through this resolver.
manualResolver *manual.Resolver
// The ClientConn to talk to the remote balancer.
ccRemoteLB *remoteBalancerCCWrapper
// backoff for calling remote balancer.
backoff backoff.Strategy
// Support client side load reporting. Each picker gets a reference to this,
// and will update its content.
clientStats *rpcStats
mu sync.Mutex // guards everything following.
// The full server list including drops, used to check if the newly received
// serverList contains anything new. Each generate picker will also have
// reference to this list to do the first layer pick.
fullServerList []*lbpb.Server
// Backend addresses. It's kept so the addresses are available when
// switching between round_robin and pickfirst.
backendAddrs []resolver.Address
// All backends addresses, with metadata set to nil. This list contains all
// backend addresses in the same order and with the same duplicates as in
// serverlist. When generating picker, a SubConn slice with the same order
// but with only READY SCs will be generated.
backendAddrsWithoutMetadata []resolver.Address
// Roundrobin functionalities.
state connectivity.State
subConns map[resolver.Address]balancer.SubConn // Used to new/shutdown SubConn.
scStates map[balancer.SubConn]connectivity.State // Used to filter READY SubConns.
picker balancer.Picker
// Support fallback to resolved backend addresses if there's no response
// from remote balancer within fallbackTimeout.
remoteBalancerConnected bool
serverListReceived bool
inFallback bool
// resolvedBackendAddrs is resolvedAddrs minus remote balancers. It's set
// when resolved address updates are received, and read in the goroutine
// handling fallback.
resolvedBackendAddrs []resolver.Address
connErr error // the last connection error
}
// regeneratePicker takes a snapshot of the balancer, and generates a picker from
// it. The picker
// - always returns ErrTransientFailure if the balancer is in TransientFailure,
// - does two layer roundrobin pick otherwise.
//
// Caller must hold lb.mu.
func (lb *lbBalancer) regeneratePicker(resetDrop bool) {
if lb.state == connectivity.TransientFailure {
lb.picker = base.NewErrPicker(fmt.Errorf("all SubConns are in TransientFailure, last connection error: %v", lb.connErr))
return
}
if lb.state == connectivity.Connecting {
lb.picker = base.NewErrPicker(balancer.ErrNoSubConnAvailable)
return
}
var readySCs []balancer.SubConn
if lb.usePickFirst {
for _, sc := range lb.subConns {
readySCs = append(readySCs, sc)
break
}
} else {
for _, a := range lb.backendAddrsWithoutMetadata {
if sc, ok := lb.subConns[a]; ok {
if st, ok := lb.scStates[sc]; ok && st == connectivity.Ready {
readySCs = append(readySCs, sc)
}
}
}
}
if len(readySCs) <= 0 {
// If there's no ready SubConns, always re-pick. This is to avoid drops
// unless at least one SubConn is ready. Otherwise we may drop more
// often than want because of drops + re-picks(which become re-drops).
//
// This doesn't seem to be necessary after the connecting check above.
// Kept for safety.
lb.picker = base.NewErrPicker(balancer.ErrNoSubConnAvailable)
return
}
if lb.inFallback {
lb.picker = newRRPicker(readySCs)
return
}
if resetDrop {
lb.picker = newLBPicker(lb.fullServerList, readySCs, lb.clientStats)
return
}
prevLBPicker, ok := lb.picker.(*lbPicker)
if !ok {
lb.picker = newLBPicker(lb.fullServerList, readySCs, lb.clientStats)
return
}
prevLBPicker.updateReadySCs(readySCs)
}
// aggregateSubConnStats calculate the aggregated state of SubConns in
// lb.SubConns. These SubConns are subconns in use (when switching between
// fallback and grpclb). lb.scState contains states for all SubConns, including
// those in cache (SubConns are cached for 10 seconds after shutdown).
//
// The aggregated state is:
// - If at least one SubConn in Ready, the aggregated state is Ready;
// - Else if at least one SubConn in Connecting or IDLE, the aggregated state is Connecting;
// - It's OK to consider IDLE as Connecting. SubConns never stay in IDLE,
// they start to connect immediately. But there's a race between the overall
// state is reported, and when the new SubConn state arrives. And SubConns
// never go back to IDLE.
// - Else the aggregated state is TransientFailure.
func (lb *lbBalancer) aggregateSubConnStates() connectivity.State {
var numConnecting uint64
for _, sc := range lb.subConns {
if state, ok := lb.scStates[sc]; ok {
switch state {
case connectivity.Ready:
return connectivity.Ready
case connectivity.Connecting, connectivity.Idle:
numConnecting++
}
}
}
if numConnecting > 0 {
return connectivity.Connecting
}
return connectivity.TransientFailure
}
// UpdateSubConnState is unused; NewSubConn's options always specifies
// updateSubConnState as the listener.
func (lb *lbBalancer) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
lb.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, scs)
}
func (lb *lbBalancer) updateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
s := scs.ConnectivityState
if lb.logger.V(2) {
lb.logger.Infof("SubConn state change: %p, %v", sc, s)
}
lb.mu.Lock()
defer lb.mu.Unlock()
oldS, ok := lb.scStates[sc]
if !ok {
if lb.logger.V(2) {
lb.logger.Infof("Received state change for an unknown SubConn: %p, %v", sc, s)
}
return
}
lb.scStates[sc] = s
switch s {
case connectivity.Idle:
sc.Connect()
case connectivity.Shutdown:
// When an address was removed by resolver, b called Shutdown but kept
// the sc's state in scStates. Remove state for this sc here.
delete(lb.scStates, sc)
case connectivity.TransientFailure:
lb.connErr = scs.ConnectionError
}
// Force regenerate picker if
// - this sc became ready from not-ready
// - this sc became not-ready from ready
lb.updateStateAndPicker((oldS == connectivity.Ready) != (s == connectivity.Ready), false)
// Enter fallback when the aggregated state is not Ready and the connection
// to remote balancer is lost.
if lb.state != connectivity.Ready {
if !lb.inFallback && !lb.remoteBalancerConnected {
// Enter fallback.
lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
}
}
}
// updateStateAndPicker re-calculate the aggregated state, and regenerate picker
// if overall state is changed.
//
// If forceRegeneratePicker is true, picker will be regenerated.
func (lb *lbBalancer) updateStateAndPicker(forceRegeneratePicker bool, resetDrop bool) {
oldAggrState := lb.state
lb.state = lb.aggregateSubConnStates()
// Regenerate picker when one of the following happens:
// - caller wants to regenerate
// - the aggregated state changed
if forceRegeneratePicker || (lb.state != oldAggrState) {
lb.regeneratePicker(resetDrop)
}
var cc balancer.ClientConn = lb.cc
if lb.usePickFirst {
// Bypass the caching layer that would wrap the picker.
cc = lb.cc.ClientConn
}
cc.UpdateState(balancer.State{ConnectivityState: lb.state, Picker: lb.picker})
}
// fallbackToBackendsAfter blocks for fallbackTimeout and falls back to use
// resolved backends (backends received from resolver, not from remote balancer)
// if no connection to remote balancers was successful.
func (lb *lbBalancer) fallbackToBackendsAfter(fallbackTimeout time.Duration) {
timer := time.NewTimer(fallbackTimeout)
defer timer.Stop()
select {
case <-timer.C:
case <-lb.doneCh:
return
}
lb.mu.Lock()
if lb.inFallback || lb.serverListReceived {
lb.mu.Unlock()
return
}
// Enter fallback.
lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
lb.mu.Unlock()
}
func (lb *lbBalancer) handleServiceConfig(gc *grpclbServiceConfig) {
lb.mu.Lock()
defer lb.mu.Unlock()
// grpclb uses the user's dial target to populate the `Name` field of the
// `InitialLoadBalanceRequest` message sent to the remote balancer. But when
// grpclb is used a child policy in the context of RLS, we want the `Name`
// field to be populated with the value received from the RLS server. To
// support this use case, an optional "target_name" field has been added to
// the grpclb LB policy's config. If specified, it overrides the name of
// the target to be sent to the remote balancer; if not, the target to be
// sent to the balancer will continue to be obtained from the target URI
// passed to the gRPC client channel. Whenever that target to be sent to the
// balancer is updated, we need to restart the stream to the balancer as
// this target is sent in the first message on the stream.
if gc != nil {
target := lb.dialTarget
if gc.ServiceName != "" {
target = gc.ServiceName
}
if target != lb.target {
lb.target = target
if lb.ccRemoteLB != nil {
lb.ccRemoteLB.cancelRemoteBalancerCall()
}
}
}
newUsePickFirst := childIsPickFirst(gc)
if lb.usePickFirst == newUsePickFirst {
return
}
if lb.logger.V(2) {
lb.logger.Infof("Switching mode. Is pick_first used for backends? %v", newUsePickFirst)
}
lb.refreshSubConns(lb.backendAddrs, lb.inFallback, newUsePickFirst)
}
func (lb *lbBalancer) ResolverError(error) {
// Ignore resolver errors. GRPCLB is not selected unless the resolver
// works at least once.
}
func (lb *lbBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
if lb.logger.V(2) {
lb.logger.Infof("UpdateClientConnState: %s", pretty.ToJSON(ccs))
}
gc, _ := ccs.BalancerConfig.(*grpclbServiceConfig)
lb.handleServiceConfig(gc)
backendAddrs := ccs.ResolverState.Addresses
var remoteBalancerAddrs []resolver.Address
if sd := grpclbstate.Get(ccs.ResolverState); sd != nil {
// Override any balancer addresses provided via
// ccs.ResolverState.Addresses.
remoteBalancerAddrs = sd.BalancerAddresses
}
if len(backendAddrs)+len(remoteBalancerAddrs) == 0 {
// There should be at least one address, either grpclb server or
// fallback. Empty address is not valid.
return balancer.ErrBadResolverState
}
if len(remoteBalancerAddrs) == 0 {
if lb.ccRemoteLB != nil {
lb.ccRemoteLB.close()
lb.ccRemoteLB = nil
}
} else if lb.ccRemoteLB == nil {
// First time receiving resolved addresses, create a cc to remote
// balancers.
if err := lb.newRemoteBalancerCCWrapper(); err != nil {
return err
}
// Start the fallback goroutine.
go lb.fallbackToBackendsAfter(lb.fallbackTimeout)
}
if lb.ccRemoteLB != nil {
// cc to remote balancers uses lb.manualResolver. Send the updated remote
// balancer addresses to it through manualResolver.
lb.manualResolver.UpdateState(resolver.State{Addresses: remoteBalancerAddrs})
}
lb.mu.Lock()
lb.resolvedBackendAddrs = backendAddrs
if len(remoteBalancerAddrs) == 0 || lb.inFallback {
// If there's no remote balancer address in ClientConn update, grpclb
// enters fallback mode immediately.
//
// If a new update is received while grpclb is in fallback, update the
// list of backends being used to the new fallback backends.
lb.refreshSubConns(lb.resolvedBackendAddrs, true, lb.usePickFirst)
}
lb.mu.Unlock()
return nil
}
func (lb *lbBalancer) Close() {
select {
case <-lb.doneCh:
return
default:
}
close(lb.doneCh)
if lb.ccRemoteLB != nil {
lb.ccRemoteLB.close()
}
lb.cc.close()
}
func (lb *lbBalancer) ExitIdle() {}

View File

@@ -0,0 +1,67 @@
/*
*
* Copyright 2019 gRPC authors.
*
* 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.
*
*/
package grpclb
import (
"encoding/json"
"google.golang.org/grpc/balancer/pickfirst"
"google.golang.org/grpc/balancer/roundrobin"
"google.golang.org/grpc/serviceconfig"
)
const (
roundRobinName = roundrobin.Name
pickFirstName = pickfirst.Name
)
type grpclbServiceConfig struct {
serviceconfig.LoadBalancingConfig
ChildPolicy *[]map[string]json.RawMessage
ServiceName string
}
func (b *lbBuilder) ParseConfig(lbConfig json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
ret := &grpclbServiceConfig{}
if err := json.Unmarshal(lbConfig, ret); err != nil {
return nil, err
}
return ret, nil
}
func childIsPickFirst(sc *grpclbServiceConfig) bool {
if sc == nil {
return false
}
childConfigs := sc.ChildPolicy
if childConfigs == nil {
return false
}
for _, childC := range *childConfigs {
// If round_robin exists before pick_first, return false
if _, ok := childC[roundRobinName]; ok {
return false
}
// If pick_first is before round_robin, return true
if _, ok := childC[pickFirstName]; ok {
return true
}
}
return false
}

View File

@@ -0,0 +1,123 @@
/*
*
* Copyright 2019 gRPC authors.
*
* 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.
*
*/
package grpclb
import (
"encoding/json"
"testing"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc/serviceconfig"
)
func (s) TestParse(t *testing.T) {
tests := []struct {
name string
sc string
want serviceconfig.LoadBalancingConfig
wantErr bool
}{
{
name: "empty",
sc: "",
want: nil,
wantErr: true,
},
{
name: "success1",
sc: `
{
"childPolicy": [
{"pick_first":{}}
],
"serviceName": "foo-service"
}`,
want: &grpclbServiceConfig{
ChildPolicy: &[]map[string]json.RawMessage{
{"pick_first": json.RawMessage("{}")},
},
ServiceName: "foo-service",
},
},
{
name: "success2",
sc: `
{
"childPolicy": [
{"round_robin":{}},
{"pick_first":{}}
],
"serviceName": "foo-service"
}`,
want: &grpclbServiceConfig{
ChildPolicy: &[]map[string]json.RawMessage{
{"round_robin": json.RawMessage("{}")},
{"pick_first": json.RawMessage("{}")},
},
ServiceName: "foo-service",
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := (&lbBuilder{}).ParseConfig(json.RawMessage(tt.sc))
if (err != nil) != (tt.wantErr) {
t.Fatalf("ParseConfig(%q) returned error: %v, wantErr: %v", tt.sc, err, tt.wantErr)
}
if diff := cmp.Diff(tt.want, got); diff != "" {
t.Fatalf("ParseConfig(%q) returned unexpected difference (-want +got):\n%s", tt.sc, diff)
}
})
}
}
func (s) TestChildIsPickFirst(t *testing.T) {
tests := []struct {
name string
s string
want bool
}{
{
name: "pickfirst_only",
s: `{"childPolicy":[{"pick_first":{}}]}`,
want: true,
},
{
name: "pickfirst_before_rr",
s: `{"childPolicy":[{"pick_first":{}},{"round_robin":{}}]}`,
want: true,
},
{
name: "rr_before_pickfirst",
s: `{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}`,
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gc, err := (&lbBuilder{}).ParseConfig(json.RawMessage(tt.s))
if err != nil {
t.Fatalf("Parse(%v) = _, %v; want _, nil", tt.s, err)
}
if got := childIsPickFirst(gc.(*grpclbServiceConfig)); got != tt.want {
t.Errorf("childIsPickFirst() = %v, want %v", got, tt.want)
}
})
}
}

View File

@@ -0,0 +1,193 @@
/*
*
* Copyright 2017 gRPC authors.
*
* 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.
*
*/
package grpclb
import (
rand "math/rand/v2"
"sync"
"sync/atomic"
"google.golang.org/grpc/balancer"
lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
// rpcStats is same as lbpb.ClientStats, except that numCallsDropped is a map
// instead of a slice.
type rpcStats struct {
// Only access the following fields atomically.
numCallsStarted int64
numCallsFinished int64
numCallsFinishedWithClientFailedToSend int64
numCallsFinishedKnownReceived int64
mu sync.Mutex
// map load_balance_token -> num_calls_dropped
numCallsDropped map[string]int64
}
func newRPCStats() *rpcStats {
return &rpcStats{
numCallsDropped: make(map[string]int64),
}
}
func isZeroStats(stats *lbpb.ClientStats) bool {
return len(stats.CallsFinishedWithDrop) == 0 &&
stats.NumCallsStarted == 0 &&
stats.NumCallsFinished == 0 &&
stats.NumCallsFinishedWithClientFailedToSend == 0 &&
stats.NumCallsFinishedKnownReceived == 0
}
// toClientStats converts rpcStats to lbpb.ClientStats, and clears rpcStats.
func (s *rpcStats) toClientStats() *lbpb.ClientStats {
stats := &lbpb.ClientStats{
NumCallsStarted: atomic.SwapInt64(&s.numCallsStarted, 0),
NumCallsFinished: atomic.SwapInt64(&s.numCallsFinished, 0),
NumCallsFinishedWithClientFailedToSend: atomic.SwapInt64(&s.numCallsFinishedWithClientFailedToSend, 0),
NumCallsFinishedKnownReceived: atomic.SwapInt64(&s.numCallsFinishedKnownReceived, 0),
}
s.mu.Lock()
dropped := s.numCallsDropped
s.numCallsDropped = make(map[string]int64)
s.mu.Unlock()
for token, count := range dropped {
stats.CallsFinishedWithDrop = append(stats.CallsFinishedWithDrop, &lbpb.ClientStatsPerToken{
LoadBalanceToken: token,
NumCalls: count,
})
}
return stats
}
func (s *rpcStats) drop(token string) {
atomic.AddInt64(&s.numCallsStarted, 1)
s.mu.Lock()
s.numCallsDropped[token]++
s.mu.Unlock()
atomic.AddInt64(&s.numCallsFinished, 1)
}
func (s *rpcStats) failedToSend() {
atomic.AddInt64(&s.numCallsStarted, 1)
atomic.AddInt64(&s.numCallsFinishedWithClientFailedToSend, 1)
atomic.AddInt64(&s.numCallsFinished, 1)
}
func (s *rpcStats) knownReceived() {
atomic.AddInt64(&s.numCallsStarted, 1)
atomic.AddInt64(&s.numCallsFinishedKnownReceived, 1)
atomic.AddInt64(&s.numCallsFinished, 1)
}
// rrPicker does roundrobin on subConns. It's typically used when there's no
// response from remote balancer, and grpclb falls back to the resolved
// backends.
//
// It guaranteed that len(subConns) > 0.
type rrPicker struct {
mu sync.Mutex
subConns []balancer.SubConn // The subConns that were READY when taking the snapshot.
subConnsNext int
}
func newRRPicker(readySCs []balancer.SubConn) *rrPicker {
return &rrPicker{
subConns: readySCs,
subConnsNext: rand.IntN(len(readySCs)),
}
}
func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
p.mu.Lock()
defer p.mu.Unlock()
sc := p.subConns[p.subConnsNext]
p.subConnsNext = (p.subConnsNext + 1) % len(p.subConns)
return balancer.PickResult{SubConn: sc}, nil
}
// lbPicker does two layers of picks:
//
// First layer: roundrobin on all servers in serverList, including drops and backends.
// - If it picks a drop, the RPC will fail as being dropped.
// - If it picks a backend, do a second layer pick to pick the real backend.
//
// Second layer: roundrobin on all READY backends.
//
// It's guaranteed that len(serverList) > 0.
type lbPicker struct {
mu sync.Mutex
serverList []*lbpb.Server
serverListNext int
subConns []balancer.SubConn // The subConns that were READY when taking the snapshot.
subConnsNext int
stats *rpcStats
}
func newLBPicker(serverList []*lbpb.Server, readySCs []balancer.SubConn, stats *rpcStats) *lbPicker {
return &lbPicker{
serverList: serverList,
subConns: readySCs,
subConnsNext: rand.IntN(len(readySCs)),
stats: stats,
}
}
func (p *lbPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
p.mu.Lock()
defer p.mu.Unlock()
// Layer one roundrobin on serverList.
s := p.serverList[p.serverListNext]
p.serverListNext = (p.serverListNext + 1) % len(p.serverList)
// If it's a drop, return an error and fail the RPC.
if s.Drop {
p.stats.drop(s.LoadBalanceToken)
return balancer.PickResult{}, status.Errorf(codes.Unavailable, "request dropped by grpclb")
}
// If not a drop but there's no ready subConns.
if len(p.subConns) <= 0 {
return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
}
// Return the next ready subConn in the list, also collect rpc stats.
sc := p.subConns[p.subConnsNext]
p.subConnsNext = (p.subConnsNext + 1) % len(p.subConns)
done := func(info balancer.DoneInfo) {
if !info.BytesSent {
p.stats.failedToSend()
} else if info.BytesReceived {
p.stats.knownReceived()
}
}
return balancer.PickResult{SubConn: sc, Done: done}, nil
}
func (p *lbPicker) updateReadySCs(readySCs []balancer.SubConn) {
p.mu.Lock()
defer p.mu.Unlock()
p.subConns = readySCs
p.subConnsNext = p.subConnsNext % len(readySCs)
}

View File

@@ -0,0 +1,459 @@
/*
*
* Copyright 2017 gRPC authors.
*
* 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.
*
*/
package grpclb
import (
"context"
"fmt"
"io"
"net"
"net/netip"
"sync"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/internal/backoff"
imetadata "google.golang.org/grpc/internal/metadata"
"google.golang.org/grpc/keepalive"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/resolver"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
)
func serverListEqual(a, b []*lbpb.Server) bool {
if len(a) != len(b) {
return false
}
for i := 0; i < len(a); i++ {
if !proto.Equal(a[i], b[i]) {
return false
}
}
return true
}
// processServerList updates balancer's internal state, create/remove SubConns
// and regenerates picker using the received serverList.
func (lb *lbBalancer) processServerList(l *lbpb.ServerList) {
if lb.logger.V(2) {
lb.logger.Infof("Processing server list: %#v", l)
}
lb.mu.Lock()
defer lb.mu.Unlock()
// Set serverListReceived to true so fallback will not take effect if it has
// not hit timeout.
lb.serverListReceived = true
// If the new server list == old server list, do nothing.
if serverListEqual(lb.fullServerList, l.Servers) {
if lb.logger.V(2) {
lb.logger.Infof("Ignoring new server list as it is the same as the previous one")
}
return
}
lb.fullServerList = l.Servers
var backendAddrs []resolver.Address
for i, s := range l.Servers {
if s.Drop {
continue
}
md := metadata.Pairs(lbTokenKey, s.LoadBalanceToken)
var ipStr string
if ip, ok := netip.AddrFromSlice(s.IpAddress); ok {
ipStr = ip.String()
} else {
ipStr = fmt.Sprintf("? %x", s.IpAddress)
}
addr := imetadata.Set(resolver.Address{Addr: net.JoinHostPort(ipStr, fmt.Sprintf("%d", s.Port))}, md)
if lb.logger.V(2) {
lb.logger.Infof("Server list entry:|%d|, ipStr:|%s|, port:|%d|, load balancer token:|%v|", i, ipStr, s.Port, s.LoadBalanceToken)
}
backendAddrs = append(backendAddrs, addr)
}
// Call refreshSubConns to create/remove SubConns. If we are in fallback,
// this is also exiting fallback.
lb.refreshSubConns(backendAddrs, false, lb.usePickFirst)
}
// refreshSubConns creates/removes SubConns with backendAddrs, and refreshes
// balancer state and picker.
//
// Caller must hold lb.mu.
func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback bool, pickFirst bool) {
opts := balancer.NewSubConnOptions{}
if !fallback {
opts.CredsBundle = lb.grpclbBackendCreds
}
lb.backendAddrs = backendAddrs
lb.backendAddrsWithoutMetadata = nil
fallbackModeChanged := lb.inFallback != fallback
lb.inFallback = fallback
if fallbackModeChanged && lb.inFallback {
// Clear previous received list when entering fallback, so if the server
// comes back and sends the same list again, the new addresses will be
// used.
lb.fullServerList = nil
}
balancingPolicyChanged := lb.usePickFirst != pickFirst
lb.usePickFirst = pickFirst
if fallbackModeChanged || balancingPolicyChanged {
// Remove all SubConns when switching balancing policy or switching
// fallback mode.
//
// For fallback mode switching with pickfirst, we want to recreate the
// SubConn because the creds could be different.
for a, sc := range lb.subConns {
sc.Shutdown()
delete(lb.subConns, a)
}
}
if lb.usePickFirst {
var (
scKey resolver.Address
sc balancer.SubConn
)
for scKey, sc = range lb.subConns {
break
}
if sc != nil {
if len(backendAddrs) == 0 {
sc.Shutdown()
delete(lb.subConns, scKey)
return
}
lb.cc.ClientConn.UpdateAddresses(sc, backendAddrs)
sc.Connect()
return
}
opts.StateListener = func(scs balancer.SubConnState) { lb.updateSubConnState(sc, scs) }
// This bypasses the cc wrapper with SubConn cache.
sc, err := lb.cc.ClientConn.NewSubConn(backendAddrs, opts)
if err != nil {
lb.logger.Warningf("Failed to create new SubConn: %v", err)
return
}
sc.Connect()
lb.subConns[backendAddrs[0]] = sc
lb.scStates[sc] = connectivity.Idle
return
}
// addrsSet is the set converted from backendAddrsWithoutMetadata, it's used to quick
// lookup for an address.
addrsSet := make(map[resolver.Address]struct{})
// Create new SubConns.
for _, addr := range backendAddrs {
addrWithoutAttrs := addr
addrWithoutAttrs.Attributes = nil
addrsSet[addrWithoutAttrs] = struct{}{}
lb.backendAddrsWithoutMetadata = append(lb.backendAddrsWithoutMetadata, addrWithoutAttrs)
if _, ok := lb.subConns[addrWithoutAttrs]; !ok {
// Use addrWithMD to create the SubConn.
var sc balancer.SubConn
opts.StateListener = func(scs balancer.SubConnState) { lb.updateSubConnState(sc, scs) }
sc, err := lb.cc.NewSubConn([]resolver.Address{addr}, opts)
if err != nil {
lb.logger.Warningf("Failed to create new SubConn: %v", err)
continue
}
lb.subConns[addrWithoutAttrs] = sc // Use the addr without MD as key for the map.
if _, ok := lb.scStates[sc]; !ok {
// Only set state of new sc to IDLE. The state could already be
// READY for cached SubConns.
lb.scStates[sc] = connectivity.Idle
}
sc.Connect()
}
}
for a, sc := range lb.subConns {
// a was removed by resolver.
if _, ok := addrsSet[a]; !ok {
sc.Shutdown()
delete(lb.subConns, a)
// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
// The entry will be deleted in UpdateSubConnState.
}
}
// Regenerate and update picker after refreshing subconns because with
// cache, even if SubConn was newed/removed, there might be no state
// changes (the subconn will be kept in cache, not actually
// newed/removed).
lb.updateStateAndPicker(true, true)
}
type remoteBalancerCCWrapper struct {
cc *grpc.ClientConn
lb *lbBalancer
backoff backoff.Strategy
done chan struct{}
streamMu sync.Mutex
streamCancel func()
// waitgroup to wait for all goroutines to exit.
wg sync.WaitGroup
}
func (lb *lbBalancer) newRemoteBalancerCCWrapper() error {
var dopts []grpc.DialOption
if creds := lb.opt.DialCreds; creds != nil {
dopts = append(dopts, grpc.WithTransportCredentials(creds))
} else if bundle := lb.grpclbClientConnCreds; bundle != nil {
dopts = append(dopts, grpc.WithCredentialsBundle(bundle))
} else {
dopts = append(dopts, grpc.WithTransportCredentials(insecure.NewCredentials()))
}
if lb.opt.Dialer != nil {
dopts = append(dopts, grpc.WithContextDialer(lb.opt.Dialer))
}
if lb.opt.CustomUserAgent != "" {
dopts = append(dopts, grpc.WithUserAgent(lb.opt.CustomUserAgent))
}
// Explicitly set pickfirst as the balancer.
dopts = append(dopts, grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy":"pick_first"}`))
dopts = append(dopts, grpc.WithResolvers(lb.manualResolver))
dopts = append(dopts, grpc.WithChannelzParentID(lb.opt.ChannelzParent))
// Enable Keepalive for grpclb client.
dopts = append(dopts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
Time: 20 * time.Second,
Timeout: 10 * time.Second,
PermitWithoutStream: true,
}))
// The dial target is not important.
//
// The grpclb server addresses will set field ServerName, and creds will
// receive ServerName as authority.
target := lb.manualResolver.Scheme() + ":///grpclb.subClientConn"
cc, err := grpc.NewClient(target, dopts...)
if err != nil {
return fmt.Errorf("grpc.NewClient(%s): %v", target, err)
}
cc.Connect()
ccw := &remoteBalancerCCWrapper{
cc: cc,
lb: lb,
backoff: lb.backoff,
done: make(chan struct{}),
}
lb.ccRemoteLB = ccw
ccw.wg.Add(1)
go ccw.watchRemoteBalancer()
return nil
}
// close closed the ClientConn to remote balancer, and waits until all
// goroutines to finish.
func (ccw *remoteBalancerCCWrapper) close() {
close(ccw.done)
ccw.cc.Close()
ccw.wg.Wait()
}
func (ccw *remoteBalancerCCWrapper) readServerList(s *balanceLoadClientStream) error {
for {
reply, err := s.Recv()
if err != nil {
if err == io.EOF {
return errServerTerminatedConnection
}
return fmt.Errorf("grpclb: failed to recv server list: %v", err)
}
if serverList := reply.GetServerList(); serverList != nil {
ccw.lb.processServerList(serverList)
}
if reply.GetFallbackResponse() != nil {
// Eagerly enter fallback
ccw.lb.mu.Lock()
ccw.lb.refreshSubConns(ccw.lb.resolvedBackendAddrs, true, ccw.lb.usePickFirst)
ccw.lb.mu.Unlock()
}
}
}
func (ccw *remoteBalancerCCWrapper) sendLoadReport(s *balanceLoadClientStream, interval time.Duration) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
lastZero := false
for {
select {
case <-ticker.C:
case <-s.Context().Done():
return
}
stats := ccw.lb.clientStats.toClientStats()
zero := isZeroStats(stats)
if zero && lastZero {
// Quash redundant empty load reports.
continue
}
lastZero = zero
t := time.Now()
stats.Timestamp = &timestamppb.Timestamp{
Seconds: t.Unix(),
Nanos: int32(t.Nanosecond()),
}
if err := s.Send(&lbpb.LoadBalanceRequest{
LoadBalanceRequestType: &lbpb.LoadBalanceRequest_ClientStats{
ClientStats: stats,
},
}); err != nil {
return
}
}
}
func (ccw *remoteBalancerCCWrapper) callRemoteBalancer(ctx context.Context) (backoff bool, _ error) {
lbClient := &loadBalancerClient{cc: ccw.cc}
stream, err := lbClient.BalanceLoad(ctx, grpc.WaitForReady(true))
if err != nil {
return true, fmt.Errorf("grpclb: failed to perform RPC to the remote balancer: %v", err)
}
ccw.lb.mu.Lock()
ccw.lb.remoteBalancerConnected = true
ccw.lb.mu.Unlock()
// grpclb handshake on the stream.
initReq := &lbpb.LoadBalanceRequest{
LoadBalanceRequestType: &lbpb.LoadBalanceRequest_InitialRequest{
InitialRequest: &lbpb.InitialLoadBalanceRequest{
Name: ccw.lb.target,
},
},
}
if err := stream.Send(initReq); err != nil {
return true, fmt.Errorf("grpclb: failed to send init request: %v", err)
}
reply, err := stream.Recv()
if err != nil {
return true, fmt.Errorf("grpclb: failed to recv init response: %v", err)
}
initResp := reply.GetInitialResponse()
if initResp == nil {
return true, fmt.Errorf("grpclb: reply from remote balancer did not include initial response")
}
ccw.wg.Add(1)
go func() {
defer ccw.wg.Done()
if d := convertDuration(initResp.ClientStatsReportInterval); d > 0 {
ccw.sendLoadReport(stream, d)
}
}()
// No backoff if init req/resp handshake was successful.
return false, ccw.readServerList(stream)
}
// cancelRemoteBalancerCall cancels the context used by the stream to the remote
// balancer. watchRemoteBalancer() takes care of restarting this call after the
// stream fails.
func (ccw *remoteBalancerCCWrapper) cancelRemoteBalancerCall() {
ccw.streamMu.Lock()
if ccw.streamCancel != nil {
ccw.streamCancel()
ccw.streamCancel = nil
}
ccw.streamMu.Unlock()
}
func (ccw *remoteBalancerCCWrapper) watchRemoteBalancer() {
defer func() {
ccw.wg.Done()
ccw.streamMu.Lock()
if ccw.streamCancel != nil {
// This is to make sure that we don't leak the context when we are
// directly returning from inside of the below `for` loop.
ccw.streamCancel()
ccw.streamCancel = nil
}
ccw.streamMu.Unlock()
}()
var retryCount int
var ctx context.Context
for {
ccw.streamMu.Lock()
if ccw.streamCancel != nil {
ccw.streamCancel()
ccw.streamCancel = nil
}
ctx, ccw.streamCancel = context.WithCancel(context.Background())
ccw.streamMu.Unlock()
doBackoff, err := ccw.callRemoteBalancer(ctx)
select {
case <-ccw.done:
return
default:
if err != nil {
if err == errServerTerminatedConnection {
ccw.lb.logger.Infof("Call to remote balancer failed: %v", err)
} else {
ccw.lb.logger.Warningf("Call to remote balancer failed: %v", err)
}
}
}
// Trigger a re-resolve when the stream errors.
ccw.lb.cc.ClientConn.ResolveNow(resolver.ResolveNowOptions{})
ccw.lb.mu.Lock()
ccw.lb.remoteBalancerConnected = false
ccw.lb.fullServerList = nil
// Enter fallback when connection to remote balancer is lost, and the
// aggregated state is not Ready.
if !ccw.lb.inFallback && ccw.lb.state != connectivity.Ready {
// Entering fallback.
ccw.lb.refreshSubConns(ccw.lb.resolvedBackendAddrs, true, ccw.lb.usePickFirst)
}
ccw.lb.mu.Unlock()
if !doBackoff {
retryCount = 0
continue
}
timer := time.NewTimer(ccw.backoff.Backoff(retryCount)) // Copy backoff
select {
case <-timer.C:
case <-ccw.done:
timer.Stop()
return
}
retryCount++
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,173 @@
/*
*
* Copyright 2016 gRPC authors.
*
* 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.
*
*/
package grpclb
import (
"fmt"
"sync"
"time"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/resolver"
)
const subConnCacheTime = time.Second * 10
// lbCacheClientConn is a wrapper balancer.ClientConn with a SubConn cache.
// SubConns will be kept in cache for subConnCacheTime before being shut down.
//
// Its NewSubconn and SubConn.Shutdown methods are updated to do cache first.
type lbCacheClientConn struct {
balancer.ClientConn
timeout time.Duration
mu sync.Mutex
// subConnCache only keeps subConns that are being deleted.
subConnCache map[resolver.Address]*subConnCacheEntry
subConnToAddr map[balancer.SubConn]resolver.Address
}
type subConnCacheEntry struct {
sc balancer.SubConn
cancel func()
abortDeleting bool
}
func newLBCacheClientConn(cc balancer.ClientConn) *lbCacheClientConn {
return &lbCacheClientConn{
ClientConn: cc,
timeout: subConnCacheTime,
subConnCache: make(map[resolver.Address]*subConnCacheEntry),
subConnToAddr: make(map[balancer.SubConn]resolver.Address),
}
}
func (ccc *lbCacheClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
if len(addrs) != 1 {
return nil, fmt.Errorf("grpclb calling NewSubConn with addrs of length %v", len(addrs))
}
addrWithoutAttrs := addrs[0]
addrWithoutAttrs.Attributes = nil
ccc.mu.Lock()
defer ccc.mu.Unlock()
if entry, ok := ccc.subConnCache[addrWithoutAttrs]; ok {
// If entry is in subConnCache, the SubConn was being deleted.
// cancel function will never be nil.
entry.cancel()
delete(ccc.subConnCache, addrWithoutAttrs)
return entry.sc, nil
}
scNew, err := ccc.ClientConn.NewSubConn(addrs, opts)
if err != nil {
return nil, err
}
scNew = &lbCacheSubConn{SubConn: scNew, ccc: ccc}
ccc.subConnToAddr[scNew] = addrWithoutAttrs
return scNew, nil
}
func (ccc *lbCacheClientConn) RemoveSubConn(sc balancer.SubConn) {
logger.Errorf("RemoveSubConn(%v) called unexpectedly", sc)
}
type lbCacheSubConn struct {
balancer.SubConn
ccc *lbCacheClientConn
}
func (sc *lbCacheSubConn) Shutdown() {
ccc := sc.ccc
ccc.mu.Lock()
defer ccc.mu.Unlock()
addr, ok := ccc.subConnToAddr[sc]
if !ok {
return
}
if entry, ok := ccc.subConnCache[addr]; ok {
if entry.sc != sc {
// This could happen if NewSubConn was called multiple times for
// the same address, and those SubConns are all shut down. We
// remove sc immediately here.
delete(ccc.subConnToAddr, sc)
sc.SubConn.Shutdown()
}
return
}
entry := &subConnCacheEntry{
sc: sc,
}
ccc.subConnCache[addr] = entry
timer := time.AfterFunc(ccc.timeout, func() {
ccc.mu.Lock()
defer ccc.mu.Unlock()
if entry.abortDeleting {
return
}
sc.SubConn.Shutdown()
delete(ccc.subConnToAddr, sc)
delete(ccc.subConnCache, addr)
})
entry.cancel = func() {
if !timer.Stop() {
// If stop was not successful, the timer has fired (this can only
// happen in a race). But the deleting function is blocked on ccc.mu
// because the mutex was held by the caller of this function.
//
// Set abortDeleting to true to abort the deleting function. When
// the lock is released, the deleting function will acquire the
// lock, check the value of abortDeleting and return.
entry.abortDeleting = true
}
}
}
func (ccc *lbCacheClientConn) UpdateState(s balancer.State) {
s.Picker = &lbCachePicker{Picker: s.Picker}
ccc.ClientConn.UpdateState(s)
}
func (ccc *lbCacheClientConn) close() {
ccc.mu.Lock()
defer ccc.mu.Unlock()
// Only cancel all existing timers. There's no need to shut down SubConns.
for _, entry := range ccc.subConnCache {
entry.cancel()
}
}
type lbCachePicker struct {
balancer.Picker
}
func (cp *lbCachePicker) Pick(i balancer.PickInfo) (balancer.PickResult, error) {
res, err := cp.Picker.Pick(i)
if err != nil {
return res, err
}
res.SubConn = res.SubConn.(*lbCacheSubConn).SubConn
return res, nil
}

View File

@@ -0,0 +1,267 @@
/*
*
* Copyright 2018 gRPC authors.
*
* 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.
*
*/
package grpclb
import (
"fmt"
"sync"
"testing"
"time"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/resolver"
)
type mockSubConn struct {
balancer.SubConn
mcc *mockClientConn
}
func (msc *mockSubConn) Shutdown() {
msc.mcc.mu.Lock()
defer msc.mcc.mu.Unlock()
delete(msc.mcc.subConns, msc)
}
type mockClientConn struct {
balancer.ClientConn
mu sync.Mutex
subConns map[balancer.SubConn]resolver.Address
}
func newMockClientConn() *mockClientConn {
return &mockClientConn{
subConns: make(map[balancer.SubConn]resolver.Address),
}
}
func (mcc *mockClientConn) NewSubConn(addrs []resolver.Address, _ balancer.NewSubConnOptions) (balancer.SubConn, error) {
sc := &mockSubConn{mcc: mcc}
mcc.mu.Lock()
defer mcc.mu.Unlock()
mcc.subConns[sc] = addrs[0]
return sc, nil
}
func (mcc *mockClientConn) RemoveSubConn(sc balancer.SubConn) {
panic(fmt.Sprintf("RemoveSubConn(%v) called unexpectedly", sc))
}
const testCacheTimeout = 100 * time.Millisecond
func checkMockCC(mcc *mockClientConn, scLen int) error {
mcc.mu.Lock()
defer mcc.mu.Unlock()
if len(mcc.subConns) != scLen {
return fmt.Errorf("mcc = %+v, want len(mcc.subConns) = %v", mcc.subConns, scLen)
}
return nil
}
func checkCacheCC(ccc *lbCacheClientConn, sccLen, sctaLen int) error {
ccc.mu.Lock()
defer ccc.mu.Unlock()
if len(ccc.subConnCache) != sccLen {
return fmt.Errorf("ccc = %+v, want len(ccc.subConnCache) = %v", ccc.subConnCache, sccLen)
}
if len(ccc.subConnToAddr) != sctaLen {
return fmt.Errorf("ccc = %+v, want len(ccc.subConnToAddr) = %v", ccc.subConnToAddr, sctaLen)
}
return nil
}
// Test that SubConn won't be immediately shut down.
func (s) TestLBCacheClientConnExpire(t *testing.T) {
mcc := newMockClientConn()
if err := checkMockCC(mcc, 0); err != nil {
t.Fatal(err)
}
ccc := newLBCacheClientConn(mcc)
ccc.timeout = testCacheTimeout
if err := checkCacheCC(ccc, 0, 0); err != nil {
t.Fatal(err)
}
sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
// One subconn in MockCC.
if err := checkMockCC(mcc, 1); err != nil {
t.Fatal(err)
}
// No subconn being deleted, and one in CacheCC.
if err := checkCacheCC(ccc, 0, 1); err != nil {
t.Fatal(err)
}
sc.Shutdown()
// One subconn in MockCC before timeout.
if err := checkMockCC(mcc, 1); err != nil {
t.Fatal(err)
}
// One subconn being deleted, and one in CacheCC.
if err := checkCacheCC(ccc, 1, 1); err != nil {
t.Fatal(err)
}
// Should all become empty after timeout.
var err error
for i := 0; i < 2; i++ {
time.Sleep(testCacheTimeout)
err = checkMockCC(mcc, 0)
if err != nil {
continue
}
err = checkCacheCC(ccc, 0, 0)
if err != nil {
continue
}
}
if err != nil {
t.Fatal(err)
}
}
// Test that NewSubConn with the same address of a SubConn being shut down will
// reuse the SubConn and cancel the removing.
func (s) TestLBCacheClientConnReuse(t *testing.T) {
mcc := newMockClientConn()
if err := checkMockCC(mcc, 0); err != nil {
t.Fatal(err)
}
ccc := newLBCacheClientConn(mcc)
ccc.timeout = testCacheTimeout
if err := checkCacheCC(ccc, 0, 0); err != nil {
t.Fatal(err)
}
sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
// One subconn in MockCC.
if err := checkMockCC(mcc, 1); err != nil {
t.Fatal(err)
}
// No subconn being deleted, and one in CacheCC.
if err := checkCacheCC(ccc, 0, 1); err != nil {
t.Fatal(err)
}
sc.Shutdown()
// One subconn in MockCC before timeout.
if err := checkMockCC(mcc, 1); err != nil {
t.Fatal(err)
}
// One subconn being deleted, and one in CacheCC.
if err := checkCacheCC(ccc, 1, 1); err != nil {
t.Fatal(err)
}
// Recreate the old subconn, this should cancel the deleting process.
sc, _ = ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
// One subconn in MockCC.
if err := checkMockCC(mcc, 1); err != nil {
t.Fatal(err)
}
// No subconn being deleted, and one in CacheCC.
if err := checkCacheCC(ccc, 0, 1); err != nil {
t.Fatal(err)
}
var err error
// Should not become empty after 2*timeout.
time.Sleep(2 * testCacheTimeout)
err = checkMockCC(mcc, 1)
if err != nil {
t.Fatal(err)
}
err = checkCacheCC(ccc, 0, 1)
if err != nil {
t.Fatal(err)
}
// Call Shutdown again, will delete after timeout.
sc.Shutdown()
// One subconn in MockCC before timeout.
if err := checkMockCC(mcc, 1); err != nil {
t.Fatal(err)
}
// One subconn being deleted, and one in CacheCC.
if err := checkCacheCC(ccc, 1, 1); err != nil {
t.Fatal(err)
}
// Should all become empty after timeout.
for i := 0; i < 2; i++ {
time.Sleep(testCacheTimeout)
err = checkMockCC(mcc, 0)
if err != nil {
continue
}
err = checkCacheCC(ccc, 0, 0)
if err != nil {
continue
}
}
if err != nil {
t.Fatal(err)
}
}
// Test that if the timer to shut down a SubConn fires at the same time
// NewSubConn cancels the timer, it doesn't cause deadlock.
func (s) TestLBCache_ShutdownTimer_New_Race(t *testing.T) {
mcc := newMockClientConn()
if err := checkMockCC(mcc, 0); err != nil {
t.Fatal(err)
}
ccc := newLBCacheClientConn(mcc)
ccc.timeout = time.Nanosecond
if err := checkCacheCC(ccc, 0, 0); err != nil {
t.Fatal(err)
}
sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
// One subconn in MockCC.
if err := checkMockCC(mcc, 1); err != nil {
t.Fatal(err)
}
// No subconn being deleted, and one in CacheCC.
if err := checkCacheCC(ccc, 0, 1); err != nil {
t.Fatal(err)
}
done := make(chan struct{})
go func() {
for i := 0; i < 1000; i++ {
// Shutdown starts a timer with 1 ns timeout, the NewSubConn will
// race with the timer.
sc.Shutdown()
sc, _ = ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
}
close(done)
}()
select {
case <-time.After(time.Second):
t.Fatalf("Test didn't finish within 1 second. Deadlock")
case <-done:
}
}

View File

@@ -0,0 +1,51 @@
/*
*
* Copyright 2020 gRPC authors.
*
* 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.
*
*/
// Package state declares grpclb types to be set by resolvers wishing to pass
// information to grpclb via resolver.State Attributes.
package state
import (
"google.golang.org/grpc/resolver"
)
// keyType is the key to use for storing State in Attributes.
type keyType string
const key = keyType("grpc.grpclb.state")
// State contains gRPCLB-relevant data passed from the name resolver.
type State struct {
// BalancerAddresses contains the remote load balancer address(es). If
// set, overrides any resolver-provided addresses with Type of GRPCLB.
BalancerAddresses []resolver.Address
}
// Set returns a copy of the provided state with attributes containing s. s's
// data should not be mutated after calling Set.
func Set(state resolver.State, s *State) resolver.State {
state.Attributes = state.Attributes.WithValue(key, s)
return state
}
// Get returns the grpclb State in the resolver.State, or nil if not present.
// The returned data should not be mutated.
func Get(state resolver.State) *State {
s, _ := state.Attributes.Value(key).(*State)
return s
}

157
balancer/lazy/lazy.go Normal file
View File

@@ -0,0 +1,157 @@
/*
*
* Copyright 2025 gRPC authors.
*
* 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.
*
*/
// Package lazy contains a load balancer that starts in IDLE instead of
// CONNECTING. Once it starts connecting, it instantiates its delegate.
//
// # Experimental
//
// Notice: This package is EXPERIMENTAL and may be changed or removed in a
// later release.
package lazy
import (
"fmt"
"sync"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/resolver"
internalgrpclog "google.golang.org/grpc/internal/grpclog"
)
var (
logger = grpclog.Component("lazy-lb")
)
const (
logPrefix = "[lazy-lb %p] "
)
// ChildBuilderFunc creates a new balancer with the ClientConn. It has the same
// type as the balancer.Builder.Build method.
type ChildBuilderFunc func(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer
// NewBalancer is the constructor for the lazy balancer.
func NewBalancer(cc balancer.ClientConn, bOpts balancer.BuildOptions, childBuilder ChildBuilderFunc) balancer.Balancer {
b := &lazyBalancer{
cc: cc,
buildOptions: bOpts,
childBuilder: childBuilder,
}
b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
cc.UpdateState(balancer.State{
ConnectivityState: connectivity.Idle,
Picker: &idlePicker{exitIdle: sync.OnceFunc(func() {
// Call ExitIdle in a new goroutine to avoid deadlocks while calling
// back into the channel synchronously.
go b.ExitIdle()
})},
})
return b
}
type lazyBalancer struct {
// The following fields are initialized at build time and read-only after
// that and therefore do not need to be guarded by a mutex.
cc balancer.ClientConn
buildOptions balancer.BuildOptions
logger *internalgrpclog.PrefixLogger
childBuilder ChildBuilderFunc
// The following fields are accessed while handling calls to the idlePicker
// and when handling ClientConn state updates. They are guarded by a mutex.
mu sync.Mutex
delegate balancer.Balancer
latestClientConnState *balancer.ClientConnState
latestResolverError error
}
func (lb *lazyBalancer) Close() {
lb.mu.Lock()
defer lb.mu.Unlock()
if lb.delegate != nil {
lb.delegate.Close()
lb.delegate = nil
}
}
func (lb *lazyBalancer) ResolverError(err error) {
lb.mu.Lock()
defer lb.mu.Unlock()
if lb.delegate != nil {
lb.delegate.ResolverError(err)
return
}
lb.latestResolverError = err
}
func (lb *lazyBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
lb.mu.Lock()
defer lb.mu.Unlock()
if lb.delegate != nil {
return lb.delegate.UpdateClientConnState(ccs)
}
lb.latestClientConnState = &ccs
lb.latestResolverError = nil
return nil
}
// UpdateSubConnState implements balancer.Balancer.
func (lb *lazyBalancer) UpdateSubConnState(balancer.SubConn, balancer.SubConnState) {
// UpdateSubConnState is deprecated.
}
func (lb *lazyBalancer) ExitIdle() {
lb.mu.Lock()
defer lb.mu.Unlock()
if lb.delegate != nil {
lb.delegate.ExitIdle()
return
}
lb.delegate = lb.childBuilder(lb.cc, lb.buildOptions)
if lb.latestClientConnState != nil {
if err := lb.delegate.UpdateClientConnState(*lb.latestClientConnState); err != nil {
if err == balancer.ErrBadResolverState {
lb.cc.ResolveNow(resolver.ResolveNowOptions{})
} else {
lb.logger.Warningf("Error from child policy on receiving initial state: %v", err)
}
}
lb.latestClientConnState = nil
}
if lb.latestResolverError != nil {
lb.delegate.ResolverError(lb.latestResolverError)
lb.latestResolverError = nil
}
}
// idlePicker is used when the SubConn is IDLE and kicks the SubConn into
// CONNECTING when Pick is called.
type idlePicker struct {
exitIdle func()
}
func (i *idlePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
i.exitIdle()
return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
}

View File

@@ -0,0 +1,466 @@
/*
*
* Copyright 2025 gRPC authors.
*
* 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.
*
*/
package lazy_test
import (
"context"
"errors"
"fmt"
"strings"
"testing"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/lazy"
"google.golang.org/grpc/balancer/pickfirst"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/internal/balancer/stub"
"google.golang.org/grpc/internal/grpcsync"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/internal/stubserver"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/resolver/manual"
testgrpc "google.golang.org/grpc/interop/grpc_testing"
testpb "google.golang.org/grpc/interop/grpc_testing"
)
const (
// Default timeout for tests in this package.
defaultTestTimeout = 10 * time.Second
// Default short timeout, to be used when waiting for events which are not
// expected to happen.
defaultTestShortTimeout = 100 * time.Millisecond
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
// TestExitIdle creates a lazy balancer than manages a pickfirst child. The test
// calls Connect() on the channel which in turn calls ExitIdle on the lazy
// balancer. The test verifies that the channel enters READY.
func (s) TestExitIdle(t *testing.T) {
backend1 := stubserver.StartTestService(t, nil)
defer backend1.Stop()
mr := manual.NewBuilderWithScheme("e2e-test")
defer mr.Close()
mr.InitialState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend1.Address}}},
},
})
bf := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build)
},
ExitIdle: func(bd *stub.BalancerData) {
bd.ChildBalancer.ExitIdle()
},
ResolverError: func(bd *stub.BalancerData, err error) {
bd.ChildBalancer.ResolverError(err)
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
stub.Register(t.Name(), bf)
json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": {}}]}`, t.Name())
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithDefaultServiceConfig(json),
grpc.WithResolvers(mr),
}
cc, err := grpc.NewClient(mr.Scheme()+":///", opts...)
if err != nil {
t.Fatalf("grpc.NewClient(_) failed: %v", err)
}
defer cc.Close()
cc.Connect()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
testutils.AwaitState(ctx, t, cc, connectivity.Ready)
// Send a resolver update to verify that the resolver state is correctly
// passed through to the leaf pickfirst balancer.
backend2 := stubserver.StartTestService(t, nil)
defer backend2.Stop()
mr.UpdateState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend2.Address}}},
},
})
var peer peer.Peer
client := testgrpc.NewTestServiceClient(cc)
if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
}
if got, want := peer.Addr.String(), backend2.Address; got != want {
t.Errorf("EmptyCall() went to unexpected backend: got %q, want %q", got, want)
}
}
// TestPicker creates a lazy balancer under a stub balancer which block all
// calls to ExitIdle. This ensures the only way to trigger lazy to exit idle is
// through the picker. The test makes an RPC and ensures it succeeds.
func (s) TestPicker(t *testing.T) {
backend := stubserver.StartTestService(t, nil)
defer backend.Stop()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
bf := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build)
},
ExitIdle: func(*stub.BalancerData) {
t.Log("Ignoring call to ExitIdle, calling the picker should make the lazy balancer exit IDLE state.")
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
stub.Register(name, bf)
json := fmt.Sprintf(`{"loadBalancingConfig": [{%q: {}}]}`, name)
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithDefaultServiceConfig(json),
}
cc, err := grpc.NewClient(backend.Address, opts...)
if err != nil {
t.Fatalf("grpc.NewClient(_) failed: %v", err)
}
defer cc.Close()
// The channel should remain in IDLE as the ExitIdle calls are not
// propagated to the lazy balancer from the stub balancer.
cc.Connect()
shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
defer shortCancel()
testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
// The picker from the lazy balancer should be send to the channel when the
// first resolver update is received by lazy. Making an RPC should trigger
// child creation.
client := testgrpc.NewTestServiceClient(cc)
if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
}
}
// Tests the scenario when a resolver produces a good state followed by a
// resolver error. The test verifies that the child balancer receives the good
// update followed by the error.
func (s) TestGoodUpdateThenResolverError(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
backend := stubserver.StartTestService(t, nil)
defer backend.Stop()
resolverStateReceived := false
resolverErrorReceived := grpcsync.NewEvent()
childBF := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
if resolverErrorReceived.HasFired() {
t.Error("Received resolver error before resolver state.")
}
resolverStateReceived = true
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
ResolverError: func(bd *stub.BalancerData, err error) {
if !resolverStateReceived {
t.Error("Received resolver error before resolver state.")
}
resolverErrorReceived.Fire()
bd.ChildBalancer.ResolverError(err)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
childBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_child", "/", "")
stub.Register(childBalName, childBF)
topLevelBF := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
},
ExitIdle: func(*stub.BalancerData) {
t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
},
ResolverError: func(bd *stub.BalancerData, err error) {
bd.ChildBalancer.ResolverError(err)
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
topLevelBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_top_level", "/", "")
stub.Register(topLevelBalName, topLevelBF)
json := fmt.Sprintf(`{"loadBalancingConfig": [{%q: {}}]}`, topLevelBalName)
mr := manual.NewBuilderWithScheme("e2e-test")
defer mr.Close()
mr.InitialState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend.Address}}},
},
})
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithResolvers(mr),
grpc.WithDefaultServiceConfig(json),
}
cc, err := grpc.NewClient(mr.Scheme()+":///whatever", opts...)
if err != nil {
t.Fatalf("grpc.NewClient(_) failed: %v", err)
}
defer cc.Close()
cc.Connect()
mr.CC().ReportError(errors.New("test error"))
// The channel should remain in IDLE as the ExitIdle calls are not
// propagated to the lazy balancer from the stub balancer.
shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
defer shortCancel()
testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
client := testgrpc.NewTestServiceClient(cc)
if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
}
if !resolverStateReceived {
t.Fatalf("Child balancer did not receive resolver state.")
}
select {
case <-resolverErrorReceived.Done():
case <-ctx.Done():
t.Fatal("Context timed out waiting for resolver error to be delivered to child balancer.")
}
}
// Tests the scenario when a resolver produces a list of endpoints followed by
// a resolver error. The test verifies that the child balancer receives only the
// good update.
func (s) TestResolverErrorThenGoodUpdate(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
backend := stubserver.StartTestService(t, nil)
defer backend.Stop()
childBF := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
ResolverError: func(bd *stub.BalancerData, err error) {
t.Error("Received unexpected resolver error.")
bd.ChildBalancer.ResolverError(err)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
childBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_child", "/", "")
stub.Register(childBalName, childBF)
topLevelBF := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
},
ExitIdle: func(*stub.BalancerData) {
t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
topLevelBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_top_level", "/", "")
stub.Register(topLevelBalName, topLevelBF)
json := fmt.Sprintf(`{"loadBalancingConfig": [{%q: {}}]}`, topLevelBalName)
mr := manual.NewBuilderWithScheme("e2e-test")
defer mr.Close()
mr.InitialState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend.Address}}},
},
})
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithResolvers(mr),
grpc.WithDefaultServiceConfig(json),
}
cc, err := grpc.NewClient(mr.Scheme()+":///whatever", opts...)
if err != nil {
t.Fatalf("grpc.NewClient(_) failed: %v", err)
}
defer cc.Close()
cc.Connect()
// Send an error followed by a good update.
mr.CC().ReportError(errors.New("test error"))
mr.UpdateState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend.Address}}},
},
})
// The channel should remain in IDLE as the ExitIdle calls are not
// propagated to the lazy balancer from the stub balancer.
shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
defer shortCancel()
testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
// An RPC would succeed only if the leaf pickfirst receives the endpoint
// list.
client := testgrpc.NewTestServiceClient(cc)
if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
t.Errorf("client.EmptyCall() returned unexpected error: %v", err)
}
}
// Tests that ExitIdle calls are correctly passed through to the child balancer.
// It starts a backend and ensures the channel connects to it. The test then
// stops the backend, making the channel enter IDLE. The test calls Connect on
// the channel and verifies that the child balancer exits idle.
func (s) TestExitIdlePassthrough(t *testing.T) {
backend1 := stubserver.StartTestService(t, nil)
defer backend1.Stop()
mr := manual.NewBuilderWithScheme("e2e-test")
defer mr.Close()
mr.InitialState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend1.Address}}},
},
})
bf := stub.BalancerFuncs{
Init: func(bd *stub.BalancerData) {
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirst.Name).Build)
},
ExitIdle: func(bd *stub.BalancerData) {
bd.ChildBalancer.ExitIdle()
},
ResolverError: func(bd *stub.BalancerData, err error) {
bd.ChildBalancer.ResolverError(err)
},
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
return bd.ChildBalancer.UpdateClientConnState(ccs)
},
Close: func(bd *stub.BalancerData) {
bd.ChildBalancer.Close()
},
}
stub.Register(t.Name(), bf)
json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": {}}]}`, t.Name())
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithDefaultServiceConfig(json),
grpc.WithResolvers(mr),
}
cc, err := grpc.NewClient(mr.Scheme()+":///", opts...)
if err != nil {
t.Fatalf("grpc.NewClient(_) failed: %v", err)
}
defer cc.Close()
cc.Connect()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
testutils.AwaitState(ctx, t, cc, connectivity.Ready)
// Stopping the active backend should put the channel in IDLE.
backend1.Stop()
testutils.AwaitState(ctx, t, cc, connectivity.Idle)
// Sending a new backend address should not kick the channel out of IDLE.
// On calling cc.Connect(), the channel should call ExitIdle on the lazy
// balancer which passes through the call to the leaf pickfirst.
backend2 := stubserver.StartTestService(t, nil)
defer backend2.Stop()
mr.UpdateState(resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backend2.Address}}},
},
})
shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
defer shortCancel()
testutils.AwaitNoStateChange(shortCtx, t, cc, connectivity.Idle)
cc.Connect()
testutils.AwaitState(ctx, t, cc, connectivity.Ready)
}

View File

@@ -0,0 +1,250 @@
/*
*
* Copyright 2023 gRPC authors.
*
* 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.
*
*/
// Package leastrequest implements a least request load balancer.
package leastrequest
import (
"encoding/json"
"fmt"
rand "math/rand/v2"
"sync"
"sync/atomic"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/endpointsharding"
"google.golang.org/grpc/balancer/pickfirst"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/grpclog"
internalgrpclog "google.golang.org/grpc/internal/grpclog"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/serviceconfig"
)
// Name is the name of the least request balancer.
const Name = "least_request_experimental"
var (
// randuint32 is a global to stub out in tests.
randuint32 = rand.Uint32
logger = grpclog.Component("least-request")
)
func init() {
balancer.Register(bb{})
}
// LBConfig is the balancer config for least_request_experimental balancer.
type LBConfig struct {
serviceconfig.LoadBalancingConfig `json:"-"`
// ChoiceCount is the number of random SubConns to sample to find the one
// with the fewest outstanding requests. If unset, defaults to 2. If set to
// < 2, the config will be rejected, and if set to > 10, will become 10.
ChoiceCount uint32 `json:"choiceCount,omitempty"`
}
type bb struct{}
func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
lbConfig := &LBConfig{
ChoiceCount: 2,
}
if err := json.Unmarshal(s, lbConfig); err != nil {
return nil, fmt.Errorf("least-request: unable to unmarshal LBConfig: %v", err)
}
// "If `choice_count < 2`, the config will be rejected." - A48
if lbConfig.ChoiceCount < 2 { // sweet
return nil, fmt.Errorf("least-request: lbConfig.choiceCount: %v, must be >= 2", lbConfig.ChoiceCount)
}
// "If a LeastRequestLoadBalancingConfig with a choice_count > 10 is
// received, the least_request_experimental policy will set choice_count =
// 10." - A48
if lbConfig.ChoiceCount > 10 {
lbConfig.ChoiceCount = 10
}
return lbConfig, nil
}
func (bb) Name() string {
return Name
}
func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
b := &leastRequestBalancer{
ClientConn: cc,
endpointRPCCounts: resolver.NewEndpointMap[*atomic.Int32](),
}
b.child = endpointsharding.NewBalancer(b, bOpts, balancer.Get(pickfirst.Name).Build, endpointsharding.Options{})
b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[%p] ", b))
b.logger.Infof("Created")
return b
}
type leastRequestBalancer struct {
// Embeds balancer.ClientConn because we need to intercept UpdateState
// calls from the child balancer.
balancer.ClientConn
child balancer.Balancer
logger *internalgrpclog.PrefixLogger
mu sync.Mutex
choiceCount uint32
// endpointRPCCounts holds RPC counts to keep track for subsequent picker
// updates.
endpointRPCCounts *resolver.EndpointMap[*atomic.Int32]
}
func (lrb *leastRequestBalancer) Close() {
lrb.child.Close()
lrb.endpointRPCCounts = nil
}
func (lrb *leastRequestBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
lrb.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
}
func (lrb *leastRequestBalancer) ResolverError(err error) {
// Will cause inline picker update from endpoint sharding.
lrb.child.ResolverError(err)
}
func (lrb *leastRequestBalancer) ExitIdle() {
lrb.child.ExitIdle()
}
func (lrb *leastRequestBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
lrCfg, ok := ccs.BalancerConfig.(*LBConfig)
if !ok {
logger.Errorf("least-request: received config with unexpected type %T: %v", ccs.BalancerConfig, ccs.BalancerConfig)
return balancer.ErrBadResolverState
}
lrb.mu.Lock()
lrb.choiceCount = lrCfg.ChoiceCount
lrb.mu.Unlock()
return lrb.child.UpdateClientConnState(balancer.ClientConnState{
// Enable the health listener in pickfirst children for client side health
// checks and outlier detection, if configured.
ResolverState: pickfirst.EnableHealthListener(ccs.ResolverState),
})
}
type endpointState struct {
picker balancer.Picker
numRPCs *atomic.Int32
}
func (lrb *leastRequestBalancer) UpdateState(state balancer.State) {
var readyEndpoints []endpointsharding.ChildState
for _, child := range endpointsharding.ChildStatesFromPicker(state.Picker) {
if child.State.ConnectivityState == connectivity.Ready {
readyEndpoints = append(readyEndpoints, child)
}
}
// If no ready pickers are present, simply defer to the round robin picker
// from endpoint sharding, which will round robin across the most relevant
// pick first children in the highest precedence connectivity state.
if len(readyEndpoints) == 0 {
lrb.ClientConn.UpdateState(state)
return
}
lrb.mu.Lock()
defer lrb.mu.Unlock()
if logger.V(2) {
lrb.logger.Infof("UpdateState called with ready endpoints: %v", readyEndpoints)
}
// Reconcile endpoints.
newEndpoints := resolver.NewEndpointMap[any]()
for _, child := range readyEndpoints {
newEndpoints.Set(child.Endpoint, nil)
}
// If endpoints are no longer ready, no need to count their active RPCs.
for endpoint := range lrb.endpointRPCCounts.All() {
if _, ok := newEndpoints.Get(endpoint); !ok {
lrb.endpointRPCCounts.Delete(endpoint)
}
}
// Copy refs to counters into picker.
endpointStates := make([]endpointState, 0, len(readyEndpoints))
for _, child := range readyEndpoints {
counter, ok := lrb.endpointRPCCounts.Get(child.Endpoint)
if !ok {
// Create new counts if needed.
counter = new(atomic.Int32)
lrb.endpointRPCCounts.Set(child.Endpoint, counter)
}
endpointStates = append(endpointStates, endpointState{
picker: child.State.Picker,
numRPCs: counter,
})
}
lrb.ClientConn.UpdateState(balancer.State{
Picker: &picker{
choiceCount: lrb.choiceCount,
endpointStates: endpointStates,
},
ConnectivityState: connectivity.Ready,
})
}
type picker struct {
// choiceCount is the number of random endpoints to sample for choosing the
// one with the least requests.
choiceCount uint32
endpointStates []endpointState
}
func (p *picker) Pick(pInfo balancer.PickInfo) (balancer.PickResult, error) {
var pickedEndpointState *endpointState
var pickedEndpointNumRPCs int32
for i := 0; i < int(p.choiceCount); i++ {
index := randuint32() % uint32(len(p.endpointStates))
endpointState := p.endpointStates[index]
n := endpointState.numRPCs.Load()
if pickedEndpointState == nil || n < pickedEndpointNumRPCs {
pickedEndpointState = &endpointState
pickedEndpointNumRPCs = n
}
}
result, err := pickedEndpointState.picker.Pick(pInfo)
if err != nil {
return result, err
}
// "The counter for a subchannel should be atomically incremented by one
// after it has been successfully picked by the picker." - A48
pickedEndpointState.numRPCs.Add(1)
// "the picker should add a callback for atomically decrementing the
// subchannel counter once the RPC finishes (regardless of Status code)." -
// A48.
originalDone := result.Done
result.Done = func(info balancer.DoneInfo) {
pickedEndpointState.numRPCs.Add(-1)
if originalDone != nil {
originalDone(info)
}
}
return result, nil
}

View File

@@ -0,0 +1,770 @@
/*
*
* Copyright 2023 gRPC authors.
*
* 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.
*/
package leastrequest
import (
"context"
"encoding/json"
"fmt"
"strings"
"sync"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/internal"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/internal/stubserver"
"google.golang.org/grpc/internal/testutils"
testgrpc "google.golang.org/grpc/interop/grpc_testing"
testpb "google.golang.org/grpc/interop/grpc_testing"
"google.golang.org/grpc/peer"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/resolver/manual"
"google.golang.org/grpc/serviceconfig"
)
const (
defaultTestTimeout = 5 * time.Second
defaultTestShortTimeout = 10 * time.Millisecond
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
func (s) TestParseConfig(t *testing.T) {
parser := bb{}
tests := []struct {
name string
input string
wantCfg serviceconfig.LoadBalancingConfig
wantErr string
}{
{
name: "happy-case-default",
input: `{}`,
wantCfg: &LBConfig{
ChoiceCount: 2,
},
},
{
name: "happy-case-choice-count-set",
input: `{"choiceCount": 3}`,
wantCfg: &LBConfig{
ChoiceCount: 3,
},
},
{
name: "happy-case-choice-count-greater-than-ten",
input: `{"choiceCount": 11}`,
wantCfg: &LBConfig{
ChoiceCount: 10,
},
},
{
name: "choice-count-less-than-2",
input: `{"choiceCount": 1}`,
wantErr: "must be >= 2",
},
{
name: "invalid-json",
input: "{{invalidjson{{",
wantErr: "invalid character",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
gotCfg, gotErr := parser.ParseConfig(json.RawMessage(test.input))
// Substring match makes this very tightly coupled to the
// internalserviceconfig.BalancerConfig error strings. However, it
// is important to distinguish the different types of error messages
// possible as the parser has a few defined buckets of ways it can
// error out.
if (gotErr != nil) != (test.wantErr != "") {
t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
}
if gotErr != nil && !strings.Contains(gotErr.Error(), test.wantErr) {
t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
}
if test.wantErr != "" {
return
}
if diff := cmp.Diff(gotCfg, test.wantCfg); diff != "" {
t.Fatalf("ParseConfig(%v) got unexpected output, diff (-got +want): %v", test.input, diff)
}
})
}
}
func startBackends(t *testing.T, numBackends int) []*stubserver.StubServer {
backends := make([]*stubserver.StubServer, 0, numBackends)
// Construct and start working backends.
for i := 0; i < numBackends; i++ {
backend := &stubserver.StubServer{
EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
return &testpb.Empty{}, nil
},
FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
<-stream.Context().Done()
return nil
},
}
if err := backend.StartServer(); err != nil {
t.Fatalf("Failed to start backend: %v", err)
}
t.Logf("Started good TestService backend at: %q", backend.Address)
t.Cleanup(func() { backend.Stop() })
backends = append(backends, backend)
}
return backends
}
// setupBackends spins up three test backends, each listening on a port on
// localhost. The three backends always reply with an empty response with no
// error, and for streaming receive until hitting an EOF error.
func setupBackends(t *testing.T, numBackends int) []string {
t.Helper()
addresses := make([]string, numBackends)
backends := startBackends(t, numBackends)
// Construct and start working backends.
for i := 0; i < numBackends; i++ {
addresses[i] = backends[i].Address
}
return addresses
}
// checkRoundRobinRPCs verifies that EmptyCall RPCs on the given ClientConn,
// connected to a server exposing the test.grpc_testing.TestService, are
// roundrobined across the given backend addresses.
//
// Returns a non-nil error if context deadline expires before RPCs start to get
// roundrobined across the given backends.
func checkRoundRobinRPCs(ctx context.Context, client testgrpc.TestServiceClient, addrs []resolver.Address) error {
wantAddrCount := make(map[string]int)
for _, addr := range addrs {
wantAddrCount[addr.Addr]++
}
gotAddrCount := make(map[string]int)
for ; ctx.Err() == nil; <-time.After(time.Millisecond) {
gotAddrCount = make(map[string]int)
// Perform 3 iterations.
var iterations [][]string
for i := 0; i < 3; i++ {
iteration := make([]string, len(addrs))
for c := 0; c < len(addrs); c++ {
var peer peer.Peer
client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer))
iteration[c] = peer.Addr.String()
}
iterations = append(iterations, iteration)
}
// Ensure the first iteration contains all addresses in addrs.
for _, addr := range iterations[0] {
gotAddrCount[addr]++
}
if !cmp.Equal(gotAddrCount, wantAddrCount) {
continue
}
// Ensure all three iterations contain the same addresses.
if !cmp.Equal(iterations[0], iterations[1]) || !cmp.Equal(iterations[0], iterations[2]) {
continue
}
return nil
}
return fmt.Errorf("timeout when waiting for roundrobin distribution of RPCs across addresses: %v; got: %v", addrs, gotAddrCount)
}
// TestLeastRequestE2E tests the Least Request LB policy in an e2e style. The
// Least Request balancer is configured as the top level balancer of the
// channel, and is passed three addresses. Eventually, the test creates three
// streams, which should be on certain backends according to the least request
// algorithm. The randomness in the picker is injected in the test to be
// deterministic, allowing the test to make assertions on the distribution.
func (s) TestLeastRequestE2E(t *testing.T) {
defer func(u func() uint32) {
randuint32 = u
}(randuint32)
var index int
indexes := []uint32{
0, 0, 1, 1, 2, 2, // Triggers a round robin distribution.
}
randuint32 = func() uint32 {
ret := indexes[index%len(indexes)]
index++
return ret
}
addresses := setupBackends(t, 3)
mr := manual.NewBuilderWithScheme("lr-e2e")
defer mr.Close()
// Configure least request as top level balancer of channel.
lrscJSON := `
{
"loadBalancingConfig": [
{
"least_request_experimental": {
"choiceCount": 2
}
}
]
}`
sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
firstThreeAddresses := []resolver.Address{
{Addr: addresses[0]},
{Addr: addresses[1]},
{Addr: addresses[2]},
}
mr.InitialState(resolver.State{
Addresses: firstThreeAddresses,
ServiceConfig: sc,
})
cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient() failed: %v", err)
}
defer cc.Close()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
testServiceClient := testgrpc.NewTestServiceClient(cc)
// Wait for all 3 backends to round robin across. The happens because a
// SubConn transitioning into READY causes a new picker update. Once the
// picker update with all 3 backends is present, this test can start to make
// assertions based on those backends.
if err := checkRoundRobinRPCs(ctx, testServiceClient, firstThreeAddresses); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Map ordering of READY SubConns is non deterministic. Thus, perform 3 RPCs
// mocked from the random to each index to learn the addresses of SubConns
// at each index.
index = 0
peerAtIndex := make([]string, 3)
var peer0 peer.Peer
if _, err := testServiceClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer0)); err != nil {
t.Fatalf("testServiceClient.EmptyCall failed: %v", err)
}
peerAtIndex[0] = peer0.Addr.String()
if _, err := testServiceClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer0)); err != nil {
t.Fatalf("testServiceClient.EmptyCall failed: %v", err)
}
peerAtIndex[1] = peer0.Addr.String()
if _, err := testServiceClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer0)); err != nil {
t.Fatalf("testServiceClient.EmptyCall failed: %v", err)
}
peerAtIndex[2] = peer0.Addr.String()
// Start streaming RPCs, but do not finish them. Each subsequent stream
// should be started according to the least request algorithm, and chosen
// between the indexes provided.
index = 0
indexes = []uint32{
0, 0, // Causes first stream to be on first address.
0, 1, // Compares first address (one RPC) to second (no RPCs), so choose second.
1, 2, // Compares second address (one RPC) to third (no RPCs), so choose third.
0, 3, // Causes another stream on first address.
1, 0, // Compares second address (one RPC) to first (two RPCs), so choose second.
2, 0, // Compares third address (one RPC) to first (two RPCs), so choose third.
0, 0, // Causes another stream on first address.
2, 2, // Causes a stream on third address.
2, 1, // Compares third address (three RPCs) to second (two RPCs), so choose third.
}
wantIndex := []uint32{0, 1, 2, 0, 1, 2, 0, 2, 1}
// Start streaming RPC's, but do not finish them. Each created stream should
// be started based on the least request algorithm and injected randomness
// (see indexes slice above for exact expectations).
for _, wantIndex := range wantIndex {
stream, err := testServiceClient.FullDuplexCall(ctx)
if err != nil {
t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
}
p, ok := peer.FromContext(stream.Context())
if !ok {
t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
}
if p.Addr.String() != peerAtIndex[wantIndex] {
t.Fatalf("testServiceClient.FullDuplexCall's Peer got: %v, want: %v", p.Addr.String(), peerAtIndex[wantIndex])
}
}
}
// TestLeastRequestPersistsCounts tests that the Least Request Balancer persists
// counts once it gets a new picker update. It first updates the Least Request
// Balancer with two backends, and creates a bunch of streams on them. Then, it
// updates the Least Request Balancer with three backends, including the two
// previous. Any created streams should then be started on the new backend.
func (s) TestLeastRequestPersistsCounts(t *testing.T) {
defer func(u func() uint32) {
randuint32 = u
}(randuint32)
var index int
indexes := []uint32{
0, 0, 1, 1,
}
randuint32 = func() uint32 {
ret := indexes[index%len(indexes)]
index++
return ret
}
addresses := setupBackends(t, 3)
mr := manual.NewBuilderWithScheme("lr-e2e")
defer mr.Close()
// Configure least request as top level balancer of channel.
lrscJSON := `
{
"loadBalancingConfig": [
{
"least_request_experimental": {
"choiceCount": 2
}
}
]
}`
sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
firstTwoAddresses := []resolver.Address{
{Addr: addresses[0]},
{Addr: addresses[1]},
}
mr.InitialState(resolver.State{
Addresses: firstTwoAddresses,
ServiceConfig: sc,
})
cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient() failed: %v", err)
}
defer cc.Close()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
testServiceClient := testgrpc.NewTestServiceClient(cc)
// Wait for the two backends to round robin across. The happens because a
// SubConn transitioning into READY causes a new picker update. Once the
// picker update with the two backends is present, this test can start to
// populate those backends with streams.
if err := checkRoundRobinRPCs(ctx, testServiceClient, firstTwoAddresses); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Start 50 streaming RPCs, and leave them unfinished for the duration of
// the test. This will populate the first two addresses with many active
// RPCs.
for i := 0; i < 50; i++ {
_, err := testServiceClient.FullDuplexCall(ctx)
if err != nil {
t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
}
}
// Update the least request balancer to choice count 3. Also update the
// address list adding a third address. Alongside the injected randomness,
// this should trigger the least request balancer to search all created
// SubConns. Thus, since address 3 is the new address and the first two
// addresses are populated with RPCs, once the picker update of all 3 READY
// SubConns takes effect, all new streams should be started on address 3.
index = 0
indexes = []uint32{
0, 1, 2, 3, 4, 5,
}
lrscJSON = `
{
"loadBalancingConfig": [
{
"least_request_experimental": {
"choiceCount": 3
}
}
]
}`
sc = internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
fullAddresses := []resolver.Address{
{Addr: addresses[0]},
{Addr: addresses[1]},
{Addr: addresses[2]},
}
mr.UpdateState(resolver.State{
Addresses: fullAddresses,
ServiceConfig: sc,
})
newAddress := fullAddresses[2]
// Poll for only address 3 to show up. This requires a polling loop because
// picker update with all three SubConns doesn't take into effect
// immediately, needs the third SubConn to become READY.
if err := checkRoundRobinRPCs(ctx, testServiceClient, []resolver.Address{newAddress}); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Start 25 rpcs, but don't finish them. They should all start on address 3,
// since the first two addresses both have 25 RPCs (and randomness
// injection/choiceCount causes all 3 to be compared every iteration).
for i := 0; i < 25; i++ {
stream, err := testServiceClient.FullDuplexCall(ctx)
if err != nil {
t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
}
p, ok := peer.FromContext(stream.Context())
if !ok {
t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
}
if p.Addr.String() != addresses[2] {
t.Fatalf("testServiceClient.FullDuplexCall's Peer got: %v, want: %v", p.Addr.String(), addresses[2])
}
}
// Now 25 RPC's are active on each address, the next three RPC's should
// round robin, since choiceCount is three and the injected random indexes
// cause it to search all three addresses for fewest outstanding requests on
// each iteration.
wantAddrCount := map[string]int{
addresses[0]: 1,
addresses[1]: 1,
addresses[2]: 1,
}
gotAddrCount := make(map[string]int)
for i := 0; i < len(addresses); i++ {
stream, err := testServiceClient.FullDuplexCall(ctx)
if err != nil {
t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
}
p, ok := peer.FromContext(stream.Context())
if !ok {
t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
}
if p.Addr != nil {
gotAddrCount[p.Addr.String()]++
}
}
if diff := cmp.Diff(gotAddrCount, wantAddrCount); diff != "" {
t.Fatalf("addr count (-got:, +want): %v", diff)
}
}
// TestConcurrentRPCs tests concurrent RPCs on the least request balancer. It
// configures a channel with a least request balancer as the top level balancer,
// and makes 100 RPCs asynchronously. This makes sure no race conditions happen
// in this scenario.
func (s) TestConcurrentRPCs(t *testing.T) {
addresses := setupBackends(t, 3)
mr := manual.NewBuilderWithScheme("lr-e2e")
defer mr.Close()
// Configure least request as top level balancer of channel.
lrscJSON := `
{
"loadBalancingConfig": [
{
"least_request_experimental": {
"choiceCount": 2
}
}
]
}`
sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
firstTwoAddresses := []resolver.Address{
{Addr: addresses[0]},
{Addr: addresses[1]},
}
mr.InitialState(resolver.State{
Addresses: firstTwoAddresses,
ServiceConfig: sc,
})
cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient() failed: %v", err)
}
defer cc.Close()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
testServiceClient := testgrpc.NewTestServiceClient(cc)
var wg sync.WaitGroup
for i := 0; i < 100; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 5; j++ {
testServiceClient.EmptyCall(ctx, &testpb.Empty{})
}
}()
}
wg.Wait()
}
// Test tests that the least request balancer persists RPC counts once it gets
// new picker updates and backends within an endpoint go down. It first updates
// the balancer with two endpoints having two addresses each. It verifies the
// requests are round robined across the first address of each endpoint. It then
// stops the active backend in endpoint[0]. It verified that the balancer starts
// using the second address in endpoint[0]. The test then creates a bunch of
// streams on two endpoints. Then, it updates the balancer with three endpoints,
// including the two previous. Any created streams should then be started on the
// new endpoint. The test shuts down the active backed in endpoint[1] and
// endpoint[2]. The test verifies that new RPCs are round robined across the
// active backends in endpoint[1] and endpoint[2].
func (s) TestLeastRequestEndpoints_MultipleAddresses(t *testing.T) {
defer func(u func() uint32) {
randuint32 = u
}(randuint32)
var index int
indexes := []uint32{
0, 0, 1, 1,
}
randuint32 = func() uint32 {
ret := indexes[index%len(indexes)]
index++
return ret
}
backends := startBackends(t, 6)
mr := manual.NewBuilderWithScheme("lr-e2e")
defer mr.Close()
// Configure least request as top level balancer of channel.
lrscJSON := `
{
"loadBalancingConfig": [
{
"least_request_experimental": {
"choiceCount": 2
}
}
]
}`
endpoints := []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: backends[0].Address}, {Addr: backends[1].Address}}},
{Addresses: []resolver.Address{{Addr: backends[2].Address}, {Addr: backends[3].Address}}},
{Addresses: []resolver.Address{{Addr: backends[4].Address}, {Addr: backends[5].Address}}},
}
sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
firstTwoEndpoints := []resolver.Endpoint{endpoints[0], endpoints[1]}
mr.InitialState(resolver.State{
Endpoints: firstTwoEndpoints,
ServiceConfig: sc,
})
cc, err := grpc.NewClient(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("grpc.NewClient() failed: %v", err)
}
defer cc.Close()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
testServiceClient := testgrpc.NewTestServiceClient(cc)
// Wait for the two backends to round robin across. The happens because a
// child pickfirst transitioning into READY causes a new picker update. Once
// the picker update with the two backends is present, this test can start
// to populate those backends with streams.
wantAddrs := []resolver.Address{
endpoints[0].Addresses[0],
endpoints[1].Addresses[0],
}
if err := checkRoundRobinRPCs(ctx, testServiceClient, wantAddrs); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Shut down one of the addresses in endpoints[0], the child pickfirst
// should fallback to the next address in endpoints[0].
backends[0].Stop()
wantAddrs = []resolver.Address{
endpoints[0].Addresses[1],
endpoints[1].Addresses[0],
}
if err := checkRoundRobinRPCs(ctx, testServiceClient, wantAddrs); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Start 50 streaming RPCs, and leave them unfinished for the duration of
// the test. This will populate the first two endpoints with many active
// RPCs.
for i := 0; i < 50; i++ {
_, err := testServiceClient.FullDuplexCall(ctx)
if err != nil {
t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
}
}
// Update the least request balancer to choice count 3. Also update the
// address list adding a third endpoint. Alongside the injected randomness,
// this should trigger the least request balancer to search all created
// endpoints. Thus, since endpoint 3 is the new endpoint and the first two
// endpoint are populated with RPCs, once the picker update of all 3 READY
// pickfirsts takes effect, all new streams should be started on endpoint 3.
index = 0
indexes = []uint32{
0, 1, 2, 3, 4, 5,
}
lrscJSON = `
{
"loadBalancingConfig": [
{
"least_request_experimental": {
"choiceCount": 3
}
}
]
}`
sc = internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(lrscJSON)
mr.UpdateState(resolver.State{
Endpoints: endpoints,
ServiceConfig: sc,
})
newAddress := endpoints[2].Addresses[0]
// Poll for only endpoint 3 to show up. This requires a polling loop because
// picker update with all three endpoints doesn't take into effect
// immediately, needs the third pickfirst to become READY.
if err := checkRoundRobinRPCs(ctx, testServiceClient, []resolver.Address{newAddress}); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
// Start 25 rpcs, but don't finish them. They should all start on endpoint 3,
// since the first two endpoints both have 25 RPCs (and randomness
// injection/choiceCount causes all 3 to be compared every iteration).
for i := 0; i < 25; i++ {
stream, err := testServiceClient.FullDuplexCall(ctx)
if err != nil {
t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
}
p, ok := peer.FromContext(stream.Context())
if !ok {
t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
}
if p.Addr.String() != newAddress.Addr {
t.Fatalf("testServiceClient.FullDuplexCall's Peer got: %v, want: %v", p.Addr.String(), newAddress)
}
}
// Now 25 RPC's are active on each endpoint, the next three RPC's should
// round robin, since choiceCount is three and the injected random indexes
// cause it to search all three endpoints for fewest outstanding requests on
// each iteration.
wantAddrCount := map[string]int{
endpoints[0].Addresses[1].Addr: 1,
endpoints[1].Addresses[0].Addr: 1,
endpoints[2].Addresses[0].Addr: 1,
}
gotAddrCount := make(map[string]int)
for i := 0; i < len(endpoints); i++ {
stream, err := testServiceClient.FullDuplexCall(ctx)
if err != nil {
t.Fatalf("testServiceClient.FullDuplexCall failed: %v", err)
}
p, ok := peer.FromContext(stream.Context())
if !ok {
t.Fatalf("testServiceClient.FullDuplexCall has no Peer")
}
if p.Addr != nil {
gotAddrCount[p.Addr.String()]++
}
}
if diff := cmp.Diff(gotAddrCount, wantAddrCount); diff != "" {
t.Fatalf("addr count (-got:, +want): %v", diff)
}
// Shutdown the active address for endpoint[1] and endpoint[2]. This should
// result in their streams failing. Now the requests should roundrobin b/w
// endpoint[1] and endpoint[2].
backends[2].Stop()
backends[4].Stop()
index = 0
indexes = []uint32{
0, 1, 2, 2, 1, 0,
}
wantAddrs = []resolver.Address{
endpoints[1].Addresses[1],
endpoints[2].Addresses[1],
}
if err := checkRoundRobinRPCs(ctx, testServiceClient, wantAddrs); err != nil {
t.Fatalf("error in expected round robin: %v", err)
}
}
// Test tests that the least request balancer properly surfaces resolver
// errors.
func (s) TestLeastRequestEndpoints_ResolverError(t *testing.T) {
const sc = `{"loadBalancingConfig": [{"least_request_experimental": {}}]}`
mr := manual.NewBuilderWithScheme("lr-e2e")
defer mr.Close()
cc, err := grpc.NewClient(
mr.Scheme()+":///",
grpc.WithResolvers(mr),
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithDefaultServiceConfig(sc),
)
if err != nil {
t.Fatalf("grpc.NewClient() failed: %v", err)
}
defer cc.Close()
// We need to pass an endpoint with a valid address to the resolver before
// reporting an error - otherwise endpointsharding does not report the
// error through.
lis, err := testutils.LocalTCPListener()
if err != nil {
t.Fatalf("net.Listen() failed: %v", err)
}
// Act like a server that closes the connection without sending a server
// preface.
go func() {
conn, err := lis.Accept()
if err != nil {
t.Errorf("Unexpected error when accepting a connection: %v", err)
}
conn.Close()
}()
mr.UpdateState(resolver.State{
Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: lis.Addr().String()}}}},
})
cc.Connect()
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
// Report an error through the resolver
resolverErr := fmt.Errorf("simulated resolver error")
mr.CC().ReportError(resolverErr)
// Ensure the client returns the expected resolver error.
testServiceClient := testgrpc.NewTestServiceClient(cc)
for ; ctx.Err() == nil; <-time.After(defaultTestShortTimeout) {
_, err = testServiceClient.EmptyCall(ctx, &testpb.Empty{})
if strings.Contains(err.Error(), resolverErr.Error()) {
break
}
}
if ctx.Err() != nil {
t.Fatalf("Timeout when waiting for RPCs to fail with error containing %s. Last error: %v", resolverErr, err)
}
}

View File

@@ -0,0 +1,37 @@
/*
* Copyright 2024 gRPC authors.
*
* 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.
*
*/
// Package internal contains code internal to the pickfirst package.
package internal
import (
rand "math/rand/v2"
"time"
)
var (
// RandShuffle pseudo-randomizes the order of addresses.
RandShuffle = rand.Shuffle
// RandFloat64 returns, as a float64, a pseudo-random number in [0.0,1.0).
RandFloat64 = rand.Float64
// TimeAfterFunc allows mocking the timer for testing connection delay
// related functionality.
TimeAfterFunc = func(d time.Duration, f func()) func() {
timer := time.AfterFunc(d, f)
return func() { timer.Stop() }
}
)

View File

@@ -0,0 +1,921 @@
/*
*
* Copyright 2017 gRPC authors.
*
* 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.
*
*/
// Package pickfirst contains the pick_first load balancing policy which
// is the universal leaf policy.
package pickfirst
import (
"cmp"
"encoding/json"
"errors"
"fmt"
"math"
"net"
"net/netip"
"slices"
"sync"
"time"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/pickfirst/internal"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/experimental/balancer/weight"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/internal/envconfig"
internalgrpclog "google.golang.org/grpc/internal/grpclog"
"google.golang.org/grpc/internal/pretty"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/serviceconfig"
)
func init() {
balancer.Register(pickfirstBuilder{})
}
// Name is the name of the pick_first balancer.
const Name = "pick_first"
// enableHealthListenerKeyType is a unique key type used in resolver
// attributes to indicate whether the health listener usage is enabled.
type enableHealthListenerKeyType struct{}
var logger = grpclog.Component("pick-first-leaf-lb")
const (
// TODO: change to pick-first when this becomes the default pick_first policy.
logPrefix = "[pick-first-leaf-lb %p] "
// connectionDelayInterval is the time to wait for during the happy eyeballs
// pass before starting the next connection attempt.
connectionDelayInterval = 250 * time.Millisecond
)
type ipAddrFamily int
const (
// ipAddrFamilyUnknown represents strings that can't be parsed as an IP
// address.
ipAddrFamilyUnknown ipAddrFamily = iota
ipAddrFamilyV4
ipAddrFamilyV6
)
type pickfirstBuilder struct{}
func (pickfirstBuilder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
b := &pickfirstBalancer{
cc: cc,
subConns: resolver.NewAddressMapV2[*scData](),
state: connectivity.Connecting,
cancelConnectionTimer: func() {},
}
b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
return b
}
func (b pickfirstBuilder) Name() string {
return Name
}
func (pickfirstBuilder) ParseConfig(js json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
var cfg pfConfig
if err := json.Unmarshal(js, &cfg); err != nil {
return nil, fmt.Errorf("pickfirst: unable to unmarshal LB policy config: %s, error: %v", string(js), err)
}
return cfg, nil
}
// EnableHealthListener updates the state to configure pickfirst for using a
// generic health listener.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a later
// release.
func EnableHealthListener(state resolver.State) resolver.State {
state.Attributes = state.Attributes.WithValue(enableHealthListenerKeyType{}, true)
return state
}
type pfConfig struct {
serviceconfig.LoadBalancingConfig `json:"-"`
// If set to true, instructs the LB policy to shuffle the order of the list
// of endpoints received from the name resolver before attempting to
// connect to them.
ShuffleAddressList bool `json:"shuffleAddressList"`
}
// scData keeps track of the current state of the subConn.
// It is not safe for concurrent access.
type scData struct {
// The following fields are initialized at build time and read-only after
// that.
subConn balancer.SubConn
addr resolver.Address
rawConnectivityState connectivity.State
// The effective connectivity state based on raw connectivity, health state
// and after following sticky TransientFailure behaviour defined in A62.
effectiveState connectivity.State
lastErr error
connectionFailedInFirstPass bool
}
func (b *pickfirstBalancer) newSCData(addr resolver.Address) (*scData, error) {
sd := &scData{
rawConnectivityState: connectivity.Idle,
effectiveState: connectivity.Idle,
addr: addr,
}
sc, err := b.cc.NewSubConn([]resolver.Address{addr}, balancer.NewSubConnOptions{
StateListener: func(state balancer.SubConnState) {
b.updateSubConnState(sd, state)
},
})
if err != nil {
return nil, err
}
sd.subConn = sc
return sd, nil
}
type pickfirstBalancer struct {
// The following fields are initialized at build time and read-only after
// that and therefore do not need to be guarded by a mutex.
logger *internalgrpclog.PrefixLogger
cc balancer.ClientConn
// The mutex is used to ensure synchronization of updates triggered
// from the idle picker and the already serialized resolver,
// SubConn state updates.
mu sync.Mutex
// State reported to the channel based on SubConn states and resolver
// updates.
state connectivity.State
// scData for active subonns mapped by address.
subConns *resolver.AddressMapV2[*scData]
addressList addressList
firstPass bool
numTF int
cancelConnectionTimer func()
healthCheckingEnabled bool
}
// ResolverError is called by the ClientConn when the name resolver produces
// an error or when pickfirst determined the resolver update to be invalid.
func (b *pickfirstBalancer) ResolverError(err error) {
b.mu.Lock()
defer b.mu.Unlock()
b.resolverErrorLocked(err)
}
func (b *pickfirstBalancer) resolverErrorLocked(err error) {
if b.logger.V(2) {
b.logger.Infof("Received error from the name resolver: %v", err)
}
// The picker will not change since the balancer does not currently
// report an error. If the balancer hasn't received a single good resolver
// update yet, transition to TRANSIENT_FAILURE.
if b.state != connectivity.TransientFailure && b.addressList.size() > 0 {
if b.logger.V(2) {
b.logger.Infof("Ignoring resolver error because balancer is using a previous good update.")
}
return
}
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.TransientFailure,
Picker: &picker{err: fmt.Errorf("name resolver error: %v", err)},
})
}
func (b *pickfirstBalancer) UpdateClientConnState(state balancer.ClientConnState) error {
b.mu.Lock()
defer b.mu.Unlock()
b.cancelConnectionTimer()
if len(state.ResolverState.Addresses) == 0 && len(state.ResolverState.Endpoints) == 0 {
// Cleanup state pertaining to the previous resolver state.
// Treat an empty address list like an error by calling b.ResolverError.
b.closeSubConnsLocked()
b.addressList.updateAddrs(nil)
b.resolverErrorLocked(errors.New("produced zero addresses"))
return balancer.ErrBadResolverState
}
b.healthCheckingEnabled = state.ResolverState.Attributes.Value(enableHealthListenerKeyType{}) != nil
cfg, ok := state.BalancerConfig.(pfConfig)
if state.BalancerConfig != nil && !ok {
return fmt.Errorf("pickfirst: received illegal BalancerConfig (type %T): %v: %w", state.BalancerConfig, state.BalancerConfig, balancer.ErrBadResolverState)
}
if b.logger.V(2) {
b.logger.Infof("Received new config %s, resolver state %s", pretty.ToJSON(cfg), pretty.ToJSON(state.ResolverState))
}
var newAddrs []resolver.Address
if endpoints := state.ResolverState.Endpoints; len(endpoints) != 0 {
// Perform the optional shuffling described in gRFC A62. The shuffling
// will change the order of endpoints but not touch the order of the
// addresses within each endpoint. - A61
if cfg.ShuffleAddressList {
if envconfig.PickFirstWeightedShuffling {
type weightedEndpoint struct {
endpoint resolver.Endpoint
weight float64
}
// For each endpoint, compute a key as described in A113 and
// https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf:
var weightedEndpoints []weightedEndpoint
for _, endpoint := range endpoints {
u := internal.RandFloat64() // Random number in [0.0, 1.0)
weight := weightAttribute(endpoint)
weightedEndpoints = append(weightedEndpoints, weightedEndpoint{
endpoint: endpoint,
weight: math.Pow(u, 1.0/float64(weight)),
})
}
// Sort endpoints by key in descending order and reconstruct the
// endpoints slice.
slices.SortFunc(weightedEndpoints, func(a, b weightedEndpoint) int {
return cmp.Compare(b.weight, a.weight)
})
// Here, and in the "else" block below, we clone the endpoints
// slice to avoid mutating the resolver state. Doing the latter
// would lead to data races if the caller is accessing the same
// slice concurrently.
sortedEndpoints := make([]resolver.Endpoint, len(endpoints))
for i, we := range weightedEndpoints {
sortedEndpoints[i] = we.endpoint
}
endpoints = sortedEndpoints
} else {
endpoints = slices.Clone(endpoints)
internal.RandShuffle(len(endpoints), func(i, j int) { endpoints[i], endpoints[j] = endpoints[j], endpoints[i] })
}
}
// "Flatten the list by concatenating the ordered list of addresses for
// each of the endpoints, in order." - A61
for _, endpoint := range endpoints {
newAddrs = append(newAddrs, endpoint.Addresses...)
}
} else {
// Endpoints not set, process addresses until we migrate resolver
// emissions fully to Endpoints. The top channel does wrap emitted
// addresses with endpoints, however some balancers such as weighted
// target do not forward the corresponding correct endpoints down/split
// endpoints properly. Once all balancers correctly forward endpoints
// down, can delete this else conditional.
newAddrs = state.ResolverState.Addresses
if cfg.ShuffleAddressList {
newAddrs = append([]resolver.Address{}, newAddrs...)
internal.RandShuffle(len(newAddrs), func(i, j int) { newAddrs[i], newAddrs[j] = newAddrs[j], newAddrs[i] })
}
}
// If an address appears in multiple endpoints or in the same endpoint
// multiple times, we keep it only once. We will create only one SubConn
// for the address because an AddressMap is used to store SubConns.
// Not de-duplicating would result in attempting to connect to the same
// SubConn multiple times in the same pass. We don't want this.
newAddrs = deDupAddresses(newAddrs)
newAddrs = interleaveAddresses(newAddrs)
prevAddr := b.addressList.currentAddress()
prevSCData, found := b.subConns.Get(prevAddr)
prevAddrsCount := b.addressList.size()
isPrevRawConnectivityStateReady := found && prevSCData.rawConnectivityState == connectivity.Ready
b.addressList.updateAddrs(newAddrs)
// If the previous ready SubConn exists in new address list,
// keep this connection and don't create new SubConns.
if isPrevRawConnectivityStateReady && b.addressList.seekTo(prevAddr) {
return nil
}
b.reconcileSubConnsLocked(newAddrs)
// If it's the first resolver update or the balancer was already READY
// (but the new address list does not contain the ready SubConn) or
// CONNECTING, enter CONNECTING.
// We may be in TRANSIENT_FAILURE due to a previous empty address list,
// we should still enter CONNECTING because the sticky TF behaviour
// mentioned in A62 applies only when the TRANSIENT_FAILURE is reported
// due to connectivity failures.
if isPrevRawConnectivityStateReady || b.state == connectivity.Connecting || prevAddrsCount == 0 {
// Start connection attempt at first address.
b.forceUpdateConcludedStateLocked(balancer.State{
ConnectivityState: connectivity.Connecting,
Picker: &picker{err: balancer.ErrNoSubConnAvailable},
})
b.startFirstPassLocked()
} else if b.state == connectivity.TransientFailure {
// If we're in TRANSIENT_FAILURE, we stay in TRANSIENT_FAILURE until
// we're READY. See A62.
b.startFirstPassLocked()
}
return nil
}
// UpdateSubConnState is unused as a StateListener is always registered when
// creating SubConns.
func (b *pickfirstBalancer) UpdateSubConnState(subConn balancer.SubConn, state balancer.SubConnState) {
b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", subConn, state)
}
func (b *pickfirstBalancer) Close() {
b.mu.Lock()
defer b.mu.Unlock()
b.closeSubConnsLocked()
b.cancelConnectionTimer()
b.state = connectivity.Shutdown
}
// ExitIdle moves the balancer out of idle state. It can be called concurrently
// by the idlePicker and clientConn so access to variables should be
// synchronized.
func (b *pickfirstBalancer) ExitIdle() {
b.mu.Lock()
defer b.mu.Unlock()
if b.state == connectivity.Idle {
// Move the balancer into CONNECTING state immediately. This is done to
// avoid staying in IDLE if a resolver update arrives before the first
// SubConn reports CONNECTING.
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.Connecting,
Picker: &picker{err: balancer.ErrNoSubConnAvailable},
})
b.startFirstPassLocked()
}
}
func (b *pickfirstBalancer) startFirstPassLocked() {
b.firstPass = true
b.numTF = 0
// Reset the connection attempt record for existing SubConns.
for _, sd := range b.subConns.All() {
sd.connectionFailedInFirstPass = false
}
b.requestConnectionLocked()
}
func (b *pickfirstBalancer) closeSubConnsLocked() {
for _, sd := range b.subConns.All() {
sd.subConn.Shutdown()
}
b.subConns = resolver.NewAddressMapV2[*scData]()
}
// deDupAddresses ensures that each address appears only once in the slice.
func deDupAddresses(addrs []resolver.Address) []resolver.Address {
seenAddrs := resolver.NewAddressMapV2[bool]()
retAddrs := []resolver.Address{}
for _, addr := range addrs {
if _, ok := seenAddrs.Get(addr); ok {
continue
}
seenAddrs.Set(addr, true)
retAddrs = append(retAddrs, addr)
}
return retAddrs
}
// interleaveAddresses interleaves addresses of both families (IPv4 and IPv6)
// as per RFC-8305 section 4.
// Whichever address family is first in the list is followed by an address of
// the other address family; that is, if the first address in the list is IPv6,
// then the first IPv4 address should be moved up in the list to be second in
// the list. It doesn't support configuring "First Address Family Count", i.e.
// there will always be a single member of the first address family at the
// beginning of the interleaved list.
// Addresses that are neither IPv4 nor IPv6 are treated as part of a third
// "unknown" family for interleaving.
// See: https://datatracker.ietf.org/doc/html/rfc8305#autoid-6
func interleaveAddresses(addrs []resolver.Address) []resolver.Address {
familyAddrsMap := map[ipAddrFamily][]resolver.Address{}
interleavingOrder := []ipAddrFamily{}
for _, addr := range addrs {
family := addressFamily(addr.Addr)
if _, found := familyAddrsMap[family]; !found {
interleavingOrder = append(interleavingOrder, family)
}
familyAddrsMap[family] = append(familyAddrsMap[family], addr)
}
interleavedAddrs := make([]resolver.Address, 0, len(addrs))
for curFamilyIdx := 0; len(interleavedAddrs) < len(addrs); curFamilyIdx = (curFamilyIdx + 1) % len(interleavingOrder) {
// Some IP types may have fewer addresses than others, so we look for
// the next type that has a remaining member to add to the interleaved
// list.
family := interleavingOrder[curFamilyIdx]
remainingMembers := familyAddrsMap[family]
if len(remainingMembers) > 0 {
interleavedAddrs = append(interleavedAddrs, remainingMembers[0])
familyAddrsMap[family] = remainingMembers[1:]
}
}
return interleavedAddrs
}
// addressFamily returns the ipAddrFamily after parsing the address string.
// If the address isn't of the format "ip-address:port", it returns
// ipAddrFamilyUnknown. The address may be valid even if it's not an IP when
// using a resolver like passthrough where the address may be a hostname in
// some format that the dialer can resolve.
func addressFamily(address string) ipAddrFamily {
// Parse the IP after removing the port.
host, _, err := net.SplitHostPort(address)
if err != nil {
return ipAddrFamilyUnknown
}
ip, err := netip.ParseAddr(host)
if err != nil {
return ipAddrFamilyUnknown
}
switch {
case ip.Is4() || ip.Is4In6():
return ipAddrFamilyV4
case ip.Is6():
return ipAddrFamilyV6
default:
return ipAddrFamilyUnknown
}
}
// reconcileSubConnsLocked updates the active subchannels based on a new address
// list from the resolver. It does this by:
// - closing subchannels: any existing subchannels associated with addresses
// that are no longer in the updated list are shut down.
// - removing subchannels: entries for these closed subchannels are removed
// from the subchannel map.
//
// This ensures that the subchannel map accurately reflects the current set of
// addresses received from the name resolver.
func (b *pickfirstBalancer) reconcileSubConnsLocked(newAddrs []resolver.Address) {
newAddrsMap := resolver.NewAddressMapV2[bool]()
for _, addr := range newAddrs {
newAddrsMap.Set(addr, true)
}
for oldAddr := range b.subConns.All() {
if _, ok := newAddrsMap.Get(oldAddr); ok {
continue
}
val, _ := b.subConns.Get(oldAddr)
val.subConn.Shutdown()
b.subConns.Delete(oldAddr)
}
}
// shutdownRemainingLocked shuts down remaining subConns. Called when a subConn
// becomes ready, which means that all other subConn must be shutdown.
func (b *pickfirstBalancer) shutdownRemainingLocked(selected *scData) {
b.cancelConnectionTimer()
for _, sd := range b.subConns.All() {
if sd.subConn != selected.subConn {
sd.subConn.Shutdown()
}
}
b.subConns = resolver.NewAddressMapV2[*scData]()
b.subConns.Set(selected.addr, selected)
}
// requestConnectionLocked starts connecting on the subchannel corresponding to
// the current address. If no subchannel exists, one is created. If the current
// subchannel is in TransientFailure, a connection to the next address is
// attempted until a subchannel is found.
func (b *pickfirstBalancer) requestConnectionLocked() {
if !b.addressList.isValid() {
return
}
var lastErr error
for valid := true; valid; valid = b.addressList.increment() {
curAddr := b.addressList.currentAddress()
sd, ok := b.subConns.Get(curAddr)
if !ok {
var err error
// We want to assign the new scData to sd from the outer scope,
// hence we can't use := below.
sd, err = b.newSCData(curAddr)
if err != nil {
// This should never happen, unless the clientConn is being shut
// down.
if b.logger.V(2) {
b.logger.Infof("Failed to create a subConn for address %v: %v", curAddr.String(), err)
}
// Do nothing, the LB policy will be closed soon.
return
}
b.subConns.Set(curAddr, sd)
}
switch sd.rawConnectivityState {
case connectivity.Idle:
sd.subConn.Connect()
b.scheduleNextConnectionLocked()
return
case connectivity.TransientFailure:
// The SubConn is being re-used and failed during a previous pass
// over the addressList. It has not completed backoff yet.
// Mark it as having failed and try the next address.
sd.connectionFailedInFirstPass = true
lastErr = sd.lastErr
continue
case connectivity.Connecting:
// Wait for the connection attempt to complete or the timer to fire
// before attempting the next address.
b.scheduleNextConnectionLocked()
return
default:
b.logger.Errorf("SubConn with unexpected state %v present in SubConns map.", sd.rawConnectivityState)
return
}
}
// All the remaining addresses in the list are in TRANSIENT_FAILURE, end the
// first pass if possible.
b.endFirstPassIfPossibleLocked(lastErr)
}
func (b *pickfirstBalancer) scheduleNextConnectionLocked() {
b.cancelConnectionTimer()
if !b.addressList.hasNext() {
return
}
curAddr := b.addressList.currentAddress()
cancelled := false // Access to this is protected by the balancer's mutex.
closeFn := internal.TimeAfterFunc(connectionDelayInterval, func() {
b.mu.Lock()
defer b.mu.Unlock()
// If the scheduled task is cancelled while acquiring the mutex, return.
if cancelled {
return
}
if b.logger.V(2) {
b.logger.Infof("Happy Eyeballs timer expired while waiting for connection to %q.", curAddr.Addr)
}
if b.addressList.increment() {
b.requestConnectionLocked()
}
})
// Access to the cancellation callback held by the balancer is guarded by
// the balancer's mutex, so it's safe to set the boolean from the callback.
b.cancelConnectionTimer = sync.OnceFunc(func() {
cancelled = true
closeFn()
})
}
func (b *pickfirstBalancer) updateSubConnState(sd *scData, newState balancer.SubConnState) {
b.mu.Lock()
defer b.mu.Unlock()
oldState := sd.rawConnectivityState
sd.rawConnectivityState = newState.ConnectivityState
// Previously relevant SubConns can still callback with state updates.
// To prevent pickers from returning these obsolete SubConns, this logic
// is included to check if the current list of active SubConns includes this
// SubConn.
if !b.isActiveSCData(sd) {
return
}
if newState.ConnectivityState == connectivity.Shutdown {
sd.effectiveState = connectivity.Shutdown
return
}
if newState.ConnectivityState == connectivity.TransientFailure {
sd.connectionFailedInFirstPass = true
}
if newState.ConnectivityState == connectivity.Ready {
b.shutdownRemainingLocked(sd)
if !b.addressList.seekTo(sd.addr) {
// This should not fail as we should have only one SubConn after
// entering READY. The SubConn should be present in the addressList.
b.logger.Errorf("Address %q not found address list in %v", sd.addr, b.addressList.addresses)
return
}
if !b.healthCheckingEnabled {
if b.logger.V(2) {
b.logger.Infof("SubConn %p reported connectivity state READY and the health listener is disabled. Transitioning SubConn to READY.", sd.subConn)
}
sd.effectiveState = connectivity.Ready
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.Ready,
Picker: &picker{result: balancer.PickResult{SubConn: sd.subConn}},
})
return
}
if b.logger.V(2) {
b.logger.Infof("SubConn %p reported connectivity state READY. Registering health listener.", sd.subConn)
}
// Send a CONNECTING update to take the SubConn out of sticky-TF if
// required.
sd.effectiveState = connectivity.Connecting
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.Connecting,
Picker: &picker{err: balancer.ErrNoSubConnAvailable},
})
sd.subConn.RegisterHealthListener(func(scs balancer.SubConnState) {
b.updateSubConnHealthState(sd, scs)
})
return
}
// If the LB policy is READY, and it receives a subchannel state change,
// it means that the READY subchannel has failed.
// A SubConn can also transition from CONNECTING directly to IDLE when
// a transport is successfully created, but the connection fails
// before the SubConn can send the notification for READY. We treat
// this as a successful connection and transition to IDLE.
// TODO: https://github.com/grpc/grpc-go/issues/7862 - Remove the second
// part of the if condition below once the issue is fixed.
if oldState == connectivity.Ready || (oldState == connectivity.Connecting && newState.ConnectivityState == connectivity.Idle) {
// Once a transport fails, the balancer enters IDLE and starts from
// the first address when the picker is used.
b.shutdownRemainingLocked(sd)
sd.effectiveState = newState.ConnectivityState
b.addressList.reset()
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.Idle,
Picker: &idlePicker{exitIdle: sync.OnceFunc(b.ExitIdle)},
})
return
}
if b.firstPass {
switch newState.ConnectivityState {
case connectivity.Connecting:
// The effective state can be in either IDLE, CONNECTING or
// TRANSIENT_FAILURE. If it's TRANSIENT_FAILURE, stay in
// TRANSIENT_FAILURE until it's READY. See A62.
if sd.effectiveState != connectivity.TransientFailure {
sd.effectiveState = connectivity.Connecting
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.Connecting,
Picker: &picker{err: balancer.ErrNoSubConnAvailable},
})
}
case connectivity.TransientFailure:
sd.lastErr = newState.ConnectionError
sd.effectiveState = connectivity.TransientFailure
// Since we're re-using common SubConns while handling resolver
// updates, we could receive an out of turn TRANSIENT_FAILURE from
// a pass over the previous address list. Happy Eyeballs will also
// cause out of order updates to arrive.
if curAddr := b.addressList.currentAddress(); equalAddressIgnoringBalAttributes(&curAddr, &sd.addr) {
b.cancelConnectionTimer()
if b.addressList.increment() {
b.requestConnectionLocked()
return
}
}
// End the first pass if we've seen a TRANSIENT_FAILURE from all
// SubConns once.
b.endFirstPassIfPossibleLocked(newState.ConnectionError)
}
return
}
// We have finished the first pass, keep re-connecting failing SubConns.
switch newState.ConnectivityState {
case connectivity.TransientFailure:
b.numTF = (b.numTF + 1) % b.subConns.Len()
sd.lastErr = newState.ConnectionError
if b.numTF%b.subConns.Len() == 0 {
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.TransientFailure,
Picker: &picker{err: newState.ConnectionError},
})
}
// We don't need to request re-resolution since the SubConn already
// does that before reporting TRANSIENT_FAILURE.
// TODO: #7534 - Move re-resolution requests from SubConn into
// pick_first.
case connectivity.Idle:
sd.subConn.Connect()
}
}
// endFirstPassIfPossibleLocked ends the first happy-eyeballs pass if all the
// addresses are tried and their SubConns have reported a failure.
func (b *pickfirstBalancer) endFirstPassIfPossibleLocked(lastErr error) {
// An optimization to avoid iterating over the entire SubConn map.
if b.addressList.isValid() {
return
}
// Connect() has been called on all the SubConns. The first pass can be
// ended if all the SubConns have reported a failure.
for _, sd := range b.subConns.All() {
if !sd.connectionFailedInFirstPass {
return
}
}
b.firstPass = false
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.TransientFailure,
Picker: &picker{err: lastErr},
})
// Start re-connecting all the SubConns that are already in IDLE.
for _, sd := range b.subConns.All() {
if sd.rawConnectivityState == connectivity.Idle {
sd.subConn.Connect()
}
}
}
func (b *pickfirstBalancer) isActiveSCData(sd *scData) bool {
activeSD, found := b.subConns.Get(sd.addr)
return found && activeSD == sd
}
func (b *pickfirstBalancer) updateSubConnHealthState(sd *scData, state balancer.SubConnState) {
b.mu.Lock()
defer b.mu.Unlock()
// Previously relevant SubConns can still callback with state updates.
// To prevent pickers from returning these obsolete SubConns, this logic
// is included to check if the current list of active SubConns includes
// this SubConn.
if !b.isActiveSCData(sd) {
return
}
sd.effectiveState = state.ConnectivityState
switch state.ConnectivityState {
case connectivity.Ready:
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.Ready,
Picker: &picker{result: balancer.PickResult{SubConn: sd.subConn}},
})
case connectivity.TransientFailure:
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.TransientFailure,
Picker: &picker{err: fmt.Errorf("pickfirst: health check failure: %v", state.ConnectionError)},
})
case connectivity.Connecting:
b.updateBalancerState(balancer.State{
ConnectivityState: connectivity.Connecting,
Picker: &picker{err: balancer.ErrNoSubConnAvailable},
})
default:
b.logger.Errorf("Got unexpected health update for SubConn %p: %v", state)
}
}
// updateBalancerState stores the state reported to the channel and calls
// ClientConn.UpdateState(). As an optimization, it avoids sending duplicate
// updates to the channel.
func (b *pickfirstBalancer) updateBalancerState(newState balancer.State) {
// In case of TransientFailures allow the picker to be updated to update
// the connectivity error, in all other cases don't send duplicate state
// updates.
if newState.ConnectivityState == b.state && b.state != connectivity.TransientFailure {
return
}
b.forceUpdateConcludedStateLocked(newState)
}
// forceUpdateConcludedStateLocked stores the state reported to the channel and
// calls ClientConn.UpdateState().
// A separate function is defined to force update the ClientConn state since the
// channel doesn't correctly assume that LB policies start in CONNECTING and
// relies on LB policy to send an initial CONNECTING update.
func (b *pickfirstBalancer) forceUpdateConcludedStateLocked(newState balancer.State) {
b.state = newState.ConnectivityState
b.cc.UpdateState(newState)
}
type picker struct {
result balancer.PickResult
err error
}
func (p *picker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
return p.result, p.err
}
// idlePicker is used when the SubConn is IDLE and kicks the SubConn into
// CONNECTING when Pick is called.
type idlePicker struct {
exitIdle func()
}
func (i *idlePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
i.exitIdle()
return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
}
// addressList manages sequentially iterating over addresses present in a list
// of endpoints. It provides a 1 dimensional view of the addresses present in
// the endpoints.
// This type is not safe for concurrent access.
type addressList struct {
addresses []resolver.Address
idx int
}
func (al *addressList) isValid() bool {
return al.idx < len(al.addresses)
}
func (al *addressList) size() int {
return len(al.addresses)
}
// increment moves to the next index in the address list.
// This method returns false if it went off the list, true otherwise.
func (al *addressList) increment() bool {
if !al.isValid() {
return false
}
al.idx++
return al.idx < len(al.addresses)
}
// currentAddress returns the current address pointed to in the addressList.
// If the list is in an invalid state, it returns an empty address instead.
func (al *addressList) currentAddress() resolver.Address {
if !al.isValid() {
return resolver.Address{}
}
return al.addresses[al.idx]
}
func (al *addressList) reset() {
al.idx = 0
}
func (al *addressList) updateAddrs(addrs []resolver.Address) {
al.addresses = addrs
al.reset()
}
// seekTo returns false if the needle was not found and the current index was
// left unchanged.
func (al *addressList) seekTo(needle resolver.Address) bool {
for ai, addr := range al.addresses {
if !equalAddressIgnoringBalAttributes(&addr, &needle) {
continue
}
al.idx = ai
return true
}
return false
}
// hasNext returns whether incrementing the addressList will result in moving
// past the end of the list. If the list has already moved past the end, it
// returns false.
func (al *addressList) hasNext() bool {
if !al.isValid() {
return false
}
return al.idx+1 < len(al.addresses)
}
// equalAddressIgnoringBalAttributes returns true is a and b are considered
// equal. This is different from the Equal method on the resolver.Address type
// which considers all fields to determine equality. Here, we only consider
// fields that are meaningful to the SubConn.
func equalAddressIgnoringBalAttributes(a, b *resolver.Address) bool {
return a.Addr == b.Addr && a.ServerName == b.ServerName &&
a.Attributes.Equal(b.Attributes)
}
// weightAttribute is a convenience function which returns the value of the
// weight endpoint Attribute.
//
// When used in the xDS context, the weight attribute is guaranteed to be
// non-zero. But, when used in a non-xDS context, the weight attribute could be
// unset. A Default of 1 is used in the latter case.
func weightAttribute(e resolver.Endpoint) uint32 {
w := weight.FromEndpoint(e).Weight
if w == 0 {
return 1
}
return w
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,376 @@
/*
*
* Copyright 2024 gRPC authors.
*
* 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.
*
*/
package pickfirst
import (
"context"
"errors"
"fmt"
"testing"
"time"
"google.golang.org/grpc/attributes"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/internal/grpctest"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/resolver"
)
const (
// Default timeout for tests in this package.
defaultTestTimeout = 10 * time.Second
// Default short timeout, to be used when waiting for events which are not
// expected to happen.
defaultTestShortTimeout = 100 * time.Millisecond
)
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
// TestPickFirst_InitialResolverError sends a resolver error to the balancer
// before a valid resolver update. It verifies that the clientconn state is
// updated to TRANSIENT_FAILURE.
func (s) TestPickFirst_InitialResolverError(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
cc := testutils.NewBalancerClientConn(t)
bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{})
defer bal.Close()
bal.ResolverError(errors.New("resolution failed: test error"))
if err := cc.WaitForConnectivityState(ctx, connectivity.TransientFailure); err != nil {
t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.TransientFailure, err)
}
// After sending a valid update, the LB policy should report CONNECTING.
ccState := balancer.ClientConnState{
ResolverState: resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
{Addresses: []resolver.Address{{Addr: "2.2.2.2:2"}}},
},
},
}
if err := bal.UpdateClientConnState(ccState); err != nil {
t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
}
if err := cc.WaitForConnectivityState(ctx, connectivity.Connecting); err != nil {
t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.Connecting, err)
}
}
// TestPickFirst_ResolverErrorinTF sends a resolver error to the balancer
// before when it's attempting to connect to a SubConn TRANSIENT_FAILURE. It
// verifies that the picker is updated and the SubConn is not closed.
func (s) TestPickFirst_ResolverErrorinTF(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
cc := testutils.NewBalancerClientConn(t)
bal := balancer.Get(Name).Build(cc, balancer.BuildOptions{})
defer bal.Close()
// After sending a valid update, the LB policy should report CONNECTING.
ccState := balancer.ClientConnState{
ResolverState: resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
},
},
}
if err := bal.UpdateClientConnState(ccState); err != nil {
t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
}
sc1 := <-cc.NewSubConnCh
if err := cc.WaitForConnectivityState(ctx, connectivity.Connecting); err != nil {
t.Fatalf("cc.WaitForConnectivityState(%v) returned error: %v", connectivity.Connecting, err)
}
scErr := fmt.Errorf("test error: connection refused")
sc1.UpdateState(balancer.SubConnState{
ConnectivityState: connectivity.TransientFailure,
ConnectionError: scErr,
})
if err := cc.WaitForPickerWithErr(ctx, scErr); err != nil {
t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", scErr, err)
}
bal.ResolverError(errors.New("resolution failed: test error"))
if err := cc.WaitForErrPicker(ctx); err != nil {
t.Fatalf("cc.WaitForPickerWithErr() returned error: %v", err)
}
select {
case <-time.After(defaultTestShortTimeout):
case sc := <-cc.ShutdownSubConnCh:
t.Fatalf("Unexpected SubConn shutdown: %v", sc)
}
}
// TestAddressList_Iteration verifies the behaviour of the addressList while
// iterating through the entries.
func (s) TestAddressList_Iteration(t *testing.T) {
addrs := []resolver.Address{
{
Addr: "192.168.1.1",
ServerName: "test-host-1",
Attributes: attributes.New("key-1", "val-1"),
BalancerAttributes: attributes.New("bal-key-1", "bal-val-1"),
},
{
Addr: "192.168.1.2",
ServerName: "test-host-2",
Attributes: attributes.New("key-2", "val-2"),
BalancerAttributes: attributes.New("bal-key-2", "bal-val-2"),
},
{
Addr: "192.168.1.3",
ServerName: "test-host-3",
Attributes: attributes.New("key-3", "val-3"),
BalancerAttributes: attributes.New("bal-key-3", "bal-val-3"),
},
}
addressList := addressList{}
addressList.updateAddrs(addrs)
for i := 0; i < len(addrs); i++ {
if got, want := addressList.isValid(), true; got != want {
t.Fatalf("addressList.isValid() = %t, want %t", got, want)
}
if got, want := addressList.currentAddress(), addrs[i]; !want.Equal(got) {
t.Errorf("addressList.currentAddress() = %v, want %v", got, want)
}
if got, want := addressList.increment(), i+1 < len(addrs); got != want {
t.Fatalf("addressList.increment() = %t, want %t", got, want)
}
}
if got, want := addressList.isValid(), false; got != want {
t.Fatalf("addressList.isValid() = %t, want %t", got, want)
}
// increment an invalid address list.
if got, want := addressList.increment(), false; got != want {
t.Errorf("addressList.increment() = %t, want %t", got, want)
}
if got, want := addressList.isValid(), false; got != want {
t.Errorf("addressList.isValid() = %t, want %t", got, want)
}
addressList.reset()
for i := 0; i < len(addrs); i++ {
if got, want := addressList.isValid(), true; got != want {
t.Fatalf("addressList.isValid() = %t, want %t", got, want)
}
if got, want := addressList.currentAddress(), addrs[i]; !want.Equal(got) {
t.Errorf("addressList.currentAddress() = %v, want %v", got, want)
}
if got, want := addressList.increment(), i+1 < len(addrs); got != want {
t.Fatalf("addressList.increment() = %t, want %t", got, want)
}
}
}
// TestAddressList_SeekTo verifies the behaviour of addressList.seekTo.
func (s) TestAddressList_SeekTo(t *testing.T) {
addrs := []resolver.Address{
{
Addr: "192.168.1.1",
ServerName: "test-host-1",
Attributes: attributes.New("key-1", "val-1"),
BalancerAttributes: attributes.New("bal-key-1", "bal-val-1"),
},
{
Addr: "192.168.1.2",
ServerName: "test-host-2",
Attributes: attributes.New("key-2", "val-2"),
BalancerAttributes: attributes.New("bal-key-2", "bal-val-2"),
},
{
Addr: "192.168.1.3",
ServerName: "test-host-3",
Attributes: attributes.New("key-3", "val-3"),
BalancerAttributes: attributes.New("bal-key-3", "bal-val-3"),
},
}
addressList := addressList{}
addressList.updateAddrs(addrs)
// Try finding an address in the list.
key := resolver.Address{
Addr: "192.168.1.2",
ServerName: "test-host-2",
Attributes: attributes.New("key-2", "val-2"),
BalancerAttributes: attributes.New("ignored", "bal-val-2"),
}
if got, want := addressList.seekTo(key), true; got != want {
t.Errorf("addressList.seekTo(%v) = %t, want %t", key, got, want)
}
// It should be possible to increment once more now that the pointer has advanced.
if got, want := addressList.increment(), true; got != want {
t.Errorf("addressList.increment() = %t, want %t", got, want)
}
if got, want := addressList.increment(), false; got != want {
t.Errorf("addressList.increment() = %t, want %t", got, want)
}
// Seek to the key again, it is behind the pointer now.
if got, want := addressList.seekTo(key), true; got != want {
t.Errorf("addressList.seekTo(%v) = %t, want %t", key, got, want)
}
// Seek to a key not in the list.
key = resolver.Address{
Addr: "192.168.1.5",
ServerName: "test-host-5",
Attributes: attributes.New("key-5", "val-5"),
BalancerAttributes: attributes.New("ignored", "bal-val-5"),
}
if got, want := addressList.seekTo(key), false; got != want {
t.Errorf("addressList.seekTo(%v) = %t, want %t", key, got, want)
}
// It should be possible to increment once more since the pointer has not advanced.
if got, want := addressList.increment(), true; got != want {
t.Errorf("addressList.increment() = %t, want %t", got, want)
}
if got, want := addressList.increment(), false; got != want {
t.Errorf("addressList.increment() = %t, want %t", got, want)
}
}
// TestPickFirstLeaf_TFPickerUpdate sends TRANSIENT_FAILURE SubConn state updates
// for each SubConn managed by a pickfirst balancer. It verifies that the picker
// is updated with the expected frequency.
func (s) TestPickFirstLeaf_TFPickerUpdate(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
cc := testutils.NewBalancerClientConn(t)
bal := pickfirstBuilder{}.Build(cc, balancer.BuildOptions{})
defer bal.Close()
ccState := balancer.ClientConnState{
ResolverState: resolver.State{
Endpoints: []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}},
{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}}, // duplicate, should be ignored.
{Addresses: []resolver.Address{{Addr: "2.2.2.2:2"}}},
{Addresses: []resolver.Address{{Addr: "1.1.1.1:1"}}}, // duplicate, should be ignored.
},
},
}
if err := bal.UpdateClientConnState(ccState); err != nil {
t.Fatalf("UpdateClientConnState(%v) returned error: %v", ccState, err)
}
// PF should report TRANSIENT_FAILURE only once all the sunbconns have failed
// once.
tfErr := fmt.Errorf("test err: connection refused")
sc1 := <-cc.NewSubConnCh
select {
case <-sc1.ConnectCh:
case <-ctx.Done():
t.Fatal("Context timed out waiting for Connect() to be called on sc1.")
}
sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: tfErr})
// Move the subconn back to IDLE, it should not be re-connected until the
// first pass is complete.
shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
defer shortCancel()
sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
select {
case <-sc1.ConnectCh:
t.Fatal("Connect() unexpectedly called on sc1.")
case <-shortCtx.Done():
}
if err := cc.WaitForPickerWithErr(ctx, balancer.ErrNoSubConnAvailable); err != nil {
t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", balancer.ErrNoSubConnAvailable, err)
}
sc2 := <-cc.NewSubConnCh
select {
case <-sc2.ConnectCh:
case <-ctx.Done():
t.Fatal("Context timed out waiting for Connect() to be called on sc2.")
}
sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: tfErr})
if err := cc.WaitForPickerWithErr(ctx, tfErr); err != nil {
t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", tfErr, err)
}
// Subsequent TRANSIENT_FAILUREs should be reported only after seeing "# of SubConns"
// TRANSIENT_FAILUREs.
// Both the subconns should be connected in parallel.
select {
case <-sc1.ConnectCh:
case <-ctx.Done():
t.Fatal("Context timed out waiting for Connect() to be called on sc1.")
}
shortCtx, shortCancel = context.WithTimeout(ctx, defaultTestShortTimeout)
defer shortCancel()
select {
case <-sc2.ConnectCh:
t.Fatal("Connect() called on sc2 before it completed backing-off.")
case <-shortCtx.Done():
}
sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
select {
case <-sc2.ConnectCh:
case <-ctx.Done():
t.Fatal("Context timed out waiting for Connect() to be called on sc2.")
}
newTfErr := fmt.Errorf("test err: unreachable")
sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: newTfErr})
select {
case <-time.After(defaultTestShortTimeout):
case p := <-cc.NewPickerCh:
sc, err := p.Pick(balancer.PickInfo{})
t.Fatalf("Unexpected picker update: %v, %v", sc, err)
}
sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: newTfErr})
if err := cc.WaitForPickerWithErr(ctx, newTfErr); err != nil {
t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", newTfErr, err)
}
}

View File

@@ -0,0 +1,201 @@
/*
*
* Copyright 2025 gRPC authors.
*
* 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.
*
*/
// Package randomsubsetting implements the random_subsetting LB policy specified
// here: https://github.com/grpc/proposal/blob/master/A68-random-subsetting.md
//
// To install the LB policy, import this package as:
//
// import _ "google.golang.org/grpc/balancer/randomsubsetting"
//
// # Experimental
//
// Notice: This package is EXPERIMENTAL and may be changed or removed in a
// later release.
package randomsubsetting
import (
"cmp"
"encoding/json"
"fmt"
"math/rand/v2"
"slices"
xxhash "github.com/cespare/xxhash/v2"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/grpclog"
"google.golang.org/grpc/internal/balancer/gracefulswitch"
internalgrpclog "google.golang.org/grpc/internal/grpclog"
iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/serviceconfig"
)
// Name is the name of the random subsetting load balancer.
const Name = "random_subsetting_experimental"
var (
logger = grpclog.Component(Name)
randUint64 = rand.Uint64
)
func prefixLogger(p *subsettingBalancer) *internalgrpclog.PrefixLogger {
return internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[random-subsetting-lb %p] ", p))
}
func init() {
balancer.Register(bb{})
}
type bb struct{}
func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
b := &subsettingBalancer{
Balancer: gracefulswitch.NewBalancer(cc, bOpts),
hashSeed: randUint64(),
hashDigest: xxhash.New(),
}
b.logger = prefixLogger(b)
b.logger.Infof("Created")
return b
}
type lbConfig struct {
serviceconfig.LoadBalancingConfig `json:"-"`
SubsetSize uint32 `json:"subsetSize,omitempty"`
ChildPolicy *iserviceconfig.BalancerConfig `json:"childPolicy,omitempty"`
}
func (bb) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
lbCfg := &lbConfig{}
// Ensure that the specified child policy is registered and validates its
// config, if present.
if err := json.Unmarshal(s, lbCfg); err != nil {
return nil, fmt.Errorf("randomsubsetting: json.Unmarshal failed for configuration: %s with error: %v", string(s), err)
}
if lbCfg.SubsetSize == 0 {
return nil, fmt.Errorf("randomsubsetting: SubsetSize must be greater than 0")
}
if lbCfg.ChildPolicy == nil {
return nil, fmt.Errorf("randomsubsetting: ChildPolicy must be specified")
}
return lbCfg, nil
}
func (bb) Name() string {
return Name
}
type subsettingBalancer struct {
*gracefulswitch.Balancer
logger *internalgrpclog.PrefixLogger
cfg *lbConfig
hashSeed uint64
hashDigest *xxhash.Digest
}
func (b *subsettingBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
lbCfg, ok := s.BalancerConfig.(*lbConfig)
if !ok {
b.logger.Warningf("Received config with unexpected type %T: %v", s.BalancerConfig, s.BalancerConfig)
return balancer.ErrBadResolverState
}
// Build config for the gracefulswitch balancer. It is safe to ignore
// JSON marshaling errors here, since the config was already validated
// as part of ParseConfig().
cfg := []map[string]any{{lbCfg.ChildPolicy.Name: lbCfg.ChildPolicy.Config}}
cfgJSON, _ := json.Marshal(cfg)
parsedCfg, err := gracefulswitch.ParseConfig(cfgJSON)
if err != nil {
return fmt.Errorf("randomsubsetting: error switching to child of type %q: %v", lbCfg.ChildPolicy.Name, err)
}
b.cfg = lbCfg
endpoints := resolver.State{
Endpoints: b.calculateSubset(s.ResolverState.Endpoints),
ServiceConfig: s.ResolverState.ServiceConfig,
Attributes: s.ResolverState.Attributes,
}
return b.Balancer.UpdateClientConnState(balancer.ClientConnState{
ResolverState: endpoints,
BalancerConfig: parsedCfg,
})
}
// calculateSubset implements the subsetting algorithm, as described in A68:
// https://github.com/grpc/proposal/blob/master/A68-random-subsetting.md#subsetting-algorithm
func (b *subsettingBalancer) calculateSubset(endpoints []resolver.Endpoint) []resolver.Endpoint {
// A helper struct to hold an endpoint and its hash.
type endpointWithHash struct {
hash uint64
ep resolver.Endpoint
}
// Empty endpoints cannot be hashed and are ignored by child policies.
// Remove them here so they do not count toward the requested subset size.
isEmpty := func(endpoint resolver.Endpoint) bool {
return len(endpoint.Addresses) == 0
}
if slices.ContainsFunc(endpoints, isEmpty) {
endpoints = slices.DeleteFunc(slices.Clone(endpoints), isEmpty)
}
subsetSize := b.cfg.SubsetSize
if len(endpoints) <= int(subsetSize) {
return endpoints
}
hashedEndpoints := make([]endpointWithHash, len(endpoints))
for i, endpoint := range endpoints {
// For every endpoint in the list, compute a hash with previously
// generated seed - A68.
//
// The xxhash package's Sum64() function does not allow setting a seed.
// This means that we need to reset the digest with the seed for every
// endpoint. Without this, an endpoint will not retain the same hash
// across resolver updates.
//
// Note that we only hash the first address of the endpoint, as per A68.
b.hashDigest.ResetWithSeed(b.hashSeed)
b.hashDigest.WriteString(endpoint.Addresses[0].String())
hashedEndpoints[i] = endpointWithHash{
hash: b.hashDigest.Sum64(),
ep: endpoint,
}
}
slices.SortFunc(hashedEndpoints, func(a, b endpointWithHash) int {
// Note: This uses the standard library cmp package, not the
// github.com/google/go-cmp/cmp package. The latter is intended for
// testing purposes only.
return cmp.Compare(a.hash, b.hash)
})
// Convert back to resolver.Endpoints
endpointSubset := make([]resolver.Endpoint, subsetSize)
for i, endpoint := range hashedEndpoints[:subsetSize] {
endpointSubset[i] = endpoint.ep
}
return endpointSubset
}

View File

@@ -0,0 +1,281 @@
/*
*
* Copyright 2025 gRPC authors.
*
* 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.
*
*/
package randomsubsetting
import (
"context"
"encoding/json"
"fmt"
"strings"
"testing"
"time"
xxhash "github.com/cespare/xxhash/v2"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/internal/balancer/stub"
"google.golang.org/grpc/internal/grpctest"
iserviceconfig "google.golang.org/grpc/internal/serviceconfig"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/serviceconfig"
_ "google.golang.org/grpc/balancer/roundrobin" // For round_robin LB policy in tests
)
var defaultTestTimeout = 5 * time.Second
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
func (s) TestParseConfig(t *testing.T) {
parser := bb{}
tests := []struct {
name string
input string
wantCfg serviceconfig.LoadBalancingConfig
wantErr string
}{
{
name: "invalid_json",
input: "{{invalidjson{{",
wantErr: "json.Unmarshal failed for configuration",
},
{
name: "empty_config",
input: `{}`,
wantErr: "SubsetSize must be greater than 0",
},
{
name: "subset_size_zero",
input: `{ "subsetSize": 0 }`,
wantErr: "SubsetSize must be greater than 0",
},
{
name: "child_policy_missing",
input: `{ "subsetSize": 1 }`,
wantErr: "ChildPolicy must be specified",
},
{
name: "child_policy_not_registered",
input: `{ "subsetSize": 1 , "childPolicy": [{"unregistered_lb": {}}] }`,
wantErr: "no supported policies found",
},
{
name: "success",
input: `{ "subsetSize": 3, "childPolicy": [{"round_robin": {}}]}`,
wantCfg: &lbConfig{
SubsetSize: 3,
ChildPolicy: &iserviceconfig.BalancerConfig{Name: "round_robin"},
},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
gotCfg, gotErr := parser.ParseConfig(json.RawMessage(test.input))
// Substring match makes this very tightly coupled to the
// internalserviceconfig.BalancerConfig error strings. However, it
// is important to distinguish the different types of error messages
// possible as the parser has a few defined buckets of ways it can
// error out.
if (gotErr != nil) != (test.wantErr != "") {
t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
}
if gotErr != nil && !strings.Contains(gotErr.Error(), test.wantErr) {
t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
}
if test.wantErr != "" {
return
}
if diff := cmp.Diff(test.wantCfg, gotCfg); diff != "" {
t.Fatalf("ParseConfig(%s) got unexpected output, diff (-want +got):\n%s", test.input, diff)
}
})
}
}
func makeEndpoints(n int) []resolver.Endpoint {
endpoints := make([]resolver.Endpoint, n)
for i := 0; i < n; i++ {
endpoints[i] = resolver.Endpoint{
Addresses: []resolver.Address{{Addr: fmt.Sprintf("endpoint-%d", i)}},
}
}
return endpoints
}
func (s) TestCalculateSubset_Simple(t *testing.T) {
tests := []struct {
name string
endpoints []resolver.Endpoint
subsetSize uint32
want []resolver.Endpoint
}{
{
name: "NoEndpoints",
endpoints: []resolver.Endpoint{},
subsetSize: 3,
want: []resolver.Endpoint{},
},
{
name: "SubsetSizeLargerThanNumberOfEndpoints",
endpoints: makeEndpoints(5),
subsetSize: 10,
want: makeEndpoints(5),
},
{
name: "SubsetSizeEqualToNumberOfEndpoints",
endpoints: makeEndpoints(5),
subsetSize: 5,
want: makeEndpoints(5),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
b := &subsettingBalancer{
cfg: &lbConfig{SubsetSize: tt.subsetSize},
hashSeed: 0,
hashDigest: xxhash.New(),
}
got := b.calculateSubset(tt.endpoints)
if diff := cmp.Diff(tt.want, got); diff != "" {
t.Errorf("calculateSubset() returned diff (-want +got):\n%s", diff)
}
})
}
}
func (s) TestCalculateSubset_IgnoresEndpointsWithoutAddresses(t *testing.T) {
emptyEndpoint := resolver.Endpoint{Addresses: []resolver.Address{}}
endpoints := append([]resolver.Endpoint{emptyEndpoint}, makeEndpoints(3)...)
b := &subsettingBalancer{
cfg: &lbConfig{SubsetSize: 2},
hashSeed: 0,
hashDigest: xxhash.New(),
}
got := b.calculateSubset(endpoints)
if len(got) != 2 {
t.Fatalf("calculateSubset() returned %d endpoints, want 2", len(got))
}
for _, endpoint := range got {
if len(endpoint.Addresses) == 0 {
t.Errorf("calculateSubset() returned an endpoint without addresses")
}
}
}
func (s) TestCalculateSubset_EndpointsRetainHashValues(t *testing.T) {
endpoints := makeEndpoints(10)
const subsetSize = 5
// The subset is deterministic based on the hash, so we can hardcode
// the expected output.
want := []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: "endpoint-6"}}},
{Addresses: []resolver.Address{{Addr: "endpoint-0"}}},
{Addresses: []resolver.Address{{Addr: "endpoint-1"}}},
{Addresses: []resolver.Address{{Addr: "endpoint-7"}}},
{Addresses: []resolver.Address{{Addr: "endpoint-3"}}},
}
b := &subsettingBalancer{
cfg: &lbConfig{SubsetSize: subsetSize},
hashSeed: 0,
hashDigest: xxhash.New(),
}
for range 10 {
got := b.calculateSubset(endpoints)
if diff := cmp.Diff(want, got); diff != "" {
t.Fatalf("calculateSubset() returned diff (-want +got):\n%s", diff)
}
}
}
func (s) TestSubsettingBalancer_DeterministicSubset(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
// Register the stub balancer builder, which will be used as the child
// policy in the random_subsetting balancer.
updateCh := make(chan balancer.ClientConnState, 1)
stub.Register("stub-child-balancer", stub.BalancerFuncs{
UpdateClientConnState: func(_ *stub.BalancerData, s balancer.ClientConnState) error {
select {
case <-ctx.Done():
case updateCh <- s:
}
return nil
},
})
// Create a random_subsetting balancer.
tcc := testutils.NewBalancerClientConn(t)
rsb := balancer.Get(Name).Build(tcc, balancer.BuildOptions{})
defer rsb.Close()
// Prepare the configuration and resolver state to be passed to the
// random_subsetting balancer.
endpoints := makeEndpoints(10)
state := balancer.ClientConnState{
ResolverState: resolver.State{Endpoints: endpoints},
BalancerConfig: &lbConfig{
SubsetSize: 5,
ChildPolicy: &iserviceconfig.BalancerConfig{Name: "stub-child-balancer"},
},
}
// Send the resolver state to the random_subsetting balancer and verify that
// the child policy receives the expected number of endpoints.
if err := rsb.UpdateClientConnState(state); err != nil {
t.Fatalf("UpdateClientConnState failed: %v", err)
}
var wantEndpoints []resolver.Endpoint
select {
case s := <-updateCh:
if len(s.ResolverState.Endpoints) != 5 {
t.Fatalf("Child policy received %d endpoints, want 5", len(s.ResolverState.Endpoints))
}
// Store the subset for the next comparison.
wantEndpoints = s.ResolverState.Endpoints
case <-ctx.Done():
t.Fatal("Timed out waiting for child policy to receive an update")
}
// Call UpdateClientConnState again with the same configuration.
if err := rsb.UpdateClientConnState(state); err != nil {
t.Fatalf("Second UpdateClientConnState failed: %v", err)
}
// Verify that the child policy receives the same subset of endpoints.
select {
case s := <-updateCh:
if diff := cmp.Diff(wantEndpoints, s.ResolverState.Endpoints); diff != "" {
t.Fatalf("Child policy received a different subset of endpoints on second update, diff (-want +got):\n%s", diff)
}
case <-ctx.Done():
t.Fatal("Timed out waiting for second child policy update")
}
}

View File

@@ -0,0 +1,77 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"encoding/json"
"fmt"
"strings"
"google.golang.org/grpc/internal/envconfig"
"google.golang.org/grpc/internal/metadata"
iringhash "google.golang.org/grpc/internal/ringhash"
)
const (
defaultMinSize = 1024
defaultMaxSize = 4096
ringHashSizeUpperBound = 8 * 1024 * 1024 // 8M
)
func parseConfig(c json.RawMessage) (*iringhash.LBConfig, error) {
var cfg iringhash.LBConfig
if err := json.Unmarshal(c, &cfg); err != nil {
return nil, err
}
if cfg.MinRingSize > ringHashSizeUpperBound {
return nil, fmt.Errorf("min_ring_size value of %d is greater than max supported value %d for this field", cfg.MinRingSize, ringHashSizeUpperBound)
}
if cfg.MaxRingSize > ringHashSizeUpperBound {
return nil, fmt.Errorf("max_ring_size value of %d is greater than max supported value %d for this field", cfg.MaxRingSize, ringHashSizeUpperBound)
}
if cfg.MinRingSize == 0 {
cfg.MinRingSize = defaultMinSize
}
if cfg.MaxRingSize == 0 {
cfg.MaxRingSize = defaultMaxSize
}
if cfg.MinRingSize > cfg.MaxRingSize {
return nil, fmt.Errorf("min %v is greater than max %v", cfg.MinRingSize, cfg.MaxRingSize)
}
if cfg.MinRingSize > envconfig.RingHashCap {
cfg.MinRingSize = envconfig.RingHashCap
}
if cfg.MaxRingSize > envconfig.RingHashCap {
cfg.MaxRingSize = envconfig.RingHashCap
}
if !envconfig.RingHashSetRequestHashKey {
cfg.RequestHashHeader = ""
}
if cfg.RequestHashHeader != "" {
cfg.RequestHashHeader = strings.ToLower(cfg.RequestHashHeader)
// See rules in https://github.com/grpc/proposal/blob/master/A76-ring-hash-improvements.md#explicitly-setting-the-request-hash-key
if err := metadata.ValidateKey(cfg.RequestHashHeader); err != nil {
return nil, fmt.Errorf("invalid requestHashHeader %q: %v", cfg.RequestHashHeader, err)
}
if strings.HasSuffix(cfg.RequestHashHeader, "-bin") {
return nil, fmt.Errorf("invalid requestHashHeader %q: key must not end with \"-bin\"", cfg.RequestHashHeader)
}
}
return &cfg, nil
}

View File

@@ -0,0 +1,173 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"encoding/json"
"testing"
"github.com/google/go-cmp/cmp"
"google.golang.org/grpc/internal/envconfig"
iringhash "google.golang.org/grpc/internal/ringhash"
"google.golang.org/grpc/internal/testutils"
)
func (s) TestParseConfig(t *testing.T) {
tests := []struct {
name string
js string
envConfigCap uint64
requestHeaderEnvVar bool
want *iringhash.LBConfig
wantErr bool
}{
{
name: "OK",
js: `{"minRingSize": 1, "maxRingSize": 2}`,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{MinRingSize: 1, MaxRingSize: 2},
},
{
name: "OK with default min",
js: `{"maxRingSize": 2000}`,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{MinRingSize: defaultMinSize, MaxRingSize: 2000},
},
{
name: "OK with default max",
js: `{"minRingSize": 2000}`,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{MinRingSize: 2000, MaxRingSize: defaultMaxSize},
},
{
name: "min greater than max",
js: `{"minRingSize": 10, "maxRingSize": 2}`,
requestHeaderEnvVar: true,
want: nil,
wantErr: true,
},
{
name: "min greater than max greater than global limit",
js: `{"minRingSize": 6000, "maxRingSize": 5000}`,
requestHeaderEnvVar: true,
want: nil,
wantErr: true,
},
{
name: "max greater than global limit",
js: `{"minRingSize": 1, "maxRingSize": 6000}`,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{MinRingSize: 1, MaxRingSize: 4096},
},
{
name: "min and max greater than global limit",
js: `{"minRingSize": 5000, "maxRingSize": 6000}`,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{MinRingSize: 4096, MaxRingSize: 4096},
},
{
name: "min and max less than raised global limit",
js: `{"minRingSize": 5000, "maxRingSize": 6000}`,
envConfigCap: 8000,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{MinRingSize: 5000, MaxRingSize: 6000},
},
{
name: "min and max greater than raised global limit",
js: `{"minRingSize": 10000, "maxRingSize": 10000}`,
envConfigCap: 8000,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{MinRingSize: 8000, MaxRingSize: 8000},
},
{
name: "min greater than upper bound",
js: `{"minRingSize": 8388610, "maxRingSize": 10}`,
requestHeaderEnvVar: true,
want: nil,
wantErr: true,
},
{
name: "max greater than upper bound",
js: `{"minRingSize": 10, "maxRingSize": 8388610}`,
requestHeaderEnvVar: true,
want: nil,
wantErr: true,
},
{
name: "request metadata key set",
js: `{"requestHashHeader": "x-foo"}`,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{
MinRingSize: defaultMinSize,
MaxRingSize: defaultMaxSize,
RequestHashHeader: "x-foo",
},
},
{
name: "request metadata key set with uppercase letters",
js: `{"requestHashHeader": "x-FOO"}`,
requestHeaderEnvVar: true,
want: &iringhash.LBConfig{
MinRingSize: defaultMinSize,
MaxRingSize: defaultMaxSize,
RequestHashHeader: "x-foo",
},
},
{
name: "invalid request hash header",
js: `{"requestHashHeader": "!invalid"}`,
requestHeaderEnvVar: true,
want: nil,
wantErr: true,
},
{
name: "binary request hash header",
js: `{"requestHashHeader": "header-with-bin"}`,
requestHeaderEnvVar: true,
want: nil,
wantErr: true,
},
{
name: "request hash header cleared when RingHashSetRequestHashKey env var is false",
js: `{"requestHashHeader": "x-foo"}`,
requestHeaderEnvVar: false,
want: &iringhash.LBConfig{
MinRingSize: defaultMinSize,
MaxRingSize: defaultMaxSize,
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.envConfigCap != 0 {
testutils.SetEnvConfig(t, &envconfig.RingHashCap, tt.envConfigCap)
}
testutils.SetEnvConfig(t, &envconfig.RingHashSetRequestHashKey, tt.requestHeaderEnvVar)
got, err := parseConfig(json.RawMessage(tt.js))
if (err != nil) != tt.wantErr {
t.Errorf("parseConfig() error = %v, wantErr %v", err, tt.wantErr)
return
}
if diff := cmp.Diff(got, tt.want); diff != "" {
t.Errorf("parseConfig() got unexpected output, diff (-got +want): %v", diff)
}
})
}
}

View File

@@ -0,0 +1,34 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"fmt"
"google.golang.org/grpc/grpclog"
internalgrpclog "google.golang.org/grpc/internal/grpclog"
)
const prefix = "[ring-hash-lb %p] "
var logger = grpclog.Component("xds")
func prefixLogger(p *ringhashBalancer) *internalgrpclog.PrefixLogger {
return internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(prefix, p))
}

124
balancer/ringhash/picker.go Normal file
View File

@@ -0,0 +1,124 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"fmt"
"strings"
xxhash "github.com/cespare/xxhash/v2"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
iringhash "google.golang.org/grpc/internal/ringhash"
"google.golang.org/grpc/metadata"
)
type picker struct {
ring *ring
// endpointStates is a cache of endpoint states.
// The ringhash balancer stores endpoint states in a `resolver.EndpointMap`,
// with access guarded by `ringhashBalancer.mu`. The `endpointStates` cache
// in the picker helps avoid locking the ringhash balancer's mutex when
// reading the latest state at RPC time.
endpointStates map[string]endpointState // endpointState.hashKey -> endpointState
// requestHashHeader is the header key to look for the request hash. If it's
// empty, the request hash is expected to be set in the context via xDS.
// See gRFC A76.
requestHashHeader string
// hasEndpointInConnectingState is true if any of the endpoints is in
// CONNECTING.
hasEndpointInConnectingState bool
randUint64 func() uint64
}
func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
usingRandomHash := false
var requestHash uint64
if p.requestHashHeader == "" {
var ok bool
if requestHash, ok = iringhash.XDSRequestHash(info.Ctx); !ok {
return balancer.PickResult{}, fmt.Errorf("ringhash: expected xDS config selector to set the request hash")
}
} else {
md, ok := metadata.FromOutgoingContext(info.Ctx)
if !ok || len(md.Get(p.requestHashHeader)) == 0 {
requestHash = p.randUint64()
usingRandomHash = true
} else {
values := strings.Join(md.Get(p.requestHashHeader), ",")
requestHash = xxhash.Sum64String(values)
}
}
e := p.ring.pick(requestHash)
ringSize := len(p.ring.items)
if !usingRandomHash {
// Per gRFC A61, because of sticky-TF with PickFirst's auto reconnect on TF,
// we ignore all TF subchannels and find the first ring entry in READY,
// CONNECTING or IDLE. If that entry is in IDLE, we need to initiate a
// connection. The idlePicker returned by the LazyLB or the new Pickfirst
// should do this automatically.
for i := 0; i < ringSize; i++ {
index := (e.idx + i) % ringSize
es := p.endpointState(p.ring.items[index])
switch es.state.ConnectivityState {
case connectivity.Ready, connectivity.Connecting, connectivity.Idle:
return es.state.Picker.Pick(info)
case connectivity.TransientFailure:
default:
panic(fmt.Sprintf("Found child balancer in unknown state: %v", es.state.ConnectivityState))
}
}
} else {
// If the picker has generated a random hash, it will walk the ring from
// this hash, and pick the first READY endpoint. If no endpoint is
// currently in CONNECTING state, it will trigger a connection attempt
// on at most one endpoint that is in IDLE state along the way. - A76
requestedConnection := p.hasEndpointInConnectingState
for i := 0; i < ringSize; i++ {
index := (e.idx + i) % ringSize
es := p.endpointState(p.ring.items[index])
if es.state.ConnectivityState == connectivity.Ready {
return es.state.Picker.Pick(info)
}
if !requestedConnection && es.state.ConnectivityState == connectivity.Idle {
requestedConnection = true
// If the SubChannel is in idle state, initiate a connection but
// continue to check other pickers to see if there is one in
// ready state.
es.exitIdle()
}
}
if requestedConnection {
return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
}
}
// All children are in transient failure. Return the first failure.
return p.endpointState(e).state.Picker.Pick(info)
}
func (p *picker) endpointState(e *ringEntry) endpointState {
return p.endpointStates[e.hashKey]
}

View File

@@ -0,0 +1,301 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"context"
"errors"
"fmt"
"math"
"testing"
"time"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
iringhash "google.golang.org/grpc/internal/ringhash"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/metadata"
)
var (
testSubConns []*testutils.TestSubConn
errPicker = errors.New("picker in TransientFailure")
)
func init() {
for i := 0; i < 8; i++ {
testSubConns = append(testSubConns, testutils.NewTestSubConn(fmt.Sprint(i)))
}
}
// fakeChildPicker is used to mock pickers from child pickfirst balancers.
type fakeChildPicker struct {
connectivityState connectivity.State
subConn *testutils.TestSubConn
tfError error
}
func (p *fakeChildPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
switch p.connectivityState {
case connectivity.Idle:
p.subConn.Connect()
return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
case connectivity.Connecting:
return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
case connectivity.Ready:
return balancer.PickResult{SubConn: p.subConn}, nil
default:
return balancer.PickResult{}, p.tfError
}
}
func testRingAndEndpointStates(states []connectivity.State) (*ring, map[string]endpointState) {
var items []*ringEntry
epStates := map[string]endpointState{}
for i, st := range states {
testSC := testSubConns[i]
items = append(items, &ringEntry{
idx: i,
hash: math.MaxUint64 / uint64(len(states)) * uint64(i),
hashKey: testSC.String(),
})
epState := endpointState{
state: balancer.State{
ConnectivityState: st,
Picker: &fakeChildPicker{
connectivityState: st,
tfError: fmt.Errorf("%d: %w", i, errPicker),
subConn: testSC,
},
},
exitIdle: func() { testSC.Connect() },
}
epStates[testSC.String()] = epState
}
return &ring{items: items}, epStates
}
func (s) TestPickerPickFirstTwo(t *testing.T) {
tests := []struct {
name string
connectivityStates []connectivity.State
wantSC balancer.SubConn
wantErr error
wantSCToConnect balancer.SubConn
}{
{
name: "picked is Ready",
connectivityStates: []connectivity.State{connectivity.Ready, connectivity.Idle},
wantSC: testSubConns[0],
},
{
name: "picked is connecting, queue",
connectivityStates: []connectivity.State{connectivity.Connecting, connectivity.Idle},
wantErr: balancer.ErrNoSubConnAvailable,
},
{
name: "picked is Idle, connect and queue",
connectivityStates: []connectivity.State{connectivity.Idle, connectivity.Idle},
wantErr: balancer.ErrNoSubConnAvailable,
wantSCToConnect: testSubConns[0],
},
{
name: "picked is TransientFailure, next is ready, return",
connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Ready},
wantSC: testSubConns[1],
},
{
name: "picked is TransientFailure, next is connecting, queue",
connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Connecting},
wantErr: balancer.ErrNoSubConnAvailable,
},
{
name: "picked is TransientFailure, next is Idle, connect and queue",
connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Idle},
wantErr: balancer.ErrNoSubConnAvailable,
wantSCToConnect: testSubConns[1],
},
{
name: "all are in TransientFailure, return picked failure",
connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.TransientFailure},
wantErr: errPicker,
},
}
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ring, epStates := testRingAndEndpointStates(tt.connectivityStates)
p := &picker{
ring: ring,
endpointStates: epStates,
}
got, err := p.Pick(balancer.PickInfo{
Ctx: iringhash.SetXDSRequestHash(ctx, 0), // always pick the first endpoint on the ring.
})
if (err != nil || tt.wantErr != nil) && !errors.Is(err, tt.wantErr) {
t.Errorf("Pick() error = %v, wantErr %v", err, tt.wantErr)
return
}
if got.SubConn != tt.wantSC {
t.Errorf("Pick() got = %v, want picked SubConn: %v", got, tt.wantSC)
}
if sc := tt.wantSCToConnect; sc != nil {
select {
case <-sc.(*testutils.TestSubConn).ConnectCh:
case <-time.After(defaultTestShortTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", sc)
}
}
})
}
}
func (s) TestPickerNoRequestHash(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
ring, epStates := testRingAndEndpointStates([]connectivity.State{connectivity.Ready})
p := &picker{
ring: ring,
endpointStates: epStates,
}
if _, err := p.Pick(balancer.PickInfo{Ctx: ctx}); err == nil {
t.Errorf("Pick() should have failed with no request hash")
}
}
func (s) TestPickerRequestHashKey(t *testing.T) {
tests := []struct {
name string
headerValues []string
expectedPick int
}{
{
name: "header not set",
expectedPick: 0, // Random hash set to 0, which is within (MaxUint64 / 3 * 2, 0]
},
{
name: "header empty",
headerValues: []string{""},
expectedPick: 0, // xxhash.Sum64String("value1,value2") is within (MaxUint64 / 3 * 2, 0]
},
{
name: "header set to one value",
headerValues: []string{"some-value"},
expectedPick: 1, // xxhash.Sum64String("some-value") is within (0, MaxUint64 / 3]
},
{
name: "header set to multiple values",
headerValues: []string{"value1", "value2"},
expectedPick: 2, // xxhash.Sum64String("value1,value2") is within (MaxUint64 / 3, MaxUint64 / 3 * 2]
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
ring, epStates := testRingAndEndpointStates(
[]connectivity.State{
connectivity.Ready,
connectivity.Ready,
connectivity.Ready,
})
headerName := "some-header"
p := &picker{
ring: ring,
endpointStates: epStates,
requestHashHeader: headerName,
randUint64: func() uint64 { return 0 },
}
for _, v := range tt.headerValues {
ctx = metadata.AppendToOutgoingContext(ctx, headerName, v)
}
if res, err := p.Pick(balancer.PickInfo{Ctx: ctx}); err != nil {
t.Errorf("Pick() failed: %v", err)
} else if res.SubConn != testSubConns[tt.expectedPick] {
t.Errorf("Pick() got = %v, want SubConn: %v", res.SubConn, testSubConns[tt.expectedPick])
}
})
}
}
func (s) TestPickerRandomHash(t *testing.T) {
tests := []struct {
name string
hash uint64
connectivityStates []connectivity.State
wantSC balancer.SubConn
wantErr error
wantSCToConnect balancer.SubConn
hasEndpointInConnectingState bool
}{
{
name: "header not set, picked is Ready",
connectivityStates: []connectivity.State{connectivity.Ready, connectivity.Idle},
wantSC: testSubConns[0],
},
{
name: "header not set, picked is Idle, another is Ready. Connect and pick Ready",
connectivityStates: []connectivity.State{connectivity.Idle, connectivity.Ready},
wantSC: testSubConns[1],
wantSCToConnect: testSubConns[0],
},
{
name: "header not set, picked is Idle, there is at least one Connecting",
connectivityStates: []connectivity.State{connectivity.Connecting, connectivity.Idle},
wantErr: balancer.ErrNoSubConnAvailable,
hasEndpointInConnectingState: true,
},
{
name: "header not set, all Idle or TransientFailure, connect",
connectivityStates: []connectivity.State{connectivity.TransientFailure, connectivity.Idle},
wantErr: balancer.ErrNoSubConnAvailable,
wantSCToConnect: testSubConns[1],
},
}
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ring, epStates := testRingAndEndpointStates(tt.connectivityStates)
p := &picker{
ring: ring,
endpointStates: epStates,
requestHashHeader: "some-header",
hasEndpointInConnectingState: tt.hasEndpointInConnectingState,
randUint64: func() uint64 { return 0 }, // always return the first endpoint on the ring.
}
if got, err := p.Pick(balancer.PickInfo{Ctx: ctx}); err != tt.wantErr {
t.Errorf("Pick() error = %v, wantErr %v", err, tt.wantErr)
return
} else if got.SubConn != tt.wantSC {
t.Errorf("Pick() got = %v, want picked SubConn: %v", got, tt.wantSC)
}
if sc := tt.wantSCToConnect; sc != nil {
select {
case <-sc.(*testutils.TestSubConn).ConnectCh:
case <-time.After(defaultTestShortTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", sc)
}
}
})
}
}

189
balancer/ringhash/ring.go Normal file
View File

@@ -0,0 +1,189 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"math"
"sort"
"strconv"
xxhash "github.com/cespare/xxhash/v2"
"google.golang.org/grpc/internal/grpclog"
"google.golang.org/grpc/resolver"
)
type ring struct {
items []*ringEntry
}
type endpointInfo struct {
hashKey string
scaledWeight float64
originalWeight uint32
}
type ringEntry struct {
idx int
hash uint64
hashKey string
weight uint32
}
// newRing creates a ring from the endpoints stored in the EndpointMap. The ring
// size is limited by the passed in max/min.
//
// ring entries will be created for each endpoint, and endpoints with high
// weight (specified by the endpoint) may have multiple entries.
//
// For example, for endpoints with weights {a:3, b:3, c:4}, a generated ring of
// size 10 could be:
// - {idx:0 hash:3689675255460411075 b}
// - {idx:1 hash:4262906501694543955 c}
// - {idx:2 hash:5712155492001633497 c}
// - {idx:3 hash:8050519350657643659 b}
// - {idx:4 hash:8723022065838381142 b}
// - {idx:5 hash:11532782514799973195 a}
// - {idx:6 hash:13157034721563383607 c}
// - {idx:7 hash:14468677667651225770 c}
// - {idx:8 hash:17336016884672388720 a}
// - {idx:9 hash:18151002094784932496 a}
//
// To pick from a ring, a binary search will be done for the given target hash,
// and first item with hash >= given hash will be returned.
//
// Must be called with a non-empty endpoints map.
func newRing(endpoints *resolver.EndpointMap[*endpointState], minRingSize, maxRingSize uint64, logger *grpclog.PrefixLogger) *ring {
if logger.V(2) {
logger.Infof("newRing: number of endpoints is %d, minRingSize is %d, maxRingSize is %d", endpoints.Len(), minRingSize, maxRingSize)
}
// https://github.com/envoyproxy/envoy/blob/765c970f06a4c962961a0e03a467e165b276d50f/source/common/upstream/ring_hash_lb.cc#L114
normalizedWeights, minWeight := normalizeWeights(endpoints)
if logger.V(2) {
logger.Infof("newRing: normalized endpoint weights is %v", normalizedWeights)
}
// Normalized weights for {3,3,4} is {0.3,0.3,0.4}.
// Scale up the size of the ring such that the least-weighted host gets a
// whole number of hashes on the ring.
//
// Note that size is limited by the input max/min.
scale := math.Min(math.Ceil(minWeight*float64(minRingSize))/minWeight, float64(maxRingSize))
ringSize := math.Ceil(scale)
items := make([]*ringEntry, 0, int(ringSize))
if logger.V(2) {
logger.Infof("newRing: creating new ring of size %v", ringSize)
}
// For each entry, scale*weight nodes are generated in the ring.
//
// Not all of these are whole numbers. E.g. for weights {a:3,b:3,c:4}, if
// ring size is 7, scale is 6.66. The numbers of nodes will be
// {a,a,b,b,c,c,c}.
//
// A hash is generated for each item, and later the results will be sorted
// based on the hash.
var currentHashes, targetHashes float64
for _, epInfo := range normalizedWeights {
targetHashes += scale * epInfo.scaledWeight
// This index ensures that ring entries corresponding to the same
// endpoint hash to different values. And since this index is
// per-endpoint, these entries hash to the same value across address
// updates.
idx := 0
for currentHashes < targetHashes {
h := xxhash.Sum64String(epInfo.hashKey + "_" + strconv.Itoa(idx))
items = append(items, &ringEntry{hash: h, hashKey: epInfo.hashKey, weight: epInfo.originalWeight})
idx++
currentHashes++
}
}
// Sort items based on hash, to prepare for binary search.
sort.Slice(items, func(i, j int) bool { return items[i].hash < items[j].hash })
for i, ii := range items {
ii.idx = i
}
return &ring{items: items}
}
// normalizeWeights calculates the normalized weights for each endpoint in the
// given endpoints map. It returns a slice of endpointWithState structs, where
// each struct contains the picker for an endpoint and its corresponding weight.
// The function also returns the minimum weight among all endpoints.
//
// The normalized weight of each endpoint is calculated by dividing its weight
// attribute by the sum of all endpoint weights. If the weight attribute is not
// found on the endpoint, a default weight of 1 is used.
//
// The endpoints are sorted in ascending order to ensure consistent results.
//
// Must be called with a non-empty endpoints map.
func normalizeWeights(endpoints *resolver.EndpointMap[*endpointState]) ([]endpointInfo, float64) {
var weightSum uint32
// Since attributes are explicitly ignored in the EndpointMap key, we need
// to iterate over the values to get the weights.
endpointVals := endpoints.Values()
for _, epState := range endpointVals {
weightSum += epState.weight
}
ret := make([]endpointInfo, 0, endpoints.Len())
min := 1.0
for _, epState := range endpointVals {
// (*endpointState).weight is set to 1 if the weight attribute is not
// found on the endpoint. And since this function is guaranteed to be
// called with a non-empty endpoints map, weightSum is guaranteed to be
// non-zero. So, we need not worry about divide by zero error here.
nw := float64(epState.weight) / float64(weightSum)
ret = append(ret, endpointInfo{
hashKey: epState.hashKey,
scaledWeight: nw,
originalWeight: epState.weight,
})
min = math.Min(min, nw)
}
// Sort the endpoints to return consistent results.
//
// Note: this might not be necessary, but this makes sure the ring is
// consistent as long as the endpoints are the same, for example, in cases
// where an endpoint is added and then removed, the RPCs will still pick the
// same old endpoint.
sort.Slice(ret, func(i, j int) bool {
return ret[i].hashKey < ret[j].hashKey
})
return ret, min
}
// pick does a binary search. It returns the item with smallest index i that
// r.items[i].hash >= h.
func (r *ring) pick(h uint64) *ringEntry {
i := sort.Search(len(r.items), func(i int) bool { return r.items[i].hash >= h })
if i == len(r.items) {
// If not found, and h is greater than the largest hash, return the
// first item.
i = 0
}
return r.items[i]
}
// next returns the next entry.
func (r *ring) next(e *ringEntry) *ringEntry {
return r.items[(e.idx+1)%len(r.items)]
}

View File

@@ -0,0 +1,112 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"fmt"
"math"
"testing"
xxhash "github.com/cespare/xxhash/v2"
"google.golang.org/grpc/experimental/balancer/weight"
"google.golang.org/grpc/resolver"
)
var testEndpoints []resolver.Endpoint
var testEndpointStateMap *resolver.EndpointMap[*endpointState]
func init() {
testEndpoints = []resolver.Endpoint{
testEndpoint("a", 3),
testEndpoint("b", 3),
testEndpoint("c", 4),
}
testEndpointStateMap = resolver.NewEndpointMap[*endpointState]()
testEndpointStateMap.Set(testEndpoints[0], &endpointState{hashKey: "a", weight: 3})
testEndpointStateMap.Set(testEndpoints[1], &endpointState{hashKey: "b", weight: 3})
testEndpointStateMap.Set(testEndpoints[2], &endpointState{hashKey: "c", weight: 4})
}
func testEndpoint(addr string, endpointWeight uint32) resolver.Endpoint {
ep := resolver.Endpoint{Addresses: []resolver.Address{{Addr: addr}}}
return weight.Set(ep, weight.EndpointInfo{Weight: endpointWeight})
}
func (s) TestRingNew(t *testing.T) {
var totalWeight float64 = 10
for _, min := range []uint64{3, 4, 6, 8} {
for _, max := range []uint64{20, 8} {
t.Run(fmt.Sprintf("size-min-%v-max-%v", min, max), func(t *testing.T) {
r := newRing(testEndpointStateMap, min, max, nil)
totalCount := len(r.items)
if totalCount < int(min) || totalCount > int(max) {
t.Fatalf("unexpected size %v, want min %v, max %v", totalCount, min, max)
}
for _, e := range testEndpoints {
var count int
for _, ii := range r.items {
if ii.hashKey == hashKey(e) {
count++
}
}
got := float64(count) / float64(totalCount)
want := float64(getWeightAttribute(e)) / totalWeight
if !equalApproximately(got, want) {
t.Fatalf("unexpected item weight in ring: %v != %v", got, want)
}
}
})
}
}
}
func equalApproximately(x, y float64) bool {
delta := math.Abs(x - y)
mean := math.Abs(x+y) / 2.0
return delta/mean < 0.25
}
func (s) TestRingPick(t *testing.T) {
r := newRing(testEndpointStateMap, 10, 20, nil)
for _, h := range []uint64{xxhash.Sum64String("1"), xxhash.Sum64String("2"), xxhash.Sum64String("3"), xxhash.Sum64String("4")} {
t.Run(fmt.Sprintf("picking-hash-%v", h), func(t *testing.T) {
e := r.pick(h)
var low uint64
if e.idx > 0 {
low = r.items[e.idx-1].hash
}
high := e.hash
// h should be in [low, high).
if h < low || h >= high {
t.Fatalf("unexpected item picked, hash: %v, low: %v, high: %v", h, low, high)
}
})
}
}
func (s) TestRingNext(t *testing.T) {
r := newRing(testEndpointStateMap, 10, 20, nil)
for _, e := range r.items {
ne := r.next(e)
if ne.idx != (e.idx+1)%len(r.items) {
t.Fatalf("next(%+v) returned unexpected %+v", e, ne)
}
}
}

View File

@@ -0,0 +1,410 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
// Package ringhash implements the ringhash balancer. See the following
// gRFCs for details:
// - https://github.com/grpc/proposal/blob/master/A42-xds-ring-hash-lb-policy.md
// - https://github.com/grpc/proposal/blob/master/A61-IPv4-IPv6-dualstack-backends.md#ring-hash
// - https://github.com/grpc/proposal/blob/master/A76-ring-hash-improvements.md
//
// # Experimental
//
// Notice: This package is EXPERIMENTAL and may be changed or removed in a
// later release.
package ringhash
import (
"encoding/json"
"errors"
"fmt"
"math/rand/v2"
"sort"
"sync"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/base"
"google.golang.org/grpc/balancer/endpointsharding"
"google.golang.org/grpc/balancer/lazy"
"google.golang.org/grpc/balancer/pickfirst"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/experimental/balancer/weight"
"google.golang.org/grpc/internal/grpclog"
"google.golang.org/grpc/internal/pretty"
iringhash "google.golang.org/grpc/internal/ringhash"
"google.golang.org/grpc/resolver"
"google.golang.org/grpc/resolver/ringhash"
"google.golang.org/grpc/serviceconfig"
)
// Name is the name of the ring_hash balancer.
const Name = "ring_hash_experimental"
func lazyPickFirstBuilder(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
return lazy.NewBalancer(cc, opts, balancer.Get(pickfirst.Name).Build)
}
func init() {
balancer.Register(bb{})
}
type bb struct{}
func (bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
b := &ringhashBalancer{
ClientConn: cc,
endpointStates: resolver.NewEndpointMap[*endpointState](),
}
esOpts := endpointsharding.Options{DisableAutoReconnect: true}
b.child = endpointsharding.NewBalancer(b, opts, lazyPickFirstBuilder, esOpts)
b.logger = prefixLogger(b)
b.logger.Infof("Created")
return b
}
func (bb) Name() string {
return Name
}
func (bb) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
return parseConfig(c)
}
type ringhashBalancer struct {
// The following fields are initialized at build time and read-only after
// that and therefore do not need to be guarded by a mutex.
// ClientConn is embedded to intercept UpdateState calls from the child
// endpointsharding balancer.
balancer.ClientConn
logger *grpclog.PrefixLogger
child balancer.Balancer
mu sync.Mutex
config *iringhash.LBConfig
inhibitChildUpdates bool
shouldRegenerateRing bool
endpointStates *resolver.EndpointMap[*endpointState]
// ring is always in sync with endpoints. When endpoints change, a new ring
// is generated. Note that address weights updates also regenerates the
// ring.
ring *ring
}
// hashKey returns the hash key to use for an endpoint. Per gRFC A61, each entry
// in the ring is a hash of the endpoint's hash key concatenated with a
// per-entry unique suffix.
func hashKey(endpoint resolver.Endpoint) string {
if hk := ringhash.HashKey(endpoint); hk != "" {
return hk
}
// If no hash key is set, use the endpoint's first address as the hash key.
// This is the default behavior when no hash key is set.
return endpoint.Addresses[0].Addr
}
// UpdateState intercepts child balancer state updates. It updates the
// per-endpoint state stored in the ring, and also the aggregated state based on
// the child picker. It also reconciles the endpoint list. It sets
// `b.shouldRegenerateRing` to true if the new endpoint list is different from
// the previous, i.e. any of the following is true:
// - an endpoint was added
// - an endpoint was removed
// - an endpoint's weight was updated
// - the first addresses of the endpoint has changed
func (b *ringhashBalancer) UpdateState(state balancer.State) {
b.mu.Lock()
defer b.mu.Unlock()
childStates := endpointsharding.ChildStatesFromPicker(state.Picker)
// endpointsSet is the set converted from endpoints, used for quick lookup.
endpointsSet := resolver.NewEndpointMap[bool]()
for _, childState := range childStates {
endpoint := childState.Endpoint
endpointsSet.Set(endpoint, true)
newWeight := getWeightAttribute(endpoint)
hk := hashKey(endpoint)
es, ok := b.endpointStates.Get(endpoint)
if !ok {
es := &endpointState{
hashKey: hk,
weight: newWeight,
state: childState.State,
exitIdle: childState.ExitIdle,
}
b.endpointStates.Set(endpoint, es)
b.shouldRegenerateRing = true
} else {
// We have seen this endpoint before and created a `endpointState`
// object for it. If the weight or the hash key of the endpoint has
// changed, update the endpoint state map with the new weight or
// hash key. This will be used when a new ring is created.
if oldWeight := es.weight; oldWeight != newWeight {
b.shouldRegenerateRing = true
es.weight = newWeight
}
if es.hashKey != hk {
b.shouldRegenerateRing = true
es.hashKey = hk
}
es.state = childState.State
}
}
for endpoint := range b.endpointStates.All() {
if _, ok := endpointsSet.Get(endpoint); ok {
continue
}
// endpoint was removed by resolver.
b.endpointStates.Delete(endpoint)
b.shouldRegenerateRing = true
}
b.updatePickerLocked()
}
func (b *ringhashBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
if b.logger.V(2) {
b.logger.Infof("Received update from resolver, balancer config: %+v", pretty.ToJSON(ccs.BalancerConfig))
}
newConfig, ok := ccs.BalancerConfig.(*iringhash.LBConfig)
if !ok {
return fmt.Errorf("unexpected balancer config with type: %T", ccs.BalancerConfig)
}
b.mu.Lock()
b.inhibitChildUpdates = true
b.mu.Unlock()
defer func() {
b.mu.Lock()
b.inhibitChildUpdates = false
b.updatePickerLocked()
b.mu.Unlock()
}()
if err := b.child.UpdateClientConnState(balancer.ClientConnState{
// Make pickfirst children use health listeners for outlier detection
// and health checking to work.
ResolverState: pickfirst.EnableHealthListener(ccs.ResolverState),
}); err != nil {
return err
}
b.mu.Lock()
// Ring updates can happen due to the following:
// 1. Addition or deletion of endpoints: The synchronous picker update from
// the child endpointsharding balancer would contain the list of updated
// endpoints. Updates triggered by the child after handling the
// `UpdateClientConnState` call will not change the endpoint list.
// 2. Change in the `LoadBalancerConfig`: Ring config such as max/min ring
// size.
// To avoid extra ring updates, a boolean is used to track the need for a
// ring update and the update is done only once at the end.
//
// If the ring configuration has changed, we need to regenerate the ring
// while sending a new picker.
if b.config == nil || b.config.MinRingSize != newConfig.MinRingSize || b.config.MaxRingSize != newConfig.MaxRingSize {
b.shouldRegenerateRing = true
}
b.config = newConfig
b.mu.Unlock()
return nil
}
func (b *ringhashBalancer) ResolverError(err error) {
b.child.ResolverError(err)
}
func (b *ringhashBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
}
func (b *ringhashBalancer) updatePickerLocked() {
state := b.aggregatedStateLocked()
// Start connecting to new endpoints if necessary.
if state == connectivity.Connecting || state == connectivity.TransientFailure {
// When overall state is TransientFailure, we need to make sure at least
// one endpoint is attempting to connect, otherwise this balancer may
// never get picks if the parent is priority.
//
// Because we report Connecting as the overall state when only one
// endpoint is in TransientFailure, we do the same check for Connecting
// here.
//
// Note that this check also covers deleting endpoints. E.g. if the
// endpoint attempting to connect is deleted, and the overall state is
// TF. Since there must be at least one endpoint attempting to connect,
// we need to trigger one.
//
// After calling `ExitIdle` on a child balancer, the child will send a
// picker update asynchronously. A race condition may occur if another
// picker update from endpointsharding arrives before the child's
// picker update. The received picker may trigger a re-execution of the
// loop below to find an idle child. Since map iteration order is
// non-deterministic, the list of `endpointState`s must be sorted to
// ensure `ExitIdle` is called on the same child, preventing unnecessary
// connections.
endpointStates := make([]*endpointState, 0, b.endpointStates.Len())
for _, s := range b.endpointStates.All() {
endpointStates = append(endpointStates, s)
}
sort.Slice(endpointStates, func(i, j int) bool {
return endpointStates[i].hashKey < endpointStates[j].hashKey
})
// Store the function to ExitIdle on the first IDLE endpoint.
var exitIdle func()
for _, es := range endpointStates {
connState := es.state.ConnectivityState
if connState == connectivity.Connecting {
exitIdle = nil
break
}
if exitIdle == nil && connState == connectivity.Idle {
exitIdle = es.exitIdle
}
}
if exitIdle != nil {
exitIdle()
}
}
if b.inhibitChildUpdates {
return
}
// Update the channel.
if b.endpointStates.Len() > 0 && b.shouldRegenerateRing {
// with a non-empty list of endpoints.
b.ring = newRing(b.endpointStates, b.config.MinRingSize, b.config.MaxRingSize, b.logger)
}
b.shouldRegenerateRing = false
var newPicker balancer.Picker
if b.endpointStates.Len() == 0 {
newPicker = base.NewErrPicker(errors.New("produced zero addresses"))
} else {
newPicker = b.newPickerLocked()
}
b.ClientConn.UpdateState(balancer.State{
ConnectivityState: state,
Picker: newPicker,
})
}
func (b *ringhashBalancer) Close() {
b.logger.Infof("Shutdown")
b.child.Close()
}
func (b *ringhashBalancer) ExitIdle() {
// ExitIdle implementation is a no-op because connections are either
// triggers from picks or from child balancer state changes.
}
// newPickerLocked generates a picker. The picker copies the endpoint states
// over to avoid locking the mutex at RPC time. The picker should be
// re-generated every time an endpoint state is updated.
func (b *ringhashBalancer) newPickerLocked() *picker {
states := make(map[string]endpointState)
hasEndpointConnecting := false
for _, epState := range b.endpointStates.All() {
// Copy the endpoint state to avoid races, since ring hash
// mutates the state, weight and hash key in place.
states[epState.hashKey] = *epState
if epState.state.ConnectivityState == connectivity.Connecting {
hasEndpointConnecting = true
}
}
return &picker{
ring: b.ring,
endpointStates: states,
requestHashHeader: b.config.RequestHashHeader,
hasEndpointInConnectingState: hasEndpointConnecting,
randUint64: rand.Uint64,
}
}
// aggregatedStateLocked returns the aggregated child balancers state
// based on the following rules.
// - If there is at least one endpoint in READY state, report READY.
// - If there are 2 or more endpoints in TRANSIENT_FAILURE state, report
// TRANSIENT_FAILURE.
// - If there is at least one endpoint in CONNECTING state, report CONNECTING.
// - If there is one endpoint in TRANSIENT_FAILURE and there is more than one
// endpoint, report state CONNECTING.
// - If there is at least one endpoint in Idle state, report Idle.
// - Otherwise, report TRANSIENT_FAILURE.
//
// Note that if there are 1 connecting, 2 transient failure, the overall state
// is transient failure. This is because the second transient failure is a
// fallback of the first failing endpoint, and we want to report transient
// failure to failover to the lower priority.
func (b *ringhashBalancer) aggregatedStateLocked() connectivity.State {
var nums [5]int
for _, es := range b.endpointStates.All() {
nums[es.state.ConnectivityState]++
}
if nums[connectivity.Ready] > 0 {
return connectivity.Ready
}
if nums[connectivity.TransientFailure] > 1 {
return connectivity.TransientFailure
}
if nums[connectivity.Connecting] > 0 {
return connectivity.Connecting
}
if nums[connectivity.TransientFailure] == 1 && b.endpointStates.Len() > 1 {
return connectivity.Connecting
}
if nums[connectivity.Idle] > 0 {
return connectivity.Idle
}
return connectivity.TransientFailure
}
// getWeightAttribute is a convenience function which returns the value of the
// weight endpoint Attribute.
//
// When used in the xDS context, the weight attribute is guaranteed to be
// non-zero. But, when used in a non-xDS context, the weight attribute could be
// unset. A Default of 1 is used in the latter case.
func getWeightAttribute(e resolver.Endpoint) uint32 {
w := weight.FromEndpoint(e).Weight
if w == 0 {
return 1
}
return w
}
type endpointState struct {
// hashKey is the hash key of the endpoint. Per gRFC A61, each entry in the
// ring is an endpoint, positioned based on the hash of the endpoint's first
// address by default. Per gRFC A76, the hash key of an endpoint may be
// overridden, for example based on EDS endpoint metadata.
hashKey string
weight uint32
exitIdle func()
// state is updated by the balancer while receiving resolver updates from
// the channel and picker updates from its children. Access to it is guarded
// by ringhashBalancer.mu.
state balancer.State
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,739 @@
/*
*
* Copyright 2021 gRPC authors.
*
* 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.
*
*/
package ringhash
import (
"context"
"fmt"
"testing"
"time"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/connectivity"
"google.golang.org/grpc/experimental/balancer/weight"
"google.golang.org/grpc/internal/grpctest"
iringhash "google.golang.org/grpc/internal/ringhash"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/resolver"
)
const (
defaultTestTimeout = 10 * time.Second
defaultTestShortTimeout = 10 * time.Millisecond
testBackendAddrsCount = 12
)
var (
testBackendAddrStrs []string
testConfig = &iringhash.LBConfig{MinRingSize: 1, MaxRingSize: 10}
)
func init() {
for i := 0; i < testBackendAddrsCount; i++ {
testBackendAddrStrs = append(testBackendAddrStrs, fmt.Sprintf("%d.%d.%d.%d:%d", i, i, i, i, i))
}
}
// setupTest creates the balancer, and does an initial sanity check.
func setupTest(t *testing.T, endpoints []resolver.Endpoint) (*testutils.BalancerClientConn, balancer.Balancer, balancer.Picker) {
t.Helper()
cc := testutils.NewBalancerClientConn(t)
builder := balancer.Get(Name)
b := builder.Build(cc, balancer.BuildOptions{})
if b == nil {
t.Fatalf("builder.Build(%s) failed and returned nil", Name)
}
t.Cleanup(func() { b.Close() })
if err := b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{Endpoints: endpoints},
BalancerConfig: testConfig,
}); err != nil {
t.Fatalf("UpdateClientConnState returned err: %v", err)
}
// The leaf pickfirst are created lazily, only when their endpoint is picked
// or other endpoints are in TF. No SubConns should be created immediately.
select {
case sc := <-cc.NewSubConnCh:
t.Errorf("unexpected SubConn creation: %v", sc)
case <-time.After(defaultTestShortTimeout):
}
// Should also have a picker, with all endpoints in Idle.
p1 := <-cc.NewPickerCh
ringHashPicker := p1.(*picker)
if got, want := len(ringHashPicker.endpointStates), len(endpoints); got != want {
t.Errorf("Number of child balancers = %d, want = %d", got, want)
}
for firstAddr, bs := range ringHashPicker.endpointStates {
if got, want := bs.state.ConnectivityState, connectivity.Idle; got != want {
t.Errorf("Child balancer connectivity state for address %q = %v, want = %v", firstAddr, got, want)
}
}
return cc, b, p1
}
type s struct {
grpctest.Tester
}
func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}
// TestUpdateClientConnState_NewRingSize tests the scenario where the ringhash
// LB policy receives new configuration which specifies new values for the ring
// min and max sizes. The test verifies that a new ring is created and a new
// picker is sent to the ClientConn.
func (s) TestUpdateClientConnState_NewRingSize(t *testing.T) {
origMinRingSize, origMaxRingSize := 1, 10 // Configured from `testConfig` in `setupTest`
newMinRingSize, newMaxRingSize := 20, 100
endpoints := []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}}}
cc, b, p1 := setupTest(t, endpoints)
ring1 := p1.(*picker).ring
if ringSize := len(ring1.items); ringSize < origMinRingSize || ringSize > origMaxRingSize {
t.Fatalf("Ring created with size %d, want between [%d, %d]", ringSize, origMinRingSize, origMaxRingSize)
}
if err := b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{Endpoints: endpoints},
BalancerConfig: &iringhash.LBConfig{
MinRingSize: uint64(newMinRingSize),
MaxRingSize: uint64(newMaxRingSize),
},
}); err != nil {
t.Fatalf("UpdateClientConnState returned err: %v", err)
}
var ring2 *ring
select {
case <-time.After(defaultTestTimeout):
t.Fatal("Timeout when waiting for a picker update after a configuration update")
case p2 := <-cc.NewPickerCh:
ring2 = p2.(*picker).ring
}
if ringSize := len(ring2.items); ringSize < newMinRingSize || ringSize > newMaxRingSize {
t.Fatalf("Ring created with size %d, want between [%d, %d]", ringSize, newMinRingSize, newMaxRingSize)
}
}
func (s) TestOneEndpoint(t *testing.T) {
wantAddr1 := resolver.Address{Addr: testBackendAddrStrs[0]}
cc, _, p0 := setupTest(t, []resolver.Endpoint{{Addresses: []resolver.Address{wantAddr1}}})
ring0 := p0.(*picker).ring
firstHash := ring0.items[0].hash
// firstHash-1 will pick the first (and only) SubConn from the ring.
testHash := firstHash - 1
// The first pick should be queued, and should trigger a connection to the
// only Endpoint which has a single address.
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)}); err != balancer.ErrNoSubConnAvailable {
t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
}
var sc0 *testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc0 = <-cc.NewSubConnCh:
}
if got, want := sc0.Addresses[0].Addr, wantAddr1.Addr; got != want {
t.Fatalf("SubConn.Addresses = %v, want = %v", got, want)
}
select {
case <-sc0.ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
}
// Send state updates to Ready.
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
if err := cc.WaitForConnectivityState(ctx, connectivity.Ready); err != nil {
t.Fatal(err)
}
// Test pick with one backend.
p1 := <-cc.NewPickerCh
for i := 0; i < 5; i++ {
gotSCSt, _ := p1.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
if gotSCSt.SubConn != sc0 {
t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc0)
}
}
}
// TestThreeBackendsAffinity covers that there are 3 SubConns, RPCs with the
// same hash always pick the same SubConn. When the one picked is down, another
// one will be picked.
func (s) TestThreeSubConnsAffinity(t *testing.T) {
endpoints := []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
}
remainingAddrs := map[string]bool{
testBackendAddrStrs[0]: true,
testBackendAddrStrs[1]: true,
testBackendAddrStrs[2]: true,
}
cc, _, p0 := setupTest(t, endpoints)
// This test doesn't update addresses, so this ring will be used by all the
// pickers.
ring := p0.(*picker).ring
firstHash := ring.items[0].hash
// firstHash+1 will pick the second endpoint from the ring.
testHash := firstHash + 1
// The first pick should be queued, and should trigger Connect() on the only
// SubConn.
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)}); err != balancer.ErrNoSubConnAvailable {
t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
}
// The picked endpoint should be the second in the ring.
var subConns [3]*testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case subConns[1] = <-cc.NewSubConnCh:
}
if got, want := subConns[1].Addresses[0].Addr, ring.items[1].hashKey; got != want {
t.Fatalf("SubConn.Address = %v, want = %v", got, want)
}
select {
case <-subConns[1].ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", subConns[1])
}
delete(remainingAddrs, ring.items[1].hashKey)
// Turn down the subConn in use.
subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
// This should trigger a connection to a new endpoint.
<-cc.NewPickerCh
var sc *testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc = <-cc.NewSubConnCh:
}
scAddr := sc.Addresses[0].Addr
if _, ok := remainingAddrs[scAddr]; !ok {
t.Fatalf("New SubConn created with previously used address: %q", scAddr)
}
delete(remainingAddrs, scAddr)
select {
case <-sc.ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", subConns[1])
}
if scAddr == ring.items[0].hashKey {
subConns[0] = sc
} else if scAddr == ring.items[2].hashKey {
subConns[2] = sc
}
// Turning down the SubConn should cause creation of a connection to the
// final endpoint.
sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc = <-cc.NewSubConnCh:
}
scAddr = sc.Addresses[0].Addr
if _, ok := remainingAddrs[scAddr]; !ok {
t.Fatalf("New SubConn created with previously used address: %q", scAddr)
}
delete(remainingAddrs, scAddr)
select {
case <-sc.ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", subConns[1])
}
if scAddr == ring.items[0].hashKey {
subConns[0] = sc
} else if scAddr == ring.items[2].hashKey {
subConns[2] = sc
}
sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
// All endpoints are in TransientFailure. Make the first endpoint in the
// ring report Ready. All picks should go to this endpoint which is two
// indexes away from the endpoint with the chosen hash.
subConns[0].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
subConns[0].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
subConns[0].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
if err := cc.WaitForConnectivityState(ctx, connectivity.Ready); err != nil {
t.Fatalf("Context timed out while waiting for channel to report Ready.")
}
p1 := <-cc.NewPickerCh
for i := 0; i < 5; i++ {
gotSCSt, _ := p1.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
if gotSCSt.SubConn != subConns[0] {
t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, subConns[0])
}
}
// Make the last endpoint in the ring report Ready. All picks should go to
// this endpoint since it is one index away from the chosen hash.
subConns[2].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
subConns[2].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
subConns[2].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
p2 := <-cc.NewPickerCh
for i := 0; i < 5; i++ {
gotSCSt, _ := p2.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
if gotSCSt.SubConn != subConns[2] {
t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, subConns[2])
}
}
// Make the second endpoint in the ring report Ready. All picks should go to
// this endpoint as it is the one with the chosen hash.
subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
subConns[1].UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
p3 := <-cc.NewPickerCh
for i := 0; i < 5; i++ {
gotSCSt, _ := p3.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
if gotSCSt.SubConn != subConns[1] {
t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, subConns[1])
}
}
}
// TestThreeBackendsAffinity covers that there are 3 SubConns, RPCs with the
// same hash always pick the same SubConn. Then try different hash to pick
// another backend, and verify the first hash still picks the first backend.
func (s) TestThreeBackendsAffinityMultiple(t *testing.T) {
wantEndpoints := []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
}
cc, _, p0 := setupTest(t, wantEndpoints)
// This test doesn't update addresses, so this ring will be used by all the
// pickers.
ring0 := p0.(*picker).ring
firstHash := ring0.items[0].hash
// firstHash+1 will pick the second SubConn from the ring.
testHash := firstHash + 1
// The first pick should be queued, and should trigger Connect() on the only
// SubConn.
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)}); err != balancer.ErrNoSubConnAvailable {
t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
}
// The picked SubConn should be the second in the ring.
var sc0 *testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc0 = <-cc.NewSubConnCh:
}
if got, want := sc0.Addresses[0].Addr, ring0.items[1].hashKey; got != want {
t.Fatalf("SubConn.Address = %v, want = %v", got, want)
}
select {
case <-sc0.ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
}
// Send state updates to Ready.
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
if err := cc.WaitForConnectivityState(ctx, connectivity.Ready); err != nil {
t.Fatal(err)
}
// First hash should always pick sc0.
p1 := <-cc.NewPickerCh
for i := 0; i < 5; i++ {
gotSCSt, _ := p1.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
if gotSCSt.SubConn != sc0 {
t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc0)
}
}
secondHash := ring0.items[1].hash
// secondHash+1 will pick the third SubConn from the ring.
testHash2 := secondHash + 1
if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash2)}); err != balancer.ErrNoSubConnAvailable {
t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
}
var sc1 *testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc1 = <-cc.NewSubConnCh:
}
if got, want := sc1.Addresses[0].Addr, ring0.items[2].hashKey; got != want {
t.Fatalf("SubConn.Address = %v, want = %v", got, want)
}
select {
case <-sc1.ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", sc1)
}
sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
// With the new generated picker, hash2 always picks sc1.
p2 := <-cc.NewPickerCh
for i := 0; i < 5; i++ {
gotSCSt, _ := p2.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash2)})
if gotSCSt.SubConn != sc1 {
t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc1)
}
}
// But the first hash still picks sc0.
for i := 0; i < 5; i++ {
gotSCSt, _ := p2.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, testHash)})
if gotSCSt.SubConn != sc0 {
t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc0)
}
}
}
// TestAddrWeightChange covers the following scenarios after setting up the
// balancer with 3 addresses [A, B, C]:
// - updates balancer with [A, B, C], a new Picker should not be sent.
// - updates balancer with [A, B] (C removed), a new Picker is sent and the
// ring is updated.
// - updates balancer with [A, B], but B has a weight of 2, a new Picker is
// sent. And the new ring should contain the correct number of entries
// and weights.
func (s) TestAddrWeightChange(t *testing.T) {
endpoints := []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
}
cc, b, p0 := setupTest(t, endpoints)
ring0 := p0.(*picker).ring
// Update with the same addresses, it will result in a new picker, but with
// the same ring.
if err := b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{Endpoints: endpoints},
BalancerConfig: testConfig,
}); err != nil {
t.Fatalf("UpdateClientConnState returned err: %v", err)
}
var p1 balancer.Picker
select {
case p1 = <-cc.NewPickerCh:
case <-time.After(defaultTestTimeout):
t.Fatalf("timeout waiting for picker after UpdateClientConn with same addresses")
}
ring1 := p1.(*picker).ring
if ring1 != ring0 {
t.Fatalf("new picker with same address has a different ring than before, want same")
}
// Delete an address, should send a new Picker.
if err := b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{Endpoints: endpoints[:2]},
BalancerConfig: testConfig,
}); err != nil {
t.Fatalf("UpdateClientConnState returned err: %v", err)
}
var p2 balancer.Picker
select {
case p2 = <-cc.NewPickerCh:
case <-time.After(defaultTestTimeout):
t.Fatalf("timeout waiting for picker after UpdateClientConn with different addresses")
}
ring2 := p2.(*picker).ring
if ring2 == ring0 {
t.Fatalf("new picker after removing address has the same ring as before, want different")
}
// Another update with the same addresses, but different weight.
if err := b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{Endpoints: []resolver.Endpoint{
endpoints[0],
weight.Set(endpoints[1], weight.EndpointInfo{Weight: 2}),
}},
BalancerConfig: testConfig,
}); err != nil {
t.Fatalf("UpdateClientConnState returned err: %v", err)
}
var p3 balancer.Picker
select {
case p3 = <-cc.NewPickerCh:
case <-time.After(defaultTestTimeout):
t.Fatalf("timeout waiting for picker after UpdateClientConn with different addresses")
}
if p3.(*picker).ring == ring2 {
t.Fatalf("new picker after changing address weight has the same ring as before, want different")
}
// With the new update, the ring must look like this:
// [
// {idx:0 endpoint: {addr: testBackendAddrStrs[0], weight: 1}},
// {idx:1 endpoint: {addr: testBackendAddrStrs[1], weight: 2}},
// {idx:2 endpoint: {addr: testBackendAddrStrs[2], weight: 1}},
// ].
if len(p3.(*picker).ring.items) != 3 {
t.Fatalf("new picker after changing address weight has %d entries, want 3", len(p3.(*picker).ring.items))
}
for _, i := range p3.(*picker).ring.items {
if i.hashKey == testBackendAddrStrs[0] {
if i.weight != 1 {
t.Fatalf("new picker after changing address weight has weight %d for %v, want 1", i.weight, i.hashKey)
}
}
if i.hashKey == testBackendAddrStrs[1] {
if i.weight != 2 {
t.Fatalf("new picker after changing address weight has weight %d for %v, want 2", i.weight, i.hashKey)
}
}
}
}
// TestAutoConnectEndpointOnTransientFailure covers the situation when an
// endpoint fails. It verifies that a new endpoint is automatically tried
// (without a pick) when there is no endpoint already in Connecting state.
func (s) TestAutoConnectEndpointOnTransientFailure(t *testing.T) {
wantEndpoints := []resolver.Endpoint{
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[0]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[1]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[2]}}},
{Addresses: []resolver.Address{{Addr: testBackendAddrStrs[3]}}},
}
cc, _, p0 := setupTest(t, wantEndpoints)
// ringhash won't tell SCs to connect until there is an RPC, so simulate
// one now.
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
ctx = iringhash.SetXDSRequestHash(ctx, 0)
defer cancel()
p0.Pick(balancer.PickInfo{Ctx: ctx})
// The picked SubConn should be the second in the ring.
var sc0 *testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc0 = <-cc.NewSubConnCh:
}
select {
case <-sc0.ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
}
// Turn the first subconn to transient failure. This should set the overall
// connectivity state to CONNECTING.
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
cc.WaitForConnectivityState(ctx, connectivity.Connecting)
// It will trigger the second subconn to connect since there is only one
// endpoint, which is in TF.
var sc1 *testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc1 = <-cc.NewSubConnCh:
}
select {
case <-sc1.ConnectCh:
case <-time.After(defaultTestShortTimeout):
t.Fatalf("timeout waiting for Connect() from SubConn %v", sc1)
}
// Turn the second subconn to TF. This will set the overall state to TF.
sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
cc.WaitForConnectivityState(ctx, connectivity.TransientFailure)
// It will trigger the third subconn to connect.
var sc2 *testutils.TestSubConn
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case sc2 = <-cc.NewSubConnCh:
}
select {
case <-sc2.ConnectCh:
case <-time.After(defaultTestShortTimeout):
t.Fatalf("timeout waiting for Connect() from SubConn %v", sc2)
}
sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
// Send the first SubConn into CONNECTING. To do this, first make it READY,
// then CONNECTING.
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
cc.WaitForConnectivityState(ctx, connectivity.Ready)
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Idle})
// Since one endpoint is in TF and one in CONNECTING, the aggregated state
// will be CONNECTING.
cc.WaitForConnectivityState(ctx, connectivity.Connecting)
p1 := <-cc.NewPickerCh
p1.Pick(balancer.PickInfo{Ctx: ctx})
select {
case <-sc0.ConnectCh:
case <-time.After(defaultTestTimeout):
t.Errorf("timeout waiting for Connect() from SubConn %v", sc0)
}
sc0.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
// This will not trigger any new SubCOnns to be created, because sc0 is
// still attempting to connect, and we only need one SubConn to connect.
sc2.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
select {
case sc := <-cc.NewSubConnCh:
t.Fatalf("unexpected SubConn creation: %v", sc)
case <-sc0.ConnectCh:
t.Fatalf("unexpected Connect() from SubConn %v", sc0)
case <-sc1.ConnectCh:
t.Fatalf("unexpected Connect() from SubConn %v", sc1)
case <-sc2.ConnectCh:
t.Fatalf("unexpected Connect() from SubConn %v", sc2)
case <-time.After(defaultTestShortTimeout):
}
}
func (s) TestAggregatedConnectivityState(t *testing.T) {
tests := []struct {
name string
endpointStates []connectivity.State
want connectivity.State
}{
{
name: "one ready",
endpointStates: []connectivity.State{connectivity.Ready},
want: connectivity.Ready,
},
{
name: "one connecting",
endpointStates: []connectivity.State{connectivity.Connecting},
want: connectivity.Connecting,
},
{
name: "one ready one transient failure",
endpointStates: []connectivity.State{connectivity.Ready, connectivity.TransientFailure},
want: connectivity.Ready,
},
{
name: "one connecting one transient failure",
endpointStates: []connectivity.State{connectivity.Connecting, connectivity.TransientFailure},
want: connectivity.Connecting,
},
{
name: "one connecting two transient failure",
endpointStates: []connectivity.State{connectivity.Connecting, connectivity.TransientFailure, connectivity.TransientFailure},
want: connectivity.TransientFailure,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
bal := &ringhashBalancer{endpointStates: resolver.NewEndpointMap[*endpointState]()}
for i, cs := range tt.endpointStates {
es := &endpointState{
state: balancer.State{ConnectivityState: cs},
}
ep := resolver.Endpoint{Addresses: []resolver.Address{{Addr: fmt.Sprintf("%d.%d.%d.%d:%d", i, i, i, i, i)}}}
bal.endpointStates.Set(ep, es)
}
if got := bal.aggregatedStateLocked(); got != tt.want {
t.Errorf("recordTransition() = %v, want %v", got, tt.want)
}
})
}
}
type testKeyType string
const testKey testKeyType = "grpc.lb.ringhash.testKey"
type testAttribute struct {
content string
}
func setTestAttrAddr(addr resolver.Address, content string) resolver.Address {
addr.BalancerAttributes = addr.BalancerAttributes.WithValue(testKey, testAttribute{content})
return addr
}
func setTestAttrEndpoint(endpoint resolver.Endpoint, content string) resolver.Endpoint {
endpoint.Attributes = endpoint.Attributes.WithValue(testKey, testAttribute{content})
return endpoint
}
// TestAddrBalancerAttributesChange tests the case where the ringhash balancer
// receives a ClientConnUpdate with the same config and addresses as received in
// the previous update. Although the `BalancerAttributes` and endpoint
// attributes contents are the same, the pointers are different. This test
// verifies that subConns are not recreated in this scenario.
func (s) TestAddrBalancerAttributesChange(t *testing.T) {
content := "test"
addrs1 := []resolver.Address{setTestAttrAddr(resolver.Address{Addr: testBackendAddrStrs[0]}, content)}
wantEndpoints1 := []resolver.Endpoint{
setTestAttrEndpoint(resolver.Endpoint{Addresses: addrs1}, "content"),
}
cc, b, p0 := setupTest(t, wantEndpoints1)
ring0 := p0.(*picker).ring
firstHash := ring0.items[0].hash
// The first pick should be queued, and should trigger a connection to the
// only Endpoint which has a single address.
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()
if _, err := p0.Pick(balancer.PickInfo{Ctx: iringhash.SetXDSRequestHash(ctx, firstHash)}); err != balancer.ErrNoSubConnAvailable {
t.Fatalf("first pick returned err %v, want %v", err, balancer.ErrNoSubConnAvailable)
}
select {
case <-ctx.Done():
t.Fatalf("Timed out waiting for SubConn creation.")
case <-cc.NewSubConnCh:
}
addrs2 := []resolver.Address{setTestAttrAddr(resolver.Address{Addr: testBackendAddrStrs[0]}, content)}
wantEndpoints2 := []resolver.Endpoint{setTestAttrEndpoint(resolver.Endpoint{Addresses: addrs2}, content)}
if err := b.UpdateClientConnState(balancer.ClientConnState{
ResolverState: resolver.State{Endpoints: wantEndpoints2},
BalancerConfig: testConfig,
}); err != nil {
t.Fatalf("UpdateClientConnState returned err: %v", err)
}
select {
case <-cc.NewSubConnCh:
t.Fatal("new subConn created for an update with the same addresses")
case <-time.After(defaultTestShortTimeout):
}
}

Some files were not shown because too many files have changed in this diff Show More