Skip to content

Commit b5d9fd5

Browse files
authored
fix bug with gzip in meter handler (#223)
1 parent 39ad61b commit b5d9fd5

File tree

1 file changed

+37
-17
lines changed

1 file changed

+37
-17
lines changed

handler/meter/meter.go

Lines changed: 37 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import (
44
"bytes"
55
"compress/gzip"
66
"context"
7-
"io"
7+
"fmt"
88
"strings"
99
"sync"
1010

@@ -91,32 +91,52 @@ func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb.
9191
log = logger.DefaultLogger
9292
}
9393

94+
var err error
95+
96+
if md, ok := metadata.FromIncomingContext(ctx); ok && gzipAccepted(md) && !h.Options.DisableCompress {
97+
err = h.writeMetricsGzip(ctx, rsp)
98+
} else {
99+
err = h.writeMetricsPlain(ctx, rsp)
100+
}
101+
102+
if err != nil {
103+
log.Error(ctx, "http/meter write failed", err)
104+
}
105+
106+
return nil
107+
}
108+
109+
func (h *Handler) writeMetricsGzip(ctx context.Context, rsp *codecpb.Frame) error {
110+
httpsrv.AppendResponseMetadata(ctx, metadata.Pairs(contentEncodingHeader, "gzip"))
111+
94112
buf := bufPool.Get().(*bytes.Buffer)
95113
defer bufPool.Put(buf)
96114
buf.Reset()
97115

98-
w := io.Writer(buf)
116+
gz := gzipPool.Get().(*gzip.Writer)
117+
defer gzipPool.Put(gz)
118+
gz.Reset(buf)
99119

100-
if md, ok := metadata.FromIncomingContext(ctx); ok && gzipAccepted(md) && !h.Options.DisableCompress {
101-
httpsrv.AppendResponseMetadata(ctx, metadata.Pairs(contentEncodingHeader, "gzip"))
120+
if err := h.Options.Meter.Write(gz, h.Options.MeterOptions...); err != nil {
121+
return fmt.Errorf("meter write: %w", err)
122+
}
102123

103-
gz := gzipPool.Get().(*gzip.Writer)
104-
defer gzipPool.Put(gz)
124+
if err := gz.Close(); err != nil {
125+
return fmt.Errorf("gzip close: %w", err)
126+
}
105127

106-
gz.Reset(w)
107-
defer gz.Close()
128+
rsp.Data = buf.Bytes()
108129

109-
w = gz
110-
}
130+
return nil
131+
}
111132

112-
if err := h.Options.Meter.Write(w, h.Options.MeterOptions...); err != nil {
113-
log.Error(ctx, "http/meter write failed", err)
114-
return nil
115-
}
133+
func (h *Handler) writeMetricsPlain(_ context.Context, rsp *codecpb.Frame) error {
134+
buf := bufPool.Get().(*bytes.Buffer)
135+
defer bufPool.Put(buf)
136+
buf.Reset()
116137

117-
// gz.Flush() must be called after writing metrics to ensure buffered data is written to the underlying writer.
118-
if gz, ok := w.(*gzip.Writer); ok {
119-
gz.Flush()
138+
if err := h.Options.Meter.Write(buf, h.Options.MeterOptions...); err != nil {
139+
return fmt.Errorf("meter write: %w", err)
120140
}
121141

122142
rsp.Data = buf.Bytes()

0 commit comments

Comments
 (0)