Merge pull request #35 from funbringer/unique_system_ids_for_cached_initdb

Unique system ids for cached initdb

Author: funbringer

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,23 @@ 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 ExecUtilException as e:
+            msg = "Failed to reset WAL for system id"
+            raise_from(InitNodeException(msg), e)
+
         except Exception as e:
             raise_from(InitNodeException("Failed to spawn a node"), e)

testgres/config.py

@@ -17,6 +17,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?
 
@@ -32,6 +33,7 @@ class GlobalConfig(object):
 
     cache_initdb = True
     _cached_initdb_dir = None
+    cached_initdb_unique = False
 
     cache_pg_config = True
 

testgres/consts.py

@@ -10,6 +10,9 @@
 TMP_CACHE = 'tgsc_'
 TMP_BACKUP = 'tgsb_'
 
+# 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"

testgres/utils.py

@@ -68,6 +68,30 @@ 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
+
+    date1 = datetime.datetime.utcfromtimestamp(0)
+    date2 = datetime.datetime.utcnow()
+
+    secs = int((date2 - date1).total_seconds())
+    usecs = date2.microsecond
+
+    # see pg_resetwal.c : GuessControlValues()
+    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).

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 util_exists(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,34 @@ def test_init_after_cleanup(self):
             node.cleanup()
             node.init().start().execute('select 1')
 
+    @unittest.skipUnless(util_exists('pg_resetwal'), 'might be missing')
+    @unittest.skipUnless(pg_version_ge('9.6'), 'query works on 9.6+')
+    def test_init_unique_system_id(self):
+        # this function exists in PostgreSQL 9.6+
+        query = 'select system_identifier from pg_control_system()'
+
+        with scoped_config(cache_initdb=False):
+            with get_new_node().init().start() as node0:
+                id0 = node0.execute(query)[0]
+
+        # yapf: disable
+        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:
+
+                id1 = node1.execute(query)[0]
+                id2 = node2.execute(query)[0]
+
+                # ids must increase
+                self.assertGreater(id1, id0)
+                self.assertGreater(id2, id1)
+
     def test_node_exit(self):
         base_dir = None
 
@@ -486,8 +514,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_exists('pgbench'), 'might be missing')
     def test_pgbench(self):
         with get_new_node().init().start() as node: