func (s *SketchHandler) TimerLoop(d time.Duration, r time.Duration) {
displayTicker := time.NewTicker(d)
resetTicker := time.NewTicker(r)
for {
select {
case <-resetTicker.C:
s.Lock()
fmt.Fprintf(os.Stdout, "\n-[reset sketch. next reset in %s]-\n", r)
s.Sketch = countmin.NewCountMinSketch(7, 20000)
s.HeavyHitters = make([]*HeavyHitter, 0)
s.messageCount = 0
s.Unlock()
case <-displayTicker.C:
fmt.Fprintf(os.Stdout, "\n----- %d messages ----\n", s.messageCount)
s.Lock()
sort.Sort(HeavyHittersByValue{s.HeavyHitters})
for _, h := range s.HeavyHitters {
fmt.Fprintf(os.Stdout, "%d - %s\n", h.Value, h.Key)
}
s.Unlock()
}
}
}
func main() {
flag.Parse()
if *showVersion {
fmt.Printf("nsq_sketch v%s go-nsq/%s\n", VERSION, nsq.VERSION)
return
}
if *topic == "" || *channel == "" {
log.Fatalf("--topic and --channel are required")
}
if *sketchKey == "" {
log.Fatalf("--key required to speicfy which message element to sketch")
}
if *maxInFlight < 0 {
log.Fatalf("--max-in-flight must be > 0")
}
if *topN < 1 {
log.Fatal("--top must be > 0 to specify the number of items to track")
}
if len(nsqdTCPAddrs) == 0 && len(lookupdHTTPAddrs) == 0 {
log.Fatalf("--nsqd-tcp-address or --lookupd-http-address required.")
}
if len(nsqdTCPAddrs) != 0 && len(lookupdHTTPAddrs) != 0 {
log.Fatalf("use --nsqd-tcp-address or --lookupd-http-address not both")
}
displayInterval, err := time.ParseDuration(*sketchDisplayInterval)
if err != nil {
log.Fatalf("failed parsing --interval %s. %s", *sketchDisplayInterval, err.Error())
}
resetInterval, err := time.ParseDuration(*sketchResetInterval)
if err != nil {
log.Fatalf("failed parsing --reset-every %s. %s", *sketchResetInterval, err.Error())
}
log.Printf("Sketching for %s intervals. Displaying %d top items every %s", resetInterval, *topN, displayInterval)
termChan := make(chan os.Signal, 1)
signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)
s := &SketchHandler{
Sketch: countmin.NewCountMinSketch(*sketchWidth, *sketchHeight),
HeavyHitters: make([]*HeavyHitter, 0),
}
if !*verbose {
log.SetOutput(ioutil.Discard)
}
cfg := nsq.NewConfig()
cfg.MaxInFlight = *maxInFlight
r, err := nsq.NewConsumer(*topic, *channel, cfg)
if err != nil {
log.Fatalf(err.Error())
}
r.AddHandler(s)
go s.TimerLoop(displayInterval, resetInterval)
go func() {
select {
case <-termChan:
r.Stop()
}
}()
err = r.ConnectToNSQDs(nsqdTCPAddrs)
if err != nil {
log.Fatalf("%s", err)
}
err = r.ConnectToNSQLookupds(lookupdHTTPAddrs)
if err != nil {
log.Fatalf("%s", err)
}
<-r.StopChan
}