micro-server-http/handler/meter/meter.go

package meter // import "go.unistack.org/micro-server-http/v4/handler/meter"

import (
	"bytes"
	"compress/gzip"
	"context"
	"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 := logger.DefaultLogger()
	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 {
		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
	}
	parts := strings.Split(a, ",")
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part == "gzip" || strings.HasPrefix(part, "gzip;") {
			return true
		}
	}
	return false
}