Muninのhttp_loadtimeをGolangで

munin_http_loadtime.go

package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"
)

func main() {
	var urls []string

	target, targetIsSet := os.LookupEnv("target")
	if targetIsSet {
		urls = strings.Split(target, ",")
	}

	switch len(os.Args) {
	case 1:
		fetch(urls)
	case 2:
		c := os.Args[1]
		switch c {
		case "autoconf":
			fmt.Println("yes")
		case "config":
			config(urls)
		default:
			fmt.Fprintf(os.Stderr, "Unknown command: %s\n", c)
			os.Exit(1)
		}
	default:
		fmt.Fprintln(os.Stderr, "Too many arguments.")
		os.Exit(1)
	}
}

func config(urls []string) {
	fmt.Println("graph_title HTTP loadtime of a page")
	fmt.Println("graph_args --base 1000 -l -1")
	fmt.Println("graph_vlabel Load time in seconds")
	fmt.Println("graph_category network")
	fmt.Println("graph_info This graph shows the load time in seconds")

	for _, url := range urls {
		var urlShort string
		if len(url) > 30 {
			urlShort = url[0:30] + "..."
		} else {
			urlShort = url
		}

		escapedUrl := escapeURL(url)

		fmt.Printf("%s.label %s\n", escapedUrl, urlShort)
		fmt.Printf("%s.info page load time\n", escapedUrl)
	}
}

type FetchResult struct {
	URL      string
	Duration time.Duration
	Err      bool
}

func fetch(urls []string) {
	if len(urls) < 1 {
		return
	}

	ch := make(chan FetchResult)

	for _, url := range urls {
		go fetchPage(url, ch)
	}

	for range urls {
		r := <-ch

		escapedURL := escapeURL(r.URL)

		if r.Err {
			fmt.Printf("%s.value -1\n", escapedURL)
		} else {
			fmt.Printf("%s.value %f\n", escapedURL, r.Duration.Seconds())
		}
	}
}

func fetchPage(url string, ch chan FetchResult) {
	r := FetchResult{
		URL: url,
		Err: false,
	}

	request, _ := http.NewRequest("GET", url, nil)
	request.Header.Set("Cache-Control", "no-cache, no-store")
	client := &http.Client{}

	t0 := time.Now()

	resp, getErr := client.Do(request)
	if getErr != nil {
		r.Err = true
		ch <- r
		return
	}

	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		r.Err = true
		ch <- r
		return
	}

	_, readErr := ioutil.ReadAll(resp.Body)
	if readErr != nil {
		r.Err = true
		ch <- r
		return
	}

	r.Duration = time.Since(t0)
	ch <- r
}

var escapeRe = regexp.MustCompile(`[-:/.]`)

func escapeURL(url string) string {
	return escapeRe.ReplaceAllString(url, "_")
}