feat: send events in batches

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:

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,
+		Disabled:     false,
-		Debug:       false,
+		Debug:        false,
-		QueueSize:   1000,
+		QueueSize:    1000,
-		TrackErrors: false,
+		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)
 		}
 	}
 

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"

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
+	var result authResponse
-	err := sendRequestAndParse(ctx, umamiHost+"/api/auth/login", Auth{
+	err := sendRequestAndParse(ctx, umamiHost+"/api/auth/login", authRequest{
 		Username: umamiUsername,
 		Password: umamiPassword,
 	}, nil, &result)

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

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) {
+func (h *UmamiFeeder) reportEventsToUmami(ctx context.Context, events []*SendBody) {
-	hostname := event.Hostname
+	h.debug("reporting %d events", len(events))
-	websiteId, ok := h.websites[hostname]
+	resp, err := sendRequest(ctx, h.umamiHost+"/api/batch", events, nil)
-	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)
 	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()
 	}()