feat: streaming least-conn loadbalance

Author: DMwangnima

client/client.go

@@ -283,12 +283,14 @@ func (kc *kClient) initMiddlewares(ctx context.Context) {
 	kc.mws = append(kc.mws, builderMWs...)
 	kc.mws = append(kc.mws, acl.NewACLMiddleware(kc.opt.ACLRules))
 	if kc.opt.Proxy == nil {
-		kc.mws = append(kc.mws, newResolveMWBuilder(kc.lbf)(ctx))
+		cs, _ := kc.opt.RemoteOpt.ConnPool.(remote.ConnStatistics)
+		kc.mws = append(kc.mws, newResolveMWBuilder(kc.lbf, cs)(ctx))
 		kc.mws = append(kc.mws, kc.opt.CBSuite.InstanceCBMW())
 		kc.mws = append(kc.mws, richMWsWithBuilder(ctx, kc.opt.IMWBs)...)
 	} else {
+		cs, _ := kc.opt.RemoteOpt.ConnPool.(remote.ConnStatistics)
 		if kc.opt.Resolver != nil { // customized service discovery
-			kc.mws = append(kc.mws, newResolveMWBuilder(kc.lbf)(ctx))
+			kc.mws = append(kc.mws, newResolveMWBuilder(kc.lbf, cs)(ctx))
 		}
 		kc.mws = append(kc.mws, newProxyMW(kc.opt.Proxy))
 	}

client/client_test.go

@@ -1138,6 +1138,9 @@ func Test_kClient_initStreamMiddlewares(t *testing.T) {
 		opt: &client.Options{
 			TracerCtl: ctl,
 			Streaming: internal_stream.StreamingConfig{},
+			RemoteOpt: &remote.ClientOption{
+				ConnPool: nil,
+			},
 		},
 	}
 	c.initStreamMiddlewares(context.Background())

client/middlewares.go

@@ -36,6 +36,7 @@ import (
 	"github.com/cloudwego/kitex/pkg/remote/codec/protobuf"
 	"github.com/cloudwego/kitex/pkg/rpcinfo"
 	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
+	"github.com/cloudwego/kitex/transport"
 )
 
 const maxRetry = 6
@@ -83,7 +84,7 @@ func discoveryEventHandler(name string, bus event.Bus, queue event.Queue) func(d
 // newResolveMWBuilder creates a middleware for service discovery.
 // This middleware selects an appropriate instance based on the resolver and loadbalancer given.
 // If retryable error is encountered, it will retry until timeout or an unretryable error is returned.
-func newResolveMWBuilder(lbf *lbcache.BalancerFactory) endpoint.MiddlewareBuilder {
+func newResolveMWBuilder(lbf *lbcache.BalancerFactory, cs remote.ConnStatistics) endpoint.MiddlewareBuilder {
 	return func(ctx context.Context) endpoint.Middleware {
 		return func(next endpoint.Endpoint) endpoint.Endpoint {
 			return func(ctx context.Context, request, response interface{}) error {
@@ -94,12 +95,12 @@ func newResolveMWBuilder(lbf *lbcache.BalancerFactory) endpoint.MiddlewareBuilde
 					return kerrors.ErrNoDestService
 				}
 
-				remote := remoteinfo.AsRemoteInfo(dest)
-				if remote == nil {
+				ri := remoteinfo.AsRemoteInfo(dest)
+				if ri == nil {
 					err := fmt.Errorf("unsupported target EndpointInfo type: %T", dest)
 					return kerrors.ErrInternalException.WithCause(err)
 				}
-				if remote.GetInstance() != nil {
+				if ri.GetInstance() != nil {
 					return next(ctx, request, response)
 				}
 				lb, err := lbf.Get(ctx, dest)
@@ -118,11 +119,16 @@ func newResolveMWBuilder(lbf *lbcache.BalancerFactory) endpoint.MiddlewareBuilde
 					// we always need to get a new picker every time, because when downstream update deployment,
 					// we may get an old picker that include all outdated instances which will cause connect always failed.
 					picker := lb.GetPicker()
+					cfg := rpcInfo.Config()
+					// gRPC streaming
+					if cs != nil && cfg.InteractionMode() == rpcinfo.Streaming && (cfg.TransportProtocol()&transport.GRPC != 0) {
+						ctx = remote.NewCtxWithConnStatistics(ctx, cs)
+					}
 					ins := picker.Next(ctx, request)
 					if ins == nil {
 						err = kerrors.ErrNoMoreInstance.WithCause(fmt.Errorf("last error: %w", lastErr))
 					} else {
-						remote.SetInstance(ins)
+						ri.SetInstance(ins)
 						// TODO: generalize retry strategy
 						err = next(ctx, request, response)
 					}

client/middlewares_test.go

@@ -35,10 +35,12 @@ import (
 	"github.com/cloudwego/kitex/pkg/event"
 	"github.com/cloudwego/kitex/pkg/kerrors"
 	"github.com/cloudwego/kitex/pkg/proxy"
+	"github.com/cloudwego/kitex/pkg/remote"
 	"github.com/cloudwego/kitex/pkg/remote/codec/protobuf"
 	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/status"
 	"github.com/cloudwego/kitex/pkg/rpcinfo"
 	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
+	"github.com/cloudwego/kitex/transport"
 )
 
 var (
@@ -84,7 +86,7 @@ func TestResolverMW(t *testing.T) {
 
 	var invoked bool
 	cli := newMockClient(t, ctrl).(*kcFinalizerClient)
-	mw := newResolveMWBuilder(cli.lbf)(ctx)
+	mw := newResolveMWBuilder(cli.lbf, nil)(ctx)
 	ep := func(ctx context.Context, request, response interface{}) error {
 		invoked = true
 		return nil
@@ -114,14 +116,14 @@ func TestResolverMWOutOfInstance(t *testing.T) {
 	}
 	var invoked bool
 	cli := newMockClient(t, ctrl, WithResolver(resolver)).(*kcFinalizerClient)
-	mw := newResolveMWBuilder(cli.lbf)(ctx)
+	mw := newResolveMWBuilder(cli.lbf, nil)(ctx)
 	ep := func(ctx context.Context, request, response interface{}) error {
 		invoked = true
 		return nil
 	}
 
 	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
-	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), nil, rpcinfo.NewRPCStats())
+	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
 
 	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
 	req := new(MockTStruct)
@@ -222,7 +224,7 @@ func BenchmarkResolverMW(b *testing.B) {
 	defer ctrl.Finish()
 
 	cli := newMockClient(b, ctrl).(*kcFinalizerClient)
-	mw := newResolveMWBuilder(cli.lbf)(ctx)
+	mw := newResolveMWBuilder(cli.lbf, nil)(ctx)
 	ep := func(ctx context.Context, request, response interface{}) error { return nil }
 	ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), nil, rpcinfo.NewRPCStats())
 
@@ -241,7 +243,7 @@ func BenchmarkResolverMWParallel(b *testing.B) {
 	defer ctrl.Finish()
 
 	cli := newMockClient(b, ctrl).(*kcFinalizerClient)
-	mw := newResolveMWBuilder(cli.lbf)(ctx)
+	mw := newResolveMWBuilder(cli.lbf, nil)(ctx)
 	ep := func(ctx context.Context, request, response interface{}) error { return nil }
 	ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), nil, rpcinfo.NewRPCStats())
 
@@ -279,3 +281,170 @@ func TestDiscoveryEventHandler(t *testing.T) {
 	added := extra["Added"].([]*instInfo)
 	test.Assert(t, len(added) == 1)
 }
+
+// mockConnStatistics implements remote.ConnStatistics for testing
+type mockConnStatistics struct {
+	activeStreams map[string]int
+}
+
+func (m *mockConnStatistics) ActiveStreams(addr string) int {
+	if m.activeStreams == nil {
+		return 0
+	}
+	return m.activeStreams[addr]
+}
+
+// TestResolverMW_WithConnStatistics_StreamingMode tests that ConnStatistics is passed to context
+// when in gRPC streaming mode
+func TestResolverMW_WithConnStatistics_StreamingMode(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockStats := &mockConnStatistics{
+		activeStreams: map[string]int{
+			"localhost:404": 5,
+		},
+	}
+
+	var contextPassedToEndpoint context.Context
+	cli := newMockClient(t, ctrl).(*kcFinalizerClient)
+	mw := newResolveMWBuilder(cli.lbf, mockStats)(ctx)
+	ep := func(ctx context.Context, request, response interface{}) error {
+		contextPassedToEndpoint = ctx
+		return nil
+	}
+
+	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
+
+	// Create RPC config with streaming mode and gRPC protocol
+	cfg := rpcinfo.NewRPCConfig()
+	rpcinfo.AsMutableRPCConfig(cfg).SetInteractionMode(rpcinfo.Streaming)
+	rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
+
+	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
+
+	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
+	req := new(MockTStruct)
+	res := new(MockTStruct)
+	err := mw(ep)(ctx, req, res)
+	test.Assert(t, err == nil)
+	test.Assert(t, to.GetInstance() == instance404[0])
+
+	// Verify ConnStatistics was passed to context
+	cs := remote.GetConnStatistics(contextPassedToEndpoint)
+	test.Assert(t, cs != nil, "ConnStatistics should be in context for streaming mode")
+	test.Assert(t, cs.ActiveStreams("localhost:404") == 5)
+}
+
+// TestResolverMW_WithConnStatistics_NonStreamingMode tests that ConnStatistics is NOT passed
+// to context when not in streaming mode
+func TestResolverMW_WithConnStatistics_NonStreamingMode(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockStats := &mockConnStatistics{
+		activeStreams: map[string]int{
+			"localhost:404": 5,
+		},
+	}
+
+	var contextPassedToEndpoint context.Context
+	cli := newMockClient(t, ctrl).(*kcFinalizerClient)
+	mw := newResolveMWBuilder(cli.lbf, mockStats)(ctx)
+	ep := func(ctx context.Context, request, response interface{}) error {
+		contextPassedToEndpoint = ctx
+		return nil
+	}
+
+	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
+
+	// Create RPC config with PingPong mode (not streaming)
+	cfg := rpcinfo.NewRPCConfig()
+	rpcinfo.AsMutableRPCConfig(cfg).SetInteractionMode(rpcinfo.PingPong)
+	rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
+
+	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
+
+	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
+	req := new(MockTStruct)
+	res := new(MockTStruct)
+	err := mw(ep)(ctx, req, res)
+	test.Assert(t, err == nil)
+
+	// Verify ConnStatistics was NOT passed to context for non-streaming mode
+	cs := remote.GetConnStatistics(contextPassedToEndpoint)
+	test.Assert(t, cs == nil, "ConnStatistics should not be in context for non-streaming mode")
+}
+
+// TestResolverMW_WithConnStatistics_NonGRPC tests that ConnStatistics is NOT passed
+// for non-gRPC protocols
+func TestResolverMW_WithConnStatistics_NonGRPC(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockStats := &mockConnStatistics{
+		activeStreams: map[string]int{
+			"localhost:404": 5,
+		},
+	}
+
+	var contextPassedToEndpoint context.Context
+	cli := newMockClient(t, ctrl).(*kcFinalizerClient)
+	mw := newResolveMWBuilder(cli.lbf, mockStats)(ctx)
+	ep := func(ctx context.Context, request, response interface{}) error {
+		contextPassedToEndpoint = ctx
+		return nil
+	}
+
+	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
+
+	// Create RPC config with streaming mode but non-gRPC protocol
+	cfg := rpcinfo.NewRPCConfig()
+	rpcinfo.AsMutableRPCConfig(cfg).SetInteractionMode(rpcinfo.Streaming)
+	rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader) // Not GRPC
+
+	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
+
+	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
+	req := new(MockTStruct)
+	res := new(MockTStruct)
+	err := mw(ep)(ctx, req, res)
+	test.Assert(t, err == nil)
+
+	// Verify ConnStatistics was NOT passed for non-gRPC protocol
+	cs := remote.GetConnStatistics(contextPassedToEndpoint)
+	test.Assert(t, cs == nil, "ConnStatistics should not be in context for non-gRPC protocol")
+}
+
+// TestResolverMW_WithoutConnStatistics tests behavior when ConnStatistics is nil
+func TestResolverMW_WithoutConnStatistics(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var contextPassedToEndpoint context.Context
+	cli := newMockClient(t, ctrl).(*kcFinalizerClient)
+	mw := newResolveMWBuilder(cli.lbf, nil)(ctx)
+	ep := func(ctx context.Context, request, response interface{}) error {
+		contextPassedToEndpoint = ctx
+		return nil
+	}
+
+	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
+
+	// Create RPC config with streaming mode and gRPC protocol
+	cfg := rpcinfo.NewRPCConfig()
+	rpcinfo.AsMutableRPCConfig(cfg).SetInteractionMode(rpcinfo.Streaming)
+	rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
+
+	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
+
+	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
+	req := new(MockTStruct)
+	res := new(MockTStruct)
+	err := mw(ep)(ctx, req, res)
+	test.Assert(t, err == nil)
+
+	// Verify ConnStatistics is nil when not provided
+	cs := remote.GetConnStatistics(contextPassedToEndpoint)
+	test.Assert(t, cs == nil, "ConnStatistics should be nil when not provided")
+}

pkg/loadbalance/weighted_balancer.go

@@ -30,6 +30,7 @@ const (
 	lbKindInterleaved
 	lbKindRandom
 	lbKindRandomWithAliasMethod
+	lbKindWeightedLeastConn
 )
 
 type weightedBalancer struct {
@@ -67,6 +68,11 @@ func NewWeightedRandomWithAliasMethodBalancer() Loadbalancer {
 	return lb
 }
 
+func NewWeightedLeastConnBalancer() Loadbalancer {
+	lb := &weightedBalancer{kind: lbKindWeightedLeastConn}
+	return lb
+}
+
 // GetPicker implements the Loadbalancer interface.
 func (wb *weightedBalancer) GetPicker(e discovery.Result) Picker {
 	if !e.Cacheable {
@@ -127,6 +133,8 @@ func (wb *weightedBalancer) createPicker(e discovery.Result) (picker Picker) {
 		} else {
 			picker = newAliasMethodPicker(instances, weightSum)
 		}
+	case lbKindWeightedLeastConn:
+		picker = newWeightedLeastConnPicker(instances, balance)
 	default: // random
 		if balance {
 			picker = newRandomPicker(instances)
@@ -161,6 +169,8 @@ func (wb *weightedBalancer) Name() string {
 		return "interleaved_weighted_round_robin"
 	case lbKindRandomWithAliasMethod:
 		return "weight_random_with_alias_method"
+	case lbKindWeightedLeastConn:
+		return "weight_least_conn"
 	default:
 		return "weight_random"
 	}