blob: d49112fd16e055a5002b13828f3c3667a381bf53 [file] [log] [blame]
Doug Zongker424296a2014-09-02 08:53:09 -07001# Copyright (C) 2014 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
Doug Zongkerfc44a512014-08-26 13:10:25 -070015from __future__ import print_function
16
17from collections import deque, OrderedDict
18from hashlib import sha1
Doug Zongker6ab2a502016-02-09 08:28:09 -080019import array
Tao Bao8dcf7382015-05-21 14:09:49 -070020import common
Doug Zongker6ab2a502016-02-09 08:28:09 -080021import functools
Doug Zongker62338182014-09-08 08:29:55 -070022import heapq
Doug Zongkerfc44a512014-08-26 13:10:25 -070023import itertools
24import multiprocessing
25import os
Doug Zongkerfc44a512014-08-26 13:10:25 -070026import re
27import subprocess
Doug Zongkerfc44a512014-08-26 13:10:25 -070028import threading
Doug Zongker6ab2a502016-02-09 08:28:09 -080029import time
Doug Zongkerfc44a512014-08-26 13:10:25 -070030import tempfile
31
Dan Albert8b72aef2015-03-23 19:13:21 -070032from rangelib import RangeSet
33
Doug Zongkerfc44a512014-08-26 13:10:25 -070034
Doug Zongkerab7ca1d2014-08-26 10:40:28 -070035__all__ = ["EmptyImage", "DataImage", "BlockImageDiff"]
36
Dan Albert8b72aef2015-03-23 19:13:21 -070037
Doug Zongkerfc44a512014-08-26 13:10:25 -070038def compute_patch(src, tgt, imgdiff=False):
39 srcfd, srcfile = tempfile.mkstemp(prefix="src-")
40 tgtfd, tgtfile = tempfile.mkstemp(prefix="tgt-")
41 patchfd, patchfile = tempfile.mkstemp(prefix="patch-")
42 os.close(patchfd)
43
44 try:
45 with os.fdopen(srcfd, "wb") as f_src:
46 for p in src:
47 f_src.write(p)
48
49 with os.fdopen(tgtfd, "wb") as f_tgt:
50 for p in tgt:
51 f_tgt.write(p)
52 try:
53 os.unlink(patchfile)
54 except OSError:
55 pass
56 if imgdiff:
57 p = subprocess.call(["imgdiff", "-z", srcfile, tgtfile, patchfile],
58 stdout=open("/dev/null", "a"),
59 stderr=subprocess.STDOUT)
60 else:
61 p = subprocess.call(["bsdiff", srcfile, tgtfile, patchfile])
62
63 if p:
64 raise ValueError("diff failed: " + str(p))
65
66 with open(patchfile, "rb") as f:
67 return f.read()
68 finally:
69 try:
70 os.unlink(srcfile)
71 os.unlink(tgtfile)
72 os.unlink(patchfile)
73 except OSError:
74 pass
75
Dan Albert8b72aef2015-03-23 19:13:21 -070076
77class Image(object):
78 def ReadRangeSet(self, ranges):
79 raise NotImplementedError
80
Tao Bao68658c02015-06-01 13:40:49 -070081 def TotalSha1(self, include_clobbered_blocks=False):
Dan Albert8b72aef2015-03-23 19:13:21 -070082 raise NotImplementedError
83
84
85class EmptyImage(Image):
Doug Zongkerfc44a512014-08-26 13:10:25 -070086 """A zero-length image."""
87 blocksize = 4096
88 care_map = RangeSet()
Tao Baoff777812015-05-12 11:42:31 -070089 clobbered_blocks = RangeSet()
Tao Baoe9b61912015-07-09 17:37:49 -070090 extended = RangeSet()
Doug Zongkerfc44a512014-08-26 13:10:25 -070091 total_blocks = 0
92 file_map = {}
93 def ReadRangeSet(self, ranges):
94 return ()
Tao Bao68658c02015-06-01 13:40:49 -070095 def TotalSha1(self, include_clobbered_blocks=False):
96 # EmptyImage always carries empty clobbered_blocks, so
97 # include_clobbered_blocks can be ignored.
98 assert self.clobbered_blocks.size() == 0
Doug Zongkerab7ca1d2014-08-26 10:40:28 -070099 return sha1().hexdigest()
100
101
Dan Albert8b72aef2015-03-23 19:13:21 -0700102class DataImage(Image):
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700103 """An image wrapped around a single string of data."""
104
105 def __init__(self, data, trim=False, pad=False):
106 self.data = data
107 self.blocksize = 4096
108
109 assert not (trim and pad)
110
111 partial = len(self.data) % self.blocksize
Tao Bao7589e962015-09-05 20:35:32 -0700112 padded = False
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700113 if partial > 0:
114 if trim:
115 self.data = self.data[:-partial]
116 elif pad:
117 self.data += '\0' * (self.blocksize - partial)
Tao Bao7589e962015-09-05 20:35:32 -0700118 padded = True
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700119 else:
120 raise ValueError(("data for DataImage must be multiple of %d bytes "
121 "unless trim or pad is specified") %
122 (self.blocksize,))
123
124 assert len(self.data) % self.blocksize == 0
125
126 self.total_blocks = len(self.data) / self.blocksize
127 self.care_map = RangeSet(data=(0, self.total_blocks))
Tao Bao7589e962015-09-05 20:35:32 -0700128 # When the last block is padded, we always write the whole block even for
129 # incremental OTAs. Because otherwise the last block may get skipped if
130 # unchanged for an incremental, but would fail the post-install
131 # verification if it has non-zero contents in the padding bytes.
132 # Bug: 23828506
133 if padded:
Tao Bao42206c32015-09-08 13:39:40 -0700134 clobbered_blocks = [self.total_blocks-1, self.total_blocks]
Tao Bao7589e962015-09-05 20:35:32 -0700135 else:
Tao Bao42206c32015-09-08 13:39:40 -0700136 clobbered_blocks = []
137 self.clobbered_blocks = clobbered_blocks
Tao Baoe9b61912015-07-09 17:37:49 -0700138 self.extended = RangeSet()
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700139
140 zero_blocks = []
141 nonzero_blocks = []
142 reference = '\0' * self.blocksize
143
Tao Bao7589e962015-09-05 20:35:32 -0700144 for i in range(self.total_blocks-1 if padded else self.total_blocks):
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700145 d = self.data[i*self.blocksize : (i+1)*self.blocksize]
146 if d == reference:
147 zero_blocks.append(i)
148 zero_blocks.append(i+1)
149 else:
150 nonzero_blocks.append(i)
151 nonzero_blocks.append(i+1)
152
Tao Bao42206c32015-09-08 13:39:40 -0700153 assert zero_blocks or nonzero_blocks or clobbered_blocks
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700154
Tao Bao42206c32015-09-08 13:39:40 -0700155 self.file_map = dict()
156 if zero_blocks:
157 self.file_map["__ZERO"] = RangeSet(data=zero_blocks)
158 if nonzero_blocks:
159 self.file_map["__NONZERO"] = RangeSet(data=nonzero_blocks)
160 if clobbered_blocks:
161 self.file_map["__COPY"] = RangeSet(data=clobbered_blocks)
Tao Bao7589e962015-09-05 20:35:32 -0700162
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700163 def ReadRangeSet(self, ranges):
164 return [self.data[s*self.blocksize:e*self.blocksize] for (s, e) in ranges]
165
Tao Bao68658c02015-06-01 13:40:49 -0700166 def TotalSha1(self, include_clobbered_blocks=False):
Tao Bao7589e962015-09-05 20:35:32 -0700167 if not include_clobbered_blocks:
168 ranges = self.care_map.subtract(self.clobbered_blocks)
169 return sha1(self.ReadRangeSet(ranges)).hexdigest()
170 else:
171 return sha1(self.data).hexdigest()
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700172
Doug Zongkerfc44a512014-08-26 13:10:25 -0700173
174class Transfer(object):
175 def __init__(self, tgt_name, src_name, tgt_ranges, src_ranges, style, by_id):
176 self.tgt_name = tgt_name
177 self.src_name = src_name
178 self.tgt_ranges = tgt_ranges
179 self.src_ranges = src_ranges
180 self.style = style
181 self.intact = (getattr(tgt_ranges, "monotonic", False) and
182 getattr(src_ranges, "monotonic", False))
Tao Baob8c87172015-03-19 19:42:12 -0700183
184 # We use OrderedDict rather than dict so that the output is repeatable;
185 # otherwise it would depend on the hash values of the Transfer objects.
186 self.goes_before = OrderedDict()
187 self.goes_after = OrderedDict()
Doug Zongkerfc44a512014-08-26 13:10:25 -0700188
Doug Zongker62338182014-09-08 08:29:55 -0700189 self.stash_before = []
190 self.use_stash = []
191
Doug Zongkerfc44a512014-08-26 13:10:25 -0700192 self.id = len(by_id)
193 by_id.append(self)
194
Doug Zongker62338182014-09-08 08:29:55 -0700195 def NetStashChange(self):
196 return (sum(sr.size() for (_, sr) in self.stash_before) -
197 sum(sr.size() for (_, sr) in self.use_stash))
198
Tao Bao82c47982015-08-17 09:45:13 -0700199 def ConvertToNew(self):
200 assert self.style != "new"
201 self.use_stash = []
202 self.style = "new"
203 self.src_ranges = RangeSet()
204
Doug Zongkerfc44a512014-08-26 13:10:25 -0700205 def __str__(self):
206 return (str(self.id) + ": <" + str(self.src_ranges) + " " + self.style +
207 " to " + str(self.tgt_ranges) + ">")
208
209
Doug Zongker6ab2a502016-02-09 08:28:09 -0800210@functools.total_ordering
211class HeapItem(object):
212 def __init__(self, item):
213 self.item = item
214 # Negate the score since python's heap is a min-heap and we want
215 # the maximum score.
216 self.score = -item.score
217 def clear(self):
218 self.item = None
219 def __bool__(self):
220 return self.item is None
221 def __eq__(self, other):
222 return self.score == other.score
223 def __le__(self, other):
224 return self.score <= other.score
225
226
Doug Zongkerfc44a512014-08-26 13:10:25 -0700227# BlockImageDiff works on two image objects. An image object is
228# anything that provides the following attributes:
229#
230# blocksize: the size in bytes of a block, currently must be 4096.
231#
232# total_blocks: the total size of the partition/image, in blocks.
233#
234# care_map: a RangeSet containing which blocks (in the range [0,
235# total_blocks) we actually care about; i.e. which blocks contain
236# data.
237#
238# file_map: a dict that partitions the blocks contained in care_map
239# into smaller domains that are useful for doing diffs on.
240# (Typically a domain is a file, and the key in file_map is the
241# pathname.)
242#
Tao Baoff777812015-05-12 11:42:31 -0700243# clobbered_blocks: a RangeSet containing which blocks contain data
244# but may be altered by the FS. They need to be excluded when
245# verifying the partition integrity.
246#
Doug Zongkerfc44a512014-08-26 13:10:25 -0700247# ReadRangeSet(): a function that takes a RangeSet and returns the
248# data contained in the image blocks of that RangeSet. The data
249# is returned as a list or tuple of strings; concatenating the
250# elements together should produce the requested data.
251# Implementations are free to break up the data into list/tuple
252# elements in any way that is convenient.
253#
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700254# TotalSha1(): a function that returns (as a hex string) the SHA-1
255# hash of all the data in the image (ie, all the blocks in the
Tao Bao68658c02015-06-01 13:40:49 -0700256# care_map minus clobbered_blocks, or including the clobbered
257# blocks if include_clobbered_blocks is True).
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700258#
Doug Zongkerfc44a512014-08-26 13:10:25 -0700259# When creating a BlockImageDiff, the src image may be None, in which
260# case the list of transfers produced will never read from the
261# original image.
262
263class BlockImageDiff(object):
Tao Baoeba409c2015-10-21 13:30:43 -0700264 def __init__(self, tgt, src=None, threads=None, version=4):
Doug Zongkerfc44a512014-08-26 13:10:25 -0700265 if threads is None:
266 threads = multiprocessing.cpu_count() // 2
Dan Albert8b72aef2015-03-23 19:13:21 -0700267 if threads == 0:
268 threads = 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700269 self.threads = threads
Doug Zongker62338182014-09-08 08:29:55 -0700270 self.version = version
Dan Albert8b72aef2015-03-23 19:13:21 -0700271 self.transfers = []
272 self.src_basenames = {}
273 self.src_numpatterns = {}
Tao Baob4cfca52016-02-04 14:26:02 -0800274 self._max_stashed_size = 0
Tao Baod522bdc2016-04-12 15:53:16 -0700275 self.touched_src_ranges = RangeSet()
276 self.touched_src_sha1 = None
Doug Zongker62338182014-09-08 08:29:55 -0700277
Tao Baoeba409c2015-10-21 13:30:43 -0700278 assert version in (1, 2, 3, 4)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700279
280 self.tgt = tgt
281 if src is None:
282 src = EmptyImage()
283 self.src = src
284
285 # The updater code that installs the patch always uses 4k blocks.
286 assert tgt.blocksize == 4096
287 assert src.blocksize == 4096
288
289 # The range sets in each filemap should comprise a partition of
290 # the care map.
291 self.AssertPartition(src.care_map, src.file_map.values())
292 self.AssertPartition(tgt.care_map, tgt.file_map.values())
293
Tao Baob4cfca52016-02-04 14:26:02 -0800294 @property
295 def max_stashed_size(self):
296 return self._max_stashed_size
297
Doug Zongkerfc44a512014-08-26 13:10:25 -0700298 def Compute(self, prefix):
299 # When looking for a source file to use as the diff input for a
300 # target file, we try:
301 # 1) an exact path match if available, otherwise
302 # 2) a exact basename match if available, otherwise
303 # 3) a basename match after all runs of digits are replaced by
304 # "#" if available, otherwise
305 # 4) we have no source for this target.
306 self.AbbreviateSourceNames()
307 self.FindTransfers()
308
309 # Find the ordering dependencies among transfers (this is O(n^2)
310 # in the number of transfers).
311 self.GenerateDigraph()
312 # Find a sequence of transfers that satisfies as many ordering
313 # dependencies as possible (heuristically).
314 self.FindVertexSequence()
315 # Fix up the ordering dependencies that the sequence didn't
316 # satisfy.
Doug Zongker62338182014-09-08 08:29:55 -0700317 if self.version == 1:
318 self.RemoveBackwardEdges()
319 else:
320 self.ReverseBackwardEdges()
321 self.ImproveVertexSequence()
322
Tao Bao82c47982015-08-17 09:45:13 -0700323 # Ensure the runtime stash size is under the limit.
324 if self.version >= 2 and common.OPTIONS.cache_size is not None:
325 self.ReviseStashSize()
326
Doug Zongkerfc44a512014-08-26 13:10:25 -0700327 # Double-check our work.
328 self.AssertSequenceGood()
329
330 self.ComputePatches(prefix)
331 self.WriteTransfers(prefix)
332
Dan Albert8b72aef2015-03-23 19:13:21 -0700333 def HashBlocks(self, source, ranges): # pylint: disable=no-self-use
Sami Tolvanendd67a292014-12-09 16:40:34 +0000334 data = source.ReadRangeSet(ranges)
335 ctx = sha1()
336
337 for p in data:
338 ctx.update(p)
339
340 return ctx.hexdigest()
341
Doug Zongkerfc44a512014-08-26 13:10:25 -0700342 def WriteTransfers(self, prefix):
343 out = []
344
Doug Zongkerfc44a512014-08-26 13:10:25 -0700345 total = 0
Doug Zongkerfc44a512014-08-26 13:10:25 -0700346
Doug Zongker62338182014-09-08 08:29:55 -0700347 stashes = {}
348 stashed_blocks = 0
349 max_stashed_blocks = 0
350
351 free_stash_ids = []
352 next_stash_id = 0
353
Doug Zongkerfc44a512014-08-26 13:10:25 -0700354 for xf in self.transfers:
355
Doug Zongker62338182014-09-08 08:29:55 -0700356 if self.version < 2:
357 assert not xf.stash_before
358 assert not xf.use_stash
359
360 for s, sr in xf.stash_before:
361 assert s not in stashes
362 if free_stash_ids:
363 sid = heapq.heappop(free_stash_ids)
364 else:
365 sid = next_stash_id
366 next_stash_id += 1
367 stashes[s] = sid
Sami Tolvanendd67a292014-12-09 16:40:34 +0000368 if self.version == 2:
caozhiyuan21b37d82015-10-21 15:14:03 +0800369 stashed_blocks += sr.size()
Sami Tolvanendd67a292014-12-09 16:40:34 +0000370 out.append("stash %d %s\n" % (sid, sr.to_string_raw()))
371 else:
372 sh = self.HashBlocks(self.src, sr)
373 if sh in stashes:
374 stashes[sh] += 1
375 else:
376 stashes[sh] = 1
caozhiyuan21b37d82015-10-21 15:14:03 +0800377 stashed_blocks += sr.size()
Tao Baod522bdc2016-04-12 15:53:16 -0700378 self.touched_src_ranges = self.touched_src_ranges.union(sr)
Sami Tolvanendd67a292014-12-09 16:40:34 +0000379 out.append("stash %s %s\n" % (sh, sr.to_string_raw()))
Doug Zongker62338182014-09-08 08:29:55 -0700380
381 if stashed_blocks > max_stashed_blocks:
382 max_stashed_blocks = stashed_blocks
383
Jesse Zhao7b985f62015-03-02 16:53:08 -0800384 free_string = []
caozhiyuan21b37d82015-10-21 15:14:03 +0800385 free_size = 0
Jesse Zhao7b985f62015-03-02 16:53:08 -0800386
Doug Zongker62338182014-09-08 08:29:55 -0700387 if self.version == 1:
Tao Bao4fcb77e2015-10-21 13:36:01 -0700388 src_str = xf.src_ranges.to_string_raw() if xf.src_ranges else ""
Sami Tolvanendd67a292014-12-09 16:40:34 +0000389 elif self.version >= 2:
Doug Zongker62338182014-09-08 08:29:55 -0700390
391 # <# blocks> <src ranges>
392 # OR
393 # <# blocks> <src ranges> <src locs> <stash refs...>
394 # OR
395 # <# blocks> - <stash refs...>
396
397 size = xf.src_ranges.size()
Dan Albert8b72aef2015-03-23 19:13:21 -0700398 src_str = [str(size)]
Doug Zongker62338182014-09-08 08:29:55 -0700399
400 unstashed_src_ranges = xf.src_ranges
401 mapped_stashes = []
402 for s, sr in xf.use_stash:
403 sid = stashes.pop(s)
Doug Zongker62338182014-09-08 08:29:55 -0700404 unstashed_src_ranges = unstashed_src_ranges.subtract(sr)
Sami Tolvanendd67a292014-12-09 16:40:34 +0000405 sh = self.HashBlocks(self.src, sr)
Doug Zongker62338182014-09-08 08:29:55 -0700406 sr = xf.src_ranges.map_within(sr)
407 mapped_stashes.append(sr)
Sami Tolvanendd67a292014-12-09 16:40:34 +0000408 if self.version == 2:
Dan Albert8b72aef2015-03-23 19:13:21 -0700409 src_str.append("%d:%s" % (sid, sr.to_string_raw()))
Tao Baobb625d22015-08-13 14:44:15 -0700410 # A stash will be used only once. We need to free the stash
411 # immediately after the use, instead of waiting for the automatic
412 # clean-up at the end. Because otherwise it may take up extra space
413 # and lead to OTA failures.
414 # Bug: 23119955
415 free_string.append("free %d\n" % (sid,))
caozhiyuan21b37d82015-10-21 15:14:03 +0800416 free_size += sr.size()
Sami Tolvanendd67a292014-12-09 16:40:34 +0000417 else:
418 assert sh in stashes
Dan Albert8b72aef2015-03-23 19:13:21 -0700419 src_str.append("%s:%s" % (sh, sr.to_string_raw()))
Sami Tolvanendd67a292014-12-09 16:40:34 +0000420 stashes[sh] -= 1
421 if stashes[sh] == 0:
caozhiyuan21b37d82015-10-21 15:14:03 +0800422 free_size += sr.size()
Sami Tolvanendd67a292014-12-09 16:40:34 +0000423 free_string.append("free %s\n" % (sh))
424 stashes.pop(sh)
Doug Zongker62338182014-09-08 08:29:55 -0700425 heapq.heappush(free_stash_ids, sid)
426
427 if unstashed_src_ranges:
Dan Albert8b72aef2015-03-23 19:13:21 -0700428 src_str.insert(1, unstashed_src_ranges.to_string_raw())
Doug Zongker62338182014-09-08 08:29:55 -0700429 if xf.use_stash:
430 mapped_unstashed = xf.src_ranges.map_within(unstashed_src_ranges)
Dan Albert8b72aef2015-03-23 19:13:21 -0700431 src_str.insert(2, mapped_unstashed.to_string_raw())
Doug Zongker62338182014-09-08 08:29:55 -0700432 mapped_stashes.append(mapped_unstashed)
433 self.AssertPartition(RangeSet(data=(0, size)), mapped_stashes)
434 else:
Dan Albert8b72aef2015-03-23 19:13:21 -0700435 src_str.insert(1, "-")
Doug Zongker62338182014-09-08 08:29:55 -0700436 self.AssertPartition(RangeSet(data=(0, size)), mapped_stashes)
437
Dan Albert8b72aef2015-03-23 19:13:21 -0700438 src_str = " ".join(src_str)
Doug Zongker62338182014-09-08 08:29:55 -0700439
Sami Tolvanendd67a292014-12-09 16:40:34 +0000440 # all versions:
Doug Zongker62338182014-09-08 08:29:55 -0700441 # zero <rangeset>
442 # new <rangeset>
443 # erase <rangeset>
444 #
445 # version 1:
446 # bsdiff patchstart patchlen <src rangeset> <tgt rangeset>
447 # imgdiff patchstart patchlen <src rangeset> <tgt rangeset>
448 # move <src rangeset> <tgt rangeset>
449 #
450 # version 2:
Dan Albert8b72aef2015-03-23 19:13:21 -0700451 # bsdiff patchstart patchlen <tgt rangeset> <src_str>
452 # imgdiff patchstart patchlen <tgt rangeset> <src_str>
453 # move <tgt rangeset> <src_str>
Sami Tolvanendd67a292014-12-09 16:40:34 +0000454 #
455 # version 3:
Dan Albert8b72aef2015-03-23 19:13:21 -0700456 # bsdiff patchstart patchlen srchash tgthash <tgt rangeset> <src_str>
457 # imgdiff patchstart patchlen srchash tgthash <tgt rangeset> <src_str>
458 # move hash <tgt rangeset> <src_str>
Doug Zongkerfc44a512014-08-26 13:10:25 -0700459
460 tgt_size = xf.tgt_ranges.size()
461
462 if xf.style == "new":
463 assert xf.tgt_ranges
464 out.append("%s %s\n" % (xf.style, xf.tgt_ranges.to_string_raw()))
465 total += tgt_size
466 elif xf.style == "move":
Doug Zongkerfc44a512014-08-26 13:10:25 -0700467 assert xf.tgt_ranges
468 assert xf.src_ranges.size() == tgt_size
469 if xf.src_ranges != xf.tgt_ranges:
Doug Zongker62338182014-09-08 08:29:55 -0700470 if self.version == 1:
471 out.append("%s %s %s\n" % (
472 xf.style,
473 xf.src_ranges.to_string_raw(), xf.tgt_ranges.to_string_raw()))
474 elif self.version == 2:
475 out.append("%s %s %s\n" % (
476 xf.style,
Dan Albert8b72aef2015-03-23 19:13:21 -0700477 xf.tgt_ranges.to_string_raw(), src_str))
Sami Tolvanendd67a292014-12-09 16:40:34 +0000478 elif self.version >= 3:
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100479 # take into account automatic stashing of overlapping blocks
480 if xf.src_ranges.overlaps(xf.tgt_ranges):
Tao Baoe9b61912015-07-09 17:37:49 -0700481 temp_stash_usage = stashed_blocks + xf.src_ranges.size()
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100482 if temp_stash_usage > max_stashed_blocks:
483 max_stashed_blocks = temp_stash_usage
484
Tao Baod522bdc2016-04-12 15:53:16 -0700485 self.touched_src_ranges = self.touched_src_ranges.union(
486 xf.src_ranges)
487
Sami Tolvanendd67a292014-12-09 16:40:34 +0000488 out.append("%s %s %s %s\n" % (
489 xf.style,
490 self.HashBlocks(self.tgt, xf.tgt_ranges),
Dan Albert8b72aef2015-03-23 19:13:21 -0700491 xf.tgt_ranges.to_string_raw(), src_str))
Doug Zongkerfc44a512014-08-26 13:10:25 -0700492 total += tgt_size
493 elif xf.style in ("bsdiff", "imgdiff"):
Doug Zongkerfc44a512014-08-26 13:10:25 -0700494 assert xf.tgt_ranges
495 assert xf.src_ranges
Doug Zongker62338182014-09-08 08:29:55 -0700496 if self.version == 1:
497 out.append("%s %d %d %s %s\n" % (
498 xf.style, xf.patch_start, xf.patch_len,
499 xf.src_ranges.to_string_raw(), xf.tgt_ranges.to_string_raw()))
500 elif self.version == 2:
501 out.append("%s %d %d %s %s\n" % (
502 xf.style, xf.patch_start, xf.patch_len,
Dan Albert8b72aef2015-03-23 19:13:21 -0700503 xf.tgt_ranges.to_string_raw(), src_str))
Sami Tolvanendd67a292014-12-09 16:40:34 +0000504 elif self.version >= 3:
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100505 # take into account automatic stashing of overlapping blocks
506 if xf.src_ranges.overlaps(xf.tgt_ranges):
Tao Baoe9b61912015-07-09 17:37:49 -0700507 temp_stash_usage = stashed_blocks + xf.src_ranges.size()
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100508 if temp_stash_usage > max_stashed_blocks:
509 max_stashed_blocks = temp_stash_usage
510
Tao Baod522bdc2016-04-12 15:53:16 -0700511 self.touched_src_ranges = self.touched_src_ranges.union(
512 xf.src_ranges)
513
Sami Tolvanendd67a292014-12-09 16:40:34 +0000514 out.append("%s %d %d %s %s %s %s\n" % (
515 xf.style,
516 xf.patch_start, xf.patch_len,
517 self.HashBlocks(self.src, xf.src_ranges),
518 self.HashBlocks(self.tgt, xf.tgt_ranges),
Dan Albert8b72aef2015-03-23 19:13:21 -0700519 xf.tgt_ranges.to_string_raw(), src_str))
Doug Zongkerfc44a512014-08-26 13:10:25 -0700520 total += tgt_size
521 elif xf.style == "zero":
522 assert xf.tgt_ranges
523 to_zero = xf.tgt_ranges.subtract(xf.src_ranges)
524 if to_zero:
525 out.append("%s %s\n" % (xf.style, to_zero.to_string_raw()))
526 total += to_zero.size()
527 else:
Dan Albert8b72aef2015-03-23 19:13:21 -0700528 raise ValueError("unknown transfer style '%s'\n" % xf.style)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700529
Sami Tolvanendd67a292014-12-09 16:40:34 +0000530 if free_string:
531 out.append("".join(free_string))
caozhiyuan21b37d82015-10-21 15:14:03 +0800532 stashed_blocks -= free_size
Sami Tolvanendd67a292014-12-09 16:40:34 +0000533
Tao Bao575d68a2015-08-07 19:49:45 -0700534 if self.version >= 2 and common.OPTIONS.cache_size is not None:
Tao Bao8dcf7382015-05-21 14:09:49 -0700535 # Sanity check: abort if we're going to need more stash space than
536 # the allowed size (cache_size * threshold). There are two purposes
537 # of having a threshold here. a) Part of the cache may have been
538 # occupied by some recovery logs. b) It will buy us some time to deal
539 # with the oversize issue.
540 cache_size = common.OPTIONS.cache_size
541 stash_threshold = common.OPTIONS.stash_threshold
542 max_allowed = cache_size * stash_threshold
543 assert max_stashed_blocks * self.tgt.blocksize < max_allowed, \
544 'Stash size %d (%d * %d) exceeds the limit %d (%d * %.2f)' % (
545 max_stashed_blocks * self.tgt.blocksize, max_stashed_blocks,
546 self.tgt.blocksize, max_allowed, cache_size,
547 stash_threshold)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700548
Tao Baod522bdc2016-04-12 15:53:16 -0700549 if self.version >= 3:
550 self.touched_src_sha1 = self.HashBlocks(
551 self.src, self.touched_src_ranges)
552
Tao Baoe9b61912015-07-09 17:37:49 -0700553 # Zero out extended blocks as a workaround for bug 20881595.
554 if self.tgt.extended:
555 out.append("zero %s\n" % (self.tgt.extended.to_string_raw(),))
Tao Baob32d56e2015-09-09 11:55:01 -0700556 total += self.tgt.extended.size()
Tao Baoe9b61912015-07-09 17:37:49 -0700557
558 # We erase all the blocks on the partition that a) don't contain useful
559 # data in the new image and b) will not be touched by dm-verity.
Doug Zongkerfc44a512014-08-26 13:10:25 -0700560 all_tgt = RangeSet(data=(0, self.tgt.total_blocks))
Tao Baoe9b61912015-07-09 17:37:49 -0700561 all_tgt_minus_extended = all_tgt.subtract(self.tgt.extended)
562 new_dontcare = all_tgt_minus_extended.subtract(self.tgt.care_map)
563 if new_dontcare:
564 out.append("erase %s\n" % (new_dontcare.to_string_raw(),))
Doug Zongkere985f6f2014-09-09 12:38:47 -0700565
566 out.insert(0, "%d\n" % (self.version,)) # format version number
Tao Baob32d56e2015-09-09 11:55:01 -0700567 out.insert(1, "%d\n" % (total,))
Doug Zongkere985f6f2014-09-09 12:38:47 -0700568 if self.version >= 2:
569 # version 2 only: after the total block count, we give the number
570 # of stash slots needed, and the maximum size needed (in blocks)
571 out.insert(2, str(next_stash_id) + "\n")
572 out.insert(3, str(max_stashed_blocks) + "\n")
Doug Zongkerfc44a512014-08-26 13:10:25 -0700573
574 with open(prefix + ".transfer.list", "wb") as f:
575 for i in out:
576 f.write(i)
577
Doug Zongker62338182014-09-08 08:29:55 -0700578 if self.version >= 2:
Tao Baob4cfca52016-02-04 14:26:02 -0800579 self._max_stashed_size = max_stashed_blocks * self.tgt.blocksize
Tao Bao575d68a2015-08-07 19:49:45 -0700580 OPTIONS = common.OPTIONS
581 if OPTIONS.cache_size is not None:
582 max_allowed = OPTIONS.cache_size * OPTIONS.stash_threshold
583 print("max stashed blocks: %d (%d bytes), "
584 "limit: %d bytes (%.2f%%)\n" % (
Tao Baob4cfca52016-02-04 14:26:02 -0800585 max_stashed_blocks, self._max_stashed_size, max_allowed,
586 self._max_stashed_size * 100.0 / max_allowed))
Tao Bao575d68a2015-08-07 19:49:45 -0700587 else:
588 print("max stashed blocks: %d (%d bytes), limit: <unknown>\n" % (
Tao Baob4cfca52016-02-04 14:26:02 -0800589 max_stashed_blocks, self._max_stashed_size))
Doug Zongker62338182014-09-08 08:29:55 -0700590
Tao Bao82c47982015-08-17 09:45:13 -0700591 def ReviseStashSize(self):
592 print("Revising stash size...")
593 stashes = {}
594
595 # Create the map between a stash and its def/use points. For example, for a
596 # given stash of (idx, sr), stashes[idx] = (sr, def_cmd, use_cmd).
597 for xf in self.transfers:
598 # Command xf defines (stores) all the stashes in stash_before.
599 for idx, sr in xf.stash_before:
600 stashes[idx] = (sr, xf)
601
602 # Record all the stashes command xf uses.
603 for idx, _ in xf.use_stash:
604 stashes[idx] += (xf,)
605
606 # Compute the maximum blocks available for stash based on /cache size and
607 # the threshold.
608 cache_size = common.OPTIONS.cache_size
609 stash_threshold = common.OPTIONS.stash_threshold
610 max_allowed = cache_size * stash_threshold / self.tgt.blocksize
611
612 stashed_blocks = 0
Tao Bao9a5caf22015-08-25 15:10:10 -0700613 new_blocks = 0
Tao Bao82c47982015-08-17 09:45:13 -0700614
615 # Now go through all the commands. Compute the required stash size on the
616 # fly. If a command requires excess stash than available, it deletes the
617 # stash by replacing the command that uses the stash with a "new" command
618 # instead.
619 for xf in self.transfers:
620 replaced_cmds = []
621
622 # xf.stash_before generates explicit stash commands.
623 for idx, sr in xf.stash_before:
624 if stashed_blocks + sr.size() > max_allowed:
625 # We cannot stash this one for a later command. Find out the command
626 # that will use this stash and replace the command with "new".
627 use_cmd = stashes[idx][2]
628 replaced_cmds.append(use_cmd)
Tao Bao9a5caf22015-08-25 15:10:10 -0700629 print("%10d %9s %s" % (sr.size(), "explicit", use_cmd))
Tao Bao82c47982015-08-17 09:45:13 -0700630 else:
631 stashed_blocks += sr.size()
632
633 # xf.use_stash generates free commands.
634 for _, sr in xf.use_stash:
635 stashed_blocks -= sr.size()
636
637 # "move" and "diff" may introduce implicit stashes in BBOTA v3. Prior to
638 # ComputePatches(), they both have the style of "diff".
639 if xf.style == "diff" and self.version >= 3:
640 assert xf.tgt_ranges and xf.src_ranges
641 if xf.src_ranges.overlaps(xf.tgt_ranges):
642 if stashed_blocks + xf.src_ranges.size() > max_allowed:
643 replaced_cmds.append(xf)
Tao Bao9a5caf22015-08-25 15:10:10 -0700644 print("%10d %9s %s" % (xf.src_ranges.size(), "implicit", xf))
Tao Bao82c47982015-08-17 09:45:13 -0700645
646 # Replace the commands in replaced_cmds with "new"s.
647 for cmd in replaced_cmds:
648 # It no longer uses any commands in "use_stash". Remove the def points
649 # for all those stashes.
650 for idx, sr in cmd.use_stash:
651 def_cmd = stashes[idx][1]
652 assert (idx, sr) in def_cmd.stash_before
653 def_cmd.stash_before.remove((idx, sr))
654
Tianjie Xuebe39a02016-01-14 14:12:26 -0800655 # Add up blocks that violates space limit and print total number to
656 # screen later.
657 new_blocks += cmd.tgt_ranges.size()
Tao Bao82c47982015-08-17 09:45:13 -0700658 cmd.ConvertToNew()
659
Tianjie Xuebe39a02016-01-14 14:12:26 -0800660 num_of_bytes = new_blocks * self.tgt.blocksize
661 print(" Total %d blocks (%d bytes) are packed as new blocks due to "
662 "insufficient cache size." % (new_blocks, num_of_bytes))
Tao Bao9a5caf22015-08-25 15:10:10 -0700663
Doug Zongkerfc44a512014-08-26 13:10:25 -0700664 def ComputePatches(self, prefix):
665 print("Reticulating splines...")
666 diff_q = []
667 patch_num = 0
668 with open(prefix + ".new.dat", "wb") as new_f:
669 for xf in self.transfers:
670 if xf.style == "zero":
671 pass
672 elif xf.style == "new":
673 for piece in self.tgt.ReadRangeSet(xf.tgt_ranges):
674 new_f.write(piece)
675 elif xf.style == "diff":
676 src = self.src.ReadRangeSet(xf.src_ranges)
677 tgt = self.tgt.ReadRangeSet(xf.tgt_ranges)
678
679 # We can't compare src and tgt directly because they may have
680 # the same content but be broken up into blocks differently, eg:
681 #
682 # ["he", "llo"] vs ["h", "ello"]
683 #
684 # We want those to compare equal, ideally without having to
685 # actually concatenate the strings (these may be tens of
686 # megabytes).
687
688 src_sha1 = sha1()
689 for p in src:
690 src_sha1.update(p)
691 tgt_sha1 = sha1()
692 tgt_size = 0
693 for p in tgt:
694 tgt_sha1.update(p)
695 tgt_size += len(p)
696
697 if src_sha1.digest() == tgt_sha1.digest():
698 # These are identical; we don't need to generate a patch,
699 # just issue copy commands on the device.
700 xf.style = "move"
701 else:
702 # For files in zip format (eg, APKs, JARs, etc.) we would
703 # like to use imgdiff -z if possible (because it usually
704 # produces significantly smaller patches than bsdiff).
705 # This is permissible if:
706 #
707 # - the source and target files are monotonic (ie, the
708 # data is stored with blocks in increasing order), and
709 # - we haven't removed any blocks from the source set.
710 #
711 # If these conditions are satisfied then appending all the
712 # blocks in the set together in order will produce a valid
713 # zip file (plus possibly extra zeros in the last block),
714 # which is what imgdiff needs to operate. (imgdiff is
715 # fine with extra zeros at the end of the file.)
716 imgdiff = (xf.intact and
717 xf.tgt_name.split(".")[-1].lower()
718 in ("apk", "jar", "zip"))
719 xf.style = "imgdiff" if imgdiff else "bsdiff"
720 diff_q.append((tgt_size, src, tgt, xf, patch_num))
721 patch_num += 1
722
723 else:
724 assert False, "unknown style " + xf.style
725
726 if diff_q:
727 if self.threads > 1:
728 print("Computing patches (using %d threads)..." % (self.threads,))
729 else:
730 print("Computing patches...")
731 diff_q.sort()
732
733 patches = [None] * patch_num
734
Dan Albert8b72aef2015-03-23 19:13:21 -0700735 # TODO: Rewrite with multiprocessing.ThreadPool?
Doug Zongkerfc44a512014-08-26 13:10:25 -0700736 lock = threading.Lock()
737 def diff_worker():
738 while True:
739 with lock:
Dan Albert8b72aef2015-03-23 19:13:21 -0700740 if not diff_q:
741 return
Doug Zongkerfc44a512014-08-26 13:10:25 -0700742 tgt_size, src, tgt, xf, patchnum = diff_q.pop()
743 patch = compute_patch(src, tgt, imgdiff=(xf.style == "imgdiff"))
744 size = len(patch)
745 with lock:
746 patches[patchnum] = (patch, xf)
747 print("%10d %10d (%6.2f%%) %7s %s" % (
748 size, tgt_size, size * 100.0 / tgt_size, xf.style,
749 xf.tgt_name if xf.tgt_name == xf.src_name else (
750 xf.tgt_name + " (from " + xf.src_name + ")")))
751
752 threads = [threading.Thread(target=diff_worker)
Dan Albert8b72aef2015-03-23 19:13:21 -0700753 for _ in range(self.threads)]
Doug Zongkerfc44a512014-08-26 13:10:25 -0700754 for th in threads:
755 th.start()
756 while threads:
757 threads.pop().join()
758 else:
759 patches = []
760
761 p = 0
762 with open(prefix + ".patch.dat", "wb") as patch_f:
763 for patch, xf in patches:
764 xf.patch_start = p
765 xf.patch_len = len(patch)
766 patch_f.write(patch)
767 p += len(patch)
768
769 def AssertSequenceGood(self):
770 # Simulate the sequences of transfers we will output, and check that:
771 # - we never read a block after writing it, and
772 # - we write every block we care about exactly once.
773
774 # Start with no blocks having been touched yet.
Doug Zongker6ab2a502016-02-09 08:28:09 -0800775 touched = array.array("B", "\0" * self.tgt.total_blocks)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700776
777 # Imagine processing the transfers in order.
778 for xf in self.transfers:
779 # Check that the input blocks for this transfer haven't yet been touched.
Doug Zongker62338182014-09-08 08:29:55 -0700780
781 x = xf.src_ranges
782 if self.version >= 2:
783 for _, sr in xf.use_stash:
784 x = x.subtract(sr)
785
Doug Zongker6ab2a502016-02-09 08:28:09 -0800786 for s, e in x:
Tao Baoff75c232016-03-04 15:23:34 -0800787 # Source image could be larger. Don't check the blocks that are in the
788 # source image only. Since they are not in 'touched', and won't ever
789 # be touched.
790 for i in range(s, min(e, self.tgt.total_blocks)):
Doug Zongker6ab2a502016-02-09 08:28:09 -0800791 assert touched[i] == 0
792
793 # Check that the output blocks for this transfer haven't yet
794 # been touched, and touch all the blocks written by this
795 # transfer.
796 for s, e in xf.tgt_ranges:
797 for i in range(s, e):
798 assert touched[i] == 0
799 touched[i] = 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700800
801 # Check that we've written every target block.
Doug Zongker6ab2a502016-02-09 08:28:09 -0800802 for s, e in self.tgt.care_map:
803 for i in range(s, e):
804 assert touched[i] == 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700805
Doug Zongker62338182014-09-08 08:29:55 -0700806 def ImproveVertexSequence(self):
807 print("Improving vertex order...")
808
809 # At this point our digraph is acyclic; we reversed any edges that
810 # were backwards in the heuristically-generated sequence. The
811 # previously-generated order is still acceptable, but we hope to
812 # find a better order that needs less memory for stashed data.
813 # Now we do a topological sort to generate a new vertex order,
814 # using a greedy algorithm to choose which vertex goes next
815 # whenever we have a choice.
816
817 # Make a copy of the edge set; this copy will get destroyed by the
818 # algorithm.
819 for xf in self.transfers:
820 xf.incoming = xf.goes_after.copy()
821 xf.outgoing = xf.goes_before.copy()
822
823 L = [] # the new vertex order
824
825 # S is the set of sources in the remaining graph; we always choose
826 # the one that leaves the least amount of stashed data after it's
827 # executed.
828 S = [(u.NetStashChange(), u.order, u) for u in self.transfers
829 if not u.incoming]
830 heapq.heapify(S)
831
832 while S:
833 _, _, xf = heapq.heappop(S)
834 L.append(xf)
835 for u in xf.outgoing:
836 del u.incoming[xf]
837 if not u.incoming:
838 heapq.heappush(S, (u.NetStashChange(), u.order, u))
839
840 # if this fails then our graph had a cycle.
841 assert len(L) == len(self.transfers)
842
843 self.transfers = L
844 for i, xf in enumerate(L):
845 xf.order = i
846
Doug Zongkerfc44a512014-08-26 13:10:25 -0700847 def RemoveBackwardEdges(self):
848 print("Removing backward edges...")
849 in_order = 0
850 out_of_order = 0
851 lost_source = 0
852
853 for xf in self.transfers:
Doug Zongkerfc44a512014-08-26 13:10:25 -0700854 lost = 0
855 size = xf.src_ranges.size()
856 for u in xf.goes_before:
857 # xf should go before u
858 if xf.order < u.order:
859 # it does, hurray!
Doug Zongker62338182014-09-08 08:29:55 -0700860 in_order += 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700861 else:
862 # it doesn't, boo. trim the blocks that u writes from xf's
863 # source, so that xf can go after u.
Doug Zongker62338182014-09-08 08:29:55 -0700864 out_of_order += 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700865 assert xf.src_ranges.overlaps(u.tgt_ranges)
866 xf.src_ranges = xf.src_ranges.subtract(u.tgt_ranges)
867 xf.intact = False
868
869 if xf.style == "diff" and not xf.src_ranges:
870 # nothing left to diff from; treat as new data
871 xf.style = "new"
872
873 lost = size - xf.src_ranges.size()
874 lost_source += lost
Doug Zongkerfc44a512014-08-26 13:10:25 -0700875
876 print((" %d/%d dependencies (%.2f%%) were violated; "
877 "%d source blocks removed.") %
878 (out_of_order, in_order + out_of_order,
879 (out_of_order * 100.0 / (in_order + out_of_order))
880 if (in_order + out_of_order) else 0.0,
881 lost_source))
882
Doug Zongker62338182014-09-08 08:29:55 -0700883 def ReverseBackwardEdges(self):
884 print("Reversing backward edges...")
885 in_order = 0
886 out_of_order = 0
887 stashes = 0
888 stash_size = 0
889
890 for xf in self.transfers:
Doug Zongker62338182014-09-08 08:29:55 -0700891 for u in xf.goes_before.copy():
892 # xf should go before u
893 if xf.order < u.order:
894 # it does, hurray!
895 in_order += 1
896 else:
897 # it doesn't, boo. modify u to stash the blocks that it
898 # writes that xf wants to read, and then require u to go
899 # before xf.
900 out_of_order += 1
901
902 overlap = xf.src_ranges.intersect(u.tgt_ranges)
903 assert overlap
904
905 u.stash_before.append((stashes, overlap))
906 xf.use_stash.append((stashes, overlap))
907 stashes += 1
908 stash_size += overlap.size()
909
910 # reverse the edge direction; now xf must go after u
911 del xf.goes_before[u]
912 del u.goes_after[xf]
913 xf.goes_after[u] = None # value doesn't matter
914 u.goes_before[xf] = None
915
916 print((" %d/%d dependencies (%.2f%%) were violated; "
917 "%d source blocks stashed.") %
918 (out_of_order, in_order + out_of_order,
919 (out_of_order * 100.0 / (in_order + out_of_order))
920 if (in_order + out_of_order) else 0.0,
921 stash_size))
922
Doug Zongkerfc44a512014-08-26 13:10:25 -0700923 def FindVertexSequence(self):
924 print("Finding vertex sequence...")
925
926 # This is based on "A Fast & Effective Heuristic for the Feedback
927 # Arc Set Problem" by P. Eades, X. Lin, and W.F. Smyth. Think of
928 # it as starting with the digraph G and moving all the vertices to
929 # be on a horizontal line in some order, trying to minimize the
930 # number of edges that end up pointing to the left. Left-pointing
931 # edges will get removed to turn the digraph into a DAG. In this
932 # case each edge has a weight which is the number of source blocks
933 # we'll lose if that edge is removed; we try to minimize the total
934 # weight rather than just the number of edges.
935
936 # Make a copy of the edge set; this copy will get destroyed by the
937 # algorithm.
938 for xf in self.transfers:
939 xf.incoming = xf.goes_after.copy()
940 xf.outgoing = xf.goes_before.copy()
Doug Zongker6ab2a502016-02-09 08:28:09 -0800941 xf.score = sum(xf.outgoing.values()) - sum(xf.incoming.values())
Doug Zongkerfc44a512014-08-26 13:10:25 -0700942
943 # We use an OrderedDict instead of just a set so that the output
944 # is repeatable; otherwise it would depend on the hash values of
945 # the transfer objects.
946 G = OrderedDict()
947 for xf in self.transfers:
948 G[xf] = None
949 s1 = deque() # the left side of the sequence, built from left to right
950 s2 = deque() # the right side of the sequence, built from right to left
951
Doug Zongker6ab2a502016-02-09 08:28:09 -0800952 heap = []
953 for xf in self.transfers:
954 xf.heap_item = HeapItem(xf)
955 heap.append(xf.heap_item)
956 heapq.heapify(heap)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700957
Doug Zongker6ab2a502016-02-09 08:28:09 -0800958 sinks = set(u for u in G if not u.outgoing)
959 sources = set(u for u in G if not u.incoming)
960
961 def adjust_score(iu, delta):
962 iu.score += delta
963 iu.heap_item.clear()
964 iu.heap_item = HeapItem(iu)
965 heapq.heappush(heap, iu.heap_item)
966
967 while G:
Doug Zongkerfc44a512014-08-26 13:10:25 -0700968 # Put all sinks at the end of the sequence.
Doug Zongker6ab2a502016-02-09 08:28:09 -0800969 while sinks:
970 new_sinks = set()
Doug Zongkerfc44a512014-08-26 13:10:25 -0700971 for u in sinks:
Doug Zongker6ab2a502016-02-09 08:28:09 -0800972 if u not in G: continue
Doug Zongkerfc44a512014-08-26 13:10:25 -0700973 s2.appendleft(u)
974 del G[u]
975 for iu in u.incoming:
Doug Zongker6ab2a502016-02-09 08:28:09 -0800976 adjust_score(iu, -iu.outgoing.pop(u))
977 if not iu.outgoing: new_sinks.add(iu)
978 sinks = new_sinks
Doug Zongkerfc44a512014-08-26 13:10:25 -0700979
980 # Put all the sources at the beginning of the sequence.
Doug Zongker6ab2a502016-02-09 08:28:09 -0800981 while sources:
982 new_sources = set()
Doug Zongkerfc44a512014-08-26 13:10:25 -0700983 for u in sources:
Doug Zongker6ab2a502016-02-09 08:28:09 -0800984 if u not in G: continue
Doug Zongkerfc44a512014-08-26 13:10:25 -0700985 s1.append(u)
986 del G[u]
987 for iu in u.outgoing:
Doug Zongker6ab2a502016-02-09 08:28:09 -0800988 adjust_score(iu, +iu.incoming.pop(u))
989 if not iu.incoming: new_sources.add(iu)
990 sources = new_sources
Doug Zongkerfc44a512014-08-26 13:10:25 -0700991
Doug Zongker6ab2a502016-02-09 08:28:09 -0800992 if not G: break
Doug Zongkerfc44a512014-08-26 13:10:25 -0700993
994 # Find the "best" vertex to put next. "Best" is the one that
995 # maximizes the net difference in source blocks saved we get by
996 # pretending it's a source rather than a sink.
997
Doug Zongker6ab2a502016-02-09 08:28:09 -0800998 while True:
999 u = heapq.heappop(heap)
1000 if u and u.item in G:
1001 u = u.item
1002 break
Doug Zongkerfc44a512014-08-26 13:10:25 -07001003
Doug Zongkerfc44a512014-08-26 13:10:25 -07001004 s1.append(u)
1005 del G[u]
1006 for iu in u.outgoing:
Doug Zongker6ab2a502016-02-09 08:28:09 -08001007 adjust_score(iu, +iu.incoming.pop(u))
1008 if not iu.incoming: sources.add(iu)
1009
Doug Zongkerfc44a512014-08-26 13:10:25 -07001010 for iu in u.incoming:
Doug Zongker6ab2a502016-02-09 08:28:09 -08001011 adjust_score(iu, -iu.outgoing.pop(u))
1012 if not iu.outgoing: sinks.add(iu)
Doug Zongkerfc44a512014-08-26 13:10:25 -07001013
1014 # Now record the sequence in the 'order' field of each transfer,
1015 # and by rearranging self.transfers to be in the chosen sequence.
1016
1017 new_transfers = []
1018 for x in itertools.chain(s1, s2):
1019 x.order = len(new_transfers)
1020 new_transfers.append(x)
1021 del x.incoming
1022 del x.outgoing
1023
1024 self.transfers = new_transfers
1025
1026 def GenerateDigraph(self):
1027 print("Generating digraph...")
Doug Zongker6ab2a502016-02-09 08:28:09 -08001028
1029 # Each item of source_ranges will be:
1030 # - None, if that block is not used as a source,
1031 # - a transfer, if one transfer uses it as a source, or
1032 # - a set of transfers.
1033 source_ranges = []
1034 for b in self.transfers:
1035 for s, e in b.src_ranges:
1036 if e > len(source_ranges):
1037 source_ranges.extend([None] * (e-len(source_ranges)))
1038 for i in range(s, e):
1039 if source_ranges[i] is None:
1040 source_ranges[i] = b
1041 else:
1042 if not isinstance(source_ranges[i], set):
1043 source_ranges[i] = set([source_ranges[i]])
1044 source_ranges[i].add(b)
1045
Doug Zongkerfc44a512014-08-26 13:10:25 -07001046 for a in self.transfers:
Doug Zongker6ab2a502016-02-09 08:28:09 -08001047 intersections = set()
1048 for s, e in a.tgt_ranges:
1049 for i in range(s, e):
1050 if i >= len(source_ranges): break
1051 b = source_ranges[i]
1052 if b is not None:
1053 if isinstance(b, set):
1054 intersections.update(b)
1055 else:
1056 intersections.add(b)
1057
1058 for b in intersections:
1059 if a is b: continue
Doug Zongkerfc44a512014-08-26 13:10:25 -07001060
1061 # If the blocks written by A are read by B, then B needs to go before A.
1062 i = a.tgt_ranges.intersect(b.src_ranges)
1063 if i:
Doug Zongkerab7ca1d2014-08-26 10:40:28 -07001064 if b.src_name == "__ZERO":
1065 # the cost of removing source blocks for the __ZERO domain
1066 # is (nearly) zero.
1067 size = 0
1068 else:
1069 size = i.size()
Doug Zongkerfc44a512014-08-26 13:10:25 -07001070 b.goes_before[a] = size
1071 a.goes_after[b] = size
1072
1073 def FindTransfers(self):
Tao Bao9a5caf22015-08-25 15:10:10 -07001074 """Parse the file_map to generate all the transfers."""
1075
1076 def AddTransfer(tgt_name, src_name, tgt_ranges, src_ranges, style, by_id,
1077 split=False):
1078 """Wrapper function for adding a Transfer().
1079
1080 For BBOTA v3, we need to stash source blocks for resumable feature.
1081 However, with the growth of file size and the shrink of the cache
1082 partition source blocks are too large to be stashed. If a file occupies
1083 too many blocks (greater than MAX_BLOCKS_PER_DIFF_TRANSFER), we split it
1084 into smaller pieces by getting multiple Transfer()s.
1085
Tianjie Xubb86e1d2016-01-13 16:14:10 -08001086 The downside is that after splitting, we may increase the package size
1087 since the split pieces don't align well. According to our experiments,
1088 1/8 of the cache size as the per-piece limit appears to be optimal.
1089 Compared to the fixed 1024-block limit, it reduces the overall package
1090 size by 30% volantis, and 20% for angler and bullhead."""
Tao Bao9a5caf22015-08-25 15:10:10 -07001091
1092 # We care about diff transfers only.
1093 if style != "diff" or not split:
1094 Transfer(tgt_name, src_name, tgt_ranges, src_ranges, style, by_id)
1095 return
1096
Tianjie Xubb86e1d2016-01-13 16:14:10 -08001097 pieces = 0
1098 cache_size = common.OPTIONS.cache_size
1099 split_threshold = 0.125
1100 max_blocks_per_transfer = int(cache_size * split_threshold /
1101 self.tgt.blocksize)
1102
Tao Bao9a5caf22015-08-25 15:10:10 -07001103 # Change nothing for small files.
Tianjie Xubb86e1d2016-01-13 16:14:10 -08001104 if (tgt_ranges.size() <= max_blocks_per_transfer and
1105 src_ranges.size() <= max_blocks_per_transfer):
Tao Bao9a5caf22015-08-25 15:10:10 -07001106 Transfer(tgt_name, src_name, tgt_ranges, src_ranges, style, by_id)
1107 return
1108
Tianjie Xubb86e1d2016-01-13 16:14:10 -08001109 while (tgt_ranges.size() > max_blocks_per_transfer and
1110 src_ranges.size() > max_blocks_per_transfer):
Tao Bao9a5caf22015-08-25 15:10:10 -07001111 tgt_split_name = "%s-%d" % (tgt_name, pieces)
1112 src_split_name = "%s-%d" % (src_name, pieces)
Tianjie Xubb86e1d2016-01-13 16:14:10 -08001113 tgt_first = tgt_ranges.first(max_blocks_per_transfer)
1114 src_first = src_ranges.first(max_blocks_per_transfer)
1115
Tao Bao9a5caf22015-08-25 15:10:10 -07001116 Transfer(tgt_split_name, src_split_name, tgt_first, src_first, style,
1117 by_id)
1118
1119 tgt_ranges = tgt_ranges.subtract(tgt_first)
1120 src_ranges = src_ranges.subtract(src_first)
1121 pieces += 1
1122
1123 # Handle remaining blocks.
1124 if tgt_ranges.size() or src_ranges.size():
1125 # Must be both non-empty.
1126 assert tgt_ranges.size() and src_ranges.size()
1127 tgt_split_name = "%s-%d" % (tgt_name, pieces)
1128 src_split_name = "%s-%d" % (src_name, pieces)
1129 Transfer(tgt_split_name, src_split_name, tgt_ranges, src_ranges, style,
1130 by_id)
1131
Doug Zongkerfc44a512014-08-26 13:10:25 -07001132 empty = RangeSet()
1133 for tgt_fn, tgt_ranges in self.tgt.file_map.items():
1134 if tgt_fn == "__ZERO":
1135 # the special "__ZERO" domain is all the blocks not contained
1136 # in any file and that are filled with zeros. We have a
1137 # special transfer style for zero blocks.
1138 src_ranges = self.src.file_map.get("__ZERO", empty)
Tao Bao9a5caf22015-08-25 15:10:10 -07001139 AddTransfer(tgt_fn, "__ZERO", tgt_ranges, src_ranges,
1140 "zero", self.transfers)
Doug Zongkerfc44a512014-08-26 13:10:25 -07001141 continue
1142
Tao Baoff777812015-05-12 11:42:31 -07001143 elif tgt_fn == "__COPY":
1144 # "__COPY" domain includes all the blocks not contained in any
1145 # file and that need to be copied unconditionally to the target.
Tao Bao9a5caf22015-08-25 15:10:10 -07001146 AddTransfer(tgt_fn, None, tgt_ranges, empty, "new", self.transfers)
Tao Baoff777812015-05-12 11:42:31 -07001147 continue
1148
Doug Zongkerfc44a512014-08-26 13:10:25 -07001149 elif tgt_fn in self.src.file_map:
1150 # Look for an exact pathname match in the source.
Tao Bao9a5caf22015-08-25 15:10:10 -07001151 AddTransfer(tgt_fn, tgt_fn, tgt_ranges, self.src.file_map[tgt_fn],
1152 "diff", self.transfers, self.version >= 3)
Doug Zongkerfc44a512014-08-26 13:10:25 -07001153 continue
1154
1155 b = os.path.basename(tgt_fn)
1156 if b in self.src_basenames:
1157 # Look for an exact basename match in the source.
1158 src_fn = self.src_basenames[b]
Tao Bao9a5caf22015-08-25 15:10:10 -07001159 AddTransfer(tgt_fn, src_fn, tgt_ranges, self.src.file_map[src_fn],
1160 "diff", self.transfers, self.version >= 3)
Doug Zongkerfc44a512014-08-26 13:10:25 -07001161 continue
1162
1163 b = re.sub("[0-9]+", "#", b)
1164 if b in self.src_numpatterns:
1165 # Look for a 'number pattern' match (a basename match after
1166 # all runs of digits are replaced by "#"). (This is useful
1167 # for .so files that contain version numbers in the filename
1168 # that get bumped.)
1169 src_fn = self.src_numpatterns[b]
Tao Bao9a5caf22015-08-25 15:10:10 -07001170 AddTransfer(tgt_fn, src_fn, tgt_ranges, self.src.file_map[src_fn],
1171 "diff", self.transfers, self.version >= 3)
Doug Zongkerfc44a512014-08-26 13:10:25 -07001172 continue
1173
Tao Bao9a5caf22015-08-25 15:10:10 -07001174 AddTransfer(tgt_fn, None, tgt_ranges, empty, "new", self.transfers)
Doug Zongkerfc44a512014-08-26 13:10:25 -07001175
1176 def AbbreviateSourceNames(self):
Doug Zongkerfc44a512014-08-26 13:10:25 -07001177 for k in self.src.file_map.keys():
1178 b = os.path.basename(k)
1179 self.src_basenames[b] = k
1180 b = re.sub("[0-9]+", "#", b)
1181 self.src_numpatterns[b] = k
1182
1183 @staticmethod
1184 def AssertPartition(total, seq):
1185 """Assert that all the RangeSets in 'seq' form a partition of the
1186 'total' RangeSet (ie, they are nonintersecting and their union
1187 equals 'total')."""
Doug Zongker6ab2a502016-02-09 08:28:09 -08001188
Doug Zongkerfc44a512014-08-26 13:10:25 -07001189 so_far = RangeSet()
1190 for i in seq:
1191 assert not so_far.overlaps(i)
1192 so_far = so_far.union(i)
1193 assert so_far == total