Repository: streamlit/streamlit
Issue category: Performance Issue

Updating images will increase memory usage
### Checklist

- [X] I have searched the [existing issues](https://github.com/streamlit/streamlit/issues) for similar issues.
- [X] I added a very descriptive title to this issue.
- [X] I have provided sufficient information below to help reproduce this issue.

### Summary

When updating images frequently the memory usage of the Streamlit app keeps increasing. A lot!

You don't need a lot of sessions to get this behavior. Running a single page and refreshing it after some time will not release the memory.

### Reproducible Code Example

```Python
import datetime
import time
import psutil
import os
pid = os.getpid()


import cv2
import numpy as np
import streamlit as st

img = np.random.rand(600,600,3)

@st.cache_resource(ttl=1)
def get_memory(pid):
    process = psutil.Process(pid)

    # Get the memory usage in RAM
    memory_usage = process.memory_info().rss

    # Convert the memory usage to MB
    memory_usage_mb = memory_usage / (1024 * 1024)
    return(memory_usage_mb)
    #print(f"Total memory usage of all running Python processes: {mem_usage} bytes")


def get_image():
    # Get current date and time
    img_with_date = img.copy()
    now = datetime.datetime.now()
    dt_string = now.strftime("%d/%m/%Y %H:%M:%S.%f")[:-3]
    cv2.putText(img_with_date, dt_string, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    return img_with_date

def app():
    if "page_open" not in st.session_state:
        st.session_state.page_open = True

    if "num_img" not in st.session_state:
        st.session_state.num_img = 0


    st.title("Test memory usage of st.image...")

    text_placeholder = st.empty()
    img_placeholder = st.empty()

    while "page_open" in st.session_state:
        text_placeholder.write(f"{st.session_state.num_img} {get_memory(pid):.2f}MB")
        img_placeholder.image(get_image())

        st.session_state.num_img += 1
        time.sleep(1/10)

    print("Session ended...")
    img_placeholder.empty()
    text_placeholder.empty()
    st.stop()
if __name__ == "__main__":
    app()
```


### Steps To Reproduce

1. Run streamlit app
2. Maybe open a few tabs to speed up the memory consumption
3. Close all tabs
4. Open a new tab, memory usage will have decreased a bit but it will never go down to an acceptable size anymore

### Expected Behavior

Resources should be removed when no sessions are running anymore. In specific use cases it might even be needed to have an aggressive cleanup of resources, for example when showing live video.

### Current Behavior

Current behavior is that images are stored somewhere in RAM and there doesn't seem to be a way to control it. This was already discussed [on the forum](https://discuss.streamlit.io/t/updating-images-in-a-while-loop-keeps-increasing-ram-even-after-sessions-are-closed-and-cleaned-up/38226/5) with @vdonato



### Is this a regression?

- [ ] Yes, this used to work in a previous version.

### Debug info

Streamlit version: 1.20.0
Python version: 3.10.9
Operating System: Windows 11 (WSL2)
Browser: Chrome
Virtual environment: conda

### Additional Information

The app will show both the image id and the memory usage to make it really easy to see that the RAM is indeed increasing over time.

Dependencies:
```
pip install streamlit==1.20.0
pip install opencv-python-headless
pip install psutil
```

### Are you willing to submit a PR?

- [ ] Yes, I am willing to submit a PR!


Functions to edit:
lib/streamlit/runtime/app_session.py:AppSession.shutdown
lib/streamlit/runtime/forward_msg_cache.py:ForwardMsgCache.remove_expired_session_entries
lib/streamlit/runtime/runtime.py:Runtime._send_message
Functions added: lib/streamlit/runtime/forward_msg_cache.py:ForwardMsgCache.remove_refs_for_session
lib/streamlit/runtime/forward_msg_cache.py:ForwardMsgCache.remove_expired_entries_for_session

Patch:
diff --git a/lib/streamlit/runtime/app_session.py b/lib/streamlit/runtime/app_session.py
index 6fae7653b807..36e681c344f7 100644
--- a/lib/streamlit/runtime/app_session.py
+++ b/lib/streamlit/runtime/app_session.py
@@ -225,8 +225,10 @@ def shutdown(self) -> None:
             self._uploaded_file_mgr.remove_session_files(self.id)
 
             if runtime.exists():
-                runtime.get_instance().media_file_mgr.clear_session_refs(self.id)
-                runtime.get_instance().media_file_mgr.remove_orphaned_files()
+                rt = runtime.get_instance()
+                rt.media_file_mgr.clear_session_refs(self.id)
+                rt.media_file_mgr.remove_orphaned_files()
+                rt.message_cache.remove_refs_for_session(self)
 
             # Shut down the ScriptRunner, if one is active.
             # self._state must not be set to SHUTDOWN_REQUESTED until
diff --git a/lib/streamlit/runtime/forward_msg_cache.py b/lib/streamlit/runtime/forward_msg_cache.py
index f348603ac520..13f494867569 100644
--- a/lib/streamlit/runtime/forward_msg_cache.py
+++ b/lib/streamlit/runtime/forward_msg_cache.py
@@ -217,7 +217,28 @@ def has_message_reference(
         age = entry.get_session_ref_age(session, script_run_count)
         return age <= int(config.get_option("global.maxCachedMessageAge"))
 
-    def remove_expired_session_entries(
+    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 being shut down.
+
+        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.
diff --git a/lib/streamlit/runtime/runtime.py b/lib/streamlit/runtime/runtime.py
index ddb18a1656dd..b34503bcac03 100644
--- a/lib/streamlit/runtime/runtime.py
+++ b/lib/streamlit/runtime/runtime.py
@@ -687,7 +687,7 @@ def _send_message(self, session_info: ActiveSessionInfo, msg: ForwardMsg) -> Non
                 config.get_option("global.maxCachedMessageAge"),
             )
             session_info.script_run_count += 1
-            self._message_cache.remove_expired_session_entries(
+            self._message_cache.remove_expired_entries_for_session(
                 session_info.session, session_info.script_run_count
             )
 

Test patch: diff --git a/lib/tests/streamlit/runtime/app_session_test.py b/lib/tests/streamlit/runtime/app_session_test.py
index 2ee2e5d8cd40..448118736e9a 100644
--- a/lib/tests/streamlit/runtime/app_session_test.py
+++ b/lib/tests/streamlit/runtime/app_session_test.py
@@ -101,15 +101,19 @@ def test_shutdown(self, patched_disconnect):
 
         mock_file_mgr = MagicMock(spec=UploadedFileManager)
         session._uploaded_file_mgr = mock_file_mgr
+        mock_message_cache = Runtime._instance.message_cache
 
         session.shutdown()
         self.assertEqual(AppSessionState.SHUTDOWN_REQUESTED, session._state)
         mock_file_mgr.remove_session_files.assert_called_once_with(session.id)
         patched_disconnect.assert_called_once_with(session._on_secrets_file_changed)
+        mock_message_cache.remove_refs_for_session.assert_called_once_with(session)
 
         # A 2nd shutdown call should have no effect.
         session.shutdown()
         self.assertEqual(AppSessionState.SHUTDOWN_REQUESTED, session._state)
+        mock_message_cache.remove_refs_for_session.assert_called_once_with(session)
+
         mock_file_mgr.remove_session_files.assert_called_once_with(session.id)
 
     def test_shutdown_with_running_scriptrunner(self):
diff --git a/lib/tests/streamlit/runtime/forward_msg_cache_test.py b/lib/tests/streamlit/runtime/forward_msg_cache_test.py
index 7d5122808cda..8bce96913de3 100644
--- a/lib/tests/streamlit/runtime/forward_msg_cache_test.py
+++ b/lib/tests/streamlit/runtime/forward_msg_cache_test.py
@@ -101,6 +101,42 @@ def test_clear(self):
         cache.clear()
         self.assertEqual(None, cache.get_message(msg_hash))
 
+    def test_remove_refs_for_session(self):
+        cache = ForwardMsgCache()
+
+        session1 = _create_mock_session()
+        session2 = _create_mock_session()
+
+        # Only session1 has a ref to msg1.
+        msg1 = _create_dataframe_msg([1, 2, 3])
+        populate_hash_if_needed(msg1)
+        cache.add_message(msg1, session1, 0)
+
+        # Only session2 has a ref to msg2.
+        msg2 = _create_dataframe_msg([1, 2, 3, 4])
+        populate_hash_if_needed(msg2)
+        cache.add_message(msg2, session2, 0)
+
+        # Both session1 and session2 have a ref to msg3.
+        msg3 = _create_dataframe_msg([1, 2, 3, 4, 5])
+        populate_hash_if_needed(msg2)
+        cache.add_message(msg3, session1, 0)
+        cache.add_message(msg3, session2, 0)
+
+        cache.remove_refs_for_session(session1)
+
+        cache_entries = list(cache._entries.values())
+
+        cached_msgs = [entry.msg for entry in cache_entries]
+        assert cached_msgs == [msg2, msg3]
+
+        sessions_with_refs = {
+            s
+            for entry in cache_entries
+            for s in entry._session_script_run_counts.keys()
+        }
+        assert sessions_with_refs == {session2}
+
     def test_message_expiration(self):
         """Test MessageCache's expiration logic"""
         config._set_option("global.maxCachedMessageAge", 1, "test")
@@ -131,14 +167,14 @@ def test_message_expiration(self):
 
         # Remove session1's expired entries. This should not remove the
         # entry from the cache, because session2 still has a reference to it.
-        cache.remove_expired_session_entries(session1, runcount1)
+        cache.remove_expired_entries_for_session(session1, runcount1)
         self.assertFalse(cache.has_message_reference(msg, session1, runcount1))
         self.assertTrue(cache.has_message_reference(msg, session2, runcount2))
 
         # Expire session2's reference. The message should no longer be
         # in the cache at all.
         runcount2 += 2
-        cache.remove_expired_session_entries(session2, runcount2)
+        cache.remove_expired_entries_for_session(session2, runcount2)
         self.assertIsNone(cache.get_message(msg_hash))
 
     def test_cache_stats_provider(self):