fix userver: fix streaming middlewares integration

1. Removed special error handling logic from HandleRequestStream because it is not needed; errors are correctly handled by common middleware
2. Changed so that all middleware headers are placed before the user-defined part of handling (because for streaming headers are sent after the first chunk that happens in the user-defined part). To support this, added to HttpResponse different sctructures for system and user headers, so that while handling errors user headers can be erased, while leaving system headers
commit_hash:7a3ae6ec36fea916afd4a31407e5842271a9d112

Author: nmolotilov

core/functional_tests/basic_chaos/tests/httpclient/test_httpclient_stream.py

@@ -121,3 +121,24 @@ async def mock(request):
     response = await call()
     assert response.status == 200
     assert mock.times_called > 1
+
+
+async def test_required_headers(call, gate, mockserver):
+    @mockserver.handler('/test')
+    async def mock(request):
+        return mockserver.make_response('OK!')
+
+    required_headers = ['X-YaRequestId', 'X-YaSpanId', 'X-YaTraceId']
+
+    response = await call()
+    assert response.status == 200
+    assert all(key in response.headers for key in required_headers)
+
+    await gate.stop_accepting()
+    await gate.sockets_close()  # close keepalive connections
+    assert gate.connections_count() == 0
+
+    response = await call()
+    assert response.status == 500
+    assert all(key in response.headers for key in required_headers)
+    assert gate.connections_count() == 0

core/include/userver/server/handlers/http_handler_base.hpp

@@ -97,11 +97,19 @@ class HttpHandlerBase : public HandlerBase {
 
     /// Takes the exception and formats it into response, as specified by
     /// exception.
-    void HandleCustomHandlerException(const http::HttpRequest& request, const CustomHandlerException& ex) const;
+    void HandleCustomHandlerException(
+        const http::HttpRequest& request,
+        request::RequestContext& context,
+        const CustomHandlerException& ex
+    ) const;
 
     /// Takes the exception and formats it into response as an internal server
     /// error.
-    void HandleUnknownException(const http::HttpRequest& request, const std::exception& ex) const;
+    void HandleUnknownException(
+        const http::HttpRequest& request,
+        request::RequestContext& context,
+        const std::exception& ex
+    ) const;
 
     /// Helper function to log an unknown exception
     void LogUnknownException(const std::exception& ex, std::optional<logging::Level> log_level_override = {}) const;

core/include/userver/server/http/http_request.hpp

@@ -293,7 +293,7 @@ class HttpRequest final {
     friend class HttpRequestHandler;
 
     struct Impl;
-    utils::FastPimpl<Impl, 1648, 16> pimpl_;
+    utils::FastPimpl<Impl, 1904, 16> pimpl_;
 };
 
 }  // namespace server::http

core/include/userver/server/http/http_response.hpp

