36 lines
1.0 KiB
Go
36 lines
1.0 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
|
|
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
|
)
|
|
|
|
// Timeout applies a cooperative deadline to handlers and downstream calls.
|
|
// It does not spawn a goroutine or race a handler's response. Business code must
|
|
// observe Context().Done(); an already committed response cannot be replaced.
|
|
func Timeout(duration time.Duration) httpx.Middleware {
|
|
if duration <= 0 {
|
|
panic("middleware: timeout must be positive")
|
|
}
|
|
return func(next httpx.Handler) httpx.Handler {
|
|
return func(c *httpx.Context) error {
|
|
parent := c.Context()
|
|
ctx, cancel := context.WithTimeout(parent, duration)
|
|
defer cancel()
|
|
c.WithContext(ctx)
|
|
defer c.WithContext(parent)
|
|
if ctx.Err() != nil {
|
|
return apierrors.New(504, "REQUEST_TIMEOUT", "request deadline exceeded")
|
|
}
|
|
err := next(c)
|
|
if ctx.Err() == context.DeadlineExceeded && !c.Responded() {
|
|
return apierrors.New(504, "REQUEST_TIMEOUT", "request deadline exceeded").WithCause(err)
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
}
|