Provide alternative in-memory queue for publishing events

baseplate/lib/events.py

@@ -22,7 +22,7 @@
 from baseplate import Span
 from baseplate.clients import ContextFactory
 from baseplate.lib import config
-from baseplate.lib.message_queue import MessageQueue
+from baseplate.lib.message_queue import InMemoryMessageQueue, PosixMessageQueue
 from baseplate.lib.message_queue import TimedOutError
 
 
@@ -93,10 +93,13 @@ class EventQueue(ContextFactory, config.Parser, Generic[T]):
 
     """
 
-    def __init__(self, name: str, event_serializer: Callable[[T], bytes]):
-        self.queue = MessageQueue(
-            "/events-" + name, max_messages=MAX_QUEUE_SIZE, max_message_size=MAX_EVENT_SIZE
-        )
+    def __init__(self, name: str, event_serializer: Callable[[T], bytes], use_in_memory_queue: bool = False):
+        if use_in_memory_queue: 
+            self.queue = InMemoryMessageQueue("/events-" + name, max_messages=MAX_QUEUE_SIZE)
+        else:
+            self.queue = PosixMessageQueue(
+                "/events-" + name, max_messages=MAX_QUEUE_SIZE, max_message_size=MAX_EVENT_SIZE
+            )
         self.serialize_event = event_serializer
 
     def put(self, event: T) -> None:

baseplate/lib/message_queue.py

@@ -1,4 +1,6 @@
 """A Gevent-friendly POSIX message queue."""
+from abc import abstractmethod
+import queue
 import select
 
 from typing import Optional
@@ -34,9 +36,37 @@ class MessageQueueOSError(OSError):
     def __init__(self, inner: Exception):
         super().__init__(f"{inner} (check `ulimit -q`?)")
 
-
 class MessageQueue:
-    """A Gevent-friendly (but not required) inter process message queue.
+    @abstractmethod
+    def __init__(self, name: str, max_messages: int, max_message_size: int):
+        pass
+
+    @abstractmethod
+    def get(self, timeout: Optional[float] = None) -> bytes:
+        """Read a message from the queue.
+
+        :param timeout: If the queue is empty, the call will block up to
+            ``timeout`` seconds or forever if ``None``.
+        :raises: :py:exc:`TimedOutError` The queue was empty for the allowed
+            duration of the call.
+
+        """
+        pass
+
+    @abstractmethod
+    def put(self, message: bytes, timeout: Optional[float] = None) -> None:
+        """Add a message to the queue.
+
+        :param timeout: If the queue is full, the call will block up to
+            ``timeout`` seconds or forever if ``None``.
+        :raises: :py:exc:`TimedOutError` The queue was full for the allowed
+            duration of the call.
+
+        """
+        pass
+
+class InMemoryMessageQueue(MessageQueue):
+    """An in-memory inter process message queue.
 
     ``name`` should be a string of up to 255 characters consisting of an
     initial slash, followed by one or more characters, none of which are
@@ -47,7 +77,61 @@ class MessageQueue:
     support this.
 
     """
