streamlit-nightly 1.43.2.dev20250307__py3-none-any.whl

streamlit/runtime/forward_msg_cache.py

@@ -0,0 +1,296 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+from weakref import WeakKeyDictionary
+
+from streamlit import config, util
+from streamlit.logger import get_logger
+from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
+from streamlit.runtime.stats import CacheStat, CacheStatsProvider, group_stats
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+    from streamlit.runtime.app_session import AppSession
+
+_LOGGER: Final = get_logger(__name__)
+
+
+def populate_hash_if_needed(msg: ForwardMsg) -> str:
+    """Computes and assigns the unique hash for a ForwardMsg.
+
+    If the ForwardMsg already has a hash, this is a no-op.
+
+    Parameters
+    ----------
+    msg : ForwardMsg
+
+    Returns
+    -------
+    string
+        The message's hash, returned here for convenience. (The hash
+        will also be assigned to the ForwardMsg; callers do not need
+        to do this.)
+
+    """
+    if msg.hash == "":
+        # Move the message's metadata aside. It's not part of the
+        # hash calculation.
+        metadata = msg.metadata
+        msg.ClearField("metadata")
+
+        # MD5 is good enough for what we need, which is uniqueness.
+        msg.hash = util.calc_md5(msg.SerializeToString())
+
+        # Restore metadata.
+        msg.metadata.CopyFrom(metadata)
+
+    return msg.hash
+
+
+def create_reference_msg(msg: ForwardMsg) -> ForwardMsg:
+    """Create a ForwardMsg that refers to the given message via its hash.
+
+    The reference message will also get a copy of the source message's
+    metadata.
+
+    Parameters
+    ----------
+    msg : ForwardMsg
+        The ForwardMsg to create the reference to.
+
+    Returns
+    -------
+    ForwardMsg
+        A new ForwardMsg that "points" to the original message via the
+        ref_hash field.
+
+    """
+    ref_msg = ForwardMsg()
+    ref_msg.ref_hash = populate_hash_if_needed(msg)
+    ref_msg.metadata.CopyFrom(msg.metadata)
+    return ref_msg
+
+
+class ForwardMsgCache(CacheStatsProvider):
+    """A cache of ForwardMsgs.
+
+    Large ForwardMsgs (e.g. those containing big DataFrame payloads) are
+    stored in this cache. The server can choose to send a ForwardMsg's hash,
+    rather than the message itself, to a client. Clients can then
+    request messages from this cache via another endpoint.
+
+    This cache is *not* thread safe. It's intended to only be accessed by
+    the server thread.
+
+    """
+
+    class Entry:
+        """Cache entry.
+
+        Stores the cached message, and the set of AppSessions
+        that we've sent the cached message to.
+
+        """
+
+        def __init__(self, msg: ForwardMsg | None):
+            self.msg = msg
+            self._session_script_run_counts: MutableMapping[AppSession, int] = (
+                WeakKeyDictionary()
+            )
+
+        def __repr__(self) -> str:
+            return util.repr_(self)
+
+        def add_session_ref(self, session: AppSession, script_run_count: int) -> None:
+            """Adds a reference to a AppSession that has referenced
+            this Entry's message.
+
+            Parameters
+            ----------
+            session : AppSession
+            script_run_count : int
+                The session's run count at the time of the call
+
+            """
+            prev_run_count = self._session_script_run_counts.get(session, 0)
+            if script_run_count < prev_run_count:
+                _LOGGER.error(
+                    "New script_run_count (%s) is < prev_run_count (%s). "
+                    "This should never happen!",
+                    script_run_count,
+                    prev_run_count,
+                )
+                script_run_count = prev_run_count
+            self._session_script_run_counts[session] = script_run_count
+
+        def has_session_ref(self, session: AppSession) -> bool:
+            return session in self._session_script_run_counts
+
+        def get_session_ref_age(
+            self, session: AppSession, script_run_count: int
+        ) -> int:
+            """The age of the given session's reference to the Entry,
+            given a new script_run_count.
+
+            """
+            return script_run_count - self._session_script_run_counts[session]
+
+        def remove_session_ref(self, session: AppSession) -> None:
+            del self._session_script_run_counts[session]
+
+        def has_refs(self) -> bool:
+            """True if this Entry has references from any AppSession.
+
+            If not, it can be removed from the cache.
+            """
+            return len(self._session_script_run_counts) > 0
+
+    def __init__(self):
+        self._entries: dict[str, ForwardMsgCache.Entry] = {}
+
+    def __repr__(self) -> str:
+        return util.repr_(self)
+
+    def add_message(
+        self, msg: ForwardMsg, session: AppSession, script_run_count: int
+    ) -> None:
+        """Add a ForwardMsg to the cache.
+
+        The cache will also record a reference to the given AppSession,
+        so that it can track which sessions have already received
+        each given ForwardMsg.
+
+        Parameters
+        ----------
+        msg : ForwardMsg
+        session : AppSession
+        script_run_count : int
+            The number of times the session's script has run
+
+        """
+        populate_hash_if_needed(msg)
+        entry = self._entries.get(msg.hash, None)
+        if entry is None:
+            if config.get_option("global.storeCachedForwardMessagesInMemory"):
+                entry = ForwardMsgCache.Entry(msg)
+            else:
+                entry = ForwardMsgCache.Entry(None)
+            self._entries[msg.hash] = entry
+        entry.add_session_ref(session, script_run_count)
+
+    def get_message(self, hash: str) -> ForwardMsg | None:
+        """Return the message with the given ID if it exists in the cache.
+
+        Parameters
+        ----------
+        hash : str
+            The id of the message to retrieve.
+
+        Returns
+        -------
+        ForwardMsg | None
+
+        """
+        entry = self._entries.get(hash, None)
+        return entry.msg if entry else None
+
+    def has_message_reference(
+        self, msg: ForwardMsg, session: AppSession, script_run_count: int
+    ) -> bool:
+        """Return True if a session has a reference to a message."""
+        populate_hash_if_needed(msg)
+
+        entry = self._entries.get(msg.hash, None)
+        if entry is None or not entry.has_session_ref(session):
+            return False
+
+        # Ensure we're not expired
+        age = entry.get_session_ref_age(session, script_run_count)
+        return age <= int(config.get_option("global.maxCachedMessageAge"))
+
+    def remove_refs_for_session(self, session: AppSession) -> None:
+        """Remove refs for all entries for the given session.
+
+        This should be called when an AppSession is disconnected or closed.
+
+        Parameters
+        ----------
+        session : AppSession
+        """
+
+        # Operate on a copy of our entries dict.
+        # We may be deleting from it.
+        for msg_hash, entry in self._entries.copy().items():
+            if entry.has_session_ref(session):
+                entry.remove_session_ref(session)
+
+            if not entry.has_refs():
+                # The entry has no more references. Remove it from
+                # the cache completely.
+                del self._entries[msg_hash]
+
+    def remove_expired_entries_for_session(
+        self, session: AppSession, script_run_count: int
+    ) -> None:
+        """Remove any cached messages that have expired from the given session.
+
+        This should be called each time a AppSession finishes executing.
+
+        Parameters
+        ----------
+        session : AppSession
+        script_run_count : int
+            The number of times the session's script has run
+
+        """
+        max_age = config.get_option("global.maxCachedMessageAge")
+
+        # Operate on a copy of our entries dict.
+        # We may be deleting from it.
+        for msg_hash, entry in self._entries.copy().items():
+            if not entry.has_session_ref(session):
+                continue
+
+            age = entry.get_session_ref_age(session, script_run_count)
+            if age > max_age:
+                _LOGGER.debug(
+                    "Removing expired entry [session=%s, hash=%s, age=%s]",
+                    id(session),
+                    msg_hash,
+                    age,
+                )
+                entry.remove_session_ref(session)
+                if not entry.has_refs():
+                    # The entry has no more references. Remove it from
+                    # the cache completely.
+                    del self._entries[msg_hash]
+
+    def clear(self) -> None:
+        """Remove all entries from the cache"""
+        self._entries.clear()
+
+    def get_stats(self) -> list[CacheStat]:
+        stats: list[CacheStat] = [
+            CacheStat(
+                category_name="ForwardMessageCache",
+                cache_name="",
+                byte_length=entry.msg.ByteSize() if entry.msg is not None else 0,
+            )
+            for _, entry in self._entries.items()
+        ]
+        return group_stats(stats)

