fix bug with gzip in meter handler (#223)
This commit is contained in:
		| @@ -4,7 +4,7 @@ import ( | ||||
| 	"bytes" | ||||
| 	"compress/gzip" | ||||
| 	"context" | ||||
| 	"io" | ||||
| 	"fmt" | ||||
| 	"strings" | ||||
| 	"sync" | ||||
|  | ||||
| @@ -91,32 +91,52 @@ func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb. | ||||
| 		log = logger.DefaultLogger | ||||
| 	} | ||||
|  | ||||
| 	var err error | ||||
|  | ||||
| 	if md, ok := metadata.FromIncomingContext(ctx); ok && gzipAccepted(md) && !h.Options.DisableCompress { | ||||
| 		err = h.writeMetricsGzip(ctx, rsp) | ||||
| 	} else { | ||||
| 		err = h.writeMetricsPlain(ctx, rsp) | ||||
| 	} | ||||
|  | ||||
| 	if err != nil { | ||||
| 		log.Error(ctx, "http/meter write failed", err) | ||||
| 	} | ||||
|  | ||||
| 	return nil | ||||
| } | ||||
|  | ||||
| func (h *Handler) writeMetricsGzip(ctx context.Context, rsp *codecpb.Frame) error { | ||||
| 	httpsrv.AppendResponseMetadata(ctx, metadata.Pairs(contentEncodingHeader, "gzip")) | ||||
|  | ||||
| 	buf := bufPool.Get().(*bytes.Buffer) | ||||
| 	defer bufPool.Put(buf) | ||||
| 	buf.Reset() | ||||
|  | ||||
| 	w := io.Writer(buf) | ||||
|  | ||||
| 	if md, ok := metadata.FromIncomingContext(ctx); ok && gzipAccepted(md) && !h.Options.DisableCompress { | ||||
| 		httpsrv.AppendResponseMetadata(ctx, metadata.Pairs(contentEncodingHeader, "gzip")) | ||||
|  | ||||
| 	gz := gzipPool.Get().(*gzip.Writer) | ||||
| 	defer gzipPool.Put(gz) | ||||
| 	gz.Reset(buf) | ||||
|  | ||||
| 		gz.Reset(w) | ||||
| 		defer gz.Close() | ||||
|  | ||||
| 		w = gz | ||||
| 	if err := h.Options.Meter.Write(gz, h.Options.MeterOptions...); err != nil { | ||||
| 		return fmt.Errorf("meter write: %w", err) | ||||
| 	} | ||||
|  | ||||
| 	if err := h.Options.Meter.Write(w, h.Options.MeterOptions...); err != nil { | ||||
| 		log.Error(ctx, "http/meter write failed", err) | ||||
| 	if err := gz.Close(); err != nil { | ||||
| 		return fmt.Errorf("gzip close: %w", err) | ||||
| 	} | ||||
|  | ||||
| 	rsp.Data = buf.Bytes() | ||||
|  | ||||
| 	return nil | ||||
| } | ||||
|  | ||||
| 	// gz.Flush() must be called after writing metrics to ensure buffered data is written to the underlying writer. | ||||
| 	if gz, ok := w.(*gzip.Writer); ok { | ||||
| 		gz.Flush() | ||||
| func (h *Handler) writeMetricsPlain(_ context.Context, rsp *codecpb.Frame) error { | ||||
| 	buf := bufPool.Get().(*bytes.Buffer) | ||||
| 	defer bufPool.Put(buf) | ||||
| 	buf.Reset() | ||||
|  | ||||
| 	if err := h.Options.Meter.Write(buf, h.Options.MeterOptions...); err != nil { | ||||
| 		return fmt.Errorf("meter write: %w", err) | ||||
| 	} | ||||
|  | ||||
| 	rsp.Data = buf.Bytes() | ||||
|   | ||||
		Reference in New Issue
	
	Block a user