Forem: Waldemar Quevedo

Introducing go_test.mod

Waldemar Quevedo — Fri, 24 Sep 2021 21:12:52 +0000

How to avoid meta test dependencies across Go modules

Since the Go v1.14 release, the go command now includes a little known flag called -modfile that can be used to manage multiple set of dependencies within the same repository.

The -modfile flag can also be really helpful to also manage better what are the dependencies that importers of your package end up bringing into their project, specially when dealing with test dependencies.

The problem

For example, let's say that we have the following couple of repos that do not have any dependencies:

# github.com/wallyqs/go-mod-a
package hello

func Hello() string {
    return "Hello"
}

# github.com/wallyqs/go-mod-b
package world

func World() string {
    return "World"
}

Though for testing purposes, they do depend on each other. We add hello_test.go first in this case in the go-mod-a repo:

package hello

import (
    "testing"
    "github.com/wallyqs/go-mod-b"
)

func TestHello(t *testing.T) {
    got := Hello()
    expected := "Hello"
    if got != expected {
        t.Fatalf("Expected %v, got: %v", expected, got)
    }
}

func TestHelloWorld(t *testing.T) {
    got := fmt.Sprintf("%s %s", Hello(), world.World())
    expected := "Hello World"
    if got != expected {
        t.Fatalf("Expected %v, got: %v", expected, got)
    }
}

The result of this will be something like the following in the
go.mod and go.sum in the go-mod-a repo:

$ cat go.mod 
module github.com/wallyqs/go-mod-a

go 1.17

require github.com/wallyqs/go-mod-b v0.0.0-20210924195232-cc93d43aa245 // indirect

$ cat go.sum 
github.com/wallyqs/go-mod-b v0.0.0-20210924195232-cc93d43aa245 h1:vD21YUG9esVWngkuysL7tmR4l9EdwvaghEfdGNcBDVw=
github.com/wallyqs/go-mod-b v0.0.0-20210924195232-cc93d43aa245/go.mod h1:IMoYMTpn8BaoMzQP+9DeRp7bHQv/n+jAhQfguxXc9oM=

At this point, things look like this in the go-mod-a repo:

tree .
.
├── LICENSE
├── README.md
├── go.mod
├── go.sum
├── hello.go
└── hello_test.go

Now, if we add the same test to the go-mod-b, we will have a circular testing dependency which Go modules is actually fine with:

package world

import (
    "fmt"
    "testing"
    "github.com/wallyqs/go-mod-a"
)

func TestWorld(t *testing.T) {
    got := World()
    expected := "World"
    if got != expected {
        t.Fatalf("Expected %v, got: %v", expected, got)
    }
}

func TestHelloWorld(t *testing.T) {
    got := fmt.Sprintf("%s %s", hello.Hello(), World())
    expected := "Hello World"
    if got != expected {
        t.Fatalf("Expected %v, got: %v", expected, got)
    }
}

$ ~/go/src/github.com/wallyqs/go-mod-b (main) $ go get github.com/wallyqs/go-mod-a
go: downloading github.com/wallyqs/go-mod-a v0.0.0-20210924201556-e67af2677460
go get: added github.com/wallyqs/go-mod-a v0.0.0-20210924201556-e67af2677460

$ ~/go/src/github.com/wallyqs/go-mod-b (main) $ go test ./... -v
=== RUN   TestWorld
--- PASS: TestWorld (0.00s)
=== RUN   TestHelloWorld
--- PASS: TestHelloWorld (0.00s)
PASS
ok      github.com/wallyqs/go-mod-b 0.014s

In the go-mod-b repo, now we will have something like this though:

$ tree .
.
├── LICENSE
├── README.md
├── go.mod
├── go.sum
├── world.go
└── world_test.go

$ cat go.sum 
github.com/wallyqs/go-mod-a v0.0.0-20210924201556-e67af2677460 h1:alim5qw73L9EfbCEF7TW3KDlbfWoMs580PtS4+1fF38=
github.com/wallyqs/go-mod-a v0.0.0-20210924201556-e67af2677460/go.mod h1:+86UH/9vUOVQadGIBaupYvn1FMXQIzXpXqmGtq28JPk=
github.com/wallyqs/go-mod-b v0.0.0-20210924195232-cc93d43aa245/go.mod h1:IMoYMTpn8BaoMzQP+9DeRp7bHQv/n+jAhQfguxXc9oM=

