ronantakizawa/codereview-bench huggingface dataset:
Repository: grpc/grpc-go
File: balancer/randomsubsetting/randomsubsetting_test.go (Go)

Code before:
	testgrpc "google.golang.org/grpc/interop/grpc_testing"
	testpb "google.golang.org/grpc/interop/grpc_testing"
)

var defaultTestTimeout = 5 * time.Second

type s struct {
	grpctest.Tester
}

func Test(t *testing.T) {
	grpctest.RunSubTests(t, s{})
}

func (s) TestParseConfig(t *testing.T) {
	parser := bb{}
	tests := []struct {
		name    string
		input   string
		wantCfg serviceconfig.LoadBalancingConfig
		wantErr string
	}{
		{
			name:    "invalid-json",
			input:   "{{invalidjson{{",
			wantErr: "invalid character",
		},
		{
			name:    "empty_config",
			input:   `{}`,
			wantErr: "SubsetSize must be greater than 0",
		},
		{
			name:    "subset_size_zero",
			input:   `{ "subsetSize": 0 }`,
			wantErr: "SubsetSize must be greater than 0",
		},
		{
			name:    "child_policy_missing",
			input:   `{ "subsetSize": 1 }`,
			wantErr: "ChildPolicy must be specified",
		},
		{
			name:    "child_policy_not_registered",
			input:   `{ "subsetSize": 1 , "childPolicy": [{"unregistered_lb": {}}] }`,
			wantErr: "no supported policies found",
		},
		{
			name:  "success",
			input: `{ "subsetSize": 3, "childPolicy": [{"round_robin": {}}]}`,
			wantCfg: &lbConfig{

Review comment (bug):
This has now become a change detector test for the json package. Please change the error string returned when unmarshal fails so that you can look for that here instead. So, maybe something like:

```golang
	if err := json.Unmarshal(s, lbCfg); err != nil {
		return nil, fmt.Errorf("randomsubsetting: json.Unmarshal failed for configuration: %s with error: %v", string(s), err)
	}
```

Then, your `wantErr` here could be "json.Unmarshal failed for configuration"

Code after:

type s struct {
	grpctest.Tester
}

func Test(t *testing.T) {
	grpctest.RunSubTests(t, s{})
}

func (s) TestParseConfig(t *testing.T) {
	parser := bb{}
	tests := []struct {
		name    string
		input   string
		wantCfg serviceconfig.LoadBalancingConfig
		wantErr string
	}{
		{
			name:    "invalid_json",
			input:   "{{invalidjson{{",
			wantErr: "json.Unmarshal failed for configuration",
		},
		{
			name:    "empty_config",
			input:   `{}`,
			wantErr: "SubsetSize must be greater than 0",
		},
		{
			name:    "subset_size_zero",
			input:   `{ "subsetSize": 0 }`,
			wantErr: "SubsetSize must be greater than 0",
		},
		{
			name:    "child_policy_missing",
			input:   `{ "subsetSize": 1 }`,
			wantErr: "ChildPolicy must be specified",
		},
		{
			name:    "child_policy_not_registered",
			input:   `{ "subsetSize": 1 , "childPolicy": [{"unregistered_lb": {}}] }`,
			wantErr: "no supported policies found",
		},
		{
			name:  "success",
			input: `{ "subsetSize": 3, "childPolicy": [{"round_robin": {}}]}`,
			wantCfg: &lbConfig{
				SubsetSize:  3,
				ChildPolicy: &iserviceconfig.BalancerConfig{Name: "round_robin"},
			},
		},
	}
Repository: grpc/grpc-go
File: internal/transport/http_util_test.go (Go)

Code before:
			wantData: []byte("foo"),
		},
		{
			name: "padded_zero_data_some_padding",
			w: func(fr *framer) error {
				return fr.fr.WriteDataPadded(1, false, []byte{}, []byte{0, 0})
			},
			wantData: []byte{},
		},
		{
			name: "stream_id_0",
			w: func(fr *framer) error {
				fr.fr.AllowIllegalWrites = true
				return fr.fr.WriteData(0, false, []byte("foo"))
			},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			fr, _ := testFramer()
			// Write the frame using the provided function.
			writeErr := tt.w(fr)
			if writeErr != nil {
				t.Fatalf("tt.w() returned unexpected error: %v", writeErr)
			}
			fr.writer.Flush()

			// Read the frame back.
			f, err := fr.readFrame()
			if gotErr := err != nil; gotErr != tt.wantErr {
				t.Fatalf("ReadFrame() err = %v; want %t", err, tt.wantErr)
			}
			if tt.wantErr {
				return
			}

			df, ok := f.(*parsedDataFrame)
			if !ok {
				t.Fatalf("ReadFrame() returned %T, want *parsedDataFrame", f)
			}
			if gotData := df.data.ReadOnlyData(); !bytes.Equal(gotData, tt.wantData) {
				t.Fatalf("parsedDataFrame.Data() = %q, want %q", gotData, tt.wantData)
			}
			df.data.Free()
		})
	}
}

func testFramer() (*framer, *bytes.Buffer) {

Review comment (bug):
Nit: `tt.w()` is not very readable when it shows up in the log of a failing test. Maybe something more descriptive here please.

Code after:
func (s) TestFramer_ParseDataFrame(t *testing.T) {
	tests := []struct {
		name                string
		wire                []byte // from frame header onward
		wantData            []byte
		wantErr             error
		wantErrDetailSubstr string
	}{
		{
			name: "good_padded",
			wire: buildDataFrame(http2.FrameHeader{
				Type: http2.FrameData, Length: 6, StreamID: 1, Flags: http2.FlagDataPadded,
			}, []byte{
				2,             // pad length
				'f', 'o', 'o', // data
				0, 0, // padding
			}),
			wantData: []byte("foo"),
		},
		{
			name: "good_unpadded",
			wire: buildDataFrame(http2.FrameHeader{
				Type: http2.FrameData, Length: 3, StreamID: 1, Flags: 0,
			}, []byte("foo")),
			wantData: []byte("foo"),
		},
		{
			name: "stream_id_0",
			wire: buildDataFrame(http2.FrameHeader{
				Type: http2.FrameData, Length: 1, StreamID: 0, Flags: 0,
			}, []byte{0}),
			wantErr:             http2.ConnectionError(http2.ErrCodeProtocol),
			wantErrDetailSubstr: "DATA frame with stream ID 0",
		},
		{
			name: "pad_size_bigger_than_payload",
			wire: buildDataFrame(http2.FrameHeader{
				Type: http2.FrameData, Length: 4, StreamID: 1, Flags: http2.FlagDataPadded,
			}, []byte{
				4,        // pad length of 4
				'f', 'o', // data 'fo' is 2 bytes
				0, // padding 0 is 1 byte.
			}), // pad length 4 but only 3 bytes for data+padding available in payload after pad length byte
			wantErr:             http2.ConnectionError(http2.ErrCodeProtocol),
			wantErrDetailSubstr: "pad size larger than data payload",
		},
		{
			name: "padded_zero_data_some_padding",
			wire: buildDataFrame(http2.FrameHeader{
				Type: http2.FrameData, Length: 3, StreamID: 1, Flags: http2.FlagDataPadded,
			}, []byte{
Repository: grpc/grpc-go
File: mem/buffer_slice.go (Go)

Code before:
// each underlying buffer.
func (s BufferSlice) Reader() Reader {
	s.Ref()
	return &sliceReader{
		data: s,
		len:  s.Len(),
	}
}

// Reader exposes a BufferSlice's data as an io.Reader, allowing it to interface
// with other parts systems. It also provides an additional convenience method
// Remaining(), which returns the number of unread bytes remaining in the slice.
// Buffers will be freed as they are read.
type Reader interface {
	io.Reader
	io.ByteReader
	// Close frees the underlying BufferSlice and never returns an error. Subsequent
	// calls to Read will return (0, io.EOF).
	Close() error
	// Remaining returns the number of unread bytes remaining in the slice.
	Remaining() int
	// Reset frees the currently held buffer slice and starts reading from the
	// provided slice. This allows reusing the reader object.
	Reset(s BufferSlice)
	// Peek returns up to the next n bytes from the reader's current position as
	// a slice of byte slices.
	// The reader's position is not advanced.
	// Next appends results to the provided res slice and returns the updated
	// slice.
	// The returned subslices are views into the underlying buffers and are only
	// valid until the reader is advanced past the corresponding buffer.
	Peek(n int, res [][]byte) [][]byte
	// Discard moves the read index forward by n bytes.
	// It frees buffers as they are fully consumed.
	Discard(n int)
}

type sliceReader struct {
	data BufferSlice
	len  int
	// The index into data[0].ReadOnlyData().
	bufferIdx int
}

func (r *sliceReader) Remaining() int {
	return r.len
}

func (r *sliceReader) Reset(s BufferSlice) {
	r.data.Free()
	s.Ref()

Review comment (suggestion):
Does it make sense to return the actual number of bytes returned in `res`? Most APIs to read data do that.

Code after:
// Reader returns a new Reader for the input slice after taking references to
// each underlying buffer.
func (s BufferSlice) Reader() *Reader {
	s.Ref()
	return &Reader{
		data: s,
		len:  s.Len(),
	}
}

// Reader exposes a BufferSlice's data as an io.Reader, allowing it to interface
// with other systems.
//
// Buffers will be freed as they are read.
//
// A Reader can be constructed from a BufferSlice; alternatively the zero value
// of a Reader may be used after calling Reset on it.
type Reader struct {
	data BufferSlice
	len  int
	// The index into data[0].ReadOnlyData().
	bufferIdx int
}

// Remaining returns the number of unread bytes remaining in the slice.
func (r *Reader) Remaining() int {
	return r.len
}

// Reset frees the currently held buffer slice and starts reading from the
// provided slice. This allows reusing the reader object.
func (r *Reader) Reset(s BufferSlice) {
	r.data.Free()
	s.Ref()
	r.data = s
	r.len = s.Len()
	r.bufferIdx = 0
}

// Close frees the underlying BufferSlice and never returns an error. Subsequent
// calls to Read will return (0, io.EOF).
func (r *Reader) Close() error {
	r.data.Free()
	r.data = nil
	r.len = 0
	return nil
}

func (r *Reader) freeFirstBufferIfEmpty() bool {
	if len(r.data) == 0 || r.bufferIdx != len(r.data[0].ReadOnlyData()) {
		return false
Repository: grpc/grpc-go
File: mem/buffer_slice_test.go (Go)

Code before:

// testPool is an implementation of BufferPool that allows to ensure that:
// - there are matching Put calls for all Get calls.
// - there are no unexpected Put calls.
type testPool struct {
	allocated map[*[]byte]struct{}
}

func (t *testPool) Get(length int) *[]byte {
	buf := make([]byte, length)
	t.allocated[&buf] = struct{}{}
	return &buf
}

func (t *testPool) Put(buf *[]byte) {
	if _, ok := t.allocated[buf]; !ok {
		panic("unexpected put")
	}
	delete(t.allocated, buf)
}

func (s) TestBufferSlice_Iteration(t *testing.T) {
	tests := []struct {
		name       string
		buffers    [][]byte
		operations func(t *testing.T, c mem.Reader)
	}{
		{
			name: "empty",
			operations: func(t *testing.T, c mem.Reader) {
				if c.Remaining() != 0 {
					t.Errorf("Remaining() = %v, want 0", c.Remaining())
				}
				res := c.Peek(1, nil)
				if len(res) != 0 {
					t.Errorf("Peek() got %v, want empty", res)
				}
				c.Discard(1)
				if c.Remaining() != 0 {
					t.Errorf("Remaining() after Discard = %v, want 0", c.Remaining())
				}
			},
		},
		{
			name:    "single_buffer",
			buffers: [][]byte{[]byte("0123456789")},
			operations: func(t *testing.T, c mem.Reader) {
				if c.Remaining() != 10 {
					t.Errorf("Remaining() = %v, want 10", c.Remaining())
				}


Review comment (suggestion):
It does feel like many of checks in many of the operations should result in calls to `t.Fatal` instead of `t.Error` since it doesn't make sense to carry on (for example, when you start with an empty buffer and `Remaining` doesn't return `0`).

And if we start calling `t.Fatal` at different places, we also probably need to defer the call to the `Close` method of the reader in the main test logic.

Code after:

// testPool is an implementation of BufferPool that allows to ensure that:
// - there are matching Put calls for all Get calls.
// - there are no unexpected Put calls.
type testPool struct {
	allocated map[*[]byte]struct{}
}

func (t *testPool) Get(length int) *[]byte {
	buf := make([]byte, length)
	t.allocated[&buf] = struct{}{}
	return &buf
}

func (t *testPool) Put(buf *[]byte) {
	if _, ok := t.allocated[buf]; !ok {
		panic("unexpected put")
	}
	delete(t.allocated, buf)
}

func (s) TestBufferSlice_Iteration(t *testing.T) {
	tests := []struct {
		name       string
		buffers    [][]byte
		operations func(t *testing.T, c *mem.Reader)
	}{
		{
			name: "empty",
			operations: func(t *testing.T, r *mem.Reader) {
				if r.Remaining() != 0 {
					t.Fatalf("Remaining() = %v, want 0", r.Remaining())
				}
				_, err := r.Peek(1, nil)
				if err == nil {
					t.Fatalf("Peek(1) returned error <nil>, want non-nil")
				}
				discarded, err := r.Discard(1)
				if got, want := discarded, 0; got != want {
					t.Fatalf("Discard(1) = %d, want %d", got, want)
				}
				if err == nil {
					t.Fatalf("Discard(1) returned error <nil>, want non-nil")
				}
				if r.Remaining() != 0 {
					t.Fatalf("Remaining() after Discard = %v, want 0", r.Remaining())
				}
			},
		},
		{
			name:    "single_buffer",
Repository: grpc/grpc-go
File: internal/transport/http_util.go (Go)

Code before:
	if endStream {
		flags = http2.FlagDataEndStream
	}
	length := uint32(0)
	for _, d := range data {
		length += uint32(len(d))
	}
	// TODO: Replace the header write with the framer API being added in
	// https://github.com/golang/go/issues/66655.
	f.headerBuf = append(f.headerBuf[:0],
		byte(length>>16),
		byte(length>>8),
		byte(length),
		byte(http2.FrameData),
		byte(flags),
		byte(streamID>>24),
		byte(streamID>>16),
		byte(streamID>>8),
		byte(streamID))
	if _, err := f.writer.Write(f.headerBuf); err != nil {
		return err
	}
	for _, d := range data {
		if len(d) == 0 {
			continue
		}
		if _, err := f.writer.Write(d); err != nil {
			return err
		}
	}
	return nil
}

func getWriteBufferPool(size int) *sync.Pool {
	writeBufferMutex.Lock()
	defer writeBufferMutex.Unlock()
	pool, ok := writeBufferPoolMap[size]
	if ok {
		return pool
	}
	pool = &sync.Pool{
		New: func() any {
			b := make([]byte, size)
			return &b
		},
	}
	writeBufferPoolMap[size] = pool
	return pool
}

// ParseDialTarget returns the network and address to pass to dialer.

Review comment (bug):
Would this be a bug if it were zero?  I would have expected it to be.

If it is, then we should delete it.  `Write` should handle a zero-length buffer as a nop already anyway.

Code after:
	w := newBufWriter(conn, writeBufferSize, pool)
	f := &framer{
		writer: w,
		fr:     http2.NewFramer(w, r),
		reader: r,
		pool:   memPool,
	}
	f.fr.SetMaxReadFrameSize(http2MaxFrameLen)
	// Opt-in to Frame reuse API on framer to reduce garbage.
	// Frames aren't safe to read from after a subsequent call to ReadFrame.
	f.fr.SetReuseFrames()
	f.fr.MaxHeaderListSize = maxHeaderListSize
	f.fr.ReadMetaHeaders = hpack.NewDecoder(http2InitHeaderTableSize, nil)
	return f
}

// writeData writes a DATA frame.
//
// It is the caller's responsibility not to violate the maximum frame size.
func (f *framer) writeData(streamID uint32, endStream bool, data [][]byte) error {
	var flags http2.Flags
	if endStream {
		flags = http2.FlagDataEndStream
	}
	length := uint32(0)
	for _, d := range data {
		length += uint32(len(d))
	}
	// TODO: Replace the header write with the framer API being added in
	// https://github.com/golang/go/issues/66655.
	f.headerBuf = append(f.headerBuf[:0],
		byte(length>>16),
		byte(length>>8),
		byte(length),
		byte(http2.FrameData),
		byte(flags),
		byte(streamID>>24),
		byte(streamID>>16),
		byte(streamID>>8),
		byte(streamID))
	if _, err := f.writer.Write(f.headerBuf); err != nil {
		return err
	}
	for _, d := range data {
		if _, err := f.writer.Write(d); err != nil {
			return err
		}
	}
	return nil
}
Repository: grpc/grpc-go
File: internal/xds/balancer/cdsbalancer/e2e_test/dns_impl_test.go (Go)

Code before:
                        "cluster": "test-cluster",
                        "type": "LOGICAL_DNS",
                        "dnsHostname": "%s:///target-name",
                        "outlierDetection": {}
                    }],
                    "xdsLbPolicy":[{"round_robin":{}}]
                }
            }]
        }`, dnsScheme)

	scpr := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(jsonSC)
	r.InitialState(xdsclient.SetClient(resolver.State{ServiceConfig: scpr}, xdsC))

	// Create a ClientConn and make a successful RPC.
	cc, err := grpc.NewClient(r.Scheme()+":///test.service",
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithResolvers(r),
	)
	cc.Connect()
	if err != nil {
		t.Fatalf("failed to create new client for local test server: %v", err)
	}
	defer cc.Close()

	testClient := testpb.NewTestServiceClient(cc)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var dnsClientConn resolver.ClientConn
	select {
	case dnsClientConn = <-dnsCCCh:
	case <-ctx.Done():
		t.Fatal("Timeout waiting for cluster_resolver to build the DNS resolver")
	}

	// For LOGICAL_DNS, this updates the SINGLE endpoint to have 2 IPs.
	dnsClientConn.UpdateState(resolver.State{
		Addresses: []resolver.Address{
			{Addr: server1.Address},
			{Addr: server2.Address},
		},
	})

	// Ensure the RPC is routed to the first backend.
	var peer peer.Peer
	if _, err := testClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
		t.Fatalf("RPC failed: %v", err)
	}

	if got, want := peer.Addr.String(), server1.Address; got != want {
		t.Errorf("peer.Addr = %q, want = %q", got, want)

Review comment (suggestion):
There is already a `defaultTestTimeout` defined in this package. Please use that instead of hardcoding it to `5s` here.

Code after:
			Policy:      e2e.LoadBalancingPolicyRoundRobin,
		})},
		Endpoints: nil,
	}

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

	if err := managementServer.Update(ctx, resources); err != nil {
		t.Fatalf("Failed to update management server: %v", err)
	}

	// Create a ClientConn and make a successful RPC.
	cc, err := grpc.NewClient(fmt.Sprintf("xds:///"+serviceName), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(resolverBuilder))
	if err != nil {
		t.Fatalf("Failed to create new client for local test server: %v", err)
	}
	defer cc.Close()

	// Ensure the RPC is routed to the first backend.
	testClient := testpb.NewTestServiceClient(cc)
	for i := 0; i < 10; i++ {
		var peer peer.Peer
		if _, err := testClient.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
			t.Fatalf("RPC failed: %v", err)
		}

		if got, want := peer.Addr.String(), server1.Address; got != want {
			t.Errorf("peer.Addr = %q, want = %q", got, want)
		}
	}
}
Repository: grpc/grpc-go
File: internal/transport/http2_client.go (Go)

Code before:
			if !timer.Stop() {
				<-timer.C
			}
			return
		}
	}
}

func (t *http2Client) Error() <-chan struct{} {
	return t.ctx.Done()
}

func (t *http2Client) GoAway() <-chan struct{} {
	return t.goAway
}

func (t *http2Client) socketMetrics() *channelz.EphemeralSocketMetrics {
	return &channelz.EphemeralSocketMetrics{
		LocalFlowControlWindow:  int64(t.fc.getSize()),
		RemoteFlowControlWindow: t.getOutFlowWindow(),
	}
}

func (t *http2Client) RemoteAddr() net.Addr { return t.remoteAddr }

func (t *http2Client) AuthInfo() credentials.AuthInfo { return t.authInfo }

func (t *http2Client) incrMsgSent() {
	if channelz.IsOn() {
		t.channelz.SocketMetrics.MessagesSent.Add(1)
		t.channelz.SocketMetrics.LastMessageSentTimestamp.Store(time.Now().UnixNano())
	}
}

func (t *http2Client) incrMsgRecv() {
	if channelz.IsOn() {
		t.channelz.SocketMetrics.MessagesReceived.Add(1)
		t.channelz.SocketMetrics.LastMessageReceivedTimestamp.Store(time.Now().UnixNano())
	}
}

func (t *http2Client) getOutFlowWindow() int64 {
	resp := make(chan uint32, 1)
	timer := time.NewTimer(time.Second)
	defer timer.Stop()
	t.controlBuf.put(&outFlowControlSizeRequest{resp})
	select {
	case sz := <-resp:
		return int64(sz)
	case <-t.ctxDone:
		return -1

Review comment (security):
Both, remoteAddr and authInfo, are available in the `Peer`.

I'm thinking we should get rid of `RemoteAddr` (which has a single caller at this point in time), and instead of adding `AuthInfo`, we could add a `Peer` method that returns `peer.Peer`. Wdyt?

Code after:
			if !timer.Stop() {
				<-timer.C
			}
			return
		}
	}
}

func (t *http2Client) Error() <-chan struct{} {
	return t.ctx.Done()
}

func (t *http2Client) GoAway() <-chan struct{} {
	return t.goAway
}

func (t *http2Client) socketMetrics() *channelz.EphemeralSocketMetrics {
	return &channelz.EphemeralSocketMetrics{
		LocalFlowControlWindow:  int64(t.fc.getSize()),
		RemoteFlowControlWindow: t.getOutFlowWindow(),
	}
}

func (t *http2Client) incrMsgSent() {
	if channelz.IsOn() {
		t.channelz.SocketMetrics.MessagesSent.Add(1)
		t.channelz.SocketMetrics.LastMessageSentTimestamp.Store(time.Now().UnixNano())
	}
}

func (t *http2Client) incrMsgRecv() {
	if channelz.IsOn() {
		t.channelz.SocketMetrics.MessagesReceived.Add(1)
		t.channelz.SocketMetrics.LastMessageReceivedTimestamp.Store(time.Now().UnixNano())
	}
}

func (t *http2Client) getOutFlowWindow() int64 {
	resp := make(chan uint32, 1)
	timer := time.NewTimer(time.Second)
	defer timer.Stop()
	t.controlBuf.put(&outFlowControlSizeRequest{resp})
	select {
	case sz := <-resp:
		return int64(sz)
	case <-t.ctxDone:
		return -1
	case <-timer.C:
		return -2
	}
}
Repository: grpc/grpc-go
File: internal/xds/balancer/clusterimpl/picker.go (Go)

Code before:
		}
	}

	if d.counter != nil {
		// Update Done() so that when the RPC finishes, the request count will
		// be released.
		oldDone := pr.Done
		pr.Done = func(doneInfo balancer.DoneInfo) {
			d.counter.EndRequest()
			if oldDone != nil {
				oldDone(doneInfo)
			}
		}
	}

	return pr, err
}

// autoHostRewriteKey is the context key used to store the value of
// route's autoHostRewrite in the RPC context.
type autoHostRewriteKey struct{}

// autoHostRewrite retrieves the autoHostRewrite value from the provided context.
func autoHostRewrite(ctx context.Context) bool {
	autohostRewrite, _ := ctx.Value(autoHostRewriteKey{}).(bool)
	return autohostRewrite
}

// AutoHostRewriteForTesting returns the value of autoHostRewrite field;
// to be used for testing only.
func AutoHostRewriteForTesting(ctx context.Context) bool {
	return autoHostRewrite(ctx)
}

// SetAutoHostRewrite adds the autoHostRewrite value to the context for
// the xds_cluster_impl LB policy to pick.
func SetAutoHostRewrite(ctx context.Context, autohostRewrite bool) context.Context {
	return context.WithValue(ctx, autoHostRewriteKey{}, autohostRewrite)
}


Review comment (suggestion):
Nit: maybe rename the variable with a one letter name. Something like `v`?

See: go/go-style/decisions#variable-names

TL;Dr
The general rule of thumb is that the length of a name should be proportional to the size of its scope and inversely proportional to the number of times that it is used within that scope

Code after:
		}
	}

	if d.counter != nil {
		// Update Done() so that when the RPC finishes, the request count will
		// be released.
		oldDone := pr.Done
		pr.Done = func(doneInfo balancer.DoneInfo) {
			d.counter.EndRequest()
			if oldDone != nil {
				oldDone(doneInfo)
			}
		}
	}

	return pr, err
}

// autoHostRewriteKey is the context key used to store the value of
// route's autoHostRewrite in the RPC context.
type autoHostRewriteKey struct{}

// autoHostRewrite retrieves the autoHostRewrite value from the provided context.
func autoHostRewrite(ctx context.Context) bool {
	v, _ := ctx.Value(autoHostRewriteKey{}).(bool)
	return v
}

// AutoHostRewriteForTesting returns the value of autoHostRewrite field;
// to be used for testing only.
func AutoHostRewriteForTesting(ctx context.Context) bool {
	return autoHostRewrite(ctx)
}

// SetAutoHostRewrite adds the autoHostRewrite value to the context for
// the xds_cluster_impl LB policy to pick.
func SetAutoHostRewrite(ctx context.Context, autohostRewrite bool) context.Context {
	return context.WithValue(ctx, autoHostRewriteKey{}, autohostRewrite)
}