From 850b678c134a32039662a87266d2c39d7d4fba1c Mon Sep 17 00:00:00 2001
From: Ge Song <gesong0606@gmail.com>
Date: Tue, 25 Apr 2023 13:17:10 +0800
Subject: [PATCH] =?UTF-8?q?log/=E6=89=A9=E5=B1=95Exit=E5=8A=9F=E8=83=BD;?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 log/README.md        |  1 +
 log/log.go           |  5 +++++
 log/logsdk/logger.go | 21 ++++++++++++++++++---
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/log/README.md b/log/README.md
index a8cd764..24c5005 100644
--- a/log/README.md
+++ b/log/README.md
@@ -54,6 +54,7 @@ log.Logger()
 - `SetReportStack` 设置生成调用栈;
 - `SetReportStackLevel` 当日志等级小于设定值时强制生成调用栈;
 - `Reset` 把 Logger 恢复到初始状态;
+- `AddBeforeExit` 新增调用 `Logger.Exit` 时优先执行的函数, 先增加的后执行;
 - `SetExit` 设置 `Logger` 的退出函数(`Logger.Exit`), 当日志等级为 `LevelFatal` 时调用这个函数;
 
 ### 日志生成
diff --git a/log/log.go b/log/log.go
index 8063a12..9884674 100644
--- a/log/log.go
+++ b/log/log.go
@@ -150,6 +150,11 @@ func Panicf(ctx context.Context, format string, args ...any) {
 	globalLogger.AddCallerSkip(1).Panicf(ctx, format, args...)
 }
 
+// Exit 调用全局 logger 的 Exit
+func Exit(code int) {
+	globalLogger.Exit(code)
+}
+
 // Logger 返回全局 logger, 通常用来在程序启动时对全局 logger 进行配置,
 // 业务代码处理日志时直接使用这个包里定义的日志函数.
 func Logger() *logsdk.Logger {
diff --git a/log/logsdk/logger.go b/log/logsdk/logger.go
index ab612bb..a9804a8 100644
--- a/log/logsdk/logger.go
+++ b/log/logsdk/logger.go
@@ -20,8 +20,9 @@ type entry = Entry
 // Logger 中保存了日志所需的全局配置,
 // 使用 Logger 处理日志.
 type Logger struct {
-	exit            func(code int)  // protected by lock
-	levelProcessors levelProcessors // protected by lock
+	exit            func(code int)   // protected by lock
+	beforeExit      []func(code int) // protected by lock
+	levelProcessors levelProcessors  // protected by lock
 	entry
 	lock             sync.RWMutex
 	level            atomic.Int32
@@ -98,6 +99,7 @@ func (logger *Logger) SetReportStackLevel(level Level) {
 func (logger *Logger) Reset() {
 	logger.lock.Lock()
 	logger.exit = nil
+	logger.beforeExit = nil
 	logger.levelProcessors = levelProcessors{}
 	logger.lock.Unlock()
 
@@ -112,9 +114,14 @@ func (logger *Logger) Reset() {
 func (logger *Logger) Exit(code int) {
 	logger.lock.RLock()
 	exit := logger.exit
+	beforeExit := make([]func(int), len(logger.beforeExit))
+	copy(beforeExit, logger.beforeExit)
 	logger.lock.RUnlock()
+	for i := len(beforeExit) - 1; i >= 0; i-- {
+		beforeExit[i](code)
+	}
 	if exit == nil {
-		os.Exit(code)
+		exit = os.Exit
 	}
 	exit(code)
 }
@@ -126,6 +133,14 @@ func (logger *Logger) SetExit(fn func(code int)) {
 	logger.lock.Unlock()
 }
 
+// AddBeforeExit 增加 Exit 在调用 SetExit 指定的函数前执行的函数,
+// 先增加的后执行.
+func (logger *Logger) AddBeforeExit(fn ...func(code int)) {
+	logger.lock.Lock()
+	logger.beforeExit = append(logger.beforeExit, fn...)
+	logger.lock.Unlock()
+}
+
 func (logger *Logger) getLevelProcessors(level Level) []EntryProcessor {
 	logger.lock.RLock()
 	defer logger.lock.RUnlock()