Fixed invalid connection handling on node connections

aerospike · May 17, 2015 · f846d7a · f846d7a
1 parent 604afc9
commit f846d7a
Show file tree

Hide file tree

Showing 9 changed files with 34 additions and 23 deletions.
diff --git a/admin_command.go b/admin_command.go
@@ -244,12 +244,12 @@ func (acmd *AdminCommand) executeCommand(cluster *Cluster, policy *AdminPolicy)
 	}
 
 	if _, err := conn.Write(acmd.dataBuffer[:acmd.dataOffset]); err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return err
 	}
 
 	if _, err := conn.Read(acmd.dataBuffer, _HEADER_SIZE); err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return err
 	}
 
@@ -280,12 +280,13 @@ func (acmd *AdminCommand) readUsers(cluster *Cluster, policy *AdminPolicy) ([]*U
 	}
 
 	if _, err := conn.Write(acmd.dataBuffer[:acmd.dataOffset]); err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return nil, err
 	}
 
 	status, list, err := acmd.readUserBlocks(conn)
 	if err != nil {
+		node.InvalidateConnection(conn)
 		return nil, err
 	}
 	node.PutConnection(conn)

diff --git a/client.go b/client.go
@@ -596,9 +596,10 @@ func (clnt *Client) RegisterUDF(policy *WritePolicy, udfBody []byte, serverPath
 
 	responseMap, err := RequestInfo(conn, strCmd.String())
 	if err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return nil, err
 	}
+	node.PutConnection(conn)
 
 	var response string
 	for _, v := range responseMap {
@@ -622,8 +623,6 @@ func (clnt *Client) RegisterUDF(policy *WritePolicy, udfBody []byte, serverPath
 		return nil, NewAerospikeError(COMMAND_REJECTED, fmt.Sprintf("Registration failed: %s\nFile: %s\nLine: %s\nMessage: %s",
 			res["error"], res["file"], res["line"], res["message"]))
 	}
-
-	node.PutConnection(conn)
 	return NewRegisterTask(clnt.cluster, serverPath), nil
 }
 
@@ -656,9 +655,10 @@ func (clnt *Client) RemoveUDF(policy *WritePolicy, udfName string) (*RemoveTask,
 
 	responseMap, err := RequestInfo(conn, strCmd.String())
 	if err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return nil, err
 	}
