diff options
Diffstat (limited to 'testing/mozbase/mozlog/mozlog/handlers/bufferhandler.py')
-rw-r--r-- | testing/mozbase/mozlog/mozlog/handlers/bufferhandler.py | 86 |
1 files changed, 86 insertions, 0 deletions
diff --git a/testing/mozbase/mozlog/mozlog/handlers/bufferhandler.py b/testing/mozbase/mozlog/mozlog/handlers/bufferhandler.py new file mode 100644 index 0000000000..5fcdcf967a --- /dev/null +++ b/testing/mozbase/mozlog/mozlog/handlers/bufferhandler.py @@ -0,0 +1,86 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +from .base import BaseHandler + + +class BufferHandler(BaseHandler): + """Handler that maintains a circular buffer of messages based on the + size and actions specified by a user. + + :param inner: The underlying handler used to emit messages. + :param message_limit: The maximum number of messages to retain for + context. If None, the buffer will grow without limit. + :param buffered_actions: The set of actions to include in the buffer + rather than log directly. + """ + + def __init__(self, inner, message_limit=100, buffered_actions=None): + BaseHandler.__init__(self, inner) + self.inner = inner + self.message_limit = message_limit + if buffered_actions is None: + buffered_actions = ["log", "test_status"] + self.buffered_actions = set(buffered_actions) + self._buffering = True + + if self.message_limit is not None: + self._buffer = [None] * self.message_limit + self._buffer_pos = 0 + else: + self._buffer = [] + + self.message_handler.register_message_handlers( + "buffer", + { + "on": self._enable_buffering, + "off": self._disable_buffering, + "flush": self._flush_buffered, + "clear": self._clear_buffer, + }, + ) + + def __call__(self, data): + action = data["action"] + if "bypass_mozlog_buffer" in data: + data.pop("bypass_mozlog_buffer") + self.inner(data) + return + if not self._buffering or action not in self.buffered_actions: + self.inner(data) + return + + self._add_message(data) + + def _add_message(self, data): + if self.message_limit is None: + self._buffer.append(data) + else: + self._buffer[self._buffer_pos] = data + self._buffer_pos = (self._buffer_pos + 1) % self.message_limit + + def _enable_buffering(self): + self._buffering = True + + def _disable_buffering(self): + self._buffering = False + + def _clear_buffer(self): + """Clear the buffer of unwanted messages.""" + current_size = len([m for m in self._buffer if m is not None]) + if self.message_limit is not None: + self._buffer = [None] * self.message_limit + else: + self._buffer = [] + return current_size + + def _flush_buffered(self): + """Logs the contents of the current buffer""" + for msg in self._buffer[self._buffer_pos :]: + if msg is not None: + self.inner(msg) + for msg in self._buffer[: self._buffer_pos]: + if msg is not None: + self.inner(msg) + return self._clear_buffer() |