rag/internal/cmd/rpc.go

package cmd

import (
	"context"
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/auth"
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/reflection"
	"leafdev.top/leaf/rag/api/library"
	libraryGw "leafdev.top/leaf/rag/api/library"
	grpc2 "leafdev.top/leaf/rag/internal/middleware/grpc"
	"leafdev.top/leaf/rag/internal/providers/jwks"
	"leafdev.top/leaf/rag/internal/services/libraryService"
	"net"
	"net/http"
	"strings"
	"sync"
)

var rpcCommand = &cobra.Command{
	Use: "rpc",

	Run: func(cmd *cobra.Command, args []string) {
		jwks.InitJwksRefresh()
		StartGRPC()
	},
}

func StartGRPC() {
	if config.ListenAddr.GRPC == "" {
		config.ListenAddr.GRPC = "0.0.0.0:8081"
	}
	if config.ListenAddr.HTTP == "" {
		config.ListenAddr.HTTP = "0.0.0.0:8080"
	}

	lis, err := net.Listen("tcp", config.ListenAddr.GRPC)
	if err != nil {
		panic("GRPC failed to listen: " + err.Error())
	}
	logger.Info("GRPC Server listening at " + config.ListenAddr.GRPC)

	var opts = []grpc.ServerOption{
		grpc.ChainUnaryInterceptor(
			logging.UnaryServerInterceptor(grpc2.ZapLogInterceptor()),
			auth.UnaryServerInterceptor(grpc2.JwtAuth),
		),
		grpc.ChainStreamInterceptor(
			logging.StreamServerInterceptor(grpc2.ZapLogInterceptor()),
			auth.StreamServerInterceptor(grpc2.JwtAuth),
		),
	}
	grpcServer := grpc.NewServer(opts...)

	reflection.Register(grpcServer)

	library.RegisterLibraryServiceServer(grpcServer, libraryService.LibraryService{})

	var wg = sync.WaitGroup{}

	wg.Add(1)
	// 同时启动 grpc 和 http
	go func() {
		err = grpcServer.Serve(lis)
		if err != nil {
			panic(err)
		}

		defer wg.Done()
	}()

	wg.Add(1)
	go func() {
		ctx := context.Background()
		ctx, cancel := context.WithCancel(ctx)
		defer cancel()
		defer wg.Done()

		mux := runtime.NewServeMux()
		clientOpts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
		err = libraryGw.RegisterLibraryServiceHandlerFromEndpoint(ctx, mux, "127.0.0.1:"+getPortFromAddr(config.ListenAddr.GRPC), clientOpts)
		if err != nil {
			panic(err)
		}
		logger.Info("GRPC Gateway listening at " + config.ListenAddr.HTTP)
		err = http.ListenAndServe(config.ListenAddr.HTTP, mux)
		if err != nil {
			panic(err)
			return
		}
	}()

	wg.Wait()
}

func getPortFromAddr(addr string) string {
	return addr[strings.LastIndex(addr, ":")+1:]
}