Merge pull request #1 from shaoxi/main

feat: support service use grpc port
api7 · Jan 7, 2025 · 5206274 · 5206274
2 parents 912bd55 + 4a11b3b
commit 5206274
Showing 2 changed files with 101 additions and 9 deletions.
diff --git a/internal/core/message/a6conf.go b/internal/core/message/a6conf.go
@@ -2,8 +2,12 @@ package message
 
 import (
 	"encoding/json"
+	"fmt"
 	"reflect"
+	"strconv"
 	"strings"
+
+	"github.com/api7/gopkg/pkg/log"
 )
 
 // Sample route config
@@ -12,14 +16,16 @@ import (
 // 	"discovery_args.group_name": "group_name",
 // 	"discovery_args.namespace_id": "test_name",
 // 	"discovery_type": "nacos",
-// 	"service_name": "test-service"
+// 	"service_name": "test-service",
+// 	"service_grpc_port": "10001"
 // },
 
 type Labels struct {
 	DiscoveryType            string `json:"discovery_type,omitempty"`
 	ServiceName              string `json:"service_name,omitempty"`
 	DiscoveryArgsNamespaceID string `json:"discovery_args.namespace_id,omitempty"`
 	DiscoveryArgsGroupName   string `json:"discovery_args.group_name,omitempty"`
+	ServiceGrpcPort          string `json:"service_grpc_port,omitempty"`
 }
 
 type UpstreamArg struct {
@@ -196,7 +202,8 @@ func NewUpstreams(value []byte) (A6Conf, error) {
 
 // "labels": {
 // 	"service_name":"aquaman-user",
-// 	"discovery_type":"nacos"
+// 	"discovery_type":"nacos",
+// 	"service_grpc_port":"10001"
 // },
 
 type Routes struct {
@@ -211,8 +218,26 @@ func (routes *Routes) GetAll() *map[string]interface{} {
 }
 
 func (routes *Routes) Marshal() ([]byte, error) {
-	// bytes1, _ := json.Marshal(routes.All)
-	// print("a6conf marshal 1=====", string(bytes1))
+	// If grpc port is configured, modify all nodes' port to grpc port
+	if routes.Labels.ServiceGrpcPort != "" && routes.Upstream.Nodes != nil {
+		grpcPort, err := strconv.ParseInt(routes.Labels.ServiceGrpcPort, 10, 64)
+		if err != nil {
+			log.Errorf("invalid grpc port configuration: failed to parse port %s to integer for route %s: %s", routes.Labels.ServiceGrpcPort, routes.All["id"], err)
+			return nil, fmt.Errorf("invalid grpc port configuration: failed to parse port %s to integer for route %s: %s", routes.Labels.ServiceGrpcPort, routes.All["id"], err)
+		}
+		if nodes, ok := routes.Upstream.Nodes.([]*Node); ok {
+			nodesCopy := make([]*Node, len(nodes))
+			for i, n := range nodes {
+				nodesCopy[i] = &Node{
+					Host:   n.Host,
+					Weight: n.Weight,
+					Port:   int(grpcPort),
+				}
+				log.Infof("updated gRPC port to %d for node %s in route %s", grpcPort, n.Host, routes.All["id"])
+			}
+			routes.Upstream.Nodes = nodesCopy
+		}
+	}
 
 	embedElm(reflect.ValueOf(routes), routes.All)
 

diff --git a/internal/core/message/a6conf_test.go b/internal/core/message/a6conf_test.go
@@ -124,11 +124,6 @@ func TestMarshal_Routes(t *testing.T) {
         "pass_host": "pass",
         "type": "roundrobin",
         "hash_on": "vars",
-        "_discovery_type": "nacos",
-        "_service_name": "APISIX-NACOS",
-        "discovery_args": {
-            "group_name": "DEFAULT_GROUP"
-        },
         "nodes": [
             {
                 "host": "192.168.1.1",
@@ -157,6 +152,78 @@ func TestMarshal_Routes(t *testing.T) {
 	assert.JSONEq(t, wantA6Str, string(ss))
 }
 
+func TestMarshal_Routes_With_Grpc_Port(t *testing.T) {
+	givenA6Str := `{
+    "status": 1,
+    "id": "3",
+    "uri": "/hh",
+	 "labels": {
+        "discovery_type": "nacos",
+        "discovery_args.group_name": "DEFAULT_GROUP",
+        "service_name": "test-service",
+		"service_grpc_port":"10001"
+    },
+    "upstream": {
+        "scheme": "http",
+        "pass_host": "pass",
+        "type": "roundrobin",
+        "hash_on": "vars",
+        "nodes": {
+            "192.168.1.1:10001": 1
+        }
+    },
+    "create_time": 1648871506,
+    "priority": 0,
+    "update_time": 1648871506
+}`
+	nodes := []*Node{
+		{Host: "192.168.1.1", Port: 80, Weight: 1},
+		{Host: "192.168.1.2", Port: 80, Weight: 1},
+	}
+
+	wantA6Str := `{
+    "status": 1,
+    "id": "3",
+    "uri": "/hh",
+	"labels": {
+        "discovery_type": "nacos",
+        "discovery_args.group_name": "DEFAULT_GROUP",
+        "service_name": "test-service",
+		"service_grpc_port":"10001"
+    },
+    "upstream": {
+        "scheme": "http",
+        "pass_host": "pass",
+        "type": "roundrobin",
+        "hash_on": "vars",
+        "nodes": [
+            {
+                "host": "192.168.1.1",
+                "port": 10001,
+                "weight": 1
+            },
+            {
+                "host": "192.168.1.2",
+                "port": 10001,
+                "weight": 1
+            }
+        ]
+    },
+    "create_time": 1648871506,
+    "priority": 0,
+    "update_time": 1648871506
+}`
+	caseDesc := "sanity"
+	a6, err := NewA6Conf([]byte(givenA6Str), A6RoutesConf)
+	assert.Nil(t, err, caseDesc)
+
+	a6.Inject(nodes)
+	ss, err := a6.Marshal()
+	assert.Nil(t, err, caseDesc)
+
+	assert.JSONEq(t, wantA6Str, string(ss))
+}
+
 func TestHasNodesAttr_Routes(t *testing.T) {
 	tests := []struct {
 		name  string