Recently, I was going through the documentation and codebase of a python package (specifically kafka-python) and was curious about how cleaning up resources and sending pending messages (messages that are still in-memory and haven’t been sent to the Kafka brokers) from a Kafka producer.

What I found is that the specific (and several other libraries) are using the atexit (built-in) python module (check out the official documentation).

The atexit module allows you to register handlers (functions) that run automatically upon interpreter termination. As the official documentation mentions:

The functions registered are not called when the program is killed by a signal not handled by Python, when a Python fatal internal error is detected, or when os._exit() is called

This means that the handlers registered with atexit run on normal termination of a python script (e.g. after sys.exit(0)), but not when the interpreter crashes, or a signal that is not handled by python (e.g. via a SIGTERM by doing kill -15 script_pid - for this we can use the built-in module signal, which we’ll see in a future post - stay tuned :-) ). os._exit is is similar to sys.exit, but does not run cleanup handlers.

Registering cleanup functions

So, let’s see how we can use atexit in practice:

1
import atexit
2
import time
3
import sys
4

5
def exit_handler():
6
  print('Message from exit handler...')
7

8
atexit.register(exit_handler)
9

10
print('Will sleep for 1 second...')
11
time.sleep(1)

At line 5, we create the handler (which, in this simple example, just prints out a message), and we register it with atexit.register. Alternatively, we can use atexit.register as a decorator like that:

@atexit.register
def exit_handler():
  print('Message from exit handler...')

If we run the script, we’ll get the following output:

Will sleep for 1 second...
Message from exit handler...

Once the script ends (after printing a message and sleeping for 1 second), the atexit handler kicks in and prints its message. Multiple handlers can be registered and upon script termination they will run in reverse order. Here’s an example:

1
import atexit
2
import time
3
import sys
4

5
@atexit.register
6
def exit_handler_1():
7
  print('Message from exit handler 1...')
8

9
@atexit.register
10
def exit_handler_2():
11
  print('Message from exit handler 2...')
12

13
print('Will sleep for 1 second...')
14
time.sleep(1)

Running the updated script, we’ll get the following output:

Will sleep for 1 second...
Message from exit handler 2...
Message from exit handler 1...

Unregistering a cleanup function

Building on top of the example above, we can unregister a cleanup function after it’s been registered:

1
import atexit
2
import time
3
import sys
4

5
@atexit.register
6
def exit_handler():
7
  print('Message from exit handler 1...')
8

9
atexit.unregister(exit_handler)
10

11
print('Will sleep for 1 second...')
12
time.sleep(1)

Running the updated script, we’ll get the following output (in which you can see that the cleanup function that has been unregistered doesn’t run at all):

Will sleep for 1 second...

Real-world examples of cleanup functions using atexit can be found in various libraries (e.g. check out kafka-python’s producer implementation) and usually are doing some kind of cleanup (e.g. closing Kafka producers, sending pending log messages that are stored in-memory until they reach a batch size or every X amount of time, etc).

That’s all for now!