sig-pub/internal/gateway/fast_gateway.go

package gateway

import (
	"context"
	"encoding/json"
	"net/http"
	"net/url"
	"sig-pub/pkg/grpc/generic"
	"sig-pub/pkg/grpc/session"
	"sig-pub/pkg/resp"
	"sig-pub/pkg/utils/strs"
	"sig-pub/pkg/zlog"
	"strings"
	"time"

	"github.com/bytedance/sonic"
	"github.com/fasthttp/router"
	"github.com/savsgio/gotils/strconv"
	"github.com/valyala/fasthttp"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/protojson"
)

const (
	ReadTimeout  = 20 * time.Second
	WriteTimeout = 20 * time.Second
	IdleTimeout  = 120 * time.Second
)

type FastGatewayServer struct {
	sigServerURL             *url.URL
	grpcGenericClientFactory *generic.GrpcGenericClientFactory

	svr       *fasthttp.Server
	sigClient *fasthttp.HostClient
}

func NewFastGatewayServer(
	sigServerURL *url.URL,
	grpcGenericClientFactory *generic.GrpcGenericClientFactory,
) *FastGatewayServer {
	return &FastGatewayServer{
		sigServerURL:             sigServerURL,
		grpcGenericClientFactory: grpcGenericClientFactory,
	}
}

func (g *FastGatewayServer) Init() (err error) {
	// fasthttp router
	r := router.New()
	r.ANY("/api/sig/{path:*}", g.reverseProxyHandler)
	r.POST("/api/v1/{path:*}", g.reverseProxyGrpcGenericCall)

	g.svr = &fasthttp.Server{
		Handler:            r.Handler,
		ReadTimeout:        ReadTimeout,
		WriteTimeout:       WriteTimeout,
		IdleTimeout:        IdleTimeout,
		MaxConnsPerIP:      1024 * 4,
		MaxRequestsPerConn: 0, // 0 = unlimited
		ReduceMemoryUsage:  true,
	}

	// sig http service client
	g.sigClient = &fasthttp.HostClient{
		Addr:                          g.sigServerURL.Host,
		Name:                          "gateway",
		MaxConns:                      512, // 根据实际业务调
		ReadTimeout:                   ReadTimeout,
		WriteTimeout:                  WriteTimeout,
		MaxConnDuration:               10 * time.Minute,
		DisableHeaderNamesNormalizing: true,
	}
	return
}

func (g *FastGatewayServer) Run(addr string) (err error) {
	return g.svr.ListenAndServe(addr)
}

func (g *FastGatewayServer) reverseProxyHandler(ctx *fasthttp.RequestCtx) {
	path := ctx.UserValue("path")
	zlog.Infof("path: %s", path)

	req := &ctx.Request
	resp := &ctx.Response

	req.Header.Set("Connection", "keep-alive")
	err := g.sigClient.Do(req, resp)
	if err != nil {
		zlog.Errorf("reverse proxy error: url=%s, %v", strconv.B2S(ctx.RequestURI()), err)
	}
}

func responseJSON(ctx *fasthttp.RequestCtx, code int, data any) {
	body, err := sonic.Marshal(data)
	if err != nil {
		zlog.Errorf("marshal resp data error: url=%s, %v", strconv.B2S(ctx.RequestURI()), err)
		ctx.SetStatusCode(http.StatusInternalServerError)
		return
	}

	ctx.SetContentType("application/json; charset=utf-8")
	ctx.SetStatusCode(code)
	ctx.SetBody(body)
}

// reverseProxyGrpcGenericCall 代理各个grpc服务请求
func (s *FastGatewayServer) reverseProxyGrpcGenericCall(c *fasthttp.RequestCtx) {
	path := c.UserValue("path").(string)
	paths := strings.Split(path, "/")
	if len(paths) != 2 {
		responseJSON(c, http.StatusBadRequest, resp.Error("service not specified"))
		return
	}
	svr := strs.UpperInitialLetter(paths[0])
	method := strs.UpperInitialLetter(paths[1])

	if svr == "" || method == "" {
		responseJSON(c, http.StatusBadRequest, resp.Error("service not found"))
		return
	}
	if !strings.HasSuffix(svr, "Service") {
		svr += "Service"
	}
	// todo service white list

	// get request body
	jsonBody := c.Request.Body()

	ctx1, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	grpcGenericClient, err := s.grpcGenericClientFactory.GetClient(ctx1, svr)
	if err != nil {
		zlog.Error(err)
		responseJSON(c, http.StatusForbidden, resp.Error(err.Error()))
		return
	}

	// put session
	ctx := context.Background()
	ctx = session.PutSubject(ctx, session.NewRpcSubject("123456"))
	ctx, cancel = context.WithTimeout(ctx, time.Second*20)
	defer cancel()

	// todo config call options
	var opts []grpc.CallOption
	if svr == "ExchangeService" && method == "HistoryKline" {
		opts = append(opts, grpc.UseCompressor("snappy"))
	}
	// generic call with json
	rsp, err := grpcGenericClient.InvokeUnaryJsonBytes(ctx, method, jsonBody, opts...)
	if err != nil {
		if err == generic.ErrorMethodNotExists {
			responseJSON(c, http.StatusNotFound, resp.Error(err.Error()))
			return
		}
		responseJSON(c, http.StatusInternalServerError, resp.Error(err.Error()))
		return
	}

	// encode response
	bytes, err := protojson.MarshalOptions{
		UseProtoNames:   false, // false:lowerCamelCase, true:snake_case
		EmitUnpopulated: false, // 是否包含默认值
	}.Marshal(rsp)
	if err != nil {
		responseJSON(c, http.StatusInternalServerError, resp.Error(err.Error()))
		return
	}

	r := json.RawMessage(bytes)
	responseJSON(c, http.StatusOK, resp.Success(r))
}