-
Notifications
You must be signed in to change notification settings - Fork 71
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #122 from batchcorp/blinktag/nsq
NSQ read/write/relay support
- Loading branch information
Showing
47 changed files
with
5,556 additions
and
142 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,93 @@ | ||
| package nsq | ||
|
|
||
| import ( | ||
| "crypto/tls" | ||
| "crypto/x509" | ||
| "io/ioutil" | ||
|
|
||
| "github.com/jhump/protoreflect/desc" | ||
| "github.com/nsqio/go-nsq" | ||
| "github.com/pkg/errors" | ||
| "github.com/sirupsen/logrus" | ||
|
|
||
| "github.com/batchcorp/plumber/cli" | ||
| ) | ||
|
|
||
| var ( | ||
| ErrMissingAddress = errors.New("you must specify either --nsqd-address or --lookupd-address") | ||
| ErrChooseAddress = errors.New("you must specify either --nsqd-address or --lookupd-address, not both") | ||
| ErrMissingTLSKey = errors.New("--tls-client-key-file cannot be blank if using TLS") | ||
| ErrMissingTlsCert = errors.New("--tls-client-cert-file cannot be blank if using TLS") | ||
| ErrMissingTLSCA = errors.New("--tls-ca-file cannot be blank if using TLS") | ||
| ) | ||
|
|
||
| // NSQ encapsulates options for calling Read() and Write() methods | ||
| type NSQ struct { | ||
| Options *cli.Options | ||
| MsgDesc *desc.MessageDescriptor | ||
| log *NSQLogger | ||
| } | ||
|
|
||
| // getNSQConfig returns the config needed for creating a new NSQ consumer or producer | ||
| func getNSQConfig(opts *cli.Options) (*nsq.Config, error) { | ||
| config := nsq.NewConfig() | ||
| config.ClientID = opts.NSQ.ClientID | ||
|
|
||
| if opts.NSQ.AuthSecret != "" { | ||
| config.AuthSecret = opts.NSQ.AuthSecret | ||
| } | ||
|
|
||
| if opts.NSQ.InsecureTLS || opts.NSQ.TLSClientCertFile != "" { | ||
| tlsConfig, err := generateTLSConfig(opts) | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "unable to generate TLS config") | ||
| } | ||
|
|
||
| config.TlsConfig = tlsConfig | ||
| } | ||
|
|
||
| return config, nil | ||
| } | ||
|
|
||
| // generateTLSConfig generates necessary TLS config for Dialing to an NSQ server | ||
| func generateTLSConfig(opts *cli.Options) (*tls.Config, error) { | ||
| certpool := x509.NewCertPool() | ||
|
|
||
| pemCerts, err := ioutil.ReadFile(opts.NSQ.TLSCAFile) | ||
| if err == nil { | ||
| certpool.AppendCertsFromPEM(pemCerts) | ||
| } | ||
|
|
||
| // Import client certificate/key pair | ||
| cert, err := tls.LoadX509KeyPair(opts.NSQ.TLSClientCertFile, opts.NSQ.TLSClientKeyFile) | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "unable to load ssl keypair") | ||
| } | ||
|
|
||
| // Just to print out the client certificate.. | ||
| cert.Leaf, err = x509.ParseCertificate(cert.Certificate[0]) | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "unable to parse certificate") | ||
| } | ||
|
|
||
| // Create tls.Config with desired tls properties | ||
| return &tls.Config{ | ||
| RootCAs: certpool, | ||
| ClientAuth: tls.NoClientCert, | ||
| ClientCAs: nil, | ||
| InsecureSkipVerify: opts.NSQ.InsecureTLS, | ||
| Certificates: []tls.Certificate{cert}, | ||
| }, nil | ||
| } | ||
|
|
||
| // NSQLogger wraps logrus and implements the Output() method so we can satisfy the interface | ||
| // requirements for NSQ's logger | ||
| type NSQLogger struct { | ||
| *logrus.Entry | ||
| } | ||
|
|
||
| // Output writes an NSQ log message via logrus | ||
| func (n *NSQLogger) Output(_ int, s string) error { | ||
| n.Info(s) | ||
| return nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,131 @@ | ||
| package nsq | ||
|
|
||
| import ( | ||
| "sync" | ||
|
|
||
| "github.com/jhump/protoreflect/desc" | ||
| "github.com/nsqio/go-nsq" | ||
| "github.com/pkg/errors" | ||
| "github.com/sirupsen/logrus" | ||
|
|
||
| "github.com/batchcorp/plumber/cli" | ||
| "github.com/batchcorp/plumber/pb" | ||
| "github.com/batchcorp/plumber/printer" | ||
| "github.com/batchcorp/plumber/reader" | ||
| ) | ||
|
|
||
| // Read is the entry point function for performing read operations in NSQ | ||
| func Read(opts *cli.Options) error { | ||
| if err := validateReadOptions(opts); err != nil { | ||
| return errors.Wrap(err, "unable to validate read options") | ||
| } | ||
|
|
||
| var mdErr error | ||
| var md *desc.MessageDescriptor | ||
|
|
||
| if opts.ReadProtobufRootMessage != "" { | ||
| md, mdErr = pb.FindMessageDescriptor(opts.ReadProtobufDirs, opts.ReadProtobufRootMessage) | ||
| if mdErr != nil { | ||
| return errors.Wrap(mdErr, "unable to find root message descriptor") | ||
| } | ||
| } | ||
|
|
||
| logger := &NSQLogger{} | ||
| logger.Entry = logrus.WithField("pkg", "nsq") | ||
|
|
||
| n := &NSQ{ | ||
| Options: opts, | ||
| MsgDesc: md, | ||
| log: logger, | ||
| } | ||
|
|
||
| return n.Read() | ||
| } | ||
|
|
||
| // Read will attempt to consume one or more messages from a given topic, | ||
| // optionally decode it and/or convert the returned output. | ||
| func (n *NSQ) Read() error { | ||
| config, err := getNSQConfig(n.Options) | ||
| if err != nil { | ||
| return errors.Wrap(err, "unable to create NSQ config") | ||
| } | ||
|
|
||
| consumer, err := nsq.NewConsumer(n.Options.NSQ.Topic, n.Options.NSQ.Channel, config) | ||
| if err != nil { | ||
| return errors.Wrap(err, "Could not start NSQ consumer") | ||
| } | ||
|
|
||
| logLevel := nsq.LogLevelError | ||
| if n.Options.Debug { | ||
| logLevel = nsq.LogLevelDebug | ||
| } | ||
|
|
||
| // Use logrus for NSQ logs | ||
| consumer.SetLogger(n.log, logLevel) | ||
|
|
||
| wg := &sync.WaitGroup{} | ||
| wg.Add(1) | ||
|
|
||
| count := 1 | ||
|
|
||
| consumer.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error { | ||
| data, err := reader.Decode(n.Options, n.MsgDesc, msg.Body) | ||
| if err != nil { | ||
| return errors.Wrap(err, "unable to decode msg") | ||
| } | ||
|
|
||
| printer.PrintNSQResult(n.Options, count, msg, data) | ||
|
|
||
| if !n.Options.ReadFollow { | ||
| wg.Done() | ||
| } | ||
| count++ | ||
| return nil | ||
| })) | ||
|
|
||
| // Connect to correct server. Reading can be done directly from an NSQD server | ||
| // or let lookupd find the correct one. | ||
| if n.Options.NSQ.NSQLookupDAddress != "" { | ||
| if err := consumer.ConnectToNSQLookupd(n.Options.NSQ.NSQLookupDAddress); err != nil { | ||
| return errors.Wrap(err, "could not connect to nsqlookupd") | ||
| } | ||
| } else { | ||
| if err := consumer.ConnectToNSQD(n.Options.NSQ.NSQDAddress); err != nil { | ||
| return errors.Wrap(err, "could not connect to nsqd") | ||
| } | ||
| } | ||
| defer consumer.Stop() | ||
|
|
||
| n.log.Infof("Waiting for messages...") | ||
|
|
||
| wg.Wait() | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| // validateReadOptions ensures all necessary flags have values required for reading from NSQ | ||
| func validateReadOptions(opts *cli.Options) error { | ||
| if opts.NSQ.NSQDAddress == "" && opts.NSQ.NSQLookupDAddress == "" { | ||
| return ErrMissingAddress | ||
| } | ||
|
|
||
| if opts.NSQ.NSQDAddress != "" && opts.NSQ.NSQLookupDAddress != "" { | ||
| return ErrChooseAddress | ||
| } | ||
|
|
||
| if opts.NSQ.TLSCAFile != "" || opts.NSQ.TLSClientCertFile != "" || opts.NSQ.TLSClientKeyFile != "" { | ||
| if opts.NSQ.TLSClientKeyFile == "" { | ||
| return ErrMissingTLSKey | ||
| } | ||
|
|
||
| if opts.NSQ.TLSClientCertFile == "" { | ||
| return ErrMissingTlsCert | ||
| } | ||
|
|
||
| if opts.NSQ.TLSCAFile == "" { | ||
| return ErrMissingTLSCA | ||
| } | ||
| } | ||
|
|
||
| return nil | ||
| } |
Oops, something went wrong.