github.com/aristanetworks/goarista@v0.0.0-20240514173732-cca2755bbd44/cmd/ockafka/main.go (about)

     1  // Copyright (c) 2016 Arista Networks, Inc.
     2  // Use of this source code is governed by the Apache License 2.0
     3  // that can be found in the COPYING file.
     4  
     5  // The occlient tool is a client for the gRPC service for getting and setting the
     6  // OpenConfig configuration and state of a network device.
     7  package main
     8  
     9  import (
    10  	"context"
    11  	"flag"
    12  	"fmt"
    13  	"strings"
    14  	"sync"
    15  
    16  	client "github.com/aristanetworks/goarista/gnmi"
    17  	"github.com/aristanetworks/goarista/kafka"
    18  	"github.com/aristanetworks/goarista/kafka/gnmi"
    19  	"github.com/aristanetworks/goarista/kafka/producer"
    20  
    21  	pb "github.com/openconfig/gnmi/proto/gnmi"
    22  
    23  	"github.com/IBM/sarama"
    24  	"github.com/aristanetworks/glog"
    25  )
    26  
    27  var keysFlag = flag.String("kafkakeys", "",
    28  	"Keys for kafka messages (comma-separated, default: the value of -addrs). The key '"+
    29  		client.HostnameArg+"' is replaced by the current hostname.")
    30  
    31  func newProducer(addresses []string, topic, key, dataset string) (producer.Producer, error) {
    32  	encodedKey := sarama.StringEncoder(key)
    33  	p, err := producer.New(gnmi.NewEncoder(topic, encodedKey, dataset), addresses, nil)
    34  	if err != nil {
    35  		return nil, fmt.Errorf("Failed to create Kafka brokers: %s", err)
    36  	}
    37  	glog.Infof("Connected to Kafka brokers at %s", addresses)
    38  	return p, nil
    39  }
    40  
    41  func main() {
    42  	ctx := context.Background()
    43  	config, subscriptions := client.ParseFlags()
    44  	ctx = client.NewContext(ctx, config)
    45  	grpcAddrs := strings.Split(config.Addr, ",")
    46  
    47  	var keys []string
    48  	var err error
    49  	if *keysFlag == "" {
    50  		keys = grpcAddrs
    51  	} else {
    52  		keys, err = client.ParseHostnames(*keysFlag)
    53  		if err != nil {
    54  			glog.Fatal(err)
    55  		}
    56  	}
    57  	if len(grpcAddrs) != len(keys) {
    58  		glog.Fatal("Please provide the same number of addresses and Kafka keys")
    59  	}
    60  	addresses := strings.Split(*kafka.Addresses, ",")
    61  	wg := new(sync.WaitGroup)
    62  	for i, grpcAddr := range grpcAddrs {
    63  		key := keys[i]
    64  		p, err := newProducer(addresses, *kafka.Topic, key, grpcAddr)
    65  		if err != nil {
    66  			glog.Fatal(err)
    67  		} else {
    68  			glog.Infof("Initialized Kafka producer for %s", grpcAddr)
    69  		}
    70  		wg.Add(1)
    71  		go func() {
    72  			p.Start()
    73  			defer p.Stop()
    74  			respChan := make(chan *pb.SubscribeResponse)
    75  			errChan := make(chan error)
    76  			c, err := client.Dial(config)
    77  			if err != nil {
    78  				glog.Fatal(err)
    79  			}
    80  			subscribeOptions := &client.SubscribeOptions{
    81  				Paths: client.SplitPaths(subscriptions),
    82  			}
    83  			go client.Subscribe(ctx, c, subscribeOptions, respChan, errChan)
    84  			for {
    85  				select {
    86  				case resp, open := <-respChan:
    87  					if !open {
    88  						return
    89  					}
    90  					p.Write(resp)
    91  				case err := <-errChan:
    92  					glog.Fatal(err)
    93  				}
    94  			}
    95  		}()
    96  	}
    97  	wg.Wait()
    98  }