In python, what happens when a thread is started inside of a context manager?

Question

As a simple example, take a basic socket:

import socket
import threading
import time

def stream(message, socket_):
    while True:
        socket_.sendall(message)
        time.sleep(1)

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
    sock.connect((address, port))
    threading.Thread(target=stream, args=("hello world!", sock), daemon=True).start()

I assume the flow is something like this:

1. Context is entered, sock is created.
2. The thread is created and stream is called with sock.
3. The context is exited and sock is closed.
4. The stream thread is not using a closed socket and will probably crash.

Is that right? If so, is there any way to continue the context in one or more separate threads?

Answer 1

Yes, you are correct. In the code you provided, the sock socket object will be closed when the with block is exited. Therefore, if the stream thread continues to use the sock object after it has been closed, it will raise a socket.error or socket.timeout exception.

To avoid this issue, you can create the socket outside of the with block and pass it as an argument to the stream function. This way, the socket will not be closed when the with block is exited, and the stream function can continue to use it.

To ensure that the socket gets closed when the program exits, you can use the atexit module in Python. The atexit module provides a way to register functions that will be called when the program is about to exit.

You can register a cleanup function that closes the socket using the atexit.register() function. This function takes the cleanup function as an argument and registers it to be called when the program exits.

Here's an updated version of your code that avoids closing the socket prematurely and ensure that the socket gets closed when the program exits:

import socket
import threading
import time
import atexit

def stream(message, socket_):
    while True:
        socket_.sendall(message)
        time.sleep(1)

def cleanup(sock):
    sock.close()

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect((address, port))

atexit.register(cleanup, sock)

thread = threading.Thread(target=stream, args=("hello world!", sock), daemon=True)
thread.start()

In this code, the sock socket object is created outside of the with block. Then, the stream function is called with the sock object as an argument. The context is not exited, so the socket will not be closed prematurely.

Additionally, I have used the start() method of the thread object to start the thread.

The cleanup() function is defined to close the socket. The atexit.register() function is used to register the cleanup() function with the sock object as an argument. This ensures that the cleanup() function will be called when the program exits.