Fix race condition

Author: ildus

bench/bench.py

@@ -30,24 +30,28 @@ def cwd(path):
         os.chdir(curdir)
 
 
-async def insert_data(con, files, table_name, table_name_c):
-    while True:
-        try:
-            filename = files.pop()
+async def insert_data(node, files, table_name, table_name_c):
+    sql = 'insert into {} values ($1)'
+    sql1 = sql.format(table_name)
+    sql2 = sql.format(table_name_c)
+
+    async with node.connect() as con:
+        while True:
+            try:
+                filename = files.pop()
+            except IndexError:
+                break
+
             print(filename)
-        except IndexError:
-            break
 
-        with open(filename, 'r') as f:
-            data = json.load(f)
+            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 ($1)'
-                    for i in range(100):
-                        await con.execute(sql.format(table_name), json.dumps(obj))
-                        await con.execute(sql.format(table_name_c), json.dumps(obj))
+            if isinstance(data, dict):
+                if 'rounds' in data:
+                    for obj in data['rounds']:
+                        await con.execute(sql1, json.dumps(obj))
+                        await con.execute(sql2, json.dumps(obj))
 
 
 def main(loop):
@@ -58,8 +62,6 @@ def main(loop):
 
         node.safe_psql('postgres', 'create extension jsonbd')
 
-        connections = [node.connect() for i in range(3)]
-
         for name, root_dir in sources:
             table_name = name
             table_name_c = '%s_c' % name
@@ -68,20 +70,21 @@ def main(loop):
             node.safe_psql('postgres', 'create table %s(a jsonb compression jsonbd)' % table_name_c)
             node.safe_psql('postgres', 'alter table %s alter column a set storage external' % table_name)
 
-            files = []
             with cwd(os.path.abspath(root_dir)):
+                files = []
                 for filename in glob.iglob('**/*.json', recursive=True):
                     if filename == 'package.json':
                         continue
 
                     files.append(filename)
 
-                coroutines = [insert_data(con, files, table_name, table_name_c)
-                        for con in connections]
+                print(len(files))
+                coroutines = [insert_data(node, files, table_name, table_name_c)
+                    for i in range(4)]
                 loop.run_until_complete(asyncio.gather(*coroutines))
 
-                print(node.safe_psql('postgres', "select pg_size_pretty(pg_total_relation_size('%s'))" % table_name))
-                print(node.safe_psql('postgres', "select pg_size_pretty(pg_total_relation_size('%s'))" % table_name_c))
+            print(node.safe_psql('postgres', "select pg_size_pretty(pg_total_relation_size('%s'))" % table_name))
+            print(node.safe_psql('postgres', "select pg_size_pretty(pg_total_relation_size('%s'))" % table_name_c))
 
 
 if __name__ == '__main__':

jsonbd.c

@@ -304,6 +304,7 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 
 	hdr = shm_toc_lookup(toc, 0, false);
 
+begin:
 	/*
 	 * find some not busy worker,
 	 * the backend can intercept a worker that just started by another
@@ -398,6 +399,17 @@ jsonbd_communicate(shm_mq_iovec *iov, int iov_len,
 
 comm:
 	Assert(wd != NULL);
+	Assert(LWLockHeldByMe(wd->lock));
+
+	/*
+	 * Even if we got the lock it doesn't mean that worker is free,
+	 * so try to set busy flag
+	 */
+	if (!pg_atomic_test_set_flag(&wd->busy))
+	{
+		LWLockRelease(wd->lock);
+		goto begin;
+	}
 
 	detached = false;
 

jsonbd.h

@@ -29,6 +29,7 @@ typedef struct jsonbd_shm_worker
 	volatile Oid		dboid;	/* database of the worker */
 	LWLock			   *lock;
 	Latch				latch;
+	pg_atomic_flag		busy;	/* worker is busy */
 } jsonbd_shm_worker;
 
 /* Shared memory structures */

jsonbd_worker.c

@@ -161,6 +161,7 @@ init_worker(dsm_segment *seg)
 	/* Set launcher free */
 	SetLatch(&hdr->launcher_latch);
 	InitLatch(&worker_state->latch);
+	pg_atomic_init_flag(&worker_state->busy);
 
 	/* make this worker visible in backend cycle */
 	hdr->workers_ready++;
@@ -778,12 +779,9 @@ jsonbd_worker_main(Datum arg)
 			if (resmq != SHM_MQ_SUCCESS)
 				elog(NOTICE, "jsonbd: backend detached early");
 
-			/*
-			 * it is not safe call shm_mq_detach here, since mq can be
-			 * already cleared.
-			 * */
-			/* shm_mq_detach(mqh); */
+			shm_mq_detach(mqh);
 			MemoryContextReset(worker_context);
+			pg_atomic_clear_flag(&worker_state->busy);
 		}
 	}