#153-v3 #187
| @@ -2,13 +2,35 @@ package meter // import "go.unistack.org/micro-server-http/v3/handler/meter" | |||||||
|  |  | ||||||
| import ( | import ( | ||||||
| 	"bytes" | 	"bytes" | ||||||
|  | 	"compress/gzip" | ||||||
| 	"context" | 	"context" | ||||||
|  | 	"io" | ||||||
|  | 	"strings" | ||||||
|  | 	"sync" | ||||||
|  |  | ||||||
| 	codecpb "go.unistack.org/micro-proto/v3/codec" | 	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" | 	"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 | // guard to fail early | ||||||
| var _ MeterServiceServer = &Handler{} | var _ MeterServiceServer = &Handler{} | ||||||
|  |  | ||||||
| @@ -56,12 +78,48 @@ 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.FromIncomingContext(ctx); gzipAccepted(md) && ok { | ||||||
|  | 		md.Set(contentEncodingHeader, "gzip") | ||||||
|  | 		ctx = metadata.NewIncomingContext(ctx, md) | ||||||
|  | 		gz := gzipPool.Get().(*gzip.Writer) | ||||||
|  | 		defer gzipPool.Put(gz) | ||||||
|  |  | ||||||
|  | 		gz.Reset(w) | ||||||
|  | 		defer gz.Close() | ||||||
|  |  | ||||||
|  | 		w = gz | ||||||
|  | 		gz.Flush() | ||||||
|  | 	} | ||||||
|  |  | ||||||
|  | 	if err := h.opts.Meter.Write(w, h.opts.MeterOptions...); err != nil { | ||||||
|  | 		log.Error(ctx, "http/meter: write failed", 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