@@ -115,6 +115,28 @@ def get_deterministic_priv_key(self):
115
115
]
116
116
return PRIV_KEYS [self .index ]
117
117
118
+ def get_mem_rss (self ):
119
+ """Get the memory usage (RSS) per `ps`.
120
+
121
+ If process is stopped or `ps` is unavailable, return None.
122
+ """
123
+ if not (self .running and self .process ):
124
+ self .log .warning ("Couldn't get memory usage; process isn't running." )
125
+ return None
126
+
127
+ try :
128
+ return int (subprocess .check_output (
129
+ "ps h -o rss {}" .format (self .process .pid ),
130
+ shell = True , stderr = subprocess .DEVNULL ).strip ())
131
+
132
+ # Catching `Exception` broadly to avoid failing on platforms where ps
133
+ # isn't installed or doesn't work as expected, e.g. OpenBSD.
134
+ #
135
+ # We could later use something like `psutils` to work across platforms.
136
+ except Exception :
137
+ self .log .exception ("Unable to get memory usage" )
138
+ return None
139
+
118
140
def _node_msg (self , msg : str ) -> str :
119
141
"""Return a modified msg that identifies this node by its index as a debugging aid."""
120
142
return "[node %d] %s" % (self .index , msg )
@@ -267,6 +289,29 @@ def assert_debug_log(self, expected_msgs):
267
289
if re .search (re .escape (expected_msg ), log , flags = re .MULTILINE ) is None :
268
290
self ._raise_assertion_error ('Expected message "{}" does not partially match log:\n \n {}\n \n ' .format (expected_msg , print_log ))
269
291
292
+ @contextlib .contextmanager
293
+ def assert_memory_usage_stable (self , perc_increase_allowed = 0.03 ):
294
+ """Context manager that allows the user to assert that a node's memory usage (RSS)
295
+ hasn't increased beyond some threshold percentage.
296
+ """
297
+ before_memory_usage = self .get_mem_rss ()
298
+
299
+ yield
300
+
301
+ after_memory_usage = self .get_mem_rss ()
302
+
303
+ if not (before_memory_usage and after_memory_usage ):
304
+ self .log .warning ("Unable to detect memory usage (RSS) - skipping memory check." )
305
+ return
306
+
307
+ perc_increase_memory_usage = 1 - (float (before_memory_usage ) / after_memory_usage )
308
+
309
+ if perc_increase_memory_usage > perc_increase_allowed :
310
+ self ._raise_assertion_error (
311
+ "Memory usage increased over threshold of {:.3f}% from {} to {} ({:.3f}%)" .format (
312
+ perc_increase_allowed * 100 , before_memory_usage , after_memory_usage ,
313
+ perc_increase_memory_usage * 100 ))
314
+
270
315
def assert_start_raises_init_error (self , extra_args = None , expected_msg = None , match = ErrorMatch .FULL_TEXT , * args , ** kwargs ):
271
316
"""Attempt to start the node and expect it to raise an error.
272
317
0 commit comments