Skip to content

Commit a20d318

Browse files
coadlercoadler
committed
feat: add single tailnet support to pgcoord
1 parent ccda1c5 commit a20d318

20 files changed

+671
-258
lines changed

‎coderd/coderd.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -406,7 +406,7 @@ func New(options *Options) *API {
406406
api.DERPMap,
407407
options.DeploymentValues.DERP.Config.ForceWebSockets.Value(),
408408
func(context.Context) (tailnet.MultiAgentConn, error) {
409-
return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New()), nil
409+
return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New())
410410
},
411411
wsconncache.New(api._dialWorkspaceAgentTailnet, 0),
412412
api.TracerProvider,

‎coderd/database/dbtestutil/db.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010

1111
"github.com/coder/coder/v2/coderd/database"
1212
"github.com/coder/coder/v2/coderd/database/dbfake"
13+
"github.com/coder/coder/v2/coderd/database/migrations"
1314
"github.com/coder/coder/v2/coderd/database/postgres"
1415
"github.com/coder/coder/v2/coderd/database/pubsub"
1516
)
@@ -42,6 +43,9 @@ func NewDB(t testing.TB) (database.Store, pubsub.Pubsub) {
4243
})
4344
db = database.New(sqlDB)
4445

46+
err = migrations.Up(sqlDB)
47+
require.NoError(t, err)
48+
4549
ps, err = pubsub.New(context.Background(), sqlDB, connectionURL)
4650
require.NoError(t, err)
4751
t.Cleanup(func() {

‎coderd/database/dump.sql

Lines changed: 8 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000151_pg_coordinator_single_tailnet.down.sql

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
BEGIN;
2+
3+
-- ALTER TABLE
4+
-- tailnet_clients
5+
-- ADD COLUMN
6+
-- agent_id uuid;
7+
8+
UPDATE
9+
tailnet_clients
10+
SET
11+
-- grab just the first agent_id, or default to an empty UUID.
12+
agent_id = COALESCE(agent_ids[0], '00000000-0000-0000-0000-000000000000'::uuid);
13+
14+
ALTER TABLE
15+
tailnet_clients
16+
DROP COLUMN
17+
agent_ids;
18+
19+
COMMIT;

‎coderd/database/migrations/000151_pg_coordinator_single_tailnet.up.sql

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
BEGIN;
2+
3+
ALTER TABLE
4+
tailnet_clients
5+
ADD COLUMN
6+
agent_ids uuid[];
7+
8+
UPDATE
9+
tailnet_clients
10+
SET
11+
agent_ids = ARRAY[agent_id]::uuid[];
12+
13+
ALTER TABLE
14+
tailnet_clients
15+
ALTER COLUMN
16+
agent_ids SET NOT NULL;
17+
18+
19+
CREATE INDEX idx_tailnet_clients_agent_ids ON tailnet_clients USING GIN (agent_ids);
20+
21+
CREATE OR REPLACE FUNCTION tailnet_notify_client_change() RETURNS trigger
22+
LANGUAGE plpgsql
23+
AS $$
24+
BEGIN
25+
-- check new first to get the updated agent ids.
26+
IF (NEW IS NOT NULL) THEN
27+
PERFORM pg_notify('tailnet_client_update', NEW.id || ',' || array_to_string(NEW.agent_ids, ','));
28+
RETURN NULL;
29+
END IF;
30+
IF (OLD IS NOT NULL) THEN
31+
PERFORM pg_notify('tailnet_client_update', OLD.id || ',' || array_to_string(OLD.agent_ids, ','));
32+
RETURN NULL;
33+
END IF;
34+
END;
35+
$$;
36+
37+
ALTER TABLE
38+
tailnet_clients
39+
DROP COLUMN
40+
agent_id;
41+
42+
COMMIT;

‎coderd/database/models.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

Lines changed: 13 additions & 14 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/tailnet.sql

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ INSERT INTO
33
tailnet_clients (
44
id,
55
coordinator_id,
6-
agent_id,
6+
agent_ids,
77
node,
88
updated_at
99
)
@@ -13,7 +13,7 @@ ON CONFLICT (id, coordinator_id)
1313
DO UPDATE SET
1414
id = $1,
1515
coordinator_id = $2,
16-
agent_id = $3,
16+
agent_ids = $3,
1717
node = $4,
1818
updated_at = now() at time zone 'utc'
1919
RETURNING *;
@@ -66,12 +66,11 @@ FROM tailnet_agents;
6666
-- name: GetTailnetClientsForAgent :many
6767
SELECT *
6868
FROM tailnet_clients
69-
WHERE agent_id = $1;
69+
WHERE $1::uuid = ANY(agent_ids);
7070

7171
-- name: GetAllTailnetClients :many
7272
SELECT *
73-
FROM tailnet_clients
74-
ORDER BY agent_id;
73+
FROM tailnet_clients;
7574

7675
-- name: UpsertTailnetCoordinator :one
7776
INSERT INTO

‎coderd/tailnet_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,7 @@ func setupAgent(t *testing.T, agentAddresses []netip.Prefix) (uuid.UUID, agent.A
233233
derpServer,
234234
func() *tailcfg.DERPMap { return manifest.DERPMap },
235235
false,
236-
func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()), nil },
236+
func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()) },
237237
cache,
238238
trace.NewNoopTracerProvider(),
239239
)

‎enterprise/coderd/workspaceproxycoordinate.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ func (api *API) workspaceProxyCoordinate(rw http.ResponseWriter, r *http.Request
6767
}
6868

6969
id := uuid.New()
70-
sub := (*api.AGPL.TailnetCoordinator.Load()).ServeMultiAgent(id)
70+
sub, err := (*api.AGPL.TailnetCoordinator.Load()).ServeMultiAgent(id)
7171
ctx, nc := websocketNetConn(ctx, conn, websocket.MessageText)
7272
defer nc.Close()
7373

‎enterprise/coderd/workspaceproxycoordinator_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ func Test_agentIsLegacy(t *testing.T) {
5959
defer cancel()
6060

6161
nodeID := uuid.New()
62-
ma := coordinator.ServeMultiAgent(nodeID)
62+
ma, err := coordinator.ServeMultiAgent(nodeID)
63+
require.NoError(t, err)
6364
defer ma.Close()
6465
require.NoError(t, ma.UpdateSelf(&agpl.Node{
6566
ID: 55,
@@ -123,7 +124,8 @@ func Test_agentIsLegacy(t *testing.T) {
123124
defer cancel()
124125

125126
nodeID := uuid.New()
126-
ma := coordinator.ServeMultiAgent(nodeID)
127+
ma, err := coordinator.ServeMultiAgent(nodeID)
128+
require.NoError(t, err)
127129
defer ma.Close()
128130
require.NoError(t, ma.UpdateSelf(&agpl.Node{
129131
ID: 55,

0 commit comments

Comments
 (0)