github.com/kubeshop/testkube@v1.17.23/pkg/event/bus/nats.go (about)

     1  package bus
     2  
     3  import (
     4  	"crypto/tls"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/nats-io/nats.go"
    10  
    11  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
    12  	"github.com/kubeshop/testkube/pkg/event/kind/common"
    13  	"github.com/kubeshop/testkube/pkg/log"
    14  )
    15  
    16  var (
    17  	_ Bus = (*NATSBus)(nil)
    18  )
    19  
    20  const (
    21  	SubscribeBuffer        = 1
    22  	SubscriptionName       = "events"
    23  	InternalPublishTopic   = "internal.all"
    24  	InternalSubscribeTopic = "internal.>"
    25  )
    26  
    27  type ConnectionConfig struct {
    28  	NatsURI            string
    29  	NatsSecure         bool
    30  	NatsSkipVerify     bool
    31  	NatsCertFile       string
    32  	NatsKeyFile        string
    33  	NatsCAFile         string
    34  	NatsConnectTimeout time.Duration
    35  }
    36  
    37  func optsFromConfig(cfg ConnectionConfig) (opts []nats.Option) {
    38  	opts = []nats.Option{}
    39  	if cfg.NatsSecure {
    40  		if cfg.NatsSkipVerify {
    41  			opts = append(opts, nats.Secure(&tls.Config{InsecureSkipVerify: true}))
    42  		} else {
    43  			opts = append(opts, nats.ClientCert(cfg.NatsCertFile, cfg.NatsKeyFile))
    44  			if cfg.NatsCAFile != "" {
    45  				opts = append(opts, nats.RootCAs(cfg.NatsCAFile))
    46  			}
    47  		}
    48  	}
    49  
    50  	if cfg.NatsConnectTimeout > 0 {
    51  		opts = append(opts, nats.Timeout(cfg.NatsConnectTimeout))
    52  	}
    53  
    54  	return opts
    55  }
    56  
    57  func NewNATSEncodedConnection(cfg ConnectionConfig, opts ...nats.Option) (*nats.EncodedConn, error) {
    58  	opts = append(opts, optsFromConfig(cfg)...)
    59  
    60  	nc, err := NewNATSConnection(cfg, opts...)
    61  	if err != nil {
    62  		log.DefaultLogger.Fatalw("error connecting to nats", "error", err)
    63  		return nil, err
    64  	}
    65  
    66  	// automatic NATS JSON CODEC
    67  	ec, err := nats.NewEncodedConn(nc, nats.JSON_ENCODER)
    68  	if err != nil {
    69  		log.DefaultLogger.Fatalw("error connecting to nats", "error", err)
    70  		return nil, err
    71  	}
    72  
    73  	if err != nil {
    74  		log.DefaultLogger.Errorw("error creating NATS connection", "error", err)
    75  	}
    76  
    77  	return ec, nil
    78  }
    79  
    80  func NewNATSConnection(cfg ConnectionConfig, opts ...nats.Option) (*nats.Conn, error) {
    81  	opts = append(opts, optsFromConfig(cfg)...)
    82  
    83  	nc, err := nats.Connect(cfg.NatsURI, opts...)
    84  	if err != nil {
    85  		log.DefaultLogger.Fatalw("error connecting to nats", "error", err)
    86  		return nil, err
    87  	}
    88  
    89  	return nc, nil
    90  }
    91  
    92  func NewNATSBus(nc *nats.EncodedConn) *NATSBus {
    93  	return &NATSBus{
    94  		nc: nc,
    95  	}
    96  }
    97  
    98  type NATSBus struct {
    99  	nc            *nats.EncodedConn
   100  	subscriptions sync.Map
   101  }
   102  
   103  // Publish publishes event to NATS on events topic
   104  func (n *NATSBus) Publish(event testkube.Event) error {
   105  	return n.PublishTopic(SubscriptionName, event)
   106  }
   107  
   108  // Subscribe subscribes to NATS events topic
   109  func (n *NATSBus) Subscribe(queueName string, handler Handler) error {
   110  	return n.SubscribeTopic(SubscriptionName, queueName, handler)
   111  }
   112  
   113  // PublishTopic publishes event to NATS on given topic
   114  func (n *NATSBus) PublishTopic(topic string, event testkube.Event) error {
   115  	return n.nc.Publish(topic, event)
   116  }
   117  
   118  // SubscribeTopic subscribes to NATS topic
   119  func (n *NATSBus) SubscribeTopic(topic, queueName string, handler Handler) error {
   120  	// sanitize names for NATS
   121  	queue := common.ListenerName(queueName)
   122  
   123  	// async subscribe on queue
   124  	s, err := n.nc.QueueSubscribe(topic, queue, handler)
   125  
   126  	if err == nil {
   127  		// store subscription for later unsubscribe
   128  		key := n.queueName(SubscriptionName, queue)
   129  		n.subscriptions.Store(key, s)
   130  	}
   131  
   132  	return err
   133  }
   134  
   135  func (n *NATSBus) Unsubscribe(queueName string) error {
   136  	// sanitize names for NATS
   137  	queue := common.ListenerName(queueName)
   138  
   139  	key := n.queueName(SubscriptionName, queue)
   140  	if s, ok := n.subscriptions.Load(key); ok {
   141  		return s.(*nats.Subscription).Drain()
   142  	}
   143  	return nil
   144  }
   145  
   146  func (n *NATSBus) Close() error {
   147  	n.nc.Close()
   148  	return nil
   149  }
   150  
   151  func (n *NATSBus) queueName(subscription, queue string) string {
   152  	return fmt.Sprintf("%s.%s", SubscriptionName, queue)
   153  }