python: fix racy prebuilt download
Bug: https://github.com/google/perfetto/issues/773
Change-Id: Ib44239b40abbdbde7300b61e66b223aeedad247d
diff --git a/tools/heap_profile b/tools/heap_profile
index 700cfeb..eec9508 100755
--- a/tools/heap_profile
+++ b/tools/heap_profile
@@ -216,6 +216,9 @@
import platform
import subprocess
import sys
+import threading
+
+DOWNLOAD_LOCK = threading.Lock()
def download_or_get_cached(file_name, url, sha256):
@@ -232,28 +235,36 @@
sha256_path = os.path.join(dir, file_name + '.sha256')
needs_download = True
- # Avoid recomputing the SHA-256 on each invocation. The SHA-256 of the last
- # download is cached into file_name.sha256, just check if that matches.
- if os.path.exists(bin_path) and os.path.exists(sha256_path):
- with open(sha256_path, 'rb') as f:
- digest = f.read().decode()
- if digest == sha256:
- needs_download = False
+ try:
+ # In BatchTraceProcessor, many threads can be trying to execute the below
+ # code in parallel. For this reason, protect the whole operation with a
+ # lock.
+ DOWNLOAD_LOCK.acquire()
- if needs_download:
- # Either the filed doesn't exist or the SHA256 doesn't match.
- tmp_path = bin_path + '.tmp'
- print('Downloading ' + url)
- subprocess.check_call(['curl', '-f', '-L', '-#', '-o', tmp_path, url])
- with open(tmp_path, 'rb') as fd:
- actual_sha256 = hashlib.sha256(fd.read()).hexdigest()
- if actual_sha256 != sha256:
- raise Exception('Checksum mismatch for %s (actual: %s, expected: %s)' %
- (url, actual_sha256, sha256))
- os.chmod(tmp_path, 0o755)
- os.replace(tmp_path, bin_path)
- with open(sha256_path, 'w') as f:
- f.write(sha256)
+ # Avoid recomputing the SHA-256 on each invocation. The SHA-256 of the last
+ # download is cached into file_name.sha256, just check if that matches.
+ if os.path.exists(bin_path) and os.path.exists(sha256_path):
+ with open(sha256_path, 'rb') as f:
+ digest = f.read().decode()
+ if digest == sha256:
+ needs_download = False
+
+ if needs_download:
+ # Either the filed doesn't exist or the SHA256 doesn't match.
+ tmp_path = bin_path + '.tmp'
+ print('Downloading ' + url)
+ subprocess.check_call(['curl', '-f', '-L', '-#', '-o', tmp_path, url])
+ with open(tmp_path, 'rb') as fd:
+ actual_sha256 = hashlib.sha256(fd.read()).hexdigest()
+ if actual_sha256 != sha256:
+ raise Exception('Checksum mismatch for %s (actual: %s, expected: %s)' %
+ (url, actual_sha256, sha256))
+ os.chmod(tmp_path, 0o755)
+ os.replace(tmp_path, bin_path)
+ with open(sha256_path, 'w') as f:
+ f.write(sha256)
+ finally:
+ DOWNLOAD_LOCK.release()
return bin_path
