github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/serverinterceptors/breakerinterceptor.go (about)

     1  package serverinterceptors
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/lingyao2333/mo-zero/core/breaker"
     7  	"github.com/lingyao2333/mo-zero/zrpc/internal/codes"
     8  	"google.golang.org/grpc"
     9  )
    10  
    11  // StreamBreakerInterceptor is an interceptor that acts as a circuit breaker.
    12  func StreamBreakerInterceptor(svr interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo,
    13  	handler grpc.StreamHandler) (err error) {
    14  	breakerName := info.FullMethod
    15  	return breaker.DoWithAcceptable(breakerName, func() error {
    16  		return handler(svr, stream)
    17  	}, codes.Acceptable)
    18  }
    19  
    20  // UnaryBreakerInterceptor is an interceptor that acts as a circuit breaker.
    21  func UnaryBreakerInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
    22  	handler grpc.UnaryHandler) (resp interface{}, err error) {
    23  	breakerName := info.FullMethod
    24  	err = breaker.DoWithAcceptable(breakerName, func() error {
    25  		var err error
    26  		resp, err = handler(ctx, req)
    27  		return err
    28  	}, codes.Acceptable)
    29  
    30  	return resp, err
    31  }