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

     1  // This is an NSQ client that publishes incoming messages from
     2  // stdin to the specified topic.
     3  
     4  package main
     5  
     6  import (
     7  	"bufio"
     8  	"flag"
     9  	"fmt"
    10  	"io"
    11  	"log"
    12  	"os"
    13  	"os/signal"
    14  	"sync/atomic"
    15  	"syscall"
    16  	"time"
    17  
    18  	"github.com/nsqio/go-nsq"
    19  	"github.com/nsqio/nsq/internal/app"
    20  	"github.com/nsqio/nsq/internal/version"
    21  )
    22  
    23  var (
    24  	topic     = flag.String("topic", "", "NSQ topic to publish to")
    25  	delimiter = flag.String("delimiter", "\n", "character to split input from stdin")
    26  
    27  	destNsqdTCPAddrs = app.StringArray{}
    28  )
    29  
    30  func init() {
    31  	flag.Var(&destNsqdTCPAddrs, "nsqd-tcp-address", "destination nsqd TCP address (may be given multiple times)")
    32  }
    33  
    34  func main() {
    35  	cfg := nsq.NewConfig()
    36  	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)")
    37  	rate := flag.Int64("rate", 0, "Throttle messages to n/second. 0 to disable")
    38  
    39  	flag.Parse()
    40  
    41  	if len(*topic) == 0 {
    42  		log.Fatal("--topic required")
    43  	}
    44  
    45  	if len(*delimiter) != 1 {
    46  		log.Fatal("--delimiter must be a single byte")
    47  	}
    48  
    49  	stopChan := make(chan bool)
    50  	termChan := make(chan os.Signal, 1)
    51  	signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)
    52  
    53  	cfg.UserAgent = fmt.Sprintf("to_nsq/%s go-nsq/%s", version.Binary, nsq.VERSION)
    54  
    55  	// make the producers
    56  	producers := make(map[string]*nsq.Producer)
    57  	for _, addr := range destNsqdTCPAddrs {
    58  		producer, err := nsq.NewProducer(addr, cfg)
    59  		if err != nil {
    60  			log.Fatalf("failed to create nsq.Producer - %s", err)
    61  		}
    62  		producers[addr] = producer
    63  	}
    64  
    65  	if len(producers) == 0 {
    66  		log.Fatal("--nsqd-tcp-address required")
    67  	}
    68  
    69  	throttleEnabled := *rate >= 1
    70  	balance := int64(1)
    71  	// avoid divide by 0 if !throttleEnabled
    72  	var interval time.Duration
    73  	if throttleEnabled {
    74  		interval = time.Second / time.Duration(*rate)
    75  	}
    76  	go func() {
    77  		if !throttleEnabled {
    78  			return
    79  		}
    80  		log.Printf("Throttling messages rate to max:%d/second", *rate)
    81  		// every tick increase the number of messages we can send
    82  		for range time.Tick(interval) {
    83  			n := atomic.AddInt64(&balance, 1)
    84  			// if we build up more than 1s of capacity just bound to that
    85  			if n > int64(*rate) {
    86  				atomic.StoreInt64(&balance, int64(*rate))
    87  			}
    88  		}
    89  	}()
    90  
    91  	r := bufio.NewReader(os.Stdin)
    92  	delim := (*delimiter)[0]
    93  	go func() {
    94  		for {
    95  			var err error
    96  			if throttleEnabled {
    97  				currentBalance := atomic.LoadInt64(&balance)
    98  				if currentBalance <= 0 {
    99  					time.Sleep(interval)
   100  				}
   101  				err = readAndPublish(r, delim, producers)
   102  				atomic.AddInt64(&balance, -1)
   103  			} else {
   104  				err = readAndPublish(r, delim, producers)
   105  			}
   106  			if err != nil {
   107  				if err != io.EOF {
   108  					log.Fatal(err)
   109  				}
   110  				close(stopChan)
   111  				break
   112  			}
   113  		}
   114  	}()
   115  
   116  	select {
   117  	case <-termChan:
   118  	case <-stopChan:
   119  	}
   120  
   121  	for _, producer := range producers {
   122  		producer.Stop()
   123  	}
   124  }
   125  
   126  // readAndPublish reads to the delim from r and publishes the bytes
   127  // to the map of producers.
   128  func readAndPublish(r *bufio.Reader, delim byte, producers map[string]*nsq.Producer) error {
   129  	line, readErr := r.ReadBytes(delim)
   130  
   131  	if len(line) > 0 {
   132  		// trim the delimiter
   133  		line = line[:len(line)-1]
   134  	}
   135  
   136  	if len(line) == 0 {
   137  		return readErr
   138  	}
   139  
   140  	for _, producer := range producers {
   141  		err := producer.Publish(*topic, line)
   142  		if err != nil {
   143  			return err
   144  		}
   145  	}
   146  
   147  	return readErr
   148  }