复制项目

2026-01-14 22:16:44 +08:00
parent e2577b8cee
commit e50142a3b9
691 changed files with 97009 additions and 1 deletions
--- a/pkg/common/startrpc/circuitbreaker.go
+++ b/pkg/common/startrpc/circuitbreaker.go
@@ -0,0 +1,107 @@
+package startrpc
+
+import (
+	"context"
+	"time"
+
+	"github.com/openimsdk/tools/log"
+	"github.com/openimsdk/tools/stability/circuitbreaker"
+	"github.com/openimsdk/tools/stability/circuitbreaker/sre"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type CircuitBreaker struct {
+	Enable  bool          `yaml:"enable"`
+	Success float64       `yaml:"success"` // success rate threshold (0.0-1.0)
+	Request int64         `yaml:"request"` // request threshold
+	Bucket  int           `yaml:"bucket"`  // number of buckets
+	Window  time.Duration `yaml:"window"`  // time window for statistics
+}
+
+func NewCircuitBreaker(config *CircuitBreaker) circuitbreaker.CircuitBreaker {
+	if !config.Enable {
+		return nil
+	}
+
+	return sre.NewSREBraker(
+		sre.WithWindow(config.Window),
+		sre.WithBucket(config.Bucket),
+		sre.WithSuccess(config.Success),
+		sre.WithRequest(config.Request),
+	)
+}
+
+func UnaryCircuitBreakerInterceptor(breaker circuitbreaker.CircuitBreaker) grpc.ServerOption {
+	if breaker == nil {
+		return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+			return handler(ctx, req)
+		})
+	}
+
+	return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+		if err := breaker.Allow(); err != nil {
+			log.ZWarn(ctx, "rpc circuit breaker open", err, "method", info.FullMethod)
+			return nil, status.Error(codes.Unavailable, "service unavailable due to circuit breaker")
+		}
+
+		resp, err = handler(ctx, req)
+
+		if err != nil {
+			if st, ok := status.FromError(err); ok {
+				switch st.Code() {
+				case codes.OK:
+					breaker.MarkSuccess()
+				case codes.InvalidArgument, codes.NotFound, codes.AlreadyExists, codes.PermissionDenied:
+					breaker.MarkSuccess()
+				default:
+					breaker.MarkFailed()
+				}
+			} else {
+				breaker.MarkFailed()
+			}
+		} else {
+			breaker.MarkSuccess()
+		}
+
+		return resp, err
+
+	})
+}
+
+func StreamCircuitBreakerInterceptor(breaker circuitbreaker.CircuitBreaker) grpc.ServerOption {
+	if breaker == nil {
+		return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+			return handler(srv, ss)
+		})
+	}
+
+	return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		if err := breaker.Allow(); err != nil {
+			log.ZWarn(ss.Context(), "rpc circuit breaker open", err, "method", info.FullMethod)
+			return status.Error(codes.Unavailable, "service unavailable due to circuit breaker")
+		}
+
+		err := handler(srv, ss)
+
+		if err != nil {
+			if st, ok := status.FromError(err); ok {
+				switch st.Code() {
+				case codes.OK:
+					breaker.MarkSuccess()
+				case codes.InvalidArgument, codes.NotFound, codes.AlreadyExists, codes.PermissionDenied:
+					breaker.MarkSuccess()
+				default:
+					breaker.MarkFailed()
+				}
+			} else {
+				breaker.MarkFailed()
+			}
+		} else {
+			breaker.MarkSuccess()
+		}
+
+		return err
+	})
+}
--- a/pkg/common/startrpc/mw.go
+++ b/pkg/common/startrpc/mw.go
@@ -0,0 +1,15 @@
+package startrpc
+
+import (
+	"context"
+
+	"git.imall.cloud/openim/open-im-server-deploy/pkg/authverify"
+	"google.golang.org/grpc"
+)
+
+func grpcServerIMAdminUserID(imAdminUserID []string) grpc.ServerOption {
+	return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+		ctx = authverify.WithIMAdminUserIDs(ctx, imAdminUserID)
+		return handler(ctx, req)
+	})
+}
--- a/pkg/common/startrpc/ratelimit.go
+++ b/pkg/common/startrpc/ratelimit.go
@@ -0,0 +1,70 @@
+package startrpc
+
+import (
+	"context"
+	"time"
+
+	"github.com/openimsdk/tools/log"
+	"github.com/openimsdk/tools/stability/ratelimit"
+	"github.com/openimsdk/tools/stability/ratelimit/bbr"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type RateLimiter struct {
+	Enable       bool
+	Window       time.Duration
+	Bucket       int
+	CPUThreshold int64
+}
+
+func NewRateLimiter(config *RateLimiter) ratelimit.Limiter {
+	if !config.Enable {
+		return nil
+	}
+
+	return bbr.NewBBRLimiter(
+		bbr.WithWindow(config.Window),
+		bbr.WithBucket(config.Bucket),
+		bbr.WithCPUThreshold(config.CPUThreshold),
+	)
+}
+
+func UnaryRateLimitInterceptor(limiter ratelimit.Limiter) grpc.ServerOption {
+	if limiter == nil {
+		return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+			return handler(ctx, req)
+		})
+	}
+
+	return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+		done, err := limiter.Allow()
+		if err != nil {
+			log.ZWarn(ctx, "rpc rate limited", err, "method", info.FullMethod)
+			return nil, status.Errorf(codes.ResourceExhausted, "rpc request rate limit exceeded: %v, please try again later", err)
+		}
+
+		defer done(ratelimit.DoneInfo{})
+		return handler(ctx, req)
+	})
+}
+
+func StreamRateLimitInterceptor(limiter ratelimit.Limiter) grpc.ServerOption {
+	if limiter == nil {
+		return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+			return handler(srv, ss)
+		})
+	}
+
+	return grpc.ChainStreamInterceptor(func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		done, err := limiter.Allow()
+		if err != nil {
+			log.ZWarn(ss.Context(), "rpc rate limited", err, "method", info.FullMethod)
+			return status.Errorf(codes.ResourceExhausted, "rpc request rate limit exceeded: %v, please try again later", err)
+		}
+		defer done(ratelimit.DoneInfo{})
+
+		return handler(srv, ss)
+	})
+}
--- a/pkg/common/startrpc/start.go
+++ b/pkg/common/startrpc/start.go
@@ -0,0 +1,321 @@
+// Copyright © 2023 OpenIM. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package startrpc
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"os"
+	"os/signal"
+	"reflect"
+	"strconv"
+	"syscall"
+	"time"
+
+	conf "git.imall.cloud/openim/open-im-server-deploy/pkg/common/config"
+	"github.com/openimsdk/tools/utils/datautil"
+	"github.com/openimsdk/tools/utils/jsonutil"
+	"github.com/openimsdk/tools/utils/network"
+	"google.golang.org/grpc/status"
+
+	kdisc "git.imall.cloud/openim/open-im-server-deploy/pkg/common/discovery"
+	"git.imall.cloud/openim/open-im-server-deploy/pkg/common/prommetrics"
+	"github.com/openimsdk/tools/discovery"
+	"github.com/openimsdk/tools/errs"
+	"github.com/openimsdk/tools/log"
+	grpccli "github.com/openimsdk/tools/mw/grpc/client"
+	grpcsrv "github.com/openimsdk/tools/mw/grpc/server"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+func init() {
+	prommetrics.RegistryAll()
+}
+
+func Start[T any](ctx context.Context, disc *conf.Discovery, circuitBreakerConfig *conf.CircuitBreaker, rateLimiterConfig *conf.RateLimiter, prometheusConfig *conf.Prometheus, listenIP,
+	registerIP string, autoSetPorts bool, rpcPorts []int, index int, rpcRegisterName string, notification *conf.Notification, config T,
+	watchConfigNames []string, watchServiceNames []string,
+	rpcFn func(ctx context.Context, config T, client discovery.SvcDiscoveryRegistry, server grpc.ServiceRegistrar) error,
+	options ...grpc.ServerOption) error {
+
+	if notification != nil {
+		conf.InitNotification(notification)
+	}
+
+	maxRequestBody := getConfigRpcMaxRequestBody(reflect.ValueOf(config))
+	shareConfig := getConfigShare(reflect.ValueOf(config))
+
+	log.ZDebug(ctx, "rpc start", "rpcMaxRequestBody", maxRequestBody, "rpcRegisterName", rpcRegisterName, "registerIP", registerIP, "listenIP", listenIP)
+
+	options = append(options,
+		grpcsrv.GrpcServerMetadataContext(),
+		grpcsrv.GrpcServerErrorConvert(),
+		grpcsrv.GrpcServerLogger(),
+		grpcsrv.GrpcServerRequestValidate(),
+		grpcsrv.GrpcServerPanicCapture(),
+	)
+	if shareConfig != nil && len(shareConfig.IMAdminUser.UserIDs) > 0 {
+		options = append(options, grpcServerIMAdminUserID(shareConfig.IMAdminUser.UserIDs))
+	}
+	var clientOptions []grpc.DialOption
+	if maxRequestBody != nil {
+		if maxRequestBody.RequestMaxBodySize > 0 {
+			options = append(options, grpc.MaxRecvMsgSize(maxRequestBody.RequestMaxBodySize))
+			clientOptions = append(clientOptions, grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(maxRequestBody.RequestMaxBodySize)))
+		}
+		if maxRequestBody.ResponseMaxBodySize > 0 {
+			options = append(options, grpc.MaxSendMsgSize(maxRequestBody.ResponseMaxBodySize))
+			clientOptions = append(clientOptions, grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(maxRequestBody.ResponseMaxBodySize)))
+		}
+	}
+
+	if circuitBreakerConfig != nil && circuitBreakerConfig.Enable {
+		cb := &CircuitBreaker{
+			Enable:  circuitBreakerConfig.Enable,
+			Success: circuitBreakerConfig.Success,
+			Request: circuitBreakerConfig.Request,
+			Bucket:  circuitBreakerConfig.Bucket,
+			Window:  circuitBreakerConfig.Window,
+		}
+
+		breaker := NewCircuitBreaker(cb)
+
+		options = append(options,
+			UnaryCircuitBreakerInterceptor(breaker),
+			StreamCircuitBreakerInterceptor(breaker),
+		)
+
+		log.ZInfo(ctx, "RPC circuit breaker enabled",
+			"service", rpcRegisterName,
+			"window", circuitBreakerConfig.Window,
+			"bucket", circuitBreakerConfig.Bucket,
+			"success", circuitBreakerConfig.Success,
+			"requestThreshold", circuitBreakerConfig.Request)
+	}
+
+	if rateLimiterConfig != nil && rateLimiterConfig.Enable {
+		rl := &RateLimiter{
+			Enable:       rateLimiterConfig.Enable,
+			Window:       rateLimiterConfig.Window,
+			Bucket:       rateLimiterConfig.Bucket,
+			CPUThreshold: rateLimiterConfig.CPUThreshold,
+		}
+
+		limiter := NewRateLimiter(rl)
+
+		options = append(options,
+			UnaryRateLimitInterceptor(limiter),
+			StreamRateLimitInterceptor(limiter),
+		)
+
+		log.ZInfo(ctx, "RPC rate limiter enabled",
+			"service", rpcRegisterName,
+			"window", rateLimiterConfig.Window,
+			"bucket", rateLimiterConfig.Bucket,
+			"cpuThreshold", rateLimiterConfig.CPUThreshold)
+	}
+
+	registerIP, err := network.GetRpcRegisterIP(registerIP)
+	if err != nil {
+		return err
+	}
+	var prometheusListenAddr string
+	if autoSetPorts {
+		prometheusListenAddr = net.JoinHostPort(listenIP, "0")
+	} else {
+		prometheusPort, err := datautil.GetElemByIndex(prometheusConfig.Ports, index)
+		if err != nil {
+			return err
+		}
+		prometheusListenAddr = net.JoinHostPort(listenIP, strconv.Itoa(prometheusPort))
+	}
+
+	watchConfigNames = append(watchConfigNames, conf.LogConfigFileName)
+
+	client, err := kdisc.NewDiscoveryRegister(disc, watchServiceNames)
+	if err != nil {
+		return err
+	}
+
+	defer client.Close()
+	client.AddOption(
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"LoadBalancingPolicy": "%s"}`, "round_robin")),
+
+		grpccli.GrpcClientLogger(),
+		grpccli.GrpcClientContext(),
+		grpccli.GrpcClientErrorConvert(),
+	)
+	if len(clientOptions) > 0 {
+		client.AddOption(clientOptions...)
+	}
+
+	ctx, cancel := context.WithCancelCause(ctx)
+
+	go func() {
+		sigs := make(chan os.Signal, 1)
+		signal.Notify(sigs, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)
+		select {
+		case <-ctx.Done():
+			return
+		case val := <-sigs:
+			log.ZDebug(ctx, "recv signal", "signal", val.String())
+			cancel(fmt.Errorf("signal %s", val.String()))
+		}
+	}()
+
+	if prometheusListenAddr != "" {
+		options = append(
+			options,
+			prommetricsUnaryInterceptor(rpcRegisterName),
+			prommetricsStreamInterceptor(rpcRegisterName),
+		)
+		prometheusListener, prometheusPort, err := listenTCP(prometheusListenAddr)
+		if err != nil {
+			return err
+		}
+		log.ZDebug(ctx, "prometheus start", "addr", prometheusListener.Addr(), "rpcRegisterName", rpcRegisterName)
+		target, err := jsonutil.JsonMarshal(prommetrics.BuildDefaultTarget(registerIP, prometheusPort))
+		if err != nil {
+			return err
+		}
+		if autoSetPorts {
+			if err = client.SetWithLease(ctx, prommetrics.BuildDiscoveryKey(rpcRegisterName, index), target, prommetrics.TTL); err != nil {
+				if !errors.Is(err, discovery.ErrNotSupported) {
+					return err
+				}
+			}
+		}
+		go func() {
+			err := prommetrics.Start(prometheusListener)
+			if err == nil {
+				err = fmt.Errorf("listener done")
+			}
+			cancel(fmt.Errorf("prommetrics %s %w", rpcRegisterName, err))
+		}()
+	}
+
+	var (
+		rpcServer       *grpc.Server
+		rpcGracefulStop chan struct{}
+	)
+
+	onGrpcServiceRegistrar := func(desc *grpc.ServiceDesc, impl any) {
+		if rpcServer != nil {
+			rpcServer.RegisterService(desc, impl)
+			return
+		}
+		var rpcListenAddr string
+		if autoSetPorts {
+			rpcListenAddr = net.JoinHostPort(listenIP, "0")
+		} else {
+			rpcPort, err := datautil.GetElemByIndex(rpcPorts, index)
+			if err != nil {
+				cancel(fmt.Errorf("rpcPorts index out of range %s %w", rpcRegisterName, err))
+				return
+			}
+			rpcListenAddr = net.JoinHostPort(listenIP, strconv.Itoa(rpcPort))
+		}
+		rpcListener, err := net.Listen("tcp", rpcListenAddr)
+		if err != nil {
+			cancel(fmt.Errorf("listen rpc %s %s %w", rpcRegisterName, rpcListenAddr, err))
+			return
+		}
+
+		rpcServer = grpc.NewServer(options...)
+		rpcServer.RegisterService(desc, impl)
+		rpcGracefulStop = make(chan struct{})
+		rpcPort := rpcListener.Addr().(*net.TCPAddr).Port
+		log.ZDebug(ctx, "rpc start register", "rpcRegisterName", rpcRegisterName, "registerIP", registerIP, "rpcPort", rpcPort)
+		grpcOpt := grpc.WithTransportCredentials(insecure.NewCredentials())
+		rpcGracefulStop = make(chan struct{})
+		go func() {
+			<-ctx.Done()
+			rpcServer.GracefulStop()
+			close(rpcGracefulStop)
+		}()
+		if err := client.Register(ctx, rpcRegisterName, registerIP, rpcListener.Addr().(*net.TCPAddr).Port, grpcOpt); err != nil {
+			cancel(fmt.Errorf("rpc register %s %w", rpcRegisterName, err))
+			return
+		}
+
+		go func() {
+			err := rpcServer.Serve(rpcListener)
+			if err == nil {
+				err = fmt.Errorf("serve end")
+			}
+			cancel(fmt.Errorf("rpc %s %w", rpcRegisterName, err))
+		}()
+	}
+
+	err = rpcFn(ctx, config, client, &grpcServiceRegistrar{onRegisterService: onGrpcServiceRegistrar})
+	if err != nil {
+		return err
+	}
+	<-ctx.Done()
+	log.ZDebug(ctx, "cmd wait done", "err", context.Cause(ctx))
+	if rpcGracefulStop != nil {
+		timeout := time.NewTimer(time.Second * 15)
+		defer timeout.Stop()
+		select {
+		case <-timeout.C:
+			log.ZWarn(ctx, "rcp graceful stop timeout", nil)
+		case <-rpcGracefulStop:
+			log.ZDebug(ctx, "rcp graceful stop done")
+		}
+	}
+	return context.Cause(ctx)
+}
+
+func listenTCP(addr string) (net.Listener, int, error) {
+	listener, err := net.Listen("tcp", addr)
+	if err != nil {
+		return nil, 0, errs.WrapMsg(err, "listen err", "addr", addr)
+	}
+	return listener, listener.Addr().(*net.TCPAddr).Port, nil
+}
+
+func prommetricsUnaryInterceptor(rpcRegisterName string) grpc.ServerOption {
+	getCode := func(err error) int {
+		if err == nil {
+			return 0
+		}
+		rpcErr, ok := err.(interface{ GRPCStatus() *status.Status })
+		if !ok {
+			return -1
+		}
+		return int(rpcErr.GRPCStatus().Code())
+	}
+	return grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+		resp, err := handler(ctx, req)
+		prommetrics.RPCCall(rpcRegisterName, info.FullMethod, getCode(err))
+		return resp, err
+	})
+}
+
+func prommetricsStreamInterceptor(rpcRegisterName string) grpc.ServerOption {
+	return grpc.ChainStreamInterceptor()
+}
+
+type grpcServiceRegistrar struct {
+	onRegisterService func(desc *grpc.ServiceDesc, impl any)
+}
+
+func (x *grpcServiceRegistrar) RegisterService(desc *grpc.ServiceDesc, impl any) {
+	x.onRegisterService(desc, impl)
+}
--- a/pkg/common/startrpc/tools.go
+++ b/pkg/common/startrpc/tools.go
@@ -0,0 +1,47 @@
+package startrpc
+
+import (
+	"reflect"
+
+	conf "git.imall.cloud/openim/open-im-server-deploy/pkg/common/config"
+)
+
+func getConfig[T any](value reflect.Value) *T {
+	for value.Kind() == reflect.Pointer {
+		value = value.Elem()
+	}
+	if value.Kind() == reflect.Struct {
+		num := value.NumField()
+		for i := 0; i < num; i++ {
+			field := value.Field(i)
+			for field.Kind() == reflect.Pointer {
+				field = field.Elem()
+			}
+			if field.Kind() == reflect.Struct {
+				if elem, ok := field.Interface().(T); ok {
+					return &elem
+				}
+				if elem := getConfig[T](field); elem != nil {
+					return elem
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func getConfigRpcMaxRequestBody(value reflect.Value) *conf.MaxRequestBody {
+	return getConfig[conf.MaxRequestBody](value)
+}
+
+func getConfigShare(value reflect.Value) *conf.Share {
+	return getConfig[conf.Share](value)
+}
+
+func getConfigRateLimiter(value reflect.Value) *conf.RateLimiter {
+	return getConfig[conf.RateLimiter](value)
+}
+
+func getConfigCircuitBreaker(value reflect.Value) *conf.CircuitBreaker {
+	return getConfig[conf.CircuitBreaker](value)
+}