在Go中使用defer减少重复代码

881 阅读5分钟

在Go中,error是经常被函数返回的。如果在调用函数时,返回的是err,就必须要处理,处理的err一旦多了起来,如果对资源处理不当,就容易出错。我们就来看看在实际程序中如何使用defer来简化我们的程序。

首先,我们看一个没有使用defer语句来关闭资源的例子。我们实现了一个拷贝文件的函数。我们还将管理该文件描述符的闭包,因为一个 *os.File一旦被打开准备读写时,它就必须要使用Close函数进行关闭。最后,在函数的最后,我们将使用Sync方法来刷新文件系统的缓冲区以便将内容强制写到磁盘上,使副本持久化。下面是具体的实现:

func CopyFile(srcName, dstName string) error {
    src, err := os.Open(srcName) ①
    if err != nil {
        return err
    }
		
    stat, err := src.Stat()
    if err != nil {
        src.Close()
	return err
    }
		
    if stat.IsDir() { ②
        src.Close()
        return fmt.Errorf("file %q is a directory", srcName)
    }
    
    dst, err := os.Create(dstName) ③
    if err != nil {
        src.Close()
	return err 
    }
		
    _, err = io.Copy(dst, src) ④
    if err != nil {
        src.Close()
        dst.Close()
        return err
    }
		
    err = dst.Sync() ⑤
    src.Close()
    dst.Close()
    return err
}

① 打开源文件

② 检查是否是目录

③ 创建目标文件

④ 拷贝源文件到目标文件

⑤ 刷新文件系统缓冲区

注意: 关闭*os.File将会返回一个错误。然而,在该例中该错误可以被安全的忽略,因为我们强制刷新了文件系统的缓冲区。否则,如果错误发生时,我们至少应该记录一条日志。在错误管理一章,我们将会看到在defer语句中如何优雅地处理错误。

这个实现是可以工作的。我们打开源文件,检查是否是目录,然后处理拷贝逻辑。然而,我们注意到一些模板式的代码:

  • src.Close()重复了5次
  • dst.Close()重复了2次

在代码中必须考虑源文件和目标文件在最终都要关闭,这使得我们的代码非常容易出错。幸运的是,Go通过defer关键词提供了一种解决该问题的方案,如图2.1:

2.1.jpeg

在函数返回的时候会调用defer函数。即使在主函数panics或意外终止时defer函数也能保证被执行。调用defer会被推送到栈中。当主函数返回时,defer函数会从栈中弹出(先进后出的顺序)。这里,将会先调用c(),然后b(),最后是a()。

注意:一个defer调用的时机是在函数返回时,而非在所在的块退出时。 如下:

func main() {
    fmt.Println("a")
    if true {
        defer fmt.Print("b")
    }
    fmt.Print("c")
}
该段代码打印结果是acb,而非abc。

让我们回到CopyFile函数的例子并使用defer关键词再实现一版:

func CopyFile(srcName, dstName string) error {
    src, err := os.Open(srcName)
    if err != nil {
	return err 
    }
    defer src.Close() ①
    
    stat, err := src.Stat()
    if err != nil {
        return err 
    }
		
    if stat.IsDir() {
        return fmt.Errorf("file %q is a directory", srcName)
    }
		
    dst, err := os.Create(dstName)
    if err != nil {
        return err 
    }
    defer dst.Close() ②
    
    _, err = io.Copy(dst, src)
    if err != nil {
	return err
    }
		
    return dst.Sync()
}

① 延迟调用 src.Close()

② 延迟调用 dst.Close()

在这个版本的实现中,我们通过defer关键词的使用,移除了重复的close调用。这使得函数更轻量并且更易读。我们不必在每一个代码路径的末尾都关闭src和dst,这样就不容易出错了

defer语句经常会跟成对出现的操作函数一起使用,就像 open/close,connect/disconnect,以及lock/unlock函数以确保在所有的场景下资源都能够得到释放。

这是另一个使用sync.Mutex的例子:

type Store struct {
    mutex sync.Mutex
    data map[string]int
}

func (s *Store) Set(key string, value int) {
    s.mutex.Lock() ①
    defer s.mutex.Unlock() ②
	
    s.data[key] = value
}

① Mutex lock

② 在defer语句中unlock

我们使用s.mutex.Lock函数锁定了mutex并在defer中调用s.mutex.Unlock()函数的配对操作。

注意:如果我们必须实现一个pre和post操作,比如不返回任何值的 mutex lock/unlock,我们也可以这样实现:

func (s *Store) Set(key string, value int) {
    defer s.lockUnlock()() ①
    s.data[key] = value
}

func (s *Store) lockUnlock() func() {
    s.mutex.Lock()
    return func() {
        s.mutex.Unlock()
    }
}

① 会立刻执行s.lockUnlock(),但会延迟执行s.lockUnlock()()

延迟执行的函数调用是s.lockUnlock()(),而非s.lockUnlock()。 因此,s.lockUnlock()部分会立即执行(s.mutex.Lock),但是返回 的闭包会操作被延迟执行(s.mutex.Unlock())。它添加了一些语法 糖,用一行代码来处理函数中的前/后操作,这有时非常方便。 如果使用这种模式,还需要注意的是,面对带有两组括号的 s.lockUnlock()() 可能会非常混乱,这取决于您团队的资历。

当重构代码时,我们还需要注意可能的影响。例如,假设我们需要将一个包含defer调用的main函数拆分成多个函数时,这种情况下,一旦应用程序执行完,defer语句并不会执行,但是当子函数执行完时defer调用才会调用:

// Before
func main() {
    consumer := createConsumer()
    defer consumer.Close() ①
    // ...
}

// After
func main() {
    consumer := handleConsumer()
    // ...
}

func handleConsumer() Consumer {
    consumer := createConsumer()
    defer consumer.Close() ②
    return consumer
}

① 一旦程序结束defer调用将会被执行

② 当handleConsumer函数结束defer调用才会被执行

这里,我们通过重构consumer引入了一个bug。因为一旦handleConsumer函数结束,consumer.Close()就会被执行。它看起来像一个简单的注释,但是当我们必须重构大量的代码时,有时很容易忽略defer语句。

同时也需要注意Go 1.14之前的版本,defer语句不是内联的。内联是编译器通过将函数调用直接保存在调用函数中的一种优化技术。这就是为什么在一些性能是关键因素的项目中,defer关键词很少被用到的原因。但是,在Go 1.14版本之后,defer语句可以通过内联来优化了。

总之,defer可以避免死板的代码以及减少忘记释放资源的风险,例如释放资源,断开链接,mutex解锁等等。