@utaba/deep-memory-storage-cosmosdb 0.10.0 → 0.11.0

package/dist/index.cjs

@@ -37,7 +37,7 @@ module.exports = __toCommonJS(index_exports);
 
 // src/CosmosDbProvider.ts
 var import_node_crypto = __toESM(require("crypto"), 1);
-var import_deep_memory5 = require("@utaba/deep-memory");
+var import_deep_memory6 = require("@utaba/deep-memory");
 
 // src/CosmosDbConnection.ts
 var import_gremlin = __toESM(require("gremlin"), 1);
@@ -129,24 +129,11 @@ function isTransientError(err) {
   return false;
 }
 function getRetryAfterMs(err, attempt) {
-  const exponentialMs = Math.min(500 * Math.pow(2, attempt), 1e4);
-  const topLevel = err?.["retryAfterMs"];
-  if (typeof topLevel === "number" && topLevel > 0) {
-    return Math.max(topLevel, exponentialMs);
-  }
-  const attrs = err?.statusAttributes;
-  let hintVal;
-  if (attrs instanceof Map) {
-    hintVal = attrs.get("x-ms-retry-after-ms") ?? attrs.get("x-ms-retry-after");
-  } else if (typeof attrs === "object" && attrs !== null) {
-    const record = attrs;
-    hintVal = record["x-ms-retry-after-ms"] ?? record["x-ms-retry-after"];
+  const retryAfter = err?.["retryAfterMs"];
+  if (typeof retryAfter === "number" && retryAfter > 0) {
+    return retryAfter;
   }
-  const hintMs = Number(hintVal);
-  if (Number.isFinite(hintMs) && hintMs > 0) {
-    return Math.max(hintMs, exponentialMs);
-  }
-  return exponentialMs;
+  return Math.min(500 * Math.pow(2, attempt), 1e4);
 }
 function extractRequestCharge(resultSet) {
   const attrs = resultSet.attributes;
@@ -430,11 +417,7 @@ async function updateRepository(conn, repositoryId, updates) {
   return await getRepository(conn, repositoryId);
 }
 var DELETE_BATCH_SIZE = 500;
-function sleep2(ms) {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
-async function deleteRepository(conn, repositoryId, onProgress, options) {
-  const batchDelayMs = options?.batchDelayMs ?? 0;
+async function deleteRepository(conn, repositoryId, onProgress) {
   const entityCountResult = await conn.submit(
     "g.V().has('repositoryId', rid).has('entityType').count()",
     { rid: repositoryId }
@@ -460,7 +443,6 @@ async function deleteRepository(conn, repositoryId, onProgress, options) {
     relationshipsDeleted = totalRelationships - remainingCount;
     await onProgress?.({ entitiesDeleted, relationshipsDeleted, totalEntities, totalRelationships });
     if (remainingCount === 0) break;
-    if (batchDelayMs > 0) await sleep2(batchDelayMs);
   }
   while (true) {
     await conn.submit(
@@ -475,11 +457,9 @@ async function deleteRepository(conn, repositoryId, onProgress, options) {
     entitiesDeleted = totalEntities - remainingCount;
     await onProgress?.({ entitiesDeleted, relationshipsDeleted, totalEntities, totalRelationships });
     if (remainingCount === 0) break;
-    if (batchDelayMs > 0) await sleep2(batchDelayMs);
   }
 }
-async function deleteAllContents(conn, repositoryId, onProgress, options) {
-  const batchDelayMs = options?.batchDelayMs ?? 0;
+async function deleteAllContents(conn, repositoryId, onProgress) {
   const entityCountResult = await conn.submit(
     "g.V().has('repositoryId', rid).has('entityType').count()",
     { rid: repositoryId }
@@ -505,7 +485,6 @@ async function deleteAllContents(conn, repositoryId, onProgress, options) {
     relationshipsDeleted = totalRelationships - remainingCount;
     await onProgress?.({ entitiesDeleted, relationshipsDeleted, totalEntities, totalRelationships });
     if (remainingCount === 0) break;
-    if (batchDelayMs > 0) await sleep2(batchDelayMs);
   }
   while (true) {
     await conn.submit(
@@ -520,7 +499,6 @@ async function deleteAllContents(conn, repositoryId, onProgress, options) {
     entitiesDeleted = totalEntities - remainingCount;
     await onProgress?.({ entitiesDeleted, relationshipsDeleted, totalEntities, totalRelationships });
     if (remainingCount === 0) break;
-    if (batchDelayMs > 0) await sleep2(batchDelayMs);
   }
   return { deletedEntities: totalEntities, deletedRelationships: totalRelationships };
 }
@@ -1147,26 +1125,229 @@ function isInTimeRange(timestamp, timeRange) {
   return timestamp >= timeRange.from && timestamp <= timeRange.to;
 }
 
-// src/queries/bulk.ts
-var EXPORT_BATCH_SIZE = 100;
-var IMPORT_CONCURRENCY = 20;
-async function runWithConcurrency(items, concurrency, fn) {
+// src/queries/adaptive-import.ts
+var import_deep_memory5 = require("@utaba/deep-memory");
+var DEFAULT_MIN = 2;
+var DEFAULT_START = 5;
+var DEFAULT_MAX = 32;
+var DEFAULT_INCREASE_AFTER = 50;
+var DEFAULT_COOLDOWN_MS = 1e3;
+var DEFAULT_MAX_CONSECUTIVE_THROTTLES_AT_MIN = 10;
+function resolveOptions(opts) {
+  const min = Math.max(1, opts?.min ?? DEFAULT_MIN);
+  const max = Math.max(min, opts?.max ?? DEFAULT_MAX);
+  const start = Math.min(max, Math.max(min, opts?.start ?? DEFAULT_START));
+  const increaseAfter = Math.max(1, opts?.increaseAfter ?? DEFAULT_INCREASE_AFTER);
+  const cooldownMs = Math.max(0, opts?.cooldownMs ?? DEFAULT_COOLDOWN_MS);
+  const maxConsecutiveThrottlesAtMin = Math.max(
+    1,
+    opts?.maxConsecutiveThrottlesAtMin ?? DEFAULT_MAX_CONSECUTIVE_THROTTLES_AT_MIN
+  );
+  return {
+    min,
+    start,
+    max,
+    increaseAfter,
+    cooldownMs,
+    maxConsecutiveThrottlesAtMin,
+    onAdjust: opts?.onAdjust
+  };
+}
+var AdaptiveConcurrencyController = class {
+  opts;
+  current;
+  streak = 0;
+  cooldownUntil = 0;
+  completed = 0;
+  throttled = 0;
+  startEmitted = false;
+  consecutiveThrottlesAtMin = 0;
+  constructor(opts) {
+    this.opts = resolveOptions(opts);
+    this.current = this.opts.start;
+  }
+  /** Current target concurrency. */
+  getConcurrency() {
+    return this.current;
+  }
+  /** Configured ceiling — used by the runner to size its worker pool. */
+  getMaxConcurrency() {
+    return this.opts.max;
+  }
+  /** Total tasks that have been noted as completed (success or throttle). */
+  getCompleted() {
+    return this.completed;
+  }
+  /** Total tasks that observed at least one throttle. */
+  getThrottledCount() {
+    return this.throttled;
+  }
+  /**
+   * Earliest time (ms since epoch) at which a new task may be dispatched. Zero
+   * if there is no active cooldown.
+   */
+  getCooldownUntil() {
+    return this.cooldownUntil;
+  }
+  /**
+   * Emit the initial `start` event the first time the controller is queried
+   * for adjustment events. Kept separate so construction has no side effects.
+   */
+  emitStartIfNeeded() {
+    if (this.startEmitted) return;
+    this.startEmitted = true;
+    this.notify("start", this.current);
+  }
+  /** Record a throttle-free task completion. May trigger ramp-up. */
+  noteSuccess() {
+    this.completed++;
+    this.streak++;
+    this.consecutiveThrottlesAtMin = 0;
+    if (this.streak >= this.opts.increaseAfter && this.current < this.opts.max) {
+      const previous = this.current;
+      this.current = Math.min(this.opts.max, this.current + 1);
+      this.streak = 0;
+      this.notify("ramp-up", previous);
+    }
+  }
+  /** Record a task that observed at least one throttle. Halves concurrency. */
+  noteThrottle(now = Date.now()) {
+    this.completed++;
+    this.throttled++;
+    this.streak = 0;
+    const previous = this.current;
+    const next = Math.max(this.opts.min, Math.floor(this.current / 2));
+    this.cooldownUntil = now + this.opts.cooldownMs;
+    if (next !== previous) {
+      this.current = next;
+      this.notify("throttle", previous);
+    }
+    if (this.current === this.opts.min) {
+      this.consecutiveThrottlesAtMin++;
+    } else {
+      this.consecutiveThrottlesAtMin = 0;
+    }
+  }
+  /** Configured circuit-breaker threshold. */
+  getMaxConsecutiveThrottlesAtMin() {
+    return this.opts.maxConsecutiveThrottlesAtMin;
+  }
+  /** How many consecutive throttles have occurred at min. */
+  getConsecutiveThrottlesAtMin() {
+    return this.consecutiveThrottlesAtMin;
+  }
+  /**
+   * Whether the circuit breaker has tripped — runner should stop dispatching
+   * new tasks and surface an `ImportThrottleAbortError` to the caller.
+   */
+  shouldAbort() {
+    return this.consecutiveThrottlesAtMin >= this.opts.maxConsecutiveThrottlesAtMin;
+  }
+  notify(reason, previous) {
+    const cb = this.opts.onAdjust;
+    if (!cb) return;
+    try {
+      cb({
+        concurrency: this.current,
+        previousConcurrency: previous,
+        reason,
+        tasksCompleted: this.completed,
+        throttledCount: this.throttled
+      });
+    } catch {
+    }
+  }
+};
+async function runTaskWithUsage(fn) {
+  const parent = usageScope.getStore();
+  const taskAcc = { ru: 0, calls: 0, retries: 0 };
+  try {
+    const result = await usageScope.run(taskAcc, fn);
+    return { result, retries: taskAcc.retries };
+  } finally {
+    if (parent) {
+      parent.ru += taskAcc.ru;
+      parent.calls += taskAcc.calls;
+      parent.retries += taskAcc.retries;
+    }
+  }
+}
+function sleep2(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+async function runAdaptive(items, controller, fn) {
+  if (items.length === 0) return [];
+  controller.emitStartIfNeeded();
   const results = new Array(items.length);
   let nextIndex = 0;
-  const workers = [];
-  for (let w = 0; w < Math.min(concurrency, items.length); w++) {
-    workers.push(
-      (async () => {
-        while (nextIndex < items.length) {
-          const idx = nextIndex++;
-          results[idx] = await fn(items[idx]);
+  let inFlight = 0;
+  let aborted = false;
+  let gate = createGate();
+  function pokeGate() {
+    const old = gate;
+    gate = createGate();
+    old.resolve();
+  }
+  async function worker() {
+    while (nextIndex < items.length) {
+      if (aborted) return;
+      while (true) {
+        if (aborted || nextIndex >= items.length) return;
+        const cooldownRemaining = controller.getCooldownUntil() - Date.now();
+        const slotsAvailable = inFlight < controller.getConcurrency();
+        if (cooldownRemaining <= 0 && slotsAvailable) break;
+        if (cooldownRemaining > 0) {
+          await Promise.race([sleep2(cooldownRemaining), gate.promise]);
+        } else {
+          await gate.promise;
         }
-      })()
-    );
+      }
+      if (aborted || nextIndex >= items.length) return;
+      const idx = nextIndex++;
+      inFlight++;
+      try {
+        const { result, retries } = await runTaskWithUsage(() => fn(items[idx]));
+        results[idx] = result;
+        if (retries > 0) {
+          controller.noteThrottle();
+        } else {
+          controller.noteSuccess();
+        }
+        if (controller.shouldAbort()) {
+          aborted = true;
+        }
+      } finally {
+        inFlight--;
+        pokeGate();
+      }
+    }
+  }
+  const workerCount = Math.min(items.length, controller.getMaxConcurrency());
+  const workers = [];
+  for (let w = 0; w < workerCount; w++) {
+    workers.push(worker());
   }
   await Promise.all(workers);
+  if (aborted) {
+    throw new import_deep_memory5.ImportThrottleAbortError(
+      controller.getConcurrency(),
+      controller.getConsecutiveThrottlesAtMin(),
+      controller.getCompleted(),
+      controller.getThrottledCount()
+    );
+  }
   return results;
 }
+function createGate() {
+  let resolve;
+  const promise = new Promise((r) => {
+    resolve = r;
+  });
+  return { promise, resolve };
+}
+
+// src/queries/bulk.ts
+var EXPORT_BATCH_SIZE = 100;
 async function* exportAll(conn, repositoryId) {
   let sequence = 0;
   let cursor = "";
@@ -1227,11 +1408,12 @@ async function importBulk(conn, repositoryId, data, options) {
   let relationshipsImported = 0;
   const errors = [];
   const skipCheck = options?.skipExistenceCheck ?? false;
+  const controller = new AdaptiveConcurrencyController(options?.adaptiveConcurrency);
   for (const chunk of data) {
     if (chunk.entities && chunk.entities.length > 0) {
-      const results = await runWithConcurrency(
+      const results = await runAdaptive(
         chunk.entities,
-        IMPORT_CONCURRENCY,
+        controller,
         async (entity) => {
           try {
             if (skipCheck) {
@@ -1258,9 +1440,9 @@ async function importBulk(conn, repositoryId, data, options) {
       }
     }
     if (chunk.relationships && chunk.relationships.length > 0) {
-      const results = await runWithConcurrency(
+      const results = await runAdaptive(
         chunk.relationships,
-        IMPORT_CONCURRENCY,
+        controller,
         async (rel) => {
           try {
             if (skipCheck) {
@@ -1369,11 +1551,11 @@ var META_VERTEX_ID = "_meta:schema";
 var CosmosDbProvider = class {
   conn;
   config;
-  compiler = new import_deep_memory5.GremlinCompiler();
+  compiler = new import_deep_memory6.GremlinCompiler();
   reportUsage;
   constructor(config) {
     this.config = config;
-    this.reportUsage = (0, import_deep_memory5.createSafeSink)(config.reportUsage);
+    this.reportUsage = (0, import_deep_memory6.createSafeSink)(config.reportUsage);
     this.conn = new CosmosDbConnection({
       endpoint: config.endpoint,
       key: config.key,
@@ -1427,8 +1609,8 @@ var CosmosDbProvider = class {
    * ever reaching query construction.
    */
   assertValidRepositoryId(repositoryId) {
-    if (!(0, import_deep_memory5.isValidUuid)(repositoryId)) {
-      throw new import_deep_memory5.InvalidInputError(
+    if (!(0, import_deep_memory6.isValidUuid)(repositoryId)) {
+      throw new import_deep_memory6.InvalidInputError(
         "repositoryId",
         `repositoryId must be a valid v4 UUID; got '${repositoryId}'`
       );
@@ -1508,7 +1690,7 @@ var CosmosDbProvider = class {
         schemaVersion: SCHEMA_VERSION
       };
     } catch (err) {
-      throw new import_deep_memory5.ProviderError(
+      throw new import_deep_memory6.ProviderError(
         `Failed to ensure CosmosDB schema: ${err instanceof Error ? err.message : String(err)}`,
         "Verify the CosmosDB REST endpoint is accessible and the Gremlin endpoint is reachable."
       );
@@ -1554,20 +1736,18 @@ var CosmosDbProvider = class {
   }
   async deleteRepository(repositoryId, onProgress) {
     this.assertValidRepositoryId(repositoryId);
-    const batchDelayMs = this.config.deleteBatchDelayMs ?? 0;
     return this.track(
       "deleteRepository",
       repositoryId,
-      () => deleteRepository(this.conn, repositoryId, onProgress, { batchDelayMs })
+      () => deleteRepository(this.conn, repositoryId, onProgress)
     );
   }
   async deleteAllContents(repositoryId, onProgress) {
     this.assertValidRepositoryId(repositoryId);
-    const batchDelayMs = this.config.deleteBatchDelayMs ?? 0;
     return this.track(
       "deleteAllContents",
       repositoryId,
-      () => deleteAllContents(this.conn, repositoryId, onProgress, { batchDelayMs })
+      () => deleteAllContents(this.conn, repositoryId, onProgress)
     );
   }
   async getRepositoryStats(repositoryId) {
@@ -1921,7 +2101,7 @@ var CosmosDbProvider = class {
               const props = obj;
               if (props["entityType"]) {
                 const stored = entityFromGremlin(props);
-                pathEntities.push((0, import_deep_memory5.projectEntity)(stored, detailLevel));
+                pathEntities.push((0, import_deep_memory6.projectEntity)(stored, detailLevel));
               }
             }
             if (pathEntities.length > 0) {
@@ -1936,7 +2116,7 @@ var CosmosDbProvider = class {
       } else {
         for (const item of result.items) {
           const stored = entityFromGremlin(item);
-          entities.push((0, import_deep_memory5.projectEntity)(stored, detailLevel));
+          entities.push((0, import_deep_memory6.projectEntity)(stored, detailLevel));
         }
       }
       const limit = spec.limit ?? 50;
@@ -1962,7 +2142,7 @@ var CosmosDbProvider = class {
         queryMetadata
       };
     } catch (err) {
-      throw new import_deep_memory5.ProviderError(
+      throw new import_deep_memory6.ProviderError(
         `Gremlin traversal failed: ${err instanceof Error ? err.message : String(err)}`,
         "Check the traversal spec and ensure the CosmosDB connection is healthy."
       );
@@ -2023,7 +2203,7 @@ var CosmosDbProvider = class {
         queryMetadata
       };
     } catch (err) {
-      throw new import_deep_memory5.ProviderError(
+      throw new import_deep_memory6.ProviderError(
         `Native Gremlin query failed: ${err instanceof Error ? err.message : String(err)}`,
         "Verify the Gremlin query syntax is valid for CosmosDB."
       );
@@ -2093,7 +2273,7 @@ async function cosmosRestBatchDelete(restBase, key, database, container, partiti
   });
   if (response.status !== 200 && response.status !== 207) {
     const text = await response.text();
-    throw new import_deep_memory5.ProviderError(`CosmosDB batch delete failed (${response.status}): ${text}`);
+    throw new import_deep_memory6.ProviderError(`CosmosDB batch delete failed (${response.status}): ${text}`);
   }
 }
 // Annotate the CommonJS export names for ESM import in node: