refactor(api): mark spans for sampling without copying logs

Keep standalone log output intact and simplify telemetry conditions and test fixtures.
Propagate heartbeat update contexts and defer SQL formatting until a log is emitted.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 60fb2b83-6e9d-4656-9f60-6183c61b716d

Author: AchoArnold

api/pkg/di/container.go

@@ -262,9 +262,9 @@ func (container *Container) DedicatedDB() (db *gorm.DB) {
 
 	config := &gorm.Config{
 		TranslateError: true,
-		Logger: telemetry.NewGormLoggerWithConfig(
+		Logger: telemetry.NewGormSpanLogger(
 			container.Tracer(),
-			container.Logger(6),
+			gormLogger.Default,
 			gormLogger.Config{
 				SlowThreshold: 200 * time.Millisecond,
 				LogLevel:      gormLogger.Warn,

api/pkg/repositories/gorm_heartbeat_monitor_repository.go

@@ -53,7 +53,7 @@ func (repository *gormHeartbeatMonitorRepository) UpdatePhoneOnline(ctx context.
 	ctx, cancel := context.WithTimeout(ctx, dbOperationDuration)
 	defer cancel()
 
-	err := repository.db.
+	err := repository.db.WithContext(ctx).
 		Model(&entities.HeartbeatMonitor{}).
 		Where("id = ?", monitorID).
 		Where("user_id = ?", userID).
@@ -75,7 +75,7 @@ func (repository *gormHeartbeatMonitorRepository) UpdateQueueID(ctx context.Cont
 	ctx, cancel := context.WithTimeout(ctx, dbOperationDuration)
 	defer cancel()
 
-	err := repository.db.
+	err := repository.db.WithContext(ctx).
 		Model(&entities.HeartbeatMonitor{}).
 		Where("id = ?", monitorID).
 		Updates(map[string]any{

api/pkg/repositories/gorm_heartbeat_monitor_repository_test.go

@@ -0,0 +1,47 @@
+package repositories
+
+import (
+	"context"
+	"testing"
+
+	"github.com/NdoleStudio/httpsms/pkg/telemetry"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+)
+
+func TestHeartbeatMonitorUpdatesCarryContext(t *testing.T) {
+	updates := map[string]func(HeartbeatMonitorRepository, context.Context) error{
+		"phone online": func(repository HeartbeatMonitorRepository, ctx context.Context) error {
+			return repository.UpdatePhoneOnline(ctx, "user-id", uuid.New(), true)
+		},
+		"queue ID": func(repository HeartbeatMonitorRepository, ctx context.Context) error {
+			return repository.UpdateQueueID(ctx, uuid.New(), "queue-id")
+		},
+	}
+	for name, update := range updates {
+		t.Run(name, func(t *testing.T) {
+			db, err := gorm.Open(postgres.New(postgres.Config{Conn: &messageThreadTestConnPool{}}), &gorm.Config{
+				DisableAutomaticPing:   true,
+				SkipDefaultTransaction: true,
+				DryRun:                 true,
+			})
+			require.NoError(t, err)
+			type contextKey struct{}
+			ctx := context.WithValue(context.Background(), contextKey{}, "request")
+			called := false
+			require.NoError(t, db.Callback().Update().Before("gorm:update").Register("check_context", func(db *gorm.DB) {
+				called = true
+				require.Equal(t, "request", db.Statement.Context.Value(contextKey{}))
+				_, hasDeadline := db.Statement.Context.Deadline()
+				require.True(t, hasDeadline)
+			}))
+			logger := &messageThreadTestLogger{}
+			repository := NewGormHeartbeatMonitorRepository(logger, telemetry.NewOtelLogger("test", logger), db)
+
+			require.NoError(t, update(repository, ctx))
+			require.True(t, called)
+		})
+	}
+}

api/pkg/telemetry/gorm_logger.go

@@ -14,7 +14,11 @@ import (
 type gormLogger struct {
 	tracer Tracer
 	logger Logger
-	config logger.Config
+	config gormLogConfig
+}
+
+type gormLogConfig struct {
+	logger.Config
 }
 
 // NewGormLogger creates a new instance of gormLogger
@@ -27,7 +31,7 @@ func NewGormLoggerWithConfig(tracer Tracer, log Logger, config logger.Config) lo
 	return &gormLogger{
 		tracer: tracer,
 		logger: log,
-		config: config,
+		config: gormLogConfig{Config: config},
 	}
 }
 
@@ -57,20 +61,96 @@ func (gorm *gormLogger) Error(ctx context.Context, s string, i ...any) {
 }
 
 func (gorm *gormLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
-	if gorm.config.LogLevel <= logger.Silent {
+	elapsed := time.Since(begin)
+	level := gorm.config.queryLevel(elapsed, err)
+	if level == logger.Silent {
 		return
 	}
-	elapsed := time.Since(begin)
 	l := gorm.tracer.CtxLogger(gorm.logger, gorm.tracer.Span(ctx)).WithString("latency", elapsed.String())
 	sql, rows := fc()
 	msg := fmt.Sprintf("[ROWS:%d][%s]", rows, sql)
 
-	switch {
-	case err != nil && gorm.config.LogLevel >= logger.Error && (!errors.Is(err, gormdb.ErrRecordNotFound) || !gorm.config.IgnoreRecordNotFoundError):
+	switch level {
+	case logger.Error:
 		l.Error(stacktrace.Propagatef(err, "[ROWS:%d][%s]", rows, sql))
-	case gorm.config.SlowThreshold > 0 && elapsed > gorm.config.SlowThreshold && gorm.config.LogLevel >= logger.Warn:
+	case logger.Warn:
 		l.Warn(stacktrace.NewErrorf("SLOW SQL >= %s [ROWS:%d][%s]", gorm.config.SlowThreshold, rows, sql))
-	case gorm.config.LogLevel >= logger.Info:
+	default:
 		l.Debug(msg)
 	}
 }
+
+func (config gormLogConfig) queryLevel(elapsed time.Duration, err error) logger.LogLevel {
+	if config.LogLevel <= logger.Silent {
+		return logger.Silent
+	}
+	switch {
+	case config.shouldLogError(err):
+		return logger.Error
+	case config.shouldLogSlowQuery(elapsed):
+		return logger.Warn
+	case config.LogLevel >= logger.Info:
+		return logger.Info
+	default:
+		return logger.Silent
+	}
+}
+
+func (config gormLogConfig) shouldLogError(err error) bool {
+	if err == nil || config.LogLevel < logger.Error {
+		return false
+	}
+	return !config.IgnoreRecordNotFoundError || !errors.Is(err, gormdb.ErrRecordNotFound)
+}
+
+func (config gormLogConfig) shouldLogSlowQuery(elapsed time.Duration) bool {
+	return config.SlowThreshold > 0 &&
+		elapsed > config.SlowThreshold &&
+		config.LogLevel >= logger.Warn
+}
+
+type gormSpanLogger struct {
+	logger.Interface
+	tracer Tracer
+	config gormLogConfig
+}
+
+// NewGormSpanLogger marks warning/error spans without changing the existing logger's output.
+func NewGormSpanLogger(tracer Tracer, delegate logger.Interface, config logger.Config) logger.Interface {
+	return &gormSpanLogger{
+		Interface: delegate,
+		tracer:    tracer,
+		config:    gormLogConfig{Config: config},
+	}
+}
+
+func (gorm *gormSpanLogger) LogMode(level logger.LogLevel) logger.Interface {
+	clone := *gorm
+	clone.Interface = gorm.Interface.LogMode(level)
+	clone.config.LogLevel = level
+	return &clone
+}
+
+func (gorm *gormSpanLogger) Warn(ctx context.Context, message string, args ...any) {
+	if gorm.config.LogLevel >= logger.Warn {
+		gorm.tracer.Span(ctx).SetAttributes(mustSampleAttribute.Bool(true))
+	}
+	gorm.Interface.Warn(ctx, message, args...)
+}
+
+func (gorm *gormSpanLogger) Error(ctx context.Context, message string, args ...any) {
+	if gorm.config.LogLevel >= logger.Error {
+		gorm.tracer.Span(ctx).SetAttributes(mustSampleAttribute.Bool(true))
+	}
+	gorm.Interface.Error(ctx, message, args...)
+}
+
+func (gorm *gormSpanLogger) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
+	gorm.Interface.Trace(ctx, begin, func() (string, int64) {
+		switch gorm.config.queryLevel(time.Since(begin), err) {
+		case logger.Error, logger.Warn:
+			gorm.tracer.Span(ctx).SetAttributes(mustSampleAttribute.Bool(true))
+		}
+		return fc()
+	}, err)
+}

api/pkg/telemetry/gorm_logger_test.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	stdlog "log"
 	"testing"
 	"time"
 
@@ -40,26 +41,103 @@ func TestGormWarningLogger(t *testing.T) {
 				SlowThreshold: 200 * time.Millisecond,
 				LogLevel:      gormlogger.Warn,
 			}).LogMode(test.logLevel)
+			formatted := false
 			queryLogger.Trace(context.Background(), time.Now().Add(-test.elapsed), func() (string, int64) {
+				formatted = true
 				return "SELECT 1", 1
 			}, test.err)
 
-			if test.want == "" {
-				require.Empty(t, output.String())
-				return
-			}
-			var record struct {
-				Level string `json:"level"`
-				Error string `json:"error"`
-			}
-			require.NoError(t, json.Unmarshal(output.Bytes(), &record))
-			require.Equal(t, test.want, record.Level)
-			require.Contains(t, record.Error, "SELECT 1")
-			if test.err != nil {
-				require.Contains(t, record.Error, test.err.Error())
-			} else {
-				require.Contains(t, record.Error, "SLOW SQL")
-			}
+			require.Equal(t, test.want != "", formatted)
+			assertGormQueryLog(t, output.Bytes(), test.want, test.err)
 		})
 	}
 }
+
+func assertGormQueryLog(t *testing.T, output []byte, wantLevel string, err error) {
+	t.Helper()
+	if wantLevel == "" {
+		require.Empty(t, output)
+		return
+	}
+	var record struct {
+		Level string `json:"level"`
+		Error string `json:"error"`
+	}
+	require.NoError(t, json.Unmarshal(output, &record))
+	require.Equal(t, wantLevel, record.Level)
+	require.Contains(t, record.Error, "SELECT 1")
+	if err != nil {
+		require.Contains(t, record.Error, err.Error())
+		return
+	}
+	require.Contains(t, record.Error, "SLOW SQL")
+}
+
+func TestGormSpanLoggerPreservesDelegate(t *testing.T) {
+	tests := []struct {
+		name    string
+		elapsed time.Duration
+		err     error
+		level   gormlogger.LogLevel
+		message string
+		spans   []string
+	}{
+		{name: "fast query", elapsed: time.Millisecond, level: gormlogger.Warn},
+		{name: "slow query", elapsed: time.Second, level: gormlogger.Warn, message: "SLOW SQL", spans: []string{"test"}},
+		{name: "query error", err: errors.New("original query error"), level: gormlogger.Warn, message: "original query error", spans: []string{"test"}},
+		{name: "info query", elapsed: time.Millisecond, level: gormlogger.Info, message: "SELECT 1"},
+		{name: "silent error", err: errors.New("query error"), level: gormlogger.Silent},
+	}
+	for _, spec := range tests {
+		t.Run(spec.name, func(t *testing.T) {
+			test := newSamplingTestSpan(t, false)
+			queryLogger, output := newDelegatedGormLogger(test)
+			formatted := false
+			queryLogger.LogMode(spec.level).Trace(test.ctx, time.Now().Add(-spec.elapsed), func() (string, int64) {
+				formatted = true
+				return "SELECT 1", 1
+			}, spec.err)
+
+			test.finish(t, spec.spans...)
+			require.Equal(t, spec.message != "", formatted)
+			assertDelegatedGormOutput(t, output.String(), spec.message)
+		})
+	}
+}
+
+func TestGormSpanLoggerDirectWarningsAndErrors(t *testing.T) {
+	logs := map[string]func(gormlogger.Interface, context.Context){
+		"warning": func(logger gormlogger.Interface, ctx context.Context) {
+			logger.Warn(ctx, "original %s", "message")
+		},
+		"error": func(logger gormlogger.Interface, ctx context.Context) {
+			logger.Error(ctx, "original %s", "message")
+		},
+	}
+	for name, log := range logs {
+		t.Run(name, func(t *testing.T) {
+			test := newSamplingTestSpan(t, false)
+			logger, output := newDelegatedGormLogger(test)
+			log(logger, test.ctx)
+			test.finish(t, "test")
+			require.Contains(t, output.String(), "original message")
+		})
+	}
+}
+
+func newDelegatedGormLogger(test *samplingTestSpan) (gormlogger.Interface, *bytes.Buffer) {
+	output := &bytes.Buffer{}
+	config := gormlogger.Config{SlowThreshold: 200 * time.Millisecond, LogLevel: gormlogger.Warn}
+	delegate := gormlogger.New(stdlog.New(output, "", 0), config)
+	return NewGormSpanLogger(test.tracer, delegate, config), output
+}
+
+func assertDelegatedGormOutput(t *testing.T, output string, message string) {
+	t.Helper()
+	if message == "" {
+		require.Empty(t, output)
+		return
+	}
+	require.Contains(t, output, message)
+	require.Contains(t, output, "SELECT 1")
+}

api/pkg/telemetry/sampling_span_processor.go

@@ -6,6 +6,8 @@ import (
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 )
 
+const mustSampleAttribute attribute.Key = "httpsms.must_sample"
+
 type samplingSpanProcessor struct {
 	sdktrace.SpanProcessor
 	sampler sdktrace.Sampler
@@ -31,28 +33,41 @@ func (processor *samplingSpanProcessor) OnEnd(span sdktrace.ReadOnlySpan) {
 
 func shouldRetainSpan(span sdktrace.ReadOnlySpan) bool {
 	// Drop counters survive buffer limits, unlike severity attributes and events.
-	if span.DroppedEvents() > 0 || span.DroppedAttributes() > 0 {
-		return true
-	}
-	if span.Status().Code == codes.Error || hasWarningOrErrorSeverity(span.Attributes()) {
-		return true
-	}
-	for _, event := range span.Events() {
-		if event.Name == "exception" || hasWarningOrErrorSeverity(event.Attributes) {
+	return span.DroppedEvents() > 0 ||
+		span.DroppedAttributes() > 0 ||
+		span.Status().Code == codes.Error ||
+		hasRetentionAttributes(span.Attributes()) ||
+		hasExceptionalEvents(span.Events())
+}
+
+func hasExceptionalEvents(events []sdktrace.Event) bool {
+	for _, event := range events {
+		if event.Name == "exception" || event.DroppedAttributeCount > 0 || hasRetentionAttributes(event.Attributes) {
 			return true
 		}
 	}
 	return false
 }
 
-func hasWarningOrErrorSeverity(attributes []attribute.KeyValue) bool {
+func hasRetentionAttributes(attributes []attribute.KeyValue) bool {
 	for _, value := range attributes {
-		if value.Key == "log.severity" {
-			switch value.Value.AsString() {
-			case "warn", "warning", "error", "fatal":
-				return true
-			}
+		if isRetentionAttribute(value) {
+			return true
 		}
 	}
 	return false
 }
+
+func isRetentionAttribute(value attribute.KeyValue) bool {
+	if value.Key == mustSampleAttribute {
+		return value.Value.AsBool()
+	}
+	if value.Key != "log.severity" {
+		return false
+	}
+	switch value.Value.AsString() {
+	case "warn", "warning", "error", "fatal":
+		return true
+	}
+	return false
+}