102 lines
2.5 KiB
Go
102 lines
2.5 KiB
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:]
|
|
}
|