Notice how go-mod-b includes a meta test dependency of itself, if left as is now all importers of go-mod-b will be actually importing two different dependencies of the same module which is not really necessary.

Let's consider that this is now v0.1.0 version of both packages, so we then update our original test dependencies to be a tagged version:

$ go get github.com/wallyqs/go-mod-b@v0.1.0
go: downloading github.com/wallyqs/go-mod-b v0.1.0
go get: upgraded github.com/wallyqs/go-mod-b v0.0.0-20210924202312-2deb2c1c1be4 => v0.1.0

$ ~/go/src/github.com/wallyqs/go-mod-a (main) $ cat go.mod
module github.com/wallyqs/go-mod-a

go 1.17

require github.com/wallyqs/go-mod-b v0.1.0

Now after go mod tidy on the repo, the library will still be
bringing the last untagged version of itself as well:

$ cd ~/go/src/githubsrc/github.com/wallyqs/go-mod-a
$ cat go.sum
github.com/wallyqs/go-mod-a v0.0.0-20210924201556-e67af2677460/go.mod h1:+86UH/9vUOVQadGIBaupYvn1FMXQIzXpXqmGtq28JPk=
github.com/wallyqs/go-mod-b v0.0.0-20210924195232-cc93d43aa245/go.mod h1:IMoYMTpn8BaoMzQP+9DeRp7bHQv/n+jAhQfguxXc9oM=
github.com/wallyqs/go-mod-b v0.1.0 h1:3Qz5nQ0B98wlrwiXv7dEW2PhPfRqPB+dTmHPI08Dr/8=
github.com/wallyqs/go-mod-b v0.1.0/go.mod h1:BtD/yaPmz5Vf2OMu9+VcCNvR9I60VpeZTKvBQsUD5wg=

How to solve this?

With the -modfile flag, we can better separate what dependencies are being imported if we change a bit how we use the testing dependencies. We will follow these steps then:

Make a test folder
Include the tests with test dependencies in the test folder. This means that the HelloWorld test is now moved to its own test file like:

// test/hello_test.go
package hello

import (
    "fmt"
    "testing"
    "github.com/wallyqs/go-mod-a"
    "github.com/wallyqs/go-mod-b"
)

func TestHelloWorld(t *testing.T) {
    got := fmt.Sprintf("%s %s", hello.Hello(), world.World())
    expected := "Hello World"
    if got != expected {
        t.Fatalf("Expected %v, got: %v", expected, got)
    }
}

$ tree ~/go/src/github.com/wallyqs/go-mod-a
├── LICENSE
├── README.md
├── go.mod
├── go_test.mod
├── go_test.sum
├── hello.go
├── hello_test.go
└── test
    └── hello_test.go

Create a go_test.mod file:

cp go.mod go_test.mod && go mod tidy -modfile go_test.mod

Then use go test -modfile go_test.mod to run the tests with the extra deps:

$ mkdir test
$ cp go.mod go_test.mod && go mod tidy -modfile go_test.mod

$ go test -modfile=go_test.mod ./... -v
=== RUN   TestHello
--- PASS: TestHello (0.00s)
PASS
ok      github.com/wallyqs/go-mod-a (cached)
=== RUN   TestHelloWorld
--- PASS: TestHelloWorld (0.00s)
PASS
ok      github.com/wallyqs/go-mod-a/test    (cached)

Next, to generate a minimal go.mod for your library without the testing dependencies, just test the package:

$ cd ~/go/src/github.com/wallyqs/go-mod-a
$ rm go.mod
$ go mod init
$ go test github.com/wallyqs/go-mod-a

To run the tests with the tests with extra dependencies you will need to use go_test.mod instead:

$ go test -modfile=go_test.mod -v github.com/wallyqs/go-mod-a/test
=== RUN   TestHelloWorld
--- PASS: TestHelloWorld (0.00s)
PASS

And if you want to test all you can do:

