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

     1  package serverinterceptors
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/lingyao2333/mo-zero/core/lang"
    10  	"github.com/lingyao2333/mo-zero/core/logx"
    11  	"github.com/lingyao2333/mo-zero/core/stat"
    12  	"github.com/lingyao2333/mo-zero/core/syncx"
    13  	"github.com/lingyao2333/mo-zero/core/timex"
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/peer"
    16  )
    17  
    18  const defaultSlowThreshold = time.Millisecond * 500
    19  
    20  var (
    21  	notLoggingContentMethods sync.Map
    22  	slowThreshold            = syncx.ForAtomicDuration(defaultSlowThreshold)
    23  )
    24  
    25  // DontLogContentForMethod disable logging content for given method.
    26  func DontLogContentForMethod(method string) {
    27  	notLoggingContentMethods.Store(method, lang.Placeholder)
    28  }
    29  
    30  // SetSlowThreshold sets the slow threshold.
    31  func SetSlowThreshold(threshold time.Duration) {
    32  	slowThreshold.Set(threshold)
    33  }
    34  
    35  // UnaryStatInterceptor returns a func that uses given metrics to report stats.
    36  func UnaryStatInterceptor(metrics *stat.Metrics) grpc.UnaryServerInterceptor {
    37  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
    38  		handler grpc.UnaryHandler) (resp interface{}, err error) {
    39  		startTime := timex.Now()
    40  		defer func() {
    41  			duration := timex.Since(startTime)
    42  			metrics.Add(stat.Task{
    43  				Duration: duration,
    44  			})
    45  			logDuration(ctx, info.FullMethod, req, duration)
    46  		}()
    47  
    48  		return handler(ctx, req)
    49  	}
    50  }
    51  
    52  func logDuration(ctx context.Context, method string, req interface{}, duration time.Duration) {
    53  	var addr string
    54  	client, ok := peer.FromContext(ctx)
    55  	if ok {
    56  		addr = client.Addr.String()
    57  	}
    58  
    59  	logger := logx.WithContext(ctx).WithDuration(duration)
    60  	_, ok = notLoggingContentMethods.Load(method)
    61  	if ok {
    62  		if duration > slowThreshold.Load() {
    63  			logger.Slowf("[RPC] slowcall - %s - %s", addr, method)
    64  		}
    65  	} else {
    66  		content, err := json.Marshal(req)
    67  		if err != nil {
    68  			logx.WithContext(ctx).Errorf("%s - %s", addr, err.Error())
    69  		} else if duration > slowThreshold.Load() {
    70  			logger.Slowf("[RPC] slowcall - %s - %s - %s", addr, method, string(content))
    71  		} else {
    72  			logger.Infof("%s - %s - %s", addr, method, string(content))
    73  		}
    74  	}
    75  }