Accelerate SRX address-book construction with an optional Rust extension

.github/workflows/native.yml

@@ -0,0 +1,61 @@
+name: Native SRX address book
+
+on:
+  push:
+    paths:
+      - native/**
+      - capirca/**
+      - tests/**
+      - requirements.txt
+      - setup.py
+      - tools/benchmark_srx.py
+      - .github/workflows/native.yml
+  pull_request:
+    paths:
+      - native/**
+      - capirca/**
+      - tests/**
+      - requirements.txt
+      - setup.py
+      - tools/benchmark_srx.py
+      - .github/workflows/native.yml
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    timeout-minutes: 15
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - os: ubuntu-latest
+            python: '3.9'
+            tests: tests/lib/junipersrx_test.py tests/lib/srx_native_test.py
+          - os: ubuntu-latest
+            python: '3.14'
+            tests: tests
+          - os: macos-latest
+            python: '3.14'
+            tests: tests
+          - os: windows-latest
+            python: '3.14'
+            tests: tests
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5
+        with:
+          python-version: ${{ matrix.python }}
+      - run: rustup toolchain install stable --profile minimal --component rustfmt --component clippy
+      - run: rustup default stable
+      - run: python -m pip install -e . -r requirements.txt pytest ./native
+      - run: cargo fmt --manifest-path native/Cargo.toml --check
+      - run: cargo clippy --locked --manifest-path native/Cargo.toml --all-targets -- -D warnings
+      # Upstream OpenConfig uses datetime.UTC, which requires Python 3.11.
+      - run: python -m pytest -q ${{ matrix.tests }}
+      - if: matrix.python != '3.9'
+        run: python tools/benchmark_srx.py --rounds 1 --output benchmark.json
+      - run: python -m pip uninstall -y capirca-native
+      - run: python -m pytest -q ${{ matrix.tests }}

README.md

@@ -638,6 +638,10 @@ pip install capirca --user
 
 ### Basic Usage
 
+An optional [Rust extension](./native/README.md) accelerates Juniper SRX
+address-book construction. Install it separately with `pip install ./native`
+from a source checkout.
+
 There are a number of command-line arguments that can be used with `aclgen`.
 
 ```

capirca/lib/junipersrx.py

@@ -27,6 +27,13 @@
 from capirca.lib import nacaddr
 import six
 
+try:
+  import _capirca_native
+except ModuleNotFoundError as error:
+  if error.name != '_capirca_native':
+    raise
+  _capirca_native = None
+
 
 ICMP_TERM_LIMIT = 8
 
@@ -569,8 +576,7 @@ def _TranslatePolicy(self, pol, exp_info):
                 'WARNING: Term %s has 0 valid source IPs, skipping.', term.name)
             continue
           term.source_address = valid_addrs
-          for addr in term.source_address:
-            self._BuildAddressBook(self.from_zone, addr)
+          self._BuildAddressBookFromList(self.from_zone, term.source_address)
 
         # Filter destination_address based on filter_type & add to address book
         if term.destination_address:
@@ -584,8 +590,7 @@ def _TranslatePolicy(self, pol, exp_info):
                 term.name)
             continue
           term.destination_address = valid_addrs
-          for addr in term.destination_address:
-            self._BuildAddressBook(self.to_zone, addr)
+          self._BuildAddressBookFromList(self.to_zone, term.destination_address)
 
         new_term = Term(term, self.from_zone, self.to_zone, self.expresspath,
                         verbose)
@@ -696,6 +701,35 @@ def Chunks(l):
       del terms[:]
       terms.extend(expanded_terms)
 
+  def _BuildAddressBookFromList(self, zone, addresses):
+    """Insert addresses in order, retaining the selected Python objects."""
+    if _capirca_native is None:
+      for address in addresses:
+        self._BuildAddressBook(zone, address)
+      return
+
+    groups = collections.OrderedDict()
+    for address in addresses:
+      groups.setdefault(address.parent_token, []).append(address)
+    if not groups:
+      return
+    if zone not in self.addressbook:
+      self.addressbook[zone] = collections.defaultdict(list)
+    for name, incoming in groups.items():
+      current = self.addressbook[zone][name]
+      candidates = current + incoming
+      # Scoped IPv6 equality and custom subnet methods are Python contracts.
+      if any(type(ip) not in (nacaddr.IPv4, nacaddr.IPv6) or
+             getattr(ip.network_address, 'scope_id', None) is not None
+             for ip in candidates):
+        for address in incoming:
+          self._BuildAddressBook(zone, address)
+        continue
+      networks = [(ip.version, int(ip.network_address), ip.prefixlen)
+                  for ip in candidates]
+      selected = _capirca_native.address_book_indices(networks, len(current))
+      current[:] = [candidates[index] for index in selected]
+
   def _BuildAddressBook(self, zone, address):
     """Create the address book configuration entries.
 

native/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "capirca-native"
+version = "0.1.0"
+edition = "2021"
+rust-version = "1.83"
+license = "Apache-2.0"
+publish = false
+
+[lib]
+name = "_capirca_native"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = "0.29.2"