From 92475a9383f336e1929a9957ab2d6d51208622ad Mon Sep 17 00:00:00 2001
From: snower <sujian199@gmail.com>
Date: Sat, 13 Mar 2021 11:08:16 +0800
Subject: [PATCH 1/2] Fix the socket read/write fatal deadlock problem when a
 large number of concurrent requests or a large amount of data is written at
 the same time

---
 socket.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/socket.go b/socket.go
index 8891dd5d7..59562e1b4 100644
--- a/socket.go
+++ b/socket.go
@@ -53,6 +53,7 @@ type mongoSocket struct {
 	gotNonce      sync.Cond
 	dead          error
 	serverInfo    *mongoServerInfo
+	writeLock     sync.Mutex
 }
 
 type queryOpFlags uint32
@@ -522,7 +523,27 @@ func (socket *mongoSocket) Query(ops ...interface{}) (err error) {
 	stats.sentOps(len(ops))
 
 	socket.updateDeadline(writeDeadline)
-	_, err = socket.conn.Write(buf)
+	socket.Unlock()
+
+	bufLen := len(buf)
+	writeCount := 0
+	socket.writeLock.Lock()
+	for writeCount < bufLen {
+		n, err := socket.conn.Write(buf[writeCount:])
+		if err != nil {
+			socket.writeLock.Unlock()
+			socket.Lock()
+			if !wasWaiting && requestCount > 0 {
+				socket.updateDeadline(readDeadline)
+			}
+			socket.Unlock()
+			return err
+		}
+		writeCount += n
+	}
+	socket.writeLock.Unlock()
+
+	socket.Lock()
 	if !wasWaiting && requestCount > 0 {
 		socket.updateDeadline(readDeadline)
 	}

From 0d5ac4d904c42b0ac1c89ba30e0384d01c4c80d4 Mon Sep 17 00:00:00 2001
From: snower <sujian199@gmail.com>
Date: Sat, 13 Mar 2021 11:29:19 +0800
Subject: [PATCH 2/2] optimization socket write lock

---
 socket.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/socket.go b/socket.go
index 59562e1b4..a55d77d59 100644
--- a/socket.go
+++ b/socket.go
@@ -532,22 +532,22 @@ func (socket *mongoSocket) Query(ops ...interface{}) (err error) {
 		n, err := socket.conn.Write(buf[writeCount:])
 		if err != nil {
 			socket.writeLock.Unlock()
-			socket.Lock()
 			if !wasWaiting && requestCount > 0 {
+				socket.Lock()
 				socket.updateDeadline(readDeadline)
+				socket.Unlock()
 			}
-			socket.Unlock()
 			return err
 		}
 		writeCount += n
 	}
 	socket.writeLock.Unlock()
 
-	socket.Lock()
 	if !wasWaiting && requestCount > 0 {
+		socket.Lock()
 		socket.updateDeadline(readDeadline)
+		socket.Unlock()
 	}
-	socket.Unlock()
 	return err
 }