Make periodic pg_query_state calls to backend running TPC-DS bench

Also refactor common functions for python tests

Author: maksm90

tests/common.py

@@ -6,13 +6,7 @@
 import psycopg2
 import psycopg2.extensions
 import select
-
-# Some queries from TPC-DS may freeze or be even broken,
-# so we allow some sort of failure, since we do not test
-# Postgres, but rather that pg_query_state do not crash
-# anything under stress load.
-MAX_PG_QS_RETRIES = 50
-
+import time
 
 def wait(conn):
 	"""wait for some event on connection to postgres"""
@@ -47,8 +41,7 @@ def n_close(conns):
 		conn.close()
 
 def pg_query_state(config, pid, verbose=False, costs=False, timing=False, \
-								buffers=False, triggers=False, format='text', \
-								stress_in_progress=False):
+								buffers=False, triggers=False, format='text'):
 	"""
 	Get query state from backend with specified pid and optional parameters.
 	Save any warning, info, notice and log data in global variable 'notices'
@@ -57,53 +50,48 @@ def pg_query_state(config, pid, verbose=False, costs=False, timing=False, \
 	conn = psycopg2.connect(**config)
 	curs = conn.cursor()
 
-	if stress_in_progress:
-		set_guc(conn, 'statement_timeout', TPC_DS_STATEMENT_TIMEOUT)
-		n_retries = 0
-
-	result = []
-	while not result:
-		curs.callproc('pg_query_state', (pid, verbose, costs, timing, buffers, triggers, format))
-		result = curs.fetchall()
-
-		if stress_in_progress:
-			n_retries += 1
-			if n_retries >= MAX_PG_QS_RETRIES:
-				print('\npg_query_state tried %s times with no effect, giving up' % MAX_PG_QS_RETRIES)
-				break
-
+	curs.callproc('pg_query_state', (pid, verbose, costs, timing, buffers, triggers, format))
+	result = curs.fetchall()
 	notices = conn.notices[:]
 	conn.close()
+
 	return result, notices
 
-def query_state(config, async_conn, query, args={}, num_workers=0, stress_in_progress=False):
+def onetime_query_state(config, async_conn, query, args={}, num_workers=0):
 	"""
 	Get intermediate state of 'query' on connection 'async_conn' after number of 'steps'
 	of node executions from start of query
 	"""
 
 	acurs = async_conn.cursor()
-	conn = psycopg2.connect(**config)
-	curs = conn.cursor()
 
 	set_guc(async_conn, 'enable_mergejoin', 'off')
 	set_guc(async_conn, 'max_parallel_workers_per_gather', num_workers)
 	acurs.execute(query)
 
 	# extract current state of query progress
+	MAX_PG_QS_RETRIES = 10
+	DELAY_BETWEEN_RETRIES = 0.1
 	pg_qs_args = {
 			'config': config,
 			'pid': async_conn.get_backend_pid()
 			}
 	for k, v in args.items():
 		pg_qs_args[k] = v
-	result, notices = pg_query_state(**pg_qs_args)
+	n_retries = 0
+	while True:
+		result, notices = pg_query_state(**pg_qs_args)
+		n_retries += 1
+		if len(result) > 0:
+			break
+		if n_retries >= MAX_PG_QS_RETRIES:
+			# pg_query_state callings don't return any result, more likely run
+			# query has completed
+			break
+		time.sleep(DELAY_BETWEEN_RETRIES)
 	wait(async_conn)
 
-	set_guc(async_conn, 'pg_query_state.executor_trace', 'off')
 	set_guc(async_conn, 'enable_mergejoin', 'on')
-
-	conn.close()
 	return result, notices
 
 def set_guc(async_conn, param, value):

tests/test_cases.py

@@ -50,7 +50,7 @@ def test_simple_query(config):
               Buckets: \d+  Batches: \d+  Memory Usage: \d+kB
               ->  Seq Scan on bar \(Current loop: actual rows=\d+, loop number=1\)"""
 
-	qs, _ = common.query_state(config, acon, query)
+	qs, _ = common.onetime_query_state(config, acon, query)
 	assert qs[0][0] == acon.get_backend_pid() and qs[0][1] == 0 \
 		and qs[0][2] == query and re.match(expected, qs[0][3]) and qs[0][4] == None
 
