github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/kafka/init.go (about)

     1  package kafka
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"strconv"
     7  	"time"
     8  
     9  	"github.com/rcrowley/go-metrics"
    10  	"github.com/sirupsen/logrus"
    11  	"github.com/wvanbergen/kazoo-go"
    12  	"gopkg.in/Shopify/sarama.v1"
    13  
    14  	"github.com/artisanhe/tools/kafka/consumergroup"
    15  )
    16  
    17  func PartitionerConstructor(topic string) sarama.Partitioner {
    18  	return sarama.NewHashPartitioner(topic)
    19  }
    20  
    21  func (kafkaConfig KafkaConfig) initKafkaConfig(clientConfig *sarama.Config) {
    22  	// init net config
    23  	clientConfig.Net.DialTimeout = time.Duration(kafkaConfig.Net.DialTimeout) * time.Millisecond
    24  	clientConfig.Net.KeepAlive = time.Duration(kafkaConfig.Net.KeepAlive) * time.Millisecond
    25  	clientConfig.Net.MaxOpenRequests = kafkaConfig.Net.MaxOpenRequests
    26  	clientConfig.Net.ReadTimeout = time.Duration(kafkaConfig.Net.ReadTimeout) * time.Millisecond
    27  	clientConfig.Net.SASL.Enable = kafkaConfig.Net.SASL.Enable
    28  	clientConfig.Net.SASL.Handshake = kafkaConfig.Net.SASL.Handshake
    29  	clientConfig.Net.SASL.Password = kafkaConfig.Net.SASL.Password
    30  	clientConfig.Net.SASL.User = kafkaConfig.Net.SASL.User
    31  	clientConfig.Net.TLS.Enable = false
    32  	clientConfig.Net.WriteTimeout = time.Duration(kafkaConfig.Net.WriteTimeout) * time.Millisecond
    33  
    34  	// init meta config
    35  	clientConfig.Metadata.RefreshFrequency = time.Duration(kafkaConfig.MetaData.RefreshFrequency) * time.Millisecond
    36  	clientConfig.Metadata.Retry.Backoff = time.Duration(kafkaConfig.MetaData.Retry.Backoff) * time.Millisecond
    37  	clientConfig.Metadata.Retry.Max = kafkaConfig.MetaData.Retry.Max
    38  
    39  	// init producer config
    40  	clientConfig.Producer.MaxMessageBytes = kafkaConfig.Producer.MaxMessageBytes
    41  	clientConfig.Producer.Flush.MaxMessages = kafkaConfig.Producer.Flush.MaxMessages
    42  	clientConfig.Producer.RequiredAcks = sarama.RequiredAcks(kafkaConfig.Producer.RequiredAcks)
    43  	clientConfig.Producer.Retry.Backoff = time.Duration(kafkaConfig.Producer.Retry.Backoff) * time.Millisecond
    44  	clientConfig.Producer.Retry.Max = kafkaConfig.Producer.Retry.Max
    45  	clientConfig.Producer.Return.Errors = kafkaConfig.Producer.Return.Errors
    46  	clientConfig.Producer.Return.Successes = kafkaConfig.Producer.Return.Successes
    47  	clientConfig.Producer.Timeout = time.Duration(kafkaConfig.Producer.Timeout) * time.Millisecond
    48  	clientConfig.Producer.Compression = sarama.CompressionCodec(kafkaConfig.Producer.Compression)
    49  	clientConfig.Producer.Partitioner = PartitionerConstructor
    50  
    51  	// init consumer config
    52  	clientConfig.Consumer.Fetch.Default = kafkaConfig.Consumer.Fetch.Default
    53  	clientConfig.Consumer.Fetch.Max = kafkaConfig.Consumer.Fetch.Max
    54  	clientConfig.Consumer.Fetch.Min = kafkaConfig.Consumer.Fetch.Min
    55  	clientConfig.Consumer.MaxProcessingTime = time.Duration(kafkaConfig.Consumer.MaxProcessingTime) * time.Millisecond
    56  	clientConfig.Consumer.MaxWaitTime = time.Duration(kafkaConfig.Consumer.MaxWaitTime) * time.Millisecond
    57  	clientConfig.Consumer.Offsets.CommitInterval = time.Duration(kafkaConfig.Consumer.Offsets.CommitInterval) * time.Millisecond
    58  	clientConfig.Consumer.Offsets.Initial = kafkaConfig.Consumer.Offsets.Initial
    59  	clientConfig.Consumer.Offsets.Retention = time.Duration(kafkaConfig.Consumer.Offsets.Retention) * time.Millisecond
    60  	clientConfig.Consumer.Retry.Backoff = time.Duration(kafkaConfig.Consumer.Retry.Backoff) * time.Millisecond
    61  	clientConfig.Consumer.Return.Errors = kafkaConfig.Consumer.Return.Errors
    62  
    63  	// misc
    64  	clientConfig.ClientID = strconv.FormatInt(int64(os.Getpid()), 10)
    65  	clientConfig.MetricRegistry = metrics.NewRegistry()
    66  	switch kafkaConfig.Version {
    67  	case "0.8.2.0":
    68  		clientConfig.Version = sarama.V0_8_2_0
    69  	case "0.8.2.1":
    70  		clientConfig.Version = sarama.V0_8_2_1
    71  	case "0.8.2.2":
    72  		clientConfig.Version = sarama.V0_8_2_2
    73  	case "0.9.0.0":
    74  		clientConfig.Version = sarama.V0_9_0_0
    75  	case "0.9.0.1":
    76  		clientConfig.Version = sarama.V0_9_0_1
    77  	case "0.10.0.0":
    78  		clientConfig.Version = sarama.V0_10_0_0
    79  	case "0.10.0.1":
    80  		clientConfig.Version = sarama.V0_10_0_1
    81  	case "0.10.1.0":
    82  		clientConfig.Version = sarama.V0_10_1_0
    83  	default:
    84  		panic(fmt.Sprintf("invalid kafka version %s", kafkaConfig.Version))
    85  	}
    86  }
    87  
    88  func (kafkaConfig KafkaConfig) NewKafkaClient() (kafkaClient sarama.Client, err error) {
    89  	clientConfig := sarama.Config{}
    90  	kafkaConfig.initKafkaConfig(&clientConfig)
    91  	metrics.UseNilMetrics = true
    92  
    93  	kafkaClient, err = sarama.NewClient(kafkaConfig.Addrs, &clientConfig)
    94  	return
    95  }
    96  
    97  type logger struct {
    98  }
    99  
   100  func (l *logger) Printf(str string, values ...interface{}) {
   101  	logrus.WithField("tag", "kafka").Warningf("%s[%+v]", str, values)
   102  }
   103  func initKazooConfig(kazooConfig *kazoo.Config, zooConfig ZookeeperConfig) {
   104  	kazooConfig.Chroot = zooConfig.Chroot
   105  	kazooConfig.Timeout = time.Duration(zooConfig.Timeout) * time.Millisecond
   106  	kazooConfig.Logger = &logger{}
   107  }
   108  
   109  func (consumerKafkaConfig ConsumerKafkaConfig) NewKafkaConsumer() (cg *consumergroup.ConsumerGroup, err error) {
   110  	// init kafka config
   111  	clientConfig := sarama.Config{}
   112  	consumerKafkaConfig.ConsumerKafka.initKafkaConfig(&clientConfig)
   113  	metrics.UseNilMetrics = true
   114  
   115  	// init zookeeper config
   116  	kazooConfig := kazoo.Config{}
   117  	initKazooConfig(&kazooConfig, consumerKafkaConfig.Zookeeper)
   118  
   119  	// init consume config
   120  	consumeConfig := consumergroup.NewConfig()
   121  	consumeConfig.Config = &clientConfig
   122  	consumeConfig.Zookeeper = &kazooConfig
   123  	consumeConfig.Offsets.CommitInterval = clientConfig.Consumer.Offsets.CommitInterval
   124  	consumeConfig.Offsets.Initial = clientConfig.Consumer.Offsets.Initial
   125  	consumeConfig.Offsets.ProcessingTimeout = clientConfig.Consumer.MaxProcessingTime
   126  
   127  	cg, err = consumergroup.JoinConsumerGroup(consumerKafkaConfig.ConsumerGroupName, consumerKafkaConfig.Topics, consumerKafkaConfig.Zookeeper.Addrs, consumeConfig)
   128  	if err != nil {
   129  		return
   130  	}
   131  	return
   132  }
   133  func (producerKafkaConfig *ProducerKafkaConfig) Init() {
   134  	if producerKafkaConfig.kafkaClient == nil {
   135  		var err error
   136  		if producerKafkaConfig.kafkaClient, err = producerKafkaConfig.ProducerKafka.NewKafkaClient(); err != nil {
   137  			panic(fmt.Sprintf("producerKafka init fail[err:%v]", err))
   138  		}
   139  	}
   140  }
   141  
   142  func (consumerKafkaConfig *ConsumerKafkaConfig) Init() {
   143  	if consumerKafkaConfig.kafkaClient == nil {
   144  		var err error
   145  		if consumerKafkaConfig.kafkaClient, err = consumerKafkaConfig.ConsumerKafka.NewKafkaClient(); err != nil {
   146  			panic(fmt.Sprintf("Kafka client fail[err:%v]", err))
   147  		}
   148  	}
   149  	if consumerKafkaConfig.cG == nil {
   150  		var err error
   151  		consumerKafkaConfig.cG, err = consumerKafkaConfig.NewKafkaConsumer()
   152  		if err != nil {
   153  			panic(fmt.Sprintf("NewKafkaConsumer fail[err:%v]", err))
   154  		}
   155  	}
   156  }
   157  
   158  func (producerKafkaConfig *ProducerKafkaConfig) GetKafkaClient() sarama.Client {
   159  	return producerKafkaConfig.kafkaClient
   160  }
   161  
   162  func (consumerKafkaConfig *ConsumerKafkaConfig) GetKafkaClient() sarama.Client {
   163  	return consumerKafkaConfig.kafkaClient
   164  }
   165  
   166  func (consumerKafkaConfig *ConsumerKafkaConfig) GetCG() *consumergroup.ConsumerGroup {
   167  	return consumerKafkaConfig.cG
   168  }