$ go test ./... -modfile=go_test.mod -v
=== RUN   TestHello
--- PASS: TestHello (0.00s)
PASS
ok      github.com/wallyqs/go-mod-a (cached)
=== RUN   TestHelloWorld
--- PASS: TestHelloWorld (0.00s)
PASS
ok      github.com/wallyqs/go-mod-a/test    (cached)

Look, no dependencies!

Since both packages didn't actually depend on each other other than for the tests this will mean that the go.mod is just empty:

cd ~/go/src/github.com/wallyqs/go-mod-a (main) $ cat go.mod
module github.com/wallyqs/go-mod-a

go 1.17

Bonus: This will also mean that it does not depend on previous untagged versions of itself for testing!

cd ~/go/src/github.com/wallyqs/go-mod-a (main) $ cat go_test.mod 
module github.com/wallyqs/go-mod-a

go 1.17

require github.com/wallyqs/go-mod-b v0.1.0

Upgrading a test dependency

Let's say now that we tag the next version of wallyqs/go-mod-a, we can then upgrade our test dependency in the wallyqs/go-mod-b as follows:

$ cd ~/go/src/github.com/wallyqs/go-mod-b (main) $ 
go get -modfile=go_test.mod github.com/wallyqs/go-mod-a@v0.2.0

The go_test.mod as a result is just the latest tagged version:

head go*
==> go.mod <==
module github.com/wallyqs/go-mod-b

go 1.17

==> go_test.mod <==
module github.com/wallyqs/go-mod-b

go 1.17

require github.com/wallyqs/go-mod-a v0.2.0

==> go_test.sum <==
github.com/wallyqs/go-mod-a v0.2.0 h1:6o2qaAI3ssJhSx89cHvvY449+V3h/1Z09Z1ySoUPNwI=
github.com/wallyqs/go-mod-a v0.2.0/go.mod h1:N2dbiJxS0vlXnYKNv9J5JE6ULgOhcHA81/VkivvCr08=

Summary

If possible, avoid including test dependencies to your library and instead move them to a subfolder of our package. Being able to do this is good for the ecosystem around your package since will mean less interactions with goproxy and simplified dependencies overall.

The example repos in this post can be found at:

We have been using this approach in the nats.go client for some time already where both the server and the client depend on each other for some of the tests. As a result, the go.mod of the client is very small and users do not require to download dependencies of the server:

https://github.com/nats-io/nats.go/blob/main/go.mod

module github.com/nats-io/nats.go

go 1.16

require (
    github.com/nats-io/nkeys v0.3.0
    github.com/nats-io/nuid v1.0.1
)

Hope this helps!

Wally

Wrapping callbacks dynamically with a Goroutine using reflect

Waldemar Quevedo — Thu, 04 Apr 2019 06:09:16 +0000

Recently a NATS user was looking for a generic way to be able to
increase the concurrency of the async NATS handler subscriptions which
by default process the messages sequentially in order. This means that
if a message takes very long to be processed, then the next message
will be processed with some delay even if the server has already
delivered it to the client:

nc, err := nats.Connect("demo.nats.io")
if err != nil {
        log.Fatalf("Failed to create a connection: %v\n", err)
}

nc.Subscribe("foo", func(m *nats.Msg) {
        log.Printf("[Received on '>'  ] %+v", string(msg.Data))
        time.Sleep(1 * time.Second) // Blocks each message for 1 second
})

An easy way to work around the head of line blocking, would be to
simply add a Goroutine:

nc.Subscribe("foo", func(m *nats.Msg) {
    log.Printf("[Received on '>'  ] %+v", string(msg.Data))
        go func(){
        // ... Processing ...
        time.Sleep(1 * time.Second) 
    }()
})

This works for regular NATS connections and subscriptions so it is
often the approach I suggest. There is one exception to this though,
which is when using NATS encoded connections, the type of the function
can be of a multiple number of types, with the NATS client internally
handling the serialization/deserialization of the payload:

package main

