roachprod: add A records during cluster create

Patch the lookup logic to handle A records in addition to SRV records
and make the update/create record logic work with public records.

Epic: none
Release note: None

Author: golgeek

pkg/roachprod/install/services.go

@@ -154,7 +154,7 @@ func (c *SyncedCluster) discoverServices(
 	mu := syncutil.Mutex{}
 	records := make([]vm.DNSRecord, 0)
 	err := vm.FanOutDNS(c.VMs, func(dnsProvider vm.DNSProvider, _ vm.List) error {
-		r, lookupErr := dnsProvider.LookupSRVRecords(ctx, serviceDNSName(dnsProvider, virtualClusterName, serviceType, c.Name))
+		r, lookupErr := dnsProvider.LookupRecords(ctx, vm.SRV, serviceDNSName(dnsProvider, virtualClusterName, serviceType, c.Name))
 		if lookupErr != nil {
 			return lookupErr
 		}

pkg/roachprod/vm/dns.go

@@ -51,7 +51,7 @@ type DNSProvider interface {
 	// subdomain. The protocol is usually "tcp" and the subdomain is usually the
 	// cluster name. The service is a combination of the virtual cluster name and
 	// type of service.
-	LookupSRVRecords(ctx context.Context, name string) ([]DNSRecord, error)
+	LookupRecords(ctx context.Context, recordType DNSType, name string) ([]DNSRecord, error)
 	// ListRecords lists all DNS records managed for the zone.
 	ListRecords(ctx context.Context) ([]DNSRecord, error)
 	// DeleteSRVRecordsBySubdomain deletes all DNS SRV records with the given subdomain.

pkg/roachprod/vm/gce/dns.go

@@ -131,8 +131,12 @@ func (n *dnsProvider) CreateRecords(ctx context.Context, records ...vm.DNSRecord
 	}
 
 	for name, recordGroup := range recordsByName {
+		// We assume that all records in a group have the same name, type, and ttl.
+		// TODO(herko): Add error checking to ensure that the above is the case.
+		firstRecord := recordGroup[0]
+
 		err := n.withRecordLock(name, func() error {
-			existingRecords, err := n.lookupSRVRecords(ctx, name)
+			existingRecords, err := n.lookupRecords(ctx, firstRecord.Type, name)
 			if err != nil {
 				return err
 			}
@@ -151,9 +155,6 @@ func (n *dnsProvider) CreateRecords(ctx context.Context, records ...vm.DNSRecord
 				combinedRecords[record.Data] = record
 			}
 
-			// We assume that all records in a group have the same name, type, and ttl.
-			// TODO(herko): Add error checking to ensure that the above is the case.
-			firstRecord := recordGroup[0]
 			data := maps.Keys(combinedRecords)
 			sort.Strings(data)
 			zone := n.managedZone
@@ -194,24 +195,26 @@ func (n *dnsProvider) CreateRecords(ctx context.Context, records ...vm.DNSRecord
 }
 
 // LookupSRVRecords implements the vm.DNSProvider interface.
-func (n *dnsProvider) LookupSRVRecords(ctx context.Context, name string) ([]vm.DNSRecord, error) {
+func (n *dnsProvider) LookupRecords(
+	ctx context.Context, recordType vm.DNSType, name string,
+) ([]vm.DNSRecord, error) {
 	var records []vm.DNSRecord
 	var err error
 	err = n.withRecordLock(name, func() error {
-		if config.FastDNS {
+		if config.FastDNS && recordType == vm.SRV {
 			rIdx := randutil.FastUint32() % uint32(len(n.resolvers))
 			records, err = n.fastLookupSRVRecords(ctx, n.resolvers[rIdx], name, true)
 			return err
 		}
-		records, err = n.lookupSRVRecords(ctx, name)
+		records, err = n.lookupRecords(ctx, recordType, name)
 		return err
 	})
 	return records, err
 }
 
 // ListRecords implements the vm.DNSProvider interface.
 func (n *dnsProvider) ListRecords(ctx context.Context) ([]vm.DNSRecord, error) {
-	return n.listSRVRecords(ctx, "", dnsMaxResults)
+	return n.listRecords(ctx, vm.SRV, "", dnsMaxResults)
 }
 
 func (n *dnsProvider) deleteRecords(
@@ -253,7 +256,7 @@ func (n *dnsProvider) DeletePublicRecordsByName(ctx context.Context, names ...st
 // DeleteRecordsBySubdomain implements the vm.DNSProvider interface.
 func (n *dnsProvider) DeleteSRVRecordsBySubdomain(ctx context.Context, subdomain string) error {
 	suffix := fmt.Sprintf("%s.%s.", subdomain, n.Domain())
-	records, err := n.listSRVRecords(ctx, suffix, dnsMaxResults)
+	records, err := n.listRecords(ctx, vm.SRV, suffix, dnsMaxResults)
 	if err != nil {
 		return err
 	}
@@ -287,13 +290,15 @@ func (n *dnsProvider) Domain() string {
 // network problems. For lookups, we prefer this to using the gcloud command as
 // it is faster, and preferable when service information is being queried
 // regularly.
-func (n *dnsProvider) lookupSRVRecords(ctx context.Context, name string) ([]vm.DNSRecord, error) {
+func (n *dnsProvider) lookupRecords(
+	ctx context.Context, recordType vm.DNSType, name string,
+) ([]vm.DNSRecord, error) {
 	// Check the cache first.
 	if cachedRecords, ok := n.getCache(name); ok {
 		return cachedRecords, nil
 	}
 	// Lookup the records, if no records are found in the cache.
-	records, err := n.listSRVRecords(ctx, name, dnsMaxResults)
+	records, err := n.listRecords(ctx, recordType, name, dnsMaxResults)
 	if err != nil {
 		return nil, err
 	}
@@ -310,16 +315,21 @@ func (n *dnsProvider) lookupSRVRecords(ctx context.Context, name string) ([]vm.D
 	return filteredRecords, nil
 }
 
-// listSRVRecords returns all SRV records that match the given filter from Google Cloud DNS.
+// listRecords returns all records that match the given filter from Google Cloud DNS.
 // The data field of the records could be a comma-separated list of values if multiple
 // records are returned for the same name.
-func (n *dnsProvider) listSRVRecords(
-	ctx context.Context, filter string, limit int,
+func (n *dnsProvider) listRecords(
+	ctx context.Context, recordType vm.DNSType, filter string, limit int,
 ) ([]vm.DNSRecord, error) {
+	zone := n.managedZone
+	if recordType == vm.A {
+		zone = n.publicZone
+	}
+
 	args := []string{"--project", n.dnsProject, "dns", "record-sets", "list",
 		"--limit", strconv.Itoa(limit),
 		"--page-size", strconv.Itoa(limit),
-		"--zone", n.managedZone,
+		"--zone", zone,
 		"--format", "json",
 	}
 	if filter != "" {
@@ -348,11 +358,11 @@ func (n *dnsProvider) listSRVRecords(
 		if record.Kind != "dns#resourceRecordSet" {
 			continue
 		}
-		if record.RecordType != string(vm.SRV) {
+		if record.RecordType != string(recordType) {
 			continue
 		}
 		for _, data := range record.RRDatas {
-			records = append(records, vm.CreateDNSRecord(record.Name, vm.SRV, data, record.TTL))
+			records = append(records, vm.CreateDNSRecord(record.Name, recordType, data, record.TTL))
 		}
 	}
 	return records, nil

pkg/roachprod/vm/local/dns.go

@@ -58,15 +58,17 @@ func (n *dnsProvider) CreateRecords(_ context.Context, records ...vm.DNSRecord)
 	return n.saveRecords(entries)
 }
 
-// LookupSRVRecords is part of the vm.DNSProvider interface.
-func (n *dnsProvider) LookupSRVRecords(_ context.Context, name string) ([]vm.DNSRecord, error) {
+// LookupRecords is part of the vm.DNSProvider interface.
+func (n *dnsProvider) LookupRecords(
+	_ context.Context, recordType vm.DNSType, name string,
+) ([]vm.DNSRecord, error) {
 	records, err := n.loadRecords()
 	if err != nil {
 		return nil, err
 	}
 	var matchingRecords []vm.DNSRecord
 	for _, record := range records {
-		if record.Name == name && record.Type == vm.SRV {
+		if record.Name == name && record.Type == recordType {
 			matchingRecords = append(matchingRecords, record)
 		}
 	}

pkg/roachprod/vm/local/dns_test.go

@@ -48,7 +48,7 @@ func TestLookupRecords(t *testing.T) {
 	}...)
 
 	t.Run("lookup system", func(t *testing.T) {
-		records, err := p.LookupSRVRecords(ctx, "_system-sql._tcp.local.local-zone")
+		records, err := p.LookupRecords(ctx, vm.SRV, "_system-sql._tcp.local.local-zone")
 		require.NoError(t, err)
 		require.Equal(t, 3, len(records))
 		for _, r := range records {
@@ -58,7 +58,7 @@ func TestLookupRecords(t *testing.T) {
 	})
 
 	t.Run("parse SRV data", func(t *testing.T) {
-		records, err := p.LookupSRVRecords(ctx, "_tenant-1-sql._tcp.local.local-zone")
+		records, err := p.LookupRecords(ctx, vm.SRV, "_tenant-1-sql._tcp.local.local-zone")
 		require.NoError(t, err)
 		require.Equal(t, 1, len(records))
 		data, err := records[0].ParseSRVRecord()