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

package/dist/index.js

@@ -1088,26 +1088,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
+import { ImportThrottleAbortError } from "@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 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 = "";
@@ -1168,11 +1371,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) {
@@ -1199,9 +1403,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) {