add gzip rsp #153 (#182)

## Pull Request template
Please, go through these steps before clicking submit on this PR.

1. Give a descriptive title to your PR.
2. Provide a description of your changes.
3. Make sure you have some relevant tests.
4. Put `closes #XXXX` in your comment to auto-close the issue that your PR fixes (if applicable).

**PLEASE REMOVE THIS TEMPLATE BEFORE SUBMITTING**

Co-authored-by: Gorbunov Kirill Andreevich <kgorbunov@mtsbank.ru>
Reviewed-on: #182
Co-authored-by: Кирилл Горбунов <kirya_gorbunov_2015@mail.ru>
Co-committed-by: Кирилл Горбунов <kirya_gorbunov_2015@mail.ru>
This commit is contained in:
Кирилл Горбунов 2024-03-11 12:31:19 +03:00 committed by Василий Толстов
parent 28eeca30c0
commit a483e08ffa

View File

@ -2,14 +2,37 @@ package meter // import "go.unistack.org/micro-server-http/v4/handler/meter"
import ( import (
"bytes" "bytes"
"compress/gzip"
"context" "context"
"io"
"strings"
"sync"
codecpb "go.unistack.org/micro-proto/v4/codec" codecpb "go.unistack.org/micro-proto/v4/codec"
"go.unistack.org/micro/v4/errors" "go.unistack.org/micro/v4/errors"
"go.unistack.org/micro/v4/logger"
"go.unistack.org/micro/v4/metadata"
"go.unistack.org/micro/v4/meter" "go.unistack.org/micro/v4/meter"
options "go.unistack.org/micro/v4/options" "go.unistack.org/micro/v4/options"
) )
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 // guard to fail early
var _ MeterServiceServer = (*Handler)(nil) var _ MeterServiceServer = (*Handler)(nil)
@ -57,12 +80,46 @@ func NewHandler(opts ...Option) *Handler {
} }
func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb.Frame) error { func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb.Frame) error {
buf := bytes.NewBuffer(nil) log, ok := logger.FromContext(ctx)
if err := h.opts.Meter.Write(buf, h.opts.MeterOptions...); err != nil { if !ok {
return errors.InternalServerError(h.opts.Name, "%v", err) log = logger.DefaultLogger()
}
buf := bufPool.Get().(*bytes.Buffer)
defer bufPool.Put(buf)
buf.Reset()
w := io.Writer(buf)
if md, ok := metadata.FromContext(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() rsp.Data = buf.Bytes()
return nil 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
}