熟悉 Python 的同學(xué)應(yīng)該知道 Python 中的上下文管理器非常好用，在對數(shù)據(jù)庫進行讀寫、訪問文件等操作時，上下文管理器能夠確保資源在使用后得到釋放。在 Go 中是否也能實現(xiàn)上下文管理器呢？這便是本文所要探討的話題。

Python 上下文管理器

以操作文件為例，為了保證操作文件完成后資源能被正確關(guān)閉，在 Python 中我們可以編寫出如下代碼：

try:
    f = open('foo.txt', 'r')
    print(f.readlines())
finally:
    f.close()

不過這種寫法顯然不夠 Pythonic，Python 在語法層面提供了 with 語句實現(xiàn)上下文管理，用法如下：

with open('foo.txt', 'r') as f:
    print(f.readlines())

這段使用 with 語句實現(xiàn)的代碼，才更符合 Python 哲學(xué)。

如果你對 Python with 語法不熟悉，可以參閱我的文章《Python 上下文管理器實現(xiàn)》。

Go 中資源釋放問題

我們知道，在 Go 語言中訪問數(shù)據(jù)庫、文件等資源時，可以使用 defer 語句完成資源釋放操作。

如下定義一個 ReadFile 函數(shù)用來讀取文件：

func ReadFile(paths []string) error {
	for _, path := range paths {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		content, err := io.ReadAll(file)
		if err != nil {
			return err
		}
		fmt.Printf("%s content: %s\n", file.Name(), content)
	}
	return nil
}

這個函數(shù)使用循環(huán)遍歷傳進來的文件路徑列表，依次打開文件并輸出文件內(nèi)容。

為了保證即使在遇到錯誤時，資源也能夠被釋放，我們往往會使用 defer file.Close() 來關(guān)閉文件。

不過，這段代碼其實是存在問題的，我們知道 defer 的調(diào)用實際上并不會立即執(zhí)行，而是等到函數(shù)退出時才會執(zhí)行。

所以，代碼中的 defer 調(diào)用并不會在本輪循環(huán)中處理完當前文件時被執(zhí)行，而是直到所有循環(huán)執(zhí)行完成，函數(shù)退出時才會執(zhí)行。

我們可以對以上示例稍作修改，來驗證下這個問題：

func ReadFile(paths []string) error {
	for _, path := range paths {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer func() {
			file.Close()
			fmt.Printf("close %s\n", file.Name())
		}()
		content, err := io.ReadAll(file)
		if err != nil {
			return err
		}
		fmt.Printf("%s content: %s\n", file.Name(), content)
	}
	return nil
}

我們將原來的 defer 語句改成：

defer func() {
    file.Close()
    fmt.Printf("close %s\n", file.Name())
}()

以此來顯示 defer 調(diào)用時機。

針對以上示例，我們使用如下代碼來調(diào)用：

func main() {
	err := ReadFile([]string{"foo.txt", "bar.txt"})
	fmt.Printf("ReadFile err: %v\n", err)
}

注意：foo.txt、bar.txt 兩個文件我已經(jīng)提前準備好了，foo.txt 文件內(nèi)容為 foo，bar.txt 文件內(nèi)容為 bar。

執(zhí)行以上示例，得到如下輸出：

$ go run main.go
foo.txt content: foo
bar.txt content: bar
close bar.txt
close foo.txt
ReadFile err: <nil>

根據(jù)輸出內(nèi)容可以驗證，defer 語句的調(diào)用，的確在 for 循環(huán)退出以后才開始執(zhí)行。

如果打開資源過多，而沒有及時關(guān)閉，勢必會造成資源的浪費，甚至因此而意外終止程序。

所以切記，不要在循環(huán)中使用 defer。

我們可以使用匿名函數(shù)來解決這個問題：

func ReadFile(paths []string) error {
	for _, path := range paths {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		err = func() error {
			defer func() {
				file.Close()
				fmt.Printf("close %s\n", file.Name())
			}()
			content, err := io.ReadAll(file)
			if err != nil {
				return err
			}
			fmt.Printf("%s content: %s\n", file.Name(), content)
			return nil
		}()
		if err != nil {
			return err
		}
	}
	return nil
}

現(xiàn)在，將 defer 語句放入到一個立即執(zhí)行的匿名函數(shù)中，就可以解決問題了。

執(zhí)行以上示例，得到如下輸出：

$ go run main.go
foo.txt content: foo
close foo.txt
bar.txt content: bar
close bar.txt
ReadFile err: <nil>

可以發(fā)現(xiàn)，現(xiàn)在 defer 語句不再是等到 for 循環(huán)退出才會執(zhí)行，而是在匿名函數(shù)退出時即可執(zhí)行。

這樣，就達到了在本輪循環(huán)中盡早釋放不再使用的文件資源的目的。

此外，為了代碼的可讀性，我們可以將匿名函數(shù)提取出來，單獨封裝一個函數(shù)：

func ReadFile(paths []string) error {
	for _, path := range paths {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		err = processFile(file)
		if err != nil {
			return err
		}
	}
	return nil
}
func processFile(file *os.File) error {
	defer func() {
		file.Close()
		fmt.Printf("close %s\n", file.Name())
	}()
	content, err := io.ReadAll(file)
	if err != nil {
		return err
	}
	fmt.Printf("%s content: %s\n", file.Name(), content)
	return nil
}

processFile 函數(shù)專門用來處理打開的文件，ReadFile 函數(shù)可讀性也得到了提高。

執(zhí)行以上示例，得到如下輸出：

go run main.go
foo.txt content: foo
close foo.txt
bar.txt content: bar
close bar.txt
ReadFile err: <nil>

這個輸出符合預(yù)期。

以上我們介紹了兩種方式，能夠解決 defer 語句延遲調(diào)用的問題。

