Async rabbitmq consumer with socket connect_ex - python

I want to write a python program that gets ip and tcp port from a rabbitmq server and scans to check if the port is open, as these scans sometimes come in bulk (maybe 100 port, ip pairs are added to the queue at a time) I need to do the scans asynchronously to get all the results in time, and even if I lower the timeout to 1 second, 30 closed ports will hold the scan for 30 seconds each time!
I tried asyncio and aio_pika to reach my goal but still the scans are being performed synchronously.
import asyncio
import aio_pika
import socket
async def tcp_check(host, port):
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
await asyncio.sleep(1)
result = sock.connect_ex((host,port))
print (str(result))
async def main(loop):
connection = await aio_pika.connect_robust("amqp://user:password#192.168.1.100/")
async with connection:
queue_name = "tcp_scans"
channel = await connection.channel()
queue = await channel.declare_queue(queue_name, auto_delete=False, durable=True)
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
context = message.body.decode("utf-8").split(',')
await tcp_check(context[0], int(context[1]))
if __name__ == "__main__":
loop = asyncio.get_event_loop()
loop.run_until_complete(main(loop))
loop.close()
UPDATE:
I used asyncio.open_connection too:
async def tcp_check(host, port):
con = asyncio.open_connection(host, port, loop=loop)
try:
await asyncio.wait_for(con, timeout=1)
print("{}:{} Connected".format(host, port))
except asyncio.TimeoutError:
print ("{}:{} Closeed".format(host, port))
Still it takes each item from the and test one by one...

Calling of synchronous long running functions inside asynchronous coroutines should be avoided. I'd suggest to use asyncio alternative to connect_ex, e.g.:
try:
await asyncio.open_connection(host, port)
except Exception as e:
print(e)
In order to execute some coroutines simultaneously "on the fly" you can use create_task that "wrap the coroutine into a Task and schedule its execution" as it is written in doc. And after this, coroutine will be executed soon, e.g. after next await or async for iteration, when control flow returns to the event-loop.
create_task return Task object which you can add to list and wait them all finished using asyncio.gather with flag return_exceptions=True.
But in your case i think it will be sufficient replace await tcp_check() to create_task(tcp_check()) and use gather at the end of your main() to guarantee all coro is finished.
...
asyncio.create_task(tcp_check(context[0], int(context[1])))
...

Related

Keep indefinite connections open on TCP server with asyncio streams?

I'm trying to understand how to use asyncio streams for multiple connections that will keep sending messages until a predefined condition or a socket timeout. Looking at Python docs, they provide the following example for a TCP server based on asyncio streams:
import asyncio
async def handle_echo(reader, writer):
data = await reader.read(100)
message = data.decode()
addr = writer.get_extra_info('peername')
print(f"Received {message!r} from {addr!r}")
print(f"Send: {message!r}")
writer.write(data)
await writer.drain()
print("Close the connection")
writer.close()
async def main():
server = await asyncio.start_server(
handle_echo, '127.0.0.1', 8888)
addrs = ', '.join(str(sock.getsockname()) for sock in server.sockets)
print(f'Serving on {addrs}')
async with server:
await server.serve_forever()
asyncio.run(main())
What I'm trying to do is more complex and it looks more like so (a lot of it is pseudocode, written in capital letters or with implementation omitted):
import asyncio
async def io_control(queue):
while true:
...
# do I/O control in this function ...
async def data_processing(queue):
while true:
...
# perform data handling
async def handle_data(reader, writer):
data = await reader.read()
message = data.decode()
addr = writer.get_extra_info('peername')
print(f"Received {message!r} from {addr!r}")
#do stuff with a queue - pass messages to other two async functions as needed
#keep open until something happens
if(ERROR or SOCKET_TIMEOUT):
writer.close()
async def server(queue):
server = await asyncio.start_server(
handle_data, '127.0.0.1', 8888)
addrs = ', '.join(str(sock.getsockname()) for sock in server.sockets)
print(f'Serving on {addrs}')
async with server:
await server.serve_forever()
async def main():
queue_io = asyncio.Queue()
queue_data = asyncio.Queue()
asyncio.run(server(queue_data))
asyncio.run(data_handling(queue_data))
asyncio.run(io_control(queue_io))
asyncio.run(main())
Does this look feasible? I'm not used to working with co-routines (I'm coming from more of a multi-threading paradigm), so I'm not sure if what I'm doing is right or if I have to explicitly include yields or do any extra stuff.
If I understand correctly, you just need the TCP server to be able to handle multiple concurrent connections. The start_server function should already give you everything you need.
The first parameter client_connected_cb is a coroutine function called whenever a client establishes a connection. If you introduce a loop into that function (in your example code handle_data), you can keep the connection open until some criterion is met. What conditions exactly should lead to closing the connection is up to you, and the implementation details will obviously depend on that. The simplest approach I can imagine is something like this:
import asyncio
import logging
log = logging.getLogger(__name__)
async def handle_data(reader, writer):
while True:
data = (await reader.readline()).decode().strip()
if not data:
log.debug("client disconnected")
break
response = await your_data_processing_function(data)
writer.write(response.encode())
await writer.drain()
...
async def main():
server = await asyncio.start_server(handle_data, '127.0.0.1', 8888)
async with server:
await server.serve_forever()
if __name__ == '__main__':
asyncio.run(main())
There is theoretically no limit for the number of concurrent connections.
If your client_connected_cb is a coroutine function, each new connection will schedule a new task for the event loop. That is where the concurrency comes from. The magic then happens at the point of awaiting new data from the client; that is where the event loop can switch execution to another coroutine. All this happens behind the scenes, so to speak.
If you want to introduce a timeout, you could wrap the awaitable readline coroutine in a wait_for for example and then catch the TimeoutError exiting the loop.
Hope this helps.

