Provide alternative in-memory queue for publishing events

baseplate/lib/events.py

@@ -22,7 +22,8 @@
 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 RemoteMessageQueue
+from baseplate.lib.message_queue import PosixMessageQueue
 from baseplate.lib.message_queue import TimedOutError
 
 
@@ -93,10 +94,15 @@ 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 = RemoteMessageQueue("/events-" + name, max_messages=MAX_QUEUE_SIZE)
+        else:
+            self.queue = PosixMessageQueue(  # type: ignore
+                "/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,12 +1,19 @@
 """A Gevent-friendly POSIX message queue."""
+import abc
+import queue as q
 import select
 
 from typing import Optional
+from baseplate.thrift import RemoteMessageQueueService
 
 import posix_ipc
 
 from baseplate.lib.retry import RetryPolicy
 
+from thrift.Thrift import TBinaryProtocol
+from thrift.Thrift import TSocket
+from thrift.transport import TTransport
+
 
 class MessageQueueError(Exception):
     """Base exception for message queue related errors."""
@@ -35,7 +42,51 @@ def __init__(self, inner: Exception):
         super().__init__(f"{inner} (check `ulimit -q`?)")
 
 
-class MessageQueue:
+class MessageQueue(abc.ABC):
+    """Abstract class for an inter-process message queue."""
+
+    @abc.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.
+
+        """
+
+    @abc.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.
+
+        """
+
+    @abc.abstractmethod
+    def unlink(self) -> None:
+        """Remove the queue from the system.
+
+        The queue will not leave until the last active user closes it.
+
+        """
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """Close the queue, freeing related resources.
+
+        This must be called explicitly if queues are created/destroyed on the
+        fly. It is not automatically called when the object is reclaimed by
+        Python.
+
+        """
+
+
+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
@@ -62,16 +113,9 @@ 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.
-
-        :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.
-
-        """
         for time_remaining in RetryPolicy.new(budget=timeout):
             try:
                 message, _ = self.queue.receive()
@@ -84,14 +128,6 @@ def get(self, timeout: Optional[float] = None) -> bytes:
         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.
-
-        """
         for time_remaining in RetryPolicy.new(budget=timeout):
             try:
                 return self.queue.send(message=message)
@@ -103,22 +139,64 @@ def put(self, message: bytes, timeout: Optional[float] = None) -> None:
         raise TimedOutError
 
     def unlink(self) -> None:
-        """Remove the queue from the system.
+        self.queue.unlink()
 
-        The queue will not leave until the last active user closes it.
+    def close(self) -> None:
+        self.queue.close()
 
-        """
-        self.queue.unlink()
+
+class InMemoryMessageQueue(MessageQueue):
+    """An in-memory inter process message queue.""" # Used by the sidecar
+
+    def __init__(self, name: str, max_messages: int):
+        self.queue: q.Queue = q.Queue(max_messages)
+        self.max_messages = max_messages
+        self.name = name
+
+    def get(self, timeout: Optional[float] = None) -> bytes:
+        try:
+            message = self.queue.get(timeout=timeout)
+            self.queue.task_done()
+            return message
+        except q.Empty:
+            raise TimedOutError
+
+    def put(self, message: bytes, timeout: Optional[float] = None) -> None:
+        try:
+            self.queue.put(message, timeout=timeout)
+        except q.Full:
+            raise TimedOutError
+
+    def unlink(self) -> None:
+        """Not implemented for in-memory queue"""
 
     def close(self) -> None:
-        """Close the queue, freeing related resources.
+        """Not implemented for in-memory queue"""
 
-        This must be called explicitly if queues are created/destroyed on the
-        fly. It is not automatically called when the object is reclaimed by
-        Python.
 
-        """
-        self.queue.close()
+class RemoteMessageQueue(MessageQueue): 
+    def __init__(self, max_messages: int): # Connect to the sidecar and instantiate a remote queue
+        transport = TSocket.TSocket( "localhost" , 9090) # todo: how do I connect to the sidecar? I dont think this is right
+        transport = TTransport.TBufferedTransport(transport)
+        protocol = TBinaryProtocol.TBinaryProtocol(transport)
+
+        self.client = RemoteMessageQueueService.Client(protocol)
+        # todo: create queue?
+
+        # Connect to server
+        self.transport.open()
+
+    def get(self, timeout: Optional[float] = None) -> bytes: 
+        return self.client.get(timeout)
+
+    def put(self, message: bytes, timeout: Optional[float] = None) -> None: 
+        return self.client.put(message, timeout)
+
+    def unlink(self) -> None:
+        """Not implemented for remote queue"""
+
+    def close(self) -> None: 
+        self.transport.close()
 
 
 def queue_tool() -> None:
@@ -140,6 +218,11 @@ 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 +249,12 @@ 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 = RemoteMessageQueue(args.queue_name, args.max_messages)
+    else:
+        queue = PosixMessageQueue(  # type: ignore
+            args.queue_name, args.max_messages, args.max_message_size
+        )
 
     if args.mode == "read":
         while True:

baseplate/observers/tracing.py

@@ -29,7 +29,8 @@
 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 RemoteMessageQueue
+from baseplate.lib.message_queue import PosixMessageQueue
 from baseplate.lib.message_queue import TimedOutError
 from baseplate.observers.timeout import ServerTimeout
 
@@ -539,18 +540,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 = RemoteMessageQueue(
+                "/traces-" + queue_name,
+                max_messages=MAX_QUEUE_SIZE,
+            )
+        else:
+            self.queue = PosixMessageQueue(  # type: ignore
+                "/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,17 +568,15 @@ 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),
             )
             return
         try:
             self.queue.put(serialized_str, timeout=0)
         except TimedOutError:
-            logger.warning(
-                "Trace queue %s is full. Is trace sidecar healthy?", self.queue.queue.name
-            )
+            logger.warning("Trace queue %s is full. Is trace sidecar healthy?", self.queue.name)
 
 
 def tracing_client_from_config(

baseplate/sidecars/event_publisher.py

@@ -12,21 +12,28 @@
 
 import requests
 
-from baseplate import __version__ as baseplate_version
+from baseplate import RequestContext, __version__ as baseplate_version
 from baseplate.lib import config
 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
+from baseplate.lib.message_queue import PosixMessageQueue
 from baseplate.lib.message_queue import TimedOutError
 from baseplate.lib.metrics import metrics_client_from_config
 from baseplate.lib.retry import RetryPolicy
 from baseplate.server import EnvironmentInterpolation
+from baseplate.thrift import RemoteMessageQueueService
 from baseplate.sidecars import Batch
 from baseplate.sidecars import BatchFull
 from baseplate.sidecars import SerializedBatch
 from baseplate.sidecars import TimeLimitedBatch
 
+from thrift.Thrift import TServer
+from thrift.Thrift import TBinaryProtocol
+from thrift.Thrift import TSocket
+from thrift.transport import TTransport
+
 
 logger = logging.getLogger(__name__)
 
@@ -161,6 +168,21 @@ def publish(self, payload: SerializedBatch) -> None:
         raise MaxRetriesError("could not sent batch")
 
 
+
+class RemoteMessageQueueHandler: # From the sidecar, create the queue and define get/put using the InMemoryQueue implementation
+    def is_healthy(self, context: RequestContext) -> bool:
+        return True
+
+    def __init__(self, name: str, max_messages: int):
+        self.queue = InMemoryMessageQueue(name, max_messages)
+
+    def get(self, timeout: Optional[float] = None) -> bytes: 
+        return self.queue.get(timeout)
+
+    def put(self, message: bytes, timeout: Optional[float] = None) -> None:
+        self.queue.put(message, timeout)
+
+
 SERIALIZER_BY_VERSION = {"2": V2Batch, "2j": V2JBatch}
 
 
@@ -174,6 +196,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 +230,26 @@ 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:
+        # start a thrift server that will be used to communicate with the main baseplate app
+        # this should not happen here, where are sidecars running?
+        processor = RemoteMessageQueueService.processor(RemoteMessageQueueHandler())
+        transport = TSocket.TServerSocket(host='127.0.0.1', port=9090)
+        tfactory = TTransport.TBufferedTransportFactory()
+        pfactory = TBinaryProtocol.TBinaryProtocolFactory()
+
+        server = TServer.TSimpleServer(processor, transport, tfactory, pfactory)
+        print('Starting the Message Queue server...')
+        server.serve()
+        print('done.')
+
+
+    else:
+        event_queue = PosixMessageQueue(  # type: ignore
+            "/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]()