github/goreplay
4
0
Fork 0
mirror of https://github.com/buger/goreplay.git synced 2024-04-21 12:32:02 +00:00
Code Issues Releases Wiki Activity

Add Kafka tests and raw format mode (#415)

Browse Source 
* Add tests and send messages in raw format by default
This commit is contained in:
Leonid Bugaev
2017-01-20 19:13:30 +04:00
committed by GitHub
parent 4ea306b5a9
commit def3c91898
3 changed files with 100 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
output_kafka.go
+46 -22
View File
@@ -3,6 +3,7 @@ package main
import (
"encoding/json"
"github.com/Shopify/sarama"
"github.com/Shopify/sarama/mocks"
"github.com/buger/gor/proto"
"io"
"log"
@@ -13,8 +14,10 @@ import (
// KafkaConfig should contains required information to
// build producers.
type KafkaConfig struct {
host string
topic string
host string
topic string
producer sarama.AsyncProducer
useJSON bool
}
// KafkaOutput should make producer client.
@@ -27,6 +30,9 @@ type KafkaOutput struct {
// passed as Json to Apache Kafka.
type KafkaMessage struct {
ReqURL string `json:"Req_URL"`
ReqType string `json:"Req_Type"`
ReqID string `json:"Req_ID"`
ReqTs string `json:"Req_Ts"`
ReqMethod string `json:"Req_Method"`
ReqBody string `json:"Req_Body,omitempty"`
ReqHeaders map[string]string `json:"Req_Headers,omitempty"`
@@ -38,15 +44,23 @@ const KafkaOutputFrequency = 500
// NewKafkaOutput creates instance of kafka producer client.
func NewKafkaOutput(address string, config *KafkaConfig) io.Writer {
c := sarama.NewConfig()
c.Producer.RequiredAcks = sarama.WaitForLocal
c.Producer.Compression = sarama.CompressionSnappy
c.Producer.Flush.Frequency = KafkaOutputFrequency * time.Millisecond
brokerList := strings.Split(config.host, ",")
var producer sarama.AsyncProducer
producer, err := sarama.NewAsyncProducer(brokerList, c)
if err != nil {
log.Fatalln("Failed to start Sarama(Kafka) producer:", err)
if config.producer.(*mocks.AsyncProducer) != nil {
producer = config.producer
} else {
c.Producer.RequiredAcks = sarama.WaitForLocal
c.Producer.Compression = sarama.CompressionSnappy
c.Producer.Flush.Frequency = KafkaOutputFrequency * time.Millisecond
brokerList := strings.Split(config.host, ",")
var err error
producer, err = sarama.NewAsyncProducer(brokerList, c)
if err != nil {
log.Fatalln("Failed to start Sarama(Kafka) producer:", err)
}
}
o := &KafkaOutput{
@@ -70,22 +84,32 @@ func (o *KafkaOutput) ErrorHandler() {
}
func (o *KafkaOutput) Write(data []byte) (n int, err error) {
headers := make(map[string]string)
proto.ParseHeaders([][]byte{data}, func(header []byte, value []byte) bool {
headers[string(header)] = string(value)
return true
})
var message sarama.StringEncoder
req := payloadBody(data)
if !o.config.useJSON {
message = sarama.StringEncoder(data)
} else {
headers := make(map[string]string)
proto.ParseHeaders([][]byte{data}, func(header []byte, value []byte) bool {
headers[string(header)] = string(value)
return true
})
kafkaMessage := KafkaMessage{
ReqURL: string(proto.Path(req)),
ReqMethod: string(proto.Method(req)),
ReqBody: string(proto.Body(req)),
ReqHeaders: headers,
meta := payloadMeta(data)
req := payloadBody(data)
kafkaMessage := KafkaMessage{
ReqURL: string(proto.Path(req)),
ReqType: string(meta[0]),
ReqID: string(meta[1]),
ReqTs: string(meta[2]),
ReqMethod: string(proto.Method(req)),
ReqBody: string(proto.Body(req)),
ReqHeaders: headers,
}
jsonMessage, _ := json.Marshal(&kafkaMessage)
message = sarama.StringEncoder(jsonMessage)
}
jsonMessage, _ := json.Marshal(&kafkaMessage)
message := sarama.StringEncoder(jsonMessage)
o.producer.Input() <- &sarama.ProducerMessage{
Topic: o.config.topic,
output_kafka_test.go
+53
View File
@@ -0,0 +1,53 @@
package main
import (
"github.com/Shopify/sarama"
"github.com/Shopify/sarama/mocks"
"testing"
)
func TestOutputKafkaRAW(t *testing.T) {
config := sarama.NewConfig()
config.Producer.Return.Successes = true
producer := mocks.NewAsyncProducer(t, config)
producer.ExpectInputAndSucceed()
output := NewKafkaOutput("", &KafkaConfig{
producer: producer,
topic: "test",
useJSON: false,
})
output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n"))
resp := <-producer.Successes()
data, _ := resp.Value.Encode()
if string(data) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" {
t.Error("Message not properly encoded: ", string(data))
}
}
func TestOutputKafkaJSON(t *testing.T) {
config := sarama.NewConfig()
config.Producer.Return.Successes = true
producer := mocks.NewAsyncProducer(t, config)
producer.ExpectInputAndSucceed()
output := NewKafkaOutput("", &KafkaConfig{
producer: producer,
topic: "test",
useJSON: true,
})
output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n"))
resp := <-producer.Successes()
data, _ := resp.Value.Encode()
if string(data) != `{"Req_URL":"/","Req_Type":"1","Req_ID":"2","Req_Ts":"3","Req_Method":"GET","Req_Headers":{"Header":"1"}}` {
t.Error("Message not properly encoded: ", string(data))
}
}
settings.go
+1
View File
@@ -130,6 +130,7 @@ func init() {
flag.StringVar(&Settings.outputKafkaConfig.host, "output-kafka-host", "", "Send request and response stats to Kafka:\n\tgor --input-raw :8080 --output-kafka-host '192.168.0.1:9092,192.168.0.2:9092'")
flag.StringVar(&Settings.outputKafkaConfig.topic, "output-kafka-topic", "", "Send request and response stats to Kafka:\n\tgor --input-raw :8080 --output-kafka-topic 'kafka-log'")
flag.BoolVar(&Settings.outputKafkaConfig.useJSON, "output-kafka-json-format", false, "If turned on, it will serialize messages from GoReplay text format to JSON.")
flag.Var(&Settings.modifierConfig.headers, "http-set-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --http-set-header 'User-Agent: Gor'")
flag.Var(&Settings.modifierConfig.headers, "output-http-header", "WARNING: `--output-http-header` DEPRECATED, use `--http-set-header` instead")