54 lines
1.1 KiB
Go
54 lines
1.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/go-kratos/kratos/v2/errors"
|
|
"github.com/go-kratos/kratos/v2/middleware"
|
|
)
|
|
|
|
// Timeout 超时中间件
|
|
func Timeout(timeout time.Duration) middleware.Middleware {
|
|
return func(handler middleware.Handler) middleware.Handler {
|
|
return func(ctx context.Context, req interface{}) (interface{}, error) {
|
|
ctx, cancel := context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
|
|
// 使用 buffered channel 避免 goroutine 泄漏
|
|
type result struct {
|
|
reply interface{}
|
|
err error
|
|
}
|
|
done := make(chan result, 1)
|
|
panicChan := make(chan interface{}, 1)
|
|
|
|
go func() {
|
|
defer func() {
|
|
if p := recover(); p != nil {
|
|
select {
|
|
case panicChan <- p:
|
|
default:
|
|
}
|
|
}
|
|
}()
|
|
|
|
reply, err := handler(ctx, req)
|
|
select {
|
|
case done <- result{reply: reply, err: err}:
|
|
default:
|
|
}
|
|
}()
|
|
|
|
select {
|
|
case p := <-panicChan:
|
|
panic(p)
|
|
case r := <-done:
|
|
return r.reply, r.err
|
|
case <-ctx.Done():
|
|
return nil, errors.GatewayTimeout("TIMEOUT", "请求超时")
|
|
}
|
|
}
|
|
}
|
|
}
|