+    def __init__(self, name: str, max_messages: int):
+        self.queue = queue.Queue(max_messages)
+        self.max_messages = max_messages
+        self.name = name
+
+    def get(self, timeout: Optional[float] = None) -> bytes:
+        """Read a message from the queue.
 
+        :param timeout: If the queue is empty, the call will block up to
+            ``timeout`` seconds or forever if ``None``.
+        :raises: :py:exc:`TimedOutError` The queue was empty for the allowed
+            duration of the call.
+
+        """
+        try: 
+            message = self.queue.get(timeout=timeout)
+            self.queue.task_done()
+            return message
+        except queue.Empty:
+            raise TimedOutError
+
+    def put(self, message: bytes, timeout: Optional[float] = None) -> None:
+        """Add a message to the queue.
+
+        :param timeout: If the queue is full, the call will block up to
+            ``timeout`` seconds or forever if ``None``.
+        :raises: :py:exc:`TimedOutError` The queue was full for the allowed
+            duration of the call.
+
+        """
+        try: 
+            return self.queue.put(message, timeout=timeout)
+        except queue.Full:
+            raise TimedOutError
+
+    def unlink(self) -> None:
+        """Not implemented for in-memory queue"""
+        pass
+
+    def close(self) -> None:
+        """Not implemented for in-memory queue"""
+        pass
+
+class PosixMessageQueue(MessageQueue):
+    """A Gevent-friendly (but not required) inter process message queue.
+
+    ``name`` should be a string of up to 255 characters consisting of an
+    initial slash, followed by one or more characters, none of which are
+    slashes.
+
+    Note: This relies on POSIX message queues being available and
+    select(2)-able like other file descriptors. Not all operating systems
+    support this.
+
+    """
     def __init__(self, name: str, max_messages: int, max_message_size: int):
         try:
             self.queue = posix_ipc.MessageQueue(
@@ -62,6 +146,7 @@ def __init__(self, name: str, max_messages: int, max_message_size: int):
         except OSError as exc:
             raise MessageQueueOSError(exc)
         self.queue.block = False
+        self.name = name
 
     def get(self, timeout: Optional[float] = None) -> bytes:
         """Read a message from the queue.
@@ -101,7 +186,7 @@ def put(self, message: bytes, timeout: Optional[float] = None) -> None:
                 select.select([], [self.queue.mqd], [], time_remaining)
 
         raise TimedOutError
-
+    
     def unlink(self) -> None:
         """Remove the queue from the system.
 
@@ -119,6 +204,7 @@ def close(self) -> None:
 
         """
         self.queue.close()
+
 
 
 def queue_tool() -> None:
@@ -140,6 +226,7 @@ def queue_tool() -> None:
         help="if creating the queue, what to set the maximum message size to",
     )
     parser.add_argument("queue_name", help="the name of the queue to consume")
+    parser.add_argument("use_in_memory_queue", default=False, help="whether to use an in-memory queue or a posix queue")
 
     group = parser.add_mutually_exclusive_group(required=True)
     group.add_argument(
@@ -166,7 +253,10 @@ def queue_tool() -> None:
 
     args = parser.parse_args()
 
-    queue = MessageQueue(args.queue_name, args.max_messages, args.max_message_size)
+    if args.use_in_memory_queue:
+        queue = InMemoryMessageQueue(args.queue_name, args.max_messages)
+    else: 
+        queue = PosixMessageQueue(args.queue_name, args.max_messages, args.max_message_size)
 
     if args.mode == "read":
         while True:

baseplate/observers/tracing.py

@@ -29,7 +29,7 @@
 from baseplate import SpanObserver
 from baseplate.lib import config
 from baseplate.lib import warn_deprecated
-from baseplate.lib.message_queue import MessageQueue
+from baseplate.lib.message_queue import InMemoryMessageQueue, PosixMessageQueue
 from baseplate.lib.message_queue import TimedOutError
 from baseplate.observers.timeout import ServerTimeout
 
@@ -539,18 +539,24 @@ class TraceQueueFullError(Exception):
 
 
 class SidecarRecorder(Recorder):
-    """Interface for recording spans to a POSIX message queue.
+    """Interface for recording spans to a message queue.
 
     The SidecarRecorder serializes spans to a string representation before
     adding them to the queue.
     """
 
-    def __init__(self, queue_name: str):
-        self.queue = MessageQueue(
-            "/traces-" + queue_name,
-            max_messages=MAX_QUEUE_SIZE,
-            max_message_size=MAX_SPAN_SIZE,
-        )
+    def __init__(self, queue_name: str, use_in_memory_queue: bool = False):
+        if use_in_memory_queue:
+            self.queue = InMemoryMessageQueue(
+                "/traces-" + queue_name,
+                max_messages=MAX_QUEUE_SIZE,
+            )
+        else:
+            self.queue = PosixMessageQueue(
+                "/traces-" + queue_name,
+                max_messages=MAX_QUEUE_SIZE,
+                max_message_size=MAX_SPAN_SIZE,
+            )
 
     def send(self, span: TraceSpanObserver) -> None:
         # Don't raise exceptions from here. This is called in the
@@ -561,7 +567,7 @@ def send(self, span: TraceSpanObserver) -> None:
                 "Trace too big. Traces published to %s are not allowed to be larger "
                 "than %d bytes. Received trace is %d bytes. This can be caused by "
                 "an excess amount of tags or a large amount of child spans.",
-                self.queue.queue.name,
+                self.queue.name,
                 MAX_SPAN_SIZE,
                 len(serialized_str),
             )
@@ -570,7 +576,7 @@ def send(self, span: TraceSpanObserver) -> None:
             self.queue.put(serialized_str, timeout=0)
         except TimedOutError:
             logger.warning(
-                "Trace queue %s is full. Is trace sidecar healthy?", self.queue.queue.name
+                "Trace queue %s is full. Is trace sidecar healthy?", self.queue.name
             )
 
 

baseplate/sidecars/event_publisher.py

@@ -17,7 +17,7 @@
 from baseplate.lib import metrics
 from baseplate.lib.events import MAX_EVENT_SIZE
 from baseplate.lib.events import MAX_QUEUE_SIZE
-from baseplate.lib.message_queue import MessageQueue
+from baseplate.lib.message_queue import InMemoryMessageQueue, PosixMessageQueue
 from baseplate.lib.message_queue import TimedOutError
 from baseplate.lib.metrics import metrics_client_from_config
 from baseplate.lib.retry import RetryPolicy
@@ -174,6 +174,11 @@ def publish_events() -> None:
         default="main",
         help="name of event queue / publisher config (default: main)",
     )
+    arg_parser.add_argument(
+        "--use-in-memory-queue",
+        default=False,
+        help="use an in memory queue instead of a posix queue",
+    )
     arg_parser.add_argument(
         "--debug", default=False, action="store_true", help="enable debug logging"
     )
