improve session close synchronization (northwesternmutual#18)

* improve session close synchronization

* CR

client.go

@@ -21,6 +21,7 @@
 package grammes
 
 import (
+	"golang.org/x/sync/errgroup"
 	"sync"
 	"time"
 
@@ -72,6 +73,7 @@ type Client struct {
 	// requestTimeout is used for time-outing requests that a response is not received for
 	requestTimeout   time.Duration
 	requestSemaphore *semaphore.Weighted
+	commRoutines 	 errgroup.Group
 }
 
 // setupClient default values some fields in the client.

connection.go

@@ -22,7 +22,6 @@ package grammes
 
 import (
 	"errors"
-
 	"github.com/northwesternmutual/grammes/gremconnect"
 	"github.com/northwesternmutual/grammes/gremerror"
 )
@@ -43,9 +42,18 @@ func (c *Client) launchConnection() error {
 	quit := c.conn.GetQuit()
 
 	// Launch processes to keep track of connection & data
-	go c.writeWorker(c.err, quit) // Initiates message writing to the Gremlin-server
-	go c.readWorker(c.err, quit)  // Initiates message reading from the Gremlin-server
-	go c.conn.Ping(c.err)         // Manages pinging and connection to the Gremlin-server
+	c.commRoutines.Go(func() error {
+		c.writeWorker(c.err, quit) // Initiates message writing to the Gremlin-server
+		return nil
+	})
+	c.commRoutines.Go(func() error {
+		c.readWorker(c.err, quit) // Initiates message reading from the Gremlin-server
+		return nil
+	})
+	c.commRoutines.Go(func() error {
+		c.conn.Ping(c.err) // Manages pinging and connection to the Gremlin-server
+		return nil
+	})
 
 	return nil
 }
@@ -55,6 +63,8 @@ func (c *Client) Close() {
 	if c.conn != nil {
 		c.conn.Close()
 	}
+
+	c.commRoutines.Wait()
 }
 
 // IsConnected returns if the client currently

gremconnect/websocket.go

@@ -145,8 +145,8 @@ func (ws *WebSocket) Read() (msg []byte, err error) {
 func (ws *WebSocket) Close() error {
 	defer func() {
 		close(ws.Quit) // close the channel to notify our pinger.
-		ws.conn.Close()
 		ws.disposed = true
+		ws.conn.Close()
 	}()
 
 	// Send the server the message that we've closed

manager/model.go

@@ -182,8 +182,8 @@ type SessionQuerier interface {
 	NewSession() Session
 	NewNoopSession() Session
 	GetSession(uuid.UUID) Session
-	WithSession(Session, func(Session) error) error
-	WithNewSession(f func(Session) error) error
+	WithSession(Session, func(Session) error, bool) error
+	WithNewSession(func(Session) error, bool) error
 }
 
 // VertexQuerier handles the vertices on the graph.

manager/session.go

@@ -49,16 +49,18 @@ func (s *sessionManager) GetSession(sessionId uuid.UUID) Session {
 	}
 }
 
-func (s *sessionManager) WithNewSession(f func(Session) error) error {
-	return s.WithSession(s.NewSession(), f)
+func (s *sessionManager) WithNewSession(f func(Session) error, closeOnDone bool) error {
+	return s.WithSession(s.NewSession(), f, closeOnDone)
 }
 
-func (s *sessionManager) WithSession(ss Session, f func(Session) error) error {
+func (s *sessionManager) WithSession(ss Session, f func(Session) error, closeOnDone bool) error {
 	var err error
 	defer func() {
-		err2 := ss.Close()
-		if err == nil { // Capture close error if everything else was ok
-			err = err2
+		if closeOnDone {
+			err2 := ss.Close()
+			if err == nil { // Capture close error if everything else was ok
+				err = err2
+			}
 		}
 	}()
 

response.go

@@ -44,7 +44,10 @@ func (c *Client) readWorker(errs chan error, quit chan struct{}) {
 		// attempt to read from the connection
 		// and store the message back into a variable.
 		if msg, err = c.conn.Read(); err != nil {
-			errs <- err
+			if !c.conn.IsDisposed() { // When disposing a connection, gorilla will return an error
+				errs <- err
+			}
+
 			c.broken = true
 			break
 		}