Merge pull request #1658 from coolwednesday/multiple-kafka-brokers

Multiple kafka brokers

Author: coolwednesday

docs/advanced-guide/using-publisher-subscriber/page.md

@@ -47,10 +47,10 @@ that are specific for the type of message broker user wants to use.
 ---
 
 - `PUBSUB_BROKER`
-- Address to connect to kafka broker.
+- Address to connect to kafka broker. Multiple brokers can be added as comma seperated values.
 - `+`
 -
-- `localhost:9092`
+- `localhost:9092` or `localhost:8087,localhost:8088,localhost:8089`
 - Not empty string
 
 ---

examples/using-subscriber/main_test.go

@@ -25,7 +25,7 @@ func (m *mockMetrics) IncrementCounter(ctx context.Context, name string, labels
 
 func initializeTest(t *testing.T) {
 	c := kafka.New(&kafka.Config{
-		Broker:       "localhost:9092",
+		Brokers:      []string{"localhost:9092"},
 		OffSet:       1,
 		BatchSize:    kafka.DefaultBatchSize,
 		BatchBytes:   kafka.DefaultBatchBytes,

pkg/gofr/container/container.go

@@ -138,8 +138,10 @@ func (c *Container) Create(conf config.Config) {
 				InsecureSkipVerify: conf.Get("KAFKA_TLS_INSECURE_SKIP_VERIFY") == "true",
 			}
 
+			pubsubBrokers := strings.Split(conf.Get("PUBSUB_BROKER"), ",")
+
 			c.PubSub = kafka.New(&kafka.Config{
-				Broker:           conf.Get("PUBSUB_BROKER"),
+				Brokers:          pubsubBrokers,
 				Partition:        partition,
 				ConsumerGroupID:  conf.Get("CONSUMER_ID"),
 				OffSet:           offSet,

pkg/gofr/container/mock_container.go

@@ -33,8 +33,8 @@ type Mocks struct {
 
 type options func(c *Container, ctrl *gomock.Controller) any
 
-//nolint:revive // WithMockHTTPService returns an exported type intentionally; options are internal and subject to change.
-func WithMockHTTPService(httpServiceNames ...string) options {
+func WithMockHTTPService(httpServiceNames ...string) options { //nolint:revive // WithMockHTTPService returns an
+	// exported type intentionally; options are internal and subject to change.
 	return func(c *Container, ctrl *gomock.Controller) any {
 		mockservice := service.NewMockHTTP(ctrl)
 		for _, s := range httpServiceNames {

pkg/gofr/datasource/pubsub/kafka/conn.go

@@ -0,0 +1,188 @@
+package kafka
+
+import (
+	"context"
+	"errors"
+	"net"
+	"strconv"
+	"sync"
+
+	"github.com/segmentio/kafka-go"
+)
+
+//nolint:unused // We need this wrap around for testing purposes.
+type Conn struct {
+	conns []*kafka.Conn
+}
+
+// initialize creates and configures all Kafka client components.
+func (k *kafkaClient) initialize(ctx context.Context) error {
+	dialer, err := setupDialer(&k.config)
+	if err != nil {
+		return err
+	}
+
+	conns, err := connectToBrokers(ctx, k.config.Brokers, dialer, k.logger)
+	if err != nil {
+		return err
+	}
+
+	multi := &multiConn{
+		conns:  conns,
+		dialer: dialer,
+	}
+
+	writer := createKafkaWriter(&k.config, dialer, k.logger)
+	reader := make(map[string]Reader)
+
+	k.logger.Logf("connected to %d Kafka brokers", len(conns))
+
+	k.dialer = dialer
+	k.conn = multi
+	k.writer = writer
+	k.reader = reader
+
+	return nil
+}
+
+func (k *kafkaClient) getNewReader(topic string) Reader {
+	reader := kafka.NewReader(kafka.ReaderConfig{
+		GroupID:     k.config.ConsumerGroupID,
+		Brokers:     k.config.Brokers,
+		Topic:       topic,
+		MinBytes:    10e3,
+		MaxBytes:    10e6,
+		Dialer:      k.dialer,
+		StartOffset: int64(k.config.OffSet),
+	})
+
+	return reader
+}
+
+func (k *kafkaClient) DeleteTopic(_ context.Context, name string) error {
+	return k.conn.DeleteTopics(name)
+}
+
+func (k *kafkaClient) Controller() (broker kafka.Broker, err error) {
+	return k.conn.Controller()
+}
+
+func (k *kafkaClient) CreateTopic(_ context.Context, name string) error {
+	topics := kafka.TopicConfig{Topic: name, NumPartitions: 1, ReplicationFactor: 1}
+
+	return k.conn.CreateTopics(topics)
+}
+
+type multiConn struct {
+	conns  []Connection
+	dialer *kafka.Dialer
+	mu     sync.RWMutex
+}
+
+func (m *multiConn) Controller() (kafka.Broker, error) {
+	if len(m.conns) == 0 {
+		return kafka.Broker{}, errNoActiveConnections
+	}
+
+	// Try all connections until we find one that works
+	for _, conn := range m.conns {
+		if conn == nil {
+			continue
+		}
+
+		controller, err := conn.Controller()
+		if err == nil {
+			return controller, nil
+		}
+	}
+
+	return kafka.Broker{}, errNoActiveConnections
+}
+
+func (m *multiConn) CreateTopics(topics ...kafka.TopicConfig) error {
+	controller, err := m.Controller()
+	if err != nil {
+		return err
+	}
+
+	controllerAddr := net.JoinHostPort(controller.Host, strconv.Itoa(controller.Port))
+
+	controllerResolvedAddr, err := net.ResolveTCPAddr("tcp", controllerAddr)
+	if err != nil {
+		return err
+	}
+
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	for _, conn := range m.conns {
+		if conn == nil {
+			continue
+		}
+
+		if tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr); ok {
+			if tcpAddr.IP.Equal(controllerResolvedAddr.IP) && tcpAddr.Port == controllerResolvedAddr.Port {
+				return conn.CreateTopics(topics...)
+			}
+		}
+	}
+
+	// If not found, create a new connection
+	conn, err := m.dialer.DialContext(context.Background(), "tcp", controllerAddr)
+	if err != nil {
+		return err
+	}
+
+	m.conns = append(m.conns, conn)
+
+	return conn.CreateTopics(topics...)
+}
+
+func (m *multiConn) DeleteTopics(topics ...string) error {
+	controller, err := m.Controller()
+	if err != nil {
+		return err
+	}
+
+	controllerAddr := net.JoinHostPort(controller.Host, strconv.Itoa(controller.Port))
+
+	controllerResolvedAddr, err := net.ResolveTCPAddr("tcp", controllerAddr)
+	if err != nil {
+		return err
+	}
+
+	for _, conn := range m.conns {
+		if conn == nil {
+			continue
+		}
+
+		if tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr); ok {
+			// Match IP (after resolution) and Port
+			if tcpAddr.IP.Equal(controllerResolvedAddr.IP) && tcpAddr.Port == controllerResolvedAddr.Port {
+				return conn.DeleteTopics(topics...)
+			}
+		}
+	}
+
+	// If not found, create a new connection
+	conn, err := m.dialer.DialContext(context.Background(), "tcp", controllerAddr)
+	if err != nil {
+		return err
+	}
+
+	m.conns = append(m.conns, conn)
+
+	return conn.DeleteTopics(topics...)
+}
+
+func (m *multiConn) Close() error {
+	var err error
+
+	for _, conn := range m.conns {
+		if conn != nil {
+			err = errors.Join(err, conn.Close())
+		}
+	}
+
+	return err
+}

pkg/gofr/datasource/pubsub/kafka/errors.go

@@ -0,0 +1,22 @@
+package kafka
+
+import "errors"
+
+var (
+	ErrConsumerGroupNotProvided    = errors.New("consumer group id not provided")
+	errFailedToConnectBrokers      = errors.New("failed to connect to any kafka brokers")
+	errBrokerNotProvided           = errors.New("kafka broker address not provided")
+	errPublisherNotConfigured      = errors.New("can't publish message. Publisher not configured or topic is empty")
+	errBatchSize                   = errors.New("KAFKA_BATCH_SIZE must be greater than 0")
+	errBatchBytes                  = errors.New("KAFKA_BATCH_BYTES must be greater than 0")
+	errBatchTimeout                = errors.New("KAFKA_BATCH_TIMEOUT must be greater than 0")
+	errClientNotConnected          = errors.New("kafka client not connected")
+	errUnsupportedSASLMechanism    = errors.New("unsupported SASL mechanism")
+	errSASLCredentialsMissing      = errors.New("SASL credentials missing")
+	errUnsupportedSecurityProtocol = errors.New("unsupported security protocol")
+	errNoActiveConnections         = errors.New("no active connections to brokers")
+	errCACertFileRead              = errors.New("failed to read CA certificate file")
+	errClientCertLoad              = errors.New("failed to load client certificate")
+	errNotController               = errors.New("not a controller")
+	errUnreachable                 = errors.New("unreachable")
+)

pkg/gofr/datasource/pubsub/kafka/health.go

@@ -2,28 +2,96 @@ package kafka
 
 import (
 	"encoding/json"
+	"fmt"
+	"net"
+	"strconv"
 
 	"gofr.dev/pkg/gofr/datasource"
 )
 
-func (k *kafkaClient) Health() (health datasource.Health) {
-	health = datasource.Health{Details: make(map[string]any)}
+func (k *kafkaClient) Health() datasource.Health {
+	health := datasource.Health{
+		Status:  datasource.StatusDown,
+		Details: make(map[string]any),
+	}
+
+	if k.conn == nil {
+		health.Details["error"] = "invalid connection type"
+		return health
+	}
 
-	health.Status = datasource.StatusUp
+	k.conn.mu.RLock()
+	defer k.conn.mu.RUnlock()
 
-	_, err := k.conn.Controller()
-	if err != nil {
-		health.Status = datasource.StatusDown
+	brokerStatus, allDown := k.evaluateBrokerHealth()
+
+	if !allDown {
+		health.Status = datasource.StatusUp
 	}
 
-	health.Details["host"] = k.config.Broker
+	health.Details["brokers"] = brokerStatus
 	health.Details["backend"] = "KAFKA"
-	health.Details["writers"] = k.getWriterStatsAsMap()
+	health.Details["writer"] = k.getWriterStatsAsMap()
 	health.Details["readers"] = k.getReaderStatsAsMap()
 
 	return health
 }
 
+func (k *kafkaClient) evaluateBrokerHealth() ([]map[string]any, bool) {
+	var (
+		statusList     = make([]map[string]any, 0)
+		controllerAddr string
+		allDown        = true
+	)
+
+	for _, conn := range k.conn.conns {
+		if conn == nil {
+			continue
+		}
+
+		status := checkBroker(conn, &controllerAddr)
+
+		if status["status"] == brokerStatusUp {
+			allDown = false
+		}
+
+		statusList = append(statusList, status)
+	}
+
+	return statusList, allDown
+}
+
+func checkBroker(conn Connection, controllerAddr *string) map[string]any {
+	brokerAddr := conn.RemoteAddr().String()
+	status := map[string]any{
+		"broker":       brokerAddr,
+		"status":       "DOWN",
+		"isController": false,
+		"error":        nil,
+	}
+
+	_, err := conn.ReadPartitions()
+	if err != nil {
+		status["error"] = err.Error()
+		return status
+	}
+
+	status["status"] = brokerStatusUp
+
+	if *controllerAddr == "" {
+		controller, err := conn.Controller()
+		if err != nil {
+			status["error"] = fmt.Sprintf("controller lookup failed: %v", err)
+		} else if controller.Host != "" {
+			*controllerAddr = net.JoinHostPort(controller.Host, strconv.Itoa(controller.Port))
+		}
+	}
+
+	status["isController"] = brokerAddr == *controllerAddr
+
+	return status
+}
+
 func (k *kafkaClient) getReaderStatsAsMap() []any {
 	readerStats := make([]any, 0)