forked from streamlit/streamlit
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathforward_msg_cache.py
More file actions
263 lines (202 loc) · 8 KB
/
Copy pathforward_msg_cache.py
File metadata and controls
263 lines (202 loc) · 8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
# Copyright 2018-2021 Streamlit Inc.
#
# 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.
import hashlib
from typing import MutableMapping, TYPE_CHECKING
from weakref import WeakKeyDictionary
from streamlit import config
from streamlit import util
from streamlit.logger import get_logger
from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
if TYPE_CHECKING:
from streamlit.report_session import ReportSession
LOGGER = get_logger(__name__)
def populate_hash_if_needed(msg):
"""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.
hasher = hashlib.md5()
hasher.update(msg.SerializeToString())
msg.hash = hasher.hexdigest()
# Restore metadata.
msg.metadata.CopyFrom(metadata)
return msg.hash
def create_reference_msg(msg):
"""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(object):
"""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(object):
"""Cache entry.
Stores the cached message, and the set of ReportSessions
that we've sent the cached message to.
"""
def __init__(self, msg):
self.msg = msg
self._session_report_run_counts = (
WeakKeyDictionary()
) # type: MutableMapping[ReportSession, int]
def __repr__(self) -> str:
return util.repr_(self)
def add_session_ref(self, session, report_run_count):
"""Adds a reference to a ReportSession that has referenced
this Entry's message.
Parameters
----------
session : ReportSession
report_run_count : int
The session's run count at the time of the call
"""
prev_run_count = self._session_report_run_counts.get(session, 0)
if report_run_count < prev_run_count:
LOGGER.error(
"New report_run_count (%s) is < prev_run_count (%s). "
"This should never happen!" % (report_run_count, prev_run_count)
)
report_run_count = prev_run_count
self._session_report_run_counts[session] = report_run_count
def has_session_ref(self, session):
return session in self._session_report_run_counts
def get_session_ref_age(self, session, report_run_count):
"""The age of the given session's reference to the Entry,
given a new report_run_count.
"""
return report_run_count - self._session_report_run_counts[session]
def remove_session_ref(self, session):
del self._session_report_run_counts[session]
def has_refs(self):
"""True if this Entry has references from any ReportSession.
If not, it can be removed from the cache.
"""
return len(self._session_report_run_counts) > 0
def __init__(self):
self._entries = {} # Map: hash -> Entry
def __repr__(self) -> str:
return util.repr_(self)
def add_message(self, msg, session, report_run_count):
"""Add a ForwardMsg to the cache.
The cache will also record a reference to the given ReportSession,
so that it can track which sessions have already received
each given ForwardMsg.
Parameters
----------
msg : ForwardMsg
session : ReportSession
report_run_count : int
The number of times the session's report has run
"""
populate_hash_if_needed(msg)
entry = self._entries.get(msg.hash, None)
if entry is None:
entry = ForwardMsgCache.Entry(msg)
self._entries[msg.hash] = entry
entry.add_session_ref(session, report_run_count)
def get_message(self, hash):
"""Return the message with the given ID if it exists in the cache.
Parameters
----------
hash : string
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, session, report_run_count):
"""Return True if a session has a reference to a message.
Parameters
----------
msg : ForwardMsg
session : ReportSession
report_run_count : int
The number of times the session's report has run
Returns
-------
bool
"""
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, report_run_count)
return age <= config.get_option("global.maxCachedMessageAge")
def remove_expired_session_entries(self, session, report_run_count):
"""Remove any cached messages that have expired from the given session.
This should be called each time a ReportSession finishes executing.
Parameters
----------
session : ReportSession
report_run_count : int
The number of times the session's report 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, report_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):
"""Remove all entries from the cache"""
self._entries.clear()