133 lines
2.4 KiB
Go
133 lines
2.4 KiB
Go
package meter_handler
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/gzip"
|
|
"context"
|
|
"io"
|
|
"strings"
|
|
"sync"
|
|
|
|
codecpb "go.unistack.org/micro-proto/v3/codec"
|
|
"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{}
|
|
|
|
type Handler struct {
|
|
Options Options
|
|
}
|
|
|
|
type Option func(*Options)
|
|
|
|
type Options struct {
|
|
Meter meter.Meter
|
|
Name string
|
|
MeterOptions []meter.Option
|
|
DisableCompress bool
|
|
}
|
|
|
|
func Meter(m meter.Meter) Option {
|
|
return func(o *Options) {
|
|
o.Meter = m
|
|
}
|
|
}
|
|
|
|
func Name(name string) Option {
|
|
return func(o *Options) {
|
|
o.Name = name
|
|
}
|
|
}
|
|
|
|
func DisableCompress(g bool) Option {
|
|
return func(o *Options) {
|
|
o.DisableCompress = g
|
|
}
|
|
}
|
|
|
|
func MeterOptions(opts ...meter.Option) Option {
|
|
return func(o *Options) {
|
|
o.MeterOptions = append(o.MeterOptions, opts...)
|
|
}
|
|
}
|
|
|
|
func NewOptions(opts ...Option) Options {
|
|
options := Options{Meter: meter.DefaultMeter, DisableCompress: false}
|
|
for _, o := range opts {
|
|
o(&options)
|
|
}
|
|
return options
|
|
}
|
|
|
|
func NewHandler(opts ...Option) *Handler {
|
|
options := NewOptions(opts...)
|
|
return &Handler{Options: options}
|
|
}
|
|
|
|
func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb.Frame) error {
|
|
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.FromOutgoingContext(ctx); gzipAccepted(md) && ok && !h.Options.DisableCompress {
|
|
omd, _ := metadata.FromOutgoingContext(ctx)
|
|
omd.Set(contentEncodingHeader, "gzip")
|
|
gz := gzipPool.Get().(*gzip.Writer)
|
|
defer gzipPool.Put(gz)
|
|
|
|
gz.Reset(w)
|
|
defer gz.Close()
|
|
|
|
w = gz
|
|
gz.Flush()
|
|
}
|
|
|
|
if err := h.Options.Meter.Write(w, h.Options.MeterOptions...); err != nil {
|
|
log.Error(ctx, "http/meter write failed", 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
|
|
}
|