6 from arvados.retry import retry_method
11 """split(path) -> streamname, filename
13 Separate the stream name and file name in a /-separated stream path.
14 If no stream name is available, assume '.'.
17 stream_name, file_name = path.rsplit('/', 1)
18 except ValueError: # No / in string
19 stream_name, file_name = '.', path
20 return stream_name, file_name
22 class ArvadosFileBase(object):
23 def __init__(self, name, mode):
29 def _before_close(orig_func):
30 @functools.wraps(orig_func)
31 def wrapper(self, *args, **kwargs):
33 raise ValueError("I/O operation on closed stream file")
34 return orig_func(self, *args, **kwargs)
40 def __exit__(self, exc_type, exc_value, traceback):
51 class ArvadosFileReaderBase(ArvadosFileBase):
52 class _NameAttribute(str):
53 # The Python file API provides a plain .name attribute.
54 # Older SDK provided a name() method.
55 # This class provides both, for maximum compatibility.
59 def __init__(self, name, mode, num_retries=None):
60 super(ArvadosFileReaderBase, self).__init__(self._NameAttribute(name), mode)
62 self.num_retries = num_retries
63 self._readline_cache = (None, None)
67 data = self.readline()
72 def decompressed_name(self):
73 return re.sub('\.(bz2|gz)$', '', self.name)
75 @ArvadosFileBase._before_close
76 def seek(self, pos, whence=os.SEEK_CUR):
77 if whence == os.SEEK_CUR:
79 elif whence == os.SEEK_END:
81 self._filepos = min(max(pos, 0L), self.size())
86 @ArvadosFileBase._before_close
88 def readall(self, size=2**20, num_retries=None):
90 data = self.read(size, num_retries=num_retries)
95 @ArvadosFileBase._before_close
97 def readline(self, size=float('inf'), num_retries=None):
98 cache_pos, cache_data = self._readline_cache
99 if self.tell() == cache_pos:
103 data_size = len(data[-1])
104 while (data_size < size) and ('\n' not in data[-1]):
105 next_read = self.read(2 ** 20, num_retries=num_retries)
108 data.append(next_read)
109 data_size += len(next_read)
112 nextline_index = data.index('\n') + 1
114 nextline_index = len(data)
115 nextline_index = min(nextline_index, size)
116 self._readline_cache = (self.tell(), data[nextline_index:])
117 return data[:nextline_index]
119 @ArvadosFileBase._before_close
121 def decompress(self, decompress, size, num_retries=None):
122 for segment in self.readall(size, num_retries):
123 data = decompress(segment)
127 @ArvadosFileBase._before_close
129 def readall_decompressed(self, size=2**20, num_retries=None):
131 if self.name.endswith('.bz2'):
132 dc = bz2.BZ2Decompressor()
133 return self.decompress(dc.decompress, size,
134 num_retries=num_retries)
135 elif self.name.endswith('.gz'):
136 dc = zlib.decompressobj(16+zlib.MAX_WBITS)
137 return self.decompress(lambda segment: dc.decompress(dc.unconsumed_tail + segment),
138 size, num_retries=num_retries)
140 return self.readall(size, num_retries=num_retries)
142 @ArvadosFileBase._before_close
144 def readlines(self, sizehint=float('inf'), num_retries=None):
147 for s in self.readall(num_retries=num_retries):
150 if data_size >= sizehint:
152 return ''.join(data).splitlines(True)
155 class StreamFileReader(ArvadosFileReaderBase):
156 def __init__(self, stream, segments, name):
157 super(StreamFileReader, self).__init__(name, 'rb', num_retries=stream.num_retries)
158 self._stream = stream
159 self.segments = segments
161 def stream_name(self):
162 return self._stream.name()
165 n = self.segments[-1]
166 return n.range_start + n.range_size
168 @ArvadosFileBase._before_close
170 def read(self, size, num_retries=None):
171 """Read up to 'size' bytes from the stream, starting at the current file position"""
176 available_chunks = locators_and_ranges(self.segments, self._filepos, size)
178 lr = available_chunks[0]
179 data = self._stream._readfrom(lr.locator+lr.segment_offset,
181 num_retries=num_retries)
183 self._filepos += len(data)
186 @ArvadosFileBase._before_close
188 def readfrom(self, start, size, num_retries=None):
189 """Read up to 'size' bytes from the stream, starting at 'start'"""
194 for lr in locators_and_ranges(self.segments, start, size):
195 data.append(self._stream._readfrom(lr.locator+lr.segment_offset, lr.segment_size,
196 num_retries=num_retries))
199 def as_manifest(self):
200 from stream import normalize_stream
202 for r in self.segments:
203 segs.extend(self._stream.locators_and_ranges(r.locator, r.range_size))
204 return " ".join(normalize_stream(".", {self.name: segs})) + "\n"
207 class BufferBlock(object):
208 def __init__(self, locator, starting_size=2**16):
209 self.locator = locator
210 self.buffer_block = bytearray(starting_size)
211 self.buffer_view = memoryview(self.buffer_block)
212 self.write_pointer = 0
214 def append(self, data):
215 while (self.write_pointer+len(data)) > len(self.buffer_block):
216 new_buffer_block = bytearray(len(self.buffer_block) * 2)
217 new_buffer_block[0:self.write_pointer] = self.buffer_block[0:self.write_pointer]
218 self.buffer_block = new_buffer_block
219 self.buffer_view = memoryview(self.buffer_block)
220 self.buffer_view[self.write_pointer:self.write_pointer+len(data)] = data
221 self.write_pointer += len(data)
224 return self.write_pointer
226 def calculate_locator(self):
227 return "%s+%i" % (hashlib.md5(self.buffer_view[0:self.write_pointer]).hexdigest(), self.size())
230 class ArvadosFile(object):
231 def __init__(self, stream=[], segments=[], keep=None):
233 stream: a list of Range objects representing a block stream
234 segments: a list of Range objects representing segments
236 self._modified = True
239 self.add_segment(stream, s.range_start, s.range_size)
240 self._current_bblock = None
241 self._bufferblocks = None
244 def set_unmodified(self):
245 self._modified = False
248 return self._modified
250 def truncate(self, size):
252 for r in self._segments:
253 range_end = r.range_start+r.range_size
254 if r.range_start >= size:
255 # segment is past the trucate size, all done
257 elif size < range_end:
258 nr = Range(r.locator, r.range_start, size - r.range_start)
259 nr.segment_offset = r.segment_offset
265 self._segments = new_segs
266 self._modified = True
268 def _keepget(self, locator, num_retries):
269 if self._bufferblocks and locator in self._bufferblocks:
270 bb = self._bufferblocks[locator]
271 return bb.buffer_view[0:bb.write_pointer].tobytes()
273 return self._keep.get(locator, num_retries=num_retries)
275 def readfrom(self, offset, size, num_retries):
276 if size == 0 or offset >= self.size():
278 if self._keep is None:
279 self._keep = KeepClient(num_retries=num_retries)
281 # TODO: initiate prefetch on all blocks in the range (offset, offset + size + config.KEEP_BLOCK_SIZE)
283 for lr in locators_and_ranges(self._segments, offset, size):
284 # TODO: if data is empty, wait on block get, otherwise only
285 # get more data if the block is already in the cache.
286 data.append(self._keepget(lr.locator, num_retries=num_retries)[lr.segment_offset:lr.segment_offset+lr.segment_size])
289 def _init_bufferblock(self):
290 if self._bufferblocks is None:
291 self._bufferblocks = {}
292 self._current_bblock = BufferBlock("bufferblock%i" % len(self._bufferblocks))
293 self._bufferblocks[self._current_bblock.locator] = self._current_bblock
295 def _repack_writes(self):
298 '''Test if the buffer block has more data than is referenced by actual segments
299 (this happens when a buffered write over-writes a file range written in
300 a previous buffered write). Re-pack the buffer block for efficiency
301 and to avoid leaking information.
303 segs = self._segments
305 # Sum up the segments to get the total bytes of the file referencing
306 # into the buffer block.
307 bufferblock_segs = [s for s in segs if s.locator == self._current_bblock.locator]
308 write_total = sum([s.range_size for s in bufferblock_segs])
310 if write_total < self._current_bblock.size():
311 # There is more data in the buffer block than is actually accounted for by segments, so
312 # re-pack into a new buffer by copying over to a new buffer block.
313 new_bb = BufferBlock(self._current_bblock.locator, starting_size=write_total)
314 for t in bufferblock_segs:
315 new_bb.append(self._current_bblock.buffer_view[t.segment_offset:t.segment_offset+t.range_size].tobytes())
316 t.segment_offset = new_bb.size() - t.range_size
318 self._current_bblock = new_bb
319 self._bufferblocks[self._current_bblock.locator] = self._current_bblock
322 def writeto(self, offset, data, num_retries):
326 if offset > self.size():
327 raise ArgumentError("Offset is past the end of the file")
329 if len(data) > config.KEEP_BLOCK_SIZE:
330 raise ArgumentError("Please append data in chunks smaller than %i bytes (config.KEEP_BLOCK_SIZE)" % (config.KEEP_BLOCK_SIZE))
332 self._modified = True
334 if self._current_bblock is None:
335 self._init_bufferblock()
337 if (self._current_bblock.write_pointer + len(data)) > config.KEEP_BLOCK_SIZE:
338 self._repack_writes()
339 if (self._current_bblock.write_pointer + len(data)) > config.KEEP_BLOCK_SIZE:
340 self._init_bufferblock()
342 self._current_bblock.append(data)
343 replace_range(self._segments, offset, len(data), self._current_bblock.locator, self._current_bblock.write_pointer - len(data))
345 def add_segment(self, blocks, pos, size):
346 self._modified = True
347 for lr in locators_and_ranges(blocks, pos, size):
348 last = self._segments[-1] if self._segments else Range(0, 0, 0)
349 r = Range(lr.locator, last.range_start+last.range_size, lr.segment_size, lr.segment_offset)
350 self._segments.append(r)
354 n = self._segments[-1]
355 return n.range_start + n.range_size
360 class ArvadosFileReader(ArvadosFileReaderBase):
361 def __init__(self, arvadosfile, name, mode="r", num_retries=None):
362 super(ArvadosFileReader, self).__init__(name, mode, num_retries=num_retries)
363 self.arvadosfile = arvadosfile
366 return self.arvadosfile.size()
368 @ArvadosFileBase._before_close
370 def read(self, size, num_retries=None):
371 """Read up to 'size' bytes from the stream, starting at the current file position"""
372 data = self.arvadosfile.readfrom(self._filepos, size, num_retries=num_retries)
373 self._filepos += len(data)
376 @ArvadosFileBase._before_close
378 def readfrom(self, offset, size, num_retries=None):
379 """Read up to 'size' bytes from the stream, starting at the current file position"""
380 return self.arvadosfile.readfrom(offset, size, num_retries)
385 class ArvadosFileWriter(ArvadosFileReader):
386 def __init__(self, arvadosfile, name, mode, num_retries=None):
387 super(ArvadosFileWriter, self).__init__(arvadosfile, name, mode, num_retries=num_retries)
389 @ArvadosFileBase._before_close
391 def write(self, data, num_retries=None):
392 if self.mode[0] == "a":
393 self.arvadosfile.writeto(self.size(), data)
395 self.arvadosfile.writeto(self._filepos, data, num_retries)
396 self._filepos += len(data)
398 @ArvadosFileBase._before_close
400 def writelines(self, seq, num_retries=None):
404 def truncate(self, size=None):
407 self.arvadosfile.truncate(size)
408 if self._filepos > self.size():
409 self._filepos = self.size()