feat rabbitmq: heartbeat and message headers support

- Add support for custom message headers
- Enable configurable heartbeat interval for connection health checks

#137

---

Pull Request resolved: #1129
commit_hash:c4879c75cbdd4d9ba225b84a8d6d13fe20746a58

Author: sav-da

.mapping.json

@@ -3825,6 +3825,7 @@
   "rabbitmq/include/userver/urabbitmq/typedefs.hpp":"taxi/uservices/userver/rabbitmq/include/userver/urabbitmq/typedefs.hpp",
   "rabbitmq/library.yaml":"taxi/uservices/userver/rabbitmq/library.yaml",
   "rabbitmq/src/tests/admin_rmqtest.cpp":"taxi/uservices/userver/rabbitmq/src/tests/admin_rmqtest.cpp",
+  "rabbitmq/src/tests/header_value_rmqtest.cpp":"taxi/uservices/userver/rabbitmq/src/tests/header_value_rmqtest.cpp",
   "rabbitmq/src/tests/invalid_auth_rmqtest.cpp":"taxi/uservices/userver/rabbitmq/src/tests/invalid_auth_rmqtest.cpp",
   "rabbitmq/src/tests/outage_rmqtest.cpp":"taxi/uservices/userver/rabbitmq/src/tests/outage_rmqtest.cpp",
   "rabbitmq/src/tests/publish_consume_rmqtest.cpp":"taxi/uservices/userver/rabbitmq/src/tests/publish_consume_rmqtest.cpp",
@@ -3860,6 +3861,8 @@
   "rabbitmq/src/urabbitmq/impl/amqp_connection_handler.hpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/amqp_connection_handler.hpp",
   "rabbitmq/src/urabbitmq/impl/deferred_wrapper.cpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/deferred_wrapper.cpp",
   "rabbitmq/src/urabbitmq/impl/deferred_wrapper.hpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/deferred_wrapper.hpp",
+  "rabbitmq/src/urabbitmq/impl/header_value.cpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/header_value.cpp",
+  "rabbitmq/src/urabbitmq/impl/header_value.hpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/header_value.hpp",
   "rabbitmq/src/urabbitmq/impl/io/socket_reader.cpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/io/socket_reader.cpp",
   "rabbitmq/src/urabbitmq/impl/io/socket_reader.hpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/io/socket_reader.hpp",
   "rabbitmq/src/urabbitmq/impl/response_awaiter.cpp":"taxi/uservices/userver/rabbitmq/src/urabbitmq/impl/response_awaiter.cpp",

rabbitmq/functional_tests/basic_chaos/rabbitmq_service.cpp

@@ -7,6 +7,8 @@
 #include <userver/components/component_context.hpp>
 #include <userver/components/minimal_server_component_list.hpp>
 #include <userver/concurrent/variable.hpp>
+#include <userver/formats/json.hpp>
+#include <userver/formats/parse/common_containers.hpp>
 #include <userver/formats/serialize/common_containers.hpp>
 #include <userver/server/handlers/http_handler_base.hpp>
 #include <userver/server/handlers/tests_control.hpp>
