本文已参与「新人创作礼」活动，一起开启掘金创作之路。

实现效果

多个客户端连接到服务器，一个客户端发送消息，服务器接收后广播到各个客户端。

Server

Proto

github.com/protocolbuf… 放到third_party/google/protobuf/timestamp.proto

编辑 .\api\helloworld\v1\greeter.proto

import "google/protobuf/timestamp.proto";

.....

// 服务方法
service Stream{
  // 双向流式rpc，同时在请求参数前和响应参数前加上stream
  rpc Conversations(stream StreamRequest) returns(stream StreamResponse){};
  rpc Chat(stream ChatRequest) returns(stream ChatResponse){};
}

message ChatRequest{
  int64 user_id = 1;
  string message = 2;
  google.protobuf.Timestamp send_time = 3;
}

message ChatResponse{
  int64 user_id = 1;
  string message = 2;
  google.protobuf.Timestamp send_time = 3;
}

生成proto对应的go文件 kratos proto client .\api\helloworld\v1\greeter.proto

Service

生成server代码 kratos proto server .\api\helloworld\v1\greeter.proto -t .\internal\service (已有不会生成)

服务端接收到客户端发送消息，通过go异步处理，交给biz层。

~~biz层处理完后发送消息到chan。~~

服务端处理批量发送:connect进来的时候将conn保存到一个全局数组；然后通过一个静态方法遍历每个conn对象再发送。

也可以在connect进来的时候创建一个chan，然后保存chan数组；service开一个go去读chan。

考虑到客户端随时会断开连接，也就是说conn随时会失效，发送前需要检查是否还存在连接。

internal/service/stream.go


package service

import (
   "fmt"
   "github.com/go-kratos/kratos/v2/log"
   "grpcStream/internal/biz"
   "io"
   "sync"

   pb "grpcStream/api/helloworld/v1"
)

type StreamService struct {
   pb.UnimplementedStreamServer
   hub       *biz.ChatHub
   chatGroup sync.Map
   log       *log.Helper
}

func NewStreamService(logger log.Logger, chatHub *biz.ChatHub) *StreamService {
   go chatHub.Run()
   return &StreamService{log: log.NewHelper(logger), hub: chatHub}
}

func (s *StreamService) Chat(conn pb.Stream_ChatServer) error {
   //ctx := context.Background()
   end := make(chan bool) //此函数不可以直接退出
   client := &biz.Client{
      Hub:     s.hub,
      Conn:    conn,
      Send:    make(chan biz.Message),
      Log:     s.log,
      EndChan: end,
   }
   s.log.Infof("客户端连接")
   go client.WritePump() // 监听写出
   go client.ReadPump()  // 监听读取
   <-end
   return nil
}

Biz

Hub

Hub用来保存Client和广播

internal/biz/chatHub.go

package biz

import (
   "github.com/go-kratos/kratos/v2/log"
   "google.golang.org/protobuf/types/known/timestamppb"
   pb "grpcStream/api/helloworld/v1"
   "io"
   "strconv"
   "time"
)

type Client struct {
   Hub *ChatHub

   Conn pb.Stream_ChatServer

   Send chan Message // 发送通道

   Name string // 客户名称

   Log *log.Helper

   EndChan chan bool // 结束信号
}

