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

Code before:
 * 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.
package internal

var (
	// SetRandShuffleForTesting pseudo-randomizes the order of addresses. n
	// is the number of elements.  swap swaps the elements with indexes i and j.
	// It returns the shuffle func before it was replaced.
	SetRandShuffleForTesting func(sf func(n int, swap func(i, j int)))

	// RevertRandShuffleFunc sets the real shuffle function back after testing.
	RevertRandShuffleFunc func()
)


Review comment (suggestion):
Alternate proposal:

```go
package internal

var RandShuffle = rand.Shuffle

package pickfirst

// when it's time to shuffle, call internal.RandShuffle

package pickfirst_test

// set internal.RandShuffle directly
```

I think this would be a bit simpler for everyone to use.

Code after:
/*
 * Copyright 2024 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

// Package internal contains code internal to the pickfirst package.
package internal

import "math/rand"

// RandShuffle pseudo-randomizes the order of addresses.
var RandShuffle = rand.Shuffle
Repository: grpc/grpc-go
File: balancer/grpclb/grpclb_picker.go (Go)

Code before:
/*
 *
 * Copyright 2017 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package grpclb

import (
	"math/rand/v2"
	"sync"
	"sync/atomic"

	"google.golang.org/grpc/balancer"
	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// rpcStats is same as lbpb.ClientStats, except that numCallsDropped is a map
// instead of a slice.
type rpcStats struct {
	// Only access the following fields atomically.
	numCallsStarted                        int64
	numCallsFinished                       int64
	numCallsFinishedWithClientFailedToSend int64
	numCallsFinishedKnownReceived          int64

	mu sync.Mutex
	// map load_balance_token -> num_calls_dropped
	numCallsDropped map[string]int64
}

func newRPCStats() *rpcStats {
	return &rpcStats{

Review comment (suggestion):
This is going to be a little painful. But could you please rename this import as `rand "math/rand/v2"`, so that we can track the package name from use sites. Thank you.

Code after:
/*
 *
 * Copyright 2017 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package grpclb

import (
	rand "math/rand/v2"
	"sync"
	"sync/atomic"

	"google.golang.org/grpc/balancer"
	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// rpcStats is same as lbpb.ClientStats, except that numCallsDropped is a map
// instead of a slice.
type rpcStats struct {
	// Only access the following fields atomically.
	numCallsStarted                        int64
	numCallsFinished                       int64
	numCallsFinishedWithClientFailedToSend int64
	numCallsFinishedKnownReceived          int64

	mu sync.Mutex
	// map load_balance_token -> num_calls_dropped
	numCallsDropped map[string]int64
}

func newRPCStats() *rpcStats {
	return &rpcStats{
Repository: grpc/grpc-go
File: xds/internal/httpfilter/fault/fault_test.go (Go)

Code before:
	if err != nil {
		t.Fatalf("testutils.LocalTCPListener() failed: %v", err)
	}

	// Initialize a test gRPC server, assign it to the stub server, and start the test service.
	stub := &stubserver.StubServer{
		Listener: lis,
		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
			return &testpb.Empty{}, nil
		},
		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
			// End RPC after client does a CloseSend.
			for {
				if _, err := stream.Recv(); err == io.EOF {
					return nil
				} else if err != nil {
					return err
				}
			}
		},
	}

	stubserver.StartTestService(t, stub)
	t.Cleanup(func() {
		stub.S.Stop()
	})
	return managementServer, nodeID, uint32(lis.Addr().(*net.TCPAddr).Port)
}

func (s) TestFaultInjection_Unary(t *testing.T) {
	type subcase struct {
		name   string
		code   codes.Code
		repeat int
		randIn []int           // Intn calls per-repeat (not per-subcase)
		delays []time.Duration // NewTimer calls per-repeat (not per-subcase)
		md     metadata.MD
	}
	testCases := []struct {
		name       string
		cfgs       []*fpb.HTTPFault
		randOutInc int
		want       []subcase
	}{{
		name: "max faults zero",
		cfgs: []*fpb.HTTPFault{{
			MaxActiveFaults: wrapperspb.UInt32(0),
			Abort: &fpb.FaultAbort{
				Percentage: &tpb.FractionalPercent{Numerator: 100, Denominator: tpb.FractionalPercent_HUNDRED},
				ErrorType:  &fpb.FaultAbort_GrpcStatus{GrpcStatus: uint32(codes.Aborted)},
			},

Review comment (suggestion):
This should also work 

```suggestion
	t.Cleanup(stub.S.Stop)
```

Code after:
	if err != nil {
		t.Fatalf("testutils.LocalTCPListener() failed: %v", err)
	}

	// Initialize a test gRPC server, assign it to the stub server, and start the test service.
	stub := &stubserver.StubServer{
		Listener: lis,
		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
			return &testpb.Empty{}, nil
		},
		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
			// End RPC after client does a CloseSend.
			for {
				if _, err := stream.Recv(); err == io.EOF {
					return nil
				} else if err != nil {
					return err
				}
			}
		},
	}

	stubserver.StartTestService(t, stub)
	t.Cleanup(stub.S.Stop)
	return managementServer, nodeID, uint32(lis.Addr().(*net.TCPAddr).Port)
}

func (s) TestFaultInjection_Unary(t *testing.T) {
	type subcase struct {
		name   string
		code   codes.Code
		repeat int
		randIn []int           // Intn calls per-repeat (not per-subcase)
		delays []time.Duration // NewTimer calls per-repeat (not per-subcase)
		md     metadata.MD
	}
	testCases := []struct {
		name       string
		cfgs       []*fpb.HTTPFault
		randOutInc int
		want       []subcase
	}{{
		name: "max faults zero",
		cfgs: []*fpb.HTTPFault{{
			MaxActiveFaults: wrapperspb.UInt32(0),
			Abort: &fpb.FaultAbort{
				Percentage: &tpb.FractionalPercent{Numerator: 100, Denominator: tpb.FractionalPercent_HUNDRED},
				ErrorType:  &fpb.FaultAbort_GrpcStatus{GrpcStatus: uint32(codes.Aborted)},
			},
		}},
		randOutInc: 5,
Repository: grpc/grpc-go
File: credentials/tls_ext_test.go (Go)

Code before:
		panic(fmt.Sprintf("Error reading CA cert file: %v", err))
	}
	certPool = x509.NewCertPool()
	if !certPool.AppendCertsFromPEM(b) {
		panic("Error appending cert from PEM")
	}
}

// Tests that the MinVersion of tls.Config is set to 1.2 if it is not already
// set by the user.
func (s) TestTLS_MinVersion12(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	serverTLS := &tls.Config{
		// MinVersion should be set to 1.2 by gRPC by default.
		Certificates: []tls.Certificate{serverCert},
	}

	testCases := []struct {
		name      string
		serverTLS func() *tls.Config
	}{
		{
			name:      "base_case",
			serverTLS: func() *tls.Config { return serverTLS },
		},
		{
			name: "fallback_to_base",
			serverTLS: func() *tls.Config {
				config := serverTLS.Clone()
				config.GetConfigForClient = func(*tls.ClientHelloInfo) (*tls.Config, error) {
					return nil, nil
				}
				return config
			},
		},
		{
			name: "dynamic_using_get_config_for_client",
			serverTLS: func() *tls.Config {
				return &tls.Config{
					GetConfigForClient: func(*tls.ClientHelloInfo) (*tls.Config, error) {
						return serverTLS, nil
					},
				}
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {

Review comment (suggestion):
Structure question - why define a `func() *tls.Config` instead of just having the field be `tls.Config`? It seems like it's just adding a layer of complication but doesn't add anything?

Same question throughout

Code after:
		panic(fmt.Sprintf("Error reading CA cert file: %v", err))
	}
	certPool = x509.NewCertPool()
	if !certPool.AppendCertsFromPEM(b) {
		panic("Error appending cert from PEM")
	}
}

// Tests that the MinVersion of tls.Config is set to 1.2 if it is not already
// set by the user.
func (s) TestTLS_MinVersion12(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	testCases := []struct {
		name      string
		serverTLS func() *tls.Config
	}{
		{
			name: "base_case",
			serverTLS: func() *tls.Config {
				return &tls.Config{
					// MinVersion should be set to 1.2 by gRPC by default.
					Certificates: []tls.Certificate{serverCert},
				}
			},
		},
		{
			name: "fallback_to_base",
			serverTLS: func() *tls.Config {
				config := &tls.Config{
					// MinVersion should be set to 1.2 by gRPC by default.
					Certificates: []tls.Certificate{serverCert},
				}
				config.GetConfigForClient = func(*tls.ClientHelloInfo) (*tls.Config, error) {
					return nil, nil
				}
				return config
			},
		},
		{
			name: "dynamic_using_get_config_for_client",
			serverTLS: func() *tls.Config {
				return &tls.Config{
					GetConfigForClient: func(*tls.ClientHelloInfo) (*tls.Config, error) {
						return &tls.Config{
							// MinVersion should be set to 1.2 by gRPC by default.
							Certificates: []tls.Certificate{serverCert},
						}, nil
					},
				}
Repository: grpc/grpc-go
File: balancer/endpointsharding/endpointsharding.go (Go)

Code before:
		return nil
	}
	return p.childStates
}

// balancerWrapper is a wrapper of a balancer. It ID's a child balancer by
// endpoint, and persists recent child balancer state.
type balancerWrapper struct {
	balancer.Balancer   // Simply forward balancer.Balancer operations.
	balancer.ClientConn // embed to intercept UpdateState, doesn't deal with SubConns

	es *endpointSharding

	childState ChildState
}

func (bw *balancerWrapper) UpdateState(state balancer.State) {
	bw.es.mu.Lock()
	bw.childState.State = state
	bw.es.mu.Unlock()
	// When a child balancer says it's IDLE, ping it to exit idle and reconnect.
	// TODO: In the future, perhaps make this a knob in configuration.
	if ei, ok := bw.Balancer.(balancer.ExitIdler); state.ConnectivityState == connectivity.Idle && ok {
		go func() {
			bw.es.childMu.Lock()
			ei.ExitIdle()
			bw.es.childMu.Unlock()
		}()
	}
	bw.es.updateState()
}

// ParseConfig parses a child config list and returns an LB config to use with
// the endpointsharding balancer.
//
// cfg is expected to be a JSON array of LB policy names + configs as the
// format of the loadBalancingConfig field in ServiceConfig.
func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
	return gracefulswitch.ParseConfig(cfg)
}

// PickFirstConfig is a pick first config without shuffling enabled.
const PickFirstConfig = "[{\"pick_first\": {}}]"


Review comment (bug):
Don't we need to avoid this racing with Close?  Not worried about a deadlock or concurrency, but this sequence:

- Child reports IDLE
- This goroutine gets scheduled
- Parent calls `Close()`
- This goroutine runs and tells the child to `ExitIdle()` after `Close()`

Code after:
	balancer.Balancer   // Simply forward balancer.Balancer operations.
	balancer.ClientConn // embed to intercept UpdateState, doesn't deal with SubConns

	es *endpointSharding

	childState ChildState
}

func (bw *balancerWrapper) UpdateState(state balancer.State) {
	bw.es.mu.Lock()
	bw.childState.State = state
	bw.es.mu.Unlock()
	// When a child balancer says it's IDLE, ping it to exit idle and reconnect.
	// TODO: In the future, perhaps make this a knob in configuration.
	if ei, ok := bw.Balancer.(balancer.ExitIdler); state.ConnectivityState == connectivity.Idle && ok {
		go func() {
			bw.es.childMu.Lock()
			if !bw.es.closed {
				ei.ExitIdle()
			}
			bw.es.childMu.Unlock()
		}()
	}
	bw.es.updateState()
}

// ParseConfig parses a child config list and returns an LB config to use with
// the endpointsharding balancer.
//
// cfg is expected to be a JSON array of LB policy names + configs as the
// format of the loadBalancingConfig field in ServiceConfig.
func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
	return gracefulswitch.ParseConfig(cfg)
}

// PickFirstConfig is a pick first config without shuffling enabled.
const PickFirstConfig = "[{\"pick_first\": {}}]"
Repository: grpc/grpc-go
File: internal/hierarchy/hierarchy.go (Go)

Code before:
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
}

// GetEndpoint returns the hierarchical path of endpoint.
func GetEndpoint(endpoint resolver.Endpoint) []string {
	attrs := endpoint.Attributes
	if attrs == nil {
		return nil
	}
	path, _ := attrs.Value(pathKey).(pathValue)
	return ([]string)(path)
}

// SetEndpoint overrides the hierarchical path in endpoint with path.
func SetEndpoint(endpoint resolver.Endpoint, path []string) resolver.Endpoint {
	endpoint.Attributes = endpoint.Attributes.WithValue(pathKey, pathValue(path))
	return endpoint
}

// Get returns the hierarchical path of addr.
func Get(addr resolver.Address) []string {
	attrs := addr.BalancerAttributes
	if attrs == nil {
		return nil
	}
	path, _ := attrs.Value(pathKey).(pathValue)
	return ([]string)(path)
}

// Set overrides the hierarchical path in addr with path.
func Set(addr resolver.Address, path []string) resolver.Address {
	addr.BalancerAttributes = addr.BalancerAttributes.WithValue(pathKey, pathValue(path))

Review comment (suggestion):
I think this function can just be:

```go
path, _ := endpoint.Attributes.Value(pathKey).(pathValue)
return path
```


Code after:
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
}

// Get returns the hierarchical path of addr.
func Get(addr resolver.Address) []string {
	attrs := addr.BalancerAttributes
	if attrs == nil {
		return nil
	}
	path, _ := attrs.Value(pathKey).(pathValue)
	return ([]string)(path)
}

// Set overrides the hierarchical path in addr with path.
func Set(addr resolver.Address, path []string) resolver.Address {
	addr.BalancerAttributes = addr.BalancerAttributes.WithValue(pathKey, pathValue(path))
	return addr
}

// Group splits a slice of addresses into groups based on
Repository: grpc/grpc-go
File: internal/hierarchy/hierarchy.go (Go)

Code before:
//	    {addr1, path: [wt1]},
//	  ],
//	  p1: [
//	    {addr2, path: [wt2]},
//	    {addr3, path: [wt3]},
//	  ],
//	}
//
// If hierarchical path is not set, or has no path in it, the address is
// dropped.
func Group(addrs []resolver.Address) map[string][]resolver.Address {
	ret := make(map[string][]resolver.Address)
	for _, addr := range addrs {
		oldPath := Get(addr)
		if len(oldPath) == 0 {
			continue
		}
		curPath := oldPath[0]
		newPath := oldPath[1:]
		newAddr := Set(addr, newPath)
		ret[curPath] = append(ret[curPath], newAddr)
	}
	return ret
}

// GroupEndpoints groups endpoints as per Group above.
func GroupEndpoints(endpoints []resolver.Endpoint) map[string][]resolver.Endpoint {
	ret := make(map[string][]resolver.Endpoint)
	for _, endpoint := range endpoints {
		oldPath := GetEndpoint(endpoint)
		if len(oldPath) == 0 {
			continue
		}
		curPath := oldPath[0]
		newPath := oldPath[1:]
		newEndpoint := SetEndpoint(endpoint, newPath)
		ret[curPath] = append(ret[curPath], newEndpoint)
	}
	return ret
}


Review comment (suggestion):
Pet peeve: please don't use "as per" in our docs. :)  Just "as" or "per" is probably what you want instead.

Except in this case, we want to just copy/paste the whole thing and update as appropriate, because we'll be deleting the other one.

Code after:
//	    {addr3, path: [wt3]},
//	  ],
//	}
//
// If hierarchical path is not set, or has no path in it, the address is
// dropped.
func Group(addrs []resolver.Address) map[string][]resolver.Address {
	ret := make(map[string][]resolver.Address)
	for _, addr := range addrs {
		oldPath := Get(addr)
		if len(oldPath) == 0 {
			continue
		}
		curPath := oldPath[0]
		newPath := oldPath[1:]
		newAddr := Set(addr, newPath)
		ret[curPath] = append(ret[curPath], newAddr)
	}
	return ret
}

// GroupEndpoints 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]},
//	  ],
//	}
//
Repository: grpc/grpc-go
File: examples/features/gracefulstop/client/main.go (Go)

Code before:
/*
 *
 * 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.
 */

// Binary client demonstrates sending multiple requests to server and observe
// graceful stop.
package main

import (
	"context"
	"flag"
	"log"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	pb "google.golang.org/grpc/examples/helloworld/helloworld"
	"google.golang.org/grpc/status"
)

var addr = flag.String("addr", "localhost:50052", "the address to connect to")

func main() {
	flag.Parse()

	// Set up a connection to the server.
	conn, err := grpc.NewClient(*addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatalf("Failed to connect: %v", err)
	}

Review comment (bug):
Is there some way this can actually prove that the stop was graceful?  E.g. by starting a stream that runs until requests start failing, and then closing the stream from the client and having it exit successfully by getting a response from the server at that time?

E.g. the server could respond with a message indicating how many requests it processed.

Code after:
/*
 *
 * 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.
 */

// Binary client demonstrates sending multiple requests to server and observe
// graceful stop.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	pb "google.golang.org/grpc/examples/features/proto/echo"
)

var addr = flag.String("addr", "localhost:50052", "the address to connect to")

func main() {
	flag.Parse()

	conn, err := grpc.NewClient(*addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatalf("Failed to create new client: %v", err)
	}
	defer conn.Close()
	c := pb.NewEchoClient(conn)