@@ -358,26 +358,29 @@ async def sock_recv(self, sock, n):
358
358
"""
359
359
if self ._debug and sock .gettimeout () != 0 :
360
360
raise ValueError ("the socket must be non-blocking" )
361
+ try :
362
+ return sock .recv (n )
363
+ except (BlockingIOError , InterruptedError ):
364
+ pass
361
365
fut = self .create_future ()
362
- self ._sock_recv (fut , None , sock , n )
366
+ fd = sock .fileno ()
367
+ self .add_reader (fd , self ._sock_recv , fut , sock , n )
368
+ fut .add_done_callback (
369
+ functools .partial (self ._sock_read_done , fd ))
363
370
return await fut
364
371
365
- def _sock_recv (self , fut , registered_fd , sock , n ):
372
+ def _sock_read_done (self , fd , fut ):
373
+ self .remove_reader (fd )
374
+
375
+ def _sock_recv (self , fut , sock , n ):
366
376
# _sock_recv() can add itself as an I/O callback if the operation can't
367
377
# be done immediately. Don't use it directly, call sock_recv().
368
- if registered_fd is not None :
369
- # Remove the callback early. It should be rare that the
370
- # selector says the fd is ready but the call still returns
371
- # EAGAIN, and I am willing to take a hit in that case in
372
- # order to simplify the common case.
373
- self .remove_reader (registered_fd )
374
- if fut .cancelled ():
378
+ if fut .done ():
375
379
return
376
380
try :
377
381
data = sock .recv (n )
378
382
except (BlockingIOError , InterruptedError ):
379
- fd = sock .fileno ()
380
- self .add_reader (fd , self ._sock_recv , fut , fd , sock , n )
383
+ return # try again next time
381
384
except Exception as exc :
382
385
fut .set_exception (exc )
383
386
else :
@@ -391,27 +394,27 @@ async def sock_recv_into(self, sock, buf):
391
394
"""
392
395
if self ._debug and sock .gettimeout () != 0 :
393
396
raise ValueError ("the socket must be non-blocking" )
397
+ try :
398
+ return sock .recv_into (buf )
399
+ except (BlockingIOError , InterruptedError ):
400
+ pass
394
401
fut = self .create_future ()
395
- self ._sock_recv_into (fut , None , sock , buf )
402
+ fd = sock .fileno ()
403
+ self .add_reader (fd , self ._sock_recv_into , fut , sock , buf )
404
+ fut .add_done_callback (
405
+ functools .partial (self ._sock_read_done , fd ))
396
406
return await fut
397
407
398
- def _sock_recv_into (self , fut , registered_fd , sock , buf ):
408
+ def _sock_recv_into (self , fut , sock , buf ):
399
409
# _sock_recv_into() can add itself as an I/O callback if the operation
400
410
# can't be done immediately. Don't use it directly, call
401
411
# sock_recv_into().
402
- if registered_fd is not None :
403
- # Remove the callback early. It should be rare that the
404
- # selector says the FD is ready but the call still returns
405
- # EAGAIN, and I am willing to take a hit in that case in
406
- # order to simplify the common case.
407
- self .remove_reader (registered_fd )
408
- if fut .cancelled ():
412
+ if fut .done ():
409
413
return
410
414
try :
411
415
nbytes = sock .recv_into (buf )
412
416
except (BlockingIOError , InterruptedError ):
413
- fd = sock .fileno ()
414
- self .add_reader (fd , self ._sock_recv_into , fut , fd , sock , buf )
417
+ return # try again next time
415
418
except Exception as exc :
416
419
fut .set_exception (exc )
417
420
else :
@@ -428,34 +431,40 @@ async def sock_sendall(self, sock, data):
428
431
"""
429
432
if self ._debug and sock .gettimeout () != 0 :
430
433
raise ValueError ("the socket must be non-blocking" )
431
- fut = self .create_future ()
432
- if data :
433
- self ._sock_sendall (fut , None , sock , data )
434
+ try :
435
+ n = sock .send (data )
436
+ except (BlockingIOError , InterruptedError ):
437
+ n = 0
438
+
439
+ if n == len (data ):
440
+ # all data sent
441
+ return
434
442
else :
435
- fut .set_result (None )
443
+ data = bytearray (memoryview (data )[n :])
444
+
445
+ fut = self .create_future ()
446
+ fd = sock .fileno ()
447
+ fut .add_done_callback (
448
+ functools .partial (self ._sock_write_done , fd ))
449
+ self .add_writer (fd , self ._sock_sendall , fut , sock , data )
436
450
return await fut
437
451
438
- def _sock_sendall (self , fut , registered_fd , sock , data ):
439
- if registered_fd is not None :
440
- self .remove_writer (registered_fd )
441
- if fut .cancelled ():
452
+ def _sock_sendall (self , fut , sock , data ):
453
+ if fut .done ():
454
+ # Future cancellation can be scheduled on previous loop iteration
442
455
return
443
-
444
456
try :
445
457
n = sock .send (data )
446
458
except (BlockingIOError , InterruptedError ):
447
- n = 0
459
+ return
448
460
except Exception as exc :
449
461
fut .set_exception (exc )
450
462
return
451
463
452
464
if n == len (data ):
453
465
fut .set_result (None )
454
466
else :
455
- if n :
456
- data = data [n :]
457
- fd = sock .fileno ()
458
- self .add_writer (fd , self ._sock_sendall , fut , fd , sock , data )
467
+ del data [:n ]
459
468
460
469
async def sock_connect (self , sock , address ):
461
470
"""Connect to a remote socket at address.
@@ -484,18 +493,18 @@ def _sock_connect(self, fut, sock, address):
484
493
# becomes writable to be notified when the connection succeed or
485
494
# fails.
486
495
fut .add_done_callback (
487
- functools .partial (self ._sock_connect_done , fd ))
496
+ functools .partial (self ._sock_write_done , fd ))
488
497
self .add_writer (fd , self ._sock_connect_cb , fut , sock , address )
489
498
except Exception as exc :
490
499
fut .set_exception (exc )
491
500
else :
492
501
fut .set_result (None )
493
502
494
- def _sock_connect_done (self , fd , fut ):
503
+ def _sock_write_done (self , fd , fut ):
495
504
self .remove_writer (fd )
496
505
497
506
def _sock_connect_cb (self , fut , sock , address ):
498
- if fut .cancelled ():
507
+ if fut .done ():
499
508
return
500
509
501
510
try :
@@ -529,7 +538,7 @@ def _sock_accept(self, fut, registered, sock):
529
538
fd = sock .fileno ()
530
539
if registered :
531
540
self .remove_reader (fd )
532
- if fut .cancelled ():
541
+ if fut .done ():
533
542
return
534
543
try :
535
544
conn , address = sock .accept ()
0 commit comments