orca/producer_test.go - third_party/github.com/grpc/grpc-go - Git at Google

 /*
  * 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 orca_test

 import (
 	"context"
 	"fmt"
 	"sync"
 	"testing"
 	"time"

 	"github.com/golang/protobuf/proto"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/roundrobin"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/testutils"
 	"google.golang.org/grpc/orca"
 	"google.golang.org/grpc/orca/internal"
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/resolver/manual"
 	"google.golang.org/grpc/status"

 	v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3"
 	v3orcaservicegrpc "github.com/cncf/xds/go/xds/service/orca/v3"
 	v3orcaservicepb "github.com/cncf/xds/go/xds/service/orca/v3"
 )

 // customLBB wraps a round robin LB policy but provides a ClientConn wrapper to
 // add an ORCA OOB report producer for all created SubConns.
 type customLBB struct{}

 func (customLBB) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
 	return balancer.Get(roundrobin.Name).Build(&ccWrapper{ClientConn: cc}, opts)
 }

 func (customLBB) Name() string { return "customLB" }

 func init() {
 	balancer.Register(customLBB{})
 }

 type ccWrapper struct {
 	balancer.ClientConn
 }

 func (w *ccWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
 	if len(addrs) != 1 {
 		panic(fmt.Sprintf("got addrs=%v; want len(addrs) == 1", addrs))
 	}
 	sc, err := w.ClientConn.NewSubConn(addrs, opts)
 	if err != nil {
 		return sc, err
 	}
 	l := getListenerInfo(addrs[0])
 	l.listener.cleanup = orca.RegisterOOBListener(sc, l.listener, l.opts)
 	l.sc = sc
 	return sc, nil
 }

 // listenerInfo is stored in an address's attributes to allow ORCA
 // listeners to be registered on subconns created for that address.
 type listenerInfo struct {
 	listener *testOOBListener
 	opts     orca.OOBListenerOptions
 	sc       balancer.SubConn // Set by the LB policy
 }

 type listenerInfoKey struct{}

 func setListenerInfo(addr resolver.Address, l *listenerInfo) resolver.Address {
 	addr.Attributes = addr.Attributes.WithValue(listenerInfoKey{}, l)
 	return addr
 }

 func getListenerInfo(addr resolver.Address) *listenerInfo {
 	return addr.Attributes.Value(listenerInfoKey{}).(*listenerInfo)
 }

 // testOOBListener is a simple listener that pushes load reports to a channel.
 type testOOBListener struct {
 	cleanup      func()
 	loadReportCh chan *v3orcapb.OrcaLoadReport
 }

 func newTestOOBListener() *testOOBListener {
 	return &testOOBListener{cleanup: func() {}, loadReportCh: make(chan *v3orcapb.OrcaLoadReport)}
 }

 func (t *testOOBListener) Stop() { t.cleanup() }

 func (t *testOOBListener) OnLoadReport(r *v3orcapb.OrcaLoadReport) {
 	t.loadReportCh <- r
 }

 // TestProducer is a basic, end-to-end style test of an LB policy with an
 // OOBListener communicating with a server with an ORCA service.
 func (s) TestProducer(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()

 	// Use a fixed backoff for stream recreation.
 	oldBackoff := internal.DefaultBackoffFunc
 	internal.DefaultBackoffFunc = func(int) time.Duration { return 10 * time.Millisecond }
 	defer func() { internal.DefaultBackoffFunc = oldBackoff }()

 	// Initialize listener for our ORCA server.
 	lis, err := testutils.LocalTCPListener()
 	if err != nil {
 		t.Fatal(err)
 	}

 	// Register the OpenRCAService with a very short metrics reporting interval.
 	const shortReportingInterval = 50 * time.Millisecond
 	opts := orca.ServiceOptions{MinReportingInterval: shortReportingInterval}
 	internal.AllowAnyMinReportingInterval.(func(*orca.ServiceOptions))(&opts)
 	s := grpc.NewServer()
 	orcaSrv, err := orca.Register(s, opts)
 	if err != nil {
 		t.Fatalf("orca.Register failed: %v", err)
 	}
 	go s.Serve(lis)
 	defer s.Stop()

 	// Create our client with an OOB listener in the LB policy it selects.
 	r := manual.NewBuilderWithScheme("whatever")
 	oobLis := newTestOOBListener()

 	lisOpts := orca.OOBListenerOptions{ReportInterval: 50 * time.Millisecond}
 	li := &listenerInfo{listener: oobLis, opts: lisOpts}
 	addr := setListenerInfo(resolver.Address{Addr: lis.Addr().String()}, li)
 	r.InitialState(resolver.State{Addresses: []resolver.Address{addr}})
 	cc, err := grpc.Dial("whatever:///whatever", grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"customLB":{}}]}`), grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		t.Fatalf("grpc.Dial failed: %v", err)
 	}
 	defer cc.Close()

 	// Ensure the OOB listener is stopped before the client is closed to avoid
 	// a potential irrelevant error in the logs.
 	defer oobLis.Stop()

 	// Set a few metrics and wait for them on the client side.
 	orcaSrv.SetCPUUtilization(10)
 	orcaSrv.SetMemoryUtilization(100)
 	orcaSrv.SetUtilization("bob", 555)
 	loadReportWant := &v3orcapb.OrcaLoadReport{
 		CpuUtilization: 10,
 		MemUtilization: 100,
 		Utilization:    map[string]float64{"bob": 555},
 	}

 testReport:
 	for {
 		select {
 		case r := <-oobLis.loadReportCh:
 			t.Log("Load report received: ", r)
 			if proto.Equal(r, loadReportWant) {
 				// Success!
 				break testReport
 			}
 		case <-ctx.Done():
 			t.Fatalf("timed out waiting for load report: %v", loadReportWant)
 		}
 	}

 	// Change and add metrics and wait for them on the client side.
 	orcaSrv.SetCPUUtilization(50)
 	orcaSrv.SetMemoryUtilization(200)
 	orcaSrv.SetUtilization("mary", 321)
 	loadReportWant = &v3orcapb.OrcaLoadReport{
 		CpuUtilization: 50,
 		MemUtilization: 200,
 		Utilization:    map[string]float64{"bob": 555, "mary": 321},
 	}

 	for {
 		select {
 		case r := <-oobLis.loadReportCh:
 			t.Log("Load report received: ", r)
 			if proto.Equal(r, loadReportWant) {
 				// Success!
 				return
 			}
 		case <-ctx.Done():
 			t.Fatalf("timed out waiting for load report: %v", loadReportWant)
 		}
 	}
 }

 // fakeORCAService is a simple implementation of an ORCA service that pushes
 // requests it receives from clients to a channel and sends responses from a
 // channel back.  This allows tests to verify the client is sending requests
 // and processing responses properly.
 type fakeORCAService struct {
 	v3orcaservicegrpc.UnimplementedOpenRcaServiceServer

 	reqCh  chan *v3orcaservicepb.OrcaLoadReportRequest
 	respCh chan interface{} // either *v3orcapb.OrcaLoadReport or error
 }

 func newFakeORCAService() *fakeORCAService {
 	return &fakeORCAService{
 		reqCh:  make(chan *v3orcaservicepb.OrcaLoadReportRequest),
 		respCh: make(chan interface{}),
 	}
 }

 func (f *fakeORCAService) close() {
 	close(f.respCh)
 }

 func (f *fakeORCAService) StreamCoreMetrics(req *v3orcaservicepb.OrcaLoadReportRequest, stream v3orcaservicegrpc.OpenRcaService_StreamCoreMetricsServer) error {
 	f.reqCh <- req
 	for resp := range f.respCh {
 		if err, ok := resp.(error); ok {
 			return err
 		}
 		if err := stream.Send(resp.(*v3orcapb.OrcaLoadReport)); err != nil {
 			// In the event that a stream error occurs, a new stream will have
 			// been created that was waiting for this response message.  Push
 			// it back onto the channel and return.
 			//
 			// This happens because we range over respCh.  If we changed to
 			// instead select on respCh + stream.Context(), the same situation
 			// could still occur due to a race between noticing the two events,
 			// so such a workaround would still be needed to prevent flakiness.
 			f.respCh <- resp
 			return err
 		}
 	}
 	return nil
 }

 // TestProducerBackoff verifies that the ORCA producer applies the proper
 // backoff after stream failures.
 func (s) TestProducerBackoff(t *testing.T) {
 	grpctest.TLogger.ExpectErrorN("injected error", 4)

 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()

 	// Provide a convenient way to expect backoff calls and return a minimal
 	// value.
 	const backoffShouldNotBeCalled = 9999 // Use to assert backoff function is not called.
 	const backoffAllowAny = -1            // Use to ignore any backoff calls.
 	expectedBackoff := backoffAllowAny
 	oldBackoff := internal.DefaultBackoffFunc
 	internal.DefaultBackoffFunc = func(got int) time.Duration {
 		if expectedBackoff == backoffShouldNotBeCalled {
 			t.Errorf("Unexpected backoff call; parameter = %v", got)
 		} else if expectedBackoff != backoffAllowAny {
 			if got != expectedBackoff {
 				t.Errorf("Unexpected backoff received; got %v want %v", got, expectedBackoff)
 			}
 		}
 		return time.Millisecond
 	}
 	defer func() { internal.DefaultBackoffFunc = oldBackoff }()

 	// Initialize listener for our ORCA server.
 	lis, err := testutils.LocalTCPListener()
 	if err != nil {
 		t.Fatal(err)
 	}

 	// Register our fake ORCA service.
 	s := grpc.NewServer()
 	fake := newFakeORCAService()
 	defer fake.close()
 	v3orcaservicegrpc.RegisterOpenRcaServiceServer(s, fake)
 	go s.Serve(lis)
 	defer s.Stop()

 	// Define the report interval and a function to wait for it to be sent to
 	// the server.
 	const reportInterval = 123 * time.Second
 	awaitRequest := func(interval time.Duration) {
 		select {
 		case req := <-fake.reqCh:
 			if got := req.GetReportInterval().AsDuration(); got != interval {
 				t.Errorf("Unexpected report interval; got %v want %v", got, interval)
 			}
 		case <-ctx.Done():
 			t.Fatalf("Did not receive client request")
 		}
 	}

 	// Create our client with an OOB listener in the LB policy it selects.
 	r := manual.NewBuilderWithScheme("whatever")
 	oobLis := newTestOOBListener()

 	lisOpts := orca.OOBListenerOptions{ReportInterval: reportInterval}
 	li := &listenerInfo{listener: oobLis, opts: lisOpts}
 	r.InitialState(resolver.State{Addresses: []resolver.Address{setListenerInfo(resolver.Address{Addr: lis.Addr().String()}, li)}})
 	cc, err := grpc.Dial("whatever:///whatever", grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"customLB":{}}]}`), grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		t.Fatalf("grpc.Dial failed: %v", err)
 	}
 	defer cc.Close()

 	// Ensure the OOB listener is stopped before the client is closed to avoid
 	// a potential irrelevant error in the logs.
 	defer oobLis.Stop()

 	// Define a load report to send and expect the client to see.
 	loadReportWant := &v3orcapb.OrcaLoadReport{
 		CpuUtilization: 10,
 		MemUtilization: 100,
 		Utilization:    map[string]float64{"bob": 555},
 	}

 	// Unblock the fake.
 	awaitRequest(reportInterval)
 	fake.respCh <- loadReportWant
 	select {
 	case r := <-oobLis.loadReportCh:
 		t.Log("Load report received: ", r)
 		if proto.Equal(r, loadReportWant) {
 			// Success!
 			break
 		}
 	case <-ctx.Done():
 		t.Fatalf("timed out waiting for load report: %v", loadReportWant)
 	}

 	// The next request should be immediate, since there was a message
 	// received.
 	expectedBackoff = backoffShouldNotBeCalled
 	fake.respCh <- status.Errorf(codes.Internal, "injected error")
 	awaitRequest(reportInterval)

 	// The next requests will need to backoff.
 	expectedBackoff = 0
 	fake.respCh <- status.Errorf(codes.Internal, "injected error")
 	awaitRequest(reportInterval)
 	expectedBackoff = 1
 	fake.respCh <- status.Errorf(codes.Internal, "injected error")
 	awaitRequest(reportInterval)
 	expectedBackoff = 2
 	fake.respCh <- status.Errorf(codes.Internal, "injected error")
 	awaitRequest(reportInterval)
 	// The next request should be immediate, since there was a message
 	// received.
 	expectedBackoff = backoffShouldNotBeCalled

 	// Send another valid response and wait for it on the client.
 	fake.respCh <- loadReportWant
 	select {
 	case r := <-oobLis.loadReportCh:
 		t.Log("Load report received: ", r)
 		if proto.Equal(r, loadReportWant) {
 			// Success!
 			break
 		}
 	case <-ctx.Done():
 		t.Fatalf("timed out waiting for load report: %v", loadReportWant)
 	}
 }

 // TestProducerMultipleListeners tests that multiple listeners works as
 // expected in a producer: requesting the proper interval and delivering the
 // update to all listeners.
 func (s) TestProducerMultipleListeners(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()

 	// Provide a convenient way to expect backoff calls and return a minimal
 	// value.
 	oldBackoff := internal.DefaultBackoffFunc
 	internal.DefaultBackoffFunc = func(got int) time.Duration {
 		return time.Millisecond
 	}
 	defer func() { internal.DefaultBackoffFunc = oldBackoff }()

 	// Initialize listener for our ORCA server.
 	lis, err := testutils.LocalTCPListener()
 	if err != nil {
 		t.Fatal(err)
 	}

 	// Register our fake ORCA service.
 	s := grpc.NewServer()
 	fake := newFakeORCAService()
 	defer fake.close()
 	v3orcaservicegrpc.RegisterOpenRcaServiceServer(s, fake)
 	go s.Serve(lis)
 	defer s.Stop()

 	// Define the report interval and a function to wait for it to be sent to
 	// the server.
 	const reportInterval1 = 123 * time.Second
 	const reportInterval2 = 234 * time.Second
 	const reportInterval3 = 56 * time.Second
 	awaitRequest := func(interval time.Duration) {
 		select {
 		case req := <-fake.reqCh:
 			if got := req.GetReportInterval().AsDuration(); got != interval {
 				t.Errorf("Unexpected report interval; got %v want %v", got, interval)
 			}
 		case <-ctx.Done():
 			t.Fatalf("Did not receive client request")
 		}
 	}

 	// Create our client with an OOB listener in the LB policy it selects.
 	r := manual.NewBuilderWithScheme("whatever")
 	oobLis1 := newTestOOBListener()
 	lisOpts1 := orca.OOBListenerOptions{ReportInterval: reportInterval1}
 	li := &listenerInfo{listener: oobLis1, opts: lisOpts1}
 	r.InitialState(resolver.State{Addresses: []resolver.Address{setListenerInfo(resolver.Address{Addr: lis.Addr().String()}, li)}})
 	cc, err := grpc.Dial("whatever:///whatever", grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"customLB":{}}]}`), grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		t.Fatalf("grpc.Dial failed: %v", err)
 	}
 	defer cc.Close()

 	// Ensure the OOB listener is stopped before the client is closed to avoid
 	// a potential irrelevant error in the logs.
 	defer oobLis1.Stop()

 	oobLis2 := newTestOOBListener()
 	lisOpts2 := orca.OOBListenerOptions{ReportInterval: reportInterval2}

 	oobLis3 := newTestOOBListener()
 	lisOpts3 := orca.OOBListenerOptions{ReportInterval: reportInterval3}

 	// Define a load report to send and expect the client to see.
 	loadReportWant := &v3orcapb.OrcaLoadReport{
 		CpuUtilization: 10,
 		MemUtilization: 100,
 		Utilization:    map[string]float64{"bob": 555},
 	}

 	// Receive reports and update counts for the three listeners.
 	var reportsMu sync.Mutex
 	var reportsReceived1, reportsReceived2, reportsReceived3 int
 	go func() {
 		for {
 			select {
 			case r := <-oobLis1.loadReportCh:
 				t.Log("Load report 1 received: ", r)
 				if !proto.Equal(r, loadReportWant) {
 					t.Errorf("Unexpected report received: %+v", r)
 				}
 				reportsMu.Lock()
 				reportsReceived1++
 				reportsMu.Unlock()
 			case r := <-oobLis2.loadReportCh:
 				t.Log("Load report 2 received: ", r)
 				if !proto.Equal(r, loadReportWant) {
 					t.Errorf("Unexpected report received: %+v", r)
 				}
 				reportsMu.Lock()
 				reportsReceived2++
 				reportsMu.Unlock()
 			case r := <-oobLis3.loadReportCh:
 				t.Log("Load report 3 received: ", r)
 				if !proto.Equal(r, loadReportWant) {
 					t.Errorf("Unexpected report received: %+v", r)
 				}
 				reportsMu.Lock()
 				reportsReceived3++
 				reportsMu.Unlock()
 			case <-ctx.Done():
 				// Test has ended; exit
 				return
 			}
 		}
 	}()

 	// checkReports is a helper function to check the report counts for the three listeners.
 	checkReports := func(r1, r2, r3 int) {
 		t.Helper()
 		for ctx.Err() == nil {
 			reportsMu.Lock()
 			if r1 == reportsReceived1 && r2 == reportsReceived2 && r3 == reportsReceived3 {
 				// Success!
 				reportsMu.Unlock()
 				return
 			}
 			if reportsReceived1 > r1 || reportsReceived2 > r2 || reportsReceived3 > r3 {
 				reportsMu.Unlock()
 				t.Fatalf("received excess reports. got %v %v %v; want %v %v %v", reportsReceived1, reportsReceived2, reportsReceived3, r1, r2, r3)
 				return
 			}
 			reportsMu.Unlock()
 			time.Sleep(10 * time.Millisecond)
 		}
 		t.Fatalf("timed out waiting for reports received. got %v %v %v; want %v %v %v", reportsReceived1, reportsReceived2, reportsReceived3, r1, r2, r3)
 	}

 	// Only 1 listener; expect reportInterval1 to be used and expect the report
 	// to be sent to the listener.
 	awaitRequest(reportInterval1)
 	fake.respCh <- loadReportWant
 	checkReports(1, 0, 0)

 	// Register listener 2 with a less frequent interval; no need to recreate
 	// stream.  Report should go to both listeners.
 	oobLis2.cleanup = orca.RegisterOOBListener(li.sc, oobLis2, lisOpts2)
 	fake.respCh <- loadReportWant
 	checkReports(2, 1, 0)

 	// Register listener 3 with a more frequent interval; stream is recreated
 	// with this interval after the next report is received.  The first report
 	// will go to all three listeners.
 	oobLis3.cleanup = orca.RegisterOOBListener(li.sc, oobLis3, lisOpts3)
 	fake.respCh <- loadReportWant
 	checkReports(3, 2, 1)
 	awaitRequest(reportInterval3)

 	// Another report without a change in listeners should go to all three listeners.
 	fake.respCh <- loadReportWant
 	checkReports(4, 3, 2)

 	// Stop listener 2.  This does not affect the interval as listener 3 is
 	// still the shortest.  The next update goes to listeners 1 and 3.
 	oobLis2.Stop()
 	fake.respCh <- loadReportWant
 	checkReports(5, 3, 3)

 	// Stop listener 3.  This makes the interval longer, with stream recreation
 	// delayed until the next report is received.  Reports should only go to
 	// listener 1 now.
 	oobLis3.Stop()
 	fake.respCh <- loadReportWant
 	checkReports(6, 3, 3)
 	awaitRequest(reportInterval1)
 	// Another report without a change in listeners should go to the first listener.
 	fake.respCh <- loadReportWant
 	checkReports(7, 3, 3)
 }
	/*
	* 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 orca_test

	import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/roundrobin"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/internal/grpctest"
	"google.golang.org/grpc/internal/testutils"
	"google.golang.org/grpc/orca"
	"google.golang.org/grpc/orca/internal"
	"google.golang.org/grpc/resolver"
	"google.golang.org/grpc/resolver/manual"
	"google.golang.org/grpc/status"

	v3orcapb "github.com/cncf/xds/go/xds/data/orca/v3"
	v3orcaservicegrpc "github.com/cncf/xds/go/xds/service/orca/v3"
	v3orcaservicepb "github.com/cncf/xds/go/xds/service/orca/v3"
	)

	// customLBB wraps a round robin LB policy but provides a ClientConn wrapper to
	// add an ORCA OOB report producer for all created SubConns.
	type customLBB struct{}

	func (customLBB) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
	return balancer.Get(roundrobin.Name).Build(&ccWrapper{ClientConn: cc}, opts)
	}

	func (customLBB) Name() string { return "customLB" }

	func init() {
	balancer.Register(customLBB{})
	}

	type ccWrapper struct {
	balancer.ClientConn
	}

	func (w *ccWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
	if len(addrs) != 1 {
	panic(fmt.Sprintf("got addrs=%v; want len(addrs) == 1", addrs))
	}
	sc, err := w.ClientConn.NewSubConn(addrs, opts)
	if err != nil {
	return sc, err
	}
	l := getListenerInfo(addrs[0])
	l.listener.cleanup = orca.RegisterOOBListener(sc, l.listener, l.opts)
	l.sc = sc
	return sc, nil
	}

	// listenerInfo is stored in an address's attributes to allow ORCA
	// listeners to be registered on subconns created for that address.
	type listenerInfo struct {
	listener *testOOBListener
	opts orca.OOBListenerOptions
	sc balancer.SubConn // Set by the LB policy
	}

	type listenerInfoKey struct{}

	func setListenerInfo(addr resolver.Address, l *listenerInfo) resolver.Address {
	addr.Attributes = addr.Attributes.WithValue(listenerInfoKey{}, l)
	return addr
	}

	func getListenerInfo(addr resolver.Address) *listenerInfo {
	return addr.Attributes.Value(listenerInfoKey{}).(*listenerInfo)
	}

	// testOOBListener is a simple listener that pushes load reports to a channel.
	type testOOBListener struct {
	cleanup func()
	loadReportCh chan *v3orcapb.OrcaLoadReport
	}

	func newTestOOBListener() *testOOBListener {
	return &testOOBListener{cleanup: func() {}, loadReportCh: make(chan *v3orcapb.OrcaLoadReport)}
	}

	func (t *testOOBListener) Stop() { t.cleanup() }

	func (t testOOBListener) OnLoadReport(r v3orcapb.OrcaLoadReport) {
	t.loadReportCh <- r
	}

	// TestProducer is a basic, end-to-end style test of an LB policy with an
	// OOBListener communicating with a server with an ORCA service.
	func (s) TestProducer(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	// Use a fixed backoff for stream recreation.
	oldBackoff := internal.DefaultBackoffFunc
	internal.DefaultBackoffFunc = func(int) time.Duration { return 10 * time.Millisecond }
	defer func() { internal.DefaultBackoffFunc = oldBackoff }()

	// Initialize listener for our ORCA server.
	lis, err := testutils.LocalTCPListener()
	if err != nil {
	t.Fatal(err)
	}

	// Register the OpenRCAService with a very short metrics reporting interval.
	const shortReportingInterval = 50 * time.Millisecond
	opts := orca.ServiceOptions{MinReportingInterval: shortReportingInterval}
	internal.AllowAnyMinReportingInterval.(func(*orca.ServiceOptions))(&opts)
	s := grpc.NewServer()
	orcaSrv, err := orca.Register(s, opts)
	if err != nil {
	t.Fatalf("orca.Register failed: %v", err)
	}
	go s.Serve(lis)
	defer s.Stop()

	// Create our client with an OOB listener in the LB policy it selects.
	r := manual.NewBuilderWithScheme("whatever")
	oobLis := newTestOOBListener()

	lisOpts := orca.OOBListenerOptions{ReportInterval: 50 * time.Millisecond}
	li := &listenerInfo{listener: oobLis, opts: lisOpts}
	addr := setListenerInfo(resolver.Address{Addr: lis.Addr().String()}, li)
	r.InitialState(resolver.State{Addresses: []resolver.Address{addr}})
	cc, err := grpc.Dial("whatever:///whatever", grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"customLB":{}}]}`), grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
	t.Fatalf("grpc.Dial failed: %v", err)
	}
	defer cc.Close()

	// Ensure the OOB listener is stopped before the client is closed to avoid
	// a potential irrelevant error in the logs.
	defer oobLis.Stop()

	// Set a few metrics and wait for them on the client side.
	orcaSrv.SetCPUUtilization(10)
	orcaSrv.SetMemoryUtilization(100)
	orcaSrv.SetUtilization("bob", 555)
	loadReportWant := &v3orcapb.OrcaLoadReport{
	CpuUtilization: 10,
	MemUtilization: 100,
	Utilization: map[string]float64{"bob": 555},
	}

	testReport:
	for {
	select {
	case r := <-oobLis.loadReportCh:
	t.Log("Load report received: ", r)
	if proto.Equal(r, loadReportWant) {
	// Success!
	break testReport
	}
	case <-ctx.Done():
	t.Fatalf("timed out waiting for load report: %v", loadReportWant)
	}
	}

	// Change and add metrics and wait for them on the client side.
	orcaSrv.SetCPUUtilization(50)
	orcaSrv.SetMemoryUtilization(200)
	orcaSrv.SetUtilization("mary", 321)
	loadReportWant = &v3orcapb.OrcaLoadReport{
	CpuUtilization: 50,
	MemUtilization: 200,
	Utilization: map[string]float64{"bob": 555, "mary": 321},
	}

	for {
	select {
	case r := <-oobLis.loadReportCh:
	t.Log("Load report received: ", r)
	if proto.Equal(r, loadReportWant) {
	// Success!
	return
	}
	case <-ctx.Done():
	t.Fatalf("timed out waiting for load report: %v", loadReportWant)
	}
	}
	}

	// fakeORCAService is a simple implementation of an ORCA service that pushes
	// requests it receives from clients to a channel and sends responses from a
	// channel back. This allows tests to verify the client is sending requests
	// and processing responses properly.
	type fakeORCAService struct {
	v3orcaservicegrpc.UnimplementedOpenRcaServiceServer

	reqCh chan *v3orcaservicepb.OrcaLoadReportRequest
	respCh chan interface{} // either *v3orcapb.OrcaLoadReport or error
	}

	func newFakeORCAService() *fakeORCAService {
	return &fakeORCAService{
	reqCh: make(chan *v3orcaservicepb.OrcaLoadReportRequest),
	respCh: make(chan interface{}),
	}
	}

	func (f *fakeORCAService) close() {
	close(f.respCh)
	}

	func (f fakeORCAService) StreamCoreMetrics(req v3orcaservicepb.OrcaLoadReportRequest, stream v3orcaservicegrpc.OpenRcaService_StreamCoreMetricsServer) error {
	f.reqCh <- req
	for resp := range f.respCh {
	if err, ok := resp.(error); ok {
	return err
	}
	if err := stream.Send(resp.(*v3orcapb.OrcaLoadReport)); err != nil {
	// In the event that a stream error occurs, a new stream will have
	// been created that was waiting for this response message. Push
	// it back onto the channel and return.
	//
	// This happens because we range over respCh. If we changed to
	// instead select on respCh + stream.Context(), the same situation
	// could still occur due to a race between noticing the two events,
	// so such a workaround would still be needed to prevent flakiness.
	f.respCh <- resp
	return err
	}
	}
	return nil
	}

	// TestProducerBackoff verifies that the ORCA producer applies the proper
	// backoff after stream failures.
	func (s) TestProducerBackoff(t *testing.T) {
	grpctest.TLogger.ExpectErrorN("injected error", 4)

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	// Provide a convenient way to expect backoff calls and return a minimal
	// value.
	const backoffShouldNotBeCalled = 9999 // Use to assert backoff function is not called.
	const backoffAllowAny = -1 // Use to ignore any backoff calls.
	expectedBackoff := backoffAllowAny
	oldBackoff := internal.DefaultBackoffFunc
	internal.DefaultBackoffFunc = func(got int) time.Duration {
	if expectedBackoff == backoffShouldNotBeCalled {
	t.Errorf("Unexpected backoff call; parameter = %v", got)
	} else if expectedBackoff != backoffAllowAny {
	if got != expectedBackoff {
	t.Errorf("Unexpected backoff received; got %v want %v", got, expectedBackoff)
	}
	}
	return time.Millisecond
	}
	defer func() { internal.DefaultBackoffFunc = oldBackoff }()

	// Initialize listener for our ORCA server.
	lis, err := testutils.LocalTCPListener()
	if err != nil {
	t.Fatal(err)
	}

	// Register our fake ORCA service.
	s := grpc.NewServer()
	fake := newFakeORCAService()
	defer fake.close()
	v3orcaservicegrpc.RegisterOpenRcaServiceServer(s, fake)
	go s.Serve(lis)
	defer s.Stop()

	// Define the report interval and a function to wait for it to be sent to
	// the server.
	const reportInterval = 123 * time.Second
	awaitRequest := func(interval time.Duration) {
	select {
	case req := <-fake.reqCh:
	if got := req.GetReportInterval().AsDuration(); got != interval {
	t.Errorf("Unexpected report interval; got %v want %v", got, interval)
	}
	case <-ctx.Done():
	t.Fatalf("Did not receive client request")
	}
	}

	// Create our client with an OOB listener in the LB policy it selects.
	r := manual.NewBuilderWithScheme("whatever")
	oobLis := newTestOOBListener()

	lisOpts := orca.OOBListenerOptions{ReportInterval: reportInterval}
	li := &listenerInfo{listener: oobLis, opts: lisOpts}
	r.InitialState(resolver.State{Addresses: []resolver.Address{setListenerInfo(resolver.Address{Addr: lis.Addr().String()}, li)}})
	cc, err := grpc.Dial("whatever:///whatever", grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"customLB":{}}]}`), grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
	t.Fatalf("grpc.Dial failed: %v", err)
	}
	defer cc.Close()

	// Ensure the OOB listener is stopped before the client is closed to avoid
	// a potential irrelevant error in the logs.
	defer oobLis.Stop()

	// Define a load report to send and expect the client to see.
	loadReportWant := &v3orcapb.OrcaLoadReport{
	CpuUtilization: 10,
	MemUtilization: 100,
	Utilization: map[string]float64{"bob": 555},
	}

	// Unblock the fake.
	awaitRequest(reportInterval)
	fake.respCh <- loadReportWant
	select {
	case r := <-oobLis.loadReportCh:
	t.Log("Load report received: ", r)
	if proto.Equal(r, loadReportWant) {
	// Success!
	break
	}
	case <-ctx.Done():
	t.Fatalf("timed out waiting for load report: %v", loadReportWant)
	}

	// The next request should be immediate, since there was a message
	// received.
	expectedBackoff = backoffShouldNotBeCalled
	fake.respCh <- status.Errorf(codes.Internal, "injected error")
	awaitRequest(reportInterval)

	// The next requests will need to backoff.
	expectedBackoff = 0
	fake.respCh <- status.Errorf(codes.Internal, "injected error")
	awaitRequest(reportInterval)
	expectedBackoff = 1
	fake.respCh <- status.Errorf(codes.Internal, "injected error")
	awaitRequest(reportInterval)
	expectedBackoff = 2
	fake.respCh <- status.Errorf(codes.Internal, "injected error")
	awaitRequest(reportInterval)
	// The next request should be immediate, since there was a message
	// received.
	expectedBackoff = backoffShouldNotBeCalled

	// Send another valid response and wait for it on the client.
	fake.respCh <- loadReportWant
	select {
	case r := <-oobLis.loadReportCh:
	t.Log("Load report received: ", r)
	if proto.Equal(r, loadReportWant) {
	// Success!
	break
	}
	case <-ctx.Done():
	t.Fatalf("timed out waiting for load report: %v", loadReportWant)
	}
	}

	// TestProducerMultipleListeners tests that multiple listeners works as
	// expected in a producer: requesting the proper interval and delivering the
	// update to all listeners.
	func (s) TestProducerMultipleListeners(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	// Provide a convenient way to expect backoff calls and return a minimal
	// value.
	oldBackoff := internal.DefaultBackoffFunc
	internal.DefaultBackoffFunc = func(got int) time.Duration {
	return time.Millisecond
	}
	defer func() { internal.DefaultBackoffFunc = oldBackoff }()

	// Initialize listener for our ORCA server.
	lis, err := testutils.LocalTCPListener()
	if err != nil {
	t.Fatal(err)
	}

	// Register our fake ORCA service.
	s := grpc.NewServer()
	fake := newFakeORCAService()
	defer fake.close()
	v3orcaservicegrpc.RegisterOpenRcaServiceServer(s, fake)
	go s.Serve(lis)
	defer s.Stop()

	// Define the report interval and a function to wait for it to be sent to
	// the server.
	const reportInterval1 = 123 * time.Second
	const reportInterval2 = 234 * time.Second
	const reportInterval3 = 56 * time.Second
	awaitRequest := func(interval time.Duration) {
	select {
	case req := <-fake.reqCh:
	if got := req.GetReportInterval().AsDuration(); got != interval {
	t.Errorf("Unexpected report interval; got %v want %v", got, interval)
	}
	case <-ctx.Done():
	t.Fatalf("Did not receive client request")
	}
	}

	// Create our client with an OOB listener in the LB policy it selects.
	r := manual.NewBuilderWithScheme("whatever")
	oobLis1 := newTestOOBListener()
	lisOpts1 := orca.OOBListenerOptions{ReportInterval: reportInterval1}
	li := &listenerInfo{listener: oobLis1, opts: lisOpts1}
	r.InitialState(resolver.State{Addresses: []resolver.Address{setListenerInfo(resolver.Address{Addr: lis.Addr().String()}, li)}})
	cc, err := grpc.Dial("whatever:///whatever", grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"customLB":{}}]}`), grpc.WithResolvers(r), grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
	t.Fatalf("grpc.Dial failed: %v", err)
	}
	defer cc.Close()

	// Ensure the OOB listener is stopped before the client is closed to avoid
	// a potential irrelevant error in the logs.
	defer oobLis1.Stop()

	oobLis2 := newTestOOBListener()
	lisOpts2 := orca.OOBListenerOptions{ReportInterval: reportInterval2}

	oobLis3 := newTestOOBListener()
	lisOpts3 := orca.OOBListenerOptions{ReportInterval: reportInterval3}

	// Define a load report to send and expect the client to see.
	loadReportWant := &v3orcapb.OrcaLoadReport{
	CpuUtilization: 10,
	MemUtilization: 100,
	Utilization: map[string]float64{"bob": 555},
	}

	// Receive reports and update counts for the three listeners.
	var reportsMu sync.Mutex
	var reportsReceived1, reportsReceived2, reportsReceived3 int
	go func() {
	for {
	select {
	case r := <-oobLis1.loadReportCh:
	t.Log("Load report 1 received: ", r)
	if !proto.Equal(r, loadReportWant) {
	t.Errorf("Unexpected report received: %+v", r)
	}
	reportsMu.Lock()
	reportsReceived1++
	reportsMu.Unlock()
	case r := <-oobLis2.loadReportCh:
	t.Log("Load report 2 received: ", r)
	if !proto.Equal(r, loadReportWant) {
	t.Errorf("Unexpected report received: %+v", r)
	}
	reportsMu.Lock()
	reportsReceived2++
	reportsMu.Unlock()
	case r := <-oobLis3.loadReportCh:
	t.Log("Load report 3 received: ", r)
	if !proto.Equal(r, loadReportWant) {
	t.Errorf("Unexpected report received: %+v", r)
	}
	reportsMu.Lock()
	reportsReceived3++
	reportsMu.Unlock()
	case <-ctx.Done():
	// Test has ended; exit
	return
	}
	}
	}()

	// checkReports is a helper function to check the report counts for the three listeners.
	checkReports := func(r1, r2, r3 int) {
	t.Helper()
	for ctx.Err() == nil {
	reportsMu.Lock()
	if r1 == reportsReceived1 && r2 == reportsReceived2 && r3 == reportsReceived3 {
	// Success!
	reportsMu.Unlock()
	return
	}
	if reportsReceived1 > r1 \|\| reportsReceived2 > r2 \|\| reportsReceived3 > r3 {
	reportsMu.Unlock()
	t.Fatalf("received excess reports. got %v %v %v; want %v %v %v", reportsReceived1, reportsReceived2, reportsReceived3, r1, r2, r3)
	return
	}
	reportsMu.Unlock()
	time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for reports received. got %v %v %v; want %v %v %v", reportsReceived1, reportsReceived2, reportsReceived3, r1, r2, r3)
	}

	// Only 1 listener; expect reportInterval1 to be used and expect the report
	// to be sent to the listener.
	awaitRequest(reportInterval1)
	fake.respCh <- loadReportWant
	checkReports(1, 0, 0)

	// Register listener 2 with a less frequent interval; no need to recreate
	// stream. Report should go to both listeners.
	oobLis2.cleanup = orca.RegisterOOBListener(li.sc, oobLis2, lisOpts2)
	fake.respCh <- loadReportWant
	checkReports(2, 1, 0)

	// Register listener 3 with a more frequent interval; stream is recreated
	// with this interval after the next report is received. The first report
	// will go to all three listeners.
	oobLis3.cleanup = orca.RegisterOOBListener(li.sc, oobLis3, lisOpts3)
	fake.respCh <- loadReportWant
	checkReports(3, 2, 1)
	awaitRequest(reportInterval3)

	// Another report without a change in listeners should go to all three listeners.
	fake.respCh <- loadReportWant
	checkReports(4, 3, 2)

	// Stop listener 2. This does not affect the interval as listener 3 is
	// still the shortest. The next update goes to listeners 1 and 3.
	oobLis2.Stop()
	fake.respCh <- loadReportWant
	checkReports(5, 3, 3)

	// Stop listener 3. This makes the interval longer, with stream recreation
	// delayed until the next report is received. Reports should only go to
	// listener 1 now.
	oobLis3.Stop()
	fake.respCh <- loadReportWant
	checkReports(6, 3, 3)
	awaitRequest(reportInterval1)
	// Another report without a change in listeners should go to the first listener.
	fake.respCh <- loadReportWant
	checkReports(7, 3, 3)
	}