github.com/Jeffail/benthos/v3@v3.65.0/lib/output/writer/nsq.go (about)

     1  package writer
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"fmt"
     7  	"io"
     8  	llog "log"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/Jeffail/benthos/v3/internal/bloblang/field"
    13  	"github.com/Jeffail/benthos/v3/internal/interop"
    14  	"github.com/Jeffail/benthos/v3/lib/log"
    15  	"github.com/Jeffail/benthos/v3/lib/metrics"
    16  	"github.com/Jeffail/benthos/v3/lib/types"
    17  	btls "github.com/Jeffail/benthos/v3/lib/util/tls"
    18  	nsq "github.com/nsqio/go-nsq"
    19  )
    20  
    21  //------------------------------------------------------------------------------
    22  
    23  // NSQConfig contains configuration fields for the NSQ output type.
    24  type NSQConfig struct {
    25  	Address     string      `json:"nsqd_tcp_address" yaml:"nsqd_tcp_address"`
    26  	Topic       string      `json:"topic" yaml:"topic"`
    27  	UserAgent   string      `json:"user_agent" yaml:"user_agent"`
    28  	TLS         btls.Config `json:"tls" yaml:"tls"`
    29  	MaxInFlight int         `json:"max_in_flight" yaml:"max_in_flight"`
    30  }
    31  
    32  // NewNSQConfig creates a new NSQConfig with default values.
    33  func NewNSQConfig() NSQConfig {
    34  	return NSQConfig{
    35  		Address:     "localhost:4150",
    36  		Topic:       "benthos_messages",
    37  		UserAgent:   "benthos_producer",
    38  		TLS:         btls.NewConfig(),
    39  		MaxInFlight: 1,
    40  	}
    41  }
    42  
    43  //------------------------------------------------------------------------------
    44  
    45  // NSQ is an output type that serves NSQ messages.
    46  type NSQ struct {
    47  	log log.Modular
    48  
    49  	topicStr *field.Expression
    50  
    51  	tlsConf  *tls.Config
    52  	connMut  sync.RWMutex
    53  	producer *nsq.Producer
    54  
    55  	conf NSQConfig
    56  }
    57  
    58  // NewNSQ creates a new NSQ output type.
    59  //
    60  // Deprecated: use the V2 API instead.
    61  func NewNSQ(conf NSQConfig, log log.Modular, stats metrics.Type) (*NSQ, error) {
    62  	return NewNSQV2(conf, types.NoopMgr(), log, stats)
    63  }
    64  
    65  // NewNSQV2 creates a new NSQ output type.
    66  func NewNSQV2(conf NSQConfig, mgr types.Manager, log log.Modular, stats metrics.Type) (*NSQ, error) {
    67  	n := NSQ{
    68  		log:  log,
    69  		conf: conf,
    70  	}
    71  	var err error
    72  	if n.topicStr, err = interop.NewBloblangField(mgr, conf.Topic); err != nil {
    73  		return nil, fmt.Errorf("failed to parse topic expression: %v", err)
    74  	}
    75  	if conf.TLS.Enabled {
    76  		if n.tlsConf, err = conf.TLS.Get(); err != nil {
    77  			return nil, err
    78  		}
    79  	}
    80  	return &n, nil
    81  }
    82  
    83  //------------------------------------------------------------------------------
    84  
    85  // ConnectWithContext attempts to establish a connection to NSQ servers.
    86  func (n *NSQ) ConnectWithContext(ctx context.Context) error {
    87  	return n.Connect()
    88  }
    89  
    90  // Connect attempts to establish a connection to NSQ servers.
    91  func (n *NSQ) Connect() error {
    92  	n.connMut.Lock()
    93  	defer n.connMut.Unlock()
    94  
    95  	cfg := nsq.NewConfig()
    96  	cfg.UserAgent = n.conf.UserAgent
    97  	if n.tlsConf != nil {
    98  		cfg.TlsV1 = true
    99  		cfg.TlsConfig = n.tlsConf
   100  	}
   101  
   102  	producer, err := nsq.NewProducer(n.conf.Address, cfg)
   103  	if err != nil {
   104  		return err
   105  	}
   106  
   107  	producer.SetLogger(llog.New(io.Discard, "", llog.Flags()), nsq.LogLevelError)
   108  
   109  	if err := producer.Ping(); err != nil {
   110  		return err
   111  	}
   112  	n.producer = producer
   113  	n.log.Infof("Sending NSQ messages to address: %s\n", n.conf.Address)
   114  	return nil
   115  }
   116  
   117  // WriteWithContext attempts to write a message.
   118  func (n *NSQ) WriteWithContext(ctx context.Context, msg types.Message) error {
   119  	return n.Write(msg)
   120  }
   121  
   122  // Write attempts to write a message.
   123  func (n *NSQ) Write(msg types.Message) error {
   124  	n.connMut.RLock()
   125  	prod := n.producer
   126  	n.connMut.RUnlock()
   127  
   128  	if prod == nil {
   129  		return types.ErrNotConnected
   130  	}
   131  
   132  	return IterateBatchedSend(msg, func(i int, p types.Part) error {
   133  		return prod.Publish(n.topicStr.String(i, msg), p.Get())
   134  	})
   135  }
   136  
   137  // CloseAsync shuts down the NSQ output and stops processing messages.
   138  func (n *NSQ) CloseAsync() {
   139  	go func() {
   140  		n.connMut.Lock()
   141  		if n.producer != nil {
   142  			n.producer.Stop()
   143  			n.producer = nil
   144  		}
   145  		n.connMut.Unlock()
   146  	}()
   147  }
   148  
   149  // WaitForClose blocks until the NSQ output has closed down.
   150  func (n *NSQ) WaitForClose(timeout time.Duration) error {
   151  	return nil
   152  }
   153  
   154  //------------------------------------------------------------------------------