feat(coderd/database): update AcquireProvisionerJob query to allow specifying exact tag match behaviour

johnstcn · johnstcn · commit 9b77b425f3a0 · 2024-02-21T11:57:49.000Z
diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -192,8 +192,9 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
 					UUID:  uuid.New(),
 					Valid: true,
 				},
-				Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
-				Tags:  []byte(`{"scope": "organization"}`),
+				Types:         []database.ProvisionerType{database.ProvisionerTypeEcho},
+				Tags:          []byte(`{"scope": "organization", "owner": ""}`),
+				ExactTagMatch: false,
 			})
 			require.NoError(b.t, err, "acquire starting job")
 			if j.ID == job.ID {
diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -417,10 +417,11 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
 	}
 	if !orig.StartedAt.Time.IsZero() {
 		job, err = db.AcquireProvisionerJob(genCtx, database.AcquireProvisionerJobParams{
-			StartedAt: orig.StartedAt,
-			Types:     []database.ProvisionerType{database.ProvisionerTypeEcho},
-			Tags:      must(json.Marshal(orig.Tags)),
-			WorkerID:  uuid.NullUUID{},
+			StartedAt:     orig.StartedAt,
+			Types:         []database.ProvisionerType{database.ProvisionerTypeEcho},
+			Tags:          must(json.Marshal(orig.Tags)),
+			WorkerID:      uuid.NullUUID{},
+			ExactTagMatch: false,
 		})
 		require.NoError(t, err)
 		// There is no easy way to make sure we acquire the correct job.
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -748,6 +748,19 @@ var deletedUserLinkError = &pq.Error{
 	Routine: "exec_stmt_raise",
 }
 
+func tagsEqual(m1, m2 map[string]string) bool {
+	return tagsSubset(m1, m2) && tagsSubset(m2, m1)
+}
+
+func tagsSubset(needle, haystack map[string]string) bool {
+	for k := range needle {
+		if needle[k] != haystack[k] {
+			return false
+		}
+	}
+	return true
+}
+
 func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
 	return xerrors.New("AcquireLock must only be called within a transaction")
 }
@@ -783,19 +796,11 @@ func (q *FakeQuerier) AcquireProvisionerJob(_ context.Context, arg database.Acqu
 			}
 		}
 
-		missing := false
-		for key, value := range provisionerJob.Tags {
-			provided, found := tags[key]
-			if !found {
-				missing = true
-				break
-			}
-			if provided != value {
-				missing = true
-				break
-			}
+		matchFunc := tagsSubset
+		if arg.ExactTagMatch {
+			matchFunc = tagsEqual
 		}
-		if missing {
+		if !matchFunc(tags, provisionerJob.Tags) {
 			continue
 		}
 		provisionerJob.StartedAt = arg.StartedAt
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/provisionerjobs.sql b/coderd/database/queries/provisionerjobs.sql
@@ -21,8 +21,12 @@ WHERE
 			nested.started_at IS NULL
 			-- Ensure the caller has the correct provisioner.
 			AND nested.provisioner = ANY(@types :: provisioner_type [ ])
-			-- Ensure the caller satisfies all job tags.
-			AND nested.tags <@ @tags :: jsonb
+			-- Ensure the caller satisfies all job tags according to the operator
+			AND CASE
+				WHEN @exact_tag_match :: boolean THEN nested.tags = @tags :: jsonb
+			ELSE
+				nested.tags = @tags :: jsonb
+			END
 		ORDER BY
 			nested.created_at
 		FOR UPDATE
diff --git a/coderd/provisionerdserver/acquirer.go b/coderd/provisionerdserver/acquirer.go
@@ -49,6 +49,8 @@ type Acquirer struct {
 	mu sync.Mutex
 	q  map[dKey]domain
 
+	exactTagMatch bool
+
 	// testing only
 	backupPollDuration time.Duration
 }
@@ -61,6 +63,12 @@ func TestingBackupPollDuration(dur time.Duration) AcquirerOption {
 	}
 }
 
+func WithExactTagMatch() AcquirerOption {
+	return func(a *Acquirer) {
+		a.exactTagMatch = true
+	}
+}
+
 // AcquirerStore is the subset of database.Store that the Acquirer needs
 type AcquirerStore interface {
 	AcquireProvisionerJob(context.Context, database.AcquireProvisionerJobParams) (database.ProvisionerJob, error)
@@ -76,6 +84,7 @@ func NewAcquirer(ctx context.Context, logger slog.Logger, store AcquirerStore, p
 		ps:                 ps,
 		q:                  make(map[dKey]domain),
 		backupPollDuration: backupPollDuration,
+		exactTagMatch:      false,
 	}
 	for _, opt := range opts {
 		opt(a)
@@ -96,7 +105,8 @@ func (a *Acquirer) AcquireJob(
 	logger := a.logger.With(
 		slog.F("worker_id", worker),
 		slog.F("provisioner_types", pt),
-		slog.F("tags", tags))
+		slog.F("tags", tags),
+		slog.F("exact_tag_match", a.exactTagMatch))
 	logger.Debug(ctx, "acquiring job")
 	dk := domainKey(pt, tags)
 	dbTags, err := tags.ToJSON()
diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -573,7 +573,8 @@ func TestUpdateJob(t *testing.T) {
 				UUID:  srvID,
 				Valid: true,
 			},
-			Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
+			Types:         []database.ProvisionerType{database.ProvisionerTypeEcho},
+			ExactTagMatch: false,
 		})
 		require.NoError(t, err)
 		return job.ID
diff --git a/enterprise/coderd/schedule/template_test.go b/enterprise/coderd/schedule/template_test.go
@@ -181,8 +181,9 @@ func TestTemplateUpdateBuildDeadlines(t *testing.T) {
 					UUID:  uuid.New(),
 					Valid: true,
 				},
-				Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
-				Tags:  json.RawMessage(fmt.Sprintf(`{%q: "yeah"}`, c.name)),
+				Types:         []database.ProvisionerType{database.ProvisionerTypeEcho},
+				Tags:          json.RawMessage(fmt.Sprintf(`{%q: "yeah"}`, c.name)),
+				ExactTagMatch: false,
 			})
 			require.NoError(t, err)
 			require.Equal(t, job.ID, acquiredJob.ID)
