blob: 411bb18906dfca5a4c16e853be715c4932c0491d [file] [log] [blame]
maruel@chromium.orgc6f90062012-11-07 18:32:22 +00001#!/usr/bin/env python
maruelea586f32016-04-05 11:11:33 -07002# Copyright 2013 The LUCI Authors. All rights reserved.
maruelf1f5e2a2016-05-25 17:10:39 -07003# Use of this source code is governed under the Apache License, Version 2.0
4# that can be found in the LICENSE file.
maruel@chromium.orgc6f90062012-11-07 18:32:22 +00005
Marc-Antoine Ruel52436aa2014-08-28 21:57:57 -04006"""Archives a set of files or directories to an Isolate Server."""
maruel@chromium.orgc6f90062012-11-07 18:32:22 +00007
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +00008__version__ = '0.9.0'
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00009
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +000010import collections
nodir90bc8dc2016-06-15 13:35:21 -070011import errno
tansell9e04a8d2016-07-28 09:31:59 -070012import functools
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000013import logging
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -040014import optparse
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000015import os
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +000016import re
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +040017import signal
tansell9e04a8d2016-07-28 09:31:59 -070018import stat
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000019import sys
tansell26de79e2016-11-13 18:41:11 -080020import tarfile
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +000021import threading
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000022import time
csharp@chromium.org59c7bcf2012-11-21 21:13:18 +000023import zlib
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000024
Marc-Antoine Ruel016c7602019-04-02 18:31:13 +000025from utils import tools
26tools.force_local_third_party()
maruel@chromium.orgfb78d432013-08-28 21:22:40 +000027
Marc-Antoine Ruel016c7602019-04-02 18:31:13 +000028# third_party/
29import colorama
30from depot_tools import fix_encoding
31from depot_tools import subcommand
Takuto Ikuta6e2ff962019-10-29 12:35:27 +000032import six
Lei Leife202df2019-06-11 17:33:34 +000033from six.moves import queue as Queue
Marc-Antoine Ruel016c7602019-04-02 18:31:13 +000034
35# pylint: disable=ungrouped-imports
36import auth
37import isolated_format
38import isolate_storage
39import local_caching
Marc-Antoine Ruel37989932013-11-19 16:28:08 -050040from utils import file_path
maruel12e30012015-10-09 11:55:35 -070041from utils import fs
Marc-Antoine Ruelf74cffe2015-07-15 15:21:34 -040042from utils import logging_utils
vadimsh@chromium.org6b706212013-08-28 15:03:46 +000043from utils import net
Marc-Antoine Ruelcfb60852014-07-02 15:22:00 -040044from utils import on_error
maruel8e4e40c2016-05-30 06:21:07 -070045from utils import subprocess42
vadimsh@chromium.orgb074b162013-08-22 17:55:46 +000046from utils import threading_utils
Vadim Shtayurae34e13a2014-02-02 11:23:26 -080047
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000048
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +000049# Version of isolate protocol passed to the server in /handshake request.
50ISOLATE_PROTOCOL_VERSION = '1.0'
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000051
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +000052
Vadim Shtayura3148e072014-09-02 18:51:52 -070053# Maximum expected delay (in seconds) between successive file fetches or uploads
54# in Storage. If it takes longer than that, a deadlock might be happening
55# and all stack frames for all threads are dumped to log.
56DEADLOCK_TIMEOUT = 5 * 60
57
58
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +000059# The number of files to check the isolate server per /pre-upload query.
vadimsh@chromium.orgeea52422013-08-21 19:35:54 +000060# All files are sorted by likelihood of a change in the file content
61# (currently file size is used to estimate this: larger the file -> larger the
62# possibility it has changed). Then first ITEMS_PER_CONTAINS_QUERIES[0] files
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +000063# are taken and send to '/pre-upload', then next ITEMS_PER_CONTAINS_QUERIES[1],
vadimsh@chromium.orgeea52422013-08-21 19:35:54 +000064# and so on. Numbers here is a trade-off; the more per request, the lower the
65# effect of HTTP round trip latency and TCP-level chattiness. On the other hand,
66# larger values cause longer lookups, increasing the initial latency to start
67# uploading, which is especially an issue for large files. This value is
68# optimized for the "few thousands files to look up with minimal number of large
69# files missing" case.
Marc-Antoine Ruel52436aa2014-08-28 21:57:57 -040070ITEMS_PER_CONTAINS_QUERIES = (20, 20, 50, 50, 50, 100)
csharp@chromium.org07fa7592013-01-11 18:19:30 +000071
maruel@chromium.org9958e4a2013-09-17 00:01:48 +000072
csharp@chromium.org59c7bcf2012-11-21 21:13:18 +000073# A list of already compressed extension types that should not receive any
74# compression before being uploaded.
75ALREADY_COMPRESSED_TYPES = [
Marc-Antoine Ruel7f234c82014-08-06 21:55:18 -040076 '7z', 'avi', 'cur', 'gif', 'h264', 'jar', 'jpeg', 'jpg', 'mp4', 'pdf',
77 'png', 'wav', 'zip',
csharp@chromium.org59c7bcf2012-11-21 21:13:18 +000078]
79
maruel@chromium.orgc6f90062012-11-07 18:32:22 +000080
maruel@chromium.org41601642013-09-18 19:40:46 +000081# The delay (in seconds) to wait between logging statements when retrieving
82# the required files. This is intended to let the user (or buildbot) know that
83# the program is still running.
84DELAY_BETWEEN_UPDATES_IN_SECS = 30
85
86
Marc-Antoine Ruelac54cb42013-11-18 14:05:35 -050087DEFAULT_BLACKLIST = (
88 # Temporary vim or python files.
89 r'^.+\.(?:pyc|swp)$',
90 # .git or .svn directory.
91 r'^(?:.+' + re.escape(os.path.sep) + r'|)\.(?:git|svn)$',
92)
93
94
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -050095class Error(Exception):
96 """Generic runtime error."""
97 pass
98
99
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400100class Aborted(Error):
101 """Operation aborted."""
102 pass
103
104
nodir90bc8dc2016-06-15 13:35:21 -0700105class AlreadyExists(Error):
106 """File already exists."""
107
108
maruel12e30012015-10-09 11:55:35 -0700109def file_read(path, chunk_size=isolated_format.DISK_FILE_CHUNK, offset=0):
Vadim Shtayuraf0cb97a2013-12-05 13:57:49 -0800110 """Yields file content in chunks of |chunk_size| starting from |offset|."""
maruel12e30012015-10-09 11:55:35 -0700111 with fs.open(path, 'rb') as f:
Vadim Shtayuraf0cb97a2013-12-05 13:57:49 -0800112 if offset:
113 f.seek(offset)
maruel@chromium.org8750e4b2013-09-18 02:37:57 +0000114 while True:
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000115 data = f.read(chunk_size)
maruel@chromium.org8750e4b2013-09-18 02:37:57 +0000116 if not data:
117 break
118 yield data
119
120
tansell9e04a8d2016-07-28 09:31:59 -0700121def fileobj_path(fileobj):
122 """Return file system path for file like object or None.
123
124 The returned path is guaranteed to exist and can be passed to file system
125 operations like copy.
126 """
127 name = getattr(fileobj, 'name', None)
128 if name is None:
Marc-Antoine Ruel793bff32019-04-18 17:50:48 +0000129 return None
tansell9e04a8d2016-07-28 09:31:59 -0700130
131 # If the file like object was created using something like open("test.txt")
132 # name will end up being a str (such as a function outside our control, like
133 # the standard library). We want all our paths to be unicode objects, so we
134 # decode it.
Takuto Ikuta95459dd2019-10-29 12:39:47 +0000135 if not isinstance(name, six.text_type):
Marc-Antoine Rueld8464b12017-12-04 15:59:41 -0500136 # We incorrectly assume that UTF-8 is used everywhere.
137 name = name.decode('utf-8')
tansell9e04a8d2016-07-28 09:31:59 -0700138
tansell26de79e2016-11-13 18:41:11 -0800139 # fs.exists requires an absolute path, otherwise it will fail with an
140 # assertion error.
141 if not os.path.isabs(name):
Takuto Ikuta523c6472019-09-18 02:53:34 +0000142 return None
tansell26de79e2016-11-13 18:41:11 -0800143
tansell9e04a8d2016-07-28 09:31:59 -0700144 if fs.exists(name):
145 return name
Marc-Antoine Ruel793bff32019-04-18 17:50:48 +0000146 return None
tansell9e04a8d2016-07-28 09:31:59 -0700147
148
149# TODO(tansell): Replace fileobj_copy with shutil.copyfileobj once proper file
150# wrappers have been created.
151def fileobj_copy(
152 dstfileobj, srcfileobj, size=-1,
153 chunk_size=isolated_format.DISK_FILE_CHUNK):
154 """Copy data from srcfileobj to dstfileobj.
155
156 Providing size means exactly that amount of data will be copied (if there
157 isn't enough data, an IOError exception is thrown). Otherwise all data until
158 the EOF marker will be copied.
159 """
160 if size == -1 and hasattr(srcfileobj, 'tell'):
161 if srcfileobj.tell() != 0:
162 raise IOError('partial file but not using size')
163
164 written = 0
165 while written != size:
166 readsize = chunk_size
167 if size > 0:
168 readsize = min(readsize, size-written)
169 data = srcfileobj.read(readsize)
170 if not data:
171 if size == -1:
172 break
173 raise IOError('partial file, got %s, wanted %s' % (written, size))
174 dstfileobj.write(data)
175 written += len(data)
176
177
178def putfile(srcfileobj, dstpath, file_mode=None, size=-1, use_symlink=False):
179 """Put srcfileobj at the given dstpath with given mode.
180
181 The function aims to do this as efficiently as possible while still allowing
182 any possible file like object be given.
183
184 Creating a tree of hardlinks has a few drawbacks:
185 - tmpfs cannot be used for the scratch space. The tree has to be on the same
186 partition as the cache.
187 - involves a write to the inode, which advances ctime, cause a metadata
188 writeback (causing disk seeking).
189 - cache ctime cannot be used to detect modifications / corruption.
190 - Some file systems (NTFS) have a 64k limit on the number of hardlink per
191 partition. This is why the function automatically fallbacks to copying the
192 file content.
193 - /proc/sys/fs/protected_hardlinks causes an additional check to ensure the
194 same owner is for all hardlinks.
195 - Anecdotal report that ext2 is known to be potentially faulty on high rate
196 of hardlink creation.
197
198 Creating a tree of symlinks has a few drawbacks:
199 - Tasks running the equivalent of os.path.realpath() will get the naked path
200 and may fail.
201 - Windows:
202 - Symlinks are reparse points:
203 https://msdn.microsoft.com/library/windows/desktop/aa365460.aspx
204 https://msdn.microsoft.com/library/windows/desktop/aa363940.aspx
205 - Symbolic links are Win32 paths, not NT paths.
206 https://googleprojectzero.blogspot.com/2016/02/the-definitive-guide-on-win32-to-nt.html
207 - Symbolic links are supported on Windows 7 and later only.
208 - SeCreateSymbolicLinkPrivilege is needed, which is not present by
209 default.
210 - SeCreateSymbolicLinkPrivilege is *stripped off* by UAC when a restricted
211 RID is present in the token;
212 https://msdn.microsoft.com/en-us/library/bb530410.aspx
213 """
214 srcpath = fileobj_path(srcfileobj)
215 if srcpath and size == -1:
216 readonly = file_mode is None or (
217 file_mode & (stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH))
218
219 if readonly:
220 # If the file is read only we can link the file
221 if use_symlink:
222 link_mode = file_path.SYMLINK_WITH_FALLBACK
223 else:
224 link_mode = file_path.HARDLINK_WITH_FALLBACK
225 else:
226 # If not read only, we must copy the file
227 link_mode = file_path.COPY
228
229 file_path.link_file(dstpath, srcpath, link_mode)
Takuto Ikuta523c6472019-09-18 02:53:34 +0000230 assert fs.exists(dstpath)
tansell9e04a8d2016-07-28 09:31:59 -0700231 else:
232 # Need to write out the file
233 with fs.open(dstpath, 'wb') as dstfileobj:
234 fileobj_copy(dstfileobj, srcfileobj, size)
235
Takuto Ikuta523c6472019-09-18 02:53:34 +0000236 if sys.platform == 'win32' and file_mode and file_mode & stat.S_IWRITE:
237 # On windows, mode other than removing stat.S_IWRITE is ignored. Returns
238 # early to skip slow/unnecessary chmod call.
239 return
tansell9e04a8d2016-07-28 09:31:59 -0700240
241 # file_mode of 0 is actually valid, so need explicit check.
242 if file_mode is not None:
243 fs.chmod(dstpath, file_mode)
244
245
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000246def zip_compress(content_generator, level=7):
247 """Reads chunks from |content_generator| and yields zip compressed chunks."""
248 compressor = zlib.compressobj(level)
249 for chunk in content_generator:
250 compressed = compressor.compress(chunk)
251 if compressed:
252 yield compressed
253 tail = compressor.flush(zlib.Z_FINISH)
254 if tail:
255 yield tail
256
257
Marc-Antoine Ruel8bee66d2014-08-28 19:02:07 -0400258def zip_decompress(
259 content_generator, chunk_size=isolated_format.DISK_FILE_CHUNK):
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000260 """Reads zipped data from |content_generator| and yields decompressed data.
261
262 Decompresses data in small chunks (no larger than |chunk_size|) so that
263 zip bomb file doesn't cause zlib to preallocate huge amount of memory.
264
265 Raises IOError if data is corrupted or incomplete.
266 """
267 decompressor = zlib.decompressobj()
268 compressed_size = 0
269 try:
270 for chunk in content_generator:
271 compressed_size += len(chunk)
272 data = decompressor.decompress(chunk, chunk_size)
273 if data:
274 yield data
275 while decompressor.unconsumed_tail:
276 data = decompressor.decompress(decompressor.unconsumed_tail, chunk_size)
277 if data:
278 yield data
279 tail = decompressor.flush()
280 if tail:
281 yield tail
282 except zlib.error as e:
283 raise IOError(
284 'Corrupted zip stream (read %d bytes) - %s' % (compressed_size, e))
285 # Ensure all data was read and decompressed.
286 if decompressor.unused_data or decompressor.unconsumed_tail:
287 raise IOError('Not all data was decompressed')
288
289
Marc-Antoine Rueldcff6462018-12-04 16:35:18 +0000290def _get_zip_compression_level(filename):
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000291 """Given a filename calculates the ideal zip compression level to use."""
292 file_ext = os.path.splitext(filename)[1].lower()
293 # TODO(csharp): Profile to find what compression level works best.
294 return 0 if file_ext in ALREADY_COMPRESSED_TYPES else 7
295
296
maruel@chromium.orgaf254852013-09-17 17:48:14 +0000297def create_directories(base_directory, files):
298 """Creates the directory structure needed by the given list of files."""
299 logging.debug('create_directories(%s, %d)', base_directory, len(files))
300 # Creates the tree of directories to create.
301 directories = set(os.path.dirname(f) for f in files)
302 for item in list(directories):
303 while item:
304 directories.add(item)
305 item = os.path.dirname(item)
306 for d in sorted(directories):
307 if d:
aludwin606aa1f2016-10-31 18:41:30 -0700308 abs_d = os.path.join(base_directory, d)
309 if not fs.isdir(abs_d):
310 fs.mkdir(abs_d)
maruel@chromium.orgaf254852013-09-17 17:48:14 +0000311
312
Marc-Antoine Rueldcff6462018-12-04 16:35:18 +0000313def _create_symlinks(base_directory, files):
Marc-Antoine Ruelccafe0e2013-11-08 16:15:36 -0500314 """Creates any symlinks needed by the given set of files."""
maruel@chromium.orgaf254852013-09-17 17:48:14 +0000315 for filepath, properties in files:
316 if 'l' not in properties:
317 continue
318 if sys.platform == 'win32':
Marc-Antoine Ruelccafe0e2013-11-08 16:15:36 -0500319 # TODO(maruel): Create symlink via the win32 api.
maruel@chromium.orgaf254852013-09-17 17:48:14 +0000320 logging.warning('Ignoring symlink %s', filepath)
321 continue
322 outfile = os.path.join(base_directory, filepath)
nodir90bc8dc2016-06-15 13:35:21 -0700323 try:
324 os.symlink(properties['l'], outfile) # pylint: disable=E1101
325 except OSError as e:
326 if e.errno == errno.EEXIST:
327 raise AlreadyExists('File %s already exists.' % outfile)
328 raise
maruel@chromium.orgaf254852013-09-17 17:48:14 +0000329
330
Marc-Antoine Ruel440eee62018-12-04 22:37:05 +0000331class _ThreadFile(object):
332 """Multithreaded fake file. Used by TarBundle."""
333 def __init__(self):
334 self._data = threading_utils.TaskChannel()
335 self._offset = 0
336
337 def __iter__(self):
338 return self._data
339
340 def tell(self):
341 return self._offset
342
343 def write(self, b):
344 self._data.send_result(b)
345 self._offset += len(b)
346
347 def close(self):
348 self._data.send_done()
349
350
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400351class FileItem(isolate_storage.Item):
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800352 """A file to push to Storage.
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000353
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800354 Its digest and size may be provided in advance, if known. Otherwise they will
355 be derived from the file content.
356 """
357
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +0000358 def __init__(self, path, algo, digest=None, size=None, high_priority=False):
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800359 super(FileItem, self).__init__(
360 digest,
maruel12e30012015-10-09 11:55:35 -0700361 size if size is not None else fs.stat(path).st_size,
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +0000362 high_priority,
363 compression_level=_get_zip_compression_level(path))
364 self._path = path
365 self._algo = algo
366 self._meta = None
367
368 @property
369 def path(self):
370 return self._path
371
372 @property
373 def digest(self):
374 if not self._digest:
375 self._digest = isolated_format.hash_file(self._path, self._algo)
376 return self._digest
377
378 @property
379 def meta(self):
380 if not self._meta:
381 # TODO(maruel): Inline.
382 self._meta = isolated_format.file_to_metadata(self.path, 0, False)
383 # We need to hash right away.
384 self._meta['h'] = self.digest
385 return self._meta
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000386
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800387 def content(self):
388 return file_read(self.path)
maruel@chromium.orgc6f90062012-11-07 18:32:22 +0000389
390
Marc-Antoine Ruel440eee62018-12-04 22:37:05 +0000391class TarBundle(isolate_storage.Item):
392 """Tarfile to push to Storage.
393
394 Its digest is the digest of all the files it contains. It is generated on the
395 fly.
396 """
397
398 def __init__(self, root, algo):
399 # 2 trailing 512 bytes headers.
400 super(TarBundle, self).__init__(size=1024)
401 self._items = []
402 self._meta = None
403 self._algo = algo
404 self._root_len = len(root) + 1
405 # Same value as for Go.
406 # https://chromium.googlesource.com/infra/luci/luci-go.git/+/master/client/archiver/tar_archiver.go
407 # https://chromium.googlesource.com/infra/luci/luci-go.git/+/master/client/archiver/upload_tracker.go
408 self._archive_max_size = int(10e6)
409
410 @property
411 def digest(self):
412 if not self._digest:
413 self._prepare()
414 return self._digest
415
416 @property
417 def size(self):
418 if self._size is None:
419 self._prepare()
420 return self._size
421
422 def try_add(self, item):
423 """Try to add this file to the bundle.
424
425 It is extremely naive but this should be just enough for
426 https://crbug.com/825418.
427
428 Future improvements should be in the Go code, and the Swarming bot should be
429 migrated to use the Go code instead.
430 """
431 if not item.size:
432 return False
433 # pylint: disable=unreachable
434 rounded = (item.size + 512) & ~511
435 if rounded + self._size > self._archive_max_size:
436 return False
437 # https://crbug.com/825418
438 return False
439 self._size += rounded
440 self._items.append(item)
441 return True
442
443 def yield_item_path_meta(self):
444 """Returns a tuple(Item, filepath, meta_dict).
445
446 If the bundle contains less than 5 items, the items are yielded.
447 """
448 if len(self._items) < 5:
449 # The tarball is too small, yield individual items, if any.
450 for item in self._items:
451 yield item, item.path[self._root_len:], item.meta
452 else:
453 # This ensures self._meta is set.
454 p = self.digest + '.tar'
455 # Yield itself as a tarball.
456 yield self, p, self._meta
457
458 def content(self):
459 """Generates the tarfile content on the fly."""
460 obj = _ThreadFile()
461 def _tar_thread():
462 try:
463 t = tarfile.open(
464 fileobj=obj, mode='w', format=tarfile.PAX_FORMAT, encoding='utf-8')
465 for item in self._items:
466 logging.info(' tarring %s', item.path)
467 t.add(item.path)
468 t.close()
469 except Exception:
470 logging.exception('Internal failure')
471 finally:
472 obj.close()
473
474 t = threading.Thread(target=_tar_thread)
475 t.start()
476 try:
477 for data in obj:
478 yield data
479 finally:
480 t.join()
481
482 def _prepare(self):
483 h = self._algo()
484 total = 0
485 for chunk in self.content():
486 h.update(chunk)
487 total += len(chunk)
488 # pylint: disable=attribute-defined-outside-init
489 # This is not true, they are defined in Item.__init__().
490 self._digest = h.hexdigest()
491 self._size = total
492 self._meta = {
493 'h': self.digest,
494 's': self.size,
495 't': u'tar',
496 }
497
498
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400499class BufferItem(isolate_storage.Item):
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000500 """A byte buffer to push to Storage."""
501
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +0000502 def __init__(self, buf, algo, high_priority=False):
503 super(BufferItem, self).__init__(
504 digest=algo(buf).hexdigest(),
505 size=len(buf),
506 high_priority=high_priority)
507 self._buffer = buf
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000508
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800509 def content(self):
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +0000510 return [self._buffer]
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000511
512
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000513class Storage(object):
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800514 """Efficiently downloads or uploads large set of files via StorageApi.
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000515
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800516 Implements compression support, parallel 'contains' checks, parallel uploads
517 and more.
518
519 Works only within single namespace (and thus hashing algorithm and compression
520 scheme are fixed).
521
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400522 Spawns multiple internal threads. Thread safe, but not fork safe. Modifies
523 signal handlers table to handle Ctrl+C.
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800524 """
525
Vadim Shtayurae0ab1902014-04-29 10:55:27 -0700526 def __init__(self, storage_api):
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000527 self._storage_api = storage_api
528 self._cpu_thread_pool = None
529 self._net_thread_pool = None
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400530 self._aborted = False
531 self._prev_sig_handlers = {}
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000532
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000533 @property
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +0000534 def server_ref(self):
535 """Shortcut to get the server_ref from storage_api.
Vadim Shtayurae0ab1902014-04-29 10:55:27 -0700536
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +0000537 This can be used to get the underlying hash_algo.
Vadim Shtayurae0ab1902014-04-29 10:55:27 -0700538 """
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +0000539 return self._storage_api.server_ref
Vadim Shtayurae0ab1902014-04-29 10:55:27 -0700540
541 @property
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000542 def cpu_thread_pool(self):
543 """ThreadPool for CPU-bound tasks like zipping."""
544 if self._cpu_thread_pool is None:
Marc-Antoine Ruelbdad1182015-02-06 16:04:35 -0500545 threads = max(threading_utils.num_processors(), 2)
Lei Leife202df2019-06-11 17:33:34 +0000546 max_size = long(2)**32 if sys.version_info.major == 2 else 2**32
547 if sys.maxsize <= max_size:
Marc-Antoine Ruelbdad1182015-02-06 16:04:35 -0500548 # On 32 bits userland, do not try to use more than 16 threads.
549 threads = min(threads, 16)
550 self._cpu_thread_pool = threading_utils.ThreadPool(2, threads, 0, 'zip')
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000551 return self._cpu_thread_pool
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000552
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000553 @property
554 def net_thread_pool(self):
555 """AutoRetryThreadPool for IO-bound tasks, retries IOError."""
556 if self._net_thread_pool is None:
Vadim Shtayura3148e072014-09-02 18:51:52 -0700557 self._net_thread_pool = threading_utils.IOAutoRetryThreadPool()
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000558 return self._net_thread_pool
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000559
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000560 def close(self):
561 """Waits for all pending tasks to finish."""
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400562 logging.info('Waiting for all threads to die...')
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000563 if self._cpu_thread_pool:
564 self._cpu_thread_pool.join()
565 self._cpu_thread_pool.close()
566 self._cpu_thread_pool = None
567 if self._net_thread_pool:
568 self._net_thread_pool.join()
569 self._net_thread_pool.close()
570 self._net_thread_pool = None
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400571 logging.info('Done.')
572
573 def abort(self):
574 """Cancels any pending or future operations."""
575 # This is not strictly theadsafe, but in the worst case the logging message
576 # will be printed twice. Not a big deal. In other places it is assumed that
577 # unprotected reads and writes to _aborted are serializable (it is true
578 # for python) and thus no locking is used.
579 if not self._aborted:
580 logging.warning('Aborting... It can take a while.')
581 self._aborted = True
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000582
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000583 def __enter__(self):
584 """Context manager interface."""
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400585 assert not self._prev_sig_handlers, self._prev_sig_handlers
586 for s in (signal.SIGINT, signal.SIGTERM):
587 self._prev_sig_handlers[s] = signal.signal(s, lambda *_args: self.abort())
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000588 return self
589
590 def __exit__(self, _exc_type, _exc_value, _traceback):
591 """Context manager interface."""
592 self.close()
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400593 while self._prev_sig_handlers:
594 s, h = self._prev_sig_handlers.popitem()
595 signal.signal(s, h)
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000596 return False
597
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000598 def upload_items(self, items):
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000599 """Uploads a generator of Item to the isolate server.
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000600
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800601 It figures out what items are missing from the server and uploads only them.
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000602
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000603 It uses 3 threads internally:
604 - One to create batches based on a timeout
605 - One to dispatch the /contains RPC and field the missing entries
606 - One to field the /push RPC
607
608 The main threads enumerates 'items' and pushes to the first thread. Then it
609 join() all the threads, waiting for them to complete.
610
611 (enumerate items of Item, this can be slow as disk is traversed)
612 |
613 v
614 _create_items_batches_thread Thread #1
615 (generates list(Item), every 3s or 20~100 items)
616 |
617 v
618 _do_lookups_thread Thread #2
619 | |
620 v v
621 (missing) (was on server)
622 |
623 v
624 _handle_missing_thread Thread #3
625 |
626 v
627 (upload Item, append to uploaded)
628
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000629 Arguments:
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400630 items: list of isolate_storage.Item instances that represents data to
631 upload.
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000632
633 Returns:
634 List of items that were uploaded. All other items are already there.
635 """
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000636 incoming = Queue.Queue()
637 batches_to_lookup = Queue.Queue()
638 missing = Queue.Queue()
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000639 uploaded = []
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800640
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000641 def _create_items_batches_thread():
642 """Creates batches for /contains RPC lookup from individual items.
643
644 Input: incoming
645 Output: batches_to_lookup
646 """
647 try:
648 batch_size_index = 0
649 batch_size = ITEMS_PER_CONTAINS_QUERIES[batch_size_index]
650 batch = []
651 while not self._aborted:
652 try:
653 item = incoming.get(True, timeout=3)
654 if item:
655 batch.append(item)
656 except Queue.Empty:
657 item = False
658 if len(batch) == batch_size or (not item and batch):
659 if len(batch) == batch_size:
660 batch_size_index += 1
661 batch_size = ITEMS_PER_CONTAINS_QUERIES[
662 min(batch_size_index, len(ITEMS_PER_CONTAINS_QUERIES)-1)]
663 batches_to_lookup.put(batch)
664 batch = []
665 if item is None:
666 break
667 finally:
668 # Unblock the next pipeline.
669 batches_to_lookup.put(None)
670
671 def _do_lookups_thread():
672 """Enqueues all the /contains RPCs and emits the missing items.
673
674 Input: batches_to_lookup
675 Output: missing, to_upload
676 """
677 try:
678 channel = threading_utils.TaskChannel()
679 def _contains(b):
680 if self._aborted:
681 raise Aborted()
682 return self._storage_api.contains(b)
683
684 pending_contains = 0
685 while not self._aborted:
686 batch = batches_to_lookup.get()
687 if batch is None:
688 break
689 self.net_thread_pool.add_task_with_channel(
690 channel, threading_utils.PRIORITY_HIGH, _contains, batch)
691 pending_contains += 1
692 while pending_contains and not self._aborted:
693 try:
Marc-Antoine Ruel4494b6c2018-11-28 21:00:41 +0000694 v = channel.next(timeout=0)
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000695 except threading_utils.TaskChannel.Timeout:
696 break
697 pending_contains -= 1
Marc-Antoine Ruel04903a32019-10-09 21:09:25 +0000698 for missing_item, push_state in v.items():
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000699 missing.put((missing_item, push_state))
700 while pending_contains and not self._aborted:
Marc-Antoine Ruel04903a32019-10-09 21:09:25 +0000701 for missing_item, push_state in channel.next().items():
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000702 missing.put((missing_item, push_state))
703 pending_contains -= 1
704 finally:
705 # Unblock the next pipeline.
706 missing.put((None, None))
707
708 def _handle_missing_thread():
709 """Sends the missing items to the uploader.
710
711 Input: missing
712 Output: uploaded
713 """
Vadim Shtayura3148e072014-09-02 18:51:52 -0700714 with threading_utils.DeadlockDetector(DEADLOCK_TIMEOUT) as detector:
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000715 channel = threading_utils.TaskChannel()
716 pending_upload = 0
717 while not self._aborted:
718 try:
719 missing_item, push_state = missing.get(True, timeout=5)
720 if missing_item is None:
721 break
722 self._async_push(channel, missing_item, push_state)
723 pending_upload += 1
724 except Queue.Empty:
725 pass
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000726 detector.ping()
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000727 while not self._aborted and pending_upload:
728 try:
Marc-Antoine Ruel4494b6c2018-11-28 21:00:41 +0000729 item = channel.next(timeout=0)
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000730 except threading_utils.TaskChannel.Timeout:
731 break
732 uploaded.append(item)
733 pending_upload -= 1
734 logging.debug(
735 'Uploaded %d; %d pending: %s (%d)',
736 len(uploaded), pending_upload, item.digest, item.size)
737 while not self._aborted and pending_upload:
Marc-Antoine Ruel4494b6c2018-11-28 21:00:41 +0000738 item = channel.next()
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000739 uploaded.append(item)
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000740 pending_upload -= 1
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000741 logging.debug(
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000742 'Uploaded %d; %d pending: %s (%d)',
743 len(uploaded), pending_upload, item.digest, item.size)
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000744
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000745 threads = [
746 threading.Thread(target=_create_items_batches_thread),
747 threading.Thread(target=_do_lookups_thread),
748 threading.Thread(target=_handle_missing_thread),
749 ]
750 for t in threads:
751 t.start()
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000752
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000753 try:
754 # For each digest keep only first isolate_storage.Item that matches it.
755 # All other items are just indistinguishable copies from the point of view
756 # of isolate server (it doesn't care about paths at all, only content and
757 # digests).
758 seen = {}
759 try:
760 # TODO(maruel): Reorder the items as a priority queue, with larger items
761 # being processed first. This is, before hashing the data.
762 # This must be done in the primary thread since items can be a
763 # generator.
764 for item in items:
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000765 if seen.setdefault(item.digest, item) is item:
766 incoming.put(item)
767 finally:
768 incoming.put(None)
769 finally:
770 for t in threads:
771 t.join()
772
773 logging.info('All %s files are uploaded', len(uploaded))
Marc-Antoine Ruel73c0ae72018-11-30 14:05:45 +0000774 if seen:
775 _print_upload_stats(seen.values(), uploaded)
vadimsh@chromium.orgf24e5c32013-10-11 21:16:21 +0000776 return uploaded
777
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000778 def _async_push(self, channel, item, push_state):
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000779 """Starts asynchronous push to the server in a parallel thread.
780
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000781 Can be used only after |item| was checked for presence on a server with a
782 /contains RPC.
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800783
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000784 Arguments:
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000785 channel: TaskChannel that receives back |item| when upload ends.
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400786 item: item to upload as instance of isolate_storage.Item class.
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000787 push_state: push state returned by storage_api.contains(). It contains
788 storage specific information describing how to upload the item (for
789 example in case of cloud storage, it is signed upload URLs).
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800790
791 Returns:
792 None, but |channel| later receives back |item| when upload ends.
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000793 """
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800794 # Thread pool task priority.
Marc-Antoine Ruel52436aa2014-08-28 21:57:57 -0400795 priority = (
Vadim Shtayura3148e072014-09-02 18:51:52 -0700796 threading_utils.PRIORITY_HIGH if item.high_priority
797 else threading_utils.PRIORITY_MED)
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800798
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000799 def _push(content):
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400800 """Pushes an isolate_storage.Item and returns it to |channel|."""
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400801 if self._aborted:
802 raise Aborted()
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800803 self._storage_api.push(item, push_state, content)
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000804 return item
805
Wei Huang1a38fbe2017-11-28 22:55:22 -0500806 # If zipping is not required, just start a push task. Don't pass 'content'
807 # so that it can create a new generator when it retries on failures.
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +0000808 if not self.server_ref.is_with_compression:
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000809 self.net_thread_pool.add_task_with_channel(channel, priority, _push, None)
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000810 return
811
812 # If zipping is enabled, zip in a separate thread.
813 def zip_and_push():
814 # TODO(vadimsh): Implement streaming uploads. Before it's done, assemble
815 # content right here. It will block until all file is zipped.
816 try:
Vadim Shtayuraf9e401b2014-10-15 18:19:37 +0400817 if self._aborted:
818 raise Aborted()
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800819 stream = zip_compress(item.content(), item.compression_level)
Lei Lei73a5f732020-03-23 20:36:14 +0000820 # In Python3, zlib.compress returns a byte object instead of str.
821 data = six.b('').join(stream)
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000822 except Exception as exc:
823 logging.error('Failed to zip \'%s\': %s', item, exc)
Vadim Shtayura0ffc4092013-11-20 17:49:52 -0800824 channel.send_exception()
vadimsh@chromium.org7cdf1c02013-09-25 00:24:16 +0000825 return
Wei Huang1a38fbe2017-11-28 22:55:22 -0500826 # Pass '[data]' explicitly because the compressed data is not same as the
827 # one provided by 'item'. Since '[data]' is a list, it can safely be
828 # reused during retries.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000829 self.net_thread_pool.add_task_with_channel(
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000830 channel, priority, _push, [data])
vadimsh@chromium.orgbcb966b2013-10-01 18:14:18 +0000831 self.cpu_thread_pool.add_task(priority, zip_and_push)
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000832
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800833 def push(self, item, push_state):
834 """Synchronously pushes a single item to the server.
835
836 If you need to push many items at once, consider using 'upload_items' or
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000837 '_async_push' with instance of TaskChannel.
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800838
839 Arguments:
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400840 item: item to upload as instance of isolate_storage.Item class.
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000841 push_state: push state returned by storage_api.contains(). It contains
842 storage specific information describing how to upload the item (for
843 example in case of cloud storage, it is signed upload URLs).
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800844
845 Returns:
846 Pushed item (same object as |item|).
847 """
848 channel = threading_utils.TaskChannel()
Vadim Shtayura3148e072014-09-02 18:51:52 -0700849 with threading_utils.DeadlockDetector(DEADLOCK_TIMEOUT):
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +0000850 self._async_push(channel, item, push_state)
Marc-Antoine Ruel4494b6c2018-11-28 21:00:41 +0000851 pushed = channel.next()
Vadim Shtayurabcff74f2014-02-27 16:19:34 -0800852 assert pushed is item
853 return item
854
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000855 def async_fetch(self, channel, priority, digest, size, sink):
856 """Starts asynchronous fetch from the server in a parallel thread.
857
858 Arguments:
859 channel: TaskChannel that receives back |digest| when download ends.
860 priority: thread pool task priority for the fetch.
861 digest: hex digest of an item to download.
862 size: expected size of the item (after decompression).
863 sink: function that will be called as sink(generator).
864 """
865 def fetch():
866 try:
867 # Prepare reading pipeline.
Adrian Ludwinb4ebc092017-09-13 07:46:24 -0400868 stream = self._storage_api.fetch(digest, size, 0)
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +0000869 if self.server_ref.is_with_compression:
Marc-Antoine Ruel8bee66d2014-08-28 19:02:07 -0400870 stream = zip_decompress(stream, isolated_format.DISK_FILE_CHUNK)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000871 # Run |stream| through verifier that will assert its size.
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +0000872 verifier = FetchStreamVerifier(
873 stream, self.server_ref.hash_algo, digest, size)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000874 # Verified stream goes to |sink|.
875 sink(verifier.run())
876 except Exception as err:
Vadim Shtayura0ffc4092013-11-20 17:49:52 -0800877 logging.error('Failed to fetch %s: %s', digest, err)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000878 raise
879 return digest
880
881 # Don't bother with zip_thread_pool for decompression. Decompression is
882 # really fast and most probably IO bound anyway.
883 self.net_thread_pool.add_task_with_channel(channel, priority, fetch)
884
vadimsh@chromium.org35122be2013-09-19 02:48:00 +0000885
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000886class FetchQueue(object):
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400887 """Fetches items from Storage and places them into ContentAddressedCache.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000888
889 It manages multiple concurrent fetch operations. Acts as a bridge between
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400890 Storage and ContentAddressedCache so that Storage and ContentAddressedCache
891 don't depend on each other at all.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000892 """
893
894 def __init__(self, storage, cache):
895 self.storage = storage
896 self.cache = cache
897 self._channel = threading_utils.TaskChannel()
898 self._pending = set()
899 self._accessed = set()
Marc-Antoine Ruel5d7606b2018-06-15 19:06:12 +0000900 self._fetched = set(cache)
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400901 # Pending digests that the caller waits for, see wait_on()/wait().
902 self._waiting_on = set()
903 # Already fetched digests the caller waits for which are not yet returned by
904 # wait().
905 self._waiting_on_ready = set()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000906
Marc-Antoine Ruel52436aa2014-08-28 21:57:57 -0400907 def add(
Vadim Shtayura3148e072014-09-02 18:51:52 -0700908 self,
909 digest,
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -0400910 size=local_caching.UNKNOWN_FILE_SIZE,
Vadim Shtayura3148e072014-09-02 18:51:52 -0700911 priority=threading_utils.PRIORITY_MED):
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000912 """Starts asynchronous fetch of item |digest|."""
913 # Fetching it now?
914 if digest in self._pending:
915 return
916
917 # Mark this file as in use, verify_all_cached will later ensure it is still
918 # in cache.
919 self._accessed.add(digest)
920
921 # Already fetched? Notify cache to update item's LRU position.
922 if digest in self._fetched:
923 # 'touch' returns True if item is in cache and not corrupted.
924 if self.cache.touch(digest, size):
925 return
Marc-Antoine Ruel5d7606b2018-06-15 19:06:12 +0000926 logging.error('%s is corrupted', digest)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000927 self._fetched.remove(digest)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000928
929 # TODO(maruel): It should look at the free disk space, the current cache
930 # size and the size of the new item on every new item:
931 # - Trim the cache as more entries are listed when free disk space is low,
932 # otherwise if the amount of data downloaded during the run > free disk
933 # space, it'll crash.
934 # - Make sure there's enough free disk space to fit all dependencies of
935 # this run! If not, abort early.
936
937 # Start fetching.
938 self._pending.add(digest)
939 self.storage.async_fetch(
940 self._channel, priority, digest, size,
941 functools.partial(self.cache.write, digest))
942
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400943 def wait_on(self, digest):
944 """Updates digests to be waited on by 'wait'."""
945 # Calculate once the already fetched items. These will be retrieved first.
946 if digest in self._fetched:
947 self._waiting_on_ready.add(digest)
948 else:
949 self._waiting_on.add(digest)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000950
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400951 def wait(self):
952 """Waits until any of waited-on items is retrieved.
953
954 Once this happens, it is remove from the waited-on set and returned.
955
956 This function is called in two waves. The first wave it is done for HIGH
957 priority items, the isolated files themselves. The second wave it is called
958 for all the files.
959
960 If the waited-on set is empty, raises RuntimeError.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000961 """
962 # Flush any already fetched items.
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400963 if self._waiting_on_ready:
964 return self._waiting_on_ready.pop()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000965
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400966 assert self._waiting_on, 'Needs items to wait on'
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000967
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400968 # Wait for one waited-on item to be fetched.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000969 while self._pending:
Marc-Antoine Ruel4494b6c2018-11-28 21:00:41 +0000970 digest = self._channel.next()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000971 self._pending.remove(digest)
972 self._fetched.add(digest)
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400973 if digest in self._waiting_on:
974 self._waiting_on.remove(digest)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000975 return digest
976
977 # Should never reach this point due to assert above.
978 raise RuntimeError('Impossible state')
979
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -0400980 @property
981 def wait_queue_empty(self):
982 """Returns True if there is no digest left for wait() to return."""
983 return not self._waiting_on and not self._waiting_on_ready
984
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000985 def inject_local_file(self, path, algo):
986 """Adds local file to the cache as if it was fetched from storage."""
maruel12e30012015-10-09 11:55:35 -0700987 with fs.open(path, 'rb') as f:
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +0000988 data = f.read()
989 digest = algo(data).hexdigest()
990 self.cache.write(digest, [data])
991 self._fetched.add(digest)
992 return digest
993
994 @property
995 def pending_count(self):
996 """Returns number of items to be fetched."""
997 return len(self._pending)
998
999 def verify_all_cached(self):
1000 """True if all accessed items are in cache."""
Marc-Antoine Ruel5d7606b2018-06-15 19:06:12 +00001001 # Not thread safe, but called after all work is done.
1002 return self._accessed.issubset(self.cache)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001003
1004
1005class FetchStreamVerifier(object):
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -04001006 """Verifies that fetched file is valid before passing it to the
1007 ContentAddressedCache.
1008 """
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001009
Adrian Ludwin6d2a8342017-08-15 19:56:54 -04001010 def __init__(self, stream, hasher, expected_digest, expected_size):
1011 """Initializes the verifier.
1012
1013 Arguments:
1014 * stream: an iterable yielding chunks of content
1015 * hasher: an object from hashlib that supports update() and hexdigest()
1016 (eg, hashlib.sha1).
1017 * expected_digest: if the entire stream is piped through hasher and then
1018 summarized via hexdigest(), this should be the result. That is, it
1019 should be a hex string like 'abc123'.
1020 * expected_size: either the expected size of the stream, or
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -04001021 local_caching.UNKNOWN_FILE_SIZE.
Adrian Ludwin6d2a8342017-08-15 19:56:54 -04001022 """
Marc-Antoine Rueldf4976d2015-04-15 19:56:21 -04001023 assert stream is not None
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001024 self.stream = stream
Adrian Ludwin6d2a8342017-08-15 19:56:54 -04001025 self.expected_digest = expected_digest
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001026 self.expected_size = expected_size
1027 self.current_size = 0
Adrian Ludwin6d2a8342017-08-15 19:56:54 -04001028 self.rolling_hash = hasher()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001029
1030 def run(self):
1031 """Generator that yields same items as |stream|.
1032
1033 Verifies |stream| is complete before yielding a last chunk to consumer.
1034
1035 Also wraps IOError produced by consumer into MappingError exceptions since
1036 otherwise Storage will retry fetch on unrelated local cache errors.
1037 """
1038 # Read one chunk ahead, keep it in |stored|.
1039 # That way a complete stream can be verified before pushing last chunk
1040 # to consumer.
1041 stored = None
1042 for chunk in self.stream:
1043 assert chunk is not None
1044 if stored is not None:
1045 self._inspect_chunk(stored, is_last=False)
1046 try:
1047 yield stored
1048 except IOError as exc:
Marc-Antoine Ruel52436aa2014-08-28 21:57:57 -04001049 raise isolated_format.MappingError(
1050 'Failed to store an item in cache: %s' % exc)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001051 stored = chunk
1052 if stored is not None:
1053 self._inspect_chunk(stored, is_last=True)
1054 try:
1055 yield stored
1056 except IOError as exc:
Marc-Antoine Ruel52436aa2014-08-28 21:57:57 -04001057 raise isolated_format.MappingError(
1058 'Failed to store an item in cache: %s' % exc)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001059
1060 def _inspect_chunk(self, chunk, is_last):
1061 """Called for each fetched chunk before passing it to consumer."""
1062 self.current_size += len(chunk)
Adrian Ludwin6d2a8342017-08-15 19:56:54 -04001063 self.rolling_hash.update(chunk)
1064 if not is_last:
1065 return
1066
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -04001067 if ((self.expected_size != local_caching.UNKNOWN_FILE_SIZE) and
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001068 (self.expected_size != self.current_size)):
Adrian Ludwin6d2a8342017-08-15 19:56:54 -04001069 msg = 'Incorrect file size: want %d, got %d' % (
1070 self.expected_size, self.current_size)
Adrian Ludwin6d2a8342017-08-15 19:56:54 -04001071 raise IOError(msg)
1072
1073 actual_digest = self.rolling_hash.hexdigest()
1074 if self.expected_digest != actual_digest:
1075 msg = 'Incorrect digest: want %s, got %s' % (
1076 self.expected_digest, actual_digest)
Adrian Ludwin21920d52017-08-22 09:34:19 -04001077 raise IOError(msg)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001078
1079
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001080class IsolatedBundle(object):
1081 """Fetched and parsed .isolated file with all dependencies."""
1082
Takuto Ikuta1e6072c2018-11-06 20:42:43 +00001083 def __init__(self, filter_cb):
1084 """
1085 filter_cb: callback function to filter downloaded content.
1086 When filter_cb is not None, Isolated file is downloaded iff
1087 filter_cb(filepath) returns True.
1088 """
1089
Vadim Shtayura3148e072014-09-02 18:51:52 -07001090 self.command = []
1091 self.files = {}
1092 self.read_only = None
1093 self.relative_cwd = None
1094 # The main .isolated file, a IsolatedFile instance.
1095 self.root = None
1096
Takuto Ikuta1e6072c2018-11-06 20:42:43 +00001097 self._filter_cb = filter_cb
1098
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001099 def fetch(self, fetch_queue, root_isolated_hash, algo):
1100 """Fetches the .isolated and all the included .isolated.
Vadim Shtayura3148e072014-09-02 18:51:52 -07001101
1102 It enables support for "included" .isolated files. They are processed in
1103 strict order but fetched asynchronously from the cache. This is important so
1104 that a file in an included .isolated file that is overridden by an embedding
1105 .isolated file is not fetched needlessly. The includes are fetched in one
1106 pass and the files are fetched as soon as all the ones on the left-side
1107 of the tree were fetched.
1108
1109 The prioritization is very important here for nested .isolated files.
1110 'includes' have the highest priority and the algorithm is optimized for both
1111 deep and wide trees. A deep one is a long link of .isolated files referenced
1112 one at a time by one item in 'includes'. A wide one has a large number of
1113 'includes' in a single .isolated file. 'left' is defined as an included
1114 .isolated file earlier in the 'includes' list. So the order of the elements
1115 in 'includes' is important.
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001116
1117 As a side effect this method starts asynchronous fetch of all data files
1118 by adding them to |fetch_queue|. It doesn't wait for data files to finish
1119 fetching though.
Vadim Shtayura3148e072014-09-02 18:51:52 -07001120 """
1121 self.root = isolated_format.IsolatedFile(root_isolated_hash, algo)
1122
1123 # Isolated files being retrieved now: hash -> IsolatedFile instance.
1124 pending = {}
1125 # Set of hashes of already retrieved items to refuse recursive includes.
1126 seen = set()
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001127 # Set of IsolatedFile's whose data files have already being fetched.
1128 processed = set()
Vadim Shtayura3148e072014-09-02 18:51:52 -07001129
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001130 def retrieve_async(isolated_file):
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -04001131 """Retrieves an isolated file included by the root bundle."""
Vadim Shtayura3148e072014-09-02 18:51:52 -07001132 h = isolated_file.obj_hash
1133 if h in seen:
1134 raise isolated_format.IsolatedError(
1135 'IsolatedFile %s is retrieved recursively' % h)
1136 assert h not in pending
1137 seen.add(h)
1138 pending[h] = isolated_file
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -04001139 # This isolated item is being added dynamically, notify FetchQueue.
1140 fetch_queue.wait_on(h)
Vadim Shtayura3148e072014-09-02 18:51:52 -07001141 fetch_queue.add(h, priority=threading_utils.PRIORITY_HIGH)
1142
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001143 # Start fetching root *.isolated file (single file, not the whole bundle).
1144 retrieve_async(self.root)
Vadim Shtayura3148e072014-09-02 18:51:52 -07001145
1146 while pending:
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001147 # Wait until some *.isolated file is fetched, parse it.
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -04001148 item_hash = fetch_queue.wait()
Vadim Shtayura3148e072014-09-02 18:51:52 -07001149 item = pending.pop(item_hash)
tansell9e04a8d2016-07-28 09:31:59 -07001150 with fetch_queue.cache.getfileobj(item_hash) as f:
1151 item.load(f.read())
Vadim Shtayura3148e072014-09-02 18:51:52 -07001152
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001153 # Start fetching included *.isolated files.
Vadim Shtayura3148e072014-09-02 18:51:52 -07001154 for new_child in item.children:
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001155 retrieve_async(new_child)
Vadim Shtayura3148e072014-09-02 18:51:52 -07001156
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001157 # Always fetch *.isolated files in traversal order, waiting if necessary
1158 # until next to-be-processed node loads. "Waiting" is done by yielding
1159 # back to the outer loop, that waits until some *.isolated is loaded.
1160 for node in isolated_format.walk_includes(self.root):
1161 if node not in processed:
1162 # Not visited, and not yet loaded -> wait for it to load.
1163 if not node.is_loaded:
1164 break
1165 # Not visited and loaded -> process it and continue the traversal.
1166 self._start_fetching_files(node, fetch_queue)
1167 processed.add(node)
Vadim Shtayura3148e072014-09-02 18:51:52 -07001168
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001169 # All *.isolated files should be processed by now and only them.
1170 all_isolateds = set(isolated_format.walk_includes(self.root))
1171 assert all_isolateds == processed, (all_isolateds, processed)
Marc-Antoine Ruel2d631542018-04-19 20:28:09 -04001172 assert fetch_queue.wait_queue_empty, 'FetchQueue should have been emptied'
Vadim Shtayura3148e072014-09-02 18:51:52 -07001173
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001174 # Extract 'command' and other bundle properties.
1175 for node in isolated_format.walk_includes(self.root):
1176 self._update_self(node)
Vadim Shtayura3148e072014-09-02 18:51:52 -07001177 self.relative_cwd = self.relative_cwd or ''
1178
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001179 def _start_fetching_files(self, isolated, fetch_queue):
1180 """Starts fetching files from |isolated| that are not yet being fetched.
Vadim Shtayura3148e072014-09-02 18:51:52 -07001181
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001182 Modifies self.files.
1183 """
maruel10bea7b2016-12-07 05:03:49 -08001184 files = isolated.data.get('files', {})
1185 logging.debug('fetch_files(%s, %d)', isolated.obj_hash, len(files))
Marc-Antoine Ruel04903a32019-10-09 21:09:25 +00001186 for filepath, properties in files.items():
Takuto Ikuta1e6072c2018-11-06 20:42:43 +00001187 if self._filter_cb and not self._filter_cb(filepath):
1188 continue
1189
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001190 # Root isolated has priority on the files being mapped. In particular,
1191 # overridden files must not be fetched.
1192 if filepath not in self.files:
1193 self.files[filepath] = properties
tansell9e04a8d2016-07-28 09:31:59 -07001194
1195 # Make sure if the isolated is read only, the mode doesn't have write
1196 # bits.
1197 if 'm' in properties and self.read_only:
1198 properties['m'] &= ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH)
1199
1200 # Preemptively request hashed files.
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001201 if 'h' in properties:
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001202 fetch_queue.add(
1203 properties['h'], properties['s'], threading_utils.PRIORITY_MED)
1204
1205 def _update_self(self, node):
1206 """Extracts bundle global parameters from loaded *.isolated file.
1207
1208 Will be called with each loaded *.isolated file in order of traversal of
1209 isolated include graph (see isolated_format.walk_includes).
1210 """
Vadim Shtayura3148e072014-09-02 18:51:52 -07001211 # Grabs properties.
1212 if not self.command and node.data.get('command'):
1213 # Ensure paths are correctly separated on windows.
1214 self.command = node.data['command']
1215 if self.command:
1216 self.command[0] = self.command[0].replace('/', os.path.sep)
Vadim Shtayura3148e072014-09-02 18:51:52 -07001217 if self.read_only is None and node.data.get('read_only') is not None:
1218 self.read_only = node.data['read_only']
1219 if (self.relative_cwd is None and
1220 node.data.get('relative_cwd') is not None):
1221 self.relative_cwd = node.data['relative_cwd']
1222
1223
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +00001224def get_storage(server_ref):
Vadim Shtayurabcff74f2014-02-27 16:19:34 -08001225 """Returns Storage class that can upload and download from |namespace|.
1226
1227 Arguments:
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +00001228 server_ref: isolate_storage.ServerRef instance.
Vadim Shtayurabcff74f2014-02-27 16:19:34 -08001229
1230 Returns:
1231 Instance of Storage.
1232 """
Lei Lei73a5f732020-03-23 20:36:14 +00001233 # Handle the specific internal use case.
1234 assert (isinstance(server_ref, isolate_storage.ServerRef) or
1235 type(server_ref).__name__ == 'ServerRef'), repr(server_ref)
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +00001236 return Storage(isolate_storage.get_storage_api(server_ref))
maruel@chromium.orgdedbf492013-09-12 20:42:11 +00001237
maruel@chromium.orgdedbf492013-09-12 20:42:11 +00001238
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001239def _map_file(dst, digest, props, cache, read_only, use_symlinks):
1240 """Put downloaded file to destination path. This function is used for multi
1241 threaded file putting.
1242 """
Takuto Ikuta523c6472019-09-18 02:53:34 +00001243 with tools.Profiler("_map_file for %s" % dst):
1244 with cache.getfileobj(digest) as srcfileobj:
1245 filetype = props.get('t', 'basic')
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001246
Takuto Ikuta523c6472019-09-18 02:53:34 +00001247 if filetype == 'basic':
1248 # Ignore all bits apart from the user.
1249 file_mode = (props.get('m') or 0o500) & 0o700
1250 if read_only:
1251 # Enforce read-only if the root bundle does.
1252 file_mode &= 0o500
1253 putfile(srcfileobj, dst, file_mode, use_symlink=use_symlinks)
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001254
Takuto Ikuta523c6472019-09-18 02:53:34 +00001255 elif filetype == 'tar':
1256 basedir = os.path.dirname(dst)
1257 with tarfile.TarFile(fileobj=srcfileobj, encoding='utf-8') as t:
1258 ensured_dirs = set()
1259 for ti in t:
1260 if not ti.isfile():
1261 logging.warning('Path(%r) is nonfile (%s), skipped', ti.name,
1262 ti.type)
1263 continue
1264 # Handle files created on Windows fetched on POSIX and the
1265 # reverse.
1266 other_sep = '/' if os.path.sep == '\\' else '\\'
1267 name = ti.name.replace(other_sep, os.path.sep)
1268 fp = os.path.normpath(os.path.join(basedir, name))
1269 if not fp.startswith(basedir):
1270 logging.error('Path(%r) is outside root directory', fp)
1271 ifd = t.extractfile(ti)
1272 fp_dir = os.path.dirname(fp)
1273 if fp_dir not in ensured_dirs:
1274 file_path.ensure_tree(fp_dir)
1275 ensured_dirs.add(fp_dir)
1276 file_mode = ti.mode & 0o700
1277 if read_only:
1278 # Enforce read-only if the root bundle does.
1279 file_mode &= 0o500
1280 putfile(ifd, fp, file_mode, ti.size)
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001281
Takuto Ikuta523c6472019-09-18 02:53:34 +00001282 else:
1283 raise isolated_format.IsolatedError('Unknown file type %r' % filetype)
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001284
1285
Takuto Ikuta1e6072c2018-11-06 20:42:43 +00001286def fetch_isolated(isolated_hash, storage, cache, outdir, use_symlinks,
1287 filter_cb=None):
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001288 """Aggressively downloads the .isolated file(s), then download all the files.
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001289
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001290 Arguments:
1291 isolated_hash: hash of the root *.isolated file.
1292 storage: Storage class that communicates with isolate storage.
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -04001293 cache: ContentAddressedCache class that knows how to store and map files
1294 locally.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001295 outdir: Output directory to map file tree to.
maruel4409e302016-07-19 14:25:51 -07001296 use_symlinks: Use symlinks instead of hardlinks when True.
Takuto Ikuta1e6072c2018-11-06 20:42:43 +00001297 filter_cb: filter that works as whitelist for downloaded files.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001298
1299 Returns:
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001300 IsolatedBundle object that holds details about loaded *.isolated file.
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001301 """
Marc-Antoine Ruel4e8cd182014-06-18 13:27:17 -04001302 logging.debug(
maruel4409e302016-07-19 14:25:51 -07001303 'fetch_isolated(%s, %s, %s, %s, %s)',
1304 isolated_hash, storage, cache, outdir, use_symlinks)
Vadim Shtayurae0ab1902014-04-29 10:55:27 -07001305 # Hash algorithm to use, defined by namespace |storage| is using.
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +00001306 algo = storage.server_ref.hash_algo
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001307 fetch_queue = FetchQueue(storage, cache)
Takuto Ikuta1e6072c2018-11-06 20:42:43 +00001308 bundle = IsolatedBundle(filter_cb)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001309
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001310 with tools.Profiler('GetIsolateds'):
1311 # Optionally support local files by manually adding them to cache.
1312 if not isolated_format.is_valid_hash(isolated_hash, algo):
1313 logging.debug('%s is not a valid hash, assuming a file '
1314 '(algo was %s, hash size was %d)',
1315 isolated_hash, algo(), algo().digest_size)
Takuto Ikuta6e2ff962019-10-29 12:35:27 +00001316 path = six.text_type(os.path.abspath(isolated_hash))
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001317 try:
1318 isolated_hash = fetch_queue.inject_local_file(path, algo)
1319 except IOError as e:
1320 raise isolated_format.MappingError(
1321 '%s doesn\'t seem to be a valid file. Did you intent to pass a '
1322 'valid hash (error: %s)?' % (isolated_hash, e))
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001323
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001324 # Load all *.isolated and start loading rest of the files.
1325 bundle.fetch(fetch_queue, isolated_hash, algo)
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001326
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001327 with tools.Profiler('GetRest'):
1328 # Create file system hierarchy.
1329 file_path.ensure_tree(outdir)
1330 create_directories(outdir, bundle.files)
Marc-Antoine Ruel04903a32019-10-09 21:09:25 +00001331 _create_symlinks(outdir, bundle.files.items())
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001332
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001333 # Ensure working directory exists.
1334 cwd = os.path.normpath(os.path.join(outdir, bundle.relative_cwd))
1335 file_path.ensure_tree(cwd)
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001336
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001337 # Multimap: digest -> list of pairs (path, props).
1338 remaining = {}
Marc-Antoine Ruel04903a32019-10-09 21:09:25 +00001339 for filepath, props in bundle.files.items():
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001340 if 'h' in props:
1341 remaining.setdefault(props['h'], []).append((filepath, props))
1342 fetch_queue.wait_on(props['h'])
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001343
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001344 # Now block on the remaining files to be downloaded and mapped.
1345 logging.info('Retrieving remaining files (%d of them)...',
1346 fetch_queue.pending_count)
1347 last_update = time.time()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001348
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001349 with threading_utils.ThreadPool(2, 32, 32) as putfile_thread_pool:
1350 with threading_utils.DeadlockDetector(DEADLOCK_TIMEOUT) as detector:
1351 while remaining:
1352 detector.ping()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001353
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001354 # Wait for any item to finish fetching to cache.
1355 digest = fetch_queue.wait()
tansell9e04a8d2016-07-28 09:31:59 -07001356
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001357 # Create the files in the destination using item in cache as the
1358 # source.
1359 for filepath, props in remaining.pop(digest):
1360 fullpath = os.path.join(outdir, filepath)
tanselle4288c32016-07-28 09:45:40 -07001361
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001362 putfile_thread_pool.add_task(threading_utils.PRIORITY_HIGH,
1363 _map_file, fullpath, digest,
1364 props, cache, bundle.read_only,
1365 use_symlinks)
tanselle4288c32016-07-28 09:45:40 -07001366
Takuto Ikutadeba39d2019-04-04 12:18:39 +00001367 # Report progress.
1368 duration = time.time() - last_update
1369 if duration > DELAY_BETWEEN_UPDATES_IN_SECS:
1370 msg = '%d files remaining...' % len(remaining)
1371 sys.stdout.write(msg + '\n')
1372 sys.stdout.flush()
1373 logging.info(msg)
1374 last_update = time.time()
1375 assert fetch_queue.wait_queue_empty, 'FetchQueue should have been emptied'
1376 putfile_thread_pool.join()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001377
Marc-Antoine Ruele9558372018-08-03 03:41:22 +00001378 # Save the cache right away to not loose the state of the new objects.
1379 cache.save()
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001380 # Cache could evict some items we just tried to fetch, it's a fatal error.
1381 if not fetch_queue.verify_all_cached():
Marc-Antoine Rueldddf6172018-01-23 14:25:43 -05001382 free_disk = file_path.get_free_space(cache.cache_dir)
1383 msg = (
1384 'Cache is too small to hold all requested files.\n'
1385 ' %s\n cache=%dbytes, %d items; %sb free_space') % (
Marc-Antoine Ruel5d7606b2018-06-15 19:06:12 +00001386 cache.policies, cache.total_size, len(cache), free_disk)
Marc-Antoine Rueldddf6172018-01-23 14:25:43 -05001387 raise isolated_format.MappingError(msg)
Vadim Shtayura7f7459c2014-09-04 13:25:10 -07001388 return bundle
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001389
1390
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +00001391def _directory_to_metadata(root, algo, blacklist):
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001392 """Yields every file and/or symlink found.
1393
1394 Yields:
1395 tuple(FileItem, relpath, metadata)
1396 For a symlink, FileItem is None.
1397 """
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001398 # Current tar file bundle, if any.
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001399 root = file_path.get_native_path_case(root)
Marc-Antoine Ruel440eee62018-12-04 22:37:05 +00001400 bundle = TarBundle(root, algo)
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001401 for relpath, issymlink in isolated_format.expand_directory_and_symlink(
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001402 root,
1403 u'.' + os.path.sep,
1404 blacklist,
1405 follow_symlinks=(sys.platform != 'win32')):
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001406
1407 filepath = os.path.join(root, relpath)
1408 if issymlink:
Marc-Antoine Ruel440eee62018-12-04 22:37:05 +00001409 # TODO(maruel): Do not call this.
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001410 meta = isolated_format.file_to_metadata(filepath, 0, False)
1411 yield None, relpath, meta
1412 continue
1413
1414 prio = relpath.endswith('.isolated')
Marc-Antoine Ruel440eee62018-12-04 22:37:05 +00001415 if bundle.try_add(FileItem(path=filepath, algo=algo, high_priority=prio)):
1416 # The file was added to the current pending tarball and won't be archived
1417 # individually.
1418 continue
1419
1420 # Flush and reset the bundle.
1421 for i, p, m in bundle.yield_item_path_meta():
1422 yield i, p, m
1423 bundle = TarBundle(root, algo)
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001424
1425 # Yield the file individually.
1426 item = FileItem(path=filepath, algo=algo, size=None, high_priority=prio)
1427 yield item, relpath, item.meta
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001428
Marc-Antoine Ruel440eee62018-12-04 22:37:05 +00001429 for i, p, m in bundle.yield_item_path_meta():
1430 yield i, p, m
1431
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001432
Marc-Antoine Ruelb69069b2018-11-28 20:50:40 +00001433def _print_upload_stats(items, missing):
1434 """Prints upload stats."""
1435 total = len(items)
1436 total_size = sum(f.size for f in items)
1437 logging.info(
1438 'Total: %6d, %9.1fkiB', total, total_size / 1024.)
1439 cache_hit = set(items).difference(missing)
1440 cache_hit_size = sum(f.size for f in cache_hit)
1441 logging.info(
1442 'cache hit: %6d, %9.1fkiB, %6.2f%% files, %6.2f%% size',
1443 len(cache_hit),
1444 cache_hit_size / 1024.,
1445 len(cache_hit) * 100. / total,
1446 cache_hit_size * 100. / total_size if total_size else 0)
1447 cache_miss = missing
1448 cache_miss_size = sum(f.size for f in cache_miss)
1449 logging.info(
1450 'cache miss: %6d, %9.1fkiB, %6.2f%% files, %6.2f%% size',
1451 len(cache_miss),
1452 cache_miss_size / 1024.,
1453 len(cache_miss) * 100. / total,
1454 cache_miss_size * 100. / total_size if total_size else 0)
1455
1456
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001457def _enqueue_dir(dirpath, blacklist, hash_algo, hash_algo_name):
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001458 """Called by archive_files_to_storage for a directory.
1459
1460 Create an .isolated file.
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001461
1462 Yields:
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001463 FileItem for every file found, plus one for the .isolated file itself.
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001464 """
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001465 files = {}
1466 for item, relpath, meta in _directory_to_metadata(
1467 dirpath, hash_algo, blacklist):
Marc-Antoine Ruel9cd5ef02018-11-29 23:47:34 +00001468 # item is None for a symlink.
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001469 files[relpath] = meta
Marc-Antoine Ruel9cd5ef02018-11-29 23:47:34 +00001470 if item:
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001471 yield item
1472
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001473 # TODO(maruel): If there' not file, don't yield an .isolated file.
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001474 data = {
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001475 'algo': hash_algo_name,
1476 'files': files,
1477 'version': isolated_format.ISOLATED_FILE_VERSION,
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001478 }
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001479 # Keep the file in memory. This is fine because .isolated files are relatively
1480 # small.
1481 yield BufferItem(
1482 tools.format_json(data, True), algo=hash_algo, high_priority=True)
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001483
1484
Vadim Shtayurae0ab1902014-04-29 10:55:27 -07001485def archive_files_to_storage(storage, files, blacklist):
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001486 """Stores every entry into remote storage and returns stats.
Marc-Antoine Ruel2283ad12014-02-09 11:14:57 -05001487
1488 Arguments:
1489 storage: a Storage object that communicates with the remote object store.
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001490 files: iterable of files to upload. If a directory is specified (with a
1491 trailing slash), a .isolated file is created and its hash is returned.
1492 Duplicates are skipped.
Marc-Antoine Ruel2283ad12014-02-09 11:14:57 -05001493 blacklist: function that returns True if a file should be omitted.
maruel064c0a32016-04-05 11:47:15 -07001494
1495 Returns:
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001496 tuple(OrderedDict(path: hash), list(FileItem cold), list(FileItem hot)).
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001497 The first file in the first item is always the .isolated file.
Marc-Antoine Ruel2283ad12014-02-09 11:14:57 -05001498 """
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001499 # Dict of path to hash.
1500 results = collections.OrderedDict()
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001501 hash_algo = storage.server_ref.hash_algo
1502 hash_algo_name = storage.server_ref.hash_algo_name
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001503 # Generator of FileItem to pass to upload_items() concurrent operation.
1504 channel = threading_utils.TaskChannel()
1505 uploaded_digests = set()
1506 def _upload_items():
1507 results = storage.upload_items(channel)
1508 uploaded_digests.update(f.digest for f in results)
1509 t = threading.Thread(target=_upload_items)
1510 t.start()
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001511
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001512 # Keep track locally of the items to determine cold and hot items.
1513 items_found = []
1514 try:
1515 for f in files:
Takuto Ikuta95459dd2019-10-29 12:39:47 +00001516 assert isinstance(f, six.text_type), repr(f)
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001517 if f in results:
1518 # Duplicate
1519 continue
1520 try:
1521 filepath = os.path.abspath(f)
1522 if fs.isdir(filepath):
1523 # Uploading a whole directory.
1524 item = None
1525 for item in _enqueue_dir(
1526 filepath, blacklist, hash_algo, hash_algo_name):
Marc-Antoine Ruelcc802b02018-11-28 21:05:01 +00001527 channel.send_result(item)
1528 items_found.append(item)
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001529 # The very last item will be the .isolated file.
1530 if not item:
1531 # There was no file in the directory.
1532 continue
1533 elif fs.isfile(filepath):
1534 item = FileItem(
1535 path=filepath,
1536 algo=hash_algo,
1537 size=None,
1538 high_priority=f.endswith('.isolated'))
1539 channel.send_result(item)
1540 items_found.append(item)
1541 else:
1542 raise Error('%s is neither a file or directory.' % f)
1543 results[f] = item.digest
1544 except OSError:
1545 raise Error('Failed to process %s.' % f)
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001546 finally:
Marc-Antoine Ruel1b2885d2018-12-04 18:30:33 +00001547 # Stops the generator, so _upload_items() can exit.
1548 channel.send_done()
1549 t.join()
1550
1551 cold = []
1552 hot = []
1553 for i in items_found:
1554 # Note that multiple FileItem may have the same .digest.
1555 if i.digest in uploaded_digests:
1556 cold.append(i)
1557 else:
1558 hot.append(i)
1559 return results, cold, hot
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001560
1561
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001562@subcommand.usage('<file1..fileN> or - to read from stdin')
1563def CMDarchive(parser, args):
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001564 """Archives data to the server.
1565
1566 If a directory is specified, a .isolated file is created the whole directory
1567 is uploaded. Then this .isolated file can be included in another one to run
1568 commands.
1569
1570 The commands output each file that was processed with its content hash. For
1571 directories, the .isolated generated for the directory is listed as the
1572 directory entry itself.
1573 """
Marc-Antoine Ruelf7d737d2014-12-10 15:36:29 -05001574 add_isolate_server_options(parser)
Marc-Antoine Ruel1f8ba352014-11-04 15:55:03 -05001575 add_archive_options(parser)
maruel@chromium.orgcb3c3d52013-03-14 18:55:30 +00001576 options, files = parser.parse_args(args)
nodir55be77b2016-05-03 09:39:57 -07001577 process_isolate_server_options(parser, options, True, True)
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +00001578 server_ref = isolate_storage.ServerRef(
1579 options.isolate_server, options.namespace)
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001580 if files == ['-']:
1581 files = (l.rstrip('\n\r') for l in sys.stdin)
1582 if not files:
1583 parser.error('Nothing to upload')
1584 files = (f.decode('utf-8') for f in files)
1585 blacklist = tools.gen_blacklist(options.blacklist)
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001586 try:
Marc-Antoine Rueld0868ec2018-11-28 20:47:29 +00001587 with get_storage(server_ref) as storage:
1588 results, _cold, _hot = archive_files_to_storage(storage, files, blacklist)
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -04001589 except (Error, local_caching.NoMoreSpace) as e:
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001590 parser.error(e.args[0])
Marc-Antoine Ruel04903a32019-10-09 21:09:25 +00001591 print('\n'.join('%s %s' % (h, f) for f, h in results.items()))
Marc-Antoine Ruelfcc3cd82013-11-19 16:31:38 -05001592 return 0
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001593
1594
1595def CMDdownload(parser, args):
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +00001596 """Download data from the server.
1597
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001598 It can either download individual files or a complete tree from a .isolated
1599 file.
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +00001600 """
Marc-Antoine Ruelf7d737d2014-12-10 15:36:29 -05001601 add_isolate_server_options(parser)
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +00001602 parser.add_option(
Marc-Antoine Ruel185ded42015-01-28 20:49:18 -05001603 '-s', '--isolated', metavar='HASH',
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001604 help='hash of an isolated file, .isolated file content is discarded, use '
1605 '--file if you need it')
1606 parser.add_option(
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +00001607 '-f', '--file', metavar='HASH DEST', default=[], action='append', nargs=2,
1608 help='hash and destination of a file, can be used multiple times')
1609 parser.add_option(
Marc-Antoine Ruelf90861c2015-03-24 20:54:49 -04001610 '-t', '--target', metavar='DIR', default='download',
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +00001611 help='destination directory')
maruel4409e302016-07-19 14:25:51 -07001612 parser.add_option(
1613 '--use-symlinks', action='store_true',
1614 help='Use symlinks instead of hardlinks')
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001615 add_cache_options(parser)
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +00001616 options, args = parser.parse_args(args)
1617 if args:
1618 parser.error('Unsupported arguments: %s' % args)
Marc-Antoine Ruel5028ba22017-08-25 17:37:51 -04001619 if not file_path.enable_symlink():
Marc-Antoine Ruel5a024272019-01-15 20:11:16 +00001620 logging.warning('Symlink support is not enabled')
Marc-Antoine Ruelf7d737d2014-12-10 15:36:29 -05001621
nodir55be77b2016-05-03 09:39:57 -07001622 process_isolate_server_options(parser, options, True, True)
maruel@chromium.org4f2ebe42013-09-19 13:09:08 +00001623 if bool(options.isolated) == bool(options.file):
1624 parser.error('Use one of --isolated or --file, and only one.')
maruel4409e302016-07-19 14:25:51 -07001625 if not options.cache and options.use_symlinks:
1626 parser.error('--use-symlinks require the use of a cache with --cache')
maruel@chromium.orgb7e79a22013-09-13 01:24:56 +00001627
John Abd-El-Maleke3a85012018-05-29 20:10:44 -07001628 cache = process_cache_options(options, trim=True)
maruel2e8d0f52016-07-16 07:51:29 -07001629 cache.cleanup()
Takuto Ikuta6e2ff962019-10-29 12:35:27 +00001630 options.target = six.text_type(os.path.abspath(options.target))
Marc-Antoine Ruelf90861c2015-03-24 20:54:49 -04001631 if options.isolated:
maruel12e30012015-10-09 11:55:35 -07001632 if (fs.isfile(options.target) or
1633 (fs.isdir(options.target) and fs.listdir(options.target))):
Marc-Antoine Ruelf90861c2015-03-24 20:54:49 -04001634 parser.error(
1635 '--target \'%s\' exists, please use another target' % options.target)
Marc-Antoine Ruelb8513132018-11-20 19:48:53 +00001636 server_ref = isolate_storage.ServerRef(
1637 options.isolate_server, options.namespace)
1638 with get_storage(server_ref) as storage:
Vadim Shtayura3172be52013-12-03 12:49:05 -08001639 # Fetching individual files.
1640 if options.file:
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001641 # TODO(maruel): Enable cache in this case too.
Vadim Shtayura3172be52013-12-03 12:49:05 -08001642 channel = threading_utils.TaskChannel()
1643 pending = {}
1644 for digest, dest in options.file:
Takuto Ikuta6e2ff962019-10-29 12:35:27 +00001645 dest = six.text_type(dest)
Vadim Shtayura3172be52013-12-03 12:49:05 -08001646 pending[digest] = dest
1647 storage.async_fetch(
1648 channel,
Vadim Shtayura3148e072014-09-02 18:51:52 -07001649 threading_utils.PRIORITY_MED,
Vadim Shtayura3172be52013-12-03 12:49:05 -08001650 digest,
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -04001651 local_caching.UNKNOWN_FILE_SIZE,
1652 functools.partial(
1653 local_caching.file_write, os.path.join(options.target, dest)))
Vadim Shtayura3172be52013-12-03 12:49:05 -08001654 while pending:
Marc-Antoine Ruel4494b6c2018-11-28 21:00:41 +00001655 fetched = channel.next()
Vadim Shtayura3172be52013-12-03 12:49:05 -08001656 dest = pending.pop(fetched)
1657 logging.info('%s: %s', fetched, dest)
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001658
Vadim Shtayura3172be52013-12-03 12:49:05 -08001659 # Fetching whole isolated tree.
1660 if options.isolated:
Marc-Antoine Ruele79ddbf2018-06-13 18:33:07 +00001661 bundle = fetch_isolated(
1662 isolated_hash=options.isolated,
1663 storage=storage,
1664 cache=cache,
1665 outdir=options.target,
1666 use_symlinks=options.use_symlinks)
1667 cache.trim()
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001668 if bundle.command:
1669 rel = os.path.join(options.target, bundle.relative_cwd)
1670 print('To run this test please run from the directory %s:' %
1671 os.path.join(options.target, rel))
1672 print(' ' + ' '.join(bundle.command))
vadimsh@chromium.org7b5dae32013-10-03 16:59:59 +00001673
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001674 return 0
1675
1676
Marc-Antoine Ruel1f8ba352014-11-04 15:55:03 -05001677def add_archive_options(parser):
1678 parser.add_option(
1679 '--blacklist',
1680 action='append', default=list(DEFAULT_BLACKLIST),
1681 help='List of regexp to use as blacklist filter when uploading '
1682 'directories')
1683
1684
Marc-Antoine Ruelf7d737d2014-12-10 15:36:29 -05001685def add_isolate_server_options(parser):
1686 """Adds --isolate-server and --namespace options to parser."""
Marc-Antoine Ruel1687b5e2014-02-06 17:47:53 -05001687 parser.add_option(
1688 '-I', '--isolate-server',
1689 metavar='URL', default=os.environ.get('ISOLATE_SERVER', ''),
Marc-Antoine Ruel8806e622014-02-12 14:15:53 -05001690 help='URL of the Isolate Server to use. Defaults to the environment '
1691 'variable ISOLATE_SERVER if set. No need to specify https://, this '
1692 'is assumed.')
Marc-Antoine Ruel1687b5e2014-02-06 17:47:53 -05001693 parser.add_option(
aludwind7b7b7e2017-06-29 16:38:50 -07001694 '--grpc-proxy', help='gRPC proxy by which to communicate to Isolate')
aludwin81178302016-11-30 17:18:49 -08001695 parser.add_option(
Marc-Antoine Ruel1687b5e2014-02-06 17:47:53 -05001696 '--namespace', default='default-gzip',
1697 help='The namespace to use on the Isolate Server, default: %default')
1698
1699
nodir55be77b2016-05-03 09:39:57 -07001700def process_isolate_server_options(
1701 parser, options, set_exception_handler, required):
1702 """Processes the --isolate-server option.
Marc-Antoine Ruelf7d737d2014-12-10 15:36:29 -05001703
1704 Returns the identity as determined by the server.
Marc-Antoine Ruel1687b5e2014-02-06 17:47:53 -05001705 """
1706 if not options.isolate_server:
nodir55be77b2016-05-03 09:39:57 -07001707 if required:
1708 parser.error('--isolate-server is required.')
1709 return
1710
aludwind7b7b7e2017-06-29 16:38:50 -07001711 if options.grpc_proxy:
1712 isolate_storage.set_grpc_proxy(options.grpc_proxy)
aludwin81178302016-11-30 17:18:49 -08001713 else:
1714 try:
1715 options.isolate_server = net.fix_url(options.isolate_server)
1716 except ValueError as e:
1717 parser.error('--isolate-server %s' % e)
Marc-Antoine Ruele290ada2014-12-10 19:48:49 -05001718 if set_exception_handler:
1719 on_error.report_on_exception_exit(options.isolate_server)
Marc-Antoine Ruelf7d737d2014-12-10 15:36:29 -05001720 try:
1721 return auth.ensure_logged_in(options.isolate_server)
1722 except ValueError as e:
1723 parser.error(str(e))
Marc-Antoine Ruel793bff32019-04-18 17:50:48 +00001724 return None
Marc-Antoine Ruel8806e622014-02-12 14:15:53 -05001725
Marc-Antoine Ruel1687b5e2014-02-06 17:47:53 -05001726
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001727def add_cache_options(parser):
1728 cache_group = optparse.OptionGroup(parser, 'Cache management')
1729 cache_group.add_option(
Marc-Antoine Ruel5aeb3bb2018-06-16 13:11:02 +00001730 '--cache', metavar='DIR', default='cache',
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001731 help='Directory to keep a local cache of the files. Accelerates download '
1732 'by reusing already downloaded files. Default=%default')
1733 cache_group.add_option(
1734 '--max-cache-size',
1735 type='int',
1736 metavar='NNN',
maruel71586102016-01-29 11:44:09 -08001737 default=50*1024*1024*1024,
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001738 help='Trim if the cache gets larger than this value, default=%default')
1739 cache_group.add_option(
1740 '--min-free-space',
1741 type='int',
1742 metavar='NNN',
1743 default=2*1024*1024*1024,
1744 help='Trim if disk free space becomes lower than this value, '
1745 'default=%default')
1746 cache_group.add_option(
1747 '--max-items',
1748 type='int',
1749 metavar='NNN',
1750 default=100000,
1751 help='Trim if more than this number of items are in the cache '
1752 'default=%default')
1753 parser.add_option_group(cache_group)
1754
1755
John Abd-El-Maleke3a85012018-05-29 20:10:44 -07001756def process_cache_options(options, trim, **kwargs):
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001757 if options.cache:
Marc-Antoine Ruel34f5f282018-05-16 16:04:31 -04001758 policies = local_caching.CachePolicies(
1759 options.max_cache_size,
1760 options.min_free_space,
1761 options.max_items,
1762 # 3 weeks.
1763 max_age_secs=21*24*60*60)
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001764
Marc-Antoine Ruel2666d9c2018-05-18 13:52:02 -04001765 # |options.cache| path may not exist until DiskContentAddressedCache()
1766 # instance is created.
1767 return local_caching.DiskContentAddressedCache(
Takuto Ikuta6e2ff962019-10-29 12:35:27 +00001768 six.text_type(os.path.abspath(options.cache)), policies, trim, **kwargs)
Marc-Antoine Ruel793bff32019-04-18 17:50:48 +00001769 return local_caching.MemoryContentAddressedCache()
Marc-Antoine Ruela57d7db2014-10-15 20:31:19 -04001770
1771
Marc-Antoine Ruelf74cffe2015-07-15 15:21:34 -04001772class OptionParserIsolateServer(logging_utils.OptionParserWithLogging):
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001773 def __init__(self, **kwargs):
Marc-Antoine Ruelf74cffe2015-07-15 15:21:34 -04001774 logging_utils.OptionParserWithLogging.__init__(
Marc-Antoine Ruelac54cb42013-11-18 14:05:35 -05001775 self,
1776 version=__version__,
1777 prog=os.path.basename(sys.modules[__name__].__file__),
1778 **kwargs)
Vadim Shtayurae34e13a2014-02-02 11:23:26 -08001779 auth.add_auth_options(self)
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001780
1781 def parse_args(self, *args, **kwargs):
Marc-Antoine Ruelf74cffe2015-07-15 15:21:34 -04001782 options, args = logging_utils.OptionParserWithLogging.parse_args(
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001783 self, *args, **kwargs)
Vadim Shtayura5d1efce2014-02-04 10:55:43 -08001784 auth.process_auth_options(self, options)
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001785 return options, args
1786
1787
1788def main(args):
1789 dispatcher = subcommand.CommandDispatcher(__name__)
Marc-Antoine Ruelcfb60852014-07-02 15:22:00 -04001790 return dispatcher.execute(OptionParserIsolateServer(), args)
maruel@chromium.orgc6f90062012-11-07 18:32:22 +00001791
1792
1793if __name__ == '__main__':
maruel8e4e40c2016-05-30 06:21:07 -07001794 subprocess42.inhibit_os_error_reporting()
maruel@chromium.orgfb78d432013-08-28 21:22:40 +00001795 fix_encoding.fix_encoding()
1796 tools.disable_buffering()
1797 colorama.init()
maruel@chromium.orgcb3c3d52013-03-14 18:55:30 +00001798 sys.exit(main(sys.argv[1:]))