zmq4: make number of retries for dial configurable

options.go

@@ -51,6 +51,14 @@ func WithLogger(msg *log.Logger) Option {
 	}
 }
 
+// WithDialerMaxRetries configures the maximum number of retries
+// when dialing an endpoint (-1 means infinite retries).
+func WithDialerMaxRetries(maxRetries int) Option {
+	return func(s *socket) {
+		s.maxRetries = maxRetries
+	}
+}
+
 /*
 // TODO(sbinet)
 

socket.go

@@ -18,8 +18,9 @@ import (
 )
 
 const (
-	defaultRetry   = 250 * time.Millisecond
-	defaultTimeout = 5 * time.Minute
+	defaultRetry      = 250 * time.Millisecond
+	defaultTimeout    = 5 * time.Minute
+	defaultMaxRetries = 10
 )
 
 var (
@@ -30,13 +31,14 @@ var (
 
 // socket implements the ZeroMQ socket interface
 type socket struct {
-	ep        string // socket end-point
-	typ       SocketType
-	id        SocketIdentity
-	retry     time.Duration
-	sec       Security
-	log       *log.Logger
-	subTopics func() []string
+	ep         string // socket end-point
+	typ        SocketType
+	id         SocketIdentity
+	retry      time.Duration
+	maxRetries int
+	sec        Security
+	log        *log.Logger
+	subTopics  func() []string
 
 	mu    sync.RWMutex
 	ids   map[string]*Conn // ZMTP connection IDs
@@ -63,6 +65,7 @@ func newDefaultSocket(ctx context.Context, sockType SocketType) *socket {
 	return &socket{
 		typ:        sockType,
 		retry:      defaultRetry,
+		maxRetries: defaultMaxRetries,
 		sec:        nullSecurity{},
 		ids:        make(map[string]*Conn),
 		conns:      nil,
@@ -247,7 +250,8 @@ connect:
 	}
 
 	if err != nil {
-		if retries < 10 {
+		// retry if retry count is lower than maximum retry count and context has not been canceled
+		if (sck.maxRetries == -1 || retries < sck.maxRetries) && sck.ctx.Err() == nil {
 			retries++
 			time.Sleep(sck.retry)
 			goto connect

socket_test.go

@@ -6,6 +6,7 @@ package zmq4_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -14,6 +15,7 @@ import (
 	"time"
 
 	"github.com/go-zeromq/zmq4"
+	"github.com/go-zeromq/zmq4/transport"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -260,3 +262,72 @@ func TestSocketSendSubscriptionOnConnect(t *testing.T) {
 		t.Fatalf("invalid message received: got '%s', wanted '%s'", msg.Frames[0], message)
 	}
 }
+
+type transportMock struct {
+	dialCalledCount int
+	errOnDial       bool
+	conn            net.Conn
+}
+
+func (t *transportMock) Dial(ctx context.Context, dialer transport.Dialer, addr string) (net.Conn, error) {
+	t.dialCalledCount++
+	if t.errOnDial {
+		return nil, errors.New("test error")
+	}
+	return t.conn, nil
+}
+
+func (t *transportMock) Listen(ctx context.Context, addr string) (net.Listener, error) {
+	return nil, nil
+}
+
+func (t *transportMock) Addr(ep string) (addr string, err error) {
+	return "", nil
+}
+
+func (t *transportMock) init() {
+	t.dialCalledCount = 0
+	t.errOnDial = false
+	t.conn = nil
+	zmq4.RegisterTransport("test", testTransport)
+}
+
+var testTransport = &transportMock{}
+
+func TestConnMaxRetries(t *testing.T) {
+	retryCount := 123
+	socket := zmq4.NewSub(context.Background(), zmq4.WithDialerRetry(time.Microsecond), zmq4.WithDialerMaxRetries(retryCount))
+	testTransport.init()
+	testTransport.errOnDial = true
+	err := socket.Dial("test://test")
+
+	if err == nil {
+		t.Fatal("expected error")
+	}
+
+	if testTransport.dialCalledCount != retryCount+1 {
+		t.Fatalf("Dial called %d times, expected %d", testTransport.dialCalledCount, retryCount+1)
+	}
+}
+
+func TestConnMaxRetriesInfinite(t *testing.T) {
+	timeout := time.Millisecond
+	retryTime := time.Nanosecond
+
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	socket := zmq4.NewSub(ctx, zmq4.WithDialerRetry(retryTime), zmq4.WithDialerMaxRetries(-1))
+	testTransport.init()
+	testTransport.errOnDial = true
+
+	err := socket.Dial("test://test")
+
+	if err == nil {
+		t.Fatal("expected error")
+	}
+
+	atLeastExpectedRetries := 100
+	if testTransport.dialCalledCount < atLeastExpectedRetries {
+		t.Fatalf("Dial called %d times, expected  at least %d", testTransport.dialCalledCount, atLeastExpectedRetries)
+	}
+}