Golang的gRPC如何支持双向流 开发实时数据推送服务案例解析

syntax = "proto3";

package realtimeservice;

option go_package = "./realtimeservice";

service DataStreamService {
  rpc PushData(stream ClientMessage) returns (stream ServerMessage);
}

message ClientMessage {
  string client_id = 1;
  string payload = 2;
}

message ServerMessage {
  string server_id = 1;
  string data = 2;
  int64 timestamp = 3;
}

package main

import (
    "context"
    "fmt"
    "io"
    "log"
    "net"
    "time"

    pb "your_module_path/realtimeservice" // 替换为你的实际模块路径

    "google.golang.org/grpc"
)

type server struct {
    pb.UnimplementedDataStreamServiceServer
}

func (s *server) PushData(stream pb.DataStreamService_PushDataServer) error {
    log.Println("New bidirectional stream established.")

    // goroutine for receiving messages from client
    go func() {
        for {
            req, err := stream.Recv()
            if err == io.EOF {
                log.Println("Client closed the stream (EOF).")
                return
            }
            if err != nil {
                log.Printf("Error receiving from client: %v", err)
                return
            }
            log.Printf("Received from client %s: %s", req.GetClientId(), req.GetPayload())
            // 可以在这里处理客户端消息，比如转发给其他客户端或存储
        }
    }()

    // goroutine for sending messages to client (simulating real-time push)
    ticker := time.NewTicker(2 * time.Second)
    defer ticker.Stop()

    for {
        select {
        case <-stream.Context().Done(): // Stream context cancelled (client disconnected or server shutting down)
            log.Println("Stream context done, stopping server send.")
            return stream.Context().Err()
        case <-ticker.C:
            msg := &pb.ServerMessage{
                ServerId: "server-node-1",
                Data:     fmt.Sprintf("实时数据更新 %d", time.Now().Unix()),
                Timestamp: time.Now().UnixMilli(),
            }
            if err := stream.Send(msg); err != nil {
                log.Printf("Error sending to client: %v", err)
                return err // Error sending, close stream
            }
            log.Printf("Sent data to client: %s", msg.GetData())
        }
    }
}

func main() {
    lis, err := net.Listen("tcp", ":50051")
    if err != nil {
        log.Fatalf("failed to listen: %v", err)
    }
    s := grpc.NewServer()
    pb.RegisterDataStreamServiceServer(s, &server{})
    log.Println("gRPC server listening on :50051")
    if err := s.Serve(lis); err != nil {
        log.Fatalf("failed to serve: %v", err)
    }
}

package main

import (
    "context"
    "fmt"
    "io"
    "log"
    "time"

    pb "your_module_path/realtimeservice" // 替换为你的实际模块路径

    "google.golang.org/grpc"
    "google.golang.org/grpc/credentials/insecure"
)

func main() {
    conn, err := grpc.Dial("localhost:50051", grpc.WithTransportCredentials(insecure.NewCredentials()))
    if err != nil {
        log.Fatalf("did not connect: %v", err)
    }
    defer conn.Close()
    c := pb.NewDataStreamServiceClient(conn)

    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    stream, err := c.PushData(ctx)
    if err != nil {
        log.Fatalf("could not create stream: %v", err)
    }

    // goroutine for receiving messages from server
    go func() {
        for {
            in, err := stream.Recv()
            if err == io.EOF {
                log.Println("Server closed the stream (EOF).")
                return
            }
            if err != nil {
                log.Printf("Error receiving from server: %v", err)
                return
            }
            log.Printf("Received from server %s: %s (timestamp: %d)", in.GetServerId(), in.GetData(), in.GetTimestamp())
        }
    }()

    // goroutine for sending messages to server (simulating client updates)
    ticker := time.NewTicker(3 * time.Second)
    defer ticker.Stop()

    clientMsgCount := 0
    for {
        select {
        case <-ctx.Done(): // Client context cancelled
            log.Println("Client context done, stopping client send.")
            return
        case <-ticker.C:
            clientMsgCount++
            msg := &pb.ClientMessage{
                ClientId: "client-go-123",
                Payload:  fmt.Sprintf("客户端心跳或更新 %d", clientMsgCount),
            }
            if err := stream.Send(msg); err != nil {
                log.Printf("Error sending to server: %v", err)
                return // Error sending, close stream
            }
            log.Printf("Sent to server: %s", msg.GetPayload())
        }
    }
}