trantor/lib/stats.go

package trantor

import (
	log "github.com/cihub/seelog"

	"net/http"
	"strings"

	"github.com/gorilla/mux"
	"gitlab.com/trantor/trantor/lib/database"
	"gitlab.com/trantor/trantor/lib/storage"
)

const (
	stats_version = 2
	statsChanSize = 100
)

type handler struct {
	w        http.ResponseWriter
	r        *http.Request
	sess     *Session
	db       database.DB
	store    storage.Store
	template *Template
	hostname string
	ro       bool
}

type StatsGatherer struct {
	db       database.DB
	store    storage.Store
	template *Template
	hostname string
	channel  chan statsRequest
	ro       bool
}

func InitStats(database database.DB, store storage.Store, hostname string, template *Template, ro bool) *StatsGatherer {
	sg := StatsGatherer{
		channel:  make(chan statsRequest, statsChanSize),
		db:       database,
		store:    store,
		template: template,
		hostname: hostname,
		ro:       ro,
	}

	go sg.worker()
	return &sg
}

func (sg StatsGatherer) Gather(function func(handler)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Info("Query ", r.Method, " ", r.RequestURI)

		h := handler{
			store:    sg.store,
			db:       sg.db,
			template: sg.template,
			hostname: sg.hostname,
			w:        w,
			r:        r,
			sess:     GetSession(r, sg.db),
			ro:       sg.ro,
		}

		function(h)
		sg.channel <- statsRequest{r}
	}
}

type statsRequest struct {
	r *http.Request
}

func (sg StatsGatherer) worker() {
	for req := range sg.channel {
		var err error
		pattern := strings.Split(req.r.URL.Path, "/")
		id := mux.Vars(req.r)["id"]
		if len(pattern) > 1 && pattern[1] != "" && id != "" {
			switch pattern[1] {
			case "download":
				id := mux.Vars(req.r)["id"]
				err = sg.db.IncDownloads(id)
			case "book":
				id := mux.Vars(req.r)["id"]
				err = sg.db.IncViews(id)
			case "read":
				id := mux.Vars(req.r)["id"]
				err = sg.db.IncViews(id)
			}
		}
		if err != nil {
			log.Warn("Problem incrementing visits: ", err)
		}
	}
}