From 18c26b749c7cf776536f302c29ed97f5696cb1f5 Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Tue, 31 Oct 2023 16:29:39 +0100
Subject: [PATCH] Add support for compression

---
 Project.toml      |  4 ++++
 src/Prometheus.jl | 48 +++++++++++++++++++++++++++++++++++++++++------
 test/runtests.jl  | 15 +++++++++++++++
 3 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index 6c7c3e5..eb86344 100644
--- a/Project.toml
+++ b/Project.toml
@@ -3,11 +3,15 @@ uuid = "f25c1797-fe98-4e0c-b252-1b4fe3b6bde6"
 version = "1.0.0"
 
 [deps]
+CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+SimpleBufferStream = "777ac1f9-54b0-4bf8-805c-2214025038e7"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 
 [compat]
+CodecZlib = "0.7"
 HTTP = "1"
+SimpleBufferStream = "1"
 Sockets = "1.6"
 julia = "1.6"
 
diff --git a/src/Prometheus.jl b/src/Prometheus.jl
index 7a0ed03..001bb5d 100644
--- a/src/Prometheus.jl
+++ b/src/Prometheus.jl
@@ -1,6 +1,8 @@
 module Prometheus
 
+using CodecZlib: GzipCompressorStream
 using HTTP: HTTP
+using SimpleBufferStream: BufferStream
 using Sockets: Sockets
 
 abstract type Collector end
@@ -50,6 +52,9 @@ if !isdefined(Base, Symbol("@atomic")) # v1.7.0
         end
     end
 end
+if !isdefined(Base, :eachsplit) # v1.8.0
+    const eachsplit = split
+end
 
 
 #####################
@@ -505,24 +510,55 @@ end
 
 const CONTENT_TYPE_LATEST = "text/plain; version=0.0.4; charset=utf-8"
 
+function gzip_accepted(http::HTTP.Stream)
+    accept_encoding = HTTP.header(http.message, "Accept-Encoding")
+    for enc in eachsplit(accept_encoding, ',')
+        if lowercase(strip(first(eachsplit(enc, ';')))) == "gzip"
+            return true
+        end
+    end
+    return false
+end
+
 """
-    expose(http::HTTP.Stream, reg::CollectorRegistry = DEFAULT_REGISTRY)
+    expose(http::HTTP.Stream, reg::CollectorRegistry = DEFAULT_REGISTRY; kwargs...)
 
 Export all metrics in `reg` by writing them to the HTTP stream `http`.
 
 The caller is responsible for checking e.g. the HTTP method and URI target. For
 HEAD requests this method do not write a body, however.
 """
-function expose(http::HTTP.Stream, reg::CollectorRegistry = DEFAULT_REGISTRY)
-    # TODO: Handle Accept request header
-    # TODO: Compression if requested/supported
+function expose(http::HTTP.Stream, reg::CollectorRegistry = DEFAULT_REGISTRY; compress::Bool=true)
+    # TODO: Handle Accept request header for different formats?
+    # Compress by default if client supports it and user haven't disabled it
+    if compress
+        compress = gzip_accepted(http)
+    end
+    # Create the response
     HTTP.setstatus(http, 200)
     HTTP.setheader(http, "Content-Type" => CONTENT_TYPE_LATEST)
+    if compress
+        HTTP.setheader(http, "Content-Encoding" => "gzip")
+    end
     HTTP.startwrite(http)
-    # The user is repsonsible for making sure that e.g. target and method is
+    # The user is responsible for making sure that e.g. target and method is
     # correct, but at least we skip writing the body for HEAD requests.
     if http.message.method != "HEAD"
-        expose_io(http, reg)
+        if compress
+            buf = BufferStream()
+            gzstream = GzipCompressorStream(buf)
+            tsk = @async try
+                expose_io(gzstream, reg)
+            finally
+                # Close the compressor stream to free resources in zlib and
+                # to let the write(http, buf) below finish.
+                close(gzstream)
+            end
+            write(http, buf)
+            wait(tsk)
+        else
+            expose_io(http, reg)
+        end
     end
     return
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 287eb7a..b7eadbe 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -323,6 +323,8 @@ end
             return Prometheus.expose(http)
         elseif http.message.target == "/metrics/reg"
             return Prometheus.expose(http, Prometheus.DEFAULT_REGISTRY)
+        elseif http.message.target == "/metrics/nogzip"
+            return Prometheus.expose(http; compress=false)
         else
             HTTP.setstatus(http, 404)
             HTTP.startwrite(http)
@@ -340,6 +342,19 @@ end
     # Bad URI
     r_bad = HTTP.request("GET", "http://localhost:8123"; status_exception=false)
     @test r_bad.status == 404
+    # Compression
+    for enc in ("gzip", "br, compress, gzip", "br;q=1.0, gzip;q=0.8, *;q=0.1")
+        r_gzip = HTTP.request(
+            "GET", "http://localhost:8123/metrics/default", ["Accept-Encoding" => enc]
+        )
+        @test HTTP.header(r_gzip, "Content-Encoding") == "gzip"
+        @test String(r_gzip.body) == reference_output # HTTP.jl decompresses gzip
+        r_nogzip = HTTP.request(
+            "GET", "http://localhost:8123/metrics/nogzip", ["Accept-Encoding" => enc]
+        )
+        @test HTTP.header(r_nogzip, "Content-Encoding") != "gzip"
+        @test String(r_nogzip.body) == reference_output # HTTP.jl decompresses gzip
+    end
     # Clean up
     close(server)
     wait(server)