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

kylecarbs · web-flow · commit 5be6c7071ea5 · 2022-11-06T15:27:09.000-06:00
This will enable displaying a graph that associates agents
to running replicas.
diff --git 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
diff --git a/coderd/database/migrations/000070_deployment_graph.down.sql 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
@@ -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
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git 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
@@ -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
@@ -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
@@ -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,

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+ALTER TABLE workspace_agents`
	`2`	`+ DROP COLUMN last_connected_replica_id;`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+ALTER TABLE workspace_agents`
	`2`	`+ ADD COLUMN last_connected_replica_id uuid;`
Original file line number	Diff line number	Diff line change
`@@ -158,6 +158,7 @@ func New(ctx context.Context, options Options) (API, error) {`
`158`	`158`	`}`
`159`	`159`	`var err error`
`160`	`160`	`api.replicaManager, err = replicasync.New(ctx, options.Logger, options.Database, options.Pubsub, &replicasync.Options{`
	`161`	`+ ID: api.AGPL.ID,`
`161`	`162`	`RelayAddress: options.DERPServerRelayAddress,`
`162`	`163`	`RegionID: int32(options.DERPServerRegionID),`
`163`	`164`	`TLSConfig: meshTLSConfig,`