micro-server-http/handler/meter/meter.go
Gorbunov Kirill Andreevich aa6c32026f update #153
2024-03-11 11:18:57 +03:00

116 lines
2.2 KiB
Go

package meter // import "go.unistack.org/micro-server-http/v4/handler/meter"
import (
"bytes"
"compress/gzip"
"context"
"io"
"strings"
"sync"
codecpb "go.unistack.org/micro-proto/v4/codec"
"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/options"
)
const (
contentEncodingHeader = "Content-Encoding"
acceptEncodingHeader = "Accept-Encoding"
)
var gzipPool = sync.Pool{
New: func() interface{} {
return gzip.NewWriter(nil)
},
}
// guard to fail early
var _ MeterServiceServer = (*Handler)(nil)
type Handler struct {
opts Options
}
type Option func(*Options)
type Options struct {
Meter meter.Meter
Name string
MeterOptions []options.Option
}
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 MeterOptions(opts ...options.Option) Option {
return func(o *Options) {
o.MeterOptions = append(o.MeterOptions, opts...)
}
}
func NewOptions(opts ...Option) Options {
options := Options{Meter: meter.DefaultMeter}
for _, o := range opts {
o(&options)
}
return options
}
func NewHandler(opts ...Option) *Handler {
options := NewOptions(opts...)
return &Handler{opts: 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 := bytes.NewBuffer(nil)
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(buf)
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()
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
}