Use poll instead of timeouts when waiting for response

This commit is contained in:
Juhani Krekelä 2021-10-17 18:58:58 +03:00
parent 5f0a735a1d
commit 0ac0006a63
1 changed files with 20 additions and 18 deletions

View File

@ -289,23 +289,26 @@ def proxy(sock, host):
# Keep sending request body, if any, until we get a response from remote # Keep sending request body, if any, until we get a response from remote
poll = select.poll() poll = select.poll()
poll.register(remote_sock, select.POLLIN) poll.register(remote_sock, select.POLLIN)
sock.settimeout(1) poll.register(sock, select.POLLIN)
while len(poll.poll(1000)) == 0: no_response = True
try: while no_response:
data = sock.recv(1024) for fd, _ in poll.poll():
except ConnectionResetError: if fd == remote_sock.fileno():
break no_response = False
except socket.timeout: break
continue else:
if data == b'': break try:
# Save the part we've sent already in case we need to re-send request data = sock.recv(1024)
request_data.extend(data) except ConnectionResetError:
try: break
remote_sock.sendall(data) if data == b'': break
except (ConnectionResetError, BrokenPipeError): # Save the part we've sent already in case we need to re-send request
print('Remote hung up', file=sys.stderr) request_data.extend(data)
return try:
sock.settimeout(None) remote_sock.sendall(data)
except (ConnectionResetError, BrokenPipeError):
print('Remote hung up', file=sys.stderr)
return
# Get response headers # Get response headers
remote_sock.settimeout(10) remote_sock.settimeout(10)
@ -382,7 +385,6 @@ def proxy(sock, host):
print('', file=sys.stderr) print('', file=sys.stderr)
sock.settimeout(60) sock.settimeout(60)
remote_sock.settimeout(60) remote_sock.settimeout(60)
poll.register(sock, select.POLLIN)
last_transfer = time.monotonic() last_transfer = time.monotonic()
while True: while True:
events = poll.poll(60_000) events = poll.poll(60_000)