You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
46 lines
1.1 KiB
46 lines
1.1 KiB
package singleflight
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/go-kratos/kratos/v2/middleware"
|
|
"github.com/go-kratos/kratos/v2/transport"
|
|
"golang.org/x/sync/singleflight"
|
|
)
|
|
|
|
var singleflightGroup singleflight.Group
|
|
|
|
// 单飞 middleware.
|
|
/*
|
|
//服务端下使用，通过传入op名称，使用单飞：
|
|
singleflight.SingleFlight(
|
|
"/service.test1/GetCityName",
|
|
"/service.test2/GetAllCityName",
|
|
)
|
|
*/
|
|
func SingleFlight(ops ...string) middleware.Middleware {
|
|
return func(handler middleware.Handler) middleware.Handler {
|
|
return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
|
|
if tr, ok := transport.FromServerContext(ctx); ok {
|
|
if Contains(ops, tr.Operation()) {
|
|
cacheKey := fmt.Sprintf("%s %s", tr.Operation(), req)
|
|
reply, err, _ = singleflightGroup.Do(cacheKey, func() (interface{}, error) {
|
|
return handler(ctx, req)
|
|
})
|
|
return reply, err
|
|
}
|
|
}
|
|
return handler(ctx, req)
|
|
}
|
|
}
|
|
}
|
|
|
|
func Contains(elems []string, elem string) bool {
|
|
for _, e := range elems {
|
|
if elem == e {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|