Fix bug in launcher

Author: ildus

.gitignore

@@ -6,3 +6,4 @@
 regression.diffs
 regression.out
 *.plist
+tests.log

jsonbd.c

@@ -103,17 +103,16 @@ static void
 jsonbd_init_worker(shm_toc *toc, jsonbd_shm_worker *wd, int worker_num,
 		size_t queue_size)
 {
+	LWLockPadded	*locks;
 	static int mqkey = MAX_JSONBD_WORKERS + 1;
 
 	/* each worker will have two mq, for input and output */
 	wd->mqin = shm_mq_create(shm_toc_allocate(toc, queue_size), queue_size);
 	wd->mqout = shm_mq_create(shm_toc_allocate(toc, queue_size), queue_size);
 
 	/* init worker context */
-	pg_atomic_init_flag(&wd->busy);
 	wd->proc = NULL;
 	wd->dboid = InvalidOid;
-	wd->dbsem = NULL;
 
 	shm_mq_clean_receiver(wd->mqin);
 	shm_mq_clean_receiver(wd->mqout);
@@ -125,6 +124,10 @@ jsonbd_init_worker(shm_toc *toc, jsonbd_shm_worker *wd, int worker_num,
 
 	shm_toc_insert(toc, mqkey++, wd->mqin);
 	shm_toc_insert(toc, mqkey++, wd->mqout);
+
+	/* initialize worker's lwlock */
+	locks = GetNamedLWLockTranche(JSONBD_LWLOCKS_TRANCHE);
+	wd->lock = &locks[worker_num].lock;
 }
 
 static void
@@ -152,7 +155,6 @@ jsonbd_shmem_startup_hook(void)
 		/* Initialize header */
 		hdr = shm_toc_allocate(toc, sizeof(jsonbd_shm_hdr));
 		hdr->workers_ready = 0;
-		sem_init(&hdr->launcher_sem, 1, 1);
 		jsonbd_init_worker(toc, &hdr->launcher, 0, shm_mq_minimum_size);
 		shm_toc_insert(toc, 0, hdr);
 
@@ -184,6 +186,8 @@ _PG_init(void)
 
 	if (jsonbd_nworkers)
 	{
+		/* jsonbd workers and one lwlock for launcher */
+		RequestNamedLWLockTranche(JSONBD_LWLOCKS_TRANCHE, MAX_JSONBD_WORKERS + 1);
 		RequestAddinShmemSpace(jsonbd_shmem_size());
 		jsonbd_register_launcher();
 	}
@@ -289,36 +293,39 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 	shm_mq_result		resmq;
 	shm_mq_handle	   *mqh;
 	jsonbd_shm_hdr	   *hdr;
-	jsonbd_shm_worker  *wd;
+	jsonbd_shm_worker  *wd = NULL;
 
 	char			   *res;
 	Size				reslen;
-	sem_t			   *cursem = NULL;
 
 	if (jsonbd_nworkers <= 0)
 		elog(ERROR, "jsonbd workers are not available");
 
 	hdr = shm_toc_lookup(toc, 0, false);
 
-	while (cursem == NULL)
+	/*
+	 * find some not busy worker,
+	 * the backend can intercept a worker that just started by another
+	 * backend, that's ok
+	 */
+	while (true)
 	{
-		/*
-		 * find some not busy worker,
-		 * the backend can intercept a worker that just started by another
-		 * backend, that's ok
-		 */
 		for (i = 0; i < hdr->workers_ready; i++)
 		{
+			bool locked;
+
 			wd = shm_toc_lookup(toc, i + 1, false);
 			if (wd->dboid != MyDatabaseId)
 				continue;
 
 			/*
 			 * we found first worker for our database, next 'jsonbd_nworkers'
-			 * workers should be ours, and one of them should be free after sem_wait
+			 * workers should be ours
 			 */
-			cursem = wd->dbsem;
-			sem_wait(cursem);
+			locked = LWLockConditionalAcquire(wd->lock, LW_EXCLUSIVE);
+			if (locked)
+				goto comm;
+
 			for (j = i; j < (i + jsonbd_nworkers); j++)
 			{
 				wd = shm_toc_lookup(toc, j + 1, false);
@@ -327,20 +334,27 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 					/* somehow not all workers started for this database, try next */
 					continue;
 
-				if (pg_atomic_test_set_flag(&wd->busy))
+				locked = LWLockConditionalAcquire(wd->lock, LW_EXCLUSIVE);
+				if (locked)
 					goto comm;
 			}
 
-			/* should never reach here if all worked correctly */
-			sem_post(cursem);
-			elog(ERROR, "jsonbd: could not make a connection with workers");
+			/* if none of the workers were free, we just wait on last one */
+			Assert(!locked);
+			LWLockAcquire(wd->lock, LW_EXCLUSIVE);
+			goto comm;
 		}
 
 		/*
-		 * there are no workers for our database,
+		 * There are no workers for our database,
 		 * so we should launch them using our jsonbd workers launcher
+		 *
+		 * But if the launcher already locked, we should check the workers
+		 * list again
 		 */
-		sem_wait(&hdr->launcher_sem);
+		if (!LWLockAcquireOrWait(hdr->launcher.lock, LW_EXCLUSIVE))
+			continue;
+
 		shm_mq_set_sender(hdr->launcher.mqin, MyProc);
 		shm_mq_set_receiver(hdr->launcher.mqout, MyProc);
 
@@ -366,7 +380,7 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 
 		shm_mq_clean_sender(hdr->launcher.mqin);
 		shm_mq_clean_receiver(hdr->launcher.mqout);
-		sem_post(&hdr->launcher_sem);
+		LWLockRelease(hdr->launcher.lock);
 
 		if (detached)
 			elog(ERROR, "jsonbd: workers launcher was detached");
@@ -376,6 +390,8 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 	}
 
 comm:
+	Assert(wd != NULL);
+
 	detached = false;
 
 	/* send data */
@@ -405,8 +421,8 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 	/* clean self as receiver and unlock mq */
 	shm_mq_clean_sender(wd->mqin);
 	shm_mq_clean_receiver(wd->mqout);
-	pg_atomic_clear_flag(&wd->busy);
-	sem_post(cursem);
+
+	LWLockRelease(wd->lock);
 
 	if (detached)
 		elog(ERROR, "jsonbd: worker has detached");

jsonbd.h

@@ -11,6 +11,7 @@
 
 #define JSONBD_SHM_MQ_MAGIC		0xAAAA
 
+#define JSONBD_LWLOCKS_TRANCHE	"jsonbd lwlocks tranche"
 #define MAX_JSONBD_WORKERS_PER_DATABASE		3
 #define MAX_DATABASES						10 /* FIXME: need more? */
 #define MAX_JSONBD_WORKERS	(MAX_DATABASES * MAX_JSONBD_WORKERS_PER_DATABASE)
@@ -22,19 +23,16 @@ typedef enum {
 
 typedef struct jsonbd_shm_worker
 {
-	sem_t			   *dbsem;
 	shm_mq			   *mqin;
 	shm_mq			   *mqout;
-	pg_atomic_flag		busy;
 	PGPROC			   *proc;
 	volatile Oid		dboid;	/* database of the worker */
+	LWLock			   *lock;
 } jsonbd_shm_worker;
 
 /* Shared memory structures */
 typedef struct jsonbd_shm_hdr
 {
-	sem_t				launcher_sem;
-	sem_t				workers_sem[MAX_DATABASES];
 	volatile int		workers_ready;
 	jsonbd_shm_worker	launcher;
 	Latch				launcher_latch;

jsonbd_worker.c

@@ -136,7 +136,6 @@ init_worker(dsm_segment *seg)
 
 	worker_state = shm_toc_lookup(toc, worker_args->worker_num, false);
 	worker_state->proc = MyProc;
-	worker_state->dbsem = &hdr->workers_sem[worker_args->database_num];
 	worker_state->dboid = worker_args->dboid;
 
 	/* input mq */
@@ -145,9 +144,6 @@ init_worker(dsm_segment *seg)
 	/* output mq */
 	shm_mq_set_sender(worker_state->mqout, MyProc);
 
-	/* not busy at start */
-	pg_atomic_clear_flag(&worker_state->busy);
-
 	/* this context will be reset after each task */
 	Assert(worker_context == NULL);
 	worker_context = AllocSetContextCreate(TopMemoryContext,
@@ -612,9 +608,6 @@ jsonbd_launcher_main(Datum arg)
 				if (started != jsonbd_nworkers)
 					elog(NOTICE, "jsonbd: not all workers for %d has started", dboid);
 
-				/* semaphore value should be equal to workers count for database */
-				sem_init(&hdr->workers_sem[database_num], 1, started);
-
 				/* we report ok if at least one worker has started */
 				resmq = shm_mq_sendv(mqh, &((shm_mq_iovec) {"y", 2}), 1, false);
 				database_num += 1;
@@ -762,7 +755,6 @@ jsonbd_worker_main(Datum arg)
 static bool
 jsonbd_register_worker(int worker_num, Oid dboid, int database_num)
 {
-	pid_t					 pid;
 	BackgroundWorker		 worker;
 	BackgroundWorkerHandle	*bgw_handle;
 	jsonbd_worker_args		*worker_args;
@@ -804,13 +796,6 @@ jsonbd_register_worker(int worker_num, Oid dboid, int database_num)
 		return false;
 	}
 
-	/* Wait till the worker starts */
-	if (WaitForBackgroundWorkerStartup(bgw_handle, &pid) == BGWH_POSTMASTER_DIED)
-	{
-		elog(LOG, "jsonbd: postmaster died during bgworker startup");
-		return false;
-	}
-
 	/* Wait to be signalled. */
 #if PG_VERSION_NUM >= 100000
 	WaitLatch(&hdr->launcher_latch, WL_LATCH_SET, 0, PG_WAIT_EXTENSION);

sql/basic.sql

@@ -1,6 +1,6 @@
 CREATE SCHEMA comp;
 CREATE EXTENSION jsonbd SCHEMA comp;
-CREATE TABLE comp.t(a JSONB COMPRESSED jsonbd);
+CREATE TABLE comp.t(a JSONB COMPRESSION jsonbd);
 \d+ comp.t;
 
 CREATE OR REPLACE FUNCTION comp.add_record()

tests/jsonbd_test.py

@@ -3,9 +3,45 @@
 import unittest
 import random
 import json
+import os.path
+import subprocess
 
 from testgres import get_new_node
 
+# set setup base logging config, it can be turned on by `use_logging`
+# parameter on node setup
+
+import logging
+import logging.config
+
+logfile = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'tests.log')
+LOG_CONFIG = {
+    'version': 1,
+    'handlers': {
+        'console': {
+            'class': 'logging.StreamHandler',
+            'formatter': 'base_format',
+            'level': logging.DEBUG,
+        },
+        'file': {
+            'class': 'logging.FileHandler',
+            'filename': logfile,
+            'formatter': 'base_format',
+            'level': logging.DEBUG,
+        },
+    },
+    'formatters': {
+        'base_format': {
+            'format': '%(node)-5s: %(message)s',
+        },
+    },
+    'root': {
+        'handlers': ('file', ),
+        'level': 'DEBUG',
+    },
+}
+
+logging.config.dictConfig(LOG_CONFIG)
 
 insert_cmd = '''
 	INSERT INTO comp.t
@@ -32,19 +68,23 @@ def generate_dict():
 
 
 class Tests(unittest.TestCase):
+    def set_trace(self, con, command="pg_debug"):
+        pid = con.execute("select pg_backend_pid()")[0][0]
+        p = subprocess.Popen([command], stdin=subprocess.PIPE)
+        p.communicate(str(pid).encode())
+
     def test_correctness(self):
         with get_new_node('node1') as node:
             node.init()
             node.append_conf("postgresql.conf", "shared_preload_libraries='jsonbd'\n")
             node.start()
 
             node.psql('postgres', 'create extension jsonbd')
-            node.psql('postgres', 'create table t1(pk serial, a jsonb compressed jsonbd);')
+            node.psql('postgres', 'create table t1(pk serial, a jsonb compression jsonbd);')
 
             data = []
             with node.connect('postgres') as con:
-                import ipdb; ipdb.set_trace()
-                for i in range(2):
+                for i in range(1000):
                     d = generate_dict()
                     data.append(d)
                     con.execute("insert into t1 (a) values ('%s');" % json.dumps(d))