twitter-prometheus/main.go

89 lines
2.2 KiB
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) {
username := r.URL.Query().Get("username")
if username == "" {
http.Error(w, "username 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, username, 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))
}