Add bench script, fix mq errors

Author: ildus

bench/.gitignore

@@ -0,0 +1 @@
+/football.json

bench/bench.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+
+import contextlib
+import glob
+import json
+import os
+
+from testgres import get_new_node
+
+sources = [
+    ('football', './football.json'),
+]
+
+conf = '''
+shared_preload_libraries='jsonbd'
+'''
+
+
+@contextlib.contextmanager
+def cwd(path):
+    print("cwd: ", path)
+    curdir = os.getcwd()
+    os.chdir(path)
+
+    try:
+        yield
+    finally:
+        print("cwd:", curdir)
+        os.chdir(curdir)
+
+
+with get_new_node('node1') as node:
+    node.init()
+    node.append_conf('postgresql.conf', conf)
+    node.start()
+
+    node.safe_psql('create extension jsonbd')
+
+    for name, root_dir in sources:
+        table_name = name
+        table_name_c = '%s_c' % name
+
+        node.safe_psql('create table %s(a jsonb)' % table_name)
+        node.safe_psql('create table %s(a jsonb compression jsonbd)' % table_name_c)
+        node.safe_psql('alter table %s alter column a set storage external' % table_name)
+
+        with node.connect() as con:
+            with cwd(os.path.abspath(root_dir)):
+                for filename in glob.iglob('**/*.json', recursive=True):
+                    if filename == 'package.json':
+                        continue
+
+                    with open(filename, 'r') as f:
+                        data = json.load(f)
+
+                    if isinstance(data, dict):
+                        if 'rounds' in data:
+                            for obj in data['rounds']:
+                                sql = 'insert into {} values (%s)'
+                                for i in range(10):
+                                    con.execute(sql.format(table_name), (json.dumps(obj), ))
+                                    con.execute(sql.format(table_name_c), (json.dumps(obj), ))
+
+        print(node.safe_psql("select pg_size_pretty(pg_total_relation_size('%s'))" % table_name))
+        print(node.safe_psql("select pg_size_pretty(pg_total_relation_size('%s'))" % table_name_c))

jsonbd.c

@@ -360,15 +360,15 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 		mqout = shm_mq_create(hdr->launcher.mqout, shm_mq_minimum_size);
 
 		/*
-		 * important that sender on mqout should be set earlier than
-		 * receiver on mqin
-		 */
-		shm_mq_set_sender(mqout, hdr->launcher.proc);
+		 * set sender, create handle and wake up launcher and wait until
+		 * it's connected
+		 * */
 		shm_mq_set_receiver(mqout, MyProc);
 		shm_mq_set_sender(mqin, MyProc);
-		shm_mq_set_receiver(mqin, hdr->launcher.proc);
-
 		mqh = shm_mq_attach(mqin, NULL, NULL);
+		SetLatch(&hdr->launcher.proc->procLatch);
+		shm_mq_wait_for_attach(mqh);
+
 		resmq = shm_mq_sendv(mqh,
 				&((shm_mq_iovec) {(char *) &MyDatabaseId, sizeof(MyDatabaseId)}), 1, false);
 		if (resmq != SHM_MQ_SUCCESS)
@@ -405,12 +405,17 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 	mqin = shm_mq_create(wd->mqin, jsonbd_total_queue_size);
 	mqout = shm_mq_create(wd->mqout, jsonbd_total_queue_size);
 
-	shm_mq_set_sender(mqin, MyProc);
-	shm_mq_set_receiver(mqin, wd->proc);
-	shm_mq_set_sender(mqout, wd->proc);
+	/*
+	 * create handle and wake up worker and wait until
+	 * it's connected
+	 * */
 	shm_mq_set_receiver(mqout, MyProc);
-
+	shm_mq_set_sender(mqin, MyProc);
 	mqh = shm_mq_attach(mqin, NULL, NULL);
+	SetLatch(&wd->latch);
+	shm_mq_wait_for_attach(mqh);
+
+	/* send data */
 	resmq = shm_mq_sendv(mqh, iov, iov_len, false);
 	if (resmq != SHM_MQ_SUCCESS)
 		detached = true;

jsonbd.h

@@ -28,6 +28,7 @@ typedef struct jsonbd_shm_worker
 	PGPROC			   *proc;
 	volatile Oid		dboid;	/* database of the worker */
 	LWLock			   *lock;
+	Latch				latch;
 } jsonbd_shm_worker;
 
 /* Shared memory structures */

jsonbd_worker.c

@@ -160,6 +160,7 @@ init_worker(dsm_segment *seg)
 
 	/* Set launcher free */
 	SetLatch(&hdr->launcher_latch);
+	InitLatch(&worker_state->latch);
 
 	/* make this worker visible in backend cycle */
 	hdr->workers_ready++;
@@ -534,7 +535,7 @@ jsonbd_launcher_main(Datum arg)
 	shm_toc			*toc;
 	jsonbd_shm_hdr	*hdr;
 
-	shm_mq_handle  *mqh = NULL;
+	shm_mq_handle  *mqh;
 	int				worker_num	= 1;
 	int				database_num = 0;
 
@@ -572,22 +573,41 @@ jsonbd_launcher_main(Datum arg)
 		if (shutdown_requested)
 			break;
 
