From d297e43b119a39b6eafac3dc013fd0e25c159ad4 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 | 60 +++++++++++++++++++++++++-----
 pkg/roachprod/install/services.go  | 11 ++++--
 2 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/pkg/roachprod/install/cockroach.go b/pkg/roachprod/install/cockroach.go
index a33dcc05b100..55feb239ba9f 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
@@ -345,28 +365,50 @@ func (c *SyncedCluster) Start(ctx context.Context, l *logger.Logger, startOpts S
 		return fmt.Errorf("start SQL proxy not implemented")
 	}
 
+	// Determine if custom ports were specified in the start options.
+	customPortsSpecified := func() bool {
+		if startOpts.SQLPort != 0 && startOpts.SQLPort != config.DefaultSQLPort {
+			return true
+		}
+		if startOpts.AdminUIPort != 0 && startOpts.AdminUIPort != config.DefaultAdminUIPort {
+			return true
+		}
+		return false
+	}
+
 	// Local clusters do not support specifying ports. An error is returned if we
 	// detect that they were set.
 	if c.IsLocal() {
+		if customPortsSpecified() {
+			return fmt.Errorf("local clusters do not support specifying ports")
+		}
 		// We don't need to return an error if the ports are the default values
 		// specified in DefaultStartOps, as these have not been specified explicitly
 		// by the user.
-		if startOpts.SQLPort != 0 && startOpts.SQLPort != config.DefaultSQLPort {
-			return fmt.Errorf("local clusters do not support specifying ports")
-		}
-		if startOpts.AdminUIPort != 0 && startOpts.AdminUIPort != config.DefaultAdminUIPort {
-			return fmt.Errorf("local clusters do not support specifying ports")
-		}
 		startOpts.SQLPort = 0
 		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 {
+		if customPortsSpecified() {
+			return fmt.Errorf("service registration is not supported for this cluster, but custom ports were specified")
+		}
+		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 {