Browse Source

update rpc example

kevin 4 years ago
parent
commit
974c19d6d3
1 changed files with 6 additions and 6 deletions
  1. 6 6
      example/rpc/client/stream/client.go

+ 6 - 6
example/rpc/client/stream/client.go

@@ -5,6 +5,7 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"sync"
 
 	"github.com/tal-tech/go-zero/core/discov"
 	"github.com/tal-tech/go-zero/example/rpc/remote/stream"
@@ -33,22 +34,21 @@ func main() {
 		log.Fatal(err)
 	}
 
-	done := make(chan struct{})
+	var wg sync.WaitGroup
 	go func() {
-		defer func() {
-			done <- struct{}{}
-		}()
-		for i := 0; i < 3; i++ {
+		for {
 			resp, err := stm.Recv()
 			if err != nil {
 				log.Fatal(err)
 			}
 
 			fmt.Println("=>", resp.Greet)
+			wg.Done()
 		}
 	}()
 
 	for i := 0; i < 3; i++ {
+		wg.Add(1)
 		fmt.Println("<=", name)
 		if err = stm.Send(&stream.StreamReq{
 			Name: name,
@@ -57,5 +57,5 @@ func main() {
 		}
 	}
 
-	<-done
+	wg.Wait()
 }