visigo.go

package visigo

import (
	"errors"
	"net/http"

	"github.com/axiomhq/hyperloglog"
	"github.com/tomasen/realip"
)

var counter map[string]*hyperloglog.Sketch

// ErrCount - error returned when you try to get count but didn't register middleware
var ErrCount = errors.New("count not found or error in HyperLogLog")

// Visits - get visits for given URL
func Visits(r *http.Request) (uint64, error) {
	if counter == nil {
		// no, you didn't ...
		panic("you need to register Visigo Counter first!")
	}
	if hll, found := counter[r.URL.String()]; found {
		return hll.Estimate(), nil
	}
	return 0, ErrCount
}

// TotalVisits gets total visits to all sites
func TotalVisits() (uint64, error) {
	hll := hyperloglog.New()
	for _, s := range counter {
		if err := hll.Merge(s); err != nil {
			return 0, err
		}
	}
	return hll.Estimate(), nil
}

// Counter - registers middleware for visits counting
func Counter(next http.Handler) http.Handler {
	counter = make(map[string]*hyperloglog.Sketch)
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if hll, found := counter[r.URL.String()]; found {
			hll.Insert([]byte(realip.RealIP(r)))
		} else {
			l := hyperloglog.New()
			l.Insert([]byte(realip.RealIP(r)))
			counter[r.URL.String()] = l
		}
		next.ServeHTTP(w, r)
	})
}