[mutianf]

Summary

Synchronizes mutation staging, flush scheduling, and shutdown in MutationsBatcher using a re-entrant lock (self._stage_lock = threading.RLock()):

1. Prevents silent batch loss on close() vs. timer flush race:

   • Protects _staged_entries swap and task submission to _sync_flush_executor under _stage_lock, and acquires _stage_lock in close() when setting _closed, canceling _flush_timer, and scheduling the final flush before shutting down _sync_flush_executor.
   • Restores _staged_entries, _staged_count, and _staged_bytes if create_task raises RuntimeError (e.g., if the executor was already shut down during atexit) so close() can flush them synchronously.
   • Initializes _stage_lock, _flush_jobs, and exception counters before starting the background _flush_timer daemon thread.

2. Prevents duplicate batch dispatch on concurrent _schedule_flush() calls:

   • Serializes append(), _timer_routine(), _schedule_flush(), close(), and _on_exit() via _stage_lock so user and background timer threads cannot read the same _staged_entries reference and dispatch duplicate MutateRows RPCs.
   • Protects _entries_processed_since_last_raise and _add_exceptions updates across concurrent flush worker threads and uses _flush_jobs.discard for idempotent task cleanup.

[gemini-code-assist]

Code Review

This pull request introduces thread safety to the MutationsBatcher class (both async and sync autogenerated versions) by using a reentrant lock (threading.RLock) to protect shared state, such as staged entries and flush metrics. It also ensures that staged entries are properly restored if a RuntimeError occurs during a flush scheduling attempt, and adds comprehensive unit tests to verify these thread-safety and error-recovery behaviors. There are no review comments to address, and we have no additional feedback to provide.

[daniel-sanche]

We also have one of these at the start of the function. Are they both needed?

[daniel-sanche]

This also modifies some of the state. Does it need to use the lock?