Skip to content

Commit

Permalink
populate initReq target name and fix IP []byte type in grpclb (#1145)
Browse files Browse the repository at this point in the history
  • Loading branch information
menghanl committed Mar 28, 2017
1 parent b3cc2b5 commit 7fc29d0
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
9 changes: 7 additions & 2 deletions grpclb/grpclb.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ import (
"errors"
"fmt"
"math/rand"
"net"
"sync"
"time"

Expand Down Expand Up @@ -96,6 +97,7 @@ type addrInfo struct {

type balancer struct {
r naming.Resolver
target string
mu sync.Mutex
seq int // a sequence number to make sure addrCh does not get stale addresses.
w naming.Watcher
Expand Down Expand Up @@ -224,7 +226,7 @@ func (b *balancer) processServerList(l *lbpb.ServerList, seq int) {
for _, s := range servers {
md := metadata.Pairs("lb-token", s.LoadBalanceToken)
addr := grpc.Address{
Addr: fmt.Sprintf("%s:%d", s.IpAddress, s.Port),
Addr: fmt.Sprintf("%s:%d", net.IP(s.IpAddress), s.Port),
Metadata: &md,
}
sl = append(sl, &addrInfo{
Expand Down Expand Up @@ -272,7 +274,9 @@ func (b *balancer) callRemoteBalancer(lbc lbpb.LoadBalancerClient, seq int) (ret
b.mu.Unlock()
initReq := &lbpb.LoadBalanceRequest{
LoadBalanceRequestType: &lbpb.LoadBalanceRequest_InitialRequest{
InitialRequest: new(lbpb.InitialLoadBalanceRequest),
InitialRequest: &lbpb.InitialLoadBalanceRequest{
Name: b.target,
},
},
}
if err := stream.Send(initReq); err != nil {
Expand Down Expand Up @@ -322,6 +326,7 @@ func (b *balancer) Start(target string, config grpc.BalancerConfig) error {
if b.r == nil {
return errors.New("there is no name resolver installed")
}
b.target = target
b.mu.Lock()
if b.done {
b.mu.Unlock()
Expand Down
18 changes: 13 additions & 5 deletions grpclb/grpclb_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,14 @@ func (b *remoteBalancer) stop() {
}

func (b *remoteBalancer) BalanceLoad(stream lbpb.LoadBalancer_BalanceLoadServer) error {
req, err := stream.Recv()
if err != nil {
return err
}
initReq := req.GetInitialRequest()
if initReq.Name != besn {
return grpc.Errorf(codes.InvalidArgument, "invalid service name: %v", initReq.Name)
}
resp := &lbpb.LoadBalanceResponse{
LoadBalanceResponseType: &lbpb.LoadBalanceResponse_InitialResponse{
InitialResponse: new(lbpb.InitialLoadBalanceResponse),
Expand Down Expand Up @@ -264,7 +272,7 @@ func TestGRPCLB(t *testing.T) {
t.Fatalf("Failed to generate the port number %v", err)
}
be := &lbpb.Server{
IpAddress: []byte(beAddr[0]),
IpAddress: beLis.Addr().(*net.TCPAddr).IP,
Port: int32(bePort),
LoadBalanceToken: lbToken,
}
Expand Down Expand Up @@ -334,12 +342,12 @@ func TestDropRequest(t *testing.T) {
}
sls := []*lbpb.ServerList{{
Servers: []*lbpb.Server{{
IpAddress: []byte(beAddr1[0]),
IpAddress: beLis1.Addr().(*net.TCPAddr).IP,
Port: int32(bePort1),
LoadBalanceToken: lbToken,
DropRequest: true,
}, {
IpAddress: []byte(beAddr2[0]),
IpAddress: beLis2.Addr().(*net.TCPAddr).IP,
Port: int32(bePort2),
LoadBalanceToken: lbToken,
DropRequest: false,
Expand Down Expand Up @@ -410,7 +418,7 @@ func TestDropRequestFailedNonFailFast(t *testing.T) {
t.Fatalf("Failed to generate the port number %v", err)
}
be := &lbpb.Server{
IpAddress: []byte(beAddr[0]),
IpAddress: beLis.Addr().(*net.TCPAddr).IP,
Port: int32(bePort),
LoadBalanceToken: lbToken,
DropRequest: true,
Expand Down Expand Up @@ -473,7 +481,7 @@ func TestServerExpiration(t *testing.T) {
t.Fatalf("Failed to generate the port number %v", err)
}
be := &lbpb.Server{
IpAddress: []byte(beAddr[0]),
IpAddress: beLis.Addr().(*net.TCPAddr).IP,
Port: int32(bePort),
LoadBalanceToken: lbToken,
}
Expand Down

0 comments on commit 7fc29d0

Please sign in to comment.