+	node.PutConnection(conn)
 
 	var response string
 	for _, v := range responseMap {
@@ -697,9 +697,10 @@ func (clnt *Client) ListUDF(policy *BasePolicy) ([]*UDF, error) {
 
 	responseMap, err := RequestInfo(conn, strCmd.String())
 	if err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return nil, err
 	}
+	node.PutConnection(conn)
 
 	var response string
 	for _, v := range responseMap {
@@ -734,8 +735,6 @@ func (clnt *Client) ListUDF(policy *BasePolicy) ([]*UDF, error) {
 		res = append(res, udf)
 	}
 
-	node.PutConnection(conn)
-
 	return res, nil
 }
 
@@ -1120,7 +1119,7 @@ func (clnt *Client) sendInfoCommand(policy *WritePolicy, command string) (map[st
 
 	info, err := newInfo(conn, command)
 	if err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return nil, err
 	}
 	node.PutConnection(conn)

diff --git a/command.go b/command.go
@@ -786,7 +786,7 @@ func (cmd *baseCommand) execute(ifc command) (err error) {
 		if err != nil {
 			// All runtime exceptions are considered fatal. Do not retry.
 			// Close socket to flush out possible garbage. Do not put back in pool.
-			cmd.conn.Close()
+			node.InvalidateConnection(cmd.conn)
 			return err
 		}
 
@@ -798,7 +798,7 @@ func (cmd *baseCommand) execute(ifc command) (err error) {
 		if err != nil {
 			// IO errors are considered temporary anomalies. Retry.
 			// Close socket to flush out possible garbage. Do not put back in pool.
-			cmd.conn.Close()
+			node.InvalidateConnection(cmd.conn)
 
 			Logger.Warn("Node " + node.String() + ": " + err.Error())
 			// IO error means connection to server node is unhealthy.
@@ -814,7 +814,7 @@ func (cmd *baseCommand) execute(ifc command) (err error) {
 			// cancelling/closing the batch/multi commands will return an error, which will
 			// close the connection to throw away its data and signal the server about the
 			// situation. We will not put back the connection in the buffer.
-			cmd.conn.Close()
+			node.InvalidateConnection(cmd.conn)
 			return err
 		}
 

diff --git a/execute_task.go b/execute_task.go
@@ -57,6 +57,7 @@ func (etsk *ExecuteTask) IsDone() (bool, error) {
 		}
 		responseMap, err := RequestInfo(conn, command)
 		if err != nil {
+			node.InvalidateConnection(conn)
 			return false, err
 		}
 

diff --git a/info.go b/info.go
@@ -43,7 +43,7 @@ func RequestNodeInfo(node *Node, name ...string) (map[string]string, error) {
 
 	response, err := RequestInfo(conn, name...)
 	if err != nil {
-		conn.Close()
+		node.InvalidateConnection(conn)
 		return nil, err
 	}
 	node.PutConnection(conn)

diff --git a/node.go b/node.go
@@ -87,22 +87,25 @@ func (nd *Node) Refresh() ([]*Host, error) {
 
 	infoMap, err := RequestInfo(conn, "node", "partition-generation", "services")
 	if err != nil {
-		conn.Close()
+		nd.InvalidateConnection(conn)
 		nd.DecreaseHealth()
 		return nil, err
 	}
 
 	if err := nd.verifyNodeName(infoMap); err != nil {
+		nd.PutConnection(conn)
 		return nil, err
 	}
 	nd.RestoreHealth()
 	nd.responded.Set(true)
 
 	if friends, err = nd.addFriends(infoMap); err != nil {
+		nd.PutConnection(conn)
 		return nil, err
 	}
 
 	if err := nd.updatePartitions(conn, infoMap); err != nil {
+		nd.InvalidateConnection(conn)
 		return nil, err
 	}
 
@@ -212,7 +215,7 @@ L:
 					return conn, nil
 				}
 			}
-			conn.Close()
+			nd.InvalidateConnection(conn)
 		}
 
 		// if connection count is limited and enough connections are already created, don't create a new one
@@ -260,11 +263,16 @@ L:
 func (nd *Node) PutConnection(conn *Connection) {
 	conn.refresh()
 	if !nd.active.Get() || !nd.connections.Offer(conn) {
-		nd.connectionCount.DecrementAndGet()
-		conn.Close()
+		nd.InvalidateConnection(conn)
 	}
 }
 
+// InvalidateConnection closes and discards a connection from the pool.
+func (nd *Node) InvalidateConnection(conn *Connection) {
+	nd.connectionCount.DecrementAndGet()
+	conn.Close()
+}
+
 // RestoreHealth marks the node as healthy.
 func (nd *Node) RestoreHealth() {
 	// There can be cases where health is full, but active is false.

diff --git a/node_test.go b/node_test.go
@@ -76,7 +76,7 @@ var _ = Describe("Aerospike", func() {
 					Expect(c).NotTo(BeNil())
 					Expect(c.IsConnected()).To(BeTrue())
 
-					c.Close()
+					node.InvalidateConnection(c)
 				}
 
 			})
@@ -104,7 +104,7 @@ var _ = Describe("Aerospike", func() {
 					Expect(c).NotTo(BeNil())
 					Expect(c.IsConnected()).To(BeTrue())
 
-					c.Close()
+					node.InvalidateConnection(c)
 				}
 
 				for i := 0; i < 4; i++ {
@@ -260,7 +260,7 @@ var _ = Describe("Aerospike", func() {
 				c3, err := node.GetConnection(0)
 				Expect(err).NotTo(HaveOccurred())
 				Expect(c3).NotTo(BeNil())
-				defer c3.Close()
+				defer node.InvalidateConnection(c3)
 				Expect(c3).ToNot(Equal(c))
 				Expect(c3.IsConnected()).To(BeTrue())
 

diff --git a/security_test.go b/security_test.go
@@ -121,7 +121,7 @@ var _ = Describe("Security tests", func() {
 					for i := 0; i < client_policy.ConnectionQueueSize; i++ {
 						conn, err := node.GetConnection(time.Second)
 						Expect(err).ToNot(HaveOccurred())
-						conn.Close()
+						node.InvalidateConnection(conn)
 					}
 				}
 

diff --git a/tools/asinfo/asinfo.go b/tools/asinfo/asinfo.go
@@ -52,8 +52,10 @@ func main() {
 			log.Fatalln(err.Error())
 		} else {
 			if infoMap, err := as.RequestInfo(conn, strings.Trim(*value, " ")); err != nil {
+				node.InvalidateConnection(conn)
 				log.Fatalln(err.Error())
 			} else {
+				node.PutConnection(conn)
 				cnt := 1
 				for k, v := range infoMap {
 					log.Printf("%d :  %s\n     %s\n\n", cnt, k, v)