github.com/argoproj/argo-events@v1.9.1/eventbus/jetstream/base/jetstream.go (about)

     1  package base
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/tls"
     6  	"fmt"
     7  
     8  	"github.com/argoproj/argo-events/common"
     9  	eventbuscommon "github.com/argoproj/argo-events/eventbus/common"
    10  	eventbusv1alpha1 "github.com/argoproj/argo-events/pkg/apis/eventbus/v1alpha1"
    11  	nats "github.com/nats-io/nats.go"
    12  	"github.com/spf13/viper"
    13  	"go.uber.org/zap"
    14  )
    15  
    16  type Jetstream struct {
    17  	url  string
    18  	auth *eventbuscommon.Auth
    19  
    20  	MgmtConnection JetstreamConnection
    21  
    22  	streamSettings string
    23  
    24  	Logger *zap.SugaredLogger
    25  }
    26  
    27  func NewJetstream(url string, streamSettings string, auth *eventbuscommon.Auth, logger *zap.SugaredLogger) (*Jetstream, error) {
    28  	js := &Jetstream{
    29  		url:            url,
    30  		auth:           auth,
    31  		Logger:         logger,
    32  		streamSettings: streamSettings,
    33  	}
    34  
    35  	return js, nil
    36  }
    37  
    38  func (stream *Jetstream) Init() error {
    39  	mgmtConnection, err := stream.MakeConnection()
    40  	if err != nil {
    41  		errStr := fmt.Sprintf("error creating Management Connection for Jetstream stream %+v: %v", stream, err)
    42  		stream.Logger.Error(errStr)
    43  		return fmt.Errorf(errStr)
    44  	}
    45  	err = stream.CreateStream(mgmtConnection)
    46  	if err != nil {
    47  		stream.Logger.Errorw("Failed to create Stream", zap.Error(err))
    48  		return err
    49  	}
    50  	stream.MgmtConnection = *mgmtConnection
    51  
    52  	return nil
    53  }
    54  
    55  func (stream *Jetstream) MakeConnection() (*JetstreamConnection, error) {
    56  	log := stream.Logger
    57  	conn := &JetstreamConnection{Logger: stream.Logger}
    58  
    59  	opts := []nats.Option{
    60  		// todo: try out Jetstream's auto-reconnection capability
    61  		nats.NoReconnect(),
    62  		nats.DisconnectErrHandler(func(nc *nats.Conn, err error) {
    63  			conn.NATSConnected = false
    64  			log.Errorw("NATS connection lost", zap.Error(err))
    65  		}),
    66  		nats.ReconnectHandler(func(nnc *nats.Conn) {
    67  			conn.NATSConnected = true
    68  			log.Info("Reconnected to NATS server")
    69  		}),
    70  		nats.Secure(&tls.Config{
    71  			InsecureSkipVerify: true,
    72  		}),
    73  	}
    74  
    75  	switch stream.auth.Strategy {
    76  	case eventbusv1alpha1.AuthStrategyToken:
    77  		log.Info("NATS auth strategy: Token")
    78  		opts = append(opts, nats.Token(stream.auth.Credential.Token))
    79  	case eventbusv1alpha1.AuthStrategyBasic:
    80  		log.Info("NATS auth strategy: Basic")
    81  		opts = append(opts, nats.UserInfo(stream.auth.Credential.Username, stream.auth.Credential.Password))
    82  	case eventbusv1alpha1.AuthStrategyNone:
    83  		log.Info("NATS auth strategy: None")
    84  	default:
    85  		return nil, fmt.Errorf("unsupported auth strategy")
    86  	}
    87  	nc, err := nats.Connect(stream.url, opts...)
    88  	if err != nil {
    89  		log.Errorw("Failed to connect to NATS server", zap.Error(err))
    90  		return nil, err
    91  	}
    92  	conn.NATSConn = nc
    93  	conn.NATSConnected = true
    94  
    95  	// Create JetStream Context
    96  	conn.JSContext, err = nc.JetStream()
    97  	if err != nil {
    98  		log.Errorw("Failed to get Jetstream context", zap.Error(err))
    99  		return nil, err
   100  	}
   101  
   102  	log.Info("Connected to NATS Jetstream server.")
   103  	return conn, nil
   104  }
   105  
   106  func (stream *Jetstream) CreateStream(conn *JetstreamConnection) error {
   107  	if conn == nil {
   108  		return fmt.Errorf("Can't create Stream on nil connection")
   109  	}
   110  	var err error
   111  
   112  	// before we add the Stream first let's check to make sure it doesn't already exist
   113  	streamInfo, err := conn.JSContext.StreamInfo(common.JetStreamStreamName)
   114  	if streamInfo != nil && err == nil {
   115  		stream.Logger.Infof("No need to create Stream '%s' as it already exists", common.JetStreamStreamName)
   116  		return nil
   117  	}
   118  	if err != nil && err != nats.ErrStreamNotFound {
   119  		stream.Logger.Warnf(`Error calling StreamInfo for Stream '%s' (this can happen if another Jetstream client "
   120  		is trying to create the Stream at the same time): %v`, common.JetStreamStreamName, err)
   121  	}
   122  
   123  	// unmarshal settings
   124  	v := viper.New()
   125  	v.SetConfigType("yaml")
   126  	if err := v.ReadConfig(bytes.NewBufferString(stream.streamSettings)); err != nil {
   127  		return err
   128  	}
   129  
   130  	streamConfig := nats.StreamConfig{
   131  		Name:       common.JetStreamStreamName,
   132  		Subjects:   []string{common.JetStreamStreamName + ".*.*"},
   133  		Retention:  nats.LimitsPolicy,
   134  		Discard:    nats.DiscardOld,
   135  		MaxMsgs:    v.GetInt64("maxMsgs"),
   136  		MaxAge:     v.GetDuration("maxAge"),
   137  		MaxBytes:   v.GetInt64("maxBytes"),
   138  		Storage:    nats.FileStorage,
   139  		Replicas:   v.GetInt("replicas"),
   140  		Duplicates: v.GetDuration("duplicates"),
   141  	}
   142  	stream.Logger.Infof("Will use this stream config:\n '%v'", streamConfig)
   143  
   144  	connectErr := common.DoWithRetry(nil, func() error { // exponential backoff if it fails the first time
   145  		_, err = conn.JSContext.AddStream(&streamConfig)
   146  		if err != nil {
   147  			errStr := fmt.Sprintf(`Failed to add Jetstream stream '%s'for connection %+v: err=%v`,
   148  				common.JetStreamStreamName, conn, err)
   149  			return fmt.Errorf(errStr)
   150  		} else {
   151  			return nil
   152  		}
   153  	})
   154  	if connectErr != nil {
   155  		return connectErr
   156  	}
   157  
   158  	stream.Logger.Infof("Created Jetstream stream '%s' for connection %+v", common.JetStreamStreamName, conn)
   159  	return nil
   160  }