github.com/psiphon-inc/goarista@v0.0.0-20160825065156-d002785f4c67/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  	"flag"
    11  	"fmt"
    12  	"strings"
    13  	"sync"
    14  
    15  	"github.com/Shopify/sarama"
    16  	"github.com/aristanetworks/glog"
    17  	"github.com/aristanetworks/goarista/kafka"
    18  	"github.com/aristanetworks/goarista/kafka/openconfig"
    19  	"github.com/aristanetworks/goarista/kafka/producer"
    20  	pb "github.com/aristanetworks/goarista/openconfig"
    21  	"github.com/aristanetworks/goarista/openconfig/client"
    22  )
    23  
    24  var keysFlag = flag.String("kafkakeys", "",
    25  	"Keys for kafka messages (comma-separated, default: the value of -addrs")
    26  
    27  func newProducer(addresses []string, topic, key string) (producer.Producer, error) {
    28  	client, err := kafka.NewClient(addresses)
    29  	if err != nil {
    30  		return nil, fmt.Errorf("Failed to create Kafka client: %s", err)
    31  	}
    32  	encodedKey := sarama.StringEncoder(key)
    33  	p, err := producer.New(topic, nil, client, encodedKey, openconfig.ElasticsearchMessageEncoder)
    34  	if err != nil {
    35  		return nil, fmt.Errorf("Failed to create Kafka producer: %s", err)
    36  	}
    37  	return p, nil
    38  }
    39  
    40  func main() {
    41  	username, password, subscriptions, grpcAddrs, opts := client.ParseFlags()
    42  
    43  	if *keysFlag == "" {
    44  		*keysFlag = strings.Join(grpcAddrs, ",")
    45  	}
    46  	keys := strings.Split(*keysFlag, ",")
    47  	if len(grpcAddrs) != len(keys) {
    48  		glog.Fatal("Please provide the same number of addresses and Kafka keys")
    49  	}
    50  	addresses := strings.Split(*kafka.Addresses, ",")
    51  	wg := new(sync.WaitGroup)
    52  	for i, grpcAddr := range grpcAddrs {
    53  		key := keys[i]
    54  		p, err := newProducer(addresses, *kafka.Topic, key)
    55  		if err != nil {
    56  			glog.Fatal(err)
    57  		} else {
    58  			glog.Infof("Initialized Kafka producer for %s", grpcAddr)
    59  		}
    60  		publish := func(notif *pb.SubscribeResponse) {
    61  			p.Write(notif)
    62  		}
    63  		wg.Add(1)
    64  		go p.Run()
    65  		go client.Run(publish, wg, username, password, grpcAddr, subscriptions, opts)
    66  	}
    67  	wg.Wait()
    68  }