github.com/mysteriumnetwork/node@v0.0.0-20240516044423-365054f76801/communication/nats/connection_wrap.go (about)

     1  /*
     2   * Copyright (C) 2019 The "MysteriumNetwork/node" Authors.
     3   *
     4   * This program is free software: you can redistribute it and/or modify
     5   * it under the terms of the GNU General Public License as published by
     6   * the Free Software Foundation, either version 3 of the License, or
     7   * (at your option) any later version.
     8   *
     9   * This program is distributed in the hope that it will be useful,
    10   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    11   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    12   * GNU General Public License for more details.
    13   *
    14   * You should have received a copy of the GNU General Public License
    15   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    16   */
    17  
    18  package nats
    19  
    20  import (
    21  	"context"
    22  	"fmt"
    23  	"net"
    24  	"net/url"
    25  	"strings"
    26  	"time"
    27  
    28  	nats_lib "github.com/nats-io/nats.go"
    29  	"github.com/pkg/errors"
    30  	"github.com/rs/zerolog/log"
    31  
    32  	"github.com/mysteriumnetwork/node/requests"
    33  )
    34  
    35  const (
    36  	// DefaultBrokerScheme broker scheme.
    37  	DefaultBrokerScheme = "nats"
    38  	// DefaultBrokerPort broker port.
    39  	DefaultBrokerPort = 4222
    40  )
    41  
    42  // ParseServerURL validates given NATS server address.
    43  func ParseServerURL(serverURI string) (*url.URL, error) {
    44  	// Add scheme first otherwise serverURL.Parse() fails.
    45  	if !strings.HasPrefix(serverURI, DefaultBrokerScheme) {
    46  		serverURI = fmt.Sprintf("%s://%s", DefaultBrokerScheme, serverURI)
    47  	}
    48  
    49  	serverURL, err := url.Parse(serverURI)
    50  	if err != nil {
    51  		return nil, errors.Wrapf(err, `failed to parse NATS server URI "%s"`, serverURI)
    52  	}
    53  
    54  	if serverURL.Port() == "" {
    55  		serverURL.Host = fmt.Sprintf("%s:%d", serverURL.Host, DefaultBrokerPort)
    56  	}
    57  
    58  	return serverURL, nil
    59  }
    60  
    61  // ParseServerURIs validates given list of NATS server addresses.
    62  func ParseServerURIs(serverURIs []string) ([]*url.URL, error) {
    63  	serverURLs := make([]*url.URL, len(serverURIs))
    64  
    65  	for i, server := range serverURIs {
    66  		natsURL, err := ParseServerURL(server)
    67  		if err != nil {
    68  			return nil, err
    69  		}
    70  
    71  		serverURLs[i] = natsURL
    72  	}
    73  
    74  	return serverURLs, nil
    75  }
    76  
    77  func newConnection(dialer requests.DialContext, serverURIs ...string) (*ConnectionWrap, error) {
    78  	return &ConnectionWrap{
    79  		servers: serverURIs,
    80  		onClose: func() {},
    81  		dialer:  dialer,
    82  	}, nil
    83  }
    84  
    85  // ConnectionWrap defines wrapped connection to NATS server(s).
    86  type ConnectionWrap struct {
    87  	*nats_lib.Conn
    88  
    89  	dialer requests.DialContext
    90  
    91  	servers []string
    92  	onClose func()
    93  }
    94  
    95  func (c *ConnectionWrap) connectOptions() nats_lib.Options {
    96  	options := nats_lib.GetDefaultOptions()
    97  	options.Servers = c.servers
    98  	options.MaxReconnect = -1
    99  	options.ReconnectWait = 1 * time.Second
   100  	options.PingInterval = 10 * time.Second
   101  	options.Timeout = 10 * time.Second
   102  	options.RetryOnFailedConnect = true
   103  
   104  	options.ClosedCB = func(conn *nats_lib.Conn) { log.Warn().Msg("NATS: connection closed") }
   105  	options.DisconnectedErrCB = func(nc *nats_lib.Conn, err error) { log.Warn().Err(err).Msg("NATS: disconnected") }
   106  	options.ReconnectedCB = func(nc *nats_lib.Conn) { log.Warn().Msg("NATS: reconnected") }
   107  
   108  	if c.dialer != nil {
   109  		options.CustomDialer = &dialer{c.dialer}
   110  	}
   111  
   112  	return options
   113  }
   114  
   115  // Open starts the connection: left for test compatibility.
   116  // Deprecated: Use nats.BrokerConnector#Connect() instead.
   117  func (c *ConnectionWrap) Open() (err error) {
   118  	c.Conn, err = c.connectOptions().Connect()
   119  	if err != nil {
   120  		log.Warn().Err(err).Msgf("Failed to connect to NATS servers %v, will reconnect again", c.connectOptions().Servers)
   121  	}
   122  
   123  	return nil
   124  }
   125  
   126  // Close destructs the connection.
   127  func (c *ConnectionWrap) Close() {
   128  	if c.Conn != nil {
   129  		c.Conn.Close()
   130  	}
   131  	c.onClose()
   132  }
   133  
   134  // Servers returns list of currently connected servers.
   135  func (c *ConnectionWrap) Servers() []string {
   136  	return c.servers
   137  }
   138  
   139  type dialer struct {
   140  	dialer requests.DialContext
   141  }
   142  
   143  func (d *dialer) Dial(network, address string) (net.Conn, error) {
   144  	ctx, cancel := context.WithTimeout(context.Background(), nats_lib.DefaultTimeout)
   145  	defer cancel()
   146  
   147  	return d.dialer(ctx, network, address)
   148  }