<commit_before>package command

import (
	"flag"
	"fmt"
	"net"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/bjtime"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/termui"
	"github.com/nsf/termbox-go"
)

type Top struct {
	Ui  cli.Ui
	Cmd string

	zone           string
	clusterPattern string

	mu sync.Mutex

	round int

	who            string
	limit          int
	topInterval    time.Duration
	batchMode      bool
	dashboardGraph bool
	topicPattern   string
	longFmt        bool
	skipIpPrefix   bool

	brokers          map[string][]string
	counters         map[string]float64 // key is cluster:topic
	lastCounters     map[string]float64
	consumerCounters map[string]float64
	partitions       map[string]int

	totalConsumerMps []float64
	totalMps         []float64 // for the dashboard graph
	maxMps           float64
}

func (this *Top) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("top", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.DurationVar(&this.topInterval, "i", time.Second*5, "refresh interval")
	cmdFlags.StringVar(&this.clusterPattern, "c", "", "")
	cmdFlags.IntVar(&this.limit, "n", 33, "")
	cmdFlags.BoolVar(&this.skipIpPrefix, "noiphead", true, "")
	cmdFlags.StringVar(&this.who, "who", "producer", "")
	cmdFlags.BoolVar(&this.dashboardGraph, "d", false, "")
	cmdFlags.BoolVar(&this.longFmt, "l", true, "")
	cmdFlags.BoolVar(&this.batchMode, "b", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.dashboardGraph {
		if this.topInterval.Seconds() < 20 {
			this.topInterval = 20 * time.Second
		}
		this.who = "both"
		go this.clusterOffsetSummary()
	}

	if this.who == "c" || this.who == "consumer" {
		if this.topInterval.Seconds() < 20 {
			this.topInterval = 20 * time.Second // consumer groups only refresh offset per minute
		}

	}

	this.brokers = make(map[string][]string)
	this.counters = make(map[string]float64)
	this.lastCounters = make(map[string]float64)
	this.partitions = make(map[string]int)
	this.consumerCounters = make(map[string]float64)
	this.totalMps = make([]float64, 0, 1000)
	this.totalConsumerMps = make([]float64, 0, 1000)

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), this.clusterPattern) {
			return
		}

		switch this.who {
		case "p", "producer":
			go this.clusterTopProducers(zkcluster)

		case "c", "consumer":
			go this.clusterTopConsumers(zkcluster)

		case "both":
			go this.clusterTopConsumers(zkcluster)
			go this.clusterTopProducers(zkcluster)

		default:
			this.Ui.Error(fmt.Sprintf("unknown type: %s", this.who))
		}
	})

	if this.dashboardGraph {
		this.drawDashboard()
		return
	}

	ticker := time.NewTicker(this.topInterval)
	defer ticker.Stop()
	keyboardPressed := make(chan struct{})
	go func() {
		var b []byte = make([]byte, 1)
		for {
			os.Stdin.Read(b)
			keyboardPressed <- struct{}{}
		}
	}()
	for {
		select {
		case <-keyboardPressed:
		case <-ticker.C:
		}

		if this.batchMode {
			this.Ui.Output(bjtime.TimeToString(bjtime.NowBj()))
		} else {
			refreshScreen()
		}

		// header
		if this.longFmt {
			this.Ui.Output(fmt.Sprintf("%-9s %15s %-30s %35s %20s %15s",
				this.who, "cluster", "brokers", "topic", "num", "mps")) // mps=msg per second
			this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 129)))
		} else {
			this.Ui.Output(fmt.Sprintf("%-9s %20s %50s %20s %15s",
				this.who, "cluster", "topic", "num", "mps")) // mps=msg per second
			this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 118)))
		}

		this.showAndResetCounters()
	}

	return
}

