fix: Properly handle recursion error when storing pickle

pull/653/head
Raphaël Vinot 2023-03-30 16:30:27 +02:00
parent 8122b9904c
commit 9281d73675
1 changed files with 18 additions and 11 deletions

View File

@ -274,6 +274,7 @@ class CapturesIndex(Mapping):
if not (har_files := sorted(capture_dir.glob('*.har'))): if not (har_files := sorted(capture_dir.glob('*.har'))):
har_files = sorted(capture_dir.glob('*.har.gz')) har_files = sorted(capture_dir.glob('*.har.gz'))
try: try:
default_recursion_limit = sys.getrecursionlimit()
with self._timeout_context(): with self._timeout_context():
tree = CrawledTree(har_files, uuid) tree = CrawledTree(har_files, uuid)
self.__resolve_dns(tree, logger) self.__resolve_dns(tree, logger)
@ -292,19 +293,25 @@ class CapturesIndex(Mapping):
except RecursionError as e: except RecursionError as e:
raise NoValidHarFile(f'Tree too deep, probably a recursive refresh: {e}.\n Append /export to the URL to get the files.') raise NoValidHarFile(f'Tree too deep, probably a recursive refresh: {e}.\n Append /export to the URL to get the files.')
else: else:
with gzip.open(capture_dir / 'tree.pickle.gz', 'wb') as _p: # Some pickles require a pretty high recursion limit, this kindof fixes it.
# Some pickles require a pretty high recursion limit, this kindof fixes it. # If the capture is really broken (generally a refresh to self), the capture
# If the capture is really broken (generally a refresh to self), the capture # is discarded in the RecursionError above.
# is discarded in the RecursionError above. sys.setrecursionlimit(int(default_recursion_limit * 1.1))
default_recursion_limit = sys.getrecursionlimit() try:
sys.setrecursionlimit(int(default_recursion_limit * 1.1)) with gzip.open(capture_dir / 'tree.pickle.gz', 'wb') as _p:
try:
_p.write(pickletools.optimize(pickle.dumps(tree, protocol=5))) _p.write(pickletools.optimize(pickle.dumps(tree, protocol=5)))
# pickle.dump(tree, _p, protocol=5) except RecursionError as e:
except RecursionError as e: logger.exception('Unable to store pickle.')
raise NoValidHarFile(f'Tree too deep, probably a recursive refresh: {e}.\n Append /export to the URL to get the files.') # unable to use the HAR files, get them out of the way
sys.setrecursionlimit(default_recursion_limit) for har_file in har_files:
har_file.rename(har_file.with_suffix('.broken'))
(capture_dir / 'tree.pickle.gz').unlink(missing_ok=True)
raise NoValidHarFile(f'Tree too deep, probably a recursive refresh: {e}.\n Append /export to the URL to get the files.')
except Exception:
(capture_dir / 'tree.pickle.gz').unlink(missing_ok=True)
logger.exception('Unable to store pickle.')
finally: finally:
sys.setrecursionlimit(default_recursion_limit)
lock_file.unlink(missing_ok=True) lock_file.unlink(missing_ok=True)
return tree return tree