fix gracefilly shutdown bug, issue valyala#958

alpha-baby · Feb 2, 2021 · aa7befc · aa7befc
1 parent 5661df8
commit aa7befc
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git 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.go b/server.go
@@ -2196,7 +2196,7 @@ func (s *Server) serveConn(c net.Conn) (err error) {
 			}
 		}
 
-		connectionClose = connectionClose || ctx.Response.ConnectionClose()
+		connectionClose = connectionClose || ctx.Response.ConnectionClose() || atomic.LoadInt32(&s.stop) == 1
 		if connectionClose {
 			ctx.Response.Header.SetCanonical(strConnection, strClose)
 		} else if !isHTTP11 {

diff --git a/server_test.go b/server_test.go
@@ -3120,7 +3120,8 @@ 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, "close")
 		clientCh <- struct{}{}
 	}()
 	time.Sleep(time.Millisecond * 100)
@@ -3472,7 +3473,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)
@@ -3482,6 +3483,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 {