diff --git a/demo/docker-compose.yml b/demo/docker-compose.yml index 536ca95..aadf351 100644 --- a/demo/docker-compose.yml +++ b/demo/docker-compose.yml @@ -1,6 +1,8 @@ services: traefik: image: traefik:3 + depends_on: + - umami labels: - "traefik.enable=true" - "traefik.http.routers.traefik.service=api@internal" @@ -47,8 +49,11 @@ services: umami: image: ghcr.io/umami-software/umami:postgresql-latest depends_on: - - postgres + postgres: + condition: service_healthy environment: + - DEBUG=umami:middleware + - ENABLE_TEST_CONSOLE=1 - DATABASE_URL=postgres://umami:umami@postgres:5432/umami ports: - target: 3000 @@ -68,6 +73,11 @@ services: - umami_postgres volumes: - ./postgres:/var/lib/postgresql/data + healthcheck: + test: [ "CMD-SHELL", "sh -c 'pg_isready -U ${POSTGRES_USER} -d ${POSTGRES_DB}'" ] + interval: 10s + timeout: 3s + retries: 3 networks: traefik: diff --git a/umami.go b/umami.go index 414f0b4..d4fe5cf 100644 --- a/umami.go +++ b/umami.go @@ -10,6 +10,7 @@ import ( "path" "regexp" "strings" + "sync" "time" ) @@ -21,6 +22,10 @@ type Config struct { Debug bool `json:"debug"` // QueueSize defines the size of queue, i.e. the amount of events that are waiting to be submitted to Umami. QueueSize int `json:"queueSize"` + // BatchSize defines the amount of events that are submitted to Umami in one request. + BatchSize int `json:"batchSize"` + // BatchMaxWait defines the maximum time to wait before submitting the batch. + BatchMaxWait time.Duration `json:"batchMaxWait"` // UmamiHost is the URL of the Umami instance. UmamiHost string `json:"umamiHost"` @@ -60,10 +65,12 @@ type Config struct { // CreateConfig creates the default plugin configuration. func CreateConfig() *Config { return &Config{ - Disabled: false, - Debug: false, - QueueSize: 1000, - TrackErrors: false, + Disabled: false, + Debug: false, + QueueSize: 1000, + BatchSize: 20, + BatchMaxWait: 5 * time.Second, + TrackErrors: false, UmamiHost: "", UmamiToken: "", @@ -93,10 +100,14 @@ type UmamiFeeder struct { logHandler *log.Logger queue chan *UmamiEvent + batchSize int + batchMaxWait time.Duration + umamiHost string umamiToken string umamiTeamId string websites map[string]string + websitesMutex sync.RWMutex createNewWebsites bool trackErrors bool @@ -121,10 +132,14 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h // Umami API does not support batching https://github.com/umami-software/umami/discussions/1473 queue: make(chan *UmamiEvent, config.QueueSize), + batchSize: config.BatchSize, + batchMaxWait: config.BatchMaxWait, + umamiHost: config.UmamiHost, umamiToken: config.UmamiToken, umamiTeamId: config.UmamiTeamId, websites: config.Websites, + websitesMutex: sync.RWMutex{}, createNewWebsites: config.CreateNewWebsites, trackErrors: config.TrackErrors, @@ -191,7 +206,7 @@ func (h *UmamiFeeder) connect(ctx context.Context, config *Config) error { } h.websites[website.Domain] = website.ID - h.debug("fetched websiteId for: %s", website.Domain) + h.debug("website fetched '%s': %s", website.Domain, website.ID) } } diff --git a/umami_test.go b/umami_test.go index 57c674d..c3408f5 100644 --- a/umami_test.go +++ b/umami_test.go @@ -8,8 +8,6 @@ import ( ) func TestTraefikUmamiFeeder(t *testing.T) { - t.Skip("requires a running Umami instance, see /demo/docker-compose.yml") - cfg := CreateConfig() cfg.UmamiHost = "http://localhost:3000" cfg.UmamiUsername = "admin" diff --git a/umami_token.go b/umami_token.go index 229c6e0..6d1ce19 100644 --- a/umami_token.go +++ b/umami_token.go @@ -2,18 +2,18 @@ package traefik_umami_feeder import "context" -type Auth struct { +type authRequest struct { Username string `json:"username"` Password string `json:"password"` } -type AuthResponse struct { +type authResponse struct { Token string `json:"token"` } func getToken(ctx context.Context, umamiHost string, umamiUsername string, umamiPassword string) (string, error) { - var result AuthResponse - err := sendRequestAndParse(ctx, umamiHost+"/api/auth/login", Auth{ + var result authResponse + err := sendRequestAndParse(ctx, umamiHost+"/api/auth/login", authRequest{ Username: umamiUsername, Password: umamiPassword, }, nil, &result) diff --git a/umami_websites.go b/umami_websites.go index 6ad4891..bf61218 100644 --- a/umami_websites.go +++ b/umami_websites.go @@ -6,7 +6,7 @@ import ( "time" ) -type WebsitesResponse struct { +type websitesResponse struct { Data []Website `json:"data"` Count int `json:"count"` Page int `json:"page"` @@ -49,7 +49,7 @@ func fetchWebsites(ctx context.Context, umamiHost string, umamiToken string, tea url = umamiHost + "/api/teams/" + teamId + "/websites?pageSize=200" } - var result WebsitesResponse + var result websitesResponse err := sendRequestAndParse(ctx, url, nil, headers, &result) if err != nil { @@ -58,3 +58,30 @@ func fetchWebsites(ctx context.Context, umamiHost string, umamiToken string, tea return &result.Data, nil } + +func getWebsiteId(h *UmamiFeeder, hostname string) string { + h.websitesMutex.RLock() + websiteId, ok := h.websites[hostname] + h.websitesMutex.RUnlock() + + if ok { + return websiteId + } + + h.websitesMutex.Lock() + defer h.websitesMutex.Unlock() + + // Create a background context for the API call + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + website, err := createWebsite(ctx, h.umamiHost, h.umamiToken, h.umamiTeamId, hostname) + if err != nil { + h.error("failed to create website: " + err.Error()) + return "" + } + + h.websites[website.Domain] = website.ID + h.debug("website created '%s': %s", website.Domain, website.ID) + return website.ID +} diff --git a/umami_worker.go b/umami_worker.go index ef7a691..44ec173 100644 --- a/umami_worker.go +++ b/umami_worker.go @@ -3,6 +3,7 @@ package traefik_umami_feeder import ( "context" "fmt" + "io" "net/http" "time" ) @@ -28,14 +29,23 @@ type SendBody struct { } func (h *UmamiFeeder) submitToFeed(req *http.Request, code int) { + hostname := parseDomainFromHost(req.Host) + websiteId := getWebsiteId(h, hostname) + + if websiteId == "" { + h.error("tracking skipped, websiteId is unknown: " + hostname) + return + } + event := &UmamiEvent{ - Hostname: parseDomainFromHost(req.Host), + Hostname: hostname, Language: parseAcceptLanguage(req.Header.Get("Accept-Language")), Referrer: req.Referer(), Url: req.URL.String(), Ip: extractRemoteIP(req), UserAgent: req.Header.Get("User-Agent"), Timestamp: time.Now().Unix(), + Website: websiteId, } select { @@ -65,50 +75,48 @@ func (h *UmamiFeeder) umamiEventFeeder(ctx context.Context) (err error) { } }() + batch := make([]*SendBody, 0, h.batchSize) + timeout := time.NewTimer(h.batchMaxWait) + for { // Wait for event. select { case <-ctx.Done(): h.debug("worker shutting down (canceled)") + if len(batch) > 0 { + h.reportEventsToUmami(ctx, batch) + } return nil case event := <-h.queue: - h.reportEventToUmami(ctx, event) + batch = append(batch, &SendBody{Payload: event, Type: "event"}) + if len(batch) >= h.batchSize { + h.reportEventsToUmami(ctx, batch) + batch = make([]*SendBody, 0, h.batchSize) + timeout.Reset(h.batchMaxWait) + } + + case <-timeout.C: + if len(batch) > 0 { + h.reportEventsToUmami(ctx, batch) + batch = make([]*SendBody, 0, h.batchSize) + } + timeout.Reset(h.batchMaxWait) } } } -func (h *UmamiFeeder) reportEventToUmami(ctx context.Context, event *UmamiEvent) { - hostname := event.Hostname - websiteId, ok := h.websites[hostname] - if !ok { - website, err := createWebsite(ctx, h.umamiHost, h.umamiToken, h.umamiTeamId, hostname) - if err != nil { - h.error("failed to create website: " + err.Error()) - return - } - - h.websites[website.Domain] = website.ID - websiteId = website.ID - h.debug("created website for: %s", website.Domain) - } - if websiteId == "" { - h.error("skip tracking, websiteId is unknown: " + hostname) - return - } - event.Website = websiteId - - body := SendBody{ - Payload: event, - Type: "event", - } - - h.debug("sending tracking request %v", event) - resp, err := sendRequest(ctx, h.umamiHost+"/api/send", body, nil) +func (h *UmamiFeeder) reportEventsToUmami(ctx context.Context, events []*SendBody) { + h.debug("reporting %d events", len(events)) + resp, err := sendRequest(ctx, h.umamiHost+"/api/batch", events, nil) if err != nil { h.error("failed to send tracking: " + err.Error()) return } + if h.isDebug { + bodyBytes, _ := io.ReadAll(resp.Body) + h.debug("%v: %s", resp.Status, string(bodyBytes)) + } defer func() { _ = resp.Body.Close() }()