From f87d646510524c85e020d84e065585b8b17b69fc Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Sun, 6 Nov 2022 17:30:01 +0000 Subject: [PATCH] feat: Associate connected workspace agents with replicas This will enable displaying a graph that associates agents to running replicas. --- coderd/coderd.go | 7 ++++ coderd/database/dump.sql | 3 +- .../000070_deployment_graph.down.sql | 2 ++ .../migrations/000070_deployment_graph.up.sql | 2 ++ coderd/database/models.go | 3 +- coderd/database/queries.sql.go | 35 ++++++++++++------- coderd/database/queries/workspaceagents.sql | 5 +-- coderd/workspaceagents.go | 4 +++ enterprise/coderd/coderd.go | 1 + enterprise/replicasync/replicasync.go | 11 +++--- 10 files changed, 52 insertions(+), 21 deletions(-) create mode 100644 coderd/database/migrations/000070_deployment_graph.down.sql create mode 100644 coderd/database/migrations/000070_deployment_graph.up.sql diff --git a/coderd/coderd.go b/coderd/coderd.go index 43ae621be144b..ad52edc401cce 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -16,6 +16,7 @@ import ( "github.com/andybalholm/brotli" "github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5/middleware" + "github.com/google/uuid" "github.com/klauspost/compress/zstd" "github.com/prometheus/client_golang/prometheus" "go.opentelemetry.io/otel/trace" @@ -165,6 +166,7 @@ func New(options *Options) *API { r := chi.NewRouter() api := &API{ + ID: uuid.New(), Options: options, RootHandler: r, siteHandler: site.Handler(site.FS(), binFS), @@ -579,6 +581,11 @@ func New(options *Options) *API { type API struct { *Options + // ID is a uniquely generated ID on initialization. + // This is used to associate objects with a specific + // Coder API instance, like workspace agents to a + // specific replica. + ID uuid.UUID Auditor atomic.Pointer[audit.Auditor] WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool] WorkspaceQuotaEnforcer atomic.Pointer[workspacequota.Enforcer] diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 3b2444d690e4c..8eed05162e176 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -390,7 +390,8 @@ CREATE TABLE workspace_agents ( instance_metadata jsonb, resource_metadata jsonb, directory character varying(4096) DEFAULT ''::character varying NOT NULL, - version text DEFAULT ''::text NOT NULL + version text DEFAULT ''::text NOT NULL, + last_connected_replica_id uuid ); COMMENT ON COLUMN workspace_agents.version IS 'Version tracks the version of the currently running workspace agent. Workspace agents register their version upon start.'; diff --git a/coderd/database/migrations/000070_deployment_graph.down.sql b/coderd/database/migrations/000070_deployment_graph.down.sql new file mode 100644 index 0000000000000..dca40cd2be12d --- /dev/null +++ b/coderd/database/migrations/000070_deployment_graph.down.sql @@ -0,0 +1,2 @@ +ALTER TABLE workspace_agents + DROP COLUMN last_connected_replica_id; diff --git a/coderd/database/migrations/000070_deployment_graph.up.sql b/coderd/database/migrations/000070_deployment_graph.up.sql new file mode 100644 index 0000000000000..ecfd9c6d6aa77 --- /dev/null +++ b/coderd/database/migrations/000070_deployment_graph.up.sql @@ -0,0 +1,2 @@ +ALTER TABLE workspace_agents + ADD COLUMN last_connected_replica_id uuid; diff --git a/coderd/database/models.go b/coderd/database/models.go index 832fcd6e18031..4f637ce5b1cec 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -660,7 +660,8 @@ type WorkspaceAgent struct { ResourceMetadata pqtype.NullRawMessage `db:"resource_metadata" json:"resource_metadata"` Directory string `db:"directory" json:"directory"` // Version tracks the version of the currently running workspace agent. Workspace agents register their version upon start. - Version string `db:"version" json:"version"` + Version string `db:"version" json:"version"` + LastConnectedReplicaID uuid.NullUUID `db:"last_connected_replica_id" json:"last_connected_replica_id"` } type WorkspaceApp struct { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 47a5cde7474c4..016110e045b05 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -4471,7 +4471,7 @@ func (q *sqlQuerier) UpdateUserStatus(ctx context.Context, arg UpdateUserStatusP const getWorkspaceAgentByAuthToken = `-- name: GetWorkspaceAgentByAuthToken :one SELECT - id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version + id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version, last_connected_replica_id FROM workspace_agents WHERE @@ -4502,13 +4502,14 @@ func (q *sqlQuerier) GetWorkspaceAgentByAuthToken(ctx context.Context, authToken &i.ResourceMetadata, &i.Directory, &i.Version, + &i.LastConnectedReplicaID, ) return i, err } const getWorkspaceAgentByID = `-- name: GetWorkspaceAgentByID :one SELECT - id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version + id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version, last_connected_replica_id FROM workspace_agents WHERE @@ -4537,13 +4538,14 @@ func (q *sqlQuerier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (W &i.ResourceMetadata, &i.Directory, &i.Version, + &i.LastConnectedReplicaID, ) return i, err } const getWorkspaceAgentByInstanceID = `-- name: GetWorkspaceAgentByInstanceID :one SELECT - id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version + id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version, last_connected_replica_id FROM workspace_agents WHERE @@ -4574,13 +4576,14 @@ func (q *sqlQuerier) GetWorkspaceAgentByInstanceID(ctx context.Context, authInst &i.ResourceMetadata, &i.Directory, &i.Version, + &i.LastConnectedReplicaID, ) return i, err } const getWorkspaceAgentsByResourceIDs = `-- name: GetWorkspaceAgentsByResourceIDs :many SELECT - id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version + id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version, last_connected_replica_id FROM workspace_agents WHERE @@ -4615,6 +4618,7 @@ func (q *sqlQuerier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids [] &i.ResourceMetadata, &i.Directory, &i.Version, + &i.LastConnectedReplicaID, ); err != nil { return nil, err } @@ -4630,7 +4634,7 @@ func (q *sqlQuerier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids [] } const getWorkspaceAgentsCreatedAfter = `-- name: GetWorkspaceAgentsCreatedAfter :many -SELECT id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version FROM workspace_agents WHERE created_at > $1 +SELECT id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version, last_connected_replica_id FROM workspace_agents WHERE created_at > $1 ` func (q *sqlQuerier) GetWorkspaceAgentsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceAgent, error) { @@ -4661,6 +4665,7 @@ func (q *sqlQuerier) GetWorkspaceAgentsCreatedAfter(ctx context.Context, created &i.ResourceMetadata, &i.Directory, &i.Version, + &i.LastConnectedReplicaID, ); err != nil { return nil, err } @@ -4694,7 +4699,7 @@ INSERT INTO resource_metadata ) VALUES - ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) RETURNING id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version + ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) RETURNING id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, startup_script, instance_metadata, resource_metadata, directory, version, last_connected_replica_id ` type InsertWorkspaceAgentParams struct { @@ -4751,6 +4756,7 @@ func (q *sqlQuerier) InsertWorkspaceAgent(ctx context.Context, arg InsertWorkspa &i.ResourceMetadata, &i.Directory, &i.Version, + &i.LastConnectedReplicaID, ) return i, err } @@ -4761,18 +4767,20 @@ UPDATE SET first_connected_at = $2, last_connected_at = $3, - disconnected_at = $4, - updated_at = $5 + last_connected_replica_id = $4, + disconnected_at = $5, + updated_at = $6 WHERE id = $1 ` type UpdateWorkspaceAgentConnectionByIDParams struct { - ID uuid.UUID `db:"id" json:"id"` - FirstConnectedAt sql.NullTime `db:"first_connected_at" json:"first_connected_at"` - LastConnectedAt sql.NullTime `db:"last_connected_at" json:"last_connected_at"` - DisconnectedAt sql.NullTime `db:"disconnected_at" json:"disconnected_at"` - UpdatedAt time.Time `db:"updated_at" json:"updated_at"` + ID uuid.UUID `db:"id" json:"id"` + FirstConnectedAt sql.NullTime `db:"first_connected_at" json:"first_connected_at"` + LastConnectedAt sql.NullTime `db:"last_connected_at" json:"last_connected_at"` + LastConnectedReplicaID uuid.NullUUID `db:"last_connected_replica_id" json:"last_connected_replica_id"` + DisconnectedAt sql.NullTime `db:"disconnected_at" json:"disconnected_at"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` } func (q *sqlQuerier) UpdateWorkspaceAgentConnectionByID(ctx context.Context, arg UpdateWorkspaceAgentConnectionByIDParams) error { @@ -4780,6 +4788,7 @@ func (q *sqlQuerier) UpdateWorkspaceAgentConnectionByID(ctx context.Context, arg arg.ID, arg.FirstConnectedAt, arg.LastConnectedAt, + arg.LastConnectedReplicaID, arg.DisconnectedAt, arg.UpdatedAt, ) diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index 541c14372d707..3c0beb01b8892 100644 --- a/coderd/database/queries/workspaceagents.sql +++ b/coderd/database/queries/workspaceagents.sql @@ -64,8 +64,9 @@ UPDATE SET first_connected_at = $2, last_connected_at = $3, - disconnected_at = $4, - updated_at = $5 + last_connected_replica_id = $4, + disconnected_at = $5, + updated_at = $6 WHERE id = $1; diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index 9e23daa6c0532..f754b7a2c162e 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -522,6 +522,10 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request LastConnectedAt: lastConnectedAt, DisconnectedAt: disconnectedAt, UpdatedAt: database.Now(), + LastConnectedReplicaID: uuid.NullUUID{ + UUID: api.ID, + Valid: true, + }, }) if err != nil { return err diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 69f09c6b9f5d6..adebd113647b0 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -158,6 +158,7 @@ func New(ctx context.Context, options *Options) (*API, error) { } var err error api.replicaManager, err = replicasync.New(ctx, options.Logger, options.Database, options.Pubsub, &replicasync.Options{ + ID: api.AGPL.ID, RelayAddress: options.DERPServerRelayAddress, RegionID: int32(options.DERPServerRegionID), TLSConfig: meshTLSConfig, diff --git a/enterprise/replicasync/replicasync.go b/enterprise/replicasync/replicasync.go index 68344c91f7315..508672cd4dff0 100644 --- a/enterprise/replicasync/replicasync.go +++ b/enterprise/replicasync/replicasync.go @@ -26,6 +26,7 @@ var ( ) type Options struct { + ID uuid.UUID CleanupInterval time.Duration UpdateInterval time.Duration PeerTimeout time.Duration @@ -40,6 +41,9 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, pubsub data if options == nil { options = &Options{} } + if options.ID == uuid.Nil { + options.ID = uuid.New() + } if options.PeerTimeout == 0 { options.PeerTimeout = 3 * time.Second } @@ -59,9 +63,8 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, pubsub data if err != nil { return nil, xerrors.Errorf("ping database: %w", err) } - id := uuid.New() replica, err := db.InsertReplica(ctx, database.InsertReplicaParams{ - ID: id, + ID: options.ID, CreatedAt: database.Now(), StartedAt: database.Now(), UpdatedAt: database.Now(), @@ -74,13 +77,13 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, pubsub data if err != nil { return nil, xerrors.Errorf("insert replica: %w", err) } - err = pubsub.Publish(PubsubEvent, []byte(id.String())) + err = pubsub.Publish(PubsubEvent, []byte(options.ID.String())) if err != nil { return nil, xerrors.Errorf("publish new replica: %w", err) } ctx, cancelFunc := context.WithCancel(ctx) manager := &Manager{ - id: id, + id: options.ID, options: options, db: db, pubsub: pubsub,