rabbit_prometheus_handler: Stream identity-encoded HTTP replies

`prometheus_text_format:format/1` produces a binary of the format for
the entire registry. For clusters with many resources, this can lead to
large replies from `/metrics/[:registry]` especially for large
registries like `per-object`. Instead of formatting the response and
then sending it, we can stream the response by taking advantage of the
new `format_into/3` callback (which needs to be added upstream to the
`prometheus` dep). This uses `cowboy_req:stream_body/3` to stream the
iodata as `prometheus` works through the registry.

This should hopefully be a nice memory improvement. The other benefit
is that results are sent eagerly. For a stress-testing example,

1. `make run-broker`
2. `rabbitmqctl import_definitions path/to/100k-classic-queues.json`
3. `curl -s localhost:15692/metrics/per-object`

Before this change `curl` would wait for around 8 seconds and then the
entire response would arrive. With this change the results start
streaming in immediately.

Author: the-mikedavis

deps/rabbitmq_prometheus/src/rabbit_prometheus_handler.erl

@@ -89,47 +89,58 @@ gen_response(_, Request) ->
     Request.
 
 gen_metrics_response(Registry, Request) ->
-    {Code, RespHeaders, Body} = reply(Registry, Request),
-
-    Headers = to_cowboy_headers(RespHeaders),
-    cowboy_req:reply(Code, maps:from_list(Headers), Body, Request).
-
-to_cowboy_headers(RespHeaders) ->
-    lists:map(fun to_cowboy_headers_/1, RespHeaders).
-
-to_cowboy_headers_({Name, Value}) ->
-    {to_cowboy_name(Name), Value}.
-
-to_cowboy_name(Name) ->
-    binary:replace(atom_to_binary(Name, utf8), <<"_">>, <<"-">>).
+    reply(Registry, Request).
 
 reply(Registry, Request) ->
     case validate_registry(Registry, registry()) of
         {true, RealRegistry} ->
             format_metrics(Request, RealRegistry);
         {registry_conflict, _ReqR, _ConfR} ->
-            {409, [], <<>>}
+            cowboy_req:reply(409, #{}, <<>>, Request)
     end.
 
 format_metrics(Request, Registry) ->
     AcceptEncoding = cowboy_req:header(<<"accept-encoding">>, Request, undefined),
     ContentType = prometheus_text_format:content_type(),
-    Scrape = render_format(ContentType, Registry),
     Encoding = accept_encoding_header:negotiate(AcceptEncoding, [<<"identity">>,
                                                                  <<"gzip">>]),
-    encode_format(ContentType, binary_to_list(Encoding), Scrape, Registry).
-
-render_format(ContentType, Registry) ->
-    Scrape = prometheus_summary:observe_duration(
-               Registry,
-               ?SCRAPE_DURATION,
-               [Registry, ContentType],
-               fun () -> prometheus_text_format:format(Registry) end),
-    prometheus_summary:observe(Registry,
-                               ?SCRAPE_SIZE,
-                               [Registry, ContentType],
-                               iolist_size(Scrape)),
-    Scrape.
+    Headers = #{<<"content-type">> => ContentType,
+                <<"content-encoding">> => Encoding},
+    case Encoding of
+        <<"gzip">> ->
+            Scrape = prometheus_summary:observe_duration(
+                       Registry,
+                       ?SCRAPE_DURATION,
+                       [Registry, ContentType],
+                       fun () -> prometheus_text_format:format(Registry) end),
+            prometheus_summary:observe(Registry,
+                                       ?SCRAPE_SIZE,
+                                       [Registry, ContentType],
+                                       iolist_size(Scrape)),
+            Encoded = zlib:gzip(Scrape),
+            prometheus_summary:observe(telemetry_registry(),
+                                       ?SCRAPE_ENCODED_SIZE,
+                                       [Registry, ContentType, Encoding],
+                                       iolist_size(Encoded)),
+            cowboy_req:reply(200, Headers, Encoded, Request);
+        <<"identity">> ->
+            Req = cowboy_req:stream_reply(200, Headers, Request),
+            Fmt = fun(Data, Size) ->
+                cowboy_req:stream_body(Data, nofin, Req),
+                Size + iolist_size(Data)
+            end,
+            Size = prometheus_summary:observe_duration(
+                     Registry,
+                     ?SCRAPE_DURATION,
+                     [Registry, ContentType],
+                     fun () -> prometheus_text_format:format_into(Registry, 0, Fmt) end),
+            cowboy_req:stream_body(<<>>, fin, Req),
+            prometheus_summary:observe(Registry,
+                                       ?SCRAPE_SIZE,
+                                       [Registry, ContentType],
+                                       Size),
+            Req
+    end.
 
 validate_registry(undefined, auto) ->
     {true, default};
@@ -146,20 +157,6 @@ telemetry_registry() ->
 registry() ->
     application:get_env(rabbitmq_prometheus, registry, auto).
 
-encode_format(ContentType, Encoding, Scrape, Registry) ->
-    Encoded = encode_format_(Encoding, Scrape),
-    prometheus_summary:observe(telemetry_registry(),
-                               ?SCRAPE_ENCODED_SIZE,
-                               [Registry, ContentType, Encoding],
-                               iolist_size(Encoded)),
-    {200, [{content_type, binary_to_list(ContentType)},
-           {content_encoding, Encoding}], Encoded}.
-
-encode_format_("gzip", Scrape) ->
-    zlib:gzip(Scrape);
-encode_format_("identity", Scrape) ->
-     Scrape.
-
 %% It's not easy to pass this information in a pure way (it'll require changing prometheus.erl)
 put_filtering_options_into_process_dictionary(Request) ->
     #{vhost := VHosts, family := Families} = cowboy_req:match_qs([{vhost, [], undefined}, {family, [], undefined}], Request),

rabbitmq-components.mk

@@ -50,7 +50,7 @@ dep_khepri = hex 0.17.2
 dep_khepri_mnesia_migration = hex 0.8.0
 dep_meck = hex 1.0.0
 dep_osiris = git https://github.com/rabbitmq/osiris v1.10.1
-dep_prometheus = hex 5.1.1
+dep_prometheus = git https://github.com/the-mikedavis/prometheus.erl ceb4c9fe14ede2876f87d4997c9825cdfbf07759
 dep_ra = hex 2.17.1
 dep_ranch = hex 2.2.0
 dep_recon = hex 2.5.6