adobe
diff --git a/‎Makefile
+1-1 b/‎Makefile
+1-1
diff --git a/‎api/util/util.go
+8 b/‎api/util/util.go
+8
diff --git a/‎api/v1beta1/kafkacluster_types.go
+16-5 b/‎api/v1beta1/kafkacluster_types.go
+16-5
diff --git a/‎api/v1beta1/kafkacluster_types_test.go
+118-6 b/‎api/v1beta1/kafkacluster_types_test.go
+118-6
diff --git a/‎config/samples/kraft/simplekafkacluster_kraft.yaml
+1-1 b/‎config/samples/kraft/simplekafkacluster_kraft.yaml
+1-1
diff --git a/‎controllers/tests/kafkacluster_controller_kafka_test.go
+78-36 b/‎controllers/tests/kafkacluster_controller_kafka_test.go
+78-36
diff --git a/‎pkg/jmxextractor/extractor.go
+10-4 b/‎pkg/jmxextractor/extractor.go
+10-4
@@ -102,7 +102,7 @@ test: generate fmt vet bin/setup-envtest
 test-e2e:
 	 IMG_E2E=${IMG_E2E} go test github.com/banzaicloud/koperator/tests/e2e \
 		-v \
-		-timeout 20m \
+		-timeout 45m \
 		-tags e2e \
 		--ginkgo.show-node-events \
 		--ginkgo.trace \
 
@@ -40,6 +40,14 @@ func LabelsForKafka(name string) map[string]string {
 	return map[string]string{"app": "kafka", "kafka_cr": name}
 }
 
