manager: implement pfx table itaration for primary pfx table

mroethke · mroethke · commit 4fd5194ce12e · 2017-02-20T15:03:19.000+01:00
This patch allows iteration in two ways, callback driven and iterator
driven. The later approach uses threads to run the callback in the
background.
diff --git a/ffi_build.py b/ffi_build.py
@@ -15,6 +15,7 @@
         extern "Python" void rtr_mgr_status_callback(const struct rtr_mgr_group *, enum rtr_mgr_status, const struct rtr_socket *, void *);
         extern "Python" void pfx_update_callback(struct pfx_table *pfx_table, const struct pfx_record record, const bool added);
         extern "Python" void spki_update_callback(struct spki_table *spki_table, const struct spki_record record, const bool added);
+        extern "Python" void pfx_table_callback(const struct pfx_record *pfx_record, void *data);
                 """)
 
 ffibuilder.set_source("_rtrlib",
diff --git a/rtrlib/callbacks.py b/rtrlib/callbacks.py
@@ -74,6 +74,17 @@ def status_callback(rtr_mgr_group, group_status, rtr_socket, object_handle):
         )
 
 
+@ffi.def_extern(name="pfx_table_callback")
+def pfx_table_callback(pfx_record, object_handle):
+    """
+    Wraps the pfx_table callback, used for iteration of the pfx table,
+    to hide cffi specifics
+    """
+    callback, data = ffi.from_handle(object_handle)
+
+    callback(PFXRecord(pfx_record), data)
+
+
 def pfx_update_callback_wrapper(func):
     """
     Wraps the given python function and wraps it to hide cffi specifics.
diff --git a/rtrlib/records.py b/rtrlib/records.py
@@ -56,6 +56,26 @@ def socket(self):
         """
         return RTRSocket(self._record.socket)
 
+    def __str__(self):
+        return "{prefix}/{min}-{max} {asn}".format(prefix=self.prefix,
+                                                   min=self.min_len,
+                                                   max=self.max_len,
+                                                   asn=self.asn
+                                                   )
+
+
+def copy_pfx_record(record):
+    """Copy a pfx record."""
+    cdata = record._record
+    new_record = ffi.new('struct pfx_record *')
+    new_record.asn = cdata.asn
+    new_record.prefix = cdata.prefix
+    new_record.min_len = cdata.min_len
+    new_record.max_len = cdata.max_len
+    new_record.socket = cdata.socket
+
+    return PFXRecord(new_record)
+
 
 class SPKIRecord(object):
     """
diff --git a/rtrlib/rtr_manager.py b/rtrlib/rtr_manager.py
@@ -11,14 +11,20 @@
 import time
 import logging
 import signal
+import threading
 
 from enum import Enum
 from _rtrlib import ffi, lib
 
-import six
 import rtrlib.callbacks as callbacks
-
-from .util import to_bytestr, is_integer, is_string, ip_str_to_addr
+import rtrlib.records as records
+
+from .util import (to_bytestr,
+                   is_integer,
+                   is_string,
+                   ip_str_to_addr,
+                   CallbackGenerator
+                   )
 from .exceptions import RTRInitError, PFXException, SyncTimeout
 
 
@@ -208,6 +214,76 @@ def validate(self, asn, prefix, mask_len):
 
         return PfxvState(result[0])
 
+    def for_each_ipv4_record(self, callback, data):
+        r"""
+        Iterate over all ipv4 records of the pfx table.
+
+        callback must take two arguments, the pfx_record and the data object.
+
+        For a more pythonic alternative see :py:meth:`ipv4_records`
+
+        :param callable callback: called for every record in the pfx table
+        :param object data: arbitrary data object \
+            that is passed to the callback function
+        """
+        data_handle = ffi.new_handle((callback, data))
+
+        lib.rtr_mgr_for_each_ipv4_record(
+            self.rtr_manager_config,
+            lib.pfx_table_callback,
+            data_handle
+            )
+
+    def ipv4_records(self):
+        r"""
+        Return iterator over all ipv4 records in the pfx table.
+
+        This iterator utilises threads to execute retrieve the records. \
+        If that is a problem for you take a look at \
+        :py:meth:`for_each_ipv4_record`.
+        """
+        def callback(record, data):
+            LOG.debug('Putting "%s" in queue', record)
+            data.put_nowait(records.copy_pfx_record(record))
+
+        generator = CallbackGenerator(self.for_each_ipv4_record, callback)
+        return generator
+
+    def for_each_ipv6_record(self, callback, data):
+        r"""
+        Iterate over all ipv6 records of the pfx table.
+
+        callback must take two arguments, the pfx_record and the data object.
+
+        For a more pythonic alternative see :py:meth:`ipv6_records`
+
+        :param callable callback: called for every record in the pfx table
+        :param object data: arbitrary data object \
+            that is passed to the callback function
+        """
+        data_handle = ffi.new_handle((callback, data))
+
+        lib.rtr_mgr_for_each_ipv6_record(
+            self.rtr_manager_config,
+            lib.pfx_table_callback,
+            data_handle
+            )
+
+    def ipv6_records(self):
+        r"""
+        Return iterator over all ipv6 records in the pfx table.
+
+        This iterator utilises threads to execute retrieve the records. \
+        If that is a problem for you take a look at \
+        :py:meth:`for_each_ipv6_record`.
+        """
+        def callback(record, data):
+            LOG.debug('Putting "%s" in queue', record)
+            data.put_nowait(records.copy_pfx_record(record))
+
+        generator = CallbackGenerator(self.for_each_ipv6_record, callback)
+        return generator
+
 
 class PfxvState(Enum):
     """
diff --git a/rtrlib/util.py b/rtrlib/util.py
@@ -11,6 +11,9 @@
 
 import logging
 import six
+import threading
+
+from six.moves import queue
 
 from _rtrlib import ffi, lib
 from .exceptions import IpConversionException
@@ -76,3 +79,68 @@ def is_string(var):
     Checks if var is a string
     """
     return isinstance(var, six.string_types)
+
+
+class StoppableThread(threading.Thread):
+    """
+    Thread class with a stop() method.
+
+    The thread itself has to check regularly for the stopped() condition.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(StoppableThread, self).__init__(*args, **kwargs)
+        self._stop_event = threading.Event()
+
+    def stop(self):
+        self._stop_event.set()
+
+    def stopped(self):
+        return self._stop_event.isSet()
+
+
+class CallbackGenerator(object):
+
+    def __init__(self, function, callback, args=()):
+        def inner_callback(pfx_record, data):
+            while not data.thread.stopped():
+                try:
+                    data.callback(pfx_record, data.queue)
+                    break
+                except queue.Full:
+                    pass
+            return
+
+        self.callback = callback
+        self.queue = queue.Queue()
+        new_args = list(args)
+        new_args.extend((inner_callback, self))
+        self.thread = StoppableThread(
+                                      target=function,
+                                      args=new_args,
+                                      daemon=True,
+                                      )
+
+        self.thread.start()
+
+    def __del__(self,):
+        self.thread.stop()
+
+    def __iter__(self,):
+        return self
+
+    def __next__(self,):
+        while(True):
+            try:
+                LOG.debug('About to take item out of queue')
+                item = self.queue.get_nowait()
+                LOG.debug('Took "%s" out of the queue', item)
+                self.queue.task_done()
+                return item
+            except queue.Empty:
+                if self.thread.is_alive() or self.queue.qsize() > 0:
+                    LOG.debug('Queue not yet filled, looping')
+                    continue
+                else:
+                    LOG.debug('Queue empty stopping iteration')
+                    raise StopIteration()
