chore: provisioner acquirer to respect organization ID of jobs

Emyrk · Emyrk · commit 026d40dfc2a3 · 2024-07-11T11:04:31.000-05:00
Prior to this the wrong provisioner was awakened on any new job
posting.
diff --git a/coderd/database/provisionerjobs/provisionerjobs.go b/coderd/database/provisionerjobs/provisionerjobs.go
@@ -3,6 +3,7 @@ package provisionerjobs
 import (
 	"encoding/json"
 
+	"github.com/google/uuid"
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/coderd/database"
@@ -12,12 +13,14 @@ import (
 const EventJobPosted = "provisioner_job_posted"
 
 type JobPosting struct {
+	OrganizationID  uuid.UUID                `json:"organization_id"`
 	ProvisionerType database.ProvisionerType `json:"type"`
 	Tags            map[string]string        `json:"tags"`
 }
 
 func PostJob(ps pubsub.Pubsub, job database.ProvisionerJob) error {
 	msg, err := json.Marshal(JobPosting{
+		OrganizationID:  job.OrganizationID,
 		ProvisionerType: job.Provisioner,
 		Tags:            job.Tags,
 	})
diff --git a/coderd/provisionerdserver/acquirer.go b/coderd/provisionerdserver/acquirer.go
@@ -163,13 +163,14 @@ func (a *Acquirer) want(organization uuid.UUID, pt []database.ProvisionerType, t
 	if !ok {
 		ctx, cancel := context.WithCancel(a.ctx)
 		d = domain{
-			ctx:       ctx,
-			cancel:    cancel,
-			a:         a,
-			key:       dk,
-			pt:        pt,
-			tags:      tags,
-			acquirees: make(map[chan<- struct{}]*acquiree),
+			ctx:            ctx,
+			cancel:         cancel,
+			a:              a,
+			key:            dk,
+			pt:             pt,
+			tags:           tags,
+			organizationID: organization,
+			acquirees:      make(map[chan<- struct{}]*acquiree),
 		}
 		a.q[dk] = d
 		go d.poll(a.backupPollDuration)
@@ -450,16 +451,22 @@ type acquiree struct {
 // tags.  Acquirees in the same domain are restricted such that only one queries
 // the database at a time.
 type domain struct {
-	ctx       context.Context
-	cancel    context.CancelFunc
-	a         *Acquirer
-	key       dKey
-	pt        []database.ProvisionerType
-	tags      Tags
-	acquirees map[chan<- struct{}]*acquiree
+	ctx            context.Context
+	cancel         context.CancelFunc
+	a              *Acquirer
+	key            dKey
+	pt             []database.ProvisionerType
+	tags           Tags
+	organizationID uuid.UUID
+	acquirees      map[chan<- struct{}]*acquiree
 }
 
 func (d domain) contains(p provisionerjobs.JobPosting) bool {
+	// If the organization ID is 'uuid.Nil', this is a legacy job posting.
+	// Ignore this check in the legacy case.
+	if p.OrganizationID != uuid.Nil && p.OrganizationID != d.organizationID {
+		return false
+	}
 	if !slices.Contains(d.pt, p.ProvisionerType) {
 		return false
 	}
diff --git a/enterprise/coderd/templates_test.go b/enterprise/coderd/templates_test.go
@@ -1855,8 +1855,11 @@ func TestMultipleOrganizationTemplates(t *testing.T) {
 		IncludeProvisionerDaemon: true,
 	})
 
-	var _, _ = third, templateAdmin
+	t.Logf("First organization: %s", first.OrganizationID.String())
+	t.Logf("Second organization: %s", second.ID.String())
+	t.Logf("Third organization: %s", third.ID.String())
 
+	t.Logf("Creating template version in second organization")
 	version := coderdtest.CreateTemplateVersion(t, templateAdmin, second.ID, nil)
 	coderdtest.AwaitTemplateVersionJobCompleted(t, ownerClient, version.ID)
 	coderdtest.CreateTemplate(t, templateAdmin, second.ID, version.ID, func(request *codersdk.CreateTemplateRequest) {
