From 43a19aca6520e644c9f75da0479e4ed761e982eb Mon Sep 17 00:00:00 2001 From: Gorbunov Kirill Andreevich Date: Mon, 11 Mar 2024 13:01:42 +0300 Subject: [PATCH] gzip for v3 #153 --- handler/meter/meter.go | 63 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 60 insertions(+), 3 deletions(-) diff --git a/handler/meter/meter.go b/handler/meter/meter.go index f92830e..76974ce 100644 --- a/handler/meter/meter.go +++ b/handler/meter/meter.go @@ -2,13 +2,36 @@ package meter // import "go.unistack.org/micro-server-http/v3/handler/meter" import ( "bytes" + "compress/gzip" "context" + "io" + "strings" + "sync" codecpb "go.unistack.org/micro-proto/v3/codec" "go.unistack.org/micro/v3/errors" + "go.unistack.org/micro/v3/logger" + "go.unistack.org/micro/v3/metadata" "go.unistack.org/micro/v3/meter" ) +const ( + contentEncodingHeader = "Content-Encoding" + acceptEncodingHeader = "Accept-Encoding" +) + +var gzipPool = sync.Pool{ + New: func() interface{} { + return gzip.NewWriter(nil) + }, +} + +var bufPool = sync.Pool{ + New: func() interface{} { + return bytes.NewBuffer(nil) + }, +} + // guard to fail early var _ MeterServiceServer = &Handler{} @@ -56,12 +79,46 @@ func NewHandler(opts ...Option) *Handler { } func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb.Frame) error { - buf := bytes.NewBuffer(nil) - if err := h.opts.Meter.Write(buf, h.opts.MeterOptions...); err != nil { - return errors.InternalServerError(h.opts.Name, "%v", err) + log, ok := logger.FromContext(ctx) + if !ok { + log = logger.DefaultLogger() + } + + buf := bufPool.Get().(*bytes.Buffer) + defer bufPool.Put(buf) + buf.Reset() + + w := io.Writer(buf) + + if md, ok := metadata.FromIncomingContext(ctx); gzipAccepted(md) && ok { + md.Set(contentEncodingHeader, "gzip") + gz := gzipPool.Get().(*gzip.Writer) + defer gzipPool.Put(gz) + + gz.Reset(w) + defer gz.Close() + + w = gz + } + + if err := h.opts.Meter.Write(w, h.opts.MeterOptions...); err != nil { + log.Error(ctx, errors.InternalServerError(h.opts.Name, "%v", err)) + return nil } rsp.Data = buf.Bytes() return nil } + +// gzipAccepted returns whether the client will accept gzip-encoded content. +func gzipAccepted(md metadata.Metadata) bool { + a, ok := md.Get(acceptEncodingHeader) + if !ok { + return false + } + if strings.Contains(a, "gzip") { + return true + } + return false +}