Skip to content

Instantly share code, notes, and snippets.

@wallyqs wallyqs/custom-connect.go
Last active Sep 5, 2018

Embed
What would you like to do?
Custom Dialer based connect
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")
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.