Custom Dialer based connect

custom-connect.go

package main

import (
	"context"
	"log"
	"net"
	"net/http/httptrace"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/nats-io/go-nats"
)

type customDialer struct {
	ctx             context.Context
	nc              *nats.Conn
	connectTimeout  time.Duration
	connectTimeWait time.Duration
}

func (cd *customDialer) Dial(network, address string) (net.Conn, error) {
	d := &net.Dialer{}
	trace := &httptrace.ClientTrace{
		ConnectStart: func(network, addr string) {
			log.Println("Attempting to connect to server at", addr)
		},
		ConnectDone: func(network, addr string, err error) {
			if err == nil {
				log.Println("Established TCP connection")
			} else {
				log.Println("Failed connecting to", addr)
			}
		},
	}
	ctx := httptrace.WithClientTrace(cd.ctx, trace)
	ctx, cancel := context.WithTimeout(ctx, cd.connectTimeout)
	defer cancel()
	for {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			if conn, err := d.DialContext(ctx, network, address); err == nil {
				log.Println("Connected to NATS successfully")
				return conn, nil
			} else {
				log.Println(err, ctx.Err())
				time.Sleep(cd.connectTimeWait)
			}
		}
	}
}

func main() {
	// Parent context cancels connecting/reconnecting altogether.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var err error
	var nc *nats.Conn
	cd := &customDialer{
		ctx:             ctx,
		nc:              nc,
		connectTimeout:  5 * time.Second,
		connectTimeWait: 1 * time.Second,
	}
	opts := []nats.Option{
		nats.SetCustomDialer(cd),
		nats.ReconnectWait(2 * time.Second),
		nats.ReconnectHandler(func(c *nats.Conn) {
			log.Println("Reconnected to", c.ConnectedUrl())
		}),
		nats.DisconnectHandler(func(c *nats.Conn) {
			log.Println("Disconnected from NATS")
		}),
		nats.ClosedHandler(func(c *nats.Conn) {
			log.Println("NATS connection is closed.")
		}),
		nats.MaxReconnects(2),
	}
	go func() {
		nc, err = nats.Connect("demo.nats.io", opts...)
	}()

	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT)
	go func() {
		for {
			select {
			case sig := <-c:
				switch sig {
				case syscall.SIGINT:
					log.Println("Shutting down...")
					cancel()
					return
				}
			}
		}
	}()

WaitForEstablishedConnection:
	for {
		if err != nil {
			log.Fatal(err)
		}
		select {
		case <-ctx.Done():
			break WaitForEstablishedConnection
		default:
		}
		if nc == nil || !nc.IsConnected() {
			log.Println("Connection not ready yet...")
			time.Sleep(500 * time.Millisecond)
			continue
		}
		break WaitForEstablishedConnection
	}
	if ctx.Err() != nil {
		log.Fatal(ctx.Err())
	}

	// Run until receiving signal
	for range time.NewTicker(1 * time.Second).C {
		select {
		case <-ctx.Done():
			// Disconnect and flush pending messages
			if err := nc.Drain(); err != nil {
				log.Println(err)
			}
			log.Println("Closed connection to NATS")
			return
		default:
		}

		if nc.IsClosed() {
			break
		}
		if err := nc.Publish("hello", []byte("world")); err != nil {
			log.Println(err)
			continue
		}
		log.Println("Published message")
	}
}