Add fork_at_current_lsn function which creates branch at current LSN (#2344)

* Add fork_at_current_lsn function which creates branch at current LSN

* Undo use of fork_at_current_lsn in test_branching because of short GC period

* Add missed return in fork_at_current_lsn

* Add missed return in fork_at_current_lsn

* Update test_runner/fixtures/neon_fixtures.py

Co-authored-by: Heikki Linnakangas <heikki@zenith.tech>

* Update test_runner/fixtures/neon_fixtures.py

Co-authored-by: Heikki Linnakangas <heikki@zenith.tech>

* Update test_runner/fixtures/neon_fixtures.py

Co-authored-by: Heikki Linnakangas <heikki@zenith.tech>

Co-authored-by: Heikki Linnakangas <heikki@zenith.tech>

test_runner/batch_others/test_twophase.py

@@ -1,7 +1,7 @@
 import os
 
 from fixtures.log_helper import log
-from fixtures.neon_fixtures import NeonEnv
+from fixtures.neon_fixtures import NeonEnv, fork_at_current_lsn
 
 
 #
@@ -55,7 +55,7 @@ def test_twophase(neon_simple_env: NeonEnv):
     assert len(twophase_files) == 2
 
     # Create a branch with the transaction in prepared state
-    env.neon_cli.create_branch("test_twophase_prepared", "test_twophase")
+    fork_at_current_lsn(env, pg, "test_twophase_prepared", "test_twophase")
 
     # Start compute on the new branch
     pg2 = env.postgres.create_start(

test_runner/batch_others/test_vm_bits.py

@@ -1,5 +1,5 @@
 from fixtures.log_helper import log
-from fixtures.neon_fixtures import NeonEnv
+from fixtures.neon_fixtures import NeonEnv, fork_at_current_lsn
 
 
 #
@@ -33,7 +33,7 @@ def test_vm_bit_clear(neon_simple_env: NeonEnv):
     cur.execute("UPDATE vmtest_update SET id = 5000 WHERE id = 1")
 
     # Branch at this point, to test that later
-    env.neon_cli.create_branch("test_vm_bit_clear_new", "test_vm_bit_clear")
+    fork_at_current_lsn(env, pg, "test_vm_bit_clear_new", "test_vm_bit_clear")
 
     # Clear the buffer cache, to force the VM page to be re-fetched from
     # the page server

test_runner/fixtures/neon_fixtures.py

@@ -2640,3 +2640,19 @@ def wait_for_last_flush_lsn(env: NeonEnv, pg: Postgres, tenant: uuid.UUID, timel
     """Wait for pageserver to catch up the latest flush LSN"""
     last_flush_lsn = lsn_from_hex(pg.safe_psql("SELECT pg_current_wal_flush_lsn()")[0][0])
     wait_for_last_record_lsn(env.pageserver.http_client(), tenant, timeline, last_flush_lsn)
+
+
+def fork_at_current_lsn(
+    env: NeonEnv,
+    pg: Postgres,
+    new_branch_name: str,
+    ancestor_branch_name: str,
+    tenant_id: Optional[uuid.UUID] = None,
+) -> uuid.UUID:
+    """
+    Create new branch at the last LSN of an existing branch.
+    The "last LSN" is taken from the given Postgres instance. The pageserver will wait for all the
+    the WAL up to that LSN to arrive in the pageserver before creating the branch.
+    """
+    current_lsn = pg.safe_psql("SELECT pg_current_wal_lsn()")[0][0]
+    return env.neon_cli.create_branch(new_branch_name, ancestor_branch_name, tenant_id, current_lsn)