Repository: gocql/gocql
Pull request #1659 (2022-10-13 20:40:35)

Rack aware host selection policy
I wrote 90% of this code, then realized that there's already a somewhat stale PR for the same issue open as #1565.

The rationale behind this PRs is the same: cassandra running on public clouds tend to use what AWS calls "region" as cassandra datacenter, and "availability zone" as rack.  A common desire is to keep traffic with-in the current AZ whenever possible from both a performance and cost standpoint.

This PR adds a `HostTierer` interface.  This is extremely similar to the `HostDistancer` interface suggested in the referenced PR, except that I feel like Tier more accurately reflects that we're grouping hosts into bins, not simply coming up with a cost metric (e.g. distance = ping RTT).

The `TokenAwareHostPolicy` now support an arbitrary number of tiers.  It will check to see if the fallback `HostSelectionPolicy` implements `HostTierer`, and if so, will use the tier values it provides for each host.  Otherwise, it assumes two tiers, the first for hosts where `IsLocal` returns `true`, and the second for the rest.

Tests and docs are included.  Please let me know if there's anything else I can provide that would help this PR get merged.


Patch:
diff --git a/AUTHORS b/AUTHORS
index 50e173a..04a97c9 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -131,3 +131,4 @@ Sergei Karetnikov <sergei.karetnikov@gmail.com>
 Stefan Miklosovic <smiklosovic@apache.org>
 Adam Burk <amburk@gmail.com>
 Valerii Ponomarov <kiparis.kh@gmail.com>
+Neal Turett <neal.turett@datadoghq.com>
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 90ae1d9..c134445 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added a RackAwareRoundRobinPolicy that attempts to keep client->server traffic in the same rack when possible.
+
 ### Changed
 
 - Supported versions of Go that we test against are now Go 1.18 and Go 1.19.
diff --git a/doc.go b/doc.go
index 489d2a3..b48cef8 100644
--- a/doc.go
+++ b/doc.go
@@ -115,6 +115,15 @@
 //
 //	session.Query("select value from mytable where pk1 = ? AND pk2 = ?", "abc", "def")
 //
+// # Rack-level awareness
+//
+// The DCAwareRoundRobinPolicy can be replaced with RackAwareRoundRobinPolicy, which takes two parameters, datacenter and rack.
+//
+// Instead of dividing hosts with two tiers (local datacenter and remote datacenters) it divides hosts into three
+// (the local rack, the rest of the local datacenter, and everything else).
+//
+// RackAwareRoundRobinPolicy can be combined with TokenAwareHostPolicy in the same way as DCAwareRoundRobinPolicy.
+//
 // # Executing queries
 //
 // Create queries with Session.Query. Query values must not be reused between different executions and must not be
diff --git a/policies.go b/policies.go
index 11bb180..f31fc70 100644
--- a/policies.go
+++ b/policies.go
@@ -259,6 +259,20 @@ type KeyspaceUpdateEvent struct {
 	Change   string
 }
 
+type HostTierer interface {
+	// HostTier returns an integer specifying how far a host is from the client.
+	// Tier must start at 0.
+	// The value is used to prioritize closer hosts during host selection.
+	// For example this could be:
+	// 0 - local rack, 1 - local DC, 2 - remote DC
+	// or:
+	// 0 - local DC, 1 - remote DC
+	HostTier(host *HostInfo) uint
+
+	// This function returns the maximum possible host tier
+	MaxHostTier() uint
+}
+
 // HostSelectionPolicy is an interface for selecting
 // the most appropriate host to execute a given query.
 // HostSelectionPolicy instances cannot be shared between sessions.
