Abhinav Gupta | About

Go Antipatterns: With* Context Managers

Python’s Context Managers are a useful feature which allow you to open and close resources cleanly.

with db.connect() as conn:
  # ...

The idiomatic way of doing the same in Go is with the use of defer statements.

conn, err := db.Connect()
if err != nil {
  // ...
}
defer conn.Close()

This suffices except when you are unable to add the cleanup method (Close) to the retuned object (conn).

An example of this that I run into often is the need to create a temporary directory and deleting it afterwards. Let’s look at the documentation for ioutil.TempDir.

TempDir creates a new temporary directory […] and returns the path of the new directory. […] the caller’s responsibility to remove the directory when no longer needed.

Since TempDir returns a string, adding a Close or Delete method is not an option. It would be easy for one to fall back to a Python-like pattern here.

// WithTempDir creates a temporary directory and calls the given
// function with the path to it. The directory is cleaned up when
// the function returns.
//
// The test is marked as failed if the temporary directory could
// not be created.
func WithTempDir(t *testing.T, f func(dir string)) {
  dir, err := ioutil.TempDir("", "test")
  if err != nil {
    t.Fatalf("failed to create temporary directory: %v", err)
  }
  defer os.RemoveAll(dir)
  f(dir)
})

Usage of the above function looks like so,

func TestFoo(t *testing.T) {
  WithTempDir(t, func(dir string) {
    // ...
  })
}

1. Problem

With the context manager, each use requires nesting your code further. Your business logic will run off the screen if you need to use more context managers.

WithTempDir(t, func(dir string) {
  WithTempFile(t, dir, func(file string) {
    var mu sync.Mutex
    WithLock(mu, func() {
      // ...
    })
  })
})

State management becomes fragile because returning values requires mutating variables in the outer scope. The functions accepted by the With* functions must have a specific signature. The inability to change the signature means that to send a value from the inner function to the caller, you must mutate variables in the outer scope.

var result *User
WithTempDir(t, func(dir string) {
  // ...

  var err error
  result, err = parseFile(filepath.Join(dir, "user"))
  if err != nil {
    // ...
  }
})

As the number of variables and context managers increases, this code will become more unreadable and more prone to bugs.

Finally, you cannot dynamically nest calls. If you need a list of N temporary directories where N is not fixed, there is no good way to use WithTempDir N times.

That is, the following is not possible.

WithTempDir(t, func(dir1 string) {
  WithTempDir(t, func(dir2 string) {
    WithTempDir(t, func(dir3 string) {
      // ?
      WithTempDir(t, func(dirN string) {
      })
    })
  })
})
📝 Note
 Python’s Context Managers solve this at the library level with contextlib.ExitStack.

2. Solution

We want to be able to defer cleanup functions but adding a Close() or Delete() method to string is not possible. How about returning an anonymous function that does the cleanup?

// CreateTempDir creates a temporary directory and returns the path
// to it, along with a function to clean up when the directory is
// no longer needed.
//
// The test is marked as failed if the temporary directory could
// not be created.
func CreateTempDir(t *testing.T) (dir string, cleanup func()) {
  dir, err := ioutil.TempDir("", "test")
  if err != nil {
    t.Fatalf("failed to create temporary directory: %v", err)
  }
  return dir, func() {
    os.RemoveAll(dir)
  }
}

That looks better! Here is how you use it:

func TestFoo(t *testing.T) {
  dir, cleanup := CreateTempDir(t)
  defer cleanup()

  // ...
}

We can call any number of similar functions without affecting readability of the business logic.

dir, cleanup := CreateTempDir(t)
defer cleanup()

file, cleanup := CreateTempFile(t, dir)
defer cleanup()

// ...

It’s easy to call them dynamically and defer cleanup in a loop.

dirs := make([]string, N)
for i := 0; i < N; i++ {
    dir, cleanup := CreateTempDir(t)
    defer cleanup() (1)
    dirs[i] = dir
}

  1. Note that deferred calls run when the surrounding function returns, not when the scope ends. In other words, the temporary directories will be cleaned up not when a single for iteration ends, but when the function returns.

3. Conclusion

This is a simple pattern and as demonstrated, it’s more readable, flexible, and idiomatic than Context Manager-style With* functions. It has found frequent use in our codebase for test setup and teardown, and sometimes outside tests as well.

Written on .