github.com/nsqio/nsq@v1.3.0/apps/nsq_tail/nsq_tail.go (about)

     1  package main
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"log"
     7  	"math/rand"
     8  	"os"
     9  	"os/signal"
    10  	"syscall"
    11  	"time"
    12  
    13  	"github.com/nsqio/go-nsq"
    14  	"github.com/nsqio/nsq/internal/app"
    15  	"github.com/nsqio/nsq/internal/version"
    16  )
    17  
    18  var (
    19  	showVersion = flag.Bool("version", false, "print version string")
    20  
    21  	channel       = flag.String("channel", "", "NSQ channel")
    22  	maxInFlight   = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")
    23  	totalMessages = flag.Int("n", 0, "total messages to show (will wait if starved)")
    24  	printTopic    = flag.Bool("print-topic", false, "print topic name where message was received")
    25  
    26  	nsqdTCPAddrs     = app.StringArray{}
    27  	lookupdHTTPAddrs = app.StringArray{}
    28  	topics           = app.StringArray{}
    29  )
    30  
    31  func init() {
    32  	flag.Var(&nsqdTCPAddrs, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
    33  	flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
    34  	flag.Var(&topics, "topic", "NSQ topic (may be given multiple times)")
    35  }
    36  
    37  type TailHandler struct {
    38  	topicName     string
    39  	totalMessages int
    40  	messagesShown int
    41  }
    42  
    43  func (th *TailHandler) HandleMessage(m *nsq.Message) error {
    44  	th.messagesShown++
    45  
    46  	if *printTopic {
    47  		_, err := os.Stdout.WriteString(th.topicName)
    48  		if err != nil {
    49  			log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
    50  		}
    51  		_, err = os.Stdout.WriteString(" | ")
    52  		if err != nil {
    53  			log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
    54  		}
    55  	}
    56  
    57  	_, err := os.Stdout.Write(m.Body)
    58  	if err != nil {
    59  		log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
    60  	}
    61  	_, err = os.Stdout.WriteString("\n")
    62  	if err != nil {
    63  		log.Fatalf("ERROR: failed to write to os.Stdout - %s", err)
    64  	}
    65  	if th.totalMessages > 0 && th.messagesShown >= th.totalMessages {
    66  		os.Exit(0)
    67  	}
    68  	return nil
    69  }
    70  
    71  func main() {
    72  	cfg := nsq.NewConfig()
    73  
    74  	flag.Var(&nsq.ConfigFlag{cfg}, "consumer-opt", "option to passthrough to nsq.Consumer (may be given multiple times, http://godoc.org/github.com/nsqio/go-nsq#Config)")
    75  	flag.Parse()
    76  
    77  	if *showVersion {
    78  		fmt.Printf("nsq_tail v%s\n", version.Binary)
    79  		return
    80  	}
    81  
    82  	if *channel == "" {
    83  		rand.Seed(time.Now().UnixNano())
    84  		*channel = fmt.Sprintf("tail%06d#ephemeral", rand.Int()%999999)
    85  	}
    86  
    87  	if len(nsqdTCPAddrs) == 0 && len(lookupdHTTPAddrs) == 0 {
    88  		log.Fatal("--nsqd-tcp-address or --lookupd-http-address required")
    89  	}
    90  	if len(nsqdTCPAddrs) > 0 && len(lookupdHTTPAddrs) > 0 {
    91  		log.Fatal("use --nsqd-tcp-address or --lookupd-http-address not both")
    92  	}
    93  	if len(topics) == 0 {
    94  		log.Fatal("--topic required")
    95  	}
    96  
    97  	sigChan := make(chan os.Signal, 1)
    98  	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
    99  
   100  	// Don't ask for more messages than we want
   101  	if *totalMessages > 0 && *totalMessages < *maxInFlight {
   102  		*maxInFlight = *totalMessages
   103  	}
   104  
   105  	cfg.UserAgent = fmt.Sprintf("nsq_tail/%s go-nsq/%s", version.Binary, nsq.VERSION)
   106  	cfg.MaxInFlight = *maxInFlight
   107  
   108  	consumers := []*nsq.Consumer{}
   109  	for i := 0; i < len(topics); i++ {
   110  		log.Printf("Adding consumer for topic: %s\n", topics[i])
   111  
   112  		consumer, err := nsq.NewConsumer(topics[i], *channel, cfg)
   113  		if err != nil {
   114  			log.Fatal(err)
   115  		}
   116  
   117  		consumer.AddHandler(&TailHandler{topicName: topics[i], totalMessages: *totalMessages})
   118  
   119  		err = consumer.ConnectToNSQDs(nsqdTCPAddrs)
   120  		if err != nil {
   121  			log.Fatal(err)
   122  		}
   123  
   124  		err = consumer.ConnectToNSQLookupds(lookupdHTTPAddrs)
   125  		if err != nil {
   126  			log.Fatal(err)
   127  		}
   128  
   129  		consumers = append(consumers, consumer)
   130  	}
   131  
   132  	<-sigChan
   133  
   134  	for _, consumer := range consumers {
   135  		consumer.Stop()
   136  	}
   137  	for _, consumer := range consumers {
   138  		<-consumer.StopChan
   139  	}
   140  }