commit c93661f5d9133256e7fc31dd2db453fc2da9dd73
author Roi Martin <jroi.martin@gmail.com> Thu Jul 14 14:39:06 2022 +0200
committer Roi Martin <jroi.martin@gmail.com> Thu Jul 14 14:39:06 2022 +0200
tree 3cbd2b21d5dfb6c3d237a9b9473ca575e8f7b5e3
parent e92d53fa0767fafa4e74f12ca84542d77d7d2976

gremlin-go: refactor writeLoop and readLoop using defer

gremlin-go/driver/gorillaTransporter.go

diff --git a/gremlin-go/driver/gorillaTransporter.go b/gremlin-go/driver/gorillaTransporter.go
index cc033c8..2a209b4 100644
--- a/gremlin-go/driver/gorillaTransporter.go
+++ b/gremlin-go/driver/gorillaTransporter.go
@@ -142,16 +142,16 @@
 }
 
 func (transporter *gorillaTransporter) writeLoop() {
+	defer transporter.wg.Done()
+
 	ticker := time.NewTicker(transporter.connSettings.keepAliveInterval)
-	defer func() {
-		ticker.Stop()
-	}()
+	defer ticker.Stop()
+
 	for {
 		select {
 		case message, ok := <-transporter.writeChannel:
 			if !ok {
 				// Channel was closed, we can disconnect and exit.
-				transporter.wg.Done()
 				return
 			}
 
@@ -159,7 +159,6 @@
 			err := transporter.connection.SetWriteDeadline(time.Now().Add(transporter.connSettings.writeDeadline))
 			if err != nil {
 				transporter.logHandler.logf(Error, failedToSetWriteDeadline, err.Error())
-				transporter.wg.Done()
 				return
 			}
 
@@ -167,7 +166,6 @@
 			err = transporter.connection.WriteMessage(websocket.BinaryMessage, message)
 			if err != nil {
 				transporter.logHandler.logf(Error, failedToWriteMessage, "BinaryMessage", err.Error())
-				transporter.wg.Done()
 				return
 			}
 		case <-ticker.C:
@@ -175,7 +173,6 @@
 			err := transporter.connection.SetWriteDeadline(time.Now().Add(transporter.connSettings.keepAliveInterval))
 			if err != nil {
 				transporter.logHandler.logf(Error, failedToSetWriteDeadline, err.Error())
-				transporter.wg.Done()
 				return
 			}
 
@@ -183,7 +180,6 @@
 			err = transporter.connection.WriteMessage(websocket.PingMessage, nil)
 			if err != nil {
 				transporter.logHandler.logf(Error, failedToWriteMessage, "PingMessage", err.Error())
-				transporter.wg.Done()
 				return
 			}
 		}

gremlin-go/driver/protocol.go

diff --git a/gremlin-go/driver/protocol.go b/gremlin-go/driver/protocol.go
index b23f5d7..d07f231 100644
--- a/gremlin-go/driver/protocol.go
+++ b/gremlin-go/driver/protocol.go
@@ -52,13 +52,14 @@
 }
 
 func (protocol *gremlinServerWSProtocol) readLoop(resultSets *synchronizedMap, errorCallback func()) {
+	defer protocol.wg.Done()
+
 	for {
 		// Read from transport layer. If the channel is closed, this will error out and exit.
 		msg, err := protocol.transporter.Read()
 		protocol.mutex.Lock()
 		if protocol.closed {
 			protocol.mutex.Unlock()
-			protocol.wg.Done()
 			return
 		}
 		protocol.mutex.Unlock()
@@ -67,7 +68,6 @@
 			_ = protocol.transporter.Close()
 			protocol.logHandler.logf(Error, readLoopError, err.Error())
 			readErrorHandler(resultSets, errorCallback, err, protocol.logHandler)
-			protocol.wg.Done()
 			return
 		}
 
@@ -76,14 +76,12 @@
 		if err != nil {
 			protocol.logHandler.logf(Error, logErrorGeneric, "gremlinServerWSProtocol.readLoop()", err.Error())
 			readErrorHandler(resultSets, errorCallback, err, protocol.logHandler)
-			protocol.wg.Done()
 			return
 		}
 
 		err = protocol.responseHandler(resultSets, resp)
 		if err != nil {
 			readErrorHandler(resultSets, errorCallback, err, protocol.logHandler)
-			protocol.wg.Done()
 			return
 		}
 	}