Ignore remote DC nodes when resolving endpoints (Astra only)

Remote DC nodes require their own secure connect bundle and are unable
to be connected to. The Astra endpoint resolver will now ignore these
nodes so that it avoid attempting connections.

Author: mpenick

astra/endpoint.go

@@ -25,13 +25,12 @@ import (
 	"sync"
 	"time"
 
-	"github.com/datastax/go-cassandra-native-protocol/primitive"
-
 	"github.com/datastax/cql-proxy/proxycore"
 )
 
 type astraResolver struct {
 	sniProxyAddress string
+	region          string
 	bundle          *Bundle
 	mu              *sync.Mutex
 }
@@ -77,6 +76,7 @@ func (r *astraResolver) Resolve() ([]proxycore.Endpoint, error) {
 
 	r.mu.Lock()
 	r.sniProxyAddress = sniProxyAddress
+	r.region = metadata.Region
 	r.mu.Unlock()
 
 	var endpoints []proxycore.Endpoint
@@ -91,31 +91,35 @@ func (r *astraResolver) Resolve() ([]proxycore.Endpoint, error) {
 	return endpoints, nil
 }
 
-func (r *astraResolver) getSNIProxyAddress() (string, error) {
+func (r *astraResolver) getSNIProxyAddressAndRegion() (string, string, error) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	if len(r.sniProxyAddress) == 0 {
-		return "", errors.New("SNI proxy address never resolved")
+		return "", "", errors.New("SNI proxy address (and region) never resolved")
 	}
-	return r.sniProxyAddress, nil
+	return r.sniProxyAddress, r.region, nil
 }
 
 func (r *astraResolver) NewEndpoint(row proxycore.Row) (proxycore.Endpoint, error) {
-	sniProxyAddress, err := r.getSNIProxyAddress()
+	sniProxyAddress, region, err := r.getSNIProxyAddressAndRegion()
 	if err != nil {
 		return nil, err
 	}
-	hostId, err := row.ByName("host_id")
+	dc, err := row.StringByName("data_center")
 	if err != nil {
 		return nil, err
 	}
-	if uuid, ok := hostId.(primitive.UUID); !ok {
-		return nil, errors.New("ignoring host because its `host_id` is not set or is invalid")
+	if len(region) > 0 && region != dc {
+		return nil, proxycore.IgnoreEndpoint
+	}
+	hostId, err := row.UUIDByName("host_id")
+	if err != nil {
+		return nil, err
 	} else {
 		return &astraEndpoint{
 			addr:      sniProxyAddress,
-			key:       fmt.Sprintf("%s:%s", sniProxyAddress, &uuid),
-			tlsConfig: copyTLSConfig(r.bundle, uuid.String()),
+			key:       fmt.Sprintf("%s:%s", sniProxyAddress, &hostId),
+			tlsConfig: copyTLSConfig(r.bundle, hostId.String()),
 		}, nil
 	}
 }

proxycore/cluster.go

@@ -240,7 +240,7 @@ func (c *Cluster) mergeHosts(hosts []*Host) error {
 	existing := make(map[string]*Host)
 
 	for _, host := range c.hosts {
-		existing[host.Endpoint().Key()] = host
+		existing[host.Key()] = host
 	}
 
 	c.currentHostIndex = -1
@@ -282,7 +282,6 @@ func (c *Cluster) sendEvent(event Event) {
 
 func (c *Cluster) queryHosts(ctx context.Context, conn *ClientConn, version primitive.ProtocolVersion) (hosts []*Host, info ClusterInfo, err error) {
 	var rs *ResultSet
-	var val interface{}
 
 	rs, err = conn.Query(ctx, version, &message.Query{
 		Query: "SELECT * FROM system.local",
@@ -299,23 +298,20 @@ func (c *Cluster) queryHosts(ctx context.Context, conn *ClientConn, version prim
 	hosts = c.addHosts(hosts, rs)
 	row := rs.Row(0)
 
-	val, err = row.ByName("partitioner")
+	partitioner, err := row.StringByName("partitioner")
 	if err != nil {
 		return nil, ClusterInfo{}, err
 	}
-	partitioner := val.(string)
 
-	val, err = row.ByName("release_version")
+	releaseVersion, err := row.StringByName("release_version")
 	if err != nil {
 		return nil, ClusterInfo{}, err
 	}
-	releaseVersion := val.(string)
 
-	val, err = row.ByName("cql_version")
+	cqlVersion, err := row.StringByName("cql_version")
 	if err != nil {
 		return nil, ClusterInfo{}, err
 	}
-	cqlVersion := val.(string)
 
 	rs, err = conn.Query(ctx, version, &message.Query{
 		Query: "SELECT * FROM system.peers",
@@ -341,7 +337,11 @@ func (c *Cluster) addHosts(hosts []*Host, rs *ResultSet) []*Host {
 		if endpoint, err := c.config.Resolver.NewEndpoint(row); err == nil {
 			if host, err := NewHostFromRow(endpoint, row); err == nil {
 				hosts = append(hosts, host)
+			} else {
+				c.logger.Error("unable to create new host", zap.Stringer("endpoint", endpoint), zap.Error(err))
 			}
+		} else if err != IgnoreEndpoint {
+			c.logger.Error("unable to create new endpoint", zap.Error(err))
 		}
 	}
 	return hosts
@@ -350,7 +350,7 @@ func (c *Cluster) addHosts(hosts []*Host, rs *ResultSet) []*Host {
 func (c *Cluster) reconnect() bool {
 	c.currentHostIndex = (c.currentHostIndex + 1) % len(c.hosts)
 	host := c.hosts[c.currentHostIndex]
-	err := c.connect(c.ctx, host.Endpoint(), false)
+	err := c.connect(c.ctx, host.Endpoint, false)
 	if err != nil {
 		c.logger.Error("error reconnecting to host", zap.Stringer("host", host), zap.Error(err))
 		return false

proxycore/endpoint.go

@@ -23,6 +23,8 @@ import (
 	"strconv"
 )
 
+var IgnoreEndpoint = errors.New("ignore endpoint")
+
 type Endpoint interface {
 	fmt.Stringer
 	Addr() string
@@ -108,19 +110,14 @@ func (r *defaultEndpointResolver) NewEndpoint(row Row) (Endpoint, error) {
 	if err != nil && !errors.Is(err, ColumnNameNotFound) {
 		return nil, err
 	}
-	rpcAddress, err := row.ByName("rpc_address")
+	rpcAddress, err := row.InetByName("rpc_address")
 	if err != nil {
-		return nil, err
-	}
-
-	var ok bool
-	var addr net.IP
-
-	if addr, ok = rpcAddress.(net.IP); !ok {
-		return nil, errors.New("ignoring host because its `rpc_address` is not set or is invalid")
+		return nil, fmt.Errorf("ignoring host because its `rpc_address` is not set or is invalid: %w", err)
 	}
 
+	addr := rpcAddress
 	if addr.Equal(net.IPv4zero) || addr.Equal(net.IPv6zero) {
+		var ok bool
 		if addr, ok = peer.(net.IP); !ok {
 			return nil, errors.New("ignoring host because its `peer` is not set or is invalid")
 		}

proxycore/host.go

@@ -15,21 +15,17 @@
 package proxycore
 
 type Host struct {
-	endpoint Endpoint
+	Endpoint
 }
 
-func NewHostFromRow(endpoint Endpoint, _ Row) (*Host, error) {
+func NewHostFromRow(endpoint Endpoint, row Row) (*Host, error) {
 	return &Host{endpoint}, nil
 }
 
-func (h *Host) Endpoint() Endpoint {
-	return h.endpoint
-}
-
 func (h *Host) Key() string {
-	return h.endpoint.Key()
+	return h.Endpoint.Key()
 }
 
 func (h *Host) String() string {
-	return h.endpoint.String()
+	return h.Endpoint.String()
 }

proxycore/host.go.orig

@@ -0,0 +1,45 @@
+// Copyright (c) DataStax, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package proxycore
+
+type Host struct {
+	endpoint Endpoint
+}
+
+<<<<<<< HEAD
+func NewHostFromRow(endpoint Endpoint, _ Row) (*Host, error) {
+	return &Host{endpoint}, nil
+=======
+func NewHostFromRow(endpoint Endpoint, row Row) (*Host, error) {
+	dc, err := row.StringByName("data_center")
+	if err != nil {
+		return nil, err
+	} else {
+		return &Host{endpoint, dc}, nil
+	}
+>>>>>>> 4a77b5a... Filter endpoints that are not in the current region
+}
+
+func (h *Host) Endpoint() Endpoint {
+	return h.endpoint
+}
+
+func (h *Host) Key() string {
+	return h.endpoint.Key()
+}
+
+func (h *Host) String() string {
+	return h.endpoint.String()
+}

proxycore/lb.go

@@ -54,7 +54,7 @@ func (l *roundRobinLoadBalancer) OnEvent(event Event) {
 	case *RemoveEvent:
 		cpy := l.copy()
 		for i, h := range cpy {
-			if h.Endpoint().Key() == evt.Host.Key() {
+			if h.Key() == evt.Host.Key() {
 				l.hosts.Store(append(cpy[:i], cpy[i+1:]...))
 				break
 			}

proxycore/resultset.go

@@ -16,6 +16,9 @@ package proxycore
 
 import (
 	"errors"
+	"fmt"
+	"net"
+
 	"github.com/datastax/go-cassandra-native-protocol/message"
 	"github.com/datastax/go-cassandra-native-protocol/primitive"
 )
@@ -72,3 +75,39 @@ func (r Row) ByName(n string) (interface{}, error) {
 		return r.ByPos(i)
 	}
 }
+
+func (r Row) StringByName(n string) (string, error) {
+	val, err := r.ByName(n)
+	if err != nil {
+		return "", err
+	}
+	if s, ok := val.(string); !ok {
+		return "", fmt.Errorf("'%s' is not a string", n)
+	} else {
+		return s, nil
+	}
+}
+
+func (r Row) InetByName(n string) (net.IP, error) {
+	val, err := r.ByName(n)
+	if err != nil {
+		return nil, err
+	}
+	if ip, ok := val.(net.IP); !ok {
+		return nil, fmt.Errorf("'%s' is not an inet", n)
+	} else {
+		return ip, nil
+	}
+}
+
+func (r Row) UUIDByName(n string) (primitive.UUID, error) {
+	val, err := r.ByName(n)
+	if err != nil {
+		return [16]byte{}, err
+	}
+	if u, ok := val.(primitive.UUID); !ok {
+		return [16]byte{}, fmt.Errorf("'%s' is not a uuid", n)
+	} else {
+		return u, nil
+	}
+}

proxycore/session.go

@@ -99,7 +99,7 @@ func (s *Session) Send(host *Host, request Request) error {
 }
 
 func (s *Session) leastBusyConn(host *Host) *ClientConn {
-	if p, ok := s.pools.Load(host.Endpoint().Key()); ok {
+	if p, ok := s.pools.Load(host.Key()); ok {
 		pool := p.(*connPool)
 		return pool.leastBusyConn()
 	}
@@ -118,7 +118,7 @@ func (s *Session) OnEvent(event Event) {
 			for _, host := range evt.Hosts {
 				go func(host *Host) {
 					pool, err := connectPool(s.ctx, connPoolConfig{
-						Endpoint:      host.Endpoint(),
+						Endpoint:      host.Endpoint,
 						SessionConfig: s.config,
 					})
 					if err != nil {
@@ -127,7 +127,7 @@ func (s *Session) OnEvent(event Event) {
 						default:
 						}
 					}
-					s.pools.Store(host.Endpoint().Key(), pool)
+					s.pools.Store(host.Key(), pool)
 					wg.Done()
 				}(host)
 			}
@@ -139,7 +139,7 @@ func (s *Session) OnEvent(event Event) {
 	case *AddEvent:
 		// There's no compute if absent for sync.Map, figure a better way to do this if the pool already exists.
 		if pool, loaded := s.pools.LoadOrStore(evt.Host.Key(), connectPoolNoFail(s.ctx, connPoolConfig{
-			Endpoint:      evt.Host.Endpoint(),
+			Endpoint:      evt.Host.Endpoint,
 			SessionConfig: s.config,
 		})); loaded {
 			p := pool.(*connPool)