schedule_converter/grpc/client.go

package grpc

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"sync"
	"time"

	pb "schedule_converter/proto/runner"

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

const defaultHeartbeatInterval = 30 * time.Second

// Client gRPC 客户端
type Client struct {
	serverAddr string
	runnerID   string
	version    string

	mu        sync.RWMutex
	conn      *grpc.ClientConn
	stream    pb.RunnerHub_ConnectClient
	connected bool

	handler           *TaskHandler
	heartbeatInterval time.Duration
}

// NewClient 创建新的 gRPC 客户端
func NewClient(serverAddr, runnerID, version string) *Client {
	return &Client{
		serverAddr:        serverAddr,
		runnerID:          runnerID,
		version:           version,
		heartbeatInterval: defaultHeartbeatInterval,
	}
}

// SetHandler 设置任务处理器
func (c *Client) SetHandler(handler *TaskHandler) {
	c.handler = handler
}

// Connect 连接到 gRPC 服务器
func (c *Client) Connect(ctx context.Context) error {
	slog.Info("正在连接 gRPC 服务器", "server", c.serverAddr)

	conn, err := grpc.Dial(c.serverAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return fmt.Errorf("连接失败: %w", err)
	}

	client := pb.NewRunnerHubClient(conn)
	stream, err := client.Connect(ctx)
	if err != nil {
		conn.Close()
		return fmt.Errorf("创建流失败: %w", err)
	}

	c.mu.Lock()
	c.conn = conn
	c.stream = stream
	c.connected = true
	c.mu.Unlock()

	if err := c.register(); err != nil {
		return err
	}

	slog.Info("gRPC 连接已建立")
	return nil
}

// register 发送注册请求
func (c *Client) register() error {
	capabilities := make(map[string]string, len(supportedTaskTypes()))
	for _, t := range supportedTaskTypes() {
		capabilities[t.String()] = "true"
	}
	msg := &pb.StreamMessage{
		Message: &pb.StreamMessage_Register{
			Register: &pb.RegisterRequest{
				RunnerId:     c.runnerID,
				Version:      c.version,
				Capabilities: capabilities,
			},
		},
	}
	return c.sendMessage(msg)
}

// Run 运行主循环，接收消息
func (c *Client) Run(ctx context.Context) error {
	defer c.markDisconnected()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		stream := c.getStream()
		if stream == nil {
			return fmt.Errorf("流未连接")
		}

		msg, err := stream.Recv()
		if err == io.EOF {
			slog.Info("服务器关闭了连接")
			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:
		slog.Info("注册成功",
			"message", m.RegisterResponse.Message,
			"session_id", m.RegisterResponse.SessionId,
			"heartbeat_interval_sec", m.RegisterResponse.HeartbeatInterval,
		)
		if m.RegisterResponse.HeartbeatInterval > 0 {
			c.heartbeatInterval = time.Duration(m.RegisterResponse.HeartbeatInterval) * time.Second
		}
	case *pb.StreamMessage_HeartbeatAck:
		slog.Debug("心跳确认", "latency_ms", m.HeartbeatAck.Latency)
	case *pb.StreamMessage_Task:
		slog.Info("收到任务", "task_id", m.Task.TaskId, "type", m.Task.Type.String())
		go c.handleTask(m.Task)
	case *pb.StreamMessage_TaskCancel:
		slog.Info("收到取消任务请求", "task_id", m.TaskCancel.TaskId, "reason", m.TaskCancel.Reason)
		if c.handler != nil {
			c.handler.CancelTask(m.TaskCancel.TaskId)
		}
	default:
		slog.Warn("收到未知消息类型", "type", fmt.Sprintf("%T", msg.Message))
	}
}

// handleTask 处理任务
func (c *Client) handleTask(task *pb.Task) {
	if c.handler == nil {
		slog.Error("未设置任务处理器")
		return
	}

	result := c.handler.Execute(task)

	err := c.sendMessage(&pb.StreamMessage{
		Message: &pb.StreamMessage_Result{Result: result},
	})

	if err != nil {
		slog.Error("发送任务结果失败", "task_id", result.TaskId, "error", err)
		return
	}
	slog.Info("任务结果已发送", "task_id", result.TaskId, "status", result.Status.String())
}

// sendHeartbeat 发送心跳
func (c *Client) sendHeartbeat() error {
	return c.sendMessage(&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 {
				slog.Warn("发送心跳失败", "error", err)
			}
		}
	}
}

// Close 关闭连接
func (c *Client) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.connected = false
	if c.stream != nil {
		_ = c.stream.CloseSend()
	}
	if c.conn != nil {
		err := c.conn.Close()
		c.conn = nil
		c.stream = nil
		return err
	}
	return nil
}

// sendMessage 加锁发送消息，确保 stream 仍然可用
func (c *Client) sendMessage(msg *pb.StreamMessage) error {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if !c.connected || c.stream == nil {
		return fmt.Errorf("流未连接")
	}
	return c.stream.Send(msg)
}

// getStream 读锁获取 stream
func (c *Client) getStream() pb.RunnerHub_ConnectClient {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.stream
}

// markDisconnected 标记连接已断开，避免后续发送
func (c *Client) markDisconnected() {
	c.mu.Lock()
	c.connected = false
	c.mu.Unlock()
}