Skip to content

Commit 48c16ce

Browse files
mruegbwplotka
andcommitted
Update prometheus/promhttp/http.go
Co-authored-by: Bartlomiej Plotka <[email protected]> Signed-off-by: Manuel Rüger <[email protected]>
1 parent 03af4da commit 48c16ce

File tree

2 files changed

+132
-53
lines changed

2 files changed

+132
-53
lines changed

prometheus/promhttp/http.go

+81-43
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,25 @@ const (
5555
processStartTimeHeader = "Process-Start-Time-Unix"
5656
)
5757

58-
var defaultEncodingOffers = []string{"gzip", "zstd"}
58+
type Compression int
59+
60+
const (
61+
Identity Compression = iota
62+
Gzip
63+
Zstd
64+
)
65+
66+
var compressions = [...]string{
67+
"identity",
68+
"gzip",
69+
"zstd",
70+
}
71+
72+
func (c Compression) String() string {
73+
return compressions[c]
74+
}
75+
76+
var defaultCompressionFormats = []Compression{Identity, Gzip, Zstd}
5977

6078
var gzipPool = sync.Pool{
6179
New: func() interface{} {
@@ -168,46 +186,14 @@ func HandlerForTransactional(reg prometheus.TransactionalGatherer, opts HandlerO
168186
} else {
169187
contentType = expfmt.Negotiate(req.Header)
170188
}
171-
header := rsp.Header()
172-
header.Set(contentTypeHeader, string(contentType))
173-
174-
w := io.Writer(rsp)
175-
if !opts.DisableCompression {
176-
offers := defaultEncodingOffers
177-
if len(opts.EncodingOffers) > 0 {
178-
offers = opts.EncodingOffers
179-
}
180-
// TODO(mrueg): Replace internal/github.com/gddo once https://github.com/golang/go/issues/19307 is implemented.
181-
compression := httputil.NegotiateContentEncoding(req, offers)
182-
switch compression {
183-
case "zstd":
184-
header.Set(contentEncodingHeader, "zstd")
185-
// TODO(mrueg): Replace klauspost/compress with stdlib implementation once https://github.com/golang/go/issues/62513 is implemented.
186-
z, err := zstd.NewWriter(rsp, zstd.WithEncoderLevel(zstd.SpeedFastest))
187-
if err != nil {
188-
return
189-
}
190-
191-
z.Reset(w)
192-
defer z.Close()
189+
rsp.Header().Set(contentTypeHeader, string(contentType))
193190

194-
w = z
195-
case "gzip":
196-
header.Set(contentEncodingHeader, "gzip")
197-
gz := gzipPool.Get().(*gzip.Writer)
198-
defer gzipPool.Put(gz)
191+
w, err := GetWriter(req, rsp, opts.DisableCompression, opts.OfferedCompressions)
199192

200-
gz.Reset(w)
201-
defer gz.Close()
202-
203-
w = gz
204-
case "identity":
205-
// This means the content is not encoded.
206-
default:
207-
// The content encoding was not implemented yet.
208-
return
193+
if err != nil {
194+
if opts.ErrorLog != nil {
195+
opts.ErrorLog.Println("error setting getting writer", err)
209196
}
210-
211197
}
212198

213199
enc := expfmt.NewEncoder(w, contentType)
@@ -373,12 +359,19 @@ type HandlerOpts struct {
373359
// no effect on the HTTP status code because ErrorHandling is set to
374360
// ContinueOnError.
375361
Registry prometheus.Registerer
376-
// If DisableCompression is true, the handler will never compress the
377-
// response, even if requested by the client.
362+
// DisableCompression disables the response encoding (compression) and
363+
// encoding negotiation. If true, the handler will
364+
// never compress the response, even if requested
365+
// by the client and the OfferedCompressions field is ignored.
378366
DisableCompression bool
379-
// If DisableCompression is false, this option will allow to define the
380-
// set of offered encoding algorithms.
381-
EncodingOffers []string
367+
// OfferedCompressions is a set of encodings (compressions) handler will
368+
// try to offer when negotiating with the client. This defaults to zstd,
369+
// gzip and identity.
370+
// NOTE: If handler can't agree on the encodings with the client or
371+
// caller using unsupported or empty encodings in OfferedCompressions,
372+
// handler always fallbacks to no compression (identity), for
373+
// compatibility reasons. In such cases ErrorLog will be used if set.
374+
OfferedCompressions []Compression
382375
// The number of concurrent HTTP requests is limited to
383376
// MaxRequestsInFlight. Additional requests are responded to with 503
384377
// Service Unavailable and a suitable message in the body. If
@@ -426,3 +419,48 @@ func httpError(rsp http.ResponseWriter, err error) {
426419
http.StatusInternalServerError,
427420
)
428421
}
422+
423+
func GetWriter(r *http.Request, rsp http.ResponseWriter, disableCompression bool, offeredCompressions []Compression) (io.Writer, error) {
424+
w := io.Writer(rsp)
425+
if !disableCompression {
426+
offers := defaultCompressionFormats
427+
if len(offeredCompressions) > 0 {
428+
offers = offeredCompressions
429+
}
430+
var compressions []string
431+
for _, comp := range offers {
432+
compressions = append(compressions, comp.String())
433+
}
434+
// TODO(mrueg): Replace internal/github.com/gddo once https://github.com/golang/go/issues/19307 is implemented.
435+
compression := httputil.NegotiateContentEncoding(r, compressions)
436+
switch compression {
437+
case "zstd":
438+
rsp.Header().Set(contentEncodingHeader, "zstd")
439+
// TODO(mrueg): Replace klauspost/compress with stdlib implementation once https://github.com/golang/go/issues/62513 is implemented.
440+
z, err := zstd.NewWriter(rsp, zstd.WithEncoderLevel(zstd.SpeedFastest))
441+
if err != nil {
442+
return nil, err
443+
}
444+
445+
z.Reset(w)
446+
defer z.Close()
447+
448+
w = z
449+
case "gzip":
450+
rsp.Header().Set(contentEncodingHeader, "gzip")
451+
gz := gzipPool.Get().(*gzip.Writer)
452+
defer gzipPool.Put(gz)
453+
454+
gz.Reset(w)
455+
defer gz.Close()
456+
457+
w = gz
458+
case "identity":
459+
// This means the content is not compressed.
460+
default:
461+
// The content encoding was not implemented yet.
462+
return w, fmt.Errorf("content compression format not recognized: %s. Valid formats are: %s", compression, defaultCompressionFormats)
463+
}
464+
}
465+
return w, nil
466+
}

prometheus/promhttp/http_test.go

+51-10
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"bytes"
1818
"errors"
1919
"fmt"
20+
"io"
2021
"log"
2122
"net/http"
2223
"net/http/httptest"
@@ -332,22 +333,62 @@ func TestHandlerTimeout(t *testing.T) {
332333
close(c.Block) // To not leak a goroutine.
333334
}
334335

335-
func BenchmarkEncoding(b *testing.B) {
336+
func TestGetWriter(t *testing.T) {
337+
testCases := []struct {
338+
name string
339+
disableCompression bool
340+
offeredCompressions []Compression
341+
req http.Request
342+
rsp http.ResponseWriter
343+
w io.Writer
344+
err error
345+
}{
346+
{
347+
name: "test without compression",
348+
disableCompression: true,
349+
offeredCompressions: defaultCompressionFormats,
350+
req: http.Request{},
351+
rsp: http.ResponseWriter{},
352+
w: io.Writer{},
353+
err: nil,
354+
},
355+
{
356+
name: "test with gzip compression requested",
357+
disableCompression: false,
358+
offeredCompressions: defaultCompressionFormats,
359+
req: http.Request{},
360+
rsp: http.ResponseWriter{},
361+
w: io.Writer{},
362+
err: nil,
363+
},
364+
}
365+
366+
for _, test := range testCases {
367+
actualW, actualErr := GetWriter(&test.req, test.rsp, test.disableCompression, test.offeredCompressions)
368+
369+
if actualErr != err {
370+
t.Errorf("got error %v, expected %v", actualErr, err)
371+
}
372+
}
373+
374+
}
375+
376+
func BenchmarkCompression(b *testing.B) {
336377
benchmarks := []struct {
337-
name string
338-
encodingType string
378+
name string
379+
compressionType string
339380
}{
340381
{
341-
name: "test with gzip encoding",
342-
encodingType: "gzip",
382+
name: "test with gzip compression",
383+
compressionType: "gzip",
343384
},
344385
{
345-
name: "test with zstd encoding",
346-
encodingType: "zstd",
386+
name: "test with zstd compression",
387+
compressionType: "zstd",
347388
},
348389
{
349-
name: "test with no encoding",
350-
encodingType: "identity",
390+
name: "test with no compression",
391+
compressionType: "identity",
351392
},
352393
}
353394
sizes := []struct {
@@ -416,7 +457,7 @@ func BenchmarkEncoding(b *testing.B) {
416457
for i := 0; i < b.N; i++ {
417458
writer := httptest.NewRecorder()
418459
request, _ := http.NewRequest("GET", "/", nil)
419-
request.Header.Add("Accept-Encoding", benchmark.encodingType)
460+
request.Header.Add("Accept-Encoding", benchmark.compressionType)
420461
handler.ServeHTTP(writer, request)
421462
}
422463
})

0 commit comments

Comments
 (0)