@@ -577,18 +591,42 @@ func (t *tokenAwareHostPolicy) Pick(qry ExecutableQuery) NextHost {
 
 	var (
 		fallbackIter NextHost
-		i, j         int
-		remote       []*HostInfo
+		i, j, k      int
+		remote       [][]*HostInfo
+		tierer       HostTierer
+		tiererOk     bool
+		maxTier      uint
 	)
 
+	if tierer, tiererOk = t.fallback.(HostTierer); tiererOk {
+		maxTier = tierer.MaxHostTier()
+	} else {
+		maxTier = 1
+	}
+
+	if t.nonLocalReplicasFallback {
+		remote = make([][]*HostInfo, maxTier)
+	}
+
 	used := make(map[*HostInfo]bool, len(replicas))
 	return func() SelectedHost {
 		for i < len(replicas) {
 			h := replicas[i]
 			i++
 
-			if !t.fallback.IsLocal(h) {
-				remote = append(remote, h)
+			var tier uint
+			if tiererOk {
+				tier = tierer.HostTier(h)
+			} else if t.fallback.IsLocal(h) {
+				tier = 0
+			} else {
+				tier = 1
+			}
+
+			if tier != 0 {
+				if t.nonLocalReplicasFallback {
+					remote[tier-1] = append(remote[tier-1], h)
+				}
 				continue
 			}
 
@@ -599,9 +637,14 @@ func (t *tokenAwareHostPolicy) Pick(qry ExecutableQuery) NextHost {
 		}
 
 		if t.nonLocalReplicasFallback {
-			for j < len(remote) {
-				h := remote[j]
-				j++
+			for j < len(remote) && k < len(remote[j]) {
+				h := remote[j][k]
+				k++
+
+				if k >= len(remote[j]) {
+					j++
+					k = 0
+				}
 
 				if h.IsUp() {
 					used[h] = true
@@ -856,6 +899,64 @@ func (d *dcAwareRR) Pick(q ExecutableQuery) NextHost {
 	return roundRobbin(int(nextStartOffset), d.localHosts.get(), d.remoteHosts.get())
 }
 
+// RackAwareRoundRobinPolicy is a host selection policies which will prioritize and
+// return hosts which are in the local rack, before hosts in the local datacenter but
+// a different rack, before hosts in all other datercentres
+
+type rackAwareRR struct {
+	localDC         string
+	localRack       string
+	hosts           []cowHostList
+	lastUsedHostIdx uint64
+}
+
+func RackAwareRoundRobinPolicy(localDC string, localRack string) HostSelectionPolicy {
+	hosts := make([]cowHostList, 3)
+	return &rackAwareRR{localDC: localDC, localRack: localRack, hosts: hosts}
+}
+
+func (d *rackAwareRR) Init(*Session)                       {}
+func (d *rackAwareRR) KeyspaceChanged(KeyspaceUpdateEvent) {}
+func (d *rackAwareRR) SetPartitioner(p string)             {}
+
+func (d *rackAwareRR) MaxHostTier() uint {
+	return 2
+}
+
+func (d *rackAwareRR) HostTier(host *HostInfo) uint {
+	if host.DataCenter() == d.localDC {
+		if host.Rack() == d.localRack {
+			return 0
+		} else {
+			return 1
+		}
+	} else {
+		return 2
+	}
+}
+
+func (d *rackAwareRR) IsLocal(host *HostInfo) bool {
+	return d.HostTier(host) == 0
+}
+
+func (d *rackAwareRR) AddHost(host *HostInfo) {
+	dist := d.HostTier(host)
+	d.hosts[dist].add(host)
+}
+
+func (d *rackAwareRR) RemoveHost(host *HostInfo) {
+	dist := d.HostTier(host)
+	d.hosts[dist].remove(host.ConnectAddress())
+}
+
+func (d *rackAwareRR) HostUp(host *HostInfo)   { d.AddHost(host) }
+func (d *rackAwareRR) HostDown(host *HostInfo) { d.RemoveHost(host) }
+
+func (d *rackAwareRR) Pick(q ExecutableQuery) NextHost {
+	nextStartOffset := atomic.AddUint64(&d.lastUsedHostIdx, 1)
+	return roundRobbin(int(nextStartOffset), d.hosts[0].get(), d.hosts[1].get(), d.hosts[2].get())
+}
+
 // ReadyPolicy defines a policy for when a HostSelectionPolicy can be used. After
 // each host connects during session initialization, the Ready method will be
 // called. If you only need a single Host to be up you can wrap a


Test patch:
diff --git a/conn_test.go b/conn_test.go
index 845ce52..916a3d1 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -241,7 +241,7 @@ func TestStartupTimeout(t *testing.T) {
 		t.Fatalf("Expected to receive no connections error - got '%s'", err)
 	}
 
-	if !strings.Contains(log.String(), "no response to connection startup within timeout") {
+	if !strings.Contains(log.String(), "no response to connection startup within timeout") && !strings.Contains(log.String(), "no response received from cassandra within timeout period") {
 		t.Fatalf("Expected to receive timeout log message  - got '%s'", log.String())
 	}
 
diff --git a/policies_test.go b/policies_test.go
index 826b2cd..119b027 100644
--- a/policies_test.go
+++ b/policies_test.go
@@ -8,6 +8,8 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"sort"
+	"strings"
 	"testing"
 	"time"
 
@@ -106,8 +108,8 @@ func TestHostPolicy_TokenAware_SimpleStrategy(t *testing.T) {
 	// now the token ring is configured
 	query.RoutingKey([]byte("20"))
 	iter = policy.Pick(query)
-	iterCheck(t, iter, "1")
-	iterCheck(t, iter, "2")
+	expectHosts(t, "hosts[0]", iter, "1")
+	expectHosts(t, "hosts[1]", iter, "2")
 }
 
 // Tests of the host pool host selection policy implementation
@@ -368,15 +370,34 @@ func TestDowngradingConsistencyRetryPolicy(t *testing.T) {
 	}
 }
 
-func iterCheck(t *testing.T, iter NextHost, hostID string) {
+// expectHosts makes sure that the next len(hostIDs) returned from iter is a permutation of hostIDs.
+func expectHosts(t *testing.T, msg string, iter NextHost, hostIDs ...string) {
 	t.Helper()
 
-	host := iter()
-	if host == nil || host.Info() == nil {
-		t.Fatalf("expected hostID %s got nil", hostID)
+	expectedHostIDs := make(map[string]struct{}, len(hostIDs))
+	for i := range hostIDs {
+		expectedHostIDs[hostIDs[i]] = struct{}{}
 	}
-	if host.Info().HostID() != hostID {
-		t.Fatalf("Expected peer %s but was %s", hostID, host.Info().HostID())
+
+	expectedStr := func() string {
+		keys := make([]string, 0, len(expectedHostIDs))
+		for k := range expectedHostIDs {
+			keys = append(keys, k)
+		}
+		sort.Strings(keys)
+		return strings.Join(keys, ", ")
+	}
+
+	for len(expectedHostIDs) > 0 {
+		host := iter()
+		if host == nil || host.Info() == nil {
+			t.Fatalf("%s: expected hostID one of {%s}, but got nil", msg, expectedStr())
+		}
+		hostID := host.Info().HostID()
+		if _, ok := expectedHostIDs[hostID]; !ok {
+			t.Fatalf("%s: expected host ID one of {%s}, but got %s", msg, expectedStr(), hostID)
+		}
+		delete(expectedHostIDs, hostID)
 	}
 }
 
@@ -522,7 +543,7 @@ func TestHostPolicy_TokenAware(t *testing.T) {
 	query.RoutingKey([]byte("23"))
 	iter = policy.Pick(query)
 	// first should be host with matching token from the local DC
-	iterCheck(t, iter, "4")
+	expectHosts(t, "matching token from local DC", iter, "4")
 	// next are in non deterministic order
 }
 
@@ -611,11 +632,207 @@ func TestHostPolicy_TokenAware_NetworkStrategy(t *testing.T) {
 	query.RoutingKey([]byte("18"))
 	iter = policy.Pick(query)
 	// first should be hosts with matching token from the local DC
-	iterCheck(t, iter, "4")
-	iterCheck(t, iter, "7")
+	expectHosts(t, "matching token from local DC", iter, "4", "7")
 	// rest should be hosts with matching token from remote DCs
-	iterCheck(t, iter, "3")
-	iterCheck(t, iter, "5")
-	iterCheck(t, iter, "6")
-	iterCheck(t, iter, "8")
+	expectHosts(t, "matching token from remote DCs", iter, "3", "5", "6", "8")
+}
+
+// Check that each element of toCheck is in valid, and no element appears twice
+func checkList(t *testing.T, msg string, toCheck []string, valid []string) {
+	data := make(map[string]bool)
+	for _, val := range valid {
+		data[val] = false
+	}
+
+	for _, val := range toCheck {
+		seen, exists := data[val]
+		if !exists || seen {
+			t.Fatal(msg)
+		}
+		data[val] = true
+	}
+}
+
+func TestHostPolicy_RackAwareRR(t *testing.T) {
+	p := RackAwareRoundRobinPolicy("local", "b")
+
+	hosts := [...]*HostInfo{
+		{hostId: "0", connectAddress: net.ParseIP("10.0.0.1"), dataCenter: "local", rack: "a"},
+		{hostId: "1", connectAddress: net.ParseIP("10.0.0.2"), dataCenter: "local", rack: "a"},
+		{hostId: "2", connectAddress: net.ParseIP("10.0.0.3"), dataCenter: "local", rack: "b"},
+		{hostId: "3", connectAddress: net.ParseIP("10.0.0.4"), dataCenter: "local", rack: "b"},
+		{hostId: "4", connectAddress: net.ParseIP("10.0.0.5"), dataCenter: "remote", rack: "a"},
+		{hostId: "5", connectAddress: net.ParseIP("10.0.0.6"), dataCenter: "remote", rack: "a"},
+		{hostId: "6", connectAddress: net.ParseIP("10.0.0.7"), dataCenter: "remote", rack: "b"},
+		{hostId: "7", connectAddress: net.ParseIP("10.0.0.8"), dataCenter: "remote", rack: "b"},
+	}
+
+	hostsOut := make([]*HostInfo, 0, len(hosts))
+
+	for _, host := range hosts {
+		p.AddHost(host)
+	}
+
+	it := p.Pick(nil)
+	for h := it(); h != nil; h = it() {
+		hostsOut = append(hostsOut, h.Info())
+	}
+
+	if len(hostsOut) != len(hosts) {
+		t.Fatalf("expected %d hosts got %d", len(hosts), len(hostsOut))
+	}
+
+	// Must start with rack-local hosts
+	checkList(
+		t,
+		"did not return correct rack-local hosts",
+		[]string{hostsOut[0].hostId, hostsOut[1].hostId},
+		[]string{"3", "2"},
+	)
+	// Then dc-local hosts
+	checkList(
+		t,
+		"did not return correct rack-local hosts",
+		[]string{hostsOut[2].hostId, hostsOut[3].hostId},
+		[]string{"0", "1"},
+	)
+	// Then the remote hosts
+	checkList(
+		t,
+		"did not return correct rack-local hosts",
+		[]string{hostsOut[4].hostId, hostsOut[5].hostId, hostsOut[6].hostId, hostsOut[7].hostId},
+		[]string{"4", "5", "6", "7"},
+	)
+}
+
+// Tests of the token-aware host selection policy implementation with a
+// DC & Rack aware round-robin host selection policy fallback
+func TestHostPolicy_TokenAware_RackAware(t *testing.T) {
+	const keyspace = "myKeyspace"
+	policy := TokenAwareHostPolicy(RackAwareRoundRobinPolicy("local", "b"))
+	policyWithFallback := TokenAwareHostPolicy(RackAwareRoundRobinPolicy("local", "b"), NonLocalReplicasFallback())
+
+	policyInternal := policy.(*tokenAwareHostPolicy)
+	policyInternal.getKeyspaceName = func() string { return keyspace }
+	policyInternal.getKeyspaceMetadata = func(ks string) (*KeyspaceMetadata, error) {
+		return nil, errors.New("not initialized")
+	}
+
+	policyWithFallbackInternal := policyWithFallback.(*tokenAwareHostPolicy)
+	policyWithFallbackInternal.getKeyspaceName = policyInternal.getKeyspaceName
+	policyWithFallbackInternal.getKeyspaceMetadata = policyInternal.getKeyspaceMetadata
+
+	query := &Query{}
+	query.getKeyspace = func() string { return keyspace }
+
+	iter := policy.Pick(nil)
+	if iter == nil {
+		t.Fatal("host iterator was nil")
+	}
+	actual := iter()
+	if actual != nil {
+		t.Fatalf("expected nil from iterator, but was %v", actual)
+	}
+
+	// set the hosts
+	hosts := [...]*HostInfo{
+		{hostId: "0", connectAddress: net.IPv4(10, 0, 0, 1), tokens: []string{"05"}, dataCenter: "remote", rack: "a"},
+		{hostId: "1", connectAddress: net.IPv4(10, 0, 0, 2), tokens: []string{"10"}, dataCenter: "remote", rack: "b"},
+		{hostId: "2", connectAddress: net.IPv4(10, 0, 0, 3), tokens: []string{"15"}, dataCenter: "local", rack: "a"},
+		{hostId: "3", connectAddress: net.IPv4(10, 0, 0, 4), tokens: []string{"20"}, dataCenter: "local", rack: "b"},
+		{hostId: "4", connectAddress: net.IPv4(10, 0, 0, 5), tokens: []string{"25"}, dataCenter: "remote", rack: "a"},
+		{hostId: "5", connectAddress: net.IPv4(10, 0, 0, 6), tokens: []string{"30"}, dataCenter: "remote", rack: "b"},
+		{hostId: "6", connectAddress: net.IPv4(10, 0, 0, 7), tokens: []string{"35"}, dataCenter: "local", rack: "a"},
+		{hostId: "7", connectAddress: net.IPv4(10, 0, 0, 8), tokens: []string{"40"}, dataCenter: "local", rack: "b"},
+		{hostId: "8", connectAddress: net.IPv4(10, 0, 0, 9), tokens: []string{"45"}, dataCenter: "remote", rack: "a"},
+		{hostId: "9", connectAddress: net.IPv4(10, 0, 0, 10), tokens: []string{"50"}, dataCenter: "remote", rack: "b"},
+		{hostId: "10", connectAddress: net.IPv4(10, 0, 0, 11), tokens: []string{"55"}, dataCenter: "local", rack: "a"},
+		{hostId: "11", connectAddress: net.IPv4(10, 0, 0, 12), tokens: []string{"60"}, dataCenter: "local", rack: "b"},
+	}
+	for _, host := range hosts {
+		policy.AddHost(host)
+		policyWithFallback.AddHost(host)
+	}
+
+	// the token ring is not setup without the partitioner, but the fallback
+	// should work
+	if actual := policy.Pick(nil)(); actual == nil {
+		t.Fatal("expected to get host from fallback got nil")
+	}
+
+	query.RoutingKey([]byte("30"))
+	if actual := policy.Pick(query)(); actual == nil {
+		t.Fatal("expected to get host from fallback got nil")
+	}
+
+	policy.SetPartitioner("OrderedPartitioner")
+	policyWithFallback.SetPartitioner("OrderedPartitioner")
+
+	policyInternal.getKeyspaceMetadata = func(keyspaceName string) (*KeyspaceMetadata, error) {
+		if keyspaceName != keyspace {
+			return nil, fmt.Errorf("unknown keyspace: %s", keyspaceName)
+		}
+		return &KeyspaceMetadata{
+			Name:          keyspace,
+			StrategyClass: "NetworkTopologyStrategy",
+			StrategyOptions: map[string]interface{}{
+				"class":  "NetworkTopologyStrategy",
+				"local":  2,
+				"remote": 2,
+			},
+		}, nil
+	}
+	policyWithFallbackInternal.getKeyspaceMetadata = policyInternal.getKeyspaceMetadata
+	policy.KeyspaceChanged(KeyspaceUpdateEvent{Keyspace: "myKeyspace"})
+	policyWithFallback.KeyspaceChanged(KeyspaceUpdateEvent{Keyspace: "myKeyspace"})
+
+	// The NetworkTopologyStrategy above should generate the following replicas.
+	// It's handy to have as reference here.
+	assertDeepEqual(t, "replicas", map[string]tokenRingReplicas{
+		"myKeyspace": {
+			{orderedToken("05"), []*HostInfo{hosts[0], hosts[1], hosts[2], hosts[3]}},
+			{orderedToken("10"), []*HostInfo{hosts[1], hosts[2], hosts[3], hosts[4]}},
+			{orderedToken("15"), []*HostInfo{hosts[2], hosts[3], hosts[4], hosts[5]}},
+			{orderedToken("20"), []*HostInfo{hosts[3], hosts[4], hosts[5], hosts[6]}},
+			{orderedToken("25"), []*HostInfo{hosts[4], hosts[5], hosts[6], hosts[7]}},
+			{orderedToken("30"), []*HostInfo{hosts[5], hosts[6], hosts[7], hosts[8]}},
+			{orderedToken("35"), []*HostInfo{hosts[6], hosts[7], hosts[8], hosts[9]}},
+			{orderedToken("40"), []*HostInfo{hosts[7], hosts[8], hosts[9], hosts[10]}},
+			{orderedToken("45"), []*HostInfo{hosts[8], hosts[9], hosts[10], hosts[11]}},
+			{orderedToken("50"), []*HostInfo{hosts[9], hosts[10], hosts[11], hosts[0]}},
+			{orderedToken("55"), []*HostInfo{hosts[10], hosts[11], hosts[0], hosts[1]}},
+			{orderedToken("60"), []*HostInfo{hosts[11], hosts[0], hosts[1], hosts[2]}},
+		},
+	}, policyInternal.getMetadataReadOnly().replicas)
+
+	query.RoutingKey([]byte("23"))
+
+	// now the token ring is configured
+	// Test the policy with fallback
+	iter = policyWithFallback.Pick(query)
+
+	// first should be host with matching token from the local DC & rack
+	expectHosts(t, "matching token from local DC and local rack", iter, "7")
+	// next should be host with matching token from local DC and other rack
+	expectHosts(t, "matching token from local DC and non-local rack", iter, "6")
+	// next should be hosts with matching token from other DC, in any order
+	expectHosts(t, "matching token from non-local DC", iter, "4", "5")
+	// then the local DC & rack that didn't match the token
+	expectHosts(t, "non-matching token from local DC and local rack", iter, "3", "11")
+	// then the local DC & other rack that didn't match the token
+	expectHosts(t, "non-matching token from local DC and non-local rack", iter, "2", "10")
+	// finally, the other DC that didn't match the token
+	expectHosts(t, "non-matching token from non-local DC", iter, "0", "1", "8", "9")
+
+	// Test the policy without fallback
+	iter = policy.Pick(query)
+
+	// first should be host with matching token from the local DC & Rack
+	expectHosts(t, "matching token from local DC and local rack", iter, "7")
+	// next should be the other two hosts from local DC & rack
+	expectHosts(t, "non-matching token local DC and local rack", iter, "3", "11")
+	// then the three hosts from the local DC but other rack
+	expectHosts(t, "local DC, non-local rack", iter, "2", "6", "10")
+	// then the 6 hosts from the other DC
+	expectHosts(t, "non-local DC", iter, "0", "1", "4", "5", "8", "9")
 }