fix(storage): does not store stats if request fails

4 files changed, 34 insertions, 27 deletions

diff --git a/backend/data.go b/backend/data.go
index 5bc5391..2f2e06d 100644
--- a/backend/data.go
+++ b/backend/data.go
@@ -111,6 +111,11 @@ func (d *data) handleGeneric(w http.ResponseWriter, r *http.Request, name string
 	if err != nil {
 		panic(err)
 	}
+	if name == "404" {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	UpdateStats(r)
 }
 
 func (d *data) handleRSS(w http.ResponseWriter, r *http.Request, custom dataUsable) {
diff --git a/backend/router.go b/backend/router.go
index 161433d..1a2f024 100644
--- a/backend/router.go
+++ b/backend/router.go
@@ -91,28 +91,11 @@ func NewRouter(debug bool, cfg *Config, db *sql.DB, assets fs.FS) *chi.Mux {
 		})
 	})
 	// context
-	setContext := func(ctx context.Context) context.Context {
-		ctx = context.WithValue(ctx, configKey, cfg)
-		ctx = context.WithValue(ctx, assetsFSKey, assets)
-		ctx = context.WithValue(ctx, debugKey, debug)
-		return context.WithValue(ctx, storage.DBKey, db)
-	}
-	r.Use(func(next http.Handler) http.Handler {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			next.ServeHTTP(w, r.WithContext(setContext(r.Context())))
-		})
-	})
-	// stats
 	r.Use(func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			go func(r *http.Request) {
-				ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
-				defer cancel()
-				if err := storage.UpdateStats(setContext(ctx), r, cfg.Domain); err != nil {
-					slog.Error("updating stats", "error", err)
-				}
-			}(r)
-			next.ServeHTTP(w, r)
+			next.ServeHTTP(w, r.WithContext(
+				setContext(r.Context(), cfg, debug, db, assets),
+			))
 		})
 	})
 	// login
@@ -179,3 +162,23 @@ func HandleStaticFiles(r *chi.Mux, path string, root fs.FS) {
 		http.StripPrefix(pathPrefix, http.FileServerFS(root)).ServeHTTP(w, req)
 	})
 }
+
+func UpdateStats(r *http.Request) {
+	ctx := r.Context()
+	cfg := ctx.Value(configKey).(*Config)
+	debug := ctx.Value(debugKey).(bool)
+	db := ctx.Value(storage.DBKey).(*sql.DB)
+	assets := ctx.Value(assetsFSKey).(fs.FS)
+
+	ctx2, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	if err := storage.UpdateStats(setContext(ctx2, cfg, debug, db, assets), r, cfg.Domain); err != nil {
+		slog.Error("updating stats", "error", err)
+	}
+}
+func setContext(ctx context.Context, cfg *Config, debug bool, db *sql.DB, assets fs.FS) context.Context {
+	ctx = context.WithValue(ctx, configKey, cfg)
+	ctx = context.WithValue(ctx, assetsFSKey, assets)
+	ctx = context.WithValue(ctx, debugKey, debug)
+	return context.WithValue(ctx, storage.DBKey, db)
+}
diff --git a/backend/storage/db.go b/backend/storage/db.go
index 7a7c6e0..0ee50da 100644
--- a/backend/storage/db.go
+++ b/backend/storage/db.go
@@ -13,6 +13,8 @@ import (
 	_ "github.com/mattn/go-sqlite3"
 )
 
+const DBKey = "db"
+
 //go:embed migrations
 var migrations embed.FS
 
@@ -66,3 +68,7 @@ func RunMigration(ctx context.Context, db *sql.DB) error {
 	}
 	return nil
 }
+
+func getDB(ctx context.Context) *sql.DB {
+	return ctx.Value(DBKey).(*sql.DB)
+}
diff --git a/backend/storage/stats.go b/backend/storage/stats.go
index 757168d..d95387d 100644
--- a/backend/storage/stats.go
+++ b/backend/storage/stats.go
@@ -2,7 +2,6 @@ package storage
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 	"log/slog"
 	"net/http"
@@ -13,8 +12,6 @@ import (
 	"time"
 )
 
-const DBKey = "db"
-
 type loaded struct {
 	data map[string]struct{}
 	mu   *sync.RWMutex
@@ -50,10 +47,6 @@ var trimRefererReg = regexp.MustCompile(`https?://([a-z-0-9.]+(:\d+)?)/.*`)
 
 var load = newLoaded()
 
-func getDB(ctx context.Context) *sql.DB {
-	return ctx.Value(DBKey).(*sql.DB)
-}
-
 func UpdateStats(ctx context.Context, r *http.Request, domain string) error {
 	target := r.URL.Path
 	if strings.HasPrefix(target, "/static") || strings.HasPrefix(target, "/admin") {