gremlin-go: fix deadlock on (*DriverRemoteConnection).Close

gremlin-go/driver/protocol.go

@@ -161,15 +161,21 @@ func (protocol *gremlinServerWSProtocol) write(request *request) error {
 	return protocol.transporter.Write(bytes)
 }
 
-func (protocol *gremlinServerWSProtocol) close() (err error) {
+func (protocol *gremlinServerWSProtocol) close() error {
 	protocol.mutex.Lock()
-	if !protocol.closed {
-		err = protocol.transporter.Close()
-		protocol.closed = true
+
+	if protocol.closed {
+		protocol.mutex.Unlock()
+		return nil
 	}
+
+	err := protocol.transporter.Close()
+	protocol.closed = true
 	protocol.mutex.Unlock()
+
 	protocol.wg.Wait()
-	return
+
+	return err
 }
 
 func newGremlinServerWSProtocol(handler *logHandler, transporterType TransporterType, url string, connSettings *connectionSettings, results *synchronizedMap,

gremlin-go/driver/protocol_test.go

@@ -20,12 +20,15 @@ under the License.
 package gremlingo
 
 import (
+	"sync"
+	"testing"
+	"time"
+
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/text/language"
-	"testing"
 )
 
-func Test(t *testing.T) {
+func TestProtocol(t *testing.T) {
 	t.Run("Test protocol connect error.", func(t *testing.T) {
 		connSettings := newDefaultConnectionSettings()
 		connSettings.authInfo, connSettings.tlsConfig = nil, nil
@@ -37,4 +40,30 @@ func Test(t *testing.T) {
 		assert.NotNil(t, err)
 		assert.Nil(t, protocol)
 	})
+
+	// protocol.closed is only modified by protocol.close(). If it is true
+	// it means that protocol.wg.Wait() has already been called, so it
+	// should not be called again.
+	t.Run("Test protocol close when closed", func(t *testing.T) {
+		wg := &sync.WaitGroup{}
+		protocol := &gremlinServerWSProtocol{
+			closed: true,
+			mutex:  sync.Mutex{},
+			wg:     wg,
+		}
+		wg.Add(1)
+
+		done := make(chan bool)
+
+		go func() {
+			protocol.close()
+			done <- true
+		}()
+
+		select {
+		case <-time.After(1 * time.Second):
+			t.Fatal("timeout")
+		case <-done:
+		}
+	})
 }