zouzehe
commented
1 year ago the main code of the timeout middleware
package gin
import (
"context"
ginfk "github.com/gin-gonic/gin"
"net/http"
"time"
"sync"
"bytes"
"fmt"
)
var (
defaultOptions Options
)
func init() {
defaultOptions = Options{
CallBack: nil,
DefaultMsg: `{"code": -1, "msg":"http: Handler timeout"}`,
Timeout: 3 * time.Second,
ErrorHttpCode: http.StatusServiceUnavailable,
CallBackResponse: nil,
}
}
func Timeout(opts ...Option) ginfk.HandlerFunc {
return func(c *ginfk.Context) {
// sync.Pool
buffer := GetBuff()
tw := &Writer{body: buffer, ResponseWriter: c.Writer,
h: make(http.Header)}
tw.Options = defaultOptions
// Loop through each option
for _, opt := range opts {
// Call the option giving the instantiated
opt(tw)
}
// wrap the request context with a timeout
ctx, cancel := context.WithTimeout(c.Request.Context(), tw.Timeout)
defer cancel()
c.Request = c.Request.WithContext(ctx)
// Channel capacity must be greater than 0.
// Otherwise, if the parent coroutine quit due to timeout,
// the child coroutine may never be able to quit.
finish := make(chan struct{}, 1)
panicChan := make(chan interface{}, 1)
go func() {
defer func() {
if p := recover(); p != nil {
panicChan <- p
}
}()
c.Next()
finish <- struct{}{}
}()
var err error
select {
case p := <-panicChan:
panic(p)
case <-ctx.Done():
tw.mu.Lock()
defer tw.mu.Unlock()
c.Writer = tw
tw.timedOut = true
tw.ResponseWriter.WriteHeader(tw.ErrorHttpCode)
if tw.CallBackResponse != nil {
body, ct := tw.CallBackResponse()
head := tw.ResponseWriter.Header()
head["Content-Type"] = ct
_, err = tw.ResponseWriter.Write(body)
} else {
_, err = tw.ResponseWriter.Write([]byte(tw.DefaultMsg))
}
if err != nil {
panic(err)
}
c.Abort()
// execute callback func
if tw.CallBack != nil {
tw.CallBack(c.Request.Clone(context.Background()))
}
// If timeout happen, the buffer cannot be cleared actively,
// but wait for the GC to recycle.
case <-finish:
tw.mu.Lock()
defer tw.mu.Unlock()
c.Writer = tw
dst := tw.ResponseWriter.Header()
for k, vv := range tw.Header() {
dst[k] = vv
}
//if !tw.wroteHeader {
// tw.code = http.StatusOK
//}
tw.ResponseWriter.WriteHeader(tw.code)
_, err = tw.ResponseWriter.Write(buffer.Bytes())
if err != nil {
panic(err)
}
PutBuff(buffer)
}
}
}
type Writer struct {
ginfk.ResponseWriter
// header
h http.Header
// body
body *bytes.Buffer
Options // TimeoutOptions in options.go
code int
mu sync.Mutex
timedOut bool
wroteHeader bool
}
func (tw *Writer) Write(b []byte) (int, error) {
tw.mu.Lock()
defer tw.mu.Unlock()
if tw.timedOut {
return 0, nil
}
return tw.body.Write(b)
}
func (tw *Writer) WriteHeader(code int) {
checkWriteHeaderCode(code)
tw.mu.Lock()
defer tw.mu.Unlock()
if tw.timedOut {
return
}
tw.writeHeader(code)
}
func (tw *Writer) writeHeader(code int) {
tw.wroteHeader = true
tw.code = code
}
func (tw *Writer) WriteHeaderNow() {}
func (tw *Writer) Header() http.Header {
return tw.h
}
func checkWriteHeaderCode(code int) {
if code < 100 || code > 999 {
panic(fmt.Sprintf("invalid WriteHeader code %v", code))
}
}
bbdshow
Description
Very randomly, I got my Gin application fatal error
How to reproduce
Expectations
Actual result
Environment