Tracing and Testing

[farazfazli]

Hi @jschaf, I saw there was some tracing code being generated via pggen at one point but it was removed from the repo. I am creating a new span before calling generated code and would like the pggen layer to be observable. How would you recommend doing this?

Additionally, can you please explain more about how you perform testing of the generated code? I see on #84 you explained how gomock can be used in source mode but interested in learning about the method you use / recommend for testing pggen-generated code in a production app. Thank you.

[jschaf]

I am creating a new span before calling generated code and would like the pggen layer to be observable. How would you recommend doing this?

Observability is lackluster pgx/v4, but much improved in pgx/v5. Since pggen currently targets pgx/v4, we created our own pool, implementing pgxpool.Pool, to read the pggen query name from context and emit spans. The main difficulty is that the Pool methods returns pgx.Rows which delay work until scanned.

Here's one method from our tracing pool implementation.

// TracingPool is a wrapper around pgxpool.Pool that just records traces.
type TracingPool struct {
	*pgxpool.Pool
}

func NewTracingPool(pool *pgxpool.Pool) *TracingPool {
	return &TracingPool{Pool: pool}
}

func (t *TracingPool) Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error) {
	waiter := newAfterReleaseWaiter()
	ctx, span := trace.OpenSpan(waiter.embedCtx(ctx), "DB Pool Query - "+ExtractPggenQueryName(ctx))
	// Span closed in tracingRows.Close.
	rows, err := t.Pool.Query(ctx, annotateQuery(ctx, sql), args...)
	if err != nil {
		return rows, err
	}
	return newTracingRows(rows, ctx, span, waiter), nil
}

The painful part is tracking the span by implementing pgx.Rows.

If I were to implement this today, I wouldn't try to track the span through Rows.scan. I'd instrument each top-level Pool method with BeginSpan/defer EndSpan and maybe use the pool lifecycle hooks to emit events:

Here's how we did the pool lifecycle hooks.

// AddPoolConfigTracing wraps existing hooks in a pgxpool.Config to add trace
// events for pool lifecycle events.
func AddPoolConfigTracing(cfg *pgxpool.Config) {
	oldBeforeConnect := cfg.BeforeConnect
	oldAfterConnect := cfg.AfterConnect
	oldBeforeAcquire := cfg.BeforeAcquire
	oldAfterRelease := cfg.AfterRelease

	oldLogger := cfg.ConnConfig.Logger
	cfg.ConnConfig.Logger = connEventLogger{next: oldLogger}
	cfg.ConnConfig.LogLevel = pgx.LogLevelInfo

	// Use a pointer to conn to relate BeforeAcquire and AfterRelease since
	// AfterRelease doesn't take a context. Using sync.Map since each go routine
	// is only modifying a single map entry, which sync.Map is specialized for.
	connCtxs := &sync.Map{}

	cfg.BeforeConnect = func(ctx context.Context, config *pgx.ConnConfig) error {
		if oldBeforeConnect != nil {
			if err := oldBeforeConnect(ctx, config); err != nil {
				return err
			}
		}
		trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: BeforeConnect"))
		return nil
	}

	cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
		if oldAfterConnect != nil {
			if err := oldAfterConnect(ctx, conn); err != nil {
				return err
			}
		}
		trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: AfterConnect"))
		return nil
	}

	cfg.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
		if oldBeforeAcquire != nil {
			if shouldContinue := oldBeforeAcquire(ctx, conn); !shouldContinue {
				return shouldContinue
			}
		}
		trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: BeforeAcquire"))
		connCtxs.Store(conn, ctx)
		return true
	}

	cfg.AfterRelease = func(conn *pgx.Conn) bool {
		if oldAfterRelease != nil {
			if shouldContinue := oldAfterRelease(conn); !shouldContinue {
				return shouldContinue
			}
		}
		ctxVal, ok := connCtxs.LoadAndDelete(conn)
		if !ok {
			return true
		}
		ctx := ctxVal.(context.Context)
		trace.SpanFromContext(ctx).AddEvent(trace.NewEvent("PgPool: AfterRelease"))
		if w := extractAfterReleaseWaiterContext(ctx); w != nil {
			w.close()
		}

		return true
	}
}

Here's what it looks like:

Additionally, can you please explain more about how you perform testing of the generated code?

We use a real database for all tests. I avoid mocking databases since database interactions are usually the most important part of a backend server and should be covered by tests.

Each test gets a fresh copy of the database. In essence, each test does the following:

• initdb
• Run all migrations
• Run the test

We employ many tricks to keep testing fast:

• Cache the data directory with Bazel so we only run migrations once for all tests.
• Run Postgres with unsafe values, like fsync=off: https://stackoverflow.com/questions/9407442/optimise-postgresql-for-fast-testing
• In each test file, run the common database setup code and snapshot the resulting data directory. Then, for each test, copy the snapshotted directory.
• Use CloneFile when available for cheap directory copies: https://www.manpagez.com/man/2/clonefile/.

We started much simpler. You can get most of the speedups by running migrations once and caching the directory.

For more Bazel-specific info, see https://news.ycombinator.com/item?id=28082628 and https://news.ycombinator.com/item?id=28470034.