How to consolidate cleanup logic using the "with" function pattern

A few of my teammates have liked the "with" function pattern and how convenient it is for ensuring proper cleanup of resources.

A "with" function takes a callback and performs some initialization, calls the callback, then performs some cleanup.

For example, instead of manually doing a COMMIT or ROLLBACK for every SQL transaction like this:

tx, err := db.Begin()
if err != nil {
    return err
}

err = tx.Exec("INSERT ...")
if err != nil {
    // 💥 Oops, forgot tx.Rollback(), tx remains open!
    return err
}

// execute more statements...

err = tx.Commit()
if err != nil {
    return err
}

We can create a withTx() function:

func withTx(db DB, callback func(Tx) error) error {
    tx, err := db.Begin()
    if err != nil {
        return err
    }
    defer tx.Rollback()

    err = callback(tx)
    if err != nil {
        return err
    }

    return tx.Commit()  
}

And use it like this:

err := withTx(db, func(tx Tx) error {
    err = tx.Exec("INSERT ...")
    if err != nil {
        // 🙂 No need to remember to call tx.Rollback()
        return err
    }
    // execute more statements...
})
if err != nil {
    return err
}

At a glance, ~5% of calls to Begin() appear to be missing a corresponding Rollback() on error (likely example). Using a withTx() function would help eliminate those bugs.

And with Go generics we can return a value from withTx:

-func withTx(db DB, callback func(Tx) error) error {
+func withTx[T any](db DB, callback func(Tx) (*T, error)) (*T, error) {
    tx, err := db.Begin()
    if err != nil {
-       return err
+       return nil, err
    }
    defer tx.Rollback()

-   err = callback(tx)
+   t, err := callback(tx)
    if err != nil {
        return err
    }

-   return tx.Commit()
+   return t, tx.Commit()
 }

Benefits of the "with" function pattern:

• Automatic initialization and cleanup
• Less code
• In some cases, there's no need to check if the operation succeeded and do something different if it didn't

One caveat to bear in mind: you have to make sure the callback does not return the resource, otherwise that resource could be used after the clean up runs and cause an error! I don't think it's possible to enforce correct use at compile time in Go. It is possible (but not ergonomic) in Haskell and it's convenient in Rust by piggy-backing on the built-in value lifetime machinery.

Comments

[Nate Finch]

I wouldn't catch and effectively swallow the panic like that. I'd instead just defer tx.Rollback() which will noop if you've already committed on the happy path, or rollback on an error or panic.

Also, you're swallowing the error from the callback when you call rollback. That's bad. Rollback will normally return nil in which case this would return nil.

[Chris Wendt]

👍 to defer tx.Rollback()! I updated the post 👌