fix: use authenticated urls for pubsub #14261
Conversation
39d9b47 to
6beb8a0
Compare
119b440 to
85f6cfe
Compare
coderd/database/pubsub/pubsub.go
Outdated
| } | ||
|
|
||
| // Set the dialer if the connector supports it. | ||
| if dc, ok := connector.(database.DialerConnector); ok { |
There was a problem hiding this comment.
We need to drop a CRITICAL log if this type assertion fails, since it means we've introduced a new driver that doesn't support a Dialer and our logs will be incomplete.
| // Create a custom connector if the database driver supports it. | ||
| connectorCreator, ok := p.db.Driver().(database.ConnectorCreator) | ||
| if ok { | ||
| connector, err = connectorCreator.Connector(connectURL) |
There was a problem hiding this comment.
This code doesn't get hit in the package unit tests. A good way to test it would be to create a pq Driver wrapper that we can control.
I'd like to see a test where we start pubsub with the wrapped driver, do some pub'ing and sub'ing, then kill the connection and verify that the pubsub / pq.Listener reconnects automatically. That would give a nice test of the pq changes you made as well.
| } | ||
|
|
||
| // DialerConnector can create a driver.Connector and set a pq.Dialer. | ||
| type DialerConnector interface { |
coderd/database/dbtestutil/driver.go
Outdated
| name string | ||
| inner driver.Driver | ||
| connections []driver.Conn | ||
| listeners map[chan struct{}]chan struct{} |
There was a problem hiding this comment.
This connections and listeners business is not threadsafe, and so can cause races in our tests.
A better design would be to just have a chan driver.Conn that we pass each new connection thru. Tests will have to know how many connections they expect, and can read from the channel. For the PubSub, it's the listener we want to test, which should use a connection when it first connects, then another one when it reconnects.
That way you don't need methods like AddConnection, WaitForConnection, or DropConnections -- the test code reads from the channel to wait for the connection, and can directly close the connection when it wants to interrupt it.
There is some complexity around the sql.DB, which we use for publishing and has a pool of connections. I suggest you sidestep that complexity by just using a second PubSub for publishing with a regular pq driver.
coderd/database/dbtestutil/driver.go
Outdated
| } | ||
|
|
||
| func Register() (*Driver, error) { | ||
| db, err := sql.Open("postgres", "") |
There was a problem hiding this comment.
This is a roundabout way to get a pq driver, which is just pq.Driver{}
coderd/database/dbtestutil/driver.go
Outdated
| listeners: make(map[chan struct{}]chan struct{}), | ||
| } | ||
|
|
||
| sql.Register(d.name, d) |
There was a problem hiding this comment.
It seems like you're doing this so that you can get a sql.DB to pass to the PubSub with this instrumented driver. But, registering it and using sql.Open() with the name is more complex than it needs to be.
If you allow test code to directly instantiate the instrumented connector, then you can get a sql.DB as:
db := sql.ConnectDB(connector)
That avoids registering the driver and worrying about a unique name.
| case <-gotChan: | ||
| case <-ctx.Done(): | ||
| t.Fatal("timeout waiting for message") | ||
| } |
There was a problem hiding this comment.
This construction comes up a lot, so we have
_ = testutil.RequireRecvContext(ctx, t, gotChan)
for this purpose.
| case <-reconnectChan: | ||
| case <-ctx.Done(): | ||
| t.Fatal("timeout waiting for reconnect") | ||
| } |
There was a problem hiding this comment.
_ = testutil.RequireRecvChan(ctx, t, reconnectChan)
| case <-gotChan: | ||
| case <-ctx.Done(): | ||
| t.Fatal("timeout waiting for message after reconnect") | ||
| } |
There was a problem hiding this comment.
_ = testutil.RequireRecvCtx(ctx, t, gotChan)
| @@ -119,3 +120,79 @@ func TestPGPubsub_Metrics(t *testing.T) { | |||
| !testutil.PromCounterGathered(t, metrics, "coder_pubsub_latency_measure_errs_total") | |||
| }, testutil.WaitShort, testutil.IntervalFast) | |||
| } | |||
|
|
|||
| func TestPGPubsubDriver(t *testing.T) { | |||
There was a problem hiding this comment.
This is a good test, but the design of the instrumented driver is racy -- see comments on dbtestutil/driver.go for how to address.
|
@spikecurtis Thanks for the feedback on the test structure, I've updated it with your suggestions and I think things are looking good. |
coderd/database/dbtestutil/driver.go
Outdated
| return conn, nil | ||
| } | ||
|
|
||
| conn, err := c.driver.Open(c.name) |
There was a problem hiding this comment.
This is circular. Driver.Open() creates a new Connector and calls into it. Fortunately, the pubsub creates a dialer, but anyone else who uses this might not and will overflow their stack.
You need to use pq.Driver{} directly here.
Fix for https://github.com/coder/customers/issues/639
The
pq.Listenerwe use for PubSub does not expose a way to hook into the lifecycle of reconnecting and that prevents us from supporting rotating passwords on reconnection. Previously, we've used a forkcoder/pqthat exposes a way to create listeners with a given connector. The new v2 branch ofcoder/pqis a copy of the1.10.9tag fromlib/pqwith our patch on top to ensure no other package changes are brought in besides ours.This now allows us to create a
*sql.DBthat usesawsiamrdsunder the hood and allows us to make a new connector that will generate new credentials on PubSub reconnect.I've updated the
awsiamrdstests to now include testing pubsub. I've tested this on an EC2 instance with the correct roles for connecting to a test RDS instance:Pubsub reconnect tests: