core/benchmarks/sharded-e2e-bench.js

import fs from 'node:fs';
import path from 'node:path';
import { ShardedSnapshot } from '../src/core/shards/ShardedSnapshot.js';
import { FileShardStorage } from '../src/core/shards/FileShardStorage.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 nowNs() {
  return process.hrtime.bigint();
}

function toMs(ns) {
  return Number(ns) / 1e6;
}

class ValueProvider {
  constructor(seed = 42) {
    this.rng = makeRng(seed);
    this._riskCache = new Map();
    this._limitCache = new Map();
  }

  risk(userId, docId) {
    const key = `${userId}:${docId}`;
    if (this._riskCache.has(key)) return this._riskCache.get(key);
    const h = (userId * 2654435761 + docId * 2246822519) >>> 0;
    const value = (h % 1000) / 1000;
    this._riskCache.set(key, value);
    return value;
  }

  riskLimit(docId) {
    if (this._limitCache.has(docId)) return this._limitCache.get(docId);
    const h = (docId * 2654435761) >>> 0;
    const value = 0.6 + ((h % 300) / 1000);
    this._limitCache.set(docId, value);
    return value;
  }
}

function buildRelationMap(relationIdToName) {
  const map = new Map();
  for (let i = 0; i < relationIdToName.length; i++) {
    map.set(relationIdToName[i], i);
  }
  return map;
}

function shardKey(snapshot, relationId, direction, nodeId) {
  return snapshot.getShardKeyForNode(relationId, direction, nodeId);
}

function directCheck(snapshot, relationId, srcId, dstId, shardSet, missing) {
  const key = shardKey(snapshot, relationId, 'out', srcId);
  if (key) shardSet.add(key);
  return snapshot.executeFindEdgeSync(srcId, relationId, dstId, missing) !== null;
}

function chainCheck(snapshot, relationIds, userId, docId, shardSet, missing, caches) {
  if (snapshot.sameComponent(userId, docId) === false) return false;
  const memberRel = relationIds.get('member');
  const viewerRel = relationIds.get('viewer');
  const memberDegree = snapshot.executeGetOutDegreeSync(userId, memberRel, missing);
  const viewerDegree = snapshot.executeGetInDegreeSync(docId, viewerRel, missing);
  if (memberDegree === null || viewerDegree === null) return null;

  if (viewerDegree < memberDegree) {
    const viewerKey = shardKey(snapshot, viewerRel, 'in', docId);
    if (viewerKey) shardSet.add(viewerKey);
    const cachedViewers = caches?.viewerByDoc?.get(docId);
    const viewerEdges = cachedViewers || snapshot.executeGetInEdgesSync(docId, viewerRel, missing);
    if (!viewerEdges) return null;
    if (caches && !cachedViewers) caches.viewerByDoc.set(docId, viewerEdges);
    for (const viewer of viewerEdges) {
      const groupId = viewer.src;
      const memberKey = shardKey(snapshot, memberRel, 'in', groupId);
      if (memberKey) shardSet.add(memberKey);
      const cachedMembers = caches?.memberByGroup?.get(groupId);
      const memberEdges = cachedMembers || snapshot.executeGetInEdgesSync(groupId, memberRel, missing);
      if (!memberEdges) return null;
      if (caches && !cachedMembers) caches.memberByGroup.set(groupId, memberEdges);
      for (const member of memberEdges) {
        if (member.src === userId) return true;
      }
    }
    return false;
  }

  const memberKey = shardKey(snapshot, memberRel, 'out', userId);
  if (memberKey) shardSet.add(memberKey);
  const cachedMembers = caches?.memberByUser?.get(userId);
  const memberEdges = cachedMembers || snapshot.executeGetOutEdgesSync(userId, memberRel, missing);
  if (!memberEdges) return null;
  if (caches && !cachedMembers) caches.memberByUser.set(userId, memberEdges);
  for (const member of memberEdges) {
    const groupId = member.dst;
    const viewerKey = shardKey(snapshot, viewerRel, 'out', groupId);
    if (viewerKey) shardSet.add(viewerKey);
    const cachedViewers = caches?.viewerByGroup?.get(groupId);
    const viewerEdges = cachedViewers || snapshot.executeGetOutEdgesSync(groupId, viewerRel, missing);
    if (!viewerEdges) return null;
    if (caches && !cachedViewers) caches.viewerByGroup.set(groupId, viewerEdges);
    for (const viewer of viewerEdges) {
      if (viewer.dst === docId) return true;
    }
  }
  return false;
}

