schedule_converter/grpc/client.go

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_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
}