maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 1 | #!/usr/bin/env python |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 2 | # Copyright 2013 The Chromium Authors. All rights reserved. |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 3 | # Use of this source code is governed by a BSD-style license that can be |
| 4 | # found in the LICENSE file. |
| 5 | |
| 6 | """Archives a set of files to a server.""" |
| 7 | |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 8 | __version__ = '0.2' |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 9 | |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 10 | import binascii |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 11 | import cStringIO |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 12 | import hashlib |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 13 | import itertools |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 14 | import logging |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 15 | import os |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 16 | import random |
| 17 | import re |
| 18 | import shutil |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 19 | import sys |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 20 | import threading |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 21 | import time |
maruel@chromium.org | e82112e | 2013-04-24 14:41:55 +0000 | [diff] [blame] | 22 | import urllib |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 23 | import zlib |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 24 | |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 25 | from third_party import colorama |
| 26 | from third_party.depot_tools import fix_encoding |
| 27 | from third_party.depot_tools import subcommand |
| 28 | |
vadimsh@chromium.org | 6b70621 | 2013-08-28 15:03:46 +0000 | [diff] [blame] | 29 | from utils import net |
vadimsh@chromium.org | b074b16 | 2013-08-22 17:55:46 +0000 | [diff] [blame] | 30 | from utils import threading_utils |
vadimsh@chromium.org | a432647 | 2013-08-24 02:05:41 +0000 | [diff] [blame] | 31 | from utils import tools |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 32 | |
| 33 | |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 34 | # Default server. |
| 35 | # TODO(maruel): Chromium-specific. |
| 36 | ISOLATE_SERVER = 'https://isolateserver-dev.appspot.com/' |
| 37 | |
| 38 | |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 39 | # The minimum size of files to upload directly to the blobstore. |
maruel@chromium.org | aef29f8 | 2012-12-12 15:00:42 +0000 | [diff] [blame] | 40 | MIN_SIZE_FOR_DIRECT_BLOBSTORE = 20 * 1024 |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 41 | |
vadimsh@chromium.org | eea5242 | 2013-08-21 19:35:54 +0000 | [diff] [blame] | 42 | # The number of files to check the isolate server per /contains query. |
| 43 | # All files are sorted by likelihood of a change in the file content |
| 44 | # (currently file size is used to estimate this: larger the file -> larger the |
| 45 | # possibility it has changed). Then first ITEMS_PER_CONTAINS_QUERIES[0] files |
| 46 | # are taken and send to '/contains', then next ITEMS_PER_CONTAINS_QUERIES[1], |
| 47 | # and so on. Numbers here is a trade-off; the more per request, the lower the |
| 48 | # effect of HTTP round trip latency and TCP-level chattiness. On the other hand, |
| 49 | # larger values cause longer lookups, increasing the initial latency to start |
| 50 | # uploading, which is especially an issue for large files. This value is |
| 51 | # optimized for the "few thousands files to look up with minimal number of large |
| 52 | # files missing" case. |
| 53 | ITEMS_PER_CONTAINS_QUERIES = [20, 20, 50, 50, 50, 100] |
csharp@chromium.org | 07fa759 | 2013-01-11 18:19:30 +0000 | [diff] [blame] | 54 | |
maruel@chromium.org | 9958e4a | 2013-09-17 00:01:48 +0000 | [diff] [blame] | 55 | |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 56 | # A list of already compressed extension types that should not receive any |
| 57 | # compression before being uploaded. |
| 58 | ALREADY_COMPRESSED_TYPES = [ |
| 59 | '7z', 'avi', 'cur', 'gif', 'h264', 'jar', 'jpeg', 'jpg', 'pdf', 'png', |
| 60 | 'wav', 'zip' |
| 61 | ] |
| 62 | |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 63 | |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 64 | # The file size to be used when we don't know the correct file size, |
| 65 | # generally used for .isolated files. |
| 66 | UNKNOWN_FILE_SIZE = None |
| 67 | |
| 68 | |
| 69 | # The size of each chunk to read when downloading and unzipping files. |
| 70 | ZIPPED_FILE_CHUNK = 16 * 1024 |
| 71 | |
| 72 | |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 73 | # Read timeout in seconds for downloads from isolate storage. If there's no |
| 74 | # response from the server within this timeout whole download will be aborted. |
| 75 | DOWNLOAD_READ_TIMEOUT = 60 |
| 76 | |
| 77 | |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 78 | class ConfigError(ValueError): |
| 79 | """Generic failure to load a .isolated file.""" |
| 80 | pass |
| 81 | |
| 82 | |
| 83 | class MappingError(OSError): |
| 84 | """Failed to recreate the tree.""" |
| 85 | pass |
| 86 | |
| 87 | |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 88 | def randomness(): |
| 89 | """Generates low-entropy randomness for MIME encoding. |
| 90 | |
| 91 | Exists so it can be mocked out in unit tests. |
| 92 | """ |
| 93 | return str(time.time()) |
| 94 | |
| 95 | |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 96 | def encode_multipart_formdata(fields, files, |
| 97 | mime_mapper=lambda _: 'application/octet-stream'): |
| 98 | """Encodes a Multipart form data object. |
| 99 | |
| 100 | Args: |
| 101 | fields: a sequence (name, value) elements for |
| 102 | regular form fields. |
| 103 | files: a sequence of (name, filename, value) elements for data to be |
| 104 | uploaded as files. |
| 105 | mime_mapper: function to return the mime type from the filename. |
| 106 | Returns: |
| 107 | content_type: for httplib.HTTP instance |
| 108 | body: for httplib.HTTP instance |
| 109 | """ |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 110 | boundary = hashlib.md5(randomness()).hexdigest() |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 111 | body_list = [] |
| 112 | for (key, value) in fields: |
| 113 | if isinstance(key, unicode): |
| 114 | value = key.encode('utf-8') |
| 115 | if isinstance(value, unicode): |
| 116 | value = value.encode('utf-8') |
| 117 | body_list.append('--' + boundary) |
| 118 | body_list.append('Content-Disposition: form-data; name="%s"' % key) |
| 119 | body_list.append('') |
| 120 | body_list.append(value) |
| 121 | body_list.append('--' + boundary) |
| 122 | body_list.append('') |
| 123 | for (key, filename, value) in files: |
| 124 | if isinstance(key, unicode): |
| 125 | value = key.encode('utf-8') |
| 126 | if isinstance(filename, unicode): |
| 127 | value = filename.encode('utf-8') |
| 128 | if isinstance(value, unicode): |
| 129 | value = value.encode('utf-8') |
| 130 | body_list.append('--' + boundary) |
| 131 | body_list.append('Content-Disposition: form-data; name="%s"; ' |
| 132 | 'filename="%s"' % (key, filename)) |
| 133 | body_list.append('Content-Type: %s' % mime_mapper(filename)) |
| 134 | body_list.append('') |
| 135 | body_list.append(value) |
| 136 | body_list.append('--' + boundary) |
| 137 | body_list.append('') |
| 138 | if body_list: |
| 139 | body_list[-2] += '--' |
| 140 | body = '\r\n'.join(body_list) |
| 141 | content_type = 'multipart/form-data; boundary=%s' % boundary |
| 142 | return content_type, body |
| 143 | |
| 144 | |
maruel@chromium.org | 7b844a6 | 2013-09-17 13:04:59 +0000 | [diff] [blame] | 145 | def is_valid_hash(value, algo): |
| 146 | """Returns if the value is a valid hash for the corresponding algorithm.""" |
| 147 | size = 2 * algo().digest_size |
| 148 | return bool(re.match(r'^[a-fA-F0-9]{%d}$' % size, value)) |
| 149 | |
| 150 | |
| 151 | def hash_file(filepath, algo): |
| 152 | """Calculates the hash of a file without reading it all in memory at once. |
| 153 | |
| 154 | |algo| should be one of hashlib hashing algorithm. |
| 155 | """ |
| 156 | digest = algo() |
maruel@chromium.org | 037758d | 2012-12-10 17:59:46 +0000 | [diff] [blame] | 157 | with open(filepath, 'rb') as f: |
| 158 | while True: |
| 159 | # Read in 1mb chunks. |
| 160 | chunk = f.read(1024*1024) |
| 161 | if not chunk: |
| 162 | break |
| 163 | digest.update(chunk) |
| 164 | return digest.hexdigest() |
| 165 | |
| 166 | |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 167 | def file_write(filepath, content_generator): |
| 168 | """Writes file content as generated by content_generator. |
| 169 | |
| 170 | Meant to be mocked out in unit tests. |
| 171 | """ |
| 172 | filedir = os.path.dirname(filepath) |
| 173 | if not os.path.isdir(filedir): |
| 174 | os.makedirs(filedir) |
| 175 | with open(filepath, 'wb') as f: |
| 176 | for d in content_generator: |
| 177 | f.write(d) |
| 178 | |
| 179 | |
maruel@chromium.org | af25485 | 2013-09-17 17:48:14 +0000 | [diff] [blame^] | 180 | def create_directories(base_directory, files): |
| 181 | """Creates the directory structure needed by the given list of files.""" |
| 182 | logging.debug('create_directories(%s, %d)', base_directory, len(files)) |
| 183 | # Creates the tree of directories to create. |
| 184 | directories = set(os.path.dirname(f) for f in files) |
| 185 | for item in list(directories): |
| 186 | while item: |
| 187 | directories.add(item) |
| 188 | item = os.path.dirname(item) |
| 189 | for d in sorted(directories): |
| 190 | if d: |
| 191 | os.mkdir(os.path.join(base_directory, d)) |
| 192 | |
| 193 | |
| 194 | def create_links(base_directory, files): |
| 195 | """Creates any links needed by the given set of files.""" |
| 196 | for filepath, properties in files: |
| 197 | if 'l' not in properties: |
| 198 | continue |
| 199 | if sys.platform == 'win32': |
| 200 | # TODO(maruel): Create junctions or empty text files similar to what |
| 201 | # cygwin do? |
| 202 | logging.warning('Ignoring symlink %s', filepath) |
| 203 | continue |
| 204 | outfile = os.path.join(base_directory, filepath) |
| 205 | # symlink doesn't exist on Windows. So the 'link' property should |
| 206 | # never be specified for windows .isolated file. |
| 207 | os.symlink(properties['l'], outfile) # pylint: disable=E1101 |
| 208 | if 'm' in properties: |
| 209 | lchmod = getattr(os, 'lchmod', None) |
| 210 | if lchmod: |
| 211 | lchmod(outfile, properties['m']) |
| 212 | |
| 213 | |
| 214 | def setup_commands(base_directory, cwd, cmd): |
| 215 | """Correctly adjusts and then returns the required working directory |
| 216 | and command needed to run the test. |
| 217 | """ |
| 218 | assert not os.path.isabs(cwd), 'The cwd must be a relative path, got %s' % cwd |
| 219 | cwd = os.path.join(base_directory, cwd) |
| 220 | if not os.path.isdir(cwd): |
| 221 | os.makedirs(cwd) |
| 222 | |
| 223 | # Ensure paths are correctly separated on windows. |
| 224 | cmd[0] = cmd[0].replace('/', os.path.sep) |
| 225 | cmd = tools.fix_python_path(cmd) |
| 226 | |
| 227 | return cwd, cmd |
| 228 | |
| 229 | |
| 230 | def generate_remaining_files(files): |
| 231 | """Generates a dictionary of all the remaining files to be downloaded.""" |
| 232 | remaining = {} |
| 233 | for filepath, props in files: |
| 234 | if 'h' in props: |
| 235 | remaining.setdefault(props['h'], []).append((filepath, props)) |
| 236 | |
| 237 | return remaining |
| 238 | |
| 239 | |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 240 | def is_valid_file(filepath, size): |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 241 | """Determines if the given files appears valid. |
| 242 | |
| 243 | Currently it just checks the file's size. |
| 244 | """ |
| 245 | if size == UNKNOWN_FILE_SIZE: |
| 246 | return True |
| 247 | actual_size = os.stat(filepath).st_size |
| 248 | if size != actual_size: |
| 249 | logging.warning( |
| 250 | 'Found invalid item %s; %d != %d', |
| 251 | os.path.basename(filepath), actual_size, size) |
| 252 | return False |
| 253 | return True |
| 254 | |
| 255 | |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 256 | def try_remove(filepath): |
| 257 | """Removes a file without crashing even if it doesn't exist.""" |
| 258 | try: |
| 259 | os.remove(filepath) |
| 260 | except OSError: |
| 261 | pass |
| 262 | |
| 263 | |
vadimsh@chromium.org | 80f7300 | 2013-07-12 14:52:44 +0000 | [diff] [blame] | 264 | def url_read(url, **kwargs): |
vadimsh@chromium.org | 6b70621 | 2013-08-28 15:03:46 +0000 | [diff] [blame] | 265 | result = net.url_read(url, **kwargs) |
vadimsh@chromium.org | 80f7300 | 2013-07-12 14:52:44 +0000 | [diff] [blame] | 266 | if result is None: |
maruel@chromium.org | ef33312 | 2013-03-12 20:36:40 +0000 | [diff] [blame] | 267 | # If we get no response from the server, assume it is down and raise an |
| 268 | # exception. |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 269 | raise MappingError('Unable to connect to server %s' % url) |
maruel@chromium.org | ef33312 | 2013-03-12 20:36:40 +0000 | [diff] [blame] | 270 | return result |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 271 | |
| 272 | |
maruel@chromium.org | dc359e6 | 2013-03-14 13:08:55 +0000 | [diff] [blame] | 273 | def upload_hash_content_to_blobstore( |
| 274 | generate_upload_url, data, hash_key, content): |
vadimsh@chromium.org | 80f7300 | 2013-07-12 14:52:44 +0000 | [diff] [blame] | 275 | """Uploads the given hash contents directly to the blobstore via a generated |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 276 | url. |
| 277 | |
| 278 | Arguments: |
| 279 | generate_upload_url: The url to get the new upload url from. |
maruel@chromium.org | dc359e6 | 2013-03-14 13:08:55 +0000 | [diff] [blame] | 280 | data: extra POST data. |
maruel@chromium.org | 7b844a6 | 2013-09-17 13:04:59 +0000 | [diff] [blame] | 281 | hash_key: hash of the uncompressed version of content. |
maruel@chromium.org | dc359e6 | 2013-03-14 13:08:55 +0000 | [diff] [blame] | 282 | content: The contents to upload. Must fit in memory for now. |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 283 | """ |
| 284 | logging.debug('Generating url to directly upload file to blobstore') |
maruel@chromium.org | 92a3d2e | 2012-12-20 16:22:29 +0000 | [diff] [blame] | 285 | assert isinstance(hash_key, str), hash_key |
| 286 | assert isinstance(content, str), (hash_key, content) |
maruel@chromium.org | d58bf5b | 2013-04-26 17:57:42 +0000 | [diff] [blame] | 287 | # TODO(maruel): Support large files. This would require streaming support. |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 288 | content_type, body = encode_multipart_formdata( |
maruel@chromium.org | d58bf5b | 2013-04-26 17:57:42 +0000 | [diff] [blame] | 289 | data, [('content', hash_key, content)]) |
vadimsh@chromium.org | 043b76d | 2013-09-12 16:15:13 +0000 | [diff] [blame] | 290 | for _ in net.retry_loop(max_attempts=net.URL_OPEN_MAX_ATTEMPTS): |
maruel@chromium.org | d58bf5b | 2013-04-26 17:57:42 +0000 | [diff] [blame] | 291 | # Retry HTTP 50x here. |
vadimsh@chromium.org | 6b70621 | 2013-08-28 15:03:46 +0000 | [diff] [blame] | 292 | upload_url = net.url_read(generate_upload_url, data=data) |
vadimsh@chromium.org | 80f7300 | 2013-07-12 14:52:44 +0000 | [diff] [blame] | 293 | if not upload_url: |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 294 | raise MappingError( |
maruel@chromium.org | d58bf5b | 2013-04-26 17:57:42 +0000 | [diff] [blame] | 295 | 'Unable to connect to server %s' % generate_upload_url) |
maruel@chromium.org | d58bf5b | 2013-04-26 17:57:42 +0000 | [diff] [blame] | 296 | |
| 297 | # Do not retry this request on HTTP 50x. Regenerate an upload url each time |
| 298 | # since uploading "consumes" the upload url. |
vadimsh@chromium.org | 6b70621 | 2013-08-28 15:03:46 +0000 | [diff] [blame] | 299 | result = net.url_read( |
maruel@chromium.org | d58bf5b | 2013-04-26 17:57:42 +0000 | [diff] [blame] | 300 | upload_url, data=body, content_type=content_type, retry_50x=False) |
vadimsh@chromium.org | 80f7300 | 2013-07-12 14:52:44 +0000 | [diff] [blame] | 301 | if result is not None: |
| 302 | return result |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 303 | raise MappingError( |
maruel@chromium.org | d58bf5b | 2013-04-26 17:57:42 +0000 | [diff] [blame] | 304 | 'Unable to connect to server %s' % generate_upload_url) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 305 | |
| 306 | |
maruel@chromium.org | 3e42ce8 | 2013-09-12 18:36:59 +0000 | [diff] [blame] | 307 | class IsolateServer(object): |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 308 | """Client class to download or upload to Isolate Server.""" |
maruel@chromium.org | 3e42ce8 | 2013-09-12 18:36:59 +0000 | [diff] [blame] | 309 | def __init__(self, base_url, namespace): |
| 310 | assert base_url.startswith('http'), base_url |
| 311 | self.content_url = base_url.rstrip('/') + '/content/' |
| 312 | self.namespace = namespace |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 313 | self._token = None |
| 314 | self._lock = threading.Lock() |
| 315 | |
| 316 | @property |
| 317 | def token(self): |
maruel@chromium.org | 3e42ce8 | 2013-09-12 18:36:59 +0000 | [diff] [blame] | 318 | # TODO(maruel): Make this request much earlier asynchronously while the |
| 319 | # files are being enumerated. |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 320 | with self._lock: |
| 321 | if not self._token: |
| 322 | self._token = urllib.quote(url_read(self.content_url + 'get_token')) |
| 323 | return self._token |
| 324 | |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 325 | def fetch(self, item, size): |
| 326 | """Fetches an object and yields its content.""" |
| 327 | zipped_url = '%sretrieve/%s/%s' % (self.content_url, self.namespace, item) |
| 328 | logging.debug('download_file(%s)', zipped_url) |
| 329 | |
| 330 | # Because the app engine DB is only eventually consistent, retry 404 errors |
| 331 | # because the file might just not be visible yet (even though it has been |
| 332 | # uploaded). |
| 333 | connection = net.url_open( |
| 334 | zipped_url, retry_404=True, read_timeout=DOWNLOAD_READ_TIMEOUT) |
| 335 | if not connection: |
| 336 | raise IOError('Unable to open connection to %s' % zipped_url) |
| 337 | |
| 338 | decompressor = zlib.decompressobj() |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 339 | try: |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 340 | compressed_size = 0 |
| 341 | decompressed_size = 0 |
| 342 | while True: |
| 343 | chunk = connection.read(ZIPPED_FILE_CHUNK) |
| 344 | if not chunk: |
| 345 | break |
| 346 | compressed_size += len(chunk) |
| 347 | decompressed = decompressor.decompress(chunk) |
| 348 | decompressed_size += len(decompressed) |
| 349 | yield decompressed |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 350 | |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 351 | # Ensure that all the data was properly decompressed. |
| 352 | uncompressed_data = decompressor.flush() |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 353 | if uncompressed_data: |
| 354 | raise IOError('Decompression failed') |
| 355 | if size != UNKNOWN_FILE_SIZE and decompressed_size != size: |
| 356 | raise IOError('File incorrect size after download of %s. Got %s and ' |
| 357 | 'expected %s' % (item, decompressed_size, size)) |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 358 | except zlib.error as e: |
| 359 | msg = 'Corrupted zlib for item %s. Processed %d of %s bytes.\n%s' % ( |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 360 | item, compressed_size, connection.content_length, e) |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 361 | logging.error(msg) |
| 362 | |
| 363 | # Testing seems to show that if a few machines are trying to download |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 364 | # the same blob, they can cause each other to fail. So if we hit a zip |
| 365 | # error, this is the most likely cause (it only downloads some of the |
| 366 | # data). Randomly sleep for between 5 and 25 seconds to try and spread |
| 367 | # out the downloads. |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 368 | sleep_duration = (random.random() * 20) + 5 |
| 369 | time.sleep(sleep_duration) |
| 370 | raise IOError(msg) |
maruel@chromium.org | c2bfef4 | 2013-08-30 21:46:26 +0000 | [diff] [blame] | 371 | |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 372 | def retrieve(self, item, dest, size): |
| 373 | """Fetches an object and save its content to |dest|.""" |
| 374 | try: |
| 375 | file_write(dest, self.fetch(item, size)) |
| 376 | except IOError as e: |
| 377 | # Remove unfinished download. |
| 378 | try_remove(dest) |
| 379 | logging.error('Failed to download %s at %s.\n%s', item, dest, e) |
| 380 | raise |
| 381 | |
| 382 | def store(self, content, hash_key, _size): |
maruel@chromium.org | 3e42ce8 | 2013-09-12 18:36:59 +0000 | [diff] [blame] | 383 | # TODO(maruel): Detect failures. |
| 384 | hash_key = str(hash_key) |
| 385 | if len(content) > MIN_SIZE_FOR_DIRECT_BLOBSTORE: |
| 386 | url = '%sgenerate_blobstore_url/%s/%s' % ( |
| 387 | self.content_url, self.namespace, hash_key) |
| 388 | # token is guaranteed to be already quoted but it is unnecessary here, and |
| 389 | # only here. |
| 390 | data = [('token', urllib.unquote(self.token))] |
| 391 | return upload_hash_content_to_blobstore(url, data, hash_key, content) |
| 392 | else: |
| 393 | url = '%sstore/%s/%s?token=%s' % ( |
| 394 | self.content_url, self.namespace, hash_key, self.token) |
| 395 | return url_read( |
| 396 | url, data=content, content_type='application/octet-stream') |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 397 | |
| 398 | |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 399 | def check_files_exist_on_server(query_url, queries): |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 400 | """Queries the server to see which files from this batch already exist there. |
| 401 | |
| 402 | Arguments: |
| 403 | queries: The hash files to potential upload to the server. |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 404 | Returns: |
| 405 | missing_files: list of files that are missing on the server. |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 406 | """ |
maruel@chromium.org | 3e42ce8 | 2013-09-12 18:36:59 +0000 | [diff] [blame] | 407 | # TODO(maruel): Move inside IsolateServer. |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 408 | logging.info('Checking existence of %d files...', len(queries)) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 409 | body = ''.join( |
maruel@chromium.org | e5c1713 | 2012-11-21 18:18:46 +0000 | [diff] [blame] | 410 | (binascii.unhexlify(meta_data['h']) for (_, meta_data) in queries)) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 411 | assert (len(body) % 20) == 0, repr(body) |
| 412 | |
vadimsh@chromium.org | 80f7300 | 2013-07-12 14:52:44 +0000 | [diff] [blame] | 413 | response = url_read( |
| 414 | query_url, data=body, content_type='application/octet-stream') |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 415 | if len(queries) != len(response): |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 416 | raise MappingError( |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 417 | 'Got an incorrect number of responses from the server. Expected %d, ' |
| 418 | 'but got %d' % (len(queries), len(response))) |
| 419 | |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 420 | missing_files = [ |
| 421 | queries[i] for i, flag in enumerate(response) if flag == chr(0) |
| 422 | ] |
| 423 | logging.info('Queried %d files, %d cache hit', |
| 424 | len(queries), len(queries) - len(missing_files)) |
| 425 | return missing_files |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 426 | |
| 427 | |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 428 | class FileSystem(object): |
| 429 | """Fetches data from the file system. |
| 430 | |
| 431 | The common use case is a NFS/CIFS file server that is mounted locally that is |
| 432 | used to fetch the file on a local partition. |
| 433 | """ |
| 434 | def __init__(self, base_path): |
| 435 | self.base_path = base_path |
| 436 | |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 437 | def fetch(self, item, size): |
| 438 | source = os.path.join(self.base_path, item) |
| 439 | if size != UNKNOWN_FILE_SIZE and not is_valid_file(source, size): |
| 440 | raise IOError('Invalid file %s' % item) |
| 441 | with open(source, 'rb') as f: |
| 442 | return [f.read()] |
| 443 | |
| 444 | def retrieve(self, item, dest, size): |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 445 | source = os.path.join(self.base_path, item) |
| 446 | if source == dest: |
| 447 | logging.info('Source and destination are the same, no action required') |
| 448 | return |
| 449 | logging.debug('copy_file(%s, %s)', source, dest) |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 450 | if size != UNKNOWN_FILE_SIZE and not is_valid_file(source, size): |
| 451 | raise IOError( |
| 452 | 'Invalid file %s, %d != %d' % (item, os.stat(source).st_size, size)) |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 453 | shutil.copy(source, dest) |
| 454 | |
| 455 | def store(self, content, hash_key): |
| 456 | raise NotImplementedError() |
| 457 | |
| 458 | |
| 459 | def get_storage_api(file_or_url, namespace): |
| 460 | """Returns an object that implements .retrieve().""" |
| 461 | if re.match(r'^https?://.+$', file_or_url): |
| 462 | return IsolateServer(file_or_url, namespace) |
| 463 | else: |
| 464 | return FileSystem(file_or_url) |
| 465 | |
| 466 | |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 467 | class RemoteOperation(object): |
| 468 | """Priority based worker queue to operate on action items. |
| 469 | |
| 470 | It execute a function with the given task items. It is specialized to download |
| 471 | files. |
| 472 | |
| 473 | When the priority of items is equals, works in strict FIFO mode. |
| 474 | """ |
| 475 | # Initial and maximum number of worker threads. |
| 476 | INITIAL_WORKERS = 2 |
| 477 | MAX_WORKERS = 16 |
| 478 | # Priorities. |
| 479 | LOW, MED, HIGH = (1<<8, 2<<8, 3<<8) |
| 480 | INTERNAL_PRIORITY_BITS = (1<<8) - 1 |
| 481 | RETRIES = 5 |
| 482 | |
| 483 | def __init__(self, do_item): |
| 484 | # Function to fetch a remote object or upload to a remote location. |
| 485 | self._do_item = do_item |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 486 | self._pool = threading_utils.ThreadPool( |
| 487 | self.INITIAL_WORKERS, self.MAX_WORKERS, 0, 'remote') |
| 488 | |
| 489 | def join(self): |
| 490 | """Blocks until the queue is empty.""" |
| 491 | return self._pool.join() |
| 492 | |
| 493 | def close(self): |
| 494 | """Terminates all worker threads.""" |
| 495 | self._pool.close() |
| 496 | |
| 497 | def add_item(self, priority, obj, dest, size): |
| 498 | """Retrieves an object from the remote data store. |
| 499 | |
| 500 | The smaller |priority| gets fetched first. |
| 501 | |
| 502 | Thread-safe. |
| 503 | """ |
| 504 | assert (priority & self.INTERNAL_PRIORITY_BITS) == 0 |
| 505 | return self._add_item(priority, obj, dest, size) |
| 506 | |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 507 | def get_one_result(self): |
| 508 | return self._pool.get_one_result() |
| 509 | |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 510 | def _add_item(self, priority, obj, dest, size): |
| 511 | assert isinstance(obj, basestring), obj |
| 512 | assert isinstance(dest, basestring), dest |
| 513 | assert size is None or isinstance(size, int), size |
| 514 | return self._pool.add_task( |
| 515 | priority, self._task_executer, priority, obj, dest, size) |
| 516 | |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 517 | def _task_executer(self, priority, obj, dest, size): |
| 518 | """Wraps self._do_item to trap and retry on IOError exceptions.""" |
| 519 | try: |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 520 | self._do_item(obj, dest, size) |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 521 | # TODO(maruel): Technically, we'd want to have an output queue to be a |
| 522 | # PriorityQueue. |
| 523 | return obj |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 524 | except IOError: |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 525 | # Retry a few times, lowering the priority. |
| 526 | if (priority & self.INTERNAL_PRIORITY_BITS) < self.RETRIES: |
| 527 | self._add_item(priority + 1, obj, dest, size) |
| 528 | return |
| 529 | raise |
| 530 | |
| 531 | |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 532 | def compression_level(filename): |
| 533 | """Given a filename calculates the ideal compression level to use.""" |
| 534 | file_ext = os.path.splitext(filename)[1].lower() |
| 535 | # TODO(csharp): Profile to find what compression level works best. |
| 536 | return 0 if file_ext in ALREADY_COMPRESSED_TYPES else 7 |
| 537 | |
| 538 | |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 539 | def read_and_compress(filepath, level): |
| 540 | """Reads a file and returns its content gzip compressed.""" |
| 541 | compressor = zlib.compressobj(level) |
| 542 | compressed_data = cStringIO.StringIO() |
| 543 | with open(filepath, 'rb') as f: |
| 544 | while True: |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 545 | chunk = f.read(ZIPPED_FILE_CHUNK) |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 546 | if not chunk: |
| 547 | break |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 548 | compressed_data.write(compressor.compress(chunk)) |
| 549 | compressed_data.write(compressor.flush(zlib.Z_FINISH)) |
| 550 | value = compressed_data.getvalue() |
| 551 | compressed_data.close() |
| 552 | return value |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 553 | |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 554 | |
| 555 | def zip_and_trigger_upload(infile, metadata, upload_function): |
| 556 | # TODO(csharp): Fix crbug.com/150823 and enable the touched logic again. |
| 557 | # if not metadata['T']: |
| 558 | compressed_data = read_and_compress(infile, compression_level(infile)) |
| 559 | priority = ( |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 560 | RemoteOperation.HIGH if metadata.get('priority', '1') == '0' |
| 561 | else RemoteOperation.MED) |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 562 | return upload_function( |
| 563 | priority, compressed_data, metadata['h'], UNKNOWN_FILE_SIZE) |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 564 | |
| 565 | |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 566 | def batch_files_for_check(infiles): |
| 567 | """Splits list of files to check for existence on the server into batches. |
maruel@chromium.org | 35fc0c8 | 2013-01-17 15:14:14 +0000 | [diff] [blame] | 568 | |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 569 | Each batch corresponds to a single 'exists?' query to the server. |
| 570 | |
| 571 | Yields: |
| 572 | batches: list of batches, each batch is a list of files. |
maruel@chromium.org | 35fc0c8 | 2013-01-17 15:14:14 +0000 | [diff] [blame] | 573 | """ |
vadimsh@chromium.org | eea5242 | 2013-08-21 19:35:54 +0000 | [diff] [blame] | 574 | batch_count = 0 |
| 575 | batch_size_limit = ITEMS_PER_CONTAINS_QUERIES[0] |
maruel@chromium.org | 35fc0c8 | 2013-01-17 15:14:14 +0000 | [diff] [blame] | 576 | next_queries = [] |
csharp@chromium.org | 90c4581 | 2013-01-23 14:27:21 +0000 | [diff] [blame] | 577 | items = ((k, v) for k, v in infiles.iteritems() if 's' in v) |
| 578 | for relfile, metadata in sorted(items, key=lambda x: -x[1]['s']): |
maruel@chromium.org | 35fc0c8 | 2013-01-17 15:14:14 +0000 | [diff] [blame] | 579 | next_queries.append((relfile, metadata)) |
vadimsh@chromium.org | eea5242 | 2013-08-21 19:35:54 +0000 | [diff] [blame] | 580 | if len(next_queries) == batch_size_limit: |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 581 | yield next_queries |
maruel@chromium.org | 35fc0c8 | 2013-01-17 15:14:14 +0000 | [diff] [blame] | 582 | next_queries = [] |
vadimsh@chromium.org | eea5242 | 2013-08-21 19:35:54 +0000 | [diff] [blame] | 583 | batch_count += 1 |
| 584 | batch_size_limit = ITEMS_PER_CONTAINS_QUERIES[ |
| 585 | min(batch_count, len(ITEMS_PER_CONTAINS_QUERIES) - 1)] |
maruel@chromium.org | 35fc0c8 | 2013-01-17 15:14:14 +0000 | [diff] [blame] | 586 | if next_queries: |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 587 | yield next_queries |
| 588 | |
| 589 | |
| 590 | def get_files_to_upload(contains_hash_url, infiles): |
| 591 | """Yields files that are missing on the server.""" |
vadimsh@chromium.org | b074b16 | 2013-08-22 17:55:46 +0000 | [diff] [blame] | 592 | with threading_utils.ThreadPool(1, 16, 0, prefix='get_files_to_upload') as tp: |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 593 | for files in batch_files_for_check(infiles): |
vadimsh@chromium.org | b074b16 | 2013-08-22 17:55:46 +0000 | [diff] [blame] | 594 | tp.add_task(0, check_files_exist_on_server, contains_hash_url, files) |
| 595 | for missing_file in itertools.chain.from_iterable(tp.iter_results()): |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 596 | yield missing_file |
maruel@chromium.org | 35fc0c8 | 2013-01-17 15:14:14 +0000 | [diff] [blame] | 597 | |
| 598 | |
maruel@chromium.org | 7b844a6 | 2013-09-17 13:04:59 +0000 | [diff] [blame] | 599 | def upload_tree(base_url, indir, infiles, namespace): |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 600 | """Uploads the given tree to the given url. |
| 601 | |
| 602 | Arguments: |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 603 | base_url: The base url, it is assume that |base_url|/has/ can be used to |
| 604 | query if an element was already uploaded, and |base_url|/store/ |
| 605 | can be used to upload a new element. |
| 606 | indir: Root directory the infiles are based in. |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 607 | infiles: dict of files to upload files from |indir| to |base_url|. |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 608 | namespace: The namespace to use on the server. |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 609 | """ |
| 610 | logging.info('upload tree(base_url=%s, indir=%s, files=%d)' % |
| 611 | (base_url, indir, len(infiles))) |
maruel@chromium.org | 034e396 | 2013-03-13 13:34:25 +0000 | [diff] [blame] | 612 | |
csharp@chromium.org | 07fa759 | 2013-01-11 18:19:30 +0000 | [diff] [blame] | 613 | # Create a pool of workers to zip and upload any files missing from |
| 614 | # the server. |
vadimsh@chromium.org | b074b16 | 2013-08-22 17:55:46 +0000 | [diff] [blame] | 615 | num_threads = threading_utils.num_processors() |
| 616 | zipping_pool = threading_utils.ThreadPool(min(2, num_threads), |
| 617 | num_threads, 0, 'zip') |
maruel@chromium.org | 3e42ce8 | 2013-09-12 18:36:59 +0000 | [diff] [blame] | 618 | remote = IsolateServer(base_url, namespace) |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 619 | remote_uploader = RemoteOperation(remote.store) |
csharp@chromium.org | 07fa759 | 2013-01-11 18:19:30 +0000 | [diff] [blame] | 620 | |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 621 | # Starts the zip and upload process for files that are missing |
| 622 | # from the server. |
maruel@chromium.org | 3e42ce8 | 2013-09-12 18:36:59 +0000 | [diff] [blame] | 623 | contains_hash_url = '%scontains/%s?token=%s' % ( |
| 624 | remote.content_url, namespace, remote.token) |
csharp@chromium.org | 20a888c | 2013-01-15 15:06:55 +0000 | [diff] [blame] | 625 | uploaded = [] |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 626 | for relfile, metadata in get_files_to_upload(contains_hash_url, infiles): |
csharp@chromium.org | 07fa759 | 2013-01-11 18:19:30 +0000 | [diff] [blame] | 627 | infile = os.path.join(indir, relfile) |
maruel@chromium.org | 831958f | 2013-01-22 15:01:46 +0000 | [diff] [blame] | 628 | zipping_pool.add_task(0, zip_and_trigger_upload, infile, metadata, |
csharp@chromium.org | 07fa759 | 2013-01-11 18:19:30 +0000 | [diff] [blame] | 629 | remote_uploader.add_item) |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 630 | uploaded.append((relfile, metadata)) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 631 | |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 632 | logging.info('Waiting for all files to finish zipping') |
| 633 | zipping_pool.join() |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 634 | zipping_pool.close() |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 635 | logging.info('All files zipped.') |
| 636 | |
| 637 | logging.info('Waiting for all files to finish uploading') |
maruel@chromium.org | 13eca0b | 2013-01-22 16:42:21 +0000 | [diff] [blame] | 638 | # Will raise if any exception occurred. |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 639 | remote_uploader.join() |
vadimsh@chromium.org | 53f8d5a | 2013-06-19 13:03:55 +0000 | [diff] [blame] | 640 | remote_uploader.close() |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 641 | logging.info('All files are uploaded') |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 642 | |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 643 | total = len(infiles) |
maruel@chromium.org | e5c1713 | 2012-11-21 18:18:46 +0000 | [diff] [blame] | 644 | total_size = sum(metadata.get('s', 0) for metadata in infiles.itervalues()) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 645 | logging.info( |
| 646 | 'Total: %6d, %9.1fkb', |
| 647 | total, |
maruel@chromium.org | e5c1713 | 2012-11-21 18:18:46 +0000 | [diff] [blame] | 648 | sum(m.get('s', 0) for m in infiles.itervalues()) / 1024.) |
csharp@chromium.org | 20a888c | 2013-01-15 15:06:55 +0000 | [diff] [blame] | 649 | cache_hit = set(infiles.iterkeys()) - set(x[0] for x in uploaded) |
maruel@chromium.org | e5c1713 | 2012-11-21 18:18:46 +0000 | [diff] [blame] | 650 | cache_hit_size = sum(infiles[i].get('s', 0) for i in cache_hit) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 651 | logging.info( |
| 652 | 'cache hit: %6d, %9.1fkb, %6.2f%% files, %6.2f%% size', |
| 653 | len(cache_hit), |
| 654 | cache_hit_size / 1024., |
| 655 | len(cache_hit) * 100. / total, |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 656 | cache_hit_size * 100. / total_size if total_size else 0) |
csharp@chromium.org | 20a888c | 2013-01-15 15:06:55 +0000 | [diff] [blame] | 657 | cache_miss = uploaded |
maruel@chromium.org | e5c1713 | 2012-11-21 18:18:46 +0000 | [diff] [blame] | 658 | cache_miss_size = sum(infiles[i[0]].get('s', 0) for i in cache_miss) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 659 | logging.info( |
| 660 | 'cache miss: %6d, %9.1fkb, %6.2f%% files, %6.2f%% size', |
| 661 | len(cache_miss), |
| 662 | cache_miss_size / 1024., |
| 663 | len(cache_miss) * 100. / total, |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 664 | cache_miss_size * 100. / total_size if total_size else 0) |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 665 | return 0 |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 666 | |
| 667 | |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 668 | @subcommand.usage('<file1..fileN> or - to read from stdin') |
| 669 | def CMDarchive(parser, args): |
| 670 | """Archives data to the server.""" |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 671 | options, files = parser.parse_args(args) |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 672 | |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 673 | if files == ['-']: |
| 674 | files = sys.stdin.readlines() |
| 675 | |
| 676 | if not files: |
| 677 | parser.error('Nothing to upload') |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 678 | if not options.isolate_server: |
| 679 | parser.error('Nowhere to send. Please specify --isolate-server') |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 680 | |
| 681 | # Load the necessary metadata. This is going to be rewritten eventually to be |
| 682 | # more efficient. |
maruel@chromium.org | 7b844a6 | 2013-09-17 13:04:59 +0000 | [diff] [blame] | 683 | algo = hashlib.sha1 |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 684 | infiles = dict( |
| 685 | ( |
| 686 | f, |
| 687 | { |
maruel@chromium.org | e5c1713 | 2012-11-21 18:18:46 +0000 | [diff] [blame] | 688 | 's': os.stat(f).st_size, |
maruel@chromium.org | 7b844a6 | 2013-09-17 13:04:59 +0000 | [diff] [blame] | 689 | 'h': hash_file(f, algo), |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 690 | } |
| 691 | ) |
| 692 | for f in files) |
| 693 | |
vadimsh@chromium.org | a432647 | 2013-08-24 02:05:41 +0000 | [diff] [blame] | 694 | with tools.Profiler('Archive'): |
maruel@chromium.org | 7b844a6 | 2013-09-17 13:04:59 +0000 | [diff] [blame] | 695 | ret = upload_tree( |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 696 | base_url=options.isolate_server, |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 697 | indir=os.getcwd(), |
csharp@chromium.org | 59c7bcf | 2012-11-21 21:13:18 +0000 | [diff] [blame] | 698 | infiles=infiles, |
| 699 | namespace=options.namespace) |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 700 | if not ret: |
| 701 | print '\n'.join('%s %s' % (infiles[f]['h'], f) for f in sorted(infiles)) |
| 702 | return ret |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 703 | |
| 704 | |
| 705 | def CMDdownload(parser, args): |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 706 | """Download data from the server. |
| 707 | |
| 708 | It can download individual files. |
| 709 | """ |
| 710 | parser.add_option( |
| 711 | '-f', '--file', metavar='HASH DEST', default=[], action='append', nargs=2, |
| 712 | help='hash and destination of a file, can be used multiple times') |
| 713 | parser.add_option( |
| 714 | '-t', '--target', metavar='DIR', default=os.getcwd(), |
| 715 | help='destination directory') |
| 716 | options, args = parser.parse_args(args) |
| 717 | if args: |
| 718 | parser.error('Unsupported arguments: %s' % args) |
| 719 | if not options.file: |
| 720 | parser.error('Use one of --file is required.') |
| 721 | |
| 722 | options.target = os.path.abspath(options.target) |
| 723 | remote = IsolateServer(options.isolate_server, options.namespace) |
| 724 | for h, dest in options.file: |
| 725 | logging.info('%s: %s', h, dest) |
maruel@chromium.org | e45728d | 2013-09-16 23:23:22 +0000 | [diff] [blame] | 726 | remote.retrieve(h, os.path.join(options.target, dest), UNKNOWN_FILE_SIZE) |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 727 | return 0 |
| 728 | |
| 729 | |
| 730 | class OptionParserIsolateServer(tools.OptionParserWithLogging): |
| 731 | def __init__(self, **kwargs): |
| 732 | tools.OptionParserWithLogging.__init__(self, **kwargs) |
| 733 | self.add_option( |
| 734 | '-I', '--isolate-server', |
| 735 | default=ISOLATE_SERVER, |
| 736 | metavar='URL', |
| 737 | help='Isolate server where data is stored. default: %default') |
| 738 | self.add_option( |
| 739 | '--namespace', default='default-gzip', |
maruel@chromium.org | b7e79a2 | 2013-09-13 01:24:56 +0000 | [diff] [blame] | 740 | help='The namespace to use on the server, default: %default') |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 741 | |
| 742 | def parse_args(self, *args, **kwargs): |
| 743 | options, args = tools.OptionParserWithLogging.parse_args( |
| 744 | self, *args, **kwargs) |
| 745 | options.isolate_server = options.isolate_server.rstrip('/') |
| 746 | if not options.isolate_server: |
| 747 | self.error('--isolate-server is required.') |
| 748 | return options, args |
| 749 | |
| 750 | |
| 751 | def main(args): |
| 752 | dispatcher = subcommand.CommandDispatcher(__name__) |
| 753 | try: |
| 754 | return dispatcher.execute( |
| 755 | OptionParserIsolateServer(version=__version__), args) |
maruel@chromium.org | dedbf49 | 2013-09-12 20:42:11 +0000 | [diff] [blame] | 756 | except (ConfigError, MappingError) as e: |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 757 | sys.stderr.write('\nError: ') |
| 758 | sys.stderr.write(str(e)) |
| 759 | sys.stderr.write('\n') |
| 760 | return 1 |
maruel@chromium.org | c6f9006 | 2012-11-07 18:32:22 +0000 | [diff] [blame] | 761 | |
| 762 | |
| 763 | if __name__ == '__main__': |
maruel@chromium.org | fb78d43 | 2013-08-28 21:22:40 +0000 | [diff] [blame] | 764 | fix_encoding.fix_encoding() |
| 765 | tools.disable_buffering() |
| 766 | colorama.init() |
maruel@chromium.org | cb3c3d5 | 2013-03-14 18:55:30 +0000 | [diff] [blame] | 767 | sys.exit(main(sys.argv[1:])) |