@@ -111,7 +111,7 @@ def test_nested_call(config):
 	util_curs.execute(create_function)
 	util_conn.commit()
 
-	qs, notices = common.query_state(config, acon, call_function)
+	qs, notices = common.onetime_query_state(config, acon, call_function)
 	assert len(qs) == 2 \
 		and qs[0][0] == qs[1][0] == acon.get_backend_pid() \
 		and qs[0][1] == 0 and qs[1][1] == 1 \
@@ -143,7 +143,7 @@ def test_insert_on_conflict(config):
 	util_curs.execute(add_field_uniqueness)
 	util_conn.commit()
 
-	qs, notices = common.query_state(config, acon, query)
+	qs, notices = common.onetime_query_state(config, acon, query)
 
 	assert qs[0][0] == acon.get_backend_pid() and qs[0][1] == 0 \
 		and qs[0][2] == query and re.match(expected, qs[0][3]) \
@@ -185,13 +185,13 @@ def test_trigger(config):
 	util_curs.execute(create_trigger)
 	util_conn.commit()
 
-	qs, notices = common.query_state(config, acon, query, {'triggers': True})
+	qs, notices = common.onetime_query_state(config, acon, query, {'triggers': True})
 	assert qs[0][0] == acon.get_backend_pid() and qs[0][1] == 0 \
 		and qs[0][2] == query and re.match(expected_upper, qs[0][3]) \
 		and qs[0][4] == None
 	assert len(notices) == 0
 
-	qs, notices = common.query_state(config, acon, query, {'triggers': False})
+	qs, notices = common.onetime_query_state(config, acon, query, {'triggers': False})
 	assert qs[0][0] == acon.get_backend_pid() and qs[0][1] == 0 \
 		and qs[0][2] == query and re.match(expected_upper, qs[0][3]) \
 		and qs[0][4] == None
@@ -215,7 +215,7 @@ def test_costs(config):
               Buckets: \d+  Batches: \d+  Memory Usage: \d+kB
               ->  Seq Scan on bar  \(cost=0.00..\d+.\d+ rows=\d+ width=4\) \(Current loop: actual rows=\d+, loop number=1\)"""
 
-	qs, notices = common.query_state(config, acon, query, {'costs': True})
+	qs, notices = common.onetime_query_state(config, acon, query, {'costs': True})
 	assert len(qs) == 1 and re.match(expected, qs[0][3])
 	assert len(notices) == 0
 
@@ -238,7 +238,7 @@ def test_buffers(config):
 
 	common.set_guc(acon, 'pg_query_state.enable_buffers', 'on')
 
-	qs, notices = common.query_state(config, acon, query, {'buffers': True})
+	qs, notices = common.onetime_query_state(config, acon, query, {'buffers': True})
 	assert len(qs) == 1 and re.match(expected, qs[0][3])
 	assert len(notices) == 0
 
@@ -259,7 +259,7 @@ def test_timing(config):
 
 	common.set_guc(acon, 'pg_query_state.enable_timing', 'on')
 
-	qs, notices = common.query_state(config, acon, query, {'timing': True})
+	qs, notices = common.onetime_query_state(config, acon, query, {'timing': True})
 	assert len(qs) == 1 and re.match(expected, qs[0][3])
 	assert len(notices) == 0
 
@@ -298,11 +298,11 @@ def test_formats(config):
               Buckets: \d+  Batches: \d+  Memory Usage: \d+kB
               ->  Seq Scan on bar \(Current loop: actual rows=\d+, loop number=1\)"""
 
-	qs, notices = common.query_state(config, acon, query, {'format': 'text'})
+	qs, notices = common.onetime_query_state(config, acon, query, {'format': 'text'})
 	assert len(qs) == 1 and re.match(expected, qs[0][3])
 	assert len(notices) == 0
 
-	qs, notices = common.query_state(config, acon, query, {'format': 'json'})
+	qs, notices = common.onetime_query_state(config, acon, query, {'format': 'json'})
 	try:
 		js_obj = json.loads(qs[0][3])
 	except ValueError:
@@ -311,7 +311,7 @@ def test_formats(config):
 	assert len(notices) == 0
 	check_plan(js_obj['Plan'])
 
