fix: combine services (cloudwego#367)

starswiss · Mar 18, 2022 · 8a916a4 · 8a916a4
1 parent fefbc93
commit 8a916a4
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/pkg/generic/thrift/parse.go b/pkg/generic/thrift/parse.go
@@ -85,8 +85,9 @@ func Parse(tree *parser.Thrift, mode ParseMode) (*descriptor.ServiceDescriptor,
 		sDsc.Name = "CombinedServices"
 	}
 
+	visitedSvcs := make(map[*parser.Service]bool, len(tree.Services))
 	for _, svc := range svcs {
-		for p := range getAllFunctions(svc, tree) {
+		for p := range getAllFunctions(svc, tree, visitedSvcs) {
 			fn := p.data.(*parser.Function)
 			if err := addFunction(fn, p.tree, sDsc, structsCache); err != nil {
 				return nil, err
@@ -101,11 +102,17 @@ type pair struct {
 	data interface{}
 }
 
-func getAllFunctions(svc *parser.Service, tree *parser.Thrift) chan *pair {
+func getAllFunctions(svc *parser.Service, tree *parser.Thrift, visitedSvcs map[*parser.Service]bool) chan *pair {
 	ch := make(chan *pair)
 	svcs := make(chan *pair)
+	addSvc := func(tree *parser.Thrift, svc *parser.Service) {
+		if exist := visitedSvcs[svc]; !exist {
+			svcs <- &pair{tree: tree, data: svc}
+			visitedSvcs[svc] = true
+		}
+	}
 	go func() {
-		svcs <- &pair{tree: tree, data: svc}
+		addSvc(tree, svc)
 		for base := svc.Extends; base != ""; base = svc.Extends {
 			ref := svc.GetReference()
 			if ref != nil {
@@ -114,10 +121,7 @@ func getAllFunctions(svc *parser.Service, tree *parser.Thrift) chan *pair {
 				tree = tree.Includes[idx].Reference
 			}
 			svc, _ = tree.GetService(base)
-			svcs <- &pair{
-				tree: tree,
-				data: svc,
-			}
+			addSvc(tree, svc)
 		}
 		close(svcs)
 	}()

diff --git a/pkg/generic/thrift/parse_test.go b/pkg/generic/thrift/parse_test.go
@@ -253,4 +253,20 @@ func TestExtends(t *testing.T) {
 	for _, fn := range []string{"simple", "int2bin", "req2res"} {
 		test.Assert(t, dp.Functions[fn] != nil)
 	}
+
+	dp, err = Parse(demo, FirstServiceOnly)
+	test.Assert(t, err == nil)
+	test.Assert(t, dp.Name == "DemoBaseService")
+	test.Assert(t, len(dp.Functions) == 2)
+	for _, fn := range []string{"simple", "int2bin"} {
+		test.Assert(t, dp.Functions[fn] != nil)
+	}
+
+	dp, err = Parse(demo, CombineServices)
+	test.Assert(t, err == nil)
+	test.Assert(t, dp.Name == "CombinedServices")
+	test.Assert(t, len(dp.Functions) == 3)
+	for _, fn := range []string{"simple", "int2bin", "req2res"} {
+		test.Assert(t, dp.Functions[fn] != nil)
+	}
 }