-
Notifications
You must be signed in to change notification settings - Fork 0
/
nsq_tail.go
123 lines (104 loc) · 2.88 KB
/
nsq_tail.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package main
import (
"flag"
"fmt"
"log"
"math/rand"
"os"
"os/signal"
"syscall"
"time"
"github.com/bitly/go-nsq"
"github.com/bitly/nsq/util"
)
var (
showVersion = flag.Bool("version", false, "print version string")
topic = flag.String("topic", "", "nsq topic")
channel = flag.String("channel", "", "nsq channel")
maxInFlight = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")
totalMessages = flag.Int("n", 0, "total messages to show (will wait if starved)")
readerOpts = util.StringArray{}
nsqdTCPAddrs = util.StringArray{}
lookupdHTTPAddrs = util.StringArray{}
)
func init() {
flag.Var(&readerOpts, "reader-opt", "option to passthrough to nsq.Reader (may be given multiple times)")
flag.Var(&nsqdTCPAddrs, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
}
type TailHandler struct {
totalMessages int
messagesShown int
}
func (th *TailHandler) HandleMessage(m *nsq.Message) error {
th.messagesShown++
_, err := os.Stdout.Write(m.Body)
if err != nil {
log.Fatalf("ERROR: failed to write to os.Stdout - %s", err.Error())
}
_, err = os.Stdout.WriteString("\n")
if err != nil {
log.Fatalf("ERROR: failed to write to os.Stdout - %s", err.Error())
}
if th.totalMessages > 0 && th.messagesShown >= th.totalMessages {
os.Exit(0)
}
return nil
}
func main() {
flag.Parse()
if *showVersion {
fmt.Printf("nsq_tail v%s\n", util.BINARY_VERSION)
return
}
if *channel == "" {
rand.Seed(time.Now().UnixNano())
*channel = fmt.Sprintf("tail%06d#ephemeral", rand.Int()%999999)
}
if *topic == "" {
log.Fatalf("--topic is required")
}
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")
}
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
r, err := nsq.NewReader(*topic, *channel)
if err != nil {
log.Fatalf(err.Error())
}
err = util.ParseReaderOpts(r, readerOpts)
if err != nil {
log.Fatalf(err.Error())
}
// Don't ask for more messages than we want
if *totalMessages > 0 && *totalMessages < *maxInFlight {
*maxInFlight = *totalMessages
}
r.SetMaxInFlight(*maxInFlight)
r.AddHandler(&TailHandler{totalMessages: *totalMessages})
for _, addrString := range nsqdTCPAddrs {
err := r.ConnectToNSQ(addrString)
if err != nil {
log.Fatalf(err.Error())
}
}
for _, addrString := range lookupdHTTPAddrs {
log.Printf("lookupd addr %s", addrString)
err := r.ConnectToLookupd(addrString)
if err != nil {
log.Fatalf(err.Error())
}
}
for {
select {
case <-r.ExitChan:
return
case <-sigChan:
r.Stop()
}
}
}