add Overlapped.getresultex to implement timeout in sendfile

Lib/socket.py

@@ -461,7 +461,8 @@ def _sendfile_use_send(self, file, offset=0, count=None):
     if _overlapped:
         def _sendfile_use_transmitfile(self, file, offset=0, count=None):
             self._check_sendfile_params(file, offset, count)
-            if self.gettimeout() == 0:
+            timeout = self.gettimeout()
+            if timeout == 0:
                 raise ValueError("non-blocking sockets are not supported")
             ov = _overlapped.Overlapped()
             offset_low = offset & 0xffff_ffff
@@ -469,7 +470,15 @@ def _sendfile_use_transmitfile(self, file, offset=0, count=None):
             count = count or 0
             ov.TransmitFile(self.fileno(), msvcrt.get_osfhandle(file.fileno()),
                             offset_low, offset_high, count, 0, 0)
-            sent = ov.getresult(True)
+            timeout_ms = _overlapped.INFINITE
+            if timeout is not None:
+                timeout_ms = int(timeout * 1000)
+            try:
+                sent = ov.getresultex(timeout_ms, False)
+            except WindowsError as e:
+                if e.winerror == 258:
+                    raise TimeoutError('timed out')
+                raise
             if sent > 0 and hasattr(file, 'seek'):
                 file.seek(offset + sent)
             return sent

Modules/clinic/overlapped.c.h

@@ -558,6 +558,41 @@ exit:
     return return_value;
 }
 
+PyDoc_STRVAR(_overlapped_Overlapped_getresultex__doc__,
+"getresultex($self, milliseconds, alertable, /)\n"
+"--\n"
+"\n");
+
+#define _OVERLAPPED_OVERLAPPED_GETRESULTEX_METHODDEF    \
+    {"getresultex", _PyCFunction_CAST(_overlapped_Overlapped_getresultex), METH_FASTCALL, _overlapped_Overlapped_getresultex__doc__},
+
+static PyObject *
+_overlapped_Overlapped_getresultex_impl(OverlappedObject *self,
+                                        DWORD milliseconds, BOOL alertable);
+
+static PyObject *
+_overlapped_Overlapped_getresultex(OverlappedObject *self, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    DWORD milliseconds;
+    BOOL alertable;
+
+    if (!_PyArg_CheckPositional("getresultex", nargs, 2, 2)) {
+        goto exit;
+    }
+    if (!_PyLong_UnsignedLong_Converter(args[0], &milliseconds)) {
+        goto exit;
+    }
+    alertable = PyLong_AsInt(args[1]);
+    if (alertable == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = _overlapped_Overlapped_getresultex_impl(self, milliseconds, alertable);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(_overlapped_Overlapped_ReadFile__doc__,
 "ReadFile($self, handle, size, /)\n"
 "--\n"
@@ -1239,4 +1274,4 @@ exit:
 
     return return_value;
 }
-/*[clinic end generated code: output=958cbddbcc355f47 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=1f395db21c8adb3f input=a9049054013a1b77]*/

Modules/overlapped.c

@@ -862,31 +862,12 @@ is false and the operation is still pending then an error is raised.
 [clinic start generated code]*/
 
 static PyObject *
-_overlapped_Overlapped_getresult_impl(OverlappedObject *self, BOOL wait)
+check_getresult_error(OverlappedObject *self, DWORD transferred)
-/*[clinic end generated code: output=8c9bd04d08994f6c input=aa5b03e9897ca074]*/
 {
-    DWORD transferred = 0;
+    PyObject *addr = NULL;
-    BOOL ret;
+    DWORD err = self->error;
-    DWORD err;
-    PyObject *addr;
 
-    if (self->type == TYPE_NONE) {
+    switch (self->error) {
-        PyErr_SetString(PyExc_ValueError, "operation not yet attempted");
-        return NULL;
-    }
-
-    if (self->type == TYPE_NOT_STARTED) {
-        PyErr_SetString(PyExc_ValueError, "operation failed to start");
-        return NULL;
-    }
-
-    Py_BEGIN_ALLOW_THREADS
-    ret = GetOverlappedResult(self->handle, &self->overlapped, &transferred,
-                              wait);
-    Py_END_ALLOW_THREADS
-
-    self->error = err = ret ? ERROR_SUCCESS : GetLastError();
-    switch (err) {
         case ERROR_SUCCESS:
         case ERROR_MORE_DATA:
             break;
@@ -978,6 +959,70 @@ _overlapped_Overlapped_getresult_impl(OverlappedObject *self, BOOL wait)
     }
 }
 
+static PyObject *
+_overlapped_Overlapped_getresult_impl(OverlappedObject *self, BOOL wait)
+/*[clinic end generated code: output=8c9bd04d08994f6c input=aa5b03e9897ca074]*/
+{
+    DWORD transferred = 0;
+    BOOL ret;
+
+    if (self->type == TYPE_NONE) {
+        PyErr_SetString(PyExc_ValueError, "operation not yet attempted");
+        return NULL;
+    }
+
+    if (self->type == TYPE_NOT_STARTED) {
+        PyErr_SetString(PyExc_ValueError, "operation failed to start");
+        return NULL;
+    }
+
+    Py_BEGIN_ALLOW_THREADS
+    ret = GetOverlappedResult(self->handle, &self->overlapped, &transferred,
+                              wait);
+    Py_END_ALLOW_THREADS
+
+    self->error = ret ? ERROR_SUCCESS : GetLastError();
+
+    return check_getresult_error(self, transferred);
+}
+
+/*[clinic input]
+_overlapped.Overlapped.getresultex
+
+    milliseconds: DWORD
+    alertable: BOOL
+    /
+
+[clinic start generated code]*/
+
+static PyObject *
+_overlapped_Overlapped_getresultex_impl(OverlappedObject *self,
+                                        DWORD milliseconds, BOOL alertable)
+/*[clinic end generated code: output=ce0eb6ffb9618e54 input=ef4f4cab49ac1d80]*/
+{
+    DWORD transferred = 0;
+    BOOL ret;
+
+    if (self->type == TYPE_NONE) {
+        PyErr_SetString(PyExc_ValueError, "operation not yet attempted");
+        return NULL;
+    }
+
+    if (self->type == TYPE_NOT_STARTED) {
+        PyErr_SetString(PyExc_ValueError, "operation failed to start");
+        return NULL;
+    }
+
+    Py_BEGIN_ALLOW_THREADS
+    ret = GetOverlappedResultEx(self->handle, &self->overlapped, &transferred,
+                                milliseconds, alertable);
+    Py_END_ALLOW_THREADS
+
+    self->error = ret ? ERROR_SUCCESS : GetLastError();
+
+    return check_getresult_error(self, transferred);
+}
+
 static PyObject *
 do_ReadFile(OverlappedObject *self, HANDLE handle,
             char *bufstart, DWORD buflen)
@@ -1927,6 +1972,7 @@ _overlapped_Overlapped_WSARecvFromInto_impl(OverlappedObject *self,
 
 static PyMethodDef Overlapped_methods[] = {
     _OVERLAPPED_OVERLAPPED_GETRESULT_METHODDEF
+    _OVERLAPPED_OVERLAPPED_GETRESULTEX_METHODDEF
     _OVERLAPPED_OVERLAPPED_CANCEL_METHODDEF
     _OVERLAPPED_OVERLAPPED_READFILE_METHODDEF
     _OVERLAPPED_OVERLAPPED_READFILEINTO_METHODDEF