function usersetComparatorCheck(snapshot, relationIds, valueProvider, docId, shardSet, missing, caches, limits, stats) {
  const viewerRel = relationIds.get('viewer');
  const memberRel = relationIds.get('member');
  const riskRel = relationIds.get('risk');

  if (snapshot.sameComponent(docId, docId) === false) return false;
  const viewerKey = shardKey(snapshot, viewerRel, 'in', docId);
  if (viewerKey) shardSet.add(viewerKey);
  const cachedViewers = caches?.viewerByDoc?.get(docId);
  const viewerEdges = cachedViewers || snapshot.executeGetInEdgesSync(docId, viewerRel, missing);
  if (!viewerEdges) return null;
  if (caches && !cachedViewers) caches.viewerByDoc.set(docId, viewerEdges);
  let maxRisk = -Infinity;
  let groupScanned = 0;
  let userScanned = 0;

  for (const viewer of viewerEdges) {
    if (limits.maxGroups && groupScanned >= limits.maxGroups) {
      if (stats) stats.cappedComparator++;
      return false;
    }
    const groupId = viewer.src;
    const memberKey = shardKey(snapshot, memberRel, 'in', groupId);
    if (memberKey) shardSet.add(memberKey);
    const cachedMembers = caches?.memberByGroup?.get(groupId);
    const memberEdges = cachedMembers || snapshot.executeGetInEdgesSync(groupId, memberRel, missing);
    if (!memberEdges) return null;
    if (caches && !cachedMembers) caches.memberByGroup.set(groupId, memberEdges);
    groupScanned++;
    for (const member of memberEdges) {
      if (limits.maxUsersTotal && userScanned >= limits.maxUsersTotal) {
        if (stats) stats.cappedComparator++;
        return false;
      }
      if (limits.maxUsersPerGroup && userScanned >= limits.maxUsersPerGroup * groupScanned) {
        break;
      }
      const userId = member.src;
      const riskShardKey = shardKey(snapshot, riskRel, 'out', userId);
      if (riskShardKey) shardSet.add(riskShardKey);
      const riskKey = `${userId}:${docId}`;
      const cachedRisk = caches?.riskByUserDoc?.get(riskKey);
      const hasRisk = cachedRisk !== undefined
        ? cachedRisk
        : snapshot.executeFindEdgeSync(userId, riskRel, docId, missing);
      if (!hasRisk && missing && missing.size > 0) return null;
      if (hasRisk) {
        if (caches && cachedRisk === undefined) caches.riskByUserDoc.set(riskKey, hasRisk);
        const risk = valueProvider.risk(userId, docId);
        if (risk > maxRisk) maxRisk = risk;
      }
      userScanned++;
    }
  }

  if (maxRisk === -Infinity) return false;
  const limit = valueProvider.riskLimit(docId);
  return maxRisk <= limit;
}

const args = parseArgs(process.argv);
const manifestPath = args.get('manifest') || 'tmp/shards-complex-1m-16384/manifest.json';
const queries = Number(args.get('queries') || 1000);
const seed = Number(args.get('seed') || 1234);

const manifest = JSON.parse(fs.readFileSync(manifestPath, 'utf8'));
const storage = new FileShardStorage(path.dirname(manifestPath));
const snapshot = new ShardedSnapshot(manifest, storage, { cacheLimit: 32, recentLimit: 128 });
snapshot.initializeSync();

const relationIds = buildRelationMap(manifest.relationIdToName || []);
const valueProvider = new ValueProvider(seed);
const rng = makeRng(seed);
const cachePaths = args.get('cache-paths') !== 'false';
const maxGroups = Number(args.get('max-groups') || 0);
const maxUsersPerGroup = Number(args.get('max-users-per-group') || 0);
const maxUsersTotal = Number(args.get('max-users-total') || 0);
const limits = {
  maxGroups: maxGroups > 0 ? maxGroups : null,
  maxUsersPerGroup: maxUsersPerGroup > 0 ? maxUsersPerGroup : null,
  maxUsersTotal: maxUsersTotal > 0 ? maxUsersTotal : null
};
const stats = { cappedComparator: 0 };
const caches = cachePaths ? {
  memberByUser: new Map(),
  viewerByGroup: new Map(),
  viewerByDoc: new Map(),
  memberByGroup: new Map(),
  riskByUserDoc: new Map()
} : null;

const userCount = 100000;
const docCount = 500000;
const userPool = Number(args.get('user-pool') || 5000);
const docPool = Number(args.get('doc-pool') || 5000);

const userIds = [];
for (let i = 0; i < Math.min(userPool, userCount); i++) {
  const id = snapshot.resolveNodeId(`user:${i}`);
  if (id !== -1) userIds.push(id);
}

const docIds = [];
for (let i = 0; i < Math.min(docPool, docCount); i++) {
  const id = snapshot.resolveNodeId(`doc:${i}`);
  if (id !== -1) docIds.push(id);
}

const directRel = relationIds.get('owner');
const memberRel = relationIds.get('member');
const viewerRel = relationIds.get('viewer');
const riskRel = relationIds.get('risk');

function sampleDirectPair() {
  for (let i = 0; i < 200; i++) {
    const userId = userIds[Math.floor(rng() * userIds.length)];
    const edges = snapshot.getOutEdgesSync(userId, directRel);
    if (edges.length > 0) {
      const edge = edges[0];
      return { userId, docId: edge.dst };
    }
  }
  return null;
}

function sampleChainPair() {
  for (let i = 0; i < 200; i++) {
    const userId = userIds[Math.floor(rng() * userIds.length)];
    const members = snapshot.getOutEdgesSync(userId, memberRel);
    if (!members.length) continue;
    const groupId = members[0].dst;
    const viewers = snapshot.getOutEdgesSync(groupId, viewerRel);
    if (!viewers.length) continue;
    return { userId, docId: viewers[0].dst };
  }
  return null;
}