func (c *Client) ReadPump() {
   defer func() {
      c.Hub.unregister <- c.Name //TODO 这里应该增加退出信息
      c.EndChan <- true
   }()
   for {
      req, err := c.Conn.Recv() // 接受客户端的消息
      //c.Log.Infof("收到消息 %+v", req)
      if err == io.EOF {
         //c.Log.Infof("结束-EOF") // 没有消息时会持续触发
         //c.Hub.unregister <- c
         continue
      }
      if err != nil {
         //c.Hub.unregister <- c
         c.Log.Infof("结束-错误") // 用户断开连接
         return
      }
      userID := req.GetUserId()
      // 用户注册业务逻辑
      if userID != 0 && c.Name == "" {
         // 注册
         c.Name = strconv.FormatInt(userID, 10)
         c.Hub.Register <- c
      }
      //TODO 收到信息后的业务逻辑

      // 返回确认收到消息 可有可无
      m := Message{
         userID:      userID,
         messageType: 1,
         message:     "收到消息" + req.GetMessage(),
         sendTime:    time.Now(),
      }
      c.Send <- m

      //广播
      timeStamp := req.GetSendTime()
      broadcastMessage := Message{
         messageType: 0,
         message:     strconv.FormatInt(userID, 10) + ":" + req.GetMessage(),
         sendTime:    time.Unix(timeStamp.GetSeconds(), 0),
      }
      c.Hub.broadcast <- broadcastMessage
   }

}

func (c *Client) WritePump() {
   for {
      message, ok := <-c.Send //读取要发送的信息
      if !ok {
         break
      }
      err := c.Conn.Send(&pb.ChatResponse{
         UserId:   message.userID,
         Message:  message.message,
         SendTime: timestamppb.New(message.sendTime),
      })
      if err != nil {
         c.Hub.unregister <- c.Name
         return
      }
   }
}

Client

Client用来对Connect进行操作包括读和写

internal/biz/chatClient.go

package biz

import (
   "github.com/go-kratos/kratos/v2/log"
   "google.golang.org/protobuf/types/known/timestamppb"
   pb "grpcStream/api/helloworld/v1"
   "io"
   "strconv"
   "time"
)

type Client struct {
   Hub *ChatHub

   Conn pb.Stream_ChatServer

   Send chan Message // 发送通道

   Name string // 客户名称

   Log *log.Helper

   EndChan chan bool // 结束信号
}

func (c *Client) ReadPump() {
   defer func() {
      c.Hub.unregister <- c.Name //TODO 这里应该增加退出信息
      c.EndChan <- true
   }()
   for {
      req, err := c.Conn.Recv() // 接受客户端的消息
      //c.Log.Infof("收到消息 %+v", req)
      if err == io.EOF {
         //c.Log.Infof("结束-EOF") // 没有消息时会持续触发
         //c.Hub.unregister <- c
         continue
      }
      if err != nil {
         //c.Hub.unregister <- c
         c.Log.Infof("结束-错误") // 用户断开连接
         return
      }
      userID := req.GetUserId()
      // 用户注册业务逻辑
      if userID != 0 && c.Name == "" {
         // 注册
         c.Name = strconv.FormatInt(userID, 10)
         c.Hub.Register <- c
      }
      //TODO 收到信息后的业务逻辑

      // 返回确认收到消息 可有可无
      m := Message{
         userID:      userID,
         messageType: 1,
         message:     "收到消息" + req.GetMessage(),
         sendTime:    time.Now(),
      }
      c.Send <- m

      //广播
      timeStamp := req.GetSendTime()
      broadcastMessage := Message{
         messageType: 0,
         message:     strconv.FormatInt(userID, 10) + ":" + req.GetMessage(),
         sendTime:    time.Unix(timeStamp.GetSeconds(), 0),
      }
      c.Hub.broadcast <- broadcastMessage
   }

}

func (c *Client) WritePump() {
   for {
      message, ok := <-c.Send //读取要发送的信息
      if !ok {
         break
      }
      err := c.Conn.Send(&pb.ChatResponse{
         UserId:   message.userID,
         Message:  message.message,
         SendTime: timestamppb.New(message.sendTime),
      })
      if err != nil {
         c.Hub.unregister <- c.Name
         return
      }
   }
}

Wire

执行go generate ./... 生成:

