github.com/jonaz/heapster@v1.3.0-beta.0.0.20170208112634-cd3c15ca3d29/common/kafka/kafka.go (about)

     1  // Copyright 2015 Google Inc. All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package kafka
    16  
    17  import (
    18  	"encoding/json"
    19  	"fmt"
    20  	"net/url"
    21  	"time"
    22  
    23  	"github.com/golang/glog"
    24  	"github.com/optiopay/kafka"
    25  	"github.com/optiopay/kafka/proto"
    26  )
    27  
    28  const (
    29  	brokerClientID           = "kafka-sink"
    30  	brokerDialTimeout        = 10 * time.Second
    31  	brokerDialRetryLimit     = 1
    32  	brokerDialRetryWait      = 0
    33  	brokerAllowTopicCreation = true
    34  	brokerLeaderRetryLimit   = 1
    35  	brokerLeaderRetryWait    = 0
    36  	metricsTopic             = "heapster-metrics"
    37  	eventsTopic              = "heapster-events"
    38  )
    39  
    40  const (
    41  	TimeSeriesTopic = "timeseriestopic"
    42  	EventsTopic     = "eventstopic"
    43  )
    44  
    45  type KafkaClient interface {
    46  	Name() string
    47  	Stop()
    48  	ProduceKafkaMessage(msgData interface{}) error
    49  }
    50  
    51  type kafkaSink struct {
    52  	producer  kafka.DistributingProducer
    53  	dataTopic string
    54  }
    55  
    56  func (sink *kafkaSink) ProduceKafkaMessage(msgData interface{}) error {
    57  	start := time.Now()
    58  	msgJson, err := json.Marshal(msgData)
    59  	if err != nil {
    60  		return fmt.Errorf("failed to transform the items to json : %s", err)
    61  	}
    62  
    63  	message := &proto.Message{Value: []byte(string(msgJson))}
    64  	_, err = sink.producer.Distribute(sink.dataTopic, message)
    65  	if err != nil {
    66  		return fmt.Errorf("failed to produce message to %s: %s", sink.dataTopic, err)
    67  	}
    68  	end := time.Now()
    69  	glog.V(4).Infof("Exported %d data to kafka in %s", len([]byte(string(msgJson))), end.Sub(start))
    70  	return nil
    71  }
    72  
    73  func (sink *kafkaSink) Name() string {
    74  	return "Apache Kafka Sink"
    75  }
    76  
    77  func (sink *kafkaSink) Stop() {
    78  	// nothing needs to be done.
    79  }
    80  
    81  // setupProducer returns a producer of kafka server
    82  func setupProducer(sinkBrokerHosts []string, topic string, brokerConf kafka.BrokerConf) (kafka.DistributingProducer, error) {
    83  	glog.V(3).Infof("attempting to setup kafka sink")
    84  	broker, err := kafka.Dial(sinkBrokerHosts, brokerConf)
    85  	if err != nil {
    86  		return nil, fmt.Errorf("failed to connect to kafka cluster: %s", err)
    87  	}
    88  	defer broker.Close()
    89  
    90  	//create kafka producer
    91  	conf := kafka.NewProducerConf()
    92  	conf.RequiredAcks = proto.RequiredAcksLocal
    93  	producer := broker.Producer(conf)
    94  
    95  	// create RoundRobinProducer with the default producer.
    96  	count, err := broker.PartitionCount(topic)
    97  	if err != nil {
    98  		count = 1
    99  		glog.Warningf("Failed to get partition count of topic %q: %s", topic, err)
   100  	}
   101  	sinkProducer := kafka.NewRoundRobinProducer(producer, count)
   102  	glog.V(3).Infof("kafka sink setup successfully")
   103  	return sinkProducer, nil
   104  }
   105  
   106  func getTopic(opts map[string][]string, topicType string) (string, error) {
   107  	var topic string
   108  	switch topicType {
   109  	case TimeSeriesTopic:
   110  		topic = metricsTopic
   111  	case EventsTopic:
   112  		topic = eventsTopic
   113  	default:
   114  		return "", fmt.Errorf("Topic type '%s' is illegal.", topicType)
   115  	}
   116  
   117  	if len(opts[topicType]) > 0 {
   118  		topic = opts[topicType][0]
   119  	}
   120  
   121  	return topic, nil
   122  }
   123  
   124  func NewKafkaClient(uri *url.URL, topicType string) (KafkaClient, error) {
   125  	opts, err := url.ParseQuery(uri.RawQuery)
   126  	if err != nil {
   127  		return nil, fmt.Errorf("failed to parser url's query string: %s", err)
   128  	}
   129  	glog.V(3).Infof("kafka sink option: %v", opts)
   130  
   131  	topic, err := getTopic(opts, topicType)
   132  	if err != nil {
   133  		return nil, err
   134  	}
   135  
   136  	var kafkaBrokers []string
   137  	if len(opts["brokers"]) < 1 {
   138  		return nil, fmt.Errorf("There is no broker assigned for connecting kafka")
   139  	}
   140  	kafkaBrokers = append(kafkaBrokers, opts["brokers"]...)
   141  	glog.V(2).Infof("initializing kafka sink with brokers - %v", kafkaBrokers)
   142  
   143  	//structure the config of broker
   144  	brokerConf := kafka.NewBrokerConf(brokerClientID)
   145  	brokerConf.DialTimeout = brokerDialTimeout
   146  	brokerConf.DialRetryLimit = brokerDialRetryLimit
   147  	brokerConf.DialRetryWait = brokerDialRetryWait
   148  	brokerConf.LeaderRetryLimit = brokerLeaderRetryLimit
   149  	brokerConf.LeaderRetryWait = brokerLeaderRetryWait
   150  	brokerConf.AllowTopicCreation = brokerAllowTopicCreation
   151  
   152  	// set up producer of kafka server.
   153  	sinkProducer, err := setupProducer(kafkaBrokers, topic, brokerConf)
   154  	if err != nil {
   155  		return nil, fmt.Errorf("Failed to setup Producer: - %v", err)
   156  	}
   157  
   158  	return &kafkaSink{
   159  		producer:  sinkProducer,
   160  		dataTopic: topic,
   161  	}, nil
   162  }