lib/net: add function WaitAlive

WaitAlive try to connect to network at address until timeout reached.
If connection cannot established it will return an error.

Unlike [net.DialTimeout], this function will retry not returning an error
immediately if the address has not ready yet.

2 files changed, 89 insertions, 0 deletions

diff --git a/lib/net/net.go b/lib/net/net.go
index 73c951bc..5cc2f223 100644
--- a/lib/net/net.go
+++ b/lib/net/net.go
@@ -7,8 +7,10 @@ package net
 
 import (
 	"errors"
+	"fmt"
 	"net"
 	"strings"
+	"time"
 )
 
 const (
@@ -146,3 +148,31 @@ func ToDotIPv6(ip net.IP) (out []byte) {
 
 	return out
 }
+
+// WaitAlive try to connect to network at address until timeout reached.
+// If connection cannot established it will return an error.
+//
+// Unlike [net.DialTimeout], this function will retry not returning an error
+// immediately if the address has not ready yet.
+func WaitAlive(network, address string, timeout time.Duration) (err error) {
+	var (
+		logp        = `WaitAlive`
+		dialTimeout = 100 * time.Millisecond
+		total       = dialTimeout
+		dialer      = net.Dialer{Timeout: timeout}
+
+		conn net.Conn
+	)
+
+	for total < timeout {
+		conn, err = dialer.Dial(network, address)
+		if err != nil {
+			total += dialTimeout
+			continue
+		}
+		// Connection successfully established.
+		_ = conn.Close()
+		return nil
+	}
+	return fmt.Errorf(`%s: timeout connecting to %s after %s`, logp, address, timeout)
+}
diff --git a/lib/net/net_test.go b/lib/net/net_test.go
index 2444704b..e36295c5 100644
--- a/lib/net/net_test.go
+++ b/lib/net/net_test.go
@@ -7,6 +7,7 @@ package net
 import (
 	"net"
 	"testing"
+	"time"
 
 	"github.com/shuLhan/share/lib/test"
 )
@@ -208,3 +209,61 @@ func TestToDotIPv6(t *testing.T) {
 		test.Assert(t, "ToDotIPv6", c.exp, got)
 	}
 }
+
+func TestWaitAlive_tcp(t *testing.T) {
+	var (
+		address = `127.0.0.1:27392`
+		tcpAddr = net.TCPAddr{
+			IP:   net.ParseIP(`127.0.0.1`),
+			Port: 27392,
+		}
+		errq = make(chan error, 1)
+
+		tcpListener *net.TCPListener
+	)
+
+	go func() {
+		var err2 error
+		tcpListener, err2 = net.ListenTCP(`tcp`, &tcpAddr)
+		errq <- err2
+	}()
+
+	var err = WaitAlive(`tcp`, address, 5*time.Second)
+	if err != nil {
+		t.Fatalf(`want no error, got %s`, err)
+	}
+	err = <-errq
+	if err != nil {
+		t.Fatal(err)
+	}
+	_ = tcpListener.Close()
+}
+
+func TestWaitAlive_udp(t *testing.T) {
+	var (
+		address = `127.0.0.1:27392`
+		udpAddr = net.UDPAddr{
+			IP:   net.ParseIP(`127.0.0.1`),
+			Port: 27392,
+		}
+		errq = make(chan error, 1)
+
+		udpConn *net.UDPConn
+	)
+
+	go func() {
+		var err2 error
+		udpConn, err2 = net.ListenUDP(`udp`, &udpAddr)
+		errq <- err2
+	}()
+
+	var err = WaitAlive(`udp`, address, 5*time.Second)
+	if err != nil {
+		t.Fatalf(`want no error, got %s`, err)
+	}
+	err = <-errq
+	if err != nil {
+		t.Fatal(err)
+	}
+	_ = udpConn.Close()
+}