Skip to content

Commit 789178e

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 789178e

File tree

2 files changed

+135
-53
lines changed

2 files changed

+135
-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 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

+54-10
Original file line numberDiff line numberDiff line change
@@ -332,22 +332,66 @@ func TestHandlerTimeout(t *testing.T) {
332332
close(c.Block) // To not leak a goroutine.
333333
}
334334

335-
func BenchmarkEncoding(b *testing.B) {
335+
func TestGetWriter(t *testing.T) {
336+
testCases := []struct {
337+
name string
338+
disableCompression bool
339+
offeredCompressions []Compression
340+
acceptEncoding string
341+
expectedCompression string
342+
err error
343+
}{
344+
{
345+
name: "test without compression",
346+
disableCompression: true,
347+
offeredCompressions: defaultCompressionFormats,
348+
acceptEncoding: "",
349+
expectedCompression: "identity",
350+
err: nil,
351+
},
352+
{
353+
name: "test with gzip compression requested",
354+
disableCompression: false,
355+
offeredCompressions: defaultCompressionFormats,
356+
acceptEncoding: "gzip",
357+
expectedCompression: "gzip",
358+
err: nil,
359+
},
360+
}
361+
362+
for _, test := range testCases {
363+
request, _ := http.NewRequest("GET", "/", nil)
364+
request.Header.Add(acceptEncodingHeader, test.acceptEncoding)
365+
rr := httptest.NewRecorder()
366+
_, err := GetWriter(request, rr, test.disableCompression, test.offeredCompressions)
367+
368+
if err != test.err {
369+
t.Errorf("got error: %v, expected: %v", err, test.err)
370+
}
371+
372+
if rr.Header().Get(contentEncodingHeader) != test.expectedCompression {
373+
t.Errorf("got different compression type: %v, expected: %v", rr.Header().Get(contentEncodingHeader), test.expectedCompression)
374+
}
375+
}
376+
377+
}
378+
379+
func BenchmarkCompression(b *testing.B) {
336380
benchmarks := []struct {
337-
name string
338-
encodingType string
381+
name string
382+
compressionType string
339383
}{
340384
{
341-
name: "test with gzip encoding",
342-
encodingType: "gzip",
385+
name: "test with gzip compression",
386+
compressionType: "gzip",
343387
},
344388
{
345-
name: "test with zstd encoding",
346-
encodingType: "zstd",
389+
name: "test with zstd compression",
390+
compressionType: "zstd",
347391
},
348392
{
349-
name: "test with no encoding",
350-
encodingType: "identity",
393+
name: "test with no compression",
394+
compressionType: "identity",
351395
},
352396
}
353397
sizes := []struct {
@@ -416,7 +460,7 @@ func BenchmarkEncoding(b *testing.B) {
416460
for i := 0; i < b.N; i++ {
417461
writer := httptest.NewRecorder()
418462
request, _ := http.NewRequest("GET", "/", nil)
419-
request.Header.Add("Accept-Encoding", benchmark.encodingType)
463+
request.Header.Add("Accept-Encoding", benchmark.compressionType)
420464
handler.ServeHTTP(writer, request)
421465
}
422466
})

0 commit comments

Comments
 (0)