diff --git a/socket.go b/socket.go index 8891dd5d7..a55d77d59 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,11 +523,31 @@ 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() + if !wasWaiting && requestCount > 0 { + socket.Lock() + socket.updateDeadline(readDeadline) + socket.Unlock() + } + return err + } + writeCount += n + } + socket.writeLock.Unlock() + if !wasWaiting && requestCount > 0 { + socket.Lock() socket.updateDeadline(readDeadline) + socket.Unlock() } - socket.Unlock() return err }