Added in rate limited KytosEventBuffers

CHANGELOG.rst

@@ -6,6 +6,12 @@ All notable changes to the kytos project will be documented in this file.
 UNRELEASED - Under development
 ******************************
 
+Added
+=====
+
+- Added ``RateLimitedBuffer``, an alternative to ``KytosEventBuffer``
+- New configuration parameters for ``KytosEventBuffers``, accessible through the ``event_buffer_conf`` parameter.
+
 [2023.1.0] - 2023-06-05
 ***********************
 

kytos/core/buffers/__init__.py

@@ -0,0 +1,5 @@
+"""Module for buffers."""
+from .buffers import KytosEventBuffer
+from .manager import KytosBuffers
+
+__all__ = ('KytosEventBuffer', 'KytosBuffers')

kytos/core/buffers.py → kytos/core/buffers/buffers.py

@@ -1,12 +1,13 @@
 """Kytos Buffer Classes, based on Python Queue."""
+import asyncio
 import logging
+import time
+from typing import Callable, Hashable, Iterable
 
-from janus import PriorityQueue, Queue
+import limits
+from janus import Queue
 
 from kytos.core.events import KytosEvent
-from kytos.core.helpers import get_thread_pool_max_workers
-
-__all__ = ('KytosBuffers', )
 
 LOG = logging.getLogger(__name__)
 
@@ -123,53 +124,38 @@ def full(self):
         return self._queue.sync_q.full()
 
 
-class KytosBuffers:
-    """Set of KytosEventBuffer used in Kytos."""
-
-    def __init__(self):
-        """Build four KytosEventBuffers.
-
-        :attr:`conn`: :class:`~kytos.core.buffers.KytosEventBuffer` with events
-        received from connection events.
-
-        :attr:`raw`: :class:`~kytos.core.buffers.KytosEventBuffer` with events
-        received from network.
+class RateLimitedBuffer(KytosEventBuffer):
+    """
+    Extension of KytosEventBuffer with ratelimiting capabilities.
+    """
+    def __init__(
+        self, *args,
+        strategy: limits.strategies.RateLimiter,
+        limit: limits.RateLimitItem,
+        gen_identifiers: Callable[[KytosEvent], Iterable[Hashable]],
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.strategy = strategy
+        self.limit = limit
+        self.gen_identifiers = gen_identifiers
 
-        :attr:`msg_in`: :class:`~kytos.core.buffers.KytosEventBuffer` with
-        events to be received.
-
-        :attr:`msg_out`: :class:`~kytos.core.buffers.KytosEventBuffer` with
-        events to be sent.
+    def get(self):
+        val = super().get()
+        identifiers = self.limit, *self.gen_identifiers(val)
+        while not self.strategy.hit(*identifiers):
+            window_reset, _ = self.strategy.get_window_stats(*identifiers)
+            sleep_time = window_reset - time.time()
+            if sleep_time > 0:
+                time.sleep(sleep_time)
+        return val
 
-        :attr:`app`: :class:`~kytos.core.buffers.KytosEventBuffer` with events
-        sent to NApps.
-        """
-        self._pool_max_workers = get_thread_pool_max_workers()
-        self.conn = KytosEventBuffer("conn")
-        self.raw = KytosEventBuffer("raw", maxsize=self._get_maxsize("sb"))
-        self.msg_in = KytosEventBuffer("msg_in",
-                                       maxsize=self._get_maxsize("sb"),
-                                       queue_cls=PriorityQueue)
-        self.msg_out = KytosEventBuffer("msg_out",
-                                        maxsize=self._get_maxsize("sb"),
-                                        queue_cls=PriorityQueue)
-        self.app = KytosEventBuffer("app", maxsize=self._get_maxsize("app"))
-
-    def get_all_buffers(self):
-        """Get all KytosEventBuffer instances."""
-        return [
-            event_buffer for event_buffer in self.__dict__.values()
-            if isinstance(event_buffer, KytosEventBuffer)
-        ]
-
-    def _get_maxsize(self, queue_name):
-        """Get queue maxsize if it's been set."""
-        return self._pool_max_workers.get(queue_name, 0)
-
-    def send_stop_signal(self):
-        """Send a ``kytos/core.shutdown`` event to each buffer."""
-        LOG.info('Stop signal received by Kytos buffers.')
-        LOG.info('Sending KytosShutdownEvent to all apps.')
-        event = KytosEvent(name='kytos/core.shutdown')
-        for buffer in self.get_all_buffers():
-            buffer.put(event)
+    async def aget(self):
+        val = await super().aget()
+        identifiers = self.limit, *self.gen_identifiers(val)
+        while not self.strategy.hit(*identifiers):
+            window_reset, _ = self.strategy.get_window_stats(*identifiers)
+            sleep_time = window_reset - time.time()
+            if sleep_time > 0:
+                await asyncio.sleep(sleep_time)
+        return val

kytos/core/buffers/factory.py

