go atexit源码分析

atexit源码解析

当我们在执行程序的时候如果想要在退出程序的时候，执行一些清理函数或者日志输出函数等，那么atexit将会是一个很好选择！可以方便地在程序结束之前执行对应的函数。本文将会对go语言中的atexit的源码进行详细的分析，并且给出了对应的使用示例方便大家学习！

源码地址: atexit

UML类图

样例一: 程序退出之前执行注册函数

1.1 流程图

1.2 代码分析

package main

import (
	"fmt"

	"github.com/tebeka/atexit"
)

func handler() {
	fmt.Println("Exiting")
}

func main() {
	atexit.Register(handler)
	atexit.Exit(0)
}
// output:
// Exiting

首先，我们首先来解析atexit.Register(handler)这个函数都干了些什么：

var (
    handlersLock  sync.RWMutex  // protects the above two
    nextHandlerID uint   // 注册函数唯一的标识符
    handlers      = make(map[HandlerID]func())  // 使用 make 函数创建了一个空的映射。这个初始化会在程序启动时执行，确保了 handlers 在使用前已经被正确地分配和初始化
)
type HandlerID uint 
/****** Register函数 *****/
func Register(handler func()) HandlerID {
    // 由于handlersLock是一个全局互斥锁 (handlersLock)，用于在并发执行中保护nextHandlerID这个全局变量，确保不会发生竞争条件
        handlersLock.Lock()
    // 确保函数结束的时候释放锁，避免死锁的情况
	defer handlersLock.Unlock()
    // nextHandlerID代表注册函数唯一的标识符，要确保其唯一性，因此要加锁
	nextHandlerID++
    // 将id转换为HandlerID类型
	id := HandlerID(nextHandlerID)
    // handlers是一个全局映射，键是注册函数的唯一ID，而值为注册函数，该map用于存储所有的注册函数
	handlers[id] = handler
     // 返回handler的唯一ID
	return id
}

​ 由此可以得知，当我们使用atexit.Register(handler)之后，就会将handler函数成功注册到全局映射handlers中了，之后就可以通过全局handlers来处理注册函数了。

​ 然后，让我们继续看atexit.Exit(0)做了些什么：

var (
    ...
    once sync.Once // sync.Once可以确保在并发程序中某个函数只执行一次，无论它被多次调用
)

func runHandler(handler func()) {
	defer func() {
        // 使用了 recover() 函数来捕获可能的 panic。如果发生 panic，将错误信息输出到标准错误流（os.Stderr），之后会立刻结束这个goroutine，但是不会结束整个程序的运行，这样做可以避免整个程序崩溃.
		if err := recover(); err != nil {
			fmt.Fprintln(os.Stderr, "error: atexit handler error:", err)
		}
	}()
	// 调用实际传入的函数，在本例中会直接调用func handler()
	handler()
}

func executeHandlers() {
    // 使用了读锁（RLock）和读锁解除（RUnlock），用于在并发执行中保护对全局变量 handlers 的读取操作。
	handlersLock.RLock()
	defer handlersLock.RUnlock()
    // 读取已经注册过的所有handler并且执行它们，这个操作是并发安全的
	for _, handler := range handlers {
		runHandler(handler)
	}
}

func runHandlers() {
	// 无论runHandlers函数被调用多少次，在同一个程序运行周期内，executeHandlers函数只会被执行一次。
    // 确保当并发调用runHandlers函数的时候，所有的注册函数只会执行一次
	once.Do(executeHandlers)
}

/****** Exit函数 *****/
// Exit runs all the atexit handlers and then terminates the program using
// os.Exit(code)
func Exit(code int) {
	runHandlers()
	os.Exit(code)
}

​ 由此可知，当我们调用atexit.Exit(0)的时候，程序会先执行所有的注册函数，之后才会调用os.Exit退出整个程序。

样例二：使用cancel取消注册函数

2.1 cancel流程图

2.2 代码分析

当我们不想执行注册函数的时候，可是函数又已经注册了，那么就可以使用cancel取消注册函数的执行。

package main

import (
	"fmt"

	"github.com/tebeka/atexit"
)

func handler() {
	fmt.Println("handler Exiting")
}
func handler2() {
	fmt.Println("handler2 Exiting")
}
func main() {
	id := atexit.Register(handler)
	err := id.Cancel()
	if err != nil {
		fmt.Println("Error cancel")
		return
	}
    atexit.Register(handler2)
	atexit.Exit(0)
}
// Output:
// handler2 Exiting

​ 让我们看看 err := id.Cancel()干了些什么：

// Cancel cancels the handler associated with id
func (id HandlerID) Cancel() error {
	handlersLock.Lock()
	defer handlersLock.Unlock()
	// 检查是否存在对应id的注册函数。如果handlers中不存在该 id，则返回一个包含错误信息的 error 类型。也就是说只能够取消已经注册函数，不能够取消未被注册的函数。
	_, ok := handlers[id]
	if !ok {
		return fmt.Errorf("handler %d not found", id)
	}
	// 删除handlers对应id的handler注册函数
	delete(handlers, id)
	return nil
}

​ 可以看到，其实cancel函数也只是把对应id的注册函数从hanlders中移除而已，之后执行注册函数的时候就不会执行该函数了。

样例三：使用Fatal/Fatalln/Fatal执行注册函数

3.1 Fatal/Fatalln/Fatal流程图

3.2 代码分析

package main

import (
	"fmt"

	"github.com/tebeka/atexit"
)

func handler() {
	fmt.Println("Exiting")
}

func main() {
	atexit.Register(handler)
	// 以下三个语句只能够执行一个其中一个，因为执行完对应的语句就会退出程序。
	// atexit.Fatal("this is a Fatal message")
	// atexit.Fatalf("this is a Fatalf message:%s", "fatalf")
	atexit.Fatalln("this is a Fatalln message")
}

​ 当执行atexit.Fatal或者atexit.Fatalf或者atexit.Fatalln会首先执行注册函数，之后才会执行对应的Fatal/Fatalf/Fatalln函数来退出程序。

​ 让我们看看这三个函数都干了些什么:

// Fatal runs all the atexit handler then calls log.Fatal (which will terminate
// the program)
func Fatal(v ...interface{}) {
	runHandlers()
	log.Fatal(v...)
}

// Fatalf runs all the atexit handler then calls log.Fatalf (which will
// terminate the program)
func Fatalf(format string, v ...interface{}) {
	runHandlers()
	log.Fatalf(format, v...)
}

// Fatalln runs all the atexit handler then calls log.Fatalln (which will
// terminate the program)
func Fatalln(v ...interface{}) {
	runHandlers()
	log.Fatalln(v...)
}

​ 可以看到，它们首先通过runHandlers()调用了所有的已经注册函数，之后再通过log.Fatal/Fatalf/Fatalln来输出日志，并且最终退出程序。
注：log.Fatal/Fatalf/Fatalln在输出日志之后会自动终止程序的运行。

到此，atexit源码分析就结束了！若文章中出现任何纰漏，欢迎大家指正批评哦！如果觉得写得还不错的话，麻烦大家点赞收藏加关注哦！