core/benchmarks/sharded-delta-compact.js

import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { CondensedGraph } from '../src/core/CondensedGraph.js';
import { ShardedSnapshotBuilder } from '../src/core/shards/ShardedSnapshotBuilder.js';
import { ShardedSnapshot } from '../src/core/shards/ShardedSnapshot.js';
import { FileShardStorage } from '../src/core/shards/FileShardStorage.js';
import { DeltaShardBinary } from '../src/core/shards/DeltaShardBinary.js';
import { WaveletShardBinary } from '../src/core/shards/WaveletShardBinary.js';

function parseArgs(argv) {
  const args = new Map();
  for (let i = 2; i < argv.length; i++) {
    const value = argv[i];
    if (!value.startsWith('--')) continue;
    const [key, inline] = value.slice(2).split('=');
    if (inline !== undefined) {
      args.set(key, inline);
      continue;
    }
    const next = argv[i + 1];
    if (next && !next.startsWith('--')) {
      args.set(key, next);
      i++;
    } else {
      args.set(key, true);
    }
  }
  return args;
}

function makeRng(seed = 1337) {
  let state = seed >>> 0;
  return () => {
    state = (state * 1664525 + 1013904223) >>> 0;
    return state / 0xffffffff;
  };
}

function buildGraph({ edges, users, groups, docs, seed }) {
  const graph = new CondensedGraph();
  const rng = makeRng(seed);
  for (let i = 0; i < edges; i++) {
    const userId = Math.floor(rng() * users);
    const groupId = Math.floor(rng() * groups);
    const docId = Math.floor(rng() * docs);
    graph.addEdge(`user:${userId}`, 'member', `group:${groupId}`);
    graph.addEdge(`group:${groupId}`, 'viewer', `doc:${docId}`);
  }
  graph.finalizePerfectHash();
  graph.finalizeWaveletAdjacency({ dropAdjacencyList: true });
  return graph;
}

function pickDeltaEdges(snapshot, relationId, count, seed) {
  const rng = makeRng(seed);
  const nodeCount = snapshot.nodeCount;
  const adds = [];
  const removes = [];
  for (let i = 0; i < count; i++) {
    const srcId = Math.floor(rng() * nodeCount);
    const edges = snapshot.getOutEdgesSync(srcId, relationId);
    if (edges.length && rng() < 0.5) {
      const edge = edges[Math.floor(rng() * edges.length)];
      removes.push({ srcId, dstId: edge.dst });
    } else {
      const dstId = Math.floor(rng() * nodeCount);
      adds.push({ srcId, dstId, possBits: 65535, relBits: 65535 });
    }
  }
  return { adds, removes };
}

function localizeDelta(snapshot, shard, srcId) {
  const localSource = snapshot._localSource(srcId, shard);
  return localSource;
}

function shardKeyForEdge(snapshot, relationId, direction, srcId) {
  const shardMeta = snapshot._selectShardMeta(relationId, direction, srcId);
  return shardMeta ? shardMeta.cacheKey : null;
}

function buildDeltaLayer(snapshot, storage, layerDir, relationId, delta, name) {
  fs.mkdirSync(layerDir, { recursive: true });
  const layerShards = new Map();

  for (const add of delta.adds) {
    const key = shardKeyForEdge(snapshot, relationId, 'out', add.srcId);
    if (!key) continue;
    const shardMeta = snapshot._cacheIndex.get(key);
    const entry = layerShards.get(key) || { shardMeta, additions: [], removals: [] };
    const localSource = localizeDelta(snapshot, shardMeta, add.srcId);
    entry.additions.push({ srcLocal: localSource, otherId: add.dstId, possBits: add.possBits, relBits: add.relBits });
    layerShards.set(key, entry);
  }

  for (const rem of delta.removes) {
    const key = shardKeyForEdge(snapshot, relationId, 'out', rem.srcId);
    if (!key) continue;
    const shardMeta = snapshot._cacheIndex.get(key);
    const entry = layerShards.get(key) || { shardMeta, additions: [], removals: [] };
    const localSource = localizeDelta(snapshot, shardMeta, rem.srcId);
    entry.removals.push({ srcLocal: localSource, otherId: rem.dstId });
    layerShards.set(key, entry);
  }

  const layerManifest = { name, shards: [] };
  for (const [cacheKey, entry] of layerShards.entries()) {
    const shardKey = `delta-${name}-${entry.shardMeta.key}`;
    const buffer = DeltaShardBinary.serialize({
      relationId: entry.shardMeta.relationId,
      direction: entry.shardMeta.direction,
      rangeStart: entry.shardMeta.rangeStart,
      rangeEnd: entry.shardMeta.rangeEnd,
      nodeCount: snapshot.nodeCount,
      additions: entry.additions,
      removals: entry.removals
    });
    fs.writeFileSync(path.join(layerDir, shardKey), new Uint8Array(buffer));
    layerManifest.shards.push({
      key: shardKey,
      relationId: entry.shardMeta.relationId,
      direction: entry.shardMeta.direction,
      rangeStart: entry.shardMeta.rangeStart,
      rangeEnd: entry.shardMeta.rangeEnd,
      cacheKey
    });
  }

  return layerManifest;
}