Python Asyncio & Websocket how to avoid global variable?

I have a python socket server using asyncio and websockets. When the websocket is active 100+ devices will connect and hold their connection waiting for commands/messages.
There are two threads the first thread accepts connections and adds their details to a global variable then waits for messages from the device:
async def thread1(websocket, path):
client_address = await websocket.recv()
CONNECTIONS[client_address] = websocket
async for message in websocket:
... do something with message
start_server = websockets.serve(thread1, host, port)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.ensure_future(thread2())
asyncio.get_event_loop().run_forever()
The second thread processes some user data and once it needs to send a command it accesses a global variable to get the websocket info:
thread2()
...some data processing
soc = CONNECTIONS[ipaddress]
await soc.send("some message")
My question: What's the best way to allow another thread to send messages?
I can keep the global variable safe using thread locking and a function made only to process that data, however global variables aren't ideal. I cannot send information between threads since thread1 is stuck waiting to receive messages.
The first thing I would like to say is the incorrect use of the term thread. You use asyncio and here the concept is used - coroutine (coroutine is wrapped into a asyncio task). How it differs from threads can be found, for example, here.
The websockets server spawns a new task for each incoming connection (there are the same number of connections and spawned tasks). I don't see anything wrong with the global object, at least in a small script. However, below I gave an example where I placed this in a separate class.
Also, in this case, special synchronization between coroutines is not required, since they are implemented through cooperative multitasking (in fact, all are executed in one thread, transferring control at certain points.)
Here is a simple example in which the server stores a dictionary of incoming connections and starts a task that every 2 seconds, notifies all clients and sends them the current time. The server also prints confirmation from clients to the console.
# ws_server.py
import asyncio
import websockets
import datetime
class Server:
def __init__(self, host, port):
self.host = host
self.port = port
self.connections = {}
self.is_active = False
self.server = None
async def start(self):
self.is_active = True
self.server = await websockets.serve(self.handler, self.host, self.port)
asyncio.create_task(self.periodic_notifier())
async def stop(self):
self.is_active = False
self.server.close()
await self.wait_closed()
async def wait_closed(self):
await self.server.wait_closed()
async def handler(self, websocket, path):
self.connections[websocket.remote_address] = websocket
try:
async for message in websocket:
print(message)
except ConnectionClosedError as e:
pass
del self.connections[websocket.remote_address]
print(f"Connection {websocket.remote_address} is closed")
async def periodic_notifier(self):
while self.is_active:
await asyncio.gather(
*[ws.send(f"Hello time {datetime.datetime.now()}") for ws in self.connections.values()],
return_exceptions=True)
await asyncio.sleep(2)
async def main():
server = Server("localhost", 8080)
await server.start()
await server.wait_closed()
asyncio.run(main())
# ws_client.py
import asyncio
import websockets
async def client():
uri = "ws://localhost:8080"
async with websockets.connect(uri) as websocket:
async for message in websocket:
print(message)
await websocket.send(f"ACK {message}")
asyncio.run(client())

asyncio tcp socket: How to cancel asyncio.sleep() when socket is closed?