在 Go 中實現(xiàn)上下文管理器

最近為了寫《Go 語言中 database/sql 是如何設(shè)計的》一文，我閱讀了下 database/sql 的源碼。在這個過程中，*sql.DB.queryDC 方法中一小段代碼激起了我的興趣：

func (db *DB) queryDC(ctx, txctx context.Context, dc *driverConn, releaseConn func(error), query string, args []any) (*Rows, error) {
	...
	if ok {
		var nvdargs []driver.NamedValue
		var rowsi driver.Rows
		var err error
		withLock(dc, func() {
			nvdargs, err = driverArgsConnLocked(dc.ci, nil, args)
			if err != nil {
				return
			}
			rowsi, err = ctxDriverQuery(ctx, queryerCtx, queryer, query, nvdargs)
		})
		...
	}
	...
}

在 *sql.DB.queryDC 方法中有一個 withLock 函數(shù)的調(diào)用，withLock 函數(shù)定義如下：

func withLock(lk sync.Locker, fn func()) {
	lk.Lock()
	defer lk.Unlock()
	fn()
}

當看到 withLock 函數(shù)定義時，我瞬間就想到了 Python 中的 with 上下文管理器。

withLock 接收一個 sync.Locker 接口，定義如下：

type Locker interface {
	Lock()
	Unlock()
}

它只有兩個方法，加鎖和釋放鎖。

withLock 能夠用于所有實現(xiàn) sync.Locker 接口的對象，在執(zhí)行 fn() 前加鎖，執(zhí)行之后釋放鎖。

這與 Python 的上下文管理器功能如出一轍，就是這么一個只有三行的小函數(shù)，實現(xiàn)卻相當精妙，真可謂短小精悍。

于是，參考 withLock 函數(shù)實現(xiàn)，解決 for 循環(huán)中defer 語句延遲調(diào)用的問題，就有了第三種解法。

我們可以模仿 withLock 實現(xiàn)一個 WithClose 函數(shù)：

func WithClose(closer io.Closer, fn func()) {
	defer func() {
		closer.Close()
		fmt.Printf("close %s\n", closer.(*os.File).Name())
	}()
	fn()
}

WithClose 接收一個 io.Closer 接口，定義如下：

type Closer interface {
	Close() error
}

我們可以在執(zhí)行 fn() 函數(shù)之前，使用 defer 語句來調(diào)用 io.Closer 的 Close 方法釋放資源。

現(xiàn)在，我們可以在 ReadFile 函數(shù)中使用這個小函數(shù)了：

func ReadFile(paths []string) error {
	for _, path := range paths {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		WithClose(file, func() {
			var content []byte
			content, err = io.ReadAll(file)
			if err != nil {
				return
			}
			fmt.Printf("%s content: %s\n", file.Name(), content)
		})
		if err != nil {
			return err
		}
	}
	return nil
}

這個用法同 *sql.DB.queryDC 中調(diào)用 withLock 函數(shù)一樣，并且因為閉包的存在，我們可以拿到 WithClose 內(nèi)部執(zhí)行的 fn() 函數(shù)所產(chǎn)生的錯誤對象。

執(zhí)行以上示例，得到如下輸出：

$ go run main.go
foo.txt content: foo
close foo.txt
bar.txt content: bar
close bar.txt
ReadFile err: <nil>

這個輸出依然符合預(yù)期。

我們可以測試下遇到錯誤的情況，修改 main 函數(shù)，調(diào)用 ReadFile 時最后傳入一個不存在的文件 baz.txt：

func main() {
	err := ReadFile([]string{"foo.txt", "bar.txt", "baz.txt"})
	fmt.Printf("ReadFile err: %v\n", err)
}

執(zhí)行以上示例，得到如下輸出：

$ go run main.go
foo.txt content: foo
close foo.txt
bar.txt content: bar
close bar.txt
ReadFile err: open baz.txt: no such file or directory

遇到錯誤能夠被正常捕獲。

現(xiàn)在，我們就在 Go 中實現(xiàn)類了似 Python 中的 with 上下文管理器，為解決 for 循環(huán)中defer 語句延遲調(diào)用的問題提供了新思路。

總結(jié)

本文靈感來自于 database/sql 源碼中的一小段代碼，為大家講解了如何在 Go 中實現(xiàn)類似 Python 中的 with 上下文管理器。

切記，不要在循環(huán)中使用 defer。為了解決這個問題，我們可以使用匿名函數(shù)、函數(shù)封裝以及 WithClose 三種方案。

希望此文能對你有所幫助。

P.S.

database/sql 源碼中的這一小段代碼，找回了我在開始用 Go 作為主力語言后，很久沒有在編程語言語法層面上體會過快感。相較于我最近寫的幾篇長篇大論型文章，本文顯得微不足道，但我還是很樂于為這一小段代碼寫一篇文章分享出來，畢竟這久違的感覺又回來了。

從把 Go 作為主力編程語言開始，寫代碼的思路都是“平鋪直敘”，很少思考怎么寫出更加優(yōu)雅且有趣的代碼。盡管我也分享過幾篇 Go 編程模式的文章，但相較于用 Python 作為主力編程語言時，還是少了很多“花哨”的小技巧在里面，更多的是遵循套路的樣板代碼。

盡管 Go 語言的哲學(xué)更適合工程化，但 Go 代碼寫多了，有時不免會略感乏味，懷念 Python 的靈活。我無意于討論哪種編程語言的好壞，只是，愿在編程的道路上，你我都能找到屬于自己的樂趣所在。

以上就是詳解Go語言如何實現(xiàn)類似Python中的with上下文管理器的詳細內(nèi)容，更多關(guān)于Go語言上下文管理器的資料請關(guān)注腳本之家其它相關(guān)文章！