feat: added connection class to manage connections to Kafka

Author: Gerfey

examples/kafka_transport/kafka_transport.go

@@ -3,13 +3,18 @@ package main
 import (
 	"context"
 	"log/slog"
+	"time"
 
 	"github.com/gerfey/messenger/builder"
 	"github.com/gerfey/messenger/config"
 	"github.com/gerfey/messenger/examples/kafka_transport/handler"
 	"github.com/gerfey/messenger/examples/kafka_transport/message"
 )
 
+const (
+	waitDurationSeconds = 20
+)
+
 func main() {
 	ctx := context.Background()
 
@@ -57,5 +62,5 @@ func main() {
 		return
 	}
 
-	<-ctx.Done()
+	time.Sleep(waitDurationSeconds * time.Second)
 }

transport/kafka/connection.go

@@ -0,0 +1,69 @@
+package kafka
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/segmentio/kafka-go"
+)
+
+const (
+	connectionTimeout = 5 * time.Second
+)
+
+type Connection struct {
+	brokers []string
+	dialer  *kafka.Dialer
+}
+
+func NewConnection(brokers []string) (*Connection, error) {
+	conn := &Connection{
+		brokers: brokers,
+		dialer: &kafka.Dialer{
+			Timeout:   connectionTimeout,
+			DualStack: true,
+		},
+	}
+
+	if err := conn.Check(connectionTimeout); err != nil {
+		return nil, err
+	}
+
+	return conn, nil
+}
+
+func (c *Connection) Check(timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	for _, broker := range c.brokers {
+		conn, connErr := c.dialer.DialContext(ctx, "tcp", broker)
+		if connErr != nil {
+			return fmt.Errorf("failed to connect to Kafka broker at '%s': %w", broker, connErr)
+		}
+
+		if closeErr := conn.Close(); closeErr != nil {
+			return fmt.Errorf("failed to close connection to Kafka broker at '%s': %w", broker, closeErr)
+		}
+	}
+
+	return nil
+}
+
+func (c *Connection) CreateReader(config kafka.ReaderConfig) *kafka.Reader {
+	config.Brokers = c.brokers
+	config.Dialer = c.dialer
+
+	return kafka.NewReader(config)
+}
+
+func (c *Connection) CreateWriter(topic string) *kafka.Writer {
+	return &kafka.Writer{
+		Addr:         kafka.TCP(c.brokers...),
+		Topic:        topic,
+		RequiredAcks: kafka.RequireAll,
+		Balancer:     &kafka.LeastBytes{},
+		Async:        false,
+	}
+}

transport/kafka/consumer.go

@@ -23,18 +23,19 @@ const (
 type Consumer struct {
 	cfg        TransportConfig
 	serializer api.Serializer
+	conn       *Connection
 }
 
-func NewConsumer(cfg TransportConfig, ser api.Serializer) *Consumer {
+func NewConsumer(cfg TransportConfig, ser api.Serializer, conn *Connection) *Consumer {
 	return &Consumer{
 		cfg:        cfg,
 		serializer: ser,
+		conn:       conn,
 	}
 }
 
 func (c *Consumer) Consume(ctx context.Context, handler func(context.Context, api.Envelope) error) error {
-	r := kafka.NewReader(kafka.ReaderConfig{
-		Brokers:         c.cfg.Options.Brokers,
+	readerConfig := kafka.ReaderConfig{
 		GroupID:         c.cfg.Options.Group,
 		Topic:           c.cfg.Options.Topic,
 		StartOffset:     c.startOffset(c.cfg.Options.Offset),
@@ -47,7 +48,9 @@ func (c *Consumer) Consume(ctx context.Context, handler func(context.Context, ap
 		RebalanceTimeout:  rebalanceTimeout,
 		HeartbeatInterval: heartbeatInterval,
 		MaxWait:           time.Second,
-	})
+	}
+
+	r := c.conn.CreateReader(readerConfig)
 	defer r.Close()
 
 	jobs := make(chan job)

transport/kafka/factory_test.go

@@ -97,7 +97,7 @@ func TestTransportFactory_Create(t *testing.T) {
 
 	transport, err := factory.Create(name, dsn, optionsBytes)
 
-	require.NoError(t, err)
-	assert.NotNil(t, transport)
-	assert.IsType(t, &kafka.Transport{}, transport)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "kafka")
+	assert.Nil(t, transport)
 }

transport/kafka/producer.go

@@ -16,17 +16,11 @@ type Producer struct {
 	serializer api.Serializer
 }
 
-func NewProducer(cfg TransportConfig, ser api.Serializer) (*Producer, error) {
+func NewProducer(cfg TransportConfig, ser api.Serializer, conn *Connection) (*Producer, error) {
 	return &Producer{
 		cfg:        cfg,
 		serializer: ser,
-		writer: &kafka.Writer{
-			Addr:         kafka.TCP(cfg.Options.Brokers...),
-			Topic:        cfg.Options.Topic,
-			RequiredAcks: kafka.RequireAll,
-			Balancer:     &kafka.LeastBytes{},
-			Async:        false,
-		},
+		writer:     conn.CreateWriter(cfg.Options.Topic),
 	}, nil
 }
 

transport/kafka/transport.go

@@ -15,24 +15,33 @@ type Transport struct {
 	consumer   *Consumer
 	serializer api.Serializer
 	logger     *slog.Logger
+	conn       *Connection
 }
 
 func NewTransport(cfg TransportConfig, resolver api.TypeResolver, logger *slog.Logger) (api.Transport, error) {
+	conn, err := NewConnection(cfg.Options.Brokers)
+	if err != nil {
+		logger.Error("failed to connect to Kafka brokers", "error", err)
+
+		return nil, fmt.Errorf("kafka: %w", err)
+	}
+
 	ser := serializer.NewSerializer(resolver)
 
-	producer, err := NewProducer(cfg, ser)
+	producer, err := NewProducer(cfg, ser, conn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create kafka producer: %w", err)
 	}
 
-	consumer := NewConsumer(cfg, ser)
+	consumer := NewConsumer(cfg, ser, conn)
 
 	return &Transport{
 		cfg:        cfg,
 		producer:   producer,
 		consumer:   consumer,
 		serializer: ser,
 		logger:     logger,
+		conn:       conn,
 	}, nil
 }