add subscriber subsystem

coder · coadler · Sep 21, 2023 · Aug 25, 2023 · Aug 25, 2023 · Aug 25, 2023
commit dcb007d87064449e95134dbebb9c7644077ee4cb
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -405,7 +405,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())
+				return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New()), nil
 			},
 			wsconncache.New(api._dialWorkspaceAgentTailnet, 0),
 			api.TracerProvider,

diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -774,13 +774,20 @@ func (q *querier) DeleteTailnetClient(ctx context.Context, arg database.DeleteTa
 	return q.db.DeleteTailnetClient(ctx, arg)
 }
 
-func (q *querier) DeleteTailnetClientSubscription(ctx context.Context, arg database.DeleteTailnetClientSubscriptionParams) (database.DeleteTailnetClientSubscriptionRow, error) {
+func (q *querier) DeleteTailnetClientSubscription(ctx context.Context, arg database.DeleteTailnetClientSubscriptionParams) error {
 	if err := q.authorizeContext(ctx, rbac.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
-		return database.DeleteTailnetClientSubscriptionRow{}, err
+		return err
 	}
 	return q.db.DeleteTailnetClientSubscription(ctx, arg)
 }
 
+func (q *querier) DeleteAllTailnetClientSubscriptions(ctx context.Context, arg database.DeleteAllTailnetClientSubscriptionsParams) error {
+	if err := q.authorizeContext(ctx, rbac.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
+		return err
+	}
+	return q.db.DeleteAllTailnetClientSubscriptions(ctx, arg)
+}
+
 func (q *querier) GetAPIKeyByID(ctx context.Context, id string) (database.APIKey, error) {
 	return fetch(q.log, q.auth, q.db.GetAPIKeyByID)(ctx, id)
 }

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -844,6 +844,15 @@ func (q *FakeQuerier) DeleteAPIKeysByUserID(_ context.Context, userID uuid.UUID)
 	return nil
 }
 
+func (q *FakeQuerier) DeleteAllTailnetClientSubscriptions(ctx context.Context, arg database.DeleteAllTailnetClientSubscriptionsParams) error {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return err
+	}
+
+	panic("not implemented")
+}
+
 func (q *FakeQuerier) DeleteApplicationConnectAPIKeysByUserID(_ context.Context, userID uuid.UUID) error {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
@@ -977,8 +986,8 @@ func (*FakeQuerier) DeleteTailnetClient(context.Context, database.DeleteTailnetC
 	return database.DeleteTailnetClientRow{}, ErrUnimplemented
 }
 
-func (*FakeQuerier) DeleteTailnetClientSubscription(context.Context, database.DeleteTailnetClientSubscriptionParams) (database.DeleteTailnetClientSubscriptionRow, error) {
-	return database.DeleteTailnetClientSubscriptionRow{}, ErrUnimplemented
+func (*FakeQuerier) DeleteTailnetClientSubscription(context.Context, database.DeleteTailnetClientSubscriptionParams) error {
+	return ErrUnimplemented
 }
 
 func (q *FakeQuerier) GetAPIKeyByID(_ context.Context, id string) (database.APIKey, error) {

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go
diff --git a/coderd/database/migrations/000155_pg_coordinator_single_tailnet.up.sql b/coderd/database/migrations/000155_pg_coordinator_single_tailnet.up.sql
@@ -40,47 +40,42 @@ CREATE OR REPLACE FUNCTION tailnet_notify_client_change() RETURNS trigger
 	AS $$
 DECLARE
 	var_client_id uuid;
+	var_coordinator_id uuid;
 	var_agent_ids uuid[];
+	var_agent_id uuid;
 BEGIN
 	IF (NEW.id IS NOT NULL) THEN
 		var_client_id = NEW.id;
-		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;
+		var_coordinator_id = NEW.coordinator_id;
 	ELSIF (OLD.id IS NOT NULL) THEN
-		-- if new is null and old is not null, that means the row was deleted.
-		-- simulate a foreign key by deleting all of the subscriptions.
 		var_client_id = OLD.id;
-		WITH agent_ids AS (
-			DELETE FROM
-				tailnet_client_subscriptions subs
-			WHERE
-				subs.client_id = OLD.id AND
-				subs.coordinator_id = OLD.coordinator_id
-			RETURNING
-				subs.agent_id
-		)
-		SELECT
-			array_agg(agent_id)
-		INTO
-			var_agent_ids
-		FROM
-			agent_ids;
+		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;
 
-	PERFORM pg_notify('tailnet_client_update', var_client_id || ',' || array_to_string(var_agent_ids, ','));
+	-- 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;
 $$;

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
@@ -58,11 +58,15 @@ FROM tailnet_clients
 WHERE id = $1 and coordinator_id = $2
 RETURNING id, coordinator_id;
 
--- name: DeleteTailnetClientSubscription :one
+-- name: DeleteTailnetClientSubscription :exec
 DELETE
 FROM tailnet_client_subscriptions
-WHERE client_id = $1 and agent_id = $2 and coordinator_id = $3
-RETURNING client_id, agent_id, coordinator_id;
+WHERE client_id = $1 and agent_id = $2 and coordinator_id = $3;
+
+-- name: DeleteAllTailnetClientSubscriptions :exec
+DELETE
+FROM tailnet_client_subscriptions
+WHERE client_id = $1 and coordinator_id = $2;
 
 -- name: DeleteTailnetAgent :one
 DELETE

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()) },
+		func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()), nil },
 		cache,
 		trace.NewNoopTracerProvider(),
 	)

