Implement reconnect mechanism on the proxy

Move SQL commit closer to the insert command
Use executor.submit instead of map

Author: mostafa

client_test.py

@@ -14,12 +14,12 @@ def worker(id):
         conn.execute(
             "CREATE TABLE IF NOT EXISTS test (id serial PRIMARY KEY, num integer, data varchar);")
         conn.execute("INSERT INTO test (num, data) VALUES (%s, %s)", (id, "abc'def"))
+        conn.execute("COMMIT")
 
         for row in conn.execute("SELECT * FROM test;"):
             print("ID=%s, NUM=%s, DATA=%s" % row)
 
         # conn.execute("DROP TABLE test;")
-        conn.execute("COMMIT")
         conn.close()
     except KeyboardInterrupt:
         if conn:
@@ -32,7 +32,8 @@ def worker(id):
 
 
 if __name__ == '__main__':
-    with ThreadPoolExecutor(max_workers=2) as executor:
+    with ThreadPoolExecutor(max_workers=10) as executor:
         # Create 11 connections to the server and run queries in parallel
         # This will cause the server to crash
-        executor.map(worker, range(1, 12, 1))
+        for i in range(11):
+            executor.submit(worker, i)

network/client.go

@@ -22,6 +22,8 @@ type Client struct {
 // unexpected EOF on client connection with an open transaction
 
 func NewClient(network, address string, receiveBufferSize int) *Client {
+	// TODO: Resolve the address and return an error if it can't be resolved
+
 	c := Client{
 		Network: network,
 		Address: address,

network/proxy.go

@@ -13,6 +13,7 @@ type Proxy interface {
 	Connect(c gnet.Conn) error
 	Disconnect(c gnet.Conn) error
 	PassThrough(c gnet.Conn) error
+	Reconnect(cl *Client) *Client
 	Shutdown()
 	Size() int
 }
@@ -92,25 +93,23 @@ func (pr *ProxyImpl) Connect(c gnet.Conn) error {
 
 func (pr *ProxyImpl) Disconnect(c gnet.Conn) error {
 	var client *Client
-	if c, ok := pr.connClients.Load(c); ok {
-		client = c.(*Client)
+	if cl, ok := pr.connClients.Load(c); ok {
+		client = cl.(*Client)
 	}
 	pr.connClients.Delete(c)
 
 	// TODO: The connection is unstable when I put the client back in the pool
 	// If the client is not in the pool, put it back
-	if client != nil && client.ID != "" {
-		if pr.Elastic && pr.ReuseElasticClients {
+
+	if pr.Elastic && pr.ReuseElasticClients || !pr.Elastic {
+		client = pr.Reconnect(client)
+		if client != nil && client.ID != "" {
 			if err := pr.pool.Put(client); err != nil {
 				return err
 			}
-		} else {
-			// FIXME: Close the client connection, as it is not reusable???
-			pr.pool.Put(client)
-			// pr.pool.Put(NewClient("tcp", "localhost:5432", 4096))
 		}
 	} else {
-		return errors.New("client is not connected (disconnect)")
+		client.Close()
 	}
 
 	logrus.Infof("[D] There are %d clients in the pool", len(pr.pool.ClientIDs()))
@@ -144,11 +143,17 @@ func (pr *ProxyImpl) PassThrough(c gnet.Conn) error {
 	// Receive the response from the server
 	size, response, err := client.Receive()
 	if err != nil {
-		return err
-	}
-
-	if size == 0 {
-		return errors.New("no response from the server")
+		// FIXME: Is this the right way to handle this error?
+		if err.Error() == "EOF" {
+			logrus.Error("The client is not connected to the server anymore")
+			// Either the client is not connected to the server anymore or
+			// server forceful closed the connection
+			// Reconnect the client
+			client = pr.Reconnect(client)
+			// Store the client in the map, replacing the old one
+			pr.connClients.Store(c, client)
+		}
+		// return err
 	}
 
 	// Write the response to the incoming connection
@@ -157,6 +162,14 @@ func (pr *ProxyImpl) PassThrough(c gnet.Conn) error {
 	return nil
 }
 
+func (pr *ProxyImpl) Reconnect(cl *Client) *Client {
+	// Close the client
+	if cl != nil && cl.ID != "" {
+		cl.Close()
+	}
+	return NewClient("tcp", "localhost:5432", 4096)
+}
+
 func (pr *ProxyImpl) Shutdown() {
 	pr.pool.Shutdown()
 	logrus.Info("All busy client connections have been closed")