Skip to content

Commit

Permalink
feat: support load balance for streaming connection creation (go-krat…
Browse files Browse the repository at this point in the history
…os#2669)

* feat: support load balance for streaming connection creation

* fix lint
  • Loading branch information
shenqidebaozi committed Feb 27, 2023
1 parent 19f008b commit 834b781
Showing 1 changed file with 23 additions and 3 deletions.
26 changes: 23 additions & 3 deletions transport/grpc/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -146,16 +146,22 @@ func dial(ctx context.Context, insecure bool, opts ...ClientOption) (*grpc.Clien
ints := []grpc.UnaryClientInterceptor{
unaryClientInterceptor(options.middleware, options.timeout, options.filters),
}
sints := []grpc.StreamClientInterceptor{
streamClientInterceptor(options.filters),
}

if len(options.ints) > 0 {
ints = append(ints, options.ints...)
}
if len(options.streamInts) > 0 {
sints = append(sints, options.streamInts...)
}
grpcOpts := []grpc.DialOption{
grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, options.balancerName)),
grpc.WithChainUnaryInterceptor(ints...),
grpc.WithChainStreamInterceptor(sints...),
}
if len(options.streamInts) > 0 {
grpcOpts = append(grpcOpts, grpc.WithChainStreamInterceptor(options.streamInts...))
}

if options.discovery != nil {
grpcOpts = append(grpcOpts,
grpc.WithResolvers(
Expand Down Expand Up @@ -211,3 +217,17 @@ func unaryClientInterceptor(ms []middleware.Middleware, timeout time.Duration, f
return err
}
}

func streamClientInterceptor(filters []selector.NodeFilter) grpc.StreamClientInterceptor {
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { // nolint
ctx = transport.NewClientContext(ctx, &Transport{
endpoint: cc.Target(),
operation: method,
reqHeader: headerCarrier{},
nodeFilters: filters,
})
var p selector.Peer
ctx = selector.NewPeerContext(ctx, &p)
return streamer(ctx, desc, cc, method, opts...)
}
}

0 comments on commit 834b781

Please sign in to comment.