<commit_before>package command

import (
	"flag"
	"fmt"
	"net"
	"sort"
	"strings"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
)

type Topology struct {
	Ui   cli.Ui
	Cmd  string
	zone string
}

func (this *Topology) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("topology", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.zone == "" {
		forAllZones(func(zkzone *zk.ZkZone) {
			this.displayZoneTopology(zkzone)
		})

		return
	}

	// a single zone
	ensureZoneValid(this.zone)
	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZonePath(this.zone)))
	this.displayZoneTopology(zkzone)

	return
}

type brokerHostInfo struct {
	ports    []int
	leadingN int // being leader of how many partitions
}

func newBrokerHostInfo() *brokerHostInfo {
	return &brokerHostInfo{
		ports: make([]int, 0),
	}
}

func (this *brokerHostInfo) addPort(port int) {
	this.ports = append(this.ports, port)
}

func (this *Topology) displayZoneTopology(zkzone *zk.ZkZone) {
	this.Ui.Output(zkzone.Name())

	instances := make(map[string]*brokerHostInfo)
	zkzone.WithinBrokers(func(cluster string, brokers map[string]*zk.BrokerZnode) {
		if len(brokers) == 0 {
			return
		}

		for _, broker := range brokers {
			if _, present := instances[broker.Host]; !present {
				instances[broker.Host] = newBrokerHostInfo()
			}
			instances[broker.Host].addPort(broker.Port)
		}

		// find how many partitions a broker is leading
		zkcluster := zkzone.NewCluster(cluster)
		brokerList := zkcluster.BrokerList()
		if len(brokerList) == 0 {
			return
		}
		kfk, err := sarama.NewClient(zkcluster.BrokerList(), sarama.NewConfig())
		if err != nil {
			this.Ui.Error(fmt.Sprintf("    %s", err.Error()))
			return
		}
		topics, err := kfk.Topics()<commit_msg>more detailed err msg<commit_after>package command

import (
	"flag"
	"fmt"
	"net"
	"sort"
	"strings"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
)

type Topology struct {
	Ui   cli.Ui
	Cmd  string
	zone string
}

func (this *Topology) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("topology", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.zone == "" {
		forAllZones(func(zkzone *zk.ZkZone) {
			this.displayZoneTopology(zkzone)
		})

		return
	}

	// a single zone
	ensureZoneValid(this.zone)
	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZonePath(this.zone)))
	this.displayZoneTopology(zkzone)

	return
}

type brokerHostInfo struct {
	ports    []int
	leadingN int // being leader of how many partitions
}

func newBrokerHostInfo() *brokerHostInfo {
	return &brokerHostInfo{
		ports: make([]int, 0),
	}
}

func (this *brokerHostInfo) addPort(port int) {
	this.ports = append(this.ports, port)
}

func (this *Topology) displayZoneTopology(zkzone *zk.ZkZone) {
	this.Ui.Output(zkzone.Name())

	instances := make(map[string]*brokerHostInfo)
	zkzone.WithinBrokers(func(cluster string, brokers map[string]*zk.BrokerZnode) {
		if len(brokers) == 0 {
			return
		}

		for _, broker := range brokers {
			if _, present := instances[broker.Host]; !present {
				instances[broker.Host] = newBrokerHostInfo()
			}
			instances[broker.Host].addPort(broker.Port)
		}

		// find how many partitions a broker is leading
		zkcluster := zkzone.NewCluster(cluster)
		brokerList := zkcluster.BrokerList()
		if len(brokerList) == 0 {
			return
		}
		kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
		if err != nil {
			this.Ui.Error(color.Red("    %+v %s", brokerList, err.Error()))
			return
		}

		topics, err := kfk.Topics()
{'fraction_non_alphanumeric': 0.14836942992282798, 'fraction_numerical': 0.0019915359721184964, 'mean_word_length': 4.098984771573604, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '89531', 'n_tokens_mistral': 1640, 'n_tokens_neox': 1518, 'n_words': 294}