lone-services/bff/internal/wsproxy/proxy.go

package wsproxy

import (
	"context"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/zrpc"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/reflect/protoreflect"
)

type StreamCreator func(ctx context.Context) (grpc.ClientStream, error)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func NewWsProxyHandler(createStream StreamCreator, newFrame func() protoreflect.ProtoMessage) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()

		stream, err := createStream(r.Context())
		if err != nil {
			return
		}
		defer stream.CloseSend()

		// 上行 ws -> grpc stream
		go func() {
			for {
				_, rawBytes, err := conn.ReadMessage()
				if err != nil {
					return
				}
				msg := newFrame()
				fd := msg.ProtoReflect().Descriptor().Fields().ByName("raw_body")
				msg.ProtoReflect().Set(fd, protoreflect.ValueOfBytes(rawBytes))
				_ = stream.SendMsg(msg)
			}
		}()

		// 下行 grpc stream -> ws
		for {
			msg := newFrame()
			err := stream.RecvMsg(msg)
			if err != nil {
				break
			}
			fd := msg.ProtoReflect().Descriptor().Fields().ByName("raw_body")
			val := msg.ProtoReflect().Get(fd)
			_ = conn.WriteMessage(websocket.TextMessage, val.Bytes())
		}
	}
}

// 全局rpc客户端管理器，线程安全
type RpcClientManager struct {
	mu         sync.RWMutex
	client     zrpc.Client
	ready      bool
	cfg        zrpc.RpcClientConf
	streamOpts []zrpc.ClientOption
}

func NewRpcClientManager(cfg zrpc.RpcClientConf, streamOpts ...zrpc.ClientOption) *RpcClientManager {
	m := &RpcClientManager{
		cfg:        cfg,
		streamOpts: streamOpts,
	}
	go m.loopReconnect()
	return m
}

// 尝试建立一次连接
func (m *RpcClientManager) tryConnect() error {
	cli, err := zrpc.NewClient(m.cfg, m.streamOpts...)
	if err != nil {
		logx.Errorf("tryConnect: create client failed, err=%v", err)
		return err
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.client = cli
	m.ready = true
	return nil
}

// 后台循环重试
func (m *RpcClientManager) loopReconnect() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		m.mu.RLock()
		ready := m.ready
		m.mu.RUnlock()
		if ready {
			logx.Debug("rpc client already ready, skip reconnect")
			continue
		}
		if err := m.tryConnect(); err != nil {
			logx.Errorf("loopReconnect reconnect failed: %v", err)
		}
	}
}

// 获取client，判断是否就绪
func (m *RpcClientManager) GetClient() (zrpc.Client, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.client, m.ready
}