more refactoring, seperate client hooks for each package

Author: rikimaru0345

e2e/client_hooks.go

@@ -0,0 +1,93 @@
+package e2e
+
+import (
+	"net"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/twmb/franz-go/pkg/kgo"
+	"go.uber.org/zap"
+)
+
+// in e2e we only use client hooks for logging connect/disconnect messages
+type clientHooks struct {
+	logger *zap.Logger
+
+	requestSentCount prometheus.Counter
+	bytesSent        prometheus.Counter
+
+	requestsReceivedCount prometheus.Counter
+	bytesReceived         prometheus.Counter
+}
+
+func newEndToEndClientHooks(logger *zap.Logger, metricsNamespace string) *clientHooks {
+	requestSentCount := promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: metricsNamespace,
+		Subsystem: "kafka",
+		Name:      "requests_sent_total"})
+	bytesSent := promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: metricsNamespace,
+		Subsystem: "kafka",
+		Name:      "sent_bytes",
+	})
+
+	requestsReceivedCount := promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: metricsNamespace,
+		Subsystem: "kafka",
+		Name:      "requests_received_total"})
+	bytesReceived := promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: metricsNamespace,
+		Subsystem: "kafka",
+		Name:      "received_bytes",
+	})
+
+	return &clientHooks{
+		logger: logger,
+
+		requestSentCount: requestSentCount,
+		bytesSent:        bytesSent,
+
+		requestsReceivedCount: requestsReceivedCount,
+		bytesReceived:         bytesReceived,
+	}
+}
+
+func (c clientHooks) OnConnect(meta kgo.BrokerMetadata, dialDur time.Duration, _ net.Conn, err error) {
+	if err != nil {
+		c.logger.Debug("kafka connection failed", zap.String("broker_host", meta.Host), zap.Error(err))
+		return
+	}
+	c.logger.Debug("kafka connection succeeded",
+		zap.String("host", meta.Host),
+		zap.Duration("dial_duration", dialDur))
+}
+
+func (c clientHooks) OnDisconnect(meta kgo.BrokerMetadata, _ net.Conn) {
+	c.logger.Debug("kafka broker disconnected",
+		zap.String("host", meta.Host))
+}
+
+// OnRead is passed the broker metadata, the key for the response that
+// was read, the number of bytes read, how long the Client waited
+// before reading the response, how long it took to read the response,
+// and any error.
+//
+// The bytes written does not count any tls overhead.
+// OnRead is called after a read from a broker.
+func (c clientHooks) OnRead(_ kgo.BrokerMetadata, _ int16, bytesRead int, _, _ time.Duration, _ error) {
+	c.requestsReceivedCount.Inc()
+	c.bytesReceived.Add(float64(bytesRead))
+}
+
+// OnWrite is passed the broker metadata, the key for the request that
+// was written, the number of bytes written, how long the request
+// waited before being written, how long it took to write the request,
+// and any error.
+//
+// The bytes written does not count any tls overhead.
+// OnWrite is called after a write to a broker.
+func (c clientHooks) OnWrite(_ kgo.BrokerMetadata, _ int16, bytesWritten int, _, _ time.Duration, _ error) {
+	c.requestSentCount.Inc()
+	c.bytesSent.Add(float64(bytesWritten))
+}

e2e/consumer.go

