coder · kylecarbs · Nov 6, 2022 · Nov 6, 2022 · coadler · Nov 6, 2022
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 @@
		ALTER TABLE workspace_agents
		DROP COLUMN last_connected_replica_id;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		ALTER TABLE workspace_agents
		ADD COLUMN last_connected_replica_id uuid;
Comment on lines +1 to +2 Copy link Contributor coadler Nov 6, 2022 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. I'd prefer to make this not null with a default `'00000000-0000-0000-0000-000000000000'::uuid` since in practice this should never be null. Copy link Member Author kylecarbs Nov 6, 2022 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. It can be null if the agent hasn't connected yet Copy link Contributor coadler Nov 6, 2022 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. oh good point