2
0

Fix Windows non-blocking I/O for CopyFrom

Created based on discussion here: https://github.com/jackc/pgx/pull/1525#pullrequestreview-1344511991

Fixes https://github.com/jackc/pgx/issues/1552
This commit is contained in:
Dmitry K
2023-03-18 21:42:55 +03:00
committed by Jack Christensen
parent 9ae852eb58
commit 8b5e8d9d89
4 changed files with 63 additions and 16 deletions
+3
View File
@@ -96,6 +96,9 @@ type NetConn struct {
writeDeadlineLock sync.Mutex
writeDeadline time.Time
// Indicates that underlying socket connection mode set to be non-blocking
isNonBlocking bool
}
func NewNetConn(conn net.Conn, fakeNonBlockingIO bool) *NetConn {
+5
View File
@@ -79,3 +79,8 @@ func (c *NetConn) realNonblockingRead(b []byte) (n int, err error) {
return n, nil
}
func (c *NetConn) SetBlockingMode(blocking bool) error {
// for UNIX do nothing
return nil
}
@@ -43,10 +43,12 @@ func setSockMode(fd uintptr, mode sockMode) error {
func (c *NetConn) realNonblockingWrite(b []byte) (n int, err error) {
if c.nonblockWriteFunc == nil {
c.nonblockWriteFunc = func(fd uintptr) (done bool) {
// Make sock non-blocking
if err := setSockMode(fd, sockModeNonBlocking); err != nil {
c.nonblockWriteErr = err
return true
if !c.isNonBlocking {
// Make sock non-blocking
if err := setSockMode(fd, sockModeNonBlocking); err != nil {
c.nonblockWriteErr = err
return true
}
}
var written uint32
@@ -56,10 +58,12 @@ func (c *NetConn) realNonblockingWrite(b []byte) (n int, err error) {
c.nonblockWriteErr = syscall.WSASend(syscall.Handle(fd), &buf, 1, &written, 0, nil, nil)
c.nonblockWriteN = int(written)
// Make sock blocking again
if err := setSockMode(fd, sockModeBlocking); err != nil {
c.nonblockWriteErr = err
return true
if !c.isNonBlocking {
// Make sock blocking again
if err := setSockMode(fd, sockModeBlocking); err != nil {
c.nonblockWriteErr = err
return true
}
}
return true
@@ -94,10 +98,12 @@ func (c *NetConn) realNonblockingWrite(b []byte) (n int, err error) {
func (c *NetConn) realNonblockingRead(b []byte) (n int, err error) {
if c.nonblockReadFunc == nil {
c.nonblockReadFunc = func(fd uintptr) (done bool) {
// Make sock non-blocking
if err := setSockMode(fd, sockModeNonBlocking); err != nil {
c.nonblockReadErr = err
return true
if !c.isNonBlocking {
// Make sock non-blocking
if err := setSockMode(fd, sockModeNonBlocking); err != nil {
c.nonblockReadErr = err
return true
}
}
var read uint32
@@ -108,10 +114,12 @@ func (c *NetConn) realNonblockingRead(b []byte) (n int, err error) {
c.nonblockReadErr = syscall.WSARecv(syscall.Handle(fd), &buf, 1, &read, &flags, nil, nil)
c.nonblockReadN = int(read)
// Make sock blocking again
if err := setSockMode(fd, sockModeBlocking); err != nil {
c.nonblockReadErr = err
return true
if !c.isNonBlocking {
// Make sock blocking again
if err := setSockMode(fd, sockModeBlocking); err != nil {
c.nonblockReadErr = err
return true
}
}
return true
@@ -147,3 +155,22 @@ func (c *NetConn) realNonblockingRead(b []byte) (n int, err error) {
return n, nil
}
func (c *NetConn) SetBlockingMode(blocking bool) error {
mode := sockModeNonBlocking
if blocking {
mode = sockModeBlocking
}
var err error
c.rawConn.Control(func(fd uintptr) {
err = setSockMode(fd, mode)
})
if err == nil {
c.isNonBlocking = !blocking
}
return err
}