main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/hexdigest/gowrap/generator"
	"github.com/pkg/errors"
)

const version = "1.3"

type (
	interfaceInfo struct {
		Name       string
		Package    string
		OutputFile string
	}
)

func main() {
	interfaces := processFlags()

	for _, iface := range interfaces {
		opts := generator.Options{
			InterfaceName:  iface.Name,
			SourcePackage:  iface.Package,
			OutputFile:     iface.OutputFile,
			HeaderTemplate: headerTemplate,
			BodyTemplate:   bodyTemplate,
			Vars: map[string]interface{}{
				"DecoratorName": iface.Name + "Metrics",
			},
			HeaderVars: map[string]interface{}{
				"Version": version,
			},
		}

		if err := generate(opts); err != nil {
			die("failed to generate %s: %v", opts.OutputFile, err)
		}

		fmt.Printf("Generated file: %s\n", opts.OutputFile)
	}
}

func generate(o generator.Options) error {
	g, err := generator.NewGenerator(o)
	if err != nil {
		return err
	}

	buf := bytes.NewBuffer([]byte{})

	if err = g.Generate(buf); err != nil {
		return errors.Wrap(err, "failed to generate decorator")
	}

	return ioutil.WriteFile(o.OutputFile, buf.Bytes(), 0644)
}

const (
	headerTemplate = `
		package {{$.Package.Name}}

		// Code generated by http://github.com/gojuno/metricsgen ({{$.Options.HeaderVars.Version}}). DO NOT EDIT.
	`

	bodyTemplate = `
import (
  "sync"
  "time"
  "github.com/prometheus/client_golang/prometheus"
)

{{ $decorator := .Vars.DecoratorName }}

// {{$decorator}} implements {{.Interface.Type}} interface with all methods wrapped
// with Prometheus metrics
type {{$decorator}} struct {
  base {{.Interface.Type}}
  instanceName string
  summary *prometheus.SummaryVec
}

func (_d {{$decorator}}) _observe(method string, startedAt time.Time) {
	duration := time.Since(startedAt)
	_d.summary.WithLabelValues(_d.instanceName, method).Observe(duration.Seconds())
}

func New{{$decorator}}Summary(metricName string) *prometheus.SummaryVec {
	sv := prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Name: metricName,
			Help: metricName,
		},
		[]string{"instance_name", "method"},
	)

	prometheus.MustRegister(sv)

	return sv
}

// New{{$decorator}}WithSummary returns an instance of the {{.Interface.Type}} decorated with prometheus summary metric
func New{{$decorator}}WithSummary(base {{.Interface.Type}}, instanceName string, sv *prometheus.SummaryVec) {{$decorator}} {
  return {{$decorator}} {
    base: base,
    instanceName: instanceName,
		summary: sv,
  }
}

{{range $method := .Interface.Methods}}
  // {{$method.Name}} implements {{$.Interface.Type}}
  func (_d {{$decorator}}) {{$method.Declaration}} {
		defer _d._observe("{{$method.Name}}", time.Now())
    {{$method.Pass "_d.base."}}
  }
{{end}}
`
)

func processFlags() []interfaceInfo {
	var (
		help       = flag.Bool("h", false, "show this help message")
		interfaces = flag.String("i", "", "comma-separated names of interfaces to wrap, i.e fmt.Stringer,io.Reader, use io.* notation to generate metric decorators for all interfaces in an io package")
		output     = flag.String("o", "", "directory to place generated files to")
		suffix     = flag.String("s", "_metrics.go", "output file name suffix which is added to file names when multiple interfaces are given")
		v          = flag.Bool("version", false, "show metricsgen version")
	)

	flag.Parse()

	if *help {
		flag.Usage()
		os.Exit(0)
	}

	if *v {
		fmt.Printf("metricsgen version: %s\n", version)
		os.Exit(0)
	}

	if *interfaces == "" {
		die("missing required parameter: -i, use -h flag for help")
	}

	if *output == "" {
		die("missing required parameter: -o, use -h flag for help")
	}

	result := []interfaceInfo{}
	for _, i := range strings.Split(*interfaces, ",") {
		chunks := strings.Split(i, ".")
		if len(chunks) < 2 {
			die("invalid interface name: %s\nname should be in the form <import path>.<interface type>, i.e. io.Reader\n", i)
		}

		interfaceName := chunks[len(chunks)-1]

		result = append(result, interfaceInfo{
			Name:       interfaceName,
			Package:    strings.Join(chunks[0:len(chunks)-1], "."),
			OutputFile: filepath.Join(*output, camelToSnake(interfaceName)+*suffix),
		})
	}

	return result
}

func camelToSnake(s string) string {
	b := buffer{
		r: make([]byte, 0, len(s)),
	}
	var m rune
	var w bool
	for _, ch := range s {
		if unicode.IsUpper(ch) {
			if m != 0 {
				if !w {
					b.indent()
					w = true
				}
				b.write(m)
			}
			m = unicode.ToLower(ch)
		} else {
			if m != 0 {
				b.indent()
				b.write(m)
				m = 0
				w = false
			}
			b.write(ch)
		}
	}
	if m != 0 {
		if !w {
			b.indent()
		}
		b.write(m)
	}
	return string(b.r)
}

type buffer struct {
	r         []byte
	runeBytes [utf8.UTFMax]byte
}

func (b *buffer) write(r rune) {
	if r < utf8.RuneSelf {
		b.r = append(b.r, byte(r))
		return
	}
	n := utf8.EncodeRune(b.runeBytes[0:], r)
	b.r = append(b.r, b.runeBytes[0:n]...)
}

func (b *buffer) indent() {
	if len(b.r) > 0 {
		b.r = append(b.r, '_')
	}
}

func die(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "metricsgen: "+format+"\n", args...)
	os.Exit(1)
}