Using Go Defer and Rust Drop to Defer the Call to Rollback

My extremely slow journey to learn Rust continues, delayed by other projects. My attention in 2021 has been primarily on Go and PostgreSQL. I'm coming to appreciate and respect the database overall (and PostgreSQL specifically) a lot more and wrote up some of my thoughts in Thick Databases.

One thing that has me very interested in Rust is the tools it gives me to write code that works in exactly the way I expect it to. It helps me to enforce that behaviour on other developers and to avoid situations where I (or others on my team) forget to do something important, such as initialising a value, closing a file, or an HTTP request, or a database transaction.

Forgetting to close things off in Go is one of those places that can potentially come back to bite you in hard-to-find ways. For a database connection, for example, it's very important that you remember to always rollback or commit a transaction. If you forget to do this, you can run into situations where you've starved yourself of connections, and any further requests fail -- your service grinds to a halt.

There's a handful of ways to do this. The most basic and straightforward method is to call rollback or commit every time you return:

func someWork() error {
	tx, err := db.Begin()

	err := foo(tx)

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

	err = bar(tx)

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

	return tx.Commit()
}

This runs the risk that we forget to add tx.Rollback() upon returning some error -- something that can happen very easily, especially as we refactor code and move code here from elsewhere that had an error check that didn't previously need to be accompanied with a rollback.

A safer option is to defer the call to rollback. This ensures that it's always called since it doesn't matter if a deferred rollback call follows a successful commit:

func someWork() error {
	tx, err := db.Begin()
	defer tx.Rollback()

	err := foo(tx)

	if err != nil {
		return err
	}

	err = bar(tx)

	if err != nil {
		return err
	}

	return tx.Commit()
}

This is better because we now only need to remember to call rollback once and make sure we commit when everything is in order. However, it's still not perfect, and this is where another footgun can appear. Suppose that we have a loop and we are creating new transactions on each iteration of the loop.

Here's a loop, where we start a new transaction each time, but for whatever reason, we do not want to commit the work we've done (or alternatively, we do want to commit the work from some iterations of the loop but not all -- e.g., with an error, we continue to the next iteration without committing):

func deferInLoop() {
	for i := 0; i < loops; i++ {
		var result bool
		tx, err := db.Begin()
		defer tx.Rollback()

		if err != nil {
			fmt.Println(err.Error())
			continue
		}
		err = tx.QueryRow("SELECT true").Scan(&result)
		if err != nil {
			fmt.Println(err.Error())
			continue
		}
		log.Printf("loop count %d.  Result: %t", i, result)
	}
}

If we try and execute this, we'll find our connections run out and the service panics:

2021/11/10 00:36:08 loop count 92.  Result: true
2021/11/10 00:36:08 loop count 93.  Result: true
2021/11/10 00:36:08 loop count 94.  Result: true
2021/11/10 00:36:08 loop count 95.  Result: true
2021/11/10 00:36:08 loop count 96.  Result: true
2021/11/10 00:36:08 loop count 97.  Result: true
2021/11/10 00:36:08 loop count 98.  Result: true
2021/11/10 00:36:08 loop count 99.  Result: true
pq: sorry, too many clients already
pq: sorry, too many clients already
pq: sorry, too many clients already
pq: sorry, too many clients already
pq: sorry, too many clients already
...
pq: sorry, too many clients already
pq: sorry, too many clients already
panic: runtime error: invalid memory address or nil pointer dereference
panic: runtime error: invalid memory address or nil pointer dereference
panic: runtime error: invalid memory address or nil pointer dereference
panic: runtime error: invalid memory address or nil pointer dereference
panic: runtime error: invalid memory address or nil pointer dereference
...
[signal SIGSEGV: segmentation violation code=0x1 addr=0x40 pc=0x10c3e62]

goroutine 1 [running]:
database/sql.(*Tx).rollback(0xc000028068, 0x0)
        /usr/local/go/src/database/sql/sql.go:2263 +0x22
database/sql.(*Tx).Rollback(0x0)
        /usr/local/go/src/database/sql/sql.go:2295 +0x1b
panic({0x120c840, 0x13e6ec0})

We can find ourselves in a similar situation where we have long-lived functions that begin a transaction, but don't return until much later -- long enough that concurrent calls of that function add up to eventually starving ourselves again because the deferred rollback hasn't yet been called (Rust doesn't help here either if the transaction variable won't leave scope until the end of the function).

Recapping our earlier lesson, we needed to remember to rollback every time we return. To avoid having to call rollback every time there's an error, we deferred the rollback so that it runs every time the function returns. However, there are some circumstances in which the defer won't be called early enough to let us avoid starving ourselves of connections. What can we do? We can either revert back in these situations to calling rollback each time we return/continue:

func rollbackInLoop() {
	for i := 0; i < loops; i++ {
		var result bool
		tx, err := db.Begin()

		if err != nil {
			fmt.Println(err.Error())
			tx.Rollback()
			continue
		}
		err = tx.QueryRow("SELECT true").Scan(&result)
		if err != nil {
			fmt.Println(err.Error())
			tx.Rollback()
			continue
		}
		log.Printf("loop count %d.  Result: %t", i, result)
		tx.Rollback()
	}
}

Or, we can perform the work for that transaction in a separate function that will return for every iteration of the loop. The defer, therefore, gets called before the next iteration of the loop begins a transaction:

func deferInFunc() {
	for i := 0; i < loops; i++ {
		err := deferInFuncFetch(db, i)
		if err != nil {
			fmt.Println(err.Error())
			continue
		}
	}
}

func deferInFuncFetch(db *sql.DB, i int) error {
	var result bool
	tx, err := db.Begin()
	defer tx.Rollback()

	if err != nil {
		return err
	}
	err = tx.QueryRow("SELECT true").Scan(&result)
	if err != nil {
		return err
	}
	log.Printf("loop count %d", i)
	err = tx.Commit()

	if err != nil {
		return err
	}
	return nil
}

These solutions work, but they're things we have to remember to be careful about -- it's easy for us to forget these things, as there's nothing that warns us that we've done something that causes problems until the point where we find our production service failing in strange and unpredicable ways. It would be fantastic if we could write our Go code in such a way that you can't forget to do these things. For example, by having rollback be called automatically in good time by default, or by the compiler throwing an error if we missed a case.

This led me to wonder how this would be handled in rust.

Rust Drop

Rust provides the Drop trait that you can implement on structs. The drop function will get called much like a destructor in C++, so that you can clean things up. This happens when the owner goes away, so it gives us the chance to perform actions earlier than just the point where the function returns. E.g., Drop might be called at the end of every iteration of the loop if the variable falls out of scope then.

This gives a great place for us to ensure that rollback is called if we forget. As a result, we can guarantee, in a way that prevents us from forgetting, that a transaction will eventually rollback.

Looking at the Postgres library's transaction method, we learn:

The transaction will rollback by default - use the commit method to commit it.

Investigating a little further, we can see that the Transaction struct implements the Drop trait. Specifically:

impl<'a> Drop for Transaction<'a> {
    fn drop(&mut self) {
        if let Some(transaction) = self.transaction.take() {
            let _ = self.connection.block_on(transaction.rollback());
        }
    }
}

Therefore, we don't even need to implement anything ourselves to ensure that a rollback is called -- no need to call rollback, and no need to schedule a deferred rollback. When using this library, if we neglect to commit the transaction, then it will rollback and in good time Suppose then we implement a similar loop function to the one we had in Go:

fn loop_() -> Result<(), Box<dyn std::error::Error>> {
    let mut client = Client::connect("host=localhost user=postgres", NoTls)?;

    for i in 0..200 {
        let mut transaction = client.transaction()?;

        let row = transaction.query_one("SELECT true", &[])?;
        let result: bool = row.get(0);

        println!("loop count {} result: {}", i, result);
    }

    Ok(())
}

This works exactly as we'd hope -- 200 iterations and no issues. This is because each time the loop iteration ends, the transaction struct gets dropped and rollback is called, all before the next iteration begins.

Looking at another Rust library, sqlx, we find the same approach with drop is used:

A transaction should end with a call to commit or rollback. If neither are called before the transaction goes out-of-scope, rollback is called. In other words, rollback is called on drop if the transaction is still in-progress.

Thus helping to ensure that connections don't remain around forever when we neglect to rollback. The useful thing about drop is that we don't have to rely on the function returning, or on running the transaction in a separate function. It's enough that a block ends, and the struct is due to be freed, for the default rollback to be called.

When I first started learning Go, I had some disappointment with the tools it gave me to ensure that code could be used exactly and only as intended. A standout example of this is the lack of a way to force users of a library to only get new instances of a struct via a function whose job is to ensure the struct is initialised in the ways it needs to be (e.g., initialising an internal map).

Eventually, I stopped wishing for those things, and just got on with the work of building things. However, now that I'm exploring Rust, I'm thinking again about the things that I had once hoped to be able to do in Go.

Missing a rollback is a footgun that I've fired off before, and I'll bet plenty of others have as well. I'm pleased to see that Rust gives the tools to prevent these kinds of issues.

If you'd like to explore more, I've uploaded the code I used to play around with this to github.com/saward/footgun-defer.