Fixed some issues around closed connections

Author: dancannon

CHANGELOG.md

@@ -17,7 +17,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### Changed
  - Changed `Timeout` connect argument to only configure the connection timeout.
  - Replaced `Db` with `DB` in `RunOpts` and `ExecOpts` (`Db` still works for now)
- - Made `Cursor` safe for concurrent use (`Session` was already safe).
+ - Made `Cursor` and `Session` safe for concurrent use
+ - Replaced `ErrClusterClosed` with `ErrConnectionClosed`
 
 ## Deprecated
  - Deprecated `UseOutdated` optional argument

cluster.go

@@ -326,7 +326,7 @@ func (c *Cluster) getSeeds() []Host {
 // TODO(dancannon) replace with hostpool
 func (c *Cluster) GetRandomNode() (*Node, error) {
 	if !c.IsConnected() {
-		return nil, ErrClusterClosed
+		return nil, ErrNoConnections
 	}
 	// Must copy array reference for copy on write semantics to work.
 	nodeArray := c.GetNodes()

connection.go

@@ -106,12 +106,12 @@ func (c *Connection) Query(q Query) (*Response, *Cursor, error) {
 	c.mu.Lock()
 	if c == nil {
 		c.mu.Unlock()
-		return nil, nil, nil
+		return nil, nil, ErrConnectionClosed
 	}
 	if c.conn == nil {
 		c.bad = true
 		c.mu.Unlock()
-		return nil, nil, nil
+		return nil, nil, ErrConnectionClosed
 	}
 
 	// Add token if query is a START/NOREPLY_WAIT

errors.go

@@ -18,9 +18,6 @@ var (
 	// ErrInvalidNode is returned when attempting to connect to a node which
 	// returns an invalid response.
 	ErrInvalidNode = errors.New("invalid node")
-	// ErrClusterClosed is returned when a query is executed after the connection
-	// to the cluster has been closed.
-	ErrClusterClosed = errors.New("cluster closed")
 	// ErrNoConnections is returned when there are no active connections in the
 	// clusters connection pool.
 	ErrNoConnections = errors.New("gorethink: no connections were available")

session.go

@@ -2,6 +2,7 @@ package gorethink
 
 import (
 	"crypto/tls"
+	"sync"
 	"time"
 
 	p "github.com/dancannon/gorethink/ql2"
@@ -10,8 +11,10 @@ import (
 // A Session represents a connection to a RethinkDB cluster and should be used
 // when executing queries.
 type Session struct {
-	hosts   []Host
-	opts    *ConnectOpts
+	hosts []Host
+	opts  *ConnectOpts
+
+	mu      sync.RWMutex
 	cluster *Cluster
 	closed  bool
 }
@@ -114,25 +117,32 @@ func (s *Session) Reconnect(optArgs ...CloseOpts) error {
 		return err
 	}
 
+	s.mu.Lock()
 	s.cluster, err = NewCluster(s.hosts, s.opts)
 	if err != nil {
 		return err
 	}
 
 	s.closed = false
+	s.mu.Unlock()
 
 	return nil
 }
 
 // Close closes the session
 func (s *Session) Close(optArgs ...CloseOpts) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	if s.closed {
 		return nil
 	}
 
 	if len(optArgs) >= 1 {
 		if optArgs[0].NoReplyWait {
+			s.mu.Unlock()
 			s.NoReplyWait()
+			s.mu.Lock()
 		}
 	}
 
@@ -148,12 +158,18 @@ func (s *Session) Close(optArgs ...CloseOpts) error {
 // SetMaxIdleConns sets the maximum number of connections in the idle
 // connection pool.
 func (s *Session) SetMaxIdleConns(n int) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	s.opts.MaxIdle = n
 	s.cluster.SetMaxIdleConns(n)
 }
 
 // SetMaxOpenConns sets the maximum number of open connections to the database.
 func (s *Session) SetMaxOpenConns(n int) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	s.opts.MaxOpen = n
 	s.cluster.SetMaxOpenConns(n)
 }
@@ -162,28 +178,55 @@ func (s *Session) SetMaxOpenConns(n int) {
 // processed by the server. Note that this guarantee only applies to queries
 // run on the given connection
 func (s *Session) NoReplyWait() error {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if s.closed {
+		return ErrConnectionClosed
+	}
+
 	return s.cluster.Exec(Query{
 		Type: p.Query_NOREPLY_WAIT,
 	})
 }
 
 // Use changes the default database used
 func (s *Session) Use(database string) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
 	s.opts.Database = database
 }
 
 // Query executes a ReQL query using the session to connect to the database
 func (s *Session) Query(q Query) (*Cursor, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if s.closed {
+		return nil, ErrConnectionClosed
+	}
+
 	return s.cluster.Query(q)
 }
 
 // Exec executes a ReQL query using the session to connect to the database
 func (s *Session) Exec(q Query) error {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if s.closed {
+		return ErrConnectionClosed
+	}
+
 	return s.cluster.Exec(q)
 }
 
 // SetHosts resets the hosts used when connecting to the RethinkDB cluster
 func (s *Session) SetHosts(hosts []Host) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	s.hosts = hosts
 }
 

session_test.go

@@ -58,6 +58,23 @@ func (s *RethinkSuite) TestSessionConnectError(c *test.C) {
 	c.Assert(err, test.NotNil)
 }
 
+func (s *RethinkSuite) TestSessionClose(c *test.C) {
+	session, err := Connect(ConnectOpts{
+		Address: url,
+		AuthKey: os.Getenv("RETHINKDB_AUTHKEY"),
+	})
+	c.Assert(err, test.IsNil)
+
+	_, err = Expr("Hello World").Run(session)
+	c.Assert(err, test.IsNil)
+
+	err = session.Close()
+	c.Assert(err, test.IsNil)
+
+	_, err = Expr("Hello World").Run(session)
+	c.Assert(err, test.NotNil)
+}
+
 func (s *RethinkSuite) TestSessionConnectDatabase(c *test.C) {
 	session, err := Connect(ConnectOpts{
 		Address:  url,