streamlit/runtime/forward_msg_queue.py

@@ -0,0 +1,240 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable
+
+from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
+
+if TYPE_CHECKING:
+    from streamlit.proto.Delta_pb2 import Delta
+
+
+class ForwardMsgQueue:
+    """Accumulates a session's outgoing ForwardMsgs.
+
+    Each AppSession adds messages to its queue, and the Server periodically
+    flushes all session queues and delivers their messages to the appropriate
+    clients.
+
+    ForwardMsgQueue is not thread-safe - a queue should only be used from
+    a single thread.
+    """
+
+    _before_enqueue_msg: Callable[[ForwardMsg], None] | None = None
+
+    @staticmethod
+    def on_before_enqueue_msg(
+        before_enqueue_msg: Callable[[ForwardMsg], None] | None,
+    ) -> None:
+        """Set a callback to be called before a message is enqueued.
+        Used in static streamlit app generation."""
+        ForwardMsgQueue._before_enqueue_msg = before_enqueue_msg
+
+    def __init__(self):
+        self._queue: list[ForwardMsg] = []
+        # A mapping of (delta_path -> _queue.indexof(msg)) for each
+        # Delta message in the queue. We use this for coalescing
+        # redundant outgoing Deltas (where a newer Delta supersedes
+        # an older Delta, with the same delta_path, that's still in the
+        # queue).
+        self._delta_index_map: dict[tuple[int, ...], int] = {}
+
+    def get_debug(self) -> dict[str, Any]:
+        from google.protobuf.json_format import MessageToDict
+
+        return {
+            "queue": [MessageToDict(m) for m in self._queue],
+            "ids": list(self._delta_index_map.keys()),
+        }
+
+    def is_empty(self) -> bool:
+        return len(self._queue) == 0
+
+    def enqueue(self, msg: ForwardMsg) -> None:
+        """Add message into queue, possibly composing it with another message."""
+
+        if ForwardMsgQueue._before_enqueue_msg:
+            ForwardMsgQueue._before_enqueue_msg(msg)
+
+        if not _is_composable_message(msg):
+            self._queue.append(msg)
+            return
+
+        # If there's a Delta message with the same delta_path already in
+        # the queue - meaning that it refers to the same location in
+        # the app - we attempt to combine this new Delta into the old
+        # one. This is an optimization that prevents redundant Deltas
+        # from being sent to the frontend.
+        delta_key = tuple(msg.metadata.delta_path)
+        if delta_key in self._delta_index_map:
+            index = self._delta_index_map[delta_key]
+            old_msg = self._queue[index]
+            composed_delta = _maybe_compose_deltas(old_msg.delta, msg.delta)
+            if composed_delta is not None:
+                new_msg = ForwardMsg()
+                new_msg.delta.CopyFrom(composed_delta)
+                new_msg.metadata.CopyFrom(msg.metadata)
+                self._queue[index] = new_msg
+                return
+
+        # No composition occurred. Append this message to the queue, and
+        # store its index for potential future composition.
+        self._delta_index_map[delta_key] = len(self._queue)
+        self._queue.append(msg)
+
+    def clear(
+        self,
+        retain_lifecycle_msgs: bool = False,
+        fragment_ids_this_run: list[str] | None = None,
+    ) -> None:
+        """Clear the queue, potentially retaining lifecycle messages.
+
+        The retain_lifecycle_msgs argument exists because in some cases (in particular
+        when a currently running script is interrupted by a new BackMsg), we don't want
+        to remove certain messages from the queue as doing so may cause the client to
+        not hear about important script lifecycle events (such as the script being
+        stopped early in order to be rerun).
+
+        If fragment_ids_this_run is provided, delta messages not belonging to any
+        fragment or belonging to a fragment not in fragment_ids_this_run will be
+        preserved to prevent clearing messages unrelated to the running fragments.
+        """
+
+        if not retain_lifecycle_msgs:
+            self._queue = []
+        else:
+            self._queue = [
+                _update_script_finished_message(msg, fragment_ids_this_run is not None)
+                for msg in self._queue
+                if msg.WhichOneof("type")
+                in {
+                    "new_session",
+                    "script_finished",
+                    "session_status_changed",
+                    "parent_message",
+                }
+                or (
+                    # preserve all messages if this is a fragment rerun and...
+                    fragment_ids_this_run is not None
+                    and (
+                        # the message is not a delta message
+                        # (not associated with a fragment) or...
+                        msg.delta is None
+                        or (
+                            # it is a delta but not associated with any of the passed
+                            # fragments
+                            msg.delta is not None
+                            and (
+                                msg.delta.fragment_id is None
+                                or msg.delta.fragment_id not in fragment_ids_this_run
+                            )
+                        )
+                    )
+                )
+            ]
+
+        self._delta_index_map = {}
+
+    def flush(self) -> list[ForwardMsg]:
+        """Clear the queue and return a list of the messages it contained
+        before being cleared.
+        """
+        queue = self._queue
+        self.clear()
+        return queue
+
+    def __len__(self) -> int:
+        return len(self._queue)
+
+
+def _is_composable_message(msg: ForwardMsg) -> bool:
+    """True if the ForwardMsg is potentially composable with other ForwardMsgs."""
+    if not msg.HasField("delta"):
+        # Non-delta messages are never composable.
+        return False
+
+    # We never compose add_rows messages in Python, because the add_rows
+    # operation can raise errors, and we don't have a good way of handling
+    # those errors in the message queue.
+    delta_type = msg.delta.WhichOneof("type")
+    return delta_type != "add_rows" and delta_type != "arrow_add_rows"
+
+
+def _maybe_compose_deltas(old_delta: Delta, new_delta: Delta) -> Delta | None:
+    """Combines new_delta onto old_delta if possible.
+
+    If the combination takes place, the function returns a new Delta that
+    should replace old_delta in the queue.
+
+    If the new_delta is incompatible with old_delta, the function returns None.
+    In this case, the new_delta should just be appended to the queue as normal.
+    """
+    old_delta_type = old_delta.WhichOneof("type")
+    if old_delta_type == "add_block":
+        # We never replace add_block deltas, because blocks can have
+        # other dependent deltas later in the queue. For example:
+        #
+        #   placeholder = st.empty()
+        #   placeholder.columns(1)
+        #   placeholder.empty()
+        #
+        # The call to "placeholder.columns(1)" creates two blocks, a parent
+        # container with delta_path (0, 0), and a column child with
+        # delta_path (0, 0, 0). If the final "placeholder.empty()" Delta
+        # is composed with the parent container Delta, the frontend will
+        # throw an error when it tries to add that column child to what is
+        # now just an element, and not a block.
+        return None
+
+    new_delta_type = new_delta.WhichOneof("type")
+    if new_delta_type == "new_element":
+        return new_delta
+
+    if new_delta_type == "add_block":
+        return new_delta
+
+    return None
+
+
+def _update_script_finished_message(
+    msg: ForwardMsg, is_fragment_run: bool
+) -> ForwardMsg:
+    """
+    When we are here, the message queue is cleared from non-lifecycle messages
+    before they were flushed to the browser.
+
+    If there were no non-lifecycle messages in the queue, changing the type here
+    should not matter for the frontend anyways, so we optimistically change the
+    `script_finished` message to `FINISHED_EARLY_FOR_RERUN`. This indicates to
+    the frontend that the previous run was interrupted by a new script start.
+    Otherwise, a `FINISHED_SUCCESSFULLY` message might trigger a reset of widget
+    states on the frontend.
+    """
+    if msg.WhichOneof("type") == "script_finished" and (
+        # If this is not a fragment run (= full app run), its okay to change the
+        # script_finished type to FINISHED_EARLY_FOR_RERUN because another full app run
+        # is about to start.
+        # If this is a fragment run, it is allowed to change the state of
+        # all script_finished states except for FINISHED_SUCCESSFULLY, which we use to
+        # indicate that a full app run has finished successfully (in other words, a
+        # fragment should not modify the finished status of a full app run, because
+        # the fragment finished state is different and the frontend might not trigger
+        # cleanups etc. correctly).
+        is_fragment_run is False
+        or msg.script_finished != ForwardMsg.ScriptFinishedStatus.FINISHED_SUCCESSFULLY
+    ):
+        msg.script_finished = ForwardMsg.ScriptFinishedStatus.FINISHED_EARLY_FOR_RERUN
+    return msg