Repository: k8ssandra/k8ssandra-operator (go)

Issue:
New token allocator is disabled with DSE
The new token allocator uses the `allocate_tokens_for_local_replication_factor` setting to perform smart token allocations during bootstrap.

While this setting exists by default in Cassandra 4, it's not the case in DSE 6.8.
We need to set it explicitly to 3 (using the same default as Cassandra) unless the setting is overridden in the manifest.

Currently we end up with balance issues when using a small number of tokens.

Pull request description:
Enable smart token allocation by default for DSE
<!--  Thanks for sending a pull request!  Before submitting:

1. Read our CONTRIBUTING.md guide
2. Rebase your PR if it gets out of sync with main
-->

**What this PR does**:
DSE cassandra.yaml file doesn't set a default value for `allocate_tokens_for_local_replication_factor` and thus disables smart allocations by default.
This PR will enable it for DSE, using the same defaults as C* 4.x, while preserving any override from users.

**Which issue(s) this PR fixes**:
Fixes #796 

**Checklist**
- [x] Changes manually tested
- [x] Automated Tests added/updated
- [ ] Documentation added/updated
- [x] CHANGELOG.md updated (not required for documentation PRs)
- [x] CLA Signed:  [DataStax CLA](https://cla.datastax.com/)


Patch:
diff --git a/CHANGELOG/CHANGELOG-1.5.md b/CHANGELOG/CHANGELOG-1.5.md
index 37ccd503..c624afc3 100644
--- a/CHANGELOG/CHANGELOG-1.5.md
+++ b/CHANGELOG/CHANGELOG-1.5.md
@@ -17,6 +17,7 @@ When cutting a new release, update the `unreleased` heading to the tag being gen
 
 * [FEATURE] [#775](https://github.com/k8ssandra/k8ssandra-operator/issues/775) Add the ability to inject and configure a Vector agent sidecar in the Cassandra pods
 * [FEATURE] [#600](https://github.com/k8ssandra/k8ssandra-operator/issues/600) Disable secrets management and replication with the external secrets provider
+* [ENHANCEMENT] [#796](https://github.com/k8ssandra/k8ssandra-operator/issues/796) Enable smart token allocation by default for DSE
 * [ENHANCEMENT] [#765](https://github.com/k8ssandra/k8ssandra-operator/issues/765) Add GitHub workflow to test various k8s versions
 * [ENHANCEMENT] [#323](https://github.com/k8ssandra/k8ssandra/issues/323) Use Cassandra internals for JMX authentication
 * [ENHANCEMENT] [#770](https://github.com/k8ssandra/k8ssandra-operator/issues/770) Update to Go 1.19 and Operator SDK 1.25.2
diff --git a/controllers/k8ssandra/dcconfigs.go b/controllers/k8ssandra/dcconfigs.go
index 14a297fc..0b7a08d1 100644
--- a/controllers/k8ssandra/dcconfigs.go
+++ b/controllers/k8ssandra/dcconfigs.go
@@ -83,6 +83,7 @@ func (r *K8ssandraClusterReconciler) createDatacenterConfigs(
 		cassandra.AddNumTokens(dcConfig)
 		cassandra.AddStartRpc(dcConfig)
 		cassandra.HandleDeprecatedJvmOptions(&dcConfig.CassandraConfig.JvmOptions)
+		cassandra.EnableSmartTokenAllocation(dcConfig)
 
 		if err := cassandra.ValidateDatacenterConfig(dcConfig); err != nil {
 			return nil, err
diff --git a/pkg/cassandra/config.go b/pkg/cassandra/config.go
index ed3705d6..bbe8d6f4 100644
--- a/pkg/cassandra/config.go
+++ b/pkg/cassandra/config.go
@@ -125,3 +125,14 @@ func ApplySystemReplication(dcConfig *DatacenterConfig, replication SystemReplic
 func AllowAlterRfDuringRangeMovement(dcConfig *DatacenterConfig) {
 	addOptionIfMissing(dcConfig, allowAlterRf)
 }
+
+// EnableSmartTokenAllocation adds the allocate_tokens_for_local_replication_factor option to
+// cassandra.yaml if it is not already present when running DSE.
+// This option is enabled by default in Cassandra but not DSE.
+func EnableSmartTokenAllocation(template *DatacenterConfig) {
+	// Note: we put int64 values because even if int values can be marshaled just fine,
+	// Unstructured.DeepCopy() would reject them since int is not a supported json type.
+	if template.ServerType == api.ServerDistributionDse {
+		template.CassandraConfig.CassandraYaml.PutIfAbsent("allocate_tokens_for_local_replication_factor", int64(3))
+	}
+}


Test patch:
diff --git a/pkg/cassandra/config_test.go b/pkg/cassandra/config_test.go
index 4b67d654..c16215cd 100644
--- a/pkg/cassandra/config_test.go
+++ b/pkg/cassandra/config_test.go
@@ -272,3 +272,36 @@ func parseQuantity(quantity string) *resource.Quantity {
 	parsed := resource.MustParse(quantity)
 	return &parsed
 }
+
+func TestEnableSmartTokenAllocDse(t *testing.T) {
+	dcConfig := &DatacenterConfig{
+		ServerType: api.ServerDistributionDse,
+	}
+
+	EnableSmartTokenAllocation(dcConfig)
+	assert.Equal(t, int64(3), dcConfig.CassandraConfig.CassandraYaml["allocate_tokens_for_local_replication_factor"].(int64), "allocate_tokens_for_local_replication_factor should be set to 3 by default for DSE")
+}
+
+func TestOverrideSmartTokenAllocDse(t *testing.T) {
+	dcConfig := &DatacenterConfig{
+		ServerType: api.ServerDistributionDse,
+		CassandraConfig: api.CassandraConfig{
+			CassandraYaml: unstructured.Unstructured{
+				"allocate_tokens_for_local_replication_factor": int64(5),
+			},
+		},
+	}
+
+	EnableSmartTokenAllocation(dcConfig)
+	assert.Equal(t, int64(5), dcConfig.CassandraConfig.CassandraYaml["allocate_tokens_for_local_replication_factor"].(int64), "allocate_tokens_for_local_replication_factor should retain configured value")
+}
+
+func TestSmartTokenAllocCassandra(t *testing.T) {
+	dcConfig := &DatacenterConfig{
+		ServerType: api.ServerDistributionCassandra,
+	}
+
+	EnableSmartTokenAllocation(dcConfig)
+	_, exists := dcConfig.CassandraConfig.CassandraYaml["allocate_tokens_for_local_replication_factor"]
+	assert.False(t, exists, "allocate_tokens_for_local_replication_factor should not be set for Cassandra")
+}
diff --git a/test/e2e/dse_test.go b/test/e2e/dse_test.go
index 0a4ee278..e1d08372 100644
--- a/test/e2e/dse_test.go
+++ b/test/e2e/dse_test.go
@@ -14,6 +14,7 @@ import (
 	gremlingo "github.com/apache/tinkerpop/gremlin-go/driver"
 	cassdcapi "github.com/k8ssandra/cass-operator/apis/cassandra/v1beta1"
 	api "github.com/k8ssandra/k8ssandra-operator/apis/k8ssandra/v1alpha1"
+	"github.com/k8ssandra/k8ssandra-operator/pkg/utils"
 	"github.com/k8ssandra/k8ssandra-operator/test/framework"
 	"github.com/stretchr/testify/require"
 	"k8s.io/apimachinery/pkg/types"
@@ -45,6 +46,16 @@ func createSingleDseDatacenterCluster(t *testing.T, ctx context.Context, namespa
 	err = f.Client.Patch(ctx, kc, patch)
 	require.NoError(t, err, "failed to patch K8ssandraCluster %v", kcKey)
 
+	// Check that the smart token allocation is enabled
+	t.Log("Check that the smart token allocation is enabled")
+	dc := &cassdcapi.CassandraDatacenter{}
+	err = f.Client.Get(ctx, dcKey.NamespacedName, dc)
+	require.NoError(t, err, "failed to get CassandraDatacenter %v", dcKey)
+	dcConfig, err := utils.UnmarshalToMap(dc.Spec.Config)
+	require.NoError(t, err, "failed to unmarshall CassandraDatacenter config")
+	allocateTokensSetting := dcConfig["cassandra-yaml"].(map[string]interface{})["allocate_tokens_for_local_replication_factor"].(float64)
+	require.Equal(t, float64(3), allocateTokensSetting, "expected smart token allocation to be enabled by default for DSE")
+
 	checkDatacenterReady(t, ctx, dcKey, f)
 	assertCassandraDatacenterK8cStatusReady(ctx, t, f, kcKey, dcKey.Name)