func wireApp(confServer *conf.Server, confData *conf.Data, logger log.Logger) (*kratos.App, func(), error) {
   dataData, cleanup, err := data.NewData(confData, logger)
   if err != nil {
      return nil, nil, err
   }
   greeterRepo := data.NewGreeterRepo(dataData, logger)
   greeterUsecase := biz.NewGreeterUsecase(greeterRepo, logger)
   greeterService := service.NewGreeterService(greeterUsecase)
   httpServer := server.NewHTTPServer(confServer, greeterService, logger)
   chatHub := biz.NewHub(logger)
   streamService := service.NewStreamService(logger, chatHub)
   grpcServer := server.NewGRPCServer(confServer, greeterService, streamService, logger)
   app := newApp(logger, httpServer, grpcServer)
   return app, func() {
      cleanup()
   }, nil
}

Client 测试

test/main.go

package main

import (
   "bufio"
   "context"
   "fmt"
   "github.com/go-kratos/kratos/v2/errors"
   "github.com/go-kratos/kratos/v2/middleware/recovery"
   transgrpc "github.com/go-kratos/kratos/v2/transport/grpc"
   "google.golang.org/grpc"
   "google.golang.org/protobuf/types/known/timestamppb"
   v1 "grpcStream/api/helloworld/v1"
   "io"
   "log"
   "math/rand"
   "os"
   "time"
)

var conn *grpc.ClientConn

// 测试客户端连接
func main() {
   var err error
   conn, err = transgrpc.DialInsecure(
      context.Background(),
      transgrpc.WithEndpoint("127.0.0.1:9000"),
      transgrpc.WithMiddleware(
         recovery.Recovery(),
      ),
   )
   if err != nil {
      panic(err)
   }
   defer conn.Close()
   chat()
}


func chat() {
   client := v1.NewStreamClient(conn)
   chatClient, err := client.Chat(context.TODO())
   if err != nil {
      log.Fatalf("Failed  : %v", err)
      return
   }
   waitc := make(chan struct{})
   go func() {
      for {
         in, err := chatClient.Recv()
         if err == io.EOF {
            // read done.
            close(waitc)
            return
         }
         if err != nil {
            log.Fatalf("Failed to receive a note : %v", err)
         }
         log.Printf("Got message (%v+)", in)
      }
   }()
   // 客户端模拟发送
   s1 := rand.NewSource(time.Now().UnixNano())
   r1 := rand.New(s1)
   randnum := r1.Intn(10000)
   for n := 0; n < 5; n++ {
      if err := chatClient.Send(&v1.ChatRequest{
         UserId:   int64(randnum),
         Message:  fmt.Sprintf("this is %d repeat %d", randnum, n),
         SendTime: timestamppb.Now(),
      }); err != nil {
         log.Fatalf("Failed to send a note: %v", err)
      }
   }

   //最后关闭流 (EOF)
   //err = chatClient.CloseSend()
   //if err != nil {
   // log.Fatalf("Conversations close stream err: %v", err)
   //}
   // 启动一个 goroutine 接收命令行输入的指令
   go func() {
      log.Println("请输入消息...")
      输入 := bufio.NewReader(os.Stdin)
      for {
         // 获取 命令行输入的字符串， 以回车 \n 作为结束标志
         msg, _ := 输入.ReadString('\n')

         // 向服务端发送 指令
         if err := chatClient.Send(&v1.ChatRequest{UserId: int64(randnum),
            Message:  fmt.Sprintf("%s", msg),
            SendTime: timestamppb.Now(),
         }); err != nil {
            return
         }
      }
   }()
   <-waitc
}

结果

图片.png

下边是开了两个Client测试

一边输入另一边可以收到消息

结语

代码有点问题，有待完善

一是注册的逻辑没有写，加个Token进来校验应该就可以了。（如果使用WebSocket的话可以在请求进来的时候校验请求头的Token是否正确，不正确就直接关闭连接。用gRPC应该也可以，可以直接设置Context来传递。总之要把注册的逻辑提前。）

二是业务逻辑没有分离，可以设一个函数作为业务逻辑，在chatClient里面调用，取返回结果在发出。（应该算解耦了）

三是c.Conn.Recv()的两个错误处理得不太好

Kratos gRPC双向数据流 聊天Demo