@@ -66,12 +66,14 @@ class HttpResponse final : public request::ResponseBase {
 
     /// @brief Add a new response header or rewrite an existing one.
     /// @returns true if the header was set. Returns false if headers
-    /// were already sent for stream'ed response and the new header was not set.
+    /// were already sent for stream'ed response and the new header was not set,
+    /// or the header overrides system header.
     bool SetHeader(std::string name, std::string value);
 
     /// @brief Add a new response header or rewrite an existing one.
     /// @returns true if the header was set. Returns false if headers
-    /// were already sent for stream'ed response and the new header was not set.
+    /// were already sent for stream'ed response and the new header was not set,
+    /// or the header overrides system header.
     bool SetHeader(std::string_view name, std::string value);
 
     /// @overload
@@ -88,10 +90,15 @@ class HttpResponse final : public request::ResponseBase {
     /// were already sent for stream'ed response and the new status was not set.
     bool SetStatus(HttpStatus status);
 
-    /// @brief Remove all headers from response.
+    /// @brief Set the end of system headers.
+    /// All headers written before this call are considered system; after - user.
+    /// User headers can't overwrite system headers.
+    void SetSystemHeadersEnd();
+
+    /// @brief Remove all headers from response, except system headers.
     /// @returns true if the headers were cleared. Returns false if headers
     /// were already sent for stream'ed response and the headers were not cleared.
-    bool ClearHeaders();
+    bool ClearUserHeaders();
 
     /// @brief Sets a cookie if it was not set before.
     void SetCookie(Cookie cookie);
@@ -102,8 +109,11 @@ class HttpResponse final : public request::ResponseBase {
     /// @return HTTP response status
     HttpStatus GetStatus() const { return status_; }
 
-    /// @return List of HTTP headers names.
-    HeadersMapKeys GetHeaderNames() const;
+    /// @return List of HTTP system headers names.
+    HeadersMapKeys GetSystemHeaderNames() const;
+
+    /// @return List of HTTP user headers names.
+    HeadersMapKeys GetUserHeaderNames() const;
 
     /// @return Value of the header with case insensitive name header_name, or an
     /// empty string if no such header.
@@ -155,9 +165,11 @@ class HttpResponse final : public request::ResponseBase {
 
     const HttpRequest& request_;
     HttpStatus status_ = HttpStatus::kOk;
-    HeadersMap headers_;
+    HeadersMap system_headers_;
+    HeadersMap user_headers_;
     CookiesMap cookies_;
 
+    bool system_headers_ended_ = false;
     engine::SingleConsumerEvent headers_end_{engine::SingleConsumerEvent::NoAutoReset()};
     std::optional<Queue::Consumer> body_stream_;
     Producer body_stream_producer_;

core/include/userver/server/http/http_response_body_stream.hpp

@@ -20,8 +20,13 @@ class ResponseBodyStream final {
 
     // Send a chunk of response data. It may NOT generate
     // exactly one HTTP chunk per call to PushBodyChunk().
+    // May not be called after SetBody().
     void PushBodyChunk(std::string&& chunk, engine::Deadline deadline);
 
+    // Set full response body instead of sending chunks.
+    // May not be called after PushBodyChunk().
+    void SetBody(std::string&& body);
+
     void SetHeader(const std::string&, const std::string&);
 
     void SetHeader(std::string_view, const std::string&);
@@ -38,6 +43,7 @@ class ResponseBodyStream final {
     ResponseBodyStream(HttpResponse::Producer&& queue_producer, HttpResponse& http_response);
 
     bool headers_ended_{false};
+    bool headers_end_sent_{false};
     HttpResponse::Producer queue_producer_;
     HttpResponse& http_response_;
 };

core/src/server/handlers/http_handler_base.cpp

@@ -74,21 +74,6 @@ void SetFormattedErrorResponse(http::HttpResponse& http_response, FormattedError
     }
 }
 
-void SetFormattedErrorResponse(
-    server::http::ResponseBodyStream& response_body_stream,
-    FormattedErrorData&& formatted_error,
-    engine::Deadline deadline
-) {
-    if (formatted_error.content_type) {
-        response_body_stream.SetHeader(
-            USERVER_NAMESPACE::http::headers::kContentType,
-            std::move(formatted_error.content_type->ToString())
-        );
-    }
-    response_body_stream.SetEndOfHeaders();
-    response_body_stream.PushBodyChunk(std::move(formatted_error.external_body), deadline);
-}
-
 std::unordered_map<int, logging::Level> ParseStatusCodesLogLevel(
     const std::unordered_map<std::string, std::string>& codes
 ) {
@@ -211,51 +196,15 @@ HttpHandlerBase::HttpHandlerBase(
 HttpHandlerBase::~HttpHandlerBase() { statistics_holder_.Unregister(); }
 
 void HttpHandlerBase::HandleRequestStream(http::HttpRequest& http_request, request::RequestContext& context) const {
-    auto& response = http_request.GetHttpResponse();
-    const utils::ScopeGuard scope([&response] { response.SetHeadersEnd(); });
+    auto& http_response = http_request.GetHttpResponse();
+    server::http::ResponseBodyStream response_body_stream(http_response.GetBodyProducer(), http_response);
 
-    server::http::ResponseBodyStream response_body_stream{response.GetBodyProducer(), response};
-
-    // Just in case HandleStreamRequest() throws an exception.
-    // Though it can be changed in HandleStreamRequest().
-    response_body_stream.SetStatusCode(500);
-
-    try {
-        HandleStreamRequest(http_request, context, response_body_stream);
-    } catch (const CustomHandlerException& e) {
-        response_body_stream.SetStatusCode(http::GetHttpStatus(e));
-
-        for (const auto& [name, value] : e.GetExtraHeaders()) {
-            response_body_stream.SetHeader(name, value);
-        }
-        if (e.IsExternalErrorBodyFormatted()) {
-            response_body_stream.SetEndOfHeaders();
-            response_body_stream.PushBodyChunk(std::string{e.GetExternalErrorBody()}, engine::Deadline());
-        } else {
-            auto formatted_error = GetFormattedExternalErrorBody(e);
-            SetFormattedErrorResponse(response_body_stream, std::move(formatted_error), engine::Deadline());
-        }
-    } catch (const std::exception& e) {
-        if (engine::current_task::ShouldCancel()) {
-            LOG_WARNING()
-                << "request task cancelled, exception in '" << HandlerName() << "' handler in handle_request: " << e;
-            response_body_stream.SetStatusCode(http::HttpStatus::kClientClosedRequest);
-        } else {
-            LOG_ERROR() << "exception in '" << HandlerName() << "' handler in handle_request: " << e;
-            response_body_stream.SetStatusCode(500);
-            SetFormattedErrorResponse(
-                response,
-                GetFormattedExternalErrorBody({
-                    HandlerErrorCode::kServerSideError,
-                    ExternalBody{response.GetData()},
-                })
-            );
-        }
-    }
+    HandleStreamRequest(http_request, context, response_body_stream);
 }
 
 void HttpHandlerBase::HandleHttpRequest(http::HttpRequest& http_request, request::RequestContext& context) const {
     auto& response = http_request.GetHttpResponse();
+    response.SetSystemHeadersEnd();
 
     // Don't hold the config snapshot for too long, especially with streaming.
     context.GetInternalContext().ResetConfigSnapshot();
@@ -273,6 +222,7 @@ void HttpHandlerBase::PrepareAndHandleRequest(http::HttpRequest& http_request, r
     auto& response = http_request.GetHttpResponse();
 
     context.GetInternalContext().SetConfigSnapshot(config_source_.GetSnapshot());
+    SetResponseServerHostname(response);
     try {
         UASSERT(first_middleware_);
         first_middleware_->HandleRequest(http_request, context);
@@ -286,7 +236,6 @@ void HttpHandlerBase::PrepareAndHandleRequest(http::HttpRequest& http_request, r
         response.SetStatus(http::HttpStatus::kInternalServerError);
     }
 
-    SetResponseServerHostname(response);
     response.SetHeadersEnd();
 }
 
@@ -434,8 +383,11 @@ std::string HttpHandlerBase::GetUrlForLoggingChecked(const http::HttpRequest& re
     }
 }
 
-void HttpHandlerBase::HandleCustomHandlerException(const http::HttpRequest& request, const CustomHandlerException& ex)
-    const {
+void HttpHandlerBase::HandleCustomHandlerException(
+    const http::HttpRequest& request,
+    request::RequestContext&,
+    const CustomHandlerException& ex
+) const {
     auto http_status = http::GetHttpStatus(ex);
     const auto level = GetLogLevelForResponseStatus(http_status);
     LOG(level) << "custom handler exception in '" << HandlerName() << "' handler: msg=" << ex;
@@ -452,7 +404,11 @@ void HttpHandlerBase::HandleCustomHandlerException(const http::HttpRequest& requ
     }
 }
 
-void HttpHandlerBase::HandleUnknownException(const http::HttpRequest& request, const std::exception& ex) const {
+void HttpHandlerBase::HandleUnknownException(
+    const http::HttpRequest& request,
+    request::RequestContext&,
+    const std::exception& ex
+) const {
     LogUnknownException(ex);
 
     auto& response = request.GetHttpResponse();

core/src/server/http/http2_writer.cpp

@@ -144,23 +144,29 @@ class Http2ResponseWriter final {
 private:
     Http2HeaderWriter GetHeaders() const {
         // Preallocate space for all headers
-        Http2HeaderWriter header_writer{response_.headers_.size() + response_.cookies_.size() + 3};
+        Http2HeaderWriter header_writer{
+            response_.system_headers_.size() + response_.user_headers_.size() + response_.cookies_.size() + 3
+        };
 
         header_writer.AddKeyValue(
             USERVER_NAMESPACE::http::headers::k2::kStatus,
             fmt::to_string(static_cast<std::uint16_t>(response_.status_))
         );
 
-        const auto& headers = response_.headers_;
-        const auto end = headers.cend();
-        if (headers.find(USERVER_NAMESPACE::http::headers::kDate) == end) {
+        if (!response_.HasHeader(USERVER_NAMESPACE::http::headers::kDate)) {
             header_writer.AddKeyValue(USERVER_NAMESPACE::http::headers::kDate, std::string{impl::GetCachedDate()});
         }
-        if (headers.find(USERVER_NAMESPACE::http::headers::kContentType) == end) {
+        if (!response_.HasHeader(USERVER_NAMESPACE::http::headers::kContentType)) {
             header_writer.AddKeyValue(USERVER_NAMESPACE::http::headers::kContentType, kDefaultContentType);
         }
         const std::string_view content_length_header{USERVER_NAMESPACE::http::headers::kContentLength};
-        for (const auto& [key, value] : headers) {
+        for (const auto& [key, value] : response_.system_headers_) {
+            if (key == content_length_header) {
+                continue;
+            }
+            header_writer.AddKeyValue(key, value);
+        }
+        for (const auto& [key, value] : response_.user_headers_) {
             if (key == content_length_header) {
                 continue;
             }

core/src/server/http/http_request.cpp

@@ -323,12 +323,7 @@ void HttpRequest::MarkAsInternalServerError() const {
     // TODO : refactor, this being here is a bit ridiculous
     pimpl_->response.SetStatus(http::HttpStatus::kInternalServerError);
     pimpl_->response.SetData({});
-
-    std::string server_header = pimpl_->response.GetHeader(USERVER_NAMESPACE::http::headers::kServer);
-    pimpl_->response.ClearHeaders();
-    if (!server_header.empty()) {
-        pimpl_->response.SetHeader(USERVER_NAMESPACE::http::headers::kServer, std::move(server_header));
-    }
+    pimpl_->response.ClearUserHeaders();
 }
 
 void HttpRequest::SetHttpHandler(const handlers::HttpHandlerBase& handler) { pimpl_->handler = &handler; }