diff --git a/server.go b/server.go index e89c5ac6136c..d2840b715d22 100644 --- a/server.go +++ b/server.go @@ -31,6 +31,7 @@ import ( "strings" "sync" "sync/atomic" + "testing" "time" "golang.org/x/net/trace" @@ -66,6 +67,8 @@ const ( listenerAddressForServeHTTP = "listenerAddressForServeHTTP" ) +var goAwayTimeout = 1 * time.Second + func init() { internal.GetServerCredentials = func(srv *Server) credentials.TransportCredentials { return srv.opts.creds @@ -83,6 +86,10 @@ func init() { internal.BinaryLogger = binaryLogger internal.JoinServerOptions = newJoinServerOption internal.RecvBufferPool = recvBufferPool + + if testing.Testing() { + goAwayTimeout = 0 + } } var statusOK = status.New(codes.OK, "") @@ -1013,6 +1020,11 @@ func (s *Server) serveStreams(ctx context.Context, st transport.ServerTransport, } defer func() { + // The client might not be done reading all the data, which at this point + // is stored in the kernel TCP buffer, yet. If we close the connection right away + // the client will get RST and the request will fail. Delay closing for 1 sec. + // See more details in https://github.com/grpc/grpc-go/pull/6957 + time.Sleep(goAwayTimeout) st.Close(errors.New("finished serving streams for the server transport")) for _, sh := range s.opts.statsHandlers { sh.HandleConn(ctx, &stats.ConnEnd{})