├── .gitignore ├── LICENSE ├── README.md ├── call_sync_code.py ├── concurrent_blocking_requests.py ├── factorial_iter.py ├── future_example1.py ├── future_example2.py ├── future_example3.py ├── mutually_recursive.py ├── parallel_examples.py ├── parallel_http_get.py ├── pubsub_robust ├── README.md ├── pubsub_client.py └── pubsub_server.py ├── run_in_executor.py ├── simple_coroutines.py ├── tail_recursion_with_asyncio.py └── tcp_pubsub_server.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Distribution / packaging 9 | .Python 10 | env/ 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | *.egg-info/ 23 | .installed.cfg 24 | *.egg 25 | 26 | # PyInstaller 27 | # Usually these files are written by a python script from a template 28 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 29 | *.manifest 30 | *.spec 31 | 32 | # Installer logs 33 | pip-log.txt 34 | pip-delete-this-directory.txt 35 | 36 | # Unit test / coverage reports 37 | htmlcov/ 38 | .tox/ 39 | .coverage 40 | .coverage.* 41 | .cache 42 | nosetests.xml 43 | coverage.xml 44 | *,cover 45 | 46 | # Translations 47 | *.mo 48 | *.pot 49 | 50 | # Django stuff: 51 | *.log 52 | 53 | # Sphinx documentation 54 | docs/_build/ 55 | 56 | # PyBuilder 57 | target/ 58 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 Caleb Madrigal 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # asyncio-examples 2 | A few examples of how to use asyncio 3 | -------------------------------------------------------------------------------- /call_sync_code.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import time 3 | from urllib.request import urlopen 4 | 5 | @asyncio.coroutine 6 | def count_to_10(): 7 | for i in range(11): 8 | print("Counter: {}".format(i)) 9 | yield from asyncio.sleep(.5) 10 | 11 | def get_page_len(url): 12 | # This is the blocking sleep (not the async-friendly one) 13 | time.sleep(2) 14 | page = urlopen(url).read() 15 | return len(page) 16 | 17 | @asyncio.coroutine 18 | def run_get_page_len(): 19 | loop = asyncio.get_event_loop() 20 | 21 | future1 = loop.run_in_executor(None, get_page_len, 'http://calebmadrigal.com') 22 | 23 | #data1 = yield from future1 24 | return future1 25 | 26 | @asyncio.coroutine 27 | def print_data_size(): 28 | data = yield from run_get_page_len() 29 | print("Data size: {}".format(data)) 30 | 31 | 32 | loop = asyncio.get_event_loop() 33 | tasks = [ 34 | asyncio.async(count_to_10()), 35 | asyncio.async(print_data_size())] 36 | loop.run_until_complete(asyncio.wait(tasks)) 37 | loop.close() 38 | 39 | -------------------------------------------------------------------------------- /concurrent_blocking_requests.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import time 3 | import random 4 | import urllib.request 5 | 6 | RAND_VARIATION = 5 7 | URL = 'http://xkcd.com' 8 | 9 | 10 | # Blocking HTTP GET 11 | def download_url(url): 12 | return urllib.request.urlopen(url).read() 13 | 14 | 15 | @asyncio.coroutine 16 | def worker_loop(hostid, interval): 17 | count = 0 18 | while True: 19 | count += 1 20 | print("Agent {} - loop count: {}".format(hostid, count)) 21 | 22 | # Call download_url in the thread pool. 23 | loop = asyncio.get_event_loop() 24 | fut = loop.run_in_executor(None, download_url, URL) 25 | result = yield from fut 26 | print("Agent {} Got data len: {}".format(hostid, len(result))) 27 | 28 | yield from asyncio.sleep(interval + random.randint(0, RAND_VARIATION)) 29 | 30 | 31 | @asyncio.coroutine 32 | def run_agents(): 33 | num_agents = 5 34 | interval = 10 35 | all_agents = [worker_loop(i, interval) for i in range(1, num_agents+1)] 36 | yield from asyncio.gather(*all_agents) 37 | 38 | 39 | def main(): 40 | loop = asyncio.get_event_loop() 41 | loop.run_until_complete(run_agents()) 42 | loop.close() 43 | 44 | if __name__ == '__main__': 45 | main() 46 | 47 | -------------------------------------------------------------------------------- /factorial_iter.py: -------------------------------------------------------------------------------- 1 | # For time comparison with tail_recursion_with_asyncio.py 2 | # 3 | # Results: 4 | # * This (iterative): 1.3s 5 | # * Tail-recursive with asyncio: 2.6s (about 2x longer) 6 | 7 | def factorial(n): 8 | acc = 1 9 | for i in range(1, n+1): 10 | acc *= i 11 | return acc 12 | 13 | print("Result: {}".format(factorial(50000))) 14 | 15 | -------------------------------------------------------------------------------- /future_example1.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | @asyncio.coroutine 4 | def slow_operation(future): 5 | yield from asyncio.sleep(1) 6 | future.set_result(42) 7 | 8 | @asyncio.coroutine 9 | def call_slow_operation(): 10 | fut = asyncio.Future() 11 | yield from slow_operation(fut) 12 | result = fut.result() 13 | print("The answer is: {}".format(result)) 14 | 15 | loop = asyncio.get_event_loop() 16 | loop.run_until_complete(call_slow_operation()) 17 | loop.close() 18 | 19 | -------------------------------------------------------------------------------- /future_example2.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | @asyncio.coroutine 4 | def slow_operation(future): 5 | yield from asyncio.sleep(1) 6 | future.set_result(42) 7 | 8 | @asyncio.coroutine 9 | def call_slow_operation(): 10 | def future_callback(future): 11 | result = fut.result() 12 | print("The answer is: {}".format(result)) 13 | 14 | fut = asyncio.Future() 15 | fut.add_done_callback(future_callback) 16 | yield from slow_operation(fut) 17 | 18 | loop = asyncio.get_event_loop() 19 | loop.run_until_complete(call_slow_operation()) 20 | loop.close() 21 | 22 | -------------------------------------------------------------------------------- /future_example3.py: -------------------------------------------------------------------------------- 1 | # This doesn't actually explicitly use a Future, but shows how to get the same 2 | # behavior with just coroutines. 3 | 4 | import asyncio 5 | 6 | @asyncio.coroutine 7 | def slow_operation(): 8 | yield from asyncio.sleep(1) 9 | return 42 10 | 11 | @asyncio.coroutine 12 | def call_slow_operation(): 13 | result = yield from slow_operation() 14 | print("The answer is: {}".format(result)) 15 | 16 | loop = asyncio.get_event_loop() 17 | loop.run_until_complete(call_slow_operation()) 18 | loop.close() 19 | 20 | -------------------------------------------------------------------------------- /mutually_recursive.py: -------------------------------------------------------------------------------- 1 | """ Simple mutually-recursive coroutines with asyncio. Using asyncio.ensure_future 2 | instead of yield from allows the coroutine to exit and merely schedules the next 3 | call with the event loop, allowing infinite mutual recursion. """ 4 | 5 | import asyncio 6 | 7 | @asyncio.coroutine 8 | def a(n): 9 | print("A: {}".format(n)) 10 | asyncio.async(b(n+1)) # asyncio.ensure_future in Python 3.4.4 11 | 12 | @asyncio.coroutine 13 | def b(n): 14 | print("B: {}".format(n)) 15 | asyncio.async(a(n+1)) # asyncio.ensure_future in Python 3.4.4 16 | 17 | loop = asyncio.get_event_loop() 18 | asyncio.async(a(0)) 19 | loop.run_forever() 20 | loop.close() 21 | 22 | -------------------------------------------------------------------------------- /parallel_examples.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | @asyncio.coroutine 4 | def waitn(n): 5 | asyncio.sleep(n) 6 | return "I waited {}".format(n) 7 | 8 | @asyncio.coroutine 9 | def run_parallel(): 10 | # Results will be in order called 11 | results = yield from asyncio.gather(waitn(3), waitn(1), waitn(2)) 12 | print("Results: {}".format(results)) 13 | 14 | @asyncio.coroutine 15 | def run_parallel2(): 16 | tasks = [waitn(i) for i in (3,1,2)] 17 | # Results will be in order called 18 | results = yield from asyncio.gather(*tasks) 19 | print("Results2: {}".format(results)) 20 | 21 | @asyncio.coroutine 22 | def run_parallel3(): 23 | tasks = [asyncio.async(waitn(i)) for i in (3,1,2)] 24 | done, pending = yield from asyncio.wait(tasks) 25 | # Results will NOT necessarily be in the order called 26 | results = [future.result() for future in done] 27 | print("Results3: {}".format(results)) 28 | 29 | loop = asyncio.get_event_loop() 30 | loop.run_until_complete(run_parallel()) 31 | loop.run_until_complete(run_parallel2()) 32 | loop.run_until_complete(run_parallel3()) 33 | loop.close() 34 | -------------------------------------------------------------------------------- /parallel_http_get.py: -------------------------------------------------------------------------------- 1 | """ Downloads a few web pages in parallel, and counts how many times a specified 2 | word is used in each of them. """ 3 | 4 | import asyncio 5 | import aiohttp 6 | 7 | @asyncio.coroutine 8 | def download_and_count_word(word, url): 9 | response = yield from aiohttp.request('GET', url) 10 | text = yield from response.read() 11 | return text.decode().count(word) 12 | 13 | @asyncio.coroutine 14 | def count_word_in_pages(word, urls): 15 | tasks = [download_and_count_word(word, url) for url in urls] 16 | counts = yield from asyncio.gather(*tasks) 17 | 18 | for i in range(len(urls)): 19 | url = urls[i] 20 | count = counts[i] 21 | print("{} appears {} times in {}".format(word, count, url)) 22 | 23 | word = "the" 24 | pages = ["http://calebmadrigal.com", 25 | "http://yahoo.com", 26 | "http://xkcd.com", 27 | "http://reddit.com", 28 | "http://news.ycombinator.com"] 29 | 30 | loop = asyncio.get_event_loop() 31 | loop.run_until_complete(count_word_in_pages(word, pages)) 32 | loop.close() 33 | 34 | -------------------------------------------------------------------------------- /pubsub_robust/README.md: -------------------------------------------------------------------------------- 1 | # pub-sub server 2 | 3 | Pub/sub tcp server with asyncio. 4 | 5 | ## Usage 6 | 7 | Run server: 8 | 9 | python3 pubsub_server.py 10 | 11 | Run subscriber client (to topic 1): 12 | 13 | python3 pubsub_client.py sub 1 14 | 15 | Run publisher client (to topic 1): 16 | 17 | python3 pubsub_client.py pub 1 18 | 19 | The publisher currently sends 100,000 32-byte messages. 20 | 21 | ## Description 22 | 23 | This is a simple pub-sub server that works over TCP. 24 | 25 | When a client first connects, it sends the list of topics it wants to subscribe to 26 | in this format: 27 | 28 | * body size (little-endian unsigned int) 29 | * list of little-endian unsigned ints, one for each topic. 0 is the only excluded value, 30 | as 0 indicates a close-connection command. 31 | 32 | Messages are in this format: 33 | 34 | * topic (little-endian unsigned int) 35 | * body size (little-endian unsigned int) 36 | * body (raw bytes in any format) 37 | 38 | When a client wishes to close the connection, it sends: 39 | 40 | * topic = 0 41 | * body size = 0 42 | * no body 43 | 44 | ## Performance 45 | 46 | Initial performance testing: 47 | 48 | 1-byte messages 49 | * 40k messages/second - 1 subscriber 50 | * 25k messages/second - 2 subscribers 51 | * 20k messages/second - 3 subscribers 52 | 53 | 32-byte messages 54 | * 20k messages/second - 1 subscriber 55 | * 16k messages/second - 2 subscribers 56 | * 15k messages/second - 3 subscribers 57 | 58 | -------------------------------------------------------------------------------- /pubsub_robust/pubsub_client.py: -------------------------------------------------------------------------------- 1 | """ pubub_client.py 2 | 3 | Simple client that sends messages on pubsub_server. 4 | """ 5 | 6 | import struct 7 | import socket 8 | import sys 9 | 10 | CONNECT_HEADER_FORMAT = ' ') 68 | sys.exit(1) 69 | 70 | try: 71 | mode = sys.argv[1] 72 | topic = int(sys.argv[2]) 73 | except IndexError: 74 | usage() 75 | if mode == 'pub': 76 | publish(host, port, topic) 77 | elif mode == 'sub': 78 | subscribe(host, port, topic) 79 | else: 80 | usage() 81 | -------------------------------------------------------------------------------- /pubsub_robust/pubsub_server.py: -------------------------------------------------------------------------------- 1 | """ pubsub_server.py 2 | 3 | This is a simple pub-sub server that works over TCP. 4 | 5 | When a client first connects, it sends the list of topics it wants to subscribe to 6 | in this format: 7 | 8 | * body size (little-endian unsigned int) 9 | * list of little-endian unsigned ints, one for each topic. 0 is the only excluded value, 10 | as 0 indicates a close-connection command. 11 | 12 | Messages are in this format: 13 | 14 | * topic (little-endian unsigned int) 15 | * body size (little-endian unsigned int) 16 | * body (raw bytes in any format) 17 | 18 | When a client wishes to close the connection, it sends: 19 | 20 | * topic = 0 21 | * body size = 0 22 | * no body 23 | 24 | """ 25 | 26 | import asyncio 27 | import struct 28 | 29 | CONNECT_HEADER_FORMAT = ' [sub1_writer, sub2_writer, ...] 41 | self.subscriber_to_topics = {} # subscriber_writer -> [topic1, topic2, ...] 42 | 43 | def add_subscriber(self, topic, subscriber_writer): 44 | if topic in self.topic_to_subscribers: 45 | self.topic_to_subscribers[topic].append(subscriber_writer) 46 | else: 47 | self.topic_to_subscribers[topic] = [subscriber_writer] 48 | 49 | if subscriber_writer in self.subscriber_to_topics: 50 | self.subscriber_to_topics[subscriber_writer].append(topic) 51 | else: 52 | self.subscriber_to_topics[subscriber_writer] = [topic] 53 | 54 | def remove_subscriber(self, subscriber_writer, client_addr): 55 | if subscriber_writer in self.subscriber_to_topics: 56 | subscriber_topics = self.subscriber_to_topics[subscriber_writer] 57 | print('Removing subscriber {} from topics: {}'.format(client_addr, subscriber_topics)) 58 | for topic in subscriber_topics: 59 | self.topic_to_subscribers[topic].remove(subscriber_writer) 60 | del self.subscriber_to_topics[subscriber_writer] 61 | 62 | @asyncio.coroutine 63 | def read_exact(self, client_reader, bytes_to_read): 64 | bytes_read = 0 65 | buf = b'' 66 | while bytes_read < bytes_to_read: 67 | just_read = yield from client_reader.read(bytes_to_read - bytes_read) 68 | if not just_read: 69 | raise IOError('Client closed before reading full message') 70 | buf += just_read 71 | bytes_read += len(just_read) 72 | return buf 73 | 74 | @asyncio.coroutine 75 | def read_message(self, client_reader): 76 | raw_header = yield from self.read_exact(client_reader, HEADER_SIZE) 77 | try: 78 | topic, body_size = struct.unpack(HEADER_FORMAT, raw_header) 79 | except struct.error: 80 | raise IOError('Invalid message header. Should be 2 little-endian ' + 81 | 'unsigned ints representing (1) topic (2) body size') 82 | 83 | body = yield from self.read_exact(client_reader, body_size) 84 | full_msg = raw_header + body 85 | 86 | return topic, full_msg 87 | 88 | @asyncio.coroutine 89 | def read_connect_message(self, client_reader): 90 | raw_header = yield from self.read_exact(client_reader, CONNECT_HEADER_SIZE) 91 | try: 92 | num_topics = struct.unpack(CONNECT_HEADER_FORMAT, raw_header)[0] 93 | except (struct.error, IndexError): 94 | raise IOError('Invalid connect message header. Should be a little-endian unsigned int') 95 | 96 | body_format = '<' + 'I' * num_topics 97 | body_size = struct.calcsize(body_format) 98 | body = yield from self.read_exact(client_reader, body_size) 99 | 100 | try: 101 | topic_list = struct.unpack(body_format, body) 102 | except struct.error: 103 | raise IOError('Invalid connect message body. Should be a sequence ' + 104 | 'of little-endian unsigned ints') 105 | 106 | return topic_list 107 | 108 | @asyncio.coroutine 109 | def handle_client(self, client_reader, client_writer): 110 | client_addr = client_writer.get_extra_info('peername') 111 | print('New connection from {}'.format(client_addr)) 112 | 113 | # Get topics - the initial message the client should send 114 | try: 115 | topic_list = yield from self.read_connect_message(client_reader) 116 | except IOError as e: 117 | print('ERROR: {}'.format(e)) 118 | client_writer.close() 119 | return 120 | 121 | for topic in topic_list: 122 | self.add_subscriber(topic, client_writer) 123 | 124 | print('Client {} subscribes to topics: {}'.format(client_addr, topic_list)) 125 | 126 | def close_connection(): 127 | self.remove_subscriber(client_writer, client_addr) 128 | client_writer.close() 129 | 130 | try: 131 | while True: 132 | topic, msg = yield from self.read_message(client_reader) 133 | 134 | if topic == CLOSE_TOPIC: 135 | break 136 | 137 | if topic in self.topic_to_subscribers: 138 | for subscriber_writer in self.topic_to_subscribers[topic]: 139 | subscriber_writer.write(msg) 140 | 141 | yield from client_writer.drain() 142 | 143 | except IOError as e: 144 | print('ERROR: {}'.format(e)) 145 | 146 | close_connection() 147 | 148 | def start(self, loop): 149 | coro = asyncio.start_server(self.handle_client, '127.0.0.1', 8888, loop=loop) 150 | self.server = loop.run_until_complete(coro) 151 | print('Serving on {}'.format(self.server.sockets[0].getsockname())) 152 | 153 | def stop(self, loop): 154 | if self.server is not None: 155 | self.server.close() 156 | loop.run_until_complete(self.server.wait_closed()) 157 | self.server = None 158 | 159 | 160 | def main(): 161 | loop = asyncio.get_event_loop() 162 | 163 | # creates a server and starts listening to TCP connections 164 | server = PubSubBus() 165 | server.start(loop) 166 | 167 | try: 168 | loop.run_forever() 169 | except KeyboardInterrupt: 170 | server.stop(loop) 171 | loop.close() 172 | 173 | if __name__ == '__main__': 174 | main() 175 | -------------------------------------------------------------------------------- /run_in_executor.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | from urllib.request import urlopen 3 | 4 | @asyncio.coroutine 5 | def print_data_size(): 6 | data = yield from get_data_size() 7 | print("Data size: {}".format(data)) 8 | 9 | # Note that this is a synchronous function 10 | def sync_get_url(url): 11 | return urlopen(url).read() 12 | 13 | @asyncio.coroutine 14 | def get_data_size(): 15 | loop = asyncio.get_event_loop() 16 | 17 | # These each run in their own thread (in parallel) 18 | future1 = loop.run_in_executor(None, sync_get_url, 'http://xkcd.com') 19 | future2 = loop.run_in_executor(None, sync_get_url, 'http://google.com') 20 | 21 | # While the synchronous code above is running in other threads, the event loop 22 | # can go do other things. 23 | data1 = yield from future1 24 | data2 = yield from future2 25 | return len(data1) + len(data2) 26 | 27 | loop = asyncio.get_event_loop() 28 | loop.run_until_complete(print_data_size()) 29 | 30 | -------------------------------------------------------------------------------- /simple_coroutines.py: -------------------------------------------------------------------------------- 1 | """ Simple mutually-recursive coroutines with asyncio. Note that these recursive calls 2 | continue to grow the stack (and will eventually hit the maximum recursion depth 3 | exception if too many recursive calls are made. """ 4 | 5 | import asyncio 6 | 7 | @asyncio.coroutine 8 | def a(n): 9 | print("A: {}".format(n)) 10 | if n > 10: return n 11 | else: yield from b(n+1) 12 | 13 | @asyncio.coroutine 14 | def b(n): 15 | print("B: {}".format(n)) 16 | yield from a(n+1) 17 | 18 | loop = asyncio.get_event_loop() 19 | loop.run_until_complete(a(0)) 20 | 21 | -------------------------------------------------------------------------------- /tail_recursion_with_asyncio.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | # Tail-recursive factorial using asyncio event loop as a trampoline to 4 | # keep the stack from growing. 5 | @asyncio.coroutine 6 | def factorial(n, callback, acc=1): 7 | if n == 0: 8 | callback(acc) 9 | else: 10 | asyncio.async(factorial(n-1, callback, acc*n)) # async -> ensure_future in Python 3.4.4 11 | 12 | def done_callback(result): 13 | print("Result: {}".format(result)) 14 | loop = asyncio.get_event_loop() 15 | loop.stop() 16 | 17 | 18 | loop = asyncio.get_event_loop() 19 | asyncio.async(factorial(50000, done_callback)) 20 | loop.run_forever() # Blocking call interrupted by loop.stop() 21 | loop.close() 22 | 23 | -------------------------------------------------------------------------------- /tcp_pubsub_server.py: -------------------------------------------------------------------------------- 1 | """ tcp_pubub_server.py 2 | 3 | This is a simple pub-sub server that works over TCP and has text-based messages. 4 | Note that this code is not robust. The 2 main problems are: 5 | 6 | * Messages may not contain commas 7 | * The receive code can break if the full message is not in the socket recv buffer 8 | by the time readline() is called. 9 | 10 | But this serves for a very simple pub-sub asyncio demo which can be connected to by telnet: 11 | 12 | (py)cmbpr:~ caleb$ telnet localhost 8888 13 | Trying 127.0.0.1... 14 | Connected to localhost. 15 | Escape character is '^]'. 16 | food,wine 17 | wine,red 18 | wine,red 19 | goodbye 20 | Connection closed by foreign host. 21 | 22 | What's happening in this example: 23 | 24 | * Client connects and subscribes to 'food' and 'wine' topics 25 | * Client sends a message ('red') under the topic ('wine') 26 | * Client receives the message because it subscribed to the 'wine' topic 27 | - Any other clients connected which subscribed to 'wine' would also receive it 28 | * Client closes the connection by the command, 'goodbye' 29 | 30 | """ 31 | 32 | import asyncio 33 | 34 | 35 | class BusServer: 36 | def __init__(self): 37 | self.server = None 38 | self.topic_to_subscribers = {} # topic -> [sub1_writer, sub2_writer, ...] 39 | self.subscriber_to_topics = {} # subscriber_writer -> [topic1, topic2, ...] 40 | 41 | def add_subscriber(self, topic, subscriber_writer): 42 | if topic in self.topic_to_subscribers: 43 | self.topic_to_subscribers[topic].append(subscriber_writer) 44 | else: 45 | self.topic_to_subscribers[topic] = [subscriber_writer] 46 | 47 | if subscriber_writer in self.subscriber_to_topics: 48 | self.subscriber_to_topics[subscriber_writer].append(topic) 49 | else: 50 | self.subscriber_to_topics[subscriber_writer] = [topic] 51 | 52 | def remove_subscriber(self, subscriber_writer, client_addr): 53 | subscriber_topics = self.subscriber_to_topics[subscriber_writer] 54 | print('Removing subscriber {} from topics: {}'.format(client_addr, subscriber_topics)) 55 | for topic in subscriber_topics: 56 | self.topic_to_subscribers[topic].remove(subscriber_writer) 57 | del self.subscriber_to_topics[subscriber_writer] 58 | 59 | @asyncio.coroutine 60 | def handle_client(self, client_reader, client_writer): 61 | client_addr = client_writer.get_extra_info('peername') 62 | # Get topics 63 | topics_raw = (yield from client_reader.readline()) 64 | topics = topics_raw.decode("utf-8").rstrip().split(',') 65 | print('New client {} subscribes to topics: {}'.format(client_addr, topics)) 66 | for topic in topics: 67 | self.add_subscriber(topic, client_writer) 68 | 69 | while True: 70 | full_msg_bytes = (yield from client_reader.readline()) 71 | full_msg = full_msg_bytes.decode("utf-8").rstrip() 72 | 73 | if not full_msg or full_msg == 'goodbye': # an empty string means the client disconnected 74 | break 75 | topic, msg = full_msg.split(',') 76 | 77 | if topic in self.topic_to_subscribers: 78 | for subscriber_writer in self.topic_to_subscribers[topic]: 79 | subscriber_writer.write(full_msg_bytes) 80 | 81 | yield from client_writer.drain() 82 | 83 | self.remove_subscriber(client_writer, client_addr) 84 | client_writer.close() 85 | 86 | def start(self, loop): 87 | coro = asyncio.start_server(self.handle_client, '127.0.0.1', 8888, loop=loop) 88 | self.server = loop.run_until_complete(coro) 89 | print('Serving on {}'.format(self.server.sockets[0].getsockname())) 90 | 91 | def stop(self, loop): 92 | if self.server is not None: 93 | self.server.close() 94 | loop.run_until_complete(self.server.wait_closed()) 95 | self.server = None 96 | 97 | 98 | def main(): 99 | loop = asyncio.get_event_loop() 100 | 101 | # creates a server and starts listening to TCP connections 102 | server = BusServer() 103 | server.start(loop) 104 | 105 | try: 106 | loop.run_forever() 107 | except KeyboardInterrupt: 108 | server.stop(loop) 109 | loop.close() 110 | 111 | if __name__ == '__main__': 112 | main() 113 | --------------------------------------------------------------------------------