Python模块atexit的应用。

Python中atexit模块可以注册一些列函数，让Python解析器在恰当的方式退出前调用这些函数。这个功能很实用，例如我们实现“文件锁”功能，通过在文件系统中创建一个文件来表示进程正在运行且确保进行的唯一性。当进程退出时，删除该文件，以告知其他需要创建进行的服务。

另外还可以用于处理异常、日志、销毁进行运行过程的中间文件、临时文件、数据库连接等。

从简单的示例开始

atexit的注册函数的签名如下：

1
atexit.register(func, *args, **kwargs)

使用atexit.register注册一个进程退出时的回调。

1
2
3
4
5
6
7
8
9
10
import atexit

def all_done:
    print("process exit after execute atexit funcs")

def main():
    print("run main function")

if __name__ == '__main__'：
    atexit.register(all_done)

在命令行下运行：

1
$python3 atexit_1.py

输出：

1
2
3
$python atexit_1.py
run main function
process exit after execute atexit funcs

可以看到在main函数运行完退出前解析器调用了通过atexit.register处注册的函数。

接下来一个稍微复杂的例子。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import atexit

def atexit_func_1():
    print("func_1")

def atexit_func_2():
    print("func_2")

def atexit_func_3():
    print("func_3")

def atexit_unregister():
    print("func_unregister")

def main():
    print("main function")

if __name__ == '__main__':
    atexit.register(atexit_func_1)
    atexit.register(atexit_func_2)
    atexit.register(atexit_func_3)
    atexit.register(atexit_unregister)
    print("unregister func:", atexit_unregister.__name__)
    atexit.unregister(atexit_unregister)
    main()

该段代码注册了多个函数，然后取消注册func_unregister函数，使其从回调栈中删除。

程序执行和输出：

1
2
3
4
5
6
$ python3 atexit_2.py
unregister func: atexit_unregister
main function
func_3
func_2
func_1

通过输出发现，被注册的回调函数的执行次序是注册顺序的反序。如果在程序退出时要执行类事务的操作，需要注意这一点。同一个函数注册多次也会服从这个规律。

不调用atexit注册的函数

很多情况下确保atexit注册的函数被调用很重要，因此我们需要搞明白哪些情况下被注册的函数不会被调用。

它可以为函数func出入参数。不过下面的例子出于简洁性，并不会构造给指定函数出参数的例子。当然，也可以通过偏函数指定参数。

在程序退出时执行atexit模块注册的函数时有条件的。如果满足如下条件中的任意一条，则atexit模块不会执行回调函数。

  1. 程序调用os._exit()退出
  2. 程序被一个信号终止
  3. 检测到解析器的一个致命错误

接下来我们一一举例子。

程序调用os._exit()退出

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import atexit
import os

def all_done():
    print("this func should not be called")

def main():
    print("main func, I call os._exit(1) here")
    os._exit(1)

if __name__ == '__main__':
    print("register all_done func")
    atexit.register(all_done)
    main()

main函数中调用了os._exit(1)阻止了解析器正常退出。

运行和输出：

1
2
3
$python3 atexit_os._exit.py
register all_done func
main func, I call os._exit(1) here

如果要确保程序终止也能调用atexit函数，可以使用sys.exit函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import atexit
import sys

def all_done():
    print("this func should not be called")

def main():
    print("main func, I call os._exit(1) here")
    sys.exit(1)

if __name__ == '__main__':
    print("register all_done func")
    atexit.register(all_done)
    main()

程序被一个信号终止

程序在运行过程中可以通过信号来终止程序，典型的操作是使用kill命令杀掉程序。接下来这个例子是：父进程创建子进程，然后子进程暂停，等待接收父进程的信号。

我们的子进程代码（atexit_signal.py）如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import atexit
import time
import sys

def all_done():
    print("all done func but not call")

def main():
    print("run main func")
    sys.stdout.flush()
    time.sleep(5)

if __name__ == '__main__':
    atexit.register(all_done)
    main()

父进程代码（atexit_signal_parent.py）如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import os
import subprocess
import time
import signal

func main():
    child = subprocess.Popen('python atexit_signal.py')
    print("parent pausing before sending signal...")
    time.sleep(2)
    print("signaling to child")
    os.kill(child.pid, signal.SIGTERM)

if __name__ == '__main__':
    main()

运行和输出如下：

1
2
3
4
$python3 atexit_signal_parent.py
parent pausing before sending signal...
signaling to child
run main func

因此，子进程用于接收到父进程的信号而退出并没有调用注册的all_done函数。

检测到解析器的一个致命错误

致命的错误有很多情况。这些错误都导致解析器无法按照正常的执行流程退出，比如执行到操作系统调用的代码段，程序却没有拿到缺陷而退出。

atexit处理异常

atexit注册的回调函数中产生的异常会输出到控制台上（如果没有捕捉），最后参数的异常会重新抛出。一个简单的例子。

1
2
3
4
5
6
7
8
9
10
import atexit

def quit_exception():
    print()
    raise Exception("atexit with exception")

if __name__ == '__main__':
    atexit.register(quit_exception)
    atexit.register(quit_exception)
    print("start...")

运行和输出：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
$python3 atexit_traceback.py
start...

Error in atexit._run_exitfuncs:
Traceback (most recent call last):
  File "atexit_traceback.py", line 5, in quit_exception
    raise Exception("atexit with exception")
Exception: atexit with exception

Error in atexit._run_exitfuncs:
Traceback (most recent call last):
  File "atexit_traceback.py", line 5, in quit_exception
    raise Exception("atexit with exception")