diff --git a/enterprise/coderd/workspaceproxycoordinate.go b/enterprise/coderd/workspaceproxycoordinate.go
@@ -67,14 +67,7 @@ func (api *API) workspaceProxyCoordinate(rw http.ResponseWriter, r *http.Request
 	}
 
 	id := uuid.New()
-	sub, err := (*api.AGPL.TailnetCoordinator.Load()).ServeMultiAgent(id)
-	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Failed to serve multi agent.",
-			Detail:  err.Error(),
-		})
-		return
-	}
+	sub := (*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,8 +59,7 @@ func Test_agentIsLegacy(t *testing.T) {
 		defer cancel()
 
 		nodeID := uuid.New()
-		ma, err := coordinator.ServeMultiAgent(nodeID)
-		require.NoError(t, err)
+		ma := coordinator.ServeMultiAgent(nodeID)
 		defer ma.Close()
 		require.NoError(t, ma.UpdateSelf(&agpl.Node{
 			ID:            55,
@@ -124,8 +123,7 @@ func Test_agentIsLegacy(t *testing.T) {
 		defer cancel()
 
 		nodeID := uuid.New()
-		ma, err := coordinator.ServeMultiAgent(nodeID)
-		require.NoError(t, err)
+		ma := coordinator.ServeMultiAgent(nodeID)
 		defer ma.Close()
 		require.NoError(t, ma.UpdateSelf(&agpl.Node{
 			ID:            55,

diff --git a/enterprise/tailnet/coordinator.go b/enterprise/tailnet/coordinator.go
@@ -52,7 +52,7 @@ func NewCoordinator(logger slog.Logger, ps pubsub.Pubsub) (agpl.Coordinator, err
 	return coord, nil
 }
 
-func (c *haCoordinator) ServeMultiAgent(id uuid.UUID) (agpl.MultiAgentConn, error) {
+func (c *haCoordinator) ServeMultiAgent(id uuid.UUID) agpl.MultiAgentConn {
 	m := (&agpl.MultiAgent{
 		ID:                id,
 		AgentIsLegacyFunc: c.agentIsLegacy,
@@ -61,7 +61,7 @@ func (c *haCoordinator) ServeMultiAgent(id uuid.UUID) (agpl.MultiAgentConn, erro
 		OnRemove:          func(enq agpl.Queue) { c.clientDisconnected(enq.UniqueID()) },
 	}).Init()
 	c.addClient(id, m)
-	return m, nil
+	return m
 }
 
 func (c *haCoordinator) addClient(id uuid.UUID, q agpl.Queue) {