function sampleComparatorDoc() {
  for (let i = 0; i < 200; i++) {
    const docId = docIds[Math.floor(rng() * docIds.length)];
    const viewers = snapshot.getInEdgesSync(docId, viewerRel);
    if (!viewers.length) continue;
    const groupId = viewers[0].src;
    const members = snapshot.getInEdgesSync(groupId, memberRel);
    if (!members.length) continue;
    const userId = members[0].src;
    const riskEdge = snapshot.findEdgeSync(userId, riskRel, docId);
    if (riskEdge) return { docId };
  }
  return null;
}

const directPairs = [];
const chainPairs = [];
const comparatorDocs = [];
while (directPairs.length < queries) {
  const pair = sampleDirectPair();
  if (!pair) break;
  directPairs.push(pair);
}
while (chainPairs.length < queries) {
  const pair = sampleChainPair();
  if (!pair) break;
  chainPairs.push(pair);
}
while (comparatorDocs.length < queries) {
  const pair = sampleComparatorDoc();
  if (!pair) break;
  comparatorDocs.push(pair);
}

let directHits = 0;
let chainHits = 0;
let comparatorHits = 0;

let directShardCount = 0;
let chainShardCount = 0;
let comparatorShardCount = 0;

let directStart = nowNs();
for (let i = 0; i < directPairs.length; i++) {
  let ok = null;
  let replans = 0;
  while (ok === null && replans < 4) {
    const missing = new Set();
    const shards = new Set();
    ok = directCheck(snapshot, directRel, directPairs[i].userId, directPairs[i].docId, shards, missing);
    if (missing.size > 0) {
      snapshot.prefetchPlanSync(missing);
      replans++;
      ok = null;
      continue;
    }
    if (ok) directHits++;
    directShardCount += shards.size;
  }
}
let directTime = nowNs() - directStart;

let chainStart = nowNs();
for (let i = 0; i < chainPairs.length; i++) {
  let ok = null;
  let replans = 0;
  while (ok === null && replans < 4) {
    const missing = new Set();
    const shards = new Set();
    ok = chainCheck(snapshot, relationIds, chainPairs[i].userId, chainPairs[i].docId, shards, missing, caches);
    if (missing.size > 0) {
      snapshot.prefetchPlanSync(missing);
      replans++;
      ok = null;
      continue;
    }
    if (ok) chainHits++;
    chainShardCount += shards.size;
  }
}
let chainTime = nowNs() - chainStart;

let comparatorStart = nowNs();
for (let i = 0; i < comparatorDocs.length; i++) {
  let ok = null;
  let replans = 0;
  while (ok === null && replans < 4) {
    const missing = new Set();
    const shards = new Set();
    ok = usersetComparatorCheck(snapshot, relationIds, valueProvider, comparatorDocs[i].docId, shards, missing, caches, limits, stats);
    if (missing.size > 0) {
      snapshot.prefetchPlanSync(missing);
      replans++;
      ok = null;
      continue;
    }
    if (ok) comparatorHits++;
    comparatorShardCount += shards.size;
  }
}
let comparatorTime = nowNs() - comparatorStart;

console.log('Sharded E2E benchmark');
console.log(`  manifest: ${manifestPath}`);
console.log(`  queries requested: ${queries}`);
console.log(`  user pool: ${userPool}`);
console.log(`  doc pool: ${docPool}`);
console.log(`  cache paths: ${cachePaths}`);
console.log(`  max groups: ${limits.maxGroups || 'none'}`);
console.log(`  max users/group: ${limits.maxUsersPerGroup || 'none'}`);
console.log(`  max users total: ${limits.maxUsersTotal || 'none'}`);
console.log('Results');
console.log(`  direct samples: ${directPairs.length}`);
console.log(`  direct avg: ${(toMs(directTime) / Math.max(1, directPairs.length) * 1000).toFixed(3)} µs`);
console.log(`  direct hits: ${directHits}`);
console.log(`  direct shards/query: ${(directShardCount / Math.max(1, directPairs.length)).toFixed(2)}`);
console.log(`  chain samples: ${chainPairs.length}`);
console.log(`  chain avg: ${(toMs(chainTime) / Math.max(1, chainPairs.length) * 1000).toFixed(3)} µs`);
console.log(`  chain hits: ${chainHits}`);
console.log(`  chain shards/query: ${(chainShardCount / Math.max(1, chainPairs.length)).toFixed(2)}`);
console.log(`  comparator samples: ${comparatorDocs.length}`);
console.log(`  comparator avg: ${(toMs(comparatorTime) / Math.max(1, comparatorDocs.length) * 1000).toFixed(3)} µs`);
console.log(`  comparator hits: ${comparatorHits}`);
console.log(`  comparator shards/query: ${(comparatorShardCount / Math.max(1, comparatorDocs.length)).toFixed(2)}`);
console.log(`  comparator capped: ${stats.cappedComparator}`);
console.log('Cache stats');
console.log(snapshot.getStats());