add gzip rsp #153

This commit is contained in:
Gorbunov Kirill Andreevich 2024-03-11 00:06:45 +03:00
parent 41f7bdf182
commit bee12a2f9c

View File

@ -2,14 +2,31 @@ package meter // import "go.unistack.org/micro-server-http/v4/handler/meter"
import ( import (
"bytes" "bytes"
"compress/gzip"
"context" "context"
"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 (
contentTypeHeader = "Content-Type"
contentEncodingHeader = "Content-Encoding"
acceptEncodingHeader = "Accept-Encoding"
)
var gzipPool = sync.Pool{
New: func() interface{} {
return gzip.NewWriter(nil)
},
}
// guard to fail early // guard to fail early
var _ MeterServiceServer = (*Handler)(nil) var _ MeterServiceServer = (*Handler)(nil)
@ -57,12 +74,44 @@ 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 {
log := logger.DefaultLogger()
buf := bytes.NewBuffer(nil) buf := bytes.NewBuffer(nil)
if md, ok := metadata.FromContext(ctx); gzipAccepted(md) && ok {
md.Set(contentEncodingHeader, "gzip")
gz := gzipPool.Get().(*gzip.Writer)
defer gzipPool.Put(gz)
gz.Reset(buf)
defer gz.Close()
zw := gzip.NewWriter(buf)
defer zw.Close()
*zw = *gz
}
if err := h.opts.Meter.Write(buf, h.opts.MeterOptions...); err != nil { if err := h.opts.Meter.Write(buf, h.opts.MeterOptions...); err != nil {
return errors.InternalServerError(h.opts.Name, "%v", err) 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
}
parts := strings.Split(a, ",")
for _, part := range parts {
part = strings.TrimSpace(part)
if part == "gzip" || strings.HasPrefix(part, "gzip;") {
return true
}
}
return false
}