Ejemplo n.º 1
0
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()
		}
	}
}
Ejemplo n.º 2
0
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
}