feat: Rate limit for Kafka and initial offset support and refactor co…

…nfig (#829) * add a rate limit to kafka event source * refactor sarama config and move kafka version up a level * fix most linter issues * fix last lint issue by bypassing * update example * make requested changes * update kafka example for new version * add +optional back in after removal of comments * add more optional based on settings that can be defaulted * remove context based cancelable sleep * remove unused context
argoproj · Aug 20, 2020 · 24caf93 · 24caf93
1 parent 44eccc3
commit 24caf93
Show file tree

Hide file tree

Showing 9 changed files with 567 additions and 378 deletions.
diff --git a/api/event-source.html b/api/event-source.html
diff --git a/api/event-source.md b/api/event-source.md
diff --git a/api/openapi-spec/swagger.json b/api/openapi-spec/swagger.json
diff --git a/eventsources/sources/kafka/start.go b/eventsources/sources/kafka/start.go
@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"strconv"
 	"sync"
+	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/argoproj/argo-events/common"
@@ -81,14 +82,10 @@ func (el *EventListener) StartListening(ctx context.Context, dispatch func([]byt
 }
 
 func (listener *EventListener) consumerGroupConsumer(ctx context.Context, log *zap.SugaredLogger, kafkaEventSource *v1alpha1.KafkaEventSource, dispatch func([]byte) error) error {
-	config := sarama.NewConfig()
-
-	version, err := sarama.ParseKafkaVersion(kafkaEventSource.ConsumerGroup.KafkaVersion)
+	config, err := getSaramaConfig(kafkaEventSource, log)
 	if err != nil {
-		log.Errorf("Error parsing Kafka version: %v", err)
 		return err
 	}
-	config.Version = version
 
 	switch kafkaEventSource.ConsumerGroup.RebalanceStrategy {
 	case "sticky":
@@ -102,15 +99,6 @@ func (listener *EventListener) consumerGroupConsumer(ctx context.Context, log *z
 		config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRange
 	}
 
-	if kafkaEventSource.TLS != nil {
-		tlsConfig, err := common.GetTLSConfig(kafkaEventSource.TLS.CACertPath, kafkaEventSource.TLS.ClientCertPath, kafkaEventSource.TLS.ClientKeyPath)
-		if err != nil {
-			return errors.Wrap(err, "failed to get the tls configuration")
-		}
-		config.Net.TLS.Config = tlsConfig
-		config.Net.TLS.Enable = true
-	}
-
 	consumer := Consumer{
 		ready:            make(chan bool),
 		dispatch:         dispatch,
@@ -169,20 +157,10 @@ func (el *EventListener) partitionConsumer(ctx context.Context, log *zap.Sugared
 	log.Info("connecting to Kafka cluster...")
 	if err := sources.Connect(common.GetConnectionBackoff(kafkaEventSource.ConnectionBackoff), func() error {
 		var err error
-		config := sarama.NewConfig()
 
-		if kafkaEventSource.TLS != nil {
-			tlsConfig, err := common.GetTLSConfig(kafkaEventSource.TLS.CACertPath, kafkaEventSource.TLS.ClientCertPath, kafkaEventSource.TLS.ClientKeyPath)
-			if err != nil {
-				return errors.Wrap(err, "failed to get the tls configuration")
-			}
-			config.Net.TLS.Config = tlsConfig
-			config.Net.TLS.Enable = true
-		} else {
-			consumer, err = sarama.NewConsumer([]string{kafkaEventSource.URL}, nil)
-			if err != nil {
-				return err
-			}
+		config, err := getSaramaConfig(kafkaEventSource, log)
+		if err != nil {
+			return err
 		}
 
 		consumer, err = sarama.NewConsumer([]string{kafkaEventSource.URL}, config)
@@ -259,6 +237,37 @@ func (el *EventListener) partitionConsumer(ctx context.Context, log *zap.Sugared
 	}
 }
 
+func getSaramaConfig(kafkaEventSource *v1alpha1.KafkaEventSource, log *zap.SugaredLogger) (*sarama.Config, error) { //nolint:interfacer
+	config := sarama.NewConfig()
+
+	if kafkaEventSource.Version == "" {
+		config.Version = sarama.V1_0_0_0
+	} else {
+		version, err := sarama.ParseKafkaVersion(kafkaEventSource.Version)
+		if err != nil {
+			log.Errorf("Error parsing Kafka version: %v", err)
+			return nil, err
+		}
+		config.Version = version
+	}
+
+	if kafkaEventSource.TLS != nil {
+		tlsConfig, err := common.GetTLSConfig(kafkaEventSource.TLS.CACertPath, kafkaEventSource.TLS.ClientCertPath, kafkaEventSource.TLS.ClientKeyPath)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to get the tls configuration")
+		}
+		config.Net.TLS.Config = tlsConfig
+		config.Net.TLS.Enable = true
+	}
+
+	if kafkaEventSource.ConsumerGroup != nil {
+		if kafkaEventSource.ConsumerGroup.Oldest {
+			config.Consumer.Offsets.Initial = sarama.OffsetOldest
+		}
+	}
+	return config, nil
+}
+
 // Consumer represents a Sarama consumer group consumer
 type Consumer struct {
 	ready            chan bool
@@ -310,6 +319,12 @@ func (consumer *Consumer) ConsumeClaim(session sarama.ConsumerGroupSession, clai
 		} else {
 			session.MarkMessage(message, "")
 		}
+		if consumer.kafkaEventSource.LimitEventsPerSecond > 0 {
+			//1000000000 is 1 second in nanoseconds
+			d := (1000000000 / time.Duration(consumer.kafkaEventSource.LimitEventsPerSecond) * time.Nanosecond) * time.Nanosecond
+			consumer.logger.Infof("Sleeping for: %v.", d)
+			time.Sleep(d)
+		}
 	}
 
 	return nil

diff --git a/examples/event-sources/kafka.yaml b/examples/event-sources/kafka.yaml
@@ -28,8 +28,10 @@ spec:
 #      Use a consumer group, if this is used you do not need to specify a "partition: <id>"
 #      consumerGroup:
 #        groupName: test-group
-#        kafkaVersion: "2.5.0"
+#        oldest: false
 #        rebalanceStrategy: range
+#      limitEventsPerSecond: 1
+#      version: "2.5.0"
 
 #    example-tls:
 #      url: "kafka.argo-events:9092"