github.com/erda-project/erda-infra@v1.0.9/providers/kafka/kafka.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  	"reflect"
    19  
    20  	"github.com/confluentinc/confluent-kafka-go/kafka"
    21  	"github.com/erda-project/erda-infra/base/logs"
    22  	"github.com/erda-project/erda-infra/base/servicehub"
    23  	writer "github.com/erda-project/erda-infra/pkg/parallel-writer"
    24  )
    25  
    26  // Interface .
    27  type Interface interface {
    28  	NewBatchReader(c *BatchReaderConfig, options ...BatchReaderOption) (BatchReader, error)
    29  	NewConsumer(c *ConsumerConfig, handler ConsumerFunc, options ...ConsumerOption) error
    30  	NewConsumerWitchCreator(c *ConsumerConfig, creator func(i int) (ConsumerFunc, error), options ...ConsumerOption) error
    31  	NewProducer(c *ProducerConfig, options ...ProducerOption) (writer.Writer, error)
    32  	Servers() string
    33  	ProduceChannelSize() int
    34  	ProduceEventsChannelSize() int
    35  	NewAdminClient() (*kafka.AdminClient, error)
    36  }
    37  
    38  // Producer .
    39  type Producer interface {
    40  	writer.Writer
    41  	ProduceChannelSize() int
    42  	EventsChannelSize() int
    43  }
    44  
    45  type config struct {
    46  	Servers  string `file:"servers" env:"BOOTSTRAP_SERVERS" default:"localhost:9092" desc:"kafka servers"`
    47  	Producer struct {
    48  		Options map[string]interface{} `file:"options"`
    49  	} `file:"producer"`
    50  	Comsumer struct {
    51  		Options map[string]interface{} `file:"options"`
    52  	} `file:"comsumer"`
    53  	Admin struct {
    54  		Options map[string]interface{} `file:"options"`
    55  	} `file:"admin"`
    56  }
    57  
    58  // provider .
    59  type provider struct {
    60  	Cfg      *config
    61  	Log      logs.Logger
    62  	producer sharedProducer
    63  }
    64  
    65  // Init .
    66  func (p *provider) Init(ctx servicehub.Context) error {
    67  	p.producer.log = p.Log
    68  	return nil
    69  }
    70  
    71  // Provide .
    72  func (p *provider) Provide(ctx servicehub.DependencyContext, options ...interface{}) interface{} {
    73  	return &service{
    74  		p:    p,
    75  		log:  p.Log.Sub(ctx.Caller()),
    76  		name: ctx.Caller(),
    77  	}
    78  }
    79  
    80  type service struct {
    81  	p    *provider
    82  	log  logs.Logger
    83  	name string
    84  }
    85  
    86  var _ Interface = (*service)(nil)
    87  
    88  func (s *service) Servers() string { return s.p.Cfg.Servers }
    89  
    90  func (s *service) NewAdminClient() (*kafka.AdminClient, error) {
    91  	kc := kafka.ConfigMap{}
    92  	if s.p.Cfg.Admin.Options != nil {
    93  		for k, v := range s.p.Cfg.Admin.Options {
    94  			kc[k] = v
    95  		}
    96  	}
    97  	kc["bootstrap.servers"] = s.p.Cfg.Servers
    98  	return kafka.NewAdminClient(&kc)
    99  }
   100  
   101  func init() {
   102  	servicehub.Register("kafka", &servicehub.Spec{
   103  		Services: []string{"kafka", "kafka-producer", "kafka-consumer"},
   104  		Types: []reflect.Type{
   105  			reflect.TypeOf((*Interface)(nil)).Elem(),
   106  		},
   107  		Description: "kafka、kafka-producer、kafka-consumer",
   108  		ConfigFunc:  func() interface{} { return &config{} },
   109  		Creator: func() servicehub.Provider {
   110  			return &provider{}
   111  		},
   112  	})
   113  }