Add kafka output

Author: oivoodoo

emitter.go

@@ -15,7 +15,7 @@ func Start(stop chan int) {
 			middleware.ReadFrom(in)
 		}
 
-		// We going only to read responses, so using same ReadFrom method
+		// We are going only to read responses, so using same ReadFrom method
 		for _, out := range Plugins.Outputs {
 			if r, ok := out.(io.Reader); ok {
 				middleware.ReadFrom(r)

output_kafka.go

@@ -0,0 +1,68 @@
+package main
+
+import (
+	"github.com/Shopify/sarama"
+	"log"
+	"strings"
+	"time"
+)
+
+// KafkaConfig should contains required information to
+// build producers.
+type KafkaConfig struct {
+	zookeeper string
+	topic     string
+}
+
+// KafkaOutput should make producer client.
+type KafkaOutput struct {
+	address  string
+	config   *KafkaConfig
+	producer sarama.AsyncProducer
+}
+
+// NewKafkaOutput creates instance of kafka producer client.
+func NewKafkaOutput(address string, config *KafkaConfig) *KafkaOutput {
+	c := sarama.NewConfig()
+	c.Producer.RequiredAcks = sarama.WaitForLocal
+	c.Producer.Compression = sarama.CompressionSnappy
+	c.Producer.Flush.Frequency = 500 * time.Millisecond
+
+	brokerList := strings.Split(config.zookeeper, ",")
+
+	producer, err := sarama.NewAsyncProducer(brokerList, c)
+	if err != nil {
+		log.Fatalln("Failed to start Sarama(Kafka) producer:", err)
+	}
+
+	o := &KafkaOutput{
+		address:  address,
+		config:   config,
+		producer: producer,
+	}
+
+	// Start infinite loop for tracking errors for kafka producer.
+	go o.ErrorHandler()
+
+	return o
+}
+
+// ErrorHandler should receive errors
+func (o *KafkaOutput) ErrorHandler() {
+	for err := range o.producer.Errors() {
+		log.Println("Failed to write access log entry:", err)
+	}
+}
+
+func (o *KafkaOutput) Write(data []byte) (n int, err error) {
+	buf := make(sarama.ByteEncoder, len(data))
+	copy(buf, data)
+
+	o.producer.Input() <- &sarama.ProducerMessage{
+		Topic: o.config.topic,
+		Key:   sarama.StringEncoder(o.address),
+		Value: buf,
+	}
+
+	return len(data), nil
+}

plugins.go

@@ -142,4 +142,6 @@ func InitPlugins() {
 	for _, options := range Settings.outputHTTP {
 		registerPlugin(NewHTTPOutput, options, &Settings.outputHTTPConfig)
 	}
+
+	registerPlugin(NewKafkaOutput, &Settings.outputKafkaConfig)
 }

settings.go

@@ -58,6 +58,8 @@ type AppSettings struct {
 
 	outputHTTPConfig HTTPOutputConfig
 	modifierConfig   HTTPModifierConfig
+
+	outputKafkaConfig KafkaConfig
 }
 
 // Settings holds Gor configuration
@@ -118,13 +120,15 @@ func init() {
 	flag.IntVar(&Settings.outputHTTPConfig.BufferSize, "output-http-response-buffer", 0, "HTTP response buffer size, all data after this size will be discarded.")
 	flag.IntVar(&Settings.outputHTTPConfig.workers, "output-http-workers", 0, "Gor uses dynamic worker scaling by default.  Enter a number to run a set number of workers.")
 	flag.IntVar(&Settings.outputHTTPConfig.redirectLimit, "output-http-redirects", 0, "Enable how often redirects should be followed.")
-	flag.DurationVar(&Settings.outputHTTPConfig.Timeout, "output-http-timeout", 5 * time.Second, "Specify HTTP request/response timeout. By default 5s. Example: --output-http-timeout 30s")
+	flag.DurationVar(&Settings.outputHTTPConfig.Timeout, "output-http-timeout", 5*time.Second, "Specify HTTP request/response timeout. By default 5s. Example: --output-http-timeout 30s")
 
 	flag.BoolVar(&Settings.outputHTTPConfig.stats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.")
 	flag.BoolVar(&Settings.outputHTTPConfig.OriginalHost, "http-original-host", false, "Normally gor replaces the Host http header with the host supplied with --output-http.  This option disables that behavior, preserving the original Host header.")
 	flag.BoolVar(&Settings.outputHTTPConfig.Debug, "output-http-debug", false, "Enables http debug output.")
 
 	flag.StringVar(&Settings.outputHTTPConfig.elasticSearch, "output-http-elasticsearch", "", "Send request and response stats to ElasticSearch:\n\tgor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'")
+	flag.StringVar(&Settings.outputKafkaConfig.zookeeper, "output-kafka-zookeeper", "", "Send request and response stats to Kafka:\n\tgor --input-raw :8080 --output-kafka-zookeeper '192.168.0.1:2181,192.168.0.2:2181'")
+	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.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")