Adding support for specify input kafka offset for each partitions (#1242)

Adding a new flag `--input-kafka-offset` for supporting consumes from Kafka by specified offset.

input_kafka.go

@@ -3,7 +3,9 @@ package goreplay
 import (
 	"encoding/json"
 	"log"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/Shopify/sarama/mocks"
@@ -12,14 +14,24 @@ import (
 // KafkaInput is used for receiving Kafka messages and
 // transforming them into HTTP payloads.
 type KafkaInput struct {
-	config    *InputKafkaConfig
-	consumers []sarama.PartitionConsumer
-	messages  chan *sarama.ConsumerMessage
-	quit      chan struct{}
+	config      *InputKafkaConfig
+	consumers   []sarama.PartitionConsumer
+	messages    chan *sarama.ConsumerMessage
+	speedFactor float64
+	quit        chan struct{}
+	kafkaTimer  *kafkaTimer
+}
+
+func getOffsetOfPartitions(offsetCfg string) int64 {
+	offset, err := strconv.ParseInt(offsetCfg, 10, 64)
+	if err != nil || offset < -2 {
+		log.Fatalln("Failed to parse offset: "+offsetCfg, err)
+	}
+	return offset
 }
 
 // NewKafkaInput creates instance of kafka consumer client with TLS config
-func NewKafkaInput(_ string, config *InputKafkaConfig, tlsConfig *KafkaTLSConfig) *KafkaInput {
+func NewKafkaInput(offsetCfg string, config *InputKafkaConfig, tlsConfig *KafkaTLSConfig) *KafkaInput {
 	c := NewKafkaConfig(&config.SASLConfig, tlsConfig)
 
 	var con sarama.Consumer
@@ -41,14 +53,17 @@ func NewKafkaInput(_ string, config *InputKafkaConfig, tlsConfig *KafkaTLSConfig
 	}
 
 	i := &KafkaInput{
-		config:    config,
-		consumers: make([]sarama.PartitionConsumer, len(partitions)),
-		messages:  make(chan *sarama.ConsumerMessage, 256),
-		quit:      make(chan struct{}),
+		config:      config,
+		consumers:   make([]sarama.PartitionConsumer, len(partitions)),
+		messages:    make(chan *sarama.ConsumerMessage, 256),
+		speedFactor: 1,
+		quit:        make(chan struct{}),
+		kafkaTimer:  new(kafkaTimer),
 	}
+	i.config.Offset = offsetCfg
 
 	for index, partition := range partitions {
-		consumer, err := con.ConsumePartition(config.Topic, partition, sarama.OffsetNewest)
+		consumer, err := con.ConsumePartition(config.Topic, partition, getOffsetOfPartitions(offsetCfg))
 		if err != nil {
 			log.Fatalln("Failed to start Sarama(Kafka) partition consumer:", err)
 		}
@@ -86,12 +101,15 @@ func (i *KafkaInput) PluginRead() (*Message, error) {
 	case message = <-i.messages:
 	}
 
+	inputTs := ""
+
 	msg.Data = message.Value
 	if i.config.UseJSON {
 
 		var kafkaMessage KafkaMessage
 		json.Unmarshal(message.Value, &kafkaMessage)
 
+		inputTs = kafkaMessage.ReqTs
 		var err error
 		msg.Data, err = kafkaMessage.Dump()
 		if err != nil {
@@ -103,8 +121,11 @@ func (i *KafkaInput) PluginRead() (*Message, error) {
 	// does it have meta
 	if isOriginPayload(msg.Data) {
 		msg.Meta, msg.Data = payloadMetaWithBody(msg.Data)
+		inputTs = string(payloadMeta(msg.Meta)[2])
 	}
 
+	i.timeWait(inputTs)
+
 	return &msg, nil
 
 }
@@ -118,3 +139,45 @@ func (i *KafkaInput) Close() error {
 	close(i.quit)
 	return nil
 }
+
+func (i *KafkaInput) timeWait(curInputTs string) {
+	if i.config.Offset == "-1" || curInputTs == "" {
+		return
+	}
+
+	// implement for Kafka input showdown or speedup emitting
+	timer := i.kafkaTimer
+	curTs := time.Now().UnixNano()
+
+	curInput, err := strconv.ParseInt(curInputTs, 10, 64)
+
+	if timer.latestInputTs == 0 || timer.latestOutputTs == 0 {
+		timer.latestInputTs = curInput
+		timer.latestOutputTs = curTs
+		return
+	}
+
+	if err != nil {
+		log.Fatalln("Fatal to parse timestamp err: ", err)
+	}
+
+	diffTs := curInput - timer.latestInputTs
+	pastTs := curTs - timer.latestOutputTs
+
+	diff := diffTs - pastTs
+	if i.speedFactor != 1 {
+		diff = int64(float64(diff) / i.speedFactor)
+	}
+
+	if diff > 0 {
+		time.Sleep(time.Duration(diff))
+	}
+
+	timer.latestInputTs = curInput
+	timer.latestOutputTs = curTs
+}
+
+type kafkaTimer struct {
+	latestInputTs  int64
+	latestOutputTs int64
+}