level-blobs/index.js

Save binary blobs in level and stream then back

var Writable = require('readable-stream/writable');
var Readable = require('readable-stream/readable');
var peek = require('level-peek');
var util = require('util');
var once = require('once');

var EMPTY = new Buffer(0);
var ENCODER = {
	encode: function(data) {
		return typeof data === 'string' ? data = new Buffer(data) : data;
	},
	decode: function(data) {
		return Buffer.isBuffer(data) ? data : new Buffer(data);
	},
	buffer: true,
	type: 'raw'
};

var noop = function() {};

var pad = function(n) {
	n = n.toString(16);
	return '00000000'.slice(0, -n.length)+n;
};

var expand = function(buf, len) {
	var tmp = new Buffer(len);
	buf.copy(tmp);
	return tmp;
};

module.exports = function(db, opts) {
	if (!opts) opts = {};

	var blobs = {};

	var blockSize = opts.blockSize || 65536;
	var maxBatch = opts.batch || 100;
	var blank = new Buffer(blockSize);

	db.put('\x00', 'ignore', noop); // memdown#12 workaround

	var reservations = {};
	var mutateBlock = function(key, offset, block, append, cb) {
		var release = function() {
			if (!--reservations[key].locks) delete reservations[key];
		};

		var onreservation = function(r) {
			r.locks++;

			if (!r.block && !offset) {
				r.block = block;
				cb(null, r.block, release);
				return;
			}

			if (!r.block) r.block = new Buffer(blockSize);
			if (r.block.length < offset + block.length) r.block = expand(r.block, offset + block.length);

			block.copy(r.block, offset);

			if (!append && offset + block.length < r.block.length) r.block = r.block.slice(0, offset+block.length);
			cb(null, r.block, release);
		};

		if (reservations[key]) return onreservation(reservations[key]);

		db.get(key, {valueEncoding:ENCODER}, function(err, block) {
			if (err && !err.notFound) return cb(err);
			if (!reservations[key]) reservations[key] = {locks:0, block:block};
			onreservation(reservations[key]);
		});
	};

	var WriteStream = function(name, opts) {
		if (!(this instanceof WriteStream)) return new WriteStream(name, opts);
		if (!opts) opts = {};

		this.name = name;
		this.blocks = [];
		this.batch = [];
		this.bytesWritten = 0;
		this.truncate = !opts.append;
		this.append = opts.append;

		this._shouldInitAppend = this.append && opts.start === undefined;
		this._destroyed = false;
		this._init(opts.start || 0);

		Writable.call(this);
	};

	util.inherits(WriteStream, Writable);

	WriteStream.prototype._init = function(start) {
		this.blockIndex = (start / blockSize) | 0;
		this.blockOffset = start - this.blockIndex * blockSize;
		this.blockLength = this.blockOffset;
	};

	WriteStream.prototype._flush = function(cb) {
		if (!this.batch.length) return cb();

		var key = this.batch[this.batch.length-1].key;
		var batch = this.batch;
		this.batch = [];

		if (!this.truncate) return db.batch(batch, cb);
		this.truncate = false;
		this._truncate(batch, key, cb);
	};

	WriteStream.prototype._truncate = function(batch, key, cb) {
		cb = once(cb);

		var dels = [];
		var keys = db.createKeyStream({
			start: key,
			end: this.name+'\xff\xff'
		});

		keys.on('error', cb);

		keys.on('data', function(key) {
			dels.push({type:'del', key:key});
		});

		keys.on('end', function() {
			dels.push.apply(dels, batch);
			db.batch(dels, cb);
		});
	};

	WriteStream.prototype._writeBlock = function(cb) {
		var block = this.blocks.length === 1 ? this.blocks[0] : Buffer.concat(this.blocks, this.blockLength - this.blockOffset);
		var index = this.blockIndex;
		var offset = this.blockOffset;
		var self = this;

		this.blockOffset = 0;
		this.blockLength = 0;
		this.blockIndex++;
		this.blocks = [];

		var key = this.name+'\xff'+pad(index);

		var append = function(block, force, cb) {
			if (block.length) {
				self.batch.push({
					type: 'put',
					key: key,
					value: block,
					valueEncoding: ENCODER
				});
			}

			if (!force && self.batch.length < maxBatch) return cb();
			return self._flush(cb);
		};

		if (!offset && block.length === blockSize) return append(block, false, cb);
		if (!offset && !this.append) return append(block, false, cb);

		// partial write
		mutateBlock(key, offset, block, this.append, function(err, block, release) {
			if (err) return cb(err);
			append(block, true, function(err) {
				release();
				cb(err);
			});
		});
	};

	WriteStream.prototype._initAppend = function(data, enc, cb) {
		var self = this;
		this._shouldInitAppend = false;
		blobs.size(this.name, function(err, size) {
			if (err) return cb(err);
			self._init(size);
			self._write(data, enc, cb);
		});
	};

	WriteStream.prototype._write = function(data, enc, cb) {
		if (!data.length || this._destroyed) return cb();
		if (this._shouldInitAppend) return this._initAppend(data, enc, cb);

		var self = this;
		var overflow;
		var free = blockSize - this.blockLength;

		var done = function(err) {
			if (err) return cb(err);
			if (overflow) return self._write(overflow, enc, cb);
			cb();
		};

		if (data.length > free) {
			overflow = data.slice(free);
			data = data.slice(0, free);
		}

		this.bytesWritten += data.length;
		this.blockLength += data.length;
		this.blocks.push(data);

		if (data.length < free) return done();
		this._writeBlock(done);
	};

	WriteStream.prototype.destroy = function() {
		if (this._destroyed) return;
		this._destroyed = true;
		process.nextTick(this.emit.bind(this, 'close'));
	};

	WriteStream.prototype.end = function(data) {
		var self = this;
		var args = arguments;

		if (data && typeof data !== 'function') {
			this.write(data);
			data = EMPTY;
		}

		this.write(EMPTY, function() {
			self._writeBlock(function(err) {
				if (err) return self.emit('error', err);
				self._flush(function(err) {
					if (err) return self.emit('error', err);
					Writable.prototype.end.apply(self, args);
				});
			});
		});
	};

	var ReadStream = function(name, opts) {
		if (!opts) opts = {};

		var self = this;

		var start = opts.start || 0;
		var blockIndex = (start / blockSize) | 0;
		var blockOffset = start - blockIndex * blockSize;
		var key = name+'\xff'+pad(blockIndex);

		this.name = name;
		this._missing = (typeof opts.end === 'number' ? opts.end : Infinity) - start + 1;
		this._paused = false;
		this._destroyed = false;

		this._reader = db.createReadStream({
			start: key,
			end: name+'\xff\xff',
			valueEncoding: ENCODER
		});

		var onblock = function(val) {
			key = name+'\xff'+pad(++blockIndex);

			if (!self._missing) return false;

			if (blockOffset) {
				val = val.slice(blockOffset);
				blockOffset = 0;
				if (!val.length) return true;
			}

			if (val.length > self._missing) val = val.slice(0, self._missing);

			self._missing -= val.length;
			self._pause(!self.push(val));

			return !!self._missing;
		};

		this._reader.on('data', function(data) {
			while (data.key > key) {
				if (!onblock(blank)) return;
			}

			onblock(data.value);
		});

		this._reader.on('error', function(err) {
			self.emit('error', err);
		});

		this._reader.on('end', function() {
			self.push(null);
		});

		Readable.call(this);
	};

	util.inherits(ReadStream, Readable);

	ReadStream.prototype.destroy = function() {
		if (this._destroyed) return;
		this._destroyed = true;
		this._reader.destroy();
		process.nextTick(this.emit.bind(this, 'close'));
	};

	ReadStream.prototype._pause = function(paused) {
		if (this._paused === paused) return;
		this._paused = paused;
		if (this._paused) this._reader.pause();
		else this._reader.resume();
	};

	ReadStream.prototype._read = function() {
		this._pause(false);
	};

	blobs.remove = function(name, cb) {
		cb = once(cb || noop);

		var batch = [];
		var keys = db.createKeyStream({
			start: name+'\xff',
			end: name+'\xff\xff'
		});

		keys.on('error', cb);

		keys.on('data', function(key) {
			batch.push({type:'del', key:key});
		});

		keys.on('end', function() {
			db.batch(batch, cb);
		});
	};

	blobs.size = function(name, cb) {
		peek.last(db, {
			start: name+'\xff',
			end: name+'\xff\xff',
			valueEncoding:ENCODER
		}, function(err, latest, val) {
			if (err && err.message === 'range not found') return cb(null, 0);
			if (err) return cb(err);
			if (latest.slice(0, name.length+1) !== name+'\xff') return cb(null, 0);

			cb(null, parseInt(latest.toString().slice(name.length+1), 16) * blockSize + val.length);
		});
	};

	blobs.write = function(name, data, opts, cb) {
		if (typeof opts === 'function') return blobs.write(name, data, null, opts);
		if (!opts) opts = {};
		if (!cb) cb = noop;

		var ws = blobs.createWriteStream(name, opts);

		ws.on('error', cb);
		ws.on('finish', function() {
			cb();
		});

		ws.write(data);
		ws.end();
	}

	blobs.read = function(name, opts, cb) {
		if (typeof opts === 'function') return blobs.read(name, null, opts);
		if (!opts) opts = {};

		var rs = blobs.createReadStream(name, opts);
		var list = [];

		rs.on('error', cb);
		rs.on('data', function(data) {
			list.push(data);
		});
		rs.on('end', function() {
			cb(null, list.length === 1 ? list[0] : Buffer.concat(list));
		});
	};

	blobs.createReadStream = function(name, opts) {
		return new ReadStream(name, opts);
	};

	blobs.createWriteStream = function(name, opts) {
		return new WriteStream(name, opts);
	};

	return blobs;
};