snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/admin/admin.go

@@ -0,0 +1,60 @@
+/*
+ *
+ * Copyright 2021 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package admin contains internal implementation for admin service.
+package admin
+
+import "google.golang.org/grpc"
+
+// services is a map from name to service register functions.
+var services []func(grpc.ServiceRegistrar) (func(), error)
+
+// AddService adds a service to the list of admin services.
+//
+// NOTE: this function must only be called during initialization time (i.e. in
+// an init() function), and is not thread-safe.
+//
+// If multiple services with the same service name are added (e.g. two services
+// for `grpc.channelz.v1.Channelz`), the server will panic on `Register()`.
+func AddService(f func(grpc.ServiceRegistrar) (func(), error)) {
+	services = append(services, f)
+}
+
+// Register registers the set of admin services to the given server.
+func Register(s grpc.ServiceRegistrar) (cleanup func(), _ error) {
+	var cleanups []func()
+	for _, f := range services {
+		cleanup, err := f(s)
+		if err != nil {
+			callFuncs(cleanups)
+			return nil, err
+		}
+		if cleanup != nil {
+			cleanups = append(cleanups, cleanup)
+		}
+	}
+	return func() {
+		callFuncs(cleanups)
+	}, nil
+}
+
+func callFuncs(fs []func()) {
+	for _, f := range fs {
+		f()
+	}
+}

internal/backoff/backoff.go

@@ -0,0 +1,109 @@
+/*
+ *
+ * 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 backoff implement the backoff strategy for gRPC.
+//
+// This is kept in internal until the gRPC project decides whether or not to
+// allow alternative backoff strategies.
+package backoff
+
+import (
+	"context"
+	"errors"
+	rand "math/rand/v2"
+	"time"
+
+	grpcbackoff "google.golang.org/grpc/backoff"
+)
+
+// Strategy defines the methodology for backing off after a grpc connection
+// failure.
+type Strategy interface {
+	// Backoff returns the amount of time to wait before the next retry given
+	// the number of consecutive failures.
+	Backoff(retries int) time.Duration
+}
+
+// DefaultExponential is an exponential backoff implementation using the
+// default values for all the configurable knobs defined in
+// https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
+var DefaultExponential = Exponential{Config: grpcbackoff.DefaultConfig}
+
+// Exponential implements exponential backoff algorithm as defined in
+// https://github.com/grpc/grpc/blob/master/doc/connection-backoff.md.
+type Exponential struct {
+	// Config contains all options to configure the backoff algorithm.
+	Config grpcbackoff.Config
+}
+
+// Backoff returns the amount of time to wait before the next retry given the
+// number of retries.
+func (bc Exponential) Backoff(retries int) time.Duration {
+	if retries == 0 {
+		return bc.Config.BaseDelay
+	}
+	backoff, max := float64(bc.Config.BaseDelay), float64(bc.Config.MaxDelay)
+	for backoff < max && retries > 0 {
+		backoff *= bc.Config.Multiplier
+		retries--
+	}
+	if backoff > max {
+		backoff = max
+	}
+	// Randomize backoff delays so that if a cluster of requests start at
+	// the same time, they won't operate in lockstep.
+	backoff *= 1 + bc.Config.Jitter*(rand.Float64()*2-1)
+	if backoff < 0 {
+		return 0
+	}
+	return time.Duration(backoff)
+}
+
+// ErrResetBackoff is the error to be returned by the function executed by RunF,
+// to instruct the latter to reset its backoff state.
+var ErrResetBackoff = errors.New("reset backoff state")
+
+// RunF provides a convenient way to run a function f repeatedly until the
+// context expires or f returns a non-nil error that is not ErrResetBackoff.
+// When f returns ErrResetBackoff, RunF continues to run f, but resets its
+// backoff state before doing so. backoff accepts an integer representing the
+// number of retries, and returns the amount of time to backoff.
+func RunF(ctx context.Context, f func() error, backoff func(int) time.Duration) {
+	attempt := 0
+	timer := time.NewTimer(0)
+	for ctx.Err() == nil {
+		select {
+		case <-timer.C:
+		case <-ctx.Done():
+			timer.Stop()
+			return
+		}
+
+		err := f()
+		if errors.Is(err, ErrResetBackoff) {
+			timer.Reset(0)
+			attempt = 0
+			continue
+		}
+		if err != nil {
+			return
+		}
+		timer.Reset(backoff(attempt))
+		attempt++
+	}
+}

internal/balancer/gracefulswitch/config.go

@@ -0,0 +1,84 @@
+/*
+ *
+ * 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 gracefulswitch
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+type lbConfig struct {
+	serviceconfig.LoadBalancingConfig
+
+	childBuilder balancer.Builder
+	childConfig  serviceconfig.LoadBalancingConfig
+}
+
+// ChildName returns the name of the child balancer of the gracefulswitch
+// Balancer.
+func ChildName(l serviceconfig.LoadBalancingConfig) string {
+	return l.(*lbConfig).childBuilder.Name()
+}
+
+// ParseConfig parses a child config list and returns a LB config for the
+// gracefulswitch Balancer.
+//
+// cfg is expected to be a json.RawMessage containing a JSON array of LB policy
+// names + configs as the format of the "loadBalancingConfig" field in
+// ServiceConfig.  It returns a type that should be passed to
+// UpdateClientConnState in the BalancerConfig field.
+func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	var lbCfg []map[string]json.RawMessage
+	if err := json.Unmarshal(cfg, &lbCfg); err != nil {
+		return nil, err
+	}
+	for i, e := range lbCfg {
+		if len(e) != 1 {
+			return nil, fmt.Errorf("expected a JSON struct with one entry; received entry %v at index %d", e, i)
+		}
+
+		var name string
+		var jsonCfg json.RawMessage
+		for name, jsonCfg = range e {
+		}
+
+		builder := balancer.Get(name)
+		if builder == nil {
+			// Skip unregistered balancer names.
+			continue
+		}
+
+		parser, ok := builder.(balancer.ConfigParser)
+		if !ok {
+			// This is a valid child with no config.
+			return &lbConfig{childBuilder: builder}, nil
+		}
+
+		cfg, err := parser.ParseConfig(jsonCfg)
+		if err != nil {
+			return nil, fmt.Errorf("error parsing config for policy %q: %v", name, err)
+		}
+		return &lbConfig{childBuilder: builder, childConfig: cfg}, nil
+	}
+
+	return nil, fmt.Errorf("no supported policies found in config: %v", string(cfg))
+}

internal/balancer/gracefulswitch/gracefulswitch.go

@@ -0,0 +1,421 @@
+/*
+ *
+ * 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 gracefulswitch implements a graceful switch load balancer.
+package gracefulswitch
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/resolver"
+)
+
+var errBalancerClosed = errors.New("gracefulSwitchBalancer is closed")
+var _ balancer.Balancer = (*Balancer)(nil)
+
+// NewBalancer returns a graceful switch Balancer.
+func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions) *Balancer {
+	return &Balancer{
+		cc:    cc,
+		bOpts: opts,
+	}
+}
+
+// Balancer is a utility to gracefully switch from one balancer to
+// a new balancer. It implements the balancer.Balancer interface.
+type Balancer struct {
+	bOpts balancer.BuildOptions
+	cc    balancer.ClientConn
+
+	// mu protects the following fields and all fields within balancerCurrent
+	// and balancerPending. mu does not need to be held when calling into the
+	// child balancers, as all calls into these children happen only as a direct
+	// result of a call into the gracefulSwitchBalancer, which are also
+	// guaranteed to be synchronous. There is one exception: an UpdateState call
+	// from a child balancer when current and pending are populated can lead to
+	// calling Close() on the current. To prevent that racing with an
+	// UpdateSubConnState from the channel, we hold currentMu during Close and
+	// UpdateSubConnState calls.
+	mu              sync.Mutex
+	balancerCurrent *balancerWrapper
+	balancerPending *balancerWrapper
+	closed          bool // set to true when this balancer is closed
+
+	// currentMu must be locked before mu. This mutex guards against this
+	// sequence of events: UpdateSubConnState() called, finds the
+	// balancerCurrent, gives up lock, updateState comes in, causes Close() on
+	// balancerCurrent before the UpdateSubConnState is called on the
+	// balancerCurrent.
+	currentMu sync.Mutex
+
+	// activeGoroutines tracks all the goroutines that this balancer has started
+	// and that should be waited on when the balancer closes.
+	activeGoroutines sync.WaitGroup
+}
+
+// swap swaps out the current lb with the pending lb and updates the ClientConn.
+// The caller must hold gsb.mu.
+func (gsb *Balancer) swap() {
+	gsb.cc.UpdateState(gsb.balancerPending.lastState)
+	cur := gsb.balancerCurrent
+	gsb.balancerCurrent = gsb.balancerPending
+	gsb.balancerPending = nil
+	gsb.activeGoroutines.Add(1)
+	go func() {
+		defer gsb.activeGoroutines.Done()
+		gsb.currentMu.Lock()
+		defer gsb.currentMu.Unlock()
+		cur.Close()
+	}()
+}
+
+// Helper function that checks if the balancer passed in is current or pending.
+// The caller must hold gsb.mu.
+func (gsb *Balancer) balancerCurrentOrPending(bw *balancerWrapper) bool {
+	return bw == gsb.balancerCurrent || bw == gsb.balancerPending
+}
+
+// SwitchTo initializes the graceful switch process, which completes based on
+// connectivity state changes on the current/pending balancer. Thus, the switch
+// process is not complete when this method returns. This method must be called
+// synchronously alongside the rest of the balancer.Balancer methods this
+// Graceful Switch Balancer implements.
+//
+// Deprecated: use ParseConfig and pass a parsed config to UpdateClientConnState
+// to cause the Balancer to automatically change to the new child when necessary.
+func (gsb *Balancer) SwitchTo(builder balancer.Builder) error {
+	_, err := gsb.switchTo(builder)
+	return err
+}
+
+func (gsb *Balancer) switchTo(builder balancer.Builder) (*balancerWrapper, error) {
+	gsb.mu.Lock()
+	if gsb.closed {
+		gsb.mu.Unlock()
+		return nil, errBalancerClosed
+	}
+	bw := &balancerWrapper{
+		ClientConn: gsb.cc,
+		builder:    builder,
+		gsb:        gsb,
+		lastState: balancer.State{
+			ConnectivityState: connectivity.Connecting,
+			Picker:            base.NewErrPicker(balancer.ErrNoSubConnAvailable),
+		},
+		subconns: make(map[balancer.SubConn]bool),
+	}
+	balToClose := gsb.balancerPending // nil if there is no pending balancer
+	if gsb.balancerCurrent == nil {
+		gsb.balancerCurrent = bw
+	} else {
+		gsb.balancerPending = bw
+	}
+	gsb.mu.Unlock()
+	balToClose.Close()
+	// This function takes a builder instead of a balancer because builder.Build
+	// can call back inline, and this utility needs to handle the callbacks.
+	newBalancer := builder.Build(bw, gsb.bOpts)
+	if newBalancer == nil {
+		// This is illegal and should never happen; we clear the balancerWrapper
+		// we were constructing if it happens to avoid a potential panic.
+		gsb.mu.Lock()
+		if gsb.balancerPending != nil {
+			gsb.balancerPending = nil
+		} else {
+			gsb.balancerCurrent = nil
+		}
+		gsb.mu.Unlock()
+		return nil, balancer.ErrBadResolverState
+	}
+
+	// This write doesn't need to take gsb.mu because this field never gets read
+	// or written to on any calls from the current or pending. Calls from grpc
+	// to this balancer are guaranteed to be called synchronously, so this
+	// bw.Balancer field will never be forwarded to until this SwitchTo()
+	// function returns.
+	bw.Balancer = newBalancer
+	return bw, nil
+}
+
+// Returns nil if the graceful switch balancer is closed.
+func (gsb *Balancer) latestBalancer() *balancerWrapper {
+	gsb.mu.Lock()
+	defer gsb.mu.Unlock()
+	if gsb.balancerPending != nil {
+		return gsb.balancerPending
+	}
+	return gsb.balancerCurrent
+}
+
+// UpdateClientConnState forwards the update to the latest balancer created.
+//
+// If the state's BalancerConfig is the config returned by a call to
+// gracefulswitch.ParseConfig, then this function will automatically SwitchTo
+// the balancer indicated by the config before forwarding its config to it, if
+// necessary.
+func (gsb *Balancer) UpdateClientConnState(state balancer.ClientConnState) error {
+	// The resolver data is only relevant to the most recent LB Policy.
+	balToUpdate := gsb.latestBalancer()
+	gsbCfg, ok := state.BalancerConfig.(*lbConfig)
+	if ok {
+		// Switch to the child in the config unless it is already active.
+		if balToUpdate == nil || gsbCfg.childBuilder.Name() != balToUpdate.builder.Name() {
+			var err error
+			balToUpdate, err = gsb.switchTo(gsbCfg.childBuilder)
+			if err != nil {
+				return fmt.Errorf("could not switch to new child balancer: %w", err)
+			}
+		}
+		// Unwrap the child balancer's config.
+		state.BalancerConfig = gsbCfg.childConfig
+	}
+
+	if balToUpdate == nil {
+		return errBalancerClosed
+	}
+
+	// Perform this call without gsb.mu to prevent deadlocks if the child calls
+	// back into the channel. The latest balancer can never be closed during a
+	// call from the channel, even without gsb.mu held.
+	return balToUpdate.UpdateClientConnState(state)
+}
+
+// ResolverError forwards the error to the latest balancer created.
+func (gsb *Balancer) ResolverError(err error) {
+	// The resolver data is only relevant to the most recent LB Policy.
+	balToUpdate := gsb.latestBalancer()
+	if balToUpdate == nil {
+		gsb.cc.UpdateState(balancer.State{
+			ConnectivityState: connectivity.TransientFailure,
+			Picker:            base.NewErrPicker(err),
+		})
+		return
+	}
+	// Perform this call without gsb.mu to prevent deadlocks if the child calls
+	// back into the channel. The latest balancer can never be closed during a
+	// call from the channel, even without gsb.mu held.
+	balToUpdate.ResolverError(err)
+}
+
+// ExitIdle forwards the call to the latest balancer created.
+//
+// If the latest balancer does not support ExitIdle, the subConns are
+// re-connected to manually.
+func (gsb *Balancer) ExitIdle() {
+	balToUpdate := gsb.latestBalancer()
+	if balToUpdate == nil {
+		return
+	}
+	// There is no need to protect this read with a mutex, as the write to the
+	// Balancer field happens in SwitchTo, which completes before this can be
+	// called.
+	balToUpdate.ExitIdle()
+}
+
+// updateSubConnState forwards the update to the appropriate child.
+func (gsb *Balancer) updateSubConnState(sc balancer.SubConn, state balancer.SubConnState, cb func(balancer.SubConnState)) {
+	gsb.currentMu.Lock()
+	defer gsb.currentMu.Unlock()
+	gsb.mu.Lock()
+	// Forward update to the appropriate child.  Even if there is a pending
+	// balancer, the current balancer should continue to get SubConn updates to
+	// maintain the proper state while the pending is still connecting.
+	var balToUpdate *balancerWrapper
+	if gsb.balancerCurrent != nil && gsb.balancerCurrent.subconns[sc] {
+		balToUpdate = gsb.balancerCurrent
+	} else if gsb.balancerPending != nil && gsb.balancerPending.subconns[sc] {
+		balToUpdate = gsb.balancerPending
+	}
+	if balToUpdate == nil {
+		// SubConn belonged to a stale lb policy that has not yet fully closed,
+		// or the balancer was already closed.
+		gsb.mu.Unlock()
+		return
+	}
+	if state.ConnectivityState == connectivity.Shutdown {
+		delete(balToUpdate.subconns, sc)
+	}
+	gsb.mu.Unlock()
+	if cb != nil {
+		cb(state)
+	} else {
+		balToUpdate.UpdateSubConnState(sc, state)
+	}
+}
+
+// UpdateSubConnState forwards the update to the appropriate child.
+func (gsb *Balancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	gsb.updateSubConnState(sc, state, nil)
+}
+
+// Close closes any active child balancers.
+func (gsb *Balancer) Close() {
+	gsb.mu.Lock()
+	gsb.closed = true
+	currentBalancerToClose := gsb.balancerCurrent
+	gsb.balancerCurrent = nil
+	pendingBalancerToClose := gsb.balancerPending
+	gsb.balancerPending = nil
+	gsb.mu.Unlock()
+
+	currentBalancerToClose.Close()
+	pendingBalancerToClose.Close()
+	gsb.activeGoroutines.Wait()
+}
+
+// balancerWrapper wraps a balancer.Balancer, and overrides some Balancer
+// methods to help cleanup SubConns created by the wrapped balancer.
+//
+// It implements the balancer.ClientConn interface and is passed down in that
+// capacity to the wrapped balancer. It maintains a set of subConns created by
+// the wrapped balancer and calls from the latter to create/update/shutdown
+// SubConns update this set before being forwarded to the parent ClientConn.
+// State updates from the wrapped balancer can result in invocation of the
+// graceful switch logic.
+type balancerWrapper struct {
+	balancer.ClientConn
+	balancer.Balancer
+	gsb     *Balancer
+	builder balancer.Builder
+
+	lastState balancer.State
+	subconns  map[balancer.SubConn]bool // subconns created by this balancer
+}
+
+// Close closes the underlying LB policy and shuts down the subconns it
+// created. bw must not be referenced via balancerCurrent or balancerPending in
+// gsb when called. gsb.mu must not be held.  Does not panic with a nil
+// receiver.
+func (bw *balancerWrapper) Close() {
+	// before Close is called.
+	if bw == nil {
+		return
+	}
+	// There is no need to protect this read with a mutex, as Close() is
+	// impossible to be called concurrently with the write in SwitchTo(). The
+	// callsites of Close() for this balancer in Graceful Switch Balancer will
+	// never be called until SwitchTo() returns.
+	bw.Balancer.Close()
+	bw.gsb.mu.Lock()
+	for sc := range bw.subconns {
+		sc.Shutdown()
+	}
+	bw.gsb.mu.Unlock()
+}
+
+func (bw *balancerWrapper) UpdateState(state balancer.State) {
+	// Hold the mutex for this entire call to ensure it cannot occur
+	// concurrently with other updateState() calls. This causes updates to
+	// lastState and calls to cc.UpdateState to happen atomically.
+	bw.gsb.mu.Lock()
+	defer bw.gsb.mu.Unlock()
+	bw.lastState = state
+
+	// If Close() acquires the mutex before UpdateState(), the balancer
+	// will already have been removed from the current or pending state when
+	// reaching this point.
+	if !bw.gsb.balancerCurrentOrPending(bw) {
+		// Returning here ensures that (*Balancer).swap() is not invoked after
+		// (*Balancer).Close() and therefore prevents "use after close".
+		return
+	}
+
+	if bw == bw.gsb.balancerCurrent {
+		// In the case that the current balancer exits READY, and there is a pending
+		// balancer, you can forward the pending balancer's cached State up to
+		// ClientConn and swap the pending into the current. This is because there
+		// is no reason to gracefully switch from and keep using the old policy as
+		// the ClientConn is not connected to any backends.
+		if state.ConnectivityState != connectivity.Ready && bw.gsb.balancerPending != nil {
+			bw.gsb.swap()
+			return
+		}
+		// Even if there is a pending balancer waiting to be gracefully switched to,
+		// continue to forward current balancer updates to the Client Conn. Ignoring
+		// state + picker from the current would cause undefined behavior/cause the
+		// system to behave incorrectly from the current LB policies perspective.
+		// Also, the current LB is still being used by grpc to choose SubConns per
+		// RPC, and thus should use the most updated form of the current balancer.
+		bw.gsb.cc.UpdateState(state)
+		return
+	}
+	// This method is now dealing with a state update from the pending balancer.
+	// If the current balancer is currently in a state other than READY, the new
+	// policy can be swapped into place immediately. This is because there is no
+	// reason to gracefully switch from and keep using the old policy as the
+	// ClientConn is not connected to any backends.
+	if state.ConnectivityState != connectivity.Connecting || bw.gsb.balancerCurrent.lastState.ConnectivityState != connectivity.Ready {
+		bw.gsb.swap()
+	}
+}
+
+func (bw *balancerWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	bw.gsb.mu.Lock()
+	if !bw.gsb.balancerCurrentOrPending(bw) {
+		bw.gsb.mu.Unlock()
+		return nil, fmt.Errorf("%T at address %p that called NewSubConn is deleted", bw, bw)
+	}
+	bw.gsb.mu.Unlock()
+
+	var sc balancer.SubConn
+	oldListener := opts.StateListener
+	opts.StateListener = func(state balancer.SubConnState) { bw.gsb.updateSubConnState(sc, state, oldListener) }
+	sc, err := bw.gsb.cc.NewSubConn(addrs, opts)
+	if err != nil {
+		return nil, err
+	}
+	bw.gsb.mu.Lock()
+	if !bw.gsb.balancerCurrentOrPending(bw) { // balancer was closed during this call
+		sc.Shutdown()
+		bw.gsb.mu.Unlock()
+		return nil, fmt.Errorf("%T at address %p that called NewSubConn is deleted", bw, bw)
+	}
+	bw.subconns[sc] = true
+	bw.gsb.mu.Unlock()
+	return sc, nil
+}
+
+func (bw *balancerWrapper) ResolveNow(opts resolver.ResolveNowOptions) {
+	// Ignore ResolveNow requests from anything other than the most recent
+	// balancer, because older balancers were already removed from the config.
+	if bw != bw.gsb.latestBalancer() {
+		return
+	}
+	bw.gsb.cc.ResolveNow(opts)
+}
+
+func (bw *balancerWrapper) RemoveSubConn(sc balancer.SubConn) {
+	// Note: existing third party balancers may call this, so it must remain
+	// until RemoveSubConn is fully removed.
+	sc.Shutdown()
+}
+
+func (bw *balancerWrapper) UpdateAddresses(sc balancer.SubConn, addrs []resolver.Address) {
+	bw.gsb.mu.Lock()
+	if !bw.gsb.balancerCurrentOrPending(bw) {
+		bw.gsb.mu.Unlock()
+		return
+	}
+	bw.gsb.mu.Unlock()
+	bw.gsb.cc.UpdateAddresses(sc, addrs)
+}

internal/balancer/nop/nop.go

@@ -0,0 +1,65 @@
+/*
+ *
+ * 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 nop implements a balancer with all of its balancer operations as
+// no-ops, other than returning a Transient Failure Picker on a Client Conn
+// update.
+package nop
+
+import (
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/connectivity"
+)
+
+// bal is a balancer with all of its balancer operations as no-ops, other than
+// returning a Transient Failure Picker on a Client Conn update.
+type bal struct {
+	cc  balancer.ClientConn
+	err error
+}
+
+// NewBalancer returns a no-op balancer.
+func NewBalancer(cc balancer.ClientConn, err error) balancer.Balancer {
+	return &bal{
+		cc:  cc,
+		err: err,
+	}
+}
+
+// UpdateClientConnState updates the bal's Client Conn with an Error Picker
+// and a Connectivity State of TRANSIENT_FAILURE.
+func (b *bal) UpdateClientConnState(_ balancer.ClientConnState) error {
+	b.cc.UpdateState(balancer.State{
+		Picker:            base.NewErrPicker(b.err),
+		ConnectivityState: connectivity.TransientFailure,
+	})
+	return nil
+}
+
+// ResolverError is a no-op.
+func (b *bal) ResolverError(_ error) {}
+
+// UpdateSubConnState is a no-op.
+func (b *bal) UpdateSubConnState(_ balancer.SubConn, _ balancer.SubConnState) {}
+
+// Close is a no-op.
+func (b *bal) Close() {}
+
+// ExitIdle is a no-op.
+func (b *bal) ExitIdle() {}

internal/balancer/stub/stub.go

@@ -0,0 +1,116 @@
+/*
+ *
+ * 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 stub implements a balancer for testing purposes.
+package stub
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// BalancerFuncs contains all balancer.Balancer functions with a preceding
+// *BalancerData parameter for passing additional instance information.  Any
+// nil functions will never be called.
+type BalancerFuncs struct {
+	// Init is called after ClientConn and BuildOptions are set in
+	// BalancerData.  It may be used to initialize BalancerData.Data.
+	Init func(*BalancerData)
+	// ParseConfig is used for parsing LB configs, if specified.
+	ParseConfig func(json.RawMessage) (serviceconfig.LoadBalancingConfig, error)
+
+	UpdateClientConnState func(*BalancerData, balancer.ClientConnState) error
+	ResolverError         func(*BalancerData, error)
+	Close                 func(*BalancerData)
+	ExitIdle              func(*BalancerData)
+}
+
+// BalancerData contains data relevant to a stub balancer.
+type BalancerData struct {
+	// ClientConn is set by the builder.
+	ClientConn balancer.ClientConn
+	// BuildOptions is set by the builder.
+	BuildOptions balancer.BuildOptions
+	// ChildBalancer holds a child balancer.
+	ChildBalancer balancer.Balancer
+}
+
+type bal struct {
+	bf BalancerFuncs
+	bd *BalancerData
+}
+
+func (b *bal) UpdateClientConnState(c balancer.ClientConnState) error {
+	if b.bf.UpdateClientConnState != nil {
+		return b.bf.UpdateClientConnState(b.bd, c)
+	}
+	return nil
+}
+
+func (b *bal) ResolverError(e error) {
+	if b.bf.ResolverError != nil {
+		b.bf.ResolverError(b.bd, e)
+	}
+}
+
+func (b *bal) UpdateSubConnState(sc balancer.SubConn, scs balancer.SubConnState) {
+	panic(fmt.Sprintf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, scs))
+}
+
+func (b *bal) Close() {
+	if b.bf.Close != nil {
+		b.bf.Close(b.bd)
+	}
+}
+
+func (b *bal) ExitIdle() {
+	if b.bf.ExitIdle != nil {
+		b.bf.ExitIdle(b.bd)
+	}
+}
+
+type bb struct {
+	name string
+	bf   BalancerFuncs
+}
+
+func (bb bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	b := &bal{bf: bb.bf, bd: &BalancerData{ClientConn: cc, BuildOptions: opts}}
+	if b.bf.Init != nil {
+		b.bf.Init(b.bd)
+	}
+	return b
+}
+
+func (bb bb) Name() string { return bb.name }
+
+func (bb bb) ParseConfig(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	if bb.bf.ParseConfig != nil {
+		return bb.bf.ParseConfig(lbCfg)
+	}
+	return nil, nil
+}
+
+// Register registers a stub balancer builder which will call the provided
+// functions.  The name used should be unique.
+func Register(name string, bf BalancerFuncs) {
+	balancer.Register(bb{name: name, bf: bf})
+}

internal/balancergroup/balancergroup.go

@@ -0,0 +1,607 @@
+/*
+ * 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 balancergroup implements a utility struct to bind multiple balancers
+// into one balancer.
+package balancergroup
+
+import (
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal/balancer/gracefulswitch"
+	"google.golang.org/grpc/internal/cache"
+	"google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// subBalancerWrapper is used to keep the configurations that will be used to start
+// the underlying balancer. It can be called to start/stop the underlying
+// balancer.
+//
+// When the config changes, it will pass the update to the underlying balancer
+// if it exists.
+//
+// TODO: move to a separate file?
+type subBalancerWrapper struct {
+	// subBalancerWrapper is passed to the sub-balancer as a ClientConn
+	// wrapper, only to keep the state and picker.  When sub-balancer is
+	// restarted while in cache, the picker needs to be resent.
+	//
+	// It also contains the sub-balancer ID, so the parent balancer group can
+	// keep track of SubConn/pickers and the sub-balancers they belong to. Some
+	// of the actions are forwarded to the parent ClientConn with no change.
+	// Some are forward to balancer group with the sub-balancer ID.
+	balancer.ClientConn
+	id    string
+	group *BalancerGroup
+
+	mu    sync.Mutex
+	state balancer.State
+
+	// The static part of sub-balancer. Keeps balancerBuilders and addresses.
+	// To be used when restarting sub-balancer.
+	builder balancer.Builder
+	// Options to be passed to sub-balancer at the time of creation.
+	buildOpts balancer.BuildOptions
+	// ccState is a cache of the addresses/balancer config, so when the balancer
+	// is restarted after close, it will get the previous update. It's a pointer
+	// and is set to nil at init, so when the balancer is built for the first
+	// time (not a restart), it won't receive an empty update. Note that this
+	// isn't reset to nil when the underlying balancer is closed.
+	ccState *balancer.ClientConnState
+	// The dynamic part of sub-balancer. Only used when balancer group is
+	// started. Gets cleared when sub-balancer is closed.
+	balancer *gracefulswitch.Balancer
+}
+
+// UpdateState overrides balancer.ClientConn, to keep state and picker.
+func (sbc *subBalancerWrapper) UpdateState(state balancer.State) {
+	sbc.mu.Lock()
+	sbc.state = state
+	sbc.group.updateBalancerState(sbc.id, state)
+	sbc.mu.Unlock()
+}
+
+// NewSubConn overrides balancer.ClientConn, so balancer group can keep track of
+// the relation between subconns and sub-balancers.
+func (sbc *subBalancerWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	return sbc.group.newSubConn(sbc, addrs, opts)
+}
+
+func (sbc *subBalancerWrapper) updateBalancerStateWithCachedPicker() {
+	sbc.mu.Lock()
+	if sbc.state.Picker != nil {
+		sbc.group.updateBalancerState(sbc.id, sbc.state)
+	}
+	sbc.mu.Unlock()
+}
+
+func (sbc *subBalancerWrapper) startBalancer() {
+	if sbc.balancer == nil {
+		sbc.balancer = gracefulswitch.NewBalancer(sbc, sbc.buildOpts)
+	}
+	sbc.group.logger.Infof("Creating child policy of type %q for child %q", sbc.builder.Name(), sbc.id)
+	sbc.balancer.SwitchTo(sbc.builder)
+	if sbc.ccState != nil {
+		sbc.balancer.UpdateClientConnState(*sbc.ccState)
+	}
+}
+
+// exitIdle invokes the ExitIdle method on the sub-balancer, a gracefulswitch
+// balancer.
+func (sbc *subBalancerWrapper) exitIdle() {
+	b := sbc.balancer
+	if b == nil {
+		return
+	}
+	b.ExitIdle()
+}
+
+func (sbc *subBalancerWrapper) updateClientConnState(s balancer.ClientConnState) error {
+	sbc.ccState = &s
+	b := sbc.balancer
+	if b == nil {
+		// A sub-balancer is closed when it is removed from the group or the
+		// group is closed as a whole, and is not expected to receive updates
+		// after that. But when used with the priority LB policy a sub-balancer
+		// (and the whole balancer group) could be closed because it's the lower
+		// priority, but it can still get address updates.
+		return nil
+	}
+	return b.UpdateClientConnState(s)
+}
+
+func (sbc *subBalancerWrapper) resolverError(err error) {
+	b := sbc.balancer
+	if b == nil {
+		// A sub-balancer is closed when it is removed from the group or the
+		// group is closed as a whole, and is not expected to receive updates
+		// after that. But when used with the priority LB policy a sub-balancer
+		// (and the whole balancer group) could be closed because it's the lower
+		// priority, but it can still get address updates.
+		return
+	}
+	b.ResolverError(err)
+}
+
+func (sbc *subBalancerWrapper) stopBalancer() {
+	if sbc.balancer == nil {
+		return
+	}
+	sbc.balancer.Close()
+	sbc.balancer = nil
+}
+
+// BalancerGroup takes a list of balancers, each behind a gracefulswitch
+// balancer, and make them into one balancer.
+//
+// Note that this struct doesn't implement balancer.Balancer, because it's not
+// intended to be used directly as a balancer. It's expected to be used as a
+// sub-balancer manager by a high level balancer.
+//
+//	Updates from ClientConn are forwarded to sub-balancers
+//	- service config update
+//	- address update
+//	- subConn state change
+//	  - find the corresponding balancer and forward
+//
+//	Actions from sub-balances are forwarded to parent ClientConn
+//	- new/remove SubConn
+//	- picker update and health states change
+//	  - sub-pickers are sent to an aggregator provided by the parent, which
+//	    will group them into a group-picker. The aggregated connectivity state is
+//	    also handled by the aggregator.
+//	- resolveNow
+//
+// Sub-balancers are only built when the balancer group is started. If the
+// balancer group is closed, the sub-balancers are also closed. And it's
+// guaranteed that no updates will be sent to parent ClientConn from a closed
+// balancer group.
+type BalancerGroup struct {
+	cc        balancer.ClientConn
+	buildOpts balancer.BuildOptions
+	logger    *grpclog.PrefixLogger
+
+	// stateAggregator is where the state/picker updates will be sent to. It's
+	// provided by the parent balancer, to build a picker with all the
+	// sub-pickers.
+	stateAggregator BalancerStateAggregator
+
+	// outgoingMu guards all operations in the direction:
+	// ClientConn-->Sub-balancer. Including start, stop, resolver updates and
+	// SubConn state changes.
+	//
+	// The corresponding boolean outgoingStarted is used to stop further updates
+	// to sub-balancers after they are closed.
+	outgoingMu         sync.Mutex
+	outgoingClosed     bool
+	idToBalancerConfig map[string]*subBalancerWrapper
+	// Cache for sub-balancers when they are removed. This is `nil` if caching
+	// is disabled by passing `0` for Options.SubBalancerCloseTimeout`.
+	deletedBalancerCache *cache.TimeoutCache
+
+	// incomingMu is to make sure this balancer group doesn't send updates to cc
+	// after it's closed.
+	//
+	// We don't share the mutex to avoid deadlocks (e.g. a call to sub-balancer
+	// may call back to balancer group inline. It causes deadlock if they
+	// require the same mutex).
+	//
+	// We should never need to hold multiple locks at the same time in this
+	// struct. The case where two locks are held can only happen when the
+	// underlying balancer calls back into balancer group inline. So there's an
+	// implicit lock acquisition order that outgoingMu is locked before
+	// incomingMu.
+
+	// incomingMu guards all operations in the direction:
+	// Sub-balancer-->ClientConn. Including NewSubConn, RemoveSubConn. It also
+	// guards the map from SubConn to balancer ID, so updateSubConnState needs
+	// to hold it shortly to potentially delete from the map.
+	//
+	// UpdateState is called by the balancer state aggregator, and it will
+	// decide when and whether to call.
+	//
+	// The corresponding boolean incomingStarted is used to stop further updates
+	// from sub-balancers after they are closed.
+	incomingMu      sync.Mutex
+	incomingClosed  bool // This boolean only guards calls back to ClientConn.
+	scToSubBalancer map[balancer.SubConn]*subBalancerWrapper
+}
+
+// Options wraps the arguments to be passed to the BalancerGroup ctor.
+type Options struct {
+	// CC is a reference to the parent balancer.ClientConn.
+	CC balancer.ClientConn
+	// BuildOpts contains build options to be used when creating sub-balancers.
+	BuildOpts balancer.BuildOptions
+	// StateAggregator is an implementation of the BalancerStateAggregator
+	// interface to aggregate picker and connectivity states from sub-balancers.
+	StateAggregator BalancerStateAggregator
+	// Logger is a group specific prefix logger.
+	Logger *grpclog.PrefixLogger
+	// SubBalancerCloseTimeout is the amount of time deleted sub-balancers spend
+	// in the idle cache. A value of zero here disables caching of deleted
+	// sub-balancers.
+	SubBalancerCloseTimeout time.Duration
+}
+
+// New creates a new BalancerGroup. Note that the BalancerGroup
+// needs to be started to work.
+func New(opts Options) *BalancerGroup {
+	var bc *cache.TimeoutCache
+	if opts.SubBalancerCloseTimeout != time.Duration(0) {
+		bc = cache.NewTimeoutCache(opts.SubBalancerCloseTimeout)
+	}
+
+	return &BalancerGroup{
+		cc:              opts.CC,
+		buildOpts:       opts.BuildOpts,
+		stateAggregator: opts.StateAggregator,
+		logger:          opts.Logger,
+
+		deletedBalancerCache: bc,
+		idToBalancerConfig:   make(map[string]*subBalancerWrapper),
+		scToSubBalancer:      make(map[balancer.SubConn]*subBalancerWrapper),
+	}
+}
+
+// AddWithClientConn adds a balancer with the given id to the group. The
+// balancer is built with a balancer builder registered with balancerName. The
+// given ClientConn is passed to the newly built balancer instead of the
+// one passed to balancergroup.New().
+//
+// TODO: Get rid of the existing Add() API and replace it with this.
+func (bg *BalancerGroup) AddWithClientConn(id, balancerName string, cc balancer.ClientConn) error {
+	bg.logger.Infof("Adding child policy of type %q for child %q", balancerName, id)
+	builder := balancer.Get(balancerName)
+	if builder == nil {
+		return fmt.Errorf("balancergroup: unregistered balancer name %q", balancerName)
+	}
+
+	// Store data in static map, and then check to see if bg is started.
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return fmt.Errorf("balancergroup: already closed")
+	}
+	var sbc *subBalancerWrapper
+	// Skip searching the cache if disabled.
+	if bg.deletedBalancerCache != nil {
+		if old, ok := bg.deletedBalancerCache.Remove(id); ok {
+			if bg.logger.V(2) {
+				bg.logger.Infof("Removing and reusing child policy of type %q for child %q from the balancer cache", balancerName, id)
+				bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
+			}
+
+			sbc, _ = old.(*subBalancerWrapper)
+			if sbc != nil && sbc.builder != builder {
+				// If the sub-balancer in cache was built with a different
+				// balancer builder, don't use it, cleanup this old-balancer,
+				// and behave as sub-balancer is not found in cache.
+				//
+				// NOTE that this will also drop the cached addresses for this
+				// sub-balancer, which seems to be reasonable.
+				sbc.stopBalancer()
+				// cleanupSubConns must be done before the new balancer starts,
+				// otherwise new SubConns created by the new balancer might be
+				// removed by mistake.
+				bg.cleanupSubConns(sbc)
+				sbc = nil
+			}
+		}
+	}
+	if sbc == nil {
+		sbc = &subBalancerWrapper{
+			ClientConn: cc,
+			id:         id,
+			group:      bg,
+			builder:    builder,
+			buildOpts:  bg.buildOpts,
+		}
+		sbc.startBalancer()
+	} else {
+		// When brining back a sub-balancer from cache, re-send the cached
+		// picker and state.
+		sbc.updateBalancerStateWithCachedPicker()
+	}
+	bg.idToBalancerConfig[id] = sbc
+	return nil
+}
+
+// Add adds a balancer built by builder to the group, with given id.
+func (bg *BalancerGroup) Add(id string, builder balancer.Builder) {
+	bg.AddWithClientConn(id, builder.Name(), bg.cc)
+}
+
+// Remove removes the balancer with id from the group.
+//
+// But doesn't close the balancer. The balancer is kept in a cache, and will be
+// closed after timeout. Cleanup work (closing sub-balancer and removing
+// subconns) will be done after timeout.
+func (bg *BalancerGroup) Remove(id string) {
+	bg.removeInternal(id, true)
+}
+
+// RemoveImmediately removes and closes the balancer with id from the group
+// immediately.
+func (bg *BalancerGroup) RemoveImmediately(id string) {
+	bg.removeInternal(id, false)
+}
+
+func (bg *BalancerGroup) removeInternal(id string, withCaching bool) {
+	bg.logger.Infof("Removing child policy for child %q", id)
+
+	bg.outgoingMu.Lock()
+	if bg.outgoingClosed {
+		bg.outgoingMu.Unlock()
+		return
+	}
+
+	sbToRemove, ok := bg.idToBalancerConfig[id]
+	if !ok {
+		bg.logger.Errorf("Child policy for child %q does not exist in the balancer group", id)
+		bg.outgoingMu.Unlock()
+		return
+	}
+
+	// Unconditionally remove the sub-balancer config from the map.
+	delete(bg.idToBalancerConfig, id)
+
+	if withCaching {
+		if bg.deletedBalancerCache != nil {
+			if bg.logger.V(2) {
+				bg.logger.Infof("Adding child policy for child %q to the balancer cache", id)
+			}
+			bg.deletedBalancerCache.Add(id, sbToRemove, func() {
+				if bg.logger.V(2) {
+					bg.logger.Infof("Removing child policy for child %q from the balancer cache after timeout", id)
+					bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
+				}
+
+				// A sub-balancer evicted from the timeout cache needs to closed
+				// and its subConns need to removed, unconditionally. There is a
+				// possibility that a sub-balancer might be removed (thereby
+				// moving it to the cache) around the same time that the
+				// balancergroup is closed, and by the time we get here the
+				// balancergroup might be closed.  Check for `outgoingStarted ==
+				// true` at that point can lead to a leaked sub-balancer.
+				bg.outgoingMu.Lock()
+				sbToRemove.stopBalancer()
+				bg.outgoingMu.Unlock()
+				bg.cleanupSubConns(sbToRemove)
+			})
+			if bg.logger.V(2) {
+				bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
+			}
+			bg.outgoingMu.Unlock()
+			return
+		}
+
+		// Fall through to remove the sub-balancer with immediate effect if we are not caching.
+		if bg.logger.V(2) {
+			bg.logger.Infof("Child policy for child %q was requested to be cached before eventual removal. No such cache exists. Removing right away.", id)
+		}
+	}
+
+	// Remove the sub-balancer with immediate effect if we are not caching.
+	sbToRemove.stopBalancer()
+	bg.outgoingMu.Unlock()
+	bg.cleanupSubConns(sbToRemove)
+}
+
+// bg.remove(id) doesn't do cleanup for the sub-balancer. This function does
+// cleanup after the timeout.
+func (bg *BalancerGroup) cleanupSubConns(config *subBalancerWrapper) {
+	bg.incomingMu.Lock()
+	defer bg.incomingMu.Unlock()
+	// Remove SubConns. This is only done after the balancer is
+	// actually closed.
+	//
+	// NOTE: if NewSubConn is called by this (closed) balancer later, the
+	// SubConn will be leaked. This shouldn't happen if the balancer
+	// implementation is correct. To make sure this never happens, we need to
+	// add another layer (balancer manager) between balancer group and the
+	// sub-balancers.
+	for sc, b := range bg.scToSubBalancer {
+		if b == config {
+			delete(bg.scToSubBalancer, sc)
+		}
+	}
+}
+
+// Following are actions from the parent grpc.ClientConn, forward to sub-balancers.
+
+// updateSubConnState forwards the update to cb and updates scToSubBalancer if
+// needed.
+func (bg *BalancerGroup) updateSubConnState(sc balancer.SubConn, state balancer.SubConnState, cb func(balancer.SubConnState)) {
+	bg.incomingMu.Lock()
+	if bg.incomingClosed {
+		bg.incomingMu.Unlock()
+		return
+	}
+	if _, ok := bg.scToSubBalancer[sc]; !ok {
+		bg.incomingMu.Unlock()
+		return
+	}
+	if state.ConnectivityState == connectivity.Shutdown {
+		// Only delete sc from the map when state changed to Shutdown.
+		delete(bg.scToSubBalancer, sc)
+	}
+	bg.incomingMu.Unlock()
+
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	if cb != nil {
+		cb(state)
+	}
+}
+
+// UpdateSubConnState handles the state for the subconn. It finds the
+// corresponding balancer and forwards the update.
+func (bg *BalancerGroup) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	bg.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+// UpdateClientConnState handles ClientState (including balancer config and
+// addresses) from resolver. It finds the balancer and forwards the update.
+func (bg *BalancerGroup) UpdateClientConnState(id string, s balancer.ClientConnState) error {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return nil
+	}
+	if config, ok := bg.idToBalancerConfig[id]; ok {
+		return config.updateClientConnState(s)
+	}
+	return nil
+}
+
+// ResolverError forwards resolver errors to all sub-balancers.
+func (bg *BalancerGroup) ResolverError(err error) {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	for _, config := range bg.idToBalancerConfig {
+		config.resolverError(err)
+	}
+}
+
+// Following are actions from sub-balancers, forward to ClientConn.
+
+// newSubConn: forward to ClientConn, and also create a map from sc to balancer,
+// so state update will find the right balancer.
+//
+// One note about removing SubConn: only forward to ClientConn, but not delete
+// from map. Delete sc from the map only when state changes to Shutdown. Since
+// it's just forwarding the action, there's no need for a removeSubConn()
+// wrapper function.
+func (bg *BalancerGroup) newSubConn(sbw *subBalancerWrapper, addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	// NOTE: if balancer with id was already removed, this should also return
+	// error. But since we call balancer.stopBalancer when removing the balancer, this
+	// shouldn't happen.
+	bg.incomingMu.Lock()
+	if bg.incomingClosed {
+		bg.incomingMu.Unlock()
+		return nil, fmt.Errorf("balancergroup: NewSubConn is called after balancer group is closed")
+	}
+	var sc balancer.SubConn
+	oldListener := opts.StateListener
+	opts.StateListener = func(state balancer.SubConnState) { bg.updateSubConnState(sc, state, oldListener) }
+	sc, err := sbw.ClientConn.NewSubConn(addrs, opts)
+	if err != nil {
+		bg.incomingMu.Unlock()
+		return nil, err
+	}
+	bg.scToSubBalancer[sc] = sbw
+	bg.incomingMu.Unlock()
+	return sc, nil
+}
+
+// updateBalancerState: forward the new state to balancer state aggregator. The
+// aggregator will create an aggregated picker and an aggregated connectivity
+// state, then forward to ClientConn.
+func (bg *BalancerGroup) updateBalancerState(id string, state balancer.State) {
+	bg.logger.Infof("Balancer state update from child %v, new state: %+v", id, state)
+
+	// Send new state to the aggregator, without holding the incomingMu.
+	// incomingMu is to protect all calls to the parent ClientConn, this update
+	// doesn't necessary trigger a call to ClientConn, and should already be
+	// protected by aggregator's mutex if necessary.
+	if bg.stateAggregator != nil {
+		bg.stateAggregator.UpdateState(id, state)
+	}
+}
+
+// Close closes the balancer. It stops sub-balancers, and removes the subconns.
+// When a BalancerGroup is closed, it can not receive further address updates.
+func (bg *BalancerGroup) Close() {
+	bg.incomingMu.Lock()
+	bg.incomingClosed = true
+	// Also remove all SubConns.
+	for sc := range bg.scToSubBalancer {
+		sc.Shutdown()
+		delete(bg.scToSubBalancer, sc)
+	}
+	bg.incomingMu.Unlock()
+
+	bg.outgoingMu.Lock()
+	// Setting `outgoingClosed` ensures that no entries are added to
+	// `deletedBalancerCache` after this point.
+	bg.outgoingClosed = true
+	bg.outgoingMu.Unlock()
+
+	// Clear(true) runs clear function to close sub-balancers in cache. It
+	// must be called out of outgoing mutex.
+	if bg.deletedBalancerCache != nil {
+		bg.deletedBalancerCache.Clear(true)
+	}
+
+	bg.outgoingMu.Lock()
+	for id, config := range bg.idToBalancerConfig {
+		config.stopBalancer()
+		delete(bg.idToBalancerConfig, id)
+	}
+	bg.outgoingMu.Unlock()
+}
+
+// ExitIdle should be invoked when the parent LB policy's ExitIdle is invoked.
+// It will trigger this on all sub-balancers, or reconnect their subconns if
+// not supported.
+func (bg *BalancerGroup) ExitIdle() {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	for _, config := range bg.idToBalancerConfig {
+		config.exitIdle()
+	}
+}
+
+// ExitIdleOne instructs the sub-balancer `id` to exit IDLE state, if
+// appropriate and possible.
+func (bg *BalancerGroup) ExitIdleOne(id string) {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	if config := bg.idToBalancerConfig[id]; config != nil {
+		config.exitIdle()
+	}
+}
+
+// ParseConfig parses a child config list and returns a LB config for the
+// gracefulswitch Balancer.
+//
+// cfg is expected to be a json.RawMessage containing a JSON array of LB policy
+// names + configs as the format of the "loadBalancingConfig" field in
+// ServiceConfig.  It returns a type that should be passed to
+// UpdateClientConnState in the BalancerConfig field.
+func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	return gracefulswitch.ParseConfig(cfg)
+}

internal/balancergroup/balancergroup_test.go

@@ -0,0 +1,930 @@
+/*
+ * 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 balancergroup_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/balancer/roundrobin"
+	"google.golang.org/grpc/balancer/weightedtarget/weightedaggregator"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/balancergroup"
+	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/resolver"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+var (
+	rrBuilder            = balancer.Get(roundrobin.Name)
+	testBalancerIDs      = []string{"b1", "b2", "b3"}
+	testBackendAddrs     []resolver.Address
+	testBackendEndpoints []resolver.Endpoint
+)
+
+const testBackendAddrsCount = 12
+
+func init() {
+	for i := 0; i < testBackendAddrsCount; i++ {
+		addr := resolver.Address{Addr: fmt.Sprintf("%d.%d.%d.%d:%d", i, i, i, i, i)}
+		testBackendAddrs = append(testBackendAddrs, addr)
+		testBackendEndpoints = append(testBackendEndpoints, resolver.Endpoint{Addresses: []resolver.Address{addr}})
+	}
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// Create a new balancer group, add balancer and backends.
+// - b1, weight 2, backends [0,1]
+// - b2, weight 1, backends [2,3]
+// Start the balancer group and check behavior.
+//
+// Close the balancer group.
+func (s) TestBalancerGroup_start_close(t *testing.T) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: time.Duration(0),
+	})
+
+	// Add two balancers to group and send two resolved addresses to both
+	// balancers.
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[2:4]}})
+
+	m1 := make(map[string]balancer.SubConn)
+	for i := 0; i < 4; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	// Test roundrobin on the last picker.
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[0].Addr],
+		m1[testBackendAddrs[1].Addr], m1[testBackendAddrs[1].Addr],
+		m1[testBackendAddrs[2].Addr], m1[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	gator.Stop()
+	bg.Close()
+	for i := 0; i < 4; i++ {
+		(<-cc.ShutdownSubConnCh).UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Shutdown})
+	}
+}
+
+// Test that balancer group start() doesn't deadlock if the balancer calls back
+// into balancer group inline when it gets an update.
+//
+// The potential deadlock can happen if we
+//   - hold a lock and send updates to balancer (e.g. update resolved addresses)
+//   - the balancer calls back (NewSubConn or update picker) in line
+//
+// The callback will try to hold the same lock again, which will cause a
+// deadlock.
+//
+// This test starts the balancer group with a test balancer, will update picker
+// whenever it gets an address update. It's expected that start() doesn't block
+// because of deadlock.
+func (s) TestBalancerGroup_start_close_deadlock(t *testing.T) {
+	const balancerName = "stub-TestBalancerGroup_start_close_deadlock"
+	stub.Register(balancerName, stub.BalancerFuncs{})
+	builder := balancer.Get(balancerName)
+
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: time.Duration(0),
+	})
+
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], builder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Addresses: testBackendAddrs[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], builder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Addresses: testBackendAddrs[2:4]}})
+}
+
+// initBalancerGroupForCachingTest creates a balancer group, and initialize it
+// to be ready for caching tests.
+//
+// Two rr balancers are added to bg, each with 2 ready subConns. A sub-balancer
+// is removed later, so the balancer group returned has one sub-balancer in its
+// own map, and one sub-balancer in cache.
+func initBalancerGroupForCachingTest(t *testing.T, idleCacheTimeout time.Duration) (*weightedaggregator.Aggregator, *balancergroup.BalancerGroup, *testutils.BalancerClientConn, map[string]*testutils.TestSubConn) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: idleCacheTimeout,
+	})
+
+	// Add two balancers to group and send two resolved addresses to both
+	// balancers.
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[2:4]}})
+
+	m1 := make(map[string]*testutils.TestSubConn)
+	for i := 0; i < 4; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	// Test roundrobin on the last picker.
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[0].Addr],
+		m1[testBackendAddrs[1].Addr], m1[testBackendAddrs[1].Addr],
+		m1[testBackendAddrs[2].Addr], m1[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	gator.Remove(testBalancerIDs[1])
+	bg.Remove(testBalancerIDs[1])
+	// Don't wait for SubConns to be removed after close, because they are only
+	// removed after close timeout.
+	for i := 0; i < 10; i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			t.Fatalf("Got request to shut down subconn %v, want no shut down subconn (because subconns were still in cache)", sc)
+		default:
+		}
+		time.Sleep(time.Millisecond)
+	}
+	// Test roundrobin on the with only sub-balancer0.
+	p2 := <-cc.NewPickerCh
+	want = []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[1].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p2)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	return gator, bg, cc, m1
+}
+
+// Test that if a sub-balancer is removed, and re-added within close timeout,
+// the subConns won't be re-created.
+func (s) TestBalancerGroup_locality_caching(t *testing.T) {
+	gator, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	// Turn down subconn for addr2, shouldn't get picker update because
+	// sub-balancer1 was removed.
+	addrToSC[testBackendAddrs[2].Addr].UpdateState(balancer.SubConnState{
+		ConnectivityState: connectivity.TransientFailure,
+		ConnectionError:   errors.New("test error"),
+	})
+	for i := 0; i < 10; i++ {
+		select {
+		case <-cc.NewPickerCh:
+			t.Fatalf("Got new picker, want no new picker (because the sub-balancer was removed)")
+		default:
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+
+	// Re-add sub-balancer-1, because subconns were in cache, no new subconns
+	// should be created. But a new picker will still be generated, with subconn
+	// states update to date.
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+
+	p3 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		addrToSC[testBackendAddrs[0].Addr], addrToSC[testBackendAddrs[0].Addr],
+		addrToSC[testBackendAddrs[1].Addr], addrToSC[testBackendAddrs[1].Addr],
+		// addr2 is down, b2 only has addr3 in READY state.
+		addrToSC[testBackendAddrs[3].Addr], addrToSC[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	for i := 0; i < 10; i++ {
+		select {
+		case <-cc.NewSubConnAddrsCh:
+			t.Fatalf("Got new subconn, want no new subconn (because subconns were still in cache)")
+		default:
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+}
+
+// Sub-balancers are put in cache when they are shut down. If balancer group is
+// closed within close timeout, all subconns should still be removed
+// immediately.
+func (s) TestBalancerGroup_locality_caching_close_group(t *testing.T) {
+	_, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	bg.Close()
+	// The balancer group is closed. The subconns should be shutdown immediately.
+	shutdownTimeout := time.After(time.Millisecond * 500)
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[0].Addr]: 1,
+		addrToSC[testBackendAddrs[1].Addr]: 1,
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-shutdownTimeout:
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+}
+
+// Sub-balancers in cache will be closed if not re-added within timeout, and
+// subConns will be shut down.
+func (s) TestBalancerGroup_locality_caching_not_read_within_timeout(t *testing.T) {
+	_, _, cc, addrToSC := initBalancerGroupForCachingTest(t, time.Second)
+
+	// The sub-balancer is not re-added within timeout. The subconns should be
+	// shut down.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-ctx.Done():
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+}
+
+// Wrap the rr builder, so it behaves the same, but has a different name.
+type noopBalancerBuilderWrapper struct {
+	balancer.Builder
+}
+
+func init() {
+	balancer.Register(&noopBalancerBuilderWrapper{Builder: rrBuilder})
+}
+
+func (*noopBalancerBuilderWrapper) Name() string {
+	return "noopBalancerBuilderWrapper"
+}
+
+// After removing a sub-balancer, re-add with same ID, but different balancer
+// builder. Old subconns should be shut down, and new subconns should be created.
+func (s) TestBalancerGroup_locality_caching_read_with_different_builder(t *testing.T) {
+	gator, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	// Re-add sub-balancer-1, but with a different balancer builder. The
+	// sub-balancer was still in cache, but can't be reused. This should cause
+	// old sub-balancer's subconns to be shut down immediately, and new
+	// subconns to be created.
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], &noopBalancerBuilderWrapper{rrBuilder})
+
+	// The cached sub-balancer should be closed, and the subconns should be
+	// shut down immediately.
+	shutdownTimeout := time.After(time.Millisecond * 500)
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-shutdownTimeout:
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[4:6]}})
+
+	newSCTimeout := time.After(time.Millisecond * 500)
+	scToAdd := map[string]int{
+		testBackendAddrs[4].Addr: 1,
+		testBackendAddrs[5].Addr: 1,
+	}
+	for i := 0; i < len(scToAdd); i++ {
+		select {
+		case addr := <-cc.NewSubConnAddrsCh:
+			c := scToAdd[addr[0].Addr]
+			if c == 0 {
+				t.Fatalf("Got newSubConn for %v when there's %d new expected", addr, c)
+			}
+			scToAdd[addr[0].Addr] = c - 1
+			sc := <-cc.NewSubConnCh
+			addrToSC[addr[0].Addr] = sc
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+		case <-newSCTimeout:
+			t.Fatalf("timeout waiting for subConns (from new sub-balancer) to be newed")
+		}
+	}
+
+	// Test roundrobin on the new picker.
+	p3 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		addrToSC[testBackendAddrs[0].Addr], addrToSC[testBackendAddrs[0].Addr],
+		addrToSC[testBackendAddrs[1].Addr], addrToSC[testBackendAddrs[1].Addr],
+		addrToSC[testBackendAddrs[4].Addr], addrToSC[testBackendAddrs[5].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+}
+
+// After removing a sub-balancer, it will be kept in cache. Make sure that this
+// sub-balancer's Close is called when the balancer group is closed.
+func (s) TestBalancerGroup_CloseStopsBalancerInCache(t *testing.T) {
+	const balancerName = "stub-TestBalancerGroup_check_close"
+	closed := make(chan struct{})
+	stub.Register(balancerName, stub.BalancerFuncs{Close: func(_ *stub.BalancerData) {
+		close(closed)
+	}})
+	builder := balancer.Get(balancerName)
+
+	gator, bg, _, _ := initBalancerGroupForCachingTest(t, time.Second)
+
+	// Add balancer, and remove
+	gator.Add(testBalancerIDs[2], 1)
+	bg.Add(testBalancerIDs[2], builder)
+	gator.Remove(testBalancerIDs[2])
+	bg.Remove(testBalancerIDs[2])
+
+	// Immediately close balancergroup, before the cache timeout.
+	bg.Close()
+
+	// Make sure the removed child balancer is closed eventually.
+	select {
+	case <-closed:
+	case <-time.After(time.Second * 2):
+		t.Fatalf("timeout waiting for the child balancer in cache to be closed")
+	}
+}
+
+// TestBalancerGroupBuildOptions verifies that the balancer.BuildOptions passed
+// to the balancergroup at creation time is passed to child policies.
+func (s) TestBalancerGroupBuildOptions(t *testing.T) {
+	const (
+		balancerName = "stubBalancer-TestBalancerGroupBuildOptions"
+		userAgent    = "ua"
+	)
+
+	// Setup the stub balancer such that we can read the build options passed to
+	// it in the UpdateClientConnState method.
+	bOpts := balancer.BuildOptions{
+		DialCreds:       insecure.NewCredentials(),
+		ChannelzParent:  channelz.RegisterChannel(nil, "test channel"),
+		CustomUserAgent: userAgent,
+	}
+	stub.Register(balancerName, stub.BalancerFuncs{
+		UpdateClientConnState: func(bd *stub.BalancerData, _ balancer.ClientConnState) error {
+			if bd.BuildOptions.DialCreds != bOpts.DialCreds || bd.BuildOptions.ChannelzParent != bOpts.ChannelzParent || bd.BuildOptions.CustomUserAgent != bOpts.CustomUserAgent {
+				return fmt.Errorf("buildOptions in child balancer: %v, want %v", bd, bOpts)
+			}
+			return nil
+		},
+	})
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       bOpts,
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	// Add the stub balancer build above as a child policy.
+	balancerBuilder := balancer.Get(balancerName)
+	bg.Add(testBalancerIDs[0], balancerBuilder)
+
+	// Send an empty clientConn state change. This should trigger the
+	// verification of the buildOptions being passed to the child policy.
+	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func (s) TestBalancerGroup_UpdateClientConnState_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	clientConnStateCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		UpdateClientConnState: func(_ *stub.BalancerData, _ balancer.ClientConnState) error {
+			clientConnStateCh <- struct{}{}
+			return nil
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+
+	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil {
+		t.Fatalf("Expected nil error, got %v", err)
+	}
+
+	select {
+	case <-clientConnStateCh:
+		t.Fatalf("UpdateClientConnState was called after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestBalancerGroup_ResolverError_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	resolveErrorCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ResolverError: func(_ *stub.BalancerData, _ error) {
+			resolveErrorCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+
+	bg.ResolverError(errors.New("test error"))
+
+	select {
+	case <-resolveErrorCh:
+		t.Fatalf("ResolverError was called on sub-balancer after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestBalancerExitIdleOne(t *testing.T) {
+	const balancerName = "stub-balancer-test-balancergroup-exit-idle-one"
+	exitIdleCh := make(chan struct{}, 1)
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(*stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	// Add the stub balancer build above as a child policy.
+	builder := balancer.Get(balancerName)
+	bg.Add(testBalancerIDs[0], builder)
+
+	// Call ExitIdleOne on the child policy.
+	bg.ExitIdleOne(testBalancerIDs[0])
+	select {
+	case <-time.After(time.Second):
+		t.Fatal("Timeout when waiting for ExitIdle to be invoked on child policy")
+	case <-exitIdleCh:
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdleOne_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{})
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			close(exitIdleCh)
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+	bg.ExitIdleOne(testBalancerIDs[0])
+
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdleOne called ExitIdle on sub-balancer after BalancerGroup was closed")
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdleOne_NonExistentID(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.ExitIdleOne("non-existent-id")
+
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdleOne called ExitIdle on wrong sub-balancer ID")
+	}
+}
+
+// TestBalancerGracefulSwitch tests the graceful switch functionality for a
+// child of the balancer group. At first, the child is configured as a round
+// robin load balancer, and thus should behave accordingly. The test then
+// gracefully switches this child to a custom type which only creates a SubConn
+// for the second passed in address and also only picks that created SubConn.
+// The new aggregated picker should reflect this change for the child.
+func (s) TestBalancerGracefulSwitch(t *testing.T) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: gator,
+		Logger:          nil,
+	})
+	gator.Add(testBalancerIDs[0], 1)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+
+	defer bg.Close()
+
+	m1 := make(map[string]balancer.SubConn)
+	scs := make(map[balancer.SubConn]bool)
+	for i := 0; i < 2; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		scs[sc] = true
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[1].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatal(err)
+	}
+
+	// The balancer type for testBalancersIDs[0] is currently Round Robin. Now,
+	// change it to a balancer that has separate behavior logically (creating
+	// SubConn for second address in address list and always picking that
+	// SubConn), and see if the downstream behavior reflects that change.
+	childPolicyName := t.Name()
+	stub.Register(childPolicyName, stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			ccs.ResolverState.Endpoints = ccs.ResolverState.Endpoints[1:]
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+	})
+	cfgJSON := json.RawMessage(fmt.Sprintf(`[{%q: {}}]`, t.Name()))
+	lbCfg, err := balancergroup.ParseConfig(cfgJSON)
+	if err != nil {
+		t.Fatalf("ParseConfig(%s) failed: %v", string(cfgJSON), err)
+	}
+	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{
+		ResolverState:  resolver.State{Endpoints: testBackendEndpoints[2:4]},
+		BalancerConfig: lbCfg,
+	}); err != nil {
+		t.Fatalf("error updating ClientConn state: %v", err)
+	}
+
+	addrs := <-cc.NewSubConnAddrsCh
+	if addrs[0].Addr != testBackendAddrs[3].Addr {
+		// Verifies forwarded to new created balancer, as the wrapped pick first
+		// balancer will delete first address.
+		t.Fatalf("newSubConn called with wrong address, want: %v, got : %v", testBackendAddrs[3].Addr, addrs[0].Addr)
+	}
+	sc := <-cc.NewSubConnCh
+
+	// Update the pick first balancers SubConn as CONNECTING. This will cause
+	// the pick first balancer to UpdateState() with CONNECTING, which shouldn't send
+	// a Picker update back, as the Graceful Switch process is not complete.
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer cancel()
+	select {
+	case <-cc.NewPickerCh:
+		t.Fatalf("No new picker should have been sent due to the Graceful Switch process not completing")
+	case <-ctx.Done():
+	}
+
+	// Update the pick first balancers SubConn as READY. This will cause
+	// the pick first balancer to UpdateState() with READY, which should send a
+	// Picker update back, as the Graceful Switch process is complete. This
+	// Picker should always pick the pick first's created SubConn which
+	// corresponds to address 3.
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	p2 := <-cc.NewPickerCh
+	pr, err := p2.Pick(balancer.PickInfo{})
+	if err != nil {
+		t.Fatalf("error picking: %v", err)
+	}
+	if pr.SubConn != sc {
+		t.Fatalf("picker.Pick(), want %v, got %v", sc, pr.SubConn)
+	}
+
+	// The Graceful Switch process completing for the child should cause the
+	// SubConns for the balancer being gracefully switched from to get deleted.
+	ctx, cancel = context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	for i := 0; i < 2; i++ {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("error waiting for Shutdown()")
+		case sc := <-cc.ShutdownSubConnCh:
+			// The SubConn shut down should have been one of the two created
+			// SubConns, and both should be deleted.
+			if ok := scs[sc]; ok {
+				delete(scs, sc)
+				continue
+			} else {
+				t.Fatalf("Shutdown called for wrong SubConn %v, want in %v", sc, scs)
+			}
+		}
+	}
+}
+
+func (s) TestBalancerExitIdle_All(t *testing.T) {
+	balancer1 := t.Name() + "-1"
+	balancer2 := t.Name() + "-2"
+
+	testID1, testID2 := testBalancerIDs[0], testBalancerIDs[1]
+
+	exitIdleCh1, exitIdleCh2 := make(chan struct{}, 1), make(chan struct{}, 1)
+
+	stub.Register(balancer1, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh1 <- struct{}{}
+		},
+	})
+
+	stub.Register(balancer2, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh2 <- struct{}{}
+		},
+	})
+
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	bg.Add(testID1, balancer.Get(balancer1))
+	bg.Add(testID2, balancer.Get(balancer2))
+
+	bg.ExitIdle()
+
+	errCh := make(chan error, 2)
+
+	go func() {
+		select {
+		case <-exitIdleCh1:
+			errCh <- nil
+		case <-time.After(defaultTestTimeout):
+			errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer1")
+		}
+	}()
+
+	go func() {
+		select {
+		case <-exitIdleCh2:
+			errCh <- nil
+		case <-time.After(defaultTestTimeout):
+			errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer2")
+		}
+	}()
+
+	for i := 0; i < 2; i++ {
+		if err := <-errCh; err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdle_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+	bg.ExitIdle()
+
+	select {
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdle was called on sub-balancer even after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+// Tests verifies that the RemoveImmediately method of balancergroup closes the
+// child balancer immediately, while the Remove method does not close the child
+// balancer immediately, when the cache is enabled.
+func (s) TestBalancerGroup_RemoveImmediately(t *testing.T) {
+	// Channels to track child balancer creation and closure.
+	childLBCreated := make(chan string, 1)
+	childLBClosed := make(chan string, 1)
+
+	childLBName1 := strings.ToLower(t.Name()) + "-child-1"
+	t.Logf("Registering a child balancer with name %q", childLBName1)
+	stub.Register(childLBName1, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName1
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName1
+		},
+	})
+
+	childLBName2 := strings.ToLower(t.Name()) + "-child-2"
+	t.Logf("Registering a child balancer with name %q", childLBName2)
+	stub.Register(childLBName2, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName2
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName2
+		},
+	})
+
+	t.Logf("Creating a balancergroup with cache enabled")
+	tcc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      tcc,
+		BuildOpts:               balancer.BuildOptions{},
+		SubBalancerCloseTimeout: defaultTestTimeout,
+	})
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-1")
+	bg.AddWithClientConn("child-1", childLBName1, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-2")
+	bg.AddWithClientConn("child-2", childLBName2, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with immediate effect", "child-1")
+	bg.RemoveImmediately("child-1")
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with caching", "child-2")
+	bg.Remove("child-2")
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	select {
+	case <-childLBClosed:
+		t.Fatalf("Child LB closed when expected to be cached")
+	case <-sCtx.Done():
+	}
+
+	t.Logf("Closing the balancergroup, which should close the second child balancer")
+	bg.Close()
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+}

internal/balancergroup/balancerstateaggregator.go

@@ -0,0 +1,37 @@
+/*
+ *
+ * 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 balancergroup
+
+import (
+	"google.golang.org/grpc/balancer"
+)
+
+// BalancerStateAggregator aggregates sub-picker and connectivity states into a
+// state.
+//
+// It takes care of merging sub-picker into one picker. The picking config is
+// passed directly from the parent to the aggregator implementation (instead
+// via balancer group).
+type BalancerStateAggregator interface {
+	// UpdateState updates the state of the id.
+	//
+	// It's up to the implementation whether this will trigger an update to the
+	// parent ClientConn.
+	UpdateState(id string, state balancer.State)
+}

internal/balancerload/load.go

@@ -0,0 +1,46 @@
+/*
+ * 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 balancerload defines APIs to parse server loads in trailers. The
+// parsed loads are sent to balancers in DoneInfo.
+package balancerload
+
+import (
+	"google.golang.org/grpc/metadata"
+)
+
+// Parser converts loads from metadata into a concrete type.
+type Parser interface {
+	// Parse parses loads from metadata.
+	Parse(md metadata.MD) any
+}
+
+var parser Parser
+
+// SetParser sets the load parser.
+//
+// Not mutex-protected, should be called before any gRPC functions.
+func SetParser(lr Parser) {
+	parser = lr
+}
+
+// Parse calls parser.Read().
+func Parse(md metadata.MD) any {
+	if parser == nil {
+		return nil
+	}
+	return parser.Parse(md)
+}

internal/binarylog/binarylog.go

@@ -0,0 +1,192 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package binarylog implementation binary logging as defined in
+// https://github.com/grpc/proposal/blob/master/A16-binary-logging.md.
+package binarylog
+
+import (
+	"fmt"
+	"os"
+
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal/grpcutil"
+)
+
+var grpclogLogger = grpclog.Component("binarylog")
+
+// Logger specifies MethodLoggers for method names with a Log call that
+// takes a context.
+//
+// This is used in the 1.0 release of gcp/observability, and thus must not be
+// deleted or changed.
+type Logger interface {
+	GetMethodLogger(methodName string) MethodLogger
+}
+
+// binLogger is the global binary logger for the binary. One of this should be
+// built at init time from the configuration (environment variable or flags).
+//
+// It is used to get a MethodLogger for each individual method.
+var binLogger Logger
+
+// SetLogger sets the binary logger.
+//
+// Only call this at init time.
+func SetLogger(l Logger) {
+	binLogger = l
+}
+
+// GetLogger gets the binary logger.
+//
+// Only call this at init time.
+func GetLogger() Logger {
+	return binLogger
+}
+
+// GetMethodLogger returns the MethodLogger for the given methodName.
+//
+// methodName should be in the format of "/service/method".
+//
+// Each MethodLogger returned by this method is a new instance. This is to
+// generate sequence id within the call.
+func GetMethodLogger(methodName string) MethodLogger {
+	if binLogger == nil {
+		return nil
+	}
+	return binLogger.GetMethodLogger(methodName)
+}
+
+func init() {
+	const envStr = "GRPC_BINARY_LOG_FILTER"
+	configStr := os.Getenv(envStr)
+	binLogger = NewLoggerFromConfigString(configStr)
+}
+
+// MethodLoggerConfig contains the setting for logging behavior of a method
+// logger. Currently, it contains the max length of header and message.
+type MethodLoggerConfig struct {
+	// Max length of header and message.
+	Header, Message uint64
+}
+
+// LoggerConfig contains the config for loggers to create method loggers.
+type LoggerConfig struct {
+	All      *MethodLoggerConfig
+	Services map[string]*MethodLoggerConfig
+	Methods  map[string]*MethodLoggerConfig
+
+	Blacklist map[string]struct{}
+}
+
+type logger struct {
+	config LoggerConfig
+}
+
+// NewLoggerFromConfig builds a logger with the given LoggerConfig.
+func NewLoggerFromConfig(config LoggerConfig) Logger {
+	return &logger{config: config}
+}
+
+// newEmptyLogger creates an empty logger. The map fields need to be filled in
+// using the set* functions.
+func newEmptyLogger() *logger {
+	return &logger{}
+}
+
+// Set method logger for "*".
+func (l *logger) setDefaultMethodLogger(ml *MethodLoggerConfig) error {
+	if l.config.All != nil {
+		return fmt.Errorf("conflicting global rules found")
+	}
+	l.config.All = ml
+	return nil
+}
+
+// Set method logger for "service/*".
+//
+// New MethodLogger with same service overrides the old one.
+func (l *logger) setServiceMethodLogger(service string, ml *MethodLoggerConfig) error {
+	if _, ok := l.config.Services[service]; ok {
+		return fmt.Errorf("conflicting service rules for service %v found", service)
+	}
+	if l.config.Services == nil {
+		l.config.Services = make(map[string]*MethodLoggerConfig)
+	}
+	l.config.Services[service] = ml
+	return nil
+}
+
+// Set method logger for "service/method".
+//
+// New MethodLogger with same method overrides the old one.
+func (l *logger) setMethodMethodLogger(method string, ml *MethodLoggerConfig) error {
+	if _, ok := l.config.Blacklist[method]; ok {
+		return fmt.Errorf("conflicting blacklist rules for method %v found", method)
+	}
+	if _, ok := l.config.Methods[method]; ok {
+		return fmt.Errorf("conflicting method rules for method %v found", method)
+	}
+	if l.config.Methods == nil {
+		l.config.Methods = make(map[string]*MethodLoggerConfig)
+	}
+	l.config.Methods[method] = ml
+	return nil
+}
+
+// Set blacklist method for "-service/method".
+func (l *logger) setBlacklist(method string) error {
+	if _, ok := l.config.Blacklist[method]; ok {
+		return fmt.Errorf("conflicting blacklist rules for method %v found", method)
+	}
+	if _, ok := l.config.Methods[method]; ok {
+		return fmt.Errorf("conflicting method rules for method %v found", method)
+	}
+	if l.config.Blacklist == nil {
+		l.config.Blacklist = make(map[string]struct{})
+	}
+	l.config.Blacklist[method] = struct{}{}
+	return nil
+}
+
+// getMethodLogger returns the MethodLogger for the given methodName.
+//
+// methodName should be in the format of "/service/method".
+//
+// Each MethodLogger returned by this method is a new instance. This is to
+// generate sequence id within the call.
+func (l *logger) GetMethodLogger(methodName string) MethodLogger {
+	s, m, err := grpcutil.ParseMethod(methodName)
+	if err != nil {
+		grpclogLogger.Infof("binarylogging: failed to parse %q: %v", methodName, err)
+		return nil
+	}
+	if ml, ok := l.config.Methods[s+"/"+m]; ok {
+		return NewTruncatingMethodLogger(ml.Header, ml.Message)
+	}
+	if _, ok := l.config.Blacklist[s+"/"+m]; ok {
+		return nil
+	}
+	if ml, ok := l.config.Services[s]; ok {
+		return NewTruncatingMethodLogger(ml.Header, ml.Message)
+	}
+	if l.config.All == nil {
+		return nil
+	}
+	return NewTruncatingMethodLogger(l.config.All.Header, l.config.All.Message)
+}

internal/binarylog/binarylog_test.go

@@ -0,0 +1,156 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package binarylog
+
+import (
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// Test that get method logger returns the one with the most exact match.
+func (s) TestGetMethodLogger(t *testing.T) {
+	testCases := []struct {
+		in       string
+		method   string
+		hdr, msg uint64
+	}{
+		// Global.
+		{
+			in:     "*{h:12;m:23}",
+			method: "/s/m",
+			hdr:    12, msg: 23,
+		},
+		// service/*.
+		{
+			in:     "*,s/*{h:12;m:23}",
+			method: "/s/m",
+			hdr:    12, msg: 23,
+		},
+		// Service/method.
+		{
+			in:     "*{h;m},s/m{h:12;m:23}",
+			method: "/s/m",
+			hdr:    12, msg: 23,
+		},
+		{
+			in:     "*{h;m},s/*{h:314;m},s/m{h:12;m:23}",
+			method: "/s/m",
+			hdr:    12, msg: 23,
+		},
+		{
+			in:     "*{h;m},s/*{h:12;m:23},s/m",
+			method: "/s/m",
+			hdr:    maxUInt, msg: maxUInt,
+		},
+
+		// service/*.
+		{
+			in:     "*{h;m},s/*{h:12;m:23},s/m1",
+			method: "/s/m",
+			hdr:    12, msg: 23,
+		},
+		{
+			in:     "*{h;m},s1/*,s/m{h:12;m:23}",
+			method: "/s/m",
+			hdr:    12, msg: 23,
+		},
+
+		// With black list.
+		{
+			in:     "*{h:12;m:23},-s/m1",
+			method: "/s/m",
+			hdr:    12, msg: 23,
+		},
+	}
+	for _, tc := range testCases {
+		l := NewLoggerFromConfigString(tc.in)
+		if l == nil {
+			t.Errorf("in: %q, failed to create logger from config string", tc.in)
+			continue
+		}
+		ml := l.GetMethodLogger(tc.method).(*TruncatingMethodLogger)
+		if ml == nil {
+			t.Errorf("in: %q, method logger is nil, want non-nil", tc.in)
+			continue
+		}
+		if ml.headerMaxLen != tc.hdr || ml.messageMaxLen != tc.msg {
+			t.Errorf("in: %q, want header: %v, message: %v, got header: %v, message: %v", tc.in, tc.hdr, tc.msg, ml.headerMaxLen, ml.messageMaxLen)
+		}
+	}
+}
+
+// expect method logger to be nil
+func (s) TestGetMethodLoggerOff(t *testing.T) {
+	testCases := []struct {
+		in     string
+		method string
+	}{
+		// method not specified.
+		{
+			in:     "s1/m",
+			method: "/s/m",
+		},
+		{
+			in:     "s/m1",
+			method: "/s/m",
+		},
+		{
+			in:     "s1/*",
+			method: "/s/m",
+		},
+		{
+			in:     "s1/*,s/m1",
+			method: "/s/m",
+		},
+
+		// blacklisted.
+		{
+			in:     "*,-s/m",
+			method: "/s/m",
+		},
+		{
+			in:     "s/*,-s/m",
+			method: "/s/m",
+		},
+		{
+			in:     "-s/m,s/*",
+			method: "/s/m",
+		},
+	}
+	for _, tc := range testCases {
+		l := NewLoggerFromConfigString(tc.in)
+		if l == nil {
+			t.Errorf("in: %q, failed to create logger from config string", tc.in)
+			continue
+		}
+		ml := l.GetMethodLogger(tc.method)
+		if ml != nil {
+			t.Errorf("in: %q, method logger is non-nil, want nil", tc.in)
+		}
+	}
+}

internal/binarylog/binarylog_testutil.go

@@ -0,0 +1,42 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// This file contains exported variables/functions that are exported for testing
+// only.
+//
+// An ideal way for this would be to put those in a *_test.go but in binarylog
+// package. But this doesn't work with staticcheck with go module. Error was:
+// "MdToMetadataProto not declared by package binarylog". This could be caused
+// by the way staticcheck looks for files for a certain package, which doesn't
+// support *_test.go files.
+//
+// Move those to binary_test.go when staticcheck is fixed.
+
+package binarylog
+
+var (
+	// AllLogger is a logger that logs all headers/messages for all RPCs. It's
+	// for testing only.
+	AllLogger = NewLoggerFromConfigString("*")
+	// MdToMetadataProto converts metadata to a binary logging proto message.
+	// It's for testing only.
+	MdToMetadataProto = mdToMetadataProto
+	// AddrToProto converts an address to a binary logging proto message. It's
+	// for testing only.
+	AddrToProto = addrToProto
+)

internal/binarylog/env_config.go

@@ -0,0 +1,208 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package binarylog
+
+import (
+	"errors"
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+// NewLoggerFromConfigString reads the string and build a logger. It can be used
+// to build a new logger and assign it to binarylog.Logger.
+//
+// Example filter config strings:
+//   - "" Nothing will be logged
+//   - "*" All headers and messages will be fully logged.
+//   - "*{h}" Only headers will be logged.
+//   - "*{m:256}" Only the first 256 bytes of each message will be logged.
+//   - "Foo/*" Logs every method in service Foo
+//   - "Foo/*,-Foo/Bar" Logs every method in service Foo except method /Foo/Bar
+//   - "Foo/*,Foo/Bar{m:256}" Logs the first 256 bytes of each message in method
+//     /Foo/Bar, logs all headers and messages in every other method in service
+//     Foo.
+//
+// If two configs exist for one certain method or service, the one specified
+// later overrides the previous config.
+func NewLoggerFromConfigString(s string) Logger {
+	if s == "" {
+		return nil
+	}
+	l := newEmptyLogger()
+	methods := strings.Split(s, ",")
+	for _, method := range methods {
+		if err := l.fillMethodLoggerWithConfigString(method); err != nil {
+			grpclogLogger.Warningf("failed to parse binary log config: %v", err)
+			return nil
+		}
+	}
+	return l
+}
+
+// fillMethodLoggerWithConfigString parses config, creates TruncatingMethodLogger and adds
+// it to the right map in the logger.
+func (l *logger) fillMethodLoggerWithConfigString(config string) error {
+	// "" is invalid.
+	if config == "" {
+		return errors.New("empty string is not a valid method binary logging config")
+	}
+
+	// "-service/method", blacklist, no * or {} allowed.
+	if config[0] == '-' {
+		s, m, suffix, err := parseMethodConfigAndSuffix(config[1:])
+		if err != nil {
+			return fmt.Errorf("invalid config: %q, %v", config, err)
+		}
+		if m == "*" {
+			return fmt.Errorf("invalid config: %q, %v", config, "* not allowed in blacklist config")
+		}
+		if suffix != "" {
+			return fmt.Errorf("invalid config: %q, %v", config, "header/message limit not allowed in blacklist config")
+		}
+		if err := l.setBlacklist(s + "/" + m); err != nil {
+			return fmt.Errorf("invalid config: %v", err)
+		}
+		return nil
+	}
+
+	// "*{h:256;m:256}"
+	if config[0] == '*' {
+		hdr, msg, err := parseHeaderMessageLengthConfig(config[1:])
+		if err != nil {
+			return fmt.Errorf("invalid config: %q, %v", config, err)
+		}
+		if err := l.setDefaultMethodLogger(&MethodLoggerConfig{Header: hdr, Message: msg}); err != nil {
+			return fmt.Errorf("invalid config: %v", err)
+		}
+		return nil
+	}
+
+	s, m, suffix, err := parseMethodConfigAndSuffix(config)
+	if err != nil {
+		return fmt.Errorf("invalid config: %q, %v", config, err)
+	}
+	hdr, msg, err := parseHeaderMessageLengthConfig(suffix)
+	if err != nil {
+		return fmt.Errorf("invalid header/message length config: %q, %v", suffix, err)
+	}
+	if m == "*" {
+		if err := l.setServiceMethodLogger(s, &MethodLoggerConfig{Header: hdr, Message: msg}); err != nil {
+			return fmt.Errorf("invalid config: %v", err)
+		}
+	} else {
+		if err := l.setMethodMethodLogger(s+"/"+m, &MethodLoggerConfig{Header: hdr, Message: msg}); err != nil {
+			return fmt.Errorf("invalid config: %v", err)
+		}
+	}
+	return nil
+}
+
+const (
+	// TODO: this const is only used by env_config now. But could be useful for
+	// other config. Move to binarylog.go if necessary.
+	maxUInt = ^uint64(0)
+
+	// For "p.s/m" plus any suffix. Suffix will be parsed again. See test for
+	// expected output.
+	longMethodConfigRegexpStr = `^([\w./]+)/((?:\w+)|[*])(.+)?$`
+
+	// For suffix from above, "{h:123,m:123}". See test for expected output.
+	optionalLengthRegexpStr      = `(?::(\d+))?` // Optional ":123".
+	headerConfigRegexpStr        = `^{h` + optionalLengthRegexpStr + `}$`
+	messageConfigRegexpStr       = `^{m` + optionalLengthRegexpStr + `}$`
+	headerMessageConfigRegexpStr = `^{h` + optionalLengthRegexpStr + `;m` + optionalLengthRegexpStr + `}$`
+)
+
+var (
+	longMethodConfigRegexp    = regexp.MustCompile(longMethodConfigRegexpStr)
+	headerConfigRegexp        = regexp.MustCompile(headerConfigRegexpStr)
+	messageConfigRegexp       = regexp.MustCompile(messageConfigRegexpStr)
+	headerMessageConfigRegexp = regexp.MustCompile(headerMessageConfigRegexpStr)
+)
+
+// Turn "service/method{h;m}" into "service", "method", "{h;m}".
+func parseMethodConfigAndSuffix(c string) (service, method, suffix string, _ error) {
+	// Regexp result:
+	//
+	// in:  "p.s/m{h:123,m:123}",
+	// out: []string{"p.s/m{h:123,m:123}", "p.s", "m", "{h:123,m:123}"},
+	match := longMethodConfigRegexp.FindStringSubmatch(c)
+	if match == nil {
+		return "", "", "", fmt.Errorf("%q contains invalid substring", c)
+	}
+	service = match[1]
+	method = match[2]
+	suffix = match[3]
+	return
+}
+
+// Turn "{h:123;m:345}" into 123, 345.
+//
+// Return maxUInt if length is unspecified.
+func parseHeaderMessageLengthConfig(c string) (hdrLenStr, msgLenStr uint64, err error) {
+	if c == "" {
+		return maxUInt, maxUInt, nil
+	}
+	// Header config only.
+	if match := headerConfigRegexp.FindStringSubmatch(c); match != nil {
+		if s := match[1]; s != "" {
+			hdrLenStr, err = strconv.ParseUint(s, 10, 64)
+			if err != nil {
+				return 0, 0, fmt.Errorf("failed to convert %q to uint", s)
+			}
+			return hdrLenStr, 0, nil
+		}
+		return maxUInt, 0, nil
+	}
+
+	// Message config only.
+	if match := messageConfigRegexp.FindStringSubmatch(c); match != nil {
+		if s := match[1]; s != "" {
+			msgLenStr, err = strconv.ParseUint(s, 10, 64)
+			if err != nil {
+				return 0, 0, fmt.Errorf("failed to convert %q to uint", s)
+			}
+			return 0, msgLenStr, nil
+		}
+		return 0, maxUInt, nil
+	}
+
+	// Header and message config both.
+	if match := headerMessageConfigRegexp.FindStringSubmatch(c); match != nil {
+		// Both hdr and msg are specified, but one or two of them might be empty.
+		hdrLenStr = maxUInt
+		msgLenStr = maxUInt
+		if s := match[1]; s != "" {
+			hdrLenStr, err = strconv.ParseUint(s, 10, 64)
+			if err != nil {
+				return 0, 0, fmt.Errorf("failed to convert %q to uint", s)
+			}
+		}
+		if s := match[2]; s != "" {
+			msgLenStr, err = strconv.ParseUint(s, 10, 64)
+			if err != nil {
+				return 0, 0, fmt.Errorf("failed to convert %q to uint", s)
+			}
+		}
+		return hdrLenStr, msgLenStr, nil
+	}
+	return 0, 0, fmt.Errorf("%q contains invalid substring", c)
+}

internal/binarylog/env_config_test.go

@@ -0,0 +1,478 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package binarylog
+
+import (
+	"fmt"
+	"testing"
+)
+
+// This tests that when multiple configs are specified, all methods loggers will
+// be set correctly. Correctness of each logger is covered by other unit tests.
+func (s) TestNewLoggerFromConfigString(t *testing.T) {
+	const (
+		s1     = "s1"
+		m1     = "m1"
+		m2     = "m2"
+		fullM1 = s1 + "/" + m1
+		fullM2 = s1 + "/" + m2
+	)
+	c := fmt.Sprintf("*{h:1;m:2},%s{h},%s{m},%s{h;m}", s1+"/*", fullM1, fullM2)
+	l := NewLoggerFromConfigString(c).(*logger)
+
+	if l.config.All.Header != 1 || l.config.All.Message != 2 {
+		t.Errorf("l.config.All = %#v, want headerLen: 1, messageLen: 2", l.config.All)
+	}
+
+	if ml, ok := l.config.Services[s1]; ok {
+		if ml.Header != maxUInt || ml.Message != 0 {
+			t.Errorf("want maxUInt header, 0 message, got header: %v, message: %v", ml.Header, ml.Message)
+		}
+	} else {
+		t.Errorf("service/* is not set")
+	}
+
+	if ml, ok := l.config.Methods[fullM1]; ok {
+		if ml.Header != 0 || ml.Message != maxUInt {
+			t.Errorf("want 0 header, maxUInt message, got header: %v, message: %v", ml.Header, ml.Message)
+		}
+	} else {
+		t.Errorf("service/method{h} is not set")
+	}
+
+	if ml, ok := l.config.Methods[fullM2]; ok {
+		if ml.Header != maxUInt || ml.Message != maxUInt {
+			t.Errorf("want maxUInt header, maxUInt message, got header: %v, message: %v", ml.Header, ml.Message)
+		}
+	} else {
+		t.Errorf("service/method{h;m} is not set")
+	}
+}
+
+func (s) TestNewLoggerFromConfigStringInvalid(t *testing.T) {
+	testCases := []string{
+		"",
+		"*{}",
+		"s/m,*{}",
+		"s/m,s/m{a}",
+
+		// Duplicate rules.
+		"s/m,-s/m",
+		"-s/m,s/m",
+		"s/m,s/m",
+		"s/m,s/m{h:1;m:1}",
+		"s/m{h:1;m:1},s/m",
+		"-s/m,-s/m",
+		"s/*,s/*{h:1;m:1}",
+		"*,*{h:1;m:1}",
+	}
+	for _, tc := range testCases {
+		l := NewLoggerFromConfigString(tc)
+		if l != nil {
+			t.Errorf("With config %q, want logger %v, got %v", tc, nil, l)
+		}
+	}
+}
+
+func (s) TestParseMethodConfigAndSuffix(t *testing.T) {
+	testCases := []struct {
+		in, service, method, suffix string
+	}{
+		{
+			in:      "p.s/m",
+			service: "p.s", method: "m", suffix: "",
+		},
+		{
+			in:      "p.s/m{h,m}",
+			service: "p.s", method: "m", suffix: "{h,m}",
+		},
+		{
+			in:      "p.s/*",
+			service: "p.s", method: "*", suffix: "",
+		},
+		{
+			in:      "p.s/*{h,m}",
+			service: "p.s", method: "*", suffix: "{h,m}",
+		},
+
+		// invalid suffix will be detected by another function.
+		{
+			in:      "p.s/m{invalidsuffix}",
+			service: "p.s", method: "m", suffix: "{invalidsuffix}",
+		},
+		{
+			in:      "p.s/*{invalidsuffix}",
+			service: "p.s", method: "*", suffix: "{invalidsuffix}",
+		},
+		{
+			in:      "s/m*",
+			service: "s", method: "m", suffix: "*",
+		},
+		{
+			in:      "s/*m",
+			service: "s", method: "*", suffix: "m",
+		},
+		{
+			in:      "s/**",
+			service: "s", method: "*", suffix: "*",
+		},
+	}
+	for _, tc := range testCases {
+		t.Logf("testing parseMethodConfigAndSuffix(%q)", tc.in)
+		s, m, suffix, err := parseMethodConfigAndSuffix(tc.in)
+		if err != nil {
+			t.Errorf("returned error %v, want nil", err)
+			continue
+		}
+		if s != tc.service {
+			t.Errorf("service = %q, want %q", s, tc.service)
+		}
+		if m != tc.method {
+			t.Errorf("method = %q, want %q", m, tc.method)
+		}
+		if suffix != tc.suffix {
+			t.Errorf("suffix = %q, want %q", suffix, tc.suffix)
+		}
+	}
+}
+
+func (s) TestParseMethodConfigAndSuffixInvalid(t *testing.T) {
+	testCases := []string{
+		"*/m",
+		"*/m{}",
+	}
+	for _, tc := range testCases {
+		s, m, suffix, err := parseMethodConfigAndSuffix(tc)
+		if err == nil {
+			t.Errorf("Parsing %q got nil error with %q, %q, %q, want non-nil error", tc, s, m, suffix)
+		}
+	}
+}
+
+func (s) TestParseHeaderMessageLengthConfig(t *testing.T) {
+	testCases := []struct {
+		in       string
+		hdr, msg uint64
+	}{
+		{
+			in:  "",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h}",
+			hdr: maxUInt, msg: 0,
+		},
+		{
+			in:  "{h:314}",
+			hdr: 314, msg: 0,
+		},
+		{
+			in:  "{m}",
+			hdr: 0, msg: maxUInt,
+		},
+		{
+			in:  "{m:213}",
+			hdr: 0, msg: 213,
+		},
+		{
+			in:  "{h;m}",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h:314;m}",
+			hdr: 314, msg: maxUInt,
+		},
+		{
+			in:  "{h;m:213}",
+			hdr: maxUInt, msg: 213,
+		},
+		{
+			in:  "{h:314;m:213}",
+			hdr: 314, msg: 213,
+		},
+	}
+	for _, tc := range testCases {
+		t.Logf("testing parseHeaderMessageLengthConfig(%q)", tc.in)
+		hdr, msg, err := parseHeaderMessageLengthConfig(tc.in)
+		if err != nil {
+			t.Errorf("returned error %v, want nil", err)
+			continue
+		}
+		if hdr != tc.hdr {
+			t.Errorf("header length = %v, want %v", hdr, tc.hdr)
+		}
+		if msg != tc.msg {
+			t.Errorf("message length = %v, want %v", msg, tc.msg)
+		}
+	}
+}
+func (s) TestParseHeaderMessageLengthConfigInvalid(t *testing.T) {
+	testCases := []string{
+		"{}",
+		"{h;a}",
+		"{h;m;b}",
+	}
+	for _, tc := range testCases {
+		_, _, err := parseHeaderMessageLengthConfig(tc)
+		if err == nil {
+			t.Errorf("Parsing %q got nil error, want non-nil error", tc)
+		}
+	}
+}
+
+func (s) TestFillMethodLoggerWithConfigStringBlacklist(t *testing.T) {
+	testCases := []string{
+		"p.s/m",
+		"service/method",
+	}
+	for _, tc := range testCases {
+		c := "-" + tc
+		t.Logf("testing fillMethodLoggerWithConfigString(%q)", c)
+		l := newEmptyLogger()
+		if err := l.fillMethodLoggerWithConfigString(c); err != nil {
+			t.Errorf("returned err %v, want nil", err)
+			continue
+		}
+		_, ok := l.config.Blacklist[tc]
+		if !ok {
+			t.Errorf("blacklist[%q] is not set", tc)
+		}
+	}
+}
+
+func (s) TestFillMethodLoggerWithConfigStringGlobal(t *testing.T) {
+	testCases := []struct {
+		in       string
+		hdr, msg uint64
+	}{
+		{
+			in:  "",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h}",
+			hdr: maxUInt, msg: 0,
+		},
+		{
+			in:  "{h:314}",
+			hdr: 314, msg: 0,
+		},
+		{
+			in:  "{m}",
+			hdr: 0, msg: maxUInt,
+		},
+		{
+			in:  "{m:213}",
+			hdr: 0, msg: 213,
+		},
+		{
+			in:  "{h;m}",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h:314;m}",
+			hdr: 314, msg: maxUInt,
+		},
+		{
+			in:  "{h;m:213}",
+			hdr: maxUInt, msg: 213,
+		},
+		{
+			in:  "{h:314;m:213}",
+			hdr: 314, msg: 213,
+		},
+	}
+	for _, tc := range testCases {
+		c := "*" + tc.in
+		t.Logf("testing fillMethodLoggerWithConfigString(%q)", c)
+		l := newEmptyLogger()
+		if err := l.fillMethodLoggerWithConfigString(c); err != nil {
+			t.Errorf("returned err %v, want nil", err)
+			continue
+		}
+		if l.config.All == nil {
+			t.Errorf("l.config.All is not set")
+			continue
+		}
+		if hdr := l.config.All.Header; hdr != tc.hdr {
+			t.Errorf("header length = %v, want %v", hdr, tc.hdr)
+
+		}
+		if msg := l.config.All.Message; msg != tc.msg {
+			t.Errorf("message length = %v, want %v", msg, tc.msg)
+		}
+	}
+}
+
+func (s) TestFillMethodLoggerWithConfigStringPerService(t *testing.T) {
+	testCases := []struct {
+		in       string
+		hdr, msg uint64
+	}{
+		{
+			in:  "",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h}",
+			hdr: maxUInt, msg: 0,
+		},
+		{
+			in:  "{h:314}",
+			hdr: 314, msg: 0,
+		},
+		{
+			in:  "{m}",
+			hdr: 0, msg: maxUInt,
+		},
+		{
+			in:  "{m:213}",
+			hdr: 0, msg: 213,
+		},
+		{
+			in:  "{h;m}",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h:314;m}",
+			hdr: 314, msg: maxUInt,
+		},
+		{
+			in:  "{h;m:213}",
+			hdr: maxUInt, msg: 213,
+		},
+		{
+			in:  "{h:314;m:213}",
+			hdr: 314, msg: 213,
+		},
+	}
+	const serviceName = "service"
+	for _, tc := range testCases {
+		c := serviceName + "/*" + tc.in
+		t.Logf("testing fillMethodLoggerWithConfigString(%q)", c)
+		l := newEmptyLogger()
+		if err := l.fillMethodLoggerWithConfigString(c); err != nil {
+			t.Errorf("returned err %v, want nil", err)
+			continue
+		}
+		ml, ok := l.config.Services[serviceName]
+		if !ok {
+			t.Errorf("l.service[%q] is not set", serviceName)
+			continue
+		}
+		if hdr := ml.Header; hdr != tc.hdr {
+			t.Errorf("header length = %v, want %v", hdr, tc.hdr)
+
+		}
+		if msg := ml.Message; msg != tc.msg {
+			t.Errorf("message length = %v, want %v", msg, tc.msg)
+		}
+	}
+}
+
+func (s) TestFillMethodLoggerWithConfigStringPerMethod(t *testing.T) {
+	testCases := []struct {
+		in       string
+		hdr, msg uint64
+	}{
+		{
+			in:  "",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h}",
+			hdr: maxUInt, msg: 0,
+		},
+		{
+			in:  "{h:314}",
+			hdr: 314, msg: 0,
+		},
+		{
+			in:  "{m}",
+			hdr: 0, msg: maxUInt,
+		},
+		{
+			in:  "{m:213}",
+			hdr: 0, msg: 213,
+		},
+		{
+			in:  "{h;m}",
+			hdr: maxUInt, msg: maxUInt,
+		},
+		{
+			in:  "{h:314;m}",
+			hdr: 314, msg: maxUInt,
+		},
+		{
+			in:  "{h;m:213}",
+			hdr: maxUInt, msg: 213,
+		},
+		{
+			in:  "{h:314;m:213}",
+			hdr: 314, msg: 213,
+		},
+	}
+	const (
+		serviceName    = "service"
+		methodName     = "method"
+		fullMethodName = serviceName + "/" + methodName
+	)
+	for _, tc := range testCases {
+		c := fullMethodName + tc.in
+		t.Logf("testing fillMethodLoggerWithConfigString(%q)", c)
+		l := newEmptyLogger()
+		if err := l.fillMethodLoggerWithConfigString(c); err != nil {
+			t.Errorf("returned err %v, want nil", err)
+			continue
+		}
+		ml, ok := l.config.Methods[fullMethodName]
+		if !ok {
+			t.Errorf("l.config.Methods[%q] is not set", fullMethodName)
+			continue
+		}
+		if hdr := ml.Header; hdr != tc.hdr {
+			t.Errorf("header length = %v, want %v", hdr, tc.hdr)
+
+		}
+		if msg := ml.Message; msg != tc.msg {
+			t.Errorf("message length = %v, want %v", msg, tc.msg)
+		}
+	}
+}
+
+func (s) TestFillMethodLoggerWithConfigStringInvalid(t *testing.T) {
+	testCases := []string{
+		"",
+		"{}",
+		"p.s/m{}",
+		"p.s/m{a}",
+		"p.s/m*",
+		"p.s/**",
+		"*/m",
+
+		"-p.s/*",
+		"-p.s/m{h}",
+	}
+	l := &logger{}
+	for _, tc := range testCases {
+		if err := l.fillMethodLoggerWithConfigString(tc); err == nil {
+			t.Errorf("fillMethodLoggerWithConfigString(%q) returned nil error, want non-nil", tc)
+		}
+	}
+}

internal/binarylog/method_logger.go

@@ -0,0 +1,446 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package binarylog
+
+import (
+	"context"
+	"net"
+	"strings"
+	"sync/atomic"
+	"time"
+
+	binlogpb "google.golang.org/grpc/binarylog/grpc_binarylog_v1"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/durationpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
+)
+
+type callIDGenerator struct {
+	id uint64
+}
+
+func (g *callIDGenerator) next() uint64 {
+	id := atomic.AddUint64(&g.id, 1)
+	return id
+}
+
+// reset is for testing only, and doesn't need to be thread safe.
+func (g *callIDGenerator) reset() {
+	g.id = 0
+}
+
+var idGen callIDGenerator
+
+// MethodLogger is the sub-logger for each method.
+//
+// This is used in the 1.0 release of gcp/observability, and thus must not be
+// deleted or changed.
+type MethodLogger interface {
+	Log(context.Context, LogEntryConfig)
+}
+
+// TruncatingMethodLogger is a method logger that truncates headers and messages
+// based on configured fields.
+type TruncatingMethodLogger struct {
+	headerMaxLen, messageMaxLen uint64
+
+	callID          uint64
+	idWithinCallGen *callIDGenerator
+
+	sink Sink // TODO(blog): make this pluggable.
+}
+
+// NewTruncatingMethodLogger returns a new truncating method logger.
+//
+// This is used in the 1.0 release of gcp/observability, and thus must not be
+// deleted or changed.
+func NewTruncatingMethodLogger(h, m uint64) *TruncatingMethodLogger {
+	return &TruncatingMethodLogger{
+		headerMaxLen:  h,
+		messageMaxLen: m,
+
+		callID:          idGen.next(),
+		idWithinCallGen: &callIDGenerator{},
+
+		sink: DefaultSink, // TODO(blog): make it pluggable.
+	}
+}
+
+// Build is an internal only method for building the proto message out of the
+// input event. It's made public to enable other library to reuse as much logic
+// in TruncatingMethodLogger as possible.
+func (ml *TruncatingMethodLogger) Build(c LogEntryConfig) *binlogpb.GrpcLogEntry {
+	m := c.toProto()
+	timestamp := timestamppb.Now()
+	m.Timestamp = timestamp
+	m.CallId = ml.callID
+	m.SequenceIdWithinCall = ml.idWithinCallGen.next()
+
+	switch pay := m.Payload.(type) {
+	case *binlogpb.GrpcLogEntry_ClientHeader:
+		m.PayloadTruncated = ml.truncateMetadata(pay.ClientHeader.GetMetadata())
+	case *binlogpb.GrpcLogEntry_ServerHeader:
+		m.PayloadTruncated = ml.truncateMetadata(pay.ServerHeader.GetMetadata())
+	case *binlogpb.GrpcLogEntry_Message:
+		m.PayloadTruncated = ml.truncateMessage(pay.Message)
+	}
+	return m
+}
+
+// Log creates a proto binary log entry, and logs it to the sink.
+func (ml *TruncatingMethodLogger) Log(_ context.Context, c LogEntryConfig) {
+	ml.sink.Write(ml.Build(c))
+}
+
+func (ml *TruncatingMethodLogger) truncateMetadata(mdPb *binlogpb.Metadata) (truncated bool) {
+	if ml.headerMaxLen == maxUInt {
+		return false
+	}
+	var (
+		bytesLimit = ml.headerMaxLen
+		index      int
+	)
+	// At the end of the loop, index will be the first entry where the total
+	// size is greater than the limit:
+	//
+	// len(entry[:index]) <= ml.hdr && len(entry[:index+1]) > ml.hdr.
+	for ; index < len(mdPb.Entry); index++ {
+		entry := mdPb.Entry[index]
+		if entry.Key == "grpc-trace-bin" {
+			// "grpc-trace-bin" is a special key. It's kept in the log entry,
+			// but not counted towards the size limit.
+			continue
+		}
+		currentEntryLen := uint64(len(entry.GetKey())) + uint64(len(entry.GetValue()))
+		if currentEntryLen > bytesLimit {
+			break
+		}
+		bytesLimit -= currentEntryLen
+	}
+	truncated = index < len(mdPb.Entry)
+	mdPb.Entry = mdPb.Entry[:index]
+	return truncated
+}
+
+func (ml *TruncatingMethodLogger) truncateMessage(msgPb *binlogpb.Message) (truncated bool) {
+	if ml.messageMaxLen == maxUInt {
+		return false
+	}
+	if ml.messageMaxLen >= uint64(len(msgPb.Data)) {
+		return false
+	}
+	msgPb.Data = msgPb.Data[:ml.messageMaxLen]
+	return true
+}
+
+// LogEntryConfig represents the configuration for binary log entry.
+//
+// This is used in the 1.0 release of gcp/observability, and thus must not be
+// deleted or changed.
+type LogEntryConfig interface {
+	toProto() *binlogpb.GrpcLogEntry
+}
+
+// ClientHeader configs the binary log entry to be a ClientHeader entry.
+type ClientHeader struct {
+	OnClientSide bool
+	Header       metadata.MD
+	MethodName   string
+	Authority    string
+	Timeout      time.Duration
+	// PeerAddr is required only when it's on server side.
+	PeerAddr net.Addr
+}
+
+func (c *ClientHeader) toProto() *binlogpb.GrpcLogEntry {
+	// This function doesn't need to set all the fields (e.g. seq ID). The Log
+	// function will set the fields when necessary.
+	clientHeader := &binlogpb.ClientHeader{
+		Metadata:   mdToMetadataProto(c.Header),
+		MethodName: c.MethodName,
+		Authority:  c.Authority,
+	}
+	if c.Timeout > 0 {
+		clientHeader.Timeout = durationpb.New(c.Timeout)
+	}
+	ret := &binlogpb.GrpcLogEntry{
+		Type: binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_HEADER,
+		Payload: &binlogpb.GrpcLogEntry_ClientHeader{
+			ClientHeader: clientHeader,
+		},
+	}
+	if c.OnClientSide {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_CLIENT
+	} else {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_SERVER
+	}
+	if c.PeerAddr != nil {
+		ret.Peer = addrToProto(c.PeerAddr)
+	}
+	return ret
+}
+
+// ServerHeader configs the binary log entry to be a ServerHeader entry.
+type ServerHeader struct {
+	OnClientSide bool
+	Header       metadata.MD
+	// PeerAddr is required only when it's on client side.
+	PeerAddr net.Addr
+}
+
+func (c *ServerHeader) toProto() *binlogpb.GrpcLogEntry {
+	ret := &binlogpb.GrpcLogEntry{
+		Type: binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_HEADER,
+		Payload: &binlogpb.GrpcLogEntry_ServerHeader{
+			ServerHeader: &binlogpb.ServerHeader{
+				Metadata: mdToMetadataProto(c.Header),
+			},
+		},
+	}
+	if c.OnClientSide {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_CLIENT
+	} else {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_SERVER
+	}
+	if c.PeerAddr != nil {
+		ret.Peer = addrToProto(c.PeerAddr)
+	}
+	return ret
+}
+
+// ClientMessage configs the binary log entry to be a ClientMessage entry.
+type ClientMessage struct {
+	OnClientSide bool
+	// Message can be a proto.Message or []byte. Other messages formats are not
+	// supported.
+	Message any
+}
+
+func (c *ClientMessage) toProto() *binlogpb.GrpcLogEntry {
+	var (
+		data []byte
+		err  error
+	)
+	if m, ok := c.Message.(proto.Message); ok {
+		data, err = proto.Marshal(m)
+		if err != nil {
+			grpclogLogger.Infof("binarylogging: failed to marshal proto message: %v", err)
+		}
+	} else if b, ok := c.Message.([]byte); ok {
+		data = b
+	} else {
+		grpclogLogger.Infof("binarylogging: message to log is neither proto.message nor []byte")
+	}
+	ret := &binlogpb.GrpcLogEntry{
+		Type: binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_MESSAGE,
+		Payload: &binlogpb.GrpcLogEntry_Message{
+			Message: &binlogpb.Message{
+				Length: uint32(len(data)),
+				Data:   data,
+			},
+		},
+	}
+	if c.OnClientSide {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_CLIENT
+	} else {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_SERVER
+	}
+	return ret
+}
+
+// ServerMessage configs the binary log entry to be a ServerMessage entry.
+type ServerMessage struct {
+	OnClientSide bool
+	// Message can be a proto.Message or []byte. Other messages formats are not
+	// supported.
+	Message any
+}
+
+func (c *ServerMessage) toProto() *binlogpb.GrpcLogEntry {
+	var (
+		data []byte
+		err  error
+	)
+	if m, ok := c.Message.(proto.Message); ok {
+		data, err = proto.Marshal(m)
+		if err != nil {
+			grpclogLogger.Infof("binarylogging: failed to marshal proto message: %v", err)
+		}
+	} else if b, ok := c.Message.([]byte); ok {
+		data = b
+	} else {
+		grpclogLogger.Infof("binarylogging: message to log is neither proto.message nor []byte")
+	}
+	ret := &binlogpb.GrpcLogEntry{
+		Type: binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_MESSAGE,
+		Payload: &binlogpb.GrpcLogEntry_Message{
+			Message: &binlogpb.Message{
+				Length: uint32(len(data)),
+				Data:   data,
+			},
+		},
+	}
+	if c.OnClientSide {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_CLIENT
+	} else {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_SERVER
+	}
+	return ret
+}
+
+// ClientHalfClose configs the binary log entry to be a ClientHalfClose entry.
+type ClientHalfClose struct {
+	OnClientSide bool
+}
+
+func (c *ClientHalfClose) toProto() *binlogpb.GrpcLogEntry {
+	ret := &binlogpb.GrpcLogEntry{
+		Type:    binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_HALF_CLOSE,
+		Payload: nil, // No payload here.
+	}
+	if c.OnClientSide {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_CLIENT
+	} else {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_SERVER
+	}
+	return ret
+}
+
+// ServerTrailer configs the binary log entry to be a ServerTrailer entry.
+type ServerTrailer struct {
+	OnClientSide bool
+	Trailer      metadata.MD
+	// Err is the status error.
+	Err error
+	// PeerAddr is required only when it's on client side and the RPC is trailer
+	// only.
+	PeerAddr net.Addr
+}
+
+func (c *ServerTrailer) toProto() *binlogpb.GrpcLogEntry {
+	st, ok := status.FromError(c.Err)
+	if !ok {
+		grpclogLogger.Info("binarylogging: error in trailer is not a status error")
+	}
+	var (
+		detailsBytes []byte
+		err          error
+	)
+	stProto := st.Proto()
+	if stProto != nil && len(stProto.Details) != 0 {
+		detailsBytes, err = proto.Marshal(stProto)
+		if err != nil {
+			grpclogLogger.Infof("binarylogging: failed to marshal status proto: %v", err)
+		}
+	}
+	ret := &binlogpb.GrpcLogEntry{
+		Type: binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_TRAILER,
+		Payload: &binlogpb.GrpcLogEntry_Trailer{
+			Trailer: &binlogpb.Trailer{
+				Metadata:      mdToMetadataProto(c.Trailer),
+				StatusCode:    uint32(st.Code()),
+				StatusMessage: st.Message(),
+				StatusDetails: detailsBytes,
+			},
+		},
+	}
+	if c.OnClientSide {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_CLIENT
+	} else {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_SERVER
+	}
+	if c.PeerAddr != nil {
+		ret.Peer = addrToProto(c.PeerAddr)
+	}
+	return ret
+}
+
+// Cancel configs the binary log entry to be a Cancel entry.
+type Cancel struct {
+	OnClientSide bool
+}
+
+func (c *Cancel) toProto() *binlogpb.GrpcLogEntry {
+	ret := &binlogpb.GrpcLogEntry{
+		Type:    binlogpb.GrpcLogEntry_EVENT_TYPE_CANCEL,
+		Payload: nil,
+	}
+	if c.OnClientSide {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_CLIENT
+	} else {
+		ret.Logger = binlogpb.GrpcLogEntry_LOGGER_SERVER
+	}
+	return ret
+}
+
+// metadataKeyOmit returns whether the metadata entry with this key should be
+// omitted.
+func metadataKeyOmit(key string) bool {
+	switch key {
+	case "lb-token", ":path", ":authority", "content-encoding", "content-type", "user-agent", "te", "authorization":
+		return true
+	case "grpc-trace-bin": // grpc-trace-bin is special because it's visible to users.
+		return false
+	}
+	return strings.HasPrefix(key, "grpc-")
+}
+
+func mdToMetadataProto(md metadata.MD) *binlogpb.Metadata {
+	ret := &binlogpb.Metadata{}
+	for k, vv := range md {
+		if metadataKeyOmit(k) {
+			continue
+		}
+		for _, v := range vv {
+			ret.Entry = append(ret.Entry,
+				&binlogpb.MetadataEntry{
+					Key:   k,
+					Value: []byte(v),
+				},
+			)
+		}
+	}
+	return ret
+}
+
+func addrToProto(addr net.Addr) *binlogpb.Address {
+	ret := &binlogpb.Address{}
+	switch a := addr.(type) {
+	case *net.TCPAddr:
+		if a.IP.To4() != nil {
+			ret.Type = binlogpb.Address_TYPE_IPV4
+		} else if a.IP.To16() != nil {
+			ret.Type = binlogpb.Address_TYPE_IPV6
+		} else {
+			ret.Type = binlogpb.Address_TYPE_UNKNOWN
+			// Do not set address and port fields.
+			break
+		}
+		ret.Address = a.IP.String()
+		ret.IpPort = uint32(a.Port)
+	case *net.UnixAddr:
+		ret.Type = binlogpb.Address_TYPE_UNIX
+		ret.Address = a.String()
+	default:
+		ret.Type = binlogpb.Address_TYPE_UNKNOWN
+	}
+	return ret
+}

internal/binarylog/method_logger_test.go

@@ -0,0 +1,562 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package binarylog
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"net"
+	"testing"
+	"time"
+
+	binlogpb "google.golang.org/grpc/binarylog/grpc_binarylog_v1"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/durationpb"
+)
+
+const defaultTestTimeout = 10 * time.Second
+
+func (s) TestLog(t *testing.T) {
+	idGen.reset()
+	ml := NewTruncatingMethodLogger(10, 10)
+	// Set sink to testing buffer.
+	buf := bytes.NewBuffer(nil)
+	ml.sink = newWriterSink(buf)
+
+	addr := "1.2.3.4"
+	port := 790
+	tcpAddr, _ := net.ResolveTCPAddr("tcp", fmt.Sprintf("%v:%d", addr, port))
+	addr6 := "2001:1db8:85a3::8a2e:1370:7334"
+	port6 := 796
+	tcpAddr6, _ := net.ResolveTCPAddr("tcp", fmt.Sprintf("[%v]:%d", addr6, port6))
+
+	testProtoMsg := &binlogpb.Message{
+		Length: 1,
+		Data:   []byte{'a'},
+	}
+	testProtoBytes, _ := proto.Marshal(testProtoMsg)
+
+	testCases := []struct {
+		config LogEntryConfig
+		want   *binlogpb.GrpcLogEntry
+	}{
+		{
+			config: &ClientHeader{
+				OnClientSide: false,
+				Header: map[string][]string{
+					"a": {"b", "bb"},
+				},
+				MethodName: "testservice/testmethod",
+				Authority:  "test.service.io",
+				Timeout:    2*time.Second + 3*time.Nanosecond,
+				PeerAddr:   tcpAddr,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_HEADER,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_SERVER,
+				Payload: &binlogpb.GrpcLogEntry_ClientHeader{
+					ClientHeader: &binlogpb.ClientHeader{
+						Metadata: &binlogpb.Metadata{
+							Entry: []*binlogpb.MetadataEntry{
+								{Key: "a", Value: []byte{'b'}},
+								{Key: "a", Value: []byte{'b', 'b'}},
+							},
+						},
+						MethodName: "testservice/testmethod",
+						Authority:  "test.service.io",
+						Timeout: &durationpb.Duration{
+							Seconds: 2,
+							Nanos:   3,
+						},
+					},
+				},
+				PayloadTruncated: false,
+				Peer: &binlogpb.Address{
+					Type:    binlogpb.Address_TYPE_IPV4,
+					Address: addr,
+					IpPort:  uint32(port),
+				},
+			},
+		},
+		{
+			config: &ClientHeader{
+				OnClientSide: false,
+				MethodName:   "testservice/testmethod",
+				Authority:    "test.service.io",
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_HEADER,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_SERVER,
+				Payload: &binlogpb.GrpcLogEntry_ClientHeader{
+					ClientHeader: &binlogpb.ClientHeader{
+						Metadata:   &binlogpb.Metadata{},
+						MethodName: "testservice/testmethod",
+						Authority:  "test.service.io",
+					},
+				},
+				PayloadTruncated: false,
+			},
+		},
+		{
+			config: &ServerHeader{
+				OnClientSide: true,
+				Header: map[string][]string{
+					"a": {"b", "bb"},
+				},
+				PeerAddr: tcpAddr6,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_HEADER,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_CLIENT,
+				Payload: &binlogpb.GrpcLogEntry_ServerHeader{
+					ServerHeader: &binlogpb.ServerHeader{
+						Metadata: &binlogpb.Metadata{
+							Entry: []*binlogpb.MetadataEntry{
+								{Key: "a", Value: []byte{'b'}},
+								{Key: "a", Value: []byte{'b', 'b'}},
+							},
+						},
+					},
+				},
+				PayloadTruncated: false,
+				Peer: &binlogpb.Address{
+					Type:    binlogpb.Address_TYPE_IPV6,
+					Address: addr6,
+					IpPort:  uint32(port6),
+				},
+			},
+		},
+		{
+			config: &ClientMessage{
+				OnClientSide: true,
+				Message:      testProtoMsg,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_MESSAGE,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_CLIENT,
+				Payload: &binlogpb.GrpcLogEntry_Message{
+					Message: &binlogpb.Message{
+						Length: uint32(len(testProtoBytes)),
+						Data:   testProtoBytes,
+					},
+				},
+				PayloadTruncated: false,
+				Peer:             nil,
+			},
+		},
+		{
+			config: &ServerMessage{
+				OnClientSide: false,
+				Message:      testProtoMsg,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_MESSAGE,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_SERVER,
+				Payload: &binlogpb.GrpcLogEntry_Message{
+					Message: &binlogpb.Message{
+						Length: uint32(len(testProtoBytes)),
+						Data:   testProtoBytes,
+					},
+				},
+				PayloadTruncated: false,
+				Peer:             nil,
+			},
+		},
+		{
+			config: &ClientHalfClose{
+				OnClientSide: false,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_HALF_CLOSE,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_SERVER,
+				Payload:              nil,
+				PayloadTruncated:     false,
+				Peer:                 nil,
+			},
+		},
+		{
+			config: &ServerTrailer{
+				OnClientSide: true,
+				Err:          status.Errorf(codes.Unavailable, "test"),
+				PeerAddr:     tcpAddr,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_TRAILER,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_CLIENT,
+				Payload: &binlogpb.GrpcLogEntry_Trailer{
+					Trailer: &binlogpb.Trailer{
+						Metadata:      &binlogpb.Metadata{},
+						StatusCode:    uint32(codes.Unavailable),
+						StatusMessage: "test",
+						StatusDetails: nil,
+					},
+				},
+				PayloadTruncated: false,
+				Peer: &binlogpb.Address{
+					Type:    binlogpb.Address_TYPE_IPV4,
+					Address: addr,
+					IpPort:  uint32(port),
+				},
+			},
+		},
+		{ // Err is nil, Log OK status.
+			config: &ServerTrailer{
+				OnClientSide: true,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_TRAILER,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_CLIENT,
+				Payload: &binlogpb.GrpcLogEntry_Trailer{
+					Trailer: &binlogpb.Trailer{
+						Metadata:      &binlogpb.Metadata{},
+						StatusCode:    uint32(codes.OK),
+						StatusMessage: "",
+						StatusDetails: nil,
+					},
+				},
+				PayloadTruncated: false,
+				Peer:             nil,
+			},
+		},
+		{
+			config: &Cancel{
+				OnClientSide: true,
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_CANCEL,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_CLIENT,
+				Payload:              nil,
+				PayloadTruncated:     false,
+				Peer:                 nil,
+			},
+		},
+		// gRPC headers should be omitted.
+		{
+			config: &ClientHeader{
+				OnClientSide: false,
+				Header: map[string][]string{
+					"lb-token":         {"to be omitted"},
+					":path":            {"to be omitted"},
+					":authority":       {"to be omitted"},
+					"content-encoding": {"to be omitted"},
+					"content-type":     {"to be omitted"},
+					"user-agent":       {"to be omitted"},
+					"te":               {"to be omitted"},
+					"authorization":    {"Bearer to-be-omitted"},
+					"grpc-reserved":    {"to be omitted"},
+
+					"a": {"b", "bb"},
+				},
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_CLIENT_HEADER,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_SERVER,
+				Payload: &binlogpb.GrpcLogEntry_ClientHeader{
+					ClientHeader: &binlogpb.ClientHeader{
+						Metadata: &binlogpb.Metadata{
+							Entry: []*binlogpb.MetadataEntry{
+								{Key: "a", Value: []byte{'b'}},
+								{Key: "a", Value: []byte{'b', 'b'}},
+							},
+						},
+					},
+				},
+				PayloadTruncated: false,
+			},
+		},
+		{
+			config: &ServerHeader{
+				OnClientSide: true,
+				Header: map[string][]string{
+					"lb-token":         {"to be omitted"},
+					":path":            {"to be omitted"},
+					":authority":       {"to be omitted"},
+					"content-encoding": {"to be omitted"},
+					"content-type":     {"to be omitted"},
+					"user-agent":       {"to be omitted"},
+					"te":               {"to be omitted"},
+					"authorization":    {"Bearer to-be-omitted"},
+					"grpc-reserved":    {"to be omitted"},
+
+					"a": {"b", "bb"},
+				},
+			},
+			want: &binlogpb.GrpcLogEntry{
+				Timestamp:            nil,
+				CallId:               1,
+				SequenceIdWithinCall: 0,
+				Type:                 binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_HEADER,
+				Logger:               binlogpb.GrpcLogEntry_LOGGER_CLIENT,
+				Payload: &binlogpb.GrpcLogEntry_ServerHeader{
+					ServerHeader: &binlogpb.ServerHeader{
+						Metadata: &binlogpb.Metadata{
+							Entry: []*binlogpb.MetadataEntry{
+								{Key: "a", Value: []byte{'b'}},
+								{Key: "a", Value: []byte{'b', 'b'}},
+							},
+						},
+					},
+				},
+				PayloadTruncated: false,
+			},
+		},
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	for i, tc := range testCases {
+		buf.Reset()
+		tc.want.SequenceIdWithinCall = uint64(i + 1)
+		ml.Log(ctx, tc.config)
+		inSink := new(binlogpb.GrpcLogEntry)
+		if err := proto.Unmarshal(buf.Bytes()[4:], inSink); err != nil {
+			t.Errorf("failed to unmarshal bytes in sink to proto: %v", err)
+			continue
+		}
+		inSink.Timestamp = nil // Strip timestamp before comparing.
+		if !proto.Equal(inSink, tc.want) {
+			t.Errorf("Log(%+v), in sink: %+v, want %+v", tc.config, inSink, tc.want)
+		}
+	}
+}
+
+func (s) TestTruncateMetadataNotTruncated(t *testing.T) {
+	testCases := []struct {
+		ml   *TruncatingMethodLogger
+		mpPb *binlogpb.Metadata
+	}{
+		{
+			ml: NewTruncatingMethodLogger(maxUInt, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1}},
+				},
+			},
+		},
+		{
+			ml: NewTruncatingMethodLogger(2, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1}},
+				},
+			},
+		},
+		{
+			ml: NewTruncatingMethodLogger(1, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: nil},
+				},
+			},
+		},
+		{
+			ml: NewTruncatingMethodLogger(2, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1, 1}},
+				},
+			},
+		},
+		{
+			ml: NewTruncatingMethodLogger(2, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1}},
+					{Key: "", Value: []byte{1}},
+				},
+			},
+		},
+		// "grpc-trace-bin" is kept in log but not counted towards the size
+		// limit.
+		{
+			ml: NewTruncatingMethodLogger(1, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1}},
+					{Key: "grpc-trace-bin", Value: []byte("some.trace.key")},
+				},
+			},
+		},
+	}
+
+	for i, tc := range testCases {
+		truncated := tc.ml.truncateMetadata(tc.mpPb)
+		if truncated {
+			t.Errorf("test case %v, returned truncated, want not truncated", i)
+		}
+	}
+}
+
+func (s) TestTruncateMetadataTruncated(t *testing.T) {
+	testCases := []struct {
+		ml   *TruncatingMethodLogger
+		mpPb *binlogpb.Metadata
+
+		entryLen int
+	}{
+		{
+			ml: NewTruncatingMethodLogger(2, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1, 1, 1}},
+				},
+			},
+			entryLen: 0,
+		},
+		{
+			ml: NewTruncatingMethodLogger(2, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1}},
+					{Key: "", Value: []byte{1}},
+					{Key: "", Value: []byte{1}},
+				},
+			},
+			entryLen: 2,
+		},
+		{
+			ml: NewTruncatingMethodLogger(2, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1, 1}},
+					{Key: "", Value: []byte{1}},
+				},
+			},
+			entryLen: 1,
+		},
+		{
+			ml: NewTruncatingMethodLogger(2, maxUInt),
+			mpPb: &binlogpb.Metadata{
+				Entry: []*binlogpb.MetadataEntry{
+					{Key: "", Value: []byte{1}},
+					{Key: "", Value: []byte{1, 1}},
+				},
+			},
+			entryLen: 1,
+		},
+	}
+
+	for i, tc := range testCases {
+		truncated := tc.ml.truncateMetadata(tc.mpPb)
+		if !truncated {
+			t.Errorf("test case %v, returned not truncated, want truncated", i)
+			continue
+		}
+		if len(tc.mpPb.Entry) != tc.entryLen {
+			t.Errorf("test case %v, entry length: %v, want: %v", i, len(tc.mpPb.Entry), tc.entryLen)
+		}
+	}
+}
+
+func (s) TestTruncateMessageNotTruncated(t *testing.T) {
+	testCases := []struct {
+		ml    *TruncatingMethodLogger
+		msgPb *binlogpb.Message
+	}{
+		{
+			ml: NewTruncatingMethodLogger(maxUInt, maxUInt),
+			msgPb: &binlogpb.Message{
+				Data: []byte{1},
+			},
+		},
+		{
+			ml: NewTruncatingMethodLogger(maxUInt, 3),
+			msgPb: &binlogpb.Message{
+				Data: []byte{1, 1},
+			},
+		},
+		{
+			ml: NewTruncatingMethodLogger(maxUInt, 2),
+			msgPb: &binlogpb.Message{
+				Data: []byte{1, 1},
+			},
+		},
+	}
+
+	for i, tc := range testCases {
+		truncated := tc.ml.truncateMessage(tc.msgPb)
+		if truncated {
+			t.Errorf("test case %v, returned truncated, want not truncated", i)
+		}
+	}
+}
+
+func (s) TestTruncateMessageTruncated(t *testing.T) {
+	testCases := []struct {
+		ml    *TruncatingMethodLogger
+		msgPb *binlogpb.Message
+
+		oldLength uint32
+	}{
+		{
+			ml: NewTruncatingMethodLogger(maxUInt, 2),
+			msgPb: &binlogpb.Message{
+				Length: 3,
+				Data:   []byte{1, 1, 1},
+			},
+			oldLength: 3,
+		},
+	}
+
+	for i, tc := range testCases {
+		truncated := tc.ml.truncateMessage(tc.msgPb)
+		if !truncated {
+			t.Errorf("test case %v, returned not truncated, want truncated", i)
+			continue
+		}
+		if len(tc.msgPb.Data) != int(tc.ml.messageMaxLen) {
+			t.Errorf("test case %v, message length: %v, want: %v", i, len(tc.msgPb.Data), tc.ml.messageMaxLen)
+		}
+		if tc.msgPb.Length != tc.oldLength {
+			t.Errorf("test case %v, message.Length field: %v, want: %v", i, tc.msgPb.Length, tc.oldLength)
+		}
+	}
+}

internal/binarylog/regexp_test.go

@@ -0,0 +1,179 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package binarylog
+
+import (
+	"reflect"
+	"testing"
+)
+
+func (s) TestLongMethodConfigRegexp(t *testing.T) {
+	testCases := []struct {
+		in  string
+		out []string
+	}{
+		{in: "", out: nil},
+		{in: "*/m", out: nil},
+
+		{
+			in:  "p.s/m{}",
+			out: []string{"p.s/m{}", "p.s", "m", "{}"},
+		},
+
+		{
+			in:  "p.s/m",
+			out: []string{"p.s/m", "p.s", "m", ""},
+		},
+		{
+			in:  "p.s/m{h}",
+			out: []string{"p.s/m{h}", "p.s", "m", "{h}"},
+		},
+		{
+			in:  "p.s/m{m}",
+			out: []string{"p.s/m{m}", "p.s", "m", "{m}"},
+		},
+		{
+			in:  "p.s/m{h:123}",
+			out: []string{"p.s/m{h:123}", "p.s", "m", "{h:123}"},
+		},
+		{
+			in:  "p.s/m{m:123}",
+			out: []string{"p.s/m{m:123}", "p.s", "m", "{m:123}"},
+		},
+		{
+			in:  "p.s/m{h:123,m:123}",
+			out: []string{"p.s/m{h:123,m:123}", "p.s", "m", "{h:123,m:123}"},
+		},
+
+		{
+			in:  "p.s/*",
+			out: []string{"p.s/*", "p.s", "*", ""},
+		},
+		{
+			in:  "p.s/*{h}",
+			out: []string{"p.s/*{h}", "p.s", "*", "{h}"},
+		},
+
+		{
+			in:  "s/m*",
+			out: []string{"s/m*", "s", "m", "*"},
+		},
+		{
+			in:  "s/**",
+			out: []string{"s/**", "s", "*", "*"},
+		},
+	}
+	for _, tc := range testCases {
+		match := longMethodConfigRegexp.FindStringSubmatch(tc.in)
+		if !reflect.DeepEqual(match, tc.out) {
+			t.Errorf("in: %q, out: %q, want: %q", tc.in, match, tc.out)
+		}
+	}
+}
+
+func (s) TestHeaderConfigRegexp(t *testing.T) {
+	testCases := []struct {
+		in  string
+		out []string
+	}{
+		{in: "{}", out: nil},
+		{in: "{a:b}", out: nil},
+		{in: "{m:123}", out: nil},
+		{in: "{h:123;m:123}", out: nil},
+
+		{
+			in:  "{h}",
+			out: []string{"{h}", ""},
+		},
+		{
+			in:  "{h:123}",
+			out: []string{"{h:123}", "123"},
+		},
+	}
+	for _, tc := range testCases {
+		match := headerConfigRegexp.FindStringSubmatch(tc.in)
+		if !reflect.DeepEqual(match, tc.out) {
+			t.Errorf("in: %q, out: %q, want: %q", tc.in, match, tc.out)
+		}
+	}
+}
+
+func (s) TestMessageConfigRegexp(t *testing.T) {
+	testCases := []struct {
+		in  string
+		out []string
+	}{
+		{in: "{}", out: nil},
+		{in: "{a:b}", out: nil},
+		{in: "{h:123}", out: nil},
+		{in: "{h:123;m:123}", out: nil},
+
+		{
+			in:  "{m}",
+			out: []string{"{m}", ""},
+		},
+		{
+			in:  "{m:123}",
+			out: []string{"{m:123}", "123"},
+		},
+	}
+	for _, tc := range testCases {
+		match := messageConfigRegexp.FindStringSubmatch(tc.in)
+		if !reflect.DeepEqual(match, tc.out) {
+			t.Errorf("in: %q, out: %q, want: %q", tc.in, match, tc.out)
+		}
+	}
+}
+
+func (s) TestHeaderMessageConfigRegexp(t *testing.T) {
+	testCases := []struct {
+		in  string
+		out []string
+	}{
+		{in: "{}", out: nil},
+		{in: "{a:b}", out: nil},
+		{in: "{h}", out: nil},
+		{in: "{h:123}", out: nil},
+		{in: "{m}", out: nil},
+		{in: "{m:123}", out: nil},
+
+		{
+			in:  "{h;m}",
+			out: []string{"{h;m}", "", ""},
+		},
+		{
+			in:  "{h:123;m}",
+			out: []string{"{h:123;m}", "123", ""},
+		},
+		{
+			in:  "{h;m:123}",
+			out: []string{"{h;m:123}", "", "123"},
+		},
+		{
+			in:  "{h:123;m:123}",
+			out: []string{"{h:123;m:123}", "123", "123"},
+		},
+	}
+	for _, tc := range testCases {
+		match := headerMessageConfigRegexp.FindStringSubmatch(tc.in)
+		if !reflect.DeepEqual(match, tc.out) {
+			t.Errorf("in: %q, out: %q, want: %q", tc.in, match, tc.out)
+		}
+	}
+}

internal/binarylog/sink.go

@@ -0,0 +1,170 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package binarylog
+
+import (
+	"bufio"
+	"encoding/binary"
+	"io"
+	"sync"
+	"time"
+
+	binlogpb "google.golang.org/grpc/binarylog/grpc_binarylog_v1"
+	"google.golang.org/protobuf/proto"
+)
+
+var (
+	// DefaultSink is the sink where the logs will be written to. It's exported
+	// for the binarylog package to update.
+	DefaultSink Sink = &noopSink{} // TODO(blog): change this default (file in /tmp).
+)
+
+// Sink writes log entry into the binary log sink.
+//
+// sink is a copy of the exported binarylog.Sink, to avoid circular dependency.
+type Sink interface {
+	// Write will be called to write the log entry into the sink.
+	//
+	// It should be thread-safe so it can be called in parallel.
+	Write(*binlogpb.GrpcLogEntry) error
+	// Close will be called when the Sink is replaced by a new Sink.
+	Close() error
+}
+
+type noopSink struct{}
+
+func (ns *noopSink) Write(*binlogpb.GrpcLogEntry) error { return nil }
+func (ns *noopSink) Close() error                       { return nil }
+
+// newWriterSink creates a binary log sink with the given writer.
+//
+// Write() marshals the proto message and writes it to the given writer. Each
+// message is prefixed with a 4 byte big endian unsigned integer as the length.
+//
+// No buffer is done, Close() doesn't try to close the writer.
+func newWriterSink(w io.Writer) Sink {
+	return &writerSink{out: w}
+}
+
+type writerSink struct {
+	out io.Writer
+}
+
+func (ws *writerSink) Write(e *binlogpb.GrpcLogEntry) error {
+	b, err := proto.Marshal(e)
+	if err != nil {
+		grpclogLogger.Errorf("binary logging: failed to marshal proto message: %v", err)
+		return err
+	}
+	hdr := make([]byte, 4)
+	binary.BigEndian.PutUint32(hdr, uint32(len(b)))
+	if _, err := ws.out.Write(hdr); err != nil {
+		return err
+	}
+	if _, err := ws.out.Write(b); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (ws *writerSink) Close() error { return nil }
+
+type bufferedSink struct {
+	mu             sync.Mutex
+	closer         io.Closer
+	out            Sink          // out is built on buf.
+	buf            *bufio.Writer // buf is kept for flush.
+	flusherStarted bool
+
+	writeTicker *time.Ticker
+	done        chan struct{}
+}
+
+func (fs *bufferedSink) Write(e *binlogpb.GrpcLogEntry) error {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+	if !fs.flusherStarted {
+		// Start the write loop when Write is called.
+		fs.startFlushGoroutine()
+		fs.flusherStarted = true
+	}
+	if err := fs.out.Write(e); err != nil {
+		return err
+	}
+	return nil
+}
+
+const (
+	bufFlushDuration = 60 * time.Second
+)
+
+func (fs *bufferedSink) startFlushGoroutine() {
+	fs.writeTicker = time.NewTicker(bufFlushDuration)
+	go func() {
+		for {
+			select {
+			case <-fs.done:
+				return
+			case <-fs.writeTicker.C:
+			}
+			fs.mu.Lock()
+			if err := fs.buf.Flush(); err != nil {
+				grpclogLogger.Warningf("failed to flush to Sink: %v", err)
+			}
+			fs.mu.Unlock()
+		}
+	}()
+}
+
+func (fs *bufferedSink) Close() error {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+	if fs.writeTicker != nil {
+		fs.writeTicker.Stop()
+	}
+	close(fs.done)
+	if err := fs.buf.Flush(); err != nil {
+		grpclogLogger.Warningf("failed to flush to Sink: %v", err)
+	}
+	if err := fs.closer.Close(); err != nil {
+		grpclogLogger.Warningf("failed to close the underlying WriterCloser: %v", err)
+	}
+	if err := fs.out.Close(); err != nil {
+		grpclogLogger.Warningf("failed to close the Sink: %v", err)
+	}
+	return nil
+}
+
+// NewBufferedSink creates a binary log sink with the given WriteCloser.
+//
+// Write() marshals the proto message and writes it to the given writer. Each
+// message is prefixed with a 4 byte big endian unsigned integer as the length.
+//
+// Content is kept in a buffer, and is flushed every 60 seconds.
+//
+// Close closes the WriteCloser.
+func NewBufferedSink(o io.WriteCloser) Sink {
+	bufW := bufio.NewWriter(o)
+	return &bufferedSink{
+		closer: o,
+		out:    newWriterSink(bufW),
+		buf:    bufW,
+		done:   make(chan struct{}),
+	}
+}

internal/buffer/unbounded.go

@@ -0,0 +1,115 @@
+/*
+ * 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 buffer provides an implementation of an unbounded buffer.
+package buffer
+
+import (
+	"errors"
+	"sync"
+)
+
+// Unbounded is an implementation of an unbounded buffer which does not use
+// extra goroutines. This is typically used for passing updates from one entity
+// to another within gRPC.
+//
+// To avoid extra memory allocations and type assertions, using any on
+// performance-critical code paths is discouraged. Use concrete types wherever
+// possible when instantiating Unbounded.
+//
+// All methods on this type are thread-safe and don't block on anything except
+// the underlying mutex used for synchronization.
+type Unbounded[T any] struct {
+	c       chan T
+	closed  bool
+	closing bool
+	mu      sync.Mutex
+	backlog []T
+}
+
+// NewUnbounded returns a new instance of Unbounded.
+func NewUnbounded[T any]() *Unbounded[T] {
+	return &Unbounded[T]{c: make(chan T, 1)}
+}
+
+var errBufferClosed = errors.New("Put called on closed buffer.Unbounded")
+
+// Put adds t to the unbounded buffer.
+func (b *Unbounded[T]) Put(t T) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.closing {
+		return errBufferClosed
+	}
+	if len(b.backlog) == 0 {
+		select {
+		case b.c <- t:
+			return nil
+		default:
+		}
+	}
+	b.backlog = append(b.backlog, t)
+	return nil
+}
+
+// Load sends the earliest buffered data, if any, onto the read channel returned
+// by Get(). Users are expected to call this every time they successfully read a
+// value from the read channel.
+func (b *Unbounded[T]) Load() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if len(b.backlog) > 0 {
+		select {
+		case b.c <- b.backlog[0]:
+			var zero T
+			b.backlog[0] = zero
+			b.backlog = b.backlog[1:]
+		default:
+		}
+	} else if b.closing && !b.closed {
+		b.closed = true
+		close(b.c)
+	}
+}
+
+// Get returns a read channel on which values added to the buffer, via Put(),
+// are sent on.
+//
+// Upon reading a value from this channel, users are expected to call Load() to
+// send the next buffered value onto the channel if there is any.
+//
+// If the unbounded buffer is closed, the read channel returned by this method
+// is closed after all data is drained.
+func (b *Unbounded[T]) Get() <-chan T {
+	return b.c
+}
+
+// Close closes the unbounded buffer. No subsequent data may be Put(), and the
+// channel returned from Get() will be closed after all the data is read and
+// Load() is called for the final time.
+func (b *Unbounded[T]) Close() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.closing {
+		return
+	}
+	b.closing = true
+	if len(b.backlog) == 0 {
+		b.closed = true
+		close(b.c)
+	}
+}

internal/buffer/unbounded_test.go

@@ -0,0 +1,148 @@
+/*
+ * 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 buffer
+
+import (
+	"reflect"
+	"sort"
+	"sync"
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+const (
+	numWriters = 10
+	numWrites  = 10
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// wantReads contains the set of values expected to be read by the reader
+// goroutine in the tests.
+var wantReads []int
+
+func init() {
+	for i := 0; i < numWriters; i++ {
+		for j := 0; j < numWrites; j++ {
+			wantReads = append(wantReads, i)
+		}
+	}
+}
+
+// TestSingleWriter starts one reader and one writer goroutine and makes sure
+// that the reader gets all the values added to the buffer by the writer.
+func (s) TestSingleWriter(t *testing.T) {
+	ub := NewUnbounded[int]()
+	reads := []int{}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		ch := ub.Get()
+		for i := 0; i < numWriters*numWrites; i++ {
+			r := <-ch
+			reads = append(reads, r)
+			ub.Load()
+		}
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numWriters; i++ {
+			for j := 0; j < numWrites; j++ {
+				ub.Put(i)
+			}
+		}
+	}()
+
+	wg.Wait()
+	if !reflect.DeepEqual(reads, wantReads) {
+		t.Errorf("reads: %#v, wantReads: %#v", reads, wantReads)
+	}
+}
+
+// TestMultipleWriters starts multiple writers and one reader goroutine and
+// makes sure that the reader gets all the data written by all writers.
+func (s) TestMultipleWriters(t *testing.T) {
+	ub := NewUnbounded[int]()
+	reads := []int{}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		ch := ub.Get()
+		for i := 0; i < numWriters*numWrites; i++ {
+			r := <-ch
+			reads = append(reads, r)
+			ub.Load()
+		}
+	}()
+
+	wg.Add(numWriters)
+	for i := 0; i < numWriters; i++ {
+		go func(index int) {
+			defer wg.Done()
+			for j := 0; j < numWrites; j++ {
+				ub.Put(index)
+			}
+		}(i)
+	}
+
+	wg.Wait()
+	sort.Ints(reads)
+	if !reflect.DeepEqual(reads, wantReads) {
+		t.Errorf("reads: %#v, wantReads: %#v", reads, wantReads)
+	}
+}
+
+// TestClose closes the buffer and makes sure that nothing is sent after the
+// buffer is closed.
+func (s) TestClose(t *testing.T) {
+	ub := NewUnbounded[int]()
+	if err := ub.Put(1); err != nil {
+		t.Fatalf("Unbounded.Put() = %v; want nil", err)
+	}
+	ub.Close()
+	if err := ub.Put(1); err == nil {
+		t.Fatalf("Unbounded.Put() = <nil>; want non-nil error")
+	}
+	if v, ok := <-ub.Get(); !ok {
+		t.Errorf("Unbounded.Get() = %v, %v, want %v, %v", v, ok, 1, true)
+	}
+	if err := ub.Put(1); err == nil {
+		t.Fatalf("Unbounded.Put() = <nil>; want non-nil error")
+	}
+	ub.Load()
+	if v, ok := <-ub.Get(); ok {
+		t.Errorf("Unbounded.Get() = %v, want closed channel", v)
+	}
+	if err := ub.Put(1); err == nil {
+		t.Fatalf("Unbounded.Put() = <nil>; want non-nil error")
+	}
+	ub.Close() // ignored
+}

internal/cache/timeoutCache.go

@@ -0,0 +1,151 @@
+/*
+ * 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 cache implements caches to be used in gRPC.
+package cache
+
+import (
+	"sync"
+	"time"
+)
+
+type cacheEntry struct {
+	item any
+	// Note that to avoid deadlocks (potentially caused by lock ordering),
+	// callback can only be called without holding cache's mutex.
+	callback func()
+	timer    *time.Timer
+	// deleted is set to true in Remove() when the call to timer.Stop() fails.
+	// This can happen when the timer in the cache entry fires around the same
+	// time that timer.stop() is called in Remove().
+	deleted bool
+}
+
+// TimeoutCache is a cache with items to be deleted after a timeout.
+type TimeoutCache struct {
+	mu      sync.Mutex
+	timeout time.Duration
+	cache   map[any]*cacheEntry
+}
+
+// NewTimeoutCache creates a TimeoutCache with the given timeout.
+func NewTimeoutCache(timeout time.Duration) *TimeoutCache {
+	return &TimeoutCache{
+		timeout: timeout,
+		cache:   make(map[any]*cacheEntry),
+	}
+}
+
+// Add adds an item to the cache, with the specified callback to be called when
+// the item is removed from the cache upon timeout. If the item is removed from
+// the cache using a call to Remove before the timeout expires, the callback
+// will not be called.
+//
+// If the Add was successful, it returns (newly added item, true). If there is
+// an existing entry for the specified key, the cache entry is not be updated
+// with the specified item and it returns (existing item, false).
+func (c *TimeoutCache) Add(key, item any, callback func()) (any, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if e, ok := c.cache[key]; ok {
+		return e.item, false
+	}
+
+	entry := &cacheEntry{
+		item:     item,
+		callback: callback,
+	}
+	entry.timer = time.AfterFunc(c.timeout, func() {
+		c.mu.Lock()
+		if entry.deleted {
+			c.mu.Unlock()
+			// Abort the delete since this has been taken care of in Remove().
+			return
+		}
+		delete(c.cache, key)
+		c.mu.Unlock()
+		entry.callback()
+	})
+	c.cache[key] = entry
+	return item, true
+}
+
+// Remove the item with the key from the cache.
+//
+// If the specified key exists in the cache, it returns (item associated with
+// key, true) and the callback associated with the item is guaranteed to be not
+// called. If the given key is not found in the cache, it returns (nil, false)
+func (c *TimeoutCache) Remove(key any) (item any, ok bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	entry, ok := c.removeInternal(key)
+	if !ok {
+		return nil, false
+	}
+	return entry.item, true
+}
+
+// removeInternal removes and returns the item with key.
+//
+// caller must hold c.mu.
+func (c *TimeoutCache) removeInternal(key any) (*cacheEntry, bool) {
+	entry, ok := c.cache[key]
+	if !ok {
+		return nil, false
+	}
+	delete(c.cache, key)
+	if !entry.timer.Stop() {
+		// If stop was not successful, the timer has fired (this can only happen
+		// in a race). But the deleting function is blocked on c.mu because the
+		// mutex was held by the caller of this function.
+		//
+		// Set deleted to true to abort the deleting function. When the lock is
+		// released, the delete function will acquire the lock, check the value
+		// of deleted and return.
+		entry.deleted = true
+	}
+	return entry, true
+}
+
+// Clear removes all entries, and runs the callbacks if runCallback is true.
+func (c *TimeoutCache) Clear(runCallback bool) {
+	var entries []*cacheEntry
+	c.mu.Lock()
+	for key := range c.cache {
+		if e, ok := c.removeInternal(key); ok {
+			entries = append(entries, e)
+		}
+	}
+	c.mu.Unlock()
+
+	if !runCallback {
+		return
+	}
+
+	// removeInternal removes entries from cache, and also stops the timer, so
+	// the callback is guaranteed to be not called. If runCallback is true,
+	// manual execute all callbacks.
+	for _, entry := range entries {
+		entry.callback()
+	}
+}
+
+// Len returns the number of entries in the cache.
+func (c *TimeoutCache) Len() int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return len(c.cache)
+}

internal/cache/timeoutCache_test.go

@@ -0,0 +1,254 @@
+/*
+ *
+ * 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 cache
+
+import (
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+const (
+	testCacheTimeout = 100 * time.Millisecond
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (c *TimeoutCache) getForTesting(key any) (*cacheEntry, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	r, ok := c.cache[key]
+	return r, ok
+}
+
+// TestCacheExpire attempts to add an entry to the cache and verifies that it
+// was added successfully. It then makes sure that on timeout, it's removed and
+// the associated callback is called.
+func (s) TestCacheExpire(t *testing.T) {
+	const k, v = 1, "1"
+	c := NewTimeoutCache(testCacheTimeout)
+
+	callbackChan := make(chan struct{})
+	c.Add(k, v, func() { close(callbackChan) })
+
+	if gotV, ok := c.getForTesting(k); !ok || gotV.item != v {
+		t.Fatalf("After Add(), before timeout, from cache got: %v, %v, want %v, %v", gotV.item, ok, v, true)
+	}
+	if l := c.Len(); l != 1 {
+		t.Fatalf("%d number of items in the cache, want 1", l)
+	}
+
+	select {
+	case <-callbackChan:
+	case <-time.After(testCacheTimeout * 2):
+		t.Fatalf("timeout waiting for callback")
+	}
+
+	if _, ok := c.getForTesting(k); ok {
+		t.Fatalf("After Add(), after timeout, from cache got: _, %v, want _, %v", ok, false)
+	}
+	if l := c.Len(); l != 0 {
+		t.Fatalf("%d number of items in the cache, want 0", l)
+	}
+}
+
+// TestCacheRemove attempts to remove an existing entry from the cache and
+// verifies that the entry is removed and the associated callback is not
+// invoked.
+func (s) TestCacheRemove(t *testing.T) {
+	const k, v = 1, "1"
+	c := NewTimeoutCache(testCacheTimeout)
+
+	callbackChan := make(chan struct{})
+	c.Add(k, v, func() { close(callbackChan) })
+
+	if got, ok := c.getForTesting(k); !ok || got.item != v {
+		t.Fatalf("After Add(), before timeout, from cache got: %v, %v, want %v, %v", got.item, ok, v, true)
+	}
+	if l := c.Len(); l != 1 {
+		t.Fatalf("%d number of items in the cache, want 1", l)
+	}
+
+	time.Sleep(testCacheTimeout / 2)
+
+	gotV, gotOK := c.Remove(k)
+	if !gotOK || gotV != v {
+		t.Fatalf("After Add(), before timeout, Remove() got: %v, %v, want %v, %v", gotV, gotOK, v, true)
+	}
+
+	if _, ok := c.getForTesting(k); ok {
+		t.Fatalf("After Add(), before timeout, after Remove(), from cache got: _, %v, want _, %v", ok, false)
+	}
+	if l := c.Len(); l != 0 {
+		t.Fatalf("%d number of items in the cache, want 0", l)
+	}
+
+	select {
+	case <-callbackChan:
+		t.Fatalf("unexpected callback after retrieve")
+	case <-time.After(testCacheTimeout * 2):
+	}
+}
+
+// TestCacheClearWithoutCallback attempts to clear all entries from the cache
+// and verifies that the associated callbacks are not invoked.
+func (s) TestCacheClearWithoutCallback(t *testing.T) {
+	var values []string
+	const itemCount = 3
+	for i := 0; i < itemCount; i++ {
+		values = append(values, strconv.Itoa(i))
+	}
+	c := NewTimeoutCache(testCacheTimeout)
+
+	done := make(chan struct{})
+	defer close(done)
+	callbackChan := make(chan struct{}, itemCount)
+
+	for i, v := range values {
+		callbackChanTemp := make(chan struct{})
+		c.Add(i, v, func() { close(callbackChanTemp) })
+		go func() {
+			select {
+			case <-callbackChanTemp:
+				callbackChan <- struct{}{}
+			case <-done:
+			}
+		}()
+	}
+
+	for i, v := range values {
+		if got, ok := c.getForTesting(i); !ok || got.item != v {
+			t.Fatalf("After Add(), before timeout, from cache got: %v, %v, want %v, %v", got.item, ok, v, true)
+		}
+	}
+	if l := c.Len(); l != itemCount {
+		t.Fatalf("%d number of items in the cache, want %d", l, itemCount)
+	}
+
+	time.Sleep(testCacheTimeout / 2)
+	c.Clear(false)
+
+	for i := range values {
+		if _, ok := c.getForTesting(i); ok {
+			t.Fatalf("After Add(), before timeout, after Remove(), from cache got: _, %v, want _, %v", ok, false)
+		}
+	}
+	if l := c.Len(); l != 0 {
+		t.Fatalf("%d number of items in the cache, want 0", l)
+	}
+
+	select {
+	case <-callbackChan:
+		t.Fatalf("unexpected callback after Clear")
+	case <-time.After(testCacheTimeout * 2):
+	}
+}
+
+// TestCacheClearWithCallback attempts to clear all entries from the cache and
+// verifies that the associated callbacks are invoked.
+func (s) TestCacheClearWithCallback(t *testing.T) {
+	var values []string
+	const itemCount = 3
+	for i := 0; i < itemCount; i++ {
+		values = append(values, strconv.Itoa(i))
+	}
+	c := NewTimeoutCache(time.Hour)
+
+	testDone := make(chan struct{})
+	defer close(testDone)
+
+	var wg sync.WaitGroup
+	wg.Add(itemCount)
+	for i, v := range values {
+		callbackChanTemp := make(chan struct{})
+		c.Add(i, v, func() { close(callbackChanTemp) })
+		go func() {
+			defer wg.Done()
+			select {
+			case <-callbackChanTemp:
+			case <-testDone:
+			}
+		}()
+	}
+
+	allGoroutineDone := make(chan struct{}, itemCount)
+	go func() {
+		wg.Wait()
+		close(allGoroutineDone)
+	}()
+
+	for i, v := range values {
+		if got, ok := c.getForTesting(i); !ok || got.item != v {
+			t.Fatalf("After Add(), before timeout, from cache got: %v, %v, want %v, %v", got.item, ok, v, true)
+		}
+	}
+	if l := c.Len(); l != itemCount {
+		t.Fatalf("%d number of items in the cache, want %d", l, itemCount)
+	}
+
+	time.Sleep(testCacheTimeout / 2)
+	c.Clear(true)
+
+	for i := range values {
+		if _, ok := c.getForTesting(i); ok {
+			t.Fatalf("After Add(), before timeout, after Remove(), from cache got: _, %v, want _, %v", ok, false)
+		}
+	}
+	if l := c.Len(); l != 0 {
+		t.Fatalf("%d number of items in the cache, want 0", l)
+	}
+
+	select {
+	case <-allGoroutineDone:
+	case <-time.After(testCacheTimeout * 2):
+		t.Fatalf("timeout waiting for all callbacks")
+	}
+}
+
+// TestCacheRetrieveTimeoutRace simulates the case where an entry's timer fires
+// around the same time that Remove() is called for it. It verifies that there
+// is no deadlock.
+func (s) TestCacheRetrieveTimeoutRace(t *testing.T) {
+	c := NewTimeoutCache(time.Nanosecond)
+
+	done := make(chan struct{})
+	go func() {
+		for i := 0; i < 1000; i++ {
+			// Add starts a timer with 1 ns timeout, then remove will race
+			// with the timer.
+			c.Add(i, strconv.Itoa(i), func() {})
+			c.Remove(i)
+		}
+		close(done)
+	}()
+
+	select {
+	case <-time.After(time.Second):
+		t.Fatalf("Test didn't finish within 1 second. Deadlock")
+	case <-done:
+	}
+}

internal/channelz/channel.go

@@ -0,0 +1,270 @@
+/*
+ *
+ * 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 channelz
+
+import (
+	"fmt"
+	"sync/atomic"
+
+	"google.golang.org/grpc/connectivity"
+)
+
+// Channel represents a channel within channelz, which includes metrics and
+// internal channelz data, such as channelz id, child list, etc.
+type Channel struct {
+	Entity
+	// ID is the channelz id of this channel.
+	ID int64
+	// RefName is the human readable reference string of this channel.
+	RefName string
+
+	closeCalled bool
+	nestedChans map[int64]string
+	subChans    map[int64]string
+	Parent      *Channel
+	trace       *ChannelTrace
+	// traceRefCount is the number of trace events that reference this channel.
+	// Non-zero traceRefCount means the trace of this channel cannot be deleted.
+	traceRefCount int32
+
+	// ChannelMetrics holds connectivity state, target and call metrics for the
+	// channel within channelz.
+	ChannelMetrics ChannelMetrics
+}
+
+// Implemented to make Channel implement the Identifier interface used for
+// nesting.
+func (c *Channel) channelzIdentifier() {}
+
+// String returns a string representation of the Channel, including its parent
+// entity and ID.
+func (c *Channel) String() string {
+	if c.Parent == nil {
+		return fmt.Sprintf("Channel #%d", c.ID)
+	}
+	return fmt.Sprintf("%s Channel #%d", c.Parent, c.ID)
+}
+
+func (c *Channel) id() int64 {
+	return c.ID
+}
+
+// SubChans returns a copy of the map of sub-channels associated with the
+// Channel.
+func (c *Channel) SubChans() map[int64]string {
+	db.mu.RLock()
+	defer db.mu.RUnlock()
+	return copyMap(c.subChans)
+}
+
+// NestedChans returns a copy of the map of nested channels associated with the
+// Channel.
+func (c *Channel) NestedChans() map[int64]string {
+	db.mu.RLock()
+	defer db.mu.RUnlock()
+	return copyMap(c.nestedChans)
+}
+
+// Trace returns a copy of the Channel's trace data.
+func (c *Channel) Trace() *ChannelTrace {
+	db.mu.RLock()
+	defer db.mu.RUnlock()
+	return c.trace.copy()
+}
+
+// ChannelMetrics holds connectivity state, target and call metrics for the
+// channel within channelz.
+type ChannelMetrics struct {
+	// The current connectivity state of the channel.
+	State atomic.Pointer[connectivity.State]
+	// The target this channel originally tried to connect to.  May be absent
+	Target atomic.Pointer[string]
+	// The number of calls started on the channel.
+	CallsStarted atomic.Int64
+	// The number of calls that have completed with an OK status.
+	CallsSucceeded atomic.Int64
+	// The number of calls that have a completed with a non-OK status.
+	CallsFailed atomic.Int64
+	// The last time a call was started on the channel.
+	LastCallStartedTimestamp atomic.Int64
+}
+
+// CopyFrom copies the metrics in o to c.  For testing only.
+func (c *ChannelMetrics) CopyFrom(o *ChannelMetrics) {
+	c.State.Store(o.State.Load())
+	c.Target.Store(o.Target.Load())
+	c.CallsStarted.Store(o.CallsStarted.Load())
+	c.CallsSucceeded.Store(o.CallsSucceeded.Load())
+	c.CallsFailed.Store(o.CallsFailed.Load())
+	c.LastCallStartedTimestamp.Store(o.LastCallStartedTimestamp.Load())
+}
+
+// Equal returns true iff the metrics of c are the same as the metrics of o.
+// For testing only.
+func (c *ChannelMetrics) Equal(o any) bool {
+	oc, ok := o.(*ChannelMetrics)
+	if !ok {
+		return false
+	}
+	if (c.State.Load() == nil) != (oc.State.Load() == nil) {
+		return false
+	}
+	if c.State.Load() != nil && *c.State.Load() != *oc.State.Load() {
+		return false
+	}
+	if (c.Target.Load() == nil) != (oc.Target.Load() == nil) {
+		return false
+	}
+	if c.Target.Load() != nil && *c.Target.Load() != *oc.Target.Load() {
+		return false
+	}
+	return c.CallsStarted.Load() == oc.CallsStarted.Load() &&
+		c.CallsFailed.Load() == oc.CallsFailed.Load() &&
+		c.CallsSucceeded.Load() == oc.CallsSucceeded.Load() &&
+		c.LastCallStartedTimestamp.Load() == oc.LastCallStartedTimestamp.Load()
+}
+
+func strFromPointer(s *string) string {
+	if s == nil {
+		return ""
+	}
+	return *s
+}
+
+// String returns a string representation of the ChannelMetrics, including its
+// state, target, and call metrics.
+func (c *ChannelMetrics) String() string {
+	return fmt.Sprintf("State: %v, Target: %s, CallsStarted: %v, CallsSucceeded: %v, CallsFailed: %v, LastCallStartedTimestamp: %v",
+		c.State.Load(), strFromPointer(c.Target.Load()), c.CallsStarted.Load(), c.CallsSucceeded.Load(), c.CallsFailed.Load(), c.LastCallStartedTimestamp.Load(),
+	)
+}
+
+// NewChannelMetricForTesting creates a new instance of ChannelMetrics with
+// specified initial values for testing purposes.
+func NewChannelMetricForTesting(state connectivity.State, target string, started, succeeded, failed, timestamp int64) *ChannelMetrics {
+	c := &ChannelMetrics{}
+	c.State.Store(&state)
+	c.Target.Store(&target)
+	c.CallsStarted.Store(started)
+	c.CallsSucceeded.Store(succeeded)
+	c.CallsFailed.Store(failed)
+	c.LastCallStartedTimestamp.Store(timestamp)
+	return c
+}
+
+func (c *Channel) addChild(id int64, e entry) {
+	switch v := e.(type) {
+	case *SubChannel:
+		c.subChans[id] = v.RefName
+	case *Channel:
+		c.nestedChans[id] = v.RefName
+	default:
+		logger.Errorf("cannot add a child (id = %d) of type %T to a channel", id, e)
+	}
+}
+
+func (c *Channel) deleteChild(id int64) {
+	delete(c.subChans, id)
+	delete(c.nestedChans, id)
+	c.deleteSelfIfReady()
+}
+
+func (c *Channel) triggerDelete() {
+	c.closeCalled = true
+	c.deleteSelfIfReady()
+}
+
+func (c *Channel) getParentID() int64 {
+	if c.Parent == nil {
+		return -1
+	}
+	return c.Parent.ID
+}
+
+// deleteSelfFromTree tries to delete the channel from the channelz entry relation tree, which means
+// deleting the channel reference from its parent's child list.
+//
+// In order for a channel to be deleted from the tree, it must meet the criteria that, removal of the
+// corresponding grpc object has been invoked, and the channel does not have any children left.
+//
+// The returned boolean value indicates whether the channel has been successfully deleted from tree.
+func (c *Channel) deleteSelfFromTree() (deleted bool) {
+	if !c.closeCalled || len(c.subChans)+len(c.nestedChans) != 0 {
+		return false
+	}
+	// not top channel
+	if c.Parent != nil {
+		c.Parent.deleteChild(c.ID)
+	}
+	return true
+}
+
+// deleteSelfFromMap checks whether it is valid to delete the channel from the map, which means
+// deleting the channel from channelz's tracking entirely. Users can no longer use id to query the
+// channel, and its memory will be garbage collected.
+//
+// The trace reference count of the channel must be 0 in order to be deleted from the map. This is
+// specified in the channel tracing gRFC that as long as some other trace has reference to an entity,
+// the trace of the referenced entity must not be deleted. In order to release the resource allocated
+// by grpc, the reference to the grpc object is reset to a dummy object.
+//
+// deleteSelfFromMap must be called after deleteSelfFromTree returns true.
+//
+// It returns a bool to indicate whether the channel can be safely deleted from map.
+func (c *Channel) deleteSelfFromMap() (delete bool) {
+	return c.getTraceRefCount() == 0
+}
+
+// deleteSelfIfReady tries to delete the channel itself from the channelz database.
+// The delete process includes two steps:
+//  1. delete the channel from the entry relation tree, i.e. delete the channel reference from its
+//     parent's child list.
+//  2. delete the channel from the map, i.e. delete the channel entirely from channelz. Lookup by id
+//     will return entry not found error.
+func (c *Channel) deleteSelfIfReady() {
+	if !c.deleteSelfFromTree() {
+		return
+	}
+	if !c.deleteSelfFromMap() {
+		return
+	}
+	db.deleteEntry(c.ID)
+	c.trace.clear()
+}
+
+func (c *Channel) getChannelTrace() *ChannelTrace {
+	return c.trace
+}
+
+func (c *Channel) incrTraceRefCount() {
+	atomic.AddInt32(&c.traceRefCount, 1)
+}
+
+func (c *Channel) decrTraceRefCount() {
+	atomic.AddInt32(&c.traceRefCount, -1)
+}
+
+func (c *Channel) getTraceRefCount() int {
+	i := atomic.LoadInt32(&c.traceRefCount)
+	return int(i)
+}
+
+func (c *Channel) getRefName() string {
+	return c.RefName
+}

internal/channelz/channelmap.go

@@ -0,0 +1,395 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package channelz
+
+import (
+	"fmt"
+	"sort"
+	"sync"
+	"time"
+)
+
+// entry represents a node in the channelz database.
+type entry interface {
+	// addChild adds a child e, whose channelz id is id to child list
+	addChild(id int64, e entry)
+	// deleteChild deletes a child with channelz id to be id from child list
+	deleteChild(id int64)
+	// triggerDelete tries to delete self from channelz database. However, if
+	// child list is not empty, then deletion from the database is on hold until
+	// the last child is deleted from database.
+	triggerDelete()
+	// deleteSelfIfReady check whether triggerDelete() has been called before,
+	// and whether child list is now empty. If both conditions are met, then
+	// delete self from database.
+	deleteSelfIfReady()
+	// getParentID returns parent ID of the entry. 0 value parent ID means no parent.
+	getParentID() int64
+	Entity
+}
+
+// channelMap is the storage data structure for channelz.
+//
+// Methods of channelMap can be divided into two categories with respect to
+// locking.
+//
+// 1. Methods acquire the global lock.
+// 2. Methods that can only be called when global lock is held.
+//
+// A second type of method need always to be called inside a first type of method.
+type channelMap struct {
+	mu               sync.RWMutex
+	topLevelChannels map[int64]struct{}
+	channels         map[int64]*Channel
+	subChannels      map[int64]*SubChannel
+	sockets          map[int64]*Socket
+	servers          map[int64]*Server
+}
+
+func newChannelMap() *channelMap {
+	return &channelMap{
+		topLevelChannels: make(map[int64]struct{}),
+		channels:         make(map[int64]*Channel),
+		subChannels:      make(map[int64]*SubChannel),
+		sockets:          make(map[int64]*Socket),
+		servers:          make(map[int64]*Server),
+	}
+}
+
+func (c *channelMap) addServer(id int64, s *Server) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	s.cm = c
+	c.servers[id] = s
+}
+
+func (c *channelMap) addChannel(id int64, cn *Channel, isTopChannel bool, pid int64) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	cn.trace.cm = c
+	c.channels[id] = cn
+	if isTopChannel {
+		c.topLevelChannels[id] = struct{}{}
+	} else if p := c.channels[pid]; p != nil {
+		p.addChild(id, cn)
+	} else {
+		logger.Infof("channel %d references invalid parent ID %d", id, pid)
+	}
+}
+
+func (c *channelMap) addSubChannel(id int64, sc *SubChannel, pid int64) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	sc.trace.cm = c
+	c.subChannels[id] = sc
+	if p := c.channels[pid]; p != nil {
+		p.addChild(id, sc)
+	} else {
+		logger.Infof("subchannel %d references invalid parent ID %d", id, pid)
+	}
+}
+
+func (c *channelMap) addSocket(s *Socket) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	s.cm = c
+	c.sockets[s.ID] = s
+	if s.Parent == nil {
+		logger.Infof("normal socket %d has no parent", s.ID)
+	}
+	s.Parent.(entry).addChild(s.ID, s)
+}
+
+// removeEntry triggers the removal of an entry, which may not indeed delete the
+// entry, if it has to wait on the deletion of its children and until no other
+// entity's channel trace references it.  It may lead to a chain of entry
+// deletion. For example, deleting the last socket of a gracefully shutting down
+// server will lead to the server being also deleted.
+func (c *channelMap) removeEntry(id int64) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.findEntry(id).triggerDelete()
+}
+
+// tracedChannel represents tracing operations which are present on both
+// channels and subChannels.
+type tracedChannel interface {
+	getChannelTrace() *ChannelTrace
+	incrTraceRefCount()
+	decrTraceRefCount()
+	getRefName() string
+}
+
+// c.mu must be held by the caller
+func (c *channelMap) decrTraceRefCount(id int64) {
+	e := c.findEntry(id)
+	if v, ok := e.(tracedChannel); ok {
+		v.decrTraceRefCount()
+		e.deleteSelfIfReady()
+	}
+}
+
+// c.mu must be held by the caller.
+func (c *channelMap) findEntry(id int64) entry {
+	if v, ok := c.channels[id]; ok {
+		return v
+	}
+	if v, ok := c.subChannels[id]; ok {
+		return v
+	}
+	if v, ok := c.servers[id]; ok {
+		return v
+	}
+	if v, ok := c.sockets[id]; ok {
+		return v
+	}
+	return &dummyEntry{idNotFound: id}
+}
+
+// c.mu must be held by the caller
+//
+// deleteEntry deletes an entry from the channelMap. Before calling this method,
+// caller must check this entry is ready to be deleted, i.e removeEntry() has
+// been called on it, and no children still exist.
+func (c *channelMap) deleteEntry(id int64) entry {
+	if v, ok := c.sockets[id]; ok {
+		delete(c.sockets, id)
+		return v
+	}
+	if v, ok := c.subChannels[id]; ok {
+		delete(c.subChannels, id)
+		return v
+	}
+	if v, ok := c.channels[id]; ok {
+		delete(c.channels, id)
+		delete(c.topLevelChannels, id)
+		return v
+	}
+	if v, ok := c.servers[id]; ok {
+		delete(c.servers, id)
+		return v
+	}
+	return &dummyEntry{idNotFound: id}
+}
+
+func (c *channelMap) traceEvent(id int64, desc *TraceEvent) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	child := c.findEntry(id)
+	childTC, ok := child.(tracedChannel)
+	if !ok {
+		return
+	}
+	childTC.getChannelTrace().append(&traceEvent{Desc: desc.Desc, Severity: desc.Severity, Timestamp: time.Now()})
+	if desc.Parent != nil {
+		parent := c.findEntry(child.getParentID())
+		var chanType RefChannelType
+		switch child.(type) {
+		case *Channel:
+			chanType = RefChannel
+		case *SubChannel:
+			chanType = RefSubChannel
+		}
+		if parentTC, ok := parent.(tracedChannel); ok {
+			parentTC.getChannelTrace().append(&traceEvent{
+				Desc:      desc.Parent.Desc,
+				Severity:  desc.Parent.Severity,
+				Timestamp: time.Now(),
+				RefID:     id,
+				RefName:   childTC.getRefName(),
+				RefType:   chanType,
+			})
+			childTC.incrTraceRefCount()
+		}
+	}
+}
+
+type int64Slice []int64
+
+func (s int64Slice) Len() int           { return len(s) }
+func (s int64Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+func (s int64Slice) Less(i, j int) bool { return s[i] < s[j] }
+
+func copyMap(m map[int64]string) map[int64]string {
+	n := make(map[int64]string)
+	for k, v := range m {
+		n[k] = v
+	}
+	return n
+}
+
+func (c *channelMap) getTopChannels(id int64, maxResults int) ([]*Channel, bool) {
+	if maxResults <= 0 {
+		maxResults = EntriesPerPage
+	}
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	l := int64(len(c.topLevelChannels))
+	ids := make([]int64, 0, l)
+
+	for k := range c.topLevelChannels {
+		ids = append(ids, k)
+	}
+	sort.Sort(int64Slice(ids))
+	idx := sort.Search(len(ids), func(i int) bool { return ids[i] >= id })
+	end := true
+	var t []*Channel
+	for _, v := range ids[idx:] {
+		if len(t) == maxResults {
+			end = false
+			break
+		}
+		if cn, ok := c.channels[v]; ok {
+			t = append(t, cn)
+		}
+	}
+	return t, end
+}
+
+func (c *channelMap) getServers(id int64, maxResults int) ([]*Server, bool) {
+	if maxResults <= 0 {
+		maxResults = EntriesPerPage
+	}
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	ids := make([]int64, 0, len(c.servers))
+	for k := range c.servers {
+		ids = append(ids, k)
+	}
+	sort.Sort(int64Slice(ids))
+	idx := sort.Search(len(ids), func(i int) bool { return ids[i] >= id })
+	end := true
+	var s []*Server
+	for _, v := range ids[idx:] {
+		if len(s) == maxResults {
+			end = false
+			break
+		}
+		if svr, ok := c.servers[v]; ok {
+			s = append(s, svr)
+		}
+	}
+	return s, end
+}
+
+func (c *channelMap) getServerSockets(id int64, startID int64, maxResults int) ([]*Socket, bool) {
+	if maxResults <= 0 {
+		maxResults = EntriesPerPage
+	}
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	svr, ok := c.servers[id]
+	if !ok {
+		// server with id doesn't exist.
+		return nil, true
+	}
+	svrskts := svr.sockets
+	ids := make([]int64, 0, len(svrskts))
+	sks := make([]*Socket, 0, min(len(svrskts), maxResults))
+	for k := range svrskts {
+		ids = append(ids, k)
+	}
+	sort.Sort(int64Slice(ids))
+	idx := sort.Search(len(ids), func(i int) bool { return ids[i] >= startID })
+	end := true
+	for _, v := range ids[idx:] {
+		if len(sks) == maxResults {
+			end = false
+			break
+		}
+		if ns, ok := c.sockets[v]; ok {
+			sks = append(sks, ns)
+		}
+	}
+	return sks, end
+}
+
+func (c *channelMap) getChannel(id int64) *Channel {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.channels[id]
+}
+
+func (c *channelMap) getSubChannel(id int64) *SubChannel {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.subChannels[id]
+}
+
+func (c *channelMap) getSocket(id int64) *Socket {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.sockets[id]
+}
+
+func (c *channelMap) getServer(id int64) *Server {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.servers[id]
+}
+
+type dummyEntry struct {
+	// dummyEntry is a fake entry to handle entry not found case.
+	idNotFound int64
+	Entity
+}
+
+func (d *dummyEntry) String() string {
+	return fmt.Sprintf("non-existent entity #%d", d.idNotFound)
+}
+
+func (d *dummyEntry) ID() int64 { return d.idNotFound }
+
+func (d *dummyEntry) addChild(id int64, e entry) {
+	// Note: It is possible for a normal program to reach here under race
+	// condition.  For example, there could be a race between ClientConn.Close()
+	// info being propagated to addrConn and http2Client. ClientConn.Close()
+	// cancel the context and result in http2Client to error. The error info is
+	// then caught by transport monitor and before addrConn.tearDown() is called
+	// in side ClientConn.Close(). Therefore, the addrConn will create a new
+	// transport. And when registering the new transport in channelz, its parent
+	// addrConn could have already been torn down and deleted from channelz
+	// tracking, and thus reach the code here.
+	logger.Infof("attempt to add child of type %T with id %d to a parent (id=%d) that doesn't currently exist", e, id, d.idNotFound)
+}
+
+func (d *dummyEntry) deleteChild(id int64) {
+	// It is possible for a normal program to reach here under race condition.
+	// Refer to the example described in addChild().
+	logger.Infof("attempt to delete child with id %d from a parent (id=%d) that doesn't currently exist", id, d.idNotFound)
+}
+
+func (d *dummyEntry) triggerDelete() {
+	logger.Warningf("attempt to delete an entry (id=%d) that doesn't currently exist", d.idNotFound)
+}
+
+func (*dummyEntry) deleteSelfIfReady() {
+	// code should not reach here. deleteSelfIfReady is always called on an existing entry.
+}
+
+func (*dummyEntry) getParentID() int64 {
+	return 0
+}
+
+// Entity is implemented by all channelz types.
+type Entity interface {
+	isEntity()
+	fmt.Stringer
+	id() int64
+}

internal/channelz/funcs.go

@@ -0,0 +1,230 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package channelz defines internal APIs for enabling channelz service, entry
+// registration/deletion, and accessing channelz data. It also defines channelz
+// metric struct formats.
+package channelz
+
+import (
+	"sync/atomic"
+	"time"
+
+	"google.golang.org/grpc/internal"
+)
+
+var (
+	// IDGen is the global channelz entity ID generator.  It should not be used
+	// outside this package except by tests.
+	IDGen IDGenerator
+
+	db = newChannelMap()
+	// EntriesPerPage defines the number of channelz entries to be shown on a web page.
+	EntriesPerPage = 50
+	curState       int32
+)
+
+// TurnOn turns on channelz data collection.
+func TurnOn() {
+	atomic.StoreInt32(&curState, 1)
+}
+
+func init() {
+	internal.ChannelzTurnOffForTesting = func() {
+		atomic.StoreInt32(&curState, 0)
+	}
+}
+
+// IsOn returns whether channelz data collection is on.
+func IsOn() bool {
+	return atomic.LoadInt32(&curState) == 1
+}
+
+// GetTopChannels returns a slice of top channel's ChannelMetric, along with a
+// boolean indicating whether there's more top channels to be queried for.
+//
+// The arg id specifies that only top channel with id at or above it will be
+// included in the result. The returned slice is up to a length of the arg
+// maxResults or EntriesPerPage if maxResults is zero, and is sorted in ascending
+// id order.
+func GetTopChannels(id int64, maxResults int) ([]*Channel, bool) {
+	return db.getTopChannels(id, maxResults)
+}
+
+// GetServers returns a slice of server's ServerMetric, along with a
+// boolean indicating whether there's more servers to be queried for.
+//
+// The arg id specifies that only server with id at or above it will be included
+// in the result. The returned slice is up to a length of the arg maxResults or
+// EntriesPerPage if maxResults is zero, and is sorted in ascending id order.
+func GetServers(id int64, maxResults int) ([]*Server, bool) {
+	return db.getServers(id, maxResults)
+}
+
+// GetServerSockets returns a slice of server's (identified by id) normal socket's
+// SocketMetrics, along with a boolean indicating whether there's more sockets to
+// be queried for.
+//
+// The arg startID specifies that only sockets with id at or above it will be
+// included in the result. The returned slice is up to a length of the arg maxResults
+// or EntriesPerPage if maxResults is zero, and is sorted in ascending id order.
+func GetServerSockets(id int64, startID int64, maxResults int) ([]*Socket, bool) {
+	return db.getServerSockets(id, startID, maxResults)
+}
+
+// GetChannel returns the Channel for the channel (identified by id).
+func GetChannel(id int64) *Channel {
+	return db.getChannel(id)
+}
+
+// GetSubChannel returns the SubChannel for the subchannel (identified by id).
+func GetSubChannel(id int64) *SubChannel {
+	return db.getSubChannel(id)
+}
+
+// GetSocket returns the Socket for the socket (identified by id).
+func GetSocket(id int64) *Socket {
+	return db.getSocket(id)
+}
+
+// GetServer returns the ServerMetric for the server (identified by id).
+func GetServer(id int64) *Server {
+	return db.getServer(id)
+}
+
+// RegisterChannel registers the given channel c in the channelz database with
+// target as its target and reference name, and adds it to the child list of its
+// parent.  parent == nil means no parent.
+//
+// Returns a unique channelz identifier assigned to this channel.
+//
+// If channelz is not turned ON, the channelz database is not mutated.
+func RegisterChannel(parent *Channel, target string) *Channel {
+	id := IDGen.genID()
+
+	if !IsOn() {
+		return &Channel{ID: id}
+	}
+
+	isTopChannel := parent == nil
+
+	cn := &Channel{
+		ID:          id,
+		RefName:     target,
+		nestedChans: make(map[int64]string),
+		subChans:    make(map[int64]string),
+		Parent:      parent,
+		trace:       &ChannelTrace{CreationTime: time.Now(), Events: make([]*traceEvent, 0, getMaxTraceEntry())},
+	}
+	cn.ChannelMetrics.Target.Store(&target)
+	db.addChannel(id, cn, isTopChannel, cn.getParentID())
+	return cn
+}
+
+// RegisterSubChannel registers the given subChannel c in the channelz database
+// with ref as its reference name, and adds it to the child list of its parent
+// (identified by pid).
+//
+// Returns a unique channelz identifier assigned to this subChannel.
+//
+// If channelz is not turned ON, the channelz database is not mutated.
+func RegisterSubChannel(parent *Channel, ref string) *SubChannel {
+	id := IDGen.genID()
+	sc := &SubChannel{
+		ID:      id,
+		RefName: ref,
+		parent:  parent,
+	}
+
+	if !IsOn() {
+		return sc
+	}
+
+	sc.sockets = make(map[int64]string)
+	sc.trace = &ChannelTrace{CreationTime: time.Now(), Events: make([]*traceEvent, 0, getMaxTraceEntry())}
+	db.addSubChannel(id, sc, parent.ID)
+	return sc
+}
+
+// RegisterServer registers the given server s in channelz database. It returns
+// the unique channelz tracking id assigned to this server.
+//
+// If channelz is not turned ON, the channelz database is not mutated.
+func RegisterServer(ref string) *Server {
+	id := IDGen.genID()
+	if !IsOn() {
+		return &Server{ID: id}
+	}
+
+	svr := &Server{
+		RefName:       ref,
+		sockets:       make(map[int64]string),
+		listenSockets: make(map[int64]string),
+		ID:            id,
+	}
+	db.addServer(id, svr)
+	return svr
+}
+
+// RegisterSocket registers the given normal socket s in channelz database
+// with ref as its reference name, and adds it to the child list of its parent
+// (identified by skt.Parent, which must be set). It returns the unique channelz
+// tracking id assigned to this normal socket.
+//
+// If channelz is not turned ON, the channelz database is not mutated.
+func RegisterSocket(skt *Socket) *Socket {
+	skt.ID = IDGen.genID()
+	if IsOn() {
+		db.addSocket(skt)
+	}
+	return skt
+}
+
+// RemoveEntry removes an entry with unique channelz tracking id to be id from
+// channelz database.
+//
+// If channelz is not turned ON, this function is a no-op.
+func RemoveEntry(id int64) {
+	if !IsOn() {
+		return
+	}
+	db.removeEntry(id)
+}
+
+// IDGenerator is an incrementing atomic that tracks IDs for channelz entities.
+type IDGenerator struct {
+	id int64
+}
+
+// Reset resets the generated ID back to zero.  Should only be used at
+// initialization or by tests sensitive to the ID number.
+func (i *IDGenerator) Reset() {
+	atomic.StoreInt64(&i.id, 0)
+}
+
+func (i *IDGenerator) genID() int64 {
+	return atomic.AddInt64(&i.id, 1)
+}
+
+// Identifier is an opaque channelz identifier used to expose channelz symbols
+// outside of grpc.  Currently only implemented by Channel since no other
+// types require exposure outside grpc.
+type Identifier interface {
+	Entity
+	channelzIdentifier()
+}

internal/channelz/logging.go

@@ -0,0 +1,75 @@
+/*
+ *
+ * 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 channelz
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+var logger = grpclog.Component("channelz")
+
+// Info logs and adds a trace event if channelz is on.
+func Info(l grpclog.DepthLoggerV2, e Entity, args ...any) {
+	AddTraceEvent(l, e, 1, &TraceEvent{
+		Desc:     fmt.Sprint(args...),
+		Severity: CtInfo,
+	})
+}
+
+// Infof logs and adds a trace event if channelz is on.
+func Infof(l grpclog.DepthLoggerV2, e Entity, format string, args ...any) {
+	AddTraceEvent(l, e, 1, &TraceEvent{
+		Desc:     fmt.Sprintf(format, args...),
+		Severity: CtInfo,
+	})
+}
+
+// Warning logs and adds a trace event if channelz is on.
+func Warning(l grpclog.DepthLoggerV2, e Entity, args ...any) {
+	AddTraceEvent(l, e, 1, &TraceEvent{
+		Desc:     fmt.Sprint(args...),
+		Severity: CtWarning,
+	})
+}
+
+// Warningf logs and adds a trace event if channelz is on.
+func Warningf(l grpclog.DepthLoggerV2, e Entity, format string, args ...any) {
+	AddTraceEvent(l, e, 1, &TraceEvent{
+		Desc:     fmt.Sprintf(format, args...),
+		Severity: CtWarning,
+	})
+}
+
+// Error logs and adds a trace event if channelz is on.
+func Error(l grpclog.DepthLoggerV2, e Entity, args ...any) {
+	AddTraceEvent(l, e, 1, &TraceEvent{
+		Desc:     fmt.Sprint(args...),
+		Severity: CtError,
+	})
+}
+
+// Errorf logs and adds a trace event if channelz is on.
+func Errorf(l grpclog.DepthLoggerV2, e Entity, format string, args ...any) {
+	AddTraceEvent(l, e, 1, &TraceEvent{
+		Desc:     fmt.Sprintf(format, args...),
+		Severity: CtError,
+	})
+}

internal/channelz/server.go

@@ -0,0 +1,121 @@
+/*
+ *
+ * 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 channelz
+
+import (
+	"fmt"
+	"sync/atomic"
+)
+
+// Server is the channelz representation of a server.
+type Server struct {
+	Entity
+	ID      int64
+	RefName string
+
+	ServerMetrics ServerMetrics
+
+	closeCalled   bool
+	sockets       map[int64]string
+	listenSockets map[int64]string
+	cm            *channelMap
+}
+
+// ServerMetrics defines a struct containing metrics for servers.
+type ServerMetrics struct {
+	// The number of incoming calls started on the server.
+	CallsStarted atomic.Int64
+	// The number of incoming calls that have completed with an OK status.
+	CallsSucceeded atomic.Int64
+	// The number of incoming calls that have a completed with a non-OK status.
+	CallsFailed atomic.Int64
+	// The last time a call was started on the server.
+	LastCallStartedTimestamp atomic.Int64
+}
+
+// NewServerMetricsForTesting returns an initialized ServerMetrics.
+func NewServerMetricsForTesting(started, succeeded, failed, timestamp int64) *ServerMetrics {
+	sm := &ServerMetrics{}
+	sm.CallsStarted.Store(started)
+	sm.CallsSucceeded.Store(succeeded)
+	sm.CallsFailed.Store(failed)
+	sm.LastCallStartedTimestamp.Store(timestamp)
+	return sm
+}
+
+// CopyFrom copies the metrics data from the provided ServerMetrics
+// instance into the current instance.
+func (sm *ServerMetrics) CopyFrom(o *ServerMetrics) {
+	sm.CallsStarted.Store(o.CallsStarted.Load())
+	sm.CallsSucceeded.Store(o.CallsSucceeded.Load())
+	sm.CallsFailed.Store(o.CallsFailed.Load())
+	sm.LastCallStartedTimestamp.Store(o.LastCallStartedTimestamp.Load())
+}
+
+// ListenSockets returns the listening sockets for s.
+func (s *Server) ListenSockets() map[int64]string {
+	db.mu.RLock()
+	defer db.mu.RUnlock()
+	return copyMap(s.listenSockets)
+}
+
+// String returns a printable description of s.
+func (s *Server) String() string {
+	return fmt.Sprintf("Server #%d", s.ID)
+}
+
+func (s *Server) id() int64 {
+	return s.ID
+}
+
+func (s *Server) addChild(id int64, e entry) {
+	switch v := e.(type) {
+	case *Socket:
+		switch v.SocketType {
+		case SocketTypeNormal:
+			s.sockets[id] = v.RefName
+		case SocketTypeListen:
+			s.listenSockets[id] = v.RefName
+		}
+	default:
+		logger.Errorf("cannot add a child (id = %d) of type %T to a server", id, e)
+	}
+}
+
+func (s *Server) deleteChild(id int64) {
+	delete(s.sockets, id)
+	delete(s.listenSockets, id)
+	s.deleteSelfIfReady()
+}
+
+func (s *Server) triggerDelete() {
+	s.closeCalled = true
+	s.deleteSelfIfReady()
+}
+
+func (s *Server) deleteSelfIfReady() {
+	if !s.closeCalled || len(s.sockets)+len(s.listenSockets) != 0 {
+		return
+	}
+	s.cm.deleteEntry(s.ID)
+}
+
+func (s *Server) getParentID() int64 {
+	return 0
+}

internal/channelz/socket.go

@@ -0,0 +1,137 @@
+/*
+ *
+ * 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 channelz
+
+import (
+	"fmt"
+	"net"
+	"sync/atomic"
+
+	"google.golang.org/grpc/credentials"
+)
+
+// SocketMetrics defines the struct that the implementor of Socket interface
+// should return from ChannelzMetric().
+type SocketMetrics struct {
+	// The number of streams that have been started.
+	StreamsStarted atomic.Int64
+	// The number of streams that have ended successfully:
+	// On client side, receiving frame with eos bit set.
+	// On server side, sending frame with eos bit set.
+	StreamsSucceeded atomic.Int64
+	// The number of streams that have ended unsuccessfully:
+	// On client side, termination without receiving frame with eos bit set.
+	// On server side, termination without sending frame with eos bit set.
+	StreamsFailed atomic.Int64
+	// The number of messages successfully sent on this socket.
+	MessagesSent     atomic.Int64
+	MessagesReceived atomic.Int64
+	// The number of keep alives sent.  This is typically implemented with HTTP/2
+	// ping messages.
+	KeepAlivesSent atomic.Int64
+	// The last time a stream was created by this endpoint.  Usually unset for
+	// servers.
+	LastLocalStreamCreatedTimestamp atomic.Int64
+	// The last time a stream was created by the remote endpoint.  Usually unset
+	// for clients.
+	LastRemoteStreamCreatedTimestamp atomic.Int64
+	// The last time a message was sent by this endpoint.
+	LastMessageSentTimestamp atomic.Int64
+	// The last time a message was received by this endpoint.
+	LastMessageReceivedTimestamp atomic.Int64
+}
+
+// EphemeralSocketMetrics are metrics that change rapidly and are tracked
+// outside of channelz.
+type EphemeralSocketMetrics struct {
+	// The amount of window, granted to the local endpoint by the remote endpoint.
+	// This may be slightly out of date due to network latency.  This does NOT
+	// include stream level or TCP level flow control info.
+	LocalFlowControlWindow int64
+	// The amount of window, granted to the remote endpoint by the local endpoint.
+	// This may be slightly out of date due to network latency.  This does NOT
+	// include stream level or TCP level flow control info.
+	RemoteFlowControlWindow int64
+}
+
+// SocketType represents the type of socket.
+type SocketType string
+
+// SocketType can be one of these.
+const (
+	SocketTypeNormal = "NormalSocket"
+	SocketTypeListen = "ListenSocket"
+)
+
+// Socket represents a socket within channelz which includes socket
+// metrics and data related to socket activity and provides methods
+// for managing and interacting with sockets.
+type Socket struct {
+	Entity
+	SocketType       SocketType
+	ID               int64
+	Parent           Entity
+	cm               *channelMap
+	SocketMetrics    SocketMetrics
+	EphemeralMetrics func() *EphemeralSocketMetrics
+
+	RefName string
+	// The locally bound address.  Immutable.
+	LocalAddr net.Addr
+	// The remote bound address.  May be absent.  Immutable.
+	RemoteAddr net.Addr
+	// Optional, represents the name of the remote endpoint, if different than
+	// the original target name.  Immutable.
+	RemoteName string
+	// Immutable.
+	SocketOptions *SocketOptionData
+	// Immutable.
+	Security credentials.ChannelzSecurityValue
+}
+
+// String returns a string representation of the Socket, including its parent
+// entity, socket type, and ID.
+func (ls *Socket) String() string {
+	return fmt.Sprintf("%s %s #%d", ls.Parent, ls.SocketType, ls.ID)
+}
+
+func (ls *Socket) id() int64 {
+	return ls.ID
+}
+
+func (ls *Socket) addChild(id int64, e entry) {
+	logger.Errorf("cannot add a child (id = %d) of type %T to a listen socket", id, e)
+}
+
+func (ls *Socket) deleteChild(id int64) {
+	logger.Errorf("cannot delete a child (id = %d) from a listen socket", id)
+}
+
+func (ls *Socket) triggerDelete() {
+	ls.cm.deleteEntry(ls.ID)
+	ls.Parent.(entry).deleteChild(ls.ID)
+}
+
+func (ls *Socket) deleteSelfIfReady() {
+	logger.Errorf("cannot call deleteSelfIfReady on a listen socket")
+}
+
+func (ls *Socket) getParentID() int64 {
+	return ls.Parent.id()
+}

internal/channelz/subchannel.go

@@ -0,0 +1,153 @@
+/*
+ *
+ * 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 channelz
+
+import (
+	"fmt"
+	"sync/atomic"
+)
+
+// SubChannel is the channelz representation of a subchannel.
+type SubChannel struct {
+	Entity
+	// ID is the channelz id of this subchannel.
+	ID int64
+	// RefName is the human readable reference string of this subchannel.
+	RefName       string
+	closeCalled   bool
+	sockets       map[int64]string
+	parent        *Channel
+	trace         *ChannelTrace
+	traceRefCount int32
+
+	ChannelMetrics ChannelMetrics
+}
+
+func (sc *SubChannel) String() string {
+	return fmt.Sprintf("%s SubChannel #%d", sc.parent, sc.ID)
+}
+
+func (sc *SubChannel) id() int64 {
+	return sc.ID
+}
+
+// Sockets returns a copy of the sockets map associated with the SubChannel.
+func (sc *SubChannel) Sockets() map[int64]string {
+	db.mu.RLock()
+	defer db.mu.RUnlock()
+	return copyMap(sc.sockets)
+}
+
+// Trace returns a copy of the ChannelTrace associated with the SubChannel.
+func (sc *SubChannel) Trace() *ChannelTrace {
+	db.mu.RLock()
+	defer db.mu.RUnlock()
+	return sc.trace.copy()
+}
+
+func (sc *SubChannel) addChild(id int64, e entry) {
+	if v, ok := e.(*Socket); ok && v.SocketType == SocketTypeNormal {
+		sc.sockets[id] = v.RefName
+	} else {
+		logger.Errorf("cannot add a child (id = %d) of type %T to a subChannel", id, e)
+	}
+}
+
+func (sc *SubChannel) deleteChild(id int64) {
+	delete(sc.sockets, id)
+	sc.deleteSelfIfReady()
+}
+
+func (sc *SubChannel) triggerDelete() {
+	sc.closeCalled = true
+	sc.deleteSelfIfReady()
+}
+
+func (sc *SubChannel) getParentID() int64 {
+	return sc.parent.ID
+}
+
+// deleteSelfFromTree tries to delete the subchannel from the channelz entry relation tree, which
+// means deleting the subchannel reference from its parent's child list.
+//
+// In order for a subchannel to be deleted from the tree, it must meet the criteria that, removal of
+// the corresponding grpc object has been invoked, and the subchannel does not have any children left.
+//
+// The returned boolean value indicates whether the channel has been successfully deleted from tree.
+func (sc *SubChannel) deleteSelfFromTree() (deleted bool) {
+	if !sc.closeCalled || len(sc.sockets) != 0 {
+		return false
+	}
+	sc.parent.deleteChild(sc.ID)
+	return true
+}
+
+// deleteSelfFromMap checks whether it is valid to delete the subchannel from the map, which means
+// deleting the subchannel from channelz's tracking entirely. Users can no longer use id to query
+// the subchannel, and its memory will be garbage collected.
+//
+// The trace reference count of the subchannel must be 0 in order to be deleted from the map. This is
+// specified in the channel tracing gRFC that as long as some other trace has reference to an entity,
+// the trace of the referenced entity must not be deleted. In order to release the resource allocated
+// by grpc, the reference to the grpc object is reset to a dummy object.
+//
+// deleteSelfFromMap must be called after deleteSelfFromTree returns true.
+//
+// It returns a bool to indicate whether the channel can be safely deleted from map.
+func (sc *SubChannel) deleteSelfFromMap() (delete bool) {
+	return sc.getTraceRefCount() == 0
+}
+
+// deleteSelfIfReady tries to delete the subchannel itself from the channelz database.
+// The delete process includes two steps:
+//  1. delete the subchannel from the entry relation tree, i.e. delete the subchannel reference from
+//     its parent's child list.
+//  2. delete the subchannel from the map, i.e. delete the subchannel entirely from channelz. Lookup
+//     by id will return entry not found error.
+func (sc *SubChannel) deleteSelfIfReady() {
+	if !sc.deleteSelfFromTree() {
+		return
+	}
+	if !sc.deleteSelfFromMap() {
+		return
+	}
+	db.deleteEntry(sc.ID)
+	sc.trace.clear()
+}
+
+func (sc *SubChannel) getChannelTrace() *ChannelTrace {
+	return sc.trace
+}
+
+func (sc *SubChannel) incrTraceRefCount() {
+	atomic.AddInt32(&sc.traceRefCount, 1)
+}
+
+func (sc *SubChannel) decrTraceRefCount() {
+	atomic.AddInt32(&sc.traceRefCount, -1)
+}
+
+func (sc *SubChannel) getTraceRefCount() int {
+	i := atomic.LoadInt32(&sc.traceRefCount)
+	return int(i)
+}
+
+func (sc *SubChannel) getRefName() string {
+	return sc.RefName
+}

internal/channelz/syscall_linux.go

@@ -0,0 +1,65 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package channelz
+
+import (
+	"syscall"
+
+	"golang.org/x/sys/unix"
+)
+
+// SocketOptionData defines the struct to hold socket option data, and related
+// getter function to obtain info from fd.
+type SocketOptionData struct {
+	Linger      *unix.Linger
+	RecvTimeout *unix.Timeval
+	SendTimeout *unix.Timeval
+	TCPInfo     *unix.TCPInfo
+}
+
+// Getsockopt defines the function to get socket options requested by channelz.
+// It is to be passed to syscall.RawConn.Control().
+func (s *SocketOptionData) Getsockopt(fd uintptr) {
+	if v, err := unix.GetsockoptLinger(int(fd), syscall.SOL_SOCKET, syscall.SO_LINGER); err == nil {
+		s.Linger = v
+	}
+	if v, err := unix.GetsockoptTimeval(int(fd), syscall.SOL_SOCKET, syscall.SO_RCVTIMEO); err == nil {
+		s.RecvTimeout = v
+	}
+	if v, err := unix.GetsockoptTimeval(int(fd), syscall.SOL_SOCKET, syscall.SO_SNDTIMEO); err == nil {
+		s.SendTimeout = v
+	}
+	if v, err := unix.GetsockoptTCPInfo(int(fd), syscall.SOL_TCP, syscall.TCP_INFO); err == nil {
+		s.TCPInfo = v
+	}
+}
+
+// GetSocketOption gets the socket option info of the conn.
+func GetSocketOption(socket any) *SocketOptionData {
+	c, ok := socket.(syscall.Conn)
+	if !ok {
+		return nil
+	}
+	data := &SocketOptionData{}
+	if rawConn, err := c.SyscallConn(); err == nil {
+		rawConn.Control(data.Getsockopt)
+		return data
+	}
+	return nil
+}

internal/channelz/syscall_nonlinux.go

@@ -0,0 +1,47 @@
+//go:build !linux
+
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package channelz
+
+import (
+	"sync"
+)
+
+var once sync.Once
+
+// SocketOptionData defines the struct to hold socket option data, and related
+// getter function to obtain info from fd.
+// Windows OS doesn't support Socket Option
+type SocketOptionData struct {
+}
+
+// Getsockopt defines the function to get socket options requested by channelz.
+// It is to be passed to syscall.RawConn.Control().
+// Windows OS doesn't support Socket Option
+func (s *SocketOptionData) Getsockopt(uintptr) {
+	once.Do(func() {
+		logger.Warning("Channelz: socket options are not supported on non-linux environments")
+	})
+}
+
+// GetSocketOption gets the socket option info of the conn.
+func GetSocketOption(any) *SocketOptionData {
+	return nil
+}

internal/channelz/syscall_test.go

@@ -0,0 +1,95 @@
+//go:build linux
+
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package channelz_test
+
+import (
+	"net"
+	"reflect"
+	"syscall"
+	"testing"
+
+	"golang.org/x/sys/unix"
+	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestGetSocketOpt(t *testing.T) {
+	network, addr := "tcp", ":0"
+	ln, err := net.Listen(network, addr)
+	if err != nil {
+		t.Fatalf("net.Listen(%s,%s) failed with err: %v", network, addr, err)
+	}
+	defer ln.Close()
+	go func() {
+		ln.Accept()
+	}()
+	conn, _ := net.Dial(network, ln.Addr().String())
+	defer conn.Close()
+	tcpc := conn.(*net.TCPConn)
+	raw, err := tcpc.SyscallConn()
+	if err != nil {
+		t.Fatalf("SyscallConn() failed due to %v", err)
+	}
+
+	l := &unix.Linger{Onoff: 1, Linger: 5}
+	recvTimeout := &unix.Timeval{Sec: 100}
+	sendTimeout := &unix.Timeval{Sec: 8888}
+	raw.Control(func(fd uintptr) {
+		err := unix.SetsockoptLinger(int(fd), syscall.SOL_SOCKET, syscall.SO_LINGER, l)
+		if err != nil {
+			t.Fatalf("failed to SetsockoptLinger(%v,%v,%v,%v) due to %v", int(fd), syscall.SOL_SOCKET, syscall.SO_LINGER, l, err)
+		}
+		err = unix.SetsockoptTimeval(int(fd), syscall.SOL_SOCKET, syscall.SO_RCVTIMEO, recvTimeout)
+		if err != nil {
+			t.Fatalf("failed to SetsockoptTimeval(%v,%v,%v,%v) due to %v", int(fd), syscall.SOL_SOCKET, syscall.SO_RCVTIMEO, recvTimeout, err)
+		}
+		err = unix.SetsockoptTimeval(int(fd), syscall.SOL_SOCKET, syscall.SO_SNDTIMEO, sendTimeout)
+		if err != nil {
+			t.Fatalf("failed to SetsockoptTimeval(%v,%v,%v,%v) due to %v", int(fd), syscall.SOL_SOCKET, syscall.SO_SNDTIMEO, sendTimeout, err)
+		}
+	})
+	sktopt := channelz.GetSocketOption(conn)
+	if !reflect.DeepEqual(sktopt.Linger, l) {
+		t.Fatalf("get socket option linger, want: %v, got %v", l, sktopt.Linger)
+	}
+	if !reflect.DeepEqual(sktopt.RecvTimeout, recvTimeout) {
+		t.Logf("get socket option recv timeout, want: %v, got %v, may be caused by system allowing non or partial setting of this value", recvTimeout, sktopt.RecvTimeout)
+	}
+	if !reflect.DeepEqual(sktopt.SendTimeout, sendTimeout) {
+		t.Logf("get socket option send timeout, want: %v, got %v, may be caused by system allowing non or partial setting of this value", sendTimeout, sktopt.SendTimeout)
+	}
+	if sktopt == nil || sktopt.TCPInfo != nil && sktopt.TCPInfo.State != 1 {
+		t.Fatalf("TCPInfo.State want 1 (TCP_ESTABLISHED), got %v", sktopt)
+	}
+
+	sktopt = channelz.GetSocketOption(ln)
+	if sktopt == nil || sktopt.TCPInfo == nil || sktopt.TCPInfo.State != 10 {
+		t.Fatalf("TCPInfo.State want 10 (TCP_LISTEN), got %v", sktopt)
+	}
+}

internal/channelz/trace.go

@@ -0,0 +1,213 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package channelz
+
+import (
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+const (
+	defaultMaxTraceEntry int32 = 30
+)
+
+var maxTraceEntry = defaultMaxTraceEntry
+
+// SetMaxTraceEntry sets maximum number of trace entries per entity (i.e.
+// channel/subchannel).  Setting it to 0 will disable channel tracing.
+func SetMaxTraceEntry(i int32) {
+	atomic.StoreInt32(&maxTraceEntry, i)
+}
+
+// ResetMaxTraceEntryToDefault resets the maximum number of trace entries per
+// entity to default.
+func ResetMaxTraceEntryToDefault() {
+	atomic.StoreInt32(&maxTraceEntry, defaultMaxTraceEntry)
+}
+
+func getMaxTraceEntry() int {
+	i := atomic.LoadInt32(&maxTraceEntry)
+	return int(i)
+}
+
+// traceEvent is an internal representation of a single trace event
+type traceEvent struct {
+	// Desc is a simple description of the trace event.
+	Desc string
+	// Severity states the severity of this trace event.
+	Severity Severity
+	// Timestamp is the event time.
+	Timestamp time.Time
+	// RefID is the id of the entity that gets referenced in the event. RefID is 0 if no other entity is
+	// involved in this event.
+	// e.g. SubChannel (id: 4[]) Created. --> RefID = 4, RefName = "" (inside [])
+	RefID int64
+	// RefName is the reference name for the entity that gets referenced in the event.
+	RefName string
+	// RefType indicates the referenced entity type, i.e Channel or SubChannel.
+	RefType RefChannelType
+}
+
+// TraceEvent is what the caller of AddTraceEvent should provide to describe the
+// event to be added to the channel trace.
+//
+// The Parent field is optional. It is used for an event that will be recorded
+// in the entity's parent trace.
+type TraceEvent struct {
+	Desc     string
+	Severity Severity
+	Parent   *TraceEvent
+}
+
+// ChannelTrace provides tracing information for a channel.
+// It tracks various events and metadata related to the channel's lifecycle
+// and operations.
+type ChannelTrace struct {
+	cm          *channelMap
+	clearCalled bool
+	// The time when the trace was created.
+	CreationTime time.Time
+	// A counter for the number of events recorded in the
+	// trace.
+	EventNum int64
+	mu       sync.Mutex
+	// A slice of traceEvent pointers representing the events recorded for
+	// this channel.
+	Events []*traceEvent
+}
+
+func (c *ChannelTrace) copy() *ChannelTrace {
+	return &ChannelTrace{
+		CreationTime: c.CreationTime,
+		EventNum:     c.EventNum,
+		Events:       append(([]*traceEvent)(nil), c.Events...),
+	}
+}
+
+func (c *ChannelTrace) append(e *traceEvent) {
+	c.mu.Lock()
+	if len(c.Events) == getMaxTraceEntry() {
+		del := c.Events[0]
+		c.Events = c.Events[1:]
+		if del.RefID != 0 {
+			// start recursive cleanup in a goroutine to not block the call originated from grpc.
+			go func() {
+				// need to acquire c.cm.mu lock to call the unlocked attemptCleanup func.
+				c.cm.mu.Lock()
+				c.cm.decrTraceRefCount(del.RefID)
+				c.cm.mu.Unlock()
+			}()
+		}
+	}
+	e.Timestamp = time.Now()
+	c.Events = append(c.Events, e)
+	c.EventNum++
+	c.mu.Unlock()
+}
+
+func (c *ChannelTrace) clear() {
+	if c.clearCalled {
+		return
+	}
+	c.clearCalled = true
+	c.mu.Lock()
+	for _, e := range c.Events {
+		if e.RefID != 0 {
+			// caller should have already held the c.cm.mu lock.
+			c.cm.decrTraceRefCount(e.RefID)
+		}
+	}
+	c.mu.Unlock()
+}
+
+// Severity is the severity level of a trace event.
+// The canonical enumeration of all valid values is here:
+// https://github.com/grpc/grpc-proto/blob/9b13d199cc0d4703c7ea26c9c330ba695866eb23/grpc/channelz/v1/channelz.proto#L126.
+type Severity int
+
+const (
+	// CtUnknown indicates unknown severity of a trace event.
+	CtUnknown Severity = iota
+	// CtInfo indicates info level severity of a trace event.
+	CtInfo
+	// CtWarning indicates warning level severity of a trace event.
+	CtWarning
+	// CtError indicates error level severity of a trace event.
+	CtError
+)
+
+// RefChannelType is the type of the entity being referenced in a trace event.
+type RefChannelType int
+
+const (
+	// RefUnknown indicates an unknown entity type, the zero value for this type.
+	RefUnknown RefChannelType = iota
+	// RefChannel indicates the referenced entity is a Channel.
+	RefChannel
+	// RefSubChannel indicates the referenced entity is a SubChannel.
+	RefSubChannel
+	// RefServer indicates the referenced entity is a Server.
+	RefServer
+	// RefListenSocket indicates the referenced entity is a ListenSocket.
+	RefListenSocket
+	// RefNormalSocket indicates the referenced entity is a NormalSocket.
+	RefNormalSocket
+)
+
+var refChannelTypeToString = map[RefChannelType]string{
+	RefUnknown:      "Unknown",
+	RefChannel:      "Channel",
+	RefSubChannel:   "SubChannel",
+	RefServer:       "Server",
+	RefListenSocket: "ListenSocket",
+	RefNormalSocket: "NormalSocket",
+}
+
+// String returns a string representation of the RefChannelType
+func (r RefChannelType) String() string {
+	return refChannelTypeToString[r]
+}
+
+// AddTraceEvent adds trace related to the entity with specified id, using the
+// provided TraceEventDesc.
+//
+// If channelz is not turned ON, this will simply log the event descriptions.
+func AddTraceEvent(l grpclog.DepthLoggerV2, e Entity, depth int, desc *TraceEvent) {
+	// Log only the trace description associated with the bottom most entity.
+	d := fmt.Sprintf("[%s] %s", e, desc.Desc)
+	switch desc.Severity {
+	case CtUnknown, CtInfo:
+		l.InfoDepth(depth+1, d)
+	case CtWarning:
+		l.WarningDepth(depth+1, d)
+	case CtError:
+		l.ErrorDepth(depth+1, d)
+	}
+
+	if getMaxTraceEntry() == 0 {
+		return
+	}
+	if IsOn() {
+		db.traceEvent(e.id(), desc)
+	}
+}

internal/credentials/credentials.go

@@ -0,0 +1,35 @@
+/*
+ * 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 credentials
+
+import (
+	"context"
+)
+
+// clientHandshakeInfoKey is a struct used as the key to store
+// ClientHandshakeInfo in a context.
+type clientHandshakeInfoKey struct{}
+
+// ClientHandshakeInfoFromContext extracts the ClientHandshakeInfo from ctx.
+func ClientHandshakeInfoFromContext(ctx context.Context) any {
+	return ctx.Value(clientHandshakeInfoKey{})
+}
+
+// NewClientHandshakeInfoContext creates a context with chi.
+func NewClientHandshakeInfoContext(ctx context.Context, chi any) context.Context {
+	return context.WithValue(ctx, clientHandshakeInfoKey{}, chi)
+}

internal/credentials/spiffe.go

@@ -0,0 +1,75 @@
+/*
+ *
+ * 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 credentials defines APIs for parsing SPIFFE ID.
+//
+// All APIs in this package are experimental.
+package credentials
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"net/url"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+var logger = grpclog.Component("credentials")
+
+// SPIFFEIDFromState parses the SPIFFE ID from State. If the SPIFFE ID format
+// is invalid, return nil with warning.
+func SPIFFEIDFromState(state tls.ConnectionState) *url.URL {
+	if len(state.PeerCertificates) == 0 || len(state.PeerCertificates[0].URIs) == 0 {
+		return nil
+	}
+	return SPIFFEIDFromCert(state.PeerCertificates[0])
+}
+
+// SPIFFEIDFromCert parses the SPIFFE ID from x509.Certificate. If the SPIFFE
+// ID format is invalid, return nil with warning.
+func SPIFFEIDFromCert(cert *x509.Certificate) *url.URL {
+	if cert == nil || cert.URIs == nil {
+		return nil
+	}
+	var spiffeID *url.URL
+	for _, uri := range cert.URIs {
+		if uri == nil || uri.Scheme != "spiffe" || uri.Opaque != "" || (uri.User != nil && uri.User.Username() != "") {
+			continue
+		}
+		// From this point, we assume the uri is intended for a SPIFFE ID.
+		if len(uri.String()) > 2048 {
+			logger.Warning("invalid SPIFFE ID: total ID length larger than 2048 bytes")
+			return nil
+		}
+		if len(uri.Host) == 0 || len(uri.Path) == 0 {
+			logger.Warning("invalid SPIFFE ID: domain or workload ID is empty")
+			return nil
+		}
+		if len(uri.Host) > 255 {
+			logger.Warning("invalid SPIFFE ID: domain length larger than 255 characters")
+			return nil
+		}
+		// A valid SPIFFE certificate can only have exactly one URI SAN field.
+		if len(cert.URIs) > 1 {
+			logger.Warning("invalid SPIFFE ID: multiple URI SANs")
+			return nil
+		}
+		spiffeID = uri
+	}
+	return spiffeID
+}

internal/credentials/spiffe/spiffe.go

@@ -0,0 +1,107 @@
+/*
+ *
+ * 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 spiffe defines APIs for working with SPIFFE Bundle Maps.
+//
+// All APIs in this package are experimental.
+package spiffe
+
+import (
+	"crypto/x509"
+	"encoding/json"
+	"fmt"
+
+	"github.com/spiffe/go-spiffe/v2/bundle/spiffebundle"
+	"github.com/spiffe/go-spiffe/v2/spiffeid"
+)
+
+type partialParsedSPIFFEBundleMap struct {
+	Bundles map[string]json.RawMessage `json:"trust_domains"`
+}
+
+// BundleMapFromBytes parses bytes into a SPIFFE Bundle Map. See the
+// SPIFFE Bundle Map spec for more detail -
+// https://github.com/spiffe/spiffe/blob/main/standards/SPIFFE_Trust_Domain_and_Bundle.md#4-spiffe-bundle-format
+// If duplicate keys are encountered in the JSON parsing, Go's default unmarshal
+// behavior occurs which causes the last processed entry to be the entry in the
+// parsed map.
+func BundleMapFromBytes(bundleMapBytes []byte) (map[string]*spiffebundle.Bundle, error) {
+	var result partialParsedSPIFFEBundleMap
+	if err := json.Unmarshal(bundleMapBytes, &result); err != nil {
+		return nil, err
+	}
+	if result.Bundles == nil {
+		return nil, fmt.Errorf("spiffe: BundleMapFromBytes() no bundles parsed from spiffe bundle map bytes")
+	}
+	bundleMap := map[string]*spiffebundle.Bundle{}
+	for td, jsonBundle := range result.Bundles {
+		trustDomain, err := spiffeid.TrustDomainFromString(td)
+		if err != nil {
+			return nil, fmt.Errorf("spiffe: BundleMapFromBytes() invalid trust domain %q found when parsing SPIFFE Bundle Map: %v", td, err)
+		}
+		bundle, err := spiffebundle.Parse(trustDomain, jsonBundle)
+		if err != nil {
+			return nil, fmt.Errorf("spiffe: BundleMapFromBytes() failed to parse bundle for trust domain %q: %v", td, err)
+		}
+		bundleMap[td] = bundle
+	}
+	return bundleMap, nil
+}
+
+// GetRootsFromSPIFFEBundleMap returns the root trust certificates from the
+// SPIFFE bundle map for the given trust domain from the leaf certificate.
+func GetRootsFromSPIFFEBundleMap(bundleMap map[string]*spiffebundle.Bundle, leafCert *x509.Certificate) (*x509.CertPool, error) {
+	// 1. Upon receiving a peer certificate, verify that it is a well-formed SPIFFE
+	//    leaf certificate.  In particular, it must have a single URI SAN containing
+	//    a well-formed SPIFFE ID ([SPIFFE ID format]).
+	spiffeID, err := idFromCert(leafCert)
+	if err != nil {
+		return nil, fmt.Errorf("spiffe: could not get spiffe ID from peer leaf cert but verification with spiffe trust map was configured: %v", err)
+	}
+
+	// 2. Use the trust domain in the peer certificate's SPIFFE ID to lookup
+	//    the SPIFFE trust bundle. If the trust domain is not contained in the
+	//    configured trust map, reject the certificate.
+	spiffeBundle, ok := bundleMap[spiffeID.TrustDomain().Name()]
+	if !ok {
+		return nil, fmt.Errorf("spiffe: no bundle found for peer certificates trust domain %q but verification with a SPIFFE trust map was configured", spiffeID.TrustDomain().Name())
+	}
+	roots := spiffeBundle.X509Authorities()
+	rootPool := x509.NewCertPool()
+	for _, root := range roots {
+		rootPool.AddCert(root)
+	}
+	return rootPool, nil
+}
+
+// idFromCert parses the SPIFFE ID from the x509.Certificate. If the certificate
+// does not have a valid SPIFFE ID, returns an error.
+func idFromCert(cert *x509.Certificate) (*spiffeid.ID, error) {
+	if cert == nil {
+		return nil, fmt.Errorf("input cert is nil")
+	}
+	// A valid SPIFFE Certificate should have exactly one URI.
+	if len(cert.URIs) != 1 {
+		return nil, fmt.Errorf("input cert has %v URIs but should have 1", len(cert.URIs))
+	}
+	id, err := spiffeid.FromURI(cert.URIs[0])
+	if err != nil {
+		return nil, fmt.Errorf("invalid spiffeid: %v", err)
+	}
+	return &id, nil
+}

internal/credentials/spiffe/spiffe_test.go

@@ -0,0 +1,246 @@
+/*
+ *
+ * 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 spiffe
+
+import (
+	"crypto/x509"
+	"encoding/pem"
+	"io"
+	"net/url"
+	"os"
+	"strings"
+	"testing"
+
+	"google.golang.org/grpc/testdata"
+)
+
+const wantURI = "spiffe://foo.bar.com/client/workload/1"
+
+func loadFileBytes(t *testing.T, filePath string) []byte {
+	bytes, err := os.ReadFile(filePath)
+	if err != nil {
+		t.Fatalf("Error reading file: %v", err)
+	}
+	return bytes
+}
+
+func TestKnownSPIFFEBundle(t *testing.T) {
+	spiffeBundleFile := testdata.Path("spiffe/spiffebundle.json")
+	spiffeBundleBytes := loadFileBytes(t, spiffeBundleFile)
+	bundles, err := BundleMapFromBytes(spiffeBundleBytes)
+	if err != nil {
+		t.Fatalf("BundleMapFromBytes(%v) error during parsing: %v", spiffeBundleFile, err)
+	}
+	const wantBundleSize = 2
+	if len(bundles) != wantBundleSize {
+		t.Fatalf("BundleMapFromBytes(%v) did not parse correct bundle length. got %v want %v", spiffeBundleFile, len(bundles), wantBundleSize)
+	}
+	if bundles["example.com"] == nil {
+		t.Fatalf("BundleMapFromBytes(%v) got no bundle for example.com", spiffeBundleFile)
+	}
+	if bundles["test.example.com"] == nil {
+		t.Fatalf("BundleMapFromBytes(%v) got no bundle for test.example.com", spiffeBundleFile)
+	}
+
+	wantExampleComCert := loadX509Cert(t, testdata.Path("spiffe/spiffe_cert.pem"))
+	wantTestExampleComCert := loadX509Cert(t, testdata.Path("spiffe/server1_spiffe.pem"))
+	if !bundles["example.com"].X509Authorities()[0].Equal(wantExampleComCert) {
+		t.Errorf("BundleMapFromBytes(%v) parsed wrong cert for example.com.", spiffeBundleFile)
+	}
+	if !bundles["test.example.com"].X509Authorities()[0].Equal(wantTestExampleComCert) {
+		t.Errorf("BundleMapFromBytes(%v) parsed wrong cert for test.example.com", spiffeBundleFile)
+	}
+
+}
+
+func loadX509Cert(t *testing.T, filePath string) *x509.Certificate {
+	t.Helper()
+	certFile, _ := os.Open(filePath)
+	certRaw, _ := io.ReadAll(certFile)
+	block, _ := pem.Decode([]byte(certRaw))
+	if block == nil {
+		t.Fatalf("pem.Decode(%v) = nil. Want a value.", certRaw)
+	}
+	cert, err := x509.ParseCertificate(block.Bytes)
+	if err != nil {
+		t.Fatalf("x509.ParseCertificate(%v) failed %v", block.Bytes, err.Error())
+	}
+	return cert
+}
+
+func TestSPIFFEBundleMapFailures(t *testing.T) {
+	filePaths := []string{
+		testdata.Path("spiffe/spiffebundle_corrupted_cert.json"),
+		testdata.Path("spiffe/spiffebundle_malformed.json"),
+		testdata.Path("spiffe/spiffebundle_wrong_kid.json"),
+		testdata.Path("spiffe/spiffebundle_wrong_kty.json"),
+		testdata.Path("spiffe/spiffebundle_wrong_multi_certs.json"),
+		testdata.Path("spiffe/spiffebundle_wrong_root.json"),
+		testdata.Path("spiffe/spiffebundle_wrong_seq_type.json"),
+		testdata.Path("spiffe/spiffebundle_invalid_trustdomain.json"),
+		testdata.Path("spiffe/spiffebundle_empty_string_key.json"),
+		testdata.Path("spiffe/spiffebundle_empty_keys.json"),
+	}
+	for _, path := range filePaths {
+		t.Run(path, func(t *testing.T) {
+			bundleBytes := loadFileBytes(t, path)
+			if _, err := BundleMapFromBytes(bundleBytes); err == nil {
+				t.Fatalf("BundleMapFromBytes(%v) did not fail but should have", path)
+			}
+		})
+	}
+}
+
+func TestSPIFFEBundleMapX509Failures(t *testing.T) {
+	// SPIFFE Bundles only support a use of x509-svid and jwt-svid. If a
+	// use other than this is specified, the parser does not fail, it
+	// just doesn't add an x509 authority or jwt authority to the bundle
+	filePath := testdata.Path("spiffe/spiffebundle_wrong_use.json")
+	bundleBytes := loadFileBytes(t, filePath)
+	bundle, err := BundleMapFromBytes(bundleBytes)
+	if err != nil {
+		t.Fatalf("BundleMapFromBytes(%v) failed with error: %v", filePath, err)
+	}
+	if len(bundle["example.com"].X509Authorities()) != 0 {
+		t.Fatalf("BundleMapFromBytes(%v) did not have empty bundle but should have", filePath)
+	}
+}
+
+func TestGetRootsFromSPIFFEBundleMapSuccess(t *testing.T) {
+	bundleMapFile := testdata.Path("spiffe/spiffebundle_match_client_spiffe.json")
+	bundleBytes := loadFileBytes(t, bundleMapFile)
+	bundle, err := BundleMapFromBytes(bundleBytes)
+	if err != nil {
+		t.Fatalf("BundleMapFromBytes(%v) failed with error: %v", bundleMapFile, err)
+	}
+
+	cert := loadX509Cert(t, testdata.Path("spiffe/client_spiffe.pem"))
+	gotRoots, err := GetRootsFromSPIFFEBundleMap(bundle, cert)
+	if err != nil {
+		t.Fatalf("GetRootsFromSPIFFEBundleMap() failed with err %v", err)
+	}
+	wantRoot := loadX509Cert(t, testdata.Path("spiffe/spiffe_cert.pem"))
+	wantRoots := x509.NewCertPool()
+	wantRoots.AddCert(wantRoot)
+	if !gotRoots.Equal(wantRoots) {
+		t.Fatalf("GetRootsFromSPIFFEBundleMap() got %v want %v", gotRoots, wantRoots)
+	}
+}
+
+func TestGetRootsFromSPIFFEBundleMapFailures(t *testing.T) {
+	bundleMapFile := testdata.Path("spiffe/spiffebundle.json")
+	bundleBytes := loadFileBytes(t, bundleMapFile)
+	bundle, err := BundleMapFromBytes(bundleBytes)
+	certWithTwoURIs := loadX509Cert(t, testdata.Path("spiffe/client_spiffe.pem"))
+	certWithTwoURIs.URIs = append(certWithTwoURIs.URIs, certWithTwoURIs.URIs[0])
+	certWithNoURIs := loadX509Cert(t, testdata.Path("spiffe/client_spiffe.pem"))
+	certWithNoURIs.URIs = nil
+	if err != nil {
+		t.Fatalf("BundleMapFromBytes(%v) failed with error: %v", bundleMapFile, err)
+	}
+	tests := []struct {
+		name          string
+		bundleMapFile string
+		leafCert      *x509.Certificate
+		wantErr       string
+	}{
+		{
+			name:     "no bundle for peer cert spiffeID",
+			leafCert: loadX509Cert(t, testdata.Path("spiffe/client_spiffe.pem")),
+			wantErr:  "no bundle found for peer certificates",
+		},
+		{
+			name:     "cert has invalid SPIFFE id",
+			leafCert: loadX509Cert(t, testdata.Path("ca.pem")),
+			wantErr:  "could not get spiffe ID from peer leaf cert",
+		},
+		{
+			name:     "nil cert",
+			leafCert: nil,
+			wantErr:  "input cert is nil",
+		},
+		{
+			name:     "cert has multiple URIs",
+			leafCert: certWithTwoURIs,
+			wantErr:  "input cert has 2 URIs but should have 1",
+		},
+		{
+			name:     "cert has no URIs",
+			leafCert: certWithNoURIs,
+			wantErr:  "input cert has 0 URIs but should have 1",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err = GetRootsFromSPIFFEBundleMap(bundle, tc.leafCert)
+			if err == nil {
+				t.Fatalf("GetRootsFromSPIFFEBundleMap() got no error but want error containing %v.", tc.wantErr)
+			}
+			if !strings.Contains(err.Error(), tc.wantErr) {
+				t.Fatalf("GetRootsFromSPIFFEBundleMap() got error: %v. want error to contain %v", err, tc.wantErr)
+			}
+		})
+	}
+}
+
+func TestIDFromCert(t *testing.T) {
+	cert := loadX509Cert(t, testdata.Path("x509/spiffe_cert.pem"))
+	uri, err := idFromCert(cert)
+	if err != nil {
+		t.Fatalf("idFromCert() failed with err: %v", err)
+	}
+	if uri != nil && uri.String() != wantURI {
+		t.Fatalf("ID not expected, got %s, want %s", uri.String(), wantURI)
+	}
+}
+
+func TestIDFromCertFileFailures(t *testing.T) {
+	certWithNoURIs := loadX509Cert(t, testdata.Path("spiffe/client_spiffe.pem"))
+	certWithNoURIs.URIs = nil
+	certWithInvalidSPIFFEID := loadX509Cert(t, testdata.Path("spiffe/client_spiffe.pem"))
+	certWithInvalidSPIFFEID.URIs = []*url.URL{{Path: "non-spiffe.bad"}}
+	tests := []struct {
+		name string
+		cert *x509.Certificate
+	}{
+		{
+			name: "certificate with multiple URIs",
+			cert: loadX509Cert(t, testdata.Path("x509/multiple_uri_cert.pem")),
+		},
+		{
+			name: "certificate with invalidSPIFFE ID",
+			cert: certWithInvalidSPIFFEID,
+		},
+		{
+			name: "nil cert",
+			cert: nil,
+		},
+		{
+			name: "cert with no URIs",
+			cert: certWithNoURIs,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if _, err := idFromCert(tt.cert); err == nil {
+				t.Fatalf("idFromCert() succeeded but want error")
+			}
+		})
+	}
+}

internal/credentials/spiffe_test.go

@@ -0,0 +1,233 @@
+/*
+ *
+ * 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 credentials
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/pem"
+	"net/url"
+	"os"
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/testdata"
+)
+
+const wantURI = "spiffe://foo.bar.com/client/workload/1"
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestSPIFFEIDFromState(t *testing.T) {
+	tests := []struct {
+		name string
+		urls []*url.URL
+		// If we expect a SPIFFE ID to be returned.
+		wantID bool
+	}{
+		{
+			name:   "empty URIs",
+			urls:   []*url.URL{},
+			wantID: false,
+		},
+		{
+			name: "good SPIFFE ID",
+			urls: []*url.URL{
+				{
+					Scheme:  "spiffe",
+					Host:    "foo.bar.com",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+			},
+			wantID: true,
+		},
+		{
+			name: "invalid host",
+			urls: []*url.URL{
+				{
+					Scheme:  "spiffe",
+					Host:    "",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+			},
+			wantID: false,
+		},
+		{
+			name: "invalid path",
+			urls: []*url.URL{
+				{
+					Scheme:  "spiffe",
+					Host:    "foo.bar.com",
+					Path:    "",
+					RawPath: "",
+				},
+			},
+			wantID: false,
+		},
+		{
+			name: "large path",
+			urls: []*url.URL{
+				{
+					Scheme:  "spiffe",
+					Host:    "foo.bar.com",
+					Path:    string(make([]byte, 2050)),
+					RawPath: string(make([]byte, 2050)),
+				},
+			},
+			wantID: false,
+		},
+		{
+			name: "large host",
+			urls: []*url.URL{
+				{
+					Scheme:  "spiffe",
+					Host:    string(make([]byte, 256)),
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+			},
+			wantID: false,
+		},
+		{
+			name: "multiple URI SANs",
+			urls: []*url.URL{
+				{
+					Scheme:  "spiffe",
+					Host:    "foo.bar.com",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+				{
+					Scheme:  "spiffe",
+					Host:    "bar.baz.com",
+					Path:    "workload/wl2",
+					RawPath: "workload/wl2",
+				},
+				{
+					Scheme:  "https",
+					Host:    "foo.bar.com",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+			},
+			wantID: false,
+		},
+		{
+			name: "multiple URI SANs without SPIFFE ID",
+			urls: []*url.URL{
+				{
+					Scheme:  "https",
+					Host:    "foo.bar.com",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+				{
+					Scheme:  "ssh",
+					Host:    "foo.bar.com",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+			},
+			wantID: false,
+		},
+		{
+			name: "multiple URI SANs with one SPIFFE ID",
+			urls: []*url.URL{
+				{
+					Scheme:  "spiffe",
+					Host:    "foo.bar.com",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+				{
+					Scheme:  "https",
+					Host:    "foo.bar.com",
+					Path:    "workload/wl1",
+					RawPath: "workload/wl1",
+				},
+			},
+			wantID: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			state := tls.ConnectionState{PeerCertificates: []*x509.Certificate{{URIs: tt.urls}}}
+			id := SPIFFEIDFromState(state)
+			if got, want := id != nil, tt.wantID; got != want {
+				t.Errorf("want wantID = %v, but SPIFFE ID is %v", want, id)
+			}
+		})
+	}
+}
+
+func (s) TestSPIFFEIDFromCert(t *testing.T) {
+	tests := []struct {
+		name     string
+		dataPath string
+		// If we expect a SPIFFE ID to be returned.
+		wantID bool
+	}{
+		{
+			name:     "good certificate with SPIFFE ID",
+			dataPath: "x509/spiffe_cert.pem",
+			wantID:   true,
+		},
+		{
+			name:     "bad certificate with SPIFFE ID and another URI",
+			dataPath: "x509/multiple_uri_cert.pem",
+			wantID:   false,
+		},
+		{
+			name:     "certificate without SPIFFE ID",
+			dataPath: "x509/client1_cert.pem",
+			wantID:   false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			data, err := os.ReadFile(testdata.Path(tt.dataPath))
+			if err != nil {
+				t.Fatalf("os.ReadFile(%s) failed: %v", testdata.Path(tt.dataPath), err)
+			}
+			block, _ := pem.Decode(data)
+			if block == nil {
+				t.Fatalf("Failed to parse the certificate: byte block is nil")
+			}
+			cert, err := x509.ParseCertificate(block.Bytes)
+			if err != nil {
+				t.Fatalf("x509.ParseCertificate(%b) failed: %v", block.Bytes, err)
+			}
+			uri := SPIFFEIDFromCert(cert)
+			if (uri != nil) != tt.wantID {
+				t.Fatalf("wantID got and want mismatch, got %t, want %t", uri != nil, tt.wantID)
+			}
+			if uri != nil && uri.String() != wantURI {
+				t.Fatalf("SPIFFE ID not expected, got %s, want %s", uri.String(), wantURI)
+			}
+		})
+	}
+}

internal/credentials/syscallconn.go

@@ -0,0 +1,58 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package credentials
+
+import (
+	"net"
+	"syscall"
+)
+
+type sysConn = syscall.Conn
+
+// syscallConn keeps reference of rawConn to support syscall.Conn for channelz.
+// SyscallConn() (the method in interface syscall.Conn) is explicitly
+// implemented on this type,
+//
+// Interface syscall.Conn is implemented by most net.Conn implementations (e.g.
+// TCPConn, UnixConn), but is not part of net.Conn interface. So wrapper conns
+// that embed net.Conn don't implement syscall.Conn. (Side note: tls.Conn
+// doesn't embed net.Conn, so even if syscall.Conn is part of net.Conn, it won't
+// help here).
+type syscallConn struct {
+	net.Conn
+	// sysConn is a type alias of syscall.Conn. It's necessary because the name
+	// `Conn` collides with `net.Conn`.
+	sysConn
+}
+
+// WrapSyscallConn tries to wrap rawConn and newConn into a net.Conn that
+// implements syscall.Conn. rawConn will be used to support syscall, and newConn
+// will be used for read/write.
+//
+// This function returns newConn if rawConn doesn't implement syscall.Conn.
+func WrapSyscallConn(rawConn, newConn net.Conn) net.Conn {
+	sysConn, ok := rawConn.(syscall.Conn)
+	if !ok {
+		return newConn
+	}
+	return &syscallConn{
+		Conn:    newConn,
+		sysConn: sysConn,
+	}
+}

internal/credentials/syscallconn_test.go

@@ -0,0 +1,56 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package credentials
+
+import (
+	"net"
+	"syscall"
+	"testing"
+)
+
+func (*syscallConn) SyscallConn() (syscall.RawConn, error) {
+	return nil, nil
+}
+
+type nonSyscallConn struct {
+	net.Conn
+}
+
+func (s) TestWrapSyscallConn(t *testing.T) {
+	sc := &syscallConn{}
+	nsc := &nonSyscallConn{}
+
+	wrapConn := WrapSyscallConn(sc, nsc)
+	if _, ok := wrapConn.(syscall.Conn); !ok {
+		t.Errorf("returned conn (type %T) doesn't implement syscall.Conn, want implement", wrapConn)
+	}
+}
+
+func (s) TestWrapSyscallConnNoWrap(t *testing.T) {
+	nscRaw := &nonSyscallConn{}
+	nsc := &nonSyscallConn{}
+
+	wrapConn := WrapSyscallConn(nscRaw, nsc)
+	if _, ok := wrapConn.(syscall.Conn); ok {
+		t.Errorf("returned conn (type %T) implements syscall.Conn, want not implement", wrapConn)
+	}
+	if wrapConn != nsc {
+		t.Errorf("returned conn is %p, want %p (the passed-in newConn)", wrapConn, nsc)
+	}
+}

internal/credentials/util.go

@@ -0,0 +1,52 @@
+/*
+ *
+ * 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 credentials
+
+import (
+	"crypto/tls"
+)
+
+const alpnProtoStrH2 = "h2"
+
+// AppendH2ToNextProtos appends h2 to next protos.
+func AppendH2ToNextProtos(ps []string) []string {
+	for _, p := range ps {
+		if p == alpnProtoStrH2 {
+			return ps
+		}
+	}
+	ret := make([]string, 0, len(ps)+1)
+	ret = append(ret, ps...)
+	return append(ret, alpnProtoStrH2)
+}
+
+// CloneTLSConfig returns a shallow clone of the exported
+// fields of cfg, ignoring the unexported sync.Once, which
+// contains a mutex and must not be copied.
+//
+// If cfg is nil, a new zero tls.Config is returned.
+//
+// TODO: inline this function if possible.
+func CloneTLSConfig(cfg *tls.Config) *tls.Config {
+	if cfg == nil {
+		return &tls.Config{}
+	}
+
+	return cfg.Clone()
+}

internal/credentials/util_test.go

@@ -0,0 +1,60 @@
+/*
+ *
+ * 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 credentials
+
+import (
+	"reflect"
+	"testing"
+)
+
+func (s) TestAppendH2ToNextProtos(t *testing.T) {
+	tests := []struct {
+		name string
+		ps   []string
+		want []string
+	}{
+		{
+			name: "empty",
+			ps:   nil,
+			want: []string{"h2"},
+		},
+		{
+			name: "only h2",
+			ps:   []string{"h2"},
+			want: []string{"h2"},
+		},
+		{
+			name: "with h2",
+			ps:   []string{"alpn", "h2"},
+			want: []string{"alpn", "h2"},
+		},
+		{
+			name: "no h2",
+			ps:   []string{"alpn"},
+			want: []string{"alpn", "h2"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := AppendH2ToNextProtos(tt.ps); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("AppendH2ToNextProtos() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

internal/credentials/xds/handshake_info.go

@@ -0,0 +1,469 @@
+/*
+ *
+ * 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 xds contains non-user facing functionality of the xds credentials.
+package xds
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"fmt"
+	"strings"
+	"sync/atomic"
+
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/credentials/tls/certprovider"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/credentials/spiffe"
+	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/xds/matcher"
+	"google.golang.org/grpc/resolver"
+)
+
+func init() {
+	internal.GetXDSHandshakeInfoForTesting = HandshakeInfoFromAttributes
+}
+
+// handshakeAttrKey is the type used as the key to store HandshakeInfo in
+// the Attributes field of resolver.Address.
+type handshakeAttrKey struct{}
+
+// hostnameKey is the type used as the key to store the hostname in the
+// Attributes field of resolver.Address.
+type hostnameKey struct{}
+
+// SetAddressHostname returns a copy of addr in which the Attributes field is
+// updated with the provided hostname.
+func SetAddressHostname(addr resolver.Address, hostname string) resolver.Address {
+	addr.Attributes = addr.Attributes.WithValue(hostnameKey{}, hostname)
+	return addr
+}
+
+// Hostname returns the endpoint hostname stored in attr.
+func Hostname(attr *attributes.Attributes) string {
+	if attr == nil {
+		return ""
+	}
+	v := attr.Value(hostnameKey{})
+	hn, _ := v.(string)
+	return hn
+}
+
+// SetHandshakeInfo returns a copy of addr in which the Attributes field is
+// updated with hiPtr.
+func SetHandshakeInfo(addr resolver.Address, hiPtr *atomic.Pointer[grpcsync.RefCounted[*HandshakeInfo]]) resolver.Address {
+	addr.Attributes = addr.Attributes.WithValue(handshakeAttrKey{}, hiPtr)
+	return addr
+}
+
+// HandshakeInfoFromAttributes returns the atomic pointer to the
+// reference-counted HandshakeInfo stored in attr.
+func HandshakeInfoFromAttributes(attr *attributes.Attributes) *atomic.Pointer[grpcsync.RefCounted[*HandshakeInfo]] {
+	v := attr.Value(handshakeAttrKey{})
+	hi, _ := v.(*atomic.Pointer[grpcsync.RefCounted[*HandshakeInfo]])
+	return hi
+}
+
+// HandshakeInfo wraps all the security configuration required by client and
+// server handshake methods in xds credentials. The xDS implementation will be
+// responsible for populating these fields.
+type HandshakeInfo struct {
+	// The fields below are initialized when the HandshakeInfo is created, and are
+	// not changed thereafter.
+	sanMatchers         []matcher.StringMatcher // Only on the client side.
+	requireClientCert   bool                    // Only on server side.
+	sni                 string                  // Only on client side, used for Server Name Indication in TLS handshake.
+	validateSANUsingSNI bool                    // Only on client side, indicates whether to perform validation of SANs based on SNI value.
+	useAutoHostSNI      bool                    // Only on client side, indicates whether to use endpoint hostname as SNI.
+
+	rootProvider     certprovider.Provider
+	identityProvider certprovider.Provider
+}
+
+// NewHandshakeInfo returns a new reference counted HandshakeInfo configured
+// with the provided options.
+func NewHandshakeInfo(rootProvider, identityProvider certprovider.Provider, sanMatchers []matcher.StringMatcher, sni string, requireClientCert, validateSANUsingSNI, useAutoHostSNI bool) *grpcsync.RefCounted[*HandshakeInfo] {
+	hi := &HandshakeInfo{
+		rootProvider:        rootProvider,
+		identityProvider:    identityProvider,
+		sanMatchers:         sanMatchers,
+		sni:                 sni,
+		requireClientCert:   requireClientCert,
+		validateSANUsingSNI: validateSANUsingSNI,
+		useAutoHostSNI:      useAutoHostSNI,
+	}
+
+	return grpcsync.NewRefCounted(hi, hi.close)
+}
+
+func (hi *HandshakeInfo) close() {
+	if hi.rootProvider != nil {
+		hi.rootProvider.Close()
+	}
+	if hi.identityProvider != nil {
+		hi.identityProvider.Close()
+	}
+}
+
+// ClientSideTLSConfig loads the HandshakeInfo from hiPtr, marks it as in-use,
+// and returns the tls.Config along with a done callback that MUST be invoked
+// when the handshake completes. If no HandshakeInfo is stored in hiPtr or if
+// fallback credentials should be used, useFallback returns true.
+func ClientSideTLSConfig(ctx context.Context, hiPtr *atomic.Pointer[grpcsync.RefCounted[*HandshakeInfo]], hostname string) (cfg *tls.Config, useFallback bool, done func(), err error) {
+	if hiPtr == nil {
+		return nil, true, func() {}, nil
+	}
+	for {
+		hiRC := hiPtr.Load()
+		if hiRC == nil {
+			return nil, false, func() {}, errors.New("xds: connection closed or HandshakeInfo dead")
+		}
+		if !hiRC.TryIncrement() {
+			if hiPtr.Load() != hiRC {
+				continue
+			}
+			return nil, false, func() {}, errors.New("xds: connection closed or HandshakeInfo dead")
+		}
+
+		hi := hiRC.Value()
+		if hi == nil || hi.UseFallbackCreds() {
+			hiRC.Decrement()
+			return nil, true, func() {}, nil
+		}
+		cfg, err := hi.clientSideTLSConfigInternal(ctx, hostname)
+		if err != nil {
+			hiRC.Decrement()
+			return nil, false, func() {}, err
+		}
+		return cfg, false, hiRC.Decrement, nil
+	}
+}
+
+// ServerSideTLSConfig checks if hi is configured and marks it as in-use,
+// returning the tls.Config along with a done callback that MUST be invoked when
+// the handshake completes. If hi is nil or fallback credentials should be used,
+// useFallback returns true.
+func ServerSideTLSConfig(ctx context.Context, hiRC *grpcsync.RefCounted[*HandshakeInfo]) (cfg *tls.Config, useFallback bool, done func(), err error) {
+	if hiRC == nil {
+		return nil, true, func() {}, nil
+	}
+	hi := hiRC.Value()
+	if hi == nil || hi.UseFallbackCreds() {
+		return nil, true, func() {}, nil
+	}
+	if !hiRC.TryIncrement() {
+		return nil, false, func() {}, errors.New("xds: connection closed or HandshakeInfo dead")
+	}
+	cfg, err = hi.serverSideTLSConfigInternal(ctx)
+	if err != nil {
+		hiRC.Decrement()
+		return nil, false, func() {}, err
+	}
+	return cfg, false, hiRC.Decrement, nil
+}
+
+// UseFallbackCreds returns true when fallback credentials are to be used based
+// on the contents of the HandshakeInfo.
+func (hi *HandshakeInfo) UseFallbackCreds() bool {
+	if hi == nil {
+		return true
+	}
+	return hi.identityProvider == nil && hi.rootProvider == nil
+}
+
+// GetSANMatchersForTesting returns the SAN matchers stored in HandshakeInfo.
+// To be used only for testing purposes.
+func (hi *HandshakeInfo) GetSANMatchersForTesting() []matcher.StringMatcher {
+	return append([]matcher.StringMatcher{}, hi.sanMatchers...)
+}
+
+// clientSideTLSConfigInternal constructs a tls.Config to be used in a
+// client-side handshake based on the contents of the HandshakeInfo.
+//
+// hostname is passed as a parameter here instead of being part of the
+// HandshakeInfo because HandshakeInfo contains cluster-level security
+// configuration that applies to all endpoints in the cluster, while hostname is
+// specific to each endpoint. This allows sharing a single HandshakeInfo
+// instance across multiple endpoints in the same cluster.
+func (hi *HandshakeInfo) clientSideTLSConfigInternal(ctx context.Context, hostname string) (*tls.Config, error) {
+	// On the client side, rootProvider is mandatory. IdentityProvider is
+	// optional based on whether the client is doing TLS or mTLS.
+	if hi.rootProvider == nil {
+		return nil, errors.New("xds: CertificateProvider to fetch trusted roots is missing, cannot perform TLS handshake. Please check configuration on the management server")
+	}
+
+	// InsecureSkipVerify needs to be set to true because we need to perform
+	// custom verification to check the SAN on the received certificate.
+	// Currently the Go stdlib does complete verification of the cert (which
+	// includes hostname verification) or none. We are forced to go with the
+	// latter and perform the normal cert validation ourselves.
+	cfg := &tls.Config{
+		InsecureSkipVerify: true,
+		NextProtos:         []string{"h2"},
+	}
+
+	km, err := hi.rootProvider.KeyMaterial(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("xds: fetching trusted roots from CertificateProvider failed: %v", err)
+	}
+	cfg.RootCAs = km.Roots
+
+	// If AutoHostSNI is true, and the endpoint hostname is present, we use the
+	// endpoint hostname as the SNI value and also for SAN validation.
+	// Otherwise, we use the SNI value from HandshakeInfo (which is configured
+	// by the control plane) and validating SANs based on that.
+	sni := hi.sni
+	if hi.useAutoHostSNI && hostname != "" {
+		sni = hostname
+	}
+
+	cfg.VerifyPeerCertificate = hi.buildVerifyFunc(km, true, sni)
+
+	if hi.identityProvider != nil {
+		km, err := hi.identityProvider.KeyMaterial(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("xds: fetching identity certificates from CertificateProvider failed: %v", err)
+		}
+		cfg.Certificates = km.Certs
+	}
+
+	if envconfig.XDSSNIEnabled && sni != "" {
+		cfg.ServerName = sni
+	}
+	return cfg, nil
+}
+
+func (hi *HandshakeInfo) buildVerifyFunc(km *certprovider.KeyMaterial, isClient bool, sni string) func(rawCerts [][]byte, _ [][]*x509.Certificate) error {
+	return func(rawCerts [][]byte, _ [][]*x509.Certificate) error {
+		if len(rawCerts) == 0 {
+			return fmt.Errorf("xds: no peer certificates presented")
+		}
+		// Parse all raw certificates presented by the peer.
+		var certs []*x509.Certificate
+		for _, rc := range rawCerts {
+			cert, err := x509.ParseCertificate(rc)
+			if err != nil {
+				return err
+			}
+			certs = append(certs, cert)
+		}
+
+		// Build the intermediates list and verify that the leaf certificate is
+		// signed by one of the root certificates. If a SPIFFE Bundle Map is
+		// configured, it is used to get the root certs. Otherwise, the
+		// configured roots in the root provider are used.
+		intermediates := x509.NewCertPool()
+		for _, cert := range certs[1:] {
+			intermediates.AddCert(cert)
+		}
+		roots := km.Roots
+		// If a SPIFFE Bundle Map is configured, find the roots for the trust
+		// domain of the leaf certificate.
+		if km.SPIFFEBundleMap != nil {
+			var err error
+			roots, err = spiffe.GetRootsFromSPIFFEBundleMap(km.SPIFFEBundleMap, certs[0])
+			if err != nil {
+				return err
+			}
+		}
+		opts := x509.VerifyOptions{
+			Roots:         roots,
+			Intermediates: intermediates,
+			KeyUsages:     []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
+		}
+		if isClient {
+			opts.KeyUsages = []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}
+		} else {
+			opts.KeyUsages = []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth}
+		}
+		if _, err := certs[0].Verify(opts); err != nil {
+			return err
+		}
+
+		// If XDSSNIEnabled and AutoSNISANValidation are both true and the SNI is
+		// non-empty, validate only DNS SANs against the SNI. Otherwise, fallback to
+		// validating all received SANs against the control plane provided SAN
+		// matchers.
+		if envconfig.XDSSNIEnabled && hi.validateSANUsingSNI && sni != "" {
+			// Verify SAN of leaf certificate with SNI using exact DNS matcher.
+			for _, san := range certs[0].DNSNames {
+				if dnsMatch(sni, san) {
+					return nil
+				}
+			}
+			return fmt.Errorf("xds: received DNS SANs: %v do not match the SNI: %s", certs[0].DNSNames, sni)
+		}
+		// The SANs sent by the xDS control plane are encoded as SPIFFE IDs. We need to
+		// only look at the SANs on the leaf cert.
+		if cert := certs[0]; !hi.MatchingSANExists(cert) {
+			// TODO: Print the complete certificate once the x509 package
+			// supports a String() method on the Certificate type.
+			return fmt.Errorf("xds: received SANs {DNSNames: %v, EmailAddresses: %v, IPAddresses: %v, URIs: %v} do not match any of the accepted SANs", cert.DNSNames, cert.EmailAddresses, cert.IPAddresses, cert.URIs)
+		}
+		return nil
+	}
+}
+
+// serverSideTLSConfigInternal constructs a tls.Config to be used in a
+// server-side handshake based on the contents of the HandshakeInfo.
+func (hi *HandshakeInfo) serverSideTLSConfigInternal(ctx context.Context) (*tls.Config, error) {
+	cfg := &tls.Config{
+		ClientAuth: tls.NoClientCert,
+		NextProtos: []string{"h2"},
+	}
+	// On the server side, identityProvider is mandatory. RootProvider is
+	// optional based on whether the server is doing TLS or mTLS.
+	if hi.identityProvider == nil {
+		return nil, errors.New("xds: CertificateProvider to fetch identity certificate is missing, cannot perform TLS handshake. Please check configuration on the management server")
+	}
+	if hi.requireClientCert {
+		cfg.ClientAuth = tls.RequireAndVerifyClientCert
+	}
+
+	// identityProvider is mandatory on the server side.
+	km, err := hi.identityProvider.KeyMaterial(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("xds: fetching identity certificates from CertificateProvider failed: %v", err)
+	}
+	cfg.Certificates = km.Certs
+
+	if hi.rootProvider != nil {
+		km, err := hi.rootProvider.KeyMaterial(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("xds: fetching trusted roots from CertificateProvider failed: %v", err)
+		}
+		if km.SPIFFEBundleMap != nil && hi.requireClientCert {
+			// ClientAuth, if set greater than tls.RequireAnyClientCert, must be
+			// dropped to tls.RequireAnyClientCert so that custom verification
+			// to use SPIFFE Bundles is done.
+			cfg.ClientAuth = tls.RequireAnyClientCert
+			cfg.VerifyPeerCertificate = hi.buildVerifyFunc(km, false, "")
+		} else {
+			cfg.ClientCAs = km.Roots
+		}
+	}
+	return cfg, nil
+}
+
+// MatchingSANExists returns true if the SANs contained in cert match the
+// criteria enforced by the list of SAN matchers in HandshakeInfo.
+//
+// If the list of SAN matchers in the HandshakeInfo is empty, this function
+// returns true for all input certificates.
+func (hi *HandshakeInfo) MatchingSANExists(cert *x509.Certificate) bool {
+	if len(hi.sanMatchers) == 0 {
+		return true
+	}
+
+	// SANs can be specified in any of these four fields on the parsed cert.
+	for _, san := range cert.DNSNames {
+		if hi.matchSAN(san, true) {
+			return true
+		}
+	}
+	for _, san := range cert.EmailAddresses {
+		if hi.matchSAN(san, false) {
+			return true
+		}
+	}
+	for _, san := range cert.IPAddresses {
+		if hi.matchSAN(san.String(), false) {
+			return true
+		}
+	}
+	for _, san := range cert.URIs {
+		if hi.matchSAN(san.String(), false) {
+			return true
+		}
+	}
+	return false
+}
+
+// Caller must hold mu.
+func (hi *HandshakeInfo) matchSAN(san string, isDNS bool) bool {
+	for _, matcher := range hi.sanMatchers {
+		if em := matcher.ExactMatch(); em != "" && isDNS {
+			// This is a special case which is documented in the xDS protos.
+			// If the DNS SAN is a wildcard entry, and the match criteria is
+			// `exact`, then we need to perform DNS wildcard matching
+			// instead of regular string comparison.
+			if dnsMatch(em, san) {
+				return true
+			}
+			continue
+		}
+		if matcher.Match(san) {
+			return true
+		}
+	}
+	return false
+}
+
+// dnsMatch implements a DNS wildcard matching algorithm based on RFC2828 and
+// grpc-java's implementation in `OkHostnameVerifier` class.
+//
+// NOTE: Here the `host` argument is the one from the set of string matchers in
+// the xDS proto and the `san` argument is a DNS SAN from the certificate, and
+// this is the one which can potentially contain a wildcard pattern.
+func dnsMatch(host, san string) bool {
+	// Add trailing "." and turn them into absolute domain names.
+	if !strings.HasSuffix(host, ".") {
+		host += "."
+	}
+	if !strings.HasSuffix(san, ".") {
+		san += "."
+	}
+	// Domain names are case-insensitive.
+	host = strings.ToLower(host)
+	san = strings.ToLower(san)
+
+	// If san does not contain a wildcard, do exact match.
+	if !strings.Contains(san, "*") {
+		return host == san
+	}
+
+	// Wildcard dns matching rules
+	// - '*' is only permitted in the left-most label and must be the only
+	//   character in that label. For example, *.example.com is permitted, while
+	//   *a.example.com, a*.example.com, a*b.example.com, a.*.example.com are
+	//   not permitted.
+	// - '*' matches a single domain name component. For example, *.example.com
+	//   matches test.example.com but does not match sub.test.example.com.
+	// - Wildcard patterns for single-label domain names are not permitted.
+	if san == "*." || !strings.HasPrefix(san, "*.") || strings.Contains(san[1:], "*") {
+		return false
+	}
+	// Optimization: at this point, we know that the san contains a '*' and
+	// is the first domain component of san. So, the host name must be at
+	// least as long as the san to be able to match.
+	if len(host) < len(san) {
+		return false
+	}
+	// Hostname must end with the non-wildcard portion of san.
+	if !strings.HasSuffix(host, san[1:]) {
+		return false
+	}
+	// At this point we know that the hostName and san share the same suffix
+	// (the non-wildcard portion of san). Now, we just need to make sure
+	// that the '*' does not match across domain components.
+	hostPrefix := strings.TrimSuffix(host, san[1:])
+	return !strings.Contains(hostPrefix, ".")
+}

internal/credentials/xds/handshake_info_test.go

@@ -0,0 +1,668 @@
+/*
+ *
+ * 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 xds
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/tls"
+	"crypto/x509"
+	"fmt"
+	"net"
+	"net/netip"
+	"net/url"
+	"os"
+	"regexp"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/testdata"
+
+	"google.golang.org/grpc/credentials/tls/certprovider"
+	"google.golang.org/grpc/internal/credentials/spiffe"
+	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/internal/xds/matcher"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestDNSMatch(t *testing.T) {
+	tests := []struct {
+		desc      string
+		host      string
+		pattern   string
+		wantMatch bool
+	}{
+		{
+			desc:      "invalid wildcard 1",
+			host:      "aa.example.com",
+			pattern:   "*a.example.com",
+			wantMatch: false,
+		},
+		{
+			desc:      "invalid wildcard 2",
+			host:      "aa.example.com",
+			pattern:   "a*.example.com",
+			wantMatch: false,
+		},
+		{
+			desc:      "invalid wildcard 3",
+			host:      "abc.example.com",
+			pattern:   "a*c.example.com",
+			wantMatch: false,
+		},
+		{
+			desc:      "wildcard in one of the middle components",
+			host:      "abc.test.example.com",
+			pattern:   "abc.*.example.com",
+			wantMatch: false,
+		},
+		{
+			desc:      "single component wildcard",
+			host:      "a.example.com",
+			pattern:   "*",
+			wantMatch: false,
+		},
+		{
+			desc:      "short host name",
+			host:      "a.com",
+			pattern:   "*.example.com",
+			wantMatch: false,
+		},
+		{
+			desc:      "suffix mismatch",
+			host:      "a.notexample.com",
+			pattern:   "*.example.com",
+			wantMatch: false,
+		},
+		{
+			desc:      "wildcard match across components",
+			host:      "sub.test.example.com",
+			pattern:   "*.example.com.",
+			wantMatch: false,
+		},
+		{
+			desc:      "host doesn't end in period",
+			host:      "test.example.com",
+			pattern:   "test.example.com.",
+			wantMatch: true,
+		},
+		{
+			desc:      "pattern doesn't end in period",
+			host:      "test.example.com.",
+			pattern:   "test.example.com",
+			wantMatch: true,
+		},
+		{
+			desc:      "case insensitive",
+			host:      "TEST.EXAMPLE.COM.",
+			pattern:   "test.example.com.",
+			wantMatch: true,
+		},
+		{
+			desc:      "simple match",
+			host:      "test.example.com",
+			pattern:   "test.example.com",
+			wantMatch: true,
+		},
+		{
+			desc:      "good wildcard",
+			host:      "a.example.com",
+			pattern:   "*.example.com",
+			wantMatch: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			gotMatch := dnsMatch(test.host, test.pattern)
+			if gotMatch != test.wantMatch {
+				t.Fatalf("dnsMatch(%s, %s) = %v, want %v", test.host, test.pattern, gotMatch, test.wantMatch)
+			}
+		})
+	}
+}
+
+func (s) TestMatchingSANExists_FailureCases(t *testing.T) {
+	url1, err := url.Parse("http://golang.org")
+	if err != nil {
+		t.Fatalf("url.Parse() failed: %v", err)
+	}
+	url2, err := url.Parse("https://github.com/grpc/grpc-go")
+	if err != nil {
+		t.Fatalf("url.Parse() failed: %v", err)
+	}
+	inputCert := &x509.Certificate{
+		DNSNames:       []string{"foo.bar.example.com", "bar.baz.test.com", "*.example.com"},
+		EmailAddresses: []string{"foobar@example.com", "barbaz@test.com"},
+		IPAddresses: []net.IP{
+			netip.MustParseAddr("192.0.0.1").AsSlice(),
+			netip.MustParseAddr("2001:db8::68").AsSlice(),
+		},
+		URIs: []*url.URL{url1, url2},
+	}
+
+	tests := []struct {
+		desc        string
+		sanMatchers []matcher.StringMatcher
+	}{
+		{
+			desc: "exact match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("abcd.test.com", false),
+				matcher.NewExactStringMatcher("http://golang", false),
+				matcher.NewExactStringMatcher("HTTP://GOLANG.ORG", false),
+			},
+		},
+		{
+			desc: "prefix match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewPrefixStringMatcher("i-aint-the-one", false),
+				matcher.NewPrefixStringMatcher("192.168.1.1", false),
+				matcher.NewPrefixStringMatcher("FOO.BAR", false),
+			},
+		},
+		{
+			desc: "suffix match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewSuffixStringMatcher("i-aint-the-one", false),
+				matcher.NewSuffixStringMatcher("1::68", false),
+				matcher.NewSuffixStringMatcher(".COM", false),
+			},
+		},
+		{
+			desc: "regex match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewRegexStringMatcher(regexp.MustCompile(`.*\.examples\.com`)),
+				matcher.NewRegexStringMatcher(regexp.MustCompile(`192\.[0-9]{1,3}\.1\.1`)),
+			},
+		},
+		{
+			desc: "contains match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewContainsStringMatcher("i-aint-the-one", false),
+				matcher.NewContainsStringMatcher("2001:db8:1:1::68", false),
+				matcher.NewContainsStringMatcher("GRPC", false),
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			hi := NewHandshakeInfo(nil, nil, test.sanMatchers, "", false, false, false)
+			defer hi.Decrement()
+
+			if hi.Value().MatchingSANExists(inputCert) {
+				t.Fatalf("hi.MatchingSANExists(%+v) with SAN matchers +%v succeeded when expected to fail", inputCert, test.sanMatchers)
+			}
+		})
+	}
+}
+
+func (s) TestMatchingSANExists_Success(t *testing.T) {
+	url1, err := url.Parse("http://golang.org")
+	if err != nil {
+		t.Fatalf("url.Parse() failed: %v", err)
+	}
+	url2, err := url.Parse("https://github.com/grpc/grpc-go")
+	if err != nil {
+		t.Fatalf("url.Parse() failed: %v", err)
+	}
+	inputCert := &x509.Certificate{
+		DNSNames:       []string{"baz.test.com", "*.example.com"},
+		EmailAddresses: []string{"foobar@example.com", "barbaz@test.com"},
+		IPAddresses: []net.IP{
+			netip.MustParseAddr("192.0.0.1").AsSlice(),
+			netip.MustParseAddr("2001:db8::68").AsSlice(),
+		},
+		URIs: []*url.URL{url1, url2},
+	}
+
+	tests := []struct {
+		desc        string
+		sanMatchers []matcher.StringMatcher
+	}{
+		{
+			desc: "no san matchers",
+		},
+		{
+			desc: "exact match dns wildcard",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewPrefixStringMatcher("192.168.1.1", false),
+				matcher.NewExactStringMatcher("https://github.com/grpc/grpc-java", false),
+				matcher.NewExactStringMatcher("abc.example.com", false),
+			},
+		},
+		{
+			desc: "exact match ignore case",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("FOOBAR@EXAMPLE.COM", true),
+			},
+		},
+		{
+			desc: "prefix match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewSuffixStringMatcher(".co.in", false),
+				matcher.NewPrefixStringMatcher("192.168.1.1", false),
+				matcher.NewPrefixStringMatcher("baz.test", false),
+			},
+		},
+		{
+			desc: "prefix match ignore case",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewPrefixStringMatcher("BAZ.test", true),
+			},
+		},
+		{
+			desc: "suffix  match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewRegexStringMatcher(regexp.MustCompile(`192\.[0-9]{1,3}\.1\.1`)),
+				matcher.NewSuffixStringMatcher("192.168.1.1", false),
+				matcher.NewSuffixStringMatcher("@test.com", false),
+			},
+		},
+		{
+			desc: "suffix  match ignore case",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewSuffixStringMatcher("@test.COM", true),
+			},
+		},
+		{
+			desc: "regex match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewContainsStringMatcher("https://github.com/grpc/grpc-java", false),
+				matcher.NewRegexStringMatcher(regexp.MustCompile(`192\.[0-9]{1,3}\.1\.1`)),
+				matcher.NewRegexStringMatcher(regexp.MustCompile(`.*\.test\.com`)),
+			},
+		},
+		{
+			desc: "contains match",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("https://github.com/grpc/grpc-java", false),
+				matcher.NewContainsStringMatcher("2001:68::db8", false),
+				matcher.NewContainsStringMatcher("192.0.0", false),
+			},
+		},
+		{
+			desc: "contains match ignore case",
+			sanMatchers: []matcher.StringMatcher{
+				matcher.NewContainsStringMatcher("GRPC", true),
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			hi := NewHandshakeInfo(nil, nil, test.sanMatchers, "", false, false, false)
+			defer hi.Decrement()
+
+			if !hi.Value().MatchingSANExists(inputCert) {
+				t.Fatalf("hi.MatchingSANExists(%+v) with SAN matchers +%v failed when expected to succeed", inputCert, test.sanMatchers)
+			}
+		})
+	}
+}
+
+// Equal reports whether the handshake info structs are identical.
+func (hi *HandshakeInfo) Equal(other *HandshakeInfo) bool {
+	if hi == nil && other == nil {
+		return true
+	}
+	if hi == nil || other == nil {
+		return false
+	}
+
+	r1, i1 := hi.rootProvider, hi.identityProvider
+	r2, i2 := other.rootProvider, other.identityProvider
+
+	switch {
+	case r1 != r2 || i1 != i2:
+		return false
+	case hi.requireClientCert != other.requireClientCert:
+		return false
+	case hi.sni != other.sni:
+		return false
+	case hi.validateSANUsingSNI != other.validateSANUsingSNI:
+		return false
+	case hi.useAutoHostSNI != other.useAutoHostSNI:
+		return false
+	case len(hi.sanMatchers) != len(other.sanMatchers):
+		return false
+	}
+	for i := range hi.sanMatchers {
+		if !hi.sanMatchers[i].Equal(other.sanMatchers[i]) {
+			return false
+		}
+	}
+	return true
+}
+
+type testCertProvider struct {
+	certprovider.Provider
+}
+
+func (s) TestEqual(t *testing.T) {
+	tests := []struct {
+		desc      string
+		hi1       *grpcsync.RefCounted[*HandshakeInfo]
+		hi2       *grpcsync.RefCounted[*HandshakeInfo]
+		wantMatch bool
+	}{
+		{
+			desc: "both HandshakeInfo are nil",
+			hi1: func() *grpcsync.RefCounted[*HandshakeInfo] {
+				return grpcsync.NewRefCounted((*HandshakeInfo)(nil), func() {})
+			}(),
+			hi2: func() *grpcsync.RefCounted[*HandshakeInfo] {
+				return grpcsync.NewRefCounted((*HandshakeInfo)(nil), func() {})
+			}(),
+			wantMatch: true,
+		},
+		{
+			desc: "one HandshakeInfo is nil",
+			hi1: func() *grpcsync.RefCounted[*HandshakeInfo] {
+				return grpcsync.NewRefCounted((*HandshakeInfo)(nil), func() {})
+			}(),
+			hi2:       NewHandshakeInfo(&testCertProvider{}, nil, nil, "", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc:      "different root providers",
+			hi1:       NewHandshakeInfo(&testCertProvider{}, nil, nil, "", false, false, false),
+			hi2:       NewHandshakeInfo(&testCertProvider{}, nil, nil, "", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc: "same providers, same SAN matchers",
+			hi1: NewHandshakeInfo(testCertProvider{}, testCertProvider{}, []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("foo.com", false),
+			}, "", false, false, false),
+			hi2: NewHandshakeInfo(testCertProvider{}, testCertProvider{}, []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("foo.com", false),
+			}, "", false, false, false),
+			wantMatch: true,
+		},
+		{
+			desc: "same providers, different SAN matchers",
+			hi1: NewHandshakeInfo(testCertProvider{}, testCertProvider{}, []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("foo.com", false),
+			}, "", false, false, false),
+			hi2: NewHandshakeInfo(testCertProvider{}, testCertProvider{}, []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("bar.com", false),
+			}, "", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc: "same SAN matchers with different content",
+			hi1: NewHandshakeInfo(testCertProvider{}, testCertProvider{}, []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("foo.com", false),
+			}, "", false, false, false),
+			hi2: NewHandshakeInfo(testCertProvider{}, testCertProvider{}, []matcher.StringMatcher{
+				matcher.NewExactStringMatcher("foo.com", false),
+				matcher.NewExactStringMatcher("bar.com", false),
+			}, "", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc:      "different requireClientCert flags",
+			hi1:       NewHandshakeInfo(testCertProvider{}, testCertProvider{}, nil, "", true, false, false),
+			hi2:       NewHandshakeInfo(testCertProvider{}, testCertProvider{}, nil, "", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc:      "same identity provider, different root provider",
+			hi1:       NewHandshakeInfo(&testCertProvider{}, testCertProvider{}, nil, "", false, false, false),
+			hi2:       NewHandshakeInfo(&testCertProvider{}, testCertProvider{}, nil, "", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc:      "different identity provider, same root provider",
+			hi1:       NewHandshakeInfo(testCertProvider{}, &testCertProvider{}, nil, "", false, false, false),
+			hi2:       NewHandshakeInfo(testCertProvider{}, &testCertProvider{}, nil, "", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc:      "same sni and validateSANUsingSNI",
+			hi1:       NewHandshakeInfo(nil, nil, nil, "sni", false, true, false),
+			hi2:       NewHandshakeInfo(nil, nil, nil, "sni", false, true, false),
+			wantMatch: true,
+		},
+		{
+			desc:      "different sni",
+			hi1:       NewHandshakeInfo(nil, nil, nil, "sni1", false, false, false),
+			hi2:       NewHandshakeInfo(nil, nil, nil, "sni2", false, false, false),
+			wantMatch: false,
+		},
+		{
+			desc:      "same sni, different validateSANUsingSNI",
+			hi1:       NewHandshakeInfo(nil, nil, nil, "sni", false, false, false),
+			hi2:       NewHandshakeInfo(nil, nil, nil, "sni", false, true, false),
+			wantMatch: false,
+		},
+		{
+			desc:      "different autoHostSNI",
+			hi1:       NewHandshakeInfo(nil, nil, nil, "", false, false, true),
+			hi2:       NewHandshakeInfo(nil, nil, nil, "", false, false, false),
+			wantMatch: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			if gotMatch := test.hi1.Value().Equal(test.hi2.Value()); gotMatch != test.wantMatch {
+				t.Errorf("hi1.Equal(hi2) = %v; wantMatch %v", gotMatch, test.wantMatch)
+			}
+		})
+	}
+}
+
+type testCertProviderWithKeyMaterial struct{}
+
+func (p *testCertProviderWithKeyMaterial) KeyMaterial(_ context.Context) (*certprovider.KeyMaterial, error) {
+	km := &certprovider.KeyMaterial{}
+	spiffeBundleMapContents, err := os.ReadFile(testdata.Path("spiffe_end2end/client_spiffebundle.json"))
+	if err != nil {
+		return nil, err
+	}
+	bundleMap, err := spiffe.BundleMapFromBytes(spiffeBundleMapContents)
+	if err != nil {
+		return nil, err
+	}
+	km.SPIFFEBundleMap = bundleMap
+	rootFileContents, err := os.ReadFile(testdata.Path("spiffe_end2end/ca.pem"))
+	if err != nil {
+		return nil, err
+	}
+	trustPool := x509.NewCertPool()
+	if !trustPool.AppendCertsFromPEM(rootFileContents) {
+		return nil, fmt.Errorf("Failed to parse root certificate")
+	}
+	km.Roots = trustPool
+
+	certFileContents, err := os.ReadFile(testdata.Path("spiffe_end2end/client_spiffe.pem"))
+	if err != nil {
+		return nil, err
+	}
+	keyFileContents, err := os.ReadFile(testdata.Path("spiffe_end2end/client.key"))
+	if err != nil {
+		return nil, err
+	}
+	cert, err := tls.X509KeyPair(certFileContents, keyFileContents)
+	if err != nil {
+		return nil, err
+	}
+	km.Certs = []tls.Certificate{cert}
+	return km, nil
+}
+
+func (p *testCertProviderWithKeyMaterial) Close() {}
+
+func (s) TestBuildVerifyFuncFailures(t *testing.T) {
+	tests := []struct {
+		desc          string
+		peerCertChain [][]byte
+		wantErr       string
+	}{
+		{
+			desc:          "invalid x509",
+			peerCertChain: [][]byte{[]byte("NOT_A_CERT")},
+			wantErr:       "x509: malformed certificate",
+		},
+		{
+			desc: "invalid SPIFFE ID in peer cert",
+			// server1.pem doesn't have a valid SPIFFE ID, so attempted to get a
+			// root from the SPIFFE Bundle Map will fail
+			peerCertChain: loadCert(t, testdata.Path("server1.pem"), testdata.Path("server1.key")),
+			wantErr:       "spiffe: could not get spiffe ID from peer leaf cert but verification with spiffe trust map was configure",
+		},
+	}
+	testProvider := testCertProviderWithKeyMaterial{}
+	hi := NewHandshakeInfo(&testProvider, &testProvider, nil, "", true, false, false)
+	defer hi.Decrement()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	hiPtr := new(atomic.Pointer[grpcsync.RefCounted[*HandshakeInfo]])
+	hiPtr.Store(hi)
+	cfg, _, done, err := ClientSideTLSConfig(ctx, hiPtr, "")
+	if err != nil {
+		t.Fatalf("GetClientSideTLSConfig() failed with err %v", err)
+	}
+	defer done()
+	for _, tc := range tests {
+		t.Run(tc.desc, func(t *testing.T) {
+			err = cfg.VerifyPeerCertificate(tc.peerCertChain, nil)
+			if !strings.Contains(err.Error(), tc.wantErr) {
+				t.Errorf("VerifyPeerCertificate got err %v, want: %v", err, tc.wantErr)
+			}
+		})
+	}
+}
+
+func loadCert(t *testing.T, certPath, keyPath string) [][]byte {
+	cert, err := tls.LoadX509KeyPair(certPath, keyPath)
+	if err != nil {
+		t.Fatalf("LoadX509KeyPair(%s, %s) failed: %v", certPath, keyPath, err)
+	}
+	return cert.Certificate
+
+}
+
+type testProviderWithRoots struct {
+	roots *x509.CertPool
+}
+
+func (m *testProviderWithRoots) KeyMaterial(context.Context) (*certprovider.KeyMaterial, error) {
+	return &certprovider.KeyMaterial{
+		Roots: m.roots,
+	}, nil
+}
+
+func (m *testProviderWithRoots) Close() {}
+
+// Tests the scenario where the SNI provided by the xDS server is only present
+// in the URI SAN of the certificate presented by the server, and not in the DNS
+// SAN. It verifies that if the SNI-based SAN validation is turned on, the TLS
+// handshake fails due to SNI mismatch.
+func (s) TestAutoHostSNI_DNS_SANValidation_Failures(t *testing.T) {
+	testutils.SetEnvConfig(t, &envconfig.XDSSNIEnabled, true)
+
+	priv, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		t.Fatalf("rsa.GenerateKey() failed: %v", err)
+	}
+
+	testURL, err := url.Parse("http://wrong.sni.domain")
+	if err != nil {
+		t.Fatalf("url.Parse() failed: %v", err)
+	}
+
+	template := &x509.Certificate{
+		NotBefore: time.Now(),
+		NotAfter:  time.Now().Add(time.Hour),
+		DNSNames:  []string{"unrelated.example.com"},
+		URIs:      []*url.URL{testURL}, // URI matching SNI
+	}
+
+	derBytes, err := x509.CreateCertificate(rand.Reader, template, template, &priv.PublicKey, priv)
+	if err != nil {
+		t.Fatalf("x509.CreateCertificate() failed: %v", err)
+	}
+
+	cert, err := x509.ParseCertificate(derBytes)
+	if err != nil {
+		t.Fatalf("x509.ParseCertificate() failed: %v", err)
+	}
+
+	roots := x509.NewCertPool()
+	roots.AddCert(cert)
+
+	provider := &testProviderWithRoots{roots: roots}
+	hi := NewHandshakeInfo(provider, nil, nil, "wrong.sni.domain", true, true, false)
+	defer hi.Decrement()
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	hiPtr := new(atomic.Pointer[grpcsync.RefCounted[*HandshakeInfo]])
+	hiPtr.Store(hi)
+	cfg, _, done, err := ClientSideTLSConfig(ctx, hiPtr, "")
+	if err != nil {
+		t.Fatalf("GetClientSideTLSConfig() failed: %v", err)
+	}
+	defer done()
+
+	err = cfg.VerifyPeerCertificate([][]byte{derBytes}, nil)
+	if err == nil {
+		t.Fatalf("VerifyPeerCertificate() unexpectedly succeeded when SNI was only in URI SAN")
+	}
+	if !strings.Contains(err.Error(), "received DNS SANs: [unrelated.example.com] do not match the SNI: wrong.sni.domain") {
+		t.Fatalf("VerifyPeerCertificate() failed with %v, want error containing 'do not match the SNI'", err)
+	}
+}
+
+func (s) TestVerifyPeerCertificateZeroCerts(t *testing.T) {
+	provider := &testProviderWithRoots{roots: x509.NewCertPool()}
+	hi := NewHandshakeInfo(provider, nil, nil, "", true, false, false)
+	defer hi.Decrement()
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	hiPtr := new(atomic.Pointer[grpcsync.RefCounted[*HandshakeInfo]])
+	hiPtr.Store(hi)
+	cfg, _, done, err := ClientSideTLSConfig(ctx, hiPtr, "")
+	if err != nil {
+		t.Fatalf("GetClientSideTLSConfig() failed: %v", err)
+	}
+	defer done()
+
+	const wantErr = "no peer certificates presented"
+	if err := cfg.VerifyPeerCertificate(nil, nil); err == nil || !strings.Contains(err.Error(), wantErr) {
+		t.Errorf("VerifyPeerCertificate(nil, nil) = %v, want error containing %q", err, wantErr)
+	}
+}

internal/envconfig/envconfig.go

@@ -0,0 +1,244 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package envconfig contains grpc settings configured by environment variables.
+package envconfig
+
+import (
+	"os"
+	"strconv"
+	"strings"
+)
+
+var (
+	// EnableTXTServiceConfig is set if the DNS resolver should perform TXT
+	// lookups for service config ("GRPC_ENABLE_TXT_SERVICE_CONFIG" is not
+	// "false").
+	EnableTXTServiceConfig = boolFromEnv("GRPC_ENABLE_TXT_SERVICE_CONFIG", true)
+
+	// TXTErrIgnore is set if TXT errors should be ignored
+	// ("GRPC_GO_IGNORE_TXT_ERRORS" is not "false").
+	TXTErrIgnore = boolFromEnv("GRPC_GO_IGNORE_TXT_ERRORS", true)
+
+	// RingHashCap indicates the maximum ring size which defaults to 4096
+	// entries but may be overridden by setting the environment variable
+	// "GRPC_RING_HASH_CAP".  This does not override the default bounds
+	// checking which NACKs configs specifying ring sizes > 8*1024*1024 (~8M).
+	RingHashCap = uint64FromEnv("GRPC_RING_HASH_CAP", 4096, 1, 8*1024*1024)
+
+	// ALTSMaxConcurrentHandshakes is the maximum number of concurrent ALTS
+	// handshakes that can be performed.
+	ALTSMaxConcurrentHandshakes = uint64FromEnv("GRPC_ALTS_MAX_CONCURRENT_HANDSHAKES", 100, 1, 100)
+
+	// EnforceALPNEnabled is set if TLS connections to servers with ALPN disabled
+	// should be rejected. The HTTP/2 protocol requires ALPN to be enabled, this
+	// option is present for backward compatibility. This option may be overridden
+	// by setting the environment variable "GRPC_ENFORCE_ALPN_ENABLED" to "true"
+	// or "false".
+	EnforceALPNEnabled = boolFromEnv("GRPC_ENFORCE_ALPN_ENABLED", true)
+
+	// XDSEndpointHashKeyBackwardCompat controls the parsing of the endpoint hash
+	// key from EDS LbEndpoint metadata. Endpoint hash keys can be disabled by
+	// setting "GRPC_XDS_ENDPOINT_HASH_KEY_BACKWARD_COMPAT" to "true". A future
+	// release will remove this environment variable, enabling the new behavior
+	// unconditionally.
+	XDSEndpointHashKeyBackwardCompat = boolFromEnv("GRPC_XDS_ENDPOINT_HASH_KEY_BACKWARD_COMPAT", false)
+
+	// LabelServerGoroutines controls setting [runtime/pprof.Labels] on the
+	// goroutines spawned by [grpc.Server] type.
+	// For now, this is limited to the goroutines spawned to handle incoming
+	// requests on the server.
+	// Set "GRPC_GO_SERVER_GOROUTINE_LABELS" to "grpc.method=true" to
+	// enable this grpc.method label, or "all" to enable all valid labels.
+	// This variable is a bit-field.
+	LabelServerGoroutines = goroutineLabelsFromEnv("GRPC_GO_SERVER_GOROUTINE_LABELS", 0)
+
+	// RingHashSetRequestHashKey is set if the ring hash balancer can get the
+	// request hash header by setting the "requestHashHeader" field, according
+	// to gRFC A76. It can be disabled by setting the environment variable
+	// "GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY" to "false".
+	RingHashSetRequestHashKey = boolFromEnv("GRPC_EXPERIMENTAL_RING_HASH_SET_REQUEST_HASH_KEY", true)
+
+	// ALTSHandshakerKeepaliveParams is set if we should add the
+	// KeepaliveParams when dial the ALTS handshaker service.
+	ALTSHandshakerKeepaliveParams = boolFromEnv("GRPC_EXPERIMENTAL_ALTS_HANDSHAKER_KEEPALIVE_PARAMS", false)
+
+	// EnableDefaultPortForProxyTarget controls whether the resolver adds a default port 443
+	// to a target address that lacks one. This flag only has an effect when all of
+	// the following conditions are met:
+	//   - A connect proxy is being used.
+	//   - Target resolution is disabled.
+	//   - The DNS resolver is being used.
+	EnableDefaultPortForProxyTarget = boolFromEnv("GRPC_EXPERIMENTAL_ENABLE_DEFAULT_PORT_FOR_PROXY_TARGET", true)
+
+	// CaseSensitiveBalancerRegistries is set if the balancer registry should be
+	// case-sensitive. This is enabled by default, but can be disabled by setting
+	// the env variable "GRPC_GO_EXPERIMENTAL_CASE_SENSITIVE_BALANCER_REGISTRIES"
+	// to "false".
+	//
+	// This env varible will be removed in release v1.82.0.
+	CaseSensitiveBalancerRegistries = boolFromEnv("GRPC_GO_EXPERIMENTAL_CASE_SENSITIVE_BALANCER_REGISTRIES", true)
+
+	// XDSAuthorityRewrite indicates whether xDS authority rewriting is enabled.
+	// This feature is defined in gRFC A81 and is enabled by setting the
+	// environment variable GRPC_EXPERIMENTAL_XDS_AUTHORITY_REWRITE to "true".
+	XDSAuthorityRewrite = boolFromEnv("GRPC_EXPERIMENTAL_XDS_AUTHORITY_REWRITE", false)
+
+	// PickFirstWeightedShuffling indicates whether weighted endpoint shuffling
+	// is enabled in the pick_first LB policy, as defined in gRFC A113. This
+	// feature can be disabled by setting the environment variable
+	// GRPC_EXPERIMENTAL_PF_WEIGHTED_SHUFFLING to "false".
+	PickFirstWeightedShuffling = boolFromEnv("GRPC_EXPERIMENTAL_PF_WEIGHTED_SHUFFLING", true)
+
+	// XDSRecoverPanicInResourceParsing indicates whether the xdsclient should
+	// recover from panics while parsing xDS resources.
+	//
+	// This feature can be disabled (e.g. for fuzz testing) by setting the
+	// environment variable "GRPC_GO_EXPERIMENTAL_XDS_RESOURCE_PANIC_RECOVERY"
+	// to "false".
+	XDSRecoverPanicInResourceParsing = boolFromEnv("GRPC_GO_EXPERIMENTAL_XDS_RESOURCE_PANIC_RECOVERY", true)
+
+	// EnablePriorityLBChildPolicyCache controls whether the priority balancer
+	// should cache child balancers that are removed from the LB policy config,
+	// for a period of 15 minutes. This is disabled by default, but can be
+	// enabled by setting the env variable
+	// GRPC_EXPERIMENTAL_ENABLE_PRIORITY_LB_CHILD_POLICY_CACHE to true.
+	EnablePriorityLBChildPolicyCache = boolFromEnv("GRPC_EXPERIMENTAL_ENABLE_PRIORITY_LB_CHILD_POLICY_CACHE", false)
+
+	// Enable8KBDefaultHeaderListSize indicates that default maximum header list
+	// size is restricted to 8KB. This is disabled by default, but can be enabled
+	// by setting the environment variable
+	// "GRPC_GO_EXPERIMENTAL_ENABLE_8KB_DEFAULT_HEADER_LIST_SIZE" to "true".
+	// When disabled, the default maximum header list size of 16MB is used.
+	//
+	// When enabled, RPCs with a total size of headers exceeding 8KB will fail
+	// unless explicitly configured otherwise by the user.
+	//
+	// TODO: In release v1.82.0, env var will be enabled by default.
+	Enable8KBDefaultHeaderListSize = boolFromEnv("GRPC_GO_EXPERIMENTAL_ENABLE_8KB_DEFAULT_HEADER_LIST_SIZE", false)
+
+	// EnableHTTPFramerReadBufferPooling enables the use of the
+	// readyreader.Reader interface to perform non-memory-pinning reads,
+	// provided the underlying net.Conn supports it. This reduces memory usage
+	// when subchannels are idle.
+	//
+	// This environment variable serves as an escape hatch to disable the
+	// feature if unforeseen issues arise, and it will be removed in a future
+	// release.
+	EnableHTTPFramerReadBufferPooling = boolFromEnv("GRPC_GO_EXPERIMENTAL_HTTP_FRAMER_READ_BUFFER_POOLING", true)
+
+	// ControlBufferThrottleLimit is the maximum number of control frames that can
+	// be queued in the control buffer before throttling is applied. The value
+	// must be between 1 and 10,000, and is set to 100 by default.
+	//
+	// This environment variable serves as an escape hatch to increase the
+	// throttling limit if unforeseen issues arise, and it will be removed in a
+	// future release.
+	//
+	// TODO: Remove this env var once v1.83.0 is released.
+	ControlBufferThrottleLimit = uint64FromEnv("GRPC_GO_EXPERIMENTAL_CONTROL_BUFFER_THROTTLE_LIMIT", 100, 1, 10000)
+
+	// ALTSMaxFrameSize is the maximum frame size for ALTS in bytes.
+	// This can be overridden by setting the environment variable
+	// "GRPC_GO_EXPERIMENTAL_ALTS_MAX_FRAME_SIZE" (min: 4096, max: 512*1024, capped at
+	// 512KiB to match altsWriteBufferMaxSize in
+	// credentials/alts/internal/conn/record.go).
+	ALTSMaxFrameSize = uint64FromEnv("GRPC_GO_EXPERIMENTAL_ALTS_MAX_FRAME_SIZE", 4096, 4096, 512*1024)
+
+	// EnableReceiveBufferCompaction enables the compaction of data buffers
+	// to reduce the number of buffers in the receive buffer.
+	//
+	// This environment variable serves as an escape hatch to disable the
+	// feature if unforeseen issues arise, and it will be removed in a future
+	// release.
+	//
+	// TODO: Remove this env var once v1.85.0 is released.
+	EnableReceiveBufferCompaction = boolFromEnv("GRPC_GO_EXPERIMENTAL_ENABLE_RECEIVE_BUFFER_COMPACTION", true)
+)
+
+func boolFromEnv(envVar string, def bool) bool {
+	if def {
+		// The default is true; return true unless the variable is "false".
+		return !strings.EqualFold(os.Getenv(envVar), "false")
+	}
+	// The default is false; return false unless the variable is "true".
+	return strings.EqualFold(os.Getenv(envVar), "true")
+}
+
+func uint64FromEnv(envVar string, def, min, max uint64) uint64 {
+	v, err := strconv.ParseUint(os.Getenv(envVar), 10, 64)
+	if err != nil {
+		return def
+	}
+	if v < min {
+		return min
+	}
+	if v > max {
+		return max
+	}
+	return v
+}
+
+// GoroutineLabels is a bitfield indicating which goroutine labels are enabled.
+type GoroutineLabels uint16
+
+func goroutineLabelsFromEnv(envVar string, def GoroutineLabels) GoroutineLabels {
+	val := def
+	v := os.Getenv(envVar)
+	if strings.EqualFold(v, "all") {
+		return AllGoroutineLabels
+	} else if strings.EqualFold(v, "none") {
+		return 0
+	}
+	for s := range strings.SplitSeq(v, ",") {
+		s = strings.TrimSpace(s)
+		if len(s) == 0 {
+			continue
+		}
+		pre, post, ok := strings.Cut(s, "=")
+		if !ok {
+			// no equals sign
+			continue
+		}
+		post = strings.TrimSpace(post)
+		pre = strings.TrimSpace(pre)
+		bitDesignator := GoroutineLabels(0)
+		switch {
+		case strings.EqualFold(pre, "grpc.method"):
+			bitDesignator = GoroutineLabelServerMethod
+		default:
+			continue
+		}
+		if strings.EqualFold(post, "true") {
+			val |= bitDesignator
+		} else if strings.EqualFold(post, "false") {
+			val &^= bitDesignator
+		}
+	}
+	return val
+}
+
+const (
+	// GoroutineLabelServerMethod sets the grpc.method label on new
+	// server-side gRPC streams.
+	GoroutineLabelServerMethod GoroutineLabels = 1 << iota
+)
+
+// AllGoroutineLabels is an or'd together bitfield of all valid GoroutineLabels
+// constant values (above).
+const AllGoroutineLabels = GoroutineLabelServerMethod

internal/envconfig/envconfig_test.go

@@ -0,0 +1,218 @@
+/*
+ *
+ * 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 envconfig
+
+import (
+	"os"
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestUint64FromEnv(t *testing.T) {
+	var testCases = []struct {
+		name          string
+		val           string
+		def, min, max uint64
+		want          uint64
+	}{
+		{
+			name: "error parsing",
+			val:  "asdf", def: 5, want: 5,
+		}, {
+			name: "unset",
+			val:  "", def: 5, want: 5,
+		}, {
+			name: "too low",
+			val:  "5", min: 10, want: 10,
+		}, {
+			name: "too high",
+			val:  "5", max: 2, want: 2,
+		}, {
+			name: "in range",
+			val:  "17391", def: 13000, min: 12000, max: 18000, want: 17391,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			const testVar = "testvar"
+			if tc.val == "" {
+				os.Unsetenv(testVar)
+			} else {
+				os.Setenv(testVar, tc.val)
+			}
+			if got := uint64FromEnv(testVar, tc.def, tc.min, tc.max); got != tc.want {
+				t.Errorf("uint64FromEnv(%q(=%q), %v, %v, %v) = %v; want %v", testVar, tc.val, tc.def, tc.min, tc.max, got, tc.want)
+			}
+		})
+	}
+}
+
+func (s) TestBoolFromEnv(t *testing.T) {
+	var testCases = []struct {
+		val  string
+		def  bool
+		want bool
+	}{
+		{val: "", def: true, want: true},
+		{val: "", def: false, want: false},
+		{val: "true", def: true, want: true},
+		{val: "true", def: false, want: true},
+		{val: "false", def: true, want: false},
+		{val: "false", def: false, want: false},
+		{val: "asdf", def: true, want: true},
+		{val: "asdf", def: false, want: false},
+	}
+	for _, tc := range testCases {
+		t.Run("", func(t *testing.T) {
+			const testVar = "testvar"
+			if tc.val == "" {
+				os.Unsetenv(testVar)
+			} else {
+				os.Setenv(testVar, tc.val)
+			}
+			if got := boolFromEnv(testVar, tc.def); got != tc.want {
+				t.Errorf("boolFromEnv(%q(=%q), %v) = %v; want %v", testVar, tc.val, tc.def, got, tc.want)
+			}
+		})
+	}
+}
+
+func (s) TestGoroutineLabelsFromEnv(t *testing.T) {
+	var testCases = []struct {
+		name string
+		val  string
+		def  GoroutineLabels
+		want GoroutineLabels
+	}{
+		{
+			name: "unset_env_non-zero_default",
+			val:  "",
+			def:  GoroutineLabelServerMethod,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "unset_env_zero_default",
+			val:  "",
+			def:  0,
+			want: 0,
+		}, {
+			name: "force-enable_zero_default",
+			val:  "grpc.method=true",
+			def:  0,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "force-enable_zero_default_all_caps",
+			val:  "grpc.method=TRUE",
+			def:  0,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "force-enable_zero_default_with_whitespace",
+			val:  " grpc.method\t= true",
+			def:  0,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "force-enable_zero_default_with_other_garbage",
+			val:  "grpc.method=true,foobar",
+			def:  0,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "force-enable_mixed_case_zero_default_with_other_garbage",
+			val:  "grpc.method=tRuE,foobar",
+			def:  0,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "force-disable_zero_default",
+			val:  "grpc.method=false",
+			def:  0,
+			want: 0,
+		}, {
+			name: "force-disable_non-zero_default",
+			val:  "grpc.method=false",
+			def:  GoroutineLabelServerMethod,
+			want: 0,
+		}, {
+			name: "force-disable_non-zero_default_all_caps",
+			val:  "grpc.method=FALSE",
+			def:  GoroutineLabelServerMethod,
+			want: 0,
+		}, {
+			name: "force-disable_non-zero_default_mixed_case",
+			val:  "grpc.method=fAlSe",
+			def:  GoroutineLabelServerMethod,
+			want: 0,
+		}, {
+			name: "unknown_val_no_equal",
+			val:  "grpc.unknown.garbage",
+			def:  GoroutineLabelServerMethod,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "unknown_val",
+			val:  "grpc.unknown.garbage=fooble",
+			def:  GoroutineLabelServerMethod,
+			want: GoroutineLabelServerMethod,
+		}, {
+			name: "all_with_empty_default",
+			val:  "all",
+			def:  0,
+			want: AllGoroutineLabels,
+		}, {
+			name: "all_with_server_method_default",
+			val:  "all",
+			def:  GoroutineLabelServerMethod,
+			want: AllGoroutineLabels,
+		}, {
+			name: "none_with_empty_default",
+			val:  "none",
+			def:  0,
+			want: 0,
+		}, {
+			name: "none_with_server_method_default",
+			val:  "none",
+			def:  GoroutineLabelServerMethod,
+			want: 0,
+		}, {
+			name: "unparseable_rhs",
+			val:  "grpc.method=quux",
+			def:  GoroutineLabelServerMethod,
+			want: GoroutineLabelServerMethod,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			const testVar = "testvar"
+			if tc.val == "" {
+				os.Unsetenv(testVar)
+			} else {
+				os.Setenv(testVar, tc.val)
+			}
+			if got := goroutineLabelsFromEnv(testVar, tc.def); got != tc.want {
+				t.Errorf("goroutineLabelsFromEnv(%q(=%q), %v) = %v; want %v", testVar, tc.val, tc.def, got, tc.want)
+			}
+		})
+	}
+
+}

internal/envconfig/observability.go

@@ -0,0 +1,42 @@
+/*
+ *
+ * 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 envconfig
+
+import "os"
+
+const (
+	envObservabilityConfig     = "GRPC_GCP_OBSERVABILITY_CONFIG"
+	envObservabilityConfigFile = "GRPC_GCP_OBSERVABILITY_CONFIG_FILE"
+)
+
+var (
+	// ObservabilityConfig is the json configuration for the gcp/observability
+	// package specified directly in the envObservabilityConfig env var.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	ObservabilityConfig = os.Getenv(envObservabilityConfig)
+	// ObservabilityConfigFile is the json configuration for the
+	// gcp/observability specified in a file with the location specified in
+	// envObservabilityConfigFile env var.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	ObservabilityConfigFile = os.Getenv(envObservabilityConfigFile)
+)

internal/envconfig/xds.go

@@ -0,0 +1,106 @@
+/*
+ *
+ * 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 envconfig
+
+import (
+	"os"
+)
+
+const (
+	// XDSBootstrapFileNameEnv is the env variable to set bootstrap file name.
+	// Do not use this and read from env directly. Its value is read and kept in
+	// variable XDSBootstrapFileName.
+	//
+	// When both bootstrap FileName and FileContent are set, FileName is used.
+	XDSBootstrapFileNameEnv = "GRPC_XDS_BOOTSTRAP"
+	// XDSBootstrapFileContentEnv is the env variable to set bootstrap file
+	// content. Do not use this and read from env directly. Its value is read
+	// and kept in variable XDSBootstrapFileContent.
+	//
+	// When both bootstrap FileName and FileContent are set, FileName is used.
+	XDSBootstrapFileContentEnv = "GRPC_XDS_BOOTSTRAP_CONFIG"
+)
+
+var (
+	// XDSBootstrapFileName holds the name of the file which contains xDS
+	// bootstrap configuration. Users can specify the location of the bootstrap
+	// file by setting the environment variable "GRPC_XDS_BOOTSTRAP".
+	//
+	// When both bootstrap FileName and FileContent are set, FileName is used.
+	XDSBootstrapFileName = os.Getenv(XDSBootstrapFileNameEnv)
+	// XDSBootstrapFileContent holds the content of the xDS bootstrap
+	// configuration. Users can specify the bootstrap config by setting the
+	// environment variable "GRPC_XDS_BOOTSTRAP_CONFIG".
+	//
+	// When both bootstrap FileName and FileContent are set, FileName is used.
+	XDSBootstrapFileContent = os.Getenv(XDSBootstrapFileContentEnv)
+
+	// C2PResolverTestOnlyTrafficDirectorURI is the TD URI for testing.
+	C2PResolverTestOnlyTrafficDirectorURI = os.Getenv("GRPC_TEST_ONLY_GOOGLE_C2P_RESOLVER_TRAFFIC_DIRECTOR_URI")
+
+	// XDSDualstackEndpointsEnabled is true if gRPC should read the
+	// "additional addresses" in the xDS endpoint resource.
+	XDSDualstackEndpointsEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_DUALSTACK_ENDPOINTS", true)
+
+	// XDSSystemRootCertsEnabled is true when xDS enabled gRPC clients can use
+	// the system's default root certificates for TLS certificate validation.
+	// For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A82-xds-system-root-certs.md.
+	XDSSystemRootCertsEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_SYSTEM_ROOT_CERTS", false)
+
+	// XDSSPIFFEEnabled controls if SPIFFE Bundle Maps can be used as roots of
+	// trust.  For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A87-mtls-spiffe-support.md
+	XDSSPIFFEEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_MTLS_SPIFFE", false)
+
+	// XDSHTTPConnectEnabled controls support for dynamic HTTP CONNECT proxying
+	// configured via the xDS control plane. For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A86-xds-http-connect.md
+	XDSHTTPConnectEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_HTTP_CONNECT", false)
+
+	// XDSBootstrapCallCredsEnabled controls if call credentials can be used in
+	// xDS bootstrap configuration via the `call_creds` field. For more details,
+	// see: https://github.com/grpc/proposal/blob/master/A97-xds-jwt-call-creds.md
+	XDSBootstrapCallCredsEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_BOOTSTRAP_CALL_CREDS", false)
+
+	// XDSSNIEnabled controls if gRPC should send SNI information in xDS
+	// configured TLS handshakes. For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A101-SNI-setting-and-SNI-SAN-validation.md
+	XDSSNIEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_SNI", false)
+
+	// XDSORCAToLRSPropEnabled controls whether ORCA metrics are explicitly
+	// filtered and prefix-propagated to the LRS server. For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A85-lrs-custom-metrics-changes.md
+	XDSORCAToLRSPropEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_ORCA_LRS_PROPAGATION", true)
+
+	// XDSClientExtProcEnabled indicates whether ExtProc filter is enabled on
+	// the client side. For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A93-xds-ext-proc.md
+	XDSClientExtProcEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", false)
+
+	// GCPAuthenticationFilterEnabled enables the xDS GCP Authentication
+	// filter. For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A83-xds-gcp-authn-filter.md
+	GCPAuthenticationFilterEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_GCP_AUTHENTICATION_FILTER", false)
+
+	// XDSClientExtAuthzEnabled indicates whether the external authorization
+	// filter is enabled on the client side. For more details, see:
+	// https://github.com/grpc/proposal/blob/master/A92-xds-ext-authz.md
+	XDSClientExtAuthzEnabled = boolFromEnv("GRPC_EXPERIMENTAL_XDS_EXT_AUTHZ_ON_CLIENT", false)
+)

internal/experimental.go

@@ -0,0 +1,35 @@
+/*
+ * 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 internal
+
+var (
+	// WithBufferPool is implemented by the grpc package and returns a dial
+	// option to configure a shared buffer pool for a grpc.ClientConn.
+	WithBufferPool any // func (grpc.SharedBufferPool) grpc.DialOption
+
+	// BufferPool is implemented by the grpc package and returns a server
+	// option to configure a shared buffer pool for a grpc.Server.
+	BufferPool any // func (grpc.SharedBufferPool) grpc.ServerOption
+
+	// SetDefaultBufferPool updates the default buffer pool.
+	SetDefaultBufferPool any // func(mem.BufferPool)
+
+	// AcceptCompressors is implemented by the grpc package and returns
+	// a call option that restricts the grpc-accept-encoding header for a call.
+	AcceptCompressors any // func(...string) grpc.CallOption
+)

internal/googlecloud/googlecloud.go

@@ -0,0 +1,72 @@
+/*
+ *
+ * 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 googlecloud contains internal helpful functions for google cloud.
+package googlecloud
+
+import (
+	"runtime"
+	"strings"
+	"sync"
+
+	"google.golang.org/grpc/grpclog"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+)
+
+const logPrefix = "[googlecloud]"
+
+var (
+	vmOnGCEOnce sync.Once
+	vmOnGCE     bool
+
+	logger = internalgrpclog.NewPrefixLogger(grpclog.Component("googlecloud"), logPrefix)
+)
+
+// OnGCE returns whether the client is running on GCE.
+//
+// It provides similar functionality as metadata.OnGCE from the cloud library
+// package. We keep this to avoid depending on the cloud library module.
+func OnGCE() bool {
+	vmOnGCEOnce.Do(func() {
+		mf, err := manufacturer()
+		if err != nil {
+			logger.Infof("failed to read manufacturer, setting onGCE=false: %v")
+			return
+		}
+		vmOnGCE = isRunningOnGCE(mf, runtime.GOOS)
+	})
+	return vmOnGCE
+}
+
+// isRunningOnGCE checks whether the local system, without doing a network request, is
+// running on GCP.
+func isRunningOnGCE(manufacturer []byte, goos string) bool {
+	name := string(manufacturer)
+	switch goos {
+	case "linux":
+		name = strings.TrimSpace(name)
+		return name == "Google" || name == "Google Compute Engine"
+	case "windows":
+		name = strings.ReplaceAll(name, " ", "")
+		name = strings.ReplaceAll(name, "\n", "")
+		name = strings.ReplaceAll(name, "\r", "")
+		return name == "Google"
+	default:
+		return false
+	}
+}

internal/googlecloud/googlecloud_test.go

@@ -0,0 +1,46 @@
+/*
+ *
+ * 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 googlecloud
+
+import (
+	"testing"
+)
+
+func TestIsRunningOnGCE(t *testing.T) {
+	for _, tc := range []struct {
+		description      string
+		testOS           string
+		testManufacturer string
+		out              bool
+	}{
+		// Linux tests.
+		{"linux: not a GCP platform", "linux", "not GCP", false},
+		{"Linux: GCP platform (Google)", "linux", "Google", true},
+		{"Linux: GCP platform (Google Compute Engine)", "linux", "Google Compute Engine", true},
+		{"Linux: GCP platform (Google Compute Engine) with extra spaces", "linux", "  Google Compute Engine        ", true},
+		// Windows tests.
+		{"windows: not a GCP platform", "windows", "not GCP", false},
+		{"windows: GCP platform (Google)", "windows", "Google", true},
+		{"windows: GCP platform (Google) with extra spaces", "windows", "  Google     ", true},
+	} {
+		if got, want := isRunningOnGCE([]byte(tc.testManufacturer), tc.testOS), tc.out; got != want {
+			t.Errorf("%v: isRunningOnGCE()=%v, want %v", tc.description, got, want)
+		}
+	}
+}

internal/googlecloud/manufacturer.go

@@ -0,0 +1,26 @@
+//go:build !(linux || windows)
+// +build !linux,!windows
+
+/*
+ *
+ * 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 googlecloud
+
+func manufacturer() ([]byte, error) {
+	return nil, nil
+}

internal/googlecloud/manufacturer_linux.go

@@ -0,0 +1,27 @@
+/*
+ *
+ * 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 googlecloud
+
+import "os"
+
+const linuxProductNameFile = "/sys/class/dmi/id/product_name"
+
+func manufacturer() ([]byte, error) {
+	return os.ReadFile(linuxProductNameFile)
+}

internal/googlecloud/manufacturer_windows.go

@@ -0,0 +1,50 @@
+/*
+ *
+ * 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 googlecloud
+
+import (
+	"errors"
+	"os/exec"
+	"regexp"
+	"strings"
+)
+
+const (
+	windowsCheckCommand      = "powershell.exe"
+	windowsCheckCommandArgs  = "Get-WmiObject -Class Win32_BIOS"
+	powershellOutputFilter   = "Manufacturer"
+	windowsManufacturerRegex = ":(.*)"
+)
+
+func manufacturer() ([]byte, error) {
+	cmd := exec.Command(windowsCheckCommand, windowsCheckCommandArgs)
+	out, err := cmd.Output()
+	if err != nil {
+		return nil, err
+	}
+	for _, line := range strings.Split(strings.TrimSuffix(string(out), "\n"), "\n") {
+		if strings.HasPrefix(line, powershellOutputFilter) {
+			re := regexp.MustCompile(windowsManufacturerRegex)
+			name := re.FindString(line)
+			name = strings.TrimLeft(name, ":")
+			return []byte(name), nil
+		}
+	}
+	return nil, errors.New("cannot determine the machine's manufacturer")
+}

internal/grpclog/prefix_logger.go

@@ -0,0 +1,79 @@
+/*
+ *
+ * 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 grpclog provides logging functionality for internal gRPC packages,
+// outside of the functionality provided by the external `grpclog` package.
+package grpclog
+
+import (
+	"fmt"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+// PrefixLogger does logging with a prefix.
+//
+// Logging method on a nil logs without any prefix.
+type PrefixLogger struct {
+	logger grpclog.DepthLoggerV2
+	prefix string
+}
+
+// Infof does info logging.
+func (pl *PrefixLogger) Infof(format string, args ...any) {
+	if pl != nil {
+		// Handle nil, so the tests can pass in a nil logger.
+		format = pl.prefix + format
+		pl.logger.InfoDepth(1, fmt.Sprintf(format, args...))
+		return
+	}
+	grpclog.InfoDepth(1, fmt.Sprintf(format, args...))
+}
+
+// Warningf does warning logging.
+func (pl *PrefixLogger) Warningf(format string, args ...any) {
+	if pl != nil {
+		format = pl.prefix + format
+		pl.logger.WarningDepth(1, fmt.Sprintf(format, args...))
+		return
+	}
+	grpclog.WarningDepth(1, fmt.Sprintf(format, args...))
+}
+
+// Errorf does error logging.
+func (pl *PrefixLogger) Errorf(format string, args ...any) {
+	if pl != nil {
+		format = pl.prefix + format
+		pl.logger.ErrorDepth(1, fmt.Sprintf(format, args...))
+		return
+	}
+	grpclog.ErrorDepth(1, fmt.Sprintf(format, args...))
+}
+
+// V reports whether verbosity level l is at least the requested verbose level.
+func (pl *PrefixLogger) V(l int) bool {
+	if pl != nil {
+		return pl.logger.V(l)
+	}
+	return true
+}
+
+// NewPrefixLogger creates a prefix logger with the given prefix.
+func NewPrefixLogger(logger grpclog.DepthLoggerV2, prefix string) *PrefixLogger {
+	return &PrefixLogger{logger: logger, prefix: prefix}
+}

internal/grpcsync/callback_serializer.go

@@ -0,0 +1,124 @@
+/*
+ *
+ * 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 grpcsync
+
+import (
+	"context"
+	"errors"
+
+	"google.golang.org/grpc/internal/buffer"
+)
+
+// ErrSerializerClosed is returned by ScheduleAndWait if the CallbackSerializer
+// was closed before the callback could be scheduled.
+var ErrSerializerClosed = errors.New("callback serializer is closed")
+
+// CallbackSerializer provides a mechanism to schedule callbacks in a
+// synchronized manner. It provides a FIFO guarantee on the order of execution
+// of scheduled callbacks. New callbacks can be scheduled by invoking the
+// Schedule() method.
+//
+// This type is safe for concurrent access.
+type CallbackSerializer struct {
+	// done is closed once the serializer is shut down completely, i.e all
+	// scheduled callbacks are executed and the serializer has deallocated all
+	// its resources.
+	done chan struct{}
+
+	callbacks *buffer.Unbounded[func(context.Context)]
+}
+
+// NewCallbackSerializer returns a new CallbackSerializer instance. The provided
+// context will be passed to the scheduled callbacks. Users should cancel the
+// provided context to shutdown the CallbackSerializer. It is guaranteed that no
+// callbacks will be added once this context is canceled, and any pending un-run
+// callbacks will be executed before the serializer is shut down.
+func NewCallbackSerializer(ctx context.Context) *CallbackSerializer {
+	cs := &CallbackSerializer{
+		done:      make(chan struct{}),
+		callbacks: buffer.NewUnbounded[func(context.Context)](),
+	}
+	go cs.run(ctx)
+	return cs
+}
+
+// TrySchedule tries to schedule the provided callback function f to be
+// executed in the order it was added. This is a best-effort operation. If the
+// context passed to NewCallbackSerializer was canceled before this method is
+// called, the callback will not be scheduled.
+//
+// Callbacks are expected to honor the context when performing any blocking
+// operations, and should return early when the context is canceled.
+func (cs *CallbackSerializer) TrySchedule(f func(ctx context.Context)) {
+	cs.callbacks.Put(f)
+}
+
+// ScheduleOr schedules the provided callback function f to be executed in the
+// order it was added. If the context passed to NewCallbackSerializer has been
+// canceled before this method is called, the onFailure callback will be
+// executed inline instead.
+//
+// Callbacks are expected to honor the context when performing any blocking
+// operations, and should return early when the context is canceled.
+func (cs *CallbackSerializer) ScheduleOr(f func(ctx context.Context), onFailure func()) {
+	if cs.callbacks.Put(f) != nil {
+		onFailure()
+	}
+}
+
+// ScheduleAndWait schedules the provided callback function f to be executed in
+// the order it was added and blocks until f has run. If the context passed to
+// NewCallbackSerializer was canceled before this method is called, f is not run
+// and ScheduleAndWait returns ErrSerializerClosed.
+//
+// Callbacks are expected to honor the context when performing any blocking
+// operations, and should return early when the context is canceled.
+func (cs *CallbackSerializer) ScheduleAndWait(f func(ctx context.Context)) error {
+	done := make(chan struct{})
+	var err error
+	cs.ScheduleOr(func(ctx context.Context) {
+		f(ctx)
+		close(done)
+	}, func() {
+		err = ErrSerializerClosed
+		close(done)
+	})
+	<-done
+	return err
+}
+
+func (cs *CallbackSerializer) run(ctx context.Context) {
+	defer close(cs.done)
+
+	// Close the buffer when the context is canceled
+	// to prevent new callbacks from being added.
+	context.AfterFunc(ctx, cs.callbacks.Close)
+
+	// Run all callbacks.
+	for cb := range cs.callbacks.Get() {
+		cs.callbacks.Load()
+		cb(ctx)
+	}
+}
+
+// Done returns a channel that is closed after the context passed to
+// NewCallbackSerializer is canceled and all callbacks have been executed.
+func (cs *CallbackSerializer) Done() <-chan struct{} {
+	return cs.done
+}

internal/grpcsync/callback_serializer_test.go

@@ -0,0 +1,247 @@
+/*
+ *
+ * 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 grpcsync
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 10 * time.Millisecond // For events expected to *not* happen.
+)
+
+// TestCallbackSerializer_Schedule_FIFO verifies that callbacks are executed in
+// the same order in which they were scheduled.
+func (s) TestCallbackSerializer_Schedule_FIFO(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	cs := NewCallbackSerializer(ctx)
+	defer cancel()
+
+	// We have two channels, one to record the order of scheduling, and the
+	// other to record the order of execution. We spawn a bunch of goroutines
+	// which record the order of scheduling and call the actual Schedule()
+	// method as well.  The callbacks record the order of execution.
+	//
+	// We need to grab a lock to record order of scheduling to guarantee that
+	// the act of recording and the act of calling Schedule() happen atomically.
+	const numCallbacks = 100
+	var mu sync.Mutex
+	scheduleOrderCh := make(chan int, numCallbacks)
+	executionOrderCh := make(chan int, numCallbacks)
+	for i := 0; i < numCallbacks; i++ {
+		go func(id int) {
+			mu.Lock()
+			defer mu.Unlock()
+			scheduleOrderCh <- id
+			cs.TrySchedule(func(ctx context.Context) {
+				select {
+				case <-ctx.Done():
+					return
+				case executionOrderCh <- id:
+				}
+			})
+		}(i)
+	}
+
+	// Spawn a couple of goroutines to capture the order or scheduling and the
+	// order of execution.
+	scheduleOrder := make([]int, numCallbacks)
+	executionOrder := make([]int, numCallbacks)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numCallbacks; i++ {
+			select {
+			case <-ctx.Done():
+				return
+			case id := <-scheduleOrderCh:
+				scheduleOrder[i] = id
+			}
+		}
+	}()
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numCallbacks; i++ {
+			select {
+			case <-ctx.Done():
+				return
+			case id := <-executionOrderCh:
+				executionOrder[i] = id
+			}
+		}
+	}()
+	wg.Wait()
+
+	if diff := cmp.Diff(executionOrder, scheduleOrder); diff != "" {
+		t.Fatalf("Callbacks are not executed in scheduled order. diff(-want, +got):\n%s", diff)
+	}
+}
+
+// TestCallbackSerializer_Schedule_Concurrent verifies that all concurrently
+// scheduled callbacks get executed.
+func (s) TestCallbackSerializer_Schedule_Concurrent(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	cs := NewCallbackSerializer(ctx)
+	defer cancel()
+
+	// Schedule callbacks concurrently by calling Schedule() from goroutines.
+	// The execution of the callbacks call Done() on the waitgroup, which
+	// eventually unblocks the test and allows it to complete.
+	const numCallbacks = 100
+	var wg sync.WaitGroup
+	wg.Add(numCallbacks)
+	for i := 0; i < numCallbacks; i++ {
+		go func() {
+			cs.TrySchedule(func(context.Context) {
+				wg.Done()
+			})
+		}()
+	}
+
+	// We call Wait() on the waitgroup from a goroutine so that we can select on
+	// the Wait() being unblocked and the overall test deadline expiring.
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for all scheduled callbacks to be executed")
+	case <-done:
+	}
+}
+
+// TestCallbackSerializer_Schedule_Close verifies that callbacks in the queue
+// are not executed once Close() returns.
+func (s) TestCallbackSerializer_Schedule_Close(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	serializerCtx, serializerCancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	cs := NewCallbackSerializer(serializerCtx)
+
+	// Schedule a callback which blocks until the context passed to it is
+	// canceled. It also closes a channel to signal that it has started.
+	firstCallbackStartedCh := make(chan struct{})
+	cs.TrySchedule(func(ctx context.Context) {
+		close(firstCallbackStartedCh)
+		<-ctx.Done()
+	})
+
+	// Schedule a bunch of callbacks. These should be executed since they are
+	// scheduled before the serializer is closed.
+	const numCallbacks = 10
+	callbackCh := make(chan int, numCallbacks)
+	for i := 0; i < numCallbacks; i++ {
+		num := i
+		callback := func(context.Context) { callbackCh <- num }
+		onFailure := func() { t.Fatal("Schedule failed to accept a callback when the serializer is yet to be closed") }
+		cs.ScheduleOr(callback, onFailure)
+	}
+
+	// Ensure that none of the newer callbacks are executed at this point.
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-callbackCh:
+		t.Fatal("Newer callback executed when older one is still executing")
+	}
+
+	// Wait for the first callback to start before closing the scheduler.
+	<-firstCallbackStartedCh
+
+	// Cancel the context which will unblock the first callback. All of the
+	// other callbacks (which have not started executing at this point) should
+	// be executed after this.
+	serializerCancel()
+
+	// Ensure that the newer callbacks are executed.
+	for i := 0; i < numCallbacks; i++ {
+		select {
+		case <-ctx.Done():
+			t.Fatal("Timeout when waiting for callback scheduled before close to be executed")
+		case num := <-callbackCh:
+			if num != i {
+				t.Fatalf("Executing callback %d, want %d", num, i)
+			}
+		}
+	}
+	<-cs.Done()
+
+	// Ensure that a callback cannot be scheduled after the serializer is
+	// closed.
+	done := make(chan struct{})
+	callback := func(context.Context) { t.Fatal("Scheduled a callback after closing the serializer") }
+	onFailure := func() { close(done) }
+	cs.ScheduleOr(callback, onFailure)
+	select {
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Successfully scheduled callback after serializer is closed")
+	case <-done:
+	}
+}
+
+// TestCallbackSerializer_ScheduleAndWait verifies that ScheduleAndWait runs the
+// provided callback and blocks until it has finished executing, returning a nil
+// error.
+func (s) TestCallbackSerializer_ScheduleAndWait(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	cs := NewCallbackSerializer(ctx)
+
+	executed := false
+	if err := cs.ScheduleAndWait(func(context.Context) {
+		executed = true
+	}); err != nil {
+		t.Fatalf("ScheduleAndWait() returned error: %v, want nil", err)
+	}
+
+	// The callback must have run by the time ScheduleAndWait returns. No
+	// synchronization is needed here because ScheduleAndWait is expected to
+	// block until the callback completes.
+	if !executed {
+		t.Fatal("ScheduleAndWait() returned before the callback was executed")
+	}
+}
+
+// TestCallbackSerializer_ScheduleAndWait_Closed verifies that ScheduleAndWait
+// does not run the callback and returns ErrSerializerClosed when the serializer
+// has already been closed.
+func (s) TestCallbackSerializer_ScheduleAndWait_Closed(t *testing.T) {
+	serializerCtx, serializerCancel := context.WithCancel(context.Background())
+	cs := NewCallbackSerializer(serializerCtx)
+
+	// Close the serializer and wait for it to finish shutting down.
+	serializerCancel()
+	<-cs.Done()
+
+	if err := cs.ScheduleAndWait(func(context.Context) {
+		t.Fatal("Callback executed on a closed serializer")
+	}); err != ErrSerializerClosed {
+		t.Fatalf("ScheduleAndWait() returned error: %v, want %v", err, ErrSerializerClosed)
+	}
+}

internal/grpcsync/event.go

@@ -0,0 +1,58 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package grpcsync implements additional synchronization primitives built upon
+// the sync package.
+package grpcsync
+
+import (
+	"sync/atomic"
+)
+
+// Event represents a one-time event that may occur in the future.
+type Event struct {
+	fired atomic.Bool
+	c     chan struct{}
+}
+
+// Fire causes e to complete.  It is safe to call multiple times, and
+// concurrently.  It returns true iff this call to Fire caused the signaling
+// channel returned by Done to close. If Fire returns false, it is possible
+// the Done channel has not been closed yet.
+func (e *Event) Fire() bool {
+	if e.fired.CompareAndSwap(false, true) {
+		close(e.c)
+		return true
+	}
+	return false
+}
+
+// Done returns a channel that will be closed when Fire is called.
+func (e *Event) Done() <-chan struct{} {
+	return e.c
+}
+
+// HasFired returns true if Fire has been called.
+func (e *Event) HasFired() bool {
+	return e.fired.Load()
+}
+
+// NewEvent returns a new, ready-to-use Event.
+func NewEvent() *Event {
+	return &Event{c: make(chan struct{})}
+}

internal/grpcsync/event_test.go

@@ -0,0 +1,81 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestEventHasFired(t *testing.T) {
+	e := NewEvent()
+	if e.HasFired() {
+		t.Fatal("e.HasFired() = true; want false")
+	}
+	if !e.Fire() {
+		t.Fatal("e.Fire() = false; want true")
+	}
+	if !e.HasFired() {
+		t.Fatal("e.HasFired() = false; want true")
+	}
+}
+
+func (s) TestEventDoneChannel(t *testing.T) {
+	e := NewEvent()
+	select {
+	case <-e.Done():
+		t.Fatal("e.HasFired() = true; want false")
+	default:
+	}
+	if !e.Fire() {
+		t.Fatal("e.Fire() = false; want true")
+	}
+	select {
+	case <-e.Done():
+	default:
+		t.Fatal("e.HasFired() = false; want true")
+	}
+}
+
+func (s) TestEventMultipleFires(t *testing.T) {
+	e := NewEvent()
+	if e.HasFired() {
+		t.Fatal("e.HasFired() = true; want false")
+	}
+	if !e.Fire() {
+		t.Fatal("e.Fire() = false; want true")
+	}
+	for i := 0; i < 3; i++ {
+		if !e.HasFired() {
+			t.Fatal("e.HasFired() = false; want true")
+		}
+		if e.Fire() {
+			t.Fatal("e.Fire() = true; want false")
+		}
+	}
+}

internal/grpcsync/pubsub.go

@@ -0,0 +1,121 @@
+/*
+ *
+ * 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 grpcsync
+
+import (
+	"context"
+	"sync"
+)
+
+// Subscriber represents an entity that is subscribed to messages published on
+// a PubSub. It wraps the callback to be invoked by the PubSub when a new
+// message is published.
+type Subscriber interface {
+	// OnMessage is invoked when a new message is published. Implementations
+	// must not block in this method.
+	OnMessage(msg any)
+}
+
+// PubSub is a simple one-to-many publish-subscribe system that supports
+// messages of arbitrary type. It guarantees that messages are delivered in
+// the same order in which they were published.
+//
+// Publisher invokes the Publish() method to publish new messages, while
+// subscribers interested in receiving these messages register a callback
+// via the Subscribe() method.
+//
+// Once a PubSub is stopped, no more messages can be published, but any pending
+// published messages will be delivered to the subscribers.  Done may be used
+// to determine when all published messages have been delivered.
+type PubSub struct {
+	cs *CallbackSerializer
+
+	// Access to the below fields are guarded by this mutex.
+	mu          sync.Mutex
+	msg         any
+	subscribers map[Subscriber]bool
+}
+
+// NewPubSub returns a new PubSub instance.  Users should cancel the
+// provided context to shutdown the PubSub.
+func NewPubSub(ctx context.Context) *PubSub {
+	return &PubSub{
+		cs:          NewCallbackSerializer(ctx),
+		subscribers: map[Subscriber]bool{},
+	}
+}
+
+// Subscribe registers the provided Subscriber to the PubSub.
+//
+// If the PubSub contains a previously published message, the Subscriber's
+// OnMessage() callback will be invoked asynchronously with the existing
+// message to begin with, and subsequently for every newly published message.
+//
+// The caller is responsible for invoking the returned cancel function to
+// unsubscribe itself from the PubSub.
+func (ps *PubSub) Subscribe(sub Subscriber) (cancel func()) {
+	ps.mu.Lock()
+	defer ps.mu.Unlock()
+
+	ps.subscribers[sub] = true
+
+	if ps.msg != nil {
+		msg := ps.msg
+		ps.cs.TrySchedule(func(context.Context) {
+			ps.mu.Lock()
+			defer ps.mu.Unlock()
+			if !ps.subscribers[sub] {
+				return
+			}
+			sub.OnMessage(msg)
+		})
+	}
+
+	return func() {
+		ps.mu.Lock()
+		defer ps.mu.Unlock()
+		delete(ps.subscribers, sub)
+	}
+}
+
+// Publish publishes the provided message to the PubSub, and invokes
+// callbacks registered by subscribers asynchronously.
+func (ps *PubSub) Publish(msg any) {
+	ps.mu.Lock()
+	defer ps.mu.Unlock()
+
+	ps.msg = msg
+	for sub := range ps.subscribers {
+		s := sub
+		ps.cs.TrySchedule(func(context.Context) {
+			ps.mu.Lock()
+			defer ps.mu.Unlock()
+			if !ps.subscribers[s] {
+				return
+			}
+			s.OnMessage(msg)
+		})
+	}
+}
+
+// Done returns a channel that is closed after the context passed to NewPubSub
+// is canceled and all updates have been sent to subscribers.
+func (ps *PubSub) Done() <-chan struct{} {
+	return ps.cs.Done()
+}

internal/grpcsync/pubsub_test.go

@@ -0,0 +1,196 @@
+/*
+ *
+ * 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 grpcsync
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+type testSubscriber struct {
+	onMsgCh chan int
+}
+
+func newTestSubscriber(chSize int) *testSubscriber {
+	return &testSubscriber{onMsgCh: make(chan int, chSize)}
+}
+
+func (ts *testSubscriber) OnMessage(msg any) {
+	select {
+	case ts.onMsgCh <- msg.(int):
+	default:
+	}
+}
+
+func (s) TestPubSub_PublishNoMsg(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	pubsub := NewPubSub(ctx)
+
+	ts := newTestSubscriber(1)
+	pubsub.Subscribe(ts)
+
+	select {
+	case <-ts.onMsgCh:
+		t.Fatal("Subscriber callback invoked when no message was published")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestPubSub_PublishMsgs_RegisterSubs_And_Stop(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	pubsub := NewPubSub(ctx)
+
+	const numPublished = 10
+
+	ts1 := newTestSubscriber(numPublished)
+	pubsub.Subscribe(ts1)
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+	// Publish ten messages on the pubsub and ensure that they are received in order by the subscriber.
+	go func() {
+		for i := 0; i < numPublished; i++ {
+			pubsub.Publish(i)
+		}
+		wg.Done()
+	}()
+
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numPublished; i++ {
+			select {
+			case m := <-ts1.onMsgCh:
+				if m != i {
+					t.Errorf("Received unexpected message: %q; want: %q", m, i)
+					return
+				}
+			case <-time.After(defaultTestTimeout):
+				t.Error("Timeout when expecting the onMessage() callback to be invoked")
+				return
+			}
+		}
+	}()
+	wg.Wait()
+	if t.Failed() {
+		t.FailNow()
+	}
+
+	// Register another subscriber and ensure that it receives the last published message.
+	ts2 := newTestSubscriber(numPublished)
+	pubsub.Subscribe(ts2)
+
+	select {
+	case m := <-ts2.onMsgCh:
+		if m != numPublished-1 {
+			t.Fatalf("Received unexpected message: %q; want: %q", m, numPublished-1)
+		}
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatal("Timeout when expecting the onMessage() callback to be invoked")
+	}
+
+	wg.Add(3)
+	// Publish ten messages on the pubsub and ensure that they are received in order by the subscribers.
+	go func() {
+		for i := 0; i < numPublished; i++ {
+			pubsub.Publish(i)
+		}
+		wg.Done()
+	}()
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numPublished; i++ {
+			select {
+			case m := <-ts1.onMsgCh:
+				if m != i {
+					t.Errorf("Received unexpected message: %q; want: %q", m, i)
+					return
+				}
+			case <-time.After(defaultTestTimeout):
+				t.Error("Timeout when expecting the onMessage() callback to be invoked")
+				return
+			}
+		}
+
+	}()
+	go func() {
+		defer wg.Done()
+		for i := 0; i < numPublished; i++ {
+			select {
+			case m := <-ts2.onMsgCh:
+				if m != i {
+					t.Errorf("Received unexpected message: %q; want: %q", m, i)
+					return
+				}
+			case <-time.After(defaultTestTimeout):
+				t.Error("Timeout when expecting the onMessage() callback to be invoked")
+				return
+			}
+		}
+	}()
+	wg.Wait()
+	if t.Failed() {
+		t.FailNow()
+	}
+
+	cancel()
+	<-pubsub.Done()
+
+	go func() {
+		pubsub.Publish(99)
+	}()
+	// Ensure that the subscriber callback is not invoked as instantiated
+	// pubsub has already closed.
+	select {
+	case <-ts1.onMsgCh:
+		t.Fatal("The callback was invoked after pubsub being stopped")
+	case <-ts2.onMsgCh:
+		t.Fatal("The callback was invoked after pubsub being stopped")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestPubSub_PublishMsgs_BeforeRegisterSub(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	pubsub := NewPubSub(ctx)
+
+	const numPublished = 3
+	for i := 0; i < numPublished; i++ {
+		pubsub.Publish(i)
+	}
+
+	ts := newTestSubscriber(numPublished)
+	pubsub.Subscribe(ts)
+
+	// Ensure that the subscriber callback is invoked with a previously
+	// published message.
+	select {
+	case d := <-ts.onMsgCh:
+		if d != numPublished-1 {
+			t.Fatalf("Unexpected message received: %q; %q", d, numPublished-1)
+		}
+
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatal("Timeout when expecting the onMessage() callback to be invoked")
+	}
+}

internal/grpcsync/refcounted.go

@@ -0,0 +1,109 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"sync/atomic"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+var logger = grpclog.Component("grpcsync")
+
+// RefCounted is a reference counted wrapper of type T. It tracks the number of
+// active references and runs a cleanup when the last reference is released.
+type RefCounted[T any] struct {
+	val      T
+	refCount atomic.Int32
+	onZero   func()
+}
+
+// NewRefCounted creates a new RefCounted instance wrapping the given value with
+// initial refcount of one.
+//
+// The value should typically be a pointer, interface, or handle rather than a
+// plain value type (such as a struct or primitive value).
+//
+// The provided onZero callback must not be nil, and is executed exactly once
+// when the reference count drops to zero.  Panics if onZero is nil.
+//
+// WARNING: onZero runs synchronously inside Decrement; it must not acquire
+// locks held by Decrement callers.
+func NewRefCounted[T any](val T, onZero func()) *RefCounted[T] {
+	if onZero == nil {
+		panic("grpcsync: onZero callback cannot be nil")
+	}
+	rc := &RefCounted[T]{
+		val:    val,
+		onZero: onZero,
+	}
+	rc.refCount.Store(1)
+	return rc
+}
+
+// Value returns the encapsulated resource.
+func (rc *RefCounted[T]) Value() T {
+	return rc.val
+}
+
+// TryIncrement attempts to increment the reference count, returning true if
+// successful. It returns false if the count has already reached 0, indicating
+// the resource has been cleaned up and cannot be resurrected.
+//
+// WARNING: Avoid calling TryIncrement on hot paths when an active reference is
+// already guaranteed by the caller; use Increment instead to bypass the
+// CompareAndSwap loop overhead. TryIncrement should be reserved for speculative
+// lookups (such as fetching from a cache or map) where the resource might be
+// dead.
+func (rc *RefCounted[T]) TryIncrement() bool {
+	// Utilize a CompareAndSwap loop to prevent race conditions where a concurrent
+	// decrement could drop the count to zero between the read and the increment
+	// operation, which would otherwise inadvertently resurrect a closed resource.
+	for {
+		count := rc.refCount.Load()
+		if count <= 0 {
+			return false // Already dead or dying
+		}
+		if rc.refCount.CompareAndSwap(count, count+1) {
+			return true
+		}
+	}
+}
+
+// Increment increments the reference count.
+//
+// WARNING: Call Increment only when there is a guarantee that an active
+// reference is already present, ensuring the resource is not dead. If there is
+// a possibility that the resource is dead or its reference count might have
+// reached zero, call TryIncrement instead.
+func (rc *RefCounted[T]) Increment() {
+	if rc.refCount.Add(1) <= 1 {
+		logger.Errorf("Resource already closed or dead")
+	}
+}
+
+// Decrement decrements the reference count. If it drops to zero, the onZero
+// callback is executed synchronously before this method returns.
+func (rc *RefCounted[T]) Decrement() {
+	if v := rc.refCount.Add(-1); v < 0 {
+		logger.Errorf("Refcount cannot be negative")
+	} else if v == 0 {
+		rc.onZero()
+	}
+}

internal/grpcsync/refcounted_test.go

@@ -0,0 +1,155 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcsync
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+// Test verifies the scenario where a RefCounted instance is created and its
+// reference count is decremented. It verifies that the encapsulated value is
+// correctly retrieved and that the onZero callback is invoked when the
+// reference count drops to zero.
+func (s) TestRefCounted(t *testing.T) {
+	val := 666
+	var onZeroCalled atomic.Bool
+	rc := NewRefCounted(&val, func() { onZeroCalled.Store(true) })
+	if got := rc.Value(); *got != val {
+		t.Fatalf("Value() = %v, want %v", *got, val)
+	}
+
+	if got := onZeroCalled.Load(); got {
+		t.Fatalf("Before decrement, onZeroCalled = %v, want false", got)
+	}
+
+	rc.Decrement()
+
+	if got := onZeroCalled.Load(); !got {
+		t.Fatalf("After decrementing count to zero, onZeroCalled = %v, want true", got)
+	}
+}
+
+// Test verifies the scenario where the reference count of a resource is
+// explicitly incremented and decremented multiple times. It verifies that the
+// onZero callback is only executed when the count drops to zero, and not
+// beforehand.
+func (s) TestRefCounted_IncrementDecrement(t *testing.T) {
+	val := 666
+	var onZeroCount atomic.Int32
+	rc := NewRefCounted(&val, func() { onZeroCount.Add(1) })
+	rc.Increment()
+	rc.Decrement()
+
+	if got := onZeroCount.Load(); got != 0 {
+		t.Fatal("OnZero called unexpectedly after decrementing refcount from 2 to 1")
+	}
+
+	rc.Decrement()
+
+	if got := onZeroCount.Load(); got != 1 {
+		t.Fatalf("After decrementing count from 1 to 0, onZero called = %v times, want 1", got)
+	}
+}
+
+// Test verifies the scenario where TryIncrement is called on both an active and
+// a dead resource. It verifies that TryIncrement successfully increments active
+// resources but fails on resources whose reference count has already dropped to
+// zero.
+func (s) TestRefCounted_TryIncrement(t *testing.T) {
+	val := 666
+	var onZeroCount atomic.Int32
+	rc := NewRefCounted(&val, func() { onZeroCount.Add(1) })
+	if got := rc.TryIncrement(); !got {
+		t.Fatalf("TryIncrement() on active resource = %v, want true", got)
+	}
+
+	rc.Decrement()
+	rc.Decrement()
+
+	if got := onZeroCount.Load(); got != 1 {
+		t.Fatalf("After decrementing count to zero, onZero called = %v times, want 1", got)
+	}
+
+	if got := rc.TryIncrement(); got {
+		t.Fatalf("TryIncrement() on dead resource = %v, want false", got)
+	}
+}
+
+// Test verifies the scenario where multiple goroutines concurrently increment
+// and decrement the reference count. It verifies that the reference counting is
+// thread-safe and the onZero callback is invoked exactly once.
+func (s) TestRefCounted_Concurrent(t *testing.T) {
+	val := 666
+	var onZeroCount atomic.Int32
+	rc := NewRefCounted(&val, func() { onZeroCount.Add(1) })
+
+	var wg sync.WaitGroup
+	const numGoroutines = 100
+	for i := 0; i < numGoroutines; i++ {
+		wg.Go(func() {
+			if rc.TryIncrement() {
+				rc.Decrement()
+			}
+		})
+	}
+	rc.Decrement()
+	wg.Wait()
+	if got := onZeroCount.Load(); got != 1 {
+		t.Fatalf("After concurrent increments/decrements and final decrement, onZeroCount = %v, want 1", got)
+	}
+}
+
+// Test verifies that Decrementing a resource whose reference count has already
+// dropped to zero results in a negative refcount log error.
+func (s) TestRefCounted_DecrementNegative(_ *testing.T) {
+	val := 666
+	rc := NewRefCounted(&val, func() {})
+	rc.Decrement()
+
+	grpctest.ExpectError("Refcount cannot be negative")
+	rc.Decrement()
+}
+
+// Test verifies that Incrementing a dead resource results in a closed resource
+// log error.
+func (s) TestRefCounted_IncrementDead(_ *testing.T) {
+	val := 666
+	rc := NewRefCounted(&val, func() {})
+	rc.Decrement()
+
+	grpctest.ExpectError("Resource already closed or dead")
+	rc.Increment()
+}
+
+// Test verifies that NewRefCounted returns an error if the provided onZero
+// callback is nil.
+func (s) TestNilOnZero(t *testing.T) {
+	const wantErr = "grpcsync: onZero callback cannot be nil"
+	defer func() {
+		if r := recover(); r == nil || r != wantErr {
+			t.Fatalf("NewRefCounted(_, nil) panic = %v, want: %q", r, wantErr)
+		}
+	}()
+	val := 666
+	NewRefCounted(&val, nil)
+}

internal/grpctest/example_test.go

@@ -0,0 +1,61 @@
+/*
+ *
+ * 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 grpctest_test
+
+import (
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	i int
+}
+
+func (s *s) Setup(t *testing.T) {
+	t.Log("Per-test setup code")
+	s.i = 5
+}
+
+func (s *s) TestSomething(t *testing.T) {
+	t.Log("TestSomething")
+	if s.i != 5 {
+		t.Errorf("s.i = %v; want 5", s.i)
+	}
+	s.i = 3
+}
+
+func (s *s) TestSomethingElse(t *testing.T) {
+	t.Log("TestSomethingElse")
+	if got, want := s.i%4, 1; got != want {
+		t.Errorf("s.i %% 4 = %v; want %v", got, want)
+	}
+	s.i = 3
+}
+
+func (s *s) Teardown(t *testing.T) {
+	t.Log("Per-test teardown code")
+	if s.i != 3 {
+		t.Fatalf("s.i = %v; want 3", s.i)
+	}
+}
+
+func TestExample(t *testing.T) {
+	grpctest.RunSubTests(t, &s{})
+}

internal/grpctest/grpctest.go

@@ -0,0 +1,135 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package grpctest implements testing helpers.
+package grpctest
+
+import (
+	"context"
+	"reflect"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/leakcheck"
+)
+
+var lcFailed uint32
+
+type logger struct {
+	t *testing.T
+}
+
+func (e logger) Logf(format string, args ...any) {
+	e.t.Logf(format, args...)
+}
+
+func (e logger) Errorf(format string, args ...any) {
+	atomic.StoreUint32(&lcFailed, 1)
+	e.t.Errorf(format, args...)
+}
+
+// Tester is an implementation of the x interface parameter to
+// grpctest.RunSubTests with default Setup and Teardown behavior. Setup updates
+// the tlogger and Teardown performs a leak check. Embed in a struct with tests
+// defined to use.
+type Tester struct{}
+
+// Setup updates the tlogger.
+func (Tester) Setup(t *testing.T) {
+	tLogr.update(t)
+	// TODO: There is one final leak around closing connections without completely
+	//  draining the recvBuffer that has yet to be resolved. All other leaks have been
+	//  completely addressed, and this can be turned back on as soon as this issue is
+	//  fixed.
+	leakcheck.SetTrackingBufferPool(logger{t: t})
+	leakcheck.TrackTimers()
+	leakcheck.TrackAsyncReporters()
+}
+
+// Teardown performs a leak check.
+func (Tester) Teardown(t *testing.T) {
+	leakcheck.CheckTrackingBufferPool()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	leakcheck.CheckTimers(ctx, logger{t: t})
+	if atomic.LoadUint32(&lcFailed) == 1 {
+		return
+	}
+	leakcheck.CheckGoroutines(ctx, logger{t: t})
+	if atomic.LoadUint32(&lcFailed) == 1 {
+		t.Log("Goroutine leak check disabled for future tests")
+	}
+	leakcheck.CheckAsyncReporters(logger{t: t})
+	if atomic.LoadUint32(&lcFailed) == 1 {
+		return
+	}
+	tLogr.endTest(t)
+}
+
+// Interface defines Tester's methods for use in this package.
+type Interface interface {
+	Setup(*testing.T)
+	Teardown(*testing.T)
+}
+
+func getTestFunc(t *testing.T, xv reflect.Value, name string) func(*testing.T) {
+	if m := xv.MethodByName(name); m.IsValid() {
+		if f, ok := m.Interface().(func(*testing.T)); ok {
+			return f
+		}
+		// Method exists but has the wrong type signature.
+		t.Fatalf("grpctest: function %v has unexpected signature (%T)", name, m.Interface())
+	}
+	return func(*testing.T) {}
+}
+
+// RunSubTests runs all "Test___" functions that are methods of x as subtests
+// of the current test.  Setup is run before the test function and Teardown is
+// run after.
+//
+// For example usage, see example_test.go.  Run it using:
+//
+//	$ go test -v -run TestExample .
+//
+// To run a specific test/subtest:
+//
+//	$ go test -v -run 'TestExample/^Something$' .
+func RunSubTests(t *testing.T, x Interface) {
+	xt := reflect.TypeOf(x)
+	xv := reflect.ValueOf(x)
+
+	for i := 0; i < xt.NumMethod(); i++ {
+		methodName := xt.Method(i).Name
+		if !strings.HasPrefix(methodName, "Test") {
+			continue
+		}
+		tfunc := getTestFunc(t, xv, methodName)
+		t.Run(strings.TrimPrefix(methodName, "Test"), func(t *testing.T) {
+			// Run leakcheck in t.Cleanup() to guarantee it is run even if tfunc
+			// or setup uses t.Fatal().
+			//
+			// Note that a defer would run before t.Cleanup, so if a goroutine
+			// is closed by a test's t.Cleanup, a deferred leakcheck would fail.
+			t.Cleanup(func() { x.Teardown(t) })
+			x.Setup(t)
+			tfunc(t)
+		})
+	}
+}

internal/grpctest/grpctest_test.go

@@ -0,0 +1,46 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpctest
+
+import (
+	"reflect"
+	"testing"
+)
+
+type tRunST struct {
+	setup, test, teardown bool
+}
+
+func (t *tRunST) Setup(*testing.T) {
+	t.setup = true
+}
+func (t *tRunST) TestSubTest(*testing.T) {
+	t.test = true
+}
+func (t *tRunST) Teardown(*testing.T) {
+	t.teardown = true
+}
+
+func TestRunSubTests(t *testing.T) {
+	x := &tRunST{}
+	RunSubTests(t, x)
+	if want := (&tRunST{setup: true, test: true, teardown: true}); !reflect.DeepEqual(x, want) {
+		t.Fatalf("x = %v; want all fields true", x)
+	}
+}

internal/grpctest/tlogger.go

@@ -0,0 +1,300 @@
+/*
+ *
+ * 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 grpctest
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path"
+	"regexp"
+	"runtime"
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+// tLogr serves as the grpclog logger and is the interface through which
+// expected errors are declared in tests.
+var tLogr *tLogger
+
+const callingFrame = 4
+
+type logType int
+
+func (l logType) String() string {
+	switch l {
+	case infoLog:
+		return "INFO"
+	case warningLog:
+		return "WARNING"
+	case errorLog:
+		return "ERROR"
+	case fatalLog:
+		return "FATAL"
+	}
+	return "UNKNOWN"
+}
+
+const (
+	infoLog logType = iota
+	warningLog
+	errorLog
+	fatalLog
+)
+
+type tLogger struct {
+	v           int
+	initialized bool
+
+	mu    sync.Mutex
+	t     *testing.T
+	start time.Time
+	logs  map[logType]map[*regexp.Regexp]int
+}
+
+func init() {
+	vLevel := 0 // Default verbosity level
+
+	if vLevelEnv, found := os.LookupEnv("GRPC_GO_LOG_VERBOSITY_LEVEL"); found {
+		// If found, attempt to convert. If conversion is successful, update vLevel.
+		// If conversion fails, log a warning, but vLevel remains its default of 0.
+		if val, err := strconv.Atoi(vLevelEnv); err == nil {
+			vLevel = val
+		} else {
+			// Log the error if the environment variable is not a valid integer.
+			fmt.Printf("Warning: GRPC_GO_LOG_VERBOSITY_LEVEL environment variable '%s' is not a valid integer. "+
+				"Using default verbosity level 0. Error: %v\n", vLevelEnv, err)
+		}
+	}
+	// Initialize tLogr with the determined verbosity level.
+	logsMap := map[logType]map[*regexp.Regexp]int{
+		errorLog:   {},
+		warningLog: {},
+	}
+	tLogr = &tLogger{logs: logsMap, v: vLevel}
+}
+
+// getCallingPrefix returns the <file:line> at the given depth from the stack.
+func getCallingPrefix(depth int) (string, error) {
+	_, file, line, ok := runtime.Caller(depth)
+	if !ok {
+		return "", errors.New("frame request out-of-bounds")
+	}
+	return fmt.Sprintf("%s:%d", path.Base(file), line), nil
+}
+
+// log logs the message with the specified parameters to the tLogger.
+func (tl *tLogger) log(ltype logType, depth int, format string, args ...any) {
+	tl.mu.Lock()
+	defer tl.mu.Unlock()
+	prefix, err := getCallingPrefix(callingFrame + depth)
+	if err != nil {
+		tl.t.Error(err)
+		return
+	}
+	args = append([]any{ltype.String() + " " + prefix}, args...)
+	args = append(args, fmt.Sprintf(" (t=+%s)", time.Since(tl.start)))
+
+	if format == "" {
+		switch ltype {
+		case errorLog:
+			// fmt.Sprintln is used rather than fmt.Sprint because tl.Log uses fmt.Sprintln behavior.
+			if tl.expected(fmt.Sprintln(args...), errorLog) {
+				tl.t.Log(args...)
+			} else {
+				tl.t.Error(args...)
+			}
+		case warningLog:
+			tl.expected(fmt.Sprintln(args...), warningLog)
+			tl.t.Log(args...)
+		case fatalLog:
+			panic(fmt.Sprint(args...))
+		default:
+			tl.t.Log(args...)
+		}
+	} else {
+		// Add formatting directives for the callingPrefix and timeSuffix.
+		format = "%v " + format + "%s"
+		switch ltype {
+		case errorLog:
+			if tl.expected(fmt.Sprintf(format, args...), errorLog) {
+				tl.t.Logf(format, args...)
+			} else {
+				tl.t.Errorf(format, args...)
+			}
+		case warningLog:
+			tl.expected(fmt.Sprintln(args...), warningLog)
+			tl.t.Log(args...)
+		case fatalLog:
+			panic(fmt.Sprintf(format, args...))
+		default:
+			tl.t.Logf(format, args...)
+		}
+	}
+}
+
+// update updates the testing.T that the testing logger logs to. Should be done
+// before every test. It also initializes the tLogger if it has not already.
+func (tl *tLogger) update(t *testing.T) {
+	tl.mu.Lock()
+	defer tl.mu.Unlock()
+	if !tl.initialized {
+		grpclog.SetLoggerV2(tl)
+		tl.initialized = true
+	}
+	tl.t = t
+	tl.start = time.Now()
+	tl.logs[errorLog] = map[*regexp.Regexp]int{}
+	tl.logs[warningLog] = map[*regexp.Regexp]int{}
+}
+
+// ExpectError declares an error to be expected. For the next test, the first
+// error log matching the expression (using FindString) will not cause the test
+// to fail. "For the next test" includes all the time until the next call to
+// Update(). Note that if an expected error is not encountered, this will cause
+// the test to fail.
+func ExpectError(expr string) {
+	expectLogsN(expr, 1, errorLog)
+}
+
+// ExpectErrorN declares an error to be expected n times.
+func ExpectErrorN(expr string, n int) {
+	expectLogsN(expr, n, errorLog)
+}
+
+// ExpectWarning declares a warning to be expected.
+func ExpectWarning(expr string) {
+	expectLogsN(expr, 1, warningLog)
+}
+
+func expectLogsN(expr string, n int, logType logType) {
+	tLogr.mu.Lock()
+	defer tLogr.mu.Unlock()
+	re, err := regexp.Compile(expr)
+	if err != nil {
+		tLogr.t.Error(err)
+		return
+	}
+	tLogr.logs[logType][re] += n
+}
+
+// endTest checks if expected errors were not encountered.
+func (tl *tLogger) endTest(t *testing.T) {
+	tl.mu.Lock()
+	defer tl.mu.Unlock()
+	for re, count := range tl.logs[errorLog] {
+		if count > 0 {
+			t.Errorf("Expected error '%v' not encountered", re.String())
+		}
+	}
+	for re, count := range tl.logs[warningLog] {
+		if count > 0 {
+			t.Errorf("Expected warning '%v' not encountered", re.String())
+		}
+	}
+	tl.logs[errorLog] = map[*regexp.Regexp]int{}
+	tl.logs[warningLog] = map[*regexp.Regexp]int{}
+}
+
+// expected determines if the log string of the particular type is protected or
+// not.
+func (tl *tLogger) expected(s string, logType logType) bool {
+	for re, count := range tl.logs[logType] {
+		if re.FindStringIndex(s) != nil {
+			tl.logs[logType][re]--
+			if count <= 1 {
+				delete(tl.logs[logType], re)
+			}
+			return true
+		}
+	}
+	return false
+}
+
+func (tl *tLogger) Info(args ...any) {
+	tl.log(infoLog, 0, "", args...)
+}
+
+func (tl *tLogger) Infoln(args ...any) {
+	tl.log(infoLog, 0, "", args...)
+}
+
+func (tl *tLogger) Infof(format string, args ...any) {
+	tl.log(infoLog, 0, format, args...)
+}
+
+func (tl *tLogger) InfoDepth(depth int, args ...any) {
+	tl.log(infoLog, depth, "", args...)
+}
+
+func (tl *tLogger) Warning(args ...any) {
+	tl.log(warningLog, 0, "", args...)
+}
+
+func (tl *tLogger) Warningln(args ...any) {
+	tl.log(warningLog, 0, "", args...)
+}
+
+func (tl *tLogger) Warningf(format string, args ...any) {
+	tl.log(warningLog, 0, format, args...)
+}
+
+func (tl *tLogger) WarningDepth(depth int, args ...any) {
+	tl.log(warningLog, depth, "", args...)
+}
+
+func (tl *tLogger) Error(args ...any) {
+	tl.log(errorLog, 0, "", args...)
+}
+
+func (tl *tLogger) Errorln(args ...any) {
+	tl.log(errorLog, 0, "", args...)
+}
+
+func (tl *tLogger) Errorf(format string, args ...any) {
+	tl.log(errorLog, 0, format, args...)
+}
+
+func (tl *tLogger) ErrorDepth(depth int, args ...any) {
+	tl.log(errorLog, depth, "", args...)
+}
+
+func (tl *tLogger) Fatal(args ...any) {
+	tl.log(fatalLog, 0, "", args...)
+}
+
+func (tl *tLogger) Fatalln(args ...any) {
+	tl.log(fatalLog, 0, "", args...)
+}
+
+func (tl *tLogger) Fatalf(format string, args ...any) {
+	tl.log(fatalLog, 0, format, args...)
+}
+
+func (tl *tLogger) FatalDepth(depth int, args ...any) {
+	tl.log(fatalLog, depth, "", args...)
+}
+
+func (tl *tLogger) V(l int) bool {
+	return l <= tl.v
+}

internal/grpctest/tlogger_test.go

@@ -0,0 +1,79 @@
+/*
+ *
+ * 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 grpctest
+
+import (
+	"testing"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+type s struct {
+	Tester
+}
+
+func Test(t *testing.T) {
+	RunSubTests(t, s{})
+}
+
+func (s) TestInfo(*testing.T) {
+	grpclog.Info("Info", "message.")
+}
+
+func (s) TestInfoln(*testing.T) {
+	grpclog.Infoln("Info", "message.")
+}
+
+func (s) TestInfof(*testing.T) {
+	grpclog.Infof("%v %v.", "Info", "message")
+}
+
+func (s) TestInfoDepth(*testing.T) {
+	grpclog.InfoDepth(0, "Info", "depth", "message.")
+}
+
+func (s) TestWarning(*testing.T) {
+	grpclog.Warning("Warning", "message.")
+}
+
+func (s) TestWarningln(*testing.T) {
+	grpclog.Warningln("Warning", "message.")
+}
+
+func (s) TestWarningf(*testing.T) {
+	grpclog.Warningf("%v %v.", "Warning", "message")
+}
+
+func (s) TestWarningDepth(*testing.T) {
+	grpclog.WarningDepth(0, "Warning", "depth", "message.")
+}
+
+func (s) TestError(*testing.T) {
+	const numErrors = 10
+	ExpectError("Expected error")
+	ExpectError("Expected ln error")
+	ExpectError("Expected formatted error")
+	ExpectErrorN("Expected repeated error", numErrors)
+	grpclog.Error("Expected", "error")
+	grpclog.Errorln("Expected", "ln", "error")
+	grpclog.Errorf("%v %v %v", "Expected", "formatted", "error")
+	for i := 0; i < numErrors; i++ {
+		grpclog.Error("Expected repeated error")
+	}
+}

internal/grpcutil/compressor.go

@@ -0,0 +1,42 @@
+/*
+ *
+ * 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 grpcutil
+
+import (
+	"strings"
+)
+
+// RegisteredCompressorNames holds names of the registered compressors.
+var RegisteredCompressorNames []string
+
+// IsCompressorNameRegistered returns true when name is available in registry.
+func IsCompressorNameRegistered(name string) bool {
+	for _, compressor := range RegisteredCompressorNames {
+		if compressor == name {
+			return true
+		}
+	}
+	return false
+}
+
+// RegisteredCompressors returns a string of registered compressor names
+// separated by comma.
+func RegisteredCompressors() string {
+	return strings.Join(RegisteredCompressorNames, ",")
+}

internal/grpcutil/compressor_test.go

@@ -0,0 +1,32 @@
+/*
+ *
+ * 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 grpcutil
+
+import (
+	"testing"
+)
+
+func TestRegisteredCompressors(t *testing.T) {
+	defer func(c []string) { RegisteredCompressorNames = c }(RegisteredCompressorNames)
+	RegisteredCompressorNames = []string{"gzip", "snappy"}
+	if got, want := RegisteredCompressors(), "gzip,snappy"; got != want {
+		t.Fatalf("Unexpected compressors got:%s, want:%s", got, want)
+
+	}
+}

internal/grpcutil/encode_duration.go

@@ -0,0 +1,62 @@
+/*
+ *
+ * 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 grpcutil
+
+import (
+	"strconv"
+	"time"
+)
+
+const maxTimeoutValue int64 = 100000000 - 1
+
+// div does integer division and round-up the result. Note that this is
+// equivalent to (d+r-1)/r but has less chance to overflow.
+func div(d, r time.Duration) int64 {
+	if d%r > 0 {
+		return int64(d/r + 1)
+	}
+	return int64(d / r)
+}
+
+// EncodeDuration encodes the duration to the format grpc-timeout header
+// accepts.
+//
+// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
+func EncodeDuration(t time.Duration) string {
+	if t <= 0 {
+		return "0n"
+	}
+	if d := div(t, time.Nanosecond); d <= maxTimeoutValue {
+		return strconv.FormatInt(d, 10) + "n"
+	}
+	if d := div(t, time.Microsecond); d <= maxTimeoutValue {
+		return strconv.FormatInt(d, 10) + "u"
+	}
+	if d := div(t, time.Millisecond); d <= maxTimeoutValue {
+		return strconv.FormatInt(d, 10) + "m"
+	}
+	if d := div(t, time.Second); d <= maxTimeoutValue {
+		return strconv.FormatInt(d, 10) + "S"
+	}
+	if d := div(t, time.Minute); d <= maxTimeoutValue {
+		return strconv.FormatInt(d, 10) + "M"
+	}
+	// Note that maxTimeoutValue * time.Hour > MaxInt64.
+	return strconv.FormatInt(div(t, time.Hour), 10) + "H"
+}

internal/grpcutil/encode_duration_test.go

@@ -0,0 +1,51 @@
+/*
+ *
+ * Copyright 2020 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcutil
+
+import (
+	"testing"
+	"time"
+)
+
+func TestEncodeDuration(t *testing.T) {
+	for _, test := range []struct {
+		in  string
+		out string
+	}{
+		{"12345678ns", "12345678n"},
+		{"123456789ns", "123457u"},
+		{"12345678us", "12345678u"},
+		{"123456789us", "123457m"},
+		{"12345678ms", "12345678m"},
+		{"123456789ms", "123457S"},
+		{"12345678s", "12345678S"},
+		{"123456789s", "2057614M"},
+		{"12345678m", "12345678M"},
+		{"123456789m", "2057614H"},
+	} {
+		d, err := time.ParseDuration(test.in)
+		if err != nil {
+			t.Fatalf("failed to parse duration string %s: %v", test.in, err)
+		}
+		out := EncodeDuration(d)
+		if out != test.out {
+			t.Fatalf("timeoutEncode(%s) = %s, want %s", test.in, out, test.out)
+		}
+	}
+}

internal/grpcutil/grpcutil.go

@@ -0,0 +1,20 @@
+/*
+ *
+ * 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 grpcutil provides utility functions used across the gRPC codebase.
+package grpcutil

internal/grpcutil/metadata.go

@@ -0,0 +1,40 @@
+/*
+ *
+ * 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 grpcutil
+
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+)
+
+type mdExtraKey struct{}
+
+// WithExtraMetadata creates a new context with incoming md attached.
+func WithExtraMetadata(ctx context.Context, md metadata.MD) context.Context {
+	return context.WithValue(ctx, mdExtraKey{}, md)
+}
+
+// ExtraMetadata returns the incoming metadata in ctx if it exists.  The
+// returned MD should not be modified. Writing to it may cause races.
+// Modification should be made to copies of the returned MD.
+func ExtraMetadata(ctx context.Context) (md metadata.MD, ok bool) {
+	md, ok = ctx.Value(mdExtraKey{}).(metadata.MD)
+	return
+}

internal/grpcutil/method.go

@@ -0,0 +1,88 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcutil
+
+import (
+	"errors"
+	"strings"
+)
+
+// ParseMethod splits service and method from the input. It expects format
+// "/service/method".
+func ParseMethod(methodName string) (service, method string, _ error) {
+	if !strings.HasPrefix(methodName, "/") {
+		return "", "", errors.New("invalid method name: should start with /")
+	}
+	methodName = methodName[1:]
+
+	pos := strings.LastIndex(methodName, "/")
+	if pos < 0 {
+		return "", "", errors.New("invalid method name: suffix /method is missing")
+	}
+	return methodName[:pos], methodName[pos+1:], nil
+}
+
+// baseContentType is the base content-type for gRPC.  This is a valid
+// content-type on its own, but can also include a content-subtype such as
+// "proto" as a suffix after "+" or ";".  See
+// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
+// for more details.
+const baseContentType = "application/grpc"
+
+// ContentSubtype returns the content-subtype for the given content-type.  The
+// given content-type must be a valid content-type that starts with
+// "application/grpc". A content-subtype will follow "application/grpc" after a
+// "+" or ";". See
+// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests for
+// more details.
+//
+// If contentType is not a valid content-type for gRPC, the boolean
+// will be false, otherwise true. If content-type == "application/grpc",
+// "application/grpc+", or "application/grpc;", the boolean will be true,
+// but no content-subtype will be returned.
+//
+// contentType is assumed to be lowercase already.
+func ContentSubtype(contentType string) (string, bool) {
+	if contentType == baseContentType {
+		return "", true
+	}
+	if !strings.HasPrefix(contentType, baseContentType) {
+		return "", false
+	}
+	// guaranteed since != baseContentType and has baseContentType prefix
+	switch contentType[len(baseContentType)] {
+	case '+', ';':
+		// this will return true for "application/grpc+" or "application/grpc;"
+		// which the previous validContentType function tested to be valid, so we
+		// just say that no content-subtype is specified in this case
+		return contentType[len(baseContentType)+1:], true
+	default:
+		return "", false
+	}
+}
+
+// ContentType builds full content type with the given sub-type.
+//
+// contentSubtype is assumed to be lowercase
+func ContentType(contentSubtype string) string {
+	if contentSubtype == "" {
+		return baseContentType
+	}
+	return baseContentType + "+" + contentSubtype
+}

internal/grpcutil/method_test.go

@@ -0,0 +1,69 @@
+/*
+ *
+ * Copyright 2018 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package grpcutil
+
+import (
+	"testing"
+)
+
+func TestParseMethod(t *testing.T) {
+	testCases := []struct {
+		methodName  string
+		wantService string
+		wantMethod  string
+		wantError   bool
+	}{
+		{methodName: "/s/m", wantService: "s", wantMethod: "m", wantError: false},
+		{methodName: "/p.s/m", wantService: "p.s", wantMethod: "m", wantError: false},
+		{methodName: "/p/s/m", wantService: "p/s", wantMethod: "m", wantError: false},
+		{methodName: "/", wantError: true},
+		{methodName: "/sm", wantError: true},
+		{methodName: "", wantError: true},
+		{methodName: "sm", wantError: true},
+	}
+	for _, tc := range testCases {
+		s, m, err := ParseMethod(tc.methodName)
+		if (err != nil) != tc.wantError || s != tc.wantService || m != tc.wantMethod {
+			t.Errorf("ParseMethod(%s) = (%s, %s, %v), want (%s, %s, %v)", tc.methodName, s, m, err, tc.wantService, tc.wantMethod, tc.wantError)
+		}
+	}
+}
+
+func TestContentSubtype(t *testing.T) {
+	tests := []struct {
+		contentType string
+		want        string
+		wantValid   bool
+	}{
+		{"application/grpc", "", true},
+		{"application/grpc+", "", true},
+		{"application/grpc+blah", "blah", true},
+		{"application/grpc;", "", true},
+		{"application/grpc;blah", "blah", true},
+		{"application/grpcd", "", false},
+		{"application/grpd", "", false},
+		{"application/grp", "", false},
+	}
+	for _, tt := range tests {
+		got, gotValid := ContentSubtype(tt.contentType)
+		if got != tt.want || gotValid != tt.wantValid {
+			t.Errorf("contentSubtype(%q) = (%v, %v); want (%v, %v)", tt.contentType, got, gotValid, tt.want, tt.wantValid)
+		}
+	}
+}

internal/hierarchy/hierarchy.go

@@ -0,0 +1,108 @@
+/*
+ *
+ * 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 hierarchy contains functions to set and get hierarchy string from
+// addresses.
+//
+// This package is experimental.
+package hierarchy
+
+import (
+	"google.golang.org/grpc/resolver"
+)
+
+type pathKeyType string
+
+const pathKey = pathKeyType("grpc.internal.address.hierarchical_path")
+
+type pathValue []string
+
+func (p pathValue) Equal(o any) bool {
+	op, ok := o.(pathValue)
+	if !ok {
+		return false
+	}
+	if len(op) != len(p) {
+		return false
+	}
+	for i, v := range p {
+		if v != op[i] {
+			return false
+		}
+	}
+	return true
+}
+
+// FromEndpoint returns the hierarchical path of endpoint.
+func FromEndpoint(endpoint resolver.Endpoint) []string {
+	path, _ := endpoint.Attributes.Value(pathKey).(pathValue)
+	return path
+}
+
+// SetInEndpoint overrides the hierarchical path in endpoint with path.
+func SetInEndpoint(endpoint resolver.Endpoint, path []string) resolver.Endpoint {
+	endpoint.Attributes = endpoint.Attributes.WithValue(pathKey, pathValue(path))
+	return endpoint
+}
+
+// Group splits a slice of endpoints into groups based on
+// the first hierarchy path. The first hierarchy path will be removed from the
+// result.
+//
+// Input:
+// [
+//
+//	{endpoint0, path: [p0, wt0]}
+//	{endpoint1, path: [p0, wt1]}
+//	{endpoint2, path: [p1, wt2]}
+//	{endpoint3, path: [p1, wt3]}
+//
+// ]
+//
+// Endpoints will be split into p0/p1, and the p0/p1 will be removed from the
+// path.
+//
+// Output:
+//
+//	{
+//	  p0: [
+//	    {endpoint0, path: [wt0]},
+//	    {endpoint1, path: [wt1]},
+//	  ],
+//	  p1: [
+//	    {endpoint2, path: [wt2]},
+//	    {endpoint3, path: [wt3]},
+//	  ],
+//	}
+//
+// If hierarchical path is not set, or has no path in it, the endpoint is
+// dropped.
+func Group(endpoints []resolver.Endpoint) map[string][]resolver.Endpoint {
+	ret := make(map[string][]resolver.Endpoint)
+	for _, endpoint := range endpoints {
+		oldPath := FromEndpoint(endpoint)
+		if len(oldPath) == 0 {
+			continue
+		}
+		curPath := oldPath[0]
+		newPath := oldPath[1:]
+		newEndpoint := SetInEndpoint(endpoint, newPath)
+		ret[curPath] = append(ret[curPath], newEndpoint)
+	}
+	return ret
+}

internal/hierarchy/hierarchy_ext_test.go

@@ -0,0 +1,174 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package hierarchy_test
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/internal/hierarchy"
+	"google.golang.org/grpc/resolver"
+)
+
+func TestFromEndpoint(t *testing.T) {
+	tests := []struct {
+		name string
+		ep   resolver.Endpoint
+		want []string
+	}{
+		{
+			name: "not set",
+			ep:   resolver.Endpoint{},
+			want: nil,
+		},
+		{
+			name: "set",
+			ep:   hierarchy.SetInEndpoint(resolver.Endpoint{}, []string{"a", "b"}),
+			want: []string{"a", "b"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := hierarchy.FromEndpoint(tt.ep); !cmp.Equal(got, tt.want) {
+				t.Errorf("FromEndpoint() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestSetInEndpoint(t *testing.T) {
+	tests := []struct {
+		name string
+		ep   resolver.Endpoint
+		path []string
+	}{
+		{
+			name: "before is not set",
+			ep:   resolver.Endpoint{},
+			path: []string{"a", "b"},
+		},
+		{
+			name: "before is set",
+			ep:   hierarchy.SetInEndpoint(resolver.Endpoint{}, []string{"before", "a", "b"}),
+			path: []string{"a", "b"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			newEP := hierarchy.SetInEndpoint(tt.ep, tt.path)
+			newPath := hierarchy.FromEndpoint(newEP)
+			if !cmp.Equal(newPath, tt.path) {
+				t.Errorf("path after SetInEndpoint() = %v, want %v", newPath, tt.path)
+			}
+		})
+	}
+}
+
+func TestGroup(t *testing.T) {
+	tests := []struct {
+		name string
+		eps  []resolver.Endpoint
+		want map[string][]resolver.Endpoint
+	}{
+		{
+			name: "all with hierarchy",
+			eps: []resolver.Endpoint{
+				hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a0"}}}, []string{"a"}),
+				hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a1"}}}, []string{"a"}),
+				hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "b0"}}}, []string{"b"}),
+				hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "b1"}}}, []string{"b"}),
+			},
+			want: map[string][]resolver.Endpoint{
+				"a": {
+					hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a0"}}}, nil),
+					hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a1"}}}, nil),
+				},
+				"b": {
+					hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "b0"}}}, nil),
+					hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "b1"}}}, nil),
+				},
+			},
+		},
+		{
+			// Endpoints without hierarchy are ignored.
+			name: "without hierarchy",
+			eps: []resolver.Endpoint{
+				hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a0"}}}, []string{"a"}),
+				hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a1"}}}, []string{"a"}),
+				{Addresses: []resolver.Address{{Addr: "b0"}}, Attributes: nil},
+				{Addresses: []resolver.Address{{Addr: "b1"}}, Attributes: nil},
+			},
+			want: map[string][]resolver.Endpoint{
+				"a": {
+					hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a0"}}}, nil),
+					hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: "a1"}}}, nil),
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := hierarchy.Group(tt.eps); !cmp.Equal(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})) {
+				t.Errorf("Group() = %v, want %v", got, tt.want)
+				t.Errorf("diff: %v", cmp.Diff(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})))
+			}
+		})
+	}
+}
+
+func TestGroupE2E(t *testing.T) {
+	testHierarchy := map[string]map[string][]string{
+		"p0": {
+			"wt0": {"addr0", "addr1"},
+			"wt1": {"addr2", "addr3"},
+		},
+		"p1": {
+			"wt10": {"addr10", "addr11"},
+			"wt11": {"addr12", "addr13"},
+		},
+	}
+
+	var epsWithHierarchy []resolver.Endpoint
+	for p, wts := range testHierarchy {
+		path1 := []string{p}
+		for wt, addrs := range wts {
+			path2 := append([]string(nil), path1...)
+			path2 = append(path2, wt)
+			for _, addr := range addrs {
+				a := hierarchy.SetInEndpoint(resolver.Endpoint{Addresses: []resolver.Address{{Addr: addr}}}, path2)
+				epsWithHierarchy = append(epsWithHierarchy, a)
+			}
+		}
+	}
+
+	gotHierarchy := make(map[string]map[string][]string)
+	for p1, wts := range hierarchy.Group(epsWithHierarchy) {
+		gotHierarchy[p1] = make(map[string][]string)
+		for p2, eps := range hierarchy.Group(wts) {
+			for _, ep := range eps {
+				gotHierarchy[p1][p2] = append(gotHierarchy[p1][p2], ep.Addresses[0].Addr)
+			}
+		}
+	}
+
+	if !cmp.Equal(gotHierarchy, testHierarchy) {
+		t.Errorf("diff: %v", cmp.Diff(gotHierarchy, testHierarchy))
+	}
+}

internal/idle/idle.go

@@ -0,0 +1,289 @@
+/*
+ *
+ * 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 idle contains a component for managing idleness (entering and exiting)
+// based on RPC activity.
+package idle
+
+import (
+	"math"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// For overriding in unit tests.
+var timeAfterFunc = func(d time.Duration, f func()) *time.Timer {
+	return time.AfterFunc(d, f)
+}
+
+// ClientConn is the functionality provided by grpc.ClientConn to enter and exit
+// from idle mode.
+type ClientConn interface {
+	ExitIdleMode()
+	EnterIdleMode()
+}
+
+// Manager implements idleness detection and calls the ClientConn to enter/exit
+// idle mode when appropriate. Must be created by NewManager.
+type Manager struct {
+	// State accessed atomically.
+	lastCallEndTime           int64 // Unix timestamp in nanos; time when the most recent RPC completed.
+	activeCallsCount          int32 // Count of active RPCs; -math.MaxInt32 means channel is idle or is trying to get there.
+	activeSinceLastTimerCheck int32 // Boolean; True if there was an RPC since the last timer callback.
+	closed                    int32 // Boolean; True when the manager is closed.
+
+	// Can be accessed without atomics or mutex since these are set at creation
+	// time and read-only after that.
+	cc      ClientConn // Functionality provided by grpc.ClientConn.
+	timeout time.Duration
+
+	// idleMu is used to guarantee mutual exclusion in two scenarios:
+	// - Opposing intentions:
+	//   - a: Idle timeout has fired and handleIdleTimeout() is trying to put
+	//     the channel in idle mode because the channel has been inactive.
+	//   - b: At the same time an RPC is made on the channel, and OnCallBegin()
+	//     is trying to prevent the channel from going idle.
+	// - Competing intentions:
+	//   - The channel is in idle mode and there are multiple RPCs starting at
+	//     the same time, all trying to move the channel out of idle. Only one
+	//     of them should succeed in doing so, while the other RPCs should
+	//     piggyback on the first one and be successfully handled.
+	idleMu       sync.RWMutex
+	actuallyIdle bool
+	timer        *time.Timer
+}
+
+// NewManager creates a new idleness manager implementation for the
+// given idle timeout.  It begins in idle mode.
+func NewManager(cc ClientConn, timeout time.Duration) *Manager {
+	return &Manager{
+		cc:               cc,
+		timeout:          timeout,
+		actuallyIdle:     true,
+		activeCallsCount: -math.MaxInt32,
+	}
+}
+
+// resetIdleTimerLocked resets the idle timer to the given duration.  Called
+// when exiting idle mode or when the timer fires and we need to reset it.
+func (m *Manager) resetIdleTimerLocked(d time.Duration) {
+	if m.isClosed() || m.timeout == 0 || m.actuallyIdle {
+		return
+	}
+
+	// It is safe to ignore the return value from Reset() because this method is
+	// only ever called from the timer callback or when exiting idle mode.
+	if m.timer != nil {
+		m.timer.Stop()
+	}
+	m.timer = timeAfterFunc(d, m.handleIdleTimeout)
+}
+
+func (m *Manager) resetIdleTimer(d time.Duration) {
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+	m.resetIdleTimerLocked(d)
+}
+
+// handleIdleTimeout is the timer callback that is invoked upon expiry of the
+// configured idle timeout. The channel is considered inactive if there are no
+// ongoing calls and no RPC activity since the last time the timer fired.
+func (m *Manager) handleIdleTimeout() {
+	if m.isClosed() {
+		return
+	}
+
+	if atomic.LoadInt32(&m.activeCallsCount) > 0 {
+		m.resetIdleTimer(m.timeout)
+		return
+	}
+
+	// There has been activity on the channel since we last got here. Reset the
+	// timer and return.
+	if atomic.LoadInt32(&m.activeSinceLastTimerCheck) == 1 {
+		// Set the timer to fire after a duration of idle timeout, calculated
+		// from the time the most recent RPC completed.
+		atomic.StoreInt32(&m.activeSinceLastTimerCheck, 0)
+		m.resetIdleTimer(time.Duration(atomic.LoadInt64(&m.lastCallEndTime)-time.Now().UnixNano()) + m.timeout)
+		return
+	}
+
+	// Now that we've checked that there has been no activity, attempt to enter
+	// idle mode, which is very likely to succeed.
+	if m.tryEnterIdleMode(true) {
+		// Successfully entered idle mode. No timer needed until we exit idle.
+		return
+	}
+
+	// Failed to enter idle mode due to a concurrent RPC that kept the channel
+	// active, or because of an error from the channel. Undo the attempt to
+	// enter idle, and reset the timer to try again later.
+	m.resetIdleTimer(m.timeout)
+}
+
+// tryEnterIdleMode instructs the channel to enter idle mode. But before
+// that, it performs a last minute check to ensure that no new RPC has come in,
+// making the channel active.
+//
+// checkActivity controls if a check for RPC activity, since the last time the
+// idle_timeout fired, is made.
+
+// Return value indicates whether or not the channel moved to idle mode.
+//
+// Holds idleMu which ensures mutual exclusion with exitIdleMode.
+func (m *Manager) tryEnterIdleMode(checkActivity bool) bool {
+	// Setting the activeCallsCount to -math.MaxInt32 indicates to OnCallBegin()
+	// that the channel is either in idle mode or is trying to get there.
+	if !atomic.CompareAndSwapInt32(&m.activeCallsCount, 0, -math.MaxInt32) {
+		// This CAS operation can fail if an RPC started after we checked for
+		// activity in the timer handler, or one was ongoing from before the
+		// last time the timer fired, or if a test is attempting to enter idle
+		// mode without checking.  In all cases, abort going into idle mode.
+		return false
+	}
+	// N.B. if we fail to enter idle mode after this, we must re-add
+	// math.MaxInt32 to m.activeCallsCount.
+
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+
+	if atomic.LoadInt32(&m.activeCallsCount) != -math.MaxInt32 {
+		// We raced and lost to a new RPC. Very rare, but stop entering idle.
+		atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+		return false
+	}
+	if checkActivity && atomic.LoadInt32(&m.activeSinceLastTimerCheck) == 1 {
+		// A very short RPC could have come in (and also finished) after we
+		// checked for calls count and activity in handleIdleTimeout(), but
+		// before the CAS operation. So, we need to check for activity again.
+		atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+		return false
+	}
+
+	// No new RPCs have come in since we set the active calls count value to
+	// -math.MaxInt32. And since we have the lock, it is safe to enter idle mode
+	// unconditionally now.
+	m.cc.EnterIdleMode()
+	m.actuallyIdle = true
+	return true
+}
+
+// EnterIdleModeForTesting instructs the channel to enter idle mode.
+func (m *Manager) EnterIdleModeForTesting() {
+	m.tryEnterIdleMode(false)
+}
+
+// OnCallBegin is invoked at the start of every RPC.
+func (m *Manager) OnCallBegin() {
+	if m.isClosed() {
+		return
+	}
+
+	if atomic.AddInt32(&m.activeCallsCount, 1) > 0 {
+		// Channel is not idle now. Set the activity bit and allow the call.
+		atomic.StoreInt32(&m.activeSinceLastTimerCheck, 1)
+		return
+	}
+
+	// Channel is either in idle mode or is in the process of moving to idle
+	// mode. Attempt to exit idle mode to allow this RPC.
+	m.ExitIdleMode()
+	atomic.StoreInt32(&m.activeSinceLastTimerCheck, 1)
+}
+
+// ExitIdleMode instructs m to call the ClientConn's ExitIdleMode and update its
+// internal state.
+func (m *Manager) ExitIdleMode() {
+	// Holds idleMu which ensures mutual exclusion with tryEnterIdleMode.
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+
+	if m.isClosed() || !m.actuallyIdle {
+		// This can happen in three scenarios:
+		// - handleIdleTimeout() set the calls count to -math.MaxInt32 and called
+		//   tryEnterIdleMode(). But before the latter could grab the lock, an RPC
+		//   came in and OnCallBegin() noticed that the calls count is negative.
+		// - Channel is in idle mode, and multiple new RPCs come in at the same
+		//   time, all of them notice a negative calls count in OnCallBegin and get
+		//   here. The first one to get the lock would get the channel to exit idle.
+		// - Channel is not in idle mode, and the user calls Connect which calls
+		//   m.ExitIdleMode.
+		//
+		// In any case, there is nothing to do here.
+		return
+	}
+
+	m.cc.ExitIdleMode()
+
+	// Undo the idle entry process. This also respects any new RPC attempts.
+	atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+	m.actuallyIdle = false
+
+	// Start a new timer to fire after the configured idle timeout.
+	m.resetIdleTimerLocked(m.timeout)
+}
+
+// UnsafeSetNotIdle instructs the Manager to update its internal state to
+// reflect the reality that the channel is no longer in IDLE mode.
+//
+// N.B. This method is intended only for internal use by the gRPC client
+// when it exits IDLE mode **manually** from `Dial`. The callsite must ensure:
+//   - The channel was **actually in IDLE mode** immediately prior to the call.
+//   - There is **no concurrent activity** that could cause the channel to exit
+//     IDLE mode *naturally* at the same time.
+func (m *Manager) UnsafeSetNotIdle() {
+	m.idleMu.Lock()
+	defer m.idleMu.Unlock()
+
+	atomic.AddInt32(&m.activeCallsCount, math.MaxInt32)
+	m.actuallyIdle = false
+	m.resetIdleTimerLocked(m.timeout)
+}
+
+// OnCallEnd is invoked at the end of every RPC.
+func (m *Manager) OnCallEnd() {
+	if m.isClosed() {
+		return
+	}
+
+	// Record the time at which the most recent call finished.
+	atomic.StoreInt64(&m.lastCallEndTime, time.Now().UnixNano())
+
+	// Decrement the active calls count. This count can temporarily go negative
+	// when the timer callback is in the process of moving the channel to idle
+	// mode, but one or more RPCs come in and complete before the timer callback
+	// can get done with the process of moving to idle mode.
+	atomic.AddInt32(&m.activeCallsCount, -1)
+}
+
+func (m *Manager) isClosed() bool {
+	return atomic.LoadInt32(&m.closed) == 1
+}
+
+// Close stops the timer associated with the Manager, if it exists.
+func (m *Manager) Close() {
+	atomic.StoreInt32(&m.closed, 1)
+
+	m.idleMu.Lock()
+	if m.timer != nil {
+		m.timer.Stop()
+		m.timer = nil
+	}
+	m.idleMu.Unlock()
+}

internal/idle/idle_e2e_test.go

@@ -0,0 +1,589 @@
+/*
+ *
+ * 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 idle_test
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/roundrobin"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"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/resolver"
+	"google.golang.org/grpc/resolver/manual"
+	"google.golang.org/grpc/status"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+func init() {
+	channelz.TurnOn()
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+const (
+	defaultTestTimeout          = 10 * time.Second
+	defaultTestShortTimeout     = 100 * time.Millisecond
+	defaultTestShortIdleTimeout = 500 * time.Millisecond
+)
+
+// channelzTraceEventFound looks up the top-channels in channelz (expects a
+// single one), and checks if there is a trace event on the channel matching the
+// provided description string.
+func channelzTraceEventFound(ctx context.Context, wantDesc string) error {
+	for ctx.Err() == nil {
+		tcs, _ := channelz.GetTopChannels(0, 0)
+		if l := len(tcs); l != 1 {
+			return fmt.Errorf("when looking for channelz trace event with description %q, found %d top-level channels, want 1", wantDesc, l)
+		}
+		trace := tcs[0].Trace()
+		if trace == nil {
+			return fmt.Errorf("when looking for channelz trace event with description %q, no trace events found for top-level channel", wantDesc)
+		}
+
+		for _, e := range trace.Events {
+			if strings.Contains(e.Desc, wantDesc) {
+				return nil
+			}
+		}
+	}
+	return fmt.Errorf("when looking for channelz trace event with description %q, %w", wantDesc, ctx.Err())
+}
+
+// Registers a wrapped round_robin LB policy for the duration of this test that
+// retains all the functionality of the round_robin LB policy and makes the
+// balancer close event available for inspection by the test.
+//
+// Returns a channel that gets pinged when the balancer is closed.
+func registerWrappedRoundRobinPolicy(t *testing.T) chan struct{} {
+	rrBuilder := balancer.Get(roundrobin.Name)
+	closeCh := make(chan struct{}, 1)
+	stub.Register(roundrobin.Name, stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = rrBuilder.Build(bd.ClientConn, bd.BuildOptions)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			select {
+			case closeCh <- struct{}{}:
+			default:
+			}
+			bd.ChildBalancer.Close()
+		},
+	})
+	t.Cleanup(func() { balancer.Register(rrBuilder) })
+
+	return closeCh
+}
+
+// Tests the case where channel idleness is disabled by passing an idle_timeout
+// of 0. Verifies that a READY channel with no RPCs does not move to IDLE.
+func (s) TestChannelIdleness_Disabled_NoActivity(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Create a ClientConn with idle_timeout set to 0.
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(0), // Disable idleness.
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	cc.Connect()
+
+	// Start a test backend and push an address update via the resolver.
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Verify that the ClientConn stays in READY.
+	sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
+	defer sCancel()
+	testutils.AwaitNoStateChange(sCtx, t, cc, connectivity.Ready)
+
+	// Verify that the LB policy is not closed which is expected to happen when
+	// the channel enters IDLE.
+	sCtx, sCancel = context.WithTimeout(ctx, defaultTestShortIdleTimeout)
+	defer sCancel()
+	select {
+	case <-sCtx.Done():
+	case <-closeCh:
+		t.Fatal("LB policy closed when expected not to")
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that a READY channel with no RPCs moves to IDLE, and
+// the connection to the backend is closed.
+func (s) TestChannelIdleness_Enabled_NoActivity(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Create a ClientConn with a short idle_timeout.
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	cc.Connect()
+	// Start a test backend and push an address update via the resolver.
+	lis := testutils.NewListenerWrapper(t, nil)
+	backend := stubserver.StartTestService(t, &stubserver.StubServer{Listener: lis})
+	defer backend.Stop()
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Retrieve the wrapped conn from the listener.
+	v, err := lis.NewConnCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("Failed to retrieve conn from test listener: %v", err)
+	}
+	conn := v.(*testutils.ConnWrapper)
+
+	// Verify that the ClientConn moves to IDLE as there is no activity.
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+
+	// Verify idleness related channelz events.
+	if err := channelzTraceEventFound(ctx, "entering idle mode"); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify that the previously open connection is closed.
+	if _, err := conn.CloseCh.Receive(ctx); err != nil {
+		t.Fatalf("Failed when waiting for connection to be closed after channel entered IDLE: %v", err)
+	}
+
+	// Verify that the LB policy is closed.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for LB policy to be closed after the channel enters IDLE")
+	case <-closeCh:
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that a READY channel with an ongoing RPC stays READY.
+func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
+	tests := []struct {
+		name    string
+		makeRPC func(ctx context.Context, client testgrpc.TestServiceClient) error
+	}{
+		{
+			name: "unary",
+			makeRPC: func(ctx context.Context, client testgrpc.TestServiceClient) error {
+				if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+					return fmt.Errorf("EmptyCall RPC failed: %v", err)
+				}
+				return nil
+			},
+		},
+		{
+			name: "streaming",
+			makeRPC: func(ctx context.Context, client testgrpc.TestServiceClient) error {
+				stream, err := client.FullDuplexCall(ctx)
+				if err != nil {
+					t.Fatalf("FullDuplexCall RPC failed: %v", err)
+				}
+				if _, err := stream.Recv(); err != nil && err != io.EOF {
+					t.Fatalf("stream.Recv() failed: %v", err)
+				}
+				return nil
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			closeCh := registerWrappedRoundRobinPolicy(t)
+
+			// Create a ClientConn with a short idle_timeout.
+			r := manual.NewBuilderWithScheme("whatever")
+			dopts := []grpc.DialOption{
+				grpc.WithTransportCredentials(insecure.NewCredentials()),
+				grpc.WithResolvers(r),
+				grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+				grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+			}
+			cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+			if err != nil {
+				t.Fatalf("grpc.NewClient() failed: %v", err)
+			}
+			defer cc.Close()
+			cc.Connect()
+			// Start a test backend that keeps the RPC call active by blocking
+			// on a channel that is closed by the test later on.
+			blockCh := make(chan struct{})
+			backend := &stubserver.StubServer{
+				EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+					<-blockCh
+					return &testpb.Empty{}, nil
+				},
+				FullDuplexCallF: func(testgrpc.TestService_FullDuplexCallServer) error {
+					<-blockCh
+					return nil
+				},
+			}
+			if err := backend.StartServer(); err != nil {
+				t.Fatalf("Failed to start backend: %v", err)
+			}
+			defer backend.Stop()
+
+			// Push an address update containing the address of the above
+			// backend via the manual resolver.
+			r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+			// Verify that the ClientConn moves to READY.
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
+			testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+			// Spawn a goroutine to check for expected behavior while a blocking
+			// RPC all is made from the main test goroutine.
+			errCh := make(chan error, 1)
+			go func() {
+				defer close(blockCh)
+
+				// Verify that the ClientConn stays in READY.
+				sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
+				defer sCancel()
+				if cc.WaitForStateChange(sCtx, connectivity.Ready) {
+					errCh <- fmt.Errorf("state changed from %q to %q when no state change was expected", connectivity.Ready, cc.GetState())
+					return
+				}
+
+				// Verify that the LB policy is not closed which is expected to happen when
+				// the channel enters IDLE.
+				sCtx, sCancel = context.WithTimeout(ctx, defaultTestShortIdleTimeout)
+				defer sCancel()
+				select {
+				case <-sCtx.Done():
+				case <-closeCh:
+					errCh <- fmt.Errorf("LB policy closed when expected not to")
+				}
+				errCh <- nil
+			}()
+
+			if err := test.makeRPC(ctx, testgrpc.NewTestServiceClient(cc)); err != nil {
+				t.Fatalf("%s rpc failed: %v", test.name, err)
+			}
+
+			select {
+			case err := <-errCh:
+				if err != nil {
+					t.Fatal(err)
+				}
+			case <-ctx.Done():
+				t.Fatalf("Timeout when trying to verify that an active RPC keeps channel from moving to IDLE")
+			}
+		})
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that activity on a READY channel (frequent and short
+// RPCs) keeps it from moving to IDLE.
+func (s) TestChannelIdleness_Enabled_ActiveSinceLastCheck(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Create a ClientConn with a short idle_timeout.
+	r := manual.NewBuilderWithScheme("whatever")
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	cc.Connect()
+	// Start a test backend and push an address update via the resolver.
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// For a duration of three times the configured idle timeout, making RPCs
+	// every now and then and ensure that the channel does not move out of
+	// READY.
+	sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
+	defer sCancel()
+	go func() {
+		for ; sCtx.Err() == nil; <-time.After(defaultTestShortIdleTimeout / 4) {
+			client := testgrpc.NewTestServiceClient(cc)
+			if _, err := client.EmptyCall(sCtx, &testpb.Empty{}); err != nil {
+				// While iterating through this for loop, at some point in time,
+				// the context deadline will expire. It is safe to ignore that
+				// error code.
+				if status.Code(err) != codes.DeadlineExceeded {
+					t.Errorf("EmptyCall RPC failed: %v", err)
+					return
+				}
+			}
+		}
+	}()
+
+	// Verify that the ClientConn stays in READY.
+	testutils.AwaitNoStateChange(sCtx, t, cc, connectivity.Ready)
+
+	// Verify that the LB policy is not closed which is expected to happen when
+	// the channel enters IDLE.
+	select {
+	case <-sCtx.Done():
+	case <-closeCh:
+		t.Fatal("LB policy closed when expected not to")
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Verifies that a READY channel with no RPCs moves to IDLE. Also
+// verifies that a subsequent RPC on the IDLE channel kicks it out of IDLE.
+func (s) TestChannelIdleness_Enabled_ExitIdleOnRPC(t *testing.T) {
+	closeCh := registerWrappedRoundRobinPolicy(t)
+
+	// Start a test backend and set the bootstrap state of the resolver to
+	// include this address. This will ensure that when the resolver is
+	// restarted when exiting idle, it will push the same address to grpc again.
+	r := manual.NewBuilderWithScheme("whatever")
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Create a ClientConn with a short idle_timeout.
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+	cc.Connect()
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+
+	// Verify that the ClientConn moves to IDLE as there is no activity.
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+
+	// Verify idleness related channelz events.
+	if err := channelzTraceEventFound(ctx, "entering idle mode"); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify that the LB policy is closed.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for LB policy to be closed after the channel enters IDLE")
+	case <-closeCh:
+	}
+
+	// Make an RPC and ensure that it succeeds and moves the channel back to
+	// READY.
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall RPC failed: %v", err)
+	}
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+	if err := channelzTraceEventFound(ctx, "exiting idle mode"); err != nil {
+		t.Fatal(err)
+	}
+}
+
+// Tests the case where channel idleness is enabled by passing a small value for
+// idle_timeout. Simulates a race between the idle timer firing and RPCs being
+// initiated, after a period of inactivity on the channel.
+//
+// After a period of inactivity (for the configured idle timeout duration), when
+// RPCs are started, there are two possibilities:
+//   - the idle timer wins the race and puts the channel in idle. The RPCs then
+//     kick it out of idle.
+//   - the RPCs win the race, and therefore the channel never moves to idle.
+//
+// In either of these cases, all RPCs must succeed.
+func (s) TestChannelIdleness_Enabled_IdleTimeoutRacesWithRPCs(t *testing.T) {
+	// Start a test backend and set the bootstrap state of the resolver to
+	// include this address. This will ensure that when the resolver is
+	// restarted when exiting idle, it will push the same address to grpc again.
+	r := manual.NewBuilderWithScheme("whatever")
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Create a ClientConn with a short idle_timeout.
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	// Verify that the ClientConn moves to READY.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	client := testgrpc.NewTestServiceClient(cc)
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Errorf("EmptyCall RPC failed: %v", err)
+	}
+
+	// Make an RPC every defaultTestShortTimeout duration so as to race with the
+	// idle timeout. Whether the idle timeout wins the race or the RPC wins the
+	// race, RPCs must succeed.
+	for i := 0; i < 20; i++ {
+		<-time.After(defaultTestShortTimeout)
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+			t.Fatalf("EmptyCall RPC failed: %v", err)
+		}
+		t.Logf("Iteration %d succeeded", i)
+	}
+}
+
+// Tests the case where the channel is IDLE and we call cc.Connect.
+func (s) TestChannelIdleness_Connect(t *testing.T) {
+	// Start a test backend and set the bootstrap state of the resolver to
+	// include this address. This will ensure that when the resolver is
+	// restarted when exiting idle, it will push the same address to grpc again.
+	r := manual.NewBuilderWithScheme("whatever")
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+	r.InitialState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
+
+	// Create a ClientConn with a short idle_timeout.
+	dopts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithResolvers(r),
+		grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
+		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
+	}
+	cc, err := grpc.NewClient(r.Scheme()+":///test.server", dopts...)
+	if err != nil {
+		t.Fatalf("grpc.NewClient() failed: %v", err)
+	}
+	defer cc.Close()
+
+	// Verify that the ClientConn moves to IDLE.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
+
+	// Connect should exit channel idleness.
+	cc.Connect()
+
+	// Verify that the ClientConn moves back to READY.
+	testutils.AwaitState(ctx, t, cc, connectivity.Ready)
+}
+
+// TestChannelIdleness_RPCFailingBeforeAttempt_FinishesOnce verifies that an
+// RPC which fails during stream creation, before any attempt is started, is
+// reported as finished exactly once.
+func (s) TestChannelIdleness_RPCFailingBeforeAttempt_FinishesOnce(t *testing.T) {
+	backend := stubserver.StartTestService(t, nil)
+	defer backend.Stop()
+
+	cc, err := grpc.NewClient(fmt.Sprintf("dns:///%s", backend.Address), 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()
+
+	client := testgrpc.NewTestServiceClient(cc)
+
+	// Warm-up RPC: moves the channel to READY and fires the resolver's
+	// first-resolution event, so that the next RPC does not block waiting for
+	// addresses and instead proceeds to fail during stream creation.
+	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall() failed: %v", err)
+	}
+
+	// Issue an RPC with an already-canceled context so that it fails during
+	// stream creation, before any attempt is started, and count how many times
+	// gRPC reports it as finished. The contract is exactly once.
+	var finishCount atomic.Int32
+	cancel()
+	_, err = client.EmptyCall(ctx, &testpb.Empty{}, grpc.OnFinish(func(error) { finishCount.Add(1) }))
+	if got, want := status.Code(err), codes.Canceled; got != want {
+		t.Fatalf("EmptyCall() with canceled context returned code %v, want %v (err: %v)", got, want, err)
+	}
+	if got := finishCount.Load(); got != 1 {
+		t.Fatalf("OnFinish callback invoked %d times, want exactly 1", got)
+	}
+}

internal/idle/idle_test.go

@@ -0,0 +1,378 @@
+/*
+ *
+ * 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 idle
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+const (
+	defaultTestTimeout      = 10 * time.Second
+	defaultTestIdleTimeout  = 500 * time.Millisecond // A short idle_timeout for tests.
+	defaultTestShortTimeout = 10 * time.Millisecond  // A small deadline to wait for events expected to not happen.
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+type testEnforcer struct {
+	exitIdleCh  chan struct{}
+	enterIdleCh chan struct{}
+}
+
+func (ti *testEnforcer) ExitIdleMode() {
+	ti.exitIdleCh <- struct{}{}
+}
+
+func (ti *testEnforcer) EnterIdleMode() {
+	ti.enterIdleCh <- struct{}{}
+}
+
+func newTestEnforcer() *testEnforcer {
+	return &testEnforcer{
+		exitIdleCh:  make(chan struct{}, 1),
+		enterIdleCh: make(chan struct{}, 1),
+	}
+}
+
+// overrideNewTimer overrides the new timer creation function by ensuring that a
+// message is pushed on the returned channel everytime the timer fires.
+func overrideNewTimer(t *testing.T) <-chan struct{} {
+	t.Helper()
+
+	ch := make(chan struct{}, 1)
+	origTimeAfterFunc := timeAfterFunc
+	timeAfterFunc = func(d time.Duration, callback func()) *time.Timer {
+		return time.AfterFunc(d, func() {
+			select {
+			case ch <- struct{}{}:
+			default:
+			}
+			callback()
+		})
+	}
+	t.Cleanup(func() { timeAfterFunc = origTimeAfterFunc })
+	return ch
+}
+
+// TestManager_Disabled tests the case where the idleness manager is
+// disabled by passing an idle_timeout of 0. Verifies the following things:
+//   - timer callback does not fire
+//   - an RPC triggers a call to ExitIdleMode on the ClientConn
+//   - more calls to RPC termination (as compared to RPC initiation) does not
+//     result in an error log
+func (s) TestManager_Disabled(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	// Create an idleness manager that is disabled because of idleTimeout being
+	// set to `0`.
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(0))
+
+	// Ensure that the timer callback does not fire within a short deadline.
+	select {
+	case <-callbackCh:
+		t.Fatal("Idle timer callback fired when manager is disabled")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// The first invocation of OnCallBegin() should lead to a call to
+	// ExitIdleMode() on the enforcer.
+	go mgr.OnCallBegin()
+	select {
+	case <-enforcer.exitIdleCh:
+	case <-time.After(defaultTestShortTimeout):
+		t.Fatal("Timeout waiting for channel to move out of idle mode")
+	}
+
+	// If the number of calls to OnCallEnd() exceeds the number of calls to
+	// OnCallBegin(), the idleness manager is expected to throw an error log
+	// (which will cause our TestLogger to fail the test). But since the manager
+	// is disabled, this should not happen.
+	mgr.OnCallEnd()
+	mgr.OnCallEnd()
+
+	// The idleness manager is explicitly not closed here. But since the manager
+	// is disabled, it will not start the run goroutine, and hence we expect the
+	// leak checker to not find any leaked goroutines.
+}
+
+// TestManager_Enabled_TimerFires tests the case where the idle manager
+// is enabled. Ensures that when there are no RPCs, the timer callback is
+// invoked and the EnterIdleMode() method is invoked on the enforcer.
+func (s) TestManager_Enabled_TimerFires(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+	mgr.ExitIdleMode()
+
+	// Ensure that the timer callback fires within an appropriate amount of time.
+	select {
+	case <-callbackCh:
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for idle timer callback to fire")
+	}
+
+	// Ensure that the channel moves to idle mode eventually.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle")
+	}
+}
+
+// TestManager_Enabled_OngoingCall tests the case where the idle manager
+// is enabled. Ensures that when there is an ongoing RPC, the channel does not
+// enter idle mode.
+func (s) TestManager_Enabled_OngoingCall(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+	mgr.ExitIdleMode()
+
+	// Fire up a goroutine that simulates an ongoing RPC that is terminated
+	// after the timer callback fires for the first time.
+	timerFired := make(chan struct{})
+	go func() {
+		mgr.OnCallBegin()
+		<-timerFired
+		mgr.OnCallEnd()
+	}()
+
+	// Ensure that the timer callback fires and unblock the above goroutine.
+	select {
+	case <-callbackCh:
+		close(timerFired)
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for idle timer callback to fire")
+	}
+
+	// The invocation of the timer callback should not put the channel in idle
+	// mode since we had an ongoing RPC.
+	select {
+	case <-enforcer.enterIdleCh:
+		t.Fatalf("EnterIdleMode() called on enforcer when active RPC exists")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// Since we terminated the ongoing RPC and we have no other active RPCs, the
+	// channel must move to idle eventually.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle")
+	}
+}
+
+// TestManager_Enabled_ActiveSinceLastCheck tests the case where the
+// idle manager is enabled. Ensures that when there are active RPCs in the last
+// period (even though there is no active call when the timer fires), the
+// channel does not enter idle mode.
+func (s) TestManager_Enabled_ActiveSinceLastCheck(t *testing.T) {
+	callbackCh := overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+	mgr.ExitIdleMode()
+
+	// Fire up a goroutine that simulates unary RPCs until the timer callback
+	// fires.
+	timerFired := make(chan struct{})
+	go func() {
+		for ; ; <-time.After(defaultTestShortTimeout) {
+			mgr.OnCallBegin()
+			mgr.OnCallEnd()
+
+			select {
+			case <-timerFired:
+				return
+			default:
+			}
+		}
+	}()
+
+	// Ensure that the timer callback fires, and that we don't enter idle as
+	// part of this invocation of the timer callback, since we had some RPCs in
+	// this period.
+	select {
+	case <-callbackCh:
+		close(timerFired)
+	case <-time.After(2 * defaultTestIdleTimeout):
+		close(timerFired)
+		t.Fatal("Timeout waiting for idle timer callback to fire")
+	}
+	select {
+	case <-enforcer.enterIdleCh:
+		t.Fatalf("EnterIdleMode() called on enforcer when one RPC completed in the last period")
+	case <-time.After(defaultTestShortTimeout):
+	}
+
+	// Since the unary RPC terminated and we have no other active RPCs, the
+	// channel must move to idle eventually.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(defaultTestTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle")
+	}
+}
+
+// TestManager_Enabled_ExitIdleOnRPC tests the case where the idle
+// manager is enabled. Ensures that the channel moves out of idle when an RPC is
+// initiated.
+func (s) TestManager_Enabled_ExitIdleOnRPC(t *testing.T) {
+	overrideNewTimer(t)
+
+	enforcer := newTestEnforcer()
+	mgr := NewManager(enforcer, time.Duration(defaultTestIdleTimeout))
+	defer mgr.Close()
+
+	mgr.ExitIdleMode()
+	<-enforcer.exitIdleCh
+	// Ensure that the channel moves to idle since there are no RPCs.
+	select {
+	case <-enforcer.enterIdleCh:
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for channel to move to idle mode")
+	}
+
+	for i := 0; i < 100; i++ {
+		// A call to OnCallBegin and OnCallEnd simulates an RPC.
+		go func() {
+			mgr.OnCallBegin()
+			mgr.OnCallEnd()
+		}()
+	}
+
+	// Ensure that the channel moves out of idle as a result of the above RPC.
+	select {
+	case <-enforcer.exitIdleCh:
+	case <-time.After(2 * defaultTestIdleTimeout):
+		t.Fatal("Timeout waiting for channel to move out of idle mode")
+	}
+
+	// Ensure that only one call to exit idle mode is made to the CC.
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	select {
+	case <-enforcer.exitIdleCh:
+		t.Fatal("More than one call to exit idle mode on the ClientConn; only one expected")
+	case <-sCtx.Done():
+	}
+}
+
+type racyState int32
+
+const (
+	stateInitial racyState = iota
+	stateEnteredIdle
+	stateExitedIdle
+	stateActiveRPCs
+)
+
+// racyEnforcer is a test idleness enforcer used specifically to test the
+// race between idle timeout and incoming RPCs.
+type racyEnforcer struct {
+	t       *testing.T
+	state   *racyState // Accessed atomically.
+	started bool
+}
+
+// ExitIdleMode sets the internal state to stateExitedIdle. We should only ever
+// exit idle when we are currently in idle.
+func (ri *racyEnforcer) ExitIdleMode() {
+	// Set only on the initial ExitIdleMode
+	if ri.started == false {
+		if !atomic.CompareAndSwapInt32((*int32)(ri.state), int32(stateInitial), int32(stateInitial)) {
+			ri.t.Errorf("idleness enforcer's first ExitIdleMode after EnterIdleMode")
+			return
+		}
+		ri.started = true
+		return
+	}
+	if !atomic.CompareAndSwapInt32((*int32)(ri.state), int32(stateEnteredIdle), int32(stateExitedIdle)) {
+		ri.t.Errorf("idleness enforcer asked to exit idle when it did not enter idle earlier")
+		return
+	}
+}
+
+// EnterIdleMode attempts to set the internal state to stateEnteredIdle. We should only ever enter idle before RPCs start.
+func (ri *racyEnforcer) EnterIdleMode() {
+	if !atomic.CompareAndSwapInt32((*int32)(ri.state), int32(stateInitial), int32(stateEnteredIdle)) {
+		ri.t.Errorf("idleness enforcer asked to enter idle after rpcs started")
+	}
+}
+
+// TestManager_IdleTimeoutRacesWithOnCallBegin tests the case where firing of
+// the idle timeout races with an incoming RPC. The test verifies that if the
+// timer callback wins the race and puts the channel in idle, the RPCs can kick
+// it out of idle. And if the RPCs win the race and keep the channel active,
+// then the timer callback should not attempt to put the channel in idle mode.
+func (s) TestManager_IdleTimeoutRacesWithOnCallBegin(t *testing.T) {
+	// Run multiple iterations to simulate different possibilities.
+	for i := 0; i < 20; i++ {
+		t.Run(fmt.Sprintf("iteration=%d", i), func(t *testing.T) {
+			var idlenessState racyState
+			enforcer := &racyEnforcer{t: t, state: &idlenessState}
+
+			// Configure a large idle timeout so that we can control the
+			// race between the timer callback and RPCs.
+			mgr := NewManager(enforcer, time.Duration(10*time.Minute))
+			defer mgr.Close()
+			mgr.ExitIdleMode()
+
+			var wg sync.WaitGroup
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				<-time.After(defaultTestIdleTimeout / 50)
+				mgr.handleIdleTimeout()
+			}()
+			for j := 0; j < 5; j++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					// Wait for the configured idle timeout and simulate an RPC to
+					// race with the idle timeout timer callback.
+					<-time.After(defaultTestIdleTimeout / 50)
+					mgr.OnCallBegin()
+					atomic.StoreInt32((*int32)(&idlenessState), int32(stateActiveRPCs))
+					mgr.OnCallEnd()
+				}()
+			}
+			wg.Wait()
+		})
+	}
+}

internal/internal.go

@@ -0,0 +1,304 @@
+/*
+ * Copyright 2016 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package internal contains gRPC-internal code, to avoid polluting
+// the godoc of the top-level grpc package.  It must not import any grpc
+// symbols to avoid circular dependencies.
+package internal
+
+import (
+	"context"
+	"time"
+
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+var (
+	// HealthCheckFunc is used to provide client-side LB channel health checking
+	HealthCheckFunc HealthChecker
+	// RegisterClientHealthCheckListener is used to provide a listener for
+	// updates from the client-side health checking service. It returns a
+	// function that can be called to stop the health producer.
+	RegisterClientHealthCheckListener any // func(ctx context.Context, sc balancer.SubConn, serviceName string, listener func(balancer.SubConnState)) func()
+	// BalancerUnregister is exported by package balancer to unregister a balancer.
+	BalancerUnregister func(name string)
+	// KeepaliveMinPingTime is the minimum ping interval.  This must be 10s by
+	// default, but tests may wish to set it lower for convenience.
+	KeepaliveMinPingTime = 10 * time.Second
+	// KeepaliveMinServerPingTime is the minimum ping interval for servers.
+	// This must be 1s by default, but tests may wish to set it lower for
+	// convenience.
+	KeepaliveMinServerPingTime = time.Second
+	// ParseServiceConfig parses a JSON representation of the service config.
+	ParseServiceConfig any // func(string) *serviceconfig.ParseResult
+	// EqualServiceConfigForTesting is for testing service config generation and
+	// parsing. Both a and b should be returned by ParseServiceConfig.
+	// This function compares the config without rawJSON stripped, in case the
+	// there's difference in white space.
+	EqualServiceConfigForTesting func(a, b serviceconfig.Config) bool
+	// GetCertificateProviderBuilder returns the registered builder for the
+	// given name. This is set by package certprovider for use from xDS
+	// bootstrap code while parsing certificate provider configs in the
+	// bootstrap file.
+	GetCertificateProviderBuilder any // func(string) certprovider.Builder
+	// GetXDSHandshakeInfoForTesting returns a pointer to the xds.HandshakeInfo
+	// stored in the passed in attributes. This is set by
+	// credentials/xds/xds.go.
+	GetXDSHandshakeInfoForTesting any // func (*attributes.Attributes) *unsafe.Pointer
+	// GetServerCredentials returns the transport credentials configured on a
+	// gRPC server. An xDS-enabled server needs to know what type of credentials
+	// is configured on the underlying gRPC server. This is set by server.go.
+	GetServerCredentials any // func (*grpc.Server) credentials.TransportCredentials
+	// MetricsRecorderForServer returns the MetricsRecorderList derived from a
+	// server's stats handlers.
+	MetricsRecorderForServer any // func (*grpc.Server) estats.MetricsRecorder
+	// CanonicalString returns the canonical string of the code defined here:
+	// https://github.com/grpc/grpc/blob/master/doc/statuscodes.md.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	CanonicalString any // func (codes.Code) string
+	// IsRegisteredMethod returns whether the passed in method is registered as
+	// a method on the server.
+	IsRegisteredMethod any // func(*grpc.Server, string) bool
+	// ServerFromContext returns the server from the context.
+	ServerFromContext any // func(context.Context) *grpc.Server
+	// AddGlobalServerOptions adds an array of ServerOption that will be
+	// effective globally for newly created servers. The priority will be: 1.
+	// user-provided; 2. this method; 3. default values.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	AddGlobalServerOptions any // func(opt ...ServerOption)
+	// ClearGlobalServerOptions clears the array of extra ServerOption. This
+	// method is useful in testing and benchmarking.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	ClearGlobalServerOptions func()
+	// AddGlobalDialOptions adds an array of DialOption that will be effective
+	// globally for newly created client channels. The priority will be: 1.
+	// user-provided; 2. this method; 3. default values.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	AddGlobalDialOptions any // func(opt ...DialOption)
+	// DisableGlobalDialOptions returns a DialOption that prevents the
+	// ClientConn from applying the global DialOptions (set via
+	// AddGlobalDialOptions).
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	DisableGlobalDialOptions any // func() grpc.DialOption
+	// ClearGlobalDialOptions clears the array of extra DialOption. This
+	// method is useful in testing and benchmarking.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	ClearGlobalDialOptions func()
+
+	// AddGlobalPerTargetDialOptions adds a PerTargetDialOption that will be
+	// configured for newly created ClientConns.
+	AddGlobalPerTargetDialOptions any // func (opt any)
+	// ClearGlobalPerTargetDialOptions clears the slice of global late apply
+	// dial options.
+	ClearGlobalPerTargetDialOptions func()
+
+	// JoinDialOptions combines the dial options passed as arguments into a
+	// single dial option.
+	JoinDialOptions any // func(...grpc.DialOption) grpc.DialOption
+	// JoinServerOptions combines the server options passed as arguments into a
+	// single server option.
+	JoinServerOptions any // func(...grpc.ServerOption) grpc.ServerOption
+
+	// WithBinaryLogger returns a DialOption that specifies the binary logger
+	// for a ClientConn.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	WithBinaryLogger any // func(binarylog.Logger) grpc.DialOption
+	// BinaryLogger returns a ServerOption that can set the binary logger for a
+	// server.
+	//
+	// This is used in the 1.0 release of gcp/observability, and thus must not be
+	// deleted or changed.
+	BinaryLogger any // func(binarylog.Logger) grpc.ServerOption
+
+	// SubscribeToConnectivityStateChanges adds a grpcsync.Subscriber to a
+	// provided grpc.ClientConn.
+	SubscribeToConnectivityStateChanges any // func(*grpc.ClientConn, grpcsync.Subscriber)
+
+	// NewXDSResolverWithConfigForTesting creates a new xds resolver builder using
+	// the provided xds bootstrap config instead of the global configuration from
+	// the supported environment variables.  The resolver.Builder is meant to be
+	// used in conjunction with the grpc.WithResolvers DialOption.
+	//
+	// Testing Only
+	//
+	// This function should ONLY be used for testing and may not work with some
+	// other features, including the CSDS service.
+	NewXDSResolverWithConfigForTesting any // func([]byte) (resolver.Builder, error)
+
+	// NewXDSResolverWithPoolForTesting creates a new xDS resolver builder
+	// using the provided xDS pool instead of creating a new one using the
+	// bootstrap configuration specified by the supported environment variables.
+	// The resolver.Builder is meant to be used in conjunction with the
+	// grpc.WithResolvers DialOption. The resolver.Builder does not take
+	// ownership of the provided xDS client and it is the responsibility of the
+	// caller to close the client when no longer required.
+	//
+	// Testing Only
+	//
+	// This function should ONLY be used for testing and may not work with some
+	// other features, including the CSDS service.
+	NewXDSResolverWithPoolForTesting any // func(*xdsclient.Pool) (resolver.Builder, error)
+
+	// NewXDSResolverWithClientForTesting creates a new xDS resolver builder
+	// using the provided xDS client instead of creating a new one using the
+	// bootstrap configuration specified by the supported environment variables.
+	// The resolver.Builder is meant to be used in conjunction with the
+	// grpc.WithResolvers DialOption. The resolver.Builder does not take
+	// ownership of the provided xDS client and it is the responsibility of the
+	// caller to close the client when no longer required.
+	//
+	// Testing Only
+	//
+	// This function should ONLY be used for testing and may not work with some
+	// other features, including the CSDS service.
+	NewXDSResolverWithClientForTesting any // func(xdsclient.XDSClient) (resolver.Builder, error)
+
+	// ORCAAllowAnyMinReportingInterval is for examples/orca use ONLY.
+	ORCAAllowAnyMinReportingInterval any // func(so *orca.ServiceOptions)
+
+	// GRPCResolverSchemeExtraMetadata determines when gRPC will add extra
+	// metadata to RPCs.
+	GRPCResolverSchemeExtraMetadata = "xds"
+
+	// EnterIdleModeForTesting gets the ClientConn to enter IDLE mode.
+	EnterIdleModeForTesting any // func(*grpc.ClientConn)
+
+	// ExitIdleModeForTesting gets the ClientConn to exit IDLE mode.
+	ExitIdleModeForTesting any // func(*grpc.ClientConn) error
+
+	// ChannelzTurnOffForTesting disables the Channelz service for testing
+	// purposes.
+	ChannelzTurnOffForTesting func()
+
+	// TriggerXDSResourceNotFoundForTesting causes the provided xDS Client to
+	// invoke resource-not-found error for the given resource type and name.
+	TriggerXDSResourceNotFoundForTesting any // func(xdsclient.XDSClient, xdsresource.Type, string) error
+
+	// FromOutgoingContextRaw returns the un-merged, intermediary contents of
+	// metadata.rawMD.
+	FromOutgoingContextRaw any // func(context.Context) (metadata.MD, [][]string, bool)
+
+	// UserSetDefaultScheme is set to true if the user has overridden the
+	// default resolver scheme.
+	UserSetDefaultScheme = false
+
+	// SnapshotMetricRegistryForTesting snapshots the global data of the metric
+	// registry. Returns a cleanup function that sets the metric registry to its
+	// original state. Only called in testing functions.
+	SnapshotMetricRegistryForTesting func() func()
+
+	// SetBufferPoolingThresholdForTesting updates the buffer pooling threshold, for
+	// testing purposes.
+	SetBufferPoolingThresholdForTesting any // func(int)
+
+	// TimeAfterFunc is used to create timers. During tests the function is
+	// replaced to track allocated timers and fail the test if a timer isn't
+	// cancelled.
+	TimeAfterFunc = func(d time.Duration, f func()) Timer {
+		return time.AfterFunc(d, f)
+	}
+
+	// NewStreamWaitingForResolver is a test hook that is triggered when a
+	// new stream blocks while waiting for name resolution. This can be
+	// used in tests to synchronize resolver updates and avoid race conditions.
+	// When set, the function will be called before the stream enters
+	// the blocking state.
+	NewStreamWaitingForResolver = func() {}
+
+	// AddressToTelemetryLabels is an xDS-provided function to extract telemetry
+	// labels from a resolver.Address. Callers must assert its type before calling.
+	AddressToTelemetryLabels any // func(addr resolver.Address) map[string]string
+
+	// AsyncReporterCleanupDelegate is initialized to a pass-through function by
+	// default (production behavior), allowing tests to swap it with an
+	// implementation which tracks registration of async reporter and its
+	// corresponding cleanup.
+	AsyncReporterCleanupDelegate = func(cleanup func()) func() {
+		return cleanup
+	}
+
+	// XDSFilterWrapperOption returns a ServerOption that sets the internal
+	// stream wrapper used for server-side xDS HTTP filters.
+	XDSFilterWrapperOption any // func(func(grpc.ServerStream) (grpc.ServerStream, error)) grpc.ServerOption
+)
+
+// HealthChecker defines the signature of the client-side LB channel health
+// checking function.
+//
+// The implementation is expected to create a health checking RPC stream by
+// calling newStream(), watch for the health status of serviceName, and report
+// its health back by calling setConnectivityState().
+//
+// The health checking protocol is defined at:
+// https://github.com/grpc/grpc/blob/master/doc/health-checking.md
+type HealthChecker func(ctx context.Context, newStream func(string) (any, error), setConnectivityState func(connectivity.State, error), serviceName string) error
+
+const (
+	// CredsBundleModeFallback switches GoogleDefaultCreds to fallback mode.
+	CredsBundleModeFallback = "fallback"
+	// CredsBundleModeBalancer switches GoogleDefaultCreds to grpclb balancer
+	// mode.
+	CredsBundleModeBalancer = "balancer"
+	// CredsBundleModeBackendFromBalancer switches GoogleDefaultCreds to mode
+	// that supports backend returned by grpclb balancer.
+	CredsBundleModeBackendFromBalancer = "backend-from-balancer"
+)
+
+// RLSLoadBalancingPolicyName 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.
+const RLSLoadBalancingPolicyName = "rls_experimental"
+
+// EnforceSubConnEmbedding is used to enforce proper SubConn implementation
+// embedding.
+type EnforceSubConnEmbedding interface {
+	enforceSubConnEmbedding()
+}
+
+// EnforceClientConnEmbedding is used to enforce proper ClientConn implementation
+// embedding.
+type EnforceClientConnEmbedding interface {
+	enforceClientConnEmbedding()
+}
+
+// Timer is an interface to allow injecting different time.Timer implementations
+// during tests.
+type Timer interface {
+	Stop() bool
+}
+
+// EnforceMetricsRecorderEmbedding is used to enforce proper MetricsRecorder
+// implementation embedding.
+type EnforceMetricsRecorderEmbedding interface {
+	enforceMetricsRecorderEmbedding()
+}

internal/leakcheck/leakcheck.go

@@ -0,0 +1,489 @@
+/*
+ *
+ * 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 leakcheck contains functions to check leaked goroutines and buffers.
+//
+// Call the following at the beginning of test:
+//
+//	defer leakcheck.NewLeakChecker(t).Check()
+package leakcheck
+
+import (
+	"context"
+	"fmt"
+	"runtime"
+	"runtime/debug"
+	"slices"
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/mem"
+)
+
+// failTestsOnLeakedBuffers is a special flag that will cause tests to fail if
+// leaked buffers are detected, instead of simply logging them as an
+// informational failure. This can be enabled with the "checkbuffers" compile
+// flag, e.g.:
+//
+//	go test -tags=checkbuffers
+var failTestsOnLeakedBuffers = false
+
+func init() {
+	defaultPool := mem.DefaultBufferPool()
+	globalPool.Store(&defaultPool)
+	internal.SetDefaultBufferPool.(func(mem.BufferPool))(&globalPool)
+}
+
+var globalPool swappableBufferPool
+var globalTimerTracker *timerFactory
+
+type swappableBufferPool struct {
+	atomic.Pointer[mem.BufferPool]
+}
+
+func (b *swappableBufferPool) Get(length int) *[]byte {
+	return (*b.Load()).Get(length)
+}
+
+func (b *swappableBufferPool) Put(buf *[]byte) {
+	(*b.Load()).Put(buf)
+}
+
+// SetTrackingBufferPool replaces the default buffer pool in the mem package to
+// one that tracks where buffers are allocated. CheckTrackingBufferPool should
+// then be invoked at the end of the test to validate that all buffers pulled
+// from the pool were returned.
+func SetTrackingBufferPool(logger Logger) {
+	newPool := mem.BufferPool(&trackingBufferPool{
+		pool:             *globalPool.Load(),
+		logger:           logger,
+		allocatedBuffers: make(map[*[]byte][]uintptr),
+	})
+	globalPool.Store(&newPool)
+}
+
+// CheckTrackingBufferPool undoes the effects of SetTrackingBufferPool, and fails
+// unit tests if not all buffers were returned. It is invalid to invoke this
+// function without previously having invoked SetTrackingBufferPool.
+func CheckTrackingBufferPool() {
+	p := (*globalPool.Load()).(*trackingBufferPool)
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	globalPool.Store(&p.pool)
+
+	type uniqueTrace struct {
+		stack []uintptr
+		count int
+	}
+
+	var totalLeakedBuffers int
+	var uniqueTraces []uniqueTrace
+	for _, stack := range p.allocatedBuffers {
+		idx, ok := slices.BinarySearchFunc(uniqueTraces, stack, func(trace uniqueTrace, stack []uintptr) int {
+			return slices.Compare(trace.stack, stack)
+		})
+		if !ok {
+			uniqueTraces = slices.Insert(uniqueTraces, idx, uniqueTrace{stack: stack})
+		}
+		uniqueTraces[idx].count++
+		totalLeakedBuffers++
+	}
+
+	for _, ut := range uniqueTraces {
+		frames := runtime.CallersFrames(ut.stack)
+		var trace strings.Builder
+		for {
+			f, ok := frames.Next()
+			if !ok {
+				break
+			}
+			trace.WriteString(f.Function)
+			trace.WriteString("\n\t")
+			trace.WriteString(f.File)
+			trace.WriteString(":")
+			trace.WriteString(strconv.Itoa(f.Line))
+			trace.WriteString("\n")
+		}
+		format := "%d allocated buffers never freed:\n%s"
+		args := []any{ut.count, trace.String()}
+		if failTestsOnLeakedBuffers {
+			p.logger.Errorf(format, args...)
+		} else {
+			p.logger.Logf("WARNING "+format, args...)
+		}
+	}
+
+	if totalLeakedBuffers > 0 {
+		p.logger.Logf("%g%% of buffers never freed", float64(totalLeakedBuffers)/float64(p.bufferCount))
+	}
+}
+
+type trackingBufferPool struct {
+	pool   mem.BufferPool
+	logger Logger
+
+	lock             sync.Mutex
+	bufferCount      int
+	allocatedBuffers map[*[]byte][]uintptr
+}
+
+func (p *trackingBufferPool) Get(length int) *[]byte {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	p.bufferCount++
+	buf := p.pool.Get(length)
+	p.allocatedBuffers[buf] = currentStack(2)
+	return buf
+}
+
+func (p *trackingBufferPool) Put(buf *[]byte) {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	if _, ok := p.allocatedBuffers[buf]; !ok {
+		p.logger.Errorf("Unknown buffer freed:\n%s", string(debug.Stack()))
+	} else {
+		delete(p.allocatedBuffers, buf)
+	}
+	p.pool.Put(buf)
+}
+
+var goroutinesToIgnore = []string{
+	"testing.Main(",
+	"testing.tRunner(",
+	"testing.(*M).",
+	"runtime.goexit",
+	"created by runtime.gc",
+	"created by runtime/trace.Start",
+	"interestingGoroutines",
+	"runtime.MHeap_Scavenger",
+	"signal.signal_recv",
+	"sigterm.handler",
+	"runtime_mcall",
+	"(*loggingT).flushDaemon",
+	"goroutine in C code",
+	// Ignore the http read/write goroutines. gce metadata.OnGCE() was leaking
+	// these, root cause unknown.
+	//
+	// https://github.com/grpc/grpc-go/issues/5171
+	// https://github.com/grpc/grpc-go/issues/5173
+	"created by net/http.(*Transport).dialConn",
+}
+
+// RegisterIgnoreGoroutine appends s into the ignore goroutine list. The
+// goroutines whose stack trace contains s will not be identified as leaked
+// goroutines. Not thread-safe, only call this function in init().
+func RegisterIgnoreGoroutine(s string) {
+	goroutinesToIgnore = append(goroutinesToIgnore, s)
+}
+
+func ignore(g string) bool {
+	sl := strings.SplitN(g, "\n", 2)
+	if len(sl) != 2 {
+		return true
+	}
+	stack := strings.TrimSpace(sl[1])
+	if strings.HasPrefix(stack, "testing.RunTests") {
+		return true
+	}
+
+	if stack == "" {
+		return true
+	}
+
+	for _, s := range goroutinesToIgnore {
+		if strings.Contains(stack, s) {
+			return true
+		}
+	}
+
+	return false
+}
+
+// interestingGoroutines returns all goroutines we care about for the purpose of
+// leak checking. It excludes testing or runtime ones.
+func interestingGoroutines() (gs []string) {
+	buf := make([]byte, 2<<20)
+	buf = buf[:runtime.Stack(buf, true)]
+	for _, g := range strings.Split(string(buf), "\n\n") {
+		if !ignore(g) {
+			gs = append(gs, g)
+		}
+	}
+	sort.Strings(gs)
+	return
+}
+
+// Logger is the interface that wraps the Logf and Errorf method. It's a subset
+// of testing.TB to make it easy to use this package.
+type Logger interface {
+	Logf(format string, args ...any)
+	Errorf(format string, args ...any)
+}
+
+// CheckGoroutines looks at the currently-running goroutines and checks if there
+// are any interesting (created by gRPC) goroutines leaked. It waits up to 10
+// seconds in the error cases.
+func CheckGoroutines(ctx context.Context, logger Logger) {
+	// Loop, waiting for goroutines to shut down.
+	// Wait up to timeout, but finish as quickly as possible.
+	var leaked []string
+	for ctx.Err() == nil {
+		if leaked = interestingGoroutines(); len(leaked) == 0 {
+			return
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	for _, g := range leaked {
+		logger.Errorf("Leaked goroutine: %v", g)
+	}
+}
+
+// LeakChecker captures a Logger and is returned by NewLeakChecker as a
+// convenient method to set up leak check tests in a unit test.
+type LeakChecker struct {
+	logger Logger
+}
+
+// NewLeakChecker offers a convenient way to set up the leak checks for a
+// specific unit test. It can be used as follows, at the beginning of tests:
+//
+//	defer leakcheck.NewLeakChecker(t).Check()
+//
+// It initially invokes SetTrackingBufferPool to set up buffer tracking, then the
+// deferred LeakChecker.Check call will invoke CheckTrackingBufferPool and
+// CheckGoroutines with a default timeout of 10 seconds.
+func NewLeakChecker(logger Logger) *LeakChecker {
+	SetTrackingBufferPool(logger)
+	return &LeakChecker{logger: logger}
+}
+
+type timerFactory struct {
+	mu              sync.Mutex
+	allocatedTimers map[internal.Timer][]uintptr
+}
+
+func (tf *timerFactory) timeAfterFunc(d time.Duration, f func()) internal.Timer {
+	tf.mu.Lock()
+	defer tf.mu.Unlock()
+	ch := make(chan internal.Timer, 1)
+	timer := time.AfterFunc(d, func() {
+		f()
+		tf.remove(<-ch)
+	})
+	ch <- timer
+	tf.allocatedTimers[timer] = currentStack(2)
+	return &trackingTimer{
+		Timer:  timer,
+		parent: tf,
+	}
+}
+
+func (tf *timerFactory) remove(timer internal.Timer) {
+	tf.mu.Lock()
+	defer tf.mu.Unlock()
+	delete(tf.allocatedTimers, timer)
+}
+
+func (tf *timerFactory) pendingTimers() []string {
+	tf.mu.Lock()
+	defer tf.mu.Unlock()
+	leaked := []string{}
+	for _, stack := range tf.allocatedTimers {
+		leaked = append(leaked, fmt.Sprintf("Allocated timer never cancelled:\n%s", traceToString(stack)))
+	}
+	return leaked
+}
+
+type trackingTimer struct {
+	internal.Timer
+	parent *timerFactory
+}
+
+func (t *trackingTimer) Stop() bool {
+	t.parent.remove(t.Timer)
+	return t.Timer.Stop()
+}
+
+// TrackTimers replaces internal.TimerAfterFunc with one that tracks timer
+// creations, stoppages and expirations. CheckTimers should then be invoked at
+// the end of the test to validate that all timers created have either executed
+// or are cancelled.
+func TrackTimers() {
+	globalTimerTracker = &timerFactory{
+		allocatedTimers: make(map[internal.Timer][]uintptr),
+	}
+	internal.TimeAfterFunc = globalTimerTracker.timeAfterFunc
+}
+
+// CheckTimers undoes the effects of TrackTimers, and fails unit tests if not
+// all timers were cancelled or executed. It is invalid to invoke this function
+// without previously having invoked TrackTimers.
+func CheckTimers(ctx context.Context, logger Logger) {
+	tt := globalTimerTracker
+
+	// Loop, waiting for timers to be cancelled.
+	// Wait up to timeout, but finish as quickly as possible.
+	var leaked []string
+	for ctx.Err() == nil {
+		if leaked = tt.pendingTimers(); len(leaked) == 0 {
+			return
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	for _, g := range leaked {
+		logger.Errorf("Leaked timers: %v", g)
+	}
+
+	// Reset the internal function.
+	internal.TimeAfterFunc = func(d time.Duration, f func()) internal.Timer {
+		return time.AfterFunc(d, f)
+	}
+}
+
+func currentStack(skip int) []uintptr {
+	var stackBuf [16]uintptr
+	var stack []uintptr
+	skip++
+	for {
+		n := runtime.Callers(skip, stackBuf[:])
+		stack = append(stack, stackBuf[:n]...)
+		if n < len(stackBuf) {
+			break
+		}
+		skip += len(stackBuf)
+	}
+	return stack
+}
+
+func traceToString(stack []uintptr) string {
+	frames := runtime.CallersFrames(stack)
+	var trace strings.Builder
+	for {
+		f, ok := frames.Next()
+		if !ok {
+			break
+		}
+		trace.WriteString(f.Function)
+		trace.WriteString("\n\t")
+		trace.WriteString(f.File)
+		trace.WriteString(":")
+		trace.WriteString(strconv.Itoa(f.Line))
+		trace.WriteString("\n")
+	}
+	return trace.String()
+}
+
+// Async Reporter Leak Checking
+
+var asyncReporterTracker *reporterTracker
+
+type reporterTracker struct {
+	mu          sync.Mutex
+	allocations map[*int][]uintptr
+}
+
+func newReporterTracker() *reporterTracker {
+	return &reporterTracker{
+		allocations: make(map[*int][]uintptr),
+	}
+}
+
+// register records the stack trace.
+func (rt *reporterTracker) register() *int {
+	rt.mu.Lock()
+	defer rt.mu.Unlock()
+
+	id := new(int)
+	// Skip 4 frames: register -> internal.Delegate -> stats.RegisterAsyncReporter -> Caller
+	rt.allocations[id] = currentStack(4)
+	return id
+}
+
+// unregister removes the ID.
+func (rt *reporterTracker) unregister(id *int) {
+	rt.mu.Lock()
+	defer rt.mu.Unlock()
+	delete(rt.allocations, id)
+}
+
+// leakedStackTraces returns formatted stack traces for all currently registered
+// reporters.
+func (rt *reporterTracker) leakedStackTraces() []string {
+	rt.mu.Lock()
+	defer rt.mu.Unlock()
+
+	var traces []string
+	for _, pcs := range rt.allocations {
+		msg := "\n--- Leaked Async Reporter Registration ---\n" + traceToString(pcs)
+		traces = append(traces, msg)
+	}
+	return traces
+}
+
+// TrackAsyncReporters installs the tracking delegate.
+func TrackAsyncReporters() {
+	asyncReporterTracker = newReporterTracker()
+
+	// Swap the delegate: Replace the default pass-through with tracking logic.
+	internal.AsyncReporterCleanupDelegate = func(originalCleanup func()) func() {
+		// 1. Capture Stack Trace (happens during Registration)
+		token := asyncReporterTracker.register()
+
+		// 2. Return Wrapped Cleanup
+		return func() {
+			// Defer unregister to ensure we stop tracking even if the original cleanup panics.
+			defer asyncReporterTracker.unregister(token)
+
+			if originalCleanup != nil {
+				originalCleanup()
+			}
+		}
+	}
+}
+
+// CheckAsyncReporters verifies that no leaks exist and restores the default delegate.
+func CheckAsyncReporters(logger Logger) {
+	// Restore the delegate: Reset to the default pass-through behavior.
+	internal.AsyncReporterCleanupDelegate = func(cleanup func()) func() {
+		return cleanup
+	}
+
+	if asyncReporterTracker == nil {
+		return
+	}
+
+	leaks := asyncReporterTracker.leakedStackTraces()
+	if len(leaks) > 0 {
+		// Join all stack traces into one message
+		allTraces := ""
+		for _, trace := range leaks {
+			allTraces += trace
+		}
+		logger.Errorf("Found %d leaked async reporters:%s", len(leaks), allTraces)
+	}
+
+	// Clean up global state
+	asyncReporterTracker = nil
+}

internal/leakcheck/leakcheck_enabled.go

@@ -0,0 +1,25 @@
+//go:build checkbuffers
+
+/*
+ *
+ * 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 leakcheck
+
+func init() {
+	failTestsOnLeakedBuffers = true
+}

internal/leakcheck/leakcheck_test.go

@@ -0,0 +1,173 @@
+/*
+ *
+ * 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 leakcheck
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal"
+)
+
+type testLogger struct {
+	errorCount int
+	errors     []string
+}
+
+func (e *testLogger) Logf(string, ...any) {
+}
+
+func (e *testLogger) Errorf(format string, args ...any) {
+	e.errors = append(e.errors, fmt.Sprintf(format, args...))
+	e.errorCount++
+}
+
+func TestCheck(t *testing.T) {
+	const leakCount = 3
+	ch := make(chan struct{})
+	for i := 0; i < leakCount; i++ {
+		go func() { <-ch }()
+	}
+	if leaked := interestingGoroutines(); len(leaked) != leakCount {
+		t.Errorf("interestingGoroutines() = %d, want length %d", len(leaked), leakCount)
+	}
+	e := &testLogger{}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if CheckGoroutines(ctx, e); e.errorCount < leakCount {
+		t.Errorf("CheckGoroutines() = %d, want count %d", e.errorCount, leakCount)
+		t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
+	}
+	close(ch)
+	ctx, cancel = context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	CheckGoroutines(ctx, t)
+}
+
+func ignoredTestingLeak(d time.Duration) {
+	time.Sleep(d)
+}
+
+func TestCheckRegisterIgnore(t *testing.T) {
+	RegisterIgnoreGoroutine("ignoredTestingLeak")
+	go ignoredTestingLeak(3 * time.Second)
+	const leakCount = 3
+	ch := make(chan struct{})
+	for i := 0; i < leakCount; i++ {
+		go func() { <-ch }()
+	}
+	if leaked := interestingGoroutines(); len(leaked) != leakCount {
+		t.Errorf("interestingGoroutines() = %d, want length %d", len(leaked), leakCount)
+	}
+	e := &testLogger{}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if CheckGoroutines(ctx, e); e.errorCount < leakCount {
+		t.Errorf("CheckGoroutines() = %d, want count %d", e.errorCount, leakCount)
+		t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
+	}
+	close(ch)
+	ctx, cancel = context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	CheckGoroutines(ctx, t)
+}
+
+// TestTrackTimers verifies that only leaked timers are reported and expired,
+// stopped timers are ignored.
+func TestTrackTimers(t *testing.T) {
+	TrackTimers()
+	const leakCount = 3
+	for i := 0; i < leakCount; i++ {
+		internal.TimeAfterFunc(2*time.Second, func() {
+			t.Logf("Timer %d fired.", i)
+		})
+	}
+	wg := sync.WaitGroup{}
+	// Let a couple of timers expire.
+	for i := 0; i < 2; i++ {
+		wg.Add(1)
+		internal.TimeAfterFunc(time.Millisecond, func() {
+			wg.Done()
+		})
+	}
+	wg.Wait()
+
+	// Stop a couple of timers.
+	for i := 0; i < leakCount; i++ {
+		t := internal.TimeAfterFunc(time.Hour, func() {
+			t.Error("Timer fired before test ended.")
+		})
+		t.Stop()
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	e := &testLogger{}
+	CheckTimers(ctx, e)
+	if e.errorCount != leakCount {
+		t.Errorf("CheckTimers found %v leaks, want %v leaks", e.errorCount, leakCount)
+		t.Logf("leaked timers:\n%v", strings.Join(e.errors, "\n"))
+	}
+	ctx, cancel = context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	CheckTimers(ctx, t)
+}
+
+func TestLeakChecker_DetectsLeak(t *testing.T) {
+	TrackAsyncReporters()
+	// Safety defer: ensure we restore the default delegate even if the test crashes
+	// before CheckAsyncReporters is called.
+	defer func() {
+		internal.AsyncReporterCleanupDelegate = func(f func()) func() { return f }
+	}()
+
+	// We utilize the internal delegate directly to simulate stats.RegisterAsyncReporter behavior.
+	noOpCleanup := func() {}
+	wrappedCleanup := internal.AsyncReporterCleanupDelegate(noOpCleanup)
+
+	// Create a leak: We discard 'wrappedCleanup' without calling it.
+	_ = wrappedCleanup
+	tl := &testLogger{}
+	CheckAsyncReporters(tl)
+	if tl.errorCount == 0 {
+		t.Error("Expected leak checker to report a leak, but it succeeded silently.")
+	}
+	if asyncReporterTracker != nil {
+		t.Error("Expected CheckAsyncReporters to cleanup global tracker, but it was not nil.")
+	}
+}
+
+func TestLeakChecker_PassesOnCleanup(t *testing.T) {
+	TrackAsyncReporters()
+	defer func() {
+		internal.AsyncReporterCleanupDelegate = func(f func()) func() { return f }
+	}()
+	noOpCleanup := func() {}
+	wrappedCleanup := internal.AsyncReporterCleanupDelegate(noOpCleanup)
+	wrappedCleanup()
+	tl := &testLogger{}
+	CheckAsyncReporters(tl)
+	if tl.errorCount > 0 {
+		t.Errorf("Expected no leaks, but got errors: %v", tl.errors)
+	}
+}

internal/mem/buffer_pool.go

@@ -0,0 +1,363 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package mem provides utilities that facilitate memory reuse in byte slices
+// that are used as buffers.
+package mem
+
+import (
+	"fmt"
+	"math/bits"
+	"slices"
+	"sort"
+	"sync"
+
+	"google.golang.org/grpc/internal"
+)
+
+const (
+	goPageSize = 4 * 1024 // 4KiB. N.B. this must be a power of 2.
+)
+
+var (
+	// BufferPoolingThreshold is the minimum size of a buffer that can be pooled.
+	// This is used to determine whether to pool buffers or allocate them directly.
+	BufferPoolingThreshold = 1 << 10
+)
+
+func init() {
+	internal.SetBufferPoolingThresholdForTesting = func(threshold int) {
+		BufferPoolingThreshold = threshold
+	}
+}
+
+var uintSize = bits.UintSize // use a variable for mocking during tests.
+
+// bufferPool is a copy of the public bufferPool interface used to avoid
+// circular dependencies.
+type bufferPool interface {
+	// Get returns a buffer with specified length from the pool.
+	Get(length int) *[]byte
+
+	// Put returns a buffer to the pool.
+	//
+	// The provided pointer must hold a prefix of the buffer obtained via
+	// BufferPool.Get to ensure the buffer's entire capacity can be re-used.
+	Put(*[]byte)
+}
+
+// BinaryTieredBufferPool is a buffer pool that uses multiple sub-pools with
+// power-of-two sizes.
+type BinaryTieredBufferPool struct {
+	// exponentToNextLargestPoolMap maps a power-of-two exponent (e.g., 12 for
+	// 4KB) to the index of the next largest sizedBufferPool. This is used by
+	// Get() to find the smallest pool that can satisfy a request for a given
+	// size.
+	exponentToNextLargestPoolMap []int
+	// exponentToPreviousLargestPoolMap maps a power-of-two exponent to the
+	// index of the previous largest sizedBufferPool. This is used by Put()
+	// to return a buffer to the most appropriate pool based on its capacity.
+	exponentToPreviousLargestPoolMap []int
+	sizedPools                       []bufferPool
+	fallbackPool                     bufferPool
+	maxPoolCap                       int // Optimization: Cache max capacity
+}
+
+// NewBinaryTieredBufferPool returns a BufferPool backed by multiple sub-pools.
+// This structure enables O(1) lookup time for Get and Put operations.
+//
+// The arguments provided are the exponents for the buffer capacities (powers
+// of 2), not the raw byte sizes. For example, to create a pool of 16KB buffers
+// (2^14 bytes), pass 14 as the argument.
+func NewBinaryTieredBufferPool(powerOfTwoExponents ...uint8) (*BinaryTieredBufferPool, error) {
+	return newBinaryTiered(func(size int) bufferPool {
+		return newSizedBufferPool(size, true)
+	}, &SimpleBufferPool{shouldZero: true}, powerOfTwoExponents...)
+}
+
+// NewDirtyBinaryTieredBufferPool returns a BufferPool backed by multiple
+// sub-pools. It is similar to NewBinaryTieredBufferPool but it does not
+// initialize the buffers before returning them.
+func NewDirtyBinaryTieredBufferPool(powerOfTwoExponents ...uint8) (*BinaryTieredBufferPool, error) {
+	return newBinaryTiered(func(size int) bufferPool {
+		return newSizedBufferPool(size, false)
+	}, NewDirtySimplePool(), powerOfTwoExponents...)
+}
+
+func newBinaryTiered(sizedPoolFactory func(int) bufferPool, fallbackPool bufferPool, powerOfTwoExponents ...uint8) (*BinaryTieredBufferPool, error) {
+	slices.Sort(powerOfTwoExponents)
+	powerOfTwoExponents = slices.Compact(powerOfTwoExponents)
+
+	// Determine the maximum exponent we need to support. This depends on the
+	// word size (32-bit vs 64-bit).
+	maxExponent := uintSize - 2
+	indexOfNextLargestBit := slices.Repeat([]int{-1}, maxExponent+1)
+	indexOfPreviousLargestBit := slices.Repeat([]int{-1}, maxExponent+1)
+
+	maxTier := 0
+	pools := make([]bufferPool, 0, len(powerOfTwoExponents))
+
+	for i, exp := range powerOfTwoExponents {
+		// Allocating slices of size > 2^maxExponent isn't possible on
+		// maxExponent-bit machines.
+		if int(exp) > maxExponent {
+			return nil, fmt.Errorf("mem: allocating slice of size 2^%d is not possible", exp)
+		}
+		tierSize := 1 << exp
+		pools = append(pools, sizedPoolFactory(tierSize))
+		maxTier = max(maxTier, tierSize)
+
+		// Map the exact power of 2 to this pool index.
+		indexOfNextLargestBit[exp] = i
+		indexOfPreviousLargestBit[exp] = i
+	}
+
+	// Fill gaps for Get() (Next Largest)
+	// We iterate backwards. If current is empty, take the value from the right (larger).
+	for i := maxExponent - 1; i >= 0; i-- {
+		if indexOfNextLargestBit[i] == -1 {
+			indexOfNextLargestBit[i] = indexOfNextLargestBit[i+1]
+		}
+	}
+
+	// Fill gaps for Put() (Previous Largest)
+	// We iterate forwards. If current is empty, take the value from the left (smaller).
+	for i := 1; i <= maxExponent; i++ {
+		if indexOfPreviousLargestBit[i] == -1 {
+			indexOfPreviousLargestBit[i] = indexOfPreviousLargestBit[i-1]
+		}
+	}
+
+	return &BinaryTieredBufferPool{
+		exponentToNextLargestPoolMap:     indexOfNextLargestBit,
+		exponentToPreviousLargestPoolMap: indexOfPreviousLargestBit,
+		sizedPools:                       pools,
+		maxPoolCap:                       maxTier,
+		fallbackPool:                     fallbackPool,
+	}, nil
+}
+
+// Get returns a buffer with specified length from the pool.
+func (b *BinaryTieredBufferPool) Get(size int) *[]byte {
+	return b.poolForGet(size).Get(size)
+}
+
+func (b *BinaryTieredBufferPool) poolForGet(size int) bufferPool {
+	if size == 0 || size > b.maxPoolCap {
+		return b.fallbackPool
+	}
+
+	// Calculate the exponent of the smallest power of 2 >= size.
+	// We subtract 1 from size to handle exact powers of 2 correctly.
+	//
+	// Examples:
+	// size=16 (0b10000) -> size-1=15 (0b01111) -> bits.Len=4 -> Pool for 2^4
+	// size=17 (0b10001) -> size-1=16 (0b10000) -> bits.Len=5 -> Pool for 2^5
+	querySize := uint(size - 1)
+	poolIdx := b.exponentToNextLargestPoolMap[bits.Len(querySize)]
+
+	return b.sizedPools[poolIdx]
+}
+
+// Put returns a buffer to the pool.
+func (b *BinaryTieredBufferPool) Put(buf *[]byte) {
+	// We pass the capacity of the buffer, and not the size of the buffer here.
+	// If we did the latter, all buffers would eventually move to the smallest
+	// pool.
+	b.poolForPut(cap(*buf)).Put(buf)
+}
+
+func (b *BinaryTieredBufferPool) poolForPut(bCap int) bufferPool {
+	if bCap == 0 {
+		return NopBufferPool{}
+	}
+	if bCap > b.maxPoolCap {
+		return b.fallbackPool
+	}
+	// Find the pool with the largest capacity <= bCap.
+	//
+	// We calculate the exponent of the largest power of 2 <= bCap.
+	// bits.Len(x) returns the minimum number of bits required to represent x;
+	// i.e. the number of bits up to and including the most significant bit.
+	// Subtracting 1 gives the 0-based index of the most significant bit,
+	// which is the exponent of the largest power of 2 <= bCap.
+	//
+	// Examples:
+	// cap=16 (0b10000) -> Len=5 -> 5-1=4 -> 2^4
+	// cap=15 (0b01111) -> Len=4 -> 4-1=3 -> 2^3
+	largestPowerOfTwo := bits.Len(uint(bCap)) - 1
+	poolIdx := b.exponentToPreviousLargestPoolMap[largestPowerOfTwo]
+	// The buffer is smaller than the smallest power of 2, discard it.
+	if poolIdx == -1 {
+		// Buffer is smaller than our smallest pool bucket.
+		return NopBufferPool{}
+	}
+	return b.sizedPools[poolIdx]
+}
+
+// NopBufferPool is a buffer pool that returns new buffers without pooling.
+type NopBufferPool struct{}
+
+// Get returns a buffer with specified length from the pool.
+func (NopBufferPool) Get(length int) *[]byte {
+	b := make([]byte, length)
+	return &b
+}
+
+// Put returns a buffer to the pool.
+func (NopBufferPool) Put(*[]byte) {
+}
+
+// sizedBufferPool is a BufferPool implementation that is optimized for specific
+// buffer sizes. For example, HTTP/2 frames within gRPC have a default max size
+// of 16kb and a sizedBufferPool can be configured to only return buffers with a
+// capacity of 16kb. Note that however it does not support returning larger
+// buffers and in fact panics if such a buffer is requested. Because of this,
+// this BufferPool implementation is not meant to be used on its own and rather
+// is intended to be embedded in a TieredBufferPool such that Get is only
+// invoked when the required size is smaller than or equal to defaultSize.
+type sizedBufferPool struct {
+	pool        sync.Pool
+	defaultSize int
+	shouldZero  bool
+}
+
+func (p *sizedBufferPool) Get(size int) *[]byte {
+	buf, ok := p.pool.Get().(*[]byte)
+	if !ok {
+		buf := make([]byte, size, p.defaultSize)
+		return &buf
+	}
+	b := *buf
+	if p.shouldZero {
+		clear(b[:cap(b)])
+	}
+	*buf = b[:size]
+	return buf
+}
+
+func (p *sizedBufferPool) Put(buf *[]byte) {
+	if cap(*buf) < p.defaultSize {
+		// Ignore buffers that are too small to fit in the pool. Otherwise, when
+		// Get is called it will panic as it tries to index outside the bounds
+		// of the buffer.
+		return
+	}
+	p.pool.Put(buf)
+}
+
+func newSizedBufferPool(size int, zero bool) *sizedBufferPool {
+	return &sizedBufferPool{
+		defaultSize: size,
+		shouldZero:  zero,
+	}
+}
+
+// TieredBufferPool implements the BufferPool interface with multiple tiers of
+// buffer pools for different sizes of buffers.
+type TieredBufferPool struct {
+	sizedPools   []*sizedBufferPool
+	fallbackPool SimpleBufferPool
+}
+
+// NewTieredBufferPool returns a BufferPool implementation that uses multiple
+// underlying pools of the given pool sizes.
+func NewTieredBufferPool(poolSizes ...int) *TieredBufferPool {
+	sort.Ints(poolSizes)
+	pools := make([]*sizedBufferPool, len(poolSizes))
+	for i, s := range poolSizes {
+		pools[i] = newSizedBufferPool(s, true)
+	}
+	return &TieredBufferPool{
+		sizedPools:   pools,
+		fallbackPool: SimpleBufferPool{shouldZero: true},
+	}
+}
+
+// Get returns a buffer with specified length from the pool.
+func (p *TieredBufferPool) Get(size int) *[]byte {
+	return p.getPool(size).Get(size)
+}
+
+// Put returns a buffer to the pool.
+func (p *TieredBufferPool) Put(buf *[]byte) {
+	p.getPool(cap(*buf)).Put(buf)
+}
+
+func (p *TieredBufferPool) getPool(size int) bufferPool {
+	poolIdx := sort.Search(len(p.sizedPools), func(i int) bool {
+		return p.sizedPools[i].defaultSize >= size
+	})
+
+	if poolIdx == len(p.sizedPools) {
+		return &p.fallbackPool
+	}
+
+	return p.sizedPools[poolIdx]
+}
+
+// SimpleBufferPool is an implementation of the mem.BufferPool interface that
+// attempts to pool buffers with a sync.Pool. When Get is invoked, it tries to
+// acquire a buffer from the pool but if that buffer is too small, it returns it
+// to the pool and creates a new one.
+type SimpleBufferPool struct {
+	pool       sync.Pool
+	shouldZero bool
+}
+
+// NewDirtySimplePool constructs a [SimpleBufferPool]. It does not initialize
+// the buffers before returning them. Callers must ensure they don't read the
+// buffers before writing data to them.
+func NewDirtySimplePool() *SimpleBufferPool {
+	return &SimpleBufferPool{
+		shouldZero: false,
+	}
+}
+
+// Get returns a buffer with specified length from the pool.
+func (p *SimpleBufferPool) Get(size int) *[]byte {
+	bs, ok := p.pool.Get().(*[]byte)
+	if ok && cap(*bs) >= size {
+		if p.shouldZero {
+			clear((*bs)[:cap(*bs)])
+		}
+		*bs = (*bs)[:size]
+		return bs
+	}
+
+	// A buffer was pulled from the pool, but it is too small. Put it back in
+	// the pool and create one large enough.
+	if ok {
+		p.pool.Put(bs)
+	}
+
+	// If we're going to allocate, round up to the nearest page. This way if
+	// requests frequently arrive with small variation we don't allocate
+	// repeatedly if we get unlucky and they increase over time. By default we
+	// only allocate here if size > 1MiB. Because goPageSize is a power of 2, we
+	// can round up efficiently.
+	allocSize := (size + goPageSize - 1) & ^(goPageSize - 1)
+
+	b := make([]byte, size, allocSize)
+	return &b
+}
+
+// Put returns a buffer to the pool.
+func (p *SimpleBufferPool) Put(buf *[]byte) {
+	p.pool.Put(buf)
+}

internal/mem/buffer_pool_ext_test.go

@@ -0,0 +1,118 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package mem_test
+
+import (
+	"testing"
+	"unsafe"
+
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/mem"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestBufferPool_Clears(t *testing.T) {
+	poolConfigs := []struct {
+		name        string
+		factory     func() (*mem.BinaryTieredBufferPool, error)
+		wantCleared bool
+		bufferSize  int
+	}{
+		{
+			name: "regular_sized",
+			factory: func() (*mem.BinaryTieredBufferPool, error) {
+				return mem.NewBinaryTieredBufferPool(3) // 8 bytes
+			},
+			bufferSize:  8,
+			wantCleared: true,
+		},
+		{
+			name: "regular_fallback",
+			factory: func() (*mem.BinaryTieredBufferPool, error) {
+				return mem.NewBinaryTieredBufferPool(3)
+			},
+			bufferSize:  10,
+			wantCleared: true,
+		},
+		{
+			name: "dirty_sized",
+			factory: func() (*mem.BinaryTieredBufferPool, error) {
+				return mem.NewDirtyBinaryTieredBufferPool(3)
+			},
+			bufferSize:  8,
+			wantCleared: false,
+		},
+		{
+			name: "dirty_fallback",
+			factory: func() (*mem.BinaryTieredBufferPool, error) {
+				return mem.NewDirtyBinaryTieredBufferPool(3)
+			},
+			bufferSize:  10,
+			wantCleared: false,
+		},
+	}
+
+	for _, tc := range poolConfigs {
+		t.Run(tc.name, func(t *testing.T) {
+			pool, err := tc.factory()
+			if err != nil {
+				t.Fatalf("Failed to create pool: %v", err)
+			}
+
+			for {
+				buf1 := pool.Get(tc.bufferSize)
+				// Mark the buffer with data.
+				for i := range *buf1 {
+					(*buf1)[i] = 0xAA
+				}
+				pool.Put(buf1)
+
+				buf2 := pool.Get(tc.bufferSize)
+				// Check if we got the same underlying array.
+				if unsafe.SliceData(*buf1) != unsafe.SliceData(*buf2) {
+					pool.Put(buf2)
+					continue
+				}
+
+				// We have a reused buffer. Check if it's cleared.
+				gotCleared := true
+				for _, b := range *buf2 {
+					if b != 0 {
+						gotCleared = false
+						break
+					}
+				}
+
+				if tc.wantCleared != gotCleared {
+					t.Fatalf("buffer cleared state mismatch: want %t, got %v", tc.wantCleared, gotCleared)
+				}
+
+				pool.Put(buf2)
+				break
+			}
+		})
+	}
+}

internal/mem/buffer_pool_test.go

@@ -0,0 +1,130 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package mem
+
+import "testing"
+
+var defaultBufferPoolSizeExponents = []uint8{
+	8,
+	12,
+	14, // 16KB (max HTTP/2 frame size used by gRPC)
+	15, // 32KB (default buffer size for io.Copy)
+	20, // 1MB
+}
+
+func TestNewBinaryTieredBufferPool_WordSize(t *testing.T) {
+	origUintSize := uintSize
+	defer func() { uintSize = origUintSize }()
+
+	tests := []struct {
+		name      string
+		wordSize  int
+		exponents []uint8
+		wantErr   bool
+	}{
+		{
+			name:      "32-bit_valid_exponent",
+			wordSize:  32,
+			exponents: []uint8{30},
+			wantErr:   false,
+		},
+		{
+			name:      "32-bit_invalid_exponent",
+			wordSize:  32,
+			exponents: []uint8{31},
+			wantErr:   true,
+		},
+		{
+			name:      "64-bit_valid_exponent",
+			wordSize:  64,
+			exponents: []uint8{62},
+			wantErr:   false,
+		},
+		{
+			name:      "64-bit_invalid_exponent",
+			wordSize:  64,
+			exponents: []uint8{63},
+			wantErr:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			uintSize = tt.wordSize
+			pool, err := NewBinaryTieredBufferPool(tt.exponents...)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("NewBinaryTieredBufferPool() error = %t, wantErr %t", err, tt.wantErr)
+			}
+			if err != nil {
+				return
+			}
+			if len(pool.exponentToNextLargestPoolMap) != tt.wordSize-1 {
+				t.Errorf("exponentToNextLargestPoolMap length = %d, want %d", len(pool.exponentToNextLargestPoolMap), tt.wordSize)
+			}
+			if len(pool.exponentToPreviousLargestPoolMap) != tt.wordSize-1 {
+				t.Errorf("exponentToPreviousLargestPoolMap length = %d, want %d", len(pool.exponentToPreviousLargestPoolMap), tt.wordSize)
+			}
+		})
+	}
+}
+
+// BenchmarkTieredPool benchmarks the performance of the tiered buffer pool
+// implementations, specifically focusing on the overhead of selecting the
+// correct bucket for a given size.
+func BenchmarkTieredPool(b *testing.B) {
+	defaultBufferPoolSizes := make([]int, len(defaultBufferPoolSizeExponents))
+	for i, exp := range defaultBufferPoolSizeExponents {
+		defaultBufferPoolSizes[i] = 1 << exp
+	}
+
+	b.Run("pool=Tiered", func(b *testing.B) {
+		p := NewTieredBufferPool(defaultBufferPoolSizes...)
+		for b.Loop() {
+			for size := range 1 << 19 {
+				// One for get, one for put.
+				_ = p.getPool(size)
+				_ = p.getPool(size)
+			}
+		}
+	})
+
+	b.Run("pool=BinaryTiered", func(b *testing.B) {
+		pool, err := NewBinaryTieredBufferPool(defaultBufferPoolSizeExponents...)
+		if err != nil {
+			b.Fatalf("Failed to create buffer pool: %v", err)
+		}
+		for b.Loop() {
+			for size := range 1 << 19 {
+				_ = pool.poolForGet(size)
+				_ = pool.poolForPut(size)
+			}
+		}
+	})
+}
+
+func TestNewBinaryTieredBufferPool_Duplicates(t *testing.T) {
+	exponents := []uint8{1, 2, 3, 4, 5, 6, 6, 5, 4, 3, 2, 1}
+	pool, err := NewBinaryTieredBufferPool(exponents...)
+	if err != nil {
+		t.Fatalf("NewBinaryTieredBufferPool() error = %v", err)
+	}
+	if len(pool.sizedPools) != 6 {
+		t.Errorf("sized buffer pool count = %d, want %d", len(pool.sizedPools), 6)
+	}
+}

internal/metadata/metadata.go

@@ -0,0 +1,144 @@
+/*
+ *
+ * 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 metadata contains functions to set and get metadata from addresses.
+//
+// This package is experimental.
+package metadata
+
+import (
+	"fmt"
+	"strings"
+
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/resolver"
+)
+
+type mdKeyType string
+
+const mdKey = mdKeyType("grpc.internal.address.metadata")
+
+type mdValue metadata.MD
+
+func (m mdValue) Equal(o any) bool {
+	om, ok := o.(mdValue)
+	if !ok {
+		return false
+	}
+	if len(m) != len(om) {
+		return false
+	}
+	for k, v := range m {
+		ov := om[k]
+		if len(ov) != len(v) {
+			return false
+		}
+		for i, ve := range v {
+			if ov[i] != ve {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+// Get returns the metadata of addr.
+func Get(addr resolver.Address) metadata.MD {
+	attrs := addr.Attributes
+	if attrs == nil {
+		return nil
+	}
+	md, _ := attrs.Value(mdKey).(mdValue)
+	return metadata.MD(md)
+}
+
+// Set sets (overrides) the metadata in addr.
+//
+// When a SubConn is created with this address, the RPCs sent on it will all
+// have this metadata.
+func Set(addr resolver.Address, md metadata.MD) resolver.Address {
+	addr.Attributes = addr.Attributes.WithValue(mdKey, mdValue(md))
+	return addr
+}
+
+// Validate validates every pair in md with ValidatePair.
+func Validate(md metadata.MD) error {
+	for k, vals := range md {
+		if err := ValidatePair(k, vals...); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// hasNotPrintable return true if msg contains any characters which are not in %x20-%x7E
+func hasNotPrintable(msg string) bool {
+	// for i that saving a conversion if not using for range
+	for i := 0; i < len(msg); i++ {
+		if msg[i] < 0x20 || msg[i] > 0x7E {
+			return true
+		}
+	}
+	return false
+}
+
+// ValidateKey validates a key with the following rules (pseudo-headers are
+// skipped):
+// - the key must contain one or more characters.
+// - the characters in the key must be in [0-9 a-z _ - .].
+func ValidateKey(key string) error {
+	// key should not be empty
+	if key == "" {
+		return fmt.Errorf("there is an empty key in the header")
+	}
+	// pseudo-header will be ignored
+	if key[0] == ':' {
+		return nil
+	}
+	// check key, for i that saving a conversion if not using for range
+	for i := 0; i < len(key); i++ {
+		r := key[i]
+		if !(r >= 'a' && r <= 'z') && !(r >= '0' && r <= '9') && r != '.' && r != '-' && r != '_' {
+			return fmt.Errorf("header key %q contains illegal characters not in [0-9a-z-_.]", key)
+		}
+	}
+	return nil
+}
+
+// ValidatePair validates a key-value pair with the following rules
+// (pseudo-header are skipped):
+//   - the key must contain one or more characters.
+//   - the characters in the key must be in [0-9 a-z _ - .].
+//   - if the key ends with a "-bin" suffix, no validation of the corresponding
+//     value is performed.
+//   - the characters in every value must be printable (in [%x20-%x7E]).
+func ValidatePair(key string, vals ...string) error {
+	if err := ValidateKey(key); err != nil {
+		return err
+	}
+	if strings.HasSuffix(key, "-bin") {
+		return nil
+	}
+	// check value
+	for _, val := range vals {
+		if hasNotPrintable(val) {
+			return fmt.Errorf("header key %q contains value with non-printable ASCII characters", key)
+		}
+	}
+	return nil
+}

internal/metadata/metadata_test.go

@@ -0,0 +1,117 @@
+/*
+ *
+ * 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 metadata
+
+import (
+	"errors"
+	"reflect"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/resolver"
+)
+
+func TestGet(t *testing.T) {
+	tests := []struct {
+		name string
+		addr resolver.Address
+		want metadata.MD
+	}{
+		{
+			name: "not set",
+			addr: resolver.Address{},
+			want: nil,
+		},
+		{
+			name: "not set",
+			addr: resolver.Address{
+				Attributes: attributes.New(mdKey, mdValue(metadata.Pairs("k", "v"))),
+			},
+			want: metadata.Pairs("k", "v"),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := Get(tt.addr); !cmp.Equal(got, tt.want) {
+				t.Errorf("Get() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestSet(t *testing.T) {
+	tests := []struct {
+		name string
+		addr resolver.Address
+		md   metadata.MD
+	}{
+		{
+			name: "unset before",
+			addr: resolver.Address{},
+			md:   metadata.Pairs("k", "v"),
+		},
+		{
+			name: "set before",
+			addr: resolver.Address{
+				Attributes: attributes.New(mdKey, mdValue(metadata.Pairs("bef", "ore"))),
+			},
+			md: metadata.Pairs("k", "v"),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			newAddr := Set(tt.addr, tt.md)
+			newMD := Get(newAddr)
+			if !cmp.Equal(newMD, tt.md) {
+				t.Errorf("md after Set() = %v, want %v", newMD, tt.md)
+			}
+		})
+	}
+}
+
+func TestValidate(t *testing.T) {
+	for _, test := range []struct {
+		md   metadata.MD
+		want error
+	}{
+		{
+			md:   map[string][]string{string(rune(0x19)): {"testVal"}},
+			want: errors.New("header key \"\\x19\" contains illegal characters not in [0-9a-z-_.]"),
+		},
+		{
+			md:   map[string][]string{"test": {string(rune(0x19))}},
+			want: errors.New("header key \"test\" contains value with non-printable ASCII characters"),
+		},
+		{
+			md:   map[string][]string{"": {"valid"}},
+			want: errors.New("there is an empty key in the header"),
+		},
+		{
+			md:   map[string][]string{"test-bin": {string(rune(0x19))}},
+			want: nil,
+		},
+	} {
+		err := Validate(test.md)
+		if !reflect.DeepEqual(err, test.want) {
+			t.Errorf("validating metadata which is %v got err :%v, want err :%v", test.md, err, test.want)
+		}
+	}
+}

internal/optional/optional.go

@@ -0,0 +1,47 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package optional implements a generic optional type.
+package optional
+
+// Optional represents an optional value of type T. The zero value is usable and
+// indicates that no value is set. This type is not safe for concurrent access.
+type Optional[T any] struct {
+	val T
+	set bool
+}
+
+// New creates a new Optional type with the provided value.
+func New[T any](value T) Optional[T] {
+	return Optional[T]{
+		val: value,
+		set: true,
+	}
+}
+
+// Value returns the underlying value and a boolean indicating if the value is
+// set. If the value is not set, it returns the zero value of T and false.
+func (o *Optional[T]) Value() (T, bool) {
+	return o.val, o.set
+}
+
+// SetValue updates or adds the value to Optional.
+func (o *Optional[T]) SetValue(v T) {
+	o.val = v
+	o.set = true
+}

internal/optional/optional_test.go

@@ -0,0 +1,130 @@
+/*
+ *
+ * Copyright 2026 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package optional_test
+
+import (
+	"slices"
+	"testing"
+
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/optional"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// TestOptionalInt tests the scenario of using integer optional values and
+// verifies that zero value, constructors, and mutation methods return correct
+// outputs.
+func (s) TestOptionalInt(t *testing.T) {
+	var opt optional.Optional[int]
+
+	// Test unset value.
+	if v, set := opt.Value(); set || v != 0 {
+		t.Fatalf("Zero-value Option[int] = (%v, %v); want (0, false)", v, set)
+	}
+
+	opt = optional.New(42)
+	if v, set := opt.Value(); !set || v != 42 {
+		t.Fatalf("NewValue(42) = (%v, %v); want (42, true)", v, set)
+	}
+
+	opt.SetValue(100)
+	if v, set := opt.Value(); !set || v != 100 {
+		t.Fatalf("Set(100) = (%v, %v); want (100, true)", v, set)
+	}
+}
+
+// TestOptionalString tests the scenario of using string optional values and
+// verifies that zero value, constructors, and mutation methods return correct
+// outputs.
+func (s) TestOptionalString(t *testing.T) {
+	var opt optional.Optional[string]
+
+	// Test unset value.
+	if v, set := opt.Value(); set || v != "" {
+		t.Fatalf("Zero-value Option[string] = (%q, %v); want (%q, false)", v, set, "")
+	}
+
+	wantString := "test-string"
+	opt = optional.New(wantString)
+	if v, set := opt.Value(); !set || v != wantString {
+		t.Fatalf("NewValue(%q) = (%q, %v); want (%q, true)", wantString, v, set, wantString)
+	}
+
+	wantStringNew := "world"
+	opt.SetValue(wantStringNew)
+	if v, set := opt.Value(); !set || v != wantStringNew {
+		t.Fatalf("Set(%q) = (%q, %v); want (%q, true)", wantStringNew, v, set, wantStringNew)
+	}
+}
+
+// TestOptionalStruct tests the scenario of using a custom struct optional value
+// and verifies that custom struct field values are preserved, modified, and
+// cleared correctly.
+func (s) TestOptionalStruct(t *testing.T) {
+	type testStruct struct {
+		name string
+		age  int
+	}
+
+	var opt optional.Optional[testStruct]
+	if v, set := opt.Value(); set || v != (testStruct{}) {
+		t.Fatalf("Zero-value Option[struct] = (%v, %v); want (empty, false)", v, set)
+	}
+
+	want := testStruct{name: "Alice", age: 30}
+	opt = optional.New(want)
+	if v, set := opt.Value(); !set || v != want {
+		t.Fatalf("NewValue(%v) = (%v, %v); want (%v, true)", want, v, set, want)
+	}
+
+	want2 := testStruct{name: "Bob", age: 40}
+	opt.SetValue(want2)
+	if v, set := opt.Value(); !set || v != want2 {
+		t.Fatalf("Set(%v) = (%v, %v); want (%v, true)", want2, v, set, want2)
+	}
+}
+
+// TestOptionalSlice tests the scenario of using a slice optional value and
+// verifies that zero value, constructors, and mutation methods return correct
+// outputs.
+func (s) TestOptionalSlice(t *testing.T) {
+	var opt optional.Optional[[]int]
+	if v, set := opt.Value(); set || v != nil {
+		t.Fatalf("Zero-value Option[[]int] = (%v, %v); want (nil, false)", v, set)
+	}
+
+	want := []int{1, 2, 3}
+	opt = optional.New(want)
+	if v, set := opt.Value(); !set || !slices.Equal(v, want) {
+		t.Fatalf("NewValue(%v) = (%v, %v); want (%v, true)", want, v, set, want)
+	}
+
+	want2 := []int{4, 5, 6}
+	opt.SetValue(want2)
+	if v, set := opt.Value(); !set || !slices.Equal(v, want2) {
+		t.Fatalf("Set(%v) = (%v, %v); want (%v, true)", &want2, v, set, &want2)
+	}
+}

internal/pretty/pretty.go

@@ -0,0 +1,73 @@
+/*
+ *
+ * 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 pretty defines helper functions to pretty-print structs for logging.
+package pretty
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/protoadapt"
+)
+
+const jsonIndent = "  "
+
+// ToJSON marshals the input into a json string.
+//
+// If marshal fails, it falls back to fmt.Sprintf("%+v").
+func ToJSON(e any) string {
+	if ee, ok := e.(protoadapt.MessageV1); ok {
+		e = protoadapt.MessageV2Of(ee)
+	}
+
+	if ee, ok := e.(protoadapt.MessageV2); ok {
+		mm := protojson.MarshalOptions{
+			Indent:    jsonIndent,
+			Multiline: true,
+		}
+		ret, err := mm.Marshal(ee)
+		if err != nil {
+			// This may fail for proto.Anys, e.g. for xDS v2, LDS, the v2
+			// messages are not imported, and this will fail because the message
+			// is not found.
+			return fmt.Sprintf("%+v", ee)
+		}
+		return string(ret)
+	}
+
+	ret, err := json.MarshalIndent(e, "", jsonIndent)
+	if err != nil {
+		return fmt.Sprintf("%+v", e)
+	}
+	return string(ret)
+}
+
+// FormatJSON formats the input json bytes with indentation.
+//
+// If Indent fails, it returns the unchanged input as string.
+func FormatJSON(b []byte) string {
+	var out bytes.Buffer
+	err := json.Indent(&out, b, "", jsonIndent)
+	if err != nil {
+		return string(b)
+	}
+	return out.String()
+}

internal/profiling/buffer/buffer.go

@@ -0,0 +1,272 @@
+/*
+ *
+ * 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 buffer provides a high-performant lock free implementation of a
+// circular buffer used by the profiling code.
+package buffer
+
+import (
+	"errors"
+	"math/bits"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"unsafe"
+)
+
+type queue struct {
+	// An array of pointers as references to the items stored in this queue.
+	arr []unsafe.Pointer
+	// The maximum number of elements this queue may store before it wraps around
+	// and overwrites older values. Must be an exponent of 2.
+	size uint32
+	// Always size - 1. A bitwise AND is performed with this mask in place of a
+	// modulo operation by the Push operation.
+	mask uint32
+	// Each Push operation into this queue increments the acquired counter before
+	// proceeding forwarding with the actual write to arr. This counter is also
+	// used by the Drain operation's drainWait subroutine to wait for all pushes
+	// to complete.
+	acquired uint32 // Accessed atomically.
+	// After the completion of a Push operation, the written counter is
+	// incremented. Also used by drainWait to wait for all pushes to complete.
+	written uint32
+}
+
+// Allocates and returns a new *queue. size needs to be an exponent of two.
+func newQueue(size uint32) *queue {
+	return &queue{
+		arr:  make([]unsafe.Pointer, size),
+		size: size,
+		mask: size - 1,
+	}
+}
+
+// drainWait blocks the caller until all Pushes on this queue are complete.
+func (q *queue) drainWait() {
+	for atomic.LoadUint32(&q.acquired) != atomic.LoadUint32(&q.written) {
+		runtime.Gosched()
+	}
+}
+
+// A queuePair has two queues. At any given time, Pushes go into the queue
+// referenced by queuePair.q. The active queue gets switched when there's a
+// drain operation on the circular buffer.
+type queuePair struct {
+	q0 unsafe.Pointer
+	q1 unsafe.Pointer
+	q  unsafe.Pointer
+}
+
+// Allocates and returns a new *queuePair with its internal queues allocated.
+func newQueuePair(size uint32) *queuePair {
+	qp := &queuePair{}
+	qp.q0 = unsafe.Pointer(newQueue(size))
+	qp.q1 = unsafe.Pointer(newQueue(size))
+	qp.q = qp.q0
+	return qp
+}
+
+// Switches the current queue for future Pushes to proceed to the other queue
+// so that there's no blocking in Push. Returns a pointer to the old queue that
+// was in place before the switch.
+func (qp *queuePair) switchQueues() *queue {
+	// Even though we have mutual exclusion across drainers (thanks to mu.Lock in
+	// drain), Push operations may access qp.q whilst we're writing to it.
+	if atomic.CompareAndSwapPointer(&qp.q, qp.q0, qp.q1) {
+		return (*queue)(qp.q0)
+	}
+
+	atomic.CompareAndSwapPointer(&qp.q, qp.q1, qp.q0)
+	return (*queue)(qp.q1)
+}
+
+// In order to not have expensive modulo operations, we require the maximum
+// number of elements in the circular buffer (N) to be an exponent of two to
+// use a bitwise AND mask. Since a CircularBuffer is a collection of queuePairs
+// (see below), we need to divide N; since exponents of two are only divisible
+// by other exponents of two, we use floorCPUCount number of queuePairs within
+// each CircularBuffer.
+//
+// Floor of the number of CPUs (and not the ceiling) was found to be the
+// optimal number through experiments.
+func floorCPUCount() uint32 {
+	floorExponent := bits.Len32(uint32(runtime.NumCPU())) - 1
+	if floorExponent < 0 {
+		floorExponent = 0
+	}
+	return 1 << uint32(floorExponent)
+}
+
+var numCircularBufferPairs = floorCPUCount()
+
+// CircularBuffer is a lock-free data structure that supports Push and Drain
+// operations.
+//
+// Note that CircularBuffer is built for performance more than reliability.
+// That is, some Push operations may fail without retries in some situations
+// (such as during a Drain operation). Order of pushes is not maintained
+// either; that is, if A was pushed before B, the Drain operation may return an
+// array with B before A. These restrictions are acceptable within gRPC's
+// profiling, but if your use-case does not permit these relaxed constraints
+// or if performance is not a primary concern, you should probably use a
+// lock-based data structure such as internal/buffer.UnboundedBuffer.
+type CircularBuffer struct {
+	drainMutex sync.Mutex
+	qp         []*queuePair
+	// qpn is a monotonically incrementing counter that's used to determine
+	// which queuePair a Push operation should write to. This approach's
+	// performance was found to be better than writing to a random queue.
+	qpn    uint32
+	qpMask uint32
+}
+
+var errInvalidCircularBufferSize = errors.New("buffer size is not an exponent of two")
+
+// NewCircularBuffer allocates a circular buffer of size size and returns a
+// reference to the struct. Only circular buffers of size 2^k are allowed
+// (saves us from having to do expensive modulo operations).
+func NewCircularBuffer(size uint32) (*CircularBuffer, error) {
+	if size&(size-1) != 0 {
+		return nil, errInvalidCircularBufferSize
+	}
+
+	n := numCircularBufferPairs
+	if size/numCircularBufferPairs < 8 {
+		// If each circular buffer is going to hold less than a very small number
+		// of items (let's say 8), using multiple circular buffers is very likely
+		// wasteful. Instead, fallback to one circular buffer holding everything.
+		n = 1
+	}
+
+	cb := &CircularBuffer{
+		qp:     make([]*queuePair, n),
+		qpMask: n - 1,
+	}
+
+	for i := uint32(0); i < n; i++ {
+		cb.qp[i] = newQueuePair(size / n)
+	}
+
+	return cb, nil
+}
+
+// Push pushes an element in to the circular buffer. Guaranteed to complete in
+// a finite number of steps (also lock-free). Does not guarantee that push
+// order will be retained. Does not guarantee that the operation will succeed
+// if a Drain operation concurrently begins execution.
+func (cb *CircularBuffer) Push(x any) {
+	n := atomic.AddUint32(&cb.qpn, 1) & cb.qpMask
+	qptr := atomic.LoadPointer(&cb.qp[n].q)
+	q := (*queue)(qptr)
+
+	acquired := atomic.AddUint32(&q.acquired, 1) - 1
+
+	// If true, it means that we have incremented acquired before any queuePair
+	// was switched, and therefore before any drainWait completion. Therefore, it
+	// is safe to proceed with the Push operation on this queue. Otherwise, it
+	// means that a Drain operation has begun execution, but we don't know how
+	// far along the process it is. If it is past the drainWait check, it is not
+	// safe to proceed with the Push operation. We choose to drop this sample
+	// entirely instead of retrying, as retrying may potentially send the Push
+	// operation into a spin loop (we want to guarantee completion of the Push
+	// operation within a finite time). Before exiting, we increment written so
+	// that any existing drainWaits can proceed.
+	if atomic.LoadPointer(&cb.qp[n].q) != qptr {
+		atomic.AddUint32(&q.written, 1)
+		return
+	}
+
+	// At this point, we're definitely writing to the right queue. That is, one
+	// of the following is true:
+	//   1. No drainer is in execution on this queue.
+	//   2. A drainer is in execution on this queue and it is waiting at the
+	//      acquired == written barrier.
+	//
+	// Let's say two Pushes A and B happen on the same queue. Say A and B are
+	// q.size apart; i.e. they get the same index. That is,
+	//
+	//   index_A = index_B
+	//   acquired_A + q.size = acquired_B
+	//
+	// We say "B has wrapped around A" when this happens. In this case, since A
+	// occurred before B, B's Push should be the final value. However, we
+	// accommodate A being the final value because wrap-arounds are extremely
+	// rare and accounting for them requires an additional counter and a
+	// significant performance penalty. Note that the below approach never leads
+	// to any data corruption.
+	index := acquired & q.mask
+	atomic.StorePointer(&q.arr[index], unsafe.Pointer(&x))
+
+	// Allows any drainWait checks to proceed.
+	atomic.AddUint32(&q.written, 1)
+}
+
+// Dereferences non-nil pointers from arr into result. Range of elements from
+// arr that are copied is [from, to). Assumes that the result slice is already
+// allocated and is large enough to hold all the elements that might be copied.
+// Also assumes mutual exclusion on the array of pointers.
+func dereferenceAppend(result []any, arr []unsafe.Pointer, from, to uint32) []any {
+	for i := from; i < to; i++ {
+		// We have mutual exclusion on arr, there's no need for atomics.
+		x := (*any)(arr[i])
+		if x != nil {
+			result = append(result, *x)
+		}
+	}
+	return result
+}
+
+// Drain allocates and returns an array of things Pushed in to the circular
+// buffer. Push order is not maintained; that is, if B was Pushed after A,
+// drain may return B at a lower index than A in the returned array.
+func (cb *CircularBuffer) Drain() []any {
+	cb.drainMutex.Lock()
+
+	qs := make([]*queue, len(cb.qp))
+	for i := 0; i < len(cb.qp); i++ {
+		qs[i] = cb.qp[i].switchQueues()
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(len(qs))
+	for i := 0; i < len(qs); i++ {
+		go func(qi int) {
+			qs[qi].drainWait()
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+
+	result := make([]any, 0)
+	for i := 0; i < len(qs); i++ {
+		if acquired := atomic.LoadUint32(&qs[i].acquired); acquired < qs[i].size {
+			result = dereferenceAppend(result, qs[i].arr, 0, acquired)
+		} else {
+			result = dereferenceAppend(result, qs[i].arr, 0, qs[i].size)
+		}
+	}
+
+	for i := 0; i < len(qs); i++ {
+		atomic.StoreUint32(&qs[i].acquired, 0)
+		atomic.StoreUint32(&qs[i].written, 0)
+	}
+
+	cb.drainMutex.Unlock()
+	return result
+}

internal/profiling/buffer/buffer_test.go

@@ -0,0 +1,186 @@
+/*
+ *
+ * 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 buffer
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestCircularBufferSerial(t *testing.T) {
+	var size, i uint32
+	var result []any
+
+	size = 1 << 15
+	cb, err := NewCircularBuffer(size)
+	if err != nil {
+		t.Fatalf("error allocating CircularBuffer: %v", err)
+	}
+
+	for i = 0; i < size/2; i++ {
+		cb.Push(i)
+	}
+
+	result = cb.Drain()
+	if uint32(len(result)) != size/2 {
+		t.Fatalf("len(result) = %d; want %d", len(result), size/2)
+	}
+
+	// The returned result isn't necessarily sorted.
+	seen := make(map[uint32]bool)
+	for _, r := range result {
+		seen[r.(uint32)] = true
+	}
+
+	for i = 0; i < uint32(len(result)); i++ {
+		if !seen[i] {
+			t.Fatalf("seen[%d] = false; want true", i)
+		}
+	}
+
+	for i = 0; i < size; i++ {
+		cb.Push(i)
+	}
+
+	result = cb.Drain()
+	if uint32(len(result)) != size {
+		t.Fatalf("len(result) = %d; want %d", len(result), size/2)
+	}
+}
+
+func (s) TestCircularBufferOverflow(t *testing.T) {
+	var size, i uint32
+	var result []any
+
+	size = 1 << 10
+	cb, err := NewCircularBuffer(size)
+	if err != nil {
+		t.Fatalf("error allocating CircularBuffer: %v", err)
+	}
+
+	for i = 0; i < 10*size; i++ {
+		cb.Push(i)
+	}
+
+	result = cb.Drain()
+
+	if uint32(len(result)) != size {
+		t.Fatalf("len(result) = %d; want %d", len(result), size)
+	}
+
+	for idx, x := range result {
+		if x.(uint32) < size {
+			t.Fatalf("result[%d] = %d; want it to be >= %d", idx, x, size)
+		}
+	}
+}
+
+func (s) TestCircularBufferConcurrent(t *testing.T) {
+	for tn := 0; tn < 2; tn++ {
+		var size uint32
+		var result []any
+
+		size = 1 << 6
+		cb, err := NewCircularBuffer(size)
+		if err != nil {
+			t.Fatalf("error allocating CircularBuffer: %v", err)
+		}
+
+		type item struct {
+			R uint32
+			N uint32
+			T time.Time
+		}
+
+		var wg sync.WaitGroup
+		for r := uint32(0); r < 1024; r++ {
+			wg.Add(1)
+			go func(r uint32) {
+				for n := uint32(0); n < size; n++ {
+					cb.Push(item{R: r, N: n, T: time.Now()})
+				}
+				wg.Done()
+			}(r)
+		}
+
+		// Wait for all goroutines to finish only in one test. Draining
+		// concurrently while Pushes are still happening will test for races in the
+		// Draining lock.
+		if tn == 0 {
+			wg.Wait()
+		}
+
+		result = cb.Drain()
+
+		// Can't expect the buffer to be full if the Pushes aren't necessarily done.
+		if tn == 0 {
+			if uint32(len(result)) != size {
+				t.Fatalf("len(result) = %d; want %d", len(result), size)
+			}
+		}
+
+		// There can be absolutely no expectation on the order of the data returned
+		// by Drain because: (a) everything is happening concurrently (b) a
+		// round-robin is used to write to different queues (and therefore
+		// different cachelines) for less write contention.
+
+		// Wait for all goroutines to complete before moving on to other tests. If
+		// the benchmarks run after this, it might affect performance unfairly.
+		wg.Wait()
+	}
+}
+
+func BenchmarkCircularBuffer(b *testing.B) {
+	x := 1
+	for size := 1 << 16; size <= 1<<20; size <<= 1 {
+		for routines := 1; routines <= 1<<8; routines <<= 1 {
+			b.Run(fmt.Sprintf("goroutines:%d/size:%d", routines, size), func(b *testing.B) {
+				cb, err := NewCircularBuffer(uint32(size))
+				if err != nil {
+					b.Fatalf("error allocating CircularBuffer: %v", err)
+				}
+
+				perRoutine := b.N / routines
+				var wg sync.WaitGroup
+				for r := 0; r < routines; r++ {
+					wg.Add(1)
+					go func() {
+						for i := 0; i < perRoutine; i++ {
+							cb.Push(&x)
+						}
+						wg.Done()
+					}()
+				}
+				wg.Wait()
+			})
+		}
+	}
+}

internal/profiling/goid_modified.go

@@ -0,0 +1,82 @@
+//go:build grpcgoid
+// +build grpcgoid
+
+/*
+ *
+ * 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 profiling
+
+import (
+	"runtime"
+)
+
+// This stubbed function usually returns zero (see goid_regular.go); however,
+// if grpc is built with `-tags 'grpcgoid'`, a runtime.Goid function, which
+// does not exist in the Go standard library, is expected. While not necessary,
+// sometimes, visualising grpc profiling data in trace-viewer is much nicer
+// with goroutines separated from each other.
+//
+// Several other approaches were considered before arriving at this:
+//
+//  1. Using a CGO module: CGO usually has access to some things that regular
+//     Go does not. Till go1.4, CGO used to have access to the goroutine struct
+//     because the Go runtime was written in C. However, 1.5+ uses a native Go
+//     runtime; as a result, CGO does not have access to the goroutine structure
+//     anymore in modern Go. Besides, CGO interop wasn't fast enough (estimated
+//     to be ~170ns/op). This would also make building grpc require a C
+//     compiler, which isn't a requirement currently, breaking a lot of stuff.
+//
+//  2. Using runtime.Stack stacktrace: While this would remove the need for a
+//     modified Go runtime, this is ridiculously slow, thanks to the all the
+//     string processing shenanigans required to extract the goroutine ID (about
+//     ~2000ns/op).
+//
+//  3. Using Go version-specific build tags: For any given Go version, the
+//     goroutine struct has a fixed structure. As a result, the goroutine ID
+//     could be extracted if we know the offset using some assembly. This would
+//     be faster then #1 and #2, but is harder to maintain. This would require
+//     special Go code that's both architecture-specific and go version-specific
+//     (a quadratic number of variants to maintain).
+//
+//  4. This approach, which requires a simple modification [1] to the Go runtime
+//     to expose the current goroutine's ID. This is the chosen approach and it
+//     takes about ~2 ns/op, which is negligible in the face of the tens of
+//     microseconds that grpc takes to complete a RPC request.
+//
+// [1] To make the goroutine ID visible to Go programs apply the following
+// change to the runtime2.go file in your Go runtime installation:
+//
+//	diff --git a/src/runtime/runtime2.go b/src/runtime/runtime2.go
+//	--- a/src/runtime/runtime2.go
+//	+++ b/src/runtime/runtime2.go
+//	@@ -392,6 +392,10 @@ type stack struct {
+//	 	hi uintptr
+//	 }
+//
+//	+func Goid() int64 {
+//	+  return getg().goid
+//	+}
+//	+
+//	 type g struct {
+//	 	// Stack parameters.
+//	 	// stack describes the actual stack memory: [stack.lo, stack.hi).
+//
+// The exposed runtime.Goid() function will return a int64 goroutine ID.
+func goid() int64 {
+	return runtime.Goid()
+}

internal/profiling/goid_regular.go

@@ -0,0 +1,30 @@
+//go:build !grpcgoid
+// +build !grpcgoid
+
+/*
+ *
+ * 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 profiling
+
+// This dummy function always returns 0. In some modified dev environments,
+// this may be replaced with a call to a function in a modified Go runtime that
+// retrieves the goroutine ID efficiently. See goid_modified.go for a different
+// version of goId that requires a grpcgoid build tag to compile.
+func goid() int64 {
+	return 0
+}

internal/profiling/profiling.go

@@ -0,0 +1,223 @@
+/*
+ *
+ * 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 profiling contains two logical components: buffer.go and
+// profiling.go. The former implements a circular buffer (a.k.a. ring buffer)
+// in a lock-free manner using atomics. This ring buffer is used by
+// profiling.go to store various statistics. For example, StreamStats is a
+// circular buffer of Stat objects, each of which is comprised of Timers.
+//
+// This abstraction is designed to accommodate more stats in the future; for
+// example, if one wants to profile the load balancing layer, which is
+// independent of RPC queries, a separate CircularBuffer can be used.
+//
+// Note that the circular buffer simply takes any type. In the future, more
+// types of measurements (such as the number of memory allocations) could be
+// measured, which might require a different type of object being pushed into
+// the circular buffer.
+package profiling
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"google.golang.org/grpc/internal/profiling/buffer"
+)
+
+// 0 or 1 representing profiling off and on, respectively. Use IsEnabled and
+// Enable to get and set this in a safe manner.
+var profilingEnabled uint32
+
+// IsEnabled returns whether or not profiling is enabled.
+func IsEnabled() bool {
+	return atomic.LoadUint32(&profilingEnabled) > 0
+}
+
+// Enable turns profiling on and off.
+//
+// Note that it is impossible to enable profiling for one server and leave it
+// turned off for another. This is intentional and by design -- if the status
+// of profiling was server-specific, clients wouldn't be able to profile
+// themselves. As a result, Enable turns profiling on and off for all servers
+// and clients in the binary. Each stat will be, however, tagged with whether
+// it's a client stat or a server stat; so you should be able to filter for the
+// right type of stats in post-processing.
+func Enable(enabled bool) {
+	if enabled {
+		atomic.StoreUint32(&profilingEnabled, 1)
+	} else {
+		atomic.StoreUint32(&profilingEnabled, 0)
+	}
+}
+
+// A Timer represents the wall-clock beginning and ending of a logical
+// operation.
+type Timer struct {
+	// Tags is a comma-separated list of strings (usually forward-slash-separated
+	// hierarchical strings) used to categorize a Timer.
+	Tags string
+	// Begin marks the beginning of this timer. The timezone is unspecified, but
+	// must use the same timezone as End; this is so shave off the small, but
+	// non-zero time required to convert to a standard timezone such as UTC.
+	Begin time.Time
+	// End marks the end of a timer.
+	End time.Time
+	// Each Timer must be started and ended within the same goroutine; GoID
+	// captures this goroutine ID. The Go runtime does not typically expose this
+	// information, so this is set to zero in the typical case. However, a
+	// trivial patch to the runtime package can make this field useful. See
+	// goid_modified.go in this package for more details.
+	GoID int64
+}
+
+// NewTimer creates and returns a new Timer object. This is useful when you
+// don't already have a Stat object to associate this Timer with; for example,
+// before the context of a new RPC query is created, a Timer may be needed to
+// measure transport-related operations.
+//
+// Use AppendTimer to append the returned Timer to a Stat.
+func NewTimer(tags string) *Timer {
+	return &Timer{
+		Tags:  tags,
+		Begin: time.Now(),
+		GoID:  goid(),
+	}
+}
+
+// Egress sets the End field of a timer to the current time.
+func (timer *Timer) Egress() {
+	if timer == nil {
+		return
+	}
+
+	timer.End = time.Now()
+}
+
+// A Stat is a collection of Timers that represent timing information for
+// different components within this Stat. For example, a Stat may be used to
+// reference the entire lifetime of an RPC request, with Timers within it
+// representing different components such as encoding, compression, and
+// transport.
+//
+// The user is expected to use the included helper functions to do operations
+// on the Stat such as creating or appending a new timer. Direct operations on
+// the Stat's exported fields (which are exported for encoding reasons) may
+// lead to data races.
+type Stat struct {
+	// Tags is a comma-separated list of strings used to categorize a Stat.
+	Tags string
+	// Stats may also need to store other unstructured information specific to
+	// this stat. For example, a StreamStat will use these bytes to encode the
+	// connection ID and stream ID for each RPC to uniquely identify it. The
+	// encoding that must be used is unspecified.
+	Metadata []byte
+	// A collection of *Timers and a mutex for append operations on the slice.
+	mu     sync.Mutex
+	Timers []*Timer
+}
+
+// A power of two that's large enough to hold all timers within an average RPC
+// request (defined to be a unary request) without any reallocation. A typical
+// unary RPC creates 80-100 timers for various things. While this number is
+// purely anecdotal and may change in the future as the resolution of profiling
+// increases or decreases, it serves as a good estimate for what the initial
+// allocation size should be.
+const defaultStatAllocatedTimers int32 = 128
+
+// NewStat creates and returns a new Stat object.
+func NewStat(tags string) *Stat {
+	return &Stat{
+		Tags:   tags,
+		Timers: make([]*Timer, 0, defaultStatAllocatedTimers),
+	}
+}
+
+// NewTimer creates a Timer object within the given stat if stat is non-nil.
+// The value passed in tags will be attached to the newly created Timer.
+// NewTimer also automatically sets the Begin value of the Timer to the current
+// time. The user is expected to call stat.Egress with the returned index as
+// argument to mark the end.
+func (stat *Stat) NewTimer(tags string) *Timer {
+	if stat == nil {
+		return nil
+	}
+
+	timer := &Timer{
+		Tags:  tags,
+		GoID:  goid(),
+		Begin: time.Now(),
+	}
+	stat.mu.Lock()
+	stat.Timers = append(stat.Timers, timer)
+	stat.mu.Unlock()
+	return timer
+}
+
+// AppendTimer appends a given Timer object to the internal slice of timers. A
+// deep copy of the timer is made (i.e. no reference is retained to this
+// pointer) and the user is expected to lose their reference to the timer to
+// allow the Timer object to be garbage collected.
+func (stat *Stat) AppendTimer(timer *Timer) {
+	if stat == nil || timer == nil {
+		return
+	}
+
+	stat.mu.Lock()
+	stat.Timers = append(stat.Timers, timer)
+	stat.mu.Unlock()
+}
+
+// statsInitialized is 0 before InitStats has been called. Changed to 1 by
+// exactly one call to InitStats.
+var statsInitialized int32
+
+// Stats for the last defaultStreamStatsBufsize RPCs will be stored in memory.
+// This can be configured by the registering server at profiling service
+// initialization with google.golang.org/grpc/profiling/service.ProfilingConfig
+const defaultStreamStatsSize uint32 = 16 << 10
+
+// StreamStats is a CircularBuffer containing data from the last N RPC calls
+// served, where N is set by the user. This will contain both server stats and
+// client stats (but each stat will be tagged with whether it's a server or a
+// client in its Tags).
+var StreamStats *buffer.CircularBuffer
+
+var errAlreadyInitialized = errors.New("profiling may be initialized at most once")
+
+// InitStats initializes all the relevant Stat objects. Must be called exactly
+// once per lifetime of a process; calls after the first one will return an
+// error.
+func InitStats(streamStatsSize uint32) error {
+	var err error
+	if !atomic.CompareAndSwapInt32(&statsInitialized, 0, 1) {
+		return errAlreadyInitialized
+	}
+
+	if streamStatsSize == 0 {
+		streamStatsSize = defaultStreamStatsSize
+	}
+
+	StreamStats, err = buffer.NewCircularBuffer(streamStatsSize)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}