func (this *Top) drawDashboard() {
	err := termui.Init()
	if err != nil {
		panic(err)
	}
	defer termui.Close()

	termui.UseTheme("helloworld")

	maxRound := termui.TermWidth() / 2
	refreshProducerData := func(round int) []float64 {
		this.showAndResetCounters()
		if round%maxRound == (maxRound - 5) {
			this.mu.Lock()
			tailLen := len(this.totalMps) - 2
			tail := this.totalMps[tailLen:]
			this.totalMps = make([]float64, tailLen, 1000)
			copy(this.totalMps, tail)
			this.mu.Unlock()
		}
		return this.totalMps
	}

	refreshConsumerData := func(round int) []float64 {
		if round%maxRound == (maxRound - 5) {
			this.mu.Lock()
			tailLen := len(this.totalConsumerMps) - 2
			tail := this.totalConsumerMps[tailLen:]
			this.totalConsumerMps = make([]float64, tailLen, 1000)
			copy(this.totalConsumerMps, tail)
			this.mu.Unlock()
		}
		return this.totalConsumerMps
	}

	producerChart := termui.NewLineChart()
	producerChart.Mode = "dot"
	producerChart.Border.Label = fmt.Sprintf("producer mps totals: %s %s %s",
		this.zone, this.clusterPattern, this.topicPattern)
	producerChart.Data = refreshProducerData(0)
	producerChart.Width = termui.TermWidth() / 2
	producerChart.Height = termui.TermHeight()
	producerChart.X = 0<commit_msg>still on 'gk top' rewind, this time it works<commit_after>package command

import (
	"flag"
	"fmt"
	"net"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/bjtime"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/termui"
	"github.com/nsf/termbox-go"
)

type Top struct {
	Ui  cli.Ui
	Cmd string

	zone           string
	clusterPattern string

	mu sync.Mutex

	round int

	who            string
	limit          int
	topInterval    time.Duration
	batchMode      bool
	dashboardGraph bool
	topicPattern   string
	longFmt        bool
	skipIpPrefix   bool

	brokers          map[string][]string
	counters         map[string]float64 // key is cluster:topic
	lastCounters     map[string]float64
	consumerCounters map[string]float64
	partitions       map[string]int

	totalConsumerMps []float64
	totalMps         []float64 // for the dashboard graph
	maxMps           float64
}

func (this *Top) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("top", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.DurationVar(&this.topInterval, "i", time.Second*5, "refresh interval")
	cmdFlags.StringVar(&this.clusterPattern, "c", "", "")
	cmdFlags.IntVar(&this.limit, "n", 33, "")
	cmdFlags.BoolVar(&this.skipIpPrefix, "noiphead", true, "")
	cmdFlags.StringVar(&this.who, "who", "producer", "")
	cmdFlags.BoolVar(&this.dashboardGraph, "d", false, "")
	cmdFlags.BoolVar(&this.longFmt, "l", true, "")
	cmdFlags.BoolVar(&this.batchMode, "b", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.dashboardGraph {
		if this.topInterval.Seconds() < 20 {
			this.topInterval = 20 * time.Second
		}
		this.who = "both"
		go this.clusterOffsetSummary()
	}

	if this.who == "c" || this.who == "consumer" {
		if this.topInterval.Seconds() < 20 {
			this.topInterval = 20 * time.Second // consumer groups only refresh offset per minute
		}

	}

	this.brokers = make(map[string][]string)
	this.counters = make(map[string]float64)
	this.lastCounters = make(map[string]float64)
	this.partitions = make(map[string]int)
	this.consumerCounters = make(map[string]float64)
	this.totalMps = make([]float64, 0, 1000)
	this.totalConsumerMps = make([]float64, 0, 1000)

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), this.clusterPattern) {
			return
		}

		switch this.who {
		case "p", "producer":
			go this.clusterTopProducers(zkcluster)

		case "c", "consumer":
			go this.clusterTopConsumers(zkcluster)

		case "both":
			go this.clusterTopConsumers(zkcluster)
			go this.clusterTopProducers(zkcluster)

		default:
			this.Ui.Error(fmt.Sprintf("unknown type: %s", this.who))
		}
	})

	if this.dashboardGraph {
		this.drawDashboard()
		return
	}

	ticker := time.NewTicker(this.topInterval)
	defer ticker.Stop()
	keyboardPressed := make(chan struct{})
	go func() {
		var b []byte = make([]byte, 1)
		for {
			os.Stdin.Read(b)
			keyboardPressed <- struct{}{}
		}
	}()
	for {
		select {
		case <-keyboardPressed:
		case <-ticker.C:
		}

		if this.batchMode {
			this.Ui.Output(bjtime.TimeToString(bjtime.NowBj()))
		} else {
			refreshScreen()
		}

		// header
		if this.longFmt {
			this.Ui.Output(fmt.Sprintf("%-9s %15s %-30s %35s %20s %15s",
				this.who, "cluster", "brokers", "topic", "num", "mps")) // mps=msg per second
			this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 129)))
		} else {
			this.Ui.Output(fmt.Sprintf("%-9s %20s %50s %20s %15s",
				this.who, "cluster", "topic", "num", "mps")) // mps=msg per second
			this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 118)))
		}

		this.showAndResetCounters()
	}

	return
}

func (this *Top) drawDashboard() {
	err := termui.Init()
	if err != nil {
		panic(err)
	}
	defer termui.Close()

	termui.UseTheme("helloworld")

	maxRound := termui.TermWidth() / 2
	refreshProducerData := func(round int) []float64 {
		this.showAndResetCounters()
		if round%maxRound == (maxRound - 5) {
			this.mu.Lock()
			tailLen := len(this.totalMps) - 2
			tail := this.totalMps[tailLen:]
			this.totalMps = make([]float64, 2, 1000)
			copy(this.totalMps, tail)
			this.mu.Unlock()
		}
		return this.totalMps
	}

	refreshConsumerData := func(round int) []float64 {
		if round%maxRound == (maxRound - 5) {
			this.mu.Lock()
			tailLen := len(this.totalConsumerMps) - 2
			tail := this.totalConsumerMps[tailLen:]
			this.totalConsumerMps = make([]float64, 2, 1000)
			copy(this.totalConsumerMps, tail)
			this.mu.Unlock()
		}
		return this.totalConsumerMps
	}

	producerChart := termui.NewLineChart()
	producerChart.Mode = "dot"
	producerChart.Border.Label = fmt.Sprintf("producer mps totals: %s %s %s",
		this.zone, this.clusterPattern, this.topicPattern)
	producerChart.Data = refreshProducerData(0)
	producerChart.Width = termui.TermWidth() / 2
	producerChart.Height = termui.TermHeight()
	producerChart.X = 0
{'fraction_non_alphanumeric': 0.15437014295308452, 'fraction_numerical': 0.018420800153506667, 'mean_word_length': 3.8416163492800743, 'pattern_counts': {'":': 6, '<': 13, '<?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': '2318790', 'n_tokens_mistral': 4285, 'n_tokens_neox': 3869, 'n_words': 729}