import (
    "log"
    "time"

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

func init() {
    log.SetFlags(log.LstdFlags | log.Lmicroseconds)
}

func main() {
    nc, err := nats.Connect("demo.nats.io")
    if err != nil {
        log.Fatalf("Failed to create a connection: %v\n", err)
    }
    ec, err := nats.NewEncodedConn(nc, nats.JSON_ENCODER)
    if err != nil {
        log.Fatalf("Failed to create an encoded connection: %v\n", err)
    }
    defer ec.Close()

    type Event struct {
        Message string `json:"message"`
    }

    ec.Subscribe(">", func(b []byte) {
        log.Printf("[Received on '>'  ] %+v", string(b))
        time.Sleep(1 * time.Second)
    })

    ec.Subscribe("foo", func(s string) {
        log.Printf("[Received on 'foo'] %+v", s)
        time.Sleep(1 * time.Second)
    })

    ec.Subscribe("bar", func(e *Event) {
        log.Printf("[Received on 'bar'] %+v", e)
        time.Sleep(1 * time.Second)
    })

    payload := []byte("Hello World!")
    for i := 0; i < 3; i++ {
        ec.Publish("foo", payload)
        ec.Publish("bar", &Event{Message: string(payload)})
    }
    select {
    case <-time.After(2 * time.Second):
    }

}

Now the simple workaround for the Goroutine does not work since would
have to implement it for a multiple number of function types. This is
because internally the NATS client uses reflect to find out how to
properly call the callback using the decoded payload.

In order to workaround this issue, I dug my way out
using more reflection to dynamically rewrap a function so that it is
ran in its own Goroutine:

type subscriber interface {
    Subscribe(subject string, cb nats.Handler) (*nats.Subscription, error)
    QueueSubscribe(subject, group string, cb nats.Handler) (*nats.Subscription, error)
}

func Subscribe(ec subscriber, subject string, cb nats.Handler) (*nats.Subscription, error) {
    parallelize := func(f, fptr interface{}) {
        // Original function.
        oV := reflect.ValueOf(f)

        // Pointer to the function that will be replaced.
        ffn := reflect.ValueOf(fptr).Interface()
        fn := reflect.ValueOf(ffn).Elem()

        // Type of original function used to create new one
        // that will replace the original.
        typ := reflect.TypeOf(f)

        // Rewrap original function.
        handler := func(in []reflect.Value) []reflect.Value {
            go oV.Call(in)
            return nil
        }
        v := reflect.MakeFunc(typ, handler)
        fn.Set(v)
    }
    parallelize(cb, &cb)
    return ec.Subscribe(subject, cb)
}

Now a user can use this helper function that has a similar API as the
original NATS client subscribe, though internally it is handling the
messages each in its own Goroutine:

Subscribe(ec, ">", func(b []byte) {
        log.Printf("[Received on '>'  ] %+v", string(b))
        time.Sleep(1 * time.Second)
})

Subscribe(ec, "foo", func(s string) {
        log.Printf("[Received on 'foo'] %+v", s)
        time.Sleep(1 * time.Second)
})

Subscribe(ec, "bar", func(e *Event) {
        log.Printf("[Received on 'bar'] %+v", e)
        time.Sleep(1 * time.Second)
})

Full example below:

package main

import (
    "log"
    "reflect"
    "time"

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

type subscriber interface {
    Subscribe(subject string, cb nats.Handler) (*nats.Subscription, error)
    QueueSubscribe(subject, group string, cb nats.Handler) (*nats.Subscription, error)
}

func Subscribe(ec subscriber, subject string, cb nats.Handler) (*nats.Subscription, error) {
    parallelize := func(f, fptr interface{}) {
        // Original function.
        oV := reflect.ValueOf(f)

        // Pointer to the function that will be replaced.
        ffn := reflect.ValueOf(fptr).Interface()
        fn := reflect.ValueOf(ffn).Elem()

        // Type of original function used to create new one
        // that will replace the original.
        typ := reflect.TypeOf(f)

        // Rewrap original function.
        handler := func(in []reflect.Value) []reflect.Value {
            go oV.Call(in)
            return nil
        }
        v := reflect.MakeFunc(typ, handler)
        fn.Set(v)
    }
    parallelize(cb, &cb)
    return ec.Subscribe(subject, cb)
}

func init() {
    log.SetFlags(log.LstdFlags | log.Lmicroseconds)
}

func main() {
    nc, err := nats.Connect("demo.nats.io")
    if err != nil {
        log.Fatalf("Failed to create a connection: %v\n", err)
    }
    ec, err := nats.NewEncodedConn(nc, nats.JSON_ENCODER)
    if err != nil {
        log.Fatalf("Failed to create an encoded connection: %v\n", err)
    }
    defer ec.Close()

    type Event struct {
        Message string `json:"message"`
    }

    Subscribe(ec, ">", func(b []byte) {
        log.Printf("[Received on '>'  ] %+v", string(b))
        time.Sleep(1 * time.Second)
    })

    Subscribe(ec, "foo", func(s string) {
        log.Printf("[Received on 'foo'] %+v", s)
        time.Sleep(1 * time.Second)
    })

    Subscribe(ec, "bar", func(e *Event) {
        log.Printf("[Received on 'bar'] %+v", e)
        time.Sleep(1 * time.Second)
    })

    testBytes := []byte("Hello World!")
    ec.Publish("foo", testBytes)
    ec.Publish("foo", testBytes)
    ec.Publish("foo", testBytes)
    ec.Publish("bar", &Event{Message: string(testBytes)})
    ec.Publish("bar", &Event{Message: string(testBytes)})
    ec.Publish("bar", &Event{Message: string(testBytes)})
    ec.PublishRequest("bar", "inbox", &Event{Message: string(testBytes)})
    select {
    case <-time.After(2 * time.Second):
    }

}

Using go modules in K8S projects

Waldemar Quevedo — Tue, 02 Apr 2019 04:09:01 +0000

Starting to adopt go mod more and more in the projects I maintain instead of other tools, but for some reason I keep hitting into the following issue: https://github.com/kubernetes/client-go/issues/551

Some solutions are written into that issue and more notes added to client-go, but so far what seems to work for me is the following ¯_(ツ)_/¯

module github.com/wallyqs/my-k8s-module

go 1.12

require (
    k8s.io/api kubernetes-1.13.4
    k8s.io/apimachinery kubernetes-1.13.4
    k8s.io/client-go v10.0.0
    k8s.io/kubernetes v1.13.4
)

Go’s map iteration order is not that random?

Waldemar Quevedo — Fri, 08 Mar 2019 21:09:23 +0000

I have been using Go for around 5 years already so by now had already internalized the behavior for range on maps to be random. This observable behavior of how Go's for range works on maps, naturally lead to code as follows to look correct to me:

package main

import (
    "fmt"
)

func pickRandom(entries map[string]int) string {
    var key string
    for k, _ := range entries {
        // Chose the random key
        key = k

        // Mark that it has been used
        entries[k]++
        break
    }
    return key
}

func main() {
    entries := make(map[string]int)
    entries["A"] = 0
    entries["B"] = 0
    entries["C"] = 0

    for i := 0; i < 10000; i++ {
        pickRandom(entries)
    }

    for k, v := range entries {
        fmt.Printf("key=%s hit=%v\n", k, v)
    }
}

So I was very surprised once I saw that the following results being repeated constantly:

key=A hit=7511
key=B hit=1279
key=C hit=1210

Looks like it doesn't matter how many times this is run, the first entry of the map will always get at least 7000 of the hits 😱

One way to fix this would be to instead create an extra array with they keys then use rand to get a random one, for example using math/rand:

package main

import (
    "fmt"
    "math/rand"
)

func pickRandom(entries map[string]int) string {
    var key string
    choices := []string{}
    for k, _ := range entries {
        choices = append(choices, k)
    }

    // Choose random one and mark that it was hit
    key = choices[rand.Intn(len(choices))]
    entries[key]++

    return key
}

func main() {
    entries := make(map[string]int)
    entries["A"] = 0
    entries["B"] = 0
    entries["C"] = 0

    for i := 0; i < 10000; i++ {
        pickRandom(entries)
    }

    for k, v := range entries {
        fmt.Printf("key=%s hit=%v\n", k, v)
    }
}

Now the results look much better!

key=A hit=3254
key=B hit=3387
key=C hit=3359

I was pointed that this could be a good example of the Hyrum's Law:

With a sufficient number of users of an API,
it does not matter what you promise in the contract:
all observable behaviors of your system
will be depended on by somebody.

We ended up accidentally depending on the for range behavior in the Go NATS client, where it was causing an imbalance with one server often having more connections than others in the server pool (issue has since then been fixed).

Using Custom Dialers in the NATS Go Client

Waldemar Quevedo — Wed, 05 Sep 2018 16:53:33 +0000

The Go client from NATS by default has support for reconnection logic, so that whenever it disconnects from a server it will retry for a number of times before giving up and closing the connection.

nc, err := nats.Connect("demo.nats.io", nats.MaxReconnects(10))
// ...

Besides limiting the number of maximum number of reconnect attempts,
the client has a number of knobs to customize how it should reconnect,
such as backing off in between connection attempts, or whether to
disable reconnecting altogether without buffering
enabled for example.

Furthermore, it also supports event callbacks to be able to trace
whenever the client disconnects, reconnects or gives up and closes the
connection to the server.

nc, err := nats.Connect("demo.nats.io",
        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("Closed connection to NATS")
        }),
)
// ...

