feat: add single tailnet support to pgcoord (#9351)

coderd/database/dump.sql

@@ -219,13 +219,57 @@ $$;
 CREATE FUNCTION tailnet_notify_client_change() RETURNS trigger
     LANGUAGE plpgsql
     AS $$
+DECLARE
+	var_client_id uuid;
+	var_coordinator_id uuid;
+	var_agent_ids uuid[];
+	var_agent_id uuid;
 BEGIN
-	IF (OLD IS NOT NULL) THEN
-		PERFORM pg_notify('tailnet_client_update', OLD.id || ',' || OLD.agent_id);
-		RETURN NULL;
+	IF (NEW.id IS NOT NULL) THEN
+		var_client_id = NEW.id;
+		var_coordinator_id = NEW.coordinator_id;
+	ELSIF (OLD.id IS NOT NULL) THEN
+		var_client_id = OLD.id;
+		var_coordinator_id = OLD.coordinator_id;
 	END IF;
+
+	-- Read all agents the client is subscribed to, so we can notify them.
+	SELECT
+		array_agg(agent_id)
+	INTO
+		var_agent_ids
+	FROM
+		tailnet_client_subscriptions subs
+	WHERE
+		subs.client_id = NEW.id AND
+		subs.coordinator_id = NEW.coordinator_id;
+
+	-- No agents to notify
+	if (var_agent_ids IS NULL) THEN
+		return NULL;
+	END IF;
+
+	-- pg_notify is limited to 8k bytes, which is approximately 221 UUIDs.
+	-- Instead of sending all agent ids in a single update, send one for each
+	-- agent id to prevent overflow.
+	FOREACH var_agent_id IN ARRAY var_agent_ids
+	LOOP
+		PERFORM pg_notify('tailnet_client_update', var_client_id || ',' || var_agent_id);
+	END LOOP;
+
+	return NULL;
+END;
+$$;
+
+CREATE FUNCTION tailnet_notify_client_subscription_change() RETURNS trigger
+    LANGUAGE plpgsql
+    AS $$
+BEGIN
 	IF (NEW IS NOT NULL) THEN
-		PERFORM pg_notify('tailnet_client_update', NEW.id || ',' || NEW.agent_id);
+		PERFORM pg_notify('tailnet_client_update', NEW.client_id || ',' || NEW.agent_id);
+		RETURN NULL;
+	ELSIF (OLD IS NOT NULL) THEN
+		PERFORM pg_notify('tailnet_client_update', OLD.client_id || ',' || OLD.agent_id);
 		RETURN NULL;
 	END IF;
 END;
@@ -495,10 +539,16 @@ CREATE TABLE tailnet_agents (
     node jsonb NOT NULL
 );
 
+CREATE TABLE tailnet_client_subscriptions (
+    client_id uuid NOT NULL,
+    coordinator_id uuid NOT NULL,
+    agent_id uuid NOT NULL,
+    updated_at timestamp with time zone NOT NULL
+);
+
 CREATE TABLE tailnet_clients (
     id uuid NOT NULL,
     coordinator_id uuid NOT NULL,
-    agent_id uuid NOT NULL,
     updated_at timestamp with time zone NOT NULL,
     node jsonb NOT NULL
 );
@@ -1144,6 +1194,9 @@ ALTER TABLE ONLY site_configs
 ALTER TABLE ONLY tailnet_agents
     ADD CONSTRAINT tailnet_agents_pkey PRIMARY KEY (id, coordinator_id);
 
+ALTER TABLE ONLY tailnet_client_subscriptions
+    ADD CONSTRAINT tailnet_client_subscriptions_pkey PRIMARY KEY (client_id, coordinator_id, agent_id);
+
 ALTER TABLE ONLY tailnet_clients
     ADD CONSTRAINT tailnet_clients_pkey PRIMARY KEY (id, coordinator_id);
 
@@ -1248,8 +1301,6 @@ CREATE UNIQUE INDEX idx_organization_name_lower ON organizations USING btree (lo
 
 CREATE INDEX idx_tailnet_agents_coordinator ON tailnet_agents USING btree (coordinator_id);
 
-CREATE INDEX idx_tailnet_clients_agent ON tailnet_clients USING btree (agent_id);
-
 CREATE INDEX idx_tailnet_clients_coordinator ON tailnet_clients USING btree (coordinator_id);
 
 CREATE UNIQUE INDEX idx_users_email ON users USING btree (email) WHERE (deleted = false);
@@ -1284,6 +1335,8 @@ CREATE TRIGGER tailnet_notify_agent_change AFTER INSERT OR DELETE OR UPDATE ON t
 
 CREATE TRIGGER tailnet_notify_client_change AFTER INSERT OR DELETE OR UPDATE ON tailnet_clients FOR EACH ROW EXECUTE FUNCTION tailnet_notify_client_change();
 
+CREATE TRIGGER tailnet_notify_client_subscription_change AFTER INSERT OR DELETE OR UPDATE ON tailnet_client_subscriptions FOR EACH ROW EXECUTE FUNCTION tailnet_notify_client_subscription_change();
+
 CREATE TRIGGER tailnet_notify_coordinator_heartbeat AFTER INSERT OR UPDATE ON tailnet_coordinators FOR EACH ROW EXECUTE FUNCTION tailnet_notify_coordinator_heartbeat();
 
 CREATE TRIGGER trigger_insert_apikeys BEFORE INSERT ON api_keys FOR EACH ROW EXECUTE FUNCTION insert_apikey_fail_if_user_deleted();
@@ -1329,6 +1382,9 @@ ALTER TABLE ONLY provisioner_jobs
 ALTER TABLE ONLY tailnet_agents
     ADD CONSTRAINT tailnet_agents_coordinator_id_fkey FOREIGN KEY (coordinator_id) REFERENCES tailnet_coordinators(id) ON DELETE CASCADE;
 
+ALTER TABLE ONLY tailnet_client_subscriptions
+    ADD CONSTRAINT tailnet_client_subscriptions_coordinator_id_fkey FOREIGN KEY (coordinator_id) REFERENCES tailnet_coordinators(id) ON DELETE CASCADE;
+
 ALTER TABLE ONLY tailnet_clients
     ADD CONSTRAINT tailnet_clients_coordinator_id_fkey FOREIGN KEY (coordinator_id) REFERENCES tailnet_coordinators(id) ON DELETE CASCADE;