github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/sink/kafka/sarama_factory.go (about)

     1  // Copyright 2023 PingCAP, 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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package kafka
    15  
    16  import (
    17  	"context"
    18  
    19  	"github.com/IBM/sarama"
    20  	"github.com/pingcap/tiflow/cdc/model"
    21  	"github.com/pingcap/tiflow/pkg/errors"
    22  	"github.com/pingcap/tiflow/pkg/util"
    23  	"github.com/rcrowley/go-metrics"
    24  )
    25  
    26  type saramaFactory struct {
    27  	changefeedID model.ChangeFeedID
    28  	option       *Options
    29  
    30  	registry metrics.Registry
    31  }
    32  
    33  // NewSaramaFactory constructs a Factory with sarama implementation.
    34  func NewSaramaFactory(
    35  	o *Options,
    36  	changefeedID model.ChangeFeedID,
    37  ) (Factory, error) {
    38  	return &saramaFactory{
    39  		changefeedID: changefeedID,
    40  		option:       o,
    41  		registry:     metrics.NewRegistry(),
    42  	}, nil
    43  }
    44  
    45  func (f *saramaFactory) AdminClient(ctx context.Context) (ClusterAdminClient, error) {
    46  	config, err := NewSaramaConfig(ctx, f.option)
    47  	if err != nil {
    48  		return nil, err
    49  	}
    50  
    51  	client, err := sarama.NewClient(f.option.BrokerEndpoints, config)
    52  	if err != nil {
    53  		return nil, errors.Trace(err)
    54  	}
    55  
    56  	admin, err := sarama.NewClusterAdminFromClient(client)
    57  	if err != nil {
    58  		return nil, errors.Trace(err)
    59  	}
    60  	return &saramaAdminClient{
    61  		client:     client,
    62  		admin:      admin,
    63  		changefeed: f.changefeedID,
    64  	}, nil
    65  }
    66  
    67  // SyncProducer returns a Sync Producer,
    68  // it should be the caller's responsibility to close the producer
    69  func (f *saramaFactory) SyncProducer(ctx context.Context) (SyncProducer, error) {
    70  	config, err := NewSaramaConfig(ctx, f.option)
    71  	if err != nil {
    72  		return nil, err
    73  	}
    74  	config.MetricRegistry = f.registry
    75  
    76  	client, err := sarama.NewClient(f.option.BrokerEndpoints, config)
    77  	if err != nil {
    78  		return nil, errors.Trace(err)
    79  	}
    80  
    81  	p, err := sarama.NewSyncProducerFromClient(client)
    82  	if err != nil {
    83  		return nil, errors.Trace(err)
    84  	}
    85  	return &saramaSyncProducer{
    86  		id:       f.changefeedID,
    87  		client:   client,
    88  		producer: p,
    89  	}, nil
    90  }
    91  
    92  // AsyncProducer return an Async Producer,
    93  // it should be the caller's responsibility to close the producer
    94  func (f *saramaFactory) AsyncProducer(
    95  	ctx context.Context,
    96  	failpointCh chan error,
    97  ) (AsyncProducer, error) {
    98  	config, err := NewSaramaConfig(ctx, f.option)
    99  	if err != nil {
   100  		return nil, err
   101  	}
   102  	config.MetricRegistry = f.registry
   103  
   104  	client, err := sarama.NewClient(f.option.BrokerEndpoints, config)
   105  	if err != nil {
   106  		return nil, errors.Trace(err)
   107  	}
   108  	p, err := sarama.NewAsyncProducerFromClient(client)
   109  	if err != nil {
   110  		return nil, errors.Trace(err)
   111  	}
   112  	return &saramaAsyncProducer{
   113  		client:       client,
   114  		producer:     p,
   115  		changefeedID: f.changefeedID,
   116  		failpointCh:  failpointCh,
   117  	}, nil
   118  }
   119  
   120  func (f *saramaFactory) MetricsCollector(
   121  	role util.Role,
   122  	adminClient ClusterAdminClient,
   123  ) MetricsCollector {
   124  	return NewSaramaMetricsCollector(
   125  		f.changefeedID, role, adminClient, f.registry)
   126  }