@@ -10,6 +10,7 @@ import (
1010 "fmt"
1111 "log"
1212 "net"
13+ "time"
1314
1415 "inet.af/tcpproxy"
1516
@@ -27,6 +28,9 @@ import (
2728type TCP struct {
2829 tcpproxy.Proxy
2930
31+ DialTimeout time.Duration
32+ KeepAlivePeriod time.Duration
33+
3034 Logger * log.Logger
3135
3236 routes map [string ]* upstream.List
@@ -51,8 +55,10 @@ func (upstream lbUpstream) HealthCheck(ctx context.Context) error {
5155}
5256
5357type lbTarget struct {
54- list * upstream.List
55- logger * log.Logger
58+ list * upstream.List
59+ logger * log.Logger
60+ dialTimeout time.Duration
61+ keepAlivePeriod time.Duration
5662}
5763
5864func (target * lbTarget ) HandleConn (conn net.Conn ) {
@@ -69,6 +75,8 @@ func (target *lbTarget) HandleConn(conn net.Conn) {
6975 target .logger .Printf ("proxying connection %s -> %s" , conn .RemoteAddr (), upstream .upstream )
7076
7177 upstreamTarget := tcpproxy .To (upstream .upstream )
78+ upstreamTarget .DialTimeout = target .dialTimeout
79+ upstreamTarget .KeepAlivePeriod = target .keepAlivePeriod
7280 upstreamTarget .OnDialError = func (src net.Conn , dstDialErr error ) {
7381 src .Close () //nolint: errcheck
7482
@@ -78,6 +86,8 @@ func (target *lbTarget) HandleConn(conn net.Conn) {
7886 }
7987
8088 upstreamTarget .HandleConn (conn )
89+
90+ target .logger .Printf ("closing connection %s -> %s" , conn .RemoteAddr (), upstream .upstream )
8191}
8292
8393// AddRoute installs load balancer route from listen address ipAddr to list of upstreams.
@@ -109,8 +119,10 @@ func (t *TCP) AddRoute(ipPort string, upstreamAddrs []string, options ...upstrea
109119 t .routes [ipPort ] = list
110120
111121 t .Proxy .AddRoute (ipPort , & lbTarget {
112- list : list ,
113- logger : t .Logger ,
122+ list : list ,
123+ logger : t .Logger ,
124+ dialTimeout : t .DialTimeout ,
125+ keepAlivePeriod : t .KeepAlivePeriod ,
114126 })
115127
116128 return nil
0 commit comments