- multiprocessing SimpleQueue takes a context as a required arg

  lookup the context and pass it into class initalization of the SimpleQueue
- for simplification, follow example patterns of importing multiprocessing as mp
- multiprocessing.Queue can cause surprising results which are avoided using a queue manager Manager(). See https://docs.python.org/3.7/library/multiprocessing.html
- optparse passes in self to the check_existing_file, so set as _
- optparse calls the callback with many more args then originally setup for.
- fix env lookup variables which were missed when rebirthdb merged back with rethinkdb

Author: Andrew Sawyers

rethinkdb/_export.py

@@ -23,7 +23,7 @@
 import ctypes
 import datetime
 import json
-import multiprocessing
+import multiprocessing as mp
 import numbers
 import optparse
 import os
@@ -259,11 +259,12 @@ def export_table(db, table, directory, options, error_queue, progress_info, sind
         with sindex_counter.get_lock():
             sindex_counter.value += len(table_info["indexes"])
         # -- start the writer
-        task_queue = SimpleQueue()
+        ctx = mp.get_context(mp.get_start_method())
+        task_queue = SimpleQueue(ctx=ctx)
         writer = None
         if options.format == "json":
             filename = directory + "/%s/%s.json" % (db, table)
-            writer = multiprocessing.Process(
+            writer = mp.Process(
                 target=json_writer,
                 args=(
                     filename,
@@ -273,7 +274,7 @@ def export_table(db, table, directory, options, error_queue, progress_info, sind
                     options.format))
         elif options.format == "csv":
             filename = directory + "/%s/%s.csv" % (db, table)
-            writer = multiprocessing.Process(
+            writer = mp.Process(
                 target=csv_writer,
                 args=(
                     filename,
@@ -283,7 +284,7 @@ def export_table(db, table, directory, options, error_queue, progress_info, sind
                     error_queue))
         elif options.format == "ndjson":
             filename = directory + "/%s/%s.ndjson" % (db, table)
-            writer = multiprocessing.Process(
+            writer = mp.Process(
                 target=json_writer,
                 args=(
                     filename,
@@ -388,12 +389,13 @@ def update_progress(progress_info, options):
 
 def run_clients(options, workingDir, db_table_set):
     # Spawn one client for each db.table, up to options.clients at a time
-    exit_event = multiprocessing.Event()
+    exit_event = mp.Event()
     processes = []
-    error_queue = SimpleQueue()
-    interrupt_event = multiprocessing.Event()
-    sindex_counter = multiprocessing.Value(ctypes.c_longlong, 0)
-    hook_counter = multiprocessing.Value(ctypes.c_longlong, 0)
+    ctx = mp.get_context(mp.get_start_method())
+    error_queue = SimpleQueue(ctx=ctx)
+    interrupt_event = mp.Event()
+    sindex_counter = mp.Value(ctypes.c_longlong, 0)
+    hook_counter = mp.Value(ctypes.c_longlong, 0)
 
     signal.signal(signal.SIGINT, lambda a, b: abort_export(a, b, exit_event, interrupt_event))
     errors = []
@@ -405,8 +407,8 @@ def run_clients(options, workingDir, db_table_set):
 
             tableSize = int(options.retryQuery("count", query.db(db).table(table).info()['doc_count_estimates'].sum()))
 
-            progress_info.append((multiprocessing.Value(ctypes.c_longlong, 0),
-                                  multiprocessing.Value(ctypes.c_longlong, tableSize)))
+            progress_info.append((mp.Value(ctypes.c_longlong, 0),
+                                  mp.Value(ctypes.c_longlong, tableSize)))
             arg_lists.append((db, table,
                               workingDir,
                               options,
@@ -428,7 +430,7 @@ def run_clients(options, workingDir, db_table_set):
             processes = [process for process in processes if process.is_alive()]
 
             if len(processes) < options.clients and len(arg_lists) > 0:
-                newProcess = multiprocessing.Process(target=export_table, args=arg_lists.pop(0))
+                newProcess = mp.Process(target=export_table, args=arg_lists.pop(0))
                 newProcess.start()
                 processes.append(newProcess)
 

rethinkdb/_import.py

@@ -26,7 +26,7 @@
 import csv
 import ctypes
 import json
-import multiprocessing
+import multiprocessing as mp
 import optparse
 import os
 import signal
@@ -110,12 +110,12 @@ def __init__(
         self.query_runner = query_runner
 
         # reporting information
-        self._bytes_size = multiprocessing.Value(ctypes.c_longlong, -1)
-        self._bytes_read = multiprocessing.Value(ctypes.c_longlong, -1)
+        self._bytes_size = mp.Value(ctypes.c_longlong, -1)
+        self._bytes_read = mp.Value(ctypes.c_longlong, -1)
 
-        self._total_rows = multiprocessing.Value(ctypes.c_longlong, -1)
-        self._rows_read = multiprocessing.Value(ctypes.c_longlong, 0)
-        self._rows_written = multiprocessing.Value(ctypes.c_longlong, 0)
+        self._total_rows = mp.Value(ctypes.c_longlong, -1)
+        self._rows_read = mp.Value(ctypes.c_longlong, 0)
+        self._rows_written = mp.Value(ctypes.c_longlong, 0)
 
         # source
         if hasattr(source, 'read'):
@@ -957,7 +957,7 @@ def table_writer(tables, options, work_queue, error_queue, warning_queue, exit_e
                             nesting_depth=MAX_NESTING_DEPTH),
                         durability=options.durability,
                         conflict=conflict_action,
-                        ignore_write_hook=True))
+                        ))
 
                 if res["errors"] > 0:
                     raise RuntimeError("Error when importing into table '%s.%s': %s" % (db, table, res["first_error"]))
@@ -1083,13 +1083,15 @@ def import_tables(options, sources, files_ignored=None):
 
     tables = dict(((x.db, x.table), x) for x in sources)  # (db, table) => table
 
-    work_queue = Queue(options.clients * 3)
-    error_queue = SimpleQueue()
-    warning_queue = SimpleQueue()
-    exit_event = multiprocessing.Event()
-    interrupt_event = multiprocessing.Event()
+    ctx = mp.get_context(mp.get_start_method())
+    max_queue_size = options.clients * 3
+    work_queue = mp.Manager().Queue(max_queue_size)
+    error_queue = SimpleQueue(ctx=ctx)
+    warning_queue = SimpleQueue(ctx=ctx)
+    exit_event = mp.Event()
+    interrupt_event = mp.Event()
 
-    timing_queue = SimpleQueue()
+    timing_queue = SimpleQueue(ctx=ctx)
 
     errors = []
     warnings = []
@@ -1166,7 +1168,7 @@ def drain_queues():
     try:
         # - start the progress bar
         if not options.quiet:
-            progress_bar = multiprocessing.Process(
+            progress_bar = mp.Process(
                 target=update_progress,
                 name="progress bar",
                 args=(sources, options.debug, exit_event, progress_bar_sleep)
@@ -1178,7 +1180,7 @@ def drain_queues():
         writers = []
         pools.append(writers)
         for i in range(options.clients):
-            writer = multiprocessing.Process(
+            writer = mp.Process(
                 target=table_writer,
                 name="table writer %d" %
                 i,
@@ -1202,7 +1204,7 @@ def drain_queues():
                 # add a workers to fill up the readers pool
                 while len(readers) < options.clients:
                     table = next(file_iter)
-                    reader = multiprocessing.Process(
+                    reader = mp.Process(
                         target=table.read_to_queue,
                         name="table reader %s.%s" %
                         (table.db,

rethinkdb/utils_common.py

@@ -151,7 +151,6 @@ def format_epilog(self, formatter):
         return self.epilog or ''
 
     def __init__(self, *args, **kwargs):
-
         # -- Type Checkers
 
         def check_tls_option(opt_str, value):
@@ -178,7 +177,7 @@ def check_positive_int(opt_str, value):
 
             return int(value)
 
-        def check_existing_file(opt_str, value):
+        def check_existing_file(_, opt_str, value):
             if not os.path.isfile(value):
                 raise optparse.OptionValueError('%s value was not an existing file: %s' % (opt_str, value))
 
@@ -207,7 +206,10 @@ def file_contents(opt_str, value):
 
         # -- Callbacks
 
-        def combined_connect_action(value, parser):
+        def combined_connect_action(obj, opt, value, parser, *args, **kwargs):
+            """optparse.takeaction() calls the callback (which this is set as)
+               with the following args: self, opt, value, parser *args, **kwargs
+            """
             res = self.__connectRegex.match(value)
             if not res:
                 raise optparse.OptionValueError("Invalid 'host:port' format: %s" % value)
@@ -295,15 +297,15 @@ def take_action(self, action, dest, opt, value, values, parser):
             help='driver port of a rethinkdb server',
             type='int',
             default=os.environ.get(
-                'REBIRTHDB_DRIVER_PORT',
+                'RETHINKDB_DRIVER_PORT',
                 net.DEFAULT_PORT))
         connection_group.add_option(
             '--host-name',
             dest='hostname',
             metavar='HOST',
             help='host and driver port of a rethinkdb server',
             default=os.environ.get(
-                'REBIRTHDB_HOSTNAME',
+                'RETHINKDB_HOSTNAME',
                 'localhost'))
         connection_group.add_option(
             '-u',
@@ -312,7 +314,7 @@ def take_action(self, action, dest, opt, value, values, parser):
             metavar='USERNAME',
             help='user name to connect as',
             default=os.environ.get(
-                'REBIRTHDB_USER',
+                'RETHINKDB_USER',
                 'admin'))
         connection_group.add_option(
             '-p',
@@ -344,12 +346,13 @@ def parse_args(self, *args, **kwargs):
 
         # - validate ENV variables
 
-        if 'REBIRTHDB_DRIVER_PORT' in os.environ:
-            driver_port = os.environ['REBIRTHDB_DRIVER_PORT']
+        if 'RETHINKDB_DRIVER_PORT' in os.environ:
+            driver_port = os.environ['RETHINKDB_DRIVER_PORT']
 
             if not isinstance(driver_port, int) or driver_port < 1:
-                self.error('ENV variable REBIRTHDB_DRIVER_PORT is not a useable integer: %s'
-                           % os.environ['REBIRTHDB_DRIVER_PORT'])
+                self.error('ENV variable RETHINKDB_DRIVER_PORT is not a useable '
+                           'integer: %s'
+                           % os.environ['RETHINKDB_DRIVER_PORT'])
 
         # - parse options