Repository: valyala/fasthttp
Pull request #960 (2021-02-02 03:24:35)

fix gracefilly shutdown bug, issue #958
### unit test

```bash
$ go test . -v -run=^TestShutdown$
=== RUN   TestShutdown
=== PAUSE TestShutdown
=== CONT  TestShutdown
--- PASS: TestShutdown (0.51s)
PASS
ok      github.com/valyala/fasthttp     0.530s
```

erikdubbelboer: See: https://github.com/valyala/fasthttp/issues/958#issuecomment-774428181
alpha-baby: @erikdubbelboer can you fix this unit test, becase i'm not hava windwos PC.

Patch:
diff --git a/client.go b/client.go
index 3e9fd5f..9c16fcf 100644
--- a/client.go
+++ b/client.go
@@ -856,8 +856,6 @@ func clientGetURLDeadline(dst []byte, url string, deadline time.Time, c clientDo
 	}
 	ch = chv.(chan clientURLResponse)
 
-	req := AcquireRequest()
-
 	// Note that the request continues execution on ErrTimeout until
 	// client-specific ReadTimeout exceeds. This helps limiting load
 	// on slow hosts by MaxConns* concurrent requests.
@@ -865,29 +863,56 @@ func clientGetURLDeadline(dst []byte, url string, deadline time.Time, c clientDo
 	// Without this 'hack' the load on slow host could exceed MaxConns*
 	// concurrent requests, since timed out requests on client side
 	// usually continue execution on the host.
+
+	var mu sync.Mutex
+	var timedout, responded bool
+
 	go func() {
+		req := AcquireRequest()
+
 		statusCodeCopy, bodyCopy, errCopy := doRequestFollowRedirectsBuffer(req, dst, url, c)
-		ch <- clientURLResponse{
-			statusCode: statusCodeCopy,
-			body:       bodyCopy,
-			err:        errCopy,
+		mu.Lock()
+		{
+			if !timedout {
+				ch <- clientURLResponse{
+					statusCode: statusCodeCopy,
+					body:       bodyCopy,
+					err:        errCopy,
+				}
+				responded = true
+			}
 		}
+		mu.Unlock()
+
+		ReleaseRequest(req)
 	}()
 
 	tc := AcquireTimer(timeout)
 	select {
 	case resp := <-ch:
-		ReleaseRequest(req)
-		clientURLResponseChPool.Put(chv)
 		statusCode = resp.statusCode
 		body = resp.body
 		err = resp.err
 	case <-tc.C:
-		body = dst
-		err = ErrTimeout
+		mu.Lock()
+		{
+			if responded {
+				resp := <-ch
+				statusCode = resp.statusCode
+				body = resp.body
+				err = resp.err
+			} else {
+				timedout = true
+				err = ErrTimeout
+				body = dst
+			}
+		}
+		mu.Unlock()
 	}
 	ReleaseTimer(tc)
 
+	clientURLResponseChPool.Put(chv)
+
 	return statusCode, body, err
 }
 
@@ -1193,10 +1218,6 @@ func clientDoDeadline(req *Request, resp *Response, deadline time.Time, c client
 	}
 	ReleaseTimer(tc)
 
-	select {
-	case <-ch:
-	default:
-	}
 	errorChPool.Put(chv)
 
 	return err
diff --git a/server.go b/server.go
index ec83471..0cda0f8 100644
--- a/server.go
+++ b/server.go
@@ -375,11 +375,16 @@ type Server struct {
 	// which will close it when needed.
 	KeepHijackedConns bool
 
+
+	// CloseOnShutdown when true adds a `Connection: close` header when when the server is shutting down.
+	CloseOnShutdown bool
+
 	// StreamRequestBody enables request body streaming,
 	// and calls the handler sooner when given body is
 	// larger then the current limit.
 	StreamRequestBody bool
 
+
 	tlsConfig  *tls.Config
 	nextProtos map[string]ServeHandler
 
@@ -2221,6 +2226,7 @@ func (s *Server) serveConn(c net.Conn) (err error) {
 		}
 
 		connectionClose = connectionClose || ctx.Response.ConnectionClose()
+		connectionClose = connectionClose || ctx.Response.ConnectionClose() || (s.CloseOnShutdown && atomic.LoadInt32(&s.stop) == 1)
 		if connectionClose {
 			ctx.Response.Header.SetCanonical(strConnection, strClose)
 		} else if !isHTTP11 {


Test patch:
diff --git a/header_test.go b/header_test.go
index 0fcca30..f82d221 100644
--- a/header_test.go
+++ b/header_test.go
@@ -2479,6 +2479,12 @@ func verifyResponseHeader(t *testing.T, h *ResponseHeader, expectedStatusCode, e
 	}
 }
 
+func verifyResponseHeaderConnection(t *testing.T, h *ResponseHeader, expectConnection string) {
+	if string(h.Peek(HeaderConnection)) != expectConnection {
+		t.Fatalf("Unexpected Connection %q. Expected %q", h.Peek(HeaderConnection), expectConnection)
+	}
+}
+
 func verifyRequestHeader(t *testing.T, h *RequestHeader, expectedContentLength int,
 	expectedRequestURI, expectedHost, expectedReferer, expectedContentType string) {
 	if h.ContentLength() != expectedContentLength {
diff --git a/server_test.go b/server_test.go
index 91bc6a8..c7baa3e 100644
--- a/server_test.go
+++ b/server_test.go
@@ -3117,7 +3117,70 @@ func TestShutdown(t *testing.T) {
 			t.Errorf("unexpected error: %s", err)
 		}
 		br := bufio.NewReader(conn)
-		verifyResponse(t, br, StatusOK, "aaa/bbb", "real response")
+		resp := verifyResponse(t, br, StatusOK, "aaa/bbb", "real response")
+		verifyResponseHeaderConnection(t, &resp.Header, "")
+		clientCh <- struct{}{}
+	}()
+	time.Sleep(time.Millisecond * 100)
+	shutdownCh := make(chan struct{})
+	go func() {
+		if err := s.Shutdown(); err != nil {
+			t.Errorf("unexepcted error: %s", err)
+		}
+		shutdownCh <- struct{}{}
+	}()
+	done := 0
+	for {
+		select {
+		case <-time.After(time.Second):
+			t.Fatal("shutdown took too long")
+		case <-serveCh:
+			done++
+		case <-clientCh:
+			done++
+		case <-shutdownCh:
+			done++
+		}
+		if done == 3 {
+			return
+		}
+	}
+}
+
+func TestCloseOnShutdown(t *testing.T) {
+	t.Parallel()
+
+	ln := fasthttputil.NewInmemoryListener()
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			time.Sleep(time.Millisecond * 500)
+			ctx.Success("aaa/bbb", []byte("real response"))
+		},
+		CloseOnShutdown: true,
+	}
+	serveCh := make(chan struct{})
+	go func() {
+		if err := s.Serve(ln); err != nil {
+			t.Errorf("unexepcted error: %s", err)
+		}
+		_, err := ln.Dial()
+		if err == nil {
+			t.Error("server is still listening")
+		}
+		serveCh <- struct{}{}
+	}()
+	clientCh := make(chan struct{})
+	go func() {
+		conn, err := ln.Dial()
+		if err != nil {
+			t.Errorf("unexepcted error: %s", err)
+		}
+		if _, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n")); err != nil {
+			t.Errorf("unexpected error: %s", err)
+		}
+		br := bufio.NewReader(conn)
+		resp := verifyResponse(t, br, StatusOK, "aaa/bbb", "real response")
+		verifyResponseHeaderConnection(t, &resp.Header, "close")
 		clientCh <- struct{}{}
 	}()
 	time.Sleep(time.Millisecond * 100)
@@ -3580,7 +3643,7 @@ func TestIncompleteBodyReturnsUnexpectedEOF(t *testing.T) {
 	}
 }
 
-func verifyResponse(t *testing.T, r *bufio.Reader, expectedStatusCode int, expectedContentType, expectedBody string) {
+func verifyResponse(t *testing.T, r *bufio.Reader, expectedStatusCode int, expectedContentType, expectedBody string) *Response {
 	var resp Response
 	if err := resp.Read(r); err != nil {
 		t.Fatalf("Unexpected error when parsing response: %s", err)
@@ -3590,6 +3653,7 @@ func verifyResponse(t *testing.T, r *bufio.Reader, expectedStatusCode int, expec
 		t.Fatalf("Unexpected body %q. Expected %q", resp.Body(), []byte(expectedBody))
 	}
 	verifyResponseHeader(t, &resp.Header, expectedStatusCode, len(resp.Body()), expectedContentType)
+	return &resp
 }
 
 type readWriter struct {