2019-07-24 07:22:31 +00:00
|
|
|
package loadbalancer
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"errors"
|
2022-04-15 00:31:49 +00:00
|
|
|
"fmt"
|
2019-07-24 07:22:31 +00:00
|
|
|
"net"
|
2021-08-17 23:02:30 +00:00
|
|
|
"os"
|
2019-07-24 07:22:31 +00:00
|
|
|
"path/filepath"
|
|
|
|
"sync"
|
|
|
|
|
2022-03-02 23:47:27 +00:00
|
|
|
"github.com/k3s-io/k3s/pkg/version"
|
2019-07-24 07:22:31 +00:00
|
|
|
"github.com/sirupsen/logrus"
|
2021-06-21 05:52:54 +00:00
|
|
|
"inet.af/tcpproxy"
|
2019-07-24 07:22:31 +00:00
|
|
|
)
|
|
|
|
|
2023-03-31 20:51:27 +00:00
|
|
|
// server tracks the connections to a server, so that they can be closed when the server is removed.
|
|
|
|
type server struct {
|
|
|
|
mutex sync.Mutex
|
|
|
|
connections map[net.Conn]struct{}
|
|
|
|
}
|
|
|
|
|
|
|
|
// serverConn wraps a net.Conn so that it can be removed from the server's connection map when closed.
|
|
|
|
type serverConn struct {
|
|
|
|
server *server
|
|
|
|
net.Conn
|
|
|
|
}
|
|
|
|
|
|
|
|
// LoadBalancer holds data for a local listener which forwards connections to a
|
|
|
|
// pool of remote servers. It is not a proper load-balancer in that it does not
|
|
|
|
// actually balance connections, but instead fails over to a new server only
|
|
|
|
// when a connection attempt to the currently selected server fails.
|
2019-07-24 07:22:31 +00:00
|
|
|
type LoadBalancer struct {
|
2023-03-31 20:51:27 +00:00
|
|
|
mutex sync.Mutex
|
|
|
|
proxy *tcpproxy.Proxy
|
2019-07-24 07:22:31 +00:00
|
|
|
|
2022-02-16 22:19:58 +00:00
|
|
|
serviceName string
|
|
|
|
configFile string
|
|
|
|
localAddress string
|
|
|
|
localServerURL string
|
|
|
|
defaultServerAddress string
|
|
|
|
ServerURL string
|
|
|
|
ServerAddresses []string
|
|
|
|
randomServers []string
|
2023-03-31 20:51:27 +00:00
|
|
|
servers map[string]*server
|
2022-02-16 22:19:58 +00:00
|
|
|
currentServerAddress string
|
|
|
|
nextServerIndex int
|
|
|
|
Listener net.Listener
|
2019-07-24 07:22:31 +00:00
|
|
|
}
|
|
|
|
|
2021-02-12 15:35:57 +00:00
|
|
|
const RandomPort = 0
|
|
|
|
|
2020-05-05 22:09:04 +00:00
|
|
|
var (
|
|
|
|
SupervisorServiceName = version.Program + "-agent-load-balancer"
|
|
|
|
APIServerServiceName = version.Program + "-api-server-agent-load-balancer"
|
2021-02-12 15:35:57 +00:00
|
|
|
ETCDServerServiceName = version.Program + "-etcd-server-load-balancer"
|
2019-07-24 07:22:31 +00:00
|
|
|
)
|
|
|
|
|
2023-03-31 20:51:27 +00:00
|
|
|
// New contstructs a new LoadBalancer instance. The default server URL, and
|
|
|
|
// currently active servers, are stored in a file within the dataDir.
|
2022-03-31 09:49:30 +00:00
|
|
|
func New(ctx context.Context, dataDir, serviceName, serverURL string, lbServerPort int, isIPv6 bool) (_lb *LoadBalancer, _err error) {
|
2021-03-08 22:10:00 +00:00
|
|
|
config := net.ListenConfig{Control: reusePort}
|
2022-04-15 00:31:49 +00:00
|
|
|
var localAddress string
|
2022-03-31 09:49:30 +00:00
|
|
|
if isIPv6 {
|
2022-04-15 00:31:49 +00:00
|
|
|
localAddress = fmt.Sprintf("[::1]:%d", lbServerPort)
|
|
|
|
} else {
|
|
|
|
localAddress = fmt.Sprintf("127.0.0.1:%d", lbServerPort)
|
2022-03-31 09:49:30 +00:00
|
|
|
}
|
2022-04-15 00:31:49 +00:00
|
|
|
listener, err := config.Listen(ctx, "tcp", localAddress)
|
2019-07-24 07:22:31 +00:00
|
|
|
defer func() {
|
|
|
|
if _err != nil {
|
|
|
|
logrus.Warnf("Error starting load balancer: %s", _err)
|
|
|
|
if listener != nil {
|
|
|
|
listener.Close()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
2022-04-15 00:31:49 +00:00
|
|
|
|
|
|
|
// if lbServerPort was 0, the port was assigned by the OS when bound - see what we ended up with.
|
|
|
|
localAddress = listener.Addr().String()
|
2019-07-24 07:22:31 +00:00
|
|
|
|
2022-02-16 22:19:58 +00:00
|
|
|
defaultServerAddress, localServerURL, err := parseURL(serverURL, localAddress)
|
2019-07-24 07:22:31 +00:00
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
2021-03-06 10:29:57 +00:00
|
|
|
if serverURL == localServerURL {
|
2022-02-16 22:19:58 +00:00
|
|
|
logrus.Debugf("Initial server URL for load balancer %s points at local server URL - starting with empty default server address", serviceName)
|
|
|
|
defaultServerAddress = ""
|
2021-03-06 10:29:57 +00:00
|
|
|
}
|
|
|
|
|
2019-07-24 07:22:31 +00:00
|
|
|
lb := &LoadBalancer{
|
2022-02-16 22:19:58 +00:00
|
|
|
serviceName: serviceName,
|
|
|
|
configFile: filepath.Join(dataDir, "etc", serviceName+".json"),
|
|
|
|
localAddress: localAddress,
|
|
|
|
localServerURL: localServerURL,
|
|
|
|
defaultServerAddress: defaultServerAddress,
|
2023-03-31 20:51:27 +00:00
|
|
|
servers: make(map[string]*server),
|
2022-02-16 22:19:58 +00:00
|
|
|
ServerURL: serverURL,
|
2019-07-24 07:22:31 +00:00
|
|
|
}
|
|
|
|
|
2022-02-16 22:19:58 +00:00
|
|
|
lb.setServers([]string{lb.defaultServerAddress})
|
2019-07-24 07:22:31 +00:00
|
|
|
|
|
|
|
lb.proxy = &tcpproxy.Proxy{
|
|
|
|
ListenFunc: func(string, string) (net.Listener, error) {
|
|
|
|
return listener, nil
|
|
|
|
},
|
|
|
|
}
|
|
|
|
lb.proxy.AddRoute(serviceName, &tcpproxy.DialProxy{
|
|
|
|
Addr: serviceName,
|
|
|
|
DialContext: lb.dialContext,
|
2019-08-08 17:48:11 +00:00
|
|
|
OnDialError: onDialError,
|
2019-07-24 07:22:31 +00:00
|
|
|
})
|
|
|
|
|
|
|
|
if err := lb.updateConfig(); err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
if err := lb.proxy.Start(); err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
2023-03-31 20:51:27 +00:00
|
|
|
logrus.Infof("Running load balancer %s %s -> %v [default: %s]", serviceName, lb.localAddress, lb.ServerAddresses, lb.defaultServerAddress)
|
2019-07-24 07:22:31 +00:00
|
|
|
|
|
|
|
return lb, nil
|
|
|
|
}
|
|
|
|
|
2022-02-16 22:19:58 +00:00
|
|
|
func (lb *LoadBalancer) SetDefault(serverAddress string) {
|
2023-03-31 20:51:27 +00:00
|
|
|
lb.mutex.Lock()
|
|
|
|
defer lb.mutex.Unlock()
|
|
|
|
|
|
|
|
_, hasOriginalServer := sortServers(lb.ServerAddresses, lb.defaultServerAddress)
|
|
|
|
// if the old default server is not currently in use, remove it from the server map
|
|
|
|
if server := lb.servers[lb.defaultServerAddress]; server != nil && !hasOriginalServer {
|
|
|
|
defer server.closeAll()
|
|
|
|
delete(lb.servers, lb.defaultServerAddress)
|
|
|
|
}
|
|
|
|
// if the new default server doesn't have an entry in the map, add one
|
|
|
|
if _, ok := lb.servers[serverAddress]; !ok {
|
|
|
|
lb.servers[serverAddress] = &server{connections: make(map[net.Conn]struct{})}
|
|
|
|
}
|
|
|
|
|
2022-02-16 22:19:58 +00:00
|
|
|
lb.defaultServerAddress = serverAddress
|
2023-03-31 20:51:27 +00:00
|
|
|
logrus.Infof("Updated load balancer %s default server address -> %s", lb.serviceName, serverAddress)
|
2022-02-16 22:19:58 +00:00
|
|
|
}
|
|
|
|
|
2019-07-24 07:22:31 +00:00
|
|
|
func (lb *LoadBalancer) Update(serverAddresses []string) {
|
|
|
|
if lb == nil {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
if !lb.setServers(serverAddresses) {
|
|
|
|
return
|
|
|
|
}
|
2023-03-31 20:51:27 +00:00
|
|
|
logrus.Infof("Updated load balancer %s server addresses -> %v [default: %s]", lb.serviceName, lb.ServerAddresses, lb.defaultServerAddress)
|
2019-07-24 07:22:31 +00:00
|
|
|
|
|
|
|
if err := lb.writeConfig(); err != nil {
|
2022-02-16 22:19:58 +00:00
|
|
|
logrus.Warnf("Error updating load balancer %s config: %s", lb.serviceName, err)
|
2019-07-24 07:22:31 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (lb *LoadBalancer) LoadBalancerServerURL() string {
|
|
|
|
if lb == nil {
|
|
|
|
return ""
|
|
|
|
}
|
|
|
|
return lb.localServerURL
|
|
|
|
}
|
|
|
|
|
|
|
|
func (lb *LoadBalancer) dialContext(ctx context.Context, network, address string) (net.Conn, error) {
|
|
|
|
startIndex := lb.nextServerIndex
|
|
|
|
for {
|
|
|
|
targetServer := lb.currentServerAddress
|
|
|
|
|
2023-03-31 20:51:27 +00:00
|
|
|
server := lb.servers[targetServer]
|
|
|
|
if server == nil || targetServer == "" {
|
|
|
|
logrus.Debugf("Nil server for load balancer %s: %s", lb.serviceName, targetServer)
|
|
|
|
} else {
|
|
|
|
conn, err := server.dialContext(ctx, network, targetServer)
|
|
|
|
if err == nil {
|
|
|
|
return conn, nil
|
|
|
|
}
|
|
|
|
logrus.Debugf("Dial error from load balancer %s: %s", lb.serviceName, err)
|
2019-07-24 07:22:31 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
newServer, err := lb.nextServer(targetServer)
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
if targetServer != newServer {
|
2023-03-31 20:51:27 +00:00
|
|
|
logrus.Debugf("Failed over to new server for load balancer %s: %s", lb.serviceName, newServer)
|
2019-07-24 07:22:31 +00:00
|
|
|
}
|
|
|
|
if ctx.Err() != nil {
|
|
|
|
return nil, ctx.Err()
|
|
|
|
}
|
|
|
|
|
|
|
|
maxIndex := len(lb.randomServers)
|
|
|
|
if startIndex > maxIndex {
|
|
|
|
startIndex = maxIndex
|
|
|
|
}
|
|
|
|
if lb.nextServerIndex == startIndex {
|
|
|
|
return nil, errors.New("all servers failed")
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2019-08-08 17:48:11 +00:00
|
|
|
|
|
|
|
func onDialError(src net.Conn, dstDialErr error) {
|
2023-03-31 20:51:27 +00:00
|
|
|
logrus.Debugf("Incoming conn %s, error dialing load balancer servers: %v", src.RemoteAddr(), dstDialErr)
|
2019-08-08 17:48:11 +00:00
|
|
|
src.Close()
|
|
|
|
}
|
2021-08-17 23:02:30 +00:00
|
|
|
|
2022-02-16 22:19:58 +00:00
|
|
|
// ResetLoadBalancer will delete the local state file for the load balancer on disk
|
2021-08-17 23:02:30 +00:00
|
|
|
func ResetLoadBalancer(dataDir, serviceName string) error {
|
|
|
|
stateFile := filepath.Join(dataDir, "etc", serviceName+".json")
|
|
|
|
if err := os.Remove(stateFile); err != nil {
|
|
|
|
logrus.Warn(err)
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|