Generate unique system ids in cached_initdb()

funbringer · funbringer · commit e10520a93900 · 2018-02-21T16:23:33.000+03:00
diff --git a/testgres/cache.py b/testgres/cache.py
@@ -1,19 +1,23 @@
 # coding: utf-8
 
+import io
 import os
 import shutil
 
 from six import raise_from
 
 from .config import testgres_config
 
+from .consts import XLOG_CONTROL_FILE
+
 from .exceptions import \
     InitNodeException, \
     ExecUtilException
 
 from .utils import \
     get_bin_path, \
-    execute_utility
+    execute_utility, \
+    generate_system_id
 
 
 def cached_initdb(data_dir, logfile=None, params=None):
@@ -42,5 +46,18 @@ def call_initdb(initdb_dir, log=None):
         try:
             # Copy cached initdb to current data dir
             shutil.copytree(cached_data_dir, data_dir)
+
+            # Assign this node a unique system id if asked to
+            if testgres_config.cached_initdb_unique:
+                # XXX: write new unique system id to control file
+                # Some users might rely upon unique system ids, but
+                # our initdb caching mechanism breaks this contract.
+                pg_control = os.path.join(data_dir, XLOG_CONTROL_FILE)
+                with io.open(pg_control, "r+b") as f:
+                    f.write(generate_system_id())    # overwrite id
+
+                # XXX: build new WAL segment with our system id
+                _params = [get_bin_path("pg_resetwal"), "-D", data_dir, "-f"]
+                execute_utility(_params, logfile)
         except Exception as e:
             raise_from(InitNodeException("Failed to spawn a node"), e)
diff --git a/testgres/config.py b/testgres/config.py
@@ -15,6 +15,7 @@ class GlobalConfig(object):
     Attributes:
         cache_initdb:           shall we use cached initdb instance?
         cached_initdb_dir:      shall we create a temp dir for cached initdb?
+        cached_initdb_unique:   shall we assign new node a unique system id?
 
         cache_pg_config:        shall we cache pg_config results?
 
@@ -30,6 +31,7 @@ class GlobalConfig(object):
 
     cache_initdb = True
     _cached_initdb_dir = None
+    cached_initdb_unique = False
 
     cache_pg_config = True
 
diff --git a/testgres/consts.py b/testgres/consts.py
@@ -4,6 +4,9 @@
 DATA_DIR = "data"
 LOGS_DIR = "logs"
 
+# path to control file
+XLOG_CONTROL_FILE = "global/pg_control"
+
 # names for config files
 RECOVERY_CONF_FILE = "recovery.conf"
 PG_AUTO_CONF_FILE = "postgresql.auto.conf"
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -68,6 +68,27 @@ def generate_app_name():
     return 'testgres-{}'.format(str(uuid.uuid4()))
 
 
+def generate_system_id():
+    """
+    Generate a new 64-bit unique system identifier for node.
+    """
+
+    import datetime
+    import struct
+
+    date = datetime.datetime.now()
+    secs = int(date.timestamp())
+    usecs = date.microsecond
+
+    system_id = 0
+    system_id |= (secs << 32)
+    system_id |= (usecs << 12)
+    system_id |= (os.getpid() & 0xFFF)
+
+    # pack ULL in native byte order
+    return struct.pack('=Q', system_id)
+
+
 def execute_utility(args, logfile=None):
     """
     Execute utility (pg_ctl, pg_dump etc).
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -38,15 +38,15 @@
     get_pg_config
 
 from testgres import bound_ports
+from testgres.utils import pg_version_ge
 
 
 def util_is_executable(util):
     def good_properties(f):
-        return (
-            os.path.exists(f) and
-            os.path.isfile(f) and
-            os.access(f, os.X_OK)
-        )
+        # yapf: disable
+        return (os.path.exists(f) and
+                os.path.isfile(f) and
+                os.access(f, os.X_OK))
 
     # try to resolve it
     if good_properties(get_bin_path(util)):
@@ -91,6 +91,27 @@ def test_init_after_cleanup(self):
             node.cleanup()
             node.init().start().execute('select 1')
 
+    @unittest.skipUnless(pg_version_ge('9.6'), 'query works on 9.6+')
+    def test_init_unique_system_id(self):
+        with scoped_config(
+                cache_initdb=True, cached_initdb_unique=True) as config:
+
+            self.assertTrue(config.cache_initdb)
+            self.assertTrue(config.cached_initdb_unique)
+
+            # spawn two nodes; ids must be different
+            with get_new_node().init().start() as node1, \
+                    get_new_node().init().start() as node2:
+
+                # this function exists in PostgreSQL 9.6+
+                query = 'select system_identifier from pg_control_system()'
+
+                id1 = node1.execute(query)[0]
+                id2 = node2.execute(query)[0]
+
+                # ids must increase
+                self.assertGreater(id2, id1)
+
     def test_node_exit(self):
         base_dir = None
 
@@ -486,8 +507,7 @@ def test_logging(self):
                 master.restart()
                 self.assertTrue(master._logger.is_alive())
 
-    @unittest.skipUnless(
-        util_is_executable("pgbench"), "pgbench may be missing")
+    @unittest.skipUnless(util_is_executable('pgbench'), 'might be missing')
     def test_pgbench(self):
         with get_new_node().init().start() as node:
 