@@ -26,8 +28,7 @@ class ChaosProducer final : public components::LoggableComponentBase {
 
     ChaosProducer(const components::ComponentConfig& config, const components::ComponentContext& context)
         : components::LoggableComponentBase{config, context},
-          rabbit_client_{context.FindComponent<components::RabbitMQ>("chaos-rabbit").GetClient()}
-    {
+          rabbit_client_{context.FindComponent<components::RabbitMQ>("chaos-rabbit").GetClient()} {
         const auto setup_deadline = engine::Deadline::FromDuration(kDefaultOperationTimeout);
 
         auto admin_channel = rabbit_client_->GetAdminChannel(setup_deadline);
@@ -77,9 +78,7 @@ class ChaosConsumer final : public components::ComponentBase {
     static constexpr std::string_view kName{"chaos-consumer"};
 
     ChaosConsumer(const components::ComponentConfig& config, const components::ComponentContext& context)
-        : components::ComponentBase{config, context},
-          consumer_{config, context, messages_}
-    {
+        : components::ComponentBase{config, context}, consumer_{config, context, messages_} {
         Start();
     }
 
@@ -119,8 +118,7 @@ class ChaosConsumer final : public components::ComponentBase {
         )
             : urabbitmq::
                   ConsumerBase{context.FindComponent<components::RabbitMQ>(config["rabbit_name"].As<std::string>()).GetClient(), ParseSettings(config)},
-              messages_{messages}
-        {}
+              messages_{messages} {}
 
     protected:
         void Process(urabbitmq::ConsumedMessage msg) override {
@@ -150,8 +148,7 @@ class ChaosHandler final : public server::handlers::HttpHandlerBase {
     ChaosHandler(const components::ComponentConfig& config, const components::ComponentContext& context)
         : server::handlers::HttpHandlerBase{config, context},
           producer_{context.FindComponent<ChaosProducer>()},
-          consumer_{context.FindComponent<ChaosConsumer>()}
-    {}
+          consumer_{context.FindComponent<ChaosConsumer>()} {}
 
     std::string HandleRequestThrow(const server::http::HttpRequest& request, server::request::RequestContext&)
         const override {
@@ -178,6 +175,13 @@ class ChaosHandler final : public server::handlers::HttpHandlerBase {
             throw server::handlers::ClientError{server::handlers::ExternalBody{"No 'message' query argument"}};
         }
         urabbitmq::Envelope envelope{message, urabbitmq::MessageType::kTransient, {}, {}, {}};
+        if (!request.RequestBody().empty()) {
+            const auto request_json = formats::json::FromString(request.RequestBody());
+            if (request_json.HasMember("headers")) {
+                envelope
+                    .headers = request_json["headers"].As<std::unordered_map<std::string, urabbitmq::HeaderValue>>();
+            }
+        }
         const auto& correlation_id = request.GetArg("correlation_id");
         if (!correlation_id.empty()) {
             envelope.correlation_id = correlation_id;
@@ -219,16 +223,17 @@ class ChaosHandler final : public server::handlers::HttpHandlerBase {
     }
 
     std::string HandleGet() const {
-        formats::json::ValueBuilder messages_builder;
+        urabbitmq::HeaderValue::Builder messages_builder;
         for (const auto& item : consumer_.GetMessages()) {
-            formats::json::ValueBuilder item_builder;
+            urabbitmq::HeaderValue::Builder item_builder;
             item_builder["message"] = item.message;
             if (item.correlation_id.has_value()) {
                 item_builder["correlation_id"] = item.correlation_id;
             }
             if (item.reply_to.has_value()) {
                 item_builder["reply_to"] = item.reply_to;
             }
+            item_builder["headers"] = item.headers;
             messages_builder.PushBack(std::move(item_builder));
         }
         return formats::json::ToString(messages_builder.ExtractValue());

rabbitmq/functional_tests/basic_chaos/static_config.yaml

@@ -17,6 +17,7 @@ components_manager:
             min_pool_size: 1
             max_pool_size: 1
             max_in_flight_requests: 5
+            heartbeat_interval_seconds: 1
             use_secure_connection: false
 
         secdist: {}  # Component that stores configuration of hosts and passwords

rabbitmq/functional_tests/basic_chaos/tests/test_rabbitmq.py

@@ -37,6 +37,10 @@ async def _clear_messages(service_client):
     assert response.status_code == 200
 
 
+def _strip_headers(messages):
+    return [{key: value for key, value in message.items() if key != 'headers'} for message in messages]
+
+
 async def _publish_and_consume(testpoint, client):
     @testpoint('message_consumed')
     def message_consumed(data):
@@ -51,7 +55,7 @@ def message_consumed(data):
     response = await client.get('/v1/messages')
     assert response.status_code == 200
 
-    assert response.json() == MESSAGES
+    assert _strip_headers(response.json()) == MESSAGES
 
 
 async def test_rabbitmq_happy(testpoint, service_client, gate):
@@ -60,6 +64,55 @@ async def test_rabbitmq_happy(testpoint, service_client, gate):
     await _publish_and_consume(testpoint, service_client)
 
 
+async def test_rabbitmq_headers(testpoint, service_client, gate):
+    await _clear_messages(service_client)
+
+    @testpoint('message_consumed')
+    def message_consumed(data):
+        pass
+
+    expected_headers = {
+        'x-bool': True,
+        'x-int': -10,
+        'x-uint': 10,
+        'x-double': 2.5,
+        'x-array': [-7, 'array-value', {'enabled': False, 'nullable': None}],
+        'x-object': {
+            'count': 42,
+            'name': 'nested-object',
+            'array': [-7, 'array-value', {'enabled': False, 'nullable': None}],
+        },
+        'x-null': None,
+    }
+
+    response = await service_client.post(
+        '/v1/messages?message=headers&reliable=1&reply_to=reply&correlation_id=corr-id',
+        json={'headers': expected_headers},
+    )
+    assert response.status_code == 200
+
+    await message_consumed.wait_call()
+
+    response = await service_client.get('/v1/messages')
+    assert response.status_code == 200
+    messages = response.json()
+    assert len(messages) == 1
+
+    consumed = messages[0]
+    assert consumed['message'] == 'headers'
+    assert consumed['reply_to'] == 'reply'
+    assert consumed['correlation_id'] == 'corr-id'
+    assert consumed['headers']['x-bool'] is True
+    assert consumed['headers']['x-int'] == -10
+    assert consumed['headers']['x-uint'] == 10
+    assert consumed['headers']['x-double'] == 2.5
+    assert consumed['headers']['x-array'] == expected_headers['x-array']
+    assert consumed['headers']['x-object'] == expected_headers['x-object']
+    assert consumed['headers']['x-null'] is None
+    assert consumed['headers']['u-trace-id']
+    assert consumed['headers']['u-parent-span-id']
+
+
 @pytest.mark.skip(reason='std::terminate is called, fix in TAXICOMMON-6995')
 async def test_consumer_reconnects(testpoint, service_client, gate):
     await _clear_messages(service_client)
@@ -85,4 +138,42 @@ def message_consumed(data):
     response = await service_client.get('/v1/messages')
     assert response.status_code == 200
 
-    assert response.json() == MESSAGES
+    assert _strip_headers(response.json()) == MESSAGES
+
+
+async def test_rabbitmq_heartbeat_reconnects(testpoint, service_client, gate):
+    await _clear_messages(service_client)
+
+    @testpoint('message_consumed')
+    def message_consumed(data):
+        pass
+
+    response = await service_client.post('/v1/messages?message=before-heartbeat')
+    assert response.status_code == 200
+    await message_consumed.wait_call()
+
+    async with service_client.capture_logs(log_level='INFO') as capture:
+
+        @capture.subscribe(text="Consumer for queue 'chaos-queue' is broken, trying to restart")
+        def consumer_broken(**kwargs):
+            pass
+
+        @capture.subscribe(text='Restarted successfully')
+        def consumer_restarted(**kwargs):
+            pass
+
+        await gate.to_server_drop()
+        await asyncio.sleep(3.0)
+        await gate.to_server_pass()
+
+        await consumer_broken.wait_call()
+        await consumer_restarted.wait_call()
+
+        response = await service_client.post('/v1/messages?message=after-heartbeat')
+        assert response.status_code == 200
+
+        await message_consumed.wait_call()
+
+    response = await service_client.get('/v1/messages')
+    assert response.status_code == 200
+    assert any(message['message'] == 'after-heartbeat' for message in response.json())

rabbitmq/include/userver/urabbitmq/client_settings.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <chrono>
 #include <cstddef>
 #include <optional>
 #include <string>
@@ -74,6 +75,10 @@ struct PoolSettings final {
     /// (tcp error/protocol error/write timeout) leads to a errors burst:
     /// all outstanding request will fails at once
     size_t max_in_flight_requests = 5;
+
+    /// Requested AMQP heartbeat interval in seconds.
+    /// Set to 0 to disable heartbeats.
+    size_t heartbeat_interval_seconds = 60;
 };
 
 class TestsHelper;

rabbitmq/include/userver/urabbitmq/typedefs.hpp

@@ -4,7 +4,12 @@
 /// @brief Convenient typedefs for RabbitMQ entities.
 
 #include <chrono>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <unordered_map>
 
+#include <userver/formats/json/value.hpp>
 #include <userver/utils/strong_typedef.hpp>
 
 USERVER_NAMESPACE_BEGIN
@@ -63,6 +68,10 @@ enum class MessageType {
     kTransient,
 };
 
+/// JSON-like representation of an AMQP header value.
+/// This is not JSON, but a convenient tree representation for AMQP field values.
+using HeaderValue = formats::json::Value;
+
 /// @brief Structure holding an AMQP message body along with some of its
 /// metadata fields. This struct is used to pass messages to the end user,
 /// hiding the actual AMQP message object implementation.
@@ -75,6 +84,7 @@ struct ConsumedMessage {
     Metadata metadata;
     std::optional<std::string> reply_to{};
     std::optional<std::string> correlation_id{};
+    std::unordered_map<std::string, HeaderValue> headers{};
 };
 
 /// @brief Structure holding an AMQP message body along with some of its
@@ -86,6 +96,7 @@ struct Envelope {
     std::optional<std::string> reply_to{};
     std::optional<std::string> correlation_id{};
     std::optional<std::chrono::milliseconds> expiration{};
+    std::optional<std::unordered_map<std::string, HeaderValue>> headers{};
 };
 
 }  // namespace urabbitmq