Merge pull request #788 from jehiah/to_nsq_rate_788

apps/to_nsq: add --rate to throttle message creation

apps/to_nsq/README.md

@@ -1,33 +1,33 @@
 # to_nsq
 
-A tool for publishing to an nsq topic data from `stdin`.
+A tool for publishing to an nsq topic with data from `stdin`.
 
 ## Usage
 
-Publish each line of a file:
-
 ```
-cat source.txt | to_nsq -topic="topic" -nsqd-tcp-address="127.0.0.1:4150"
+Usage of ./to_nsq:
+  -delimiter string
+    	character to split input from stdin (default "\n")
+  -nsqd-tcp-address value
+    	destination nsqd TCP address (may be given multiple times)
+  -producer-opt value
+    	option to passthrough to nsq.Producer (may be given multiple times, http://godoc.org/github.com/nsqio/go-nsq#Config)
+  -rate int
+    	Throttle messages to n/second. 0 to disable
+  -topic string
+    	NSQ topic to publish to
 ```
+
+### Examples
 
-Publish manually entered lines in a shell:
-
-```
-to_nsq -topic="topic" -nsqd-tcp-address="127.0.0.1:4150"
-one
-two
-three
-(Ctrl+C to stop)
-```
-
-Publish comma separated values from a source file:
+Publish each line of a file:
 
-```
-cat source.txt | to_nsq -delimiter="," -topic="topic" -nsqd-tcp-address="127.0.0.1:4150"
+```bash
+$ cat source.txt | to_nsq -topic="topic" -nsqd-tcp-address="127.0.0.1:4150"
 ```
 
 Publish three messages, in one go:
 
-```
-echo "one,two,three" | to_nsq -delimiter="," -topic="topic" -nsqd-tcp-address="127.0.0.1:4150"
+```bash
+$ echo "one,two,three" | to_nsq -delimiter="," -topic="topic" -nsqd-tcp-address="127.0.0.1:4150"
 ```

apps/to_nsq/to_nsq.go

@@ -11,7 +11,9 @@ import (
 	"log"
 	"os"
 	"os/signal"
+	"sync/atomic"
 	"syscall"
+	"time"
 
 	"github.com/nsqio/go-nsq"
 	"github.com/nsqio/nsq/internal/app"
@@ -20,7 +22,7 @@ import (
 
 var (
 	topic     = flag.String("topic", "", "NSQ topic to publish to")
-	delimiter = flag.String("delimiter", "\n", "character to split input from stdin (defaults to '\n')")
+	delimiter = flag.String("delimiter", "\n", "character to split input from stdin")
 
 	destNsqdTCPAddrs = app.StringArray{}
 )
@@ -32,6 +34,7 @@ func init() {
 func main() {
 	cfg := nsq.NewConfig()
 	flag.Var(&nsq.ConfigFlag{cfg}, "producer-opt", "option to passthrough to nsq.Producer (may be given multiple times, http://godoc.org/github.com/nsqio/go-nsq#Config)")
+	rate := flag.Int64("rate", 0, "Throttle messages to n/second. 0 to disable")
 
 	flag.Parse()
 
@@ -63,11 +66,39 @@ func main() {
 		log.Fatal("--nsqd-tcp-address required")
 	}
 
+	throttleEnabled := *rate >= 1
+	balance := int64(1)
+	interval := time.Second / time.Duration(*rate)
+	go func() {
+		if !throttleEnabled {
+			return
+		}
+		log.Printf("Throttling messages rate to max:%d/second", *rate)
+		// every tick increase the number of messages we can send
+		for _ = range time.Tick(interval) {
+			n := atomic.AddInt64(&balance, 1)
+			// if we build up more than 1s of capacity just bound to that
+			if n > int64(*rate) {
+				atomic.StoreInt64(&balance, int64(*rate))
+			}
+		}
+	}()
+
 	r := bufio.NewReader(os.Stdin)
 	delim := (*delimiter)[0]
 	go func() {
 		for {
-			err := readAndPublish(r, delim, producers)
+			var err error
+			if throttleEnabled {
+				currentBalance := atomic.LoadInt64(&balance)
+				if currentBalance <= 0 {
+					time.Sleep(interval)
+				}
+				err = readAndPublish(r, delim, producers)
+				atomic.AddInt64(&balance, -1)
+			} else {
+				err = readAndPublish(r, delim, producers)
+			}
 			if err != nil {
 				if err != io.EOF {
 					log.Fatal(err)