fix e2e producer

e2e/service.go

@@ -155,14 +155,6 @@ func (s *Service) Start(ctx context.Context) error {
 		return fmt.Errorf("could not validate end-to-end topic: %w", err)
 	}
 
-	// Get up-to-date metadata and inform our custom partitioner about the partition count
-	topicMetadata, err := s.getTopicMetadata(ctx)
-	if err != nil {
-		return fmt.Errorf("could not get topic metadata after validation: %w", err)
-	}
-	partitions := len(topicMetadata.Topics[0].Partitions)
-	s.partitionCount = partitions
-
 	// finally start everything else (producing, consuming, continuous validation, consumer group tracking)
 	go s.startReconciliation(ctx)
 

e2e/topic.go

@@ -2,6 +2,7 @@ package e2e
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math"
 	"time"
@@ -13,12 +14,14 @@ import (
 
 // Check our end-to-end test topic and adapt accordingly if something does not match our expectations.
 // - does it exist?
+//
 // - is it configured correctly?
-//     - does it have enough partitions?
-//     - is the replicationFactor correct?
+//   - does it have enough partitions?
+//   - is the replicationFactor correct?
+//
 // - are assignments good?
-//     - is each broker leading at least one partition?
-//     - are replicas distributed correctly?
+//   - is each broker leading at least one partition?
+//   - are replicas distributed correctly?
 func (s *Service) validateManagementTopic(ctx context.Context) error {
 	s.logger.Debug("validating end-to-end topic...")
 
@@ -29,10 +32,10 @@ func (s *Service) validateManagementTopic(ctx context.Context) error {
 
 	typedErr := kerr.TypedErrorForCode(meta.Topics[0].ErrorCode)
 	topicExists := false
-	switch typedErr {
-	case nil:
+	switch {
+	case typedErr == nil:
 		topicExists = true
-	case kerr.UnknownTopicOrPartition:
+	case errors.Is(typedErr, kerr.UnknownTopicOrPartition):
 		// UnknownTopicOrPartition (Error code 3) means that the topic does not exist.
 		// When the topic doesn't exist, continue to create it further down in the code.
 		topicExists = false
@@ -68,7 +71,43 @@ func (s *Service) validateManagementTopic(ctx context.Context) error {
 		return fmt.Errorf("failed to create partitions: %w", err)
 	}
 
-	return nil
+	return s.updatePartitionCount(ctx)
+}
+
+// updatePartitionCount retrieves metadata to inform kminion about the updated
+// partition count of its e2e topic. It must be updated after topic validation
+// because the validation process may lead to the creation of new partitions.
+// This can occur when new brokers are added to the cluster.
+func (s *Service) updatePartitionCount(ctx context.Context) error {
+	retryTicker := time.NewTicker(1 * time.Second)
+	defer retryTicker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-retryTicker.C:
+			meta, err := s.getTopicMetadata(ctx)
+			if err != nil {
+				return fmt.Errorf("could not get topic metadata while updating partition count: %w", err)
+			}
+
+			typedErr := kerr.TypedErrorForCode(meta.Topics[0].ErrorCode)
+			if typedErr == nil {
+				s.partitionCount = len(meta.Topics[0].Partitions)
+				s.logger.Debug("updatePartitionCount: successfully updated partition count", zap.Int("partition_count", s.partitionCount))
+				return nil
+			}
+			if !errors.Is(typedErr, kerr.UnknownTopicOrPartition) {
+				return fmt.Errorf("unexpected error while updating partition count: %w", typedErr)
+			}
+			s.logger.Warn("updatePartitionCount: received UNKNOWN_TOPIC_OR_PARTITION error, possibly due to timing issue. Retrying...")
+			// The UNKNOWN_TOPIC_OR_PARTITION error occurs occasionally even though the topic is created
+			// in the validateManagementTopic function. It appears to be a timing issue where the topic metadata
+			// is not immediately available after creation. In practice, waiting for a short period and then retrying
+			// the operation resolves the issue.
+		}
+	}
 }
 
 func (s *Service) executeCreatePartitions(ctx context.Context, req *kmsg.CreatePartitionsRequest) error {

kafka/client_config_helper.go

@@ -34,6 +34,9 @@ func NewKgoConfig(cfg Config, logger *zap.Logger) ([]kgo.Opt, error) {
 		kgo.ClientID(cfg.ClientID),
 		kgo.FetchMaxBytes(5 * 1000 * 1000), // 5MB
 		kgo.MaxConcurrentFetches(10),
+		// Limit metadata age to mitigate issues with unknown partition/topic
+		// requests shortly after the creation or changing the partitions.
+		kgo.MetadataMaxAge(time.Second),
 	}
 
 	// Create Logger