Merge #123089

123089: roachprod: implement azure addLabels and removeLabels r=srosenberg a=DarrylWong

These were previously noops, which prevented adding the test_run_id and test_name as tags.

Release note: none
Epic: none
Fixes: #111885

----------

This doesn't do anything for us yet as we don't have prom set up for azure, but seemed like a quick win to slightly reduce some of the log spam in the test runner.

Roachprod now also starts updating the local cluster cache with labels when added/removed. This is because the Azure API doesn't support adding/removing individual tags like GCE does, you can only redeclare all the tags a VM should have 😢. So we have to keep track of old tags.

Co-authored-by: DarrylWong <darryl@cockroachlabs.com>

pkg/roachprod/roachprod.go

@@ -1371,9 +1371,27 @@ func AddLabels(l *logger.Logger, clusterName string, labels map[string]string) e
 		return err
 	}
 
-	return vm.FanOut(c.VMs, func(p vm.Provider, vms vm.List) error {
+	err = vm.FanOut(c.VMs, func(p vm.Provider, vms vm.List) error {
 		return p.AddLabels(l, vms, labels)
 	})
+	if err != nil {
+		return err
+	}
+
+	// Adding labels is not supported for local clusters, we don't
+	// need to update the local cluster cache.
+	if config.IsLocalClusterName(clusterName) {
+		return nil
+	}
+
+	// Update the tags in the local cluster cache.
+	for _, m := range c.Cluster.VMs {
+		for k, v := range labels {
+			m.Labels[k] = v
+		}
+	}
+
+	return saveCluster(l, &c.Cluster)
 }
 
 func RemoveLabels(l *logger.Logger, clusterName string, labels []string) error {
@@ -1382,9 +1400,20 @@ func RemoveLabels(l *logger.Logger, clusterName string, labels []string) error {
 		return err
 	}
 
-	return vm.FanOut(c.VMs, func(p vm.Provider, vms vm.List) error {
+	err = vm.FanOut(c.VMs, func(p vm.Provider, vms vm.List) error {
 		return p.RemoveLabels(l, vms, labels)
 	})
+	if err != nil {
+		return err
+	}
+
+	// Update the tags in the local cluster cache.
+	for _, m := range c.Cluster.VMs {
+		for _, label := range labels {
+			delete(m.Labels, label)
+		}
+	}
+	return saveCluster(l, &c.Cluster)
 }
 
 // Create TODO

pkg/roachprod/vm/azure/azure.go

@@ -199,12 +199,74 @@ func getAzureDefaultLabelMap(opts vm.CreateOpts) map[string]string {
 }
 
 func (p *Provider) AddLabels(l *logger.Logger, vms vm.List, labels map[string]string) error {
-	l.Printf("adding labels to Azure VMs not yet supported")
-	return nil
+	return p.editLabels(l, vms, labels, false /*removeLabels*/)
 }
 
 func (p *Provider) RemoveLabels(l *logger.Logger, vms vm.List, labels []string) error {
-	l.Printf("removing labels from Azure VMs not yet supported")
+	labelsMap := make(map[string]string, len(labels))
+	for _, label := range labels {
+		labelsMap[label] = ""
+	}
+	return p.editLabels(l, vms, labelsMap, true /*removeLabels*/)
+}
+
+func (p *Provider) editLabels(
+	l *logger.Logger, vms vm.List, labels map[string]string, removeLabels bool,
+) error {
+	ctx, cancel := context.WithTimeout(context.Background(), p.OperationTimeout)
+	defer cancel()
+
+	sub, err := p.getSubscription(ctx)
+	if err != nil {
+		return err
+	}
+	client := compute.NewVirtualMachinesClient(sub)
+	if client.Authorizer, err = p.getAuthorizer(); err != nil {
+		return err
+	}
+
+	futures := make([]compute.VirtualMachinesUpdateFuture, len(vms))
+	for idx, m := range vms {
+		vmParts, err := parseAzureID(m.ProviderID)
+		if err != nil {
+			return err
+		}
+		// N.B. VirtualMachineUpdate below overwrites _all_ VM tags. Hence, we must copy all unmodified tags.
+		tags := make(map[string]*string)
+		// Copy all known VM tags.
+		for k, v := range m.Labels {
+			tags[k] = to.StringPtr(v)
+		}
+
+		if removeLabels {
+			// Remove the matching VM tags.
+			for k := range labels {
+				delete(tags, k)
+			}
+		} else {
+			// Add the new VM tags.
+			for k, v := range labels {
+				tags[k] = to.StringPtr(v)
+			}
+		}
+
+		update := compute.VirtualMachineUpdate{
+			Tags: tags,
+		}
+		futures[idx], err = client.Update(ctx, vmParts.resourceGroup, vmParts.resourceName, update)
+		if err != nil {
+			return err
+		}
+	}
+
+	for _, future := range futures {
+		if err := future.WaitForCompletionRef(ctx, client.Client); err != nil {
+			return err
+		}
+		if _, err := future.Result(client); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
@@ -406,50 +468,9 @@ func (p *Provider) DeleteCluster(l *logger.Logger, name string) error {
 
 // Extend implements the vm.Provider interface.
 func (p *Provider) Extend(l *logger.Logger, vms vm.List, lifetime time.Duration) error {
-	ctx, cancel := context.WithTimeout(context.Background(), p.OperationTimeout)
-	defer cancel()
-
-	sub, err := p.getSubscription(ctx)
-	if err != nil {
-		return err
-	}
-	client := compute.NewVirtualMachinesClient(sub)
-	if client.Authorizer, err = p.getAuthorizer(); err != nil {
-		return err
-	}
-
-	futures := make([]compute.VirtualMachinesUpdateFuture, len(vms))
-	for idx, m := range vms {
-		vmParts, err := parseAzureID(m.ProviderID)
-		if err != nil {
-			return err
-		}
-		// N.B. VirtualMachineUpdate below overwrites _all_ VM tags. Hence, we must copy all unmodified tags.
-		tags := make(map[string]*string)
-		// Copy all known VM tags.
-		for k, v := range m.Labels {
-			tags[k] = to.StringPtr(v)
-		}
-		// Overwrite Lifetime tag.
-		tags[vm.TagLifetime] = to.StringPtr(lifetime.String())
-		update := compute.VirtualMachineUpdate{
-			Tags: tags,
-		}
-		futures[idx], err = client.Update(ctx, vmParts.resourceGroup, vmParts.resourceName, update)
-		if err != nil {
-			return err
-		}
-	}
-
-	for _, future := range futures {
-		if err := future.WaitForCompletionRef(ctx, client.Client); err != nil {
-			return err
-		}
-		if _, err := future.Result(client); err != nil {
-			return err
-		}
-	}
-	return nil
+	return p.AddLabels(l, vms, map[string]string{
+		vm.TagLifetime: lifetime.String(),
+	})
 }
 
 // FindActiveAccount implements vm.Provider.

pkg/roachprod/vm/gce/gcloud.go

@@ -153,7 +153,7 @@ func (jsonVM *jsonVM) toVM(
 
 	// Check "lifetime" label.
 	var lifetime time.Duration
-	if lifetimeStr, ok := jsonVM.Labels["lifetime"]; ok {
+	if lifetimeStr, ok := jsonVM.Labels[vm.TagLifetime]; ok {
 		if lifetime, err = time.ParseDuration(lifetimeStr); err != nil {
 			vmErrors = append(vmErrors, vm.ErrNoExpiration)
 		}
@@ -2303,7 +2303,7 @@ func (p *Provider) Reset(l *logger.Logger, vms vm.List) error {
 // Extend TODO(peter): document
 func (p *Provider) Extend(l *logger.Logger, vms vm.List, lifetime time.Duration) error {
 	return p.AddLabels(l, vms, map[string]string{
-		"lifetime": lifetime.String(),
+		vm.TagLifetime: lifetime.String(),
 	})
 }