package grpc import ( "context" "fmt" "io" "log" "sync" "time" pb "schedule_converter/proto/runner" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" ) // Client gRPC 客户端 type Client struct { serverAddr string runnerID string version string conn *grpc.ClientConn stream pb.RunnerHub_ConnectClient mu sync.Mutex handler *TaskHandler // 配置 heartbeatInterval time.Duration } // NewClient 创建新的 gRPC 客户端 func NewClient(serverAddr, runnerID, version string) *Client { return &Client{ serverAddr: serverAddr, runnerID: runnerID, version: version, heartbeatInterval: 30 * time.Second, } } // SetHandler 设置任务处理器 func (c *Client) SetHandler(handler *TaskHandler) { c.handler = handler } // Connect 连接到 gRPC 服务器 func (c *Client) Connect(ctx context.Context) error { log.Printf("正在连接到 gRPC 服务器: %s", c.serverAddr) conn, err := grpc.Dial(c.serverAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) if err != nil { return fmt.Errorf("连接失败: %w", err) } c.conn = conn client := pb.NewRunnerHubClient(conn) stream, err := client.Connect(ctx) if err != nil { return fmt.Errorf("创建流失败: %w", err) } c.stream = stream // 发送注册请求 if err := c.register(); err != nil { return err } log.Println("gRPC 连接成功") return nil } // register 发送注册请求 func (c *Client) register() error { msg := &pb.StreamMessage{ Message: &pb.StreamMessage_Register{ Register: &pb.RegisterRequest{ RunnerId: c.runnerID, Version: c.version, Capabilities: map[string]string{ "TASK_TYPE_GET_SCHEDULE": "true", "TASK_TYPE_GET_GRADES": "true", "TASK_TYPE_GET_EXAMS": "true", "TASK_TYPE_LOGIN": "true", }, }, }, } return c.stream.Send(msg) } // Run 运行主循环,接收消息 func (c *Client) Run(ctx context.Context) error { for { select { case <-ctx.Done(): return ctx.Err() default: msg, err := c.stream.Recv() if err == io.EOF { log.Println("服务器关闭了连接") return nil } if err != nil { return fmt.Errorf("接收消息失败: %w", err) } c.handleMessage(msg) } } } // handleMessage 处理接收到的消息 func (c *Client) handleMessage(msg *pb.StreamMessage) { switch m := msg.Message.(type) { case *pb.StreamMessage_RegisterResponse: log.Printf("注册成功: %s, 会话ID: %s, 心跳间隔: %d秒", m.RegisterResponse.Message, m.RegisterResponse.SessionId, m.RegisterResponse.HeartbeatInterval) if m.RegisterResponse.HeartbeatInterval > 0 { c.heartbeatInterval = time.Duration(m.RegisterResponse.HeartbeatInterval) * time.Second } case *pb.StreamMessage_HeartbeatAck: // 心跳确认,可以用于计算延迟 log.Printf("心跳确认,延迟: %dms", m.HeartbeatAck.Latency) case *pb.StreamMessage_Task: log.Printf("收到任务: %s, 类型: %v", m.Task.TaskId, m.Task.Type) go c.handleTask(m.Task) case *pb.StreamMessage_TaskCancel: log.Printf("收到取消任务请求: %s, 原因: %s", m.TaskCancel.TaskId, m.TaskCancel.Reason) if c.handler != nil { c.handler.CancelTask(m.TaskCancel.TaskId) } default: log.Printf("收到未知消息类型: %T", msg.Message) } } // handleTask 处理任务 func (c *Client) handleTask(task *pb.Task) { if c.handler == nil { log.Println("错误: 未设置任务处理器") return } result := c.handler.Execute(task) c.mu.Lock() err := c.stream.Send(&pb.StreamMessage{ Message: &pb.StreamMessage_Result{ Result: result, }, }) c.mu.Unlock() if err != nil { log.Printf("发送任务结果失败: %v", err) } else { log.Printf("任务结果已发送: %s, 状态: %v", result.TaskId, result.Status) } } // sendHeartbeat 发送心跳 func (c *Client) sendHeartbeat() error { c.mu.Lock() defer c.mu.Unlock() if c.stream == nil { return fmt.Errorf("流未连接") } return c.stream.Send(&pb.StreamMessage{ Message: &pb.StreamMessage_Heartbeat{ Heartbeat: &pb.Heartbeat{ Timestamp: time.Now().UnixMilli(), }, }, }) } // HeartbeatLoop 心跳循环 func (c *Client) HeartbeatLoop(ctx context.Context, interval time.Duration) { ticker := time.NewTicker(interval) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: if err := c.sendHeartbeat(); err != nil { log.Printf("发送心跳失败: %v", err) } } } } // Close 关闭连接 func (c *Client) Close() error { if c.stream != nil { c.stream.CloseSend() } if c.conn != nil { return c.conn.Close() } return nil }