Fix #45: Support Connection Load balancing (#85)

* Add load balancing

* Update README.md

Author: sitingren

README.md

@@ -95,6 +95,8 @@ Currently supported query arguments are:
 |----------------|-------------|--------|
 | use_prepared_statements    | whether to use client-side query interpolation or server-side argument binding | 1 = (default) use server-side bindings |
 |                |             | 0 = user client side interpolation **(LESS SECURE)** |
+| connection_load_balance    | whether to enable connection load balancing on the client side | 0 = (default) disable load balancing |
+|                |             | 1 = enable load balancing |
 | tlsmode            | the ssl/tls policy for this connection | 'none' (default) = don't use SSL/TLS for this connection |
 |                |                                    | 'server' = server must support SSL/TLS, but skip verification **(INSECURE!)** |
 |                |                                    | 'server-strict' = server must support SSL/TLS |

connection.go

@@ -180,6 +180,9 @@ func newConnection(connString string) (*connection, error) {
 		result.usePreparedStmts = iFlag == "1"
 	}
 
+	// Read connection load balance flag.
+	loadBalanceFlag := result.connURL.Query().Get("connection_load_balance")
+
 	sslFlag := strings.ToLower(result.connURL.Query().Get("tlsmode"))
 	if sslFlag == "" {
 		sslFlag = "none"
@@ -191,6 +194,13 @@ func newConnection(connString string) (*connection, error) {
 		return nil, fmt.Errorf("cannot connect to %s (%s)", result.connURL.Host, err.Error())
 	}
 
+	// Load Balancing
+	if loadBalanceFlag == "1" {
+		if err = result.balanceLoad(); err != nil {
+			return nil, err
+		}
+	}
+
 	if sslFlag != "none" {
 		if err = result.initializeSSL(sslFlag); err != nil {
 			return nil, err
@@ -432,6 +442,67 @@ func (v *connection) readAll(buf []byte) error {
 	}
 }
 
+func (v *connection) balanceLoad() error {
+	v.sendMessage(&msgs.FELoadBalanceMsg{})
+	response := v.scratch[:1]
+
+	var err error
+	if err = v.readAll(response); err != nil {
+		return err
+	}
+
+	if response[0] == 'N' {
+		// keep existing connection
+		connectionLogger.Debug("<- LoadBalanceResponse: N")
+		connectionLogger.Warn("Load balancing requested but not supported by server")
+		return nil
+	}
+
+	if response[0] != 'Y' {
+		connectionLogger.Debug("<- LoadBalanceResponse: %c", response[0])
+		return fmt.Errorf("Load balancing request gave unknown response: %c", response[0])
+	}
+
+	header := v.scratch[1:5]
+	if err = v.readAll(header); err != nil {
+		return err
+	}
+	msgSize := int(binary.BigEndian.Uint32(header) - 4)
+	msgBytes := v.scratch[5:]
+
+	var y []byte
+	if msgSize > 0 {
+		if msgSize <= len(msgBytes) {
+			y = msgBytes[:msgSize]
+		} else {
+			y = make([]byte, msgSize)
+		}
+		if err = v.readAll(y); err != nil {
+			return err
+		}
+	}
+
+	bem, err := msgs.CreateBackEndMsg(response[0], y)
+	if err != nil {
+		return err
+	}
+	connectionLogger.Debug("<- " + bem.String())
+	msg := bem.(*msgs.BELoadBalanceMsg)
+
+	// v.connURL.Hostname() is used by initializeSSL(), so load balancing info should not write into v.connURL
+	loadBalanceAddr := fmt.Sprintf("%s:%d", msg.Host, msg.Port)
+
+	// Connect to new host
+	v.conn.Close()
+	v.conn, err = net.Dial("tcp", loadBalanceAddr)
+
+	if err != nil {
+		return fmt.Errorf("cannot redirect to %s (%s)", loadBalanceAddr, err.Error())
+	}
+
+	return nil
+}
+
 func (v *connection) initializeSSL(sslFlag string) error {
 	v.sendMessage(&msgs.FESSLMsg{})
 

msgs/beloadbalancemsg.go

@@ -0,0 +1,59 @@
+package msgs
+
+// Copyright (c) 2020 Micro Focus or one of its affiliates.
+//
+// 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.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+import (
+	"fmt"
+)
+
+type BELoadBalanceMsg struct {
+	Port uint32
+	Host string
+}
+
+// CreateFromMsgBody
+func (m *BELoadBalanceMsg) CreateFromMsgBody(buf *msgBuffer) (BackEndMsg, error) {
+	msg := &BELoadBalanceMsg{
+		Port: buf.readUint32(),
+		Host: buf.readString(),
+	}
+	return msg, nil
+}
+
+func (m *BELoadBalanceMsg) String() string {
+	return fmt.Sprintf("LoadBalanceResponse: host=%s, port=%d", m.Host, m.Port)
+}
+
+func init() {
+	registerBackEndMsgType('Y', &BELoadBalanceMsg{})
+}

msgs/feloadbalancemsg.go

@@ -0,0 +1,51 @@
+package msgs
+
+// Copyright (c) 2020 Micro Focus or one of its affiliates.
+//
+// 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.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+// FELoadBalanceMsg docs
+type FELoadBalanceMsg struct {
+}
+
+// Flatten docs
+func (m *FELoadBalanceMsg) Flatten() ([]byte, byte) {
+
+	buf := newMsgBuffer()
+
+	buf.appendUint32(80936960)
+
+	return buf.bytes(), 0
+}
+
+func (m *FELoadBalanceMsg) String() string {
+	return "LoadBalanceRequest"
+}