From 6e296c91227e279b11c0eab6db91732d39e5357d Mon Sep 17 00:00:00 2001
From: Herko Lategan <herko@cockroachlabs.com>
Date: Fri, 8 Mar 2024 17:04:33 +0000
Subject: [PATCH] roachprod: only register services when applicable

Previously, DNS services for clusters would be registered in all cases. This can
lead to some scenarios where `roachprod` is unable to correctly determine the
port of a service.

When a cluster is created in a custom GCP project (non-default project) the
`roachprod` garbage collector will destroy the records as it determines it has
nothing to tie the resources to since it has no visibility of the custom
project.

As we currently do not test virtual clusters on AWS or Azure, these have also
been excluded, because these providers do not specify using the Google Cloud DNS
Provider yet. For that reason, there is no point in trying to register services.

Epic: None
Release Note: None
---
 pkg/roachprod/install/cockroach.go | 37 +++++++++++++++++++++++++++---
 pkg/roachprod/install/services.go  | 11 ++++++---
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/pkg/roachprod/install/cockroach.go b/pkg/roachprod/install/cockroach.go
index a33dcc05b100..d5991da299b5 100644
--- a/pkg/roachprod/install/cockroach.go
+++ b/pkg/roachprod/install/cockroach.go
@@ -207,6 +207,26 @@ func (so StartOpts) GetJoinTargets() []Node {
 	return nodes
 }
 
+// allowServiceRegistration is a gating function that prevents the usage of
+// service registration, with DNS services, in scenarios where it is not
+// supported. This is currently the case for non-GCE clusters and for clusters
+// that are not part of the default GCE project. For custom projects the DNS
+// services get garbage collected, hence the registration is not allowed.
+func (c *SyncedCluster) allowServiceRegistration() bool {
+	if c.IsLocal() {
+		return true
+	}
+	for _, cVM := range c.VMs {
+		if cVM.Provider != gce.ProviderName {
+			return false
+		}
+		if cVM.Project != gce.DefaultProject() {
+			return false
+		}
+	}
+	return true
+}
+
 // maybeRegisterServices registers the SQL and Admin UI DNS services
 // for the cluster if no previous services for the virtual or storage
 // cluster are found. Any ports specified in the startOpts are used
@@ -361,12 +381,23 @@ func (c *SyncedCluster) Start(ctx context.Context, l *logger.Logger, startOpts S
 		startOpts.AdminUIPort = 0
 	}
 
-	err := c.maybeRegisterServices(ctx, l, startOpts)
-	if err != nil {
-		return err
+	if c.allowServiceRegistration() {
+		err := c.maybeRegisterServices(ctx, l, startOpts)
+		if err != nil {
+			return err
+		}
+	} else {
+		l.Printf(strings.Join([]string{
+			"WARNING: Service registration and custom ports are not supported for this cluster.",
+			fmt.Sprintf("Setting ports to default SQL Port: %d, and Admin UI Port: %d.", config.DefaultSQLPort, config.DefaultAdminUIPort),
+			"Attempting to start any additional external SQL processes will fail.",
+		}, "\n"))
+		startOpts.SQLPort = config.DefaultSQLPort
+		startOpts.AdminUIPort = config.DefaultAdminUIPort
 	}
 
 	if startOpts.IsVirtualCluster() {
+		var err error
 		startOpts.VirtualClusterID, err = c.upsertVirtualClusterMetadata(ctx, l, startOpts)
 		if err != nil {
 			return err
diff --git a/pkg/roachprod/install/services.go b/pkg/roachprod/install/services.go
index b0d28f93d4d2..3f996606b8df 100644
--- a/pkg/roachprod/install/services.go
+++ b/pkg/roachprod/install/services.go
@@ -215,9 +215,14 @@ func (c *SyncedCluster) DiscoverService(
 	}
 
 	// Finally, fall back to the default ports if no services are found. This is
-	// useful for backwards compatibility with clusters that were created before
-	// the introduction of service discovery, or without a DNS provider.
-	// TODO(Herko): Remove this once DNS support is fully functional.
+	// required for scenarios where the services were not registered with a DNS
+	// provider (Google DNS). Currently, services will not be registered with DNS
+	// for clusters not on GCP, and it will also not be registered for GCP
+	// clusters that specify a custom project. The fall back is also useful for
+	// backwards compatibility with clusters that were created before the
+	// introduction of service discovery, or without a DNS provider.
+	// TODO(Herko): Remove this once DNS support is fully
+	// functional.
 	if len(services) == 0 {
 		var port int
 		switch serviceType {