Skip to content

Commit

Permalink
proxy: provide keepalive helpers (#231)
Browse files Browse the repository at this point in the history
  • Loading branch information
xhebox authored Mar 5, 2023
1 parent 513b8f2 commit 67e2c78
Show file tree
Hide file tree
Showing 7 changed files with 196 additions and 13 deletions.
20 changes: 15 additions & 5 deletions lib/config/proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import (
"bytes"
"os"
"path/filepath"
"time"

"github.com/BurntSushi/toml"
"github.com/pingcap/TiProxy/lib/util/errors"
Expand All @@ -43,11 +44,19 @@ type Metrics struct {
MetricsInterval uint `toml:"metrics-interval" json:"metrics-interval"`
}

type KeepAlive struct {
Enabled bool `yaml:"enabled,omitempty" toml:"enabled,omitempty" json:"enabled,omitempty"`
Cnt int `yaml:"cnt,omitempty" toml:"cnt,omitempty" json:"cnt,omitempty"`
Idle time.Duration `yaml:"idle,omitempty" toml:"idle,omitempty" json:"idle,omitempty"`
Intvl time.Duration `yaml:"intvl,omitempty" toml:"intvl,omitempty" json:"intvl,omitempty"`
}

type ProxyServerOnline struct {
MaxConnections uint64 `yaml:"max-connections,omitempty" toml:"max-connections,omitempty" json:"max-connections,omitempty"`
TCPKeepAlive bool `yaml:"tcp-keep-alive,omitempty" toml:"tcp-keep-alive,omitempty" json:"tcp-keep-alive,omitempty"`
ProxyProtocol string `yaml:"proxy-protocol,omitempty" toml:"proxy-protocol,omitempty" json:"proxy-protocol,omitempty"`
GracefulWaitBeforeShutdown int `yaml:"graceful-wait-before-shutdown,omitempty" toml:"graceful-wait-before-shutdown,omitempty" json:"graceful-wait-before-shutdown,omitempty"`
MaxConnections uint64 `yaml:"max-connections,omitempty" toml:"max-connections,omitempty" json:"max-connections,omitempty"`
FrontendKeepalive KeepAlive `yaml:"frontend-keepalive" toml:"frontend-keepalive" json:"frontend-keepalive"`
BackendKeepalive KeepAlive `yaml:"backend-keepalive" toml:"backend-keepalive" json:"backend-keepalive"`
ProxyProtocol string `yaml:"proxy-protocol,omitempty" toml:"proxy-protocol,omitempty" json:"proxy-protocol,omitempty"`
GracefulWaitBeforeShutdown int `yaml:"graceful-wait-before-shutdown,omitempty" toml:"graceful-wait-before-shutdown,omitempty" json:"graceful-wait-before-shutdown,omitempty"`
}

type ProxyServer struct {
Expand Down Expand Up @@ -114,7 +123,8 @@ func NewConfig() *Config {
var cfg Config

cfg.Proxy.Addr = "0.0.0.0:6000"
cfg.Proxy.TCPKeepAlive = true
cfg.Proxy.FrontendKeepalive.Enabled = true
cfg.Proxy.BackendKeepalive.Enabled = true
cfg.Proxy.RequireBackendTLS = true
cfg.Proxy.PDAddrs = "127.0.0.1:2379"

Expand Down
2 changes: 1 addition & 1 deletion lib/config/proxy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ var testProxyConfig = Config{
RequireBackendTLS: true,
ProxyServerOnline: ProxyServerOnline{
MaxConnections: 1,
TCPKeepAlive: true,
FrontendKeepalive: KeepAlive{Enabled: true},
ProxyProtocol: "v2",
GracefulWaitBeforeShutdown: 10,
},
Expand Down
47 changes: 47 additions & 0 deletions pkg/proxy/keepalive/keepalive.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
// Copyright 2023 PingCAP, Inc.
//
// 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 keepalive

import (
"net"

"github.com/pingcap/TiProxy/lib/config"
"github.com/pingcap/TiProxy/lib/util/errors"
)

var (
ErrKeepAlive = errors.New("failed to set keepalive")
)

func SetKeepalive(conn net.Conn, cfg config.KeepAlive) error {
tcpcn, ok := conn.(*net.TCPConn)
if !ok {
return errors.Wrapf(ErrKeepAlive, "not net.TCPConn")
}

if err := tcpcn.SetKeepAlive(cfg.Enabled); err != nil {
return errors.Wrap(ErrKeepAlive, err)
}
if !cfg.Enabled {
return nil
}

syscn, err := tcpcn.SyscallConn()
if err != nil {
return errors.Wrap(ErrKeepAlive, err)
}

return setKeepalive(syscn, cfg)
}
54 changes: 54 additions & 0 deletions pkg/proxy/keepalive/keepalive_darwin.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
//go:build darwin

// Copyright 2023 PingCAP, Inc.
//
// 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 keepalive

import (
"syscall"

"github.com/pingcap/TiProxy/lib/config"
"github.com/pingcap/TiProxy/lib/util/errors"
)

const (
// missing in older darwin
_TCP_KEEPINTVL = 0x101
_TCP_KEEPCNT = 0x102
)

func setKeepalive(syscn syscall.RawConn, cfg config.KeepAlive) error {
var serr error
return errors.Collect(ErrKeepAlive, serr, syscn.Control(func(fd uintptr) {
if val := cfg.Idle.Seconds(); val > 0 {
serr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_KEEPALIVE, int(val))
if serr != nil {
return
}
}
if cfg.Cnt > 0 {
serr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, _TCP_KEEPCNT, cfg.Cnt)
if serr != nil {
return
}
}
if val := cfg.Intvl.Seconds(); val > 0 {
serr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, _TCP_KEEPINTVL, int(val))
if serr != nil {
return
}
}
}))
}
27 changes: 27 additions & 0 deletions pkg/proxy/keepalive/keepalive_default.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
//go:build !(linux || netbsd || freebsd || dragonfly || aix || darwin)

// Copyright 2023 PingCAP, Inc.
//
// 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 keepalive

import (
"syscall"

"github.com/pingcap/TiProxy/lib/config"
)

func setKeepalive(syscn syscall.RawConn, cfg config.KeepAlive) error {
return nil
}
48 changes: 48 additions & 0 deletions pkg/proxy/keepalive/keepalive_unix.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
//go:build linux || netbsd || freebsd || dragonfly || aix

// Copyright 2023 PingCAP, Inc.
//
// 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 keepalive

import (
"syscall"

"github.com/pingcap/TiProxy/lib/config"
"github.com/pingcap/TiProxy/lib/util/errors"
)

func setKeepalive(syscn syscall.RawConn, cfg config.KeepAlive) error {
var serr error
return errors.Collect(ErrKeepAlive, serr, syscn.Control(func(fd uintptr) {
if val := cfg.Idle.Seconds(); val > 0 {
serr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_KEEPIDLE, int(val))
if serr != nil {
return
}
}
if cfg.Cnt > 0 {
serr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_KEEPCNT, cfg.Cnt)
if serr != nil {
return
}
}
if val := cfg.Intvl.Seconds(); val > 0 {
serr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_KEEPINTVL, int(val))
if serr != nil {
return
}
}
}))
}
11 changes: 4 additions & 7 deletions pkg/proxy/proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"github.com/pingcap/TiProxy/pkg/metrics"
"github.com/pingcap/TiProxy/pkg/proxy/backend"
"github.com/pingcap/TiProxy/pkg/proxy/client"
"github.com/pingcap/TiProxy/pkg/proxy/keepalive"
pnet "github.com/pingcap/TiProxy/pkg/proxy/net"
"go.uber.org/zap"
)
Expand Down Expand Up @@ -81,7 +82,7 @@ func NewSQLServer(logger *zap.Logger, cfg config.ProxyServer, certMgr *cert.Cert

func (s *SQLServer) reset(cfg *config.ProxyServerOnline) {
s.mu.Lock()
s.mu.tcpKeepAlive = cfg.TCPKeepAlive
s.mu.tcpKeepAlive = cfg.FrontendKeepalive.Enabled
s.mu.maxConnections = cfg.MaxConnections
s.mu.proxyProtocol = cfg.ProxyProtocol != ""
s.mu.gracefulWait = cfg.GracefulWaitBeforeShutdown
Expand Down Expand Up @@ -172,12 +173,8 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn) {
metrics.ConnGauge.Dec()
}()

if tcpKeepAlive {
if tcpConn, ok := conn.(*net.TCPConn); ok {
if err := tcpConn.SetKeepAlive(true); err != nil {
logger.Warn("failed to set tcp keep alive option", zap.Error(err))
}
}
if err := keepalive.SetKeepalive(conn, config.KeepAlive{Enabled: tcpKeepAlive}); err != nil {
logger.Warn("failed to set tcp keep alive option", zap.Error(err))
}

clientConn.Run(ctx)
Expand Down

0 comments on commit 67e2c78

Please sign in to comment.