Repository: k8ssandra/k8ssandra-operator (go)

Issue:
Set defaults for num_tokens and make it configurable
Even though we use the stock defaults for `num_tokens` we need to explicitly set it because the config builder uses a default of `num_tokens: 1`. It should be configurable as well in case users want to use a non-default value.



┆Issue is synchronized with this [Jira Task](https://k8ssandra.atlassian.net/browse/K8SSAND-923) by [Unito](https://www.unito.io)
┆Reviewer: 557058:376e6d8d-84b1-46a1-a227-0887055aa65a
┆Fix Versions: k8ssandra-operator-1.0.0-alpha.1
┆Issue Number: K8SSAND-923
┆Priority: Medium


Pull request description:
configure num_tokens [K8SSAND-923]
<!--  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**:
Configures `num_tokens` and `allocate_tokens_for_local_replication_factor`

**Which issue(s) this PR fixes**:
Fixes #151 

**Checklist**
- [ ] Changes manually tested
- [x] Automated Tests added/updated
- [ ] Documentation added/updated
- [ ] CHANGELOG.md updated (not required for documentation PRs)
- [x] CLA Signed:  [DataStax CLA](https://cla.datastax.com/)


Patch:
diff --git a/api/v1alpha1/k8ssandracluster_types.go b/api/v1alpha1/k8ssandracluster_types.go
index eb617ec1..36be404b 100644
--- a/api/v1alpha1/k8ssandracluster_types.go
+++ b/api/v1alpha1/k8ssandracluster_types.go
@@ -247,6 +247,12 @@ type CassandraYaml struct {
 	//
 	//PermissionValidityMillis *int64 `json:"permissions_validity_in_ms,omitempty"`
 
+	// +optional
+	NumTokens *int `json:"num_tokens,omitempty"`
+
+	// +optional
+	AllocateTokensForLocalReplicationFactor *int `json:"allocate_tokens_for_local_replication_factor,omitempty"`
+
 	// +optional
 	ConcurrentReads *int `json:"concurrent_reads,omitempty"`
 
diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go
index 3bb8d272..6f05d467 100644
--- a/api/v1alpha1/zz_generated.deepcopy.go
+++ b/api/v1alpha1/zz_generated.deepcopy.go
@@ -187,6 +187,16 @@ func (in *CassandraDatacenterTemplate) DeepCopy() *CassandraDatacenterTemplate {
 // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil.
 func (in *CassandraYaml) DeepCopyInto(out *CassandraYaml) {
 	*out = *in
+	if in.NumTokens != nil {
+		in, out := &in.NumTokens, &out.NumTokens
+		*out = new(int)
+		**out = **in
+	}
+	if in.AllocateTokensForLocalReplicationFactor != nil {
+		in, out := &in.AllocateTokensForLocalReplicationFactor, &out.AllocateTokensForLocalReplicationFactor
+		*out = new(int)
+		**out = **in
+	}
 	if in.ConcurrentReads != nil {
 		in, out := &in.ConcurrentReads, &out.ConcurrentReads
 		*out = new(int)
diff --git a/config/crd/bases/k8ssandra.io_k8ssandraclusters.yaml b/config/crd/bases/k8ssandra.io_k8ssandraclusters.yaml
index 64562077..f33684a2 100644
--- a/config/crd/bases/k8ssandra.io_k8ssandraclusters.yaml
+++ b/config/crd/bases/k8ssandra.io_k8ssandraclusters.yaml
@@ -54,6 +54,8 @@ spec:
                     properties:
                       cassandraYaml:
                         properties:
+                          allocate_tokens_for_local_replication_factor:
+                            type: integer
                           auto_snapshot:
                             type: boolean
                           commitlog_segment_size_in_mb:
@@ -76,6 +78,8 @@ spec:
                             type: integer
                           memtable_flush_writers:
                             type: integer
+                          num_tokens:
+                            type: integer
                           prepared_statements_cache_size_mb:
                             type: integer
                           row_cache_size_in_mb:
@@ -120,6 +124,8 @@ spec:
                           properties:
                             cassandraYaml:
                               properties:
+                                allocate_tokens_for_local_replication_factor:
+                                  type: integer
                                 auto_snapshot:
                                   type: boolean
                                 commitlog_segment_size_in_mb:
@@ -142,6 +148,8 @@ spec:
                                   type: integer
                                 memtable_flush_writers:
                                   type: integer
+                                num_tokens:
+                                  type: integer
                                 prepared_statements_cache_size_mb:
                                   type: integer
                                 row_cache_size_in_mb:
diff --git a/pkg/cassandra/config.go b/pkg/cassandra/config.go
index 415a1b73..2b08c6a3 100644
--- a/pkg/cassandra/config.go
+++ b/pkg/cassandra/config.go
@@ -48,6 +48,19 @@ func (c config) MarshalJSON() ([]byte, error) {
 			c.StartRpc = nil
 			c.ThriftPreparedStatementCacheSizeMb = nil
 		}
+
+		// Even though we default to Cassandra's stock defaults for num_tokens, we need to
+		// explicitly set it because the config builder defaults to num_tokens: 1
+		if c.NumTokens == nil {
+			if isCassandra4(c.cassandraVersion) {
+				numTokens := 16
+				c.NumTokens = &numTokens
+			} else {
+				numTokens := 256
+				c.NumTokens = &numTokens
+			}
+		}
+
 		config["cassandra-yaml"] = c.CassandraYaml
 	}
 
@@ -65,7 +78,9 @@ func (c config) MarshalJSON() ([]byte, error) {
 func newConfig(apiConfig *api.CassandraConfig, cassandraVersion string) config {
 	config := config{cassandraVersion: cassandraVersion}
 
-	if apiConfig.CassandraYaml != nil {
+	if apiConfig.CassandraYaml == nil {
+		config.CassandraYaml = &api.CassandraYaml{}
+	} else {
 		config.CassandraYaml = apiConfig.CassandraYaml
 	}
 


Test patch:
diff --git a/pkg/cassandra/config_test.go b/pkg/cassandra/config_test.go
index c997aae5..94c8a53d 100644
--- a/pkg/cassandra/config_test.go
+++ b/pkg/cassandra/config_test.go
@@ -3,8 +3,9 @@ package cassandra
 import (
 	"github.com/Jeffail/gabs"
 	api "github.com/k8ssandra/k8ssandra-operator/api/v1alpha1"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"k8s.io/apimachinery/pkg/api/resource"
-	"reflect"
 	"testing"
 )
 
@@ -106,10 +107,10 @@ func TestApplySystemReplication(t *testing.T) {
 	}
 
 	for _, tc := range tests {
-		ApplySystemReplication(tc.dcConfig, tc.replication)
-		if !reflect.DeepEqual(tc.want, tc.dcConfig) {
-			t.Errorf("%s - expected: %+v, got: %+v", tc.name, *tc.want, *tc.dcConfig)
-		}
+		t.Run(tc.name, func(t *testing.T) {
+			ApplySystemReplication(tc.dcConfig, tc.replication)
+			require.Equal(t, tc.want, tc.dcConfig)
+		})
 	}
 }
 
@@ -126,8 +127,8 @@ func TestCreateJsonConfig(t *testing.T) {
 
 	tests := []test{
 		{
-			name:             "concurrent_reads, concurrent_writes, concurrent_counter_writes",
-			cassandraVersion: "4.0",
+			name:             "[4.0.0] concurrent_reads, concurrent_writes, concurrent_counter_writes",
+			cassandraVersion: "4.0.0",
 			config: &api.CassandraConfig{
 				CassandraYaml: &api.CassandraYaml{
 					ConcurrentReads:         intPtr(8),
@@ -137,6 +138,7 @@ func TestCreateJsonConfig(t *testing.T) {
 			},
 			want: `{
               "cassandra-yaml": {
+                "num_tokens": 16,
                 "concurrent_reads": 8,
                 "concurrent_writes": 16,
                 "concurrent_counter_writes": 4
@@ -144,14 +146,35 @@ func TestCreateJsonConfig(t *testing.T) {
             }`,
 		},
 		{
-			name:             "heap size - Cassandra 4.0",
-			cassandraVersion: "4.0",
+			name:             "[3.11.11] heap size",
+			cassandraVersion: "3.11.11",
 			config: &api.CassandraConfig{
 				JvmOptions: &api.JvmOptions{
 					HeapSize: &heapSize,
 				},
 			},
 			want: `{
+              "cassandra-yaml": {
+                "num_tokens": 256
+              },
+              "jvm-options": {
+                "initial_heap_size": 1073741824,
+                "max_heap_size": 1073741824
+              }
+            }`,
+		},
+		{
+			name:             "[4.0.0] heap size",
+			cassandraVersion: "4.0.0",
+			config: &api.CassandraConfig{
+				JvmOptions: &api.JvmOptions{
+					HeapSize: &heapSize,
+				},
+			},
+			want: `{
+              "cassandra-yaml": {
+                "num_tokens": 16
+              },
               "jvm-server-options": {
                 "initial_heap_size": 1073741824,
                 "max_heap_size": 1073741824
@@ -159,8 +182,8 @@ func TestCreateJsonConfig(t *testing.T) {
             }`,
 		},
 		{
-			name:             "concurrent_reads and concurrent_writes with system replication",
-			cassandraVersion: "4.0",
+			name:             "[4.0.0] concurrent_reads and concurrent_writes with system replication",
+			cassandraVersion: "4.0.0",
 			config: &api.CassandraConfig{
 				CassandraYaml: &api.CassandraYaml{
 					ConcurrentReads:  intPtr(8),
@@ -175,6 +198,7 @@ func TestCreateJsonConfig(t *testing.T) {
 			},
 			want: `{
               "cassandra-yaml": {
+                "num_tokens": 16,
                 "concurrent_reads": 8,
                 "concurrent_writes": 16
               },
@@ -187,8 +211,8 @@ func TestCreateJsonConfig(t *testing.T) {
             }`,
 		},
 		{
-			name:             "auto_snapshot, memtable_flush_writers, commitlog_segment_size_in_mb",
-			cassandraVersion: "4.0",
+			name:             "[4.0.0] auto_snapshot, memtable_flush_writers, commitlog_segment_size_in_mb",
+			cassandraVersion: "4.0.0",
 			config: &api.CassandraConfig{
 				CassandraYaml: &api.CassandraYaml{
 					AutoSnapshot:           boolPtr(true),
@@ -198,6 +222,7 @@ func TestCreateJsonConfig(t *testing.T) {
 			},
 			want: `{
               "cassandra-yaml": {
+                "num_tokens": 16,
                 "auto_snapshot": true,
                 "memtable_flush_writers": 10,
                 "commitlog_segment_size_in_mb": 8192
@@ -205,8 +230,8 @@ func TestCreateJsonConfig(t *testing.T) {
             }`,
 		},
 		{
-			name:             "concurrent_compactors, compaction_throughput_mb_per_sec, sstable_preemptive_open_interval_in_mb",
-			cassandraVersion: "4.0",
+			name:             "[4.0.0] concurrent_compactors, compaction_throughput_mb_per_sec, sstable_preemptive_open_interval_in_mb",
+			cassandraVersion: "4.0.0",
 			config: &api.CassandraConfig{
 				CassandraYaml: &api.CassandraYaml{
 					ConcurrentCompactors:            intPtr(4),
@@ -216,6 +241,7 @@ func TestCreateJsonConfig(t *testing.T) {
 			},
 			want: `{
               "cassandra-yaml": {
+				"num_tokens": 16,
                 "concurrent_compactors": 4,
                 "compaction_throughput_mb_per_sec": 64,
                 "sstable_preemptive_open_interval_in_mb": 0
@@ -223,8 +249,8 @@ func TestCreateJsonConfig(t *testing.T) {
             }`,
 		},
 		{
-			name:             "key_cache_size_in_mb, counter_cache_size_in_mb, prepared_statements_cache_size_mb, slow_query_log_timeout_in_ms",
-			cassandraVersion: "4.0",
+			name:             "[4.0.0] key_cache_size_in_mb, counter_cache_size_in_mb, prepared_statements_cache_size_mb, slow_query_log_timeout_in_ms",
+			cassandraVersion: "4.0.0",
 			config: &api.CassandraConfig{
 				CassandraYaml: &api.CassandraYaml{
 					KeyCacheSizeMb:                intPtr(100),
@@ -235,6 +261,7 @@ func TestCreateJsonConfig(t *testing.T) {
 			},
 			want: `{
               "cassandra-yaml": {
+				"num_tokens": 16,
                 "key_cache_size_in_mb": 100,
                 "counter_cache_size_in_mb": 50,
                 "prepared_statements_cache_size_mb": 180,
@@ -243,8 +270,8 @@ func TestCreateJsonConfig(t *testing.T) {
             }`,
 		},
 		{
-			name:             "file_cache_size_in_mb, row_cache_size_in_mb",
-			cassandraVersion: "4.0",
+			name:             "[4.0.0] file_cache_size_in_mb, row_cache_size_in_mb",
+			cassandraVersion: "4.0.0",
 			config: &api.CassandraConfig{
 				CassandraYaml: &api.CassandraYaml{
 					FileCacheSizeMb: intPtr(500),
@@ -253,6 +280,7 @@ func TestCreateJsonConfig(t *testing.T) {
 			},
 			want: `{
               "cassandra-yaml": {
+				"num_tokens": 16,
                 "file_cache_size_in_mb": 500,
                 "row_cache_size_in_mb": 100
               }
@@ -269,14 +297,15 @@ func TestCreateJsonConfig(t *testing.T) {
 			},
 			want: `{
               "cassandra-yaml": {
+				"num_tokens": 256,
                 "start_rpc": false,
                 "thrift_prepared_statements_cache_size_mb": 1
               }
             }`,
 		},
 		{
-			name:             "[4.0] start_rpc, thrift_prepared_statements_cache_size_mb",
-			cassandraVersion: "4.0",
+			name:             "[4.0.0] start_rpc, thrift_prepared_statements_cache_size_mb",
+			cassandraVersion: "4.0.0",
 			config: &api.CassandraConfig{
 				CassandraYaml: &api.CassandraYaml{
 					StartRpc:                           boolPtr(false),
@@ -284,7 +313,51 @@ func TestCreateJsonConfig(t *testing.T) {
 				},
 			},
 			want: `{
-              "cassandra-yaml": { 
+              "cassandra-yaml": {
+				"num_tokens": 16
+              }
+            }`,
+		},
+		{
+			name:             "[3.11.11] num_tokens",
+			cassandraVersion: "3.11.11",
+			config: &api.CassandraConfig{
+				CassandraYaml: &api.CassandraYaml{
+					NumTokens: intPtr(32),
+				},
+			},
+			want: `{
+              "cassandra-yaml": {
+                "num_tokens": 32
+              }
+            }`,
+		},
+		{
+			name:             "[4.0.0] num_tokens",
+			cassandraVersion: "4.0.0",
+			config: &api.CassandraConfig{
+				CassandraYaml: &api.CassandraYaml{
+					NumTokens: intPtr(32),
+				},
+			},
+			want: `{
+              "cassandra-yaml": {
+                "num_tokens": 32
+              }
+            }`,
+		},
+		{
+			name:             "[4.0.0] allocate_tokens_for_local_replication_factor",
+			cassandraVersion: "4.0.0",
+			config: &api.CassandraConfig{
+				CassandraYaml: &api.CassandraYaml{
+					AllocateTokensForLocalReplicationFactor: intPtr(5),
+				},
+			},
+			want: `{
+              "cassandra-yaml": {
+                "allocate_tokens_for_local_replication_factor": 5,
+				"num_tokens": 16
               }
             }`,
 		},
@@ -307,29 +380,18 @@ func TestCreateJsonConfig(t *testing.T) {
 	}
 
 	for _, tc := range tests {
-		var err error
-		tc.got, err = CreateJsonConfig(tc.config, tc.cassandraVersion)
-		if err != nil {
-			t.Errorf("%s - failed to create json dcConfig: %s", tc.name, err)
-			continue
-		}
-
-		expected, err := gabs.ParseJSON([]byte(tc.want))
-		if err != nil {
-			t.Errorf("%s - failed to parse expected value: %s", tc.name, err)
-			continue
-		}
-
-		actual, err := gabs.ParseJSON(tc.got)
-		if err != nil {
-			t.Errorf("%s - failed to parse actual value: %s", tc.name, err)
-			continue
-		}
-
-		if !reflect.DeepEqual(expected, actual) {
-			t.Errorf("%s - wanted: %s, got: %s", tc.name, expected, actual)
-		}
+		t.Run(tc.name, func(t *testing.T) {
+			var err error
+			tc.got, err = CreateJsonConfig(tc.config, tc.cassandraVersion)
+			require.NoError(t, err, "failed to create json dcConfig")
+			expected, err := gabs.ParseJSON([]byte(tc.want))
+			require.NoError(t, err, "failed to parse expected value")
+			actual, err := gabs.ParseJSON(tc.got)
+			require.NoError(t, err, "failed to parse actual value")
+			assert.Equal(t, expected, actual)
+		})
 	}
+
 }
 
 func intPtr(n int) *int {
diff --git a/pkg/cassandra/datacenter_test.go b/pkg/cassandra/datacenter_test.go
index dc91504b..79081d27 100644
--- a/pkg/cassandra/datacenter_test.go
+++ b/pkg/cassandra/datacenter_test.go
@@ -1,7 +1,7 @@
 package cassandra
 
 import (
-	"reflect"
+	"github.com/stretchr/testify/require"
 	"testing"
 
 	cassdcapi "github.com/k8ssandra/cass-operator/apis/cassandra/v1beta1"
@@ -232,9 +232,9 @@ func TestCoalesce(t *testing.T) {
 	}
 
 	for _, tc := range tests {
-		tc.got = Coalesce(tc.clusterTemplate, tc.dcTemplate)
-		if !reflect.DeepEqual(tc.got, tc.want) {
-			t.Errorf("%s - expected: %+v, got: %+v", tc.name, tc.want, tc.got)
-		}
+		t.Run(tc.name, func(t *testing.T) {
+			tc.got = Coalesce(tc.clusterTemplate, tc.dcTemplate)
+			require.Equal(t, tc.want, tc.got)
+		})
 	}
 }
diff --git a/pkg/cassandra/util_test.go b/pkg/cassandra/util_test.go
index 76836e0b..ad45da76 100644
--- a/pkg/cassandra/util_test.go
+++ b/pkg/cassandra/util_test.go
@@ -2,7 +2,7 @@ package cassandra
 
 import (
 	api "github.com/k8ssandra/k8ssandra-operator/api/v1alpha1"
-	"reflect"
+	"github.com/stretchr/testify/require"
 	"testing"
 )
 
@@ -94,9 +94,9 @@ func TestComputeSystemReplication(t *testing.T) {
 	}
 
 	for _, tc := range tests {
-		tc.got = ComputeSystemReplication(tc.kluster)
-		if !reflect.DeepEqual(tc.want, tc.got) {
-			t.Errorf("%s - expected: %+v, got: %+v", tc.name, tc.want, tc.got)
-		}
+		t.Run(tc.name, func(t *testing.T) {
+			tc.got = ComputeSystemReplication(tc.kluster)
+			require.Equal(t, tc.want, tc.got)
+		})
 	}
 }