refactor: multiprocess helper class

Author: NguyenHoangSon96

influxdb_client_3/write_client/client/util/multiprocessing_helper.py

@@ -7,8 +7,9 @@
 import logging
 import multiprocessing
 
-from influxdb_client_3.write_client import WriteOptions
+from influxdb_client_3 import write_client_options
 from influxdb_client_3.exceptions import InfluxDBError
+from influxdb_client_3.write_client import WriteOptions, WriteApi
 
 logger = logging.getLogger('influxdb_client.client.util.multiprocessing_helper')
 
@@ -118,17 +119,17 @@ def main():
     __started__ = False
     __disposed__ = False
 
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, start_method='spawn', **kwargs) -> None:
         """
         Initialize defaults.
 
-        For more information how to initialize the writer see the examples above.
+        For more information on how to initialize the writer, see the examples above.
 
         :param kwargs: arguments are passed into ``__init__`` function of ``InfluxDBClient`` and ``write_api``.
         """
+        multiprocessing.set_start_method(start_method, force=True)
         multiprocessing.Process.__init__(self)
         self.kwargs = kwargs
-        self.client = None
         self.write_api = None
         self.queue_ = multiprocessing.Manager().Queue()
 
@@ -146,13 +147,22 @@ def write(self, **kwargs) -> None:
         self.queue_.put(kwargs)
 
     def run(self):
-        """Initialize ``InfluxDBClient`` and waits for data to writes into InfluxDB."""
+        """Initialize ``InfluxDBClient3`` and wait for data to write into InfluxDB."""
         # Initialize Client and Write API
-        self.client = InfluxDBClient(**self.kwargs)
-        self.write_api = self.client.write_api(write_options=self.kwargs.get('write_options', WriteOptions()),
-                                               success_callback=self.kwargs.get('success_callback', _success_callback),
-                                               error_callback=self.kwargs.get('error_callback', _error_callback),
-                                               retry_callback=self.kwargs.get('retry_callback', _retry_callback))
+        wco = write_client_options(write_options=self.kwargs.get('write_options', WriteOptions()),
+                                   success_callback=self.kwargs.get('success_callback', _success_callback),
+                                   error_callback=self.kwargs.get('error_callback', _error_callback),
+                                   retry_callback=self.kwargs.get('retry_callback', _retry_callback)
+                                   )
+
+        self.write_api = WriteApi(
+            bucket=self.kwargs.get('database'),
+            org=self.kwargs.get('org'),
+            default_header=self.kwargs.get('default_header'),
+            rest_client=self.kwargs.get('rest_client'),
+            write_client_options=wco,
+
+        )
         # Infinite loop - until poison pill
         while True:
             next_record = self.queue_.get()
@@ -180,9 +190,6 @@ def terminate(self) -> None:
             logger.info("flushing data...")
             self.write_api.__del__()
             self.write_api = None
-        if self.client:
-            self.client.__del__()
-            self.client = None
             logger.info("closed")
 
     def __enter__(self):

influxdb_client_3/write_client/client/write_api.py

