github.com/Jeffail/benthos/v3@v3.65.0/internal/impl/nats/output_jetstream.go (about)

     1  package nats
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"strings"
     7  	"sync"
     8  
     9  	"github.com/Jeffail/benthos/v3/internal/impl/nats/auth"
    10  	"github.com/Jeffail/benthos/v3/internal/shutdown"
    11  	"github.com/Jeffail/benthos/v3/lib/output"
    12  	"github.com/Jeffail/benthos/v3/public/service"
    13  	"github.com/nats-io/nats.go"
    14  )
    15  
    16  func natsJetStreamOutputConfig() *service.ConfigSpec {
    17  	return service.NewConfigSpec().
    18  		// Stable(). TODO
    19  		Categories("Services").
    20  		Version("3.46.0").
    21  		Summary("Write messages to a NATS JetStream subject.").
    22  		Description(auth.Description()).
    23  		Field(service.NewStringListField("urls").
    24  			Description("A list of URLs to connect to. If an item of the list contains commas it will be expanded into multiple URLs.").
    25  			Example([]string{"nats://127.0.0.1:4222"}).
    26  			Example([]string{"nats://username:password@127.0.0.1:4222"})).
    27  		Field(service.NewInterpolatedStringField("subject").
    28  			Description("A subject to write to.").
    29  			Example("foo.bar.baz").
    30  			Example(`${! meta("kafka_topic") }`).
    31  			Example(`foo.${! json("meta.type") }`)).
    32  		Field(service.NewIntField("max_in_flight").
    33  			Description("The maximum number of messages to have in flight at a given time. Increase this to improve throughput.").
    34  			Default(1024)).
    35  		Field(service.NewTLSToggledField("tls")).
    36  		Field(service.NewInternalField(auth.FieldSpec()))
    37  }
    38  
    39  func init() {
    40  	err := service.RegisterOutput(
    41  		output.TypeNATSJetStream, natsJetStreamOutputConfig(),
    42  		func(conf *service.ParsedConfig, mgr *service.Resources) (service.Output, int, error) {
    43  			maxInFlight, err := conf.FieldInt("max_in_flight")
    44  			if err != nil {
    45  				return nil, 0, err
    46  			}
    47  			w, err := newJetStreamWriterFromConfig(conf, mgr.Logger())
    48  			return w, maxInFlight, err
    49  		})
    50  
    51  	if err != nil {
    52  		panic(err)
    53  	}
    54  }
    55  
    56  //------------------------------------------------------------------------------
    57  
    58  type jetStreamOutput struct {
    59  	urls       string
    60  	conf       output.NATSJetStreamConfig
    61  	subjectStr *service.InterpolatedString
    62  	authConf   auth.Config
    63  	tlsConf    *tls.Config
    64  
    65  	log *service.Logger
    66  
    67  	connMut  sync.Mutex
    68  	natsConn *nats.Conn
    69  	jCtx     nats.JetStreamContext
    70  
    71  	shutSig *shutdown.Signaller
    72  }
    73  
    74  func newJetStreamWriterFromConfig(conf *service.ParsedConfig, log *service.Logger) (*jetStreamOutput, error) {
    75  	j := jetStreamOutput{
    76  		log:     log,
    77  		shutSig: shutdown.NewSignaller(),
    78  	}
    79  
    80  	urlList, err := conf.FieldStringList("urls")
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  	j.urls = strings.Join(urlList, ",")
    85  
    86  	if j.subjectStr, err = conf.FieldInterpolatedString("subject"); err != nil {
    87  		return nil, err
    88  	}
    89  
    90  	tlsConf, tlsEnabled, err := conf.FieldTLSToggled("tls")
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  	if tlsEnabled {
    95  		j.tlsConf = tlsConf
    96  	}
    97  
    98  	if j.authConf, err = AuthFromParsedConfig(conf.Namespace("auth")); err != nil {
    99  		return nil, err
   100  	}
   101  	return &j, nil
   102  }
   103  
   104  //------------------------------------------------------------------------------
   105  
   106  func (j *jetStreamOutput) Connect(ctx context.Context) error {
   107  	j.connMut.Lock()
   108  	defer j.connMut.Unlock()
   109  
   110  	if j.natsConn != nil {
   111  		return nil
   112  	}
   113  
   114  	var natsConn *nats.Conn
   115  	var jCtx nats.JetStreamContext
   116  	var err error
   117  
   118  	defer func() {
   119  		if err != nil && natsConn != nil {
   120  			natsConn.Close()
   121  		}
   122  	}()
   123  
   124  	var opts []nats.Option
   125  	if j.tlsConf != nil {
   126  		opts = append(opts, nats.Secure(j.tlsConf))
   127  	}
   128  	opts = append(opts, auth.GetOptions(j.authConf)...)
   129  	if natsConn, err = nats.Connect(j.urls, opts...); err != nil {
   130  		return err
   131  	}
   132  
   133  	if jCtx, err = natsConn.JetStream(); err != nil {
   134  		return err
   135  	}
   136  
   137  	j.log.Infof("Sending NATS messages to JetStream subject: %v", j.conf.Subject)
   138  
   139  	j.natsConn = natsConn
   140  	j.jCtx = jCtx
   141  	return nil
   142  }
   143  
   144  func (j *jetStreamOutput) disconnect() {
   145  	j.connMut.Lock()
   146  	defer j.connMut.Unlock()
   147  
   148  	if j.natsConn != nil {
   149  		j.natsConn.Close()
   150  		j.natsConn = nil
   151  	}
   152  	j.jCtx = nil
   153  }
   154  
   155  //------------------------------------------------------------------------------
   156  
   157  func (j *jetStreamOutput) Write(ctx context.Context, msg *service.Message) error {
   158  	j.connMut.Lock()
   159  	jCtx := j.jCtx
   160  	j.connMut.Unlock()
   161  	if jCtx == nil {
   162  		return service.ErrNotConnected
   163  	}
   164  
   165  	subject := j.subjectStr.String(msg)
   166  	msgBytes, err := msg.AsBytes()
   167  	if err != nil {
   168  		return err
   169  	}
   170  
   171  	_, err = jCtx.Publish(subject, msgBytes)
   172  	return err
   173  }
   174  
   175  func (j *jetStreamOutput) Close(ctx context.Context) error {
   176  	go func() {
   177  		j.disconnect()
   178  		j.shutSig.ShutdownComplete()
   179  	}()
   180  	select {
   181  	case <-j.shutSig.HasClosedChan():
   182  	case <-ctx.Done():
   183  		return ctx.Err()
   184  	}
   185  	return nil
   186  }