-	qs, notices = common.query_state(config, acon, query, {'format': 'xml'})
+	qs, notices = common.onetime_query_state(config, acon, query, {'format': 'xml'})
 	assert len(qs) == 1
 	assert len(notices) == 0
 	try:
@@ -320,7 +320,7 @@ def test_formats(config):
 		assert False, 'Invalid xml format'
 	check_xml(xml_root)
 
-	qs, _ = common.query_state(config, acon, query, {'format': 'yaml'})
+	qs, _ = common.onetime_query_state(config, acon, query, {'format': 'yaml'})
 	try:
 		yaml_doc = yaml.load(qs[0][3], Loader=yaml.FullLoader)
 	except:
@@ -339,15 +339,15 @@ def test_timing_buffers_conflicts(config):
 	timing_pattern = '(?:running time=\d+.\d+)|(?:actual time=\d+.\d+..\d+.\d+)'
 	buffers_pattern = 'Buffers:'
 
-	qs, notices = common.query_state(config, acon, query, {'timing': True, 'buffers': False})
+	qs, notices = common.onetime_query_state(config, acon, query, {'timing': True, 'buffers': False})
 	assert len(qs) == 1 and not re.search(timing_pattern, qs[0][3])
 	assert notices == ['WARNING:  timing statistics disabled\n']
 
-	qs, notices = common.query_state(config, acon, query, {'timing': False, 'buffers': True})
+	qs, notices = common.onetime_query_state(config, acon, query, {'timing': False, 'buffers': True})
 	assert len(qs) == 1 and not re.search(buffers_pattern, qs[0][3])
 	assert notices == ['WARNING:  buffers statistics disabled\n']
 
-	qs, notices = common.query_state(config, acon, query, {'timing': True, 'buffers': True})
+	qs, notices = common.onetime_query_state(config, acon, query, {'timing': True, 'buffers': True})
 	assert len(qs) == 1 and not re.search(timing_pattern, qs[0][3]) \
 						 and not re.search(buffers_pattern, qs[0][3])
 	assert len(notices) == 2 and 'WARNING:  timing statistics disabled\n' in notices \

tests/tpcds.py

@@ -8,6 +8,7 @@
 import progressbar
 import psycopg2.extensions
 import subprocess
+import time
 
 class DataLoadException(Exception): pass
 class StressTestException(Exception): pass
@@ -58,6 +59,7 @@ def run_tpcds(config):
 			queries.append(f.read())
 
 	acon, = common.n_async_connect(config)
+	pid = acon.get_backend_pid()
 
 	print('Starting TPC-DS queries...')
 	timeout_list = []
@@ -69,7 +71,36 @@ def run_tpcds(config):
 		try:
 			# Set query timeout to TPC_DS_STATEMENT_TIMEOUT / 1000 seconds
 			common.set_guc(acon, 'statement_timeout', TPC_DS_STATEMENT_TIMEOUT)
-			qs = common.query_state(config, acon, query, stress_in_progress=True)
+
+			# run query
+			acurs = acon.cursor()
+			acurs.execute(query)
+
+			# periodically run pg_query_state on running backend trying to get
+			# crash of PostgreSQL
+			MAX_PG_QS_RETRIES = 10
+			PG_QS_DELAY, BEFORE_GOT_QS_DELAY = 0.1, 0.1
+			BEFORE_GOT_QS, GOT_QS = range(2)
+			state, n_retries = BEFORE_GOT_QS, 0
+			while True:
+				result, _ = common.pg_query_state(config, pid)
+				if state == BEFORE_GOT_QS:
+					if len(result) > 0:
+						state = GOT_QS
+						continue
+					n_retries += 1
+					if n_retries >= MAX_PG_QS_RETRIES:
+						# pg_query_state callings don't return any result, more likely run
+						# query has completed
+						break
+					time.sleep(BEFORE_GOT_QS_DELAY)
+				if state == GOT_QS:
+					if len(result) == 0:
+						break
+				time.sleep(PG_QS_DELAY)
+
+			# wait for real query completion
+			common.wait(acon)
 
 		except psycopg2.extensions.QueryCanceledError:
 			timeout_list.append(i + 1)