chore: provisioner acquirer to respect organization ID of jobs (coder#13874)

Emyrk · web-flow · commit 687d9538de6f · 2024-07-11T11:26:47.000-05:00
* test: add unit test to verify creation of templates in multiple orgs
* chore: provisioner acquirer to respect organization ID of jobs

Prior to this the wrong provisioner was awakened on any new job
posting.

* add comment and stricter check
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
@@ -1829,3 +1829,51 @@ func TestTemplateAccess(t *testing.T) {
 		}
 	})
 }
+
+func TestMultipleOrganizationTemplates(t *testing.T) {
+	t.Parallel()
+
+	ownerClient, first := coderdenttest.New(t, &coderdenttest.Options{
+		Options: &coderdtest.Options{
+			// This only affects the first org.
+			IncludeProvisionerDaemon: true,
+		},
+		LicenseOptions: &coderdenttest.LicenseOptions{
+			Features: license.Features{
+				codersdk.FeatureExternalProvisionerDaemons: 1,
+			},
+		},
+	})
+
+	templateAdmin, _ := coderdtest.CreateAnotherUser(t, ownerClient, first.OrganizationID, rbac.RoleTemplateAdmin())
+
+	second := coderdtest.CreateOrganization(t, ownerClient, coderdtest.CreateOrganizationOptions{
+		IncludeProvisionerDaemon: true,
+	})
+
+	third := coderdtest.CreateOrganization(t, ownerClient, coderdtest.CreateOrganizationOptions{
+		IncludeProvisionerDaemon: true,
+	})
+
+	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")
+
+	start := time.Now()
+	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) {
+		request.Name = "random"
+	})
+
+	if time.Since(start) > time.Second*10 {
+		// The test can sometimes pass because 'AwaitTemplateVersionJobCompleted'
+		// allows 25s, and the provisioner will check every 30s if not awakened
+		// from the pubsub. So there is a chance it will pass. If it takes longer
+		// than 10s, then it's a problem. The provisioner is not getting clearance.
+		t.Error("Creating template version in second organization took too long")
+		t.FailNow()
+	}
+}
