Shutdown close channel returned by RequestCtx.Done

RequestCtx.Done() now returns a channel that will be closed when
Server.Shutdown() is called.

server.go

@@ -353,6 +353,7 @@ type Server struct {
 	mu   sync.Mutex
 	open int32
 	stop int32
+	done chan struct{}
 }
 
 // TimeoutHandler creates RequestHandler, which returns StatusRequestTimeout
@@ -1526,6 +1527,7 @@ func (s *Server) Serve(ln net.Listener) error {
 		}
 
 		s.ln = ln
+		s.done = make(chan struct{})
 	}
 	s.mu.Unlock()
 
@@ -1607,6 +1609,10 @@ func (s *Server) Shutdown() error {
 		return err
 	}
 
+	if s.done != nil {
+		close(s.done)
+	}
+
 	// Closing the listener will make Serve() call Stop on the worker pool.
 	// Setting .stop to 1 will make serveConn() break out of its loop.
 	// Now we just have to wait until all workers are done.
@@ -1824,11 +1830,6 @@ func (s *Server) serveConn(c net.Conn) error {
 		isHTTP11        bool
 	)
 	for {
-		if atomic.LoadInt32(&s.stop) == 1 {
-			err = nil
-			break
-		}
-
 		connRequestNum++
 		ctx.time = currentTime
 
@@ -2013,6 +2014,11 @@ func (s *Server) serveConn(c net.Conn) error {
 
 		currentTime = time.Now()
 		s.setState(c, StateIdle)
+
+		if atomic.LoadInt32(&s.stop) == 1 {
+			err = nil
+			break
+		}
 	}
 
 	if br != nil {
@@ -2297,11 +2303,8 @@ func (ctx *RequestCtx) Deadline() (deadline time.Time, ok bool) {
 // Done returns a channel that's closed when work done on behalf of this
 // context should be canceled. Done may return nil if this context can
 // never be canceled. Successive calls to Done return the same value.
-//
-// This method always returns nil and is only present to make
-// RequestCtx implement the context interface.
 func (ctx *RequestCtx) Done() <-chan struct{} {
-	return nil
+	return ctx.s.done
 }
 
 // Err returns a non-nil error value after Done is closed,

server_test.go

@@ -2722,12 +2722,11 @@ func TestServeConnMultiRequests(t *testing.T) {
 
 func TestShutdown(t *testing.T) {
 	ln := fasthttputil.NewInmemoryListener()
-	h := func(ctx *RequestCtx) {
-		time.Sleep(time.Millisecond * 500)
-		ctx.Success("aaa/bbb", []byte("real response"))
-	}
 	s := &Server{
-		Handler: h,
+		Handler: func(ctx *RequestCtx) {
+			time.Sleep(time.Millisecond * 500)
+			ctx.Success("aaa/bbb", []byte("real response"))
+		},
 	}
 	serveCh := make(chan struct{})
 	go func() {
@@ -2778,13 +2777,13 @@ func TestShutdown(t *testing.T) {
 		}
 	}
 }
+
 func TestShutdownReuse(t *testing.T) {
 	ln := fasthttputil.NewInmemoryListener()
-	h := func(ctx *RequestCtx) {
-		ctx.Success("aaa/bbb", []byte("real response"))
-	}
 	s := &Server{
-		Handler:     h,
+		Handler: func(ctx *RequestCtx) {
+			ctx.Success("aaa/bbb", []byte("real response"))
+		},
 		ReadTimeout: time.Second,
 		Logger:      &customLogger{}, // Ignore log output.
 	}
@@ -2825,6 +2824,39 @@ func TestShutdownReuse(t *testing.T) {
 	}
 }
 
+func TestShutdownDone(t *testing.T) {
+	ln := fasthttputil.NewInmemoryListener()
+	s := &Server{
+		Handler: func(ctx *RequestCtx) {
+			<-ctx.Done()
+			ctx.Success("aaa/bbb", []byte("real response"))
+		},
+	}
+	go func() {
+		if err := s.Serve(ln); err != nil {
+			t.Fatalf("unexepcted error: %s", err)
+		}
+	}()
+	conn, err := ln.Dial()
+	if err != nil {
+		t.Fatalf("unexepcted error: %s", err)
+	}
+	if _, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n")); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	go func() {
+		// Shutdown won't return if the connection doesn't close,
+		// which doesn't happen until we read the response.
+		if err := s.Shutdown(); err != nil {
+			t.Fatalf("unexepcted error: %s", err)
+		}
+	}()
+	// We can only reach this point and get a valid response
+	// if reading from ctx.Done() returned.
+	br := bufio.NewReader(conn)
+	verifyResponse(t, br, StatusOK, "aaa/bbb", "real response")
+}
+
 func verifyResponse(t *testing.T, r *bufio.Reader, expectedStatusCode int, expectedContentType, expectedBody string) {
 	var resp Response
 	if err := resp.Read(r); err != nil {