twitter-prometheus/main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	log "github.com/sirupsen/logrus"
)

func probeTwitter(ctx context.Context, target string, registry *prometheus.Registry) (success bool) {
	followersGauge := prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "twitter_followers",
		Help: "The number of followers of the twitter account.",
	},
		[]string{"screen_name"})

	registry.MustRegister(followersGauge)

	res, err := http.Get(fmt.Sprintf("https://cdn.syndication.twimg.com/widgets/followbutton/info.json?screen_names=%s", url.QueryEscape(target)))
	if err != nil {
		log.Errorf("fetching url: %s", err)
		return false
	}

	data := []struct {
		FollowersCount int `json:"followers_count"`
	}{}
	if err := json.NewDecoder(res.Body).Decode(&data); err != nil {
		log.Errorf("decoding json: %s", err)
		return false
	}
	switch l := len(data); {
	case l == 0:
		log.Errorf("No data returned. Does screen name exist?")
		return false
	case l > 1:
		log.Errorf("Unexpected (>1) results returned. WTF twitter??")
		return false
	}
	followersGauge.With(prometheus.Labels{"screen_name": target}).Set(float64(data[0].FollowersCount))
	return true
}

func twitterHandler(w http.ResponseWriter, r *http.Request) {
	screenName := r.URL.Query().Get("screen_name")
	if screenName == "" {
		http.Error(w, "screen_name parameter is missing", http.StatusBadRequest)
		return
	}

	reg := prometheus.NewRegistry()

	ctx, cancel := context.WithTimeout(r.Context(), time.Duration(5*time.Second))
	defer cancel()

	if success := probeTwitter(ctx, screenName, reg); !success {
		log.Error("Probe failed!")
	}

	h := promhttp.HandlerFor(reg, promhttp.HandlerOpts{})
	h.ServeHTTP(w, r)

}

func bail(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func main() {
	var listenAddr string
	flag.StringVar(&listenAddr, "listen", "0.0.0.0:9700", "Address to listen on")
	flag.Parse()
	http.Handle("/metrics", promhttp.Handler())
	http.HandleFunc("/probe", func(w http.ResponseWriter, r *http.Request) {
		twitterHandler(w, r)
	})
	log.Infof("Listening on %s", listenAddr)
	bail(http.ListenAndServe(listenAddr, nil))
}