github.com/prysmaticlabs/prysm@v1.4.4/beacon-chain/db/kafka/export_wrapper.go (about)

     1  // Package kafka defines an implementation of Database interface
     2  // which exports streaming data using Kafka for data analysis.
     3  package kafka
     4  
     5  import (
     6  	"context"
     7  
     8  	fssz "github.com/ferranbt/fastssz"
     9  	"github.com/prysmaticlabs/prysm/beacon-chain/db/iface"
    10  	"github.com/prysmaticlabs/prysm/proto/interfaces"
    11  	"github.com/prysmaticlabs/prysm/shared/featureconfig"
    12  	"github.com/prysmaticlabs/prysm/shared/traceutil"
    13  	"go.opencensus.io/trace"
    14  	jsonpb "google.golang.org/protobuf/encoding/protojson"
    15  	"gopkg.in/confluentinc/confluent-kafka-go.v1/kafka"
    16  	_ "gopkg.in/confluentinc/confluent-kafka-go.v1/kafka/librdkafka" // Required for c++ kafka library.
    17  	"gopkg.in/errgo.v2/fmt/errors"
    18  )
    19  
    20  var _ iface.Database = (*Exporter)(nil)
    21  var marshaler = jsonpb.MarshalOptions{}
    22  
    23  // Exporter wraps a database interface and exports certain objects to kafka topics.
    24  type Exporter struct {
    25  	db iface.Database
    26  	p  *kafka.Producer
    27  }
    28  
    29  // Wrap the db with kafka exporter. If the feature flag is not enabled, this service does not wrap
    30  // the database, but returns the underlying database pointer itself.
    31  func Wrap(db iface.Database) (iface.Database, error) {
    32  	if featureconfig.Get().KafkaBootstrapServers == "" {
    33  		log.Debug("Empty Kafka bootstrap servers list, database was not wrapped with Kafka exporter")
    34  		return db, nil
    35  	}
    36  	p, err := kafka.NewProducer(&kafka.ConfigMap{"bootstrap.servers": featureconfig.Get().KafkaBootstrapServers})
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  
    41  	return &Exporter{db: db, p: p}, nil
    42  }
    43  
    44  func (e Exporter) publish(ctx context.Context, topic string, msg interfaces.SignedBeaconBlock) error {
    45  	ctx, span := trace.StartSpan(ctx, "kafka.publish")
    46  	defer span.End()
    47  
    48  	var err error
    49  	var buf []byte
    50  	if buf, err = marshaler.Marshal(msg.Proto()); err != nil {
    51  		traceutil.AnnotateError(span, err)
    52  		return err
    53  	}
    54  
    55  	var key [32]byte
    56  	if v, ok := msg.(fssz.HashRoot); ok {
    57  		key, err = v.HashTreeRoot()
    58  	} else {
    59  		err = errors.New("object does not follow hash tree root interface")
    60  	}
    61  	if err != nil {
    62  		traceutil.AnnotateError(span, err)
    63  		return err
    64  	}
    65  
    66  	if err := e.p.Produce(&kafka.Message{
    67  		TopicPartition: kafka.TopicPartition{
    68  			Topic: &topic,
    69  		},
    70  		Value: buf,
    71  		Key:   key[:],
    72  	}, nil); err != nil {
    73  		traceutil.AnnotateError(span, err)
    74  		return err
    75  	}
    76  	return nil
    77  }
    78  
    79  // Close closes kafka producer and underlying db.
    80  func (e Exporter) Close() error {
    81  	e.p.Close()
    82  	return e.db.Close()
    83  }
    84  
    85  // SaveBlock publishes to the kafka topic for beacon blocks.
    86  func (e Exporter) SaveBlock(ctx context.Context, block interfaces.SignedBeaconBlock) error {
    87  	go func() {
    88  		if err := e.publish(ctx, "beacon_block", block); err != nil {
    89  			log.WithError(err).Error("Failed to publish block")
    90  		}
    91  	}()
    92  
    93  	return e.db.SaveBlock(ctx, block)
    94  }
    95  
    96  // SaveBlocks publishes to the kafka topic for beacon blocks.
    97  func (e Exporter) SaveBlocks(ctx context.Context, blocks []interfaces.SignedBeaconBlock) error {
    98  	go func() {
    99  		for _, block := range blocks {
   100  			if err := e.publish(ctx, "beacon_block", block); err != nil {
   101  				log.WithError(err).Error("Failed to publish block")
   102  			}
   103  		}
   104  	}()
   105  
   106  	return e.db.SaveBlocks(ctx, blocks)
   107  }