unitech-golib/metric/exporter.go

package metric

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"os"
	"strconv"
	"time"

	"code.infininov.com/Infini/unitech-golib/tlog"
	consulapi "github.com/hashicorp/consul/api"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const (
	DefaultListenerAddr = "0.0.0.0:9166"
)

type ExporterCfg struct {
	Listener       net.Listener
	RegisterConsul bool
	ProductName    string
	ServiceRegion  string
	AppName        string
}

// Exporter init function must be executed asynchronously
func InitExporter(cfg ExporterCfg) (err error) {
	lis := cfg.Listener
	if cfg.Listener == nil {
		if lis, err = net.Listen("tcp", DefaultListenerAddr); err != nil {
			tlog.Error("fail to create listener with default address")
			return err
		}
	}

	_, port, err := net.SplitHostPort(lis.Addr().String())
	if err != nil {
		tlog.Error("fail to get port from listener, err: ", err)
		return err
	}

	h := http.NewServeMux()
	h.Handle("/metrics", promhttp.Handler())
	h.HandleFunc("/health_checker", healthChecker)
	server := &http.Server{
		Handler:           h,
		ReadTimeout:       30 * time.Second,
		ReadHeaderTimeout: 30 * time.Second,
		IdleTimeout:       30 * time.Second,
		WriteTimeout:      10 * time.Second,
	}

	if cfg.RegisterConsul {
		var tags []string
		tags = append(tags, fmt.Sprintf("app=%s", cfg.ProductName))
		tags = append(tags, fmt.Sprintf("country=%s", cfg.ServiceRegion))
		err = consulRegister(port, "", tags, cfg.AppName)
		if err != nil {
			tlog.Error("consul register failed, err: ", err)
			return err
		}
	}

	if err = server.Serve(lis); err != nil {
		tlog.Error("fail to init prometheus exporter, err: ", err)
		return err
	}

	return nil
}

func consulRegister(checkPort string, regAddr string, tags []string, name string) (err error) {
	config := consulapi.DefaultConfig()
	if regAddr != "" {
		config.Address = regAddr
	}

	client, err := consulapi.NewClient(config)
	if err != nil {
		tlog.Error("fail to new consul client, err: ", err)
		return err
	}

	intPort, _ := strconv.Atoi(checkPort)
	hostName, _ := os.Hostname()

	registration := &consulapi.AgentServiceRegistration{
		Port: intPort,
		Name: name,
		ID:   name + ":" + hostName,
		Tags: append(tags, "lang=go"),
		Check: &consulapi.AgentServiceCheck{
			HTTP:     fmt.Sprintf("http://%s:%s%s", "127.0.0.1", checkPort, "/health_checker"),
			Interval: "10s",
			Timeout:  "10s",
		},
	}

	return client.Agent().ServiceRegister(registration)
}

func healthChecker(resp http.ResponseWriter, req *http.Request) {
	data := map[string]interface{}{
		"code":        0,
		"message":     "ok",
		"server_time": time.Now().Unix(),
	}
	res, _ := json.Marshal(data)
	resp.Write(res)
}