@@ -203,11 +208,17 @@ def publish_events() -> None:
 
     metrics_client = metrics_client_from_config(raw_config)
 
-    event_queue = MessageQueue(
-        "/events-" + args.queue_name,
-        max_messages=cfg.max_queue_size,
-        max_message_size=MAX_EVENT_SIZE,
-    )
+    if args.use_in_memory_queue:
+        event_queue = InMemoryMessageQueue(
+            "/events-" + args.queue_name,
+            max_messages=cfg.max_queue_size,
+        )
+    else: 
+        event_queue = PosixMessageQueue(
+            "/events-" + args.queue_name,
+            max_messages=cfg.max_queue_size,
+            max_message_size=MAX_EVENT_SIZE,
+        )
 
     # pylint: disable=maybe-no-member
     serializer = SERIALIZER_BY_VERSION[cfg.collector.version]()

baseplate/sidecars/trace_publisher.py

@@ -10,7 +10,7 @@
 from baseplate import __version__ as baseplate_version
 from baseplate.lib import config
 from baseplate.lib import metrics
-from baseplate.lib.message_queue import MessageQueue
+from baseplate.lib.message_queue import InMemoryMessageQueue, PosixMessageQueue
 from baseplate.lib.message_queue import TimedOutError
 from baseplate.lib.metrics import metrics_client_from_config
 from baseplate.lib.retry import RetryPolicy
@@ -128,6 +128,11 @@ def publish_traces() -> None:
         default="main",
         help="name of trace queue / publisher config (default: main)",
     )
+    arg_parser.add_argument(
+        "--use-in-memory-queue",
+        default=False,
+        help="use in memory queue instead of a posix queue",
+    )
     arg_parser.add_argument(
         "--debug", default=False, action="store_true", help="enable debug logging"
     )
@@ -160,11 +165,17 @@ def publish_traces() -> None:
         },
     )
 
-    trace_queue = MessageQueue(
-        "/traces-" + args.queue_name,
-        max_messages=publisher_cfg.max_queue_size,
-        max_message_size=MAX_SPAN_SIZE,
-    )
+    if args.use_in_memory_queue:
+        trace_queue = InMemoryMessageQueue(
+            "/traces-" + args.queue_name,
+            max_messages=publisher_cfg.max_queue_size,
+        )
+    else:
+        trace_queue = PosixMessageQueue(
+            "/traces-" + args.queue_name,
+            max_messages=publisher_cfg.max_queue_size,
+            max_message_size=MAX_SPAN_SIZE,
+        )
 
     # pylint: disable=maybe-no-member
     inner_batch = TraceBatch(max_size=publisher_cfg.max_batch_size)

docker-compose.yml

@@ -14,6 +14,13 @@ services:
       - "redis"
       - "zookeeper"
       - "redis-cluster-node"
+  test-baseplate:
+    build:
+      context: "."
+      dockerfile: "Dockerfile"
+    command: "pytest -v tests/"
+    volumes:
+      - ".:/src"
   cassandra:
     image: "cassandra:3.11"
     environment:

docs/api/baseplate/lib/message_queue.rst

@@ -27,10 +27,10 @@ process pair (run the producer then the consumer):
 .. testcode::
 
     # producer.py
-    from baseplate.lib.message_queue import MessageQueue
+    from baseplate.lib.message_queue import PosixMessageQueue
 
     # If the queue doesn't already exist, we'll create it.
-    mq = MessageQueue(
+    mq = PosixMessageQueue(
         "/baseplate-testing", max_messages=1, max_message_size=1)
     message = "1"
     mq.put(message)
@@ -48,9 +48,9 @@ POSIX message queue. Next up, run the consumer:
 .. testcode::
 
     # consumer.py
-    from baseplate.lib.message_queue import MessageQueue
+    from baseplate.lib.message_queue import PosixMessageQueue
 
-    mq = MessageQueue(
+    mq = PosixMessageQueue(
         "/baseplate-testing", max_messages=1, max_message_size=1)
     # Unless a `timeout` kwarg is passed, this will block until
     # we can pop a message from the queue.
@@ -97,8 +97,8 @@ up seeing a vague ``ValueError`` exception. Here's an example:
 
 .. code-block:: pycon
 
-    >>> from baseplate.lib.message_queue import MessageQueue
-    >>> mq = MessageQueue(
+    >>> from baseplate.lib.message_queue import PosixMessageQueue
+    >>> mq = PosixMessageQueue(
             "/over-the-limit", max_messages=11, max_message_size=8096)
     Traceback (most recent call last):
       File "<input>", line 2, in <module>
@@ -138,7 +138,10 @@ See ``--help`` for more info.
 
 .. automodule:: baseplate.lib.message_queue
 
-.. autoclass:: MessageQueue
+.. autoclass:: PosixMessageQueue
+   :members:
+
+.. autoclass:: InMemoryMessageQueue
    :members: