Allow configuration of Cassandra's host selection policy (#4069)

Signed-off-by: ChinYing-Li <[email protected]>

Author: ChinYing-Li

CHANGELOG.md

@@ -24,6 +24,7 @@
 * [ENHANCEMENT] Distributor: reduce CPU and memory when an high number of errors are returned by the distributor on the write path. #3990
 * [ENHANCEMENT] Put metric before label value in the "label value too long" error message. #4018
 * [ENHANCEMENT] Allow use of `y|w|d` suffixes for duration related limits and per-tenant limits. #4044
+* [ENHANCEMENT] Allow configuration of Cassandra's host selection policy. (#4069)
 * [BUGFIX] Ruler-API: fix bug where `/api/v1/rules/<namespace>/<group_name>` endpoint return `400` instead of `404`. #4013
 * [BUGFIX] Distributor: reverted changes done to rate limiting in #3825. #3948
 * [BUGFIX] Ingester: Fix race condition when opening and closing tsdb concurrently. #3959

docs/configuration/config-file-reference.md

@@ -2961,6 +2961,11 @@ cassandra:
   # CLI flag: -cassandra.host-verification
   [host_verification: <boolean> | default = true]
 
+  # Policy for selecting Cassandra host. Supported values are: round-robin,
+  # token-aware.
+  # CLI flag: -cassandra.host-selection-policy
+  [host_selection_policy: <string> | default = "round-robin"]
+
   # Path to certificate file to verify the peer.
   # CLI flag: -cassandra.ca-path
   [CA_path: <string> | default = ""]

pkg/chunk/cassandra/storage_client.go

@@ -33,6 +33,7 @@ type Config struct {
 	DisableInitialHostLookup bool                `yaml:"disable_initial_host_lookup"`
 	SSL                      bool                `yaml:"SSL"`
 	HostVerification         bool                `yaml:"host_verification"`
+	HostSelectionPolicy      string              `yaml:"host_selection_policy"`
 	CAPath                   string              `yaml:"CA_path"`
 	CertPath                 string              `yaml:"tls_cert_path"`
 	KeyPath                  string              `yaml:"tls_key_path"`
@@ -53,6 +54,11 @@ type Config struct {
 	TableOptions             string              `yaml:"table_options"`
 }
 
+const (
+	HostPolicyRoundRobin = "round-robin"
+	HostPolicyTokenAware = "token-aware"
+)
+
 // RegisterFlags adds the flags required to config this to the given FlagSet
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Addresses, "cassandra.addresses", "", "Comma-separated hostnames or IPs of Cassandra instances.")
@@ -63,6 +69,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&cfg.DisableInitialHostLookup, "cassandra.disable-initial-host-lookup", false, "Instruct the cassandra driver to not attempt to get host info from the system.peers table.")
 	f.BoolVar(&cfg.SSL, "cassandra.ssl", false, "Use SSL when connecting to cassandra instances.")
 	f.BoolVar(&cfg.HostVerification, "cassandra.host-verification", true, "Require SSL certificate validation.")
+	f.StringVar(&cfg.HostSelectionPolicy, "cassandra.host-selection-policy", HostPolicyRoundRobin, "Policy for selecting Cassandra host. Supported values are: round-robin, token-aware.")
 	f.StringVar(&cfg.CAPath, "cassandra.ca-path", "", "Path to certificate file to verify the peer.")
 	f.StringVar(&cfg.CertPath, "cassandra.tls-cert-path", "", "Path to certificate file used by TLS.")
 	f.StringVar(&cfg.KeyPath, "cassandra.tls-key-path", "", "Path to private key file used by TLS.")
@@ -180,6 +187,15 @@ func (cfg *Config) setClusterConfig(cluster *gocql.ClusterConfig) error {
 			}
 		}
 	}
+
+	if cfg.HostSelectionPolicy == HostPolicyRoundRobin {
+		cluster.PoolConfig.HostSelectionPolicy = gocql.RoundRobinHostPolicy()
+	} else if cfg.HostSelectionPolicy == HostPolicyTokenAware {
+		cluster.PoolConfig.HostSelectionPolicy = gocql.TokenAwareHostPolicy(gocql.RoundRobinHostPolicy())
+	} else {
+		return errors.New("Unknown host selection policy")
+	}
+
 	if cfg.Auth {
 		password := cfg.Password.Value
 		if cfg.PasswordFile != "" {