Merge "Protect SparseImage._GetRangeData() with lock"

This commit is contained in:
Tianjie Xu 2018-01-30 23:34:50 +00:00 committed by Gerrit Code Review
commit 5edd2821de
2 changed files with 36 additions and 36 deletions

View file

@ -776,15 +776,13 @@ class BlockImageDiff(object):
src_ranges = xf.src_ranges src_ranges = xf.src_ranges
tgt_ranges = xf.tgt_ranges tgt_ranges = xf.tgt_ranges
# Needs lock since WriteRangeDataToFd() is stateful (calling seek). src_file = common.MakeTempFile(prefix="src-")
with lock: with open(src_file, "wb") as fd:
src_file = common.MakeTempFile(prefix="src-") self.src.WriteRangeDataToFd(src_ranges, fd)
with open(src_file, "wb") as fd:
self.src.WriteRangeDataToFd(src_ranges, fd)
tgt_file = common.MakeTempFile(prefix="tgt-") tgt_file = common.MakeTempFile(prefix="tgt-")
with open(tgt_file, "wb") as fd: with open(tgt_file, "wb") as fd:
self.tgt.WriteRangeDataToFd(tgt_ranges, fd) self.tgt.WriteRangeDataToFd(tgt_ranges, fd)
message = [] message = []
try: try:
@ -1430,11 +1428,10 @@ class BlockImageDiff(object):
src_file = common.MakeTempFile(prefix="src-") src_file = common.MakeTempFile(prefix="src-")
tgt_file = common.MakeTempFile(prefix="tgt-") tgt_file = common.MakeTempFile(prefix="tgt-")
with transfer_lock: with open(src_file, "wb") as src_fd:
with open(src_file, "wb") as src_fd: self.src.WriteRangeDataToFd(src_ranges, src_fd)
self.src.WriteRangeDataToFd(src_ranges, src_fd) with open(tgt_file, "wb") as tgt_fd:
with open(tgt_file, "wb") as tgt_fd: self.tgt.WriteRangeDataToFd(tgt_ranges, tgt_fd)
self.tgt.WriteRangeDataToFd(tgt_ranges, tgt_fd)
patch_file = common.MakeTempFile(prefix="patch-") patch_file = common.MakeTempFile(prefix="patch-")
patch_info_file = common.MakeTempFile(prefix="split_info-") patch_info_file = common.MakeTempFile(prefix="split_info-")

View file

@ -15,6 +15,7 @@
import bisect import bisect
import os import os
import struct import struct
import threading
from hashlib import sha1 from hashlib import sha1
import rangelib import rangelib
@ -111,6 +112,8 @@ class SparseImage(object):
raise ValueError("Unknown chunk type 0x%04X not supported" % raise ValueError("Unknown chunk type 0x%04X not supported" %
(chunk_type,)) (chunk_type,))
self.generator_lock = threading.Lock()
self.care_map = rangelib.RangeSet(care_data) self.care_map = rangelib.RangeSet(care_data)
self.offset_index = [i[0] for i in offset_map] self.offset_index = [i[0] for i in offset_map]
@ -173,39 +176,39 @@ class SparseImage(object):
particular is not necessarily equal to the number of ranges in particular is not necessarily equal to the number of ranges in
'ranges'. 'ranges'.
This generator is stateful -- it depends on the open file object Use a lock to protect the generator so that we will not run two
contained in this SparseImage, so you should not try to run two
instances of this generator on the same object simultaneously.""" instances of this generator on the same object simultaneously."""
f = self.simg_f f = self.simg_f
for s, e in ranges: with self.generator_lock:
to_read = e-s for s, e in ranges:
idx = bisect.bisect_right(self.offset_index, s) - 1 to_read = e-s
chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx] idx = bisect.bisect_right(self.offset_index, s) - 1
# for the first chunk we may be starting partway through it.
remain = chunk_len - (s - chunk_start)
this_read = min(remain, to_read)
if filepos is not None:
p = filepos + ((s - chunk_start) * self.blocksize)
f.seek(p, os.SEEK_SET)
yield f.read(this_read * self.blocksize)
else:
yield fill_data * (this_read * (self.blocksize >> 2))
to_read -= this_read
while to_read > 0:
# continue with following chunks if this range spans multiple chunks.
idx += 1
chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx] chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx]
this_read = min(chunk_len, to_read)
# for the first chunk we may be starting partway through it.
remain = chunk_len - (s - chunk_start)
this_read = min(remain, to_read)
if filepos is not None: if filepos is not None:
f.seek(filepos, os.SEEK_SET) p = filepos + ((s - chunk_start) * self.blocksize)
f.seek(p, os.SEEK_SET)
yield f.read(this_read * self.blocksize) yield f.read(this_read * self.blocksize)
else: else:
yield fill_data * (this_read * (self.blocksize >> 2)) yield fill_data * (this_read * (self.blocksize >> 2))
to_read -= this_read to_read -= this_read
while to_read > 0:
# continue with following chunks if this range spans multiple chunks.
idx += 1
chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx]
this_read = min(chunk_len, to_read)
if filepos is not None:
f.seek(filepos, os.SEEK_SET)
yield f.read(this_read * self.blocksize)
else:
yield fill_data * (this_read * (self.blocksize >> 2))
to_read -= this_read
def LoadFileBlockMap(self, fn, clobbered_blocks): def LoadFileBlockMap(self, fn, clobbered_blocks):
remaining = self.care_map remaining = self.care_map
self.file_map = out = {} self.file_map = out = {}