feat: add single tailnet support to pgcoord

coder · coadler · Sep 21, 2023 · Aug 25, 2023 · Aug 25, 2023 · Aug 25, 2023
commit a20d318c7358936d18698e255dab317e4583c3cb
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -406,7 +406,7 @@ func New(options *Options) *API {
 			api.DERPMap,
 			options.DeploymentValues.DERP.Config.ForceWebSockets.Value(),
 			func(context.Context) (tailnet.MultiAgentConn, error) {
-				return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New()), nil
+				return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New())
 			},
 			wsconncache.New(api._dialWorkspaceAgentTailnet, 0),
 			api.TracerProvider,

diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/coderd/database/migrations"
 	"github.com/coder/coder/v2/coderd/database/postgres"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 )
@@ -42,6 +43,9 @@ func NewDB(t testing.TB) (database.Store, pubsub.Pubsub) {
 		})
 		db = database.New(sqlDB)
 
+		err = migrations.Up(sqlDB)
+		require.NoError(t, err)
+
 		ps, err = pubsub.New(context.Background(), sqlDB, connectionURL)
 		require.NoError(t, err)
 		t.Cleanup(func() {

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000151_pg_coordinator_single_tailnet.down.sql b/coderd/database/migrations/000151_pg_coordinator_single_tailnet.down.sql
@@ -0,0 +1,19 @@
+BEGIN;
+
+-- ALTER TABLE
+-- 	tailnet_clients
+-- ADD COLUMN
+-- 	agent_id uuid;
+
+UPDATE
+	tailnet_clients
+SET
+	-- grab just the first agent_id, or default to an empty UUID.
+	agent_id = COALESCE(agent_ids[0], '00000000-0000-0000-0000-000000000000'::uuid);
+
+ALTER TABLE
+	tailnet_clients
+DROP COLUMN
+	agent_ids;
+
+COMMIT;
diff --git a/coderd/database/migrations/000151_pg_coordinator_single_tailnet.up.sql b/coderd/database/migrations/000151_pg_coordinator_single_tailnet.up.sql
@@ -0,0 +1,42 @@
+BEGIN;
+
+ALTER TABLE
+	tailnet_clients
+ADD COLUMN
+	agent_ids uuid[];
+
+UPDATE
+	tailnet_clients
+SET
+	agent_ids = ARRAY[agent_id]::uuid[];
+
+ALTER TABLE
+	tailnet_clients
+ALTER COLUMN
+	agent_ids SET NOT NULL;
+
+
+CREATE INDEX idx_tailnet_clients_agent_ids ON tailnet_clients USING GIN (agent_ids);
+
+CREATE OR REPLACE FUNCTION tailnet_notify_client_change() RETURNS trigger
+	LANGUAGE plpgsql
+	AS $$
+BEGIN
+	-- check new first to get the updated agent ids.
+	IF (NEW IS NOT NULL) THEN
+		PERFORM pg_notify('tailnet_client_update', NEW.id || ',' || array_to_string(NEW.agent_ids, ','));
+		RETURN NULL;
+	END IF;
+	IF (OLD IS NOT NULL) THEN
+		PERFORM pg_notify('tailnet_client_update', OLD.id || ',' || array_to_string(OLD.agent_ids, ','));
+		RETURN NULL;
+	END IF;
+END;
+$$;
+
+ALTER TABLE
+	tailnet_clients
+DROP COLUMN
+	agent_id;
+
+COMMIT;
diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/tailnet.sql b/coderd/database/queries/tailnet.sql
@@ -3,7 +3,7 @@ INSERT INTO
 	tailnet_clients (
 	id,
 	coordinator_id,
-	agent_id,
+	agent_ids,
 	node,
 	updated_at
 )
@@ -13,7 +13,7 @@ ON CONFLICT (id, coordinator_id)
 DO UPDATE SET
 	id = $1,
 	coordinator_id = $2,
-	agent_id = $3,
+	agent_ids = $3,
 	node = $4,
 	updated_at = now() at time zone 'utc'
 RETURNING *;
@@ -66,12 +66,11 @@ FROM tailnet_agents;
 -- name: GetTailnetClientsForAgent :many
 SELECT *
 FROM tailnet_clients
-WHERE agent_id = $1;
+WHERE $1::uuid = ANY(agent_ids);
 
 -- name: GetAllTailnetClients :many
 SELECT *
-FROM tailnet_clients
-ORDER BY agent_id;
+FROM tailnet_clients;
 
 -- name: UpsertTailnetCoordinator :one
 INSERT INTO

diff --git a/coderd/tailnet_test.go b/coderd/tailnet_test.go
@@ -233,7 +233,7 @@ func setupAgent(t *testing.T, agentAddresses []netip.Prefix) (uuid.UUID, agent.A
 		derpServer,
 		func() *tailcfg.DERPMap { return manifest.DERPMap },
 		false,
-		func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()), nil },
+		func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()) },
 		cache,
 		trace.NewNoopTracerProvider(),
 	)

diff --git a/enterprise/coderd/workspaceproxycoordinate.go b/enterprise/coderd/workspaceproxycoordinate.go
@@ -67,7 +67,7 @@ func (api *API) workspaceProxyCoordinate(rw http.ResponseWriter, r *http.Request
 	}
 
 	id := uuid.New()
-	sub := (*api.AGPL.TailnetCoordinator.Load()).ServeMultiAgent(id)
+	sub, err := (*api.AGPL.TailnetCoordinator.Load()).ServeMultiAgent(id)
 	ctx, nc := websocketNetConn(ctx, conn, websocket.MessageText)
 	defer nc.Close()
 

diff --git a/enterprise/coderd/workspaceproxycoordinator_test.go b/enterprise/coderd/workspaceproxycoordinator_test.go
@@ -59,7 +59,8 @@ func Test_agentIsLegacy(t *testing.T) {
 		defer cancel()
 
 		nodeID := uuid.New()
-		ma := coordinator.ServeMultiAgent(nodeID)
+		ma, err := coordinator.ServeMultiAgent(nodeID)
+		require.NoError(t, err)
 		defer ma.Close()
 		require.NoError(t, ma.UpdateSelf(&agpl.Node{
 			ID:            55,
@@ -123,7 +124,8 @@ func Test_agentIsLegacy(t *testing.T) {
 		defer cancel()
 
 		nodeID := uuid.New()
-		ma := coordinator.ServeMultiAgent(nodeID)
+		ma, err := coordinator.ServeMultiAgent(nodeID)
+		require.NoError(t, err)
 		defer ma.Close()
 		require.NoError(t, ma.UpdateSelf(&agpl.Node{
 			ID:            55,