Come posso aggiungere un timeout della connessione con asyncio?

Question

Voglio collegarmi a un elenco di molti siti diversi molto velocemente. Sto usando asyncio per farlo in modo asincrono e ora voglio aggiungere un timeout per quando le connessioni devono essere ignorate se impiegano troppo tempo a rispondere.

Come implementarlo?

import ssl 
import asyncio 
from contextlib import suppress 
from concurrent.futures import ThreadPoolExecutor 
import time 


@asyncio.coroutine 
def run(): 
    while True: 
     host = yield from q.get() 
     if not host: 
      break 

     with suppress(ssl.CertificateError): 
      reader, writer = yield from asyncio.open_connection(host[1], 443, ssl=True) #timout option? 
      reader.close() 
      writer.close() 


@asyncio.coroutine 
def load_q(): 
    # only 3 entries for debugging reasons 
    for host in [[1, 'python.org'], [2, 'qq.com'], [3, 'google.com']]: 
     yield from q.put(host) 
    for _ in range(NUM): 
     q.put(None) 


if __name__ == "__main__": 
    NUM = 1000 
    q = asyncio.Queue() 

    loop = asyncio.get_event_loop() 
    loop.set_default_executor(ThreadPoolExecutor(NUM)) 

    start = time.time() 
    coros = [asyncio.async(run()) for i in range(NUM)] 
    loop.run_until_complete(load_q()) 
    loop.run_until_complete(asyncio.wait(coros)) 
    end = time.time() 
    print(end-start) 

(Su un sidenote: Ha qualcuno un'idea di come ottimizzare questo?)

Answer 1

Puoi avvolgere la chiamata a open_connection in asyncio.wait_for, che permette di specificare un timeout:

with suppress(ssl.CertificateError): 
     fut = asyncio.open_connection(host[1], 443, ssl=True) 
     try: 
      # Wait for 3 seconds, then raise TimeoutError 
      reader, writer = yield from asyncio.wait_for(fut, timeout=3) 
     except asyncio.TimeoutError: 
      print("Timeout, skipping {}".format(host[1])) 
      continue 

Si noti che quando viene sollevato TimeoutError, viene annullata anche la coroutine open_connection. Se non vuoi che venga cancellato (anche se penso che tu sia do vuoi che venga annullato in questo caso), hai terminato la chiamata in asyncio.shield.