job_risk_third/internal/services/instrument.go

package services

import (
	"context"
	"fmt"
	"github.com/go-kit/kit/endpoint"
	"github.com/go-kit/kit/metrics"
	"github.com/juju/ratelimit"
	"golang.org/x/time/rate"
	"job_risk_third/internal/proto"
	"time"
)

var ErrLimitExceed = fmt.Errorf("Rate limit exceed!")

// NewTokenBucketLimitterWithJuju 使用juju/ratelimit创建限流中间件
func NewTokenBucketLimitterWithJuju(bkt *ratelimit.Bucket) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			if bkt.TakeAvailable(1) == 0 {
				return nil, ErrLimitExceed
			}
			return next(ctx, request)
		}
	}
}

// NewTokenBucketLimitterWithBuildIn 使用x/time/rate创建限流中间件
func NewTokenBucketLimitterWithBuildIn(bkt *rate.Limiter) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			if !bkt.Allow() {
				return nil, ErrLimitExceed
			}
			return next(ctx, request)
		}
	}
}

// metricMiddleware 定义监控中间件，嵌入Service
// 新增监控指标项：requestCount和requestLatency
type metricMiddleware struct {
	JobRiskThirdServices
	requestCount   metrics.Counter
	requestLatency metrics.Histogram
}

// Metrics 封装监控方法
func Metrics(requestCount metrics.Counter, requestLatency metrics.Histogram) ServiceMiddleware {
	return func(next JobRiskThirdServices) JobRiskThirdServices {
		return metricMiddleware{
			next,
			requestCount,
			requestLatency}
	}
}

func (mw metricMiddleware) PushFeatureDeal(req proto.PushFeatureDealRequest) error {
	defer func(beign time.Time) {
		lvs := []string{"method", "PushFeatureDeal"}
		mw.requestCount.With(lvs...).Add(1)
		mw.requestLatency.With(lvs...).Observe(time.Since(beign).Seconds())
	}(time.Now())

	b := mw.JobRiskThirdServices.PushFeatureDeal(req)
	return b
}

func (mw metricMiddleware) PushUserInfo(req proto.PushUserInfoRequest) error {
	defer func(beign time.Time) {
		lvs := []string{"method", "PushUserInfo"}
		mw.requestCount.With(lvs...).Add(1)
		mw.requestLatency.With(lvs...).Observe(time.Since(beign).Seconds())
	}(time.Now())
	return mw.JobRiskThirdServices.PushUserInfo(req)
}

func (mw metricMiddleware) PushFeatureInfo(req proto.PushFeatureInfoRequest) error {
	defer func(beign time.Time) {
		lvs := []string{"method", "PushFeatureInfo"}
		mw.requestCount.With(lvs...).Add(1)
		mw.requestLatency.With(lvs...).Observe(time.Since(beign).Seconds())
	}(time.Now())

	return mw.JobRiskThirdServices.PushFeatureInfo(req)
}