Ver código fonte

Merge pull request #136 from Zariel/dont-hide-write-error

Dont hide the error from write
Ben Hood 11 anos atrás
pai
commit
36fcdadad8
2 arquivos alterados com 11 adições e 9 exclusões
  1. 9 4
      cluster.go
  2. 2 5
      conn.go

+ 9 - 4
cluster.go

@@ -182,9 +182,8 @@ func (c *clusterImpl) addConn(conn *Conn, keyspace string) {
 	c.conns[conn] = struct{}{}
 }
 
-func (c *clusterImpl) removeConn(conn *Conn) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
+// Should only be called if c.mu is locked
+func (c *clusterImpl) removeConnLocked(conn *Conn) {
 	conn.Close()
 	connPool := c.connPool[conn.addr]
 	if connPool == nil {
@@ -198,6 +197,12 @@ func (c *clusterImpl) removeConn(conn *Conn) {
 	delete(c.conns, conn)
 }
 
+func (c *clusterImpl) removeConn(conn *Conn) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.removeConnLocked(conn)
+}
+
 func (c *clusterImpl) HandleError(conn *Conn, err error, closed bool) {
 	if !closed {
 		// ignore all non-fatal errors
@@ -242,7 +247,7 @@ func (c *clusterImpl) Close() {
 		c.quit = true
 		close(c.quitWait)
 		for conn := range c.conns {
-			conn.Close()
+			c.removeConnLocked(conn)
 		}
 	})
 }

+ 2 - 5
conn.go

@@ -283,14 +283,11 @@ func (c *Conn) exec(op operation, trace Tracer) (interface{}, error) {
 	atomic.AddInt32(&c.nwait, 1)
 	atomic.StoreInt32(&call.active, 1)
 
-	if n, err := c.conn.Write(req); err != nil {
+	if _, err = c.conn.Write(req); err != nil {
 		c.conn.Close()
 		c.cluster.HandleError(c, err, true)
 		c.uniq <- id
-		if n > 0 {
-			return nil, ErrProtocol
-		}
-		return nil, ErrUnavailable
+		return nil, err
 	}
 
 	reply := <-call.resp