Make bench script async

Author: ildus

bench/bench.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import asyncio
 import contextlib
 import glob
 import json
@@ -29,37 +30,61 @@ def cwd(path):
         os.chdir(curdir)
 
 
-with get_new_node('node1') as node:
-    node.init()
-    node.append_conf('postgresql.conf', conf)
-    node.start()
+async def insert_data(con, files, table_name, table_name_c):
+    while True:
+        try:
+            filename = files.pop()
+            print(filename)
+        except IndexError:
+            break
 
-    node.safe_psql('create extension jsonbd')
+        with open(filename, 'r') as f:
+            data = json.load(f)
 
-    for name, root_dir in sources:
-        table_name = name
-        table_name_c = '%s_c' % name
+        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))
 
-        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:
+def main(loop):
+    with get_new_node('node1') as node:
+        node.init()
+        node.append_conf('postgresql.conf', conf)
+        node.start()
+
+        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
+
+            node.safe_psql('postgres', 'create table %s(a jsonb)' % table_name)
+            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)):
                 for filename in glob.iglob('**/*.json', recursive=True):
                     if filename == 'package.json':
                         continue
 
-                    with open(filename, 'r') as f:
-                        data = json.load(f)
+                    files.append(filename)
+
+                coroutines = [insert_data(con, files, table_name, table_name_c)
+                        for con in connections]
+                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))
 
-                    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))
+if __name__ == '__main__':
+    loop = asyncio.get_event_loop()
+    main(loop)
+    loop.close()