+func LabelsForBroker(name string) map[string]string {
+	return map[string]string{"isBrokerNode": "true", "app": "kafka", "kafka_cr": name}
+}
+
+func LabelsForController(name string) map[string]string {
+	return map[string]string{"isControllerNode": "true", "app": "kafka", "kafka_cr": name}
+}
+
 // StringSliceContains returns true if list contains s
 func StringSliceContains(list []string, s string) bool {
 	for _, v := range list {
 
@@ -45,6 +45,12 @@ const (
 	// ProcessRolesKey is used to identify which process roles the Kafka pod has
 	ProcessRolesKey = "processRoles"
 
+	// IsBrokerNodeKey is used to identify if the kafka pod is either a broker or a broker_controller
+	IsBrokerNodeKey = "isBrokerNode"
+
+	// IsControllerNodeKey is used to identify if the kafka pod is a controller or broker_controller
+	IsControllerNodeKey = "isControllerNode"
+
 	// DefaultCruiseControlImage is the default CC image used when users don't specify it in CruiseControlConfig.Image
 	DefaultCruiseControlImage = "ghcr.io/banzaicloud/cruise-control:2.5.123"
 
@@ -1095,14 +1101,19 @@ func (bConfig *BrokerConfig) GetBrokerAnnotations() map[string]string {
 }
 
 // GetBrokerLabels returns the labels that are applied to broker pods
-func (bConfig *BrokerConfig) GetBrokerLabels(kafkaClusterName string, brokerId int32) map[string]string {
+func (bConfig *BrokerConfig) GetBrokerLabels(kafkaClusterName string, brokerId int32, kRaftMode bool) map[string]string {
+	kraftLabels := make(map[string]string, 0)
+	if kRaftMode {
+		kraftLabels = map[string]string{
+			ProcessRolesKey:     strings.Join(bConfig.Roles, "_"),
+			IsControllerNodeKey: fmt.Sprintf("%t", bConfig.IsControllerNode()),
+			IsBrokerNodeKey:     fmt.Sprintf("%t", bConfig.IsBrokerNode()),
+		}
+	}
 	return util.MergeLabels(
 		bConfig.BrokerLabels,
 		util.LabelsForKafka(kafkaClusterName),
-		map[string]string{
-			BrokerIdLabelKey: fmt.Sprintf("%d", brokerId),
-			ProcessRolesKey:  strings.Join(bConfig.Roles, "_"),
-		},
+		kraftLabels, map[string]string{BrokerIdLabelKey: fmt.Sprintf("%d", brokerId)},
 	)
 }
 
 
@@ -436,13 +436,125 @@ func TestGetBrokerLabels(t *testing.T) {
 
 		expectedBrokerId = 0
 	)
+	testCases := []struct {
+		testName       string
+		brokerConfig   *BrokerConfig
+		expectedLabels map[string]string
+		kRaftMode      bool
+	}{
+		{
+			testName: "Labels in zookeeper mode",
+			expectedLabels: map[string]string{
+				AppLabelKey:      expectedDefaultLabelApp,
+				BrokerIdLabelKey: strconv.Itoa(expectedBrokerId),
+				KafkaCRLabelKey:  expectedKafkaCRName,
+				"test_label_key": "test_label_value",
+			},
+			brokerConfig: &BrokerConfig{
+				Roles: nil,
+				BrokerLabels: map[string]string{
+					AppLabelKey:      "test_app",
+					BrokerIdLabelKey: "test_id",
+					KafkaCRLabelKey:  "test_cr_name",
+					"test_label_key": "test_label_value",
+				},
+			},
+			kRaftMode: false,
+		},
+		{
+			testName: "Labels for broker in kraft mode",
+			expectedLabels: map[string]string{
+				AppLabelKey:         expectedDefaultLabelApp,
+				BrokerIdLabelKey:    strconv.Itoa(expectedBrokerId),
+				KafkaCRLabelKey:     expectedKafkaCRName,
+				"test_label_key":    "test_label_value",
+				ProcessRolesKey:     "broker",
+				IsBrokerNodeKey:     "true",
+				IsControllerNodeKey: "false",
+			},
+			brokerConfig: &BrokerConfig{
+				Roles: []string{"broker"},
+				BrokerLabels: map[string]string{
+					AppLabelKey:      "test_app",
+					BrokerIdLabelKey: "test_id",
+					KafkaCRLabelKey:  "test_cr_name",
+					"test_label_key": "test_label_value",
+				},
+			},
+			kRaftMode: true,
+		},
+		{
+			testName: "Labels for controller in kraft mode",
+			expectedLabels: map[string]string{
+				AppLabelKey:         expectedDefaultLabelApp,
+				BrokerIdLabelKey:    strconv.Itoa(expectedBrokerId),
+				KafkaCRLabelKey:     expectedKafkaCRName,
+				"test_label_key":    "test_label_value",
+				ProcessRolesKey:     "controller",
+				IsBrokerNodeKey:     "false",
+				IsControllerNodeKey: "true",
+			},
+			brokerConfig: &BrokerConfig{
+				Roles: []string{"controller"},
+				BrokerLabels: map[string]string{
+					AppLabelKey:      "test_app",
+					BrokerIdLabelKey: "test_id",
+					KafkaCRLabelKey:  "test_cr_name",
+					"test_label_key": "test_label_value",
+				},
+			},
+			kRaftMode: true,
+		},
+		{
+			testName: "Labels for controller/broker in kraft mode",
+			expectedLabels: map[string]string{
+				AppLabelKey:         expectedDefaultLabelApp,
+				BrokerIdLabelKey:    strconv.Itoa(expectedBrokerId),
+				KafkaCRLabelKey:     expectedKafkaCRName,
+				"test_label_key":    "test_label_value",
+				ProcessRolesKey:     "controller_broker",
+				IsBrokerNodeKey:     "true",
+				IsControllerNodeKey: "true",
+			},
+			brokerConfig: &BrokerConfig{
+				Roles: []string{"controller", "broker"},
+				BrokerLabels: map[string]string{
+					AppLabelKey:      "test_app",
+					BrokerIdLabelKey: "test_id",
+					KafkaCRLabelKey:  "test_cr_name",
+					"test_label_key": "test_label_value",
+				},
+			},
+			kRaftMode: true,
+		},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.testName, func(t *testing.T) {
+			result := test.brokerConfig.GetBrokerLabels(expectedKafkaCRName, expectedBrokerId, test.kRaftMode)
+			if !reflect.DeepEqual(result, test.expectedLabels) {
+				t.Error("Expected:", test.expectedLabels, "Got:", result)
+			}
+		})
+	}
+}
+
+func TestGetBrokerLabelKraft(t *testing.T) {
+	const (
+		expectedDefaultLabelApp = "kafka"
+		expectedKafkaCRName     = "kafka"
+
+		expectedBrokerId = 0
+	)
 
 	expected := map[string]string{
-		AppLabelKey:      expectedDefaultLabelApp,
-		BrokerIdLabelKey: strconv.Itoa(expectedBrokerId),
-		KafkaCRLabelKey:  expectedKafkaCRName,
-		"test_label_key": "test_label_value",
-		ProcessRolesKey:  "broker",
+		AppLabelKey:         expectedDefaultLabelApp,
+		BrokerIdLabelKey:    strconv.Itoa(expectedBrokerId),
+		KafkaCRLabelKey:     expectedKafkaCRName,
+		"test_label_key":    "test_label_value",
+		ProcessRolesKey:     "broker",
+		IsBrokerNodeKey:     "true",
+		IsControllerNodeKey: "false",
 	}
 
 	brokerConfig := &BrokerConfig{
@@ -455,7 +567,7 @@ func TestGetBrokerLabels(t *testing.T) {
 		},
 	}
 
-	result := brokerConfig.GetBrokerLabels(expectedKafkaCRName, expectedBrokerId)
+	result := brokerConfig.GetBrokerLabels(expectedKafkaCRName, expectedBrokerId, true)
 
 	if !reflect.DeepEqual(result, expected) {
 		t.Error("Expected:", expected, "Got:", result)
 
@@ -53,7 +53,7 @@ spec:
       brokerConfig:
         processRoles:
           - controller
-          - broker
+          # - broker
     - id: 4
       brokerConfigGroup: "default"
       brokerConfig:
 
@@ -525,48 +525,90 @@ func expectKafkaCRStatus(ctx context.Context, kafkaCluster *v1beta1.KafkaCluster
 	Expect(kafkaCluster.Status.State).To(Equal(v1beta1.KafkaClusterRunning))
 	Expect(kafkaCluster.Status.AlertCount).To(Equal(0))
 
-	Expect(kafkaCluster.Status.ListenerStatuses).To(Equal(v1beta1.ListenerStatuses{
-		InternalListeners: map[string]v1beta1.ListenerStatusList{
-			"internal": {
-				{
-					Name:    "any-broker",
-					Address: fmt.Sprintf("%s-all-broker.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
-				},
-				{
-					Name:    "broker-0",
-					Address: fmt.Sprintf("%s-0.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
-				},
-				{
-					Name:    "broker-1",
-					Address: fmt.Sprintf("%s-1.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
-				},
-				{
-					Name:    "broker-2",
-					Address: fmt.Sprintf("%s-2.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+	if kafkaCluster.Spec.KRaftMode == false {
+		Expect(kafkaCluster.Status.ListenerStatuses).To(Equal(v1beta1.ListenerStatuses{
+			InternalListeners: map[string]v1beta1.ListenerStatusList{
+				"internal": {
+					{
+						Name:    "any-broker",
+						Address: fmt.Sprintf("%s-all-broker.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+					},
+					{
+						Name:    "broker-0",
+						Address: fmt.Sprintf("%s-0.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+					},
+					{
+						Name:    "broker-1",
+						Address: fmt.Sprintf("%s-1.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+					},
+					{
+						Name:    "broker-2",
+						Address: fmt.Sprintf("%s-2.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+					},
 				},
 			},
-		},
-		ExternalListeners: map[string]v1beta1.ListenerStatusList{
-			"test": {
-				{
-					Name:    "any-broker",
-					Address: "test.host.com:29092",
-				},
-				{
-					Name:    "broker-0",
-					Address: "test.host.com:19090",
+			ExternalListeners: map[string]v1beta1.ListenerStatusList{
+				"test": {
+					{
+						Name:    "any-broker",
+						Address: "test.host.com:29092",
+					},
+					{
+						Name:    "broker-0",
+						Address: "test.host.com:19090",
+					},
+					{
+						Name:    "broker-1",
+						Address: "test.host.com:19091",
+					},
+					{
+						Name:    "broker-2",
+						Address: "test.host.com:19092",
+					},
 				},
-				{
-					Name:    "broker-1",
-					Address: "test.host.com:19091",
+			},
+		}))
+	}
+	if kafkaCluster.Spec.KRaftMode == true {
+		Expect(kafkaCluster.Status.ListenerStatuses).To(Equal(v1beta1.ListenerStatuses{
+			InternalListeners: map[string]v1beta1.ListenerStatusList{
+				"internal": {
+					{
+						Name:    "any-broker",
+						Address: fmt.Sprintf("%s-all-broker.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+					},
+					{
+						Name:    "broker-0",
+						Address: fmt.Sprintf("%s-0.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+					},
+					{
+						Name:    "broker-2",
+						Address: fmt.Sprintf("%s-2.%s.svc.cluster.local:29092", kafkaCluster.Name, kafkaCluster.Namespace),
+					},
 				},
-				{
-					Name:    "broker-2",
-					Address: "test.host.com:19092",
+			},
+			ExternalListeners: map[string]v1beta1.ListenerStatusList{
+				"test": {
+					{
+						Name:    "any-broker",
+						Address: "test.host.com:29092",
+					},
+					{
+						Name:    "broker-0",
+						Address: "test.host.com:19090",
+					},
+					{
+						Name:    "broker-1",
+						Address: "test.host.com:19091",
+					},
+					{
+						Name:    "broker-2",
+						Address: "test.host.com:19092",
+					},
 				},
 			},
-		},
-	}))
+		}))
+	}
 	for _, brokerState := range kafkaCluster.Status.BrokersState {
 		Expect(brokerState.Version).To(Equal("3.4.1"))
 		Expect(brokerState.Image).To(Equal(kafkaCluster.Spec.GetClusterImage()))
 
@@ -29,9 +29,10 @@ import (
 )
 
 const (
-	headlessServiceJMXTemplate = "http://%s-%d." + kafka.HeadlessServiceTemplate + ".%s.svc.%s:%d"
-	serviceJMXTemplate         = "http://%s-%d.%s.svc.%s:%d"
-	versionRegexGroup          = "version"
+	headlessServiceJMXTemplate           = "http://%s-%d." + kafka.HeadlessServiceTemplate + ".%s.svc.%s:%d"
+	headlessControllerServiceJMXTemplate = "http://%s-%d." + kafka.HeadlessControllerServiceTemplate + ".%s.svc.%s:%d"
+	serviceJMXTemplate                   = "http://%s-%d.%s.svc.%s:%d"
+	versionRegexGroup                    = "version"
 )
 
 var newJMXExtractor = createNewJMXExtractor
@@ -74,9 +75,14 @@ func NewMockJMXExtractor() {
 func (exp *jmxExtractor) ExtractDockerImageAndVersion(brokerId int32, brokerConfig *v1beta1.BrokerConfig,
 	clusterImage string, headlessServiceEnabled bool) (*v1beta1.KafkaVersion, error) {
 	var requestURL string
+
 	if headlessServiceEnabled {
+		var jmxTemplate = headlessServiceJMXTemplate
+		if brokerConfig.IsControllerNode() {
+			jmxTemplate = headlessControllerServiceJMXTemplate
+		}
 		requestURL =
-			fmt.Sprintf(headlessServiceJMXTemplate,
+			fmt.Sprintf(jmxTemplate,
 				exp.clusterName, brokerId, exp.clusterName, exp.clusterNamespace,
 				exp.kubernetesClusterDomain, 9020)
 	} else {