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

     1  package serverinterceptors
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/load"
     8  	"github.com/lingyao2333/mo-zero/core/stat"
     9  	"google.golang.org/grpc"
    10  )
    11  
    12  const serviceType = "rpc"
    13  
    14  var (
    15  	sheddingStat *load.SheddingStat
    16  	lock         sync.Mutex
    17  )
    18  
    19  // UnarySheddingInterceptor returns a func that does load shedding on processing unary requests.
    20  func UnarySheddingInterceptor(shedder load.Shedder, metrics *stat.Metrics) grpc.UnaryServerInterceptor {
    21  	ensureSheddingStat()
    22  
    23  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
    24  		handler grpc.UnaryHandler) (val interface{}, err error) {
    25  		sheddingStat.IncrementTotal()
    26  		var promise load.Promise
    27  		promise, err = shedder.Allow()
    28  		if err != nil {
    29  			metrics.AddDrop()
    30  			sheddingStat.IncrementDrop()
    31  			return
    32  		}
    33  
    34  		defer func() {
    35  			if err == context.DeadlineExceeded {
    36  				promise.Fail()
    37  			} else {
    38  				sheddingStat.IncrementPass()
    39  				promise.Pass()
    40  			}
    41  		}()
    42  
    43  		return handler(ctx, req)
    44  	}
    45  }
    46  
    47  func ensureSheddingStat() {
    48  	lock.Lock()
    49  	if sheddingStat == nil {
    50  		sheddingStat = load.NewSheddingStat(serviceType)
    51  	}
    52  	lock.Unlock()
    53  }