feat: Associate connected workspace agents with replicas (#4914)

This will enable displaying a graph that associates agents
to running replicas.
This commit is contained in:
Kyle Carberry
2022-11-06 13:27:09 -08:00
committed by GitHub
parent 267b81af83
commit 5be6c7071e
10 changed files with 52 additions and 21 deletions

View File

@ -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]

View File

@ -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.';

View File

@ -0,0 +1,2 @@
ALTER TABLE workspace_agents
DROP COLUMN last_connected_replica_id;

View File

@ -0,0 +1,2 @@
ALTER TABLE workspace_agents
ADD COLUMN last_connected_replica_id uuid;

View File

@ -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 {

View File

@ -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,
)

View File

@ -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;

View File

@ -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