repository.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594
  1. from configparser import RawConfigParser
  2. from binascii import hexlify
  3. from itertools import islice
  4. import errno
  5. import os
  6. import struct
  7. import sys
  8. from zlib import crc32
  9. from .hashindex import NSIndex
  10. from .helpers import Error, IntegrityError, read_msgpack, write_msgpack, unhexlify, UpgradableLock
  11. from .lrucache import LRUCache
  12. MAX_OBJECT_SIZE = 20 * 1024 * 1024
  13. MAGIC = b'ATTICSEG'
  14. TAG_PUT = 0
  15. TAG_DELETE = 1
  16. TAG_COMMIT = 2
  17. class Repository:
  18. """Filesystem based transactional key value store
  19. On disk layout:
  20. dir/README
  21. dir/config
  22. dir/data/<X / SEGMENTS_PER_DIR>/<X>
  23. dir/index.X
  24. dir/hints.X
  25. """
  26. DEFAULT_MAX_SEGMENT_SIZE = 5 * 1024 * 1024
  27. DEFAULT_SEGMENTS_PER_DIR = 10000
  28. class DoesNotExist(Error):
  29. """Repository {} does not exist."""
  30. class AlreadyExists(Error):
  31. """Repository {} already exists."""
  32. class InvalidRepository(Error):
  33. """{} is not a valid repository."""
  34. class CheckNeeded(Error):
  35. """Inconsistency detected. Please run "attic check {}"."""
  36. class ObjectNotFound(Error):
  37. """Object with key {} not found in repository {}."""
  38. def __init__(self, path, create=False, exclusive=False):
  39. self.path = path
  40. self.io = None
  41. self.lock = None
  42. self.index = None
  43. self._active_txn = False
  44. if create:
  45. self.create(path)
  46. self.open(path, exclusive)
  47. def __del__(self):
  48. self.close()
  49. def create(self, path):
  50. """Create a new empty repository at `path`
  51. """
  52. if os.path.exists(path) and (not os.path.isdir(path) or os.listdir(path)):
  53. raise self.AlreadyExists(path)
  54. if not os.path.exists(path):
  55. os.mkdir(path)
  56. with open(os.path.join(path, 'README'), 'w') as fd:
  57. fd.write('This is an Attic repository\n')
  58. os.mkdir(os.path.join(path, 'data'))
  59. config = RawConfigParser()
  60. config.add_section('repository')
  61. config.set('repository', 'version', '1')
  62. config.set('repository', 'segments_per_dir', self.DEFAULT_SEGMENTS_PER_DIR)
  63. config.set('repository', 'max_segment_size', self.DEFAULT_MAX_SEGMENT_SIZE)
  64. config.set('repository', 'id', hexlify(os.urandom(32)).decode('ascii'))
  65. with open(os.path.join(path, 'config'), 'w') as fd:
  66. config.write(fd)
  67. def destroy(self):
  68. """Destroy the repository at `self.path`
  69. """
  70. self.close()
  71. os.remove(os.path.join(self.path, 'config')) # kill config first
  72. shutil.rmtree(self.path)
  73. def get_index_transaction_id(self):
  74. indices = sorted((int(name[6:]) for name in os.listdir(self.path) if name.startswith('index.') and name[6:].isdigit()))
  75. if indices:
  76. return indices[-1]
  77. else:
  78. return None
  79. def get_transaction_id(self):
  80. index_transaction_id = self.get_index_transaction_id()
  81. segments_transaction_id = self.io.get_segments_transaction_id()
  82. if index_transaction_id is not None and segments_transaction_id is None:
  83. raise self.CheckNeeded(self.path)
  84. # Attempt to automatically rebuild index if we crashed between commit
  85. # tag write and index save
  86. if index_transaction_id != segments_transaction_id:
  87. if index_transaction_id is not None and index_transaction_id > segments_transaction_id:
  88. replay_from = None
  89. else:
  90. replay_from = index_transaction_id
  91. self.replay_segments(replay_from, segments_transaction_id)
  92. return self.get_index_transaction_id()
  93. def open(self, path, exclusive):
  94. self.path = path
  95. if not os.path.isdir(path):
  96. raise self.DoesNotExist(path)
  97. self.config = RawConfigParser()
  98. self.config.read(os.path.join(self.path, 'config'))
  99. if 'repository' not in self.config.sections() or self.config.getint('repository', 'version') != 1:
  100. raise self.InvalidRepository(path)
  101. self.lock = UpgradableLock(os.path.join(path, 'config'), exclusive)
  102. self.max_segment_size = self.config.getint('repository', 'max_segment_size')
  103. self.segments_per_dir = self.config.getint('repository', 'segments_per_dir')
  104. self.id = unhexlify(self.config.get('repository', 'id').strip())
  105. self.io = LoggedIO(self.path, self.max_segment_size, self.segments_per_dir)
  106. def close(self):
  107. if self.lock:
  108. if self.io:
  109. self.io.close()
  110. self.io = None
  111. self.lock.release()
  112. self.lock = None
  113. def commit(self):
  114. """Commit transaction
  115. """
  116. self.io.write_commit()
  117. self.compact_segments()
  118. self.write_index()
  119. self.rollback()
  120. def open_index(self, transaction_id):
  121. if transaction_id is None:
  122. return NSIndex()
  123. return NSIndex.read((os.path.join(self.path, 'index.%d') % transaction_id).encode('utf-8'))
  124. def prepare_txn(self, transaction_id, do_cleanup=True):
  125. self._active_txn = True
  126. self.lock.upgrade()
  127. if not self.index:
  128. self.index = self.open_index(transaction_id)
  129. if transaction_id is None:
  130. self.segments = {}
  131. self.compact = set()
  132. else:
  133. if do_cleanup:
  134. self.io.cleanup(transaction_id)
  135. hints = read_msgpack(os.path.join(self.path, 'hints.%d' % transaction_id))
  136. if hints[b'version'] != 1:
  137. raise ValueError('Unknown hints file version: %d' % hints['version'])
  138. self.segments = hints[b'segments']
  139. self.compact = set(hints[b'compact'])
  140. def write_index(self):
  141. hints = {b'version': 1,
  142. b'segments': self.segments,
  143. b'compact': list(self.compact)}
  144. transaction_id = self.io.get_segments_transaction_id()
  145. write_msgpack(os.path.join(self.path, 'hints.%d' % transaction_id), hints)
  146. self.index.write(os.path.join(self.path, 'index.tmp'))
  147. os.rename(os.path.join(self.path, 'index.tmp'),
  148. os.path.join(self.path, 'index.%d' % transaction_id))
  149. # Remove old indices
  150. current = '.%d' % transaction_id
  151. for name in os.listdir(self.path):
  152. if not name.startswith('index.') and not name.startswith('hints.'):
  153. continue
  154. if name.endswith(current):
  155. continue
  156. os.unlink(os.path.join(self.path, name))
  157. self.index = None
  158. def compact_segments(self):
  159. """Compact sparse segments by copying data into new segments
  160. """
  161. if not self.compact:
  162. return
  163. index_transaction_id = self.get_index_transaction_id()
  164. segments = self.segments
  165. for segment in sorted(self.compact):
  166. if self.io.segment_exists(segment):
  167. for tag, key, offset, data in self.io.iter_objects(segment, include_data=True):
  168. if tag == TAG_PUT and self.index.get(key, (-1, -1)) == (segment, offset):
  169. new_segment, offset = self.io.write_put(key, data)
  170. self.index[key] = new_segment, offset
  171. segments.setdefault(new_segment, 0)
  172. segments[new_segment] += 1
  173. segments[segment] -= 1
  174. elif tag == TAG_DELETE:
  175. if index_transaction_id is None or segment > index_transaction_id:
  176. self.io.write_delete(key)
  177. assert segments[segment] == 0
  178. self.io.write_commit()
  179. for segment in sorted(self.compact):
  180. assert self.segments.pop(segment) == 0
  181. self.io.delete_segment(segment)
  182. self.compact = set()
  183. def replay_segments(self, index_transaction_id, segments_transaction_id):
  184. self.prepare_txn(index_transaction_id, do_cleanup=False)
  185. for segment, filename in self.io.segment_iterator():
  186. if index_transaction_id is not None and segment <= index_transaction_id:
  187. continue
  188. if segment > segments_transaction_id:
  189. break
  190. self.segments[segment] = 0
  191. for tag, key, offset in self.io.iter_objects(segment):
  192. if tag == TAG_PUT:
  193. try:
  194. s, _ = self.index[key]
  195. self.compact.add(s)
  196. self.segments[s] -= 1
  197. except KeyError:
  198. pass
  199. self.index[key] = segment, offset
  200. self.segments[segment] += 1
  201. elif tag == TAG_DELETE:
  202. try:
  203. s, _ = self.index.pop(key)
  204. self.segments[s] -= 1
  205. self.compact.add(s)
  206. except KeyError:
  207. pass
  208. self.compact.add(segment)
  209. elif tag == TAG_COMMIT:
  210. continue
  211. else:
  212. raise self.CheckNeeded(self.path)
  213. if self.segments[segment] == 0:
  214. self.compact.add(segment)
  215. self.write_index()
  216. self.rollback()
  217. def check(self, repair=False):
  218. """Check repository consistency
  219. This method verifies all segment checksums and makes sure
  220. the index is consistent with the data stored in the segments.
  221. """
  222. error_found = False
  223. def report_error(msg):
  224. nonlocal error_found
  225. error_found = True
  226. print(msg, file=sys.stderr)
  227. assert not self._active_txn
  228. try:
  229. transaction_id = self.get_transaction_id()
  230. current_index = self.open_index(transaction_id)
  231. except Exception:
  232. transaction_id = self.io.get_segments_transaction_id()
  233. current_index = None
  234. if transaction_id is None:
  235. transaction_id = self.get_index_transaction_id()
  236. if transaction_id is None:
  237. transaction_id = self.io.get_latest_segment()
  238. if repair:
  239. self.io.cleanup(transaction_id)
  240. segments_transaction_id = self.io.get_segments_transaction_id()
  241. self.prepare_txn(None)
  242. for segment, filename in self.io.segment_iterator():
  243. if segment > transaction_id:
  244. continue
  245. try:
  246. objects = list(self.io.iter_objects(segment))
  247. except (IntegrityError, struct.error):
  248. report_error('Error reading segment {}'.format(segment))
  249. objects = []
  250. if repair:
  251. self.io.recover_segment(segment, filename)
  252. objects = list(self.io.iter_objects(segment))
  253. self.segments[segment] = 0
  254. for tag, key, offset in objects:
  255. if tag == TAG_PUT:
  256. try:
  257. s, _ = self.index[key]
  258. self.compact.add(s)
  259. self.segments[s] -= 1
  260. except KeyError:
  261. pass
  262. self.index[key] = segment, offset
  263. self.segments[segment] += 1
  264. elif tag == TAG_DELETE:
  265. try:
  266. s, _ = self.index.pop(key)
  267. self.segments[s] -= 1
  268. self.compact.add(s)
  269. except KeyError:
  270. pass
  271. self.compact.add(segment)
  272. elif tag == TAG_COMMIT:
  273. continue
  274. else:
  275. report_error('Unexpected tag {} in segment {}'.format(tag, segment))
  276. # We might need to add a commit tag if no committed segment is found
  277. if repair and segments_transaction_id is None:
  278. report_error('Adding commit tag to segment {}'.format(transaction_id))
  279. self.io.segment = transaction_id + 1
  280. self.io.write_commit()
  281. self.io.close_segment()
  282. if current_index and not repair:
  283. if len(current_index) != len(self.index):
  284. report_error('Index object count mismatch. {} != {}'.format(len(current_index), len(self.index)))
  285. elif current_index:
  286. for key, value in self.index.iteritems():
  287. if current_index.get(key, (-1, -1)) != value:
  288. report_error('Index mismatch for key {}. {} != {}'.format(key, value, current_index.get(key, (-1, -1))))
  289. if repair:
  290. self.compact_segments()
  291. self.write_index()
  292. self.rollback()
  293. return not error_found or repair
  294. def rollback(self):
  295. """
  296. """
  297. self.index = None
  298. self._active_txn = False
  299. def __len__(self):
  300. if not self.index:
  301. self.index = self.open_index(self.get_transaction_id())
  302. return len(self.index)
  303. def list(self, limit=None, marker=None):
  304. if not self.index:
  305. self.index = self.open_index(self.get_transaction_id())
  306. return [id_ for id_, _ in islice(self.index.iteritems(marker=marker), limit)]
  307. def get(self, id_):
  308. if not self.index:
  309. self.index = self.open_index(self.get_transaction_id())
  310. try:
  311. segment, offset = self.index[id_]
  312. return self.io.read(segment, offset, id_)
  313. except KeyError:
  314. raise self.ObjectNotFound(id_, self.path)
  315. def get_many(self, ids, is_preloaded=False):
  316. for id_ in ids:
  317. yield self.get(id_)
  318. def put(self, id, data, wait=True):
  319. if not self._active_txn:
  320. self.prepare_txn(self.get_transaction_id())
  321. try:
  322. segment, _ = self.index[id]
  323. self.segments[segment] -= 1
  324. self.compact.add(segment)
  325. segment = self.io.write_delete(id)
  326. self.segments.setdefault(segment, 0)
  327. self.compact.add(segment)
  328. except KeyError:
  329. pass
  330. segment, offset = self.io.write_put(id, data)
  331. self.segments.setdefault(segment, 0)
  332. self.segments[segment] += 1
  333. self.index[id] = segment, offset
  334. def delete(self, id, wait=True):
  335. if not self._active_txn:
  336. self.prepare_txn(self.get_transaction_id())
  337. try:
  338. segment, offset = self.index.pop(id)
  339. except KeyError:
  340. raise self.ObjectNotFound(id, self.path)
  341. self.segments[segment] -= 1
  342. self.compact.add(segment)
  343. segment = self.io.write_delete(id)
  344. self.compact.add(segment)
  345. self.segments.setdefault(segment, 0)
  346. def preload(self, ids):
  347. """Preload objects (only applies to remote repositories
  348. """
  349. class LoggedIO:
  350. header_fmt = struct.Struct('<IIB')
  351. assert header_fmt.size == 9
  352. put_header_fmt = struct.Struct('<IIB32s')
  353. assert put_header_fmt.size == 41
  354. header_no_crc_fmt = struct.Struct('<IB')
  355. assert header_no_crc_fmt.size == 5
  356. crc_fmt = struct.Struct('<I')
  357. assert crc_fmt.size == 4
  358. _commit = header_no_crc_fmt.pack(9, TAG_COMMIT)
  359. COMMIT = crc_fmt.pack(crc32(_commit)) + _commit
  360. def __init__(self, path, limit, segments_per_dir, capacity=90):
  361. self.path = path
  362. self.fds = LRUCache(capacity)
  363. self.segment = 0
  364. self.limit = limit
  365. self.segments_per_dir = segments_per_dir
  366. self.offset = 0
  367. self._write_fd = None
  368. def close(self):
  369. for segment in list(self.fds.keys()):
  370. self.fds.pop(segment).close()
  371. self.close_segment()
  372. self.fds = None # Just to make sure we're disabled
  373. def segment_iterator(self, reverse=False):
  374. data_path = os.path.join(self.path, 'data')
  375. dirs = sorted((dir for dir in os.listdir(data_path) if dir.isdigit()), key=int, reverse=reverse)
  376. for dir in dirs:
  377. filenames = os.listdir(os.path.join(data_path, dir))
  378. sorted_filenames = sorted((filename for filename in filenames
  379. if filename.isdigit()), key=int, reverse=reverse)
  380. for filename in sorted_filenames:
  381. yield int(filename), os.path.join(data_path, dir, filename)
  382. def get_latest_segment(self):
  383. for segment, filename in self.segment_iterator(reverse=True):
  384. return segment
  385. return None
  386. def get_segments_transaction_id(self):
  387. """Verify that the transaction id is consistent with the index transaction id
  388. """
  389. for segment, filename in self.segment_iterator(reverse=True):
  390. if self.is_committed_segment(filename):
  391. return segment
  392. return None
  393. def cleanup(self, transaction_id):
  394. """Delete segment files left by aborted transactions
  395. """
  396. self.segment = transaction_id + 1
  397. for segment, filename in self.segment_iterator(reverse=True):
  398. if segment > transaction_id:
  399. os.unlink(filename)
  400. else:
  401. break
  402. def is_committed_segment(self, filename):
  403. """Check if segment ends with a COMMIT_TAG tag
  404. """
  405. with open(filename, 'rb') as fd:
  406. try:
  407. fd.seek(-self.header_fmt.size, os.SEEK_END)
  408. except OSError as e:
  409. # return False if segment file is empty or too small
  410. if e.errno == errno.EINVAL:
  411. return False
  412. raise e
  413. return fd.read(self.header_fmt.size) == self.COMMIT
  414. def segment_filename(self, segment):
  415. return os.path.join(self.path, 'data', str(segment // self.segments_per_dir), str(segment))
  416. def get_write_fd(self, no_new=False):
  417. if not no_new and self.offset and self.offset > self.limit:
  418. self.close_segment()
  419. if not self._write_fd:
  420. if self.segment % self.segments_per_dir == 0:
  421. dirname = os.path.join(self.path, 'data', str(self.segment // self.segments_per_dir))
  422. if not os.path.exists(dirname):
  423. os.mkdir(dirname)
  424. self._write_fd = open(self.segment_filename(self.segment), 'ab')
  425. self._write_fd.write(MAGIC)
  426. self.offset = 8
  427. return self._write_fd
  428. def get_fd(self, segment):
  429. try:
  430. return self.fds[segment]
  431. except KeyError:
  432. fd = open(self.segment_filename(segment), 'rb')
  433. self.fds[segment] = fd
  434. return fd
  435. def delete_segment(self, segment):
  436. try:
  437. os.unlink(self.segment_filename(segment))
  438. except OSError:
  439. pass
  440. def segment_exists(self, segment):
  441. return os.path.exists(self.segment_filename(segment))
  442. def iter_objects(self, segment, include_data=False):
  443. fd = self.get_fd(segment)
  444. fd.seek(0)
  445. if fd.read(8) != MAGIC:
  446. raise IntegrityError('Invalid segment header')
  447. offset = 8
  448. header = fd.read(self.header_fmt.size)
  449. while header:
  450. crc, size, tag = self.header_fmt.unpack(header)
  451. if size > MAX_OBJECT_SIZE:
  452. raise IntegrityError('Invalid segment object size')
  453. rest = fd.read(size - self.header_fmt.size)
  454. if crc32(rest, crc32(memoryview(header)[4:])) & 0xffffffff != crc:
  455. raise IntegrityError('Segment checksum mismatch')
  456. if tag not in (TAG_PUT, TAG_DELETE, TAG_COMMIT):
  457. raise IntegrityError('Invalid segment entry header')
  458. key = None
  459. if tag in (TAG_PUT, TAG_DELETE):
  460. key = rest[:32]
  461. if include_data:
  462. yield tag, key, offset, rest[32:]
  463. else:
  464. yield tag, key, offset
  465. offset += size
  466. header = fd.read(self.header_fmt.size)
  467. def recover_segment(self, segment, filename):
  468. self.fds.pop(segment).close()
  469. # FIXME: save a copy of the original file
  470. with open(filename, 'rb') as fd:
  471. data = memoryview(fd.read())
  472. os.rename(filename, filename + '.beforerecover')
  473. print('attempting to recover ' + filename, file=sys.stderr)
  474. with open(filename, 'wb') as fd:
  475. fd.write(MAGIC)
  476. while len(data) >= self.header_fmt.size:
  477. crc, size, tag = self.header_fmt.unpack(data[:self.header_fmt.size])
  478. if size < self.header_fmt.size or size > len(data):
  479. data = data[1:]
  480. continue
  481. if crc32(data[4:size]) & 0xffffffff != crc:
  482. data = data[1:]
  483. continue
  484. fd.write(data[:size])
  485. data = data[size:]
  486. def read(self, segment, offset, id):
  487. if segment == self.segment and self._write_fd:
  488. self._write_fd.flush()
  489. fd = self.get_fd(segment)
  490. fd.seek(offset)
  491. header = fd.read(self.put_header_fmt.size)
  492. crc, size, tag, key = self.put_header_fmt.unpack(header)
  493. if size > MAX_OBJECT_SIZE:
  494. raise IntegrityError('Invalid segment object size')
  495. data = fd.read(size - self.put_header_fmt.size)
  496. if crc32(data, crc32(memoryview(header)[4:])) & 0xffffffff != crc:
  497. raise IntegrityError('Segment checksum mismatch')
  498. if tag != TAG_PUT or id != key:
  499. raise IntegrityError('Invalid segment entry header')
  500. return data
  501. def write_put(self, id, data):
  502. size = len(data) + self.put_header_fmt.size
  503. fd = self.get_write_fd()
  504. offset = self.offset
  505. header = self.header_no_crc_fmt.pack(size, TAG_PUT)
  506. crc = self.crc_fmt.pack(crc32(data, crc32(id, crc32(header))) & 0xffffffff)
  507. fd.write(b''.join((crc, header, id, data)))
  508. self.offset += size
  509. return self.segment, offset
  510. def write_delete(self, id):
  511. fd = self.get_write_fd()
  512. header = self.header_no_crc_fmt.pack(self.put_header_fmt.size, TAG_DELETE)
  513. crc = self.crc_fmt.pack(crc32(id, crc32(header)) & 0xffffffff)
  514. fd.write(b''.join((crc, header, id)))
  515. self.offset += self.put_header_fmt.size
  516. return self.segment
  517. def write_commit(self):
  518. fd = self.get_write_fd(no_new=True)
  519. header = self.header_no_crc_fmt.pack(self.header_fmt.size, TAG_COMMIT)
  520. crc = self.crc_fmt.pack(crc32(header) & 0xffffffff)
  521. fd.write(b''.join((crc, header)))
  522. self.close_segment()
  523. def close_segment(self):
  524. if self._write_fd:
  525. self.segment += 1
  526. self.offset = 0
  527. self._write_fd.flush()
  528. os.fsync(self._write_fd.fileno())
  529. if hasattr(os, 'posix_fadvise'): # python >= 3.3, only on UNIX
  530. # tell the OS that it does not need to cache what we just wrote,
  531. # avoids spoiling the cache for the OS and other processes.
  532. os.posix_fadvise(self._write_fd.fileno(), 0, 0, os.POSIX_FADV_DONTNEED)
  533. self._write_fd.close()
  534. self._write_fd = None