Add marine_pool

marine/__init__.py

@@ -1,4 +1,5 @@
 from .marine import Marine
+from .marine_pool import MarinePool
 
 
 # TODO: create simple functions for using Marine instead of exposing the class

marine/marine_pool.py

@@ -0,0 +1,64 @@
+import math
+import multiprocessing
+from itertools import repeat
+from typing import List, Dict, Optional, Tuple, ClassVar
+
+from marine import Marine
+
+
+class MarinePool:
+    _marine_instance: ClassVar[Optional[Marine]] = None
+
+    def __init__(
+        self, lib_path: str, epan_auto_reset_count: int, process_count: int = 4
+    ):
+        self._lib_path = lib_path
+        self._epan_auto_reset_count = epan_auto_reset_count
+        self._process_count = process_count
+
+    def __enter__(self):
+        ctx = multiprocessing.get_context("spawn")
+        # Using spawn so child processes won't get the already initialized marine from the parent process.
+        # We do that because initializing marine more than one time in a process causes SIGTRAP
+        self.pool = ctx.Pool(
+            self._process_count,
+            initializer=self._init_marine,
+            initargs=[self._lib_path, self._epan_auto_reset_count],
+        )
+        return self
+
+    def filter_and_parse(
+        self,
+        packets: List[bytes],
+        bpf: Optional[str] = None,
+        display_filter: Optional[str] = None,
+        fields: Optional[List[str]] = None,
+    ) -> List[Tuple[bool, Dict[str, str]]]:
+        if len(packets) == 0:
+            return []
+
+        chunk_size = int(math.ceil(len(packets) / float(self._process_count)))
+        return self.pool.starmap(
+            self._filter_and_parse,
+            zip(packets, repeat(bpf), repeat(display_filter), repeat(fields)),
+            chunksize=chunk_size,
+        )
+
+    @classmethod
+    def _init_marine(cls, lib_path: str, epan_auto_reset_count: int) -> None:
+        cls._marine_instance = Marine(lib_path, epan_auto_reset_count)
+
+    @classmethod
+    def _filter_and_parse(
+        cls,
+        packet: bytes,
+        bpf: Optional[str] = None,
+        display_filter: Optional[str] = None,
+        fields: Optional[list] = None,
+    ) -> (bool, Dict[str, str]):
+        return cls._marine_instance.filter_and_parse(
+            packet, bpf, display_filter, fields
+        )
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.pool.close()

tests/fixtures/marine/marine_fixtures.py

@@ -1,12 +1,12 @@
-import os
-from typing import List
+from pathlib import Path
+from typing import List, Union
 
 import pytest
 from pypacker.layer12 import ethernet
 from pypacker.layer3 import ip
 from pypacker.layer4 import tcp
 
-from marine import Marine
+from marine import Marine, MarinePool
 
 
 @pytest.fixture
@@ -25,8 +25,9 @@ def tcp_packet() -> bytes:
 
 
 @pytest.fixture(scope="session")
-def libmarine_path() -> str:
-    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "libmarine.so")
+def marine_so_path() -> str:
+    path = Path(__file__).parent / "libmarine.so"
+    return str(path.resolve())
 
 
 @pytest.fixture(scope="session")
@@ -35,5 +36,16 @@ def epan_auto_reset_count() -> int:
 
 
 @pytest.fixture(scope="session")
-def marine_instance(libmarine_path: str, epan_auto_reset_count: int) -> Marine:
-    return Marine(libmarine_path, epan_auto_reset_count=epan_auto_reset_count)
+def marine_instance(marine_so_path: str, epan_auto_reset_count: int) -> Marine:
+    return Marine(marine_so_path, epan_auto_reset_count=epan_auto_reset_count)
+
+
+@pytest.fixture(scope="session")
+def marine_pool_instance(marine_so_path: str, epan_auto_reset_count: int) -> MarinePool:
+    with MarinePool(marine_so_path, epan_auto_reset_count) as mp:
+        yield mp
+
+
+@pytest.fixture(scope="session", params=["marine_instance", "marine_pool_instance"])
+def marine_or_marine_pool(request) -> Union[Marine, MarinePool]:
+    return request.getfixturevalue(request.param)