rethinkdb dedup

This commit is contained in:
Noah Levitt 2015-08-20 21:46:40 +00:00
parent 0e7a7fdd69
commit e66dc3a9fb
5 changed files with 88 additions and 12 deletions

View File

@ -47,7 +47,7 @@ setuptools.setup(name='warcprox',
license='GPL',
packages=['warcprox'],
package_data={'warcprox':['version.txt']},
install_requires=['certauth>=1.1.0', 'warctools>=4.8.3'], # gdbm not in pip :(
install_requires=['certauth>=1.1.0', 'warctools>=4.8.3', 'rethinkdb'], # gdbm not in pip :(
dependency_links=['git+https://github.com/internetarchive/warctools.git#egg=warctools-4.8.3'],
tests_require=['requests>=2.0.1', 'pytest'], # >=2.0.1 for https://github.com/kennethreitz/requests/pull/1636
cmdclass = {'test': PyTest},

View File

@ -15,6 +15,9 @@ import os
import json
from hanzo import warctools
import warcprox
import rethinkdb
r = rethinkdb
import random
class DedupDb(object):
logger = logging.getLogger("warcprox.dedup.DedupDb")
@ -36,12 +39,12 @@ class DedupDb(object):
except:
pass
def save(self, key, response_record, offset):
def save(self, key, response_record):
record_id = response_record.get_header(warctools.WarcRecord.ID).decode('latin1')
url = response_record.get_header(warctools.WarcRecord.URL).decode('latin1')
date = response_record.get_header(warctools.WarcRecord.DATE).decode('latin1')
py_value = {'i':record_id, 'u':url, 'd':date}
py_value = {'id':record_id, 'url':url, 'date':date}
json_value = json.dumps(py_value, separators=(',',':'))
self.db[key] = json_value.encode('utf-8')
@ -52,9 +55,9 @@ class DedupDb(object):
if key in self.db:
json_result = self.db[key]
result = json.loads(json_result.decode('utf-8'))
result['i'] = result['i'].encode('latin1')
result['u'] = result['u'].encode('latin1')
result['d'] = result['d'].encode('latin1')
result['id'] = result['id'].encode('latin1')
result['url'] = result['url'].encode('latin1')
result['date'] = result['date'].encode('latin1')
self.logger.debug('dedup db lookup of key=%s returning %s', key, result)
return result
@ -63,3 +66,62 @@ def decorate_with_dedup_info(dedup_db, recorded_url, base32=False):
key = warcprox.digest_str(recorded_url.response_recorder.payload_digest, base32)
recorded_url.dedup_info = dedup_db.lookup(key)
class RethinkDedupDb:
logger = logging.getLogger("warcprox.dedup.RethinkDedupDb")
def __init__(self, servers=["localhost"], db="warcprox", table="dedup", shards=3, replicas=3):
self.servers = servers
self.db = db
self.table = table
self.shards = shards
self.replicas = replicas
self._ensure_db_table()
# https://github.com/rethinkdb/rethinkdb-example-webpy-blog/blob/master/model.py
# "Best practices: Managing connections: a connection per request"
def _random_server_connection(self):
server = random.choice(self.servers)
try:
host, port = server.split(":")
return r.connect(host=host, port=port)
except ValueError:
return r.connect(host=server)
def _ensure_db_table(self):
with self._random_server_connection() as conn:
dbs = r.db_list().run(conn)
if not self.db in dbs:
self.logger.info("creating rethinkdb database %s", repr(self.db))
r.db_create(self.db).run(conn)
tables = r.db(self.db).table_list().run(conn)
if not self.table in tables:
self.logger.info("creating rethinkdb table %s in database %s", repr(self.table), repr(self.db))
r.db(self.db).table_create(self.table, primary_key="key", shards=self.shards, replicas=self.replicas).run(conn)
def close(self):
pass
def sync(self):
pass
def save(self, key, response_record):
k = key.decode("utf-8") if isinstance(key, bytes) else key
record_id = response_record.get_header(warctools.WarcRecord.ID).decode('latin1')
url = response_record.get_header(warctools.WarcRecord.URL).decode('latin1')
date = response_record.get_header(warctools.WarcRecord.DATE).decode('latin1')
record = {'key':k,'url':url,'date':date,'id':record_id}
with self._random_server_connection() as conn:
result = r.db(self.db).table(self.table).insert(record,conflict="replace").run(conn)
if sorted(result.values()) != [0,0,0,0,0,1] or [result["deleted"],result["skipped"],result["errors"]] != [0,0,0]:
raise Exception("unexpected result %s saving %s", result, record)
self.logger.debug('dedup db saved %s:%s', key, record)
def lookup(self, key):
k = key.decode("utf-8") if isinstance(key, bytes) else key
with self._random_server_connection() as conn:
result = r.db(self.db).table(self.table).get(k).run(conn)
if result:
for x in result:
result[x] = result[x].encode("utf-8")
self.logger.debug('dedup db lookup of key=%s returning %s', key, result)
return result

View File

@ -20,6 +20,7 @@ import signal
import threading
import certauth.certauth
import warcprox
import re
def _build_arg_parser(prog=os.path.basename(sys.argv[0])):
arg_parser = argparse.ArgumentParser(prog=prog,
@ -55,7 +56,10 @@ def _build_arg_parser(prog=os.path.basename(sys.argv[0])):
default='sha1', help='digest algorithm, one of {}'.format(', '.join(hash_algos)))
arg_parser.add_argument('--base32', dest='base32', action='store_true',
default=False, help='write digests in Base32 instead of hex')
arg_parser.add_argument('-j', '--dedup-db-file', dest='dedup_db_file',
group = arg_parser.add_mutually_exclusive_group()
group.add_argument('--dedup-rethinkdb-url', dest='dedup_rethinkdb_url',
help='persistent deduplication rethink db url, e.g. rethinkdb://db0.foo.org,db0.foo.org:38015,db1.foo.org/warcprox/dedup')
group.add_argument('-j', '--dedup-db-file', dest='dedup_db_file',
default='./warcprox-dedup.db', help='persistent deduplication database file; empty string or /dev/null disables deduplication')
arg_parser.add_argument('--stats-db-file', dest='stats_db_file',
default='./warcprox-stats.db', help='persistent statistics database file; empty string or /dev/null disables statistics tracking')
@ -108,7 +112,17 @@ def main(argv=sys.argv):
logging.fatal(e)
exit(1)
if args.dedup_db_file in (None, '', '/dev/null'):
if args.dedup_rethinkdb_url:
m = re.fullmatch(r"rethinkdb://([^/]+)/([^/]+)/([^/]+)", args.dedup_rethinkdb_url)
if m:
servers = m.group(1).split(",")
db = m.group(2)
table = m.group(3)
dedup_db = warcprox.dedup.RethinkDedupDb(servers, db, table)
else:
logging.fatal("failed to parse dedup rethinkdb url %s", args.dedup_rethinkdb_url)
exit(1)
elif args.dedup_db_file in (None, '', '/dev/null'):
logging.info('deduplication disabled')
dedup_db = None
else:

View File

@ -32,9 +32,9 @@ class WarcRecordBuilder:
url=recorded_url.url, warc_date=warc_date,
data=response_header_block,
warc_type=warctools.WarcRecord.REVISIT,
refers_to=recorded_url.dedup_info['i'],
refers_to_target_uri=recorded_url.dedup_info['u'],
refers_to_date=recorded_url.dedup_info['d'],
refers_to=recorded_url.dedup_info['id'],
refers_to_target_uri=recorded_url.dedup_info['url'],
refers_to_date=recorded_url.dedup_info['date'],
payload_digest=warcprox.digest_str(recorded_url.response_recorder.payload_digest, self.base32),
profile=warctools.WarcRecord.PROFILE_IDENTICAL_PAYLOAD_DIGEST,
content_type=hanzo.httptools.ResponseMessage.CONTENT_TYPE,

View File

@ -70,7 +70,7 @@ class WarcWriterThread(threading.Thread):
and recorded_url.response_recorder.payload_size() > 0):
key = warcprox.digest_str(recorded_url.response_recorder.payload_digest,
self.writer_pool.default_warc_writer.record_builder.base32)
self.dedup_db.save(key, records[0], records[0].offset)
self.dedup_db.save(key, records[0])
def _save_playback_info(self, recorded_url, records):
if self.playback_index_db is not None: