Add OTel tracing support (#774)

go.mod

@@ -19,6 +19,8 @@ require (
 	github.com/lib/pq v1.10.9
 	github.com/minio/minio-go/v7 v7.0.80
 	github.com/pkg/errors v0.9.1
+	github.com/ravilushqa/otelgqlgen v0.17.0
+	github.com/riandyrn/otelchi v0.11.0
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/rs/cors v1.11.1
 	github.com/sirupsen/logrus v1.9.3
@@ -28,9 +30,15 @@ require (
 	github.com/vektah/gqlparser/v2 v2.5.19
 	github.com/wneessen/go-mail v0.5.2
 	go.deanishe.net/favicon v0.1.0
+	go.nhat.io/otelsql v0.14.0
+	go.opentelemetry.io/otel v1.33.0
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.33.0
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.33.0
+	go.opentelemetry.io/otel/sdk v1.33.0
+	go.opentelemetry.io/otel/trace v1.33.0
 	golang.org/x/crypto v0.31.0
 	golang.org/x/image v0.22.0
-	golang.org/x/net v0.31.0
+	golang.org/x/net v0.32.0
 	golang.org/x/sync v0.10.0
 	gotest.tools/v3 v3.5.1
 )
@@ -39,17 +47,22 @@ require (
 	github.com/PuerkitoBio/goquery v1.9.3 // indirect
 	github.com/agnivade/levenshtein v1.1.1 // indirect
 	github.com/andybalholm/cascadia v1.3.2 // indirect
+	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.5 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
+	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/friendsofgo/errors v0.9.2 // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
 	github.com/go-ini/ini v1.67.0 // indirect
+	github.com/go-logr/logr v1.4.2 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-viper/mapstructure/v2 v2.2.1 // indirect
 	github.com/goccy/go-json v0.10.3 // indirect
 	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/gorilla/securecookie v1.1.2 // indirect
-	github.com/gorilla/websocket v1.5.0 // indirect
+	github.com/gorilla/websocket v1.5.3 // indirect
+	github.com/grpc-ecosystem/grpc-gateway/v2 v2.24.0 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
@@ -70,6 +83,10 @@ require (
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/urfave/cli/v2 v2.27.5 // indirect
 	github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 // indirect
+	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
+	go.opentelemetry.io/contrib v1.29.0 // indirect
+	go.opentelemetry.io/otel/metric v1.33.0 // indirect
+	go.opentelemetry.io/proto/otlp v1.4.0 // indirect
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.9.0 // indirect
 	golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
@@ -78,6 +95,10 @@ require (
 	golang.org/x/text v0.21.0 // indirect
 	golang.org/x/tools v0.24.0 // indirect
 	golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20241209162323-e6fa225c2576 // indirect
+	google.golang.org/grpc v1.68.1 // indirect
+	google.golang.org/protobuf v1.35.2 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

main.go

@@ -2,11 +2,13 @@
 package main
 
 import (
+	"context"
 	"embed"
 
 	"github.com/stashapp/stash-box/pkg/api"
 	"github.com/stashapp/stash-box/pkg/database"
 	"github.com/stashapp/stash-box/pkg/image"
+	"github.com/stashapp/stash-box/pkg/logger"
 	"github.com/stashapp/stash-box/pkg/manager"
 	"github.com/stashapp/stash-box/pkg/manager/config"
 	"github.com/stashapp/stash-box/pkg/manager/cron"
@@ -19,12 +21,17 @@ var ui embed.FS
 
 func main() {
 	manager.Initialize()
+
+	cleanup := logger.InitTracer()
+	//nolint:errcheck
+	defer cleanup(context.Background())
+
 	api.InitializeSession()
 
 	const databaseProvider = "postgres"
 	db := database.Initialize(databaseProvider, config.GetDatabasePath())
 	txnMgr := sqlx.NewTxnMgr(db)
-	user.CreateSystemUsers(txnMgr.Repo())
+	user.CreateSystemUsers(txnMgr.Repo(context.Background()))
 	api.Start(txnMgr, ui)
 	cron.Init(txnMgr)
 

pkg/api/factory.go

@@ -10,14 +10,15 @@ import (
 type RepoProvider interface {
 	// IMPORTANT: the returned Repo object MUST NOT be shared between goroutines.
 	// that is: call Repo for each new request/goroutine
-	Repo() models.Repo
+	Repo(ctx context.Context) models.Repo
 }
 
 // creates a new Repo (with its own transaction boundary) for each incoming request
 func repoMiddleware(provider RepoProvider) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			r = r.WithContext(context.WithValue(r.Context(), ContextRepo, provider.Repo()))
+			ctx := r.Context()
+			r = r.WithContext(context.WithValue(ctx, ContextRepo, provider.Repo(ctx)))
 
 			next.ServeHTTP(w, r)
 		})

pkg/api/server.go

@@ -14,13 +14,19 @@ import (
 	"strings"
 
 	"github.com/klauspost/compress/flate"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
 
+	"github.com/ravilushqa/otelgqlgen"
+
+	"github.com/99designs/gqlgen/graphql"
 	gqlHandler "github.com/99designs/gqlgen/graphql/handler"
 	gqlExtension "github.com/99designs/gqlgen/graphql/handler/extension"
 	gqlTransport "github.com/99designs/gqlgen/graphql/handler/transport"
 	gqlPlayground "github.com/99designs/gqlgen/graphql/playground"
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/chi/v5/middleware"
+	"github.com/riandyrn/otelchi"
 	"github.com/rs/cors"
 	"github.com/stashapp/stash-box/pkg/dataloader"
 	"github.com/stashapp/stash-box/pkg/logger"
@@ -93,6 +99,11 @@ func authenticateHandler() func(http.Handler) http.Handler {
 			ctx = context.WithValue(ctx, user.ContextUser, u)
 			ctx = context.WithValue(ctx, user.ContextRoles, roles)
 
+			span := trace.SpanFromContext(ctx)
+			if span.SpanContext().IsValid() && u != nil {
+				span.SetAttributes(attribute.String("user.id", u.ID.String()))
+			}
+
 			r = r.WithContext(ctx)
 
 			next.ServeHTTP(w, r)
@@ -110,6 +121,7 @@ func redirect(w http.ResponseWriter, req *http.Request) {
 
 func Start(rfp RepoProvider, ui embed.FS) {
 	r := chi.NewRouter()
+	r.Use(otelchi.Middleware("", otelchi.WithChiRoutes(r)))
 
 	var corsConfig *cors.Cors
 	if config.GetIsProduction() {
@@ -154,8 +166,9 @@ func Start(rfp RepoProvider, ui embed.FS) {
 	gqlSrv.AddTransport(gqlTransport.POST{})
 	gqlSrv.AddTransport(gqlTransport.MultipartForm{})
 	gqlSrv.Use(gqlExtension.Introspection{})
+	gqlSrv.Use(otelgqlgen.Middleware(otelgqlgen.WithCreateSpanFromFields(func(fieldCtx *graphql.FieldContext) bool { return fieldCtx.IsResolver })))
 
-	r.Handle("/graphql", dataloader.Middleware(rfp.Repo())(gqlSrv))
+	r.Handle("/graphql", dataloader.Middleware(getRepo)(gqlSrv))
 
 	if !config.GetIsProduction() {
 		r.Handle("/playground", gqlPlayground.Handler("GraphQL playground", "/graphql"))

pkg/database/databasetest/database_test_utils.go

@@ -1,6 +1,7 @@
 package databasetest
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"fmt"
@@ -59,7 +60,7 @@ func initPostgres(connString string) func() {
 
 	db = database.Initialize(databaseType, connString)
 	txnMgr := sqlxx.NewTxnMgr(db)
-	repo = txnMgr.Repo()
+	repo = txnMgr.Repo(context.TODO())
 
 	return teardownPostgres
 }

pkg/database/postgres.go

@@ -1,6 +1,7 @@
 package database
 
 import (
+	"database/sql"
 	"embed"
 	"fmt"
 	"time"
@@ -9,11 +10,13 @@ import (
 	"github.com/jmoiron/sqlx"
 	"github.com/stashapp/stash-box/pkg/logger"
 	"github.com/stashapp/stash-box/pkg/manager/config"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 
 	// Driver used here only
 	_ "github.com/golang-migrate/migrate/v4/database/postgres"
 	"github.com/golang-migrate/migrate/v4/source/iofs"
 	_ "github.com/lib/pq"
+	"go.nhat.io/otelsql"
 )
 
 const postgresDriver = "postgres"
@@ -30,7 +33,25 @@ type PostgresProvider struct{}
 func (p *PostgresProvider) Open(databasePath string) *sqlx.DB {
 	p.runMigrations(databasePath)
 
-	conn, err := sqlx.Open(postgresDriver, "postgres://"+databasePath)
+	driverName, err := otelsql.Register("postgres",
+		otelsql.TraceQueryWithoutArgs(),
+		otelsql.WithSystem(semconv.DBSystemPostgreSQL),
+	)
+
+	if err != nil {
+		logger.Fatalf("db.Open(): %q\n", err)
+	}
+
+	db, err := sql.Open(driverName, "postgres://"+databasePath)
+	if err != nil {
+		logger.Fatalf("db.Open(): %q\n", err)
+	}
+
+	if err := otelsql.RecordStats(db); err != nil {
+		logger.Fatalf("db.Open(): %q\n", err)
+	}
+
+	conn := sqlx.NewDb(db, "postgres")
 	conn.SetMaxOpenConns(config.GetMaxOpenConns())
 	conn.SetMaxIdleConns(config.GetMaxIdleConns())
 	conn.SetConnMaxLifetime(time.Duration(config.GetConnMaxLifetime()) * time.Minute)

pkg/dataloader/loaders.go

@@ -46,10 +46,12 @@ type Loaders struct {
 	EditCommentByID                EditCommentLoader
 }
 
-func Middleware(fac models.Repo) func(next http.Handler) http.Handler {
+func Middleware(getRepo func(ctx context.Context) models.Repo) func(next http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			ctx := context.WithValue(r.Context(), loadersKey, GetLoaders(r.Context(), fac))
+			ctx := r.Context()
+			fac := getRepo(ctx)
+			ctx = context.WithValue(ctx, loadersKey, GetLoaders(r.Context(), fac))
 			r = r.WithContext(ctx)
 			next.ServeHTTP(w, r)
 		})

pkg/logger/otel.go

@@ -0,0 +1,56 @@
+package logger
+
+import (
+	"context"
+	"log"
+
+	"github.com/stashapp/stash-box/pkg/manager/config"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
+	"go.opentelemetry.io/otel/sdk/resource"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+)
+
+func InitTracer() func(context.Context) error {
+	otelConfig := config.GetOTelConfig()
+	if otelConfig == nil {
+		return nil
+	}
+
+	exporter, err := otlptrace.New(
+		context.Background(),
+		otlptracegrpc.NewClient(
+			otlptracegrpc.WithInsecure(),
+			otlptracegrpc.WithEndpoint(otelConfig.Endpoint),
+		),
+	)
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	resources, err := resource.New(
+		context.Background(),
+		resource.WithAttributes(
+			attribute.String("service.name", config.GetTitle()),
+			attribute.String("library.language", "go"),
+		),
+	)
+	if err != nil {
+		log.Print("Could not set resources: ", err)
+	}
+
+	otel.SetTracerProvider(
+		sdktrace.NewTracerProvider(
+			sdktrace.WithSampler(sdktrace.TraceIDRatioBased(otelConfig.TraceRatio)),
+			sdktrace.WithBatcher(exporter),
+			sdktrace.WithResource(resources),
+		),
+	)
+
+	logger.Infof("otel initialized with collector: %s, ratio: %f", otelConfig.Endpoint, otelConfig.TraceRatio)
+
+	return exporter.Shutdown
+}

pkg/manager/config/config.go

@@ -25,6 +25,11 @@ type PostgresConfig struct {
 	ConnMaxLifetime int `mapstructure:"conn_max_lifetime"`
 }
 
+type OTelConfig struct {
+	Endpoint   string  `mapstructure:"endpoint"`
+	TraceRatio float64 `mapstructure:"trace_ratio"`
+}
+
 type config struct {
 	Host         string `mapstructure:"host"`
 	Port         int    `mapstructure:"port"`
@@ -91,6 +96,10 @@ type config struct {
 		PostgresConfig `mapstructure:",squash"`
 	}
 
+	OTel struct {
+		OTelConfig `mapstructure:",squash"`
+	}
+
 	PHashDistance int `mapstructure:"phash_distance"`
 
 	Title string `mapstructure:"title"`
@@ -246,6 +255,13 @@ func GetS3Config() *S3Config {
 	return &C.S3.S3Config
 }
 
+func GetOTelConfig() *OTelConfig {
+	if C.OTel.Endpoint != "" {
+		return &C.OTel.OTelConfig
+	}
+	return nil
+}
+
 // ValidateImageLocation returns an error is image_location is not set.
 func ValidateImageLocation() error {
 	if C.ImageLocation == "" {