github.com/erda-project/erda-infra@v1.0.9/providers/kafka/consumer.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     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  	"time"
    19  
    20  	"github.com/confluentinc/confluent-kafka-go/kafka"
    21  )
    22  
    23  // ConsumerConfig .
    24  type ConsumerConfig struct {
    25  	Topics      []string               `file:"topics" desc:"topics"`
    26  	Group       string                 `file:"group" desc:"consumer group id"`
    27  	Parallelism uint64                 `file:"parallelism" desc:"parallelism"`
    28  	Options     map[string]interface{} `file:"options" desc:"options"`
    29  }
    30  
    31  // ConsumerFunc .
    32  type ConsumerFunc func(key []byte, value []byte, topic *string, timestamp time.Time) error
    33  
    34  func mergeMap(a, b map[string]interface{}) map[string]interface{} {
    35  	if a == nil || len(a) == 0 {
    36  		return b
    37  	}
    38  	if b == nil || len(b) == 0 {
    39  		return a
    40  	}
    41  	c := make(map[string]interface{}, len(a)+len(b))
    42  	for k, v := range a {
    43  		c[k] = v
    44  	}
    45  	for k, v := range b {
    46  		c[k] = v
    47  	}
    48  	return c
    49  }
    50  
    51  func (s *service) NewConsumer(cfg *ConsumerConfig, handler ConsumerFunc, options ...ConsumerOption) error {
    52  	return s.NewConsumerWitchCreator(cfg, func(int) (ConsumerFunc, error) { return handler, nil })
    53  }
    54  
    55  func (s *service) NewConsumerWitchCreator(cfg *ConsumerConfig, creator func(i int) (ConsumerFunc, error), opts ...ConsumerOption) error {
    56  	options := mergeMap(s.p.Cfg.Comsumer.Options, cfg.Options)
    57  	parallelism := int(cfg.Parallelism)
    58  	var consumerListener func(i int, c *kafka.Consumer)
    59  	for _, opt := range opts {
    60  		switch v := opt.(type) {
    61  		case func(i int, c *kafka.Consumer):
    62  			consumerListener = v
    63  		}
    64  	}
    65  	for i := 0; i < parallelism; i++ {
    66  		var kc kafka.ConfigMap
    67  		if options != nil {
    68  			kc = convertToConfigMap(options)
    69  		} else {
    70  			kc = kafka.ConfigMap{}
    71  		}
    72  		kc["bootstrap.servers"] = s.p.Cfg.Servers
    73  		kc["group.id"] = cfg.Group
    74  
    75  		handler, err := creator(i)
    76  		if err != nil {
    77  			return err
    78  		}
    79  		go func(i int, handler ConsumerFunc) {
    80  		loop:
    81  			for {
    82  				consumer, err := kafka.NewConsumer(&kc)
    83  				if err != nil {
    84  					s.log.Errorf("failed to create kafka consumer with config %#v: %v", cfg, err)
    85  					time.Sleep(3 * time.Second)
    86  					continue
    87  				}
    88  				if err = consumer.SubscribeTopics(cfg.Topics, nil); err != nil {
    89  					s.log.Errorf("failed to subscribe kafka topics %v: %#v", cfg.Topics, err)
    90  					s.logError(consumer.Close())
    91  					time.Sleep(3 * time.Second)
    92  					continue
    93  				}
    94  				if consumerListener != nil {
    95  					consumerListener(i, consumer)
    96  				}
    97  				s.log.Infof("create kafka consumer %d with topics: %v, config: %#v", i, cfg.Topics, kc)
    98  				var errors int
    99  				for {
   100  					message, err := consumer.ReadMessage(-1)
   101  					if err != nil {
   102  						var topic string
   103  						if message != nil {
   104  							topic = *message.TopicPartition.Topic
   105  						}
   106  						s.log.Errorf("topic: %s .fail to read message from kafka: %v", topic, err)
   107  						errors++
   108  						if errors > 10 {
   109  							s.logError(consumer.Close())
   110  							time.Sleep(1 * time.Second)
   111  							continue loop
   112  						}
   113  						continue
   114  					}
   115  					err = handler(message.Key, message.Value, message.TopicPartition.Topic, message.Timestamp)
   116  					if err != nil {
   117  						s.log.Errorf("fail to process message: %v", err)
   118  					}
   119  					errors = 0
   120  				}
   121  			}
   122  		}(i, handler)
   123  	}
   124  	return nil
   125  }
   126  
   127  func convertToConfigMap(m map[string]interface{}) kafka.ConfigMap {
   128  	cm := make(kafka.ConfigMap, len(m))
   129  	for k, v := range m {
   130  		cm[k] = v
   131  	}
   132  	return cm
   133  }
   134  
   135  func (s *service) logError(err error) {
   136  	if err != nil {
   137  		s.log.Error(err)
   138  	}
   139  }
   140  
   141  // ConsumerOption .
   142  type ConsumerOption interface{}
   143  
   144  // WithConsumerListener .
   145  func WithConsumerListener(fn func(i int, c *kafka.Consumer)) ConsumerOption {
   146  	return ConsumerOption(fn)
   147  }