Exception: atexit with exception

类比思考之Go语言

在Go语言中，实现Python的atexit模块功能，有defer功能。在Go语言中称为延时调用。但和Python不同的是，defer机制是在声明函数调用的函数执行完毕时调用。

即使函数执行出错了，延时调用也会执行。在Go编程开发中，它能保证资源回收操作被执行。由于延时操作并不是一个简单的CALL调用，会有一定的性能损耗，因此，如果代码对性能要求高，资源回收敏感，通过直接方法执行回收操作更好。

当然强大的Python就怎么能在这一点上认输呢？Python自带的sched模块可以实现定时事件调度器，一定程度上通过delay来弥补defer的不足。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
package main

import (
    "fmt"
)

func defer1() {
    fmt.Println("defer func 1")
}

func defer2() {
    fmt.Println("defer func 2")
}

func main() {
    defer defer1()
    defer defer2()
}

运行和输出：

1
2
3
$go run defer_demo.go
defer func 2
defer func 1

Go的defer特性也有一个小问题，看代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
package main

func defer1() {
    defer defer2()
}

func defer2() {
    defer defer1()
}

func main() {
    defer defer1()
}

如果在实际编程中没有注意这个细节：两个函数互相通过defer机制注册的对方。程序在时就会出现栈溢出。

1
2
3
4
5
6
7
8
9
10
11
12
$ go run defer_wrap_defer.go
runtime: goroutine stack exceeds 1000000000-byte limit
fatal error: stack overflow

runtime stack:
runtime.throw(0x46c030, 0xe)
        ~/go1.9/src/runtime/panic.go:605 +0x9c
        .
        .
        .
...additional frames elided...
exit status 2

和Python一样，被defer注册的函数是逆序调用的。Go语言的defer机制的比Python的atexit好的地方是其粒度更小：前者注册的函数精确到注册所在函数结束时调用；后者只能在程序退出时调用。

一个应用例子

接下来以一个在实际开发中使用的例子来体验atexit的应用。

Linux下编程常常需要程序以守护进程的方式启动，根据UNIX环境编程知识可知通常的做法（代码层面）就是fork以及会话、权限、I/O设置等。写多了感到挺无聊的，于是采用一个类的方式把这个过程封装一次，程序代码只要重写run方法即可，就像多线程编程中继承Thread类重写run方法。大致流程如下：

  1. fork 退出父进程 （在此前可以处理好缓冲区数据以免带到子进程中）
  2. 更改子进程目录、修改文件掩码、更新会话ID设置为首领
  3. fork 退出父进程
  4. I/O重定向设置
  5. 管理当前进程的pid（确保其唯一性）
  6. 登记进程终止时的处理函数（通常是删除pid文件、释放可能引起安全问题的资源）
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import os
import atexit
import signal
import sys

class Daemon:

    def __init__(self, *args, pidfile='./center.pid', stdin='/dev/null',
                                   stdout='/dev/null',
                                   stderr='/dev/null', **kwargs):
        self.pidfile = pidfile
        self.stdin = stdin
        self.stdout = stdout
        self.stderr = stderr
        self.args = args
        self.kwargs = kwargs

    def set_pidfile(self):
        # 保存pid到文件
        if os.path.exists(self.pidfile):
            with open(self.pidfile, 'r') as fp:
                pid = fp.read()
            raise RuntimeError("Process[{}] already running".format(pid))
        pid = os.getpid()
        with open(self.pid, 'w') as fp:
            print(pid, file=fp)

        atexit.register(lambda: os.remote(self.pidfile))

    def fork(self):
        try:
            if os.fork() > 0:
                raise SystemExit(0)
        except OSError as err:
            raise RuntimeError("fork #1 failed.")

        os.chdir("/")
        os.umask()
        # 父进程退出后，子进程称为孤儿。调用os.setsid()创建一个全新的会话。
        # 把子进程设置为leadership，确保没有与之关联的终端。
        # 
        os.setsid() 

        try:
            if os.fork() > 0: # 第二次fork使守护进程放弃获取新终端的能力
                raise SystemExit(0)
        except Exception as err:
            raise RuntimeError("fork #2 failed.")

    def set_stdio(self):
        # 重定向IO
        sys.stdout.flush()
        sys.stderr.flush()

        with open(self.stdin, 'rb', 0) as fp:
            os.dup2(fp.fileno(), sys.stdin.fileno())
        with open(self.stdout, 'ab', 0) as fp:
            os.dup2(fp.fileno(), sys.stdout.fileno())
        with open(self.stderr, 'ab', 0) as fp:
            os.dup2(fp.fileno(), sys.stderr.fileno())

    def _sigterm_handler(self, signo, frame):
        # 设置退出信号
        raise SystemExit(0)

    def set_signal(self):
        signal.signal(signal.SIGTERM, self._sigterm_handler)

    def start(self):
        self.set_pidfile()
        self.fork()
        self.set_stdio()
        self.set_signal()
        self.run()

    def run(self):
        # 继承的类重写该方法，运行需要的程序逻辑
        pass

该类的I/O重定向默认写到黑洞上。atexit函数会在程序退出时删掉daemon进程的pid文件。确保统一操作系统下只有一个这样的daemon进程。

这一部分可参看旧文Linux下Python创建守护进程通用类。

转载请包括本文地址：https://allenwind.github.io/blog/5417
更多文章请参考：https://allenwind.github.io/blog/archives/