-
Notifications
You must be signed in to change notification settings - Fork 3.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
57431: ccl/sqlproxyccl: idle connection timeout support r=darinpp a=darinpp Previusly, the connection from the end user to the backend can be idle for any period of time without disconnecting. In certain cases, we want the ability for idle connections to disconnect when the idle time exceeds redefined timeout. To address this, this patch adds tracking of how long the connection to the backend has been idle. If there is a timeout specified, the connection will be disconnected and an error message will be sent back to the end user. Release note: None Co-authored-by: Darin Peshev <[email protected]>
- Loading branch information
Showing
11 changed files
with
536 additions
and
139 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,79 @@ | ||
// Copyright 2020 The Cockroach Authors. | ||
// | ||
// Licensed as a CockroachDB Enterprise file under the Cockroach Community | ||
// License (the "License"); you may not use this file except in compliance with | ||
// the License. You may obtain a copy of the License at | ||
// | ||
// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt | ||
|
||
package sqlproxyccl | ||
|
||
import ( | ||
"crypto/tls" | ||
"encoding/binary" | ||
"io" | ||
"net" | ||
|
||
"github.com/jackc/pgproto3/v2" | ||
) | ||
|
||
// BackendDial is an example backend dialer that does a TCP/IP connection | ||
// to a backend, SSL and forwards the start message. | ||
func BackendDial( | ||
msg *pgproto3.StartupMessage, outgoingAddress string, tlsConfig *tls.Config, | ||
) (net.Conn, error) { | ||
conn, err := net.Dial("tcp", outgoingAddress) | ||
if err != nil { | ||
return nil, NewErrorf( | ||
CodeBackendDown, "unable to reach backend SQL server: %v", err, | ||
) | ||
} | ||
conn, err = SSLOverlay(conn, tlsConfig) | ||
if err != nil { | ||
return nil, err | ||
} | ||
err = RelayStartupMsg(conn, msg) | ||
if err != nil { | ||
return nil, NewErrorf( | ||
CodeBackendDown, "relaying StartupMessage to target server %v: %v", | ||
outgoingAddress, err) | ||
} | ||
return conn, nil | ||
} | ||
|
||
// SSLOverlay attempts to upgrade the PG connection to use SSL | ||
// if a tls.Config is specified.. | ||
func SSLOverlay(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) { | ||
if tlsConfig == nil { | ||
return conn, nil | ||
} | ||
|
||
var err error | ||
// Send SSLRequest. | ||
if err := binary.Write(conn, binary.BigEndian, pgSSLRequest); err != nil { | ||
return nil, NewErrorf( | ||
CodeBackendDown, "sending SSLRequest to target server: %v", err, | ||
) | ||
} | ||
|
||
response := make([]byte, 1) | ||
if _, err = io.ReadFull(conn, response); err != nil { | ||
return nil, | ||
NewErrorf(CodeBackendDown, "reading response to SSLRequest") | ||
} | ||
|
||
if response[0] != pgAcceptSSLRequest { | ||
return nil, NewErrorf( | ||
CodeBackendRefusedTLS, "target server refused TLS connection", | ||
) | ||
} | ||
|
||
outCfg := tlsConfig.Clone() | ||
return tls.Client(conn, outCfg), nil | ||
} | ||
|
||
// RelayStartupMsg forwards the start message on the backend connection. | ||
func RelayStartupMsg(conn net.Conn, msg *pgproto3.StartupMessage) (err error) { | ||
_, err = conn.Write(msg.Encode(nil)) | ||
return | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,97 @@ | ||
// Copyright 2020 The Cockroach Authors. | ||
// | ||
// Licensed as a CockroachDB Enterprise file under the Cockroach Community | ||
// License (the "License"); you may not use this file except in compliance with | ||
// the License. You may obtain a copy of the License at | ||
// | ||
// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt | ||
|
||
package sqlproxyccl | ||
|
||
import ( | ||
"net" | ||
"time" | ||
|
||
"github.com/cockroachdb/cockroach/pkg/util/syncutil" | ||
"github.com/cockroachdb/cockroach/pkg/util/timeutil" | ||
"github.com/cockroachdb/errors" | ||
) | ||
|
||
// IdleDisconnectConnection is a wrapper around net.Conn that disconnects if | ||
// connection is idle. The idle time is only counted while the client is | ||
// waiting, blocked on Read. | ||
type IdleDisconnectConnection struct { | ||
net.Conn | ||
timeout time.Duration | ||
mu struct { | ||
syncutil.Mutex | ||
lastDeadlineSetAt time.Time | ||
} | ||
} | ||
|
||
var errNotSupported = errors.Errorf( | ||
"Not supported for IdleDisconnectConnection", | ||
) | ||
|
||
func (c *IdleDisconnectConnection) updateDeadline() error { | ||
now := timeutil.Now() | ||
// If it has been more than 1% of the timeout duration - advance the deadline. | ||
c.mu.Lock() | ||
defer c.mu.Unlock() | ||
if now.Sub(c.mu.lastDeadlineSetAt) > c.timeout/100 { | ||
c.mu.lastDeadlineSetAt = now | ||
|
||
if err := c.Conn.SetReadDeadline(now.Add(c.timeout)); err != nil { | ||
return err | ||
} | ||
} | ||
return nil | ||
} | ||
|
||
// Read reads data from the connection with timeout. | ||
func (c *IdleDisconnectConnection) Read(b []byte) (n int, err error) { | ||
if err := c.updateDeadline(); err != nil { | ||
return 0, err | ||
} | ||
return c.Conn.Read(b) | ||
} | ||
|
||
// Write writes data to the connection and sets the read timeout. | ||
func (c *IdleDisconnectConnection) Write(b []byte) (n int, err error) { | ||
// The Write for the connection is not blocking (or can block only temporary | ||
// in case of flow control). For idle connections, the Read will be the call | ||
// that will block and stay blocked until the backend doesn't send something. | ||
// However, it is theoretically possible, that the traffic is only going in | ||
// one direction - from the proxy to the backend, in which case we will call | ||
// repeatedly Write but stay blocked on the Read. For that specific case - the | ||
// write pushes further out the read deadline so the read doesn't timeout. | ||
if err := c.updateDeadline(); err != nil { | ||
return 0, err | ||
} | ||
return c.Conn.Write(b) | ||
} | ||
|
||
// SetDeadline is unsupported as it will interfere with the reads. | ||
func (c *IdleDisconnectConnection) SetDeadline(t time.Time) error { | ||
return errNotSupported | ||
} | ||
|
||
// SetReadDeadline is unsupported as it will interfere with the reads. | ||
func (c *IdleDisconnectConnection) SetReadDeadline(t time.Time) error { | ||
return errNotSupported | ||
} | ||
|
||
// SetWriteDeadline is unsupported as it will interfere with the reads. | ||
func (c *IdleDisconnectConnection) SetWriteDeadline(t time.Time) error { | ||
return errNotSupported | ||
} | ||
|
||
// IdleDisconnectOverlay upgrades the connection to one that closes when | ||
// idle for more than timeout duration. Timeout of zero will turn off | ||
// the idle disconnect code. | ||
func IdleDisconnectOverlay(conn net.Conn, timeout time.Duration) net.Conn { | ||
if timeout != 0 { | ||
return &IdleDisconnectConnection{Conn: conn, timeout: timeout} | ||
} | ||
return conn | ||
} |
Oops, something went wrong.