From 05e96ba1c4fa332cd7ec3b7989503c44aedc1142 Mon Sep 17 00:00:00 2001 From: Robert Lin Date: Wed, 7 Aug 2024 16:14:02 -0700 Subject: [PATCH] fix/enterpriseportal/importer: disable query tracing when importing (#64352) This can generate 10k+ spans in production Also added simple handling for `pgx.NamedArgs` here, since I notice those are missing ## Test plan CI --- .../internal/database/importer/BUILD.bazel | 1 + .../internal/database/importer/importer.go | 7 ++- .../cloudsql/tracer.go | 46 ++++++++++++++----- 3 files changed, 41 insertions(+), 13 deletions(-) diff --git a/cmd/enterprise-portal/internal/database/importer/BUILD.bazel b/cmd/enterprise-portal/internal/database/importer/BUILD.bazel index 7b4f059886e..0bc0fc3864f 100644 --- a/cmd/enterprise-portal/internal/database/importer/BUILD.bazel +++ b/cmd/enterprise-portal/internal/database/importer/BUILD.bazel @@ -20,6 +20,7 @@ go_library( "//lib/background", "//lib/enterpriseportal/subscriptions/v1:subscriptions", "//lib/errors", + "//lib/managedservicesplatform/cloudsql", "//lib/pointers", "@com_github_sourcegraph_conc//pool", "@com_github_sourcegraph_log//:log", diff --git a/cmd/enterprise-portal/internal/database/importer/importer.go b/cmd/enterprise-portal/internal/database/importer/importer.go index 6c74f4d2342..90171d80dcd 100644 --- a/cmd/enterprise-portal/internal/database/importer/importer.go +++ b/cmd/enterprise-portal/internal/database/importer/importer.go @@ -25,6 +25,7 @@ import ( "github.com/sourcegraph/sourcegraph/lib/background" subscriptionsv1 "github.com/sourcegraph/sourcegraph/lib/enterpriseportal/subscriptions/v1" "github.com/sourcegraph/sourcegraph/lib/errors" + "github.com/sourcegraph/sourcegraph/lib/managedservicesplatform/cloudsql" "github.com/sourcegraph/sourcegraph/lib/pointers" ) @@ -110,7 +111,9 @@ func (i *importer) Handle(ctx context.Context) (err error) { } }() } - return i.ImportSubscriptions(ctx) + // Disable tracing on database interactions, because we could generate + // traces with 10k+ spans in production. + return i.ImportSubscriptions(cloudsql.WithoutTrace(ctx)) } func (i *importer) ImportSubscriptions(ctx context.Context) error { @@ -276,7 +279,7 @@ func (i *importer) importSubscription(ctx context.Context, dotcomSub *dotcomdb.S } func (i *importer) importLicense(ctx context.Context, subscriptionID string, dotcomLicense *dotcomdb.LicenseAttributes) (err error) { - tr, ctx := trace.New(ctx, "importSubscription", + tr, ctx := trace.New(ctx, "importLicense", attribute.String("dotcomSub.ID", subscriptionID), attribute.String("dotcomLicense.ID", dotcomLicense.ID)) defer tr.EndWithErr(&err) diff --git a/lib/managedservicesplatform/cloudsql/tracer.go b/lib/managedservicesplatform/cloudsql/tracer.go index 01eec29cc07..ec96ab243c2 100644 --- a/lib/managedservicesplatform/cloudsql/tracer.go +++ b/lib/managedservicesplatform/cloudsql/tracer.go @@ -14,6 +14,20 @@ import ( "go.opentelemetry.io/otel/trace" ) +type contextKey int + +const contextKeyWithoutTrace contextKey = iota + +// WithoutTrace disables CloudSQL connection tracing for child contexts. +func WithoutTrace(ctx context.Context) context.Context { + return context.WithValue(ctx, contextKeyWithoutTrace, true) +} + +func shouldNotTrace(ctx context.Context) bool { + withoutTrace, ok := ctx.Value(contextKeyWithoutTrace).(bool) + return ok && withoutTrace +} + var tracer = otel.GetTracerProvider().Tracer("msp/cloudsql/pgx") type pgxTracer struct{} @@ -31,6 +45,10 @@ var ( // TraceQueryStart is called at the beginning of Query, QueryRow, and Exec calls. The returned context is used for the // rest of the call and will be passed to TraceQueryEnd. func (pgxTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { + if shouldNotTrace(ctx) { + return ctx // do nothing + } + ctx, _ = tracer.Start(ctx, "pgx.Query", trace.WithAttributes( attribute.String("query", data.SQL), @@ -40,7 +58,11 @@ func (pgxTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.Trac return ctx } -func (pgxTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) { +func (pgxTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) { + if shouldNotTrace(ctx) { + return // do nothing + } + span := trace.SpanFromContext(ctx) defer span.End() @@ -48,16 +70,7 @@ func (pgxTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.Tra attribute.String("command_tag", data.CommandTag.String()), attribute.Int64("rows_affected", data.CommandTag.RowsAffected()), ) - switch { - case data.CommandTag.Insert(): - span.SetName("pgx.Query: INSERT") - case data.CommandTag.Update(): - span.SetName("pgx.Query: UPDATE") - case data.CommandTag.Delete(): - span.SetName("pgx.Query: DELETE") - case data.CommandTag.Select(): - span.SetName("pgx.Query: SELECT") - } + span.SetName("pgx.Query: " + data.CommandTag.String()) if data.Err != nil { span.SetStatus(codes.Error, data.Err.Error()) @@ -65,6 +78,10 @@ func (pgxTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.Tra } func (pgxTracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectStartData) context.Context { + if shouldNotTrace(ctx) { + return ctx // do nothing + } + ctx, _ = tracer.Start(ctx, "pgx.Connect", trace.WithAttributes( attribute.String("database", data.ConnConfig.Database), attribute.String("instance", fmt.Sprintf("%s:%d", data.ConnConfig.Host, data.ConnConfig.Port)), @@ -73,6 +90,10 @@ func (pgxTracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectSta } func (pgxTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) { + if shouldNotTrace(ctx) { + return // do nothing + } + span := trace.SpanFromContext(ctx) defer span.End() @@ -121,6 +142,9 @@ func argsAsAttributes(args []any) []attribute.KeyValue { case time.Time: attrs[i] = attribute.String(key, v.String()) + case pgx.NamedArgs: + attrs[i] = attribute.String(key, truncateStringValue(fmt.Sprintf("%+v", v))) + default: // in case we miss anything attrs[i] = attribute.String(key, fmt.Sprintf("unhandled type %T", v)) }