feat: send events in batches

This commit is contained in:
Oleh Astappiev 2025-03-31 18:41:09 +02:00
parent 07ef234cac
commit c84831d4f4
No known key found for this signature in database
GPG key ID: 37D7AC22589E9632
6 changed files with 101 additions and 43 deletions

View file

@ -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:

View file

@ -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)
}
}

View file

@ -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"

View file

@ -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)

View file

@ -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
}

View file

@ -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()
}()