throttling for go-dmux

conf.json

@@ -17,7 +17,8 @@
           "topic": "sample-topic",
           "kafka_version_major": 2,
           "force_restart": true,
-          "read_newest": true
+          "read_newest": true,
+          "throttle_threshold": 2
         },
         "pending_acks": 1000000,
         "sink": {
@@ -40,6 +41,7 @@
     }
   ],
   "metric_port": 9999,
+  "offset_polling_interval": "1s",
   "logging": {
     "type": "file",
     "config": {

config.go

@@ -2,9 +2,11 @@ package main
 
 import (
 	"encoding/json"
+	"github.com/go-dmux/core"
 	"io/ioutil"
 	"log"
 	"os"
+	"time"
 
 	"github.com/go-dmux/connection"
 	"github.com/go-dmux/logging"
@@ -37,19 +39,23 @@ func (c ConnectionType) getConfig(data []byte) interface{} {
 }
 
 //Start invokes Run of the respective connection in a go routine
-func (c ConnectionType) Start(conf interface{}, enableDebug bool) {
+func (c ConnectionType) Start(conf interface{}, enableDebug bool, name string, offsetPollinginterval time.Duration) {
 	switch c {
 	case KafkaHTTP:
 		connObj := &connection.KafkaHTTPConn{
-			EnableDebugLog: enableDebug,
-			Conf:           conf,
+			Name:            name,
+			EnableDebugLog:  enableDebug,
+			Conf:            conf,
+			PollingInterval: offsetPollinginterval,
 		}
 		log.Println("Starting ", KafkaHTTP)
 		connObj.Run()
 	case KafkaFoxtrot:
 		connObj := &connection.KafkaFoxtrotConn{
-			EnableDebugLog: enableDebug,
-			Conf:           conf,
+			Name:            name,
+			EnableDebugLog:  enableDebug,
+			Conf:            conf,
+			PollingInterval: offsetPollinginterval,
 		}
 		log.Println("Starting ", KafkaFoxtrot)
 		connObj.Run()
@@ -70,8 +76,9 @@ type DmuxConf struct {
 	Name      string     `json:"name"`
 	DMuxItems []DmuxItem `json:"dmuxItems"`
 	// DMuxMap    map[string]KafkaHTTPConnConfig `json:"dmuxMap"`
-	MetricPort int	     `json:"metric_port"`
-	Logging logging.LogConf `json:"logging"`
+	MetricPort            int             `json:"metric_port"`
+	OffsetPollingInterval core.Duration   `json:"offset_polling_interval"`
+	Logging               logging.LogConf `json:"logging"`
 }
 
 //DmuxItem struct defines name and type of connection

connection/kafka_foxtrot_conn.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/go-dmux/core"
@@ -22,8 +23,10 @@ type KafkaFoxtrotConnConfig struct {
 
 //KafkaFoxtrotConn struct to abstract this connections Run
 type KafkaFoxtrotConn struct {
-	EnableDebugLog bool
-	Conf           interface{}
+	Name            string
+	EnableDebugLog  bool
+	Conf            interface{}
+	PollingInterval time.Duration
 }
 
 //CustomURLKey  place holder name, which will be replaced by kafka key
@@ -47,11 +50,14 @@ func (c *KafkaFoxtrotConn) Run() {
 	}
 	kafkaMsgFactory := getKafkaFoxtrotFactory()
 	src := source.GetKafkaSource(conf.Source, kafkaMsgFactory)
-	offsetTracker := source.GetKafkaOffsetTracker(conf.PendingAcks, src)
+	offsetTracker := source.GetKafkaOffsetTracker(conf.PendingAcks, src, c.Name)
 	hook := GetKafkaHook(offsetTracker, c.EnableDebugLog)
 	sk := sink.GetHTTPSink(conf.Dmux.Size, conf.Sink)
 	sk.RegisterHook(hook)
+	sk.SetConnectionName(c.Name)
 	src.RegisterHook(hook)
+	src.SetConnectionName(c.Name)
+	src.SetPollinginterval(c.PollingInterval)
 
 	//hash distribution
 	h := GetKafkaMsgHasher()

connection/kafka_http_conn.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/go-dmux/core"
@@ -27,8 +28,10 @@ type KafkaHTTPConnConfig struct {
 
 //KafkaHTTPConn struct to abstract this connections Run
 type KafkaHTTPConn struct {
-	EnableDebugLog bool
-	Conf           interface{}
+	Name            string
+	EnableDebugLog  bool
+	Conf            interface{}
+	PollingInterval time.Duration
 }
 
 func (c *KafkaHTTPConn) getConfiguration() *KafkaHTTPConnConfig {
@@ -49,11 +52,14 @@ func (c *KafkaHTTPConn) Run() {
 	}
 	kafkaMsgFactory := getKafkaHTTPFactory()
 	src := source.GetKafkaSource(conf.Source, kafkaMsgFactory)
-	offsetTracker := source.GetKafkaOffsetTracker(conf.PendingAcks, src)
+	offsetTracker := source.GetKafkaOffsetTracker(conf.PendingAcks, src, c.Name)
 	hook := GetKafkaHook(offsetTracker, c.EnableDebugLog)
 	sk := sink.GetHTTPSink(conf.Dmux.Size, conf.Sink)
 	sk.RegisterHook(hook)
+	sk.SetConnectionName(c.Name)
 	src.RegisterHook(hook)
+	src.SetConnectionName(c.Name)
+	src.SetPollinginterval(c.PollingInterval)
 
 	//hash distribution
 	h := GetKafkaMsgHasher()

http/http_sink.go

@@ -20,6 +20,7 @@ type HTTPSink struct {
 	client *http.Client
 	hook   HTTPSinkHook
 	conf   HTTPSinkConf
+	name   string
 }
 
 //HTTPSinkConf  holds config to HTTPSink
@@ -92,6 +93,10 @@ func (h *HTTPSink) RegisterHook(hook HTTPSinkHook) {
 	h.hook = hook
 }
 
+func (h *HTTPSink) SetConnectionName(name string) {
+	h.name = name
+}
+
 //HTTPMsg is an interface which incoming data should implment for HttpSink to
 //work
 type HTTPMsg interface {

kafka/consumer-group/consumer_group.go

@@ -3,11 +3,13 @@ package consumergroup
 import (
 	"errors"
 	"fmt"
+	"github.com/go-dmux/metrics"
+	"github.com/prometheus/client_golang/prometheus"
 	"sync"
 	"time"
 
-    "github.com/go-dmux/kafka/kazoo-go"
 	"github.com/Shopify/sarama"
+	"github.com/go-dmux/kafka/kazoo-go"
 )
 
 var (
@@ -83,7 +85,7 @@ type ConsumerGroup struct {
 }
 
 // Connects to a consumer group, using Zookeeper for auto-discovery
-func JoinConsumerGroup(name string, topics []string, zookeeper []string, config *Config) (cg *ConsumerGroup, err error) {
+func JoinConsumerGroup(name string, topics []string, zookeeper []string, config *Config, connectionName string, brokerList *[]string) (cg *ConsumerGroup, err error) {
 
 	if name == "" {
 		return nil, sarama.ConfigurationError("Empty consumergroup name")
@@ -116,6 +118,8 @@ func JoinConsumerGroup(name string, topics []string, zookeeper []string, config
 	if err != nil {
 		kz.Close()
 		return
+	} else {
+		*brokerList = brokers
 	}
 
 	group := kz.Consumergroup(name)
@@ -176,7 +180,7 @@ func JoinConsumerGroup(name string, topics []string, zookeeper []string, config
 	offsetConfig := OffsetManagerConfig{CommitInterval: config.Offsets.CommitInterval}
 	cg.offsetManager = NewZookeeperOffsetManager(cg, &offsetConfig)
 
-	go cg.topicListConsumer(topics)
+	go cg.topicListConsumer(connectionName, topics)
 
 	return
 }
@@ -250,7 +254,7 @@ func (cg *ConsumerGroup) FlushOffsets() error {
 	return cg.offsetManager.Flush()
 }
 
-func (cg *ConsumerGroup) topicListConsumer(topics []string) {
+func (cg *ConsumerGroup) topicListConsumer(connectionName string, topics []string) {
 	for {
 		select {
 		case <-cg.stopper:
@@ -271,7 +275,7 @@ func (cg *ConsumerGroup) topicListConsumer(topics []string) {
 
 		for _, topic := range topics {
 			cg.wg.Add(1)
-			go cg.topicConsumer(topic, cg.messages, cg.errors, stopper)
+			go cg.topicConsumer(connectionName, topic, cg.messages, cg.errors, stopper)
 		}
 
 		select {
@@ -299,7 +303,7 @@ func (cg *ConsumerGroup) topicListConsumer(topics []string) {
 	}
 }
 
-func (cg *ConsumerGroup) topicConsumer(topic string, messages chan<- *sarama.ConsumerMessage, errors chan<- error, stopper <-chan struct{}) {
+func (cg *ConsumerGroup) topicConsumer(connectionName string, topic string, messages chan<- *sarama.ConsumerMessage, errors chan<- error, stopper <-chan struct{}) {
 	defer cg.wg.Done()
 
 	select {
@@ -340,7 +344,15 @@ func (cg *ConsumerGroup) topicConsumer(topic string, messages chan<- *sarama.Con
 	// Consume all the assigned partitions
 	var wg sync.WaitGroup
 	for _, pid := range myPartitions {
-
+		//Create PartitionInfo and send it for ingestion through the partition channel
+		//In case of re-balancing this function will be triggered again and the latest information will be sent
+
+		metricName := connectionName + "." + topic + "." + cg.instance.ID + "." + time.Now().Format(time.RFC850)
+		metrics.Reg.Ingest(metrics.Metric{
+			MetricType:  prometheus.GaugeValue,
+			MetricName:  "partition_owned." + metricName,
+			MetricValue: int64(pid.ID),
+		})
 		wg.Add(1)
 		go cg.partitionConsumer(topic, pid.ID, messages, errors, &wg, stopper)
 	}
@@ -510,3 +522,14 @@ partitionConsumerLoop:
 		cg.Logf("%s/%d :: %s\n", topic, partition, err)
 	}
 }
+
+//GetConsumerOffset fetches the highest offset for a partition
+func (c *ConsumerGroup) GetConsumerOffset(topic string, partition int32) (int64, error) {
+	mark := c.consumer.HighWaterMarks()
+	if partitions, ok := mark[topic]; ok {
+		if offset, ok1 := partitions[partition]; ok1 {
+			return offset, nil
+		}
+	}
+	return -1, errors.New("could not get offset")
+}