hyapp-server/services/room-service/internal/transport/grpc/owner_forwarder.go

package grpc

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"hyapp/pkg/xerr"
	roomservice "hyapp/services/room-service/internal/room/service"
	"hyapp/services/room-service/internal/router"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	roomv1 "hyapp.local/api/proto/room/v1"
)

const (
	// ownerForwardDepthHeader 防止错误注册到负载均衡地址后出现递归转发。
	ownerForwardDepthHeader = "x-hyapp-room-owner-forward-depth"
	// maxOwnerForwardDepth 允许一次中转；超过后让本节点走本地逻辑并返回真实 owner 冲突。
	maxOwnerForwardDepth = 1
	// defaultOwnerForwardTimeout 限制转发占用调用方请求预算的最长时间。
	defaultOwnerForwardTimeout = 2 * time.Second
)

type ownerForwarder struct {
	// nodeID 是当前 room-service 实例 ID，用于判断自己是否就是 owner。
	nodeID string
	// directory 保存 room_id -> owner node_id 的短租约。
	directory router.Directory
	// registry 保存 owner node_id -> 实例 gRPC 地址的短租约。
	registry router.NodeRegistry
	// timeout 是单次 owner 转发的额外 deadline 上限。
	timeout time.Duration

	mu    sync.Mutex
	conns map[string]forwardConn
}

type forwardConn struct {
	addr string
	conn *grpc.ClientConn
}

func newOwnerForwarder(nodeID string, directory router.Directory, registry router.NodeRegistry, timeout time.Duration) *ownerForwarder {
	if timeout <= 0 {
		timeout = defaultOwnerForwardTimeout
	}
	return &ownerForwarder{
		nodeID:    strings.TrimSpace(nodeID),
		directory: directory,
		registry:  registry,
		timeout:   timeout,
		conns:     make(map[string]forwardConn),
	}
}

func forwardCommandToOwner[T any](f *ownerForwarder, ctx context.Context, appCode string, roomID string, call func(context.Context, roomv1.RoomCommandServiceClient) (T, error)) (T, bool, error) {
	return forwardToOwner(ctx, f, appCode, roomID, func(callCtx context.Context, conn *grpc.ClientConn) (T, error) {
		return call(callCtx, roomv1.NewRoomCommandServiceClient(conn))
	})
}

func forwardGuardToOwner[T any](f *ownerForwarder, ctx context.Context, appCode string, roomID string, call func(context.Context, roomv1.RoomGuardServiceClient) (T, error)) (T, bool, error) {
	return forwardToOwner(ctx, f, appCode, roomID, func(callCtx context.Context, conn *grpc.ClientConn) (T, error) {
		return call(callCtx, roomv1.NewRoomGuardServiceClient(conn))
	})
}

func forwardQueryToOwner[T any](f *ownerForwarder, ctx context.Context, appCode string, roomID string, call func(context.Context, roomv1.RoomQueryServiceClient) (T, error)) (T, bool, error) {
	return forwardToOwner(ctx, f, appCode, roomID, func(callCtx context.Context, conn *grpc.ClientConn) (T, error) {
		return call(callCtx, roomv1.NewRoomQueryServiceClient(conn))
	})
}

func forwardToOwner[T any](ctx context.Context, f *ownerForwarder, appCode string, roomID string, call func(context.Context, *grpc.ClientConn) (T, error)) (T, bool, error) {
	var zero T
	if f == nil || f.directory == nil || f.registry == nil || strings.TrimSpace(roomID) == "" {
		return zero, false, nil
	}
	if ownerForwardDepth(ctx) >= maxOwnerForwardDepth {
		// 已经中转过的请求继续走本地逻辑，避免注册错误造成递归 RPC。
		return zero, false, nil
	}

	now := time.Now().UTC()
	lease, exists, err := f.directory.Lookup(ctx, roomservice.RuntimeRoomKey(appCode, roomID))
	if err != nil {
		return zero, true, xerr.ToGRPCError(xerr.New(xerr.Unavailable, "room owner lookup failed"))
	}
	if !exists || !lease.ValidAt(now) || lease.NodeID == f.nodeID {
		return zero, false, nil
	}

	node, exists, err := f.registry.LookupNode(ctx, lease.NodeID)
	if err != nil {
		return zero, true, xerr.ToGRPCError(xerr.New(xerr.Unavailable, "room owner node lookup failed"))
	}
	if !exists || !node.ValidAt(now) {
		return zero, true, xerr.ToGRPCError(xerr.New(xerr.Unavailable, "room owner node is not registered"))
	}

	conn, err := f.connForNode(ctx, node)
	if err != nil {
		return zero, true, xerr.ToGRPCError(xerr.New(xerr.Unavailable, fmt.Sprintf("room owner node dial failed: %v", err)))
	}

	callCtx, cancel := context.WithTimeout(ctx, f.timeout)
	defer cancel()
	callCtx = metadata.AppendToOutgoingContext(callCtx, ownerForwardDepthHeader, strconv.Itoa(ownerForwardDepth(ctx)+1))
	result, err := call(callCtx, conn)
	return result, true, err
}

func (f *ownerForwarder) connForNode(_ context.Context, node router.NodeRegistration) (*grpc.ClientConn, error) {
	f.mu.Lock()
	defer f.mu.Unlock()

	if cached, ok := f.conns[node.NodeID]; ok && cached.addr == node.GRPCAddr {
		return cached.conn, nil
	}
	if cached, ok := f.conns[node.NodeID]; ok && cached.conn != nil {
		// 节点地址变化通常来自实例重建；关闭旧连接，避免继续向旧进程转发。
		_ = cached.conn.Close()
	}

	conn, err := grpc.Dial(node.GRPCAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, err
	}
	f.conns[node.NodeID] = forwardConn{addr: node.GRPCAddr, conn: conn}
	return conn, nil
}

func (f *ownerForwarder) close() {
	if f == nil {
		return
	}
	f.mu.Lock()
	defer f.mu.Unlock()

	for nodeID, cached := range f.conns {
		if cached.conn != nil {
			_ = cached.conn.Close()
		}
		delete(f.conns, nodeID)
	}
}

func ownerForwardDepth(ctx context.Context) int {
	var values []string
	if incoming, ok := metadata.FromIncomingContext(ctx); ok {
		values = incoming.Get(ownerForwardDepthHeader)
	}
	if len(values) == 0 {
		if outgoing, ok := metadata.FromOutgoingContext(ctx); ok {
			values = outgoing.Get(ownerForwardDepthHeader)
		}
	}
	if len(values) == 0 {
		return 0
	}
	depth, err := strconv.Atoi(strings.TrimSpace(values[len(values)-1]))
	if err != nil || depth < 0 {
		return 0
	}
	return depth
}