Skip to content
21 changes: 18 additions & 3 deletions modules/graceful/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ var (
// PerWriteWriteTimeout timeout for writes
const PerWriteWriteTimeout = 5 * time.Second

// PerWriteWriteTimeoutKbRate is a timeout taking account of how much there is to be written
const PerWriteWriteTimeoutKbRate = 2 * time.Second

func init() {
DefaultMaxHeaderBytes = 0 // use http.DefaultMaxHeaderBytes - which currently is 1 << 20 (1MB)
}
Expand Down Expand Up @@ -249,13 +252,25 @@ func (wl *wrappedListener) File() (*os.File, error) {

type wrappedConn struct {
net.Conn
server *Server
closed *int32
server *Server
closed *int32
deadline time.Time
}

func (w wrappedConn) Write(p []byte) (n int, err error) {
if PerWriteWriteTimeout > 0 {
_ = w.Conn.SetWriteDeadline(time.Now().Add(PerWriteWriteTimeout))
minTimeout := PerWriteWriteTimeout/2 + time.Duration(len(p)/1024)*PerWriteWriteTimeoutKbRate
minDeadline := time.Now().Add(minTimeout)

if minTimeout < PerWriteWriteTimeout {
minDeadline = time.Now().Add(PerWriteWriteTimeout)
}

w.deadline = w.deadline.Add(minTimeout)
if minDeadline.After(w.deadline) {
w.deadline = minDeadline
}
_ = w.Conn.SetWriteDeadline(w.deadline)
}
return w.Conn.Write(p)
}
Expand Down