grpc 双向通道 编写聊天室
利用 grpc 双向通道的特性，编写一个简易的聊天室。
服务端: go
客户端: go 和 python
go 语言依赖 https://blog.csdn.net/sunt2018/article/details/106630815
python grpc 入门 https://blog.csdn.net/sunt2018/article/details/90176015
grpc protobuf契约协议
protoes/hello.proto
syntax = "proto3";
package protoes;
import "google/protobuf/timestamp.proto";
service OnLineChat {
rpc SayHi(stream HiRequest) returns (stream HiReply) {};
}
message HiRequest {
string message = 1;
}
message HiReply {
string message = 1;
google.protobuf.Timestamp TS = 2;
MessageType message_type = 3;
enum MessageType{
CONNECT_SUCCESS = 0;
CONNECT_FAILED = 1;
NORMAL_MESSAGE = 2;
}
}
go 服务端
package main
import (
"github.com/golang/protobuf/ptypes/timestamp"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
pb "grcp_chat_room/protoes"
"log"
"net"
"os"
"sync"
"time"
"fmt"
)
// 定义一个类，继承字典(异步,带锁的),一会存入grpc stream对象 { name : stream }
type ConnectPool struct {
sync.Map
}
//为这个 类<对象池> 添加方法，分别为 Get,Add,Del和BroadCast(广播信息,群发)
func (p *ConnectPool) Get(name string) pb.OnLineChat_SayHiServer{
if stream, ok := p.Load(name); ok {
return stream.(pb.OnLineChat_SayHiServer)
} else {
return nil
}
}
func (p *ConnectPool) Add(name string, stream pb.OnLineChat_SayHiServer) {
p.Store(name, stream)
}
func (p *ConnectPool) Del(name string) {
p.Delete(name)
}
// 聊天室内 广播
func (p *ConnectPool) BroadCast(from, message string) {
log.Printf("BroadCast from: %s, message: %s\n", from, message)
p.Range(func(username_i, stream_i interface{}) bool {
username := username_i.(string)
stream := stream_i.(pb.OnLineChat_SayHiServer)
if username == from {
return true
} else {
stream.Send(&pb.HiReply{
Message: message,
MessageType: pb.HiReply_NORMAL_MESSAGE, // 2.正常数据
TS: ×tamp.Timestamp{Seconds: time.Now().Unix()},
})
}
return true
})
}
var connect_pool *ConnectPool
// 定义服务器类
type Service struct{}
func (s *Service) SayHi(stream pb.OnLineChat_SayHiServer) error {
peer, _ := peer.FromContext(stream.Context())
log.Printf("Received new connection. %s", peer.Addr.String())
md, _ := metadata.FromIncomingContext(stream.Context())
username := md["name"][0] // 从metadata中获取用户名信息，可以理解为请求头里的数据
if connect_pool.Get(username) != nil {
stream.Send(&pb.HiReply{
Message: fmt.Sprintf("username %s already exists!", username),
MessageType: pb.HiReply_CONNECT_FAILED, // 1. 连接失败 ， 重名了 用户已经存在
})
return nil
} else { // 连接成功
connect_pool.Add(username, stream)
stream.Send(&pb.HiReply{
Message: fmt.Sprintf("Connect success!"),
MessageType: pb.HiReply_CONNECT_SUCCESS, // 0 连接成功
})
}
go func() {
// 阻塞住，等待断开连接的时候触发
<-stream.Context().Done()
connect_pool.Del(username)
connect_pool.BroadCast(username, fmt.Sprintf("%s leval room", username))
}()
// 广播，xxxx进入了聊天室直播间
connect_pool.BroadCast(username, fmt.Sprintf("Welcome %s!", username))
// 阻塞接收 该用户后续传来的消息
for {
req, err := stream.Recv()
if err != nil {
return err
}
connect_pool.BroadCast(username, fmt.Sprintf("%s: %s", username, req.Message))
}
return nil
}
func GetListen() string {
if len(os.Args) < 2 {
return ":9999"
}
return os.Args[1]
}
func main() {
connect_pool = &ConnectPool{}
// 监听一个 地址:端口
address, err := net.Listen("tcp", GetListen())
if err != nil {
panic(err)
}
// 注册 grpc
ser := grpc.NewServer()
pb.RegisterOnLineChatServer(ser, &Service{}) //必须实现protoes中定义的方法，不然这里无法通过检测
// 启动服务
if err := ser.Serve(address); err != nil {
panic(err)
}
}
go client
package main
import (
"bufio"
"context"
"flag"
"fmt"
"io"
"log"
"os"
"sync"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
pb "grcp_chat_room/protoes"
)
// client.exe -name xxxx -address xxxxx 不写使用默认值
var name *string = flag.String("name", "guess", "what's your name?")
var address *string = flag.String("address", "127.0.0.1:8881", "server address")
var mutex sync.Mutex
// 这是一个加锁的输出，防止乱序或中间插入print数据
func ConsoleLog(message string) {
mutex.Lock()
defer mutex.Unlock()
fmt.Printf("\n------ %s -----\n%s\n> ", time.Now(), message)
}
// 输入
func Input(prompt string) string {
fmt.Print(prompt)
reader := bufio.NewReader(os.Stdin)
line, _, err := reader.ReadLine()
if err != nil {
if err == io.EOF {
return ""
} else {
panic(err)
}
}
return string(line)
}
func main() {
// 接受命令行参数
flag.Parse()
// 创建连接，拨号
conn, err := grpc.Dial("localhost:9999", grpc.WithInsecure())
if err != nil {
log.Printf("连接失败: [%v] ", err)
return
}
defer conn.Close()
// 声明客户端
client := pb.NewOnLineChatClient(conn)
// 声明 context
//ctx := context.Background()
ctx, cancel := context.WithCancel(context.Background())
ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("name", *name))
// 创建双向数据流
stream, err := client.SayHi(ctx)
if err != nil {
log.Printf("创建数据流失败: [%v] ", err)
}
// 创建了一个连接管道
connected := make(chan bool)
// 接收 服务端信息
go func() {
var (
reply *pb.HiReply
err error
)
for {
reply, err = stream.Recv()
if err != nil{
panic(err)
}
ConsoleLog(reply.Message)
if reply.MessageType == pb.HiReply_CONNECT_FAILED { // code=1 连接失败
cancel()
break
}
if reply.MessageType == pb.HiReply_CONNECT_SUCCESS { // code=0 连接成功
connected <- true
}
// 基本都是两个if都不执行，去下一次循环,返回的是 code=2 正常消息
}
}()
go func() {
<-connected
var (
line string
err error
)
for {
line = Input("")
if line == "exit" {
cancel()
break
}
err = stream.Send(&pb.HiRequest{
Message: line,
})
fmt.Print("> ")
if err != nil{
panic(err)
}
}
}()
<-ctx.Done()
fmt.Println("Bye")
}
python client
import time
from time import strftime, localtime
import grpc
from protoes import hello_pb2, hello_pb2_grpc
def gen():
while 1:
i = str(input(">"))
if i == "q":
break
yield hello_pb2.HiRequest(message=i)
time.sleep(0.1)
def run():
channel = grpc.insecure_channel("127.0.0.1:9999")
stub = hello_pb2_grpc.OnLineChatStub(channel)
# 连接metadata中带着 用户名python
metadata=(("name","python"),)
it = stub.SayHi(gen(),metadata=metadata)
for r in it:
ztime = strftime("%Y-%m-%d %H:%M:%S", localtime())
print(f"\n------ {ztime} -----\n{r.message}\n>", end = "")
# print(dir(r))
if __name__ == '__main__':
run()