@@ -0,0 +1,86 @@
+"""Utilities for composing KytosEventBuffers"""
+from functools import reduce
+
+import limits
+import limits.storage as lstorage
+import limits.strategies as lstrategies
+from janus import PriorityQueue, Queue
+
+from .buffers import KytosEventBuffer, RateLimitedBuffer
+
+
+def process_default(config):
+    """
+    Create a default KytosEventBuffer from a given config dict
+    """
+    queue_classes = {
+        'default': Queue,
+        'queue': Queue,
+        'priority': PriorityQueue,
+    }
+    return {
+        'buffer_cls': KytosEventBuffer,
+        'buffer_args': {
+            'queue_cls': queue_classes[config.get('queue', 'default')]
+        },
+    }
+
+
+def process_storage(config):
+    """
+    Create a rate limit storage from a given config dict
+    """
+    storages = {
+        'memory': lstorage.MemoryStorage,
+    }
+    return storages[config.get('type', 'memory')]()
+
+
+def process_strategy(config):
+    """
+    Create a rate limit strategy from a given config dict
+    """
+    strategies = {
+        'moving_window': lstrategies.MovingWindowRateLimiter,
+        'fixed_window': lstrategies.FixedWindowRateLimiter,
+        'elastic_window': lstrategies.FixedWindowElasticExpiryRateLimiter,
+    }
+    strategy_cls = strategies[config.get('type', 'moving_window')]
+    return strategy_cls(
+        process_storage(
+            config.get('storage', {})
+        )
+    )
+
+
+def process_rate_limited(config):
+    """
+    Create a rate limited KytosEventBuffer from a given config dict
+    """
+    processed = process_default(config)
+    processed['buffer_cls'] = RateLimitedBuffer
+    args = processed['buffer_args']
+    args['strategy'] = process_strategy(config.get('strategy', {}))
+    args['limit'] = limits.parse(config.get('limit', '100/second'))
+    args['gen_identifiers'] = lambda event: reduce(
+        lambda ev, attr: getattr(ev, attr, ('unknown',)),
+        config.get('identifier', 'connection.id').split('.'),
+        event
+    )
+    return processed
+
+
+def buffer_from_config(name, config: dict) -> KytosEventBuffer:
+    """
+    Create a KytosEventBuffer from a given config dict
+    """
+    buffer_conf_processors = {
+        'default': process_default,
+        'rate_limited': process_rate_limited,
+    }
+    buffer_type = config.get('type', 'default')
+    processed_conf = buffer_conf_processors[buffer_type](config)
+    return processed_conf['buffer_cls'](
+        name,
+        **(processed_conf['buffer_args'])
+    )

kytos/core/buffers/manager.py

@@ -0,0 +1,79 @@
+"""Kytos Buffer Classes, based on Python Queue."""
+import logging
+
+from janus import PriorityQueue
+
+from kytos.core.config import KytosConfig
+from kytos.core.events import KytosEvent
+from kytos.core.helpers import get_thread_pool_max_workers
+
+from .buffers import KytosEventBuffer
+from .factory import buffer_from_config
+
+LOG = logging.getLogger(__name__)
+
+
+class KytosBuffers:
+    """Set of KytosEventBuffer used in Kytos."""
+
+    def __init__(self):
+        """Build four KytosEventBuffers.
+
+        :attr:`conn`: :class:`~kytos.core.buffers.KytosEventBuffer` with events
+        received from connection events.
+
+        :attr:`raw`: :class:`~kytos.core.buffers.KytosEventBuffer` with events
+        received from network.
+
+        :attr:`msg_in`: :class:`~kytos.core.buffers.KytosEventBuffer` with
+        events to be received.
+
+        :attr:`msg_out`: :class:`~kytos.core.buffers.KytosEventBuffer` with
+        events to be sent.
+
+        :attr:`app`: :class:`~kytos.core.buffers.KytosEventBuffer` with events
+        sent to NApps.
+        """
+
+        self._pool_max_workers = get_thread_pool_max_workers()
+
+        self.conn = KytosEventBuffer("conn")
+        self.raw = KytosEventBuffer("raw", maxsize=self._get_maxsize("sb"))
+        self.msg_in = KytosEventBuffer(
+            "msg_in",
+            maxsize=self._get_maxsize("sb"),
+            queue_cls=PriorityQueue,
+        )
+        self.msg_out = KytosEventBuffer(
+            "msg_out",
+            maxsize=self._get_maxsize("sb"),
+            queue_cls=PriorityQueue,
+        )
+        self.app = KytosEventBuffer(
+            "app",
+            maxsize=self._get_maxsize("app")
+        )
+
+        buffer_conf = KytosConfig().options['daemon'].event_buffer_conf
+
+        for name, config in buffer_conf.items():
+            setattr(self, name, buffer_from_config(name, config))
+
+    def get_all_buffers(self):
+        """Get all KytosEventBuffer instances."""
+        return [
+            event_buffer for event_buffer in self.__dict__.values()
+            if isinstance(event_buffer, KytosEventBuffer)
+        ]
+
+    def _get_maxsize(self, queue_name):
+        """Get queue maxsize if it's been set."""
+        return self._pool_max_workers.get(queue_name, 0)
+
+    def send_stop_signal(self):
+        """Send a ``kytos/core.shutdown`` event to each buffer."""
+        LOG.info('Stop signal received by Kytos buffers.')
+        LOG.info('Sending KytosShutdownEvent to all apps.')
+        event = KytosEvent(name='kytos/core.shutdown')
+        for buffer in self.get_all_buffers():
+            buffer.put(event)