I am currently implementing the TCP socket protocol. The protocol requires sending heartbeat messages every five minutes. I am implementing a protocol using asyncio in Python. The source code below is a program that connects to localhost:8889, sends hello, and disconnects the socket after 1 second. In this case, the connection is disconnected after one second (if this actually happens, the network is down or the server is disconnected). The problem is that the send_heartbeat function waits 5 minutes without knowing that the socket is down. I would like to cancel the coroutine immediately instead of waiting 5 minutes when the socket is disconnected. What's the best way to do it?
import asyncio
async def run(host: str, port: int):
while True:
try:
reader, writer = await asyncio.open_connection(host, port)
except OSError as e:
print('connection failed:', e)
await asyncio.sleep(0.5)
continue
await asyncio.wait([
handle_stream(reader, writer),
send_heartbeat(reader, writer),
], return_when=asyncio.FIRST_COMPLETED) # will stop after 1 second
writer.close() # close socket after 1 second
await writer.wait_closed()
async def handle_stream(reader, writer):
writer.write(b'hello\n') # will success because socket is alive
await writer.drain()
await asyncio.sleep(1)
async def send_heartbeat(reader, writer):
while True:
await asyncio.sleep(300)
heartbeat_message = b'heartbeat\n'
writer.write(heartbeat_message) # will fail because socket is already closed after 1 second
await writer.drain()
if __name__ == '__main__':
asyncio.run(run('127.0.0.1', 8889))
You can cancel the sleep by canceling a task that executes it. Creating send_heartbeat as a separate task ensures that it runs in parallel to handle_stream while you await the latter:
async def run(host: str, port: int):
while True:
...
heartbeat = asyncio.create_task(send_heartbeat(reader, writer))
try:
await handle_stream(reader, writer)
finally:
heartbeat.cancel()
writer.close()
await writer.wait_closed()
BTW, since you're awaiting writer.drain() inside handle_stream, there is no guarantee that handle_stream will always complete in 1 second. This might be a place where you might want to avoid the drain, or you can use asyncio.wait_for when awaiting handle_stream(...).

Make websocket callback asynchronous with asyncio

I am trying to implement a basic websocket client using asyncio and websockets with Python 3.5.2.
Basically, I want connect_to_dealer to be a blocking call, but wait for the websocket message on a different thread.
After reading some docs (I have very little exp with Python), I concluded that asyncio.ensure_future() passing a coroutine (listen_for_message) was the way to go.
Now, I get to run listen_for_message on a different thread, but from within the coroutine I can't seem to use await or any other mechanism to make the calls synchronous. If I do it, the execution waits forever (it hangs) even for a simple sleep.
I'd like to know what I'm doing wrong.
async def listen_for_message(self, future, websocket):
while (True):
try:
await asyncio.sleep(1) # It hangs here
print('Listening for a message...')
message = await websocket.recv() # If I remove the sleep, hangs here
print("< {}".format(message))
future.set_result(message)
future.done()
except websockets.ConnectionClosed as cc:
print('Connection closed')
except Exception as e:
print('Something happened')
def handle_connect_message(self, future):
# We must first remove the websocket-specific payload because we're only interested in the connect protocol msg
print(future.result)
async def connect_to_dealer(self):
print('connect to dealer')
websocket = await websockets.connect('wss://mywebsocket'))
hello_message = await websocket.recv()
print("< {}".format(hello_message))
# We need to parse the connection ID out of the message
connection_id = hello_message['connectionId']
print('Got connection id {}'.format(connection_id))
sub_response = requests.put('https://subscribetotraffic{user_id}?connection={connection_id}'.format(user_id='username', connection_id=connection_id), headers=headers)
if sub_response.status_code == 200:
print('Now we\'re observing traffic')
else:
print('Oops request failed with {code}'.format(code=sub_response.status_code))
# Now we need to handle messages but continue with the regular execution
try:
future = asyncio.get_event_loop().create_future()
future.add_done_callback(self.handle_connect_message)
asyncio.ensure_future(self.listen_for_message(future, websocket))
except Exception as e:
print(e)
Is there a specific reason you need to work with explicit futures?
With asyncio you can use a combination of coroutines and Tasks to achieve most purposes. Tasks are essentially wrapped coroutines that go about cranking themselves over in the background, independently of other async code, so you don't have to explicitly manage their flow or juggle them with other bits of code.
I am not entirely sure of your end goal, but perhaps the approach elaborated below gives you something to work with:
import asyncio
async def listen_for_message(websocket):
while True:
await asyncio.sleep(0)
try:
print('Listening for a message...')
message = await websocket.recv()
print("< {}".format(message))
except websockets.ConnectionClosed as cc:
print('Connection closed')
except Exception as e:
print('Something happened')
async def connect_to_dealer():
print('connect to dealer')
websocket = await websockets.connect('wss://mywebsocket')
hello_message = await websocket.recv()
print("< {}".format(hello_message))
# We need to parse the connection ID out of the message
connection_id = hello_message['connectionId']
print('Got connection id {}'.format(connection_id))
sub_response = requests.put('https://subscribetotraffic{user_id}?connection={connection_id}'.format(
user_id='username', connection_id=connection_id), headers=headers)
if sub_response.status_code == 200:
print('Now we\'re observing traffic')
else:
print('Oops request failed with {code}'.format(code=sub_response.status_code))
async def my_app():
# this will block until connect_to_dealer() returns
websocket = await connect_to_dealer()
# start listen_for_message() in its own task wrapper, so doing it continues in the background
asyncio.ensure_future(listen_for_message(websocket))
# you can continue with other code here that can now coexist with listen_for_message()
if __name__ == '__main__':
loop = asyncio.get_event_loop()
loop.run_until_complete(my_app())
loop.run_forever()