@@ -12,7 +12,7 @@ import (
 )
 
 func (s *Service) ConsumeFromManagementTopic(ctx context.Context) error {
-	client := s.kafkaSvc.Client
+	client := s.client
 	topicName := s.config.TopicManagement.Name
 	topic := kgo.ConsumeTopics(kgo.NewOffset().AtEnd(), topicName)
 	balancer := kgo.Balancers(kgo.CooperativeStickyBalancer()) // Default GroupBalancer

e2e/producer.go

@@ -3,11 +3,11 @@ package e2e
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/twmb/franz-go/pkg/kgo"
+	"go.uber.org/zap"
 )
 
 type EndToEndMessage struct {
@@ -32,13 +32,13 @@ func (s *Service) produceToManagementTopic(ctx context.Context) error {
 			startTime := timeNowMs()
 			s.endToEndMessagesProduced.Inc()
 
-			err = s.kafkaSvc.Client.Produce(ctx, record, func(r *kgo.Record, err error) {
+			err = s.client.Produce(ctx, record, func(r *kgo.Record, err error) {
 				endTime := timeNowMs()
 				ackDurationMs := endTime - startTime
 				ackDuration := time.Duration(ackDurationMs) * time.Millisecond
 
 				if err != nil {
-					fmt.Printf("record had a produce error: %v\n", err)
+					s.logger.Error("error producing record: %w", zap.Error(err))
 				} else {
 					s.onAck(r.Partition, ackDuration)
 				}

e2e/service.go

@@ -37,13 +37,18 @@ type Service struct {
 }
 
 // NewService creates a new instance of the e2e moinitoring service (wow)
-func NewService(cfg Config, logger *zap.Logger, kafkaSvc *kafka.Service, metricNamespace string) (*Service, error) {
+func NewService(cfg Config, logger *zap.Logger, kafkaSvc *kafka.Service, metricNamespace string, ctx context.Context) (*Service, error) {
+
+	client, err := createKafkaClient(cfg, logger, kafkaSvc, ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create kafka client for e2e: %w", err)
+	}
 
 	svc := &Service{
 		config:   cfg,
-		logger:   logger,
+		logger:   logger.With(zap.String("source", "end_to_end")),
 		kafkaSvc: kafkaSvc,
-		client:   nil,
+		client:   client,
 
 		minionID: uuid.NewString(),
 	}
@@ -83,6 +88,28 @@ func NewService(cfg Config, logger *zap.Logger, kafkaSvc *kafka.Service, metricN
 	return svc, nil
 }
 
+func createKafkaClient(cfg Config, logger *zap.Logger, kafkaSvc *kafka.Service, ctx context.Context) (*kgo.Client, error) {
+
+	// Add RequiredAcks, as options can't be altered later
+	kgoOpts := []kgo.Opt{}
+	if cfg.Enabled {
+		ack := kgo.AllISRAcks()
+		if cfg.Producer.RequiredAcks == 1 {
+			ack = kgo.LeaderAck()
+			kgoOpts = append(kgoOpts, kgo.DisableIdempotentWrite())
+		}
+		kgoOpts = append(kgoOpts, kgo.RequiredAcks(ack))
+	}
+
+	// Prepare hooks
+	hooksChildLogger := logger.With(zap.String("source", "end_to_end"))
+	e2eHooks := newEndToEndClientHooks(hooksChildLogger, "")
+	kgoOpts = append(kgoOpts, kgo.WithHooks(e2eHooks))
+
+	// Create kafka service and check if client can successfully connect to Kafka cluster
+	return kafkaSvc.CreateAndTestClient(logger, kgoOpts, ctx)
+}
+
 // Start starts the service (wow)
 func (s *Service) Start(ctx context.Context) error {
 

e2e/topic.go

@@ -63,7 +63,7 @@ func (s *Service) validateManagementTopic(ctx context.Context) error {
 
 		create := kmsg.NewCreatePartitionsRequest()
 		create.Topics = []kmsg.CreatePartitionsRequestTopic{topic}
-		_, err := create.RequestWith(ctx, s.kafkaSvc.Client)
+		_, err := create.RequestWith(ctx, s.client)
 		if err != nil {
 			return fmt.Errorf("failed to do kmsg request on creating partitions: %w", err)
 		}
@@ -114,7 +114,7 @@ func (s *Service) validateManagementTopic(ctx context.Context) error {
 		reassignment := kmsg.NewAlterPartitionAssignmentsRequest()
 		reassignment.Topics = []kmsg.AlterPartitionAssignmentsRequestTopic{managamentTopicReassignment}
 
-		_, err := reassignment.RequestWith(ctx, s.kafkaSvc.Client)
+		_, err := reassignment.RequestWith(ctx, s.client)
 		if err != nil {
 			return fmt.Errorf("failed to do kmsg request on topic reassignment: %w", err)
 		}
@@ -187,7 +187,7 @@ func (s *Service) createManagementTopic(ctx context.Context, topicMetadata *kmsg
 	req := kmsg.NewCreateTopicsRequest()
 	req.Topics = []kmsg.CreateTopicsRequestTopic{topic}
 
-	res, err := req.RequestWith(ctx, s.kafkaSvc.Client)
+	res, err := req.RequestWith(ctx, s.client)
 	// Sometimes it won't throw Error, but the Error will be abstracted to res.Topics[0].ErrorMessage
 	if res.Topics[0].ErrorMessage != nil {
 		return fmt.Errorf("failed to create topic: %s", *res.Topics[0].ErrorMessage)
@@ -209,7 +209,7 @@ func (s *Service) getTopicMetadata(ctx context.Context) (*kmsg.MetadataResponse,
 	req := kmsg.NewMetadataRequest()
 	req.Topics = []kmsg.MetadataRequestTopic{topicReq}
 
-	res, err := req.RequestWith(ctx, s.kafkaSvc.Client)
+	res, err := req.RequestWith(ctx, s.client)
 	if err != nil {
 		return nil, fmt.Errorf("failed to request metadata: %w", err)
 	}

kafka/client_config_helper.go

@@ -24,7 +24,8 @@ import (
 
 // NewKgoConfig creates a new Config for the Kafka Client as exposed by the franz-go library.
 // If TLS certificates can't be read an error will be returned.
-func NewKgoConfig(cfg Config, logger *zap.Logger, hooks kgo.Hook) ([]kgo.Opt, error) {
+// logger is only used to print warnings about TLS.
+func NewKgoConfig(cfg Config, logger *zap.Logger) ([]kgo.Opt, error) {
 	opts := []kgo.Opt{
 		kgo.SeedBrokers(cfg.Brokers...),
 		kgo.MaxVersions(kversion.V2_7_0()),
@@ -35,13 +36,10 @@ func NewKgoConfig(cfg Config, logger *zap.Logger, hooks kgo.Hook) ([]kgo.Opt, er
 
 	// Create Logger
 	kgoLogger := KgoZapLogger{
-		logger: logger.With(zap.String("source", "kafka_client")).Sugar(),
+		logger: logger.Sugar(),
 	}
 	opts = append(opts, kgo.WithLogger(kgoLogger))
 
-	// Attach hooks
-	opts = append(opts, kgo.WithHooks(hooks))
-
 	// Add Rack Awareness if configured
 	if cfg.RackID != "" {
 		opts = append(opts, kgo.Rack(cfg.RackID))

kafka/service.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/twmb/franz-go/pkg/kerr"
 	"github.com/twmb/franz-go/pkg/kgo"
@@ -14,52 +15,61 @@ import (
 
 type Service struct {
 	cfg    Config
-	Client *kgo.Client
 	logger *zap.Logger
 }
 
-func NewService(cfg Config, logger *zap.Logger, opts []kgo.Opt) (*Service, error) {
-	// Create Kafka Client
-	hooksChildLogger := logger.With(zap.String("source", "kafka_client_hooks"))
-	clientHooks := newClientHooks(hooksChildLogger, "")
-
-	kgoOpts, err := NewKgoConfig(cfg, logger, clientHooks)
-	for _, opt := range opts {
-		kgoOpts = append(kgoOpts, opt)
+func NewService(cfg Config, logger *zap.Logger) *Service {
+	return &Service{
+		cfg:    cfg,
+		logger: logger.With(zap.String("source", "kafka_service")),
 	}
+}
+
+// Create a client with the services default settings
+// logger: will be used to log connections, errors, warnings about tls config, ...
+func (s *Service) CreateAndTestClient(logger *zap.Logger, opts []kgo.Opt, ctx context.Context) (*kgo.Client, error) {
+	// Config with default options
+	kgoOpts, err := NewKgoConfig(s.cfg, logger)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create a valid kafka Client config: %w", err)
 	}
+	// Append user (the service calling this method) provided options
+	kgoOpts = append(kgoOpts, opts...)
 
-	kafkaClient, err := kgo.NewClient(kgoOpts...)
+	// Create kafka client
+	client, err := kgo.NewClient(kgoOpts...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create kafka Client: %w", err)
 	}
 
-	return &Service{
-		cfg:    cfg,
-		Client: kafkaClient,
-		logger: logger,
-	}, nil
+	// Test connection
+	connectCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
+	defer cancel()
+	err = s.testConnection(client, connectCtx)
+	if err != nil {
+		logger.Fatal("failed to test connectivity to Kafka cluster", zap.Error(err))
+	}
+
+	return client, nil
 }
 
-// TestConnection tries to fetch Broker metadata and prints some information if connection succeeds. An error will be
+// testConnection tries to fetch Broker metadata and prints some information if connection succeeds. An error will be
 // returned if connecting fails.
-func (s *Service) TestConnection(ctx context.Context) error {
+func (s *Service) testConnection(client *kgo.Client, ctx context.Context) error {
 	s.logger.Info("connecting to Kafka seed brokers, trying to fetch cluster metadata",
 		zap.String("seed_brokers", strings.Join(s.cfg.Brokers, ",")))
 
 	req := kmsg.MetadataRequest{
 		Topics: nil,
 	}
-	res, err := req.RequestWith(ctx, s.Client)
+	res, err := req.RequestWith(ctx, client)
 	if err != nil {
 		return fmt.Errorf("failed to request metadata: %w", err)
 	}
 
 	// Request versions in order to guess Kafka Cluster version
 	versionsReq := kmsg.NewApiVersionsRequest()
-	versionsRes, err := versionsReq.RequestWith(ctx, s.Client)
+	versionsRes, err := versionsReq.RequestWith(ctx, client)
 	if err != nil {
 		return fmt.Errorf("failed to request api versions: %w", err)
 	}