@@ -981,6 +981,49 @@ def _should_gzip(self, payload: str, enable_gzip: bool = False, gzip_threshold:
 
         return False
 
+    @staticmethod
+    def _on_error(ex):
+        logger.error("unexpected error during batching: %s", ex)
+
+    def _to_response(self, data: _BatchItem, delay: datetime.timedelta):
+        return rx.of(data).pipe(
+            ops.subscribe_on(self._write_options.write_scheduler),
+            # use delay if its specified
+            ops.delay(duetime=delay, scheduler=self._write_options.write_scheduler),
+            # invoke http call
+            ops.map(lambda x: self._http(x, **x.key.kwargs)),
+            # catch exception to fail batch response
+            ops.catch(handler=lambda exception, source: rx.just(_BatchResponse(exception=exception, data=data))),
+        )
+
+    def _on_next(self, response: _BatchResponse):
+        if response.exception:
+            logger.error("The batch item wasn't processed successfully because: %s", response.exception)
+            if self._error_callback:
+                try:
+                    self._error_callback(response.data.to_key_tuple(), response.data.data, response.exception)
+                except Exception as e:
+                    """
+                    Unfortunately, because callbacks are user-provided generic code, exceptions can be entirely
+                    arbitrary
+
+                    We trap it, log that it occurred and then proceed - there's not much more that we can
+                    really do.
+                    """
+                    logger.error("The configured error callback threw an exception: %s", e)
+
+        else:
+            logger.debug("The batch item: %s was processed successfully.", response)
+            if self._success_callback:
+                try:
+                    self._success_callback(response.data.to_key_tuple(), response.data.data)
+                except Exception as e:
+                    logger.error("The configured success callback threw an exception: %s", e)
+
+    def _on_complete(self):
+        self._disposable.dispose()
+        logger.debug("the batching processor was disposed")
+
     def _append_default_tag(self, key, val, record):
         from influxdb_client_3.write_client import Point
         if isinstance(record, bytes) or isinstance(record, str):
@@ -1065,49 +1108,6 @@ def __getstate__(self):
         del state['_disposable']
         return state
 
-    @staticmethod
-    def _on_error(ex):
-        logger.error("unexpected error during batching: %s", ex)
-
-    def _on_complete(self):
-        self._disposable.dispose()
-        logger.debug("the batching processor was disposed")
-
-    def _to_response(self, data: _BatchItem, delay: datetime.timedelta):
-        return rx.of(data).pipe(
-            ops.subscribe_on(self._write_options.write_scheduler),
-            # use delay if its specified
-            ops.delay(duetime=delay, scheduler=self._write_options.write_scheduler),
-            # invoke http call
-            ops.map(lambda x: self._http(x, **x.key.kwargs)),
-            # catch exception to fail batch response
-            ops.catch(handler=lambda exception, source: rx.just(_BatchResponse(exception=exception, data=data))),
-        )
-
-    def _on_next(self, response: _BatchResponse):
-        if response.exception:
-            logger.error("The batch item wasn't processed successfully because: %s", response.exception)
-            if self._error_callback:
-                try:
-                    self._error_callback(response.data.to_key_tuple(), response.data.data, response.exception)
-                except Exception as e:
-                    """
-                    Unfortunately, because callbacks are user-provided generic code, exceptions can be entirely
-                    arbitrary
-
-                    We trap it, log that it occurred and then proceed - there's not much more that we can
-                    really do.
-                    """
-                    logger.error("The configured error callback threw an exception: %s", e)
-
-        else:
-            logger.debug("The batch item: %s was processed successfully.", response)
-            if self._success_callback:
-                try:
-                    self._success_callback(response.data.to_key_tuple(), response.data.data)
-                except Exception as e:
-                    logger.error("The configured success callback threw an exception: %s", e)
-
     def __setstate__(self, state):
         """Set your object with the provided dict."""
         self.__dict__.update(state)

tests/test_influxdb_client_3_integration.py

@@ -1,10 +1,11 @@
-import asyncio
 import json
 import logging
+import multiprocessing
 import os
 import random
 import string
 import time
+import asyncio
 import unittest
 
 import pandas as pd
@@ -343,28 +344,77 @@ def test_batch_write_closed(self):
             list_results = reader.to_pylist()
             self.assertEqual(data_size, len(list_results))
 
-    @pytest.mark.skipif(running_on_posix, reason="Skipping this test in POSIX environments")
+    # @pytest.mark.skipif(running_on_posix, reason="Skipping this test in POSIX environments")
     def test_multiprocessing_helper(self):
-        org = 'my-org'
-        writer = MultiprocessingWriter(
+        default_header = {
+            'Authorization': f'Token {self.token}'
+        }
+        rest = rest_client.RestClient(
+            base_url=self.host,
+            default_header=default_header,
+        )
+
+        with MultiprocessingWriter(
+                host=self.host,
+                database=self.database,
+                token=self.token,
+                org='my-org',
+                default_header=default_header,
+                rest_client=rest,
+                write_options=WriteOptions(batch_size=1)) as mp:
+            self.assertEqual(multiprocessing.get_start_method(), 'spawn')
+
+            measurement = f'test{random_hex(6)}'.lower()
+            for x in range(1, 5):
+                time.sleep(0.5)
+                mp.write(
+                    bucket=self.database,
+                    record=f"{measurement},tag=a value=\"number{x}\" {time.time_ns()}"
+                )
+
+        time.sleep(1)
+        df = self.client.query(f'select * from {measurement}', mode="pandas")
+        self.assertEqual(4, len(df))
+
+    def test_multiprocessing_start_method_forkserver(self):
+        default_header = {
+            'Authorization': f'Token {self.token}'
+        }
+
+        MultiprocessingWriter(
             host=self.host,
             database=self.database,
             token=self.token,
-            org=org,
-            write_options=WriteOptions(batch_size=1))
-        writer.start()
-        measurement = f'test{random_hex(6)}'.lower()
-        for x in range(1, 10):
-            time.sleep(0.2)
-            writer.write(
-                bucket=self.database,
-                record=f"{measurement},tag=a value=\"number{x}\" {time.time_ns()}"
-            )
-        writer.__del__()
+            org='my-org',
+            default_header=default_header,
+            rest_client=(rest_client.RestClient(
+                base_url=self.host,
+                default_header=default_header,
+            )),
+            write_options=WriteOptions(batch_size=1),
+            start_method='forkserver'
+        )
+        self.assertEqual(multiprocessing.get_start_method(), 'forkserver')
 
-        time.sleep(1)
-        df = self.client.query(f'select * from {measurement}', mode="pandas")
-        self.assertEqual(9, len(df))
+    def test_multiprocessing_start_method_fork(self):
+        default_header = {
+            'Authorization': f'Token {self.token}'
+        }
+
+        MultiprocessingWriter(
+            host=self.host,
+            database=self.database,
+            token=self.token,
+            org='my-org',
+            default_header=default_header,
+            rest_client=(rest_client.RestClient(
+                base_url=self.host,
+                default_header=default_header,
+            )),
+            write_options=WriteOptions(batch_size=1),
+            start_method='fork'
+        )
+        self.assertEqual(multiprocessing.get_start_method(), 'fork')
 
     test_cert = """-----BEGIN CERTIFICATE-----
 MIIDUzCCAjugAwIBAgIUZB55ULutbc9gy6xLp1BkTQU7siowDQYJKoZIhvcNAQEL