repository.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760
  1. from configparser import ConfigParser
  2. from binascii import hexlify, unhexlify
  3. from datetime import datetime
  4. from itertools import islice
  5. import errno
  6. import logging
  7. logger = logging.getLogger(__name__)
  8. import os
  9. import shutil
  10. import struct
  11. from zlib import crc32
  12. import msgpack
  13. from .helpers import Error, ErrorWithTraceback, IntegrityError, Location, ProgressIndicatorPercent
  14. from .hashindex import NSIndex
  15. from .locking import UpgradableLock, LockError, LockErrorT
  16. from .lrucache import LRUCache
  17. from .platform import sync_dir
  18. MAX_OBJECT_SIZE = 20 * 1024 * 1024
  19. MAGIC = b'BORG_SEG'
  20. MAGIC_LEN = len(MAGIC)
  21. TAG_PUT = 0
  22. TAG_DELETE = 1
  23. TAG_COMMIT = 2
  24. class Repository:
  25. """Filesystem based transactional key value store
  26. On disk layout:
  27. dir/README
  28. dir/config
  29. dir/data/<X / SEGMENTS_PER_DIR>/<X>
  30. dir/index.X
  31. dir/hints.X
  32. """
  33. DEFAULT_MAX_SEGMENT_SIZE = 5 * 1024 * 1024
  34. DEFAULT_SEGMENTS_PER_DIR = 10000
  35. class DoesNotExist(Error):
  36. """Repository {} does not exist."""
  37. class AlreadyExists(Error):
  38. """Repository {} already exists."""
  39. class InvalidRepository(Error):
  40. """{} is not a valid repository. Check repo config."""
  41. class CheckNeeded(ErrorWithTraceback):
  42. """Inconsistency detected. Please run "borg check {}"."""
  43. class ObjectNotFound(ErrorWithTraceback):
  44. """Object with key {} not found in repository {}."""
  45. def __init__(self, path, create=False, exclusive=False, lock_wait=None, lock=True, append_only=False):
  46. self.path = os.path.abspath(path)
  47. self._location = Location('file://%s' % self.path)
  48. self.io = None
  49. self.lock = None
  50. self.index = None
  51. self._active_txn = False
  52. self.lock_wait = lock_wait
  53. self.do_lock = lock
  54. self.do_create = create
  55. self.exclusive = exclusive
  56. self.append_only = append_only
  57. def __del__(self):
  58. if self.lock:
  59. self.close()
  60. assert False, "cleanup happened in Repository.__del__"
  61. def __repr__(self):
  62. return '<%s %s>' % (self.__class__.__name__, self.path)
  63. def __enter__(self):
  64. if self.do_create:
  65. self.do_create = False
  66. self.create(self.path)
  67. self.open(self.path, self.exclusive, lock_wait=self.lock_wait, lock=self.do_lock)
  68. return self
  69. def __exit__(self, exc_type, exc_val, exc_tb):
  70. if exc_type is not None:
  71. no_space_left_on_device = exc_type is OSError and exc_val.errno == errno.ENOSPC
  72. # The ENOSPC could have originated somewhere else besides the Repository. The cleanup is always safe, unless
  73. # EIO or FS corruption ensues, which is why we specifically check for ENOSPC.
  74. if self._active_txn and no_space_left_on_device:
  75. logger.warning('No space left on device, cleaning up partial transaction to free space.')
  76. self.io.cleanup(self.io.get_segments_transaction_id())
  77. self.rollback()
  78. self.close()
  79. def create(self, path):
  80. """Create a new empty repository at `path`
  81. """
  82. if os.path.exists(path) and (not os.path.isdir(path) or os.listdir(path)):
  83. raise self.AlreadyExists(path)
  84. if not os.path.exists(path):
  85. os.mkdir(path)
  86. with open(os.path.join(path, 'README'), 'w') as fd:
  87. fd.write('This is a Borg repository\n')
  88. os.mkdir(os.path.join(path, 'data'))
  89. config = ConfigParser(interpolation=None)
  90. config.add_section('repository')
  91. config.set('repository', 'version', '1')
  92. config.set('repository', 'segments_per_dir', str(self.DEFAULT_SEGMENTS_PER_DIR))
  93. config.set('repository', 'max_segment_size', str(self.DEFAULT_MAX_SEGMENT_SIZE))
  94. config.set('repository', 'append_only', str(int(self.append_only)))
  95. config.set('repository', 'id', hexlify(os.urandom(32)).decode('ascii'))
  96. self.save_config(path, config)
  97. def save_config(self, path, config):
  98. config_path = os.path.join(path, 'config')
  99. with open(config_path, 'w') as fd:
  100. config.write(fd)
  101. def save_key(self, keydata):
  102. assert self.config
  103. keydata = keydata.decode('utf-8') # remote repo: msgpack issue #99, getting bytes
  104. self.config.set('repository', 'key', keydata)
  105. self.save_config(self.path, self.config)
  106. def load_key(self):
  107. keydata = self.config.get('repository', 'key')
  108. return keydata.encode('utf-8') # remote repo: msgpack issue #99, returning bytes
  109. def destroy(self):
  110. """Destroy the repository at `self.path`
  111. """
  112. if self.append_only:
  113. raise ValueError(self.path + " is in append-only mode")
  114. self.close()
  115. os.remove(os.path.join(self.path, 'config')) # kill config first
  116. shutil.rmtree(self.path)
  117. def get_index_transaction_id(self):
  118. indices = sorted(int(fn[6:])
  119. for fn in os.listdir(self.path)
  120. if fn.startswith('index.') and fn[6:].isdigit() and os.stat(os.path.join(self.path, fn)).st_size != 0)
  121. if indices:
  122. return indices[-1]
  123. else:
  124. return None
  125. def get_transaction_id(self):
  126. index_transaction_id = self.get_index_transaction_id()
  127. segments_transaction_id = self.io.get_segments_transaction_id()
  128. if index_transaction_id is not None and segments_transaction_id is None:
  129. raise self.CheckNeeded(self.path)
  130. # Attempt to automatically rebuild index if we crashed between commit
  131. # tag write and index save
  132. if index_transaction_id != segments_transaction_id:
  133. if index_transaction_id is not None and index_transaction_id > segments_transaction_id:
  134. replay_from = None
  135. else:
  136. replay_from = index_transaction_id
  137. self.replay_segments(replay_from, segments_transaction_id)
  138. return self.get_index_transaction_id()
  139. def break_lock(self):
  140. UpgradableLock(os.path.join(self.path, 'lock')).break_lock()
  141. def open(self, path, exclusive, lock_wait=None, lock=True):
  142. self.path = path
  143. if not os.path.isdir(path):
  144. raise self.DoesNotExist(path)
  145. if lock:
  146. self.lock = UpgradableLock(os.path.join(path, 'lock'), exclusive, timeout=lock_wait).acquire()
  147. else:
  148. self.lock = None
  149. self.config = ConfigParser(interpolation=None)
  150. self.config.read(os.path.join(self.path, 'config'))
  151. if 'repository' not in self.config.sections() or self.config.getint('repository', 'version') != 1:
  152. raise self.InvalidRepository(path)
  153. self.max_segment_size = self.config.getint('repository', 'max_segment_size')
  154. self.segments_per_dir = self.config.getint('repository', 'segments_per_dir')
  155. # append_only can be set in the constructor
  156. # it shouldn't be overridden (True -> False) here
  157. self.append_only = self.append_only or self.config.getboolean('repository', 'append_only', fallback=False)
  158. self.id = unhexlify(self.config.get('repository', 'id').strip())
  159. self.io = LoggedIO(self.path, self.max_segment_size, self.segments_per_dir)
  160. def close(self):
  161. if self.lock:
  162. if self.io:
  163. self.io.close()
  164. self.io = None
  165. self.lock.release()
  166. self.lock = None
  167. def commit(self, save_space=False):
  168. """Commit transaction
  169. """
  170. self.io.write_commit()
  171. if not self.append_only:
  172. self.compact_segments(save_space=save_space)
  173. self.write_index()
  174. self.rollback()
  175. def open_index(self, transaction_id):
  176. if transaction_id is None:
  177. return NSIndex()
  178. return NSIndex.read((os.path.join(self.path, 'index.%d') % transaction_id).encode('utf-8'))
  179. def prepare_txn(self, transaction_id, do_cleanup=True):
  180. self._active_txn = True
  181. try:
  182. self.lock.upgrade()
  183. except (LockError, LockErrorT):
  184. # if upgrading the lock to exclusive fails, we do not have an
  185. # active transaction. this is important for "serve" mode, where
  186. # the repository instance lives on - even if exceptions happened.
  187. self._active_txn = False
  188. raise
  189. if not self.index or transaction_id is None:
  190. self.index = self.open_index(transaction_id)
  191. if transaction_id is None:
  192. self.segments = {} # XXX bad name: usage_count_of_segment_x = self.segments[x]
  193. self.compact = set() # XXX bad name: segments_needing_compaction = self.compact
  194. else:
  195. if do_cleanup:
  196. self.io.cleanup(transaction_id)
  197. with open(os.path.join(self.path, 'hints.%d' % transaction_id), 'rb') as fd:
  198. hints = msgpack.unpack(fd)
  199. hints_version = hints[b'version']
  200. if hints_version not in (1, 2):
  201. raise ValueError('Unknown hints file version: %d' % hints_version)
  202. self.segments = hints[b'segments']
  203. if hints_version == 1:
  204. self.compact = set(hints[b'compact'])
  205. elif hints_version == 2:
  206. self.compact = set(hints[b'compact'].keys())
  207. def write_index(self):
  208. hints = {b'version': 1,
  209. b'segments': self.segments,
  210. b'compact': list(self.compact)}
  211. transaction_id = self.io.get_segments_transaction_id()
  212. hints_file = os.path.join(self.path, 'hints.%d' % transaction_id)
  213. with open(hints_file + '.tmp', 'wb') as fd:
  214. msgpack.pack(hints, fd)
  215. fd.flush()
  216. os.fsync(fd.fileno())
  217. os.rename(hints_file + '.tmp', hints_file)
  218. self.index.write(os.path.join(self.path, 'index.tmp'))
  219. os.rename(os.path.join(self.path, 'index.tmp'),
  220. os.path.join(self.path, 'index.%d' % transaction_id))
  221. if self.append_only:
  222. with open(os.path.join(self.path, 'transactions'), 'a') as log:
  223. print('transaction %d, UTC time %s' % (transaction_id, datetime.utcnow().isoformat()), file=log)
  224. # Remove old indices
  225. current = '.%d' % transaction_id
  226. for name in os.listdir(self.path):
  227. if not name.startswith('index.') and not name.startswith('hints.'):
  228. continue
  229. if name.endswith(current):
  230. continue
  231. os.unlink(os.path.join(self.path, name))
  232. self.index = None
  233. def compact_segments(self, save_space=False):
  234. """Compact sparse segments by copying data into new segments
  235. """
  236. if not self.compact:
  237. return
  238. index_transaction_id = self.get_index_transaction_id()
  239. segments = self.segments
  240. unused = [] # list of segments, that are not used anymore
  241. def complete_xfer():
  242. # complete the transfer (usually exactly when some target segment
  243. # is full, or at the very end when everything is processed)
  244. nonlocal unused
  245. # commit the new, compact, used segments
  246. self.io.write_commit()
  247. # get rid of the old, sparse, unused segments. free space.
  248. for segment in unused:
  249. assert self.segments.pop(segment) == 0
  250. self.io.delete_segment(segment)
  251. unused = []
  252. for segment in sorted(self.compact):
  253. if self.io.segment_exists(segment):
  254. for tag, key, offset, data in self.io.iter_objects(segment, include_data=True):
  255. if tag == TAG_PUT and self.index.get(key, (-1, -1)) == (segment, offset):
  256. try:
  257. new_segment, offset = self.io.write_put(key, data, raise_full=save_space)
  258. except LoggedIO.SegmentFull:
  259. complete_xfer()
  260. new_segment, offset = self.io.write_put(key, data)
  261. self.index[key] = new_segment, offset
  262. segments.setdefault(new_segment, 0)
  263. segments[new_segment] += 1
  264. segments[segment] -= 1
  265. elif tag == TAG_DELETE:
  266. if index_transaction_id is None or segment > index_transaction_id:
  267. try:
  268. self.io.write_delete(key, raise_full=save_space)
  269. except LoggedIO.SegmentFull:
  270. complete_xfer()
  271. self.io.write_delete(key)
  272. assert segments[segment] == 0
  273. unused.append(segment)
  274. complete_xfer()
  275. self.compact = set()
  276. def replay_segments(self, index_transaction_id, segments_transaction_id):
  277. self.prepare_txn(index_transaction_id, do_cleanup=False)
  278. try:
  279. segment_count = sum(1 for _ in self.io.segment_iterator())
  280. pi = ProgressIndicatorPercent(total=segment_count, msg="Replaying segments %3.0f%%", same_line=True)
  281. for i, (segment, filename) in enumerate(self.io.segment_iterator()):
  282. pi.show(i)
  283. if index_transaction_id is not None and segment <= index_transaction_id:
  284. continue
  285. if segment > segments_transaction_id:
  286. break
  287. objects = self.io.iter_objects(segment)
  288. self._update_index(segment, objects)
  289. pi.finish()
  290. self.write_index()
  291. finally:
  292. self.rollback()
  293. def _update_index(self, segment, objects, report=None):
  294. """some code shared between replay_segments and check"""
  295. self.segments[segment] = 0
  296. for tag, key, offset in objects:
  297. if tag == TAG_PUT:
  298. try:
  299. s, _ = self.index[key]
  300. self.compact.add(s)
  301. self.segments[s] -= 1
  302. except KeyError:
  303. pass
  304. self.index[key] = segment, offset
  305. self.segments[segment] += 1
  306. elif tag == TAG_DELETE:
  307. try:
  308. s, _ = self.index.pop(key)
  309. self.segments[s] -= 1
  310. self.compact.add(s)
  311. except KeyError:
  312. pass
  313. self.compact.add(segment)
  314. elif tag == TAG_COMMIT:
  315. continue
  316. else:
  317. msg = 'Unexpected tag {} in segment {}'.format(tag, segment)
  318. if report is None:
  319. raise self.CheckNeeded(msg)
  320. else:
  321. report(msg)
  322. if self.segments[segment] == 0:
  323. self.compact.add(segment)
  324. def check(self, repair=False, save_space=False):
  325. """Check repository consistency
  326. This method verifies all segment checksums and makes sure
  327. the index is consistent with the data stored in the segments.
  328. """
  329. if self.append_only and repair:
  330. raise ValueError(self.path + " is in append-only mode")
  331. error_found = False
  332. def report_error(msg):
  333. nonlocal error_found
  334. error_found = True
  335. logger.error(msg)
  336. logger.info('Starting repository check')
  337. assert not self._active_txn
  338. try:
  339. transaction_id = self.get_transaction_id()
  340. current_index = self.open_index(transaction_id)
  341. except Exception:
  342. transaction_id = self.io.get_segments_transaction_id()
  343. current_index = None
  344. if transaction_id is None:
  345. transaction_id = self.get_index_transaction_id()
  346. if transaction_id is None:
  347. transaction_id = self.io.get_latest_segment()
  348. if repair:
  349. self.io.cleanup(transaction_id)
  350. segments_transaction_id = self.io.get_segments_transaction_id()
  351. self.prepare_txn(None) # self.index, self.compact, self.segments all empty now!
  352. segment_count = sum(1 for _ in self.io.segment_iterator())
  353. pi = ProgressIndicatorPercent(total=segment_count, msg="Checking segments %3.1f%%", step=0.1, same_line=True)
  354. for i, (segment, filename) in enumerate(self.io.segment_iterator()):
  355. pi.show(i)
  356. if segment > transaction_id:
  357. continue
  358. try:
  359. objects = list(self.io.iter_objects(segment))
  360. except IntegrityError as err:
  361. report_error(str(err))
  362. objects = []
  363. if repair:
  364. self.io.recover_segment(segment, filename)
  365. objects = list(self.io.iter_objects(segment))
  366. self._update_index(segment, objects, report_error)
  367. pi.finish()
  368. # self.index, self.segments, self.compact now reflect the state of the segment files up to <transaction_id>
  369. # We might need to add a commit tag if no committed segment is found
  370. if repair and segments_transaction_id is None:
  371. report_error('Adding commit tag to segment {}'.format(transaction_id))
  372. self.io.segment = transaction_id + 1
  373. self.io.write_commit()
  374. if current_index and not repair:
  375. # current_index = "as found on disk"
  376. # self.index = "as rebuilt in-memory from segments"
  377. if len(current_index) != len(self.index):
  378. report_error('Index object count mismatch. {} != {}'.format(len(current_index), len(self.index)))
  379. elif current_index:
  380. for key, value in self.index.iteritems():
  381. if current_index.get(key, (-1, -1)) != value:
  382. report_error('Index mismatch for key {}. {} != {}'.format(key, value, current_index.get(key, (-1, -1))))
  383. if repair:
  384. self.compact_segments(save_space=save_space)
  385. self.write_index()
  386. self.rollback()
  387. if error_found:
  388. if repair:
  389. logger.info('Completed repository check, errors found and repaired.')
  390. else:
  391. logger.error('Completed repository check, errors found.')
  392. else:
  393. logger.info('Completed repository check, no problems found.')
  394. return not error_found or repair
  395. def rollback(self):
  396. """
  397. """
  398. self.index = None
  399. self._active_txn = False
  400. def __len__(self):
  401. if not self.index:
  402. self.index = self.open_index(self.get_transaction_id())
  403. return len(self.index)
  404. def __contains__(self, id):
  405. if not self.index:
  406. self.index = self.open_index(self.get_transaction_id())
  407. return id in self.index
  408. def list(self, limit=None, marker=None):
  409. if not self.index:
  410. self.index = self.open_index(self.get_transaction_id())
  411. return [id_ for id_, _ in islice(self.index.iteritems(marker=marker), limit)]
  412. def get(self, id_):
  413. if not self.index:
  414. self.index = self.open_index(self.get_transaction_id())
  415. try:
  416. segment, offset = self.index[id_]
  417. return self.io.read(segment, offset, id_)
  418. except KeyError:
  419. raise self.ObjectNotFound(id_, self.path) from None
  420. def get_many(self, ids, is_preloaded=False):
  421. for id_ in ids:
  422. yield self.get(id_)
  423. def put(self, id, data, wait=True):
  424. if not self._active_txn:
  425. self.prepare_txn(self.get_transaction_id())
  426. try:
  427. segment, _ = self.index[id]
  428. self.segments[segment] -= 1
  429. self.compact.add(segment)
  430. segment = self.io.write_delete(id)
  431. self.segments.setdefault(segment, 0)
  432. self.compact.add(segment)
  433. except KeyError:
  434. pass
  435. segment, offset = self.io.write_put(id, data)
  436. self.segments.setdefault(segment, 0)
  437. self.segments[segment] += 1
  438. self.index[id] = segment, offset
  439. def delete(self, id, wait=True):
  440. if not self._active_txn:
  441. self.prepare_txn(self.get_transaction_id())
  442. try:
  443. segment, offset = self.index.pop(id)
  444. except KeyError:
  445. raise self.ObjectNotFound(id, self.path) from None
  446. self.segments[segment] -= 1
  447. self.compact.add(segment)
  448. segment = self.io.write_delete(id)
  449. self.compact.add(segment)
  450. self.segments.setdefault(segment, 0)
  451. def preload(self, ids):
  452. """Preload objects (only applies to remote repositories)
  453. """
  454. class LoggedIO:
  455. class SegmentFull(Exception):
  456. """raised when a segment is full, before opening next"""
  457. header_fmt = struct.Struct('<IIB')
  458. assert header_fmt.size == 9
  459. put_header_fmt = struct.Struct('<IIB32s')
  460. assert put_header_fmt.size == 41
  461. header_no_crc_fmt = struct.Struct('<IB')
  462. assert header_no_crc_fmt.size == 5
  463. crc_fmt = struct.Struct('<I')
  464. assert crc_fmt.size == 4
  465. _commit = header_no_crc_fmt.pack(9, TAG_COMMIT)
  466. COMMIT = crc_fmt.pack(crc32(_commit)) + _commit
  467. def __init__(self, path, limit, segments_per_dir, capacity=90):
  468. self.path = path
  469. self.fds = LRUCache(capacity,
  470. dispose=lambda fd: fd.close())
  471. self.segment = 0
  472. self.limit = limit
  473. self.segments_per_dir = segments_per_dir
  474. self.offset = 0
  475. self._write_fd = None
  476. def close(self):
  477. self.close_segment()
  478. self.fds.clear()
  479. self.fds = None # Just to make sure we're disabled
  480. def segment_iterator(self, reverse=False):
  481. data_path = os.path.join(self.path, 'data')
  482. dirs = sorted((dir for dir in os.listdir(data_path) if dir.isdigit()), key=int, reverse=reverse)
  483. for dir in dirs:
  484. filenames = os.listdir(os.path.join(data_path, dir))
  485. sorted_filenames = sorted((filename for filename in filenames
  486. if filename.isdigit()), key=int, reverse=reverse)
  487. for filename in sorted_filenames:
  488. yield int(filename), os.path.join(data_path, dir, filename)
  489. def get_latest_segment(self):
  490. for segment, filename in self.segment_iterator(reverse=True):
  491. return segment
  492. return None
  493. def get_segments_transaction_id(self):
  494. """Return last committed segment
  495. """
  496. for segment, filename in self.segment_iterator(reverse=True):
  497. if self.is_committed_segment(segment):
  498. return segment
  499. return None
  500. def cleanup(self, transaction_id):
  501. """Delete segment files left by aborted transactions
  502. """
  503. self.segment = transaction_id + 1
  504. for segment, filename in self.segment_iterator(reverse=True):
  505. if segment > transaction_id:
  506. os.unlink(filename)
  507. else:
  508. break
  509. def is_committed_segment(self, segment):
  510. """Check if segment ends with a COMMIT_TAG tag
  511. """
  512. try:
  513. iterator = self.iter_objects(segment)
  514. except IntegrityError:
  515. return False
  516. with open(self.segment_filename(segment), 'rb') as fd:
  517. try:
  518. fd.seek(-self.header_fmt.size, os.SEEK_END)
  519. except OSError as e:
  520. # return False if segment file is empty or too small
  521. if e.errno == errno.EINVAL:
  522. return False
  523. raise e
  524. if fd.read(self.header_fmt.size) != self.COMMIT:
  525. return False
  526. seen_commit = False
  527. while True:
  528. try:
  529. tag, key, offset = next(iterator)
  530. except IntegrityError:
  531. return False
  532. except StopIteration:
  533. break
  534. if tag == TAG_COMMIT:
  535. seen_commit = True
  536. continue
  537. if seen_commit:
  538. return False
  539. return seen_commit
  540. def segment_filename(self, segment):
  541. return os.path.join(self.path, 'data', str(segment // self.segments_per_dir), str(segment))
  542. def get_write_fd(self, no_new=False, raise_full=False):
  543. if not no_new and self.offset and self.offset > self.limit:
  544. if raise_full:
  545. raise self.SegmentFull
  546. self.close_segment()
  547. if not self._write_fd:
  548. if self.segment % self.segments_per_dir == 0:
  549. dirname = os.path.join(self.path, 'data', str(self.segment // self.segments_per_dir))
  550. if not os.path.exists(dirname):
  551. os.mkdir(dirname)
  552. sync_dir(os.path.join(self.path, 'data'))
  553. self._write_fd = open(self.segment_filename(self.segment), 'ab')
  554. self._write_fd.write(MAGIC)
  555. self.offset = MAGIC_LEN
  556. return self._write_fd
  557. def get_fd(self, segment):
  558. try:
  559. return self.fds[segment]
  560. except KeyError:
  561. fd = open(self.segment_filename(segment), 'rb')
  562. self.fds[segment] = fd
  563. return fd
  564. def delete_segment(self, segment):
  565. if segment in self.fds:
  566. del self.fds[segment]
  567. try:
  568. os.unlink(self.segment_filename(segment))
  569. except FileNotFoundError:
  570. pass
  571. def segment_exists(self, segment):
  572. return os.path.exists(self.segment_filename(segment))
  573. def iter_objects(self, segment, include_data=False):
  574. fd = self.get_fd(segment)
  575. fd.seek(0)
  576. if fd.read(MAGIC_LEN) != MAGIC:
  577. raise IntegrityError('Invalid segment magic [segment {}, offset {}]'.format(segment, 0))
  578. offset = MAGIC_LEN
  579. header = fd.read(self.header_fmt.size)
  580. while header:
  581. size, tag, key, data = self._read(fd, self.header_fmt, header, segment, offset,
  582. (TAG_PUT, TAG_DELETE, TAG_COMMIT))
  583. if include_data:
  584. yield tag, key, offset, data
  585. else:
  586. yield tag, key, offset
  587. offset += size
  588. header = fd.read(self.header_fmt.size)
  589. def recover_segment(self, segment, filename):
  590. if segment in self.fds:
  591. del self.fds[segment]
  592. with open(filename, 'rb') as fd:
  593. data = memoryview(fd.read())
  594. os.rename(filename, filename + '.beforerecover')
  595. logger.info('attempting to recover ' + filename)
  596. with open(filename, 'wb') as fd:
  597. fd.write(MAGIC)
  598. while len(data) >= self.header_fmt.size:
  599. crc, size, tag = self.header_fmt.unpack(data[:self.header_fmt.size])
  600. if size < self.header_fmt.size or size > len(data):
  601. data = data[1:]
  602. continue
  603. if crc32(data[4:size]) & 0xffffffff != crc:
  604. data = data[1:]
  605. continue
  606. fd.write(data[:size])
  607. data = data[size:]
  608. def read(self, segment, offset, id):
  609. if segment == self.segment and self._write_fd:
  610. self._write_fd.flush()
  611. fd = self.get_fd(segment)
  612. fd.seek(offset)
  613. header = fd.read(self.put_header_fmt.size)
  614. size, tag, key, data = self._read(fd, self.put_header_fmt, header, segment, offset, (TAG_PUT, ))
  615. if id != key:
  616. raise IntegrityError('Invalid segment entry header, is not for wanted id [segment {}, offset {}]'.format(
  617. segment, offset))
  618. return data
  619. def _read(self, fd, fmt, header, segment, offset, acceptable_tags):
  620. # some code shared by read() and iter_objects()
  621. try:
  622. hdr_tuple = fmt.unpack(header)
  623. except struct.error as err:
  624. raise IntegrityError('Invalid segment entry header [segment {}, offset {}]: {}'.format(
  625. segment, offset, err)) from None
  626. if fmt is self.put_header_fmt:
  627. crc, size, tag, key = hdr_tuple
  628. elif fmt is self.header_fmt:
  629. crc, size, tag = hdr_tuple
  630. key = None
  631. else:
  632. raise TypeError("_read called with unsupported format")
  633. if size > MAX_OBJECT_SIZE or size < fmt.size:
  634. raise IntegrityError('Invalid segment entry size [segment {}, offset {}]'.format(
  635. segment, offset))
  636. length = size - fmt.size
  637. data = fd.read(length)
  638. if len(data) != length:
  639. raise IntegrityError('Segment entry data short read [segment {}, offset {}]: expected {}, got {} bytes'.format(
  640. segment, offset, length, len(data)))
  641. if crc32(data, crc32(memoryview(header)[4:])) & 0xffffffff != crc:
  642. raise IntegrityError('Segment entry checksum mismatch [segment {}, offset {}]'.format(
  643. segment, offset))
  644. if tag not in acceptable_tags:
  645. raise IntegrityError('Invalid segment entry header, did not get acceptable tag [segment {}, offset {}]'.format(
  646. segment, offset))
  647. if key is None and tag in (TAG_PUT, TAG_DELETE):
  648. key, data = data[:32], data[32:]
  649. return size, tag, key, data
  650. def write_put(self, id, data, raise_full=False):
  651. fd = self.get_write_fd(raise_full=raise_full)
  652. size = len(data) + self.put_header_fmt.size
  653. offset = self.offset
  654. header = self.header_no_crc_fmt.pack(size, TAG_PUT)
  655. crc = self.crc_fmt.pack(crc32(data, crc32(id, crc32(header))) & 0xffffffff)
  656. fd.write(b''.join((crc, header, id, data)))
  657. self.offset += size
  658. return self.segment, offset
  659. def write_delete(self, id, raise_full=False):
  660. fd = self.get_write_fd(raise_full=raise_full)
  661. header = self.header_no_crc_fmt.pack(self.put_header_fmt.size, TAG_DELETE)
  662. crc = self.crc_fmt.pack(crc32(id, crc32(header)) & 0xffffffff)
  663. fd.write(b''.join((crc, header, id)))
  664. self.offset += self.put_header_fmt.size
  665. return self.segment
  666. def write_commit(self):
  667. fd = self.get_write_fd(no_new=True)
  668. header = self.header_no_crc_fmt.pack(self.header_fmt.size, TAG_COMMIT)
  669. crc = self.crc_fmt.pack(crc32(header) & 0xffffffff)
  670. # first fsync(fd) here (to ensure data supposedly hits the disk before the commit tag)
  671. fd.flush()
  672. os.fsync(fd.fileno())
  673. fd.write(b''.join((crc, header)))
  674. self.close_segment() # after-commit fsync()
  675. def close_segment(self):
  676. if self._write_fd:
  677. self.segment += 1
  678. self.offset = 0
  679. self._write_fd.flush()
  680. os.fsync(self._write_fd.fileno())
  681. if hasattr(os, 'posix_fadvise'): # only on UNIX
  682. # tell the OS that it does not need to cache what we just wrote,
  683. # avoids spoiling the cache for the OS and other processes.
  684. os.posix_fadvise(self._write_fd.fileno(), 0, 0, os.POSIX_FADV_DONTNEED)
  685. self._write_fd.close()
  686. sync_dir(os.path.dirname(self._write_fd.name))
  687. self._write_fd = None