Search code examples
pythonpython-trio

How to manually exit an infinite trio-loop, like the trio's tutorial echo client


Is there a way to manually exit a trio infinite loop, like the echo client in the trio tutorial, https://trio.readthedocs.io/en/latest/tutorial.html#an-echo-client , other than using Ctrl-C or using timeouts?

My idea is to use call the echo client from another python script, and be able to close it too with the same python script, arbitrarily. I was thinking of using a flag (maybe event?) as a switch to trigger the cancel_scope.cancel() in the nursery. But I don't know how to trigger the switch. Below is my attempt at modifying the tutorial echo client code.

import sys
import trio

PORT = 12345
BUFSIZE = 16384
FLAG = 1 # FLAG is a global variable

async def sender(client_stream):
    print("sender: started")
    while FLAG:
        data = b'async can sometimes be confusing but I believe in you!'
        print(f"sender: sending {data}")
        await client_stream.send_all(data)
        await trio.sleep(1)

async def receiver(client_stream):
    print("recevier: started!")
    while FLAG:
        data = await client_stream.receive_some(BUFSIZE)
        print(f"receiver: got data {data}")
        if not data:
            print("receiver: connection closed")
            sys.exit()

async def checkflag(nursery): # function to trigger cancel()
    global FLAG
    if not FLAG:
        nursery.cancel_scope.cancel()
    else:
        # keep this task running if not triggered, but how to trigger it, 
        # without Ctrl-C or timeout?
        await trio.sleep(1) 

async def parent():
    print(f"parent: connecting to 127.0.0.1:{PORT}")
    client_stream = await trio.open_tcp_stream("127.0.0.1", PORT)
    async with client_stream:
        async with trio.open_nursery() as nursery:
            print("parent: spawning sender ...")
            nursery.start_soon(sender, client_stream)

            print("parent: spawning receiver ...")
            nursery.start_soon(receiver, client_stream)

            print("parent: spawning checkflag...")
            nursery.start_soon(checkflag, nursery)

        print('Close nursery...')
    print("Close stream...")

trio.run(parent)

I find that I am unable to input any commands into the python REPL after trio.run(), to manually change the FLAG, and am wondering if I call this echo client from another script, how exactly to trigger the cancel_scope.cancel() in the nursery? Or is there a better way? Really appreciate all help. Thanks.


Solution

  • If you want to use input from the keyboard to exit, here's a solution for Linux and Mac OS X. You could use the Python msvcrt module to do something similar on Windows.

    I copied echo-client.py from the Trio tutorial and put a "NEW" comment over the three added blocks of code. From the REPL, you can type 'q' to cancel the nursery scope and exit:

    # -- NEW
    import termios, tty
    
    import sys
    import trio
    
    PORT = 12345
    BUFSIZE = 16384
    
    # -- NEW
    async def keyboard():
        """Return an iterator of characters from stdin."""
        stashed_term = termios.tcgetattr(sys.stdin)
        try:
            tty.setcbreak(sys.stdin, termios.TCSANOW)
            while True:
                yield await trio.run_sync_in_worker_thread(
                    sys.stdin.read, 1,
                    cancellable=True
                )
        finally:
            termios.tcsetattr(sys.stdin, termios.TCSANOW, stashed_term)
    
    async def sender(client_stream):
        print("sender: started!")
        while True:
            data = b"async can sometimes be confusing, but I believe in you!"
            print("sender: sending {!r}".format(data))
            await client_stream.send_all(data)
            await trio.sleep(1)
    
    async def receiver(client_stream):
        print("receiver: started!")
        while True:
            data = await client_stream.receive_some(BUFSIZE)
            print("receiver: got data {!r}".format(data))
            if not data:
                print("receiver: connection closed")
                sys.exit()
    
    async def parent():
        print("parent: connecting to 127.0.0.1:{}".format(PORT))
        client_stream = await trio.open_tcp_stream("127.0.0.1", PORT)
        async with client_stream:
            async with trio.open_nursery() as nursery:
                print("parent: spawning sender...")
                nursery.start_soon(sender, client_stream)
    
                print("parent: spawning receiver...")
                nursery.start_soon(receiver, client_stream)
    
                # -- NEW
                async for key in keyboard():
                    if key == 'q':
                        nursery.cancel_scope.cancel()
    
    trio.run(parent)
    

    The call to tty.setcbreak puts the terminal in unbuffered mode so you don't have to press return before the program receives input. It also prevents characters from being echoed to the screen. Furthermore, as the name implies, it allows Ctrl-C to work as normal.

    In the finally block, termios.tcsetattr restores the terminal to whatever mode it was in before tty.setcbreak. So your terminal is back to normal on exit.

    sys.stdin.read is spawned in a separate thread because it needs to run in blocking mode (not good in an async context). The reason is that stdin shares its file description with stdout and stderr. Setting stdin to non-blocking would also set stdout to non-blocking as a side effect, and that might cause issues with the print function (truncation in my case).

    Inter-process communication

    Here's a basic example of cancelling one Trio process from another with a socket:

    # infinite_loop.py    
    import trio    
    
    async def task():     
        while True:       
            print("ping")    
            await trio.sleep(0.5)    
    
    async def quitter(cancel_scope):      
        async def quit(server_stream):    
            await server_stream.receive_some(1024)    
            cancel_scope.cancel()    
        await trio.serve_tcp(quit, 12346)    
    
    async def main():    
        async with trio.open_nursery() as nursery:    
            nursery.start_soon(task)    
            nursery.start_soon(quitter, nursery.cancel_scope)    
    
    trio.run(main)
    
    # slayer.py        
    import trio    
    
    async def main():    
        async with await trio.open_tcp_stream("127.0.0.1", 12346) as s:
            await trio.sleep(3)    
            await s.send_all(b'quit')    
    
    trio.run(main)