add redial socket

1 files changed, 12 insertions, 2 deletions

diff --git a/pkg/server/socket/dial.go b/pkg/server/socket/dial.go
index a31e192..aea74d2 100644
--- a/pkg/server/socket/dial.go
+++ b/pkg/server/socket/dial.go
@@ -1,6 +1,7 @@
 package socket
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net"
@@ -21,8 +22,17 @@ func (s *dialSocket) String() string {
 }
 
 func (s *dialSocket) New(e env.Env) (Conn, error) {
-	proto, addr := parseProtoAddr(s.Proto, e.Expand(s.Addr))
-	conn, err := net.DialTimeout(proto, addr, defaultTimeout)
+	return dial(context.TODO(), e, s.Proto, s.Addr)
+}
+
+func dial(ctx context.Context, e env.Env, proto, addr string) (*conn, error) {
+	ctx, cancel := context.WithTimeout(ctx, defaultTimeout)
+	defer cancel()
+
+	var d net.Dialer
+
+	proto, addr = parseProtoAddr(proto, e.Expand(addr))
+	conn, err := d.DialContext(ctx, proto, addr)
 	if err != nil {
 		return nil, err
 	}