mpackdb

19.7 KB

1. import { writeFile, readFile, open, rename } from 'fs/promises';
2. import { createWriteStream } from 'fs';
3. import { resolve } from 'path';
4. import { deserialize } from './mpack.js';
5. import { PrimaryKeyType, IndexType } from './MPackDB.js';
7. const BLOCK_SIZE = 4096; // 4KB
9. /**
10. * Manages indexes for the database (in-memory deltas + on-disk persistence)
11. */
12. export class IndexManager {
13. _dbDir;
14. _dbName;
15. _indexes;
16. _indexTypes;
17. _primaryKeyType;
18. _persistIntervalId = null;
20. _indexPaths = {};
21. _deltaIndexes = {};
22. _tombstones = {};
23. _persistThreshold = 1000;
24. _totalDeltaCount = 0;
26. /**
27. * Creates a new IndexManager
28. * @param {string} dbDir - Database directory path
29. * @param {string} dbName - Database name
30. * @param {string[]} indexes - Array of field names to index
31. * @param {Object} indexTypes - Map of field names to IndexType
32. * @param {number} primaryKeyType - Primary key type
33. */
34. constructor(dbDir, dbName, indexes, indexTypes, primaryKeyType) {
35. this._dbDir = dbDir;
36. this._dbName = dbName;
37. this._indexes = indexes;
38. this._indexTypes = indexTypes;
39. this._primaryKeyType = primaryKeyType;
40. }
42. /**
43. * Initializes the index manager (rebuilds indexes from data file)
44. * @param {string} dataPath - Path to the data file
45. * @returns {Promise<void>}
46. */
47. async init(dataPath, { forceRebuild = false } = {}) {
48. for (const field of this._indexes) {
49. this._deltaIndexes[field] = [];
50. this._tombstones[field] = new Map();
51. const indexPath = resolve(this._dbDir, `${this._dbName}.${field}.txt`);
52. this._indexPaths[field] = indexPath;
54. if (forceRebuild) {
55. await this._rebuildIndex(dataPath, field);
56. continue;
57. }
59. try {
60. const fileHandle = await open(indexPath, 'r');
61. const stats = await fileHandle.stat();
62. await fileHandle.close();
63. if (stats.size === 0) {
64. await this._rebuildIndex(dataPath, field);
65. }
66. } catch (e) {
67. if (e.code === 'ENOENT') {
68. await this._rebuildIndex(dataPath, field);
69. } else {
70. throw e;
71. }
72. }
73. }
74. }
76. /**
77. * Rebuilds an index from the data file
78. * @param {string} dataPath - Path to the data file
79. * @param {string} field - Field name to index
80. * @returns {Promise<void>}
81. * @private
82. */
83. async _rebuildIndex(dataPath, field) {
84. const tempIndexPath = `${this._indexPaths[field]}.${Date.now()}.tmp`;
85. const writeStream = createWriteStream(tempIndexPath, { flags: 'w', encoding: 'utf-8' });
87. const docIterator = this._getDocsAndLocationsFromDataFile(dataPath);
88. const entries = [];
89. for await (const { doc, loc } of docIterator) {
90. if (doc && doc[field] !== undefined) {
91. const key = doc[field];
92. entries.push({ key, loc });
93. }
94. }
96. // Sort entries before writing to the index file
97. entries.sort((a, b) => this._compareKeys(a.key, b.key));
99. for (const entry of entries) {
100. writeStream.write(`${entry.key},${entry.loc[0]},${entry.loc[1]}\n`);
101. }
103. await new Promise(resolve => writeStream.end(resolve));
104. await rename(tempIndexPath, this._indexPaths[field]);
105. }
107. /**
108. * Generator that yields documents and their locations from the data file
109. * @param {string} dataPath - Path to the data file
110. * @yields {{doc: Object, loc: Array}} Document and location tuple
111. * @private
112. */
113. async* _getDocsAndLocationsFromDataFile(dataPath) {
114. let fileHandle;
115. try {
116. fileHandle = await open(dataPath, 'r');
117. const stats = await fileHandle.stat();
118. let offset = 0;
120. while (offset < stats.size) {
121. // Read size header
122. const sizeBuffer = Buffer.alloc(4);
123. await fileHandle.read(sizeBuffer, 0, 4, offset);
124. const size = sizeBuffer.readInt32LE(0);
126. if (size <= 0 || offset + size > stats.size) {
127. break; // Corrupted or incomplete record
128. }
130. const docBuffer = Buffer.alloc(size);
131. await fileHandle.read(docBuffer, 0, size, offset);
132. const doc = deserialize(docBuffer.subarray(4));
134. yield { doc, loc: [offset, size] };
135. offset += size;
136. }
137. } catch (e) {
138. if (e.code !== 'ENOENT') {
139. throw e;
140. }
141. // if file does not exist, do nothing
142. } finally {
143. await fileHandle?.close();
144. }
145. }
147. /**
148. * Compares two keys for sorting
149. * @param {any} keyA - First key
150. * @param {any} keyB - Second key
151. * @returns {number} -1 if keyA < keyB, 0 if equal, 1 if keyA > keyB
152. * @private
153. */
154. _compareKeys(keyA, keyB) {
155. if (typeof keyA === 'number' && typeof keyB === 'number') {
156. return keyA - keyB;
157. }
158. return String(keyA).localeCompare(String(keyB));
159. }
161. /**
162. * Starts automatic periodic persistence of indexes
163. * @param {number} interval - Interval in milliseconds between auto-persists
164. * @param {number} threshold - Number of operations before triggering auto-persist
165. */
166. startAutoPersist(interval, threshold) {
167. this._persistThreshold = threshold;
169. // Start periodic persist
170. if (interval > 0) {
171. this._persistIntervalId = setInterval(async () => {
172. if (this._totalDeltaCount > 0) {
173. await this.persist();
174. }
175. }, interval);
176. // Don't keep process alive just for this timer
177. this._persistIntervalId.unref();
178. }
179. }
181. /**
182. * Closes the index manager (persists indexes and clears interval)
183. * @returns {Promise<void>}
184. */
185. async close() {
186. if (this._persistIntervalId) {
187. clearInterval(this._persistIntervalId);
188. this._persistIntervalId = null;
189. }
190. await this.persist();
191. }
193. /**
194. * Persists all in-memory index deltas to disk
195. * @returns {Promise<void>}
196. */
197. async persist() {
198. for (const field of this._indexes) {
199. let onDiskEntries = [];
200. const indexPath = this._indexPaths[field];
201. try {
202. const indexContent = await readFile(indexPath, 'utf-8');
203. onDiskEntries = indexContent.trim().split('\n').filter(Boolean).map(line => {
204. const [key, offset, length] = line.split(',');
205. return { key: this._parseKey(key, field), loc: [parseInt(offset), parseInt(length)] };
206. });
207. } catch (e) {
208. if (e.code !== 'ENOENT') throw e;
209. }
211. const fieldTombstones = this._tombstones[field];
212. const validDiskEntries = onDiskEntries.filter(e => !fieldTombstones.has(e.loc[0]));
214. const deltaEntries = this._deltaIndexes[field] || [];
215. const finalIndexData = [...validDiskEntries, ...deltaEntries];
217. finalIndexData.sort((a, b) => this._compareKeys(a.key, b.key));
219. if (finalIndexData.length > 0) {
220. const indexContent = finalIndexData.map(e => `${e.key},${e.loc[0]},${e.loc[1]}`).join('\n') + '\n';
221. await writeFile(indexPath, indexContent, 'utf-8');
222. } else {
223. await writeFile(indexPath, '', 'utf-8');
224. }
226. // Clear the in-memory changes now that they are persisted
227. this._deltaIndexes[field] = [];
228. this._tombstones[field].clear();
229. }
230. this._totalDeltaCount = 0;
231. }
233. /**
234. * Parses a key based on the field's index type
235. * @param {any} key - The key to parse
236. * @param {string} field - Field name
237. * @returns {any} Parsed key (number or string)
238. * @private
239. */
240. _parseKey(key, field) {
241. const indexType = this._indexTypes[field];
243. // Use index type if specified, otherwise fall back to primary key type logic
244. if (indexType === IndexType.NUMERIC) {
245. const num = parseInt(key, 10);
246. if (!isNaN(num)) return num;
247. } else if (indexType === IndexType.LEXICAL) {
248. return key;
249. }
251. // Fallback to primary key type for fields without explicit index type
252. if (this._primaryKeyType === PrimaryKeyType.NUMBER) {
253. const num = parseInt(key, 10);
254. if (!isNaN(num)) return num;
255. }
256. return key;
257. }
259. /**
260. * Performs a binary search on the blocks of an index file on disk.
261. * This is a direct port of Joshua Bloch's famously correct binary search
262. * algorithm, adapted for file blocks.
263. *
264. * @param {string} field The index field to search.
265. * @param {string|number} key The key to search for.
266. * @returns {Promise<number>} The starting offset in the file for the linear scan.
267. */
268. async _binarySearchOnDisk(field, key) {
269. try {
270. const fileHandle = await open(this._indexPaths[field], 'r');
271. const stats = await fileHandle.stat();
272. const totalBlocks = Math.ceil(stats.size / BLOCK_SIZE);
273. let low = 0;
274. let high = totalBlocks - 1;
275. const buffer = Buffer.alloc(1024);
277. while (low <= high) {
278. const mid = Math.floor((low + high) / 2);
279. const offset = mid * BLOCK_SIZE;
281. const { bytesRead } = await fileHandle.read(buffer, 0, 1024, offset);
282. if (bytesRead === 0) break;
284. const blockContent = buffer.subarray(0, bytesRead).toString('utf-8');
285. const lines = blockContent.split('\n');
287. const firstKeyInBlock = this._parseKey(lines[0].split(',')[0], field);
288. const lastKeyInBlock = this._parseKey(lines[lines.length - 2]?.split(',')[0] || firstKeyInBlock, field);
290. const cmpLast = this._compareKeys(lastKeyInBlock, key);
292. if (cmpLast < 0) {
293. // The entire block is smaller than the key
294. low = mid + 1;
295. } else {
296. // The key might be in this block or a previous one
297. high = mid - 1;
298. }
299. }
301. await fileHandle.close();
302. // low is the insertion point, start scan from the previous block to be safe.
303. const finalOffset = Math.max(0, low - 1) * BLOCK_SIZE;
304. return finalOffset;
305. } catch (e) {
306. if (e.code === 'ENOENT') {
307. return -1;
308. }
309. throw e;
310. }
311. }
313. /**
314. * Finds the last entry for a given field and key (checks both delta and disk)
315. * @param {string} field - Field name
316. * @param {any} key - Key value to search for
317. * @returns {Promise<Object|null>} Entry object with key and loc, or null if not found
318. */
319. async findLastEntry(field, key) {
320. const parsedKey = this._parseKey(key, field);
322. const deltaIndex = this._deltaIndexes[field] || [];
323. for (let i = deltaIndex.length - 1; i >= 0; i--) {
324. const entry = deltaIndex[i];
325. if (this._compareKeys(entry.key, parsedKey) === 0) {
326. return entry;
327. }
328. }
330. const onDiskEntry = await this._findLastEntryOnDisk(field, parsedKey);
331. if (!onDiskEntry) return null;
333. const isTombstoned = this._tombstones[field].has(onDiskEntry.loc[0]);
334. return isTombstoned ? null : onDiskEntry;
335. }
337. /**
338. * Finds the last entry for a key on disk (linear scan from binary search position)
339. * @param {string} field - Field name
340. * @param {any} key - Key to search for
341. * @returns {Promise<Object|null>} Entry object or null
342. * @private
343. */
344. async _findLastEntryOnDisk(field, key) {
345. const startOffset = await this._binarySearchOnDisk(field, key);
346. if (startOffset === -1) return null;
348. let lastMatch = null;
349. const fileHandle = await open(this._indexPaths[field], 'r');
350. const buffer = Buffer.alloc(BLOCK_SIZE);
351. let currentOffset = startOffset;
353. try {
354. const stats = await fileHandle.stat();
355. while (currentOffset < stats.size) {
356. const { bytesRead } = await fileHandle.read(buffer, 0, BLOCK_SIZE, currentOffset);
357. if (bytesRead === 0) break;
359. const blockContent = buffer.slice(0, bytesRead).toString('utf-8');
360. const lines = blockContent.split('\n');
361. let found = false;
362. for (const line of lines) {
363. if (!line) continue;
364. const [keyStr, offset, length] = line.split(',');
365. const lineKey = this._parseKey(keyStr, field);
366. if (this._compareKeys(lineKey, key) === 0) {
367. found = true;
368. lastMatch = { key: lineKey, loc: [parseInt(offset), parseInt(length)] };
369. }
370. }
371. if (found) break; // We can stop once we find the block with the key
372. currentOffset += BLOCK_SIZE;
373. }
374. } finally {
375. await fileHandle.close();
376. }
377. return lastMatch;
378. }
380. /**
381. * Adds a document to the indexes
382. * @param {Object} doc - The document to index
383. * @param {Array} loc - [offset, length] location in data file
384. */
385. insert(doc, loc) {
386. for (const field of this._indexes) {
387. if (doc[field] !== undefined) {
388. this._deltaIndexes[field].push({ key: doc[field], loc });
389. this._totalDeltaCount++;
390. }
391. }
393. // Auto-persist if threshold reached
394. if (this._totalDeltaCount >= this._persistThreshold) {
395. // Persist async without blocking
396. this.persist().catch(err => console.error('Index persist error:', err));
397. }
398. }
400. /**
401. * Removes a document from the indexes
402. * @param {Object} doc - The document to remove
403. * @param {string} primaryKeyField - Primary key field name
404. * @returns {Promise<void>}
405. */
406. async remove(doc, primaryKeyField) {
407. const pkValue = doc[primaryKeyField];
409. const deltaPkIndex = this._deltaIndexes[primaryKeyField] || [];
410. const indexInDelta = deltaPkIndex.findIndex(e => this._compareKeys(e.key, pkValue) === 0);
412. if (indexInDelta !== -1) {
413. const offsetToRemove = deltaPkIndex[indexInDelta].loc[0];
414. for (const field of this._indexes) {
415. this._deltaIndexes[field] = (this._deltaIndexes[field] || []).filter(e => e.loc[0] !== offsetToRemove);
416. }
417. } else {
418. const onDiskEntry = await this._findLastEntryOnDisk(primaryKeyField, pkValue);
419. if (onDiskEntry) {
420. for (const field of this._indexes) {
421. if (doc[field] !== undefined) {
422. this._tombstones[field].set(onDiskEntry.loc[0], true);
423. }
424. }
425. }
426. }
427. }
429. /**
430. * Gets all document locations from the first index (used for full scans)
431. * @returns {Promise<Array[]>} Array of all [offset, length] locations
432. */
433. async getAllLocations() {
434. const field = this._indexes[0];
435. const indexPath = this._indexPaths[field];
436. let onDiskEntries = [];
437. try {
438. const indexContent = await readFile(indexPath, 'utf-8');
439. onDiskEntries = indexContent.trim().split('\n').filter(Boolean).map(line => {
440. const [, offset, length] = line.split(',');
441. return { loc: [parseInt(offset), parseInt(length)] };
442. });
443. } catch (e) {
444. if (e.code !== 'ENOENT') throw e;
445. }
447. const fieldTombstones = this._tombstones[field] || new Map();
448. const validDiskEntries = onDiskEntries.filter(e => !fieldTombstones.has(e.loc[0]));
450. const deltaEntries = this._deltaIndexes[field] || [];
452. const finalEntries = new Map();
453. for (const entry of validDiskEntries) {
454. finalEntries.set(`${entry.loc[0]}:${entry.loc[1]}`, entry.loc);
455. }
456. for (const entry of deltaEntries) {
457. finalEntries.set(`${entry.loc[0]}:${entry.loc[1]}`, entry.loc);
458. }
459. return Array.from(finalEntries.values());
460. }
462. /**
463. * Gets document locations for a given field and key
464. * @param {string} field - Field name
465. * @param {any} key - Key value to search for
466. * @returns {Promise<Array[]>} Array of [offset, length] locations
467. */
468. async get(field, key) {
469. const parsedKey = this._parseKey(key, field);
470. const onDiskResults = await this._getOnDisk(field, parsedKey);
471. const deltaResults = (this._deltaIndexes[field] || []).filter(e => this._compareKeys(e.key, parsedKey) === 0);
472. const tombstonedOffsets = this._tombstones[field] || new Map();
474. const finalEntries = new Map();
476. for (const entry of onDiskResults) {
477. if (!tombstonedOffsets.has(entry.loc[0])) {
478. finalEntries.set(`${entry.loc[0]}:${entry.loc[1]}`, entry);
479. }
480. }
482. for (const entry of deltaResults) {
483. finalEntries.set(`${entry.loc[0]}:${entry.loc[1]}`, entry);
484. }
486. return Array.from(finalEntries.values());
487. }
489. /**
490. * Gets all entries for a key from disk
491. * @param {string} field - Field name
492. * @param {any} key - Key to search for
493. * @returns {Promise<Array>} Array of entry objects
494. * @private
495. */
496. async _getOnDisk(field, key) {
497. const startOffset = await this._binarySearchOnDisk(field, key);
498. if (startOffset === -1) return [];
500. const results = [];
501. const fileHandle = await open(this._indexPaths[field], 'r');
502. const buffer = Buffer.alloc(BLOCK_SIZE);
503. let currentOffset = startOffset;
505. try {
506. const stats = await fileHandle.stat();
507. while (currentOffset < stats.size) {
508. const { bytesRead } = await fileHandle.read(buffer, 0, BLOCK_SIZE, currentOffset);
509. if (bytesRead === 0) break;
511. const blockContent = buffer.subarray(0, bytesRead).toString('utf-8');
512. const lines = blockContent.split('\n');
513. let firstKeyInBlock;
514. let foundMatchInBlock = false;
516. for (const line of lines) {
517. if (!line) continue;
518. const [keyStr, locOffset, locLength] = line.split(',');
519. if (firstKeyInBlock === undefined) {
520. firstKeyInBlock = this._parseKey(keyStr, field);
521. }
522. const lineKey = this._parseKey(keyStr, field);
523. if (this._compareKeys(lineKey, key) === 0) {
524. foundMatchInBlock = true;
525. results.push({ key: lineKey, loc: [parseInt(locOffset), parseInt(locLength)] });
526. }
527. }
529. if (results.length > 0 && !foundMatchInBlock) {
530. break;
531. }
532. if (foundMatchInBlock === false && firstKeyInBlock && this._compareKeys(firstKeyInBlock, key) > 0) {
533. break;
534. }
535. currentOffset += BLOCK_SIZE;
536. }
537. } finally {
538. await fileHandle.close();
539. }
540. return results;
541. }
543. /**
544. * Yields index entries in sorted order by streaming blocks from disk.
545. * Merges with in-memory deltas, excludes tombstones.
546. * @param {string} field - Field name
547. * @param {Object} [options]
548. * @param {any} [options.from] - Start key (inclusive), uses binary search to skip ahead
549. * @param {number} [options.limit] - Max entries to yield
550. * @yields {{key: any, loc: [number, number]}}
551. */
552. async *entries(field, { from, limit } = {}) {
553. const tombstones = this._tombstones[field] || new Map();
555. // Collect and sort deltas (small, always in memory)
556. const deltas = (this._deltaIndexes[field] || [])
557. .filter(e => !tombstones.has(e.loc[0]))
558. .slice()
559. .sort((a, b) => this._compareKeys(a.key, b.key));
560. let deltaIdx = 0;
562. // Skip deltas before `from`
563. if (from !== undefined) {
564. while (deltaIdx < deltas.length && this._compareKeys(deltas[deltaIdx].key, from) < 0) {
565. deltaIdx++;
566. }
567. }
569. // Stream disk entries block by block
570. let fileHandle;
571. let fileSize = 0;
572. let startOffset = 0;
573. try {
574. fileHandle = await open(this._indexPaths[field], 'r');
575. fileSize = (await fileHandle.stat()).size;
576. if (from !== undefined) {
577. startOffset = await this._binarySearchOnDisk(field, from);
578. if (startOffset === -1) startOffset = 0;
579. }
580. } catch (e) {
581. if (e.code !== 'ENOENT') throw e;
582. }
584. const buffer = Buffer.alloc(BLOCK_SIZE);
585. let currentOffset = startOffset;
586. let count = 0;
587. const seen = new Set();
589. // Generator for disk entries
590. const self = this;
591. async function* diskEntries() {
592. if (!fileHandle) return;
593. try {
594. while (currentOffset < fileSize) {
595. const { bytesRead } = await fileHandle.read(buffer, 0, BLOCK_SIZE, currentOffset);
596. if (bytesRead === 0) break;
597. const lines = buffer.subarray(0, bytesRead).toString('utf-8').split('\n');
598. for (const line of lines) {
599. if (!line) continue;
600. const [keyStr, locOffset, locLength] = line.split(',');
601. const key = self._parseKey(keyStr, field);
602. const loc = [parseInt(locOffset), parseInt(locLength)];
603. if (tombstones.has(loc[0])) continue;
604. if (from !== undefined && self._compareKeys(key, from) < 0) continue;
605. yield { key, loc };
606. }
607. currentOffset += BLOCK_SIZE;
608. }
609. } finally {
610. await fileHandle.close();
611. }
612. }
614. // Merge disk + deltas in sorted order
615. const diskIter = diskEntries();
616. let diskNext = await diskIter.next();
618. while (true) {
619. const hasDisk = !diskNext.done;
620. const hasDelta = deltaIdx < deltas.length;
621. if (!hasDisk && !hasDelta) break;
623. let entry;
624. if (!hasDisk) {
625. entry = deltas[deltaIdx++];
626. } else if (!hasDelta) {
627. entry = diskNext.value;
628. diskNext = await diskIter.next();
629. } else if (this._compareKeys(diskNext.value.key, deltas[deltaIdx].key) <= 0) {
630. entry = diskNext.value;
631. diskNext = await diskIter.next();
632. } else {
633. entry = deltas[deltaIdx++];
634. }
636. const locKey = `${entry.loc[0]}:${entry.loc[1]}`;
637. if (seen.has(locKey)) continue;
638. seen.add(locKey);
640. yield entry;
641. count++;
642. if (limit !== undefined && count >= limit) {
643. // Clean up disk iterator
644. if (!diskNext.done) await diskIter.return();
645. return;
646. }
647. }
648. }
650. async* _getDocsByLocation(dataPath, locations) {
651. const fileHandle = await open(dataPath, 'r');
652. try {
653. for (const loc of locations) {
654. if (loc && loc.length === 2 && loc[1] > 0) {
655. const [offset, length] = loc;
656. const buffer = Buffer.alloc(length);
657. await fileHandle.read(buffer, 0, length, offset);
658. yield deserialize(buffer.subarray(4));
659. }
660. }
661. } finally {
662. await fileHandle.close();
663. }
664. }
665. }
667. export default IndexManager;