@@ -10,6 +10,7 @@ import (
1010 "fmt"
1111 "log"
1212 "net"
13+ "time"
1314
1415 "inet.af/tcpproxy"
1516
@@ -27,6 +28,10 @@ import (
2728type TCP struct {
2829 tcpproxy.Proxy
2930
31+ DialTimeout time.Duration
32+ KeepAlivePeriod time.Duration
33+ TCPUserTimeout time.Duration
34+
3035 Logger * log.Logger
3136
3237 routes map [string ]* upstream.List
@@ -51,8 +56,11 @@ func (upstream lbUpstream) HealthCheck(ctx context.Context) error {
5156}
5257
5358type lbTarget struct {
54- list * upstream.List
55- logger * log.Logger
59+ list * upstream.List
60+ logger * log.Logger
61+ dialTimeout time.Duration
62+ keepAlivePeriod time.Duration
63+ tcpUserTimeout time.Duration
5664}
5765
5866func (target * lbTarget ) HandleConn (conn net.Conn ) {
@@ -69,6 +77,9 @@ func (target *lbTarget) HandleConn(conn net.Conn) {
6977 target .logger .Printf ("proxying connection %s -> %s" , conn .RemoteAddr (), upstream .upstream )
7078
7179 upstreamTarget := tcpproxy .To (upstream .upstream )
80+ upstreamTarget .DialTimeout = target .dialTimeout
81+ upstreamTarget .KeepAlivePeriod = target .keepAlivePeriod
82+ upstreamTarget .TCPUserTimeout = target .tcpUserTimeout
7283 upstreamTarget .OnDialError = func (src net.Conn , dstDialErr error ) {
7384 src .Close () //nolint: errcheck
7485
@@ -78,6 +89,8 @@ func (target *lbTarget) HandleConn(conn net.Conn) {
7889 }
7990
8091 upstreamTarget .HandleConn (conn )
92+
93+ target .logger .Printf ("closing connection %s -> %s" , conn .RemoteAddr (), upstream .upstream )
8194}
8295
8396// AddRoute installs load balancer route from listen address ipAddr to list of upstreams.
@@ -109,8 +122,11 @@ func (t *TCP) AddRoute(ipPort string, upstreamAddrs []string, options ...upstrea
109122 t .routes [ipPort ] = list
110123
111124 t .Proxy .AddRoute (ipPort , & lbTarget {
112- list : list ,
113- logger : t .Logger ,
125+ list : list ,
126+ logger : t .Logger ,
127+ dialTimeout : t .DialTimeout ,
128+ keepAlivePeriod : t .KeepAlivePeriod ,
129+ tcpUserTimeout : t .TCPUserTimeout ,
114130 })
115131
116132 return nil
0 commit comments