On the other hand, there aren't as many knobs for the initial
connection attempt done by the client. Also, if the initial connect to
a NATS server fails, then nats.Connect will just return an error and
abort.

    nc, err = nats.Connect("127.0.0.1:4223")
    if err != nil {
        // nats: no servers available for connection
        log.Fatal(err)
    }

Instead of exposing a number of options to manage how to Dial to NATS
as part of the client, it is possible to define a CustomDialer that
can be used to replace the default dialer in order customize how to
dial without having to do further internal changes to the library.

type customDialer struct {
}

func (cd *customDialer) Dial(network, address string) (net.Conn, error) {
    // ...
    return conn, nil
}

func main() {
    cd := &customDialer{}
    opts := []nats.Option{
        nats.SetCustomDialer(cd),
        // ...
    }
    nc, err = nats.Connect("demo.nats.io", opts...)
    // ...
}

This usage of Go interfaces makes things very flexible as now the
client could add custom logic as part of the dialer implemention and
everything will work transparently.

For example, let's say that you want to make the client use the
context package to use DialContext and be able to cancel
connecting to NATS altogether with a deadline, one could then
define a Dialer implementation as follows:

package main

import (
    "context"
    "log"
    "net"
    "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) {
    ctx, cancel := context.WithTimeout(cd.ctx, cd.connectTimeout)
    defer cancel()

    for {
        log.Println("Attempting to connect to", address)
        if ctx.Err() != nil {
            return nil, ctx.Err()
        }

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

With the dialer implementation above, the NATS client will retry a
number of times to connect to the NATS server until the context is no
longer valid:

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("127.0.0.1:4222", 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")
    }
}

Using DialContext as in the example above, also opens the door to be
able to hook to some of the tracing features part of the standard
library, for example to check when a connection attempt will happen or
TCP connection is established.

In the example below, the httptrace.ClientTrace package is being used
in order to activate some of the built-in tracing instrumentation that
is part of the standard library via context and DialContext.

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("Connecting to server at", addr)
        },
        ConnectDone: func(network, addr string, err error) {
            if err == nil {
                log.Println("Established TCP connection")
            }
        },
    }
    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 {
                time.Sleep(cd.connectTimeWait)
            }
        }
    }
}

Result:

go run custom-connect.go
2018/09/05 09:02:17 Connection not ready yet...
2018/09/05 09:02:17 Attempting to connect to server at 127.0.0.1:4222
2018/09/05 09:02:17 Established TCP connection
2018/09/05 09:02:17 Connected to NATS successfully
2018/09/05 09:02:19 Published message
2018/09/05 09:02:20 Published message
2018/09/05 09:02:21 Published message
2018/09/05 09:02:22 Published message
2018/09/05 09:02:23 Published message
^C2018/09/05 09:02:23 Shutting down...

You can find the full example here. Interested in learning more? Join the NATS slack channel, or you can also find more examples of advanced NATS usage on my book Practical NATS.