-		rc = WaitLatch(&MyProc->procLatch, WL_LATCH_SET | WL_POSTMASTER_DEATH,
-			0, PG_WAIT_EXTENSION);
+		/* Wait to be signalled. */
+		rc = WaitLatch(MyLatch, WL_LATCH_SET | WL_POSTMASTER_DEATH,
+					   0, PG_WAIT_EXTENSION);
 
 		if (rc & WL_POSTMASTER_DEATH)
 			break;
 
-		ResetLatch(&MyProc->procLatch);
+		/* Reset the latch so we don't spin. */
+		ResetLatch(MyLatch);
 
-		if (shm_mq_get_receiver(worker_state->mqin) != MyProc)
+		if (shm_mq_get_sender(worker_state->mqin) == NULL)
+		{
+			CHECK_FOR_INTERRUPTS();
 			continue;
+		}
 
-		Assert(shm_mq_get_sender(worker_state->mqout) == MyProc);
+		/*
+		 * set myself as receiver on mqin and sender on mqout,
+		 * and get data from backend
+		 * */
+		if (!shm_mq_get_sender(worker_state->mqout))
+			shm_mq_set_sender(worker_state->mqout, MyProc);
+
+		if (!shm_mq_get_receiver(worker_state->mqin))
+			shm_mq_set_receiver(worker_state->mqin, MyProc);
 
 		mqh = shm_mq_attach(worker_state->mqin, NULL, NULL);
 		resmq = shm_mq_receive(mqh, &nbytes, &data, false);
 
+		if (resmq == SHM_MQ_DETACHED)
+		{
+			shm_mq_detach(mqh);
+			continue;
+		}
+
 		if (resmq == SHM_MQ_SUCCESS)
 		{
 			int		started = 0;
@@ -614,9 +634,6 @@ jsonbd_launcher_main(Datum arg)
 
 			shm_mq_detach(mqh);
 
-			/* we don't need start this cycle again after we send data */
-			shm_mq_clean_receiver(worker_state->mqin);
-
 			mqh = shm_mq_attach(worker_state->mqout, NULL, NULL);
 			if (started)
 			{
@@ -629,15 +646,11 @@ jsonbd_launcher_main(Datum arg)
 			}
 			else
 				resmq = shm_mq_sendv(mqh, &((shm_mq_iovec) {"n", 2}), 1, false);
-
-			if (resmq != SHM_MQ_SUCCESS)
-				elog(NOTICE, "jsonbd: backend detached early");
-
-			shm_mq_detach(mqh);
-
-			/* mark we need new handle */
-			mqh = NULL;
 		}
+		if (resmq != SHM_MQ_SUCCESS)
+			elog(NOTICE, "jsonbd: backend detached early");
+
+		/* shm_mq_detach(mqh); */
 	}
 
 	elog(LOG, "jsonbd launcher has ended its work");
@@ -679,22 +692,37 @@ jsonbd_worker_main(Datum arg)
 		if (shutdown_requested)
 			break;
 
-		rc = WaitLatch(&MyProc->procLatch, WL_LATCH_SET | WL_POSTMASTER_DEATH,
-			0, PG_WAIT_EXTENSION);
+		/* Wait to be signalled. */
+		rc = WaitLatch(&worker_state->latch, WL_LATCH_SET | WL_POSTMASTER_DEATH,
+					   0, PG_WAIT_EXTENSION);
 
 		if (rc & WL_POSTMASTER_DEATH)
 			break;
 
-		ResetLatch(&MyProc->procLatch);
+		/* Reset the latch so we don't spin. */
+		ResetLatch(&worker_state->latch);
 
-		if (shm_mq_get_receiver(worker_state->mqin) != MyProc)
+		if (shm_mq_get_sender(worker_state->mqin) == NULL)
+		{
+			CHECK_FOR_INTERRUPTS();
 			continue;
+		}
+
+		if (!shm_mq_get_sender(worker_state->mqout))
+			shm_mq_set_sender(worker_state->mqout, MyProc);
 
-		Assert(shm_mq_get_sender(worker_state->mqout) == MyProc);
+		if (!shm_mq_get_receiver(worker_state->mqin))
+			shm_mq_set_receiver(worker_state->mqin, MyProc);
 
 		mqh = shm_mq_attach(worker_state->mqin, NULL, NULL);
 		resmq = shm_mq_receive(mqh, &nbytes, &data, false);
 
+		if (resmq == SHM_MQ_DETACHED)
+		{
+			shm_mq_detach(mqh);
+			continue;
+		}
+
 		if (resmq == SHM_MQ_SUCCESS)
 		{
 			JsonbcCommand	cmd;
@@ -740,7 +768,6 @@ jsonbd_worker_main(Datum arg)
 			}
 
 			shm_mq_detach(mqh);
-			shm_mq_clean_receiver(worker_state->mqin);
 			mqh = shm_mq_attach(worker_state->mqout, NULL, NULL);
 
 			if (iov != NULL)
@@ -751,7 +778,11 @@ jsonbd_worker_main(Datum arg)
 			if (resmq != SHM_MQ_SUCCESS)
 				elog(NOTICE, "jsonbd: backend detached early");
 
-			shm_mq_detach(mqh);
+			/*
+			 * it is not safe call shm_mq_detach here, since mq can be
+			 * already cleared.
+			 * */
+			/* shm_mq_detach(mqh); */
 			MemoryContextReset(worker_context);
 		}
 	}