diff --git a/example/rpc/client/stream/client.go b/example/rpc/client/stream/client.go index 55f93db8..02dae928 100644 --- a/example/rpc/client/stream/client.go +++ b/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() }