Repository: gocql/gocql
Pull request #1632 (2022-06-22 19:20:12)

Identify nodes by their host_id instead broadcasted_address
Currently driver identifies nodes based on their broadcasted IP
address. In cloud case, broadcasted IP addresses are private and are not
meant to be used as a contact point, and they may change overtime.
Hence driver internals were changed to identify nodes based on their
host_id which is unique per node and it's persistant throught entire
node lifecycle.
Because CQL Events are still using broadcasted IP addresses, driver
keep mapping between already known IP addresses and host_ids.
Prepared statement cache key was also changed to host_id, to not
invalidate it upon IP change.

zimnx: > We should probably use the host ID also as the key for prepared statement cache.

Good idea, changed the key to HostID.

jameshartig: LGTM 👍 

Patch:
diff --git a/conn.go b/conn.go
index a3b4d41..1c80f3e 100644
--- a/conn.go
+++ b/conn.go
@@ -1251,7 +1251,7 @@ type inflightPrepare struct {
 }
 
 func (c *Conn) prepareStatement(ctx context.Context, stmt string, tracer Tracer) (*preparedStatment, error) {
-	stmtCacheKey := c.session.stmtsLRU.keyFor(c.addr, c.currentKeyspace, stmt)
+	stmtCacheKey := c.session.stmtsLRU.keyFor(c.host.HostID(), c.currentKeyspace, stmt)
 	flight, ok := c.session.stmtsLRU.execIfMissing(stmtCacheKey, func(lru *lru.Cache) *inflightPrepare {
 		flight := &inflightPrepare{
 			done: make(chan struct{}),
@@ -1486,7 +1486,7 @@ func (c *Conn) executeQuery(ctx context.Context, qry *Query) *Iter {
 		// is not consistent with regards to its schema.
 		return iter
 	case *RequestErrUnprepared:
-		stmtCacheKey := c.session.stmtsLRU.keyFor(c.addr, c.currentKeyspace, qry.stmt)
+		stmtCacheKey := c.session.stmtsLRU.keyFor(c.host.HostID(), c.currentKeyspace, qry.stmt)
 		c.session.stmtsLRU.evictPreparedID(stmtCacheKey, x.StatementId)
 		return c.executeQuery(ctx, qry)
 	case error:
@@ -1632,7 +1632,7 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) *Iter {
 	case *RequestErrUnprepared:
 		stmt, found := stmts[string(x.StatementId)]
 		if found {
-			key := c.session.stmtsLRU.keyFor(c.addr, c.currentKeyspace, stmt)
+			key := c.session.stmtsLRU.keyFor(c.host.HostID(), c.currentKeyspace, stmt)
 			c.session.stmtsLRU.evictPreparedID(key, x.StatementId)
 		}
 		return c.executeBatch(ctx, batch)
diff --git a/connectionpool.go b/connectionpool.go
index 5353d80..24f1c7c 100644
--- a/connectionpool.go
+++ b/connectionpool.go
@@ -85,8 +85,6 @@ type policyConnPool struct {
 
 	mu            sync.RWMutex
 	hostConnPools map[string]*hostConnPool
-
-	endpoints []string
 }
 
 func connConfig(cfg *ClusterConfig) (*ConnConfig, error) {
@@ -130,9 +128,6 @@ func newPolicyConnPool(session *Session) *policyConnPool {
 		hostConnPools: map[string]*hostConnPool{},
 	}
 
-	pool.endpoints = make([]string, len(session.cfg.Hosts))
-	copy(pool.endpoints, session.cfg.Hosts)
-
 	return pool
 }
 
@@ -141,8 +136,8 @@ func (p *policyConnPool) SetHosts(hosts []*HostInfo) {
 	defer p.mu.Unlock()
 
 	toRemove := make(map[string]struct{})
-	for addr := range p.hostConnPools {
-		toRemove[addr] = struct{}{}
+	for hostID := range p.hostConnPools {
+		toRemove[hostID] = struct{}{}
 	}
 
 	pools := make(chan *hostConnPool)
@@ -152,10 +147,10 @@ func (p *policyConnPool) SetHosts(hosts []*HostInfo) {
 			// don't create a connection pool for a down host
 			continue
 		}
-		ip := host.ConnectAddress().String()
-		if _, exists := p.hostConnPools[ip]; exists {
+		hostID := host.HostID()
+		if _, exists := p.hostConnPools[hostID]; exists {
 			// still have this host, so don't remove it
-			delete(toRemove, ip)
+			delete(toRemove, hostID)
 			continue
 		}
 
@@ -178,7 +173,7 @@ func (p *policyConnPool) SetHosts(hosts []*HostInfo) {
 		createCount--
 		if pool.Size() > 0 {
 			// add pool only if there a connections available
-			p.hostConnPools[string(pool.host.ConnectAddress())] = pool
+			p.hostConnPools[pool.host.HostID()] = pool
 		}
 	}
 
@@ -201,9 +196,9 @@ func (p *policyConnPool) Size() int {
 }
 
 func (p *policyConnPool) getPool(host *HostInfo) (pool *hostConnPool, ok bool) {
-	ip := host.ConnectAddress().String()
+	hostID := host.HostID()
 	p.mu.RLock()
-	pool, ok = p.hostConnPools[ip]
+	pool, ok = p.hostConnPools[hostID]
 	p.mu.RUnlock()
 	return
 }
@@ -220,9 +215,9 @@ func (p *policyConnPool) Close() {
 }
 
 func (p *policyConnPool) addHost(host *HostInfo) {
-	ip := host.ConnectAddress().String()
+	hostID := host.HostID()
 	p.mu.Lock()
-	pool, ok := p.hostConnPools[ip]
+	pool, ok := p.hostConnPools[hostID]
 	if !ok {
 		pool = newHostConnPool(
 			p.session,
@@ -232,41 +227,33 @@ func (p *policyConnPool) addHost(host *HostInfo) {
 			p.keyspace,
 		)
 
-		p.hostConnPools[ip] = pool
+		p.hostConnPools[hostID] = pool
 	}
 	p.mu.Unlock()
 
 	pool.fill()
 }
 
-func (p *policyConnPool) removeHost(ip net.IP) {
-	k := ip.String()
+func (p *policyConnPool) removeHost(hostID string) {
 	p.mu.Lock()
-	pool, ok := p.hostConnPools[k]
+	pool, ok := p.hostConnPools[hostID]
 	if !ok {
 		p.mu.Unlock()
 		return
 	}
 
-	delete(p.hostConnPools, k)
+	delete(p.hostConnPools, hostID)
 	p.mu.Unlock()
 
 	go pool.Close()
 }
 
-func (p *policyConnPool) hostDown(ip net.IP) {
-	// TODO(zariel): mark host as down so we can try to connect to it later, for
-	// now just treat it has removed.
-	p.removeHost(ip)
-}
-
 // hostConnPool is a connection pool for a single host.
 // Connection selection is based on a provided ConnSelectionPolicy
 type hostConnPool struct {
 	session  *Session
 	host     *HostInfo
 	port     int
-	addr     string
 	size     int
 	keyspace string
 	// protection for conns, closed, filling
@@ -293,7 +280,6 @@ func newHostConnPool(session *Session, host *HostInfo, port, size int,
 		session:  session,
 		host:     host,
 		port:     port,
-		addr:     (&net.TCPAddr{IP: host.ConnectAddress(), Port: host.Port()}).String(),
 		size:     size,
 		keyspace: keyspace,
 		conns:    make([]*Conn, 0, size),
@@ -461,11 +447,11 @@ func (pool *hostConnPool) logConnectErr(err error) {
 		// connection refused
 		// these are typical during a node outage so avoid log spam.
 		if gocqlDebug {
-			pool.logger.Printf("gocql: unable to dial %q: %v\n", pool.host.ConnectAddress(), err)
+			pool.logger.Printf("gocql: unable to dial %q: %v\n", pool.host, err)
 		}
 	} else if err != nil {
 		// unexpected error
-		pool.logger.Printf("gocql: error: failed to connect to %s due to error: %v", pool.addr, err)
+		pool.logger.Printf("error: failed to connect to %q due to error: %v", pool.host, err)
 	}
 }
 
diff --git a/events.go b/events.go
index 93ba085..086bed8 100644
--- a/events.go
+++ b/events.go
@@ -190,11 +190,11 @@ func (s *Session) handleNodeEvent(frames []frame) {
 	}
 }
 
-func (s *Session) addNewNode(ip net.IP, port int) {
+func (s *Session) addNewNode(hostID UUID) {
 	// Get host info and apply any filters to the host
-	hostInfo, err := s.hostSource.getHostInfo(ip, port)
+	hostInfo, err := s.hostSource.getHostInfo(hostID)
 	if err != nil {
-		s.logger.Printf("gocql: events: unable to fetch host info for (%s:%d): %v\n", ip, port, err)
+		s.logger.Printf("gocql: events: unable to fetch host info for hostID: %q: %v\n", hostID, err)
 		return
 	} else if hostInfo == nil {
 		// ignore if it's null because we couldn't find it
@@ -223,15 +223,14 @@ func (s *Session) handleNewNode(ip net.IP, port int) {
 		s.logger.Printf("gocql: Session.handleNewNode: %s:%d\n", ip.String(), port)
 	}
 
-	ip, port = s.cfg.translateAddressPort(ip, port)
-
-	// if we already have the host and it's already up, then do nothing
-	host := s.ring.getHost(ip)
-	if host != nil && host.IsUp() {
+	host, ok := s.ring.getHostByIP(ip.String())
+	if ok && host.IsUp() {
 		return
 	}
 
-	s.addNewNode(ip, port)
+	if err := s.hostSource.refreshRing(); err != nil && gocqlDebug {
+		s.logger.Printf("gocql: Session.handleNewNode: failed to refresh ring: %w\n", err.Error())
+	}
 }
 
 func (s *Session) handleRemovedNode(ip net.IP, port int) {
@@ -239,23 +238,22 @@ func (s *Session) handleRemovedNode(ip net.IP, port int) {
 		s.logger.Printf("gocql: Session.handleRemovedNode: %s:%d\n", ip.String(), port)
 	}
 
-	ip, port = s.cfg.translateAddressPort(ip, port)
-
 	// we remove all nodes but only add ones which pass the filter
-	host := s.ring.getHost(ip)
-	if host == nil {
-		host = &HostInfo{connectAddress: ip, port: port}
-	}
-	s.ring.removeHost(ip)
+	host, ok := s.ring.getHostByIP(ip.String())
+	hostID := host.HostID()
+	if ok {
+		s.ring.removeHost(hostID)
+
+		host.setState(NodeDown)
+		if !s.cfg.filterHost(host) {
+			s.policy.RemoveHost(host)
+			s.pool.removeHost(hostID)
+		}
 
-	host.setState(NodeDown)
-	if !s.cfg.filterHost(host) {
-		s.policy.RemoveHost(host)
-		s.pool.removeHost(ip)
 	}
 
-	if !s.cfg.IgnorePeerAddr {
-		s.hostSource.refreshRing()
+	if err := s.hostSource.refreshRing(); err != nil && gocqlDebug {
+		s.logger.Println("failed to refresh ring:", err)
 	}
 }
 
@@ -264,11 +262,8 @@ func (s *Session) handleNodeUp(eventIp net.IP, eventPort int) {
 		s.logger.Printf("gocql: Session.handleNodeUp: %s:%d\n", eventIp.String(), eventPort)
 	}
 
-	ip, port := s.cfg.translateAddressPort(eventIp, eventPort)
-
-	host := s.ring.getHost(ip)
-	if host == nil {
-		s.addNewNode(ip, port)
+	host, ok := s.ring.getHostByIP(eventIp.String())
+	if !ok {
 		return
 	}
 
@@ -305,16 +300,15 @@ func (s *Session) handleNodeDown(ip net.IP, port int) {
 		s.logger.Printf("gocql: Session.handleNodeDown: %s:%d\n", ip.String(), port)
 	}
 
-	host := s.ring.getHost(ip)
-	if host == nil {
-		host = &HostInfo{connectAddress: ip, port: port}
-	}
+	host, ok := s.ring.getHostByIP(ip.String())
+	if ok {
+		host.setState(NodeDown)
+		if s.cfg.filterHost(host) {
+			return
+		}
 
-	host.setState(NodeDown)
-	if s.cfg.filterHost(host) {
-		return
+		s.policy.HostDown(host)
+		hostID := host.HostID()
+		s.pool.removeHost(hostID)
 	}
-
-	s.policy.HostDown(host)
-	s.pool.hostDown(ip)
 }
diff --git a/host_source.go b/host_source.go
index 70cabc4..a4d98bb 100644
--- a/host_source.go
+++ b/host_source.go
@@ -174,6 +174,22 @@ func (h *HostInfo) connectAddressLocked() (net.IP, string) {
 	return net.IPv4zero, "invalid"
 }
 
+// nodeToNodeAddress returns address broadcasted between node to nodes.
+// It's either `broadcast_address` if host info is read from system.local or `peer` if read from system.peers.
+// This IP address is also part of CQL Event emitted on topology/status changes,
+// but does not uniquely identify the node in case multiple nodes use the same IP address.
+func (h *HostInfo) nodeToNodeAddress() net.IP {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	if validIpAddr(h.broadcastAddress) {
+		return h.broadcastAddress
+	} else if validIpAddr(h.peer) {
+		return h.peer
+	}
+	return net.IPv4zero
+}
+
 // Returns the address that should be used to connect to the host.
 // If you wish to override this, use an AddressTranslator or
 // use a HostFilter to SetConnectAddress()
@@ -239,6 +255,12 @@ func (h *HostInfo) HostID() string {
 	return h.hostId
 }
 
+func (h *HostInfo) SetHostID(hostID string) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.hostId = hostID
+}
+
 func (h *HostInfo) WorkLoad() string {
 	h.mu.RLock()
 	defer h.mu.RUnlock()
@@ -583,44 +605,40 @@ func (r *ringDescriber) GetHosts() ([]*HostInfo, string, error) {
 }
 
 // Given an ip/port return HostInfo for the specified ip/port
-func (r *ringDescriber) getHostInfo(ip net.IP, port int) (*HostInfo, error) {
+func (r *ringDescriber) getHostInfo(hostID UUID) (*HostInfo, error) {
 	var host *HostInfo
-	iter := r.session.control.withConnHost(func(ch *connHost) *Iter {
-		if ch.host.ConnectAddress().Equal(ip) {
-			host = ch.host
-			return nil
-		}
+	for _, table := range []string{"system.peers", "system.local"} {
+		iter := r.session.control.withConnHost(func(ch *connHost) *Iter {
+			if ch.host.HostID() == hostID.String() {
+				host = ch.host
+				return nil
+			}
 
-		return ch.conn.query(context.TODO(), "SELECT * FROM system.peers")
-	})
+			return ch.conn.query(context.TODO(), fmt.Sprintf("SELECT * FROM %s", table))
+		})
 
-	if iter != nil {
-		rows, err := iter.SliceMap()
-		if err != nil {
-			return nil, err
-		}
-
-		for _, row := range rows {
-			h, err := r.session.hostInfoFromMap(row, &HostInfo{port: port})
+		if iter != nil {
+			rows, err := iter.SliceMap()
 			if err != nil {
 				return nil, err
 			}
 
-			if h.ConnectAddress().Equal(ip) {
-				host = h
-				break
-			}
-		}
+			for _, row := range rows {
+				h, err := r.session.hostInfoFromMap(row, &HostInfo{port: r.session.cfg.Port})
+				if err != nil {
+					return nil, err
+				}
 
-		if host == nil {
-			return nil, errors.New("host not found in peers table")
+				if h.HostID() == hostID.String() {
+					host = h
+					break
+				}
+			}
 		}
 	}
 
 	if host == nil {
 		return nil, errors.New("unable to fetch host info: invalid control connection")
-	} else if host.invalidConnectAddr() {
-		return nil, fmt.Errorf("host ConnectAddress invalid ip=%v: %v", ip, host)
 	}
 
 	return host, nil
@@ -649,7 +667,7 @@ func (r *ringDescriber) refreshRing() error {
 		} else {
 			host.update(h)
 		}
-		delete(prevHosts, h.ConnectAddress().String())
+		delete(prevHosts, h.HostID())
 	}
 
 	// TODO(zariel): it may be worth having a mutex covering the overall ring state
diff --git a/prepared_cache.go b/prepared_cache.go
index 346f36f..bb7d697 100644
--- a/prepared_cache.go
+++ b/prepared_cache.go
@@ -47,9 +47,9 @@ func (p *preparedLRU) execIfMissing(key string, fn func(lru *lru.Cache) *infligh
 	return fn(p.lru), false
 }
 
-func (p *preparedLRU) keyFor(addr, keyspace, statement string) string {
+func (p *preparedLRU) keyFor(hostID, keyspace, statement string) string {
 	// TODO: we should just use a struct for the key in the map
-	return addr + keyspace + statement
+	return hostID + keyspace + statement
 }
 
 func (p *preparedLRU) evictPreparedID(key string, id []byte) {
diff --git a/ring.go b/ring.go
index c6e0034..5b77370 100644
--- a/ring.go
+++ b/ring.go
@@ -2,7 +2,6 @@ package gocql
 
 import (
 	"fmt"
-	"net"
 	"sync"
 	"sync/atomic"
 )
@@ -13,9 +12,12 @@ type ring struct {
 	// strap the initial connection.
 	endpoints []*HostInfo
 
-	// hosts are the set of all hosts in the cassandra ring that we know of
-	mu    sync.RWMutex
+	mu sync.RWMutex
+	// hosts are the set of all hosts in the cassandra ring that we know of.
+	// key of map is host_id.
 	hosts map[string]*HostInfo
+	// hostIPToUUID maps host native address to host_id.
+	hostIPToUUID map[string]string
 
 	hostList []*HostInfo
 	pos      uint32
@@ -34,9 +36,16 @@ func (r *ring) rrHost() *HostInfo {
 	return r.hostList[pos%len(r.hostList)]
 }
 
-func (r *ring) getHost(ip net.IP) *HostInfo {
+func (r *ring) getHostByIP(ip string) (*HostInfo, bool) {
 	r.mu.RLock()
-	host := r.hosts[ip.String()]
+	defer r.mu.RUnlock()
+	hi, ok := r.hostIPToUUID[ip]
+	return r.hosts[hi], ok
+}
+
+func (r *ring) getHost(hostID string) *HostInfo {
+	r.mu.RLock()
+	host := r.hosts[hostID]
 	r.mu.RUnlock()
 	return host
 }
@@ -73,16 +82,20 @@ func (r *ring) addHostIfMissing(host *HostInfo) (*HostInfo, bool) {
 	if host.invalidConnectAddr() {
 		panic(fmt.Sprintf("invalid host: %v", host))
 	}
-	ip := host.ConnectAddress().String()
+	hostID := host.HostID()
 
 	r.mu.Lock()
 	if r.hosts == nil {
 		r.hosts = make(map[string]*HostInfo)
 	}
+	if r.hostIPToUUID == nil {
+		r.hostIPToUUID = make(map[string]string)
+	}
 
-	existing, ok := r.hosts[ip]
+	existing, ok := r.hosts[hostID]
 	if !ok {
-		r.hosts[ip] = host
+		r.hosts[hostID] = host
+		r.hostIPToUUID[host.nodeToNodeAddress().String()] = hostID
 		existing = host
 		r.hostList = append(r.hostList, host)
 	}
@@ -90,23 +103,26 @@ func (r *ring) addHostIfMissing(host *HostInfo) (*HostInfo, bool) {
 	return existing, ok
 }
 
-func (r *ring) removeHost(ip net.IP) bool {
+func (r *ring) removeHost(hostID string) bool {
 	r.mu.Lock()
 	if r.hosts == nil {
 		r.hosts = make(map[string]*HostInfo)
 	}
+	if r.hostIPToUUID == nil {
+		r.hostIPToUUID = make(map[string]string)
+	}
 
-	k := ip.String()
-	_, ok := r.hosts[k]
+	h, ok := r.hosts[hostID]
 	if ok {
 		for i, host := range r.hostList {
-			if host.ConnectAddress().Equal(ip) {
+			if host.HostID() == hostID {
 				r.hostList = append(r.hostList[:i], r.hostList[i+1:]...)
 				break
 			}
 		}
+		delete(r.hostIPToUUID, h.nodeToNodeAddress().String())
 	}
-	delete(r.hosts, k)
+	delete(r.hosts, hostID)
 	r.mu.Unlock()
 	return ok
 }
diff --git a/session.go b/session.go
index 3afdc6e..fc44abd 100644
--- a/session.go
+++ b/session.go
@@ -226,13 +226,24 @@ func (s *Session) init() error {
 					filteredHosts = append(filteredHosts, host)
 				}
 			}
-			hosts = append(hosts, filteredHosts...)
+
+			hosts = filteredHosts
+		}
+	}
+
+	for _, host := range hosts {
+		// In case when host lookup is disabled and when we are in unit tests,
+		// host are not discovered, and we are missing host ID information used
+		// by internal logic.
+		// Associate random UUIDs here with all hosts missing this information.
+		if len(host.HostID()) == 0 {
+			host.SetHostID(MustRandomUUID().String())
 		}
 	}
 
 	hostMap := make(map[string]*HostInfo, len(hosts))
 	for _, host := range hosts {
-		hostMap[host.ConnectAddress().String()] = host
+		hostMap[host.HostID()] = host
 	}
 
 	hosts = hosts[:0]
@@ -512,8 +523,9 @@ func (s *Session) executeQuery(qry *Query) (it *Iter) {
 
 func (s *Session) removeHost(h *HostInfo) {
 	s.policy.RemoveHost(h)
-	s.pool.removeHost(h.ConnectAddress())
-	s.ring.removeHost(h.ConnectAddress())
+	hostID := h.HostID()
+	s.pool.removeHost(hostID)
+	s.ring.removeHost(hostID)
 }
 
 // KeyspaceMetadata returns the schema metadata for the keyspace specified. Returns an error if the keyspace does not exist.
diff --git a/uuid.go b/uuid.go
index 09727a0..acdd81f 100644
--- a/uuid.go
+++ b/uuid.go
@@ -96,6 +96,14 @@ func UUIDFromBytes(input []byte) (UUID, error) {
 	return u, nil
 }
 
+func MustRandomUUID() UUID {
+	uuid, err := RandomUUID()
+	if err != nil {
+		panic(err)
+	}
+	return uuid
+}
+
 // RandomUUID generates a totally random UUID (version 4) as described in
 // RFC 4122.
 func RandomUUID() (UUID, error) {


Test patch:
diff --git a/cassandra_test.go b/cassandra_test.go
index 8dc4a1f..501f3e9 100644
--- a/cassandra_test.go
+++ b/cassandra_test.go
@@ -1,3 +1,4 @@
+//go:build all || cassandra
 // +build all cassandra
 
 package gocql
@@ -1371,7 +1372,7 @@ func injectInvalidPreparedStatement(t *testing.T, session *Session, table string
 	conn := getRandomConn(t, session)
 
 	flight := new(inflightPrepare)
-	key := session.stmtsLRU.keyFor(conn.addr, "", stmt)
+	key := session.stmtsLRU.keyFor(conn.host.HostID(), "", stmt)
 	session.stmtsLRU.add(key, flight)
 
 	flight.preparedStatment = &preparedStatment{
@@ -1538,28 +1539,28 @@ func TestPrepare_PreparedCacheEviction(t *testing.T) {
 	}
 
 	// Walk through all the configured hosts and test cache retention and eviction
-	for _, host := range session.cfg.Hosts {
-		_, ok := session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host+":9042", session.cfg.Keyspace, "SELECT id,mod FROM prepcachetest WHERE id = 0"))
+	for _, host := range session.ring.hosts {
+		_, ok := session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host.HostID(), session.cfg.Keyspace, "SELECT id,mod FROM prepcachetest WHERE id = 0"))
 		if ok {
 			t.Errorf("expected first select to be purged but was in cache for host=%q", host)
 		}
 
-		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host+":9042", session.cfg.Keyspace, "SELECT id,mod FROM prepcachetest WHERE id = 1"))
+		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host.HostID(), session.cfg.Keyspace, "SELECT id,mod FROM prepcachetest WHERE id = 1"))
 		if !ok {
 			t.Errorf("exepected second select to be in cache for host=%q", host)
 		}
 
-		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host+":9042", session.cfg.Keyspace, "INSERT INTO prepcachetest (id,mod) VALUES (?, ?)"))
+		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host.HostID(), session.cfg.Keyspace, "INSERT INTO prepcachetest (id,mod) VALUES (?, ?)"))
 		if !ok {
 			t.Errorf("expected insert to be in cache for host=%q", host)
 		}
 
-		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host+":9042", session.cfg.Keyspace, "UPDATE prepcachetest SET mod = ? WHERE id = ?"))
+		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host.HostID(), session.cfg.Keyspace, "UPDATE prepcachetest SET mod = ? WHERE id = ?"))
 		if !ok {
 			t.Errorf("expected update to be in cached for host=%q", host)
 		}
 
-		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host+":9042", session.cfg.Keyspace, "DELETE FROM prepcachetest WHERE id = ?"))
+		_, ok = session.stmtsLRU.lru.Get(session.stmtsLRU.keyFor(host.HostID(), session.cfg.Keyspace, "DELETE FROM prepcachetest WHERE id = ?"))
 		if !ok {
 			t.Errorf("expected delete to be cached for host=%q", host)
 		}
@@ -3001,7 +3002,15 @@ func TestDiscoverViaProxy(t *testing.T) {
 
 	session.pool.mu.RLock()
 	for _, host := range clusterHosts {
-		if _, ok := session.pool.hostConnPools[host]; !ok {
+		found := false
+		for _, hi := range session.pool.hostConnPools {
+			if hi.host.ConnectAddress().String() == host {
+				found = true
+				break
+			}
+		}
+
+		if !found {
 			t.Errorf("missing host in pool after discovery: %q", host)
 		}
 	}
diff --git a/ring_test.go b/ring_test.go
index 88e593a..e32728e 100644
--- a/ring_test.go
+++ b/ring_test.go
@@ -8,7 +8,7 @@ import (
 func TestRing_AddHostIfMissing_Missing(t *testing.T) {
 	ring := &ring{}
 
-	host := &HostInfo{connectAddress: net.IPv4(1, 1, 1, 1)}
+	host := &HostInfo{hostId: MustRandomUUID().String(), connectAddress: net.IPv4(1, 1, 1, 1)}
 	h1, ok := ring.addHostIfMissing(host)
 	if ok {
 		t.Fatal("host was reported as already existing")
@@ -22,10 +22,10 @@ func TestRing_AddHostIfMissing_Missing(t *testing.T) {
 func TestRing_AddHostIfMissing_Existing(t *testing.T) {
 	ring := &ring{}
 
-	host := &HostInfo{connectAddress: net.IPv4(1, 1, 1, 1)}
+	host := &HostInfo{hostId: MustRandomUUID().String(), connectAddress: net.IPv4(1, 1, 1, 1)}
 	ring.addHostIfMissing(host)
 
-	h2 := &HostInfo{connectAddress: net.IPv4(1, 1, 1, 1)}
+	h2 := &HostInfo{hostId: host.hostId, connectAddress: net.IPv4(2, 2, 2, 2)}
 
 	h1, ok := ring.addHostIfMissing(h2)
 	if !ok {