Repository: streamlit/streamlit

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!


Discussion:
@MathijsNL Thanks for reporting this. I was able to reproduce this issue. We do have something that is actually supposed to clean up the media file storage after a session. But this somehow might not work correctly.  

I also tested with older Streamlit versions (e.g. 1.15), and this does seem to behave the same way. So it's not something that we broke recently.
In the  [media_file_manager.py](https://github.com/streamlit/streamlit/blob/1.20.0/lib/streamlit/runtime/media_file_manager.py#L213) file, I added a line of code that calls the remove_orphaned_files() function (self.remove_orphaned_files()) to maintain stable memory usage at 232 MB during one session with the provided code. 

However, I noticed that there is no existing call to remove_orphaned_files() in the media_file_manager code at all right now. I'm unsure of the best location to add the function call or if the function itself is outdated and because of that it is no longer in use.

When I tried adding the function call to the MediaFileManager's add function for multiple sessions, I started receiving missing file errors, but at least if there is a single long running session this workaround would work for me.

@vdonato I am tagging you here as requested on the forum.
Hey @MathijsNL,

I finally had a bit of time to dig into this, and I think there may not actually be a true memory leak here (at this point I'm fairly certain but not 100% confident as it's been quite awhile since I've had to reason about what may be going on at the OS level).

After inspecting the streamlit server's media file storage when sessions are cleaned up + poking around with a memory profiler + seeing what objects the `gc` is keeping track of, I couldn't find any evidence that we're failing to clean up the media files that images are saved in after a session ends (which is when all of these mechanisms get kicked off -- your workaround should be sufficient for the use case where a session lives indefinitely, but the ultra-long-lived session isn't a use case we currently support very well. Supporting this use case would best be filed as an enhancement request rather than a bug report).

What I think might be going on here is that the memory used by the test script is properly garbage collected by Python when a session closes, but the OS isn't reclaiming all of this memory back from the process unless there's sufficient memory pressure on the system that it needs to do so -- you'd likely need another memory hungry process on the same system to force the memory reallocation to actually happen. This seems to be consistent with information that I've been able to find online about how RSS is actually calculated (see https://utcc.utoronto.ca/~cks/space/blog/linux/LinuxMemoryStats).

Let me know if this seems inconsistent with any findings that you may have.
Going to change the status of this one to `unlikely` because it seems like there may not be a bug here, but I'm not 100% sure about this conclusion yet.

Resolving patch:
diff --git a/lib/streamlit/runtime/app_session.py b/lib/streamlit/runtime/app_session.py
--- 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
--- 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
--- 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
--- 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
--- 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):