xdsclient: revert #8369: delay resource cache deletion (#8527)

The change being reverted here (#8369) is a prime suspect for a race
that can show up with the following sequence of events:
- create a new gRPC channel with the `xds:///` scheme
- make an RPC
- close the channel
- repeat (possibly from multiple goroutines)

The observable behavior from the race is that the xDS client thinks that
a Listener resource is removed by the control plane when it clearly is
not. This results in the user's gRPC channel moving to TRANSIENT_FAILURE
and subsequent RPC failures.

The reason the above mentioned PR is not being rolled back using `git
revert` is because the xds directory structure has changed significantly
since the time the PR was originally merged. Manually performing the
revert seemed much easier.

RELEASE NOTES:
* xdsclient: Revert a change that introduces a race with xDS resource
processing, leading to RPC failures

.github/codecov.yml

@@ -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"

.github/mergeable.yml

@@ -1,21 +0,0 @@
-version: 2
-mergeable:
-  - when: pull_request.*
-    validate:
-      - do: label
-        must_include:
-          regex: '^Type:'
-      - do: description
-        must_include:
-          # Allow:
-          # RELEASE NOTES: none (case insensitive)
-          #
-          # RELEASE NOTES: N/A (case insensitive)
-          #
-          # RELEASE NOTES:
-          # * <text>
-          regex: '^RELEASE NOTES:\s*([Nn][Oo][Nn][Ee]|[Nn]/[Aa]|\n(\*|-)\s*.+)$'
-          regex_flag: 'm'
-      - do: milestone
-        must_include:
-          regex: 'Release$'

.github/pull_request_template.md

@@ -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.

.github/workflows/codeql-analysis.yml

@@ -3,35 +3,35 @@ name: "CodeQL"
 on:
   push:
     branches: [ master ]
-  pull_request:
-    # The branches below must be a subset of the branches above
-    branches: [ master ]
   schedule:
     - cron: '24 20 * * 3'
 
 permissions:
   contents: read
-  security-events: write
-  pull-requests: read
-  actions: 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@v2
+      uses: actions/checkout@v4
 
     # Initializes the CodeQL tools for scanning.
     - name: Initialize CodeQL
-      uses: github/codeql-action/init@v1
+      uses: github/codeql-action/init@v2
       with:
         languages: go
 
     - name: Perform CodeQL Analysis
-      uses: github/codeql-action/analyze@v1
+      uses: github/codeql-action/analyze@v2

.github/workflows/coverage.yml

@@ -0,0 +1,29 @@
+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: Run coverage with old pickfirst
+        run: GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST=false go test -coverprofile=coverage_old_pickfirst.out -coverpkg=./... ./...
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: true

.github/workflows/deps.yml

@@ -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

.github/workflows/lock.yml

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

.github/workflows/pr-validation.yml

@@ -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 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 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 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'.");
+              }
+            }

.github/workflows/release.yml

@@ -4,25 +4,31 @@ 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]
+        goarch: [386, amd64, arm64]
         exclude:
           - goos: darwin
             goarch: 386
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
 
       - name: Set up Go
-        uses: actions/setup-go@v2
+        uses: actions/setup-go@v5
 
       - name: Download dependencies
         run: |
@@ -48,14 +54,10 @@ jobs:
         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 ::set-output name=name::${PACKAGE_NAME}
+          echo "name=${PACKAGE_NAME}" >> $GITHUB_OUTPUT
 
       - name: Upload asset
-        uses: actions/upload-release-asset@v1
+        run: |
+          gh release upload ${{ github.event.release.tag_name }} ./${{ steps.package.outputs.name }}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        with:
-          upload_url: ${{ github.event.release.upload_url }}
-          asset_path: ./${{ steps.package.outputs.name }}
-          asset_name: ${{ steps.package.outputs.name }}
-          asset_content_type: application/gzip

.github/workflows/stale.yml

@@ -5,6 +5,9 @@ on:
   schedule:
   - cron: "44 */2 * * *"
 
+permissions:
+  contents: read
+
 jobs:
   stale:
     runs-on: ubuntu-latest
@@ -13,7 +16,7 @@ jobs:
       pull-requests: write
 
     steps:
-    - uses: actions/stale@v4
+    - uses: actions/stale@v8
       with:
         repo-token: ${{ secrets.GITHUB_TOKEN }}
         days-before-stale: 6

.github/workflows/testing.yml

@@ -18,51 +18,61 @@ 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@v2
+        uses: actions/setup-go@v5
         with:
-          go-version: 1.15
-      - name: Checkout repo
-        uses: actions/checkout@v2
+          go-version: '1.25'
+          cache-dependency-path: "**/go.sum"
 
-      # Run the vet checks.
-      - name: vet
-        run: ./vet.sh -install && ./vet.sh
+      # 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:
-    # Proto checks are run in the above job.
-    env:
-      VET_SKIP_PROTO: 1
-    runs-on: ubuntu-latest
+    # 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+tests
-            goversion: 1.17
-
-          - type: tests
-            goversion: 1.17
-            testflags: -race
+          - type: vet
+            goversion: '1.24'
 
           - type: extras
-            goversion: 1.17
+            goversion: '1.25'
 
           - type: tests
-            goversion: 1.17
+            goversion: '1.25'
+
+          - type: tests
+            goversion: '1.25'
+            testflags: -race
+
+          - type: tests
+            goversion: '1.25'
             goarch: 386
 
           - type: tests
-            goversion: 1.17
+            goversion: '1.25'
             goarch: arm64
+            runner: ubuntu-24.04-arm
 
           - type: tests
-            goversion: 1.16
+            goversion: '1.24'
 
           - type: tests
-            goversion: 1.15
+            goversion: '1.25'
+            testflags: -race
+            grpcenv: 'GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST=false'
 
     steps:
       # Setup the environment.
@@ -70,45 +80,52 @@ jobs:
         if: matrix.goarch != ''
         run: echo "GOARCH=${{ matrix.goarch }}" >> $GITHUB_ENV
 
-      - name: Setup qemu emulator
-        if: matrix.goarch == 'arm64'
-        # setup qemu-user-static emulator and register it with binfmt_misc so that aarch64 binaries
-        # are automatically executed using qemu.
-        run: docker run --rm --privileged multiarch/qemu-user-static:5.2.0-2 --reset --credential yes --persistent yes
-
       - 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@v2
+        uses: actions/setup-go@v5
         with:
           go-version: ${{ matrix.goversion }}
-
-      - name: Checkout repo
-        uses: actions/checkout@v2
+          cache-dependency-path: "**/*go.sum"
 
       # Only run vet for 'vet' runs.
       - name: Run vet.sh
-        if: startsWith(matrix.type, 'vet')
-        run: ./vet.sh -install && ./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: contains(matrix.type, 'tests')
+        if: matrix.type == 'tests'
         run: |
           go version
-          go test ${{ matrix.testflags }} -cpu 1,4 -timeout 7m google.golang.org/grpc/...
-          cd ${GITHUB_WORKSPACE}/security/advancedtls && go test ${{ matrix.testflags }} -timeout 2m google.golang.org/grpc/security/advancedtls/...
-          cd ${GITHUB_WORKSPACE}/security/authorization && go test ${{ matrix.testflags }} -timeout 2m google.golang.org/grpc/security/authorization/...
-
+          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

CONTRIBUTING.md

@@ -1,61 +1,159 @@
 # How to contribute
 
-We definitely 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)
-and [contribution guidelines](https://github.com/grpc/grpc-community/blob/master/CONTRIBUTING.md) before proceeding.
+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/)
+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).
+[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.
+
+## 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
-How to get your contributions merged smoothly and quickly.
+
+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 times receive PRs that are trying to fix several things at
-  a time, but only one fix is considered acceptable, nothing gets merged and
-  both author's & review's time is wasted. Create more PRs to address different
-  concerns and everyone will be happy.
+  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.
 
-- The grpc package should only depend on standard Go packages and a small number
-  of exceptions. If your contribution introduces new dependencies which are NOT
-  in the [list](https://godoc.org/google.golang.org/grpc?imports), you need a
-  discussion with gRPC-Go authors and consultants.
+- 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.
 
-- For speculative changes, 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).
+- 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.
 
-- Provide a good **PR description** as a record of **what** change is being made
-  and **why** it was made. Link to a github issue if it exists.
-
-- Don't fix code style and formatting unless you are already changing that line
-  to address an issue. PRs with irrelevant changes won't be merged. If you do
-  want to fix formatting or style, do that in a separate PR.
-
-- Unless your PR is trivial, you should expect there will be reviewer comments
-  that you'll need to address before merging. We expect you to be reasonably
-  responsive to those comments, otherwise the PR will be closed after 2-3 weeks
-  of inactivity.
-
-- Maintain **clean commit history** and use **meaningful commit messages**. PRs
-  with messy commit history are difficult to review and won't be merged. Use
-  `rebase -i upstream/master` to curate your commit history and/or to bring in
-  latest changes from master (but avoid rebasing in the middle of a code
-  review).
-
-- Keep your PR up to date with upstream/master (if there are merge conflicts, we
-  can't really merge your 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.
-  - `make all` to test everything, OR
-  - `make vet` to catch vet errors
-  - `make test` to run the tests
-  - `make testrace` to run tests in race mode
+  recommend you run tests locally before creating your PR to catch breakages
+  early on:
 
-- Exceptions to the rules can be made if there's a compelling reason for doing so.
+  - `./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.

Documentation/anti-patterns.md

@@ -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)
+```

Documentation/benchmark.md

@@ -13,9 +13,9 @@ simulate your application:
 ```bash
 $ go run google.golang.org/grpc/benchmark/benchmain/main.go \
     -workloads=streaming \
-  	-reqSizeBytes=1024 \
-  	-respSizeBytes=1024 \
-  	-compression=gzip
+    -reqSizeBytes=1024 \
+    -respSizeBytes=1024 \
+    -compression=gzip
 ```
 
 Pass the `-h` flag to the `benchmain` utility to see other flags and workloads
@@ -45,8 +45,8 @@ 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
+    -reqPayloadCurveFiles=/path/to/csv \
+    -respPayloadCurveFiles=/path/to/csv
 ```
 
 This tells the `benchmain` utility to generate unary RPC requests with a 25%
@@ -61,8 +61,8 @@ 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
+    -reqPayloadCurveFiles=/path/to/csv1,/path/to/csv2 \
+    -respPayloadCurveFiles=/path/to/csv3,/path/to/csv4
 ```
 
 You may also combine `PayloadCurveFiles` with `SizeBytes` options. For example:
@@ -70,6 +70,6 @@ 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
+    -reqPayloadCurveFiles=/path/to/csv \
+    -respSizeBytes=1
 ```

Documentation/encoding.md

@@ -22,7 +22,7 @@ package proto
 import "google.golang.org/grpc/encoding"
 
 func init() {
-	encoding.RegisterCodec(protoCodec{})
+    encoding.RegisterCodec(protoCodec{})
 }
 
 // ... implementation of protoCodec ...
@@ -50,14 +50,14 @@ 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"))
+    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.Dial(ctx, target, grpc.WithDefaultCallOptions(grpc.CallContentSubtype("mycodec")))
+    myclient := grpc.NewClient(target, grpc.WithDefaultCallOptions(grpc.CallContentSubtype("mycodec")))
 ```
 
 When specified in either of these ways, messages will be encoded using this
@@ -83,7 +83,7 @@ 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
+[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.
 
@@ -98,7 +98,7 @@ package gzip
 import "google.golang.org/grpc/encoding"
 
 func init() {
-	encoding.RegisterCompressor(compressor{})
+    encoding.RegisterCompressor(compressor{})
 }
 
 // ... implementation of compressor ...
@@ -125,14 +125,14 @@ 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"))
+    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.Dial(ctx, target, grpc.WithDefaultCallOptions(grpc.UseCompressor("gzip")))
+    myclient := grpc.NewClient(target, grpc.WithDefaultCallOptions(grpc.UseCompressor("gzip")))
 ```
 
 When specified in either of these ways, messages will be compressed using this

Documentation/grpc-auth-support.md

@@ -1,11 +1,11 @@
 # 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 an 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.
+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.Dial(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")))
+conn, err := grpc.NewClient(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")))
 ```
 
 # Enabling TLS on a gRPC server
@@ -53,7 +53,7 @@ Alternatively, a client may also use the `grpc.CallOption`
 on each invocation of an RPC.
 
 To create a `credentials.PerRPCCredentials`, use
-[oauth.NewOauthAccess](https://godoc.org/google.golang.org/grpc/credentials/oauth#NewOauthAccess).
+[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.
@@ -63,7 +63,7 @@ to prevent any insecure transmission of tokens.
 ## Google Compute Engine (GCE)
 
 ```Go
-conn, err := grpc.Dial(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), grpc.WithPerRPCCredentials(oauth.NewComputeEngine()))
+conn, err := grpc.NewClient(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), grpc.WithPerRPCCredentials(oauth.NewComputeEngine()))
 ```
 
 ## JWT
@@ -73,6 +73,6 @@ jwtCreds, err := oauth.NewServiceAccountFromFile(*serviceAccountKeyFile, *oauthS
 if err != nil {
   log.Fatalf("Failed to create JWT credentials: %v", err)
 }
-conn, err := grpc.Dial(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), grpc.WithPerRPCCredentials(jwtCreds))
+conn, err := grpc.NewClient(serverAddr, grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, "")), grpc.WithPerRPCCredentials(jwtCreds))
 ```
 

Documentation/grpc-metadata.md

@@ -12,11 +12,11 @@ Four kinds of service method:
 - [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](https://grpc.io/docs/guides/concepts.html#metadata).
+And concept of [metadata].
 
 ## Constructing metadata
 
-A metadata can be created using package [metadata](https://godoc.org/google.golang.org/grpc/metadata).
+A metadata can be created using package [metadata].
 The type MD is actually a map from string to a list of strings:
 
 ```go
@@ -64,20 +64,10 @@ md := metadata.Pairs(
 )
 ```
 
-## Retrieving metadata from context
-
-Metadata can be retrieved from context using `FromIncomingContext`:
-
-```go
-func (s *server) SomeRPC(ctx context.Context, in *pb.SomeRequest) (*pb.SomeResponse, err) {
-    md, ok := metadata.FromIncomingContext(ctx)
-    // do something with metadata
-}
-```
-
 ## Sending and receiving metadata - client side
 
-Client side metadata sending and receiving examples are available [here](../examples/features/metadata/client/main.go).
+Client side metadata sending and receiving examples are available
+[here](../examples/features/metadata/client/main.go).
 
 ### Sending metadata
 
@@ -127,7 +117,8 @@ 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](https://godoc.org/google.golang.org/grpc#Header) and [Trailer](https://godoc.org/google.golang.org/grpc#Trailer) in [CallOption](https://godoc.org/google.golang.org/grpc#CallOption):
+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
@@ -149,7 +140,8 @@ For streaming calls including:
 - Client streaming RPC
 - Bidirectional streaming RPC
 
-Header and trailer can be retrieved from the returned stream using function `Header` and `Trailer` in interface [ClientStream](https://godoc.org/google.golang.org/grpc#ClientStream):
+Header and trailer can be retrieved from the returned stream using function
+`Header` and `Trailer` in interface [ClientStream]:
 
 ```go
 stream, err := client.SomeStreamingRPC(ctx)
@@ -164,11 +156,13 @@ 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).
+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.
+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.
 
@@ -194,15 +188,16 @@ func (s *server) SomeStreamingRPC(stream pb.Service_SomeStreamingRPCServer) erro
 
 #### Unary call
 
-To send header and trailer to client in unary call, the server can call [SendHeader](https://godoc.org/google.golang.org/grpc#SendHeader) and [SetTrailer](https://godoc.org/google.golang.org/grpc#SetTrailer) functions in module [grpc](https://godoc.org/google.golang.org/grpc).
+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 send header
+    // create and set header
     header := metadata.Pairs("header-key", "val")
-    grpc.SendHeader(ctx, header)
+    grpc.SetHeader(ctx, header)
     // create and set trailer
     trailer := metadata.Pairs("trailer-key", "val")
     grpc.SetTrailer(ctx, trailer)
@@ -211,15 +206,39 @@ func (s *server) SomeRPC(ctx context.Context, in *pb.someRequest) (*pb.someRespo
 
 #### Streaming call
 
-For streaming calls, header and trailer can be sent using function `SendHeader` and `SetTrailer` in interface [ServerStream](https://godoc.org/google.golang.org/grpc#ServerStream):
+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 send header
+    // create and set header
     header := metadata.Pairs("header-key", "val")
-    stream.SendHeader(header)
+    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

Documentation/proxy.md

@@ -1,8 +1,8 @@
 # Proxy
 
 HTTP CONNECT proxies are supported by default in gRPC. The proxy address can be
-specified by the environment variables HTTP_PROXY, HTTPS_PROXY and NO_PROXY (or
-the lowercase versions thereof).
+specified by the environment variables `HTTPS_PROXY` and `NO_PROXY`.  (Note that
+these environment variables are case insensitive.)
 
 ## Custom proxy
 
@@ -12,4 +12,4 @@ 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
-[`WithDialer`](https://godoc.org/google.golang.org/grpc#WithDialer).
+[`WithContextDialer`](https://pkg.go.dev/google.golang.org/grpc#WithContextDialer).

Documentation/rpc-errors.md

@@ -65,4 +65,4 @@ exit status 1
 [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/errors
+[example]:      https://github.com/grpc/grpc-go/tree/master/examples/features/error_details

Documentation/server-reflection-tutorial.md

@@ -2,8 +2,9 @@
 
 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 gRPC CLI,
-which can be used to introspect server protos and send/receive test RPCs.
+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
 
@@ -39,36 +40,41 @@ make the following changes:
 An example server with reflection registered can be found at
 `examples/features/reflection/server`.
 
-## gRPC CLI
+## gRPCurl
 
-After enabling Server Reflection in a server application, you can use gRPC CLI
-to check its services. gRPC CLI is only available in c++. Instructions on how to
-build and use gRPC CLI can be found at
-[command_line_tool.md](https://github.com/grpc/grpc/blob/master/doc/command_line_tool.md).
+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 gRPC CLI to check services
+## Use gRPCurl to check services
 
 First, start the helloworld server in grpc-go directory:
 
 ```sh
-$ cd <grpc-go-directory>
-$ go run examples/features/reflection/server/main.go
+$ cd <grpc-go-directory>/examples
+$ go run features/reflection/server/main.go
 ```
 
-Open a new terminal and make sure you are in the directory where grpc_cli lives:
-
+output:
 ```sh
-$ cd <grpc-cpp-directory>/bins/opt
+server listening at [::]:50051
 ```
 
-### List services
+After installing gRPCurl, open a new terminal and run the commands from the new
+terminal.
 
-`grpc_cli ls` command lists services and methods exposed at a given port:
+**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
-  $ ./grpc_cli ls localhost:50051
+  $ grpcurl -plaintext localhost:50051 list
   ```
 
   output:
@@ -78,72 +84,88 @@ $ cd <grpc-cpp-directory>/bins/opt
   helloworld.Greeter
   ```
 
-- List one service with details
+- List all the methods of a service
 
-  `grpc_cli ls` command inspects a service given its full name (in the format of
-  \<package\>.\<service\>). It can print information with a long listing format
-  when `-l` flag is set. This flag can be used to get more details about a
-  service.
+  The `list` command lists methods given the full service name (in the format of
+  \<package\>.\<service\>).
 
   ```sh
-  $ ./grpc_cli ls localhost:50051 helloworld.Greeter -l
+  $ grpcurl -plaintext localhost:50051 list helloworld.Greeter
   ```
 
   output:
   ```sh
-  filename: helloworld.proto
-  package: helloworld;
+  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) {}
+    rpc SayHello ( .helloworld.HelloRequest ) returns ( .helloworld.HelloReply );
   }
-
   ```
 
-### List methods
+- Describe all methods of a service
 
-- List one method with details
-
-  `grpc_cli ls` command also inspects a method given its full name (in the
-  format of \<package\>.\<service\>.\<method\>).
+  The `describe` command inspects a method given its full name (in the format of
+  \<package\>.\<service\>.\<method\>).
 
   ```sh
-  $ ./grpc_cli ls localhost:50051 helloworld.Greeter.SayHello -l
+  $ grpcurl -plaintext localhost:50051 describe helloworld.Greeter.SayHello
   ```
 
   output:
   ```sh
-    rpc SayHello(helloworld.HelloRequest) returns (helloworld.HelloReply) {}
+  helloworld.Greeter.SayHello is a method:
+  rpc SayHello ( .helloworld.HelloRequest ) returns ( .helloworld.HelloReply );
   ```
 
 ### Inspect message types
 
-We can use`grpc_cli type` command to inspect request/response types given the
+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
-  $ ./grpc_cli type localhost:50051 helloworld.HelloRequest
+  $ grpcurl -plaintext localhost:50051 describe helloworld.HelloRequest
   ```
 
   output:
   ```sh
+  helloworld.HelloRequest is a message:
   message HelloRequest {
-    optional string name = 1[json_name = "name"];
+    string name = 1;
   }
   ```
 
 ### Call a remote method
 
-We can send RPCs to a server and get responses using `grpc_cli call` command.
+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
-  $ ./grpc_cli call localhost:50051 SayHello "name: 'gRPC CLI'"
+  $ grpcurl -plaintext -format text -d 'name: "gRPCurl"' \
+    localhost:50051 helloworld.Greeter.SayHello
   ```
 
   output:
   ```sh
-  message: "Hello gRPC CLI"
+  message: "Hello gRPCurl"
   ```

MAINTAINERS.md

@@ -9,20 +9,28 @@ for general contribution guidelines.
 
 ## Maintainers (in alphabetical order)
 
-- [cesarghali](https://github.com/cesarghali), Google LLC
+- [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
-- [menghanl](https://github.com/menghanl), Google LLC
-- [srini100](https://github.com/srini100), Google LLC
+- [gtcooke94](https://github.com/gtcooke94), Google LLC
 
 ## Emeritus Maintainers (in alphabetical order)
-- [adelez](https://github.com/adelez), Google LLC
-- [canguler](https://github.com/canguler), Google LLC
-- [iamqizhao](https://github.com/iamqizhao), Google LLC
-- [jadekler](https://github.com/jadekler), Google LLC
-- [jtattermusch](https://github.com/jtattermusch), Google LLC
-- [lyuxuan](https://github.com/lyuxuan), Google LLC
-- [makmukhi](https://github.com/makmukhi), Google LLC
-- [matt-kwong](https://github.com/matt-kwong), Google LLC
-- [nicolasnoble](https://github.com/nicolasnoble), Google LLC
-- [yongni](https://github.com/yongni), Google LLC
+- [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)

Makefile

@@ -30,17 +30,20 @@ testdeps:
 	GO111MODULE=on go get -d -v -t google.golang.org/grpc/...
 
 vet: vetdeps
-	./vet.sh
+	./scripts/vet.sh
 
 vetdeps:
-	./vet.sh -install
+	./scripts/vet.sh -install
 
 .PHONY: \
 	all \
 	build \
 	clean \
+	deps \
 	proto \
 	test \
+	testsubmodule \
 	testrace \
+	testdeps \
 	vet \
 	vetdeps

README.md

@@ -1,8 +1,8 @@
 # gRPC-Go
 
-[![Build Status](https://travis-ci.org/grpc/grpc-go.svg)](https://travis-ci.org/grpc/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
@@ -10,25 +10,18 @@ RPC framework that puts mobile and HTTP/2 first. For more information see the
 
 ## Prerequisites
 
-- **[Go][]**: any one of the **three latest major** [releases][go-releases].
+- **[Go][]**: any one of the **two latest major** [releases][go-releases].
 
 ## Installation
 
-With [Go module][] support (Go 1.11+), simply add the following import
+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"
 ```
 
-to your code, and then `go [build|run|test]` will automatically fetch the
-necessary dependencies.
-
-Otherwise, to install the `grpc-go` package, run the following command:
-
-```console
-$ go get -u google.golang.org/grpc
-```
-
 > **Note:** If you are trying to access `grpc-go` from **China**, see the
 > [FAQ](#FAQ) below.
 
@@ -39,6 +32,7 @@ $ go get -u google.golang.org/grpc
 - [Low-level technical docs](Documentation) from this repository
 - [Performance benchmark][]
 - [Examples](examples)
+- [Contribution guidelines](CONTRIBUTING.md)
 
 ## FAQ
 
@@ -56,15 +50,6 @@ To build Go code, there are several options:
 
 - Set up a VPN and access google.golang.org through that.
 
-- Without Go module support: `git clone` the repo manually:
-
-  ```sh
-  git clone https://github.com/grpc/grpc-go.git $GOPATH/src/google.golang.org/grpc
-  ```
-
-  You will need to do the same for all of grpc's dependencies in `golang.org`,
-  e.g. `golang.org/x/net`.
-
 - 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:
 
@@ -76,33 +61,13 @@ To build Go code, there are several options:
   ```
 
   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).
+  golang.org as well. For details, refer to [golang/go issue
+  #28652](https://github.com/golang/go/issues/28652).
 
 ### Compiling error, undefined: grpc.SupportPackageIsVersion
 
-#### If you are using Go modules:
-
-Ensure your gRPC-Go version is `require`d at the appropriate version in
-the same module containing the generated `.pb.go` files.  For example,
-`SupportPackageIsVersion6` needs `v1.27.0`, so in your `go.mod` file:
-
-```go
-module <your module name>
-
-require (
-    google.golang.org/grpc v1.27.0
-)
-```
-
-#### If you are *not* using Go modules:
-
-Update the `proto` package, gRPC package, and rebuild the `.proto` files:
-
-```sh
-go get -u github.com/golang/protobuf/{proto,protoc-gen-go}
-go get -u google.golang.org/grpc
-protoc --go_out=plugins=grpc:. *.proto
-```
+Please update to the latest version of gRPC-Go using
+`go get google.golang.org/grpc`.
 
 ### How to turn on logging
 
@@ -121,9 +86,11 @@ 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),
+ 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

SECURITY.md

@@ -1,3 +1,3 @@
 # Security Policy
 
-For information on gRPC Security Policy and reporting potentional security issues, please see [gRPC CVE Process](https://github.com/grpc/proposal/blob/master/P4-grpc-cve-process.md).
+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).

admin/admin.go

@@ -20,9 +20,10 @@
 // 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
+// # Experimental
 //
 // Notice: All APIs in this package are experimental and may be removed in a
 // later release.

admin/test/utils.go

@@ -26,15 +26,16 @@ import (
 	"testing"
 	"time"
 
-	v3statusgrpc "github.com/envoyproxy/go-control-plane/envoy/service/status/v3"
-	v3statuspb "github.com/envoyproxy/go-control-plane/envoy/service/status/v3"
-	"github.com/google/uuid"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/admin"
-	channelzpb "google.golang.org/grpc/channelz/grpc_channelz_v1"
 	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/internal/xds"
+	"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 (
@@ -51,17 +52,6 @@ type ExpectedStatusCodes struct {
 // RunRegisterTests makes a client, runs the RPCs, and compares the status
 // codes.
 func RunRegisterTests(t *testing.T, ec ExpectedStatusCodes) {
-	nodeID := uuid.New().String()
-	bootstrapCleanup, err := xds.SetupBootstrapFile(xds.BootstrapOptions{
-		Version:   xds.TransportV3,
-		NodeID:    nodeID,
-		ServerURI: "no.need.for.a.server",
-	})
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer bootstrapCleanup()
-
 	lis, err := net.Listen("tcp", "localhost:0")
 	if err != nil {
 		t.Fatalf("cannot create listener: %v", err)
@@ -78,9 +68,9 @@ func RunRegisterTests(t *testing.T, ec ExpectedStatusCodes) {
 		server.Serve(lis)
 	}()
 
-	conn, err := grpc.Dial(lis.Addr().String(), grpc.WithInsecure())
+	conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
-		t.Fatalf("cannot connect to server: %v", err)
+		t.Fatalf("grpc.NewClient(%q) = %v", lis.Addr().String(), err)
 	}
 
 	t.Run("channelz", func(t *testing.T) {
@@ -97,7 +87,7 @@ func RunRegisterTests(t *testing.T, ec ExpectedStatusCodes) {
 
 // RunChannelz makes a channelz RPC.
 func RunChannelz(conn *grpc.ClientConn) error {
-	c := channelzpb.NewChannelzClient(conn)
+	c := channelzgrpc.NewChannelzClient(conn)
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 	_, err := c.GetTopChannels(ctx, &channelzpb.GetTopChannelsRequest{}, grpc.WaitForReady(true))

attributes/attributes.go

@@ -19,61 +19,123 @@
 // Package attributes defines a generic key/value store used in various gRPC
 // components.
 //
-// Experimental
+// # Experimental
 //
 // Notice: This package is EXPERIMENTAL and may be changed or removed in a
 // later release.
 package attributes
 
-import "fmt"
+import (
+	"fmt"
+	"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.
+// 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 {
-	m map[interface{}]interface{}
+	m map[any]any
 }
 
-// New returns a new Attributes containing all key/value pairs in kvs.  If the
-// same key appears multiple times, the last value overwrites all previous
-// values for that key.  Panics if len(kvs) is not even.
-func New(kvs ...interface{}) *Attributes {
-	if len(kvs)%2 != 0 {
-		panic(fmt.Sprintf("attributes.New called with unexpected input: len(kvs) = %v", len(kvs)))
-	}
-	a := &Attributes{m: make(map[interface{}]interface{}, len(kvs)/2)}
-	for i := 0; i < len(kvs)/2; i++ {
-		a.m[kvs[i*2]] = kvs[i*2+1]
-	}
-	return a
+// New returns a new Attributes containing the key/value pair.
+func New(key, value any) *Attributes {
+	return &Attributes{m: map[any]any{key: value}}
 }
 
-// WithValues returns a new Attributes containing all key/value pairs in a and
-// kvs.  Panics if len(kvs) is not even.  If the same key appears multiple
-// times, the last value overwrites all previous values for that key.  To
-// remove an existing key, use a nil value.
-func (a *Attributes) WithValues(kvs ...interface{}) *Attributes {
+// 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.  To remove an
+// existing key, use a nil value.  value should not be modified later.
+func (a *Attributes) WithValue(key, value any) *Attributes {
 	if a == nil {
-		return New(kvs...)
+		return New(key, value)
 	}
-	if len(kvs)%2 != 0 {
-		panic(fmt.Sprintf("attributes.New called with unexpected input: len(kvs) = %v", len(kvs)))
-	}
-	n := &Attributes{m: make(map[interface{}]interface{}, len(a.m)+len(kvs)/2)}
+	n := &Attributes{m: make(map[any]any, len(a.m)+1)}
 	for k, v := range a.m {
 		n.m[k] = v
 	}
-	for i := 0; i < len(kvs)/2; i++ {
-		n.m[kvs[i*2]] = kvs[i*2+1]
-	}
+	n.m[key] = value
 	return n
 }
 
 // Value returns the value associated with these attributes for key, or nil if
-// no value is associated with key.
-func (a *Attributes) Value(key interface{}) interface{} {
+// no value is associated with key.  The returned value should not be modified.
+func (a *Attributes) Value(key any) any {
 	if a == nil {
 		return nil
 	}
 	return a.m[key]
 }
+
+// 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 len(a.m) != len(o.m) {
+		return false
+	}
+	for k, v := range a.m {
+		ov, ok := o.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 true
+}
+
+// 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.m {
+		if !first {
+			sb.WriteString(", ")
+		}
+		sb.WriteString(fmt.Sprintf("%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
+}

attributes/attributes_test.go

@@ -20,41 +20,109 @@ package attributes_test
 
 import (
 	"fmt"
-	"reflect"
 	"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, keyTwo{}, "two")
+	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
+	// Key two: {two}
 }
 
-func ExampleAttributes_WithValues() {
+func ExampleAttributes_WithValue() {
 	type keyOne struct{}
 	type keyTwo struct{}
 	a := attributes.New(keyOne{}, 1)
-	a = a.WithValues(keyTwo{}, "two")
+	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
+	// Key two: {two}
 }
 
-// Test that two attributes with the same content are `reflect.DeepEqual`.
-func TestDeepEqual(t *testing.T) {
+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 the same content are Equal.
+func TestEqual(t *testing.T) {
 	type keyOne struct{}
-	a1 := attributes.New(keyOne{}, 1)
-	a2 := attributes.New(keyOne{}, 1)
-	if !reflect.DeepEqual(a1, a2) {
-		t.Fatalf("reflect.DeepEqual(%+v, %+v), want true, got false", a1, a2)
+	type keyTwo struct{}
+	a1 := attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"})
+	a2 := attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"})
+	if !a1.Equal(a2) {
+		t.Fatalf("%+v.Equals(%+v) = false; want true", a1, a2)
+	}
+	if !a2.Equal(a1) {
+		t.Fatalf("%+v.Equals(%+v) = false; want true", a2, a1)
+	}
+}
+
+// Test that two attributes with different content are not Equal.
+func TestNotEqual(t *testing.T) {
+	type keyOne struct{}
+	type keyTwo struct{}
+	a1 := attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "two"})
+	a2 := attributes.New(keyOne{}, 2).WithValue(keyTwo{}, stringVal{s: "two"})
+	a3 := attributes.New(keyOne{}, 1).WithValue(keyTwo{}, stringVal{s: "one"})
+	if a1.Equal(a2) {
+		t.Fatalf("%+v.Equals(%+v) = true; want false", a1, a2)
+	}
+	if a2.Equal(a1) {
+		t.Fatalf("%+v.Equals(%+v) = true; want false", a2, a1)
+	}
+	if a3.Equal(a1) {
+		t.Fatalf("%+v.Equals(%+v) = true; want false", a3, a1)
 	}
 }

authz/audit/audit_logger.go

@@ -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
+}

authz/audit/audit_logging_test.go

@@ -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
+}

authz/audit/stdout/stdout_logger.go

@@ -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),
+	}
+}

authz/audit/stdout/stdout_logger_test.go

@@ -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,
+	}
+}

authz/grpc_authz_end2end_test.go

@@ -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, "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, "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, "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, "unauthorized RPC request rejected"),
+	},
+	"DeniesRPCRequestWithPrincipalsFieldOnUnauthenticatedConnection": {
+		authzPolicy: `{
+				"name": "authz",
+				"allow_rules":
+				[
+					{
+						"name": "allow_authenticated",
+						"source": {
+							"principals": ["*", ""]
+						}
+					}
+				]
+			}`,
+		wantStatus: status.New(codes.PermissionDenied, "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, "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())
+	}
+}

authz/grpc_authz_server_interceptors.go

@@ -0,0 +1,178 @@
+/*
+ * 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, "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, "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 {
+	internalInterceptor unsafe.Pointer // *StaticInterceptor
+	policyFile          string
+	policyContents      []byte
+	refreshDuration     time.Duration
+	cancel              context.CancelFunc
+}
+
+// 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) {
+	if file == "" {
+		return nil, fmt.Errorf("authorization policy file path is empty")
+	}
+	if duration <= time.Duration(0) {
+		return nil, fmt.Errorf("requires refresh interval(%v) greater than 0s", duration)
+	}
+	i := &FileWatcherInterceptor{policyFile: file, refreshDuration: duration}
+	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.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.policyFile)
+	if err != nil {
+		return fmt.Errorf("policyFile(%s) read failed: %v", i.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)
+	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)
+}

authz/grpc_authz_server_interceptors_test.go

@@ -0,0 +1,120 @@
+/*
+ *
+ * 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 (
+	"fmt"
+	"os"
+	"path"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/authz"
+)
+
+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("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()
+			}
+		})
+	}
+}

authz/rbac_translator.go

@@ -16,7 +16,7 @@
 
 // Package authz exposes methods to manage authorization within gRPC.
 //
-// Experimental
+// # Experimental
 //
 // Notice: This package is EXPERIMENTAL and may be changed or removed
 // in a later release.
@@ -28,11 +28,19 @@ import (
 	"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
@@ -53,11 +61,23 @@ type rule struct {
 	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"`
+	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 {
@@ -94,7 +114,8 @@ func getStringMatcher(value string) *v3matcherpb.StringMatcher {
 	switch {
 	case value == "*":
 		return &v3matcherpb.StringMatcher{
-			MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{},
+			MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{
+				SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
 		}
 	case strings.HasSuffix(value, "*"):
 		prefix := strings.TrimSuffix(value, "*")
@@ -117,8 +138,9 @@ func getHeaderMatcher(key, value string) *v3routepb.HeaderMatcher {
 	switch {
 	case value == "*":
 		return &v3routepb.HeaderMatcher{
-			Name:                 key,
-			HeaderMatchSpecifier: &v3routepb.HeaderMatcher_SafeRegexMatch{},
+			Name: key,
+			HeaderMatchSpecifier: &v3routepb.HeaderMatcher_SafeRegexMatch{
+				SafeRegexMatch: &v3matcherpb.RegexMatcher{Regex: ".+"}},
 		}
 	case strings.HasSuffix(value, "*"):
 		prefix := strings.TrimSuffix(value, "*")
@@ -154,15 +176,15 @@ func parsePrincipalNames(principalNames []string) []*v3rbacpb.Principal {
 	return ps
 }
 
-func parsePeer(source peer) (*v3rbacpb.Principal, error) {
-	if len(source.Principals) > 0 {
-		return principalOr(parsePrincipalNames(source.Principals)), nil
+func parsePeer(source peer) *v3rbacpb.Principal {
+	if len(source.Principals) == 0 {
+		return &v3rbacpb.Principal{
+			Identifier: &v3rbacpb.Principal_Any{
+				Any: true,
+			},
+		}
 	}
-	return &v3rbacpb.Principal{
-		Identifier: &v3rbacpb.Principal_Any{
-			Any: true,
-		},
-	}, nil
+	return principalOr(parsePrincipalNames(source.Principals))
 }
 
 func parsePaths(paths []string) []*v3rbacpb.Permission {
@@ -251,56 +273,126 @@ func parseRules(rules []rule, prefixName string) (map[string]*v3rbacpb.Policy, e
 		if rule.Name == "" {
 			return policies, fmt.Errorf(`%d: "name" is not present`, i)
 		}
-		principal, err := parsePeer(rule.Source)
-		if err != nil {
-			return nil, fmt.Errorf("%d: %v", i, err)
-		}
 		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{principal},
+			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. If the input policy cannot be parsed or is invalid, an error
-// will be returned.
-func translatePolicy(policyStr string) ([]*v3rbacpb.RBAC, error) {
+// 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)
+		return nil, "", fmt.Errorf("failed to unmarshal policy: %v", err)
 	}
 	if policy.Name == "" {
-		return nil, fmt.Errorf(`"name" is not present`)
+		return nil, "", fmt.Errorf(`"name" is not present`)
 	}
 	if len(policy.AllowRules) == 0 {
-		return nil, fmt.Errorf(`"allow_rules" is not present`)
+		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)
+			return nil, "", fmt.Errorf(`"deny_rules" %v`, err)
 		}
 		denyRBAC := &v3rbacpb.RBAC{
-			Action:   v3rbacpb.RBAC_DENY,
-			Policies: denyPolicies,
+			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)
+		return nil, "", fmt.Errorf(`"allow_rules" %v`, err)
 	}
-	allowRBAC := &v3rbacpb.RBAC{Action: v3rbacpb.RBAC_ALLOW, Policies: allowPolicies}
-	return append(rbacs, allowRBAC), nil
+	allowRBAC := &v3rbacpb.RBAC{Action: v3rbacpb.RBAC_ALLOW, Policies: allowPolicies, AuditLoggingOptions: allowLogger}
+	return append(rbacs, allowRBAC), policy.Name, nil
 }

authz/rbac_translator_test.go

@@ -22,9 +22,13 @@ import (
 	"strings"
 	"testing"
 
+	v1xdsudpatypepb "github.com/cncf/xds/go/udpa/type/v1"
 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/protobuf/testing/protocmp"
+	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/structpb"
 
+	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"
@@ -32,9 +36,10 @@ import (
 
 func TestTranslatePolicy(t *testing.T) {
 	tests := map[string]struct {
-		authzPolicy  string
-		wantErr      string
-		wantPolicies []*v3rbacpb.RBAC
+		authzPolicy    string
+		wantErr        string
+		wantPolicies   []*v3rbacpb.RBAC
+		wantPolicyName string
 	}{
 		"valid policy": {
 			authzPolicy: `{
@@ -117,6 +122,7 @@ func TestTranslatePolicy(t *testing.T) {
 							},
 						},
 					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{},
 				},
 				{
 					Action: v3rbacpb.RBAC_ALLOW,
@@ -127,7 +133,7 @@ func TestTranslatePolicy(t *testing.T) {
 									Ids: []*v3rbacpb.Principal{
 										{Identifier: &v3rbacpb.Principal_Authenticated_{
 											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
-												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{},
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
 											}},
 										}},
 									},
@@ -202,6 +208,690 @@ func TestTranslatePolicy(t *testing.T) {
 							},
 						},
 					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{},
+				},
+			},
+			wantPolicyName: "authz",
+		},
+		"allow authenticated": {
+			authzPolicy: `{
+						"name": "authz",
+						"allow_rules": [
+						{
+							"name": "allow_authenticated",
+							"source": {
+								"principals":["*", ""]
+							}
+						}]
+					}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{},
+				},
+			},
+		},
+		"audit_logging_ALLOW empty config": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"deny_rules": [
+				{
+					"name": "deny_policy_1",
+					"source": {
+						"principals":[
+						"spiffe://foo.abc"
+						]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "ON_ALLOW",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": {},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_DENY,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_deny_policy_1": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: "spiffe://foo.abc"},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_ON_ALLOW,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+			},
+		},
+		"audit_logging_DENY_AND_ALLOW": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"deny_rules": [
+				{
+					"name": "deny_policy_1",
+					"source": {
+						"principals":[
+						"spiffe://foo.abc"
+						]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "ON_DENY_AND_ALLOW",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": {},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_DENY,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_deny_policy_1": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: "spiffe://foo.abc"},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY_AND_ALLOW,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+			},
+		},
+		"audit_logging_NONE": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"deny_rules": [
+				{
+					"name": "deny_policy_1",
+					"source": {
+						"principals":[
+						"spiffe://foo.abc"
+						]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "NONE",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": {},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_DENY,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_deny_policy_1": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: "spiffe://foo.abc"},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+			},
+		},
+		"audit_logging_custom_config simple": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"deny_rules": [
+				{
+					"name": "deny_policy_1",
+					"source": {
+						"principals":[
+						"spiffe://foo.abc"
+						]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "NONE",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": {"abc":123, "xyz":"123"},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_DENY,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_deny_policy_1": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: "spiffe://foo.abc"},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{"abc": 123, "xyz": "123"}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{"abc": 123, "xyz": "123"}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+			},
+		},
+		"audit_logging_custom_config nested": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "NONE",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": {"abc":123, "xyz":{"abc":123}},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{"abc": 123, "xyz": map[string]any{"abc": 123}}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+			},
+		},
+		"missing audit logger config": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "NONE"
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs:  []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{},
+					},
+				},
+			},
+		},
+		"missing audit condition": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": {},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_NONE,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+			},
+		},
+		"missing custom config audit logger": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"deny_rules": [
+				{
+					"name": "deny_policy_1",
+					"source": {
+						"principals":[
+						"spiffe://foo.abc"
+						]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "ON_DENY",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantPolicies: []*v3rbacpb.RBAC{
+				{
+					Action: v3rbacpb.RBAC_DENY,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_deny_policy_1": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: "spiffe://foo.abc"},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
+				},
+				{
+					Action: v3rbacpb.RBAC_ALLOW,
+					Policies: map[string]*v3rbacpb.Policy{
+						"authz_allow_authenticated": {
+							Principals: []*v3rbacpb.Principal{
+								{Identifier: &v3rbacpb.Principal_OrIds{OrIds: &v3rbacpb.Principal_Set{
+									Ids: []*v3rbacpb.Principal{
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_SafeRegex{SafeRegex: &v3matcherpb.RegexMatcher{Regex: ".+"}},
+											}},
+										}},
+										{Identifier: &v3rbacpb.Principal_Authenticated_{
+											Authenticated: &v3rbacpb.Principal_Authenticated{PrincipalName: &v3matcherpb.StringMatcher{
+												MatchPattern: &v3matcherpb.StringMatcher_Exact{Exact: ""},
+											}},
+										}},
+									},
+								}}},
+							},
+							Permissions: []*v3rbacpb.Permission{
+								{Rule: &v3rbacpb.Permission_Any{Any: true}},
+							},
+						},
+					},
+					AuditLoggingOptions: &v3rbacpb.RBAC_AuditLoggingOptions{
+						AuditCondition: v3rbacpb.RBAC_AuditLoggingOptions_ON_DENY,
+						LoggerConfigs: []*v3rbacpb.RBAC_AuditLoggingOptions_AuditLoggerConfig{
+							{AuditLogger: &v3corepb.TypedExtensionConfig{Name: "stdout_logger", TypedConfig: anyPbHelper(t, map[string]any{}, "stdout_logger")},
+								IsOptional: false,
+							},
+						},
+					},
 				},
 			},
 		},
@@ -258,16 +948,105 @@ func TestTranslatePolicy(t *testing.T) {
 			}`,
 			wantErr: `"allow_rules" 0: "headers" 0: unsupported "key" :method`,
 		},
+		"bad audit condition": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "ABC",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": {},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantErr: `failed to parse AuditCondition ABC`,
+		},
+		"bad audit logger config": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "NONE",
+					"audit_loggers": [
+						{
+							"name": "stdout_logger",
+							"config": "abc",
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantErr: `failed to unmarshal policy`,
+		},
+		"missing audit logger name": {
+			authzPolicy: `{
+				"name": "authz",
+				"allow_rules": [
+				{
+					"name": "allow_authenticated",
+					"source": {
+						"principals":["*", ""]
+					}
+				}],
+				"audit_logging_options": {
+					"audit_condition": "NONE",
+					"audit_loggers": [
+						{
+							"name": "",
+							"config": {},
+							"is_optional": false
+						}
+					]
+				}
+			}`,
+			wantErr: `missing required field: name`,
+		},
 	}
 	for name, test := range tests {
 		t.Run(name, func(t *testing.T) {
-			gotPolicies, gotErr := translatePolicy(test.authzPolicy)
+			gotPolicies, gotPolicyName, gotErr := translatePolicy(test.authzPolicy)
 			if gotErr != nil && !strings.HasPrefix(gotErr.Error(), test.wantErr) {
 				t.Fatalf("unexpected error\nwant:%v\ngot:%v", test.wantErr, gotErr)
 			}
 			if diff := cmp.Diff(gotPolicies, test.wantPolicies, protocmp.Transform()); diff != "" {
 				t.Fatalf("unexpected policy\ndiff (-want +got):\n%s", diff)
 			}
+			if test.wantPolicyName != "" && gotPolicyName != test.wantPolicyName {
+				t.Fatalf("unexpected policy name\nwant:%v\ngot:%v", test.wantPolicyName, gotPolicyName)
+			}
 		})
 	}
 }
+
+func anyPbHelper(t *testing.T, in map[string]any, name string) *anypb.Any {
+	t.Helper()
+	pb, err := structpb.NewStruct(in)
+	typedStruct := &v1xdsudpatypepb.TypedStruct{
+		TypeUrl: typeURLPrefix + name,
+		Value:   pb,
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+	customConfig, err := anypb.New(typedStruct)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return customConfig
+}

authz/sdk_end2end_test.go

@@ -1,307 +0,0 @@
-/*
- *
- * 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"
-	"io"
-	"net"
-	"testing"
-	"time"
-
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/authz"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/status"
-	pb "google.golang.org/grpc/test/grpc_testing"
-)
-
-type testServer struct {
-	pb.UnimplementedTestServiceServer
-}
-
-func (s *testServer) UnaryCall(ctx context.Context, req *pb.SimpleRequest) (*pb.SimpleResponse, error) {
-	return &pb.SimpleResponse{}, nil
-}
-
-func (s *testServer) StreamingInputCall(stream pb.TestService_StreamingInputCallServer) error {
-	for {
-		_, err := stream.Recv()
-		if err == io.EOF {
-			return stream.SendAndClose(&pb.StreamingInputCallResponse{})
-		}
-		if err != nil {
-			return err
-		}
-	}
-}
-
-func TestSDKEnd2End(t *testing.T) {
-	tests := map[string]struct {
-		authzPolicy    string
-		md             metadata.MD
-		wantStatusCode codes.Code
-		wantErr        string
-	}{
-		"DeniesRpcRequestMatchInDenyNoMatchInAllow": {
-			authzPolicy: `{
-				"name": "authz",
-				"allow_rules": 
-				[
-					{
-						"name": "allow_StreamingOutputCall",
-						"request": {
-							"paths":
-							[
-								"/grpc.testing.TestService/StreamingOutputCall"
-							]
-						}
-					}
-				],
-				"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-abc", "val-abc"),
-			wantStatusCode: codes.PermissionDenied,
-			wantErr:        "unauthorized RPC request rejected",
-		},
-		"DeniesRpcRequestMatchInDenyAndAllow": {
-			authzPolicy: `{
-				"name": "authz",
-				"allow_rules":
-				[
-					{
-						"name": "allow_TestServiceCalls",
-						"request": {
-							"paths":
-							[
-								"/grpc.testing.TestService/*"
-							]
-						}
-					}
-				],
-				"deny_rules":
-				[
-					{
-						"name": "deny_TestServiceCalls",
-						"request": {
-							"paths":
-							[
-								"/grpc.testing.TestService/*"
-							]
-						}
-					}
-				]
-			}`,
-			wantStatusCode: codes.PermissionDenied,
-			wantErr:        "unauthorized RPC request rejected",
-		},
-		"AllowsRpcRequestNoMatchInDenyMatchInAllow": {
-			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"),
-			wantStatusCode: codes.OK,
-		},
-		"AllowsRpcRequestNoMatchInDenyAndAllow": {
-			authzPolicy: `{
-				"name": "authz",
-				"allow_rules":
-				[
-					{
-						"name": "allow_some_user",
-						"source": {
-							"principals":
-							[
-								"some_user"
-							]
-						}
-					}
-				],
-				"deny_rules":
-				[
-					{
-						"name": "deny_StreamingOutputCall",
-						"request": {
-							"paths":
-							[
-								"/grpc.testing.TestService/StreamingOutputCall"
-							]
-						}
-					}
-				]
-			}`,
-			wantStatusCode: codes.PermissionDenied,
-			wantErr:        "unauthorized RPC request rejected",
-		},
-		"AllowsRpcRequestEmptyDenyMatchInAllow": {
-			authzPolicy: `{
-				"name": "authz",
-				"allow_rules":
-				[
-					{
-						"name": "allow_UnaryCall",
-						"request":
-						{
-							"paths":
-							[
-								"/grpc.testing.TestService/UnaryCall"
-							]
-						}
-					},
-					{
-						"name": "allow_StreamingInputCall",
-						"request":
-						{
-							"paths":
-							[
-								"/grpc.testing.TestService/StreamingInputCall"
-							]
-						}
-					}
-				]
-			}`,
-			wantStatusCode: codes.OK,
-		},
-		"DeniesRpcRequestEmptyDenyNoMatchInAllow": {
-			authzPolicy: `{
-				"name": "authz",
-				"allow_rules":
-				[
-					{
-						"name": "allow_StreamingOutputCall",
-						"request": 
-						{
-							"paths":
-							[
-								"/grpc.testing.TestService/StreamingOutputCall"
-							]
-						}
-					}
-				]
-			}`,
-			wantStatusCode: codes.PermissionDenied,
-			wantErr:        "unauthorized RPC request rejected",
-		},
-	}
-	for name, test := range tests {
-		t.Run(name, func(t *testing.T) {
-			// Start a gRPC server with SDK unary and stream server interceptors.
-			i, _ := authz.NewStatic(test.authzPolicy)
-			lis, err := net.Listen("tcp", "localhost:0")
-			if err != nil {
-				t.Fatalf("error listening: %v", err)
-			}
-			s := grpc.NewServer(
-				grpc.ChainUnaryInterceptor(i.UnaryInterceptor),
-				grpc.ChainStreamInterceptor(i.StreamInterceptor))
-			pb.RegisterTestServiceServer(s, &testServer{})
-			go s.Serve(lis)
-
-			// Establish a connection to the server.
-			clientConn, err := grpc.Dial(lis.Addr().String(), grpc.WithInsecure())
-			if err != nil {
-				t.Fatalf("grpc.Dial(%v) failed: %v", lis.Addr().String(), err)
-			}
-			defer clientConn.Close()
-			client := pb.NewTestServiceClient(clientConn)
-
-			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, &pb.SimpleRequest{})
-			if got := status.Convert(err); got.Code() != test.wantStatusCode || got.Message() != test.wantErr {
-				t.Fatalf("[UnaryCall] error want:{%v %v} got:{%v %v}", test.wantStatusCode, test.wantErr, got.Code(), got.Message())
-			}
-
-			// Verifying authorization decision for Streaming RPC.
-			stream, err := client.StreamingInputCall(ctx)
-			if err != nil {
-				t.Fatalf("failed StreamingInputCall err: %v", err)
-			}
-			req := &pb.StreamingInputCallRequest{
-				Payload: &pb.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.wantStatusCode || got.Message() != test.wantErr {
-				t.Fatalf("[StreamingCall] error want:{%v %v} got:{%v %v}", test.wantStatusCode, test.wantErr, got.Code(), got.Message())
-			}
-		})
-	}
-}

authz/sdk_server_interceptors.go

@@ -1,75 +0,0 @@
-/*
- * 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 (
-	"context"
-
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/internal/xds/rbac"
-	"google.golang.org/grpc/status"
-)
-
-// 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, err := translatePolicy(authzPolicy)
-	if err != nil {
-		return nil, err
-	}
-	chainEngine, err := rbac.NewChainEngine(rbacs)
-	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 interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-	err := i.engines.IsAuthorized(ctx)
-	if err != nil {
-		if status.Code(err) == codes.PermissionDenied {
-			return nil, status.Errorf(codes.PermissionDenied, "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 interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-	err := i.engines.IsAuthorized(ss.Context())
-	if err != nil {
-		if status.Code(err) == codes.PermissionDenied {
-			return status.Errorf(codes.PermissionDenied, "unauthorized RPC request rejected")
-		}
-		return err
-	}
-	return handler(srv, ss)
-}

authz/sdk_server_interceptors_test.go

@@ -1,56 +0,0 @@
-/*
- *
- * 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 (
-	"testing"
-
-	"google.golang.org/grpc/authz"
-)
-
-func TestNewStatic(t *testing.T) {
-	tests := map[string]struct {
-		authzPolicy string
-		wantErr     bool
-	}{
-		"InvalidPolicyFailsToCreateInterceptor": {
-			authzPolicy: `{}`,
-			wantErr:     true,
-		},
-		"ValidPolicyCreatesInterceptor": {
-			authzPolicy: `{		
-				"name": "authz",
-				"allow_rules": 
-				[
-					{
-						"name": "allow_all"
-					}
-				]
-			}`,
-			wantErr: false,
-		},
-	}
-	for name, test := range tests {
-		t.Run(name, func(t *testing.T) {
-			if _, err := authz.NewStatic(test.authzPolicy); (err != nil) != test.wantErr {
-				t.Fatalf("NewStatic(%v) returned err: %v, want err: %v", test.authzPolicy, err, test.wantErr)
-			}
-		})
-	}
-}

backoff.go

@@ -48,7 +48,7 @@ type BackoffConfig struct {
 // here for more details:
 // https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
 //
-// Experimental
+// # Experimental
 //
 // Notice: This type is EXPERIMENTAL and may be changed or removed in a
 // later release.

backoff/backoff.go

@@ -39,7 +39,7 @@ type Config struct {
 	MaxDelay time.Duration
 }
 
-// DefaultConfig is a backoff configuration with the default values specfied
+// 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

balancer/balancer.go

@@ -27,8 +27,11 @@ import (
 	"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/metadata"
 	"google.golang.org/grpc/resolver"
@@ -38,6 +41,8 @@ import (
 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
@@ -50,7 +55,14 @@ var (
 // 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) {
-	m[strings.ToLower(b.Name())] = b
+	name := strings.ToLower(b.Name())
+	if name != b.Name() {
+		// TODO: Skip the use of strings.ToLower() to index the map after v1.59
+		// is released to switch to case sensitive balancer registry. Also,
+		// remove this warning and update the docstrings for Register and Get.
+		logger.Warningf("Balancer registered with name %q. grpc-go will be switching to case sensitive balancer registries soon", b.Name())
+	}
+	m[name] = b
 }
 
 // unregisterForTesting deletes the balancer with the given name from the
@@ -63,54 +75,26 @@ func unregisterForTesting(name string) {
 
 func init() {
 	internal.BalancerUnregister = unregisterForTesting
+	internal.ConnectedAddress = connectedAddress
+	internal.SetConnectedAddress = setConnectedAddress
 }
 
 // Get returns the resolver builder registered with the given name.
 // Note that the compare is done in a case-insensitive fashion.
 // If no builder is register with the name, nil will be returned.
 func Get(name string) Builder {
+	if strings.ToLower(name) != name {
+		// TODO: Skip the use of strings.ToLower() to index the map after v1.59
+		// is released to switch to case sensitive balancer registry. Also,
+		// remove this warning and update the docstrings for Register and Get.
+		logger.Warningf("Balancer retrieved for name %q. grpc-go will be switching to case sensitive balancer registries soon", name)
+	}
 	if b, ok := m[strings.ToLower(name)]; ok {
 		return b
 	}
 	return nil
 }
 
-// A SubConn represents a single connection to a gRPC backend service.
-//
-// Each SubConn contains a list of addresses.
-//
-// All SubConns start in IDLE, and will not try to connect. To trigger the
-// connecting, Balancers must call Connect.  If a connection re-enters IDLE,
-// Balancers must call Connect again to trigger a new connection attempt.
-//
-// gRPC will try to connect to the addresses in sequence, and stop trying the
-// remainder once the first connection is successful. If an attempt to connect
-// to all addresses encounters an error, the SubConn will enter
-// TRANSIENT_FAILURE for a backoff period, and then transition to IDLE.
-//
-// Once established, if a connection is lost, the SubConn will transition
-// directly to IDLE.
-//
-// This interface is to be implemented by gRPC. Users should not need their own
-// implementation of this interface. For situations like testing, any
-// implementations should embed this interface. This allows gRPC to add new
-// methods to this interface.
-type SubConn interface {
-	// UpdateAddresses updates the addresses used in this SubConn.
-	// gRPC checks if currently-connected address is still in the new list.
-	// If it's in the list, the connection will be kept.
-	// If it's not in the list, the connection will gracefully closed, and
-	// a new connection will be created.
-	//
-	// This will trigger a state transition for the SubConn.
-	//
-	// Deprecated: This method is now part of the ClientConn interface and will
-	// eventually be removed from here.
-	UpdateAddresses([]resolver.Address)
-	// Connect starts the connecting for this SubConn.
-	Connect()
-}
-
 // NewSubConnOptions contains options to create new SubConn.
 type NewSubConnOptions struct {
 	// CredsBundle is the credentials bundle that will be used in the created
@@ -123,6 +107,11 @@ type NewSubConnOptions struct {
 	// 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.
@@ -140,20 +129,35 @@ type State struct {
 // 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 will trigger a state transition for the SubConn.
+	// 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
@@ -170,29 +174,47 @@ type ClientConn interface {
 	//
 	// 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 credential the Balancer implementation can
-	// use to dial to a remote load balancer server. The Balancer implementations
-	// can ignore this if it does not need to talk to another party securely.
+	// 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 that the Balancer can use.
+	// 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 the Balancer implementation can use to dial
-	// to a remote load balancer server. The Balancer implementations
-	// can ignore this if it doesn't need to talk to remote balancer.
+	// 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)
-	// ChannelzParentID is the entity parent's channelz unique identification number.
-	ChannelzParentID int64
+	// 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 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
 }
 
@@ -236,8 +258,8 @@ type DoneInfo struct {
 	// ServerLoad is the load received from server. It's usually sent as part of
 	// trailing metadata.
 	//
-	// The only supported type now is *orca_v1.LoadReport.
-	ServerLoad interface{}
+	// The only supported type now is *orca_v3.LoadReport.
+	ServerLoad any
 }
 
 var (
@@ -266,6 +288,14 @@ type PickResult struct {
 	// 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
@@ -322,10 +352,18 @@ type Balancer interface {
 	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 required to call
-	// ClientConn.RemoveSubConn for its existing SubConns.
+	// 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
@@ -333,8 +371,8 @@ type Balancer interface {
 // the ClientConn is idle.  If unimplemented, ClientConn.Connect will cause
 // all SubConns to connect.
 //
-// Notice: it will be required for all balancers to implement this in a future
-// release.
+// 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
@@ -342,15 +380,6 @@ type ExitIdler interface {
 	ExitIdle()
 }
 
-// SubConnState describes the state of a SubConn.
-type SubConnState struct {
-	// ConnectivityState is the connectivity state of the SubConn.
-	ConnectivityState connectivity.State
-	// ConnectionError is set if the ConnectivityState is TransientFailure,
-	// describing the reason the SubConn failed.  Otherwise, it is nil.
-	ConnectionError error
-}
-
 // ClientConnState describes the state of a ClientConn relevant to the
 // balancer.
 type ClientConnState struct {
@@ -363,56 +392,3 @@ type ClientConnState struct {
 // ErrBadResolverState may be returned by UpdateClientConnState to indicate a
 // problem with the provided name resolver data.
 var ErrBadResolverState = errors.New("bad resolver state")
-
-// 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 TransientFailure, the aggregated state is Transient Failure;
-//  - Else if at least one SubConn is Idle, the aggregated state is Idle;
-//  - Else there are no subconns and 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
-		}
-	}
-
-	// Evaluate.
-	if cse.numReady > 0 {
-		return connectivity.Ready
-	}
-	if cse.numConnecting > 0 {
-		return connectivity.Connecting
-	}
-	if cse.numTransientFailure > 0 {
-		return connectivity.TransientFailure
-	}
-	if cse.numIdle > 0 {
-		return connectivity.Idle
-	}
-	return connectivity.TransientFailure
-}

balancer/base/balancer.go

@@ -22,7 +22,6 @@ import (
 	"errors"
 	"fmt"
 
-	"google.golang.org/grpc/attributes"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
@@ -37,15 +36,16 @@ type baseBuilder struct {
 	config        Config
 }
 
-func (bb *baseBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) balancer.Balancer {
+func (bb *baseBuilder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
 	bal := &baseBalancer{
 		cc:            cc,
 		pickerBuilder: bb.pickerBuilder,
 
-		subConns: make(map[resolver.Address]subConnInfo),
+		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
@@ -58,11 +58,6 @@ func (bb *baseBuilder) Name() string {
 	return bb.name
 }
 
-type subConnInfo struct {
-	subConn balancer.SubConn
-	attrs   *attributes.Attributes
-}
-
 type baseBalancer struct {
 	cc            balancer.ClientConn
 	pickerBuilder PickerBuilder
@@ -70,7 +65,7 @@ type baseBalancer struct {
 	csEvltr *balancer.ConnectivityStateEvaluator
 	state   connectivity.State
 
-	subConns map[resolver.Address]subConnInfo // `attributes` is stripped from the keys of this map (the addresses)
+	subConns *resolver.AddressMapV2[balancer.SubConn]
 	scStates map[balancer.SubConn]connectivity.State
 	picker   balancer.Picker
 	config   Config
@@ -81,7 +76,7 @@ type baseBalancer struct {
 
 func (b *baseBalancer) ResolverError(err error) {
 	b.resolverErr = err
-	if len(b.subConns) == 0 {
+	if b.subConns.Len() == 0 {
 		b.state = connectivity.TransientFailure
 	}
 
@@ -105,65 +100,48 @@ func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 	// 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 := make(map[resolver.Address]struct{})
+	addrsSet := resolver.NewAddressMapV2[any]()
 	for _, a := range s.ResolverState.Addresses {
-		// Strip attributes from addresses before using them as map keys. So
-		// that when two addresses only differ in attributes pointers (but with
-		// the same attribute content), they are considered the same address.
-		//
-		// Note that this doesn't handle the case where the attribute content is
-		// different. So if users want to set different attributes to create
-		// duplicate connections to the same backend, it doesn't work. This is
-		// fine for now, because duplicate is done by setting Metadata today.
-		//
-		// TODO: read attributes to handle duplicate connections.
-		aNoAttrs := a
-		aNoAttrs.Attributes = nil
-		addrsSet[aNoAttrs] = struct{}{}
-		if scInfo, ok := b.subConns[aNoAttrs]; !ok {
+		addrsSet.Set(a, nil)
+		if _, ok := b.subConns.Get(a); !ok {
 			// a is a new address (not existing in b.subConns).
-			//
-			// When creating SubConn, the original address with attributes is
-			// passed through. So that connection configurations in attributes
-			// (like creds) will be used.
-			sc, err := b.cc.NewSubConn([]resolver.Address{a}, balancer.NewSubConnOptions{HealthCheckEnabled: b.config.HealthCheck})
+			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[aNoAttrs] = subConnInfo{subConn: sc, attrs: a.Attributes}
+			b.subConns.Set(a, sc)
 			b.scStates[sc] = connectivity.Idle
 			b.csEvltr.RecordTransition(connectivity.Shutdown, connectivity.Idle)
 			sc.Connect()
-		} else {
-			// Always update the subconn's address in case the attributes
-			// changed.
-			//
-			// The SubConn does a reflect.DeepEqual of the new and old
-			// addresses. So this is a noop if the current address is the same
-			// as the old one (including attributes).
-			scInfo.attrs = a.Attributes
-			b.subConns[aNoAttrs] = scInfo
-			b.cc.UpdateAddresses(scInfo.subConn, []resolver.Address{a})
 		}
 	}
-	for a, scInfo := range b.subConns {
+	for _, a := range b.subConns.Keys() {
+		sc, _ := b.subConns.Get(a)
 		// a was removed by resolver.
-		if _, ok := addrsSet[a]; !ok {
-			b.cc.RemoveSubConn(scInfo.subConn)
-			delete(b.subConns, a)
+		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.
+			// 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 an resolver error, so when
+	// 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
 }
 
@@ -183,8 +161,8 @@ func (b *baseBalancer) mergeErrors() error {
 
 // 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.
+//   - 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())
@@ -193,16 +171,21 @@ func (b *baseBalancer) regeneratePicker() {
 	readySCs := make(map[balancer.SubConn]SubConnInfo)
 
 	// Filter out all ready SCs from full subConn map.
-	for addr, scInfo := range b.subConns {
-		if st, ok := b.scStates[scInfo.subConn]; ok && st == connectivity.Ready {
-			addr.Attributes = scInfo.attrs
-			readySCs[scInfo.subConn] = SubConnInfo{Address: addr}
+	for _, addr := range b.subConns.Keys() {
+		sc, _ := b.subConns.Get(addr)
+		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)
@@ -229,8 +212,8 @@ func (b *baseBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.Su
 	case connectivity.Idle:
 		sc.Connect()
 	case connectivity.Shutdown:
-		// When an address was removed by resolver, b called RemoveSubConn but
-		// kept the sc's state in scStates. Remove state for this sc here.
+		// 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.
@@ -251,7 +234,7 @@ func (b *baseBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.Su
 }
 
 // Close is a nop because base balancer doesn't have internal state to clean up,
-// and it doesn't need to call RemoveSubConn for the SubConns.
+// and it doesn't need to call Shutdown for the SubConns.
 func (b *baseBalancer) Close() {
 }
 
@@ -274,6 +257,6 @@ type errPicker struct {
 	err error // Pick() always returns this err.
 }
 
-func (p *errPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+func (p *errPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
 	return balancer.PickResult{}, p.err
 }

balancer/base/balancer_test.go

@@ -19,7 +19,9 @@
 package base
 
 import (
+	"context"
 	"testing"
+	"time"
 
 	"google.golang.org/grpc/attributes"
 	"google.golang.org/grpc/balancer"
@@ -38,12 +40,24 @@ func (c *testClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewS
 
 func (c *testClientConn) UpdateState(balancer.State) {}
 
-type testSubConn struct{}
+type testSubConn struct {
+	balancer.SubConn
+	updateState func(balancer.SubConnState)
+}
 
-func (sc *testSubConn) UpdateAddresses(addresses []resolver.Address) {}
+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)
@@ -54,36 +68,12 @@ func (p *testPickBuilder) Build(info PickerBuildInfo) balancer.Picker {
 	return nil
 }
 
-func TestBaseBalancerStripAttributes(t *testing.T) {
-	b := (&baseBuilder{}).Build(&testClientConn{
-		newSubConn: func(addrs []resolver.Address, _ balancer.NewSubConnOptions) (balancer.SubConn, error) {
-			for _, addr := range addrs {
-				if addr.Attributes == nil {
-					t.Errorf("in NewSubConn, got address %+v with nil attributes, want not nil", addr)
-				}
-			}
-			return &testSubConn{}, nil
-		},
-	}, balancer.BuildOptions{}).(*baseBalancer)
-
-	b.UpdateClientConnState(balancer.ClientConnState{
-		ResolverState: resolver.State{
-			Addresses: []resolver.Address{
-				{Addr: "1.1.1.1", Attributes: &attributes.Attributes{}},
-				{Addr: "2.2.2.2", Attributes: &attributes.Attributes{}},
-			},
-		},
-	})
-
-	for addr := range b.subConns {
-		if addr.Attributes != nil {
-			t.Errorf("in b.subConns, got address %+v with not nil attributes, want nil", addr)
-		}
-	}
-}
-
 func TestBaseBalancerReserveAttributes(t *testing.T) {
-	var v = func(info PickerBuildInfo) {
+	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 {
@@ -102,8 +92,8 @@ func TestBaseBalancerReserveAttributes(t *testing.T) {
 	}
 	pickBuilder := &testPickBuilder{validate: v}
 	b := (&baseBuilder{pickerBuilder: pickBuilder}).Build(&testClientConn{
-		newSubConn: func(addrs []resolver.Address, _ balancer.NewSubConnOptions) (balancer.SubConn, error) {
-			return &testSubConn{}, nil
+		newSubConn: func(_ []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+			return &testSubConn{updateState: opts.StateListener}, nil
 		},
 	}, balancer.BuildOptions{}).(*baseBalancer)
 
@@ -115,8 +105,18 @@ func TestBaseBalancerReserveAttributes(t *testing.T) {
 			},
 		},
 	})
+	select {
+	case <-validated:
+	case <-ctx.Done():
+		t.Fatalf("timed out waiting for UpdateClientConnState to call picker.Build")
+	}
 
 	for sc := range b.scStates {
-		b.UpdateSubConnState(sc, balancer.SubConnState{ConnectivityState: connectivity.Ready, ConnectionError: nil})
+		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")
+		}
 	}
 }

balancer/conn_state_evaluator.go

@@ -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
+}

balancer/conn_state_evaluator_test.go

@@ -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)
+			}
+		})
+	}
+}

balancer/endpointsharding/endpointsharding.go

@@ -0,0 +1,389 @@
+/*
+ *
+ * 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 balancer state of a child along with the endpoint which
+// identifies the child balancer.
+type ChildState struct {
+	Endpoint resolver.Endpoint
+	State    balancer.State
+
+	// Balancer exposes only the ExitIdler interface of the child LB policy.
+	// Other methods of the child policy are called only by endpointsharding.
+	Balancer ExitIdler
+}
+
+// ExitIdler provides access to only the ExitIdle method of the child balancer.
+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()
+}
+
+// Options are the options to 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 {
+	es := &endpointSharding{
+		cc:           cc,
+		bOpts:        opts,
+		esOpts:       esOpts,
+		childBuilder: childBuilder,
+	}
+	es.children.Store(resolver.NewEndpointMap[*balancerWrapper]())
+	return es
+}
+
+// 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
+
+	// childMu synchronizes calls to any single child. It must be held for all
+	// calls into a child. To avoid deadlocks, do not acquire childMu while
+	// holding mu.
+	childMu  sync.Mutex
+	children atomic.Pointer[resolver.EndpointMap[*balancerWrapper]]
+
+	// inhibitChildUpdates is set during UpdateClientConnState/ResolverError
+	// calls (calls to children will each produce an update, only want one
+	// update).
+	inhibitChildUpdates atomic.Bool
+
+	// mu synchronizes access to the state stored in balancerWrappers in the
+	// children field. mu must not be held during calls into a child since
+	// synchronous calls back from the child may require taking mu, causing a
+	// deadlock. To avoid deadlocks, do not acquire childMu while holding mu.
+	mu sync.Mutex
+}
+
+// rotateEndpoints returns a slice of all the input endpoints rotated a random
+// amount.
+func rotateEndpoints(es []resolver.Endpoint) []resolver.Endpoint {
+	les := len(es)
+	if les == 0 {
+		return es
+	}
+	r := randIntN(les)
+	// Make a copy to avoid mutating data beyond the end of es.
+	ret := make([]resolver.Endpoint, les)
+	copy(ret, es[r:])
+	copy(ret[les-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. If any endpoint has no
+// addresses it will ignore that endpoint. Otherwise, returns first error found
+// from a child, but fully processes the new update.
+func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState) error {
+	es.childMu.Lock()
+	defer es.childMu.Unlock()
+
+	es.inhibitChildUpdates.Store(true)
+	defer func() {
+		es.inhibitChildUpdates.Store(false)
+		es.updateState()
+	}()
+	var ret error
+
+	children := es.children.Load()
+	newChildren := resolver.NewEndpointMap[*balancerWrapper]()
+
+	// Update/Create new children.
+	for _, endpoint := range rotateEndpoints(state.ResolverState.Endpoints) {
+		if _, ok := newChildren.Get(endpoint); ok {
+			// Endpoint child was already created, continue to avoid duplicate
+			// update.
+			continue
+		}
+		childBalancer, ok := children.Get(endpoint)
+		if ok {
+			// Endpoint attributes may have changed, update the stored endpoint.
+			es.mu.Lock()
+			childBalancer.childState.Endpoint = endpoint
+			es.mu.Unlock()
+		} else {
+			childBalancer = &balancerWrapper{
+				childState: ChildState{Endpoint: endpoint},
+				ClientConn: es.cc,
+				es:         es,
+			}
+			childBalancer.childState.Balancer = childBalancer
+			childBalancer.child = es.childBuilder(childBalancer, es.bOpts)
+		}
+		newChildren.Set(endpoint, childBalancer)
+		if err := childBalancer.updateClientConnStateLocked(balancer.ClientConnState{
+			BalancerConfig: state.BalancerConfig,
+			ResolverState: resolver.State{
+				Endpoints:  []resolver.Endpoint{endpoint},
+				Attributes: state.ResolverState.Attributes,
+			},
+		}); err != nil && ret == nil {
+			// Return first error found, and always commit full processing of
+			// updating children. If desired to process more specific errors
+			// across all endpoints, caller should make these specific
+			// validations, this is a current limitation for simplicity sake.
+			ret = err
+		}
+	}
+	// Delete old children that are no longer present.
+	for _, e := range children.Keys() {
+		child, _ := children.Get(e)
+		if _, ok := newChildren.Get(e); !ok {
+			child.closeLocked()
+		}
+	}
+	es.children.Store(newChildren)
+	if newChildren.Len() == 0 {
+		return balancer.ErrBadResolverState
+	}
+	return ret
+}
+
+// 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.childMu.Lock()
+	defer es.childMu.Unlock()
+	es.inhibitChildUpdates.Store(true)
+	defer func() {
+		es.inhibitChildUpdates.Store(false)
+		es.updateState()
+	}()
+	children := es.children.Load()
+	for _, child := range children.Values() {
+		child.resolverErrorLocked(err)
+	}
+}
+
+func (es *endpointSharding) UpdateSubConnState(balancer.SubConn, balancer.SubConnState) {
+	// UpdateSubConnState is deprecated.
+}
+
+func (es *endpointSharding) Close() {
+	es.childMu.Lock()
+	defer es.childMu.Unlock()
+	children := es.children.Load()
+	for _, child := range children.Values() {
+		child.closeLocked()
+	}
+}
+
+func (es *endpointSharding) ExitIdle() {
+	es.childMu.Lock()
+	defer es.childMu.Unlock()
+	for _, bw := range es.children.Load().Values() {
+		if !bw.isClosed {
+			bw.child.ExitIdle()
+		}
+	}
+}
+
+// updateState 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.
+func (es *endpointSharding) updateState() {
+	if es.inhibitChildUpdates.Load() {
+		return
+	}
+	var readyPickers, connectingPickers, idlePickers, transientFailurePickers []balancer.Picker
+
+	es.mu.Lock()
+	defer es.mu.Unlock()
+
+	children := es.children.Load()
+	childStates := make([]ChildState, 0, children.Len())
+
+	for _, child := range children.Values() {
+		childState := child.childState
+		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).
+	p := &pickerWithChildStates{
+		pickers:     pickers,
+		childStates: childStates,
+		next:        uint32(randIntN(len(pickers))),
+	}
+	es.cc.UpdateState(balancer.State{
+		ConnectivityState: aggState,
+		Picker:            p,
+	})
+}
+
+// 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
+}
+
+// balancerWrapper is a wrapper of a balancer. It ID's a child balancer by
+// endpoint, and persists recent child balancer state.
+type balancerWrapper 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.
+
+	// child contains the wrapped balancer. Access its methods only through
+	// methods on balancerWrapper to ensure proper synchronization
+	child               balancer.Balancer
+	balancer.ClientConn // embed to intercept UpdateState, doesn't deal with SubConns
+
+	es *endpointSharding
+
+	// Access to the following fields is guarded by es.mu.
+
+	childState ChildState
+	isClosed   bool
+}
+
+func (bw *balancerWrapper) UpdateState(state balancer.State) {
+	bw.es.mu.Lock()
+	bw.childState.State = state
+	bw.es.mu.Unlock()
+	if state.ConnectivityState == connectivity.Idle && !bw.es.esOpts.DisableAutoReconnect {
+		bw.ExitIdle()
+	}
+	bw.es.updateState()
+}
+
+// ExitIdle pings an IDLE child balancer to exit idle in a new goroutine to
+// avoid deadlocks due to synchronous balancer state updates.
+func (bw *balancerWrapper) ExitIdle() {
+	go func() {
+		bw.es.childMu.Lock()
+		if !bw.isClosed {
+			bw.child.ExitIdle()
+		}
+		bw.es.childMu.Unlock()
+	}()
+}
+
+// updateClientConnStateLocked delivers the ClientConnState to the child
+// balancer. Callers must hold the child mutex of the parent endpointsharding
+// balancer.
+func (bw *balancerWrapper) updateClientConnStateLocked(ccs balancer.ClientConnState) error {
+	return bw.child.UpdateClientConnState(ccs)
+}
+
+// closeLocked closes the child balancer. Callers must hold the child mutext of
+// the parent endpointsharding balancer.
+func (bw *balancerWrapper) closeLocked() {
+	bw.child.Close()
+	bw.isClosed = true
+}
+
+func (bw *balancerWrapper) resolverErrorLocked(err error) {
+	bw.child.ResolverError(err)
+}

balancer/endpointsharding/endpointsharding_ext_test.go

@@ -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/pickfirstleaf"
+	"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(pickfirstleaf.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(pickfirstleaf.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(pickfirstleaf.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)
+}

balancer/endpointsharding/endpointsharding_test.go

@@ -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)
+			}
+		})
+	}
+}

balancer/grpclb/grpc_lb_v1/load_balancer.pb.go

@@ -19,20 +19,20 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.25.0
-// 	protoc        v3.14.0
+// 	protoc-gen-go v1.36.6
+// 	protoc        v5.27.1
 // source: grpc/lb/v1/load_balancer.proto
 
 package grpc_lb_v1
 
 import (
-	proto "github.com/golang/protobuf/proto"
 	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 (
@@ -42,28 +42,22 @@ const (
 	_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
 )
 
-// This is a compile-time assertion that a sufficiently up-to-date version
-// of the legacy proto package is being used.
-const _ = proto.ProtoPackageIsVersion4
-
 type LoadBalanceRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	// Types that are assignable to LoadBalanceRequestType:
+	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{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[0]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -74,7 +68,7 @@ func (*LoadBalanceRequest) ProtoMessage() {}
 
 func (x *LoadBalanceRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[0]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -89,23 +83,27 @@ func (*LoadBalanceRequest) Descriptor() ([]byte, []int) {
 	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{0}
 }
 
-func (m *LoadBalanceRequest) GetLoadBalanceRequestType() isLoadBalanceRequest_LoadBalanceRequestType {
-	if m != nil {
-		return m.LoadBalanceRequestType
+func (x *LoadBalanceRequest) GetLoadBalanceRequestType() isLoadBalanceRequest_LoadBalanceRequestType {
+	if x != nil {
+		return x.LoadBalanceRequestType
 	}
 	return nil
 }
 
 func (x *LoadBalanceRequest) GetInitialRequest() *InitialLoadBalanceRequest {
-	if x, ok := x.GetLoadBalanceRequestType().(*LoadBalanceRequest_InitialRequest); ok {
-		return x.InitialRequest
+	if x != nil {
+		if x, ok := x.LoadBalanceRequestType.(*LoadBalanceRequest_InitialRequest); ok {
+			return x.InitialRequest
+		}
 	}
 	return nil
 }
 
 func (x *LoadBalanceRequest) GetClientStats() *ClientStats {
-	if x, ok := x.GetLoadBalanceRequestType().(*LoadBalanceRequest_ClientStats); ok {
-		return x.ClientStats
+	if x != nil {
+		if x, ok := x.LoadBalanceRequestType.(*LoadBalanceRequest_ClientStats); ok {
+			return x.ClientStats
+		}
 	}
 	return nil
 }
@@ -130,24 +128,21 @@ func (*LoadBalanceRequest_InitialRequest) isLoadBalanceRequest_LoadBalanceReques
 func (*LoadBalanceRequest_ClientStats) isLoadBalanceRequest_LoadBalanceRequestType() {}
 
 type InitialLoadBalanceRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	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"`
+	Name          string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *InitialLoadBalanceRequest) Reset() {
 	*x = InitialLoadBalanceRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[1]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -158,7 +153,7 @@ func (*InitialLoadBalanceRequest) ProtoMessage() {}
 
 func (x *InitialLoadBalanceRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[1]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -182,23 +177,20 @@ func (x *InitialLoadBalanceRequest) GetName() string {
 
 // Contains the number of calls finished for a particular load balance token.
 type ClientStatsPerToken struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	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"`
+	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{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[2]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -209,7 +201,7 @@ func (*ClientStatsPerToken) ProtoMessage() {}
 
 func (x *ClientStatsPerToken) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[2]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -241,10 +233,7 @@ func (x *ClientStatsPerToken) GetNumCalls() int64 {
 // 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
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	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.
@@ -258,15 +247,15 @@ type ClientStats struct {
 	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{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[3]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -277,7 +266,7 @@ func (*ClientStats) ProtoMessage() {}
 
 func (x *ClientStats) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[3]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -335,24 +324,22 @@ func (x *ClientStats) GetCallsFinishedWithDrop() []*ClientStatsPerToken {
 }
 
 type LoadBalanceResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	// Types that are assignable to LoadBalanceResponseType:
+	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{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[4]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -363,7 +350,7 @@ func (*LoadBalanceResponse) ProtoMessage() {}
 
 func (x *LoadBalanceResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[4]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -378,30 +365,36 @@ func (*LoadBalanceResponse) Descriptor() ([]byte, []int) {
 	return file_grpc_lb_v1_load_balancer_proto_rawDescGZIP(), []int{4}
 }
 
-func (m *LoadBalanceResponse) GetLoadBalanceResponseType() isLoadBalanceResponse_LoadBalanceResponseType {
-	if m != nil {
-		return m.LoadBalanceResponseType
+func (x *LoadBalanceResponse) GetLoadBalanceResponseType() isLoadBalanceResponse_LoadBalanceResponseType {
+	if x != nil {
+		return x.LoadBalanceResponseType
 	}
 	return nil
 }
 
 func (x *LoadBalanceResponse) GetInitialResponse() *InitialLoadBalanceResponse {
-	if x, ok := x.GetLoadBalanceResponseType().(*LoadBalanceResponse_InitialResponse); ok {
-		return x.InitialResponse
+	if x != nil {
+		if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_InitialResponse); ok {
+			return x.InitialResponse
+		}
 	}
 	return nil
 }
 
 func (x *LoadBalanceResponse) GetServerList() *ServerList {
-	if x, ok := x.GetLoadBalanceResponseType().(*LoadBalanceResponse_ServerList); ok {
-		return x.ServerList
+	if x != nil {
+		if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_ServerList); ok {
+			return x.ServerList
+		}
 	}
 	return nil
 }
 
 func (x *LoadBalanceResponse) GetFallbackResponse() *FallbackResponse {
-	if x, ok := x.GetLoadBalanceResponseType().(*LoadBalanceResponse_FallbackResponse); ok {
-		return x.FallbackResponse
+	if x != nil {
+		if x, ok := x.LoadBalanceResponseType.(*LoadBalanceResponse_FallbackResponse); ok {
+			return x.FallbackResponse
+		}
 	}
 	return nil
 }
@@ -434,18 +427,16 @@ func (*LoadBalanceResponse_ServerList) isLoadBalanceResponse_LoadBalanceResponse
 func (*LoadBalanceResponse_FallbackResponse) isLoadBalanceResponse_LoadBalanceResponseType() {}
 
 type FallbackResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *FallbackResponse) Reset() {
 	*x = FallbackResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[5]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -456,7 +447,7 @@ func (*FallbackResponse) ProtoMessage() {}
 
 func (x *FallbackResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[5]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -472,23 +463,20 @@ func (*FallbackResponse) Descriptor() ([]byte, []int) {
 }
 
 type InitialLoadBalanceResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	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{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[6]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -499,7 +487,7 @@ func (*InitialLoadBalanceResponse) ProtoMessage() {}
 
 func (x *InitialLoadBalanceResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[6]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -522,24 +510,21 @@ func (x *InitialLoadBalanceResponse) GetClientStatsReportInterval() *durationpb.
 }
 
 type ServerList struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	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"`
+	Servers       []*Server `protobuf:"bytes,1,rep,name=servers,proto3" json:"servers,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *ServerList) Reset() {
 	*x = ServerList{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[7]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -550,7 +535,7 @@ func (*ServerList) ProtoMessage() {}
 
 func (x *ServerList) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[7]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -575,10 +560,7 @@ func (x *ServerList) GetServers() []*Server {
 // Contains server information. When the drop field is not true, use the other
 // fields.
 type Server struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	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"`
@@ -595,16 +577,16 @@ type Server struct {
 	// 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"`
+	Drop          bool `protobuf:"varint,4,opt,name=drop,proto3" json:"drop,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *Server) Reset() {
 	*x = Server{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[8]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	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 {
@@ -615,7 +597,7 @@ func (*Server) ProtoMessage() {}
 
 func (x *Server) ProtoReflect() protoreflect.Message {
 	mi := &file_grpc_lb_v1_load_balancer_proto_msgTypes[8]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -660,130 +642,62 @@ func (x *Server) GetDrop() bool {
 
 var File_grpc_lb_v1_load_balancer_proto protoreflect.FileDescriptor
 
-var file_grpc_lb_v1_load_balancer_proto_rawDesc = []byte{
-	0x0a, 0x1e, 0x67, 0x72, 0x70, 0x63, 0x2f, 0x6c, 0x62, 0x2f, 0x76, 0x31, 0x2f, 0x6c, 0x6f, 0x61,
-	0x64, 0x5f, 0x62, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x72, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f,
-	0x12, 0x0a, 0x67, 0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x1a, 0x1e, 0x67, 0x6f,
-	0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x64, 0x75,
-	0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a, 0x1f, 0x67, 0x6f,
-	0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74, 0x69,
-	0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0xc1, 0x01,
-	0x0a, 0x12, 0x4c, 0x6f, 0x61, 0x64, 0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x71,
-	0x75, 0x65, 0x73, 0x74, 0x12, 0x50, 0x0a, 0x0f, 0x69, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x6c, 0x5f,
-	0x72, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x25, 0x2e,
-	0x67, 0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x2e, 0x49, 0x6e, 0x69, 0x74, 0x69,
-	0x61, 0x6c, 0x4c, 0x6f, 0x61, 0x64, 0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x71,
-	0x75, 0x65, 0x73, 0x74, 0x48, 0x00, 0x52, 0x0e, 0x69, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x6c, 0x52,
-	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x3c, 0x0a, 0x0c, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74,
-	0x5f, 0x73, 0x74, 0x61, 0x74, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x17, 0x2e, 0x67,
-	0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74,
-	0x53, 0x74, 0x61, 0x74, 0x73, 0x48, 0x00, 0x52, 0x0b, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x53,
-	0x74, 0x61, 0x74, 0x73, 0x42, 0x1b, 0x0a, 0x19, 0x6c, 0x6f, 0x61, 0x64, 0x5f, 0x62, 0x61, 0x6c,
-	0x61, 0x6e, 0x63, 0x65, 0x5f, 0x72, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x5f, 0x74, 0x79, 0x70,
-	0x65, 0x22, 0x2f, 0x0a, 0x19, 0x49, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x6c, 0x4c, 0x6f, 0x61, 0x64,
-	0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x12,
-	0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61,
-	0x6d, 0x65, 0x22, 0x60, 0x0a, 0x13, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74,
-	0x73, 0x50, 0x65, 0x72, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x12, 0x2c, 0x0a, 0x12, 0x6c, 0x6f, 0x61,
-	0x64, 0x5f, 0x62, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x18,
-	0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x10, 0x6c, 0x6f, 0x61, 0x64, 0x42, 0x61, 0x6c, 0x61, 0x6e,
-	0x63, 0x65, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x12, 0x1b, 0x0a, 0x09, 0x6e, 0x75, 0x6d, 0x5f, 0x63,
-	0x61, 0x6c, 0x6c, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x08, 0x6e, 0x75, 0x6d, 0x43,
-	0x61, 0x6c, 0x6c, 0x73, 0x22, 0xb0, 0x03, 0x0a, 0x0b, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x53,
-	0x74, 0x61, 0x74, 0x73, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d,
-	0x70, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65,
-	0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74,
-	0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x12, 0x2a,
-	0x0a, 0x11, 0x6e, 0x75, 0x6d, 0x5f, 0x63, 0x61, 0x6c, 0x6c, 0x73, 0x5f, 0x73, 0x74, 0x61, 0x72,
-	0x74, 0x65, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x0f, 0x6e, 0x75, 0x6d, 0x43, 0x61,
-	0x6c, 0x6c, 0x73, 0x53, 0x74, 0x61, 0x72, 0x74, 0x65, 0x64, 0x12, 0x2c, 0x0a, 0x12, 0x6e, 0x75,
-	0x6d, 0x5f, 0x63, 0x61, 0x6c, 0x6c, 0x73, 0x5f, 0x66, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x64,
-	0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x10, 0x6e, 0x75, 0x6d, 0x43, 0x61, 0x6c, 0x6c, 0x73,
-	0x46, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x64, 0x12, 0x5d, 0x0a, 0x2d, 0x6e, 0x75, 0x6d, 0x5f,
-	0x63, 0x61, 0x6c, 0x6c, 0x73, 0x5f, 0x66, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x64, 0x5f, 0x77,
-	0x69, 0x74, 0x68, 0x5f, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x5f, 0x66, 0x61, 0x69, 0x6c, 0x65,
-	0x64, 0x5f, 0x74, 0x6f, 0x5f, 0x73, 0x65, 0x6e, 0x64, 0x18, 0x06, 0x20, 0x01, 0x28, 0x03, 0x52,
-	0x26, 0x6e, 0x75, 0x6d, 0x43, 0x61, 0x6c, 0x6c, 0x73, 0x46, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65,
-	0x64, 0x57, 0x69, 0x74, 0x68, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x46, 0x61, 0x69, 0x6c, 0x65,
-	0x64, 0x54, 0x6f, 0x53, 0x65, 0x6e, 0x64, 0x12, 0x48, 0x0a, 0x21, 0x6e, 0x75, 0x6d, 0x5f, 0x63,
-	0x61, 0x6c, 0x6c, 0x73, 0x5f, 0x66, 0x69, 0x6e, 0x69, 0x73, 0x68, 0x65, 0x64, 0x5f, 0x6b, 0x6e,
-	0x6f, 0x77, 0x6e, 0x5f, 0x72, 0x65, 0x63, 0x65, 0x69, 0x76, 0x65, 0x64, 0x18, 0x07, 0x20, 0x01,
-	0x28, 0x03, 0x52, 0x1d, 0x6e, 0x75, 0x6d, 0x43, 0x61, 0x6c, 0x6c, 0x73, 0x46, 0x69, 0x6e, 0x69,
-	0x73, 0x68, 0x65, 0x64, 0x4b, 0x6e, 0x6f, 0x77, 0x6e, 0x52, 0x65, 0x63, 0x65, 0x69, 0x76, 0x65,
-	0x64, 0x12, 0x58, 0x0a, 0x18, 0x63, 0x61, 0x6c, 0x6c, 0x73, 0x5f, 0x66, 0x69, 0x6e, 0x69, 0x73,
-	0x68, 0x65, 0x64, 0x5f, 0x77, 0x69, 0x74, 0x68, 0x5f, 0x64, 0x72, 0x6f, 0x70, 0x18, 0x08, 0x20,
-	0x03, 0x28, 0x0b, 0x32, 0x1f, 0x2e, 0x67, 0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31,
-	0x2e, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x73, 0x50, 0x65, 0x72, 0x54,
-	0x6f, 0x6b, 0x65, 0x6e, 0x52, 0x15, 0x63, 0x61, 0x6c, 0x6c, 0x73, 0x46, 0x69, 0x6e, 0x69, 0x73,
-	0x68, 0x65, 0x64, 0x57, 0x69, 0x74, 0x68, 0x44, 0x72, 0x6f, 0x70, 0x4a, 0x04, 0x08, 0x04, 0x10,
-	0x05, 0x4a, 0x04, 0x08, 0x05, 0x10, 0x06, 0x22, 0x90, 0x02, 0x0a, 0x13, 0x4c, 0x6f, 0x61, 0x64,
-	0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12,
-	0x53, 0x0a, 0x10, 0x69, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x6c, 0x5f, 0x72, 0x65, 0x73, 0x70, 0x6f,
-	0x6e, 0x73, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x26, 0x2e, 0x67, 0x72, 0x70, 0x63,
-	0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x2e, 0x49, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x6c, 0x4c, 0x6f,
-	0x61, 0x64, 0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73,
-	0x65, 0x48, 0x00, 0x52, 0x0f, 0x69, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x6c, 0x52, 0x65, 0x73, 0x70,
-	0x6f, 0x6e, 0x73, 0x65, 0x12, 0x39, 0x0a, 0x0b, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x5f, 0x6c,
-	0x69, 0x73, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x67, 0x72, 0x70, 0x63,
-	0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x2e, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x4c, 0x69, 0x73,
-	0x74, 0x48, 0x00, 0x52, 0x0a, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x4c, 0x69, 0x73, 0x74, 0x12,
-	0x4b, 0x0a, 0x11, 0x66, 0x61, 0x6c, 0x6c, 0x62, 0x61, 0x63, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x70,
-	0x6f, 0x6e, 0x73, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c, 0x2e, 0x67, 0x72, 0x70,
-	0x63, 0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x2e, 0x46, 0x61, 0x6c, 0x6c, 0x62, 0x61, 0x63, 0x6b,
-	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x10, 0x66, 0x61, 0x6c, 0x6c,
-	0x62, 0x61, 0x63, 0x6b, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x1c, 0x0a, 0x1a,
-	0x6c, 0x6f, 0x61, 0x64, 0x5f, 0x62, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x5f, 0x72, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x22, 0x12, 0x0a, 0x10, 0x46, 0x61,
-	0x6c, 0x6c, 0x62, 0x61, 0x63, 0x6b, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x7e,
-	0x0a, 0x1a, 0x49, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x6c, 0x4c, 0x6f, 0x61, 0x64, 0x42, 0x61, 0x6c,
-	0x61, 0x6e, 0x63, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x5a, 0x0a, 0x1c,
-	0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x5f, 0x73, 0x74, 0x61, 0x74, 0x73, 0x5f, 0x72, 0x65, 0x70,
-	0x6f, 0x72, 0x74, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x76, 0x61, 0x6c, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x0b, 0x32, 0x19, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74,
-	0x6f, 0x62, 0x75, 0x66, 0x2e, 0x44, 0x75, 0x72, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x19, 0x63,
-	0x6c, 0x69, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x73, 0x52, 0x65, 0x70, 0x6f, 0x72, 0x74,
-	0x49, 0x6e, 0x74, 0x65, 0x72, 0x76, 0x61, 0x6c, 0x4a, 0x04, 0x08, 0x01, 0x10, 0x02, 0x22, 0x40,
-	0x0a, 0x0a, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x4c, 0x69, 0x73, 0x74, 0x12, 0x2c, 0x0a, 0x07,
-	0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x12, 0x2e,
-	0x67, 0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x2e, 0x53, 0x65, 0x72, 0x76, 0x65,
-	0x72, 0x52, 0x07, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x73, 0x4a, 0x04, 0x08, 0x03, 0x10, 0x04,
-	0x22, 0x83, 0x01, 0x0a, 0x06, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x12, 0x1d, 0x0a, 0x0a, 0x69,
-	0x70, 0x5f, 0x61, 0x64, 0x64, 0x72, 0x65, 0x73, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52,
-	0x09, 0x69, 0x70, 0x41, 0x64, 0x64, 0x72, 0x65, 0x73, 0x73, 0x12, 0x12, 0x0a, 0x04, 0x70, 0x6f,
-	0x72, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x12, 0x2c,
-	0x0a, 0x12, 0x6c, 0x6f, 0x61, 0x64, 0x5f, 0x62, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x5f, 0x74,
-	0x6f, 0x6b, 0x65, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x10, 0x6c, 0x6f, 0x61, 0x64,
-	0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x12, 0x12, 0x0a, 0x04,
-	0x64, 0x72, 0x6f, 0x70, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x04, 0x64, 0x72, 0x6f, 0x70,
-	0x4a, 0x04, 0x08, 0x05, 0x10, 0x06, 0x32, 0x62, 0x0a, 0x0c, 0x4c, 0x6f, 0x61, 0x64, 0x42, 0x61,
-	0x6c, 0x61, 0x6e, 0x63, 0x65, 0x72, 0x12, 0x52, 0x0a, 0x0b, 0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63,
-	0x65, 0x4c, 0x6f, 0x61, 0x64, 0x12, 0x1e, 0x2e, 0x67, 0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e,
-	0x76, 0x31, 0x2e, 0x4c, 0x6f, 0x61, 0x64, 0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x52, 0x65,
-	0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x67, 0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e,
-	0x76, 0x31, 0x2e, 0x4c, 0x6f, 0x61, 0x64, 0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x52, 0x65,
-	0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x28, 0x01, 0x30, 0x01, 0x42, 0x57, 0x0a, 0x0d, 0x69, 0x6f,
-	0x2e, 0x67, 0x72, 0x70, 0x63, 0x2e, 0x6c, 0x62, 0x2e, 0x76, 0x31, 0x42, 0x11, 0x4c, 0x6f, 0x61,
-	0x64, 0x42, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65, 0x72, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x50, 0x01,
-	0x5a, 0x31, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x67, 0x6f, 0x6c, 0x61, 0x6e, 0x67, 0x2e,
-	0x6f, 0x72, 0x67, 0x2f, 0x67, 0x72, 0x70, 0x63, 0x2f, 0x62, 0x61, 0x6c, 0x61, 0x6e, 0x63, 0x65,
-	0x72, 0x2f, 0x67, 0x72, 0x70, 0x63, 0x6c, 0x62, 0x2f, 0x67, 0x72, 0x70, 0x63, 0x5f, 0x6c, 0x62,
-	0x5f, 0x76, 0x31, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-}
+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 = file_grpc_lb_v1_load_balancer_proto_rawDesc
+	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(file_grpc_lb_v1_load_balancer_proto_rawDescData)
+		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 = []interface{}{
+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
@@ -820,121 +734,11 @@ func file_grpc_lb_v1_load_balancer_proto_init() {
 	if File_grpc_lb_v1_load_balancer_proto != nil {
 		return
 	}
-	if !protoimpl.UnsafeEnabled {
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*LoadBalanceRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*InitialLoadBalanceRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*ClientStatsPerToken); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*ClientStats); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*LoadBalanceResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*FallbackResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*InitialLoadBalanceResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*ServerList); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_grpc_lb_v1_load_balancer_proto_msgTypes[8].Exporter = func(v interface{}, i int) interface{} {
-			switch v := v.(*Server); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-	}
-	file_grpc_lb_v1_load_balancer_proto_msgTypes[0].OneofWrappers = []interface{}{
+	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 = []interface{}{
+	file_grpc_lb_v1_load_balancer_proto_msgTypes[4].OneofWrappers = []any{
 		(*LoadBalanceResponse_InitialResponse)(nil),
 		(*LoadBalanceResponse_ServerList)(nil),
 		(*LoadBalanceResponse_FallbackResponse)(nil),
@@ -943,7 +747,7 @@ func file_grpc_lb_v1_load_balancer_proto_init() {
 	out := protoimpl.TypeBuilder{
 		File: protoimpl.DescBuilder{
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
-			RawDescriptor: file_grpc_lb_v1_load_balancer_proto_rawDesc,
+			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,
@@ -954,7 +758,6 @@ func file_grpc_lb_v1_load_balancer_proto_init() {
 		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_rawDesc = nil
 	file_grpc_lb_v1_load_balancer_proto_goTypes = nil
 	file_grpc_lb_v1_load_balancer_proto_depIdxs = nil
 }

balancer/grpclb/grpc_lb_v1/load_balancer_grpc.pb.go

@@ -1,7 +1,26 @@
+// 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.1.0
-// - protoc             v3.14.0
+// - protoc-gen-go-grpc v1.5.1
+// - protoc             v5.27.1
 // source: grpc/lb/v1/load_balancer.proto
 
 package grpc_lb_v1
@@ -15,15 +34,19 @@ import (
 
 // 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.32.0 or later.
-const _ = grpc.SupportPackageIsVersion7
+// 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) (LoadBalancer_BalanceLoadClient, error)
+	BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[LoadBalanceRequest, LoadBalanceResponse], error)
 }
 
 type loadBalancerClient struct {
@@ -34,52 +57,38 @@ func NewLoadBalancerClient(cc grpc.ClientConnInterface) LoadBalancerClient {
 	return &loadBalancerClient{cc}
 }
 
-func (c *loadBalancerClient) BalanceLoad(ctx context.Context, opts ...grpc.CallOption) (LoadBalancer_BalanceLoadClient, error) {
-	stream, err := c.cc.NewStream(ctx, &LoadBalancer_ServiceDesc.Streams[0], "/grpc.lb.v1.LoadBalancer/BalanceLoad", opts...)
+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 := &loadBalancerBalanceLoadClient{stream}
+	x := &grpc.GenericClientStream[LoadBalanceRequest, LoadBalanceResponse]{ClientStream: stream}
 	return x, nil
 }
 
-type LoadBalancer_BalanceLoadClient interface {
-	Send(*LoadBalanceRequest) error
-	Recv() (*LoadBalanceResponse, error)
-	grpc.ClientStream
-}
-
-type loadBalancerBalanceLoadClient struct {
-	grpc.ClientStream
-}
-
-func (x *loadBalancerBalanceLoadClient) Send(m *LoadBalanceRequest) error {
-	return x.ClientStream.SendMsg(m)
-}
-
-func (x *loadBalancerBalanceLoadClient) Recv() (*LoadBalanceResponse, error) {
-	m := new(LoadBalanceResponse)
-	if err := x.ClientStream.RecvMsg(m); err != nil {
-		return nil, err
-	}
-	return m, 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
+// for forward compatibility.
 type LoadBalancerServer interface {
 	// Bidirectional rpc to get a list of servers.
-	BalanceLoad(LoadBalancer_BalanceLoadServer) error
+	BalanceLoad(grpc.BidiStreamingServer[LoadBalanceRequest, LoadBalanceResponse]) error
 }
 
-// UnimplementedLoadBalancerServer should be embedded to have forward compatible implementations.
-type UnimplementedLoadBalancerServer struct {
-}
+// 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(LoadBalancer_BalanceLoadServer) error {
-	return status.Errorf(codes.Unimplemented, "method BalanceLoad not implemented")
+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
@@ -89,34 +98,22 @@ type UnsafeLoadBalancerServer interface {
 }
 
 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(&loadBalancerBalanceLoadServer{stream})
+	return srv.(LoadBalancerServer).BalanceLoad(&grpc.GenericServerStream[LoadBalanceRequest, LoadBalanceResponse]{ServerStream: stream})
 }
 
-type LoadBalancer_BalanceLoadServer interface {
-	Send(*LoadBalanceResponse) error
-	Recv() (*LoadBalanceRequest, error)
-	grpc.ServerStream
-}
-
-type loadBalancerBalanceLoadServer struct {
-	grpc.ServerStream
-}
-
-func (x *loadBalancerBalanceLoadServer) Send(m *LoadBalanceResponse) error {
-	return x.ServerStream.SendMsg(m)
-}
-
-func (x *loadBalancerBalanceLoadServer) Recv() (*LoadBalanceRequest, error) {
-	m := new(LoadBalanceRequest)
-	if err := x.ServerStream.RecvMsg(m); err != nil {
-		return nil, err
-	}
-	return m, nil
-}
+// 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,

balancer/grpclb/grpclb.go

@@ -19,7 +19,8 @@
 // Package grpclb defines a grpclb balancer.
 //
 // To install grpclb balancer, import this package as:
-//    import _ "google.golang.org/grpc/balancer/grpclb"
+//
+//	import _ "google.golang.org/grpc/balancer/grpclb"
 package grpclb
 
 import (
@@ -31,16 +32,20 @@ import (
 
 	"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"
 
-	durationpb "github.com/golang/protobuf/ptypes/duration"
 	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
 )
 
@@ -131,32 +136,38 @@ func (b *lbBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) bal
 	// 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.
-	r := &lbManualResolver{scheme: "grpclb-internal", ccb: cc}
+	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),
-		target:          opt.Target.Endpoint,
+		dialTarget:      opt.Target.Endpoint(),
+		target:          opt.Target.Endpoint(),
 		opt:             opt,
 		fallbackTimeout: b.fallbackTimeout,
 		doneCh:          make(chan struct{}),
 
-		manualResolver: r,
+		manualResolver: mr,
 		subConns:       make(map[resolver.Address]balancer.SubConn),
 		scStates:       make(map[balancer.SubConn]connectivity.State),
-		picker:         &errPicker{err: balancer.ErrNoSubConnAvailable},
+		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 {
-			logger.Warningf("lbBalancer: client connection creds NewWithMode failed: %v", err)
+			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 {
-			logger.Warningf("lbBalancer: backend creds NewWithMode failed: %v", err)
+			lb.logger.Warningf("Failed to create credentials used for connecting to backends returned by grpclb: %v", err)
 		}
 	}
 
@@ -164,9 +175,11 @@ func (b *lbBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) bal
 }
 
 type lbBalancer struct {
-	cc     *lbCacheClientConn
-	target string
-	opt    balancer.BuildOptions
+	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
 
@@ -184,8 +197,8 @@ type lbBalancer struct {
 
 	// manualResolver is used in the remote LB ClientConn inside grpclb. When
 	// resolved address updates are received by grpclb, filtered updates will be
-	// send to remote LB ClientConn through this resolver.
-	manualResolver *lbManualResolver
+	// 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.
@@ -206,11 +219,11 @@ type lbBalancer struct {
 	// 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 gerenated.
+	// 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/remove SubConn.
+	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
@@ -227,17 +240,18 @@ type lbBalancer struct {
 
 // 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.
+//   - 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 = &errPicker{err: fmt.Errorf("all SubConns are in TransientFailure, last connection error: %v", lb.connErr)}
+		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 = &errPicker{err: balancer.ErrNoSubConnAvailable}
+		lb.picker = base.NewErrPicker(balancer.ErrNoSubConnAvailable)
 		return
 	}
 
@@ -264,7 +278,7 @@ func (lb *lbBalancer) regeneratePicker(resetDrop bool) {
 		//
 		// This doesn't seem to be necessary after the connecting check above.
 		// Kept for safety.
-		lb.picker = &errPicker{err: balancer.ErrNoSubConnAvailable}
+		lb.picker = base.NewErrPicker(balancer.ErrNoSubConnAvailable)
 		return
 	}
 	if lb.inFallback {
@@ -286,16 +300,16 @@ func (lb *lbBalancer) regeneratePicker(resetDrop bool) {
 // 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 remove).
+// 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.
+//	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
 
@@ -315,18 +329,24 @@ func (lb *lbBalancer) aggregateSubConnStates() connectivity.State {
 	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 logger.V(2) {
-		logger.Infof("lbBalancer: handle SubConn state change: %p, %v", sc, s)
+	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 logger.V(2) {
-			logger.Infof("lbBalancer: got state changes for an unknown SubConn: %p, %v", sc, s)
+		if lb.logger.V(2) {
+			lb.logger.Infof("Received state change for an unknown SubConn: %p, %v", sc, s)
 		}
 		return
 	}
@@ -335,8 +355,8 @@ func (lb *lbBalancer) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubCo
 	case connectivity.Idle:
 		sc.Connect()
 	case connectivity.Shutdown:
-		// When an address was removed by resolver, b called RemoveSubConn but
-		// kept the sc's state in scStates. Remove state for this sc here.
+		// 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
@@ -369,8 +389,13 @@ func (lb *lbBalancer) updateStateAndPicker(forceRegeneratePicker bool, resetDrop
 	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
+	}
 
-	lb.cc.UpdateState(balancer.State{ConnectivityState: lb.state, Picker: lb.picker})
+	cc.UpdateState(balancer.State{ConnectivityState: lb.state, Picker: lb.picker})
 }
 
 // fallbackToBackendsAfter blocks for fallbackTimeout and falls back to use
@@ -398,12 +423,36 @@ 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 logger.V(2) {
-		logger.Infof("lbBalancer: switching mode, new usePickFirst: %+v", newUsePickFirst)
+	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)
 }
@@ -414,23 +463,15 @@ func (lb *lbBalancer) ResolverError(error) {
 }
 
 func (lb *lbBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
-	if logger.V(2) {
-		logger.Infof("lbBalancer: UpdateClientConnState: %+v", ccs)
+	if lb.logger.V(2) {
+		lb.logger.Infof("UpdateClientConnState: %s", pretty.ToJSON(ccs))
 	}
 	gc, _ := ccs.BalancerConfig.(*grpclbServiceConfig)
 	lb.handleServiceConfig(gc)
 
-	addrs := ccs.ResolverState.Addresses
+	backendAddrs := ccs.ResolverState.Addresses
 
-	var remoteBalancerAddrs, backendAddrs []resolver.Address
-	for _, a := range addrs {
-		if a.Type == resolver.GRPCLB {
-			a.Type = resolver.Backend
-			remoteBalancerAddrs = append(remoteBalancerAddrs, a)
-		} else {
-			backendAddrs = append(backendAddrs, a)
-		}
-	}
+	var remoteBalancerAddrs []resolver.Address
 	if sd := grpclbstate.Get(ccs.ResolverState); sd != nil {
 		// Override any balancer addresses provided via
 		// ccs.ResolverState.Addresses.
@@ -451,7 +492,9 @@ func (lb *lbBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error
 	} else if lb.ccRemoteLB == nil {
 		// First time receiving resolved addresses, create a cc to remote
 		// balancers.
-		lb.newRemoteBalancerCCWrapper()
+		if err := lb.newRemoteBalancerCCWrapper(); err != nil {
+			return err
+		}
 		// Start the fallback goroutine.
 		go lb.fallbackToBackendsAfter(lb.fallbackTimeout)
 	}

balancer/grpclb/grpclb_config.go

@@ -21,19 +21,20 @@ package grpclb
 import (
 	"encoding/json"
 
-	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer/pickfirst"
 	"google.golang.org/grpc/balancer/roundrobin"
 	"google.golang.org/grpc/serviceconfig"
 )
 
 const (
 	roundRobinName = roundrobin.Name
-	pickFirstName  = grpc.PickFirstBalancerName
+	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) {

balancer/grpclb/grpclb_config_test.go

@@ -20,52 +20,68 @@ package grpclb
 
 import (
 	"encoding/json"
-	"errors"
-	"fmt"
-	"reflect"
-	"strings"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"google.golang.org/grpc/serviceconfig"
 )
 
 func (s) TestParse(t *testing.T) {
 	tests := []struct {
 		name    string
-		s       string
+		sc      string
 		want    serviceconfig.LoadBalancingConfig
-		wantErr error
+		wantErr bool
 	}{
 		{
 			name:    "empty",
-			s:       "",
+			sc:      "",
 			want:    nil,
-			wantErr: errors.New("unexpected end of JSON input"),
+			wantErr: true,
 		},
 		{
 			name: "success1",
-			s:    `{"childPolicy":[{"pick_first":{}}]}`,
+			sc: `
+{
+	"childPolicy": [
+		{"pick_first":{}}
+	],
+	"serviceName": "foo-service"
+}`,
 			want: &grpclbServiceConfig{
 				ChildPolicy: &[]map[string]json.RawMessage{
 					{"pick_first": json.RawMessage("{}")},
 				},
+				ServiceName: "foo-service",
 			},
 		},
 		{
 			name: "success2",
-			s:    `{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}`,
+			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) {
-			if got, err := (&lbBuilder{}).ParseConfig(json.RawMessage(tt.s)); !reflect.DeepEqual(got, tt.want) || !strings.Contains(fmt.Sprint(err), fmt.Sprint(tt.wantErr)) {
-				t.Errorf("parseFullServiceConfig() = %+v, %+v, want %+v, <contains %q>", got, err, tt.want, tt.wantErr)
+			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)
 			}
 		})
 	}

balancer/grpclb/grpclb_picker.go

@@ -19,13 +19,13 @@
 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/internal/grpcrand"
 	"google.golang.org/grpc/status"
 )
 
@@ -98,15 +98,6 @@ func (s *rpcStats) knownReceived() {
 	atomic.AddInt64(&s.numCallsFinished, 1)
 }
 
-type errPicker struct {
-	// Pick always returns this err.
-	err error
-}
-
-func (p *errPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
-	return balancer.PickResult{}, p.err
-}
-
 // 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.
@@ -121,7 +112,7 @@ type rrPicker struct {
 func newRRPicker(readySCs []balancer.SubConn) *rrPicker {
 	return &rrPicker{
 		subConns:     readySCs,
-		subConnsNext: grpcrand.Intn(len(readySCs)),
+		subConnsNext: rand.IntN(len(readySCs)),
 	}
 }
 
@@ -156,7 +147,7 @@ func newLBPicker(serverList []*lbpb.Server, readySCs []balancer.SubConn, stats *
 	return &lbPicker{
 		serverList:   serverList,
 		subConns:     readySCs,
-		subConnsNext: grpcrand.Intn(len(readySCs)),
+		subConnsNext: rand.IntN(len(readySCs)),
 		stats:        stats,
 	}
 }

balancer/grpclb/grpclb_remote_balancer.go

@@ -26,26 +26,38 @@ import (
 	"sync"
 	"time"
 
-	"github.com/golang/protobuf/proto"
-	timestamppb "github.com/golang/protobuf/ptypes/timestamp"
-	"github.com/google/go-cmp/cmp"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
-	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
 	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal/backoff"
-	"google.golang.org/grpc/internal/channelz"
 	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 logger.V(2) {
-		logger.Infof("lbBalancer: processing server list: %+v", l)
+	if lb.logger.V(2) {
+		lb.logger.Infof("Processing server list: %#v", l)
 	}
 	lb.mu.Lock()
 	defer lb.mu.Unlock()
@@ -55,9 +67,9 @@ func (lb *lbBalancer) processServerList(l *lbpb.ServerList) {
 	lb.serverListReceived = true
 
 	// If the new server list == old server list, do nothing.
-	if cmp.Equal(lb.fullServerList, l.Servers, cmp.Comparer(proto.Equal)) {
-		if logger.V(2) {
-			logger.Infof("lbBalancer: new serverlist same as the previous one, ignoring")
+	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
 	}
@@ -70,17 +82,10 @@ func (lb *lbBalancer) processServerList(l *lbpb.ServerList) {
 		}
 
 		md := metadata.Pairs(lbTokenKey, s.LoadBalanceToken)
-		ip := net.IP(s.IpAddress)
-		ipStr := ip.String()
-		if ip.To4() == nil {
-			// Add square brackets to ipv6 addresses, otherwise net.Dial() and
-			// net.SplitHostPort() will return too many colons error.
-			ipStr = fmt.Sprintf("[%s]", ipStr)
-		}
-		addr := imetadata.Set(resolver.Address{Addr: fmt.Sprintf("%s:%d", ipStr, s.Port)}, md)
-		if logger.V(2) {
-			logger.Infof("lbBalancer: server list entry[%d]: ipStr:|%s|, port:|%d|, load balancer token:|%v|",
-				i, ipStr, s.Port, s.LoadBalanceToken)
+		ipStr := net.IP(s.IpAddress).String()
+		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)
 	}
@@ -113,7 +118,6 @@ func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback
 	}
 
 	balancingPolicyChanged := lb.usePickFirst != pickFirst
-	oldUsePickFirst := lb.usePickFirst
 	lb.usePickFirst = pickFirst
 
 	if fallbackModeChanged || balancingPolicyChanged {
@@ -123,31 +127,34 @@ func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback
 		// For fallback mode switching with pickfirst, we want to recreate the
 		// SubConn because the creds could be different.
 		for a, sc := range lb.subConns {
-			if oldUsePickFirst {
-				// If old SubConn were created for pickfirst, bypass cache and
-				// remove directly.
-				lb.cc.cc.RemoveSubConn(sc)
-			} else {
-				lb.cc.RemoveSubConn(sc)
-			}
+			sc.Shutdown()
 			delete(lb.subConns, a)
 		}
 	}
 
 	if lb.usePickFirst {
-		var sc balancer.SubConn
-		for _, sc = range lb.subConns {
+		var (
+			scKey resolver.Address
+			sc    balancer.SubConn
+		)
+		for scKey, sc = range lb.subConns {
 			break
 		}
 		if sc != nil {
-			lb.cc.cc.UpdateAddresses(sc, backendAddrs)
+			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.cc.NewSubConn(backendAddrs, opts)
+		sc, err := lb.cc.ClientConn.NewSubConn(backendAddrs, opts)
 		if err != nil {
-			logger.Warningf("grpclb: failed to create new SubConn: %v", err)
+			lb.logger.Warningf("Failed to create new SubConn: %v", err)
 			return
 		}
 		sc.Connect()
@@ -168,9 +175,11 @@ func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback
 
 		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 {
-				logger.Warningf("grpclb: failed to create new SubConn: %v", err)
+				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.
@@ -186,7 +195,7 @@ func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback
 	for a, sc := range lb.subConns {
 		// a was removed by resolver.
 		if _, ok := addrsSet[a]; !ok {
-			lb.cc.RemoveSubConn(sc)
+			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.
@@ -206,18 +215,21 @@ type remoteBalancerCCWrapper struct {
 	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() {
+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.WithInsecure())
+		dopts = append(dopts, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	}
 	if lb.opt.Dialer != nil {
 		dopts = append(dopts, grpc.WithContextDialer(lb.opt.Dialer))
@@ -228,9 +240,7 @@ func (lb *lbBalancer) newRemoteBalancerCCWrapper() {
 	// Explicitly set pickfirst as the balancer.
 	dopts = append(dopts, grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy":"pick_first"}`))
 	dopts = append(dopts, grpc.WithResolvers(lb.manualResolver))
-	if channelz.IsOn() {
-		dopts = append(dopts, grpc.WithChannelzParentID(lb.opt.ChannelzParentID))
-	}
+	dopts = append(dopts, grpc.WithChannelzParentID(lb.opt.ChannelzParent))
 
 	// Enable Keepalive for grpclb client.
 	dopts = append(dopts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
@@ -243,10 +253,12 @@ func (lb *lbBalancer) newRemoteBalancerCCWrapper() {
 	//
 	// The grpclb server addresses will set field ServerName, and creds will
 	// receive ServerName as authority.
-	cc, err := grpc.DialContext(context.Background(), lb.manualResolver.Scheme()+":///grpclb.subClientConn", dopts...)
+	target := lb.manualResolver.Scheme() + ":///grpclb.subClientConn"
+	cc, err := grpc.NewClient(target, dopts...)
 	if err != nil {
-		logger.Fatalf("failed to dial: %v", err)
+		return fmt.Errorf("grpc.NewClient(%s): %v", target, err)
 	}
+	cc.Connect()
 	ccw := &remoteBalancerCCWrapper{
 		cc:      cc,
 		lb:      lb,
@@ -256,6 +268,7 @@ func (lb *lbBalancer) newRemoteBalancerCCWrapper() {
 	lb.ccRemoteLB = ccw
 	ccw.wg.Add(1)
 	go ccw.watchRemoteBalancer()
+	return nil
 }
 
 // close closed the ClientConn to remote balancer, and waits until all
@@ -319,13 +332,11 @@ func (ccw *remoteBalancerCCWrapper) sendLoadReport(s *balanceLoadClientStream, i
 	}
 }
 
-func (ccw *remoteBalancerCCWrapper) callRemoteBalancer() (backoff bool, _ error) {
+func (ccw *remoteBalancerCCWrapper) callRemoteBalancer(ctx context.Context) (backoff bool, _ error) {
 	lbClient := &loadBalancerClient{cc: ccw.cc}
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 	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)
+		return true, fmt.Errorf("grpclb: failed to perform RPC to the remote balancer: %v", err)
 	}
 	ccw.lb.mu.Lock()
 	ccw.lb.remoteBalancerConnected = true
@@ -362,25 +373,57 @@ func (ccw *remoteBalancerCCWrapper) callRemoteBalancer() (backoff bool, _ error)
 	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 ccw.wg.Done()
+	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 {
-		doBackoff, err := ccw.callRemoteBalancer()
+		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 {
-					logger.Info(err)
+					ccw.lb.logger.Infof("Call to remote balancer failed: %v", err)
 				} else {
-					logger.Warning(err)
+					ccw.lb.logger.Warningf("Call to remote balancer failed: %v", err)
 				}
 			}
 		}
 		// Trigger a re-resolve when the stream errors.
-		ccw.lb.cc.cc.ResolveNow(resolver.ResolveNowOptions{})
+		ccw.lb.cc.ClientConn.ResolveNow(resolver.ResolveNowOptions{})
 
 		ccw.lb.mu.Lock()
 		ccw.lb.remoteBalancerConnected = false

balancer/grpclb/grpclb_test_util_test.go

@@ -1,85 +0,0 @@
-/*
- *
- * 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 (
-	"net"
-	"sync"
-)
-
-type tempError struct{}
-
-func (*tempError) Error() string {
-	return "grpclb test temporary error"
-}
-func (*tempError) Temporary() bool {
-	return true
-}
-
-type restartableListener struct {
-	net.Listener
-	addr string
-
-	mu     sync.Mutex
-	closed bool
-	conns  []net.Conn
-}
-
-func newRestartableListener(l net.Listener) *restartableListener {
-	return &restartableListener{
-		Listener: l,
-		addr:     l.Addr().String(),
-	}
-}
-
-func (l *restartableListener) Accept() (conn net.Conn, err error) {
-	conn, err = l.Listener.Accept()
-	if err == nil {
-		l.mu.Lock()
-		if l.closed {
-			conn.Close()
-			l.mu.Unlock()
-			return nil, &tempError{}
-		}
-		l.conns = append(l.conns, conn)
-		l.mu.Unlock()
-	}
-	return
-}
-
-func (l *restartableListener) Close() error {
-	return l.Listener.Close()
-}
-
-func (l *restartableListener) stopPreviousConns() {
-	l.mu.Lock()
-	l.closed = true
-	tmp := l.conns
-	l.conns = nil
-	l.mu.Unlock()
-	for _, conn := range tmp {
-		conn.Close()
-	}
-}
-
-func (l *restartableListener) restart() {
-	l.mu.Lock()
-	l.closed = false
-	l.mu.Unlock()
-}

balancer/grpclb/grpclb_util.go

@@ -27,75 +27,15 @@ import (
 	"google.golang.org/grpc/resolver"
 )
 
-// The parent ClientConn should re-resolve when grpclb loses connection to the
-// remote balancer. When the ClientConn inside grpclb gets a TransientFailure,
-// it calls lbManualResolver.ResolveNow(), which calls parent ClientConn's
-// ResolveNow, and eventually results in re-resolve happening in parent
-// ClientConn's resolver (DNS for example).
-//
-//                          parent
-//                          ClientConn
-//  +-----------------------------------------------------------------+
-//  |             parent          +---------------------------------+ |
-//  | DNS         ClientConn      |  grpclb                         | |
-//  | resolver    balancerWrapper |                                 | |
-//  | +              +            |    grpclb          grpclb       | |
-//  | |              |            |    ManualResolver  ClientConn   | |
-//  | |              |            |     +              +            | |
-//  | |              |            |     |              | Transient  | |
-//  | |              |            |     |              | Failure    | |
-//  | |              |            |     |  <---------  |            | |
-//  | |              | <--------------- |  ResolveNow  |            | |
-//  | |  <---------  | ResolveNow |     |              |            | |
-//  | |  ResolveNow  |            |     |              |            | |
-//  | |              |            |     |              |            | |
-//  | +              +            |     +              +            | |
-//  |                             +---------------------------------+ |
-//  +-----------------------------------------------------------------+
-
-// lbManualResolver is used by the ClientConn inside grpclb. It's a manual
-// resolver with a special ResolveNow() function.
-//
-// When ResolveNow() is called, it calls ResolveNow() on the parent ClientConn,
-// so when grpclb client lose contact with remote balancers, the parent
-// ClientConn's resolver will re-resolve.
-type lbManualResolver struct {
-	scheme string
-	ccr    resolver.ClientConn
-
-	ccb balancer.ClientConn
-}
-
-func (r *lbManualResolver) Build(_ resolver.Target, cc resolver.ClientConn, _ resolver.BuildOptions) (resolver.Resolver, error) {
-	r.ccr = cc
-	return r, nil
-}
-
-func (r *lbManualResolver) Scheme() string {
-	return r.scheme
-}
-
-// ResolveNow calls resolveNow on the parent ClientConn.
-func (r *lbManualResolver) ResolveNow(o resolver.ResolveNowOptions) {
-	r.ccb.ResolveNow(o)
-}
-
-// Close is a noop for Resolver.
-func (*lbManualResolver) Close() {}
-
-// UpdateState calls cc.UpdateState.
-func (r *lbManualResolver) UpdateState(s resolver.State) {
-	r.ccr.UpdateState(s)
-}
-
 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 removed.
+// SubConns will be kept in cache for subConnCacheTime before being shut down.
 //
-// Its new and remove methods are updated to do cache first.
+// Its NewSubconn and SubConn.Shutdown methods are updated to do cache first.
 type lbCacheClientConn struct {
-	cc      balancer.ClientConn
+	balancer.ClientConn
+
 	timeout time.Duration
 
 	mu sync.Mutex
@@ -113,7 +53,7 @@ type subConnCacheEntry struct {
 
 func newLBCacheClientConn(cc balancer.ClientConn) *lbCacheClientConn {
 	return &lbCacheClientConn{
-		cc:            cc,
+		ClientConn:    cc,
 		timeout:       subConnCacheTime,
 		subConnCache:  make(map[resolver.Address]*subConnCacheEntry),
 		subConnToAddr: make(map[balancer.SubConn]resolver.Address),
@@ -137,16 +77,27 @@ func (ccc *lbCacheClientConn) NewSubConn(addrs []resolver.Address, opts balancer
 		return entry.sc, nil
 	}
 
-	scNew, err := ccc.cc.NewSubConn(addrs, opts)
+	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]
@@ -156,11 +107,11 @@ func (ccc *lbCacheClientConn) RemoveSubConn(sc balancer.SubConn) {
 
 	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 removed. We remove sc
-			// immediately here.
+			// 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)
-			ccc.cc.RemoveSubConn(sc)
+			sc.SubConn.Shutdown()
 		}
 		return
 	}
@@ -176,7 +127,7 @@ func (ccc *lbCacheClientConn) RemoveSubConn(sc balancer.SubConn) {
 		if entry.abortDeleting {
 			return
 		}
-		ccc.cc.RemoveSubConn(sc)
+		sc.SubConn.Shutdown()
 		delete(ccc.subConnToAddr, sc)
 		delete(ccc.subConnCache, addr)
 	})
@@ -195,14 +146,28 @@ func (ccc *lbCacheClientConn) RemoveSubConn(sc balancer.SubConn) {
 }
 
 func (ccc *lbCacheClientConn) UpdateState(s balancer.State) {
-	ccc.cc.UpdateState(s)
+	s.Picker = &lbCachePicker{Picker: s.Picker}
+	ccc.ClientConn.UpdateState(s)
 }
 
 func (ccc *lbCacheClientConn) close() {
 	ccc.mu.Lock()
-	// Only cancel all existing timers. There's no need to remove SubConns.
+	defer ccc.mu.Unlock()
+	// Only cancel all existing timers. There's no need to shut down SubConns.
 	for _, entry := range ccc.subConnCache {
 		entry.cancel()
 	}
-	ccc.mu.Unlock()
+}
+
+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
 }

balancer/grpclb/grpclb_util_test.go

@@ -30,6 +30,13 @@ import (
 
 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 {
@@ -45,8 +52,8 @@ func newMockClientConn() *mockClientConn {
 	}
 }
 
-func (mcc *mockClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
-	sc := &mockSubConn{}
+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]
@@ -54,9 +61,7 @@ func (mcc *mockClientConn) NewSubConn(addrs []resolver.Address, opts balancer.Ne
 }
 
 func (mcc *mockClientConn) RemoveSubConn(sc balancer.SubConn) {
-	mcc.mu.Lock()
-	defer mcc.mu.Unlock()
-	delete(mcc.subConns, sc)
+	panic(fmt.Sprintf("RemoveSubConn(%v) called unexpectedly", sc))
 }
 
 const testCacheTimeout = 100 * time.Millisecond
@@ -82,7 +87,7 @@ func checkCacheCC(ccc *lbCacheClientConn, sccLen, sctaLen int) error {
 	return nil
 }
 
-// Test that SubConn won't be immediately removed.
+// Test that SubConn won't be immediately shut down.
 func (s) TestLBCacheClientConnExpire(t *testing.T) {
 	mcc := newMockClientConn()
 	if err := checkMockCC(mcc, 0); err != nil {
@@ -105,7 +110,7 @@ func (s) TestLBCacheClientConnExpire(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	ccc.RemoveSubConn(sc)
+	sc.Shutdown()
 	// One subconn in MockCC before timeout.
 	if err := checkMockCC(mcc, 1); err != nil {
 		t.Fatal(err)
@@ -133,7 +138,7 @@ func (s) TestLBCacheClientConnExpire(t *testing.T) {
 	}
 }
 
-// Test that NewSubConn with the same address of a SubConn being removed will
+// 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()
@@ -157,7 +162,7 @@ func (s) TestLBCacheClientConnReuse(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	ccc.RemoveSubConn(sc)
+	sc.Shutdown()
 	// One subconn in MockCC before timeout.
 	if err := checkMockCC(mcc, 1); err != nil {
 		t.Fatal(err)
@@ -190,8 +195,8 @@ func (s) TestLBCacheClientConnReuse(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	// Call remove again, will delete after timeout.
-	ccc.RemoveSubConn(sc)
+	// 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)
@@ -218,9 +223,9 @@ func (s) TestLBCacheClientConnReuse(t *testing.T) {
 	}
 }
 
-// Test that if the timer to remove a SubConn fires at the same time NewSubConn
-// cancels the timer, it doesn't cause deadlock.
-func (s) TestLBCache_RemoveTimer_New_Race(t *testing.T) {
+// 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)
@@ -246,9 +251,9 @@ func (s) TestLBCache_RemoveTimer_New_Race(t *testing.T) {
 
 	go func() {
 		for i := 0; i < 1000; i++ {
-			// Remove starts a timer with 1 ns timeout, the NewSubConn will race
-			// with with the timer.
-			ccc.RemoveSubConn(sc)
+			// 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)

balancer/grpclb/state/state.go

@@ -39,7 +39,7 @@ type State struct {
 // 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.WithValues(key, s)
+	state.Attributes = state.Attributes.WithValue(key, s)
 	return state
 }
 

balancer/lazy/lazy.go

@@ -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
+}

balancer/lazy/lazy_ext_test.go

@@ -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/pickfirstleaf"
+	"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(pickfirstleaf.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(pickfirstleaf.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(pickfirstleaf.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(pickfirstleaf.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(pickfirstleaf.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)
+}

balancer/leastrequest/leastrequest.go

@@ -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/pickfirstleaf"
+	"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(pickfirstleaf.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: pickfirstleaf.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.Keys() {
+		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
+}

balancer/leastrequest/leastrequest_test.go

@@ -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)
+	}
+}

balancer/pickfirst/internal/internal.go

@@ -0,0 +1,35 @@
+/*
+ * 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
+	// 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() }
+	}
+)

balancer/pickfirst/pickfirst.go

@@ -0,0 +1,291 @@
+/*
+ *
+ * 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.
+package pickfirst
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	rand "math/rand/v2"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst/internal"
+	"google.golang.org/grpc/connectivity"
+	"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"
+
+	_ "google.golang.org/grpc/balancer/pickfirst/pickfirstleaf" // For automatically registering the new pickfirst if required.
+)
+
+func init() {
+	if envconfig.NewPickFirstEnabled {
+		return
+	}
+	balancer.Register(pickfirstBuilder{})
+}
+
+var logger = grpclog.Component("pick-first-lb")
+
+const (
+	// Name is the name of the pick_first balancer.
+	Name      = "pick_first"
+	logPrefix = "[pick-first-lb %p] "
+)
+
+type pickfirstBuilder struct{}
+
+func (pickfirstBuilder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
+	b := &pickfirstBalancer{cc: cc}
+	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
+	return b
+}
+
+func (pickfirstBuilder) Name() string {
+	return Name
+}
+
+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"`
+}
+
+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
+}
+
+type pickfirstBalancer struct {
+	logger  *internalgrpclog.PrefixLogger
+	state   connectivity.State
+	cc      balancer.ClientConn
+	subConn balancer.SubConn
+}
+
+func (b *pickfirstBalancer) ResolverError(err error) {
+	if b.logger.V(2) {
+		b.logger.Infof("Received error from the name resolver: %v", err)
+	}
+	if b.subConn == nil {
+		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.cc.UpdateState(balancer.State{
+		ConnectivityState: connectivity.TransientFailure,
+		Picker:            &picker{err: fmt.Errorf("name resolver error: %v", err)},
+	})
+}
+
+// Shuffler is an interface for shuffling an address list.
+type Shuffler interface {
+	ShuffleAddressListForTesting(n int, swap func(i, j int))
+}
+
+// ShuffleAddressListForTesting pseudo-randomizes the order of addresses.  n
+// is the number of elements.  swap swaps the elements with indexes i and j.
+func ShuffleAddressListForTesting(n int, swap func(i, j int)) { rand.Shuffle(n, swap) }
+
+func (b *pickfirstBalancer) UpdateClientConnState(state balancer.ClientConnState) error {
+	if len(state.ResolverState.Addresses) == 0 && len(state.ResolverState.Endpoints) == 0 {
+		// The resolver reported an empty address list. Treat it like an error by
+		// calling b.ResolverError.
+		if b.subConn != nil {
+			// Shut down the old subConn. All addresses were removed, so it is
+			// no longer valid.
+			b.subConn.Shutdown()
+			b.subConn = nil
+		}
+		b.ResolverError(errors.New("produced zero addresses"))
+		return balancer.ErrBadResolverState
+	}
+	// We don't have to guard this block with the env var because ParseConfig
+	// already does so.
+	cfg, ok := state.BalancerConfig.(pfConfig)
+	if state.BalancerConfig != nil && !ok {
+		return fmt.Errorf("pickfirst: received illegal BalancerConfig (type %T): %v", state.BalancerConfig, state.BalancerConfig)
+	}
+
+	if b.logger.V(2) {
+		b.logger.Infof("Received new config %s, resolver state %s", pretty.ToJSON(cfg), pretty.ToJSON(state.ResolverState))
+	}
+
+	var addrs []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 {
+			endpoints = append([]resolver.Endpoint{}, 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 {
+			// "In the flattened list, interleave addresses from the two address
+			// families, as per RFC-8304 section 4." - A61
+			// TODO: support the above language.
+			addrs = append(addrs, 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.
+		addrs = state.ResolverState.Addresses
+		if cfg.ShuffleAddressList {
+			addrs = append([]resolver.Address{}, addrs...)
+			rand.Shuffle(len(addrs), func(i, j int) { addrs[i], addrs[j] = addrs[j], addrs[i] })
+		}
+	}
+
+	if b.subConn != nil {
+		b.cc.UpdateAddresses(b.subConn, addrs)
+		return nil
+	}
+
+	var subConn balancer.SubConn
+	subConn, err := b.cc.NewSubConn(addrs, balancer.NewSubConnOptions{
+		StateListener: func(state balancer.SubConnState) {
+			b.updateSubConnState(subConn, state)
+		},
+	})
+	if err != nil {
+		if b.logger.V(2) {
+			b.logger.Infof("Failed to create new SubConn: %v", err)
+		}
+		b.state = connectivity.TransientFailure
+		b.cc.UpdateState(balancer.State{
+			ConnectivityState: connectivity.TransientFailure,
+			Picker:            &picker{err: fmt.Errorf("error creating connection: %v", err)},
+		})
+		return balancer.ErrBadResolverState
+	}
+	b.subConn = subConn
+	b.state = connectivity.Idle
+	b.cc.UpdateState(balancer.State{
+		ConnectivityState: connectivity.Connecting,
+		Picker:            &picker{err: balancer.ErrNoSubConnAvailable},
+	})
+	b.subConn.Connect()
+	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) updateSubConnState(subConn balancer.SubConn, state balancer.SubConnState) {
+	if b.logger.V(2) {
+		b.logger.Infof("Received SubConn state update: %p, %+v", subConn, state)
+	}
+	if b.subConn != subConn {
+		if b.logger.V(2) {
+			b.logger.Infof("Ignored state change because subConn is not recognized")
+		}
+		return
+	}
+	if state.ConnectivityState == connectivity.Shutdown {
+		b.subConn = nil
+		return
+	}
+
+	switch state.ConnectivityState {
+	case connectivity.Ready:
+		b.cc.UpdateState(balancer.State{
+			ConnectivityState: state.ConnectivityState,
+			Picker:            &picker{result: balancer.PickResult{SubConn: subConn}},
+		})
+	case connectivity.Connecting:
+		if b.state == connectivity.TransientFailure {
+			// We stay in TransientFailure until we are Ready. See A62.
+			return
+		}
+		b.cc.UpdateState(balancer.State{
+			ConnectivityState: state.ConnectivityState,
+			Picker:            &picker{err: balancer.ErrNoSubConnAvailable},
+		})
+	case connectivity.Idle:
+		if b.state == connectivity.TransientFailure {
+			// We stay in TransientFailure until we are Ready. Also kick the
+			// subConn out of Idle into Connecting. See A62.
+			b.subConn.Connect()
+			return
+		}
+		b.cc.UpdateState(balancer.State{
+			ConnectivityState: state.ConnectivityState,
+			Picker:            &idlePicker{subConn: subConn},
+		})
+	case connectivity.TransientFailure:
+		b.cc.UpdateState(balancer.State{
+			ConnectivityState: state.ConnectivityState,
+			Picker:            &picker{err: state.ConnectionError},
+		})
+	}
+	b.state = state.ConnectivityState
+}
+
+func (b *pickfirstBalancer) Close() {
+}
+
+func (b *pickfirstBalancer) ExitIdle() {
+	if b.subConn != nil && b.state == connectivity.Idle {
+		b.subConn.Connect()
+	}
+}
+
+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 {
+	subConn balancer.SubConn
+}
+
+func (i *idlePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	i.subConn.Connect()
+	return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+}

balancer/pickfirst/pickfirst_ext_test.go

@@ -0,0 +1,965 @@
+/*
+ *
+ * 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 pickfirst_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/backoff"
+	pfinternal "google.golang.org/grpc/balancer/pickfirst/internal"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/internal/testutils/pickfirst"
+	"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"
+)
+
+const (
+	pickFirstServiceConfig = `{"loadBalancingConfig": [{"pick_first":{}}]}`
+	// 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
+)
+
+func init() {
+	channelz.TurnOn()
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// parseServiceConfig is a test helper which uses the manual resolver to parse
+// the given service config. It calls t.Fatal() if service config parsing fails.
+func parseServiceConfig(t *testing.T, r *manual.Resolver, sc string) *serviceconfig.ParseResult {
+	t.Helper()
+
+	scpr := r.CC().ParseServiceConfig(sc)
+	if scpr.Err != nil {
+		t.Fatalf("Failed to parse service config %q: %v", sc, scpr.Err)
+	}
+	return scpr
+}
+
+// setupPickFirst performs steps required for pick_first tests. It starts a
+// bunch of backends exporting the TestService, creates a ClientConn to them
+// with service config specifying the use of the pick_first LB policy.
+func setupPickFirst(t *testing.T, backendCount int, opts ...grpc.DialOption) (*grpc.ClientConn, *manual.Resolver, []*stubserver.StubServer) {
+	t.Helper()
+
+	r := manual.NewBuilderWithScheme("whatever")
+
+	backends := make([]*stubserver.StubServer, backendCount)
+	addrs := make([]resolver.Address, backendCount)
+	for i := 0; i < backendCount; i++ {
+		backend := &stubserver.StubServer{
+			EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+				return &testpb.Empty{}, nil
+			},
+		}
+		if err := backend.StartServer(); err != nil {
+			t.Fatalf("Failed to start backend: %v", err)
+		}
+		t.Logf("Started TestService backend at: %q", backend.Address)
+		t.Cleanup(func() { backend.Stop() })
+
+		backends[i] = backend
+		addrs[i] = resolver.Address{Addr: backend.Address}
+	}
+
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithDefaultServiceConfig(pickFirstServiceConfig),
+	}
+	dopts = append(dopts, opts...)
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	t.Cleanup(func() { cc.Close() })
+
+	// At this point, the resolver has not returned any addresses to the channel.
+	// This RPC must block until the context expires.
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(sCtx, &testpb.Empty{}); status.Code(err) != codes.DeadlineExceeded {
+		t.Fatalf("EmptyCall() = %s, want %s", status.Code(err), codes.DeadlineExceeded)
+	}
+	return cc, r, backends
+}
+
+// stubBackendsToResolverAddrs converts from a set of stub server backends to
+// resolver addresses. Useful when pushing addresses to the manual resolver.
+func stubBackendsToResolverAddrs(backends []*stubserver.StubServer) []resolver.Address {
+	addrs := make([]resolver.Address, len(backends))
+	for i, backend := range backends {
+		addrs[i] = resolver.Address{Addr: backend.Address}
+	}
+	return addrs
+}
+
+// TestPickFirst_OneBackend tests the most basic scenario for pick_first. It
+// brings up a single backend and verifies that all RPCs get routed to it.
+func (s) TestPickFirst_OneBackend(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 1)
+
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// TestPickFirst_MultipleBackends tests the scenario with multiple backends and
+// verifies that all RPCs get routed to the first one.
+func (s) TestPickFirst_MultipleBackends(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 2)
+
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// TestPickFirst_OneServerDown tests the scenario where we have multiple
+// backends and pick_first is working as expected. Verifies that RPCs get routed
+// to the next backend in the list when the first one goes down.
+func (s) TestPickFirst_OneServerDown(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 2)
+
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Stop the backend which is currently being used. RPCs should get routed to
+	// the next backend in the list.
+	backends[0].Stop()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// TestPickFirst_AllServersDown tests the scenario where we have multiple
+// backends and pick_first is working as expected. When all backends go down,
+// the test verifies that RPCs fail with appropriate status code.
+func (s) TestPickFirst_AllServersDown(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 2)
+
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	for _, b := range backends {
+		b.Stop()
+	}
+
+	client := testgrpc.NewTestServiceClient(cc)
+	for {
+		if ctx.Err() != nil {
+			t.Fatalf("channel failed to move to Unavailable after all backends were stopped: %v", ctx.Err())
+		}
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}); status.Code(err) == codes.Unavailable {
+			return
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+}
+
+// TestPickFirst_AddressesRemoved tests the scenario where we have multiple
+// backends and pick_first is working as expected. It then verifies that when
+// addresses are removed by the name resolver, RPCs get routed appropriately.
+func (s) TestPickFirst_AddressesRemoved(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 3)
+
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Remove the first backend from the list of addresses originally pushed.
+	// RPCs should get routed to the first backend in the new list.
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{addrs[1], addrs[2]}})
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Append the backend that we just removed to the end of the list.
+	// Nothing should change.
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{addrs[1], addrs[2], addrs[0]}})
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Remove the first backend from the existing list of addresses.
+	// RPCs should get routed to the first backend in the new list.
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{addrs[2], addrs[0]}})
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[2]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Remove the first backend from the existing list of addresses.
+	// RPCs should get routed to the first backend in the new list.
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{addrs[0]}})
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// TestPickFirst_NewAddressWhileBlocking tests the case where pick_first is
+// configured on a channel, things are working as expected and then a resolver
+// updates removes all addresses. An RPC attempted at this point in time will be
+// blocked because there are no valid backends. This test verifies that when new
+// backends are added, the RPC is able to complete.
+func (s) TestPickFirst_NewAddressWhileBlocking(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 2)
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Send a resolver update with no addresses. This should push the channel into
+	// TransientFailure.
+	r.UpdateState(resolver.State{})
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+
+	doneCh := make(chan struct{})
+	client := testgrpc.NewTestServiceClient(cc)
+	go func() {
+		// The channel is currently in TransientFailure and this RPC will block
+		// until the channel becomes Ready, which will only happen when we push a
+		// resolver update with a valid backend address.
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.WaitForReady(true)); err != nil {
+			t.Errorf("EmptyCall() = %v, want <nil>", err)
+		}
+		close(doneCh)
+	}()
+
+	// Make sure that there is one pending RPC on the ClientConn before attempting
+	// to push new addresses through the name resolver. If we don't do this, the
+	// resolver update can happen before the above goroutine gets to make the RPC.
+	for {
+		if err := ctx.Err(); err != nil {
+			t.Fatal(err)
+		}
+		tcs, _ := channelz.GetTopChannels(0, 0)
+		if len(tcs) != 1 {
+			t.Fatalf("there should only be one top channel, not %d", len(tcs))
+		}
+		started := tcs[0].ChannelMetrics.CallsStarted.Load()
+		completed := tcs[0].ChannelMetrics.CallsSucceeded.Load() + tcs[0].ChannelMetrics.CallsFailed.Load()
+		if (started - completed) == 1 {
+			break
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+
+	// Send a resolver update with a valid backend to push the channel to Ready
+	// and unblock the above RPC.
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backends[0].Address}}})
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout when waiting for blocked RPC to complete")
+	case <-doneCh:
+	}
+}
+
+// TestPickFirst_StickyTransientFailure tests the case where pick_first is
+// configured on a channel, and the backend is configured to close incoming
+// connections as soon as they are accepted. The test verifies that the channel
+// enters TransientFailure and stays there. The test also verifies that the
+// pick_first LB policy is constantly trying to reconnect to the backend.
+func (s) TestPickFirst_StickyTransientFailure(t *testing.T) {
+	// Spin up a local server which closes the connection as soon as it receives
+	// one. It also sends a signal on a channel whenever it received a connection.
+	lis, err := testutils.LocalTCPListener()
+	if err != nil {
+		t.Fatalf("Failed to create listener: %v", err)
+	}
+	t.Cleanup(func() { lis.Close() })
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	connCh := make(chan struct{}, 1)
+	go func() {
+		for {
+			conn, err := lis.Accept()
+			if err != nil {
+				return
+			}
+			select {
+			case connCh <- struct{}{}:
+				conn.Close()
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	// Dial the above server with a ConnectParams that does a constant backoff
+	// of defaultTestShortTimeout duration.
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultServiceConfig(pickFirstServiceConfig),
+		grpc.WithConnectParams(grpc.ConnectParams{
+			Backoff: backoff.Config{
+				BaseDelay:  defaultTestShortTimeout,
+				Multiplier: float64(0),
+				Jitter:     float64(0),
+				MaxDelay:   defaultTestShortTimeout,
+			},
+		}),
+	}
+	cc, err := grpc.NewClient(lis.Addr().String(), dopts...)
+	if err != nil {
+		t.Fatalf("Failed to create new client: %v", err)
+	}
+	t.Cleanup(func() { cc.Close() })
+	cc.Connect()
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+
+	// Spawn a goroutine to ensure that the channel stays in TransientFailure.
+	// The call to cc.WaitForStateChange will return false when the main
+	// goroutine exits and the context is cancelled.
+	go func() {
+		if cc.WaitForStateChange(ctx, connectivity.TransientFailure) {
+			if state := cc.GetState(); state != connectivity.Shutdown {
+				t.Errorf("Unexpected state change from TransientFailure to %s", cc.GetState())
+			}
+		}
+	}()
+
+	// Ensures that the pick_first LB policy is constantly trying to reconnect.
+	for i := 0; i < 10; i++ {
+		select {
+		case <-connCh:
+		case <-time.After(2 * defaultTestShortTimeout):
+			t.Error("Timeout when waiting for pick_first to reconnect")
+		}
+	}
+}
+
+// Tests the PF LB policy with shuffling enabled.
+func (s) TestPickFirst_ShuffleAddressList(t *testing.T) {
+	const serviceConfig = `{"loadBalancingConfig": [{"pick_first":{ "shuffleAddressList": true }}]}`
+
+	// Install a shuffler that always reverses two entries.
+	origShuf := pfinternal.RandShuffle
+	defer func() { pfinternal.RandShuffle = origShuf }()
+	pfinternal.RandShuffle = func(n int, f func(int, int)) {
+		if n != 2 {
+			t.Errorf("Shuffle called with n=%v; want 2", n)
+			return
+		}
+		f(0, 1) // reverse the two addresses
+	}
+	// Set up our backends.
+	cc, r, backends := setupPickFirst(t, 2)
+	addrs := stubBackendsToResolverAddrs(backends)
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Push an update with both addresses and shuffling disabled.  We should
+	// connect to backend 0.
+	r.UpdateState(resolver.State{Endpoints: []resolver.Endpoint{
+		{Addresses: []resolver.Address{addrs[0]}},
+		{Addresses: []resolver.Address{addrs[1]}},
+	}})
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Send a config with shuffling enabled.  This will reverse the addresses,
+	// but the channel should still be connected to backend 0.
+	shufState := resolver.State{
+		ServiceConfig: parseServiceConfig(t, r, serviceConfig),
+		Endpoints: []resolver.Endpoint{
+			{Addresses: []resolver.Address{addrs[0]}},
+			{Addresses: []resolver.Address{addrs[1]}},
+		},
+	}
+	r.UpdateState(shufState)
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Send a resolver update with no addresses. This should push the channel
+	// into TransientFailure.
+	r.UpdateState(resolver.State{})
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+
+	// Send the same config as last time with shuffling enabled.  Since we are
+	// not connected to backend 0, we should connect to backend 1.
+	r.UpdateState(shufState)
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// Test config parsing with the env var turned on and off for various scenarios.
+func (s) TestPickFirst_ParseConfig_Success(t *testing.T) {
+	// Install a shuffler that always reverses two entries.
+	origShuf := pfinternal.RandShuffle
+	defer func() { pfinternal.RandShuffle = origShuf }()
+	pfinternal.RandShuffle = func(n int, f func(int, int)) {
+		if n != 2 {
+			t.Errorf("Shuffle called with n=%v; want 2", n)
+			return
+		}
+		f(0, 1) // reverse the two addresses
+	}
+
+	tests := []struct {
+		name          string
+		serviceConfig string
+		wantFirstAddr bool
+	}{
+		{
+			name:          "empty pickfirst config",
+			serviceConfig: `{"loadBalancingConfig": [{"pick_first":{}}]}`,
+			wantFirstAddr: true,
+		},
+		{
+			name:          "empty good pickfirst config",
+			serviceConfig: `{"loadBalancingConfig": [{"pick_first":{ "shuffleAddressList": true }}]}`,
+			wantFirstAddr: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			// Set up our backends.
+			cc, r, backends := setupPickFirst(t, 2)
+			addrs := stubBackendsToResolverAddrs(backends)
+
+			r.UpdateState(resolver.State{
+				ServiceConfig: parseServiceConfig(t, r, test.serviceConfig),
+				Addresses:     addrs,
+			})
+
+			// Some tests expect address shuffling to happen, and indicate that
+			// by setting wantFirstAddr to false (since our shuffling function
+			// defined at the top of this test, simply reverses the list of
+			// addresses provided to it).
+			wantAddr := addrs[0]
+			if !test.wantFirstAddr {
+				wantAddr = addrs[1]
+			}
+
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
+			if err := pickfirst.CheckRPCsToBackend(ctx, cc, wantAddr); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
+
+// Test config parsing for a bad service config.
+func (s) TestPickFirst_ParseConfig_Failure(t *testing.T) {
+	// Service config should fail with the below config. Name resolvers are
+	// expected to perform this parsing before they push the parsed service
+	// config to the channel.
+	const sc = `{"loadBalancingConfig": [{"pick_first":{ "shuffleAddressList": 666 }}]}`
+	scpr := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(sc)
+	if scpr.Err == nil {
+		t.Fatalf("ParseConfig() succeeded and returned %+v, when expected to fail", scpr)
+	}
+}
+
+// setupPickFirstWithListenerWrapper is very similar to setupPickFirst, but uses
+// a wrapped listener that the test can use to track accepted connections.
+func setupPickFirstWithListenerWrapper(t *testing.T, backendCount int, opts ...grpc.DialOption) (*grpc.ClientConn, *manual.Resolver, []*stubserver.StubServer, []*testutils.ListenerWrapper) {
+	t.Helper()
+
+	backends := make([]*stubserver.StubServer, backendCount)
+	addrs := make([]resolver.Address, backendCount)
+	listeners := make([]*testutils.ListenerWrapper, backendCount)
+	for i := 0; i < backendCount; i++ {
+		lis := testutils.NewListenerWrapper(t, nil)
+		backend := &stubserver.StubServer{
+			Listener: lis,
+			EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+				return &testpb.Empty{}, nil
+			},
+		}
+		if err := backend.StartServer(); err != nil {
+			t.Fatalf("Failed to start backend: %v", err)
+		}
+		t.Logf("Started TestService backend at: %q", backend.Address)
+		t.Cleanup(func() { backend.Stop() })
+
+		backends[i] = backend
+		addrs[i] = resolver.Address{Addr: backend.Address}
+		listeners[i] = lis
+	}
+
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithDefaultServiceConfig(pickFirstServiceConfig),
+	}
+	dopts = append(dopts, opts...)
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	t.Cleanup(func() { cc.Close() })
+
+	// At this point, the resolver has not returned any addresses to the channel.
+	// This RPC must block until the context expires.
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(sCtx, &testpb.Empty{}); status.Code(err) != codes.DeadlineExceeded {
+		t.Fatalf("EmptyCall() = %s, want %s", status.Code(err), codes.DeadlineExceeded)
+	}
+	return cc, r, backends, listeners
+}
+
+// TestPickFirst_AddressUpdateWithAttributes tests the case where an address
+// update received by the pick_first LB policy differs in attributes. Addresses
+// which differ in attributes are considered different from the perspective of
+// subconn creation and connection establishment and the test verifies that new
+// connections are created when attributes change.
+func (s) TestPickFirst_AddressUpdateWithAttributes(t *testing.T) {
+	cc, r, backends, listeners := setupPickFirstWithListenerWrapper(t, 2)
+
+	// Add a set of attributes to the addresses before pushing them to the
+	// pick_first LB policy through the manual resolver.
+	addrs := stubBackendsToResolverAddrs(backends)
+	for i := range addrs {
+		addrs[i].Attributes = addrs[i].Attributes.WithValue("test-attribute-1", fmt.Sprintf("%d", i))
+	}
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	// Ensure that RPCs succeed to the first backend in the list.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Grab the wrapped connection from the listener wrapper. This will be used
+	// to verify the connection is closed.
+	val, err := listeners[0].NewConnCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("Failed to receive new connection from wrapped listener: %v", err)
+	}
+	conn := val.(*testutils.ConnWrapper)
+
+	// Add another set of attributes to the addresses, and push them to the
+	// pick_first LB policy through the manual resolver. Leave the order of the
+	// addresses unchanged.
+	for i := range addrs {
+		addrs[i].Attributes = addrs[i].Attributes.WithValue("test-attribute-2", fmt.Sprintf("%d", i))
+	}
+	r.UpdateState(resolver.State{Addresses: addrs})
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// A change in the address attributes results in the new address being
+	// considered different to the current address. This will result in the old
+	// connection being closed and a new connection to the same backend (since
+	// address order is not modified).
+	if _, err := conn.CloseCh.Receive(ctx); err != nil {
+		t.Fatalf("Timeout when expecting existing connection to be closed: %v", err)
+	}
+	val, err = listeners[0].NewConnCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("Failed to receive new connection from wrapped listener: %v", err)
+	}
+	conn = val.(*testutils.ConnWrapper)
+
+	// Add another set of attributes to the addresses, and push them to the
+	// pick_first LB policy through the manual resolver.  Reverse of the order
+	// of addresses.
+	for i := range addrs {
+		addrs[i].Attributes = addrs[i].Attributes.WithValue("test-attribute-3", fmt.Sprintf("%d", i))
+	}
+	addrs[0], addrs[1] = addrs[1], addrs[0]
+	r.UpdateState(resolver.State{Addresses: addrs})
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Ensure that the old connection is closed and a new connection is
+	// established to the first address in the new list.
+	if _, err := conn.CloseCh.Receive(ctx); err != nil {
+		t.Fatalf("Timeout when expecting existing connection to be closed: %v", err)
+	}
+	_, err = listeners[1].NewConnCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("Failed to receive new connection from wrapped listener: %v", err)
+	}
+}
+
+// TestPickFirst_AddressUpdateWithBalancerAttributes tests the case where an
+// address update received by the pick_first LB policy differs in balancer
+// attributes, which are meant only for consumption by LB policies. In this
+// case, the test verifies that new connections are not created when the address
+// update only changes the balancer attributes.
+func (s) TestPickFirst_AddressUpdateWithBalancerAttributes(t *testing.T) {
+	cc, r, backends, listeners := setupPickFirstWithListenerWrapper(t, 2)
+
+	// Add a set of balancer attributes to the addresses before pushing them to
+	// the pick_first LB policy through the manual resolver.
+	addrs := stubBackendsToResolverAddrs(backends)
+	for i := range addrs {
+		addrs[i].BalancerAttributes = addrs[i].BalancerAttributes.WithValue("test-attribute-1", fmt.Sprintf("%d", i))
+	}
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	// Ensure that RPCs succeed to the expected backend.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Grab the wrapped connection from the listener wrapper. This will be used
+	// to verify the connection is not closed.
+	val, err := listeners[0].NewConnCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("Failed to receive new connection from wrapped listener: %v", err)
+	}
+	conn := val.(*testutils.ConnWrapper)
+
+	// Add a set of balancer attributes to the addresses before pushing them to
+	// the pick_first LB policy through the manual resolver. Leave the order of
+	// the addresses unchanged.
+	for i := range addrs {
+		addrs[i].BalancerAttributes = addrs[i].BalancerAttributes.WithValue("test-attribute-2", fmt.Sprintf("%d", i))
+	}
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	// Ensure that no new connection is established, and ensure that the old
+	// connection is not closed.
+	for i := range listeners {
+		sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+		defer sCancel()
+		if _, err := listeners[i].NewConnCh.Receive(sCtx); err != context.DeadlineExceeded {
+			t.Fatalf("Unexpected error when expecting no new connection: %v", err)
+		}
+	}
+	sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel()
+	if _, err := conn.CloseCh.Receive(sCtx); err != context.DeadlineExceeded {
+		t.Fatalf("Unexpected error when expecting existing connection to stay active: %v", err)
+	}
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	// Add a set of balancer attributes to the addresses before pushing them to
+	// the pick_first LB policy through the manual resolver. Reverse of the
+	// order of addresses.
+	for i := range addrs {
+		addrs[i].BalancerAttributes = addrs[i].BalancerAttributes.WithValue("test-attribute-3", fmt.Sprintf("%d", i))
+	}
+	addrs[0], addrs[1] = addrs[1], addrs[0]
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	// Ensure that no new connection is established, and ensure that the old
+	// connection is not closed.
+	for i := range listeners {
+		sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+		defer sCancel()
+		if _, err := listeners[i].NewConnCh.Receive(sCtx); err != context.DeadlineExceeded {
+			t.Fatalf("Unexpected error when expecting no new connection: %v", err)
+		}
+	}
+	sCtx, sCancel = context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel()
+	if _, err := conn.CloseCh.Receive(sCtx); err != context.DeadlineExceeded {
+		t.Fatalf("Unexpected error when expecting existing connection to stay active: %v", err)
+	}
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// Tests the case where the pick_first LB policy receives an error from the name
+// resolver without previously receiving a good update. Verifies that the
+// channel moves to TRANSIENT_FAILURE and that error received from the name
+// resolver is propagated to the caller of an RPC.
+func (s) TestPickFirst_ResolverError_NoPreviousUpdate(t *testing.T) {
+	cc, r, _ := setupPickFirst(t, 0)
+
+	nrErr := errors.New("error from name resolver")
+	r.CC().ReportError(nrErr)
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+
+	client := testgrpc.NewTestServiceClient(cc)
+	_, err := client.EmptyCall(ctx, &testpb.Empty{})
+	if err == nil {
+		t.Fatalf("EmptyCall() succeeded when expected to fail with error: %v", nrErr)
+	}
+	if !strings.Contains(err.Error(), nrErr.Error()) {
+		t.Fatalf("EmptyCall() failed with error: %v, want error: %v", err, nrErr)
+	}
+}
+
+// Tests the case where the pick_first LB policy receives an error from the name
+// resolver after receiving a good update (and the channel is currently READY).
+// The test verifies that the channel continues to use the previously received
+// good update.
+func (s) TestPickFirst_ResolverError_WithPreviousUpdate_Ready(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 1)
+
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	nrErr := errors.New("error from name resolver")
+	r.CC().ReportError(nrErr)
+
+	// Ensure that RPCs continue to succeed for the next second.
+	client := testgrpc.NewTestServiceClient(cc)
+	for end := time.Now().Add(time.Second); time.Now().Before(end); <-time.After(defaultTestShortTimeout) {
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+			t.Fatalf("EmptyCall() failed: %v", err)
+		}
+	}
+}
+
+// Tests the case where the pick_first LB policy receives an error from the name
+// resolver after receiving a good update (and the channel is currently in
+// CONNECTING state). The test verifies that the channel continues to use the
+// previously received good update, and that RPCs don't fail with the error
+// received from the name resolver.
+func (s) TestPickFirst_ResolverError_WithPreviousUpdate_Connecting(t *testing.T) {
+	lis, err := testutils.LocalTCPListener()
+	if err != nil {
+		t.Fatalf("net.Listen() failed: %v", err)
+	}
+
+	// Listen on a local port and act like a server that blocks until the
+	// channel reaches CONNECTING and closes the connection without sending a
+	// server preface.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	waitForConnecting := make(chan struct{})
+	go func() {
+		conn, err := lis.Accept()
+		if err != nil {
+			t.Errorf("Unexpected error when accepting a connection: %v", err)
+		}
+		defer conn.Close()
+
+		select {
+		case <-waitForConnecting:
+		case <-ctx.Done():
+			t.Error("Timeout when waiting for channel to move to CONNECTING state")
+		}
+	}()
+
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithDefaultServiceConfig(pickFirstServiceConfig),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	t.Cleanup(func() { cc.Close() })
+	cc.Connect()
+	addrs := []resolver.Address{{Addr: lis.Addr().String()}}
+	r.UpdateState(resolver.State{Addresses: addrs})
+	testutils.AwaitState(ctx, t, cc, connectivity.Connecting)
+
+	nrErr := errors.New("error from name resolver")
+	r.CC().ReportError(nrErr)
+
+	// RPCs should fail with deadline exceed error as long as they are in
+	// CONNECTING and not the error returned by the name resolver.
+	client := testgrpc.NewTestServiceClient(cc)
+	sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+	defer sCancel()
+	if _, err := client.EmptyCall(sCtx, &testpb.Empty{}); !strings.Contains(err.Error(), context.DeadlineExceeded.Error()) {
+		t.Fatalf("EmptyCall() failed with error: %v, want error: %v", err, context.DeadlineExceeded)
+	}
+
+	// Closing this channel leads to closing of the connection by our listener.
+	// gRPC should see this as a connection error.
+	close(waitForConnecting)
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+	checkForConnectionError(ctx, t, cc)
+}
+
+// Tests the case where the pick_first LB policy receives an error from the name
+// resolver after receiving a good update. The previous good update though has
+// seen the channel move to TRANSIENT_FAILURE.  The test verifies that the
+// channel fails RPCs with the new error from the resolver.
+func (s) TestPickFirst_ResolverError_WithPreviousUpdate_TransientFailure(t *testing.T) {
+	lis, err := testutils.LocalTCPListener()
+	if err != nil {
+		t.Fatalf("net.Listen() failed: %v", err)
+	}
+
+	// Listen on a local port and 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()
+	}()
+
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithDefaultServiceConfig(pickFirstServiceConfig),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	t.Cleanup(func() { cc.Close() })
+	cc.Connect()
+	addrs := []resolver.Address{{Addr: lis.Addr().String()}}
+	r.UpdateState(resolver.State{Addresses: addrs})
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
+	checkForConnectionError(ctx, t, cc)
+
+	// An error from the name resolver should result in RPCs failing with that
+	// error instead of the old error that caused the channel to move to
+	// TRANSIENT_FAILURE in the first place.
+	nrErr := errors.New("error from name resolver")
+	r.CC().ReportError(nrErr)
+	client := testgrpc.NewTestServiceClient(cc)
+	for ; ctx.Err() == nil; <-time.After(defaultTestShortTimeout) {
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}); strings.Contains(err.Error(), nrErr.Error()) {
+			break
+		}
+	}
+	if ctx.Err() != nil {
+		t.Fatal("Timeout when waiting for RPCs to fail with error returned by the name resolver")
+	}
+}
+
+func checkForConnectionError(ctx context.Context, t *testing.T, cc *grpc.ClientConn) {
+	t.Helper()
+
+	// RPCs may fail on the client side in two ways, once the fake server closes
+	// the accepted connection:
+	// - writing the client preface succeeds, but not reading the server preface
+	// - writing the client preface fails
+	// In either case, we should see it fail with UNAVAILABLE.
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); status.Code(err) != codes.Unavailable {
+		t.Fatalf("EmptyCall() failed with error: %v, want code %v", err, codes.Unavailable)
+	}
+}
+
+// Tests the case where the pick_first LB policy receives an update from the
+// name resolver with no addresses after receiving a good update. The test
+// verifies that the channel fails RPCs with an error indicating the fact that
+// the name resolver returned no addresses.
+func (s) TestPickFirst_ResolverError_ZeroAddresses_WithPreviousUpdate(t *testing.T) {
+	cc, r, backends := setupPickFirst(t, 1)
+
+	addrs := stubBackendsToResolverAddrs(backends)
+	r.UpdateState(resolver.State{Addresses: addrs})
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
+		t.Fatal(err)
+	}
+
+	r.UpdateState(resolver.State{})
+	wantErr := "produced zero addresses"
+	client := testgrpc.NewTestServiceClient(cc)
+	for ; ctx.Err() == nil; <-time.After(defaultTestShortTimeout) {
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}); strings.Contains(err.Error(), wantErr) {
+			break
+		}
+	}
+	if ctx.Err() != nil {
+		t.Fatal("Timeout when waiting for RPCs to fail with error returned by the name resolver")
+	}
+}

balancer/pickfirst/pickfirst_test.go

@@ -0,0 +1,132 @@
+/*
+ *
+ * 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/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)
+	}
+}

balancer/pickfirst/pickfirstleaf/metrics_test.go

@@ -0,0 +1,273 @@
+/*
+ *
+ * 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 pickfirstleaf_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/stubserver"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/internal/testutils/stats"
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/serviceconfig"
+	"google.golang.org/grpc/stats/opentelemetry"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/sdk/metric"
+	"go.opentelemetry.io/otel/sdk/metric/metricdata"
+	"go.opentelemetry.io/otel/sdk/metric/metricdata/metricdatatest"
+)
+
+var pfConfig string
+
+func init() {
+	pfConfig = fmt.Sprintf(`{
+  		"loadBalancingConfig": [
+    		{
+      			%q: {
+      		}
+    	}
+  	]
+	}`, pickfirstleaf.Name)
+}
+
+// TestPickFirstMetrics tests pick first metrics. It configures a pick first
+// balancer, causes it to connect and then disconnect, and expects the
+// subsequent metrics to emit from that.
+func (s) TestPickFirstMetrics(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	ss := &stubserver.StubServer{
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			return &testpb.Empty{}, nil
+		},
+	}
+	ss.StartServer()
+	defer ss.Stop()
+
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(pfConfig)
+
+	r := manual.NewBuilderWithScheme("whatever")
+	r.InitialState(resolver.State{
+		ServiceConfig: sc,
+		Addresses:     []resolver.Address{{Addr: ss.Address}}},
+	)
+
+	tmr := stats.NewTestMetricsRecorder()
+	cc, err := grpc.NewClient(r.Scheme()+":///", grpc.WithStatsHandler(tmr), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
+	if err != nil {
+		t.Fatalf("NewClient() failed with error: %v", err)
+	}
+	defer cc.Close()
+
+	tsc := testgrpc.NewTestServiceClient(cc)
+	if _, err := tsc.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall() failed: %v", err)
+	}
+
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 1 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 1)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 0)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
+	}
+
+	ss.Stop()
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 1 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 1)
+	}
+}
+
+// TestPickFirstMetricsFailure tests the connection attempts failed metric. It
+// configures a channel and scenario that causes a pick first connection attempt
+// to fail, and then expects that metric to emit.
+func (s) TestPickFirstMetricsFailure(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(pfConfig)
+
+	r := manual.NewBuilderWithScheme("whatever")
+	r.InitialState(resolver.State{
+		ServiceConfig: sc,
+		Addresses:     []resolver.Address{{Addr: "bad address"}}},
+	)
+	grpcTarget := r.Scheme() + ":///"
+	tmr := stats.NewTestMetricsRecorder()
+	cc, err := grpc.NewClient(grpcTarget, grpc.WithStatsHandler(tmr), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
+	if err != nil {
+		t.Fatalf("NewClient() failed with error: %v", err)
+	}
+	defer cc.Close()
+
+	tsc := testgrpc.NewTestServiceClient(cc)
+	if _, err := tsc.EmptyCall(ctx, &testpb.Empty{}); err == nil {
+		t.Fatalf("EmptyCall() passed when expected to fail")
+	}
+
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 0)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 1 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 1)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
+	}
+}
+
+// TestPickFirstMetricsE2E tests the pick first metrics end to end. It
+// configures a channel with an OpenTelemetry plugin, induces all 3 pick first
+// metrics to emit, and makes sure the correct OpenTelemetry metrics atoms emit.
+func (s) TestPickFirstMetricsE2E(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	ss := &stubserver.StubServer{
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			return &testpb.Empty{}, nil
+		},
+	}
+	ss.StartServer()
+	defer ss.Stop()
+
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(pfConfig)
+	r := manual.NewBuilderWithScheme("whatever")
+	r.InitialState(resolver.State{
+		ServiceConfig: sc,
+		Addresses:     []resolver.Address{{Addr: "bad address"}}},
+	) // Will trigger connection failed.
+
+	grpcTarget := r.Scheme() + ":///"
+	reader := metric.NewManualReader()
+	provider := metric.NewMeterProvider(metric.WithReader(reader))
+	mo := opentelemetry.MetricsOptions{
+		MeterProvider: provider,
+		Metrics:       opentelemetry.DefaultMetrics().Add("grpc.lb.pick_first.disconnections", "grpc.lb.pick_first.connection_attempts_succeeded", "grpc.lb.pick_first.connection_attempts_failed"),
+	}
+
+	cc, err := grpc.NewClient(grpcTarget, opentelemetry.DialOption(opentelemetry.Options{MetricsOptions: mo}), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
+	if err != nil {
+		t.Fatalf("NewClient() failed with error: %v", err)
+	}
+	defer cc.Close()
+
+	tsc := testgrpc.NewTestServiceClient(cc)
+	if _, err := tsc.EmptyCall(ctx, &testpb.Empty{}); err == nil {
+		t.Fatalf("EmptyCall() passed when expected to fail")
+	}
+
+	r.UpdateState(resolver.State{
+		ServiceConfig: sc,
+		Addresses:     []resolver.Address{{Addr: ss.Address}},
+	}) // Will trigger successful connection metric.
+	if _, err := tsc.EmptyCall(ctx, &testpb.Empty{}, grpc.WaitForReady(true)); err != nil {
+		t.Fatalf("EmptyCall() failed: %v", err)
+	}
+
+	// Stop the server, that should send signal to disconnect, which will
+	// eventually emit disconnection metric before ClientConn goes IDLE.
+	ss.Stop()
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+	wantMetrics := []metricdata.Metrics{
+		{
+			Name:        "grpc.lb.pick_first.connection_attempts_succeeded",
+			Description: "EXPERIMENTAL. Number of successful connection attempts.",
+			Unit:        "{attempt}",
+			Data: metricdata.Sum[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget)),
+						Value:      1,
+					},
+				},
+				Temporality: metricdata.CumulativeTemporality,
+				IsMonotonic: true,
+			},
+		},
+		{
+			Name:        "grpc.lb.pick_first.connection_attempts_failed",
+			Description: "EXPERIMENTAL. Number of failed connection attempts.",
+			Unit:        "{attempt}",
+			Data: metricdata.Sum[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget)),
+						Value:      1,
+					},
+				},
+				Temporality: metricdata.CumulativeTemporality,
+				IsMonotonic: true,
+			},
+		},
+		{
+			Name:        "grpc.lb.pick_first.disconnections",
+			Description: "EXPERIMENTAL. Number of times the selected subchannel becomes disconnected.",
+			Unit:        "{disconnection}",
+			Data: metricdata.Sum[int64]{
+				DataPoints: []metricdata.DataPoint[int64]{
+					{
+						Attributes: attribute.NewSet(attribute.String("grpc.target", grpcTarget)),
+						Value:      1,
+					},
+				},
+				Temporality: metricdata.CumulativeTemporality,
+				IsMonotonic: true,
+			},
+		},
+	}
+
+	gotMetrics := metricsDataFromReader(ctx, reader)
+	for _, metric := range wantMetrics {
+		val, ok := gotMetrics[metric.Name]
+		if !ok {
+			t.Fatalf("Metric %v not present in recorded metrics", metric.Name)
+		}
+		if !metricdatatest.AssertEqual(t, metric, val, metricdatatest.IgnoreTimestamp(), metricdatatest.IgnoreExemplars()) {
+			t.Fatalf("Metrics data type not equal for metric: %v", metric.Name)
+		}
+	}
+}
+
+func metricsDataFromReader(ctx context.Context, reader *metric.ManualReader) map[string]metricdata.Metrics {
+	rm := &metricdata.ResourceMetrics{}
+	reader.Collect(ctx, rm)
+	gotMetrics := map[string]metricdata.Metrics{}
+	for _, sm := range rm.ScopeMetrics {
+		for _, m := range sm.Metrics {
+			gotMetrics[m.Name] = m
+		}
+	}
+	return gotMetrics
+}

balancer/pickfirst/pickfirstleaf/pickfirstleaf.go

@@ -0,0 +1,906 @@
+/*
+ *
+ * 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 pickfirstleaf contains the pick_first load balancing policy which
+// will be the universal leaf policy after dualstack changes are implemented.
+//
+// # Experimental
+//
+// Notice: This package is EXPERIMENTAL and may be changed or removed in a
+// later release.
+package pickfirstleaf
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net"
+	"net/netip"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst/internal"
+	"google.golang.org/grpc/connectivity"
+	expstats "google.golang.org/grpc/experimental/stats"
+	"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() {
+	if envconfig.NewPickFirstEnabled {
+		// Register as the default pick_first balancer.
+		Name = "pick_first"
+	}
+	balancer.Register(pickfirstBuilder{})
+}
+
+// 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")
+	// Name is the name of the pick_first_leaf balancer.
+	// It is changed to "pick_first" in init() if this balancer is to be
+	// registered as the default pickfirst.
+	Name                 = "pick_first_leaf"
+	disconnectionsMetric = expstats.RegisterInt64Count(expstats.MetricDescriptor{
+		Name:        "grpc.lb.pick_first.disconnections",
+		Description: "EXPERIMENTAL. Number of times the selected subchannel becomes disconnected.",
+		Unit:        "{disconnection}",
+		Labels:      []string{"grpc.target"},
+		Default:     false,
+	})
+	connectionAttemptsSucceededMetric = expstats.RegisterInt64Count(expstats.MetricDescriptor{
+		Name:        "grpc.lb.pick_first.connection_attempts_succeeded",
+		Description: "EXPERIMENTAL. Number of successful connection attempts.",
+		Unit:        "{attempt}",
+		Labels:      []string{"grpc.target"},
+		Default:     false,
+	})
+	connectionAttemptsFailedMetric = expstats.RegisterInt64Count(expstats.MetricDescriptor{
+		Name:        "grpc.lb.pick_first.connection_attempts_failed",
+		Description: "EXPERIMENTAL. Number of failed connection attempts.",
+		Unit:        "{attempt}",
+		Labels:      []string{"grpc.target"},
+		Default:     false,
+	})
+)
+
+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, bo balancer.BuildOptions) balancer.Balancer {
+	b := &pickfirstBalancer{
+		cc:              cc,
+		target:          bo.Target.String(),
+		metricsRecorder: cc.MetricsRecorder(),
+
+		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.
+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
+	target          string
+	metricsRecorder expstats.MetricsRecorder // guaranteed to be non nil
+
+	// 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 {
+			endpoints = append([]resolver.Endpoint{}, 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(endpoints), func(i, j int) { endpoints[i], endpoints[j] = endpoints[j], endpoints[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 {
+		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.Values() {
+		sd.connectionFailedInFirstPass = false
+	}
+	b.requestConnectionLocked()
+}
+
+func (b *pickfirstBalancer) closeSubConnsLocked() {
+	for _, sd := range b.subConns.Values() {
+		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[*scData]()
+	retAddrs := []resolver.Address{}
+
+	for _, addr := range addrs {
+		if _, ok := seenAddrs.Get(addr); ok {
+			continue
+		}
+		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.Keys() {
+		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.Values() {
+		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
+	}
+
+	// Record a connection attempt when exiting CONNECTING.
+	if newState.ConnectivityState == connectivity.TransientFailure {
+		sd.connectionFailedInFirstPass = true
+		connectionAttemptsFailedMetric.Record(b.metricsRecorder, 1, b.target)
+	}
+
+	if newState.ConnectivityState == connectivity.Ready {
+		connectionAttemptsSucceededMetric.Record(b.metricsRecorder, 1, b.target)
+		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
+		// READY SubConn interspliced in between CONNECTING and IDLE, need to
+		// account for that.
+		if oldState == connectivity.Connecting {
+			// A known issue (https://github.com/grpc/grpc-go/issues/7862)
+			// causes a race that prevents the READY state change notification.
+			// This works around it.
+			connectionAttemptsSucceededMetric.Record(b.metricsRecorder, 1, b.target)
+		}
+		disconnectionsMetric.Record(b.metricsRecorder, 1, b.target)
+		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.Values() {
+		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.Values() {
+		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)
+}

balancer/pickfirst/pickfirstleaf/pickfirstleaf_test.go

@@ -0,0 +1,246 @@
+/*
+ *
+ * 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 pickfirstleaf
+
+import (
+	"context"
+	"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{})
+}
+
+// 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: "2.2.2.2:2"}}},
+			},
+		},
+	}
+	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
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.TransientFailure, ConnectionError: tfErr})
+
+	if err := cc.WaitForPickerWithErr(ctx, balancer.ErrNoSubConnAvailable); err != nil {
+		t.Fatalf("cc.WaitForPickerWithErr(%v) returned error: %v", balancer.ErrNoSubConnAvailable, err)
+	}
+
+	sc2 := <-cc.NewSubConnCh
+	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.
+	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)
+	}
+}

balancer/ringhash/config.go

@@ -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
+}

balancer/ringhash/config_test.go

@@ -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)
+			}
+		})
+	}
+}

balancer/ringhash/picker.go

@@ -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.balancer.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]
+}

balancer/ringhash/picker_test.go

@@ -0,0 +1,311 @@
+/*
+ *
+ * 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
+	}
+}
+
+type fakeExitIdler struct {
+	sc *testutils.TestSubConn
+}
+
+func (ei *fakeExitIdler) ExitIdle() {
+	ei.sc.Connect()
+}
+
+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,
+				},
+			},
+			balancer: &fakeExitIdler{
+				sc: testSC,
+			},
+		}
+		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)
+				}
+			}
+		})
+	}
+}

balancer/ringhash/ring.go

@@ -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)]
+}

balancer/ringhash/ring_test.go

@@ -24,42 +24,49 @@ import (
 	"testing"
 
 	xxhash "github.com/cespare/xxhash/v2"
+	"google.golang.org/grpc/internal/balancer/weight"
 	"google.golang.org/grpc/resolver"
 )
 
-func testAddr(addr string, weight uint32) resolver.Address {
-	return resolver.Address{Addr: addr, Metadata: weight}
+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 TestRingNew(t *testing.T) {
-	testAddrs := []resolver.Address{
-		testAddr("a", 3),
-		testAddr("b", 3),
-		testAddr("c", 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
-	testSubConnMap := map[resolver.Address]*subConn{
-		testAddr("a", 3): {addr: "a"},
-		testAddr("b", 3): {addr: "b"},
-		testAddr("c", 4): {addr: "c"},
-	}
 	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(testSubConnMap, min, max)
+				r := newRing(testEndpointStateMap, min, max, nil)
 				totalCount := len(r.items)
 				if totalCount < int(min) || totalCount > int(max) {
-					t.Fatalf("unexpect size %v, want min %v, max %v", totalCount, min, max)
+					t.Fatalf("unexpected size %v, want min %v, max %v", totalCount, min, max)
 				}
-				for _, a := range testAddrs {
+				for _, e := range testEndpoints {
 					var count int
 					for _, ii := range r.items {
-						if ii.sc.addr == a.Addr {
+						if ii.hashKey == hashKey(e) {
 							count++
 						}
 					}
 					got := float64(count) / float64(totalCount)
-					want := float64(a.Metadata.(uint32)) / totalWeight
+					want := float64(getWeightAttribute(e)) / totalWeight
 					if !equalApproximately(got, want) {
 						t.Fatalf("unexpected item weight in ring: %v != %v", got, want)
 					}
@@ -75,12 +82,8 @@ func equalApproximately(x, y float64) bool {
 	return delta/mean < 0.25
 }
 
-func TestRingPick(t *testing.T) {
-	r, _ := newRing(map[resolver.Address]*subConn{
-		{Addr: "a", Metadata: uint32(3)}: {addr: "a"},
-		{Addr: "b", Metadata: uint32(3)}: {addr: "b"},
-		{Addr: "c", Metadata: uint32(4)}: {addr: "c"},
-	}, 10, 20)
+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)
@@ -97,12 +100,8 @@ func TestRingPick(t *testing.T) {
 	}
 }
 
-func TestRingNext(t *testing.T) {
-	r, _ := newRing(map[resolver.Address]*subConn{
-		{Addr: "a", Metadata: uint32(3)}: {addr: "a"},
-		{Addr: "b", Metadata: uint32(3)}: {addr: "b"},
-		{Addr: "c", Metadata: uint32(4)}: {addr: "c"},
-	}, 10, 20)
+func (s) TestRingNext(t *testing.T) {
+	r := newRing(testEndpointStateMap, 10, 20, nil)
 
 	for _, e := range r.items {
 		ne := r.next(e)

balancer/ringhash/ringhash.go

@@ -0,0 +1,408 @@
+/*
+ *
+ * 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/pickfirstleaf"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal/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(pickfirstleaf.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{
+				balancer: childState.Balancer,
+				hashKey:  hk,
+				weight:   newWeight,
+				state:    childState.State,
+			}
+			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.Keys() {
+		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: pickfirstleaf.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.
+		var endpointStates = make([]*endpointState, b.endpointStates.Len())
+		for i, s := range b.endpointStates.Values() {
+			endpointStates[i] = s
+		}
+		sort.Slice(endpointStates, func(i, j int) bool {
+			return endpointStates[i].hashKey < endpointStates[j].hashKey
+		})
+		var idleBalancer endpointsharding.ExitIdler
+		for _, es := range endpointStates {
+			connState := es.state.ConnectivityState
+			if connState == connectivity.Connecting {
+				idleBalancer = nil
+				break
+			}
+			if idleBalancer == nil && connState == connectivity.Idle {
+				idleBalancer = es.balancer
+			}
+		}
+		if idleBalancer != nil {
+			idleBalancer.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.Values() {
+		// 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.Values() {
+		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
+	balancer endpointsharding.ExitIdler
+
+	// 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
+}

balancer/ringhash/ringhash_test.go

@@ -0,0 +1,737 @@
+/*
+ *
+ * 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/internal/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)
+	}
+	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):
+	}
+}

balancer/rls/balancer.go

@@ -0,0 +1,704 @@
+/*
+ *
+ * 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 rls implements the RLS LB policy.
+package rls
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"time"
+	"unsafe"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/backoff"
+	"google.golang.org/grpc/internal/balancergroup"
+	"google.golang.org/grpc/internal/buffer"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/pretty"
+	"google.golang.org/grpc/resolver"
+)
+
+const (
+	// Name is the name of the RLS LB policy.
+	//
+	// It currently has an experimental suffix which would be removed once
+	// end-to-end testing of the policy is completed.
+	Name = internal.RLSLoadBalancingPolicyName
+	// Default frequency for data cache purging.
+	periodicCachePurgeFreq = time.Minute
+)
+
+var (
+	logger            = grpclog.Component("rls")
+	errBalancerClosed = errors.New("rls LB policy is closed")
+
+	// Below defined vars for overriding in unit tests.
+
+	// Default exponential backoff strategy for data cache entries.
+	defaultBackoffStrategy = backoff.Strategy(backoff.DefaultExponential)
+	// Ticker used for periodic data cache purging.
+	dataCachePurgeTicker = func() *time.Ticker { return time.NewTicker(periodicCachePurgeFreq) }
+	// We want every cache entry to live in the cache for at least this
+	// duration. If we encounter a cache entry whose minimum expiration time is
+	// in the future, we abort the LRU pass, which may temporarily leave the
+	// cache being too large. This is necessary to ensure that in cases where
+	// the cache is too small, when we receive an RLS Response, we keep the
+	// resulting cache entry around long enough for the pending incoming
+	// requests to be re-processed through the new Picker. If we didn't do this,
+	// then we'd risk throwing away each RLS response as we receive it, in which
+	// case we would fail to actually route any of our incoming requests.
+	minEvictDuration = 5 * time.Second
+
+	// Following functions are no-ops in actual code, but can be overridden in
+	// tests to give tests visibility into exactly when certain events happen.
+	clientConnUpdateHook = func() {}
+	dataCachePurgeHook   = func() {}
+	resetBackoffHook     = func() {}
+
+	cacheEntriesMetric = estats.RegisterInt64Gauge(estats.MetricDescriptor{
+		Name:        "grpc.lb.rls.cache_entries",
+		Description: "EXPERIMENTAL. Number of entries in the RLS cache.",
+		Unit:        "{entry}",
+		Labels:      []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.instance_uuid"},
+		Default:     false,
+	})
+	cacheSizeMetric = estats.RegisterInt64Gauge(estats.MetricDescriptor{
+		Name:        "grpc.lb.rls.cache_size",
+		Description: "EXPERIMENTAL. The current size of the RLS cache.",
+		Unit:        "By",
+		Labels:      []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.instance_uuid"},
+		Default:     false,
+	})
+	defaultTargetPicksMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:        "grpc.lb.rls.default_target_picks",
+		Description: "EXPERIMENTAL. Number of LB picks sent to the default target.",
+		Unit:        "{pick}",
+		Labels:      []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.data_plane_target", "grpc.lb.pick_result"},
+		Default:     false,
+	})
+	targetPicksMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:        "grpc.lb.rls.target_picks",
+		Description: "EXPERIMENTAL. Number of LB picks sent to each RLS target. Note that if the default target is also returned by the RLS server, RPCs sent to that target from the cache will be counted in this metric, not in grpc.rls.default_target_picks.",
+		Unit:        "{pick}",
+		Labels:      []string{"grpc.target", "grpc.lb.rls.server_target", "grpc.lb.rls.data_plane_target", "grpc.lb.pick_result"},
+		Default:     false,
+	})
+	failedPicksMetric = estats.RegisterInt64Count(estats.MetricDescriptor{
+		Name:        "grpc.lb.rls.failed_picks",
+		Description: "EXPERIMENTAL. Number of LB picks failed due to either a failed RLS request or the RLS channel being throttled.",
+		Unit:        "{pick}",
+		Labels:      []string{"grpc.target", "grpc.lb.rls.server_target"},
+		Default:     false,
+	})
+)
+
+func init() {
+	balancer.Register(&rlsBB{})
+}
+
+type rlsBB struct{}
+
+func (rlsBB) Name() string {
+	return Name
+}
+
+func (rlsBB) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	lb := &rlsBalancer{
+		closed:             grpcsync.NewEvent(),
+		done:               grpcsync.NewEvent(),
+		cc:                 cc,
+		bopts:              opts,
+		purgeTicker:        dataCachePurgeTicker(),
+		dataCachePurgeHook: dataCachePurgeHook,
+		lbCfg:              &lbConfig{},
+		pendingMap:         make(map[cacheKey]*backoffState),
+		childPolicies:      make(map[string]*childPolicyWrapper),
+		updateCh:           buffer.NewUnbounded(),
+	}
+	lb.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-experimental-lb %p] ", lb))
+	lb.dataCache = newDataCache(maxCacheSize, lb.logger, cc.MetricsRecorder(), opts.Target.String())
+	lb.bg = balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               opts,
+		StateAggregator:         lb,
+		Logger:                  lb.logger,
+		SubBalancerCloseTimeout: time.Duration(0), // Disable caching of removed child policies
+	})
+	go lb.run()
+	return lb
+}
+
+// rlsBalancer implements the RLS LB policy.
+type rlsBalancer struct {
+	closed             *grpcsync.Event // Fires when Close() is invoked. Guarded by stateMu.
+	done               *grpcsync.Event // Fires when Close() is done.
+	cc                 balancer.ClientConn
+	bopts              balancer.BuildOptions
+	purgeTicker        *time.Ticker
+	dataCachePurgeHook func()
+	logger             *internalgrpclog.PrefixLogger
+
+	// If both cacheMu and stateMu need to be acquired, the former must be
+	// acquired first to prevent a deadlock. This order restriction is due to the
+	// fact that in places where we need to acquire both the locks, we always
+	// start off reading the cache.
+
+	// cacheMu guards access to the data cache and pending requests map. We
+	// cannot use an RWMutex here since even an operation like
+	// dataCache.getEntry() modifies the underlying LRU, which is implemented as
+	// a doubly linked list.
+	cacheMu    sync.Mutex
+	dataCache  *dataCache                 // Cache of RLS data.
+	pendingMap map[cacheKey]*backoffState // Map of pending RLS requests.
+
+	// stateMu guards access to all LB policy state.
+	stateMu            sync.Mutex
+	lbCfg              *lbConfig        // Most recently received service config.
+	childPolicyBuilder balancer.Builder // Cached child policy builder.
+	resolverState      resolver.State   // Cached resolver state.
+	ctrlCh             *controlChannel  // Control channel to the RLS server.
+	bg                 *balancergroup.BalancerGroup
+	childPolicies      map[string]*childPolicyWrapper
+	defaultPolicy      *childPolicyWrapper
+	// A reference to the most recent picker sent to gRPC as part of a state
+	// update is cached in this field so that we can release the reference to the
+	// default child policy wrapper when a new picker is created. See
+	// sendNewPickerLocked() for details.
+	lastPicker *rlsPicker
+	// Set during UpdateClientConnState when pushing updates to child policies.
+	// Prevents state updates from child policies causing new pickers to be sent
+	// up the channel. Cleared after all child policies have processed the
+	// updates sent to them, after which a new picker is sent up the channel.
+	inhibitPickerUpdates bool
+
+	// Channel on which all updates are pushed. Processed in run().
+	updateCh *buffer.Unbounded
+}
+
+type resumePickerUpdates struct {
+	done chan struct{}
+}
+
+// childPolicyIDAndState wraps a child policy id and its state update.
+type childPolicyIDAndState struct {
+	id    string
+	state balancer.State
+}
+
+type controlChannelReady struct{}
+
+// run is a long-running goroutine which handles all the updates that the
+// balancer wishes to handle. The appropriate updateHandler will push the update
+// on to a channel that this goroutine will select on, thereby the handling of
+// the update will happen asynchronously.
+func (b *rlsBalancer) run() {
+	// We exit out of the for loop below only after `Close()` has been invoked.
+	// Firing the done event here will ensure that Close() returns only after
+	// all goroutines are done.
+	defer func() { b.done.Fire() }()
+
+	// Wait for purgeDataCache() goroutine to exit before returning from here.
+	doneCh := make(chan struct{})
+	defer func() {
+		<-doneCh
+	}()
+	go b.purgeDataCache(doneCh)
+
+	for {
+		select {
+		case u, ok := <-b.updateCh.Get():
+			if !ok {
+				return
+			}
+			b.updateCh.Load()
+			switch update := u.(type) {
+			case childPolicyIDAndState:
+				b.handleChildPolicyStateUpdate(update.id, update.state)
+			case controlChannelReady:
+				b.logger.Infof("Resetting backoff state after control channel getting back to READY")
+				b.cacheMu.Lock()
+				updatePicker := b.dataCache.resetBackoffState(&backoffState{bs: defaultBackoffStrategy})
+				b.cacheMu.Unlock()
+				if updatePicker {
+					b.sendNewPicker()
+				}
+				resetBackoffHook()
+			case resumePickerUpdates:
+				b.stateMu.Lock()
+				b.logger.Infof("Resuming picker updates after config propagation to child policies")
+				b.inhibitPickerUpdates = false
+				b.sendNewPickerLocked()
+				close(update.done)
+				b.stateMu.Unlock()
+			default:
+				b.logger.Errorf("Unsupported update type %T", update)
+			}
+		case <-b.closed.Done():
+			return
+		}
+	}
+}
+
+// purgeDataCache is a long-running goroutine which periodically deletes expired
+// entries. An expired entry is one for which both the expiryTime and
+// backoffExpiryTime are in the past.
+func (b *rlsBalancer) purgeDataCache(doneCh chan struct{}) {
+	defer close(doneCh)
+
+	for {
+		select {
+		case <-b.closed.Done():
+			return
+		case <-b.purgeTicker.C:
+			b.cacheMu.Lock()
+			updatePicker := b.dataCache.evictExpiredEntries()
+			b.cacheMu.Unlock()
+			if updatePicker {
+				b.sendNewPicker()
+			}
+			b.dataCachePurgeHook()
+		}
+	}
+}
+
+func (b *rlsBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	defer clientConnUpdateHook()
+
+	b.stateMu.Lock()
+	if b.closed.HasFired() {
+		b.stateMu.Unlock()
+		b.logger.Warningf("Received service config after balancer close: %s", pretty.ToJSON(ccs.BalancerConfig))
+		return errBalancerClosed
+	}
+
+	newCfg := ccs.BalancerConfig.(*lbConfig)
+	if b.lbCfg.Equal(newCfg) {
+		b.stateMu.Unlock()
+		b.logger.Infof("New service config matches existing config")
+		return nil
+	}
+
+	b.logger.Infof("Delaying picker updates until config is propagated to and processed by child policies")
+	b.inhibitPickerUpdates = true
+
+	// When the RLS server name changes, the old control channel needs to be
+	// swapped out for a new one. All state associated with the throttling
+	// algorithm is stored on a per-control-channel basis; when we swap out
+	// channels, we also swap out the throttling state.
+	b.handleControlChannelUpdate(newCfg)
+
+	// Any changes to child policy name or configuration needs to be handled by
+	// either creating new child policies or pushing updates to existing ones.
+	b.resolverState = ccs.ResolverState
+	b.handleChildPolicyConfigUpdate(newCfg, &ccs)
+
+	// Resize the cache if the size in the config has changed.
+	resizeCache := newCfg.cacheSizeBytes != b.lbCfg.cacheSizeBytes
+
+	// Update the copy of the config in the LB policy before releasing the lock.
+	b.lbCfg = newCfg
+	b.stateMu.Unlock()
+
+	// We cannot do cache operations above because `cacheMu` needs to be grabbed
+	// before `stateMu` if we are to hold both locks at the same time.
+	b.cacheMu.Lock()
+	b.dataCache.updateRLSServerTarget(newCfg.lookupService)
+	if resizeCache {
+		// If the new config changes reduces the size of the data cache, we
+		// might have to evict entries to get the cache size down to the newly
+		// specified size. If we do evict an entry with valid backoff timer,
+		// the new picker needs to be sent to the channel to re-process any
+		// RPCs queued as a result of this backoff timer.
+		b.dataCache.resize(newCfg.cacheSizeBytes)
+	}
+	b.cacheMu.Unlock()
+	// Enqueue an event which will notify us when the above update has been
+	// propagated to all child policies, and the child policies have all
+	// processed their updates, and we have sent a picker update.
+	done := make(chan struct{})
+	b.updateCh.Put(resumePickerUpdates{done: done})
+	<-done
+	return nil
+}
+
+// handleControlChannelUpdate handles updates to service config fields which
+// influence the control channel to the RLS server.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) handleControlChannelUpdate(newCfg *lbConfig) {
+	if newCfg.lookupService == b.lbCfg.lookupService && newCfg.lookupServiceTimeout == b.lbCfg.lookupServiceTimeout {
+		return
+	}
+
+	// Create a new control channel and close the existing one.
+	b.logger.Infof("Creating control channel to RLS server at: %v", newCfg.lookupService)
+	backToReadyFn := func() {
+		b.updateCh.Put(controlChannelReady{})
+	}
+	ctrlCh, err := newControlChannel(newCfg.lookupService, newCfg.controlChannelServiceConfig, newCfg.lookupServiceTimeout, b.bopts, backToReadyFn)
+	if err != nil {
+		// This is very uncommon and usually represents a non-transient error.
+		// There is not much we can do here other than wait for another update
+		// which might fix things.
+		b.logger.Errorf("Failed to create control channel to %q: %v", newCfg.lookupService, err)
+		return
+	}
+	if b.ctrlCh != nil {
+		b.ctrlCh.close()
+	}
+	b.ctrlCh = ctrlCh
+}
+
+// handleChildPolicyConfigUpdate handles updates to service config fields which
+// influence child policy configuration.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) handleChildPolicyConfigUpdate(newCfg *lbConfig, ccs *balancer.ClientConnState) {
+	// Update child policy builder first since other steps are dependent on this.
+	if b.childPolicyBuilder == nil || b.childPolicyBuilder.Name() != newCfg.childPolicyName {
+		b.logger.Infof("Child policy changed to %q", newCfg.childPolicyName)
+		b.childPolicyBuilder = balancer.Get(newCfg.childPolicyName)
+		for _, cpw := range b.childPolicies {
+			// If the child policy has changed, we need to remove the old policy
+			// from the BalancerGroup and add a new one. The BalancerGroup takes
+			// care of closing the old one in this case.
+			b.bg.Remove(cpw.target)
+			b.bg.Add(cpw.target, b.childPolicyBuilder)
+		}
+	}
+
+	configSentToDefault := false
+	if b.lbCfg.defaultTarget != newCfg.defaultTarget {
+		// If the default target has changed, create a new childPolicyWrapper for
+		// the new target if required. If a new wrapper is created, add it to the
+		// childPolicies map and the BalancerGroup.
+		b.logger.Infof("Default target in LB config changing from %q to %q", b.lbCfg.defaultTarget, newCfg.defaultTarget)
+		cpw := b.childPolicies[newCfg.defaultTarget]
+		if cpw == nil {
+			cpw = newChildPolicyWrapper(newCfg.defaultTarget)
+			b.childPolicies[newCfg.defaultTarget] = cpw
+			b.bg.Add(newCfg.defaultTarget, b.childPolicyBuilder)
+			b.logger.Infof("Child policy %q added to BalancerGroup", newCfg.defaultTarget)
+		}
+		if err := b.buildAndPushChildPolicyConfigs(newCfg.defaultTarget, newCfg, ccs); err != nil {
+			cpw.lamify(err)
+		}
+
+		// If an old default exists, release its reference. If this was the last
+		// reference, remove the child policy from the BalancerGroup and remove the
+		// corresponding entry the childPolicies map.
+		if b.defaultPolicy != nil {
+			if b.defaultPolicy.releaseRef() {
+				delete(b.childPolicies, b.lbCfg.defaultTarget)
+				b.bg.Remove(b.defaultPolicy.target)
+			}
+		}
+		b.defaultPolicy = cpw
+		configSentToDefault = true
+	}
+
+	// No change in configuration affecting child policies. Return early.
+	if b.lbCfg.childPolicyName == newCfg.childPolicyName && b.lbCfg.childPolicyTargetField == newCfg.childPolicyTargetField && childPolicyConfigEqual(b.lbCfg.childPolicyConfig, newCfg.childPolicyConfig) {
+		return
+	}
+
+	// If fields affecting child policy configuration have changed, the changes
+	// are pushed to the childPolicyWrapper which handles them appropriately.
+	for _, cpw := range b.childPolicies {
+		if configSentToDefault && cpw.target == newCfg.defaultTarget {
+			// Default target has already been taken care of.
+			continue
+		}
+		if err := b.buildAndPushChildPolicyConfigs(cpw.target, newCfg, ccs); err != nil {
+			cpw.lamify(err)
+		}
+	}
+}
+
+// buildAndPushChildPolicyConfigs builds the final child policy configuration by
+// adding the `targetField` to the base child policy configuration received in
+// RLS LB policy configuration. The `targetField` is set to target and
+// configuration is pushed to the child policy through the BalancerGroup.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) buildAndPushChildPolicyConfigs(target string, newCfg *lbConfig, ccs *balancer.ClientConnState) error {
+	jsonTarget, err := json.Marshal(target)
+	if err != nil {
+		return fmt.Errorf("failed to marshal child policy target %q: %v", target, err)
+	}
+
+	config := newCfg.childPolicyConfig
+	targetField := newCfg.childPolicyTargetField
+	config[targetField] = jsonTarget
+	jsonCfg, err := json.Marshal(config)
+	if err != nil {
+		return fmt.Errorf("failed to marshal child policy config %+v: %v", config, err)
+	}
+
+	parser, _ := b.childPolicyBuilder.(balancer.ConfigParser)
+	parsedCfg, err := parser.ParseConfig(jsonCfg)
+	if err != nil {
+		return fmt.Errorf("childPolicy config parsing failed: %v", err)
+	}
+
+	state := balancer.ClientConnState{ResolverState: ccs.ResolverState, BalancerConfig: parsedCfg}
+	b.logger.Infof("Pushing new state to child policy %q: %+v", target, state)
+	if err := b.bg.UpdateClientConnState(target, state); err != nil {
+		b.logger.Warningf("UpdateClientConnState(%q, %+v) failed : %v", target, ccs, err)
+	}
+	return nil
+}
+
+func (b *rlsBalancer) ResolverError(err error) {
+	b.bg.ResolverError(err)
+}
+
+func (b *rlsBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+func (b *rlsBalancer) Close() {
+	b.stateMu.Lock()
+	b.closed.Fire()
+	b.purgeTicker.Stop()
+	if b.ctrlCh != nil {
+		b.ctrlCh.close()
+	}
+	b.bg.Close()
+	b.stateMu.Unlock()
+
+	b.cacheMu.Lock()
+	b.dataCache.stop()
+	b.cacheMu.Unlock()
+
+	b.updateCh.Close()
+
+	<-b.done.Done()
+}
+
+func (b *rlsBalancer) ExitIdle() {
+	b.bg.ExitIdle()
+}
+
+// sendNewPickerLocked pushes a new picker on to the channel.
+//
+// Note that regardless of what connectivity state is reported, the policy will
+// return its own picker, and not a picker that unconditionally queues
+// (typically used for IDLE or CONNECTING) or a picker that unconditionally
+// fails (typically used for TRANSIENT_FAILURE). This is required because,
+// irrespective of the connectivity state, we need to able to perform RLS
+// lookups for incoming RPCs and affect the status of queued RPCs based on the
+// receipt of RLS responses.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) sendNewPickerLocked() {
+	aggregatedState := b.aggregatedConnectivityState()
+
+	// Acquire a separate reference for the picker. This is required to ensure
+	// that the wrapper held by the old picker is not closed when the default
+	// target changes in the config, and a new wrapper is created for the new
+	// default target. See handleChildPolicyConfigUpdate() for how config changes
+	// affecting the default target are handled.
+	if b.defaultPolicy != nil {
+		b.defaultPolicy.acquireRef()
+	}
+
+	picker := &rlsPicker{
+		kbm:             b.lbCfg.kbMap,
+		origEndpoint:    b.bopts.Target.Endpoint(),
+		lb:              b,
+		defaultPolicy:   b.defaultPolicy,
+		ctrlCh:          b.ctrlCh,
+		maxAge:          b.lbCfg.maxAge,
+		staleAge:        b.lbCfg.staleAge,
+		bg:              b.bg,
+		rlsServerTarget: b.lbCfg.lookupService,
+		grpcTarget:      b.bopts.Target.String(),
+		metricsRecorder: b.cc.MetricsRecorder(),
+	}
+	picker.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-picker %p] ", picker))
+	state := balancer.State{
+		ConnectivityState: aggregatedState,
+		Picker:            picker,
+	}
+
+	if !b.inhibitPickerUpdates {
+		b.logger.Infof("New balancer.State: %+v", state)
+		b.cc.UpdateState(state)
+	} else {
+		b.logger.Infof("Delaying picker update: %+v", state)
+	}
+
+	if b.lastPicker != nil {
+		if b.defaultPolicy != nil {
+			b.defaultPolicy.releaseRef()
+		}
+	}
+	b.lastPicker = picker
+}
+
+func (b *rlsBalancer) sendNewPicker() {
+	b.stateMu.Lock()
+	defer b.stateMu.Unlock()
+	if b.closed.HasFired() {
+		return
+	}
+	b.sendNewPickerLocked()
+}
+
+// The aggregated connectivity state reported is determined as follows:
+//   - If there is at least one child policy in state READY, the connectivity
+//     state is READY.
+//   - Otherwise, if there is at least one child policy in state CONNECTING, the
+//     connectivity state is CONNECTING.
+//   - Otherwise, if there is at least one child policy in state IDLE, the
+//     connectivity state is IDLE.
+//   - Otherwise, all child policies are in TRANSIENT_FAILURE, and the
+//     connectivity state is TRANSIENT_FAILURE.
+//
+// If the RLS policy has no child policies and no configured default target,
+// then we will report connectivity state IDLE.
+//
+// Caller must hold lb.stateMu.
+func (b *rlsBalancer) aggregatedConnectivityState() connectivity.State {
+	if len(b.childPolicies) == 0 && b.lbCfg.defaultTarget == "" {
+		return connectivity.Idle
+	}
+
+	var readyN, connectingN, idleN int
+	for _, cpw := range b.childPolicies {
+		state := (*balancer.State)(atomic.LoadPointer(&cpw.state))
+		switch state.ConnectivityState {
+		case connectivity.Ready:
+			readyN++
+		case connectivity.Connecting:
+			connectingN++
+		case connectivity.Idle:
+			idleN++
+		}
+	}
+
+	switch {
+	case readyN > 0:
+		return connectivity.Ready
+	case connectingN > 0:
+		return connectivity.Connecting
+	case idleN > 0:
+		return connectivity.Idle
+	default:
+		return connectivity.TransientFailure
+	}
+}
+
+// UpdateState is a implementation of the balancergroup.BalancerStateAggregator
+// interface. The actual state aggregation functionality is handled
+// asynchronously. This method only pushes the state update on to channel read
+// and dispatched by the run() goroutine.
+func (b *rlsBalancer) UpdateState(id string, state balancer.State) {
+	b.updateCh.Put(childPolicyIDAndState{id: id, state: state})
+}
+
+// handleChildPolicyStateUpdate provides the state aggregator functionality for
+// the BalancerGroup.
+//
+// This method is invoked by the BalancerGroup whenever a child policy sends a
+// state update. We cache the child policy's connectivity state and picker for
+// two reasons:
+//   - to suppress connectivity state transitions from TRANSIENT_FAILURE to states
+//     other than READY
+//   - to delegate picks to child policies
+func (b *rlsBalancer) handleChildPolicyStateUpdate(id string, newState balancer.State) {
+	b.stateMu.Lock()
+	defer b.stateMu.Unlock()
+
+	cpw := b.childPolicies[id]
+	if cpw == nil {
+		// All child policies start with an entry in the map. If ID is not in
+		// map, it's either been removed, or never existed.
+		b.logger.Warningf("Received state update %+v for missing child policy %q", newState, id)
+		return
+	}
+
+	oldState := (*balancer.State)(atomic.LoadPointer(&cpw.state))
+	if oldState.ConnectivityState == connectivity.TransientFailure && newState.ConnectivityState == connectivity.Connecting {
+		// Ignore state transitions from TRANSIENT_FAILURE to CONNECTING, and thus
+		// fail pending RPCs instead of queuing them indefinitely when all
+		// subChannels are failing, even if the subChannels are bouncing back and
+		// forth between CONNECTING and TRANSIENT_FAILURE.
+		return
+	}
+	atomic.StorePointer(&cpw.state, unsafe.Pointer(&newState))
+	b.logger.Infof("Child policy %q has new state %+v", id, newState)
+	b.sendNewPickerLocked()
+}
+
+// acquireChildPolicyReferences attempts to acquire references to
+// childPolicyWrappers corresponding to the passed in targets. If there is no
+// childPolicyWrapper corresponding to one of the targets, a new one is created
+// and added to the BalancerGroup.
+func (b *rlsBalancer) acquireChildPolicyReferences(targets []string) []*childPolicyWrapper {
+	b.stateMu.Lock()
+	var newChildPolicies []*childPolicyWrapper
+	for _, target := range targets {
+		// If the target exists in the LB policy's childPolicies map. a new
+		// reference is taken here and added to the new list.
+		if cpw := b.childPolicies[target]; cpw != nil {
+			cpw.acquireRef()
+			newChildPolicies = append(newChildPolicies, cpw)
+			continue
+		}
+
+		// If the target does not exist in the child policy map, then a new
+		// child policy wrapper is created and added to the new list.
+		cpw := newChildPolicyWrapper(target)
+		b.childPolicies[target] = cpw
+		b.bg.Add(target, b.childPolicyBuilder)
+		b.logger.Infof("Child policy %q added to BalancerGroup", target)
+		newChildPolicies = append(newChildPolicies, cpw)
+		if err := b.buildAndPushChildPolicyConfigs(target, b.lbCfg, &balancer.ClientConnState{
+			ResolverState: b.resolverState,
+		}); err != nil {
+			cpw.lamify(err)
+		}
+	}
+	b.stateMu.Unlock()
+	return newChildPolicies
+}
+
+// releaseChildPolicyReferences releases references to childPolicyWrappers
+// corresponding to the passed in targets. If the release reference was the last
+// one, the child policy is removed from the BalancerGroup.
+func (b *rlsBalancer) releaseChildPolicyReferences(targets []string) {
+	b.stateMu.Lock()
+	for _, target := range targets {
+		if cpw := b.childPolicies[target]; cpw.releaseRef() {
+			delete(b.childPolicies, cpw.target)
+			b.bg.Remove(cpw.target)
+		}
+	}
+	b.stateMu.Unlock()
+}

balancer/rls/cache.go

@@ -0,0 +1,383 @@
+/*
+ *
+ * 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 rls
+
+import (
+	"container/list"
+	"time"
+
+	"github.com/google/uuid"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/internal/backoff"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
+)
+
+// cacheKey represents the key used to uniquely identify an entry in the data
+// cache and in the pending requests map.
+type cacheKey struct {
+	// path is the full path of the incoming RPC request.
+	path string
+	// keys is a stringified version of the RLS request key map built using the
+	// RLS keyBuilder. Since maps are not a type which is comparable in Go, it
+	// cannot be part of the key for another map (entries in the data cache and
+	// pending requests map are stored in maps).
+	keys string
+}
+
+// cacheEntry wraps all the data to be stored in a data cache entry.
+type cacheEntry struct {
+	// childPolicyWrappers contains the list of child policy wrappers
+	// corresponding to the targets returned by the RLS server for this entry.
+	childPolicyWrappers []*childPolicyWrapper
+	// headerData is received in the RLS response and is to be sent in the
+	// X-Google-RLS-Data header for matching RPCs.
+	headerData string
+	// expiryTime is the absolute time at which this cache entry stops
+	// being valid. When an RLS request succeeds, this is set to the current
+	// time plus the max_age field from the LB policy config.
+	expiryTime time.Time
+	// staleTime is the absolute time after which this cache entry will be
+	// proactively refreshed if an incoming RPC matches this entry. When an RLS
+	// request succeeds, this is set to the current time plus the stale_age from
+	// the LB policy config.
+	staleTime time.Time
+	// earliestEvictTime is the absolute time before which this entry should not
+	// be evicted from the cache. When a cache entry is created, this is set to
+	// the current time plus a default value of 5 seconds. This is required to
+	// make sure that a new entry added to the cache is not evicted before the
+	// RLS response arrives (usually when the cache is too small).
+	earliestEvictTime time.Time
+
+	// status stores the RPC status of the previous RLS request for this
+	// entry. Picks for entries with a non-nil value for this field are failed
+	// with the error stored here.
+	status error
+	// backoffState contains all backoff related state. When an RLS request
+	// succeeds, backoffState is reset. This state moves between the data cache
+	// and the pending requests map.
+	backoffState *backoffState
+	// backoffTime is the absolute time at which the backoff period for this
+	// entry ends. When an RLS request fails, this is set to the current time
+	// plus the backoff value returned by the backoffState. The backoff timer is
+	// also setup with this value. No new RLS requests are sent out for this
+	// entry until the backoff period ends.
+	//
+	// Set to zero time instant upon a successful RLS response.
+	backoffTime time.Time
+	// backoffExpiryTime is the absolute time at which an entry which has gone
+	// through backoff stops being valid.  When an RLS request fails, this is
+	// set to the current time plus twice the backoff time. The cache expiry
+	// timer will only delete entries for which both expiryTime and
+	// backoffExpiryTime are in the past.
+	//
+	// Set to zero time instant upon a successful RLS response.
+	backoffExpiryTime time.Time
+
+	// size stores the size of this cache entry. Used to enforce the cache size
+	// specified in the LB policy configuration.
+	size int64
+}
+
+// backoffState wraps all backoff related state associated with a cache entry.
+type backoffState struct {
+	// retries keeps track of the number of RLS failures, to be able to
+	// determine the amount of time to backoff before the next attempt.
+	retries int
+	// bs is the exponential backoff implementation which returns the amount of
+	// time to backoff, given the number of retries.
+	bs backoff.Strategy
+	// timer fires when the backoff period ends and incoming requests after this
+	// will trigger a new RLS request.
+	timer *time.Timer
+}
+
+// lru is a cache implementation with a least recently used eviction policy.
+// Internally it uses a doubly linked list, with the least recently used element
+// at the front of the list and the most recently used element at the back of
+// the list. The value stored in this cache will be of type `cacheKey`.
+//
+// It is not safe for concurrent access.
+type lru struct {
+	ll *list.List
+
+	// A map from the value stored in the lru to its underlying list element is
+	// maintained to have a clean API. Without this, a subset of the lru's API
+	// would accept/return cacheKey while another subset would accept/return
+	// list elements.
+	m map[cacheKey]*list.Element
+}
+
+// newLRU creates a new cache with a least recently used eviction policy.
+func newLRU() *lru {
+	return &lru{
+		ll: list.New(),
+		m:  make(map[cacheKey]*list.Element),
+	}
+}
+
+func (l *lru) addEntry(key cacheKey) {
+	e := l.ll.PushBack(key)
+	l.m[key] = e
+}
+
+func (l *lru) makeRecent(key cacheKey) {
+	e := l.m[key]
+	l.ll.MoveToBack(e)
+}
+
+func (l *lru) removeEntry(key cacheKey) {
+	e := l.m[key]
+	l.ll.Remove(e)
+	delete(l.m, key)
+}
+
+func (l *lru) getLeastRecentlyUsed() cacheKey {
+	e := l.ll.Front()
+	if e == nil {
+		return cacheKey{}
+	}
+	return e.Value.(cacheKey)
+}
+
+// dataCache contains a cache of RLS data used by the LB policy to make routing
+// decisions.
+//
+// The dataCache will be keyed by the request's path and keys, represented by
+// the `cacheKey` type. It will maintain the cache keys in an `lru` and the
+// cache data, represented by the `cacheEntry` type, in a native map.
+//
+// It is not safe for concurrent access.
+type dataCache struct {
+	maxSize         int64 // Maximum allowed size.
+	currentSize     int64 // Current size.
+	keys            *lru  // Cache keys maintained in lru order.
+	entries         map[cacheKey]*cacheEntry
+	logger          *internalgrpclog.PrefixLogger
+	shutdown        *grpcsync.Event
+	rlsServerTarget string
+
+	// Read only after initialization.
+	grpcTarget      string
+	uuid            string
+	metricsRecorder estats.MetricsRecorder
+}
+
+func newDataCache(size int64, logger *internalgrpclog.PrefixLogger, metricsRecorder estats.MetricsRecorder, grpcTarget string) *dataCache {
+	return &dataCache{
+		maxSize:         size,
+		keys:            newLRU(),
+		entries:         make(map[cacheKey]*cacheEntry),
+		logger:          logger,
+		shutdown:        grpcsync.NewEvent(),
+		grpcTarget:      grpcTarget,
+		uuid:            uuid.New().String(),
+		metricsRecorder: metricsRecorder,
+	}
+}
+
+// updateRLSServerTarget updates the RLS Server Target the RLS Balancer is
+// configured with.
+func (dc *dataCache) updateRLSServerTarget(rlsServerTarget string) {
+	dc.rlsServerTarget = rlsServerTarget
+}
+
+// resize changes the maximum allowed size of the data cache.
+//
+// The return value indicates if an entry with a valid backoff timer was
+// evicted. This is important to the RLS LB policy which would send a new picker
+// on the channel to re-process any RPCs queued as a result of this backoff
+// timer.
+func (dc *dataCache) resize(size int64) (backoffCancelled bool) {
+	if dc.shutdown.HasFired() {
+		return false
+	}
+
+	backoffCancelled = false
+	for dc.currentSize > size {
+		key := dc.keys.getLeastRecentlyUsed()
+		entry, ok := dc.entries[key]
+		if !ok {
+			// This should never happen.
+			dc.logger.Errorf("cacheKey %+v not found in the cache while attempting to resize it", key)
+			break
+		}
+
+		// When we encounter a cache entry whose minimum expiration time is in
+		// the future, we abort the LRU pass, which may temporarily leave the
+		// cache being too large. This is necessary to ensure that in cases
+		// where the cache is too small, when we receive an RLS Response, we
+		// keep the resulting cache entry around long enough for the pending
+		// incoming requests to be re-processed through the new Picker. If we
+		// didn't do this, then we'd risk throwing away each RLS response as we
+		// receive it, in which case we would fail to actually route any of our
+		// incoming requests.
+		if entry.earliestEvictTime.After(time.Now()) {
+			dc.logger.Warningf("cachekey %+v is too recent to be evicted. Stopping cache resizing for now", key)
+			break
+		}
+
+		// Stop the backoff timer before evicting the entry.
+		if entry.backoffState != nil && entry.backoffState.timer != nil {
+			if entry.backoffState.timer.Stop() {
+				entry.backoffState.timer = nil
+				backoffCancelled = true
+			}
+		}
+		dc.deleteAndCleanup(key, entry)
+	}
+	dc.maxSize = size
+	return backoffCancelled
+}
+
+// evictExpiredEntries sweeps through the cache and deletes expired entries. An
+// expired entry is one for which both the `expiryTime` and `backoffExpiryTime`
+// fields are in the past.
+//
+// The return value indicates if any expired entries were evicted.
+//
+// The LB policy invokes this method periodically to purge expired entries.
+func (dc *dataCache) evictExpiredEntries() bool {
+	if dc.shutdown.HasFired() {
+		return false
+	}
+
+	evicted := false
+	for key, entry := range dc.entries {
+		// Only evict entries for which both the data expiration time and
+		// backoff expiration time fields are in the past.
+		now := time.Now()
+		if entry.expiryTime.After(now) || entry.backoffExpiryTime.After(now) {
+			continue
+		}
+		dc.deleteAndCleanup(key, entry)
+		evicted = true
+	}
+	return evicted
+}
+
+// resetBackoffState sweeps through the cache and for entries with a backoff
+// state, the backoff timer is cancelled and the backoff state is reset. The
+// return value indicates if any entries were mutated in this fashion.
+//
+// The LB policy invokes this method when the control channel moves from READY
+// to TRANSIENT_FAILURE back to READY. See `monitorConnectivityState` method on
+// the `controlChannel` type for more details.
+func (dc *dataCache) resetBackoffState(newBackoffState *backoffState) bool {
+	if dc.shutdown.HasFired() {
+		return false
+	}
+
+	backoffReset := false
+	for _, entry := range dc.entries {
+		if entry.backoffState == nil {
+			continue
+		}
+		if entry.backoffState.timer != nil {
+			entry.backoffState.timer.Stop()
+			entry.backoffState.timer = nil
+		}
+		entry.backoffState = &backoffState{bs: newBackoffState.bs}
+		entry.backoffTime = time.Time{}
+		entry.backoffExpiryTime = time.Time{}
+		backoffReset = true
+	}
+	return backoffReset
+}
+
+// addEntry adds a cache entry for the given key.
+//
+// Return value backoffCancelled indicates if a cache entry with a valid backoff
+// timer was evicted to make space for the current entry. This is important to
+// the RLS LB policy which would send a new picker on the channel to re-process
+// any RPCs queued as a result of this backoff timer.
+//
+// Return value ok indicates if entry was successfully added to the cache.
+func (dc *dataCache) addEntry(key cacheKey, entry *cacheEntry) (backoffCancelled bool, ok bool) {
+	if dc.shutdown.HasFired() {
+		return false, false
+	}
+
+	// Handle the extremely unlikely case that a single entry is bigger than the
+	// size of the cache.
+	if entry.size > dc.maxSize {
+		return false, false
+	}
+	dc.entries[key] = entry
+	dc.currentSize += entry.size
+	dc.keys.addEntry(key)
+	// If the new entry makes the cache go over its configured size, remove some
+	// old entries.
+	if dc.currentSize > dc.maxSize {
+		backoffCancelled = dc.resize(dc.maxSize)
+	}
+	cacheSizeMetric.Record(dc.metricsRecorder, dc.currentSize, dc.grpcTarget, dc.rlsServerTarget, dc.uuid)
+	cacheEntriesMetric.Record(dc.metricsRecorder, int64(len(dc.entries)), dc.grpcTarget, dc.rlsServerTarget, dc.uuid)
+	return backoffCancelled, true
+}
+
+// updateEntrySize updates the size of a cache entry and the current size of the
+// data cache. An entry's size can change upon receipt of an RLS response.
+func (dc *dataCache) updateEntrySize(entry *cacheEntry, newSize int64) {
+	dc.currentSize -= entry.size
+	entry.size = newSize
+	dc.currentSize += entry.size
+	cacheSizeMetric.Record(dc.metricsRecorder, dc.currentSize, dc.grpcTarget, dc.rlsServerTarget, dc.uuid)
+}
+
+func (dc *dataCache) getEntry(key cacheKey) *cacheEntry {
+	if dc.shutdown.HasFired() {
+		return nil
+	}
+
+	entry, ok := dc.entries[key]
+	if !ok {
+		return nil
+	}
+	dc.keys.makeRecent(key)
+	return entry
+}
+
+func (dc *dataCache) removeEntryForTesting(key cacheKey) {
+	entry, ok := dc.entries[key]
+	if !ok {
+		return
+	}
+	dc.deleteAndCleanup(key, entry)
+}
+
+// deleteAndCleanup performs actions required at the time of deleting an entry
+// from the data cache.
+// - the entry is removed from the map of entries
+// - current size of the data cache is update
+// - the key is removed from the LRU
+func (dc *dataCache) deleteAndCleanup(key cacheKey, entry *cacheEntry) {
+	delete(dc.entries, key)
+	dc.currentSize -= entry.size
+	dc.keys.removeEntry(key)
+	cacheSizeMetric.Record(dc.metricsRecorder, dc.currentSize, dc.grpcTarget, dc.rlsServerTarget, dc.uuid)
+	cacheEntriesMetric.Record(dc.metricsRecorder, int64(len(dc.entries)), dc.grpcTarget, dc.rlsServerTarget, dc.uuid)
+}
+
+func (dc *dataCache) stop() {
+	for key, entry := range dc.entries {
+		dc.deleteAndCleanup(key, entry)
+	}
+	dc.shutdown.Fire()
+}

balancer/rls/cache_test.go

@@ -0,0 +1,302 @@
+/*
+ *
+ * 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 rls
+
+import (
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"google.golang.org/grpc/internal/backoff"
+	"google.golang.org/grpc/internal/testutils/stats"
+)
+
+var (
+	cacheKeys = []cacheKey{
+		{path: "0", keys: "a"},
+		{path: "1", keys: "b"},
+		{path: "2", keys: "c"},
+		{path: "3", keys: "d"},
+		{path: "4", keys: "e"},
+	}
+
+	longDuration  = 10 * time.Minute
+	shortDuration = 1 * time.Millisecond
+	cacheEntries  []*cacheEntry
+)
+
+func initCacheEntries() {
+	// All entries have a dummy size of 1 to simplify resize operations.
+	cacheEntries = []*cacheEntry{
+		{
+			// Entry is valid and minimum expiry time has not expired.
+			expiryTime:        time.Now().Add(longDuration),
+			earliestEvictTime: time.Now().Add(longDuration),
+			size:              1,
+		},
+		{
+			// Entry is valid and is in backoff.
+			expiryTime:   time.Now().Add(longDuration),
+			backoffTime:  time.Now().Add(longDuration),
+			backoffState: &backoffState{timer: time.NewTimer(longDuration)},
+			size:         1,
+		},
+		{
+			// Entry is valid, and not in backoff.
+			expiryTime: time.Now().Add(longDuration),
+			size:       1,
+		},
+		{
+			// Entry is invalid.
+			expiryTime: time.Time{}.Add(shortDuration),
+			size:       1,
+		},
+		{
+			// Entry is invalid valid and backoff has expired.
+			expiryTime:        time.Time{}.Add(shortDuration),
+			backoffExpiryTime: time.Time{}.Add(shortDuration),
+			size:              1,
+		},
+	}
+}
+
+func (s) TestLRU_BasicOperations(t *testing.T) {
+	initCacheEntries()
+	// Create an LRU and add some entries to it.
+	lru := newLRU()
+	for _, k := range cacheKeys {
+		lru.addEntry(k)
+	}
+
+	// Get the least recent entry. This should be the first entry we added.
+	if got, want := lru.getLeastRecentlyUsed(), cacheKeys[0]; got != want {
+		t.Fatalf("lru.getLeastRecentlyUsed() = %v, want %v", got, want)
+	}
+
+	// Iterate through the slice of keys we added earlier, making them the most
+	// recent entry, one at a time. The least recent entry at that point should
+	// be the next entry from our slice of keys.
+	for i, k := range cacheKeys {
+		lru.makeRecent(k)
+
+		lruIndex := (i + 1) % len(cacheKeys)
+		if got, want := lru.getLeastRecentlyUsed(), cacheKeys[lruIndex]; got != want {
+			t.Fatalf("lru.getLeastRecentlyUsed() = %v, want %v", got, want)
+		}
+	}
+
+	// Iterate through the slice of keys we added earlier, removing them one at
+	// a time The least recent entry at that point should be the next entry from
+	// our slice of keys, except for the last one because the lru will be empty.
+	for i, k := range cacheKeys {
+		lru.removeEntry(k)
+
+		var want cacheKey
+		if i < len(cacheKeys)-1 {
+			want = cacheKeys[i+1]
+		}
+		if got := lru.getLeastRecentlyUsed(); got != want {
+			t.Fatalf("lru.getLeastRecentlyUsed() = %v, want %v", got, want)
+		}
+	}
+}
+
+func (s) TestDataCache_BasicOperations(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(5, nil, &stats.NoopMetricsRecorder{}, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+	for i, k := range cacheKeys {
+		entry := dc.getEntry(k)
+		if !cmp.Equal(entry, cacheEntries[i], cmp.AllowUnexported(cacheEntry{}, backoffState{}), cmpopts.IgnoreUnexported(time.Timer{})) {
+			t.Fatalf("Data cache lookup for key %v returned entry %v, want %v", k, entry, cacheEntries[i])
+		}
+	}
+}
+
+func (s) TestDataCache_AddForcesResize(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(1, nil, &stats.NoopMetricsRecorder{}, "")
+
+	// The first entry in cacheEntries has a minimum expiry time in the future.
+	// This entry would stop the resize operation since we do not evict entries
+	// whose minimum expiration time is in the future. So, we do not use that
+	// entry in this test. The entry being added has a running backoff timer.
+	evicted, ok := dc.addEntry(cacheKeys[1], cacheEntries[1])
+	if evicted || !ok {
+		t.Fatalf("dataCache.addEntry() returned (%v, %v) want (false, true)", evicted, ok)
+	}
+
+	// Add another entry leading to the eviction of the above entry which has a
+	// running backoff timer. The first return value is expected to be true.
+	backoffCancelled, ok := dc.addEntry(cacheKeys[2], cacheEntries[2])
+	if !backoffCancelled || !ok {
+		t.Fatalf("dataCache.addEntry() returned (%v, %v) want (true, true)", backoffCancelled, ok)
+	}
+
+	// Add another entry leading to the eviction of the above entry which does not
+	// have a running backoff timer. This should evict the above entry, but the
+	// first return value is expected to be false.
+	backoffCancelled, ok = dc.addEntry(cacheKeys[3], cacheEntries[3])
+	if backoffCancelled || !ok {
+		t.Fatalf("dataCache.addEntry() returned (%v, %v) want (false, true)", backoffCancelled, ok)
+	}
+}
+
+func (s) TestDataCache_Resize(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(5, nil, &stats.NoopMetricsRecorder{}, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+
+	// The first cache entry (with a key of cacheKeys[0]) that we added has an
+	// earliestEvictTime in the future. As part of the resize operation, we
+	// traverse the cache in least recently used order, and this will be first
+	// entry that we will encounter. And since the earliestEvictTime is in the
+	// future, the resize operation will stop, leaving the cache bigger than
+	// what was asked for.
+	if dc.resize(1) {
+		t.Fatalf("dataCache.resize() returned true, want false")
+	}
+	if dc.currentSize != 5 {
+		t.Fatalf("dataCache.size is %d, want 5", dc.currentSize)
+	}
+
+	// Remove the entry with earliestEvictTime in the future and retry the
+	// resize operation.
+	dc.removeEntryForTesting(cacheKeys[0])
+	if !dc.resize(1) {
+		t.Fatalf("dataCache.resize() returned false, want true")
+	}
+	if dc.currentSize != 1 {
+		t.Fatalf("dataCache.size is %d, want 1", dc.currentSize)
+	}
+}
+
+func (s) TestDataCache_EvictExpiredEntries(t *testing.T) {
+	initCacheEntries()
+	dc := newDataCache(5, nil, &stats.NoopMetricsRecorder{}, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+
+	// The last two entries in the cacheEntries list have expired, and will be
+	// evicted. The first three should still remain in the cache.
+	if !dc.evictExpiredEntries() {
+		t.Fatal("dataCache.evictExpiredEntries() returned false, want true")
+	}
+	if dc.currentSize != 3 {
+		t.Fatalf("dataCache.size is %d, want 3", dc.currentSize)
+	}
+	for i := 0; i < 3; i++ {
+		entry := dc.getEntry(cacheKeys[i])
+		if !cmp.Equal(entry, cacheEntries[i], cmp.AllowUnexported(cacheEntry{}, backoffState{}), cmpopts.IgnoreUnexported(time.Timer{})) {
+			t.Fatalf("Data cache lookup for key %v returned entry %v, want %v", cacheKeys[i], entry, cacheEntries[i])
+		}
+	}
+}
+
+func (s) TestDataCache_ResetBackoffState(t *testing.T) {
+	type fakeBackoff struct {
+		backoff.Strategy
+	}
+
+	initCacheEntries()
+	dc := newDataCache(5, nil, &stats.NoopMetricsRecorder{}, "")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntries[i])
+	}
+
+	newBackoffState := &backoffState{bs: &fakeBackoff{}}
+	if updatePicker := dc.resetBackoffState(newBackoffState); !updatePicker {
+		t.Fatal("dataCache.resetBackoffState() returned updatePicker is false, want true")
+	}
+
+	// Make sure that the entry with no backoff state was not touched.
+	if entry := dc.getEntry(cacheKeys[0]); cmp.Equal(entry.backoffState, newBackoffState, cmp.AllowUnexported(backoffState{})) {
+		t.Fatal("dataCache.resetBackoffState() touched entries without a valid backoffState")
+	}
+
+	// Make sure that the entry with a valid backoff state was reset.
+	entry := dc.getEntry(cacheKeys[1])
+	if diff := cmp.Diff(entry.backoffState, newBackoffState, cmp.AllowUnexported(backoffState{})); diff != "" {
+		t.Fatalf("unexpected diff in backoffState for cache entry after dataCache.resetBackoffState(): %s", diff)
+	}
+}
+
+func (s) TestDataCache_Metrics(t *testing.T) {
+	cacheEntriesMetricsTests := []*cacheEntry{
+		{size: 1},
+		{size: 2},
+		{size: 3},
+		{size: 4},
+		{size: 5},
+	}
+	tmr := stats.NewTestMetricsRecorder()
+	dc := newDataCache(50, nil, tmr, "")
+
+	dc.updateRLSServerTarget("rls-server-target")
+	for i, k := range cacheKeys {
+		dc.addEntry(k, cacheEntriesMetricsTests[i])
+	}
+
+	const cacheEntriesKey = "grpc.lb.rls.cache_entries"
+	const cacheSizeKey = "grpc.lb.rls.cache_size"
+	// 5 total entries which add up to 15 size, so should record that.
+	if got, _ := tmr.Metric(cacheEntriesKey); got != 5 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheEntriesKey, got, 5)
+	}
+	if got, _ := tmr.Metric(cacheSizeKey); got != 15 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheSizeKey, got, 15)
+	}
+
+	// Resize down the cache to 2 entries (deterministic as based of LRU).
+	dc.resize(9)
+	if got, _ := tmr.Metric(cacheEntriesKey); got != 2 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheEntriesKey, got, 2)
+	}
+	if got, _ := tmr.Metric(cacheSizeKey); got != 9 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheSizeKey, got, 9)
+	}
+
+	// Update an entry to have size 6. This should reflect in the size metrics,
+	// which will increase by 1 to 11, while the number of cache entries should
+	// stay same. This write is deterministic and writes to the last one.
+	dc.updateEntrySize(cacheEntriesMetricsTests[4], 6)
+
+	if got, _ := tmr.Metric(cacheEntriesKey); got != 2 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheEntriesKey, got, 2)
+	}
+	if got, _ := tmr.Metric(cacheSizeKey); got != 10 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheSizeKey, got, 10)
+	}
+
+	// Delete this scaled up cache key. This should scale down the cache to 1
+	// entries, and remove 6 size so cache size should be 4.
+	dc.deleteAndCleanup(cacheKeys[4], cacheEntriesMetricsTests[4])
+	if got, _ := tmr.Metric(cacheEntriesKey); got != 1 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheEntriesKey, got, 1)
+	}
+	if got, _ := tmr.Metric(cacheSizeKey); got != 4 {
+		t.Fatalf("Unexpected data for metric %v, got: %v, want: %v", cacheSizeKey, got, 4)
+	}
+}

balancer/rls/child_policy.go

@@ -0,0 +1,109 @@
+/*
+ *
+ * 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 rls
+
+import (
+	"fmt"
+	"sync/atomic"
+	"unsafe"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/connectivity"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+)
+
+// childPolicyWrapper is a reference counted wrapper around a child policy.
+//
+// The LB policy maintains a map of these wrappers keyed by the target returned
+// by RLS. When a target is seen for the first time, a child policy wrapper is
+// created for it and the wrapper is added to the child policy map. Each entry
+// in the data cache holds references to the corresponding child policy
+// wrappers. The LB policy also holds a reference to the child policy wrapper
+// for the default target specified in the LB Policy Configuration
+//
+// When a cache entry is evicted, it releases references to the child policy
+// wrappers that it contains. When all references have been released, the
+// wrapper is removed from the child policy map and is destroyed.
+//
+// The child policy wrapper also caches the connectivity state and most recent
+// picker from the child policy. Once the child policy wrapper reports
+// TRANSIENT_FAILURE, it will continue reporting that state until it goes READY;
+// transitions from TRANSIENT_FAILURE to CONNECTING are ignored.
+//
+// Whenever a child policy wrapper changes its connectivity state, the LB policy
+// returns a new picker to the channel, since the channel may need to re-process
+// the picks for queued RPCs.
+//
+// It is not safe for concurrent access.
+type childPolicyWrapper struct {
+	logger *internalgrpclog.PrefixLogger
+	target string // RLS target corresponding to this child policy.
+	refCnt int    // Reference count.
+
+	// Balancer state reported by the child policy. The RLS LB policy maintains
+	// these child policies in a BalancerGroup. The state reported by the child
+	// policy is pushed to the state aggregator (which is also implemented by the
+	// RLS LB policy) and cached here. See handleChildPolicyStateUpdate() for
+	// details on how the state aggregation is performed.
+	//
+	// While this field is written to by the LB policy, it is read by the picker
+	// at Pick time. Making this an atomic to enable the picker to read this value
+	// without a mutex.
+	state unsafe.Pointer // *balancer.State
+}
+
+// newChildPolicyWrapper creates a child policy wrapper for the given target,
+// and is initialized with one reference and starts off in CONNECTING state.
+func newChildPolicyWrapper(target string) *childPolicyWrapper {
+	c := &childPolicyWrapper{
+		target: target,
+		refCnt: 1,
+		state: unsafe.Pointer(&balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            base.NewErrPicker(balancer.ErrNoSubConnAvailable),
+		}),
+	}
+	c.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-child-policy-wrapper %s %p] ", c.target, c))
+	c.logger.Infof("Created")
+	return c
+}
+
+// acquireRef increments the reference count on the child policy wrapper.
+func (c *childPolicyWrapper) acquireRef() {
+	c.refCnt++
+}
+
+// releaseRef decrements the reference count on the child policy wrapper. The
+// return value indicates whether the released reference was the last one.
+func (c *childPolicyWrapper) releaseRef() bool {
+	c.refCnt--
+	return c.refCnt == 0
+}
+
+// lamify causes the child policy wrapper to return a picker which will always
+// fail requests. This is used when the wrapper runs into errors when trying to
+// build and parse the child policy configuration.
+func (c *childPolicyWrapper) lamify(err error) {
+	c.logger.Warningf("Entering lame mode: %v", err)
+	atomic.StorePointer(&c.state, unsafe.Pointer(&balancer.State{
+		ConnectivityState: connectivity.TransientFailure,
+		Picker:            base.NewErrPicker(err),
+	}))
+}

balancer/rls/config.go

@@ -0,0 +1,330 @@
+/*
+ *
+ * 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 rls
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/url"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/rls/internal/keys"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/pretty"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+const (
+	// Default max_age if not specified (or greater than this value) in the
+	// service config.
+	maxMaxAge = 5 * time.Minute
+	// Upper limit for cache_size since we don't fully trust the service config.
+	maxCacheSize = 5 * 1024 * 1024 * 8 // 5MB in bytes
+	// Default lookup_service_timeout if not specified in the service config.
+	defaultLookupServiceTimeout = 10 * time.Second
+	// Default value for targetNameField in the child policy config during
+	// service config validation.
+	dummyChildPolicyTarget = "target_name_to_be_filled_in_later"
+)
+
+// lbConfig is the internal representation of the RLS LB policy's config.
+type lbConfig struct {
+	serviceconfig.LoadBalancingConfig
+
+	cacheSizeBytes       int64 // Keep this field 64-bit aligned.
+	kbMap                keys.BuilderMap
+	lookupService        string
+	lookupServiceTimeout time.Duration
+	maxAge               time.Duration
+	staleAge             time.Duration
+	defaultTarget        string
+
+	childPolicyName             string
+	childPolicyConfig           map[string]json.RawMessage
+	childPolicyTargetField      string
+	controlChannelServiceConfig string
+}
+
+func (lbCfg *lbConfig) Equal(other *lbConfig) bool {
+	return lbCfg.kbMap.Equal(other.kbMap) &&
+		lbCfg.lookupService == other.lookupService &&
+		lbCfg.lookupServiceTimeout == other.lookupServiceTimeout &&
+		lbCfg.maxAge == other.maxAge &&
+		lbCfg.staleAge == other.staleAge &&
+		lbCfg.cacheSizeBytes == other.cacheSizeBytes &&
+		lbCfg.defaultTarget == other.defaultTarget &&
+		lbCfg.childPolicyName == other.childPolicyName &&
+		lbCfg.childPolicyTargetField == other.childPolicyTargetField &&
+		lbCfg.controlChannelServiceConfig == other.controlChannelServiceConfig &&
+		childPolicyConfigEqual(lbCfg.childPolicyConfig, other.childPolicyConfig)
+}
+
+func childPolicyConfigEqual(a, b map[string]json.RawMessage) bool {
+	if (b == nil) != (a == nil) {
+		return false
+	}
+	if len(b) != len(a) {
+		return false
+	}
+	for k, jsonA := range a {
+		jsonB, ok := b[k]
+		if !ok {
+			return false
+		}
+		if !bytes.Equal(jsonA, jsonB) {
+			return false
+		}
+	}
+	return true
+}
+
+// This struct resembles the JSON representation of the loadBalancing config
+// and makes it easier to unmarshal.
+type lbConfigJSON struct {
+	RouteLookupConfig                json.RawMessage
+	RouteLookupChannelServiceConfig  json.RawMessage
+	ChildPolicy                      []map[string]json.RawMessage
+	ChildPolicyConfigTargetFieldName string
+}
+
+// ParseConfig parses the JSON load balancer config provided into an
+// internal form or returns an error if the config is invalid.
+//
+//	 When parsing a config update, the following validations are performed:
+//	 - routeLookupConfig:
+//	   - grpc_keybuilders field:
+//	     - must have at least one entry
+//	     - must not have two entries with the same `Name`
+//	     - within each entry:
+//	       - must have at least one `Name`
+//	       - must not have a `Name` with the `service` field unset or empty
+//	       - within each `headers` entry:
+//	         - must not have `required_match` set
+//	         - must not have `key` unset or empty
+//	       - across all `headers`, `constant_keys` and `extra_keys` fields:
+//	         - must not have the same `key` specified twice
+//	         - no `key` must be the empty string
+//	   - `lookup_service` field must be set and must parse as a target URI
+//	   - if `max_age` > 5m, it should be set to 5 minutes
+//	   - if `stale_age` > `max_age`, ignore it
+//	   - if `stale_age` is set, then `max_age` must also be set
+//	   - ignore `valid_targets` field
+//	   - `cache_size_bytes` field must have a value greater than 0, and if its
+//	     value is greater than 5M, we cap it at 5M
+//
+//	- routeLookupChannelServiceConfig:
+//	  - if specified, must parse as valid service config
+//
+//	- childPolicy:
+//	  - must find a valid child policy with a valid config
+//
+//	- childPolicyConfigTargetFieldName:
+//	  - must be set and non-empty
+func (rlsBB) ParseConfig(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	if logger.V(2) {
+		logger.Infof("Received JSON service config: %v", pretty.ToJSON(c))
+	}
+
+	cfgJSON := &lbConfigJSON{}
+	if err := json.Unmarshal(c, cfgJSON); err != nil {
+		return nil, fmt.Errorf("rls: json unmarshal failed for service config %+v: %v", string(c), err)
+	}
+
+	m := protojson.UnmarshalOptions{DiscardUnknown: true}
+	rlsProto := &rlspb.RouteLookupConfig{}
+	if err := m.Unmarshal(cfgJSON.RouteLookupConfig, rlsProto); err != nil {
+		return nil, fmt.Errorf("rls: bad RouteLookupConfig proto %+v: %v", string(cfgJSON.RouteLookupConfig), err)
+	}
+	lbCfg, err := parseRLSProto(rlsProto)
+	if err != nil {
+		return nil, err
+	}
+
+	if sc := string(cfgJSON.RouteLookupChannelServiceConfig); sc != "" {
+		parsed := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(sc)
+		if parsed.Err != nil {
+			return nil, fmt.Errorf("rls: bad control channel service config %q: %v", sc, parsed.Err)
+		}
+		lbCfg.controlChannelServiceConfig = sc
+	}
+
+	if cfgJSON.ChildPolicyConfigTargetFieldName == "" {
+		return nil, fmt.Errorf("rls: childPolicyConfigTargetFieldName field is not set in service config %+v", string(c))
+	}
+	name, config, err := parseChildPolicyConfigs(cfgJSON.ChildPolicy, cfgJSON.ChildPolicyConfigTargetFieldName)
+	if err != nil {
+		return nil, err
+	}
+	lbCfg.childPolicyName = name
+	lbCfg.childPolicyConfig = config
+	lbCfg.childPolicyTargetField = cfgJSON.ChildPolicyConfigTargetFieldName
+	return lbCfg, nil
+}
+
+func parseRLSProto(rlsProto *rlspb.RouteLookupConfig) (*lbConfig, error) {
+	// Validations specified on the `grpc_keybuilders` field are performed here.
+	kbMap, err := keys.MakeBuilderMap(rlsProto)
+	if err != nil {
+		return nil, err
+	}
+
+	// `lookup_service` field must be set and must parse as a target URI.
+	lookupService := rlsProto.GetLookupService()
+	if lookupService == "" {
+		return nil, fmt.Errorf("rls: empty lookup_service in route lookup config %+v", rlsProto)
+	}
+	parsedTarget, err := url.Parse(lookupService)
+	if err != nil {
+		// url.Parse() fails if scheme is missing. Retry with default scheme.
+		parsedTarget, err = url.Parse(resolver.GetDefaultScheme() + ":///" + lookupService)
+		if err != nil {
+			return nil, fmt.Errorf("rls: invalid target URI in lookup_service %s", lookupService)
+		}
+	}
+	if parsedTarget.Scheme == "" {
+		parsedTarget.Scheme = resolver.GetDefaultScheme()
+	}
+	if resolver.Get(parsedTarget.Scheme) == nil {
+		return nil, fmt.Errorf("rls: unregistered scheme in lookup_service %s", lookupService)
+	}
+
+	lookupServiceTimeout, err := convertDuration(rlsProto.GetLookupServiceTimeout())
+	if err != nil {
+		return nil, fmt.Errorf("rls: failed to parse lookup_service_timeout in route lookup config %+v: %v", rlsProto, err)
+	}
+	if lookupServiceTimeout == 0 {
+		lookupServiceTimeout = defaultLookupServiceTimeout
+	}
+
+	// Validations performed here:
+	// - if `max_age` > 5m, it should be set to 5 minutes
+	//   only if stale age is not set
+	// - if `stale_age` > `max_age`, ignore it
+	// - if `stale_age` is set, then `max_age` must also be set
+	maxAgeSet := false
+	maxAge, err := convertDuration(rlsProto.GetMaxAge())
+	if err != nil {
+		return nil, fmt.Errorf("rls: failed to parse max_age in route lookup config %+v: %v", rlsProto, err)
+	}
+	if maxAge == 0 {
+		maxAge = maxMaxAge
+	} else {
+		maxAgeSet = true
+	}
+
+	staleAgeSet := false
+	staleAge, err := convertDuration(rlsProto.GetStaleAge())
+	if err != nil {
+		return nil, fmt.Errorf("rls: failed to parse staleAge in route lookup config %+v: %v", rlsProto, err)
+	}
+	if staleAge == 0 {
+		staleAge = maxMaxAge
+	} else {
+		staleAgeSet = true
+	}
+
+	if staleAgeSet && !maxAgeSet {
+		return nil, fmt.Errorf("rls: stale_age is set, but max_age is not in route lookup config %+v", rlsProto)
+	}
+	if staleAge > maxMaxAge {
+		staleAge = maxMaxAge
+	}
+	if !staleAgeSet && maxAge > maxMaxAge {
+		maxAge = maxMaxAge
+	}
+	if staleAge > maxAge {
+		staleAge = maxAge
+	}
+
+	// `cache_size_bytes` field must have a value greater than 0, and if its
+	// value is greater than 5M, we cap it at 5M
+	cacheSizeBytes := rlsProto.GetCacheSizeBytes()
+	if cacheSizeBytes <= 0 {
+		return nil, fmt.Errorf("rls: cache_size_bytes must be set to a non-zero value: %+v", rlsProto)
+	}
+	if cacheSizeBytes > maxCacheSize {
+		logger.Info("rls: cache_size_bytes %v is too large, setting it to: %v", cacheSizeBytes, maxCacheSize)
+		cacheSizeBytes = maxCacheSize
+	}
+	return &lbConfig{
+		kbMap:                kbMap,
+		lookupService:        lookupService,
+		lookupServiceTimeout: lookupServiceTimeout,
+		maxAge:               maxAge,
+		staleAge:             staleAge,
+		cacheSizeBytes:       cacheSizeBytes,
+		defaultTarget:        rlsProto.GetDefaultTarget(),
+	}, nil
+}
+
+// parseChildPolicyConfigs iterates through the list of child policies and picks
+// the first registered policy and validates its config.
+func parseChildPolicyConfigs(childPolicies []map[string]json.RawMessage, targetFieldName string) (string, map[string]json.RawMessage, error) {
+	for i, config := range childPolicies {
+		if len(config) != 1 {
+			return "", nil, fmt.Errorf("rls: invalid childPolicy: entry %v does not contain exactly 1 policy/config pair: %q", i, config)
+		}
+
+		var name string
+		var rawCfg json.RawMessage
+		for name, rawCfg = range config {
+		}
+		builder := balancer.Get(name)
+		if builder == nil {
+			continue
+		}
+		parser, ok := builder.(balancer.ConfigParser)
+		if !ok {
+			return "", nil, fmt.Errorf("rls: childPolicy %q with config %q does not support config parsing", name, string(rawCfg))
+		}
+
+		// To validate child policy configs we do the following:
+		// - unmarshal the raw JSON bytes of the child policy config into a map
+		// - add an entry with key set to `target_field_name` and a dummy value
+		// - marshal the map back to JSON and parse the config using the parser
+		// retrieved previously
+		var childConfig map[string]json.RawMessage
+		if err := json.Unmarshal(rawCfg, &childConfig); err != nil {
+			return "", nil, fmt.Errorf("rls: json unmarshal failed for child policy config %q: %v", string(rawCfg), err)
+		}
+		childConfig[targetFieldName], _ = json.Marshal(dummyChildPolicyTarget)
+		jsonCfg, err := json.Marshal(childConfig)
+		if err != nil {
+			return "", nil, fmt.Errorf("rls: json marshal failed for child policy config {%+v}: %v", childConfig, err)
+		}
+		if _, err := parser.ParseConfig(jsonCfg); err != nil {
+			return "", nil, fmt.Errorf("rls: childPolicy config validation failed: %v", err)
+		}
+		return name, childConfig, nil
+	}
+	return "", nil, fmt.Errorf("rls: invalid childPolicy config: no supported policies found in %+v", childPolicies)
+}
+
+func convertDuration(d *durationpb.Duration) (time.Duration, error) {
+	if d == nil {
+		return 0, nil
+	}
+	return d.AsDuration(), d.CheckValid()
+}

balancer/rls/config_test.go

@@ -25,27 +25,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/google/go-cmp/cmp"
-
-	"google.golang.org/grpc/balancer"
 	_ "google.golang.org/grpc/balancer/grpclb"               // grpclb for config parsing.
 	_ "google.golang.org/grpc/internal/resolver/passthrough" // passthrough resolver.
 )
 
-const balancerWithoutConfigParserName = "dummy_balancer"
-
-type dummyBB struct {
-	balancer.Builder
-}
-
-func (*dummyBB) Name() string {
-	return balancerWithoutConfigParserName
-}
-
-func init() {
-	balancer.Register(&dummyBB{})
-}
-
 // testEqual reports whether the lbCfgs a and b are equal. This is to be used
 // only from tests. This ignores the keyBuilderMap field because its internals
 // are not exported, and hence not possible to specify in the want section of
@@ -58,25 +41,29 @@ func testEqual(a, b *lbConfig) bool {
 		a.staleAge == b.staleAge &&
 		a.cacheSizeBytes == b.cacheSizeBytes &&
 		a.defaultTarget == b.defaultTarget &&
-		a.cpName == b.cpName &&
-		a.cpTargetField == b.cpTargetField &&
-		cmp.Equal(a.cpConfig, b.cpConfig)
+		a.controlChannelServiceConfig == b.controlChannelServiceConfig &&
+		a.childPolicyName == b.childPolicyName &&
+		a.childPolicyTargetField == b.childPolicyTargetField &&
+		childPolicyConfigEqual(a.childPolicyConfig, b.childPolicyConfig)
 }
 
-func TestParseConfig(t *testing.T) {
+// TestParseConfig verifies successful config parsing scenarios.
+func (s) TestParseConfig(t *testing.T) {
+	childPolicyTargetFieldVal, _ := json.Marshal(dummyChildPolicyTarget)
 	tests := []struct {
 		desc    string
 		input   []byte
 		wantCfg *lbConfig
 	}{
-		// This input validates a few cases:
-		// - A top-level unknown field should not fail.
-		// - An unknown field in routeLookupConfig proto should not fail.
-		// - lookupServiceTimeout is set to its default value, since it is not specified in the input.
-		// - maxAge is set to maxMaxAge since the value is too large in the input.
-		// - staleAge is ignore because it is higher than maxAge in the input.
 		{
-			desc: "with transformations",
+			// This input validates a few cases:
+			// - A top-level unknown field should not fail.
+			// - An unknown field in routeLookupConfig proto should not fail.
+			// - lookupServiceTimeout is set to its default value, since it is not specified in the input.
+			// - maxAge is clamped to maxMaxAge if staleAge is not set.
+			// - staleAge is ignored because it is higher than maxAge in the input.
+			// - cacheSizeBytes is greater than the hard upper limit of 5MB
+			desc: "with transformations 1",
 			input: []byte(`{
 				"top-level-unknown-field": "unknown-value",
 				"routeLookupConfig": {
@@ -85,10 +72,10 @@ func TestParseConfig(t *testing.T) {
 						"names": [{"service": "service", "method": "method"}],
 						"headers": [{"key": "k1", "names": ["v1"]}]
 					}],
-					"lookupService": "passthrough:///target",
+					"lookupService": ":///target",
 					"maxAge" : "500s",
 					"staleAge": "600s",
-					"cacheSizeBytes": 1000,
+					"cacheSizeBytes": 100000000,
 					"defaultTarget": "passthrough:///default"
 				},
 				"childPolicy": [
@@ -96,18 +83,84 @@ func TestParseConfig(t *testing.T) {
 					{"unknown-policy": {"unknown-field": "unknown-value"}},
 					{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}
 				],
-				"childPolicyConfigTargetFieldName": "service_name"
+				"childPolicyConfigTargetFieldName": "serviceName"
 			}`),
 			wantCfg: &lbConfig{
-				lookupService:        "passthrough:///target",
-				lookupServiceTimeout: 10 * time.Second, // This is the default value.
-				maxAge:               5 * time.Minute,  // This is max maxAge.
-				staleAge:             time.Duration(0), // StaleAge is ignore because it was higher than maxAge.
-				cacheSizeBytes:       1000,
-				defaultTarget:        "passthrough:///default",
-				cpName:               "grpclb",
-				cpTargetField:        "service_name",
-				cpConfig:             map[string]json.RawMessage{"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`)},
+				lookupService:          ":///target",
+				lookupServiceTimeout:   10 * time.Second,  // This is the default value.
+				maxAge:                 500 * time.Second, // Max age is not clamped when stale age is set.
+				staleAge:               300 * time.Second, // StaleAge is clamped because it was higher than maxMaxAge.
+				cacheSizeBytes:         maxCacheSize,
+				defaultTarget:          "passthrough:///default",
+				childPolicyName:        "grpclb",
+				childPolicyTargetField: "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
+			},
+		},
+		{
+			desc: "maxAge not clamped when staleAge is set",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": ":///target",
+					"maxAge" : "500s",
+					"staleAge": "200s",
+					"cacheSizeBytes": 100000000
+				},
+				"childPolicy": [
+					{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantCfg: &lbConfig{
+				lookupService:          ":///target",
+				lookupServiceTimeout:   10 * time.Second,  // This is the default value.
+				maxAge:                 500 * time.Second, // Max age is not clamped when stale age is set.
+				staleAge:               200 * time.Second, // This is stale age within maxMaxAge.
+				cacheSizeBytes:         maxCacheSize,
+				childPolicyName:        "grpclb",
+				childPolicyTargetField: "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
+			},
+		},
+		{
+			desc: "maxAge clamped when staleAge is not set",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": ":///target",
+					"maxAge" : "500s",
+					"cacheSizeBytes": 100000000
+				},
+				"childPolicy": [
+					{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantCfg: &lbConfig{
+				lookupService:          ":///target",
+				lookupServiceTimeout:   10 * time.Second,  // This is the default value.
+				maxAge:                 300 * time.Second, // Max age is clamped when stale age is not set.
+				staleAge:               300 * time.Second,
+				cacheSizeBytes:         maxCacheSize,
+				childPolicyName:        "grpclb",
+				childPolicyTargetField: "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
 			},
 		},
 		{
@@ -118,31 +171,36 @@ func TestParseConfig(t *testing.T) {
 						"names": [{"service": "service", "method": "method"}],
 						"headers": [{"key": "k1", "names": ["v1"]}]
 					}],
-					"lookupService": "passthrough:///target",
+					"lookupService": "target",
 					"lookupServiceTimeout" : "100s",
 					"maxAge": "60s",
 					"staleAge" : "50s",
 					"cacheSizeBytes": 1000,
 					"defaultTarget": "passthrough:///default"
 				},
+				"routeLookupChannelServiceConfig": {"loadBalancingConfig": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}]},
 				"childPolicy": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}],
-				"childPolicyConfigTargetFieldName": "service_name"
+				"childPolicyConfigTargetFieldName": "serviceName"
 			}`),
 			wantCfg: &lbConfig{
-				lookupService:        "passthrough:///target",
-				lookupServiceTimeout: 100 * time.Second,
-				maxAge:               60 * time.Second,
-				staleAge:             50 * time.Second,
-				cacheSizeBytes:       1000,
-				defaultTarget:        "passthrough:///default",
-				cpName:               "grpclb",
-				cpTargetField:        "service_name",
-				cpConfig:             map[string]json.RawMessage{"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`)},
+				lookupService:               "target",
+				lookupServiceTimeout:        100 * time.Second,
+				maxAge:                      60 * time.Second,
+				staleAge:                    50 * time.Second,
+				cacheSizeBytes:              1000,
+				defaultTarget:               "passthrough:///default",
+				controlChannelServiceConfig: `{"loadBalancingConfig": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}]}`,
+				childPolicyName:             "grpclb",
+				childPolicyTargetField:      "serviceName",
+				childPolicyConfig: map[string]json.RawMessage{
+					"childPolicy": json.RawMessage(`[{"pickfirst": {}}]`),
+					"serviceName": json.RawMessage(childPolicyTargetFieldVal),
+				},
 			},
 		},
 	}
 
-	builder := &rlsBB{}
+	builder := rlsBB{}
 	for _, test := range tests {
 		t.Run(test.desc, func(t *testing.T) {
 			lbCfg, err := builder.ParseConfig(test.input)
@@ -153,7 +211,8 @@ func TestParseConfig(t *testing.T) {
 	}
 }
 
-func TestParseConfigErrors(t *testing.T) {
+// TestParseConfigErrors verifies config parsing failure scenarios.
+func (s) TestParseConfigErrors(t *testing.T) {
 	tests := []struct {
 		desc    string
 		input   []byte
@@ -191,10 +250,10 @@ func TestParseConfigErrors(t *testing.T) {
 						}]
 					}
 				}`),
-			wantErr: "rls: empty lookup_service in service config",
+			wantErr: "rls: empty lookup_service in route lookup config",
 		},
 		{
-			desc: "invalid lookup service URI",
+			desc: "unregistered scheme in lookup service URI",
 			input: []byte(`{
 					"routeLookupConfig": {
 						"grpcKeybuilders": [{
@@ -204,7 +263,7 @@ func TestParseConfigErrors(t *testing.T) {
 						"lookupService": "badScheme:///target"
 					}
 				}`),
-			wantErr: "rls: invalid target URI in lookup_service",
+			wantErr: "rls: unregistered scheme in lookup_service",
 		},
 		{
 			desc: "invalid lookup service timeout",
@@ -218,7 +277,7 @@ func TestParseConfigErrors(t *testing.T) {
 					"lookupServiceTimeout" : "315576000001s"
 				}
 			}`),
-			wantErr: "bad Duration: time: invalid duration",
+			wantErr: "google.protobuf.Duration value out of range",
 		},
 		{
 			desc: "invalid max age",
@@ -233,7 +292,7 @@ func TestParseConfigErrors(t *testing.T) {
 					"maxAge" : "315576000001s"
 				}
 			}`),
-			wantErr: "bad Duration: time: invalid duration",
+			wantErr: "google.protobuf.Duration value out of range",
 		},
 		{
 			desc: "invalid stale age",
@@ -249,7 +308,7 @@ func TestParseConfigErrors(t *testing.T) {
 					"staleAge" : "315576000001s"
 				}
 			}`),
-			wantErr: "bad Duration: time: invalid duration",
+			wantErr: "google.protobuf.Duration value out of range",
 		},
 		{
 			desc: "invalid max age stale age combo",
@@ -264,10 +323,10 @@ func TestParseConfigErrors(t *testing.T) {
 					"staleAge" : "10s"
 				}
 			}`),
-			wantErr: "rls: stale_age is set, but max_age is not in service config",
+			wantErr: "rls: stale_age is set, but max_age is not in route lookup config",
 		},
 		{
-			desc: "invalid cache size",
+			desc: "cache_size_bytes field is not set",
 			input: []byte(`{
 				"routeLookupConfig": {
 					"grpcKeybuilders": [{
@@ -277,10 +336,56 @@ func TestParseConfigErrors(t *testing.T) {
 					"lookupService": "passthrough:///target",
 					"lookupServiceTimeout" : "10s",
 					"maxAge": "30s",
-					"staleAge" : "25s"
-				}
+					"staleAge" : "25s",
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicyConfigTargetFieldName": "serviceName"
 			}`),
-			wantErr: "rls: cache_size_bytes must be greater than 0 in service config",
+			wantErr: "rls: cache_size_bytes must be set to a non-zero value",
+		},
+		{
+			desc: "routeLookupChannelServiceConfig is not in service config format",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "target",
+					"lookupServiceTimeout" : "100s",
+					"maxAge": "60s",
+					"staleAge" : "50s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"routeLookupChannelServiceConfig": "unknown",
+				"childPolicy": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "cannot unmarshal string into Go value of type grpc.jsonSC",
+		},
+		{
+			desc: "routeLookupChannelServiceConfig contains unknown LB policy",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "target",
+					"lookupServiceTimeout" : "100s",
+					"maxAge": "60s",
+					"staleAge" : "50s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"routeLookupChannelServiceConfig": {
+					"loadBalancingConfig": [{"not_a_balancer1": {} }, {"not_a_balancer2": {}}]
+				},
+				"childPolicy": [{"grpclb": {"childPolicy": [{"pickfirst": {}}]}}],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "no supported policies found in config",
 		},
 		{
 			desc: "no child policy",
@@ -296,9 +401,10 @@ func TestParseConfigErrors(t *testing.T) {
 					"staleAge" : "25s",
 					"cacheSizeBytes": 1000,
 					"defaultTarget": "passthrough:///default"
-				}
+				},
+				"childPolicyConfigTargetFieldName": "serviceName"
 			}`),
-			wantErr: "rls: childPolicy is invalid in service config",
+			wantErr: "rls: invalid childPolicy config: no supported policies found",
 		},
 		{
 			desc: "no known child policy",
@@ -318,9 +424,35 @@ func TestParseConfigErrors(t *testing.T) {
 				"childPolicy": [
 					{"cds_experimental": {"Cluster": "my-fav-cluster"}},
 					{"unknown-policy": {"unknown-field": "unknown-value"}}
-				]
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
 			}`),
-			wantErr: "rls: childPolicy is invalid in service config",
+			wantErr: "rls: invalid childPolicy config: no supported policies found",
+		},
+		{
+			desc: "invalid child policy config - more than one entry in map",
+			input: []byte(`{
+				"routeLookupConfig": {
+					"grpcKeybuilders": [{
+						"names": [{"service": "service", "method": "method"}],
+						"headers": [{"key": "k1", "names": ["v1"]}]
+					}],
+					"lookupService": "passthrough:///target",
+					"lookupServiceTimeout" : "10s",
+					"maxAge": "30s",
+					"staleAge" : "25s",
+					"cacheSizeBytes": 1000,
+					"defaultTarget": "passthrough:///default"
+				},
+				"childPolicy": [
+					{
+						"cds_experimental": {"Cluster": "my-fav-cluster"},
+						"unknown-policy": {"unknown-field": "unknown-value"}
+					}
+				],
+				"childPolicyConfigTargetFieldName": "serviceName"
+			}`),
+			wantErr: "does not contain exactly 1 policy/config pair",
 		},
 		{
 			desc: "no childPolicyConfigTargetFieldName",
@@ -365,13 +497,13 @@ func TestParseConfigErrors(t *testing.T) {
 					{"unknown-policy": {"unknown-field": "unknown-value"}},
 					{"grpclb": {"childPolicy": "not-an-array"}}
 				],
-				"childPolicyConfigTargetFieldName": "service_name"
+				"childPolicyConfigTargetFieldName": "serviceName"
 			}`),
 			wantErr: "rls: childPolicy config validation failed",
 		},
 	}
 
-	builder := &rlsBB{}
+	builder := rlsBB{}
 	for _, test := range tests {
 		t.Run(test.desc, func(t *testing.T) {
 			lbCfg, err := builder.ParseConfig(test.input)
@@ -381,60 +513,3 @@ func TestParseConfigErrors(t *testing.T) {
 		})
 	}
 }
-
-func TestValidateChildPolicyConfig(t *testing.T) {
-	jsonCfg := json.RawMessage(`[{"round_robin" : {}}, {"pick_first" : {}}]`)
-	wantChildConfig := map[string]json.RawMessage{"childPolicy": jsonCfg}
-	cp := &loadBalancingConfig{
-		Name:   "grpclb",
-		Config: []byte(`{"childPolicy": [{"round_robin" : {}}, {"pick_first" : {}}]}`),
-	}
-	cpTargetField := "serviceName"
-
-	gotChildConfig, err := validateChildPolicyConfig(cp, cpTargetField)
-	if err != nil || !cmp.Equal(gotChildConfig, wantChildConfig) {
-		t.Errorf("validateChildPolicyConfig(%v, %v) = {%v, %v}, want {%v, nil}", cp, cpTargetField, gotChildConfig, err, wantChildConfig)
-	}
-}
-
-func TestValidateChildPolicyConfigErrors(t *testing.T) {
-	tests := []struct {
-		desc          string
-		cp            *loadBalancingConfig
-		wantErrPrefix string
-	}{
-		{
-			desc: "unknown child policy",
-			cp: &loadBalancingConfig{
-				Name:   "unknown",
-				Config: []byte(`{}`),
-			},
-			wantErrPrefix: "rls: balancer builder not found for child_policy",
-		},
-		{
-			desc: "balancer builder does not implement ConfigParser",
-			cp: &loadBalancingConfig{
-				Name:   balancerWithoutConfigParserName,
-				Config: []byte(`{}`),
-			},
-			wantErrPrefix: "rls: balancer builder for child_policy does not implement balancer.ConfigParser",
-		},
-		{
-			desc: "child policy config parsing failure",
-			cp: &loadBalancingConfig{
-				Name:   "grpclb",
-				Config: []byte(`{"childPolicy": "not-an-array"}`),
-			},
-			wantErrPrefix: "rls: childPolicy config validation failed",
-		},
-	}
-
-	for _, test := range tests {
-		t.Run(test.desc, func(t *testing.T) {
-			gotChildConfig, gotErr := validateChildPolicyConfig(test.cp, "")
-			if gotChildConfig != nil || !strings.HasPrefix(fmt.Sprint(gotErr), test.wantErrPrefix) {
-				t.Errorf("validateChildPolicyConfig(%v) = {%v, %v}, want {nil, %v}", test.cp, gotChildConfig, gotErr, test.wantErrPrefix)
-			}
-		})
-	}
-}

balancer/rls/control_channel.go

@@ -0,0 +1,247 @@
+/*
+ *
+ * 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 rls
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/rls/internal/adaptive"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/buffer"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/pretty"
+	rlsgrpc "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+)
+
+var newAdaptiveThrottler = func() adaptiveThrottler { return adaptive.New() }
+
+type adaptiveThrottler interface {
+	ShouldThrottle() bool
+	RegisterBackendResponse(throttled bool)
+}
+
+// controlChannel is a wrapper around the gRPC channel to the RLS server
+// specified in the service config.
+type controlChannel struct {
+	// rpcTimeout specifies the timeout for the RouteLookup RPC call. The LB
+	// policy receives this value in its service config.
+	rpcTimeout time.Duration
+	// backToReadyFunc is a callback to be invoked when the connectivity state
+	// changes from READY --> TRANSIENT_FAILURE --> READY.
+	backToReadyFunc func()
+	// throttler in an adaptive throttling implementation used to avoid
+	// hammering the RLS service while it is overloaded or down.
+	throttler adaptiveThrottler
+
+	cc                  *grpc.ClientConn
+	client              rlsgrpc.RouteLookupServiceClient
+	logger              *internalgrpclog.PrefixLogger
+	connectivityStateCh *buffer.Unbounded
+	unsubscribe         func()
+	monitorDoneCh       chan struct{}
+}
+
+// newControlChannel creates a controlChannel to rlsServerName and uses
+// serviceConfig, if non-empty, as the default service config for the underlying
+// gRPC channel.
+func newControlChannel(rlsServerName, serviceConfig string, rpcTimeout time.Duration, bOpts balancer.BuildOptions, backToReadyFunc func()) (*controlChannel, error) {
+	ctrlCh := &controlChannel{
+		rpcTimeout:          rpcTimeout,
+		backToReadyFunc:     backToReadyFunc,
+		throttler:           newAdaptiveThrottler(),
+		connectivityStateCh: buffer.NewUnbounded(),
+		monitorDoneCh:       make(chan struct{}),
+	}
+	ctrlCh.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[rls-control-channel %p] ", ctrlCh))
+
+	dopts, err := ctrlCh.dialOpts(bOpts, serviceConfig)
+	if err != nil {
+		return nil, err
+	}
+	ctrlCh.cc, err = grpc.NewClient(rlsServerName, dopts...)
+	if err != nil {
+		return nil, err
+	}
+	// Subscribe to connectivity state before connecting to avoid missing initial
+	// updates, which are only delivered to active subscribers.
+	ctrlCh.unsubscribe = internal.SubscribeToConnectivityStateChanges.(func(cc *grpc.ClientConn, s grpcsync.Subscriber) func())(ctrlCh.cc, ctrlCh)
+	ctrlCh.cc.Connect()
+	ctrlCh.client = rlsgrpc.NewRouteLookupServiceClient(ctrlCh.cc)
+	ctrlCh.logger.Infof("Control channel created to RLS server at: %v", rlsServerName)
+	go ctrlCh.monitorConnectivityState()
+	return ctrlCh, nil
+}
+
+func (cc *controlChannel) OnMessage(msg any) {
+	st, ok := msg.(connectivity.State)
+	if !ok {
+		panic(fmt.Sprintf("Unexpected message type %T , wanted connectectivity.State type", msg))
+	}
+	cc.connectivityStateCh.Put(st)
+}
+
+// dialOpts constructs the dial options for the control plane channel.
+func (cc *controlChannel) dialOpts(bOpts balancer.BuildOptions, serviceConfig string) ([]grpc.DialOption, error) {
+	// The control plane channel will use the same authority as the parent
+	// channel for server authorization. This ensures that the identity of the
+	// RLS server and the identity of the backends is the same, so if the RLS
+	// config is injected by an attacker, it cannot cause leakage of private
+	// information contained in headers set by the application.
+	dopts := []grpc.DialOption{grpc.WithAuthority(bOpts.Authority)}
+	if bOpts.Dialer != nil {
+		dopts = append(dopts, grpc.WithContextDialer(bOpts.Dialer))
+	}
+	// The control channel will use the channel credentials from the parent
+	// channel, including any call creds associated with the channel creds.
+	var credsOpt grpc.DialOption
+	switch {
+	case bOpts.DialCreds != nil:
+		credsOpt = grpc.WithTransportCredentials(bOpts.DialCreds.Clone())
+	case bOpts.CredsBundle != nil:
+		// The "fallback" mode in google default credentials (which is the only
+		// type of credentials we expect to be used with RLS) uses TLS/ALTS
+		// creds for transport and uses the same call creds as that on the
+		// parent bundle.
+		bundle, err := bOpts.CredsBundle.NewWithMode(internal.CredsBundleModeFallback)
+		if err != nil {
+			return nil, err
+		}
+		credsOpt = grpc.WithCredentialsBundle(bundle)
+	default:
+		cc.logger.Warningf("no credentials available, using Insecure")
+		credsOpt = grpc.WithTransportCredentials(insecure.NewCredentials())
+	}
+	dopts = append(dopts, credsOpt)
+
+	// If the RLS LB policy's configuration specified a service config for the
+	// control channel, use that and disable service config fetching via the name
+	// resolver for the control channel.
+	if serviceConfig != "" {
+		cc.logger.Infof("Disabling service config from the name resolver and instead using: %s", serviceConfig)
+		dopts = append(dopts, grpc.WithDisableServiceConfig(), grpc.WithDefaultServiceConfig(serviceConfig))
+	}
+
+	return dopts, nil
+}
+
+func (cc *controlChannel) monitorConnectivityState() {
+	cc.logger.Infof("Starting connectivity state monitoring goroutine")
+	defer close(cc.monitorDoneCh)
+
+	// Since we use two mechanisms to deal with RLS server being down:
+	//   - adaptive throttling for the channel as a whole
+	//   - exponential backoff on a per-request basis
+	// we need a way to avoid double-penalizing requests by counting failures
+	// toward both mechanisms when the RLS server is unreachable.
+	//
+	// To accomplish this, we monitor the state of the control plane channel. If
+	// the state has been TRANSIENT_FAILURE since the last time it was in state
+	// READY, and it then transitions into state READY, we push on a channel
+	// which is being read by the LB policy.
+	//
+	// The LB the policy will iterate through the cache to reset the backoff
+	// timeouts in all cache entries. Specifically, this means that it will
+	// reset the backoff state and cancel the pending backoff timer. Note that
+	// when cancelling the backoff timer, just like when the backoff timer fires
+	// normally, a new picker is returned to the channel, to force it to
+	// re-process any wait-for-ready RPCs that may still be queued if we failed
+	// them while we were in backoff. However, we should optimize this case by
+	// returning only one new picker, regardless of how many backoff timers are
+	// cancelled.
+
+	// Wait for the control channel to become READY for the first time.
+	for s, ok := <-cc.connectivityStateCh.Get(); s != connectivity.Ready; s, ok = <-cc.connectivityStateCh.Get() {
+		if !ok {
+			return
+		}
+
+		cc.connectivityStateCh.Load()
+		if s == connectivity.Shutdown {
+			return
+		}
+	}
+	cc.connectivityStateCh.Load()
+	cc.logger.Infof("Connectivity state is READY")
+
+	for {
+		s, ok := <-cc.connectivityStateCh.Get()
+		if !ok {
+			return
+		}
+		cc.connectivityStateCh.Load()
+
+		if s == connectivity.Shutdown {
+			return
+		}
+		if s == connectivity.Ready {
+			cc.logger.Infof("Control channel back to READY")
+			cc.backToReadyFunc()
+		}
+
+		cc.logger.Infof("Connectivity state is %s", s)
+	}
+}
+
+func (cc *controlChannel) close() {
+	cc.unsubscribe()
+	cc.connectivityStateCh.Close()
+	<-cc.monitorDoneCh
+	cc.cc.Close()
+	cc.logger.Infof("Shutdown")
+}
+
+type lookupCallback func(targets []string, headerData string, err error)
+
+// lookup starts a RouteLookup RPC in a separate goroutine and returns the
+// results (and error, if any) in the provided callback.
+//
+// The returned boolean indicates whether the request was throttled by the
+// client-side adaptive throttling algorithm in which case the provided callback
+// will not be invoked.
+func (cc *controlChannel) lookup(reqKeys map[string]string, reason rlspb.RouteLookupRequest_Reason, staleHeaders string, cb lookupCallback) (throttled bool) {
+	if cc.throttler.ShouldThrottle() {
+		cc.logger.Infof("RLS request throttled by client-side adaptive throttling")
+		return true
+	}
+	go func() {
+		req := &rlspb.RouteLookupRequest{
+			TargetType:      "grpc",
+			KeyMap:          reqKeys,
+			Reason:          reason,
+			StaleHeaderData: staleHeaders,
+		}
+		if cc.logger.V(2) {
+			cc.logger.Infof("Sending RLS request %+v", pretty.ToJSON(req))
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), cc.rpcTimeout)
+		defer cancel()
+		resp, err := cc.client.RouteLookup(ctx, req)
+		cb(resp.GetTargets(), resp.GetHeaderData(), err)
+	}()
+	return false
+}

balancer/rls/control_channel_test.go

@@ -0,0 +1,465 @@
+/*
+ *
+ * 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 rls
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"fmt"
+	"os"
+	"regexp"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/internal"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	rlstest "google.golang.org/grpc/internal/testutils/rls"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/testdata"
+	"google.golang.org/protobuf/proto"
+)
+
+// TestControlChannelThrottled tests the case where the adaptive throttler
+// indicates that the control channel needs to be throttled.
+func (s) TestControlChannelThrottled(t *testing.T) {
+	// Start an RLS server and set the throttler to always throttle requests.
+	rlsServer, rlsReqCh := rlstest.SetupFakeRLSServer(t, nil)
+	overrideAdaptiveThrottler(t, alwaysThrottlingThrottler())
+
+	// Create a control channel to the fake RLS server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, balancer.BuildOptions{}, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the attempt to be throttled.
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, nil)
+
+	select {
+	case <-rlsReqCh:
+		t.Fatal("RouteLookup RPC invoked when control channel is throttled")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+// TestLookupFailure tests the case where the RLS server responds with an error.
+func (s) TestLookupFailure(t *testing.T) {
+	// Start an RLS server and set the throttler to never throttle requests.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil)
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Setup the RLS server to respond with errors.
+	rlsServer.SetResponseCallback(func(context.Context, *rlspb.RouteLookupRequest) *rlstest.RouteLookupResponse {
+		return &rlstest.RouteLookupResponse{Err: errors.New("rls failure")}
+	})
+
+	// Create a control channel to the fake RLS server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, balancer.BuildOptions{}, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the callback to be invoked with an error.
+	errCh := make(chan error, 1)
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(_ []string, _ string, err error) {
+		if err == nil {
+			errCh <- errors.New("rlsClient.lookup() succeeded, should have failed")
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// TestLookupDeadlineExceeded tests the case where the RLS server does not
+// respond within the configured rpc timeout.
+func (s) TestLookupDeadlineExceeded(t *testing.T) {
+	// A unary interceptor which returns a status error with DeadlineExceeded.
+	interceptor := func(context.Context, any, *grpc.UnaryServerInfo, grpc.UnaryHandler) (resp any, err error) {
+		return nil, status.Error(codes.DeadlineExceeded, "deadline exceeded")
+	}
+
+	// Start an RLS server and set the throttler to never throttle.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil, grpc.UnaryInterceptor(interceptor))
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Create a control channel with a small deadline.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestShortTimeout, balancer.BuildOptions{}, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the callback to be invoked with an error.
+	errCh := make(chan error)
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(_ []string, _ string, err error) {
+		if st, ok := status.FromError(err); !ok || st.Code() != codes.DeadlineExceeded {
+			errCh <- fmt.Errorf("rlsClient.lookup() returned error: %v, want %v", err, codes.DeadlineExceeded)
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// testCredsBundle wraps a test call creds and real transport creds.
+type testCredsBundle struct {
+	transportCreds credentials.TransportCredentials
+	callCreds      credentials.PerRPCCredentials
+}
+
+func (f *testCredsBundle) TransportCredentials() credentials.TransportCredentials {
+	return f.transportCreds
+}
+
+func (f *testCredsBundle) PerRPCCredentials() credentials.PerRPCCredentials {
+	return f.callCreds
+}
+
+func (f *testCredsBundle) NewWithMode(mode string) (credentials.Bundle, error) {
+	if mode != internal.CredsBundleModeFallback {
+		return nil, fmt.Errorf("unsupported mode: %v", mode)
+	}
+	return &testCredsBundle{
+		transportCreds: f.transportCreds,
+		callCreds:      f.callCreds,
+	}, nil
+}
+
+var (
+	// Call creds sent by the testPerRPCCredentials on the client, and verified
+	// by an interceptor on the server.
+	perRPCCredsData = map[string]string{
+		"test-key":     "test-value",
+		"test-key-bin": string([]byte{1, 2, 3}),
+	}
+)
+
+type testPerRPCCredentials struct {
+	callCreds map[string]string
+}
+
+func (f *testPerRPCCredentials) GetRequestMetadata(context.Context, ...string) (map[string]string, error) {
+	return f.callCreds, nil
+}
+
+func (f *testPerRPCCredentials) RequireTransportSecurity() bool {
+	return true
+}
+
+// Unary server interceptor which validates if the RPC contains call credentials
+// which match `perRPCCredsData
+func callCredsValidatingServerInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return nil, status.Error(codes.PermissionDenied, "didn't find metadata in context")
+	}
+	for k, want := range perRPCCredsData {
+		got, ok := md[k]
+		if !ok {
+			return ctx, status.Errorf(codes.PermissionDenied, "didn't find call creds key %v in context", k)
+		}
+		if got[0] != want {
+			return ctx, status.Errorf(codes.PermissionDenied, "for key %v, got value %v, want %v", k, got, want)
+		}
+	}
+	return handler(ctx, req)
+}
+
+// makeTLSCreds is a test helper which creates a TLS based transport credentials
+// from files specified in the arguments.
+func makeTLSCreds(t *testing.T, certPath, keyPath, rootsPath string) credentials.TransportCredentials {
+	cert, err := tls.LoadX509KeyPair(testdata.Path(certPath), testdata.Path(keyPath))
+	if err != nil {
+		t.Fatalf("tls.LoadX509KeyPair(%q, %q) failed: %v", certPath, keyPath, err)
+	}
+	b, err := os.ReadFile(testdata.Path(rootsPath))
+	if err != nil {
+		t.Fatalf("os.ReadFile(%q) failed: %v", rootsPath, err)
+	}
+	roots := x509.NewCertPool()
+	if !roots.AppendCertsFromPEM(b) {
+		t.Fatal("failed to append certificates")
+	}
+	return credentials.NewTLS(&tls.Config{
+		Certificates: []tls.Certificate{cert},
+		RootCAs:      roots,
+	})
+}
+
+const (
+	wantHeaderData  = "headerData"
+	staleHeaderData = "staleHeaderData"
+)
+
+var (
+	keyMap = map[string]string{
+		"k1": "v1",
+		"k2": "v2",
+	}
+	wantTargets   = []string{"us_east_1.firestore.googleapis.com"}
+	lookupRequest = &rlspb.RouteLookupRequest{
+		TargetType:      "grpc",
+		KeyMap:          keyMap,
+		Reason:          rlspb.RouteLookupRequest_REASON_MISS,
+		StaleHeaderData: staleHeaderData,
+	}
+	lookupResponse = &rlstest.RouteLookupResponse{
+		Resp: &rlspb.RouteLookupResponse{
+			Targets:    wantTargets,
+			HeaderData: wantHeaderData,
+		},
+	}
+)
+
+func testControlChannelCredsSuccess(t *testing.T, sopts []grpc.ServerOption, bopts balancer.BuildOptions) {
+	// Start an RLS server and set the throttler to never throttle requests.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil, sopts...)
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Setup the RLS server to respond with a valid response.
+	rlsServer.SetResponseCallback(func(context.Context, *rlspb.RouteLookupRequest) *rlstest.RouteLookupResponse {
+		return lookupResponse
+	})
+
+	// Verify that the request received by the RLS matches the expected one.
+	rlsServer.SetRequestCallback(func(got *rlspb.RouteLookupRequest) {
+		if diff := cmp.Diff(lookupRequest, got, cmp.Comparer(proto.Equal)); diff != "" {
+			t.Errorf("RouteLookupRequest diff (-want, +got):\n%s", diff)
+		}
+	})
+
+	// Create a control channel to the fake server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, bopts, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect a successful callback invocation.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	errCh := make(chan error, 1)
+	ctrlCh.lookup(keyMap, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(targets []string, headerData string, err error) {
+		if err != nil {
+			errCh <- fmt.Errorf("rlsClient.lookup() failed with err: %v", err)
+			return
+		}
+		if !cmp.Equal(targets, wantTargets) || headerData != wantHeaderData {
+			errCh <- fmt.Errorf("rlsClient.lookup() = (%v, %s), want (%v, %s)", targets, headerData, wantTargets, wantHeaderData)
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// TestControlChannelCredsSuccess tests creation of the control channel with
+// different credentials, which are expected to succeed.
+func (s) TestControlChannelCredsSuccess(t *testing.T) {
+	serverCreds := makeTLSCreds(t, "x509/server1_cert.pem", "x509/server1_key.pem", "x509/client_ca_cert.pem")
+	clientCreds := makeTLSCreds(t, "x509/client1_cert.pem", "x509/client1_key.pem", "x509/server_ca_cert.pem")
+
+	tests := []struct {
+		name  string
+		sopts []grpc.ServerOption
+		bopts balancer.BuildOptions
+	}{
+		{
+			name:  "insecure",
+			sopts: nil,
+			bopts: balancer.BuildOptions{},
+		},
+		{
+			name:  "transport creds only",
+			sopts: []grpc.ServerOption{grpc.Creds(serverCreds)},
+			bopts: balancer.BuildOptions{
+				DialCreds: clientCreds,
+				Authority: "x.test.example.com",
+			},
+		},
+		{
+			name: "creds bundle",
+			sopts: []grpc.ServerOption{
+				grpc.Creds(serverCreds),
+				grpc.UnaryInterceptor(callCredsValidatingServerInterceptor),
+			},
+			bopts: balancer.BuildOptions{
+				CredsBundle: &testCredsBundle{
+					transportCreds: clientCreds,
+					callCreds:      &testPerRPCCredentials{callCreds: perRPCCredsData},
+				},
+				Authority: "x.test.example.com",
+			},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			testControlChannelCredsSuccess(t, test.sopts, test.bopts)
+		})
+	}
+}
+
+func testControlChannelCredsFailure(t *testing.T, sopts []grpc.ServerOption, bopts balancer.BuildOptions, wantCode codes.Code, wantErrRegex *regexp.Regexp) {
+	// StartFakeRouteLookupServer a fake server.
+	//
+	// Start an RLS server and set the throttler to never throttle requests. The
+	// creds failures happen before the RPC handler on the server is invoked.
+	// So, there is need to setup the request and responses on the fake server.
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil, sopts...)
+	overrideAdaptiveThrottler(t, neverThrottlingThrottler())
+
+	// Create the control channel to the fake server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, bopts, nil)
+	if err != nil {
+		t.Fatalf("Failed to create control channel to RLS server: %v", err)
+	}
+	defer ctrlCh.close()
+
+	// Perform the lookup and expect the callback to be invoked with an error.
+	errCh := make(chan error)
+	ctrlCh.lookup(nil, rlspb.RouteLookupRequest_REASON_MISS, staleHeaderData, func(_ []string, _ string, err error) {
+		if st, ok := status.FromError(err); !ok || st.Code() != wantCode || !wantErrRegex.MatchString(st.String()) {
+			errCh <- fmt.Errorf("rlsClient.lookup() returned error: %v, wantCode: %v, wantErr: %s", err, wantCode, wantErrRegex.String())
+			return
+		}
+		errCh <- nil
+	})
+
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("timeout when waiting for lookup callback to be invoked")
+	case err := <-errCh:
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// TestControlChannelCredsFailure tests creation of the control channel with
+// different credentials, which are expected to fail.
+func (s) TestControlChannelCredsFailure(t *testing.T) {
+	serverCreds := makeTLSCreds(t, "x509/server1_cert.pem", "x509/server1_key.pem", "x509/client_ca_cert.pem")
+	clientCreds := makeTLSCreds(t, "x509/client1_cert.pem", "x509/client1_key.pem", "x509/server_ca_cert.pem")
+
+	tests := []struct {
+		name         string
+		sopts        []grpc.ServerOption
+		bopts        balancer.BuildOptions
+		wantCode     codes.Code
+		wantErrRegex *regexp.Regexp
+	}{
+		{
+			name:  "transport creds authority mismatch",
+			sopts: []grpc.ServerOption{grpc.Creds(serverCreds)},
+			bopts: balancer.BuildOptions{
+				DialCreds: clientCreds,
+				Authority: "authority-mismatch",
+			},
+			wantCode:     codes.Unavailable,
+			wantErrRegex: regexp.MustCompile(`transport: authentication handshake failed: .* \*\.test\.example\.com.*authority-mismatch`),
+		},
+		{
+			name:  "transport creds handshake failure",
+			sopts: nil, // server expects insecure connection
+			bopts: balancer.BuildOptions{
+				DialCreds: clientCreds,
+				Authority: "x.test.example.com",
+			},
+			wantCode:     codes.Unavailable,
+			wantErrRegex: regexp.MustCompile("transport: authentication handshake failed: .*"),
+		},
+		{
+			name: "call creds mismatch",
+			sopts: []grpc.ServerOption{
+				grpc.Creds(serverCreds),
+				grpc.UnaryInterceptor(callCredsValidatingServerInterceptor), // server expects call creds
+			},
+			bopts: balancer.BuildOptions{
+				CredsBundle: &testCredsBundle{
+					transportCreds: clientCreds,
+					callCreds:      &testPerRPCCredentials{}, // sends no call creds
+				},
+				Authority: "x.test.example.com",
+			},
+			wantCode:     codes.PermissionDenied,
+			wantErrRegex: regexp.MustCompile("didn't find call creds"),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			testControlChannelCredsFailure(t, test.sopts, test.bopts, test.wantCode, test.wantErrRegex)
+		})
+	}
+}
+
+type unsupportedCredsBundle struct {
+	credentials.Bundle
+}
+
+func (*unsupportedCredsBundle) NewWithMode(mode string) (credentials.Bundle, error) {
+	return nil, fmt.Errorf("unsupported mode: %v", mode)
+}
+
+// TestNewControlChannelUnsupportedCredsBundle tests the case where the control
+// channel is configured with a bundle which does not support the mode we use.
+func (s) TestNewControlChannelUnsupportedCredsBundle(t *testing.T) {
+	rlsServer, _ := rlstest.SetupFakeRLSServer(t, nil)
+
+	// Create the control channel to the fake server.
+	ctrlCh, err := newControlChannel(rlsServer.Address, "", defaultTestTimeout, balancer.BuildOptions{CredsBundle: &unsupportedCredsBundle{}}, nil)
+	if err == nil {
+		ctrlCh.close()
+		t.Fatal("newControlChannel succeeded when expected to fail")
+	}
+}

balancer/rls/helpers_test.go

@@ -0,0 +1,299 @@
+/*
+ *
+ * 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 rls
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer/rls/internal/test/e2e"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/grpctest"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
+	"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/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/serviceconfig"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// fakeBackoffStrategy is a fake implementation of the backoff.Strategy
+// interface, for tests to inject the backoff duration.
+type fakeBackoffStrategy struct {
+	backoff time.Duration
+}
+
+func (f *fakeBackoffStrategy) Backoff(int) time.Duration {
+	return f.backoff
+}
+
+// fakeThrottler is a fake implementation of the adaptiveThrottler interface.
+type fakeThrottler struct {
+	throttleFunc func() bool   // Fake throttler implementation.
+	throttleCh   chan struct{} // Invocation of ShouldThrottle signals here.
+}
+
+func (f *fakeThrottler) ShouldThrottle() bool {
+	select {
+	case <-f.throttleCh:
+	default:
+	}
+	f.throttleCh <- struct{}{}
+
+	return f.throttleFunc()
+}
+
+func (f *fakeThrottler) RegisterBackendResponse(bool) {}
+
+// alwaysThrottlingThrottler returns a fake throttler which always throttles.
+func alwaysThrottlingThrottler() *fakeThrottler {
+	return &fakeThrottler{
+		throttleFunc: func() bool { return true },
+		throttleCh:   make(chan struct{}, 1),
+	}
+}
+
+// neverThrottlingThrottler returns a fake throttler which never throttles.
+func neverThrottlingThrottler() *fakeThrottler {
+	return &fakeThrottler{
+		throttleFunc: func() bool { return false },
+		throttleCh:   make(chan struct{}, 1),
+	}
+}
+
+// oneTimeAllowingThrottler returns a fake throttler which does not throttle
+// requests until the client RPC succeeds, but throttles everything that comes
+// after. This is useful for tests which need to set up a valid cache entry
+// before testing other cases.
+func oneTimeAllowingThrottler(firstRPCDone *grpcsync.Event) *fakeThrottler {
+	return &fakeThrottler{
+		throttleFunc: firstRPCDone.HasFired,
+		throttleCh:   make(chan struct{}, 1),
+	}
+}
+
+func overrideAdaptiveThrottler(t *testing.T, f *fakeThrottler) {
+	origAdaptiveThrottler := newAdaptiveThrottler
+	newAdaptiveThrottler = func() adaptiveThrottler { return f }
+	t.Cleanup(func() { newAdaptiveThrottler = origAdaptiveThrottler })
+}
+
+// buildBasicRLSConfig constructs a basic service config for the RLS LB policy
+// with header matching rules. This expects the passed child policy name to
+// have been registered by the caller.
+func buildBasicRLSConfig(childPolicyName, rlsServerAddress string) *e2e.RLSConfig {
+	return &e2e.RLSConfig{
+		RouteLookupConfig: &rlspb.RouteLookupConfig{
+			GrpcKeybuilders: []*rlspb.GrpcKeyBuilder{
+				{
+					Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "grpc.testing.TestService"}},
+					Headers: []*rlspb.NameMatcher{
+						{Key: "k1", Names: []string{"n1"}},
+						{Key: "k2", Names: []string{"n2"}},
+					},
+				},
+			},
+			LookupService:        rlsServerAddress,
+			LookupServiceTimeout: durationpb.New(defaultTestTimeout),
+			CacheSizeBytes:       1024,
+		},
+		RouteLookupChannelServiceConfig:  `{"loadBalancingConfig": [{"pick_first": {}}]}`,
+		ChildPolicy:                      &internalserviceconfig.BalancerConfig{Name: childPolicyName},
+		ChildPolicyConfigTargetFieldName: e2e.RLSChildPolicyTargetNameField,
+	}
+}
+
+// buildBasicRLSConfigWithChildPolicy constructs a very basic service config for
+// the RLS LB policy. It also registers a test LB policy which is capable of
+// being a child of the RLS LB policy.
+func buildBasicRLSConfigWithChildPolicy(t *testing.T, childPolicyName, rlsServerAddress string) *e2e.RLSConfig {
+	childPolicyName = "test-child-policy" + childPolicyName
+	e2e.RegisterRLSChildPolicy(childPolicyName, nil)
+	t.Logf("Registered child policy with name %q", childPolicyName)
+
+	return &e2e.RLSConfig{
+		RouteLookupConfig: &rlspb.RouteLookupConfig{
+			GrpcKeybuilders:      []*rlspb.GrpcKeyBuilder{{Names: []*rlspb.GrpcKeyBuilder_Name{{Service: "grpc.testing.TestService"}}}},
+			LookupService:        rlsServerAddress,
+			LookupServiceTimeout: durationpb.New(defaultTestTimeout),
+			CacheSizeBytes:       1024,
+		},
+		RouteLookupChannelServiceConfig:  `{"loadBalancingConfig": [{"pick_first": {}}]}`,
+		ChildPolicy:                      &internalserviceconfig.BalancerConfig{Name: childPolicyName},
+		ChildPolicyConfigTargetFieldName: e2e.RLSChildPolicyTargetNameField,
+	}
+}
+
+// startBackend starts a backend implementing the TestService on a local port.
+// It returns a channel for tests to get notified whenever an RPC is invoked on
+// the backend. This allows tests to ensure that RPCs reach expected backends.
+// Also returns the address of the backend.
+func startBackend(t *testing.T, sopts ...grpc.ServerOption) (rpcCh chan struct{}, address string) {
+	t.Helper()
+
+	rpcCh = make(chan struct{}, 1)
+	backend := &stubserver.StubServer{
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			select {
+			case rpcCh <- struct{}{}:
+			default:
+			}
+			return &testpb.Empty{}, nil
+		},
+	}
+	if err := backend.StartServer(sopts...); err != nil {
+		t.Fatalf("Failed to start backend: %v", err)
+	}
+	t.Logf("Started TestService backend at: %q", backend.Address)
+	t.Cleanup(func() { backend.Stop() })
+	return rpcCh, backend.Address
+}
+
+// startManualResolverWithConfig registers and returns a manual resolver which
+// pushes the RLS LB policy's service config on the channel.
+func startManualResolverWithConfig(t *testing.T, rlsConfig *e2e.RLSConfig) *manual.Resolver {
+	t.Helper()
+
+	scJSON, err := rlsConfig.ServiceConfigJSON()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(scJSON)
+	r := manual.NewBuilderWithScheme("rls-e2e")
+	r.InitialState(resolver.State{ServiceConfig: sc})
+	t.Cleanup(r.Close)
+	return r
+}
+
+// makeTestRPCAndExpectItToReachBackend is a test helper function which makes
+// the EmptyCall RPC on the given ClientConn and verifies that it reaches a
+// backend. The latter is accomplished by listening on the provided channel
+// which gets pushed to whenever the backend in question gets an RPC.
+//
+// There are many instances where it can take a while before the attempted RPC
+// reaches the expected backend. Examples include, but are not limited to:
+//   - control channel is changed in a config update. The RLS LB policy creates a
+//     new control channel, and sends a new picker to gRPC. But it takes a while
+//     before gRPC actually starts using the new picker.
+//   - test is waiting for a cache entry to expire after which we expect a
+//     different behavior because we have configured the fake RLS server to return
+//     different backends.
+//
+// Therefore, we do not return an error when the RPC fails. Instead, we wait for
+// the context to expire before failing.
+func makeTestRPCAndExpectItToReachBackend(ctx context.Context, t *testing.T, cc *grpc.ClientConn, ch chan struct{}) {
+	t.Helper()
+
+	// Drain the backend channel before performing the RPC to remove any
+	// notifications from previous RPCs.
+	select {
+	case <-ch:
+	default:
+	}
+
+	for {
+		if err := ctx.Err(); err != nil {
+			t.Fatalf("Timeout when waiting for RPCs to be routed to the given target: %v", err)
+		}
+		sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+		client := testgrpc.NewTestServiceClient(cc)
+		client.EmptyCall(sCtx, &testpb.Empty{})
+
+		select {
+		case <-sCtx.Done():
+		case <-ch:
+			sCancel()
+			return
+		}
+	}
+}
+
+// makeTestRPCAndVerifyError is a test helper function which makes the EmptyCall
+// RPC on the given ClientConn and verifies that the RPC fails with the given
+// status code and error.
+//
+// Similar to makeTestRPCAndExpectItToReachBackend, retries until expected
+// outcome is reached or the provided context has expired.
+func makeTestRPCAndVerifyError(ctx context.Context, t *testing.T, cc *grpc.ClientConn, wantCode codes.Code, wantErr error) {
+	t.Helper()
+
+	for {
+		if err := ctx.Err(); err != nil {
+			t.Fatalf("Timeout when waiting for RPCs to fail with given error: %v", err)
+		}
+		sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
+		client := testgrpc.NewTestServiceClient(cc)
+		_, err := client.EmptyCall(sCtx, &testpb.Empty{})
+
+		// If the RPC fails with the expected code and expected error message (if
+		// one was provided), we return. Else we retry after blocking for a little
+		// while to ensure that we don't keep blasting away with RPCs.
+		if code := status.Code(err); code == wantCode {
+			if wantErr == nil || strings.Contains(err.Error(), wantErr.Error()) {
+				sCancel()
+				return
+			}
+		}
+		<-sCtx.Done()
+	}
+}
+
+// verifyRLSRequest is a test helper which listens on a channel to see if an RLS
+// request was received by the fake RLS server. Based on whether the test
+// expects a request to be sent out or not, it uses a different timeout.
+func verifyRLSRequest(t *testing.T, ch chan struct{}, wantRequest bool) {
+	t.Helper()
+
+	if wantRequest {
+		select {
+		case <-time.After(defaultTestTimeout):
+			t.Fatalf("Timeout when waiting for an RLS request to be sent out")
+		case <-ch:
+		}
+	} else {
+		select {
+		case <-time.After(defaultTestShortTimeout):
+		case <-ch:
+			t.Fatalf("RLS request sent out when not expecting one")
+		}
+	}
+}

balancer/rls/internal/adaptive/adaptive.go

@@ -20,16 +20,15 @@
 package adaptive
 
 import (
+	rand "math/rand/v2"
 	"sync"
 	"time"
-
-	"google.golang.org/grpc/internal/grpcrand"
 )
 
 // For overriding in unittests.
 var (
-	timeNowFunc = func() time.Time { return time.Now() }
-	randFunc    = func() float64 { return grpcrand.Float64() }
+	timeNowFunc = time.Now
+	randFunc    = rand.Float64
 )
 
 const (
@@ -45,21 +44,21 @@ const (
 // The throttler has the following knobs for which we will use defaults for
 // now. If there is a need to make them configurable at a later point in time,
 // support for the same will be added.
-// * Duration: amount of recent history that will be taken into account for
-//   making client-side throttling decisions. A default of 30 seconds is used.
-// * Bins: number of bins to be used for bucketing historical data. A default
-//   of 100 is used.
-// * RatioForAccepts: ratio by which accepts are multiplied, typically a value
-//   slightly larger than 1.0. This is used to make the throttler behave as if
-//   the backend had accepted more requests than it actually has, which lets us
-//   err on the side of sending to the backend more requests than we think it
-//   will accept for the sake of speeding up the propagation of state. A
-//   default of 2.0 is used.
-// * RequestsPadding: is used to decrease the (client-side) throttling
-//   probability in the low QPS regime (to speed up propagation of state), as
-//   well as to safeguard against hitting a client-side throttling probability
-//   of 100%. The weight of this value decreases as the number of requests in
-//   recent history grows. A default of 8 is used.
+//   - Duration: amount of recent history that will be taken into account for
+//     making client-side throttling decisions. A default of 30 seconds is used.
+//   - Bins: number of bins to be used for bucketing historical data. A default
+//     of 100 is used.
+//   - RatioForAccepts: ratio by which accepts are multiplied, typically a value
+//     slightly larger than 1.0. This is used to make the throttler behave as if
+//     the backend had accepted more requests than it actually has, which lets us
+//     err on the side of sending to the backend more requests than we think it
+//     will accept for the sake of speeding up the propagation of state. A
+//     default of 2.0 is used.
+//   - RequestsPadding: is used to decrease the (client-side) throttling
+//     probability in the low QPS regime (to speed up propagation of state), as
+//     well as to safeguard against hitting a client-side throttling probability
+//     of 100%. The weight of this value decreases as the number of requests in
+//     recent history grows. A default of 8 is used.
 //
 // The adaptive throttler attempts to estimate the probability that a request
 // will be throttled using recent history. Server requests (both throttled and

balancer/rls/internal/adaptive/adaptive_test.go

@@ -25,13 +25,13 @@ import (
 )
 
 // stats returns a tuple with accepts, throttles for the current time.
-func (th *Throttler) stats() (int64, int64) {
+func (t *Throttler) stats() (int64, int64) {
 	now := timeNowFunc()
 
-	th.mu.Lock()
-	a, t := th.accepts.sum(now), th.throttles.sum(now)
-	th.mu.Unlock()
-	return a, t
+	t.mu.Lock()
+	a, th := t.accepts.sum(now), t.throttles.sum(now)
+	t.mu.Unlock()
+	return a, th
 }
 
 // Enums for responses.
@@ -156,7 +156,7 @@ func TestShouldThrottleOptions(t *testing.T) {
 	for _, test := range testcases {
 		t.Run(test.desc, func(t *testing.T) {
 			m.SetNanos(0)
-			th := newWithArgs(time.Duration(time.Nanosecond), 1, test.ratioForAccepts, test.requestsPadding)
+			th := newWithArgs(time.Nanosecond, 1, test.ratioForAccepts, test.requestsPadding)
 			for i, response := range responses {
 				if response != E {
 					th.RegisterBackendResponse(response == T)

balancer/rls/internal/adaptive/lookback.go

@@ -82,10 +82,3 @@ func (l *lookback) advance(t time.Time) int64 {
 	l.head = nh
 	return nh
 }
-
-func min(x int64, y int64) int64 {
-	if x < y {
-		return x
-	}
-	return y
-}

balancer/rls/internal/balancer.go

@@ -1,201 +0,0 @@
-/*
- *
- * 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 rls
-
-import (
-	"sync"
-
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/grpclog"
-	"google.golang.org/grpc/internal/grpcsync"
-)
-
-var (
-	_ balancer.Balancer = (*rlsBalancer)(nil)
-
-	// For overriding in tests.
-	newRLSClientFunc = newRLSClient
-	logger           = grpclog.Component("rls")
-)
-
-// rlsBalancer implements the RLS LB policy.
-type rlsBalancer struct {
-	done *grpcsync.Event
-	cc   balancer.ClientConn
-	opts balancer.BuildOptions
-
-	// Mutex protects all the state maintained by the LB policy.
-	// TODO(easwars): Once we add the cache, we will also have another lock for
-	// the cache alone.
-	mu    sync.Mutex
-	lbCfg *lbConfig        // Most recently received service config.
-	rlsCC *grpc.ClientConn // ClientConn to the RLS server.
-	rlsC  *rlsClient       // RLS client wrapper.
-
-	ccUpdateCh chan *balancer.ClientConnState
-}
-
-// run is a long running goroutine which handles all the updates that the
-// balancer wishes to handle. The appropriate updateHandler will push the update
-// on to a channel that this goroutine will select on, thereby the handling of
-// the update will happen asynchronously.
-func (lb *rlsBalancer) run() {
-	for {
-		// TODO(easwars): Handle other updates like subConn state changes, RLS
-		// responses from the server etc.
-		select {
-		case u := <-lb.ccUpdateCh:
-			lb.handleClientConnUpdate(u)
-		case <-lb.done.Done():
-			return
-		}
-	}
-}
-
-// handleClientConnUpdate handles updates to the service config.
-// If the RLS server name or the RLS RPC timeout changes, it updates the control
-// channel accordingly.
-// TODO(easwars): Handle updates to other fields in the service config.
-func (lb *rlsBalancer) handleClientConnUpdate(ccs *balancer.ClientConnState) {
-	logger.Infof("rls: service config: %+v", ccs.BalancerConfig)
-	lb.mu.Lock()
-	defer lb.mu.Unlock()
-
-	if lb.done.HasFired() {
-		logger.Warning("rls: received service config after balancer close")
-		return
-	}
-
-	newCfg := ccs.BalancerConfig.(*lbConfig)
-	if lb.lbCfg.Equal(newCfg) {
-		logger.Info("rls: new service config matches existing config")
-		return
-	}
-
-	lb.updateControlChannel(newCfg)
-	lb.lbCfg = newCfg
-}
-
-// UpdateClientConnState pushes the received ClientConnState update on the
-// update channel which will be processed asynchronously by the run goroutine.
-// Implements balancer.Balancer interface.
-func (lb *rlsBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
-	select {
-	case lb.ccUpdateCh <- &ccs:
-	case <-lb.done.Done():
-	}
-	return nil
-}
-
-// ResolverErr implements balancer.Balancer interface.
-func (lb *rlsBalancer) ResolverError(error) {
-	// ResolverError is called by gRPC when the name resolver reports an error.
-	// TODO(easwars): How do we handle this?
-	logger.Fatal("rls: ResolverError is not yet unimplemented")
-}
-
-// UpdateSubConnState implements balancer.Balancer interface.
-func (lb *rlsBalancer) UpdateSubConnState(_ balancer.SubConn, _ balancer.SubConnState) {
-	logger.Fatal("rls: UpdateSubConnState is not yet implemented")
-}
-
-// Cleans up the resources allocated by the LB policy including the clientConn
-// to the RLS server.
-// Implements balancer.Balancer.
-func (lb *rlsBalancer) Close() {
-	lb.mu.Lock()
-	defer lb.mu.Unlock()
-
-	lb.done.Fire()
-	if lb.rlsCC != nil {
-		lb.rlsCC.Close()
-	}
-}
-
-func (lb *rlsBalancer) ExitIdle() {
-	// TODO: are we 100% sure this should be a nop?
-}
-
-// updateControlChannel updates the RLS client if required.
-// Caller must hold lb.mu.
-func (lb *rlsBalancer) updateControlChannel(newCfg *lbConfig) {
-	oldCfg := lb.lbCfg
-	if newCfg.lookupService == oldCfg.lookupService && newCfg.lookupServiceTimeout == oldCfg.lookupServiceTimeout {
-		return
-	}
-
-	// Use RPC timeout from new config, if different from existing one.
-	timeout := oldCfg.lookupServiceTimeout
-	if timeout != newCfg.lookupServiceTimeout {
-		timeout = newCfg.lookupServiceTimeout
-	}
-
-	if newCfg.lookupService == oldCfg.lookupService {
-		// This is the case where only the timeout has changed. We will continue
-		// to use the existing clientConn. but will create a new rlsClient with
-		// the new timeout.
-		lb.rlsC = newRLSClientFunc(lb.rlsCC, lb.opts.Target.Endpoint, timeout)
-		return
-	}
-
-	// This is the case where the RLS server name has changed. We need to create
-	// a new clientConn and close the old one.
-	var dopts []grpc.DialOption
-	if dialer := lb.opts.Dialer; dialer != nil {
-		dopts = append(dopts, grpc.WithContextDialer(dialer))
-	}
-	dopts = append(dopts, dialCreds(lb.opts))
-
-	cc, err := grpc.Dial(newCfg.lookupService, dopts...)
-	if err != nil {
-		logger.Errorf("rls: dialRLS(%s, %v): %v", newCfg.lookupService, lb.opts, err)
-		// An error from a non-blocking dial indicates something serious. We
-		// should continue to use the old control channel if one exists, and
-		// return so that the rest of the config updates can be processes.
-		return
-	}
-	if lb.rlsCC != nil {
-		lb.rlsCC.Close()
-	}
-	lb.rlsCC = cc
-	lb.rlsC = newRLSClientFunc(cc, lb.opts.Target.Endpoint, timeout)
-}
-
-func dialCreds(opts balancer.BuildOptions) grpc.DialOption {
-	// The control channel should use the same authority as that of the parent
-	// channel. This ensures that the identify of the RLS server and that of the
-	// backend is the same, so if the RLS config is injected by an attacker, it
-	// cannot cause leakage of private information contained in headers set by
-	// the application.
-	server := opts.Target.Authority
-	switch {
-	case opts.DialCreds != nil:
-		if err := opts.DialCreds.OverrideServerName(server); err != nil {
-			logger.Warningf("rls: OverrideServerName(%s) = (%v), using Insecure", server, err)
-			return grpc.WithInsecure()
-		}
-		return grpc.WithTransportCredentials(opts.DialCreds)
-	case opts.CredsBundle != nil:
-		return grpc.WithTransportCredentials(opts.CredsBundle.TransportCredentials())
-	default:
-		logger.Warning("rls: no credentials available, using Insecure")
-		return grpc.WithInsecure()
-	}
-}

balancer/rls/internal/balancer_test.go

@@ -1,238 +0,0 @@
-/*
- *
- * 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 rls
-
-import (
-	"context"
-	"net"
-	"testing"
-	"time"
-
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/balancer/rls/internal/testutils/fakeserver"
-	"google.golang.org/grpc/credentials"
-	"google.golang.org/grpc/internal/grpctest"
-	"google.golang.org/grpc/internal/testutils"
-	"google.golang.org/grpc/testdata"
-)
-
-const defaultTestTimeout = 1 * time.Second
-
-type s struct {
-	grpctest.Tester
-}
-
-func Test(t *testing.T) {
-	grpctest.RunSubTests(t, s{})
-}
-
-type listenerWrapper struct {
-	net.Listener
-	connCh *testutils.Channel
-}
-
-// Accept waits for and returns the next connection to the listener.
-func (l *listenerWrapper) Accept() (net.Conn, error) {
-	c, err := l.Listener.Accept()
-	if err != nil {
-		return nil, err
-	}
-	l.connCh.Send(c)
-	return c, nil
-}
-
-func setupwithListener(t *testing.T, opts ...grpc.ServerOption) (*fakeserver.Server, *listenerWrapper, func()) {
-	t.Helper()
-
-	l, err := net.Listen("tcp", "localhost:0")
-	if err != nil {
-		t.Fatalf("net.Listen(tcp, localhost:0): %v", err)
-	}
-	lw := &listenerWrapper{
-		Listener: l,
-		connCh:   testutils.NewChannel(),
-	}
-
-	server, cleanup, err := fakeserver.Start(lw, opts...)
-	if err != nil {
-		t.Fatalf("fakeserver.Start(): %v", err)
-	}
-	t.Logf("Fake RLS server started at %s ...", server.Address)
-
-	return server, lw, cleanup
-}
-
-type testBalancerCC struct {
-	balancer.ClientConn
-}
-
-// TestUpdateControlChannelFirstConfig tests the scenario where the LB policy
-// receives its first service config and verifies that a control channel to the
-// RLS server specified in the serviceConfig is established.
-func (s) TestUpdateControlChannelFirstConfig(t *testing.T) {
-	server, lis, cleanup := setupwithListener(t)
-	defer cleanup()
-
-	bb := balancer.Get(rlsBalancerName)
-	if bb == nil {
-		t.Fatalf("balancer.Get(%s) = nil", rlsBalancerName)
-	}
-	rlsB := bb.Build(&testBalancerCC{}, balancer.BuildOptions{})
-	defer rlsB.Close()
-	t.Log("Built RLS LB policy ...")
-
-	lbCfg := &lbConfig{lookupService: server.Address}
-	t.Logf("Sending service config %+v to RLS LB policy ...", lbCfg)
-	rlsB.UpdateClientConnState(balancer.ClientConnState{BalancerConfig: lbCfg})
-
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-	defer cancel()
-	if _, err := lis.connCh.Receive(ctx); err != nil {
-		t.Fatal("Timeout expired when waiting for LB policy to create control channel")
-	}
-
-	// TODO: Verify channel connectivity state once control channel connectivity
-	// state monitoring is in place.
-
-	// TODO: Verify RLS RPC can be made once we integrate with the picker.
-}
-
-// TestUpdateControlChannelSwitch tests the scenario where a control channel
-// exists and the LB policy receives a new serviceConfig with a different RLS
-// server name. Verifies that the new control channel is created and the old one
-// is closed (the leakchecker takes care of this).
-func (s) TestUpdateControlChannelSwitch(t *testing.T) {
-	server1, lis1, cleanup1 := setupwithListener(t)
-	defer cleanup1()
-
-	server2, lis2, cleanup2 := setupwithListener(t)
-	defer cleanup2()
-
-	bb := balancer.Get(rlsBalancerName)
-	if bb == nil {
-		t.Fatalf("balancer.Get(%s) = nil", rlsBalancerName)
-	}
-	rlsB := bb.Build(&testBalancerCC{}, balancer.BuildOptions{})
-	defer rlsB.Close()
-	t.Log("Built RLS LB policy ...")
-
-	lbCfg := &lbConfig{lookupService: server1.Address}
-	t.Logf("Sending service config %+v to RLS LB policy ...", lbCfg)
-	rlsB.UpdateClientConnState(balancer.ClientConnState{BalancerConfig: lbCfg})
-
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-	defer cancel()
-	if _, err := lis1.connCh.Receive(ctx); err != nil {
-		t.Fatal("Timeout expired when waiting for LB policy to create control channel")
-	}
-
-	lbCfg = &lbConfig{lookupService: server2.Address}
-	t.Logf("Sending service config %+v to RLS LB policy ...", lbCfg)
-	rlsB.UpdateClientConnState(balancer.ClientConnState{BalancerConfig: lbCfg})
-
-	if _, err := lis2.connCh.Receive(ctx); err != nil {
-		t.Fatal("Timeout expired when waiting for LB policy to create control channel")
-	}
-
-	// TODO: Verify channel connectivity state once control channel connectivity
-	// state monitoring is in place.
-
-	// TODO: Verify RLS RPC can be made once we integrate with the picker.
-}
-
-// TestUpdateControlChannelTimeout tests the scenario where the LB policy
-// receives a service config update with a different lookupServiceTimeout, but
-// the lookupService itself remains unchanged. It verifies that the LB policy
-// does not create a new control channel in this case.
-func (s) TestUpdateControlChannelTimeout(t *testing.T) {
-	server, lis, cleanup := setupwithListener(t)
-	defer cleanup()
-
-	bb := balancer.Get(rlsBalancerName)
-	if bb == nil {
-		t.Fatalf("balancer.Get(%s) = nil", rlsBalancerName)
-	}
-	rlsB := bb.Build(&testBalancerCC{}, balancer.BuildOptions{})
-	defer rlsB.Close()
-	t.Log("Built RLS LB policy ...")
-
-	lbCfg := &lbConfig{lookupService: server.Address, lookupServiceTimeout: 1 * time.Second}
-	t.Logf("Sending service config %+v to RLS LB policy ...", lbCfg)
-	rlsB.UpdateClientConnState(balancer.ClientConnState{BalancerConfig: lbCfg})
-
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-	defer cancel()
-	if _, err := lis.connCh.Receive(ctx); err != nil {
-		t.Fatal("Timeout expired when waiting for LB policy to create control channel")
-	}
-
-	lbCfg = &lbConfig{lookupService: server.Address, lookupServiceTimeout: 2 * time.Second}
-	t.Logf("Sending service config %+v to RLS LB policy ...", lbCfg)
-	rlsB.UpdateClientConnState(balancer.ClientConnState{BalancerConfig: lbCfg})
-	if _, err := lis.connCh.Receive(ctx); err != context.DeadlineExceeded {
-		t.Fatal("LB policy created new control channel when only lookupServiceTimeout changed")
-	}
-
-	// TODO: Verify channel connectivity state once control channel connectivity
-	// state monitoring is in place.
-
-	// TODO: Verify RLS RPC can be made once we integrate with the picker.
-}
-
-// TestUpdateControlChannelWithCreds tests the scenario where the control
-// channel is to established with credentials from the parent channel.
-func (s) TestUpdateControlChannelWithCreds(t *testing.T) {
-	sCreds, err := credentials.NewServerTLSFromFile(testdata.Path("x509/server1_cert.pem"), testdata.Path("x509/server1_key.pem"))
-	if err != nil {
-		t.Fatalf("credentials.NewServerTLSFromFile(server1.pem, server1.key) = %v", err)
-	}
-	cCreds, err := credentials.NewClientTLSFromFile(testdata.Path("x509/server_ca_cert.pem"), "")
-	if err != nil {
-		t.Fatalf("credentials.NewClientTLSFromFile(ca.pem) = %v", err)
-	}
-
-	server, lis, cleanup := setupwithListener(t, grpc.Creds(sCreds))
-	defer cleanup()
-
-	bb := balancer.Get(rlsBalancerName)
-	if bb == nil {
-		t.Fatalf("balancer.Get(%s) = nil", rlsBalancerName)
-	}
-	rlsB := bb.Build(&testBalancerCC{}, balancer.BuildOptions{
-		DialCreds: cCreds,
-	})
-	defer rlsB.Close()
-	t.Log("Built RLS LB policy ...")
-
-	lbCfg := &lbConfig{lookupService: server.Address}
-	t.Logf("Sending service config %+v to RLS LB policy ...", lbCfg)
-	rlsB.UpdateClientConnState(balancer.ClientConnState{BalancerConfig: lbCfg})
-
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-	defer cancel()
-	if _, err := lis.connCh.Receive(ctx); err != nil {
-		t.Fatal("Timeout expired when waiting for LB policy to create control channel")
-	}
-
-	// TODO: Verify channel connectivity state once control channel connectivity
-	// state monitoring is in place.
-
-	// TODO: Verify RLS RPC can be made once we integrate with the picker.
-}

balancer/rls/internal/builder.go

@@ -1,53 +0,0 @@
-/*
- *
- * 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 rls implements the RLS LB policy.
-package rls
-
-import (
-	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/internal/grpcsync"
-)
-
-const rlsBalancerName = "rls"
-
-func init() {
-	balancer.Register(&rlsBB{})
-}
-
-// rlsBB helps build RLS load balancers and parse the service config to be
-// passed to the RLS load balancer.
-type rlsBB struct{}
-
-// Name returns the name of the RLS LB policy and helps implement the
-// balancer.Balancer interface.
-func (*rlsBB) Name() string {
-	return rlsBalancerName
-}
-
-func (*rlsBB) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
-	lb := &rlsBalancer{
-		done:       grpcsync.NewEvent(),
-		cc:         cc,
-		opts:       opts,
-		lbCfg:      &lbConfig{},
-		ccUpdateCh: make(chan *balancer.ClientConnState),
-	}
-	go lb.run()
-	return lb
-}