Asyncio: Start a non-blocking listening server

This is the basic tcp server from asyncio tutotial:
import asyncio
class EchoServerClientProtocol(asyncio.Protocol):
def connection_made(self, transport):
peername = transport.get_extra_info('peername')
print('Connection from {}'.format(peername))
self.transport = transport
def data_received(self, data):
message = data.decode()
print('Data received: {!r}'.format(message))
print('Send: {!r}'.format(message))
self.transport.write(data)
print('Close the client socket')
self.transport.close()
loop = asyncio.get_event_loop()
# Each client connection will create a new protocol instance
coro = loop.create_server(EchoServerClientProtocol, '127.0.0.1', 8888)
server = loop.run_until_complete(coro)
# Serve requests until CTRL+c is pressed
print('Serving on {}'.format(server.sockets[0].getsockname()))
try:
loop.run_forever()
except KeyboardInterrupt:
pass
# Close the server
server.close()
loop.run_until_complete(server.wait_closed())
loop.close()
Like all (i found) other examples it uses blocking loop.run_forever().
How do i start listeting server and do something else in the time?
I have tried to outsource starting server in a function and start this function with asyncio.async(), but with no success.
What i'm missing here?
You can schedule several concurrent asyncio tasks before calling loop.run_forever().
#asyncio.coroutine
def other_task_coroutine():
pass # do something
start_tcp_server_task = loop.create_task(loop.create_server(
EchoServerClientProtocol, '127.0.0.1', 8888))
other_task = loop.create_task(other_task_coroutine())
self.run_forever()
When you call loop.create_task(loop.create_server()) or loop.create_task(other_task_coroutine()), nothing is actually executed: a coroutine object is created and wrapped in a task (consider a task to be a shell and the coroutine an instance of the code that will be executed in the task). The tasks are scheduled on the loop when created.
The loop will execute start_tcp_server_task first (as it's scheduled first) until a blocking IO event is pending or the passive socket is ready to listen for incoming connections.
You can see asyncio as a non-preemptible scheduler running on one CPU: once the first task interrupts itself or is done, the second task will be executed. Hence, when one task is executed, the other one has to wait until the running task finishes or yields (or "awaits" with Python 3.5). "yielding" (yield from client.read()) or "awaiting" (await client.read()) means that the task gives back the hand to the loop's scheduler, until client.read() can be executed (data is available on the socket).
Once the task gave back the control to the loop, it can schedule the other pending tasks, process incoming events and schedule the tasks which were waiting for those events. Once there is nothing left to do, the loop will perform the only blocking call of the process: sleep until the kernel notifies it that events are ready to be processed.
In this context, you must understand that when using asyncio, everything running in the process must run asynchronously so the loop can do its work. You can not use multiprocessing objects in the loop.
Note that asyncio.async(coroutine(), loop=loop) is equivalent to loop.create_task(coroutine()).
Additionally, you can consider running what you want in an executor.
For example.
coro = loop.create_server(EchoServerClientProtocol, '127.0.0.1', 8888)
server = loop.run_until_complete(coro)
async def execute(self, loop):
await loop.run_in_executor(None, your_func_here, args:
asyncio.async(execute(loop))
loop.run_forever()
An executor will run whatever function you want in an executor, which wont block your server.

Categories

Resources