package broker import ( "context" "net" "time" "github.com/eclipse/paho.mqtt.golang/packets" pb "github.com/fhmq/hmq/grpc" "go.uber.org/zap" "google.golang.org/grpc" "google.golang.org/grpc/keepalive" "google.golang.org/grpc/reflection" ) func (b *Broker) initRPCService() { lis, err := net.Listen("tcp", ":"+b.config.RpcPort) if err != nil { log.Error("failed to listen: ", zap.Error(err)) return } s := grpc.NewServer(grpc.KeepaliveParams(keepalive.ServerParameters{ Time: 30 * time.Minute, })) pb.RegisterHMQServiceServer(s, &HMQ{b: b}) reflection.Register(s) if err := s.Serve(lis); err != nil { log.Error("failed to serve: ", zap.Error(err)) } } func (b *Broker) initRPCClient(id, url string) { conn, err := grpc.Dial(url, grpc.WithInsecure(), grpc.WithKeepaliveParams(keepalive.ClientParameters{ Time: 30 * time.Minute, })) if err != nil { log.Error("create connect rpc service failed", zap.String("url", url), zap.Error(err)) } cli := pb.NewHMQServiceClient(conn) b.rpcClient[id] = cli } type HMQ struct { b *Broker } func (h *HMQ) QuerySubscribe(ctx context.Context, in *pb.QuerySubscribeRequest) (*pb.Response, error) { resp := &pb.Response{ RetCode: 0, } topic := in.Topic qos := in.Qos if qos > 1 { resp.RetCode = 404 return resp, nil } b := h.b var subs []interface{} var qoss []byte err := b.topicsMgr.Subscribers([]byte(topic), byte(qos), &subs, &qoss) if err != nil { log.Error("search sub client error, ", zap.Error(err)) resp.RetCode = 404 } if len(subs) == 0 { resp.RetCode = 404 } return resp, nil } func (h *HMQ) QueryConnect(ctx context.Context, in *pb.QueryConnectRequest) (*pb.Response, error) { resp := &pb.Response{ RetCode: 0, } b := h.b cli, exist := b.clients.Load(in.ClientID) if exist { client := cli.(*client) client.Close() } return resp, nil } func (h *HMQ) DeliverMessage(ctx context.Context, in *pb.DeliverMessageRequest) (*pb.Response, error) { b := h.b p := packets.NewControlPacket(packets.Publish).(*packets.PublishPacket) p.TopicName = in.Topic p.Payload = in.Payload p.Retain = false b.PublishMessage(p, false) resp := &pb.Response{ RetCode: 0, } return resp, nil } func (h *HMQ) QueryShareSubscribe(ctx context.Context, in *pb.QueryShareSubscribeRequest) (*pb.ShareSubscribeResponse, error) { resp := &pb.ShareSubscribeResponse{ RetCode: 0, } topic := in.Topic qos := in.Qos if qos > 1 { resp.RetCode = 404 return resp, nil } b := h.b var qoss []byte var subs []interface{} err := b.topicsMgr.Subscribers([]byte(topic), byte(qos), &subs, &qoss) if err != nil { log.Error("search sub client error, ", zap.Error(err)) resp.RetCode = 404 } if len(subs) == 0 { resp.RetCode = 404 } var qsub int32 for _, sub := range subs { s, ok := sub.(*subscription) if ok { if s.share { qsub++ } } } resp.ShareSubCount = qsub return resp, nil } func (b *Broker) DeliverMessage(packet *packets.PublishPacket) { for _, client := range b.rpcClient { resp, err := client.QuerySubscribe(context.Background(), &pb.QuerySubscribeRequest{Topic: packet.TopicName, Qos: int32(packet.Qos)}) if err != nil { log.Error("rpc request error:", zap.Error(err)) continue } if resp.RetCode == 0 { client.DeliverMessage(context.Background(), &pb.DeliverMessageRequest{Topic: packet.TopicName, Payload: packet.Payload}) } } } func (b *Broker) QueryConnect(clientID string) { for _, client := range b.rpcClient { client.QueryConnect(context.Background(), &pb.QueryConnectRequest{ClientID: clientID}) } } func (b *Broker) QueryShareSubscribe(topic string, qos byte) map[string]int32 { result := make(map[string]int32) for id, client := range b.rpcClient { resp, err := client.QueryShareSubscribe(context.Background(), &pb.QueryShareSubscribeRequest{Topic: topic, Qos: int32(qos)}) if err != nil { log.Error("rpc request error:", zap.Error(err)) continue } if resp.ShareSubCount > 0 { result[id] = resp.ShareSubCount } } return result }