## 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 was merged in pull request #182.
	This commit is contained in:
		| @@ -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 | ||||||
|  | } | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user