fix(tee): fix race condition in batch locking

[pbeza]

What ❔

After scaling zksync-tee-prover to two instances/replicas on Azure for azure-stage2, azure-testnet2, and azure-mainnet2, we started experiencing duplicated proving for some batches.

While this is not an erroneous situation, it is wasteful from a resource perspective. This was due to a race condition in batch locking. This PR fixes the issue by adding atomic batch locking.

Why ❔

To fix the bug that only activates after running zksync-tee-prover on multiple instances.

Checklist

• PR title corresponds to the body of PR (we generate changelog entries from PRs).
• Tests for the changes have been added / updated.
• Documentation comments have been added / updated.
• Code has been formatted via zkstack dev fmt and zkstack dev lint.

[slowli]

Dumb question: How is the locking made atomic in this PR? AFAIU, the first SELECT statement, if queried concurrently, can still return the same L1 batch number unless some kind of row-level locking is implemented (cf. SELECT FOR UPDATE SKIP LOCKED in this contract verifier query). I'm not even sure the UPDATE query will fail for the transaction committed last in case of a race (maybe it would with serialization isolation level, but I'd argue that erroring is not the best cause of action here; row-level locks seem to work better).

[slowli]

There are seemingly no changes in the transaction at this point. Is the commit here just to not clean up DB resources faster?

[pbeza]

To be honest, I just presumed every .start_transaction() needed to either end with .commit() or .rollback(). I'm not sure how you determine that it's OK to skip it altogether, but I trust your judgment here. I removed this line here. ✅

[slowli]

A transaction that is dropped w/o being committed is rolled back implicitly. You can add an explicit call to be more precise (but IMO it should be rollback so that it's clear that no modifying actions have been performed).

[pbeza]

Dumb question: How is the locking made atomic in this PR? (...)

Not a dumb question at all! The dumb one here was me! ;P I totally misunderstood what SQL transactions can actually handle in this context. Had to brush up on the finer details of SQL locking. Thanks for steering me in the right direction! These two links were super helpful:

• https://stackoverflow.com/q/6998372/1321680
• https://www.postgresql.org/docs/current/sql-select.html#SQL-FOR-UPDATE-SHARE

[pbeza]

I had to explicitly use OF p here. If you leave it out or switch it to OF tee, you will run into this error:

error: error returned from database: FOR UPDATE cannot be applied to the nullable side of an outer join

[slowli]

The error is understandable. Just to clarify, since it's the tee_proof_generation_details table that is updated below: Do I understand correctly that this new query can actually have tee_proof_generation_details side of the join absent (I guess if the query is invoked for a certain batch for the first time?). If so, I'd suggest to add a comment explaining why rows are locked on a table that's not modified in the transaction.

[pbeza]

@slowli, I’ve addressed your code review comments. Take a look when you get a chance.

It’s kinda hard to test properly without deploying it to stage and letting it run for a while. Specifically, let me know if locking rows in the proof_generation_details table is okay (instead of just locking tee_proof_generation_details rows).