Handle custom types with differing OIDs across shards

[sgrif]

Handle custom types with differing OIDs across shards

At the wire protocol level, type information is entirely handled by OID, not type name. An OID is a unique identifier that PG assigns for virtually every schema related entity. They are extremely dependent on the order that plugins are loaded, DDL is run, etc. Because of this it's quite likely that a user or extension defined type will not have the same OID across multiple shards.

To solve this, we need to do three things:

• Ensure clients receive a single canonical answer when they query for
  type information
• Rewrite any messages being sent to the server to ensure they have the
  correct OIDs for a given type for that sever
• Rewrite any messages being sent to the client to ensure they have the
  canonical OIDs for a given type

This commit's solution to ensuring the client gets canoincal type information is straightforward, but not ideal. We look for any select statements that could be loading type information (referencing the relevant tables in pg_catalog, or doing certain casts), and route them to shard 0.

This means that we have a single point of failure for schema information. Ideally we would be responding to these queries ourselves rather than sending them to the server, but the breadth of different ways different client libraries will load this type information would mean we essentially have to implement our own shitty SQL server. It's possible, and probably necessary in the long run, but outside of the scope of this specific fix.

It's not clear that there'd be any reason one shard would be better than another for the canonical information, but if there is one, we could probably surface this as a configuration option. But since there's no clear reason for it, we can just do the reasonable default.

For rewriting messages, we need to keep track of the mappings for each shard on our end. This was actually somewhat tricky to make work, as schema loading happens on Shard, but the message rewriting needs to happen all the way down in Server, and these are not types that share an API bounary.

To rectify this, we use a new primitive, which is the async equivalent of std::sync::OnceLock. It ensures that only one writer will ever attempt to write to it, while also allowing waiters to wait for a value to be set without attempting to initialize themselves.

With this primitive, all we need to do is stick it in an Arc to have it shared in as many disparate places as we need. We use the SchemaCache that was previously introduced to ensure we only have one copy of this mapping per database, regardless of how many user/clusters we have, but its use is relatively minimal.

Since this sits in a very hot path, I've done as much as I could to avoid excess allocations or loops if they're not strictly necessary. I'm assuming that the majority of users won't have custom types that vary between shards, so this change should ideally be free for those users, short of the query to determine that there are no mappings at startup.

Tests were a little bit trickier. Resolving the OIDs happens at cluster launch, and the code that needs to do the rewriting will never run without the cluster launching and loading first. But many of our tests were constructing clusters without launching them.

I worked around this by making Oids::default return an instance that is already resolved to an empty set of mappings. Since the normal construction requires canonical type information to be passed in, we don't need to worry about production code accidentally calling this incorrectly.

[sgrif]

This is now a working implementation. It needs some cleanup and backfill of unit tests, but should be able to make it into tomorrow's release

[sgrif]

Prior art: tokio-rs/tokio#4788 https://doc.rust-lang.org/stable/std/sync/struct.OnceLock.html

(Likely going to try to upstream this to tokio since it was previously rejected as a new feature on OnceCell, but it was mentioned it may be accepted as a new type)

[codecov]

Codecov Report

❌ Patch coverage is 78.57143% with 102 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/backend/pool/shard/oids.rs	44.44%	45 Missing ⚠️
pgdog/src/backend/pool/cluster/schema_loader.rs	18.18%	18 Missing ⚠️
pgdog/src/backend/prepared_statements.rs	85.33%	11 Missing ⚠️
pgdog/src/sync.rs	80.76%	10 Missing ⚠️
pgdog/src/net/messages/row_description.rs	0.00%	9 Missing ⚠️
pgdog/src/net/messages/parameter_description.rs	0.00%	6 Missing ⚠️
pgdog/src/net/messages/mod.rs	0.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

[sgrif]

Between making sure we have adequate test coverage for RowDescription on the simple protocol, and that python test failure, this needs a bit more work.

[levkk]

LGTM!

Might be worth considering using fnv (already in Cargo.toml) instead of hashbrown for mappings. It's faster and also we don't need dos protection since oids are database-generated.

Bonus points: include the code our friends gave us as repro into our Ruby test suite to make sure it works as expected and doesn't break in the future.