Using one balancer for all clients

Author: nexcode

client.go

@@ -17,6 +17,7 @@
 package rpcplatform
 
 import (
+	"github.com/nexcode/rpcplatform/internal/config"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/resolver/manual"
 )
@@ -26,4 +27,5 @@ type Client struct {
 	target   string
 	client   *grpc.ClientConn
 	resolver *manual.Resolver
+	config   *config.Client
 }

client_updatestate.go

@@ -17,19 +17,20 @@
 package rpcplatform
 
 import (
-	"google.golang.org/grpc/attributes"
+	"github.com/nexcode/rpcplatform/internal/grpcattrs"
 	"google.golang.org/grpc/resolver"
 )
 
 func (c *Client) updateState(init bool, serverInfoTree map[string]*ServerInfo) {
 	state := resolver.State{
 		Addresses: make([]resolver.Address, 0, len(serverInfoTree)),
+		// Attributes: grpcattrs.SetClientConfig(nil, c.config), // https://github.com/grpc/grpc-go/pull/8696
 	}
 
 	for _, value := range serverInfoTree {
 		state.Addresses = append(state.Addresses, resolver.Address{
 			Addr:       value.Address,
-			Attributes: attributes.New(struct{}{}, value.Attributes),
+			Attributes: grpcattrs.SetClientConfig(grpcattrs.SetAttributes(nil, value.Attributes), c.config),
 		})
 	}
 

init.go

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2025 RPCPlatform 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 rpcplatform
+
+import (
+	"github.com/nexcode/rpcplatform/internal/balancer"
+)
+
+func init() {
+	balancer.Register()
+}

internal/balancer/pickerbuilder.go

@@ -16,6 +16,4 @@
 
 package balancer
 
-type pickerBuilder struct {
-	maxActiveServers int
-}
+type pickerBuilder struct{}

internal/balancer/pickerbuilder_makeconninfo.go

@@ -21,7 +21,8 @@ import (
 	"math"
 	"slices"
 
-	"github.com/nexcode/rpcplatform/internal/attributes"
+	"github.com/nexcode/rpcplatform/internal/config"
+	"github.com/nexcode/rpcplatform/internal/grpcattrs"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/base"
 )
@@ -34,12 +35,18 @@ type connInfo struct {
 	count    int
 }
 
-func (pb *pickerBuilder) makeConnInfo(pickerInfo base.PickerBuildInfo) ([]*connInfo, int) {
+func (*pickerBuilder) makeConnInfo(pickerInfo base.PickerBuildInfo) ([]*connInfo, int) {
 	connInfoArr := make([]*connInfo, 0, len(pickerInfo.ReadySCs))
+
 	var totalWeight int
+	var config *config.Client
 
 	for subConn, subConnInfo := range pickerInfo.ReadySCs {
-		attributes := subConnInfo.Address.Attributes.Value(struct{}{}).(*attributes.Attributes)
+		if config == nil {
+			config = grpcattrs.GetClientConfig(subConnInfo.Address.Attributes)
+		}
+
+		attributes := grpcattrs.GetAttributes(subConnInfo.Address.Attributes)
 
 		if attributes.BalancerWeight <= 0 {
 			continue
@@ -56,8 +63,8 @@ func (pb *pickerBuilder) makeConnInfo(pickerInfo base.PickerBuildInfo) ([]*connI
 		return cmp.Compare(b.priority, a.priority)
 	})
 
-	if pb.maxActiveServers > 0 && pb.maxActiveServers < len(connInfoArr) {
-		connInfoArr = connInfoArr[:pb.maxActiveServers]
+	if config.MaxActiveServers > 0 && config.MaxActiveServers < len(connInfoArr) {
+		connInfoArr = connInfoArr[:config.MaxActiveServers]
 	}
 
 	for _, connInfo := range connInfoArr {

internal/balancer/register.go

@@ -21,10 +21,8 @@ import (
 	"google.golang.org/grpc/balancer/base"
 )
 
-func Register(name string, maxActiveServers int) {
-	balancer.Register(base.NewBalancerBuilder(name, &pickerBuilder{
-		maxActiveServers: maxActiveServers,
-	}, base.Config{
+func Register() {
+	balancer.Register(base.NewBalancerBuilder("rpcplatform", &pickerBuilder{}, base.Config{
 		HealthCheck: true,
 	}))
 }

internal/grpcattrs/attributes.go

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2025 RPCPlatform 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 grpcattrs
+
+import (
+	"github.com/nexcode/rpcplatform/internal/attributes"
+	grpcattrs "google.golang.org/grpc/attributes"
+)
+
+func GetAttributes(attrs *grpcattrs.Attributes) *attributes.Attributes {
+	value, _ := attrs.Value(keyAttributes).(*attributes.Attributes)
+	return value
+}
+
+func SetAttributes(attrs *grpcattrs.Attributes, value *attributes.Attributes) *grpcattrs.Attributes {
+	return attrs.WithValue(keyAttributes, value)
+}

internal/grpcattrs/attrkey.go

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2025 RPCPlatform 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 grpcattrs
+
+type attrKey uint8
+
+const (
+	keyAttributes attrKey = iota
+	keyClientConfig
+)

internal/grpcattrs/clientconfig.go

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2025 RPCPlatform 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 grpcattrs
+
+import (
+	"github.com/nexcode/rpcplatform/internal/config"
+	grpcattrs "google.golang.org/grpc/attributes"
+)
+
+func GetClientConfig(attrs *grpcattrs.Attributes) *config.Client {
+	value, _ := attrs.Value(keyClientConfig).(*config.Client)
+	return value
+}
+
+func SetClientConfig(attrs *grpcattrs.Attributes, value *config.Client) *grpcattrs.Attributes {
+	return attrs.WithValue(keyClientConfig, value)
+}

internal/resolver/resolver.go

@@ -20,6 +20,6 @@ import (
 	"google.golang.org/grpc/resolver/manual"
 )
 
-func NewResolver() *manual.Resolver {
-	return manual.NewBuilderWithScheme("dns")
+func New() *manual.Resolver {
+	return manual.NewBuilderWithScheme("rpcplatform")
 }