2016-05-13 06:51:48 +03:00
// Copyright 2016 The etcd Authors
2016-04-14 19:07:21 +03:00
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package tcpproxy
import (
2017-05-05 07:57:54 +03:00
"fmt"
2016-04-14 19:07:21 +03:00
"io"
2017-05-05 07:57:54 +03:00
"math/rand"
2016-04-14 19:07:21 +03:00
"net"
"sync"
"time"
2016-05-13 03:03:14 +03:00
"github.com/coreos/pkg/capnslog"
2018-04-15 08:52:39 +03:00
"go.uber.org/zap"
2016-05-13 03:03:14 +03:00
)
2018-08-29 03:13:25 +03:00
var plog = capnslog . NewPackageLogger ( "go.etcd.io/etcd" , "proxy/tcpproxy" )
2016-04-14 19:07:21 +03:00
type remote struct {
mu sync . Mutex
2017-05-05 07:57:54 +03:00
srv * net . SRV
2016-04-14 19:07:21 +03:00
addr string
inactive bool
}
func ( r * remote ) inactivate ( ) {
r . mu . Lock ( )
defer r . mu . Unlock ( )
r . inactive = true
}
2016-05-13 03:03:14 +03:00
func ( r * remote ) tryReactivate ( ) error {
2016-04-14 19:07:21 +03:00
conn , err := net . Dial ( "tcp" , r . addr )
if err != nil {
2016-05-13 03:03:14 +03:00
return err
2016-04-14 19:07:21 +03:00
}
conn . Close ( )
r . mu . Lock ( )
defer r . mu . Unlock ( )
r . inactive = false
2016-05-13 03:03:14 +03:00
return nil
2016-04-14 19:07:21 +03:00
}
func ( r * remote ) isActive ( ) bool {
r . mu . Lock ( )
defer r . mu . Unlock ( )
return ! r . inactive
}
2016-04-15 07:00:03 +03:00
type TCPProxy struct {
2018-04-15 08:52:39 +03:00
Logger * zap . Logger
2016-04-15 07:00:03 +03:00
Listener net . Listener
2017-05-05 07:57:54 +03:00
Endpoints [ ] * net . SRV
2016-04-15 07:00:03 +03:00
MonitorInterval time . Duration
donec chan struct { }
2017-05-05 07:57:54 +03:00
mu sync . Mutex // guards the following fields
remotes [ ] * remote
pickCount int // for round robin
2016-04-15 07:00:03 +03:00
}
func ( tp * TCPProxy ) Run ( ) error {
tp . donec = make ( chan struct { } )
if tp . MonitorInterval == 0 {
tp . MonitorInterval = 5 * time . Minute
}
2017-05-05 07:57:54 +03:00
for _ , srv := range tp . Endpoints {
addr := fmt . Sprintf ( "%s:%d" , srv . Target , srv . Port )
tp . remotes = append ( tp . remotes , & remote { srv : srv , addr : addr } )
2016-04-15 07:00:03 +03:00
}
2017-05-17 20:51:35 +03:00
eps := [ ] string { }
for _ , ep := range tp . Endpoints {
eps = append ( eps , fmt . Sprintf ( "%s:%d" , ep . Target , ep . Port ) )
}
2018-04-15 08:52:39 +03:00
if tp . Logger != nil {
tp . Logger . Info ( "ready to proxy client requests" , zap . Strings ( "endpoints" , eps ) )
} else {
plog . Printf ( "ready to proxy client requests to %+v" , eps )
}
2017-05-17 20:51:35 +03:00
2016-04-14 19:07:21 +03:00
go tp . runMonitor ( )
for {
2016-04-15 07:00:03 +03:00
in , err := tp . Listener . Accept ( )
2016-04-14 19:07:21 +03:00
if err != nil {
return err
}
go tp . serve ( in )
}
}
2017-05-05 07:57:54 +03:00
func ( tp * TCPProxy ) pick ( ) * remote {
var weighted [ ] * remote
var unweighted [ ] * remote
bestPr := uint16 ( 65535 )
w := 0
// find best priority class
for _ , r := range tp . remotes {
switch {
case ! r . isActive ( ) :
case r . srv . Priority < bestPr :
bestPr = r . srv . Priority
w = 0
2017-10-03 12:14:22 +03:00
weighted = nil
2017-05-05 07:57:54 +03:00
unweighted = [ ] * remote { r }
fallthrough
case r . srv . Priority == bestPr :
if r . srv . Weight > 0 {
weighted = append ( weighted , r )
w += int ( r . srv . Weight )
} else {
unweighted = append ( unweighted , r )
}
}
}
if weighted != nil {
if len ( unweighted ) > 0 && rand . Intn ( 100 ) == 1 {
// In the presence of records containing weights greater
// than 0, records with weight 0 should have a very small
// chance of being selected.
r := unweighted [ tp . pickCount % len ( unweighted ) ]
tp . pickCount ++
return r
}
// choose a uniform random number between 0 and the sum computed
// (inclusive), and select the RR whose running sum value is the
// first in the selected order
choose := rand . Intn ( w )
for i := 0 ; i < len ( weighted ) ; i ++ {
choose -= int ( weighted [ i ] . srv . Weight )
if choose <= 0 {
return weighted [ i ]
}
}
}
if unweighted != nil {
for i := 0 ; i < len ( tp . remotes ) ; i ++ {
picked := tp . remotes [ tp . pickCount % len ( tp . remotes ) ]
tp . pickCount ++
if picked . isActive ( ) {
return picked
}
}
}
return nil
2016-04-14 19:07:21 +03:00
}
2016-04-15 07:00:03 +03:00
func ( tp * TCPProxy ) serve ( in net . Conn ) {
2016-04-14 19:07:21 +03:00
var (
err error
out net . Conn
)
2017-05-05 07:57:54 +03:00
for {
tp . mu . Lock ( )
2016-04-14 19:07:21 +03:00
remote := tp . pick ( )
2017-05-05 07:57:54 +03:00
tp . mu . Unlock ( )
if remote == nil {
break
2016-04-14 19:07:21 +03:00
}
// TODO: add timeout
out , err = net . Dial ( "tcp" , remote . addr )
if err == nil {
break
}
remote . inactivate ( )
2018-04-15 08:52:39 +03:00
if tp . Logger != nil {
tp . Logger . Warn ( "deactivated endpoint" , zap . String ( "address" , remote . addr ) , zap . Duration ( "interval" , tp . MonitorInterval ) , zap . Error ( err ) )
} else {
plog . Warningf ( "deactivated endpoint [%s] due to %v for %v" , remote . addr , err , tp . MonitorInterval )
}
2016-04-14 19:07:21 +03:00
}
if out == nil {
in . Close ( )
return
}
go func ( ) {
io . Copy ( in , out )
in . Close ( )
out . Close ( )
} ( )
io . Copy ( out , in )
out . Close ( )
in . Close ( )
}
2016-04-15 07:00:03 +03:00
func ( tp * TCPProxy ) runMonitor ( ) {
2016-04-14 19:07:21 +03:00
for {
select {
2016-04-15 07:00:03 +03:00
case <- time . After ( tp . MonitorInterval ) :
2016-04-14 19:07:21 +03:00
tp . mu . Lock ( )
2017-02-21 23:39:49 +03:00
for _ , rem := range tp . remotes {
if rem . isActive ( ) {
continue
2016-04-14 19:07:21 +03:00
}
2017-02-21 23:39:49 +03:00
go func ( r * remote ) {
if err := r . tryReactivate ( ) ; err != nil {
2018-04-15 08:52:39 +03:00
if tp . Logger != nil {
tp . Logger . Warn ( "failed to activate endpoint (stay inactive for another interval)" , zap . String ( "address" , r . addr ) , zap . Duration ( "interval" , tp . MonitorInterval ) , zap . Error ( err ) )
} else {
plog . Warningf ( "failed to activate endpoint [%s] due to %v (stay inactive for another %v)" , r . addr , err , tp . MonitorInterval )
}
2017-02-21 23:39:49 +03:00
} else {
2018-04-15 08:52:39 +03:00
if tp . Logger != nil {
tp . Logger . Info ( "activated" , zap . String ( "address" , r . addr ) )
} else {
plog . Printf ( "activated %s" , r . addr )
}
2017-02-21 23:39:49 +03:00
}
} ( rem )
2016-04-14 19:07:21 +03:00
}
tp . mu . Unlock ( )
case <- tp . donec :
return
}
}
}
2016-04-15 07:00:03 +03:00
func ( tp * TCPProxy ) Stop ( ) {
2016-04-14 19:07:21 +03:00
// graceful shutdown?
// shutdown current connections?
2016-04-15 07:00:03 +03:00
tp . Listener . Close ( )
2016-04-14 19:07:21 +03:00
close ( tp . donec )
}