coder · johnstcn · Feb 21, 2024 · Feb 21, 2024 · Feb 21, 2024 · mafredri
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"}`),
+				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,22 @@ var deletedUserLinkError = &pq.Error{
 	Routine: "exec_stmt_raise",
 }
 
+// m1 and m2 are equal iff |m1| = |m2| ^ m2 ⊆ m1
+func tagsEqual(m1, m2 map[string]string) bool {
+	return len(m1) == len(m2) && tagsSubset(m1, m2)
+}
+
+// m2 is a subset of m1 if each key in m1 exists in m2
+// with the same value
+func tagsSubset(m1, m2 map[string]string) bool {
+	for k, v1 := range m1 {
+		if v2, found := m2[k]; !found || v1 != v2 {
+			return false
+		}
+	}
+	return true
+}
+
 func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
 	return xerrors.New("AcquireLock must only be called within a transaction")
 }
@@ -783,19 +799,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(provisionerJob.Tags, 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,13 @@ 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 if requested,
+			AND CASE
+				WHEN @exact_tag_match :: boolean THEN nested.tags = @tags :: jsonb
+			-- Otherwise, ensure caller satisfies a subset of tags.
+			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()
@@ -128,8 +138,9 @@ func (a *Acquirer) AcquireJob(
 					UUID:  worker,
 					Valid: true,
 				},
-				Types: pt,
-				Tags:  dbTags,
+				Types:         pt,
+				Tags:          dbTags,
+				ExactTagMatch: a.exactTagMatch,
 			})
 			if xerrors.Is(err, sql.ErrNoRows) {
 				logger.Debug(ctx, "no job available")

diff --git a/coderd/provisionerdserver/acquirer_test.go b/coderd/provisionerdserver/acquirer_test.go
@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/sqlc-dev/pqtype"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
@@ -18,6 +19,8 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbmem"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/database/provisionerjobs"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/provisionerdserver"
@@ -315,6 +318,91 @@ func TestAcquirer_UnblockOnCancel(t *testing.T) {
 	require.Equal(t, jobID, job.ID)
 }
 
+func TestAcquirer_ExactTagMatch(t *testing.T) {
+	t.Parallel()
+	if testing.Short() {
+		t.Skip("skipping this test due to -short")
+	}
+
+	for _, tt := range []struct {
+		name               string
+		provisionerJobTags map[string]string
+		acquireJobTags     map[string]string
+		expectAcquire      bool
+	}{
+		{
+			name:               "match",
+			provisionerJobTags: map[string]string{"scope": "organization", "owner": "", "foo": "bar"},
+			acquireJobTags:     map[string]string{"scope": "organization", "owner": "", "foo": "bar"},
+			expectAcquire:      true,
+		},
+		{
+			name:               "subset",
+			provisionerJobTags: map[string]string{"scope": "organization", "owner": "", "foo": "bar"},
+			acquireJobTags:     map[string]string{"scope": "organization", "owner": ""},
+			expectAcquire:      false,
+		},
+		{
+			name:               "key mismatch",
+			provisionerJobTags: map[string]string{"scope": "organization", "owner": "", "fop": "bar"},
+			acquireJobTags:     map[string]string{"scope": "organization", "owner": "", "foo": "bar"},
+			expectAcquire:      false,
+		},
+		{
+			name:               "value mismatch",
+			provisionerJobTags: map[string]string{"scope": "organization", "owner": "", "foo": "baz"},
+			acquireJobTags:     map[string]string{"scope": "organization", "owner": "", "foo": "bar"},
+			expectAcquire:      false,
+		},
+	} {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := testutil.Context(t, testutil.WaitShort)
+			// NOTE: explicitly not using fake store for this test.
+			db, ps := dbtestutil.NewDB(t)
+			log := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+			org, err := db.InsertOrganization(ctx, database.InsertOrganizationParams{
+				ID:          uuid.New(),
+				Name:        "test org",
+				Description: "the organization of testing",
+				CreatedAt:   dbtime.Now(),
+				UpdatedAt:   dbtime.Now(),
+			})
+			require.NoError(t, err)
+			pj, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
+				ID:             uuid.New(),
+				CreatedAt:      dbtime.Now(),
+				UpdatedAt:      dbtime.Now(),
+				OrganizationID: org.ID,
+				InitiatorID:    uuid.New(),
+				Provisioner:    database.ProvisionerTypeEcho,
+				StorageMethod:  database.ProvisionerStorageMethodFile,
+				FileID:         uuid.New(),
+				Type:           database.ProvisionerJobTypeWorkspaceBuild,
+				Input:          []byte("{}"),
+				Tags:           tt.provisionerJobTags,
+				TraceMetadata:  pqtype.NullRawMessage{},
+			})
+			require.NoError(t, err)
+			ptypes := []database.ProvisionerType{database.ProvisionerTypeEcho}
+			opts := []provisionerdserver.AcquirerOption{
+				provisionerdserver.WithExactTagMatch(),
+			}
+			acq := provisionerdserver.NewAcquirer(ctx, log, db, ps, opts...)
+			aj, err := acq.AcquireJob(ctx, uuid.New(), ptypes, tt.acquireJobTags)
+			if tt.expectAcquire {
+				require.NoError(t, err)
+				require.Equal(t, pj.ID, aj.ID)
+			} else {
+				require.ErrorIs(t, err, context.DeadlineExceeded, "should have timed out")
+				require.Empty(t, aj, "should not have acquired job")
+			}
+		})
+	}
+}
+
 func postJob(t *testing.T, ps pubsub.Pubsub, pt database.ProvisionerType, tags provisionerdserver.Tags) {
 	t.Helper()
 	msg, err := json.Marshal(provisionerjobs.JobPosting{

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)