function compactLayer(snapshot, layerStorage, layer, outputDir) {
  fs.mkdirSync(outputDir, { recursive: true });
  let compacted = 0;
  for (const shardMeta of layer.shards) {
    const base = snapshot._cacheIndex.get(shardMeta.cacheKey);
    if (!base) continue;
    const shard = snapshot._loadShardSync(base.relationId, base.direction, base.rangeStart);
    if (!shard) continue;

    const deltaBuffer = layerStorage.getSync(shardMeta.key);
    if (!deltaBuffer) continue;
    const deltaShard = DeltaShardBinary.deserialize(deltaBuffer);
    const rangeSize = shard.rangeEnd - shard.rangeStart;
    const sources = new Array(rangeSize);
    for (let localSource = 0; localSource < rangeSize; localSource++) {
      const range = snapshot._rangeForSource(shard, localSource);
      const list = [];
      if (range) {
        for (let pos = range.start; pos < range.end; pos++) {
          list.push({ otherId: shard.dstIds[pos], possBits: shard.possBits[pos], relBits: shard.relBits[pos] });
        }
      }
      sources[localSource] = list;
    }

    for (const removal of deltaShard.removals) {
      const list = sources[removal.srcLocal];
      if (!list) continue;
      const idx = list.findIndex((item) => item.otherId === removal.otherId);
      if (idx !== -1) list.splice(idx, 1);
    }

    for (const addition of deltaShard.additions) {
      const list = sources[addition.srcLocal] || (sources[addition.srcLocal] = []);
      list.push({ otherId: addition.otherId, possBits: addition.possBits, relBits: addition.relBits });
    }

    const buffer = WaveletShardBinary.serialize({
      relationId: shard.relationId,
      direction: shard.direction,
      rangeStart: shard.rangeStart,
      rangeEnd: shard.rangeEnd,
      nodeCount: snapshot.nodeCount,
      sources
    });
    fs.writeFileSync(path.join(outputDir, base.key), new Uint8Array(buffer));
    compacted++;
  }
  return compacted;
}

const args = parseArgs(process.argv);
const edges = Number(args.get('edges') || 200000);
const users = Number(args.get('users') || 50000);
const groups = Number(args.get('groups') || 2000);
const docs = Number(args.get('docs') || 100000);
const bucketSize = Number(args.get('bucket') || 4096);
const deltaEdges = Number(args.get('delta-edges') || 10000);
const seed = Number(args.get('seed') || 1337);

console.log('Sharded delta compact');
console.log(`  edges: ${edges}`);
console.log(`  delta edges: ${deltaEdges}`);

const graph = buildGraph({ edges, users, groups, docs, seed });
const baseDir = fs.mkdtempSync(path.join(os.tmpdir(), 'sharded-delta-'));
const baseShardDir = path.join(baseDir, 'base');
const deltaDir = path.join(baseDir, 'delta');
const compactDir = path.join(baseDir, 'compact');

const builder = new ShardedSnapshotBuilder({ bucketSize, includeDirections: ['out', 'in'], shardMode: 'range' });
const manifest = builder.build(graph, baseShardDir);
const storage = new FileShardStorage(baseShardDir);
const snapshot = new ShardedSnapshot(manifest, storage, { cacheLimit: 32, recentLimit: 128 });
snapshot.initializeSync();

const relationId = snapshot.relationIdToName.indexOf('member');
const delta = pickDeltaEdges(snapshot, relationId, deltaEdges, seed + 1);
const layerManifest = buildDeltaLayer(snapshot, storage, deltaDir, relationId, delta, 'l1');
const deltaStorage = new FileShardStorage(deltaDir);
snapshot.setDeltaLayers([{ shards: layerManifest.shards, storage: deltaStorage }]);

const compacted = compactLayer(snapshot, deltaStorage, layerManifest, compactDir);
console.log(`Compacted shards: ${compacted}`);