fast-md5-web/dist/index.js

A TypeScript project with tsup bundler for Rust WASM MD5 calculation

"use client"

// src/index.ts
import WasmInit, { Md5Calculator } from "../wasm/pkg";
import { v4 as uuidv4 } from "uuid";
var DEFAULT_CHUNK_SIZE = 8 * 1024 * 1024;
var DEFAULT_SHARED_MEMORY_SIZE = 512 * 1024 * 1024;
var MEMORY_CLEANUP_THRESHOLD = 0.8;
var Md5CalculatorPool = class {
  constructor(poolSize = 4, sharedMemoryConfig, maxConcurrentTasks) {
    this.workers = [];
    this.availableWorkers = [];
    this.pendingTasks = [];
    this.activeTasks = /* @__PURE__ */ new Map();
    this.taskCallbacks = /* @__PURE__ */ new Map();
    this.sharedMemory = null;
    this.sharedMemoryView = null;
    this.memoryBlocks = [];
    this.poolSize = poolSize;
    this.maxConcurrentTasks = maxConcurrentTasks || poolSize;
    this.sharedMemoryConfig = sharedMemoryConfig || {
      enabled: false,
      memorySize: DEFAULT_SHARED_MEMORY_SIZE,
      chunkSize: DEFAULT_CHUNK_SIZE
    };
    if (this.sharedMemoryConfig.enabled && this.isSharedArrayBufferSupported()) {
      this.initializeSharedMemory();
    }
    this.initializeWorkers();
  }
  isSharedArrayBufferSupported() {
    return typeof SharedArrayBuffer !== "undefined";
  }
  initializeSharedMemory() {
    try {
      this.sharedMemory = new SharedArrayBuffer(
        this.sharedMemoryConfig.memorySize
      );
      this.sharedMemoryView = new Uint8Array(this.sharedMemory);
      this.memoryBlocks = [
        {
          offset: 0,
          size: this.sharedMemoryConfig.memorySize,
          inUse: false
        }
      ];
    } catch (error) {
      this.sharedMemoryConfig.enabled = false;
    }
  }
  initializeWorkers() {
    for (let i = 0; i < this.poolSize; i++) {
      const worker = this.createWorker();
      this.workers.push(worker);
      this.availableWorkers.push(worker);
      if (this.sharedMemoryConfig.enabled && this.sharedMemory) {
        worker.postMessage({
          id: `init-${i}`,
          type: "init_shared_memory",
          data: {
            sharedMemory: this.sharedMemory
          }
        });
      }
    }
  }
  createWorker() {
    const worker = new Worker(new URL("./md5-worker.js", import.meta.url), {
      type: "module"
    });
    worker.onmessage = (e) => {
      const { id, type, data } = e.data;
      if (type === "result" || type === "error") {
        const callbacks = this.taskCallbacks.get(id);
        const task = this.activeTasks.get(id);
        if (callbacks) {
          this.taskCallbacks.delete(id);
          this.activeTasks.delete(id);
          if (task?.isLargeFile) {
            this.releaseSharedMemory(id);
          }
          if (type === "result") {
            callbacks.resolve(data.result);
          } else {
            callbacks.reject(new Error(data.error));
          }
        }
        this.availableWorkers.push(worker);
        this.processNextTask();
      } else if (type === "progress" && data?.progress !== void 0) {
        const callbacks = this.taskCallbacks.get(id);
        if (callbacks?.onProgress) {
          callbacks.onProgress(data.progress);
        }
      }
    };
    worker.onerror = (error) => {
      for (const [taskId, callbacks] of this.taskCallbacks.entries()) {
        callbacks.reject(
          new Error(`Worker error: ${error.message || "Unknown worker error"}`)
        );
        this.taskCallbacks.delete(taskId);
      }
      const index = this.workers.indexOf(worker);
      if (index !== -1) {
        this.workers[index] = this.createWorker();
        const availableIndex = this.availableWorkers.indexOf(worker);
        if (availableIndex !== -1) {
          this.availableWorkers[availableIndex] = this.workers[index];
        } else {
          this.availableWorkers.push(this.workers[index]);
        }
      }
    };
    return worker;
  }
  processNextTask() {
    if (this.activeTasks.size >= this.maxConcurrentTasks) {
      return;
    }
    if (this.pendingTasks.length > 0 && this.availableWorkers.length > 0) {
      this.pendingTasks.sort((a, b) => b.priority - a.priority);
      const task = this.pendingTasks.shift();
      const worker = this.availableWorkers.shift();
      this.activeTasks.set(task.id, task);
      this.taskCallbacks.set(task.id, {
        resolve: task.resolve,
        reject: task.reject,
        onProgress: task.onProgress
      });
      if (task.isLargeFile && task.data instanceof File) {
        this.processLargeFile(task, worker);
      } else {
        this.processSmallFile(task, worker);
      }
    }
  }
  async processLargeFile(task, worker) {
    const file = task.data;
    const chunkSize = this.sharedMemoryConfig.chunkSize;
    const totalChunks = Math.ceil(file.size / chunkSize);
    const memoryOffset = this.allocateSharedMemory(chunkSize, task.id);
    if (memoryOffset === -1) {
      const uint8Array = new Uint8Array(await file.arrayBuffer());
      this.processSmallFile({ ...task, data: uint8Array }, worker);
      return;
    }
    worker.postMessage({
      id: task.id,
      type: "calculate",
      data: {
        dataOffset: memoryOffset,
        dataLength: file.size,
        md5Length: task.md5Length,
        isStreamMode: true,
        totalChunks
      }
    });
    for (let i = 0; i < totalChunks; i++) {
      const start = i * chunkSize;
      const end = Math.min(start + chunkSize, file.size);
      const chunk = file.slice(start, end);
      const chunkData = new Uint8Array(await chunk.arrayBuffer());
      if (this.sharedMemoryView) {
        this.sharedMemoryView.set(chunkData, memoryOffset);
        worker.postMessage({
          id: task.id,
          type: "calculate_chunk",
          data: {
            chunkIndex: i,
            dataLength: chunkData.length,
            dataOffset: memoryOffset
          }
        });
      }
      await new Promise((resolve) => setTimeout(resolve, 0));
    }
  }
  async processSmallFile(task, worker) {
    let data;
    if (task.data instanceof File) {
      const arrayBuffer = await task.data.arrayBuffer();
      data = new Uint8Array(arrayBuffer);
    } else {
      data = task.data;
    }
    if (this.sharedMemoryConfig.enabled && this.sharedMemoryView && data.length <= this.sharedMemoryConfig.memorySize) {
      const dataOffset = this.allocateSharedMemory(data.length, task.id);
      if (dataOffset !== -1) {
        this.sharedMemoryView.set(data, dataOffset);
        worker.postMessage({
          id: task.id,
          type: "calculate",
          data: {
            dataOffset,
            dataLength: data.length,
            md5Length: task.md5Length,
            isStreamMode: false
          }
        });
        return;
      }
    }
    let buffer;
    if (data.byteLength === 0) {
      buffer = new ArrayBuffer(0);
    } else if (data.byteOffset === 0 && data.byteLength === data.buffer.byteLength) {
      if (data.buffer instanceof ArrayBuffer) {
        buffer = data.buffer;
      } else {
        buffer = new ArrayBuffer(data.byteLength);
        new Uint8Array(buffer).set(data);
      }
    } else {
      buffer = new ArrayBuffer(data.byteLength);
      new Uint8Array(buffer).set(data);
    }
    worker.postMessage(
      {
        id: task.id,
        type: "calculate",
        data: {
          fileData: buffer,
          md5Length: task.md5Length,
          isStreamMode: false
        }
      },
      [buffer]
    );
  }
  async calculateMd5(data, md5Length = 32, timeout = 6e4, onProgress, priority = 0) {
    return new Promise((resolve, reject) => {
      const taskId = uuidv4();
      const isLargeFile = data instanceof File ? data.size > DEFAULT_CHUNK_SIZE : data.length > DEFAULT_CHUNK_SIZE;
      const timeoutId = setTimeout(() => {
        const callbacks = this.taskCallbacks.get(taskId);
        if (callbacks) {
          this.taskCallbacks.delete(taskId);
          this.activeTasks.delete(taskId);
          this.releaseSharedMemory(taskId);
          reject(new Error(`MD5 calculation timeout after ${timeout}ms`));
        }
      }, timeout);
      const wrappedResolve = (result) => {
        clearTimeout(timeoutId);
        resolve(result);
      };
      const wrappedReject = (error) => {
        clearTimeout(timeoutId);
        reject(error);
      };
      const task = {
        id: taskId,
        data,
        md5Length,
        resolve: wrappedResolve,
        reject: wrappedReject,
        priority,
        isLargeFile,
        onProgress
      };
      if (this.activeTasks.size >= this.maxConcurrentTasks) {
        this.pendingTasks.push(task);
      } else if (this.availableWorkers.length > 0) {
        const worker = this.availableWorkers.shift();
        this.activeTasks.set(taskId, task);
        this.taskCallbacks.set(taskId, {
          resolve: wrappedResolve,
          reject: wrappedReject,
          onProgress
        });
        if (task.isLargeFile && task.data instanceof File) {
          this.processLargeFile(task, worker);
        } else {
          this.processSmallFile(task, worker).catch(wrappedReject);
        }
      } else {
        this.pendingTasks.push(task);
      }
    });
  }
  destroy() {
    this.workers.forEach((worker) => {
      worker.terminate();
    });
    this.workers = [];
    this.availableWorkers = [];
    this.pendingTasks = [];
    this.taskCallbacks.clear();
  }
  allocateSharedMemory(size, taskId) {
    if (!this.sharedMemoryView) {
      return -1;
    }
    for (let i = 0; i < this.memoryBlocks.length; i++) {
      const block = this.memoryBlocks[i];
      if (!block.inUse && block.size >= size) {
        block.inUse = true;
        block.taskId = taskId;
        if (block.size > size) {
          this.memoryBlocks.splice(i + 1, 0, {
            offset: block.offset + size,
            size: block.size - size,
            inUse: false
          });
          block.size = size;
        }
        return block.offset;
      }
    }
    this.defragmentMemory();
    for (let i = 0; i < this.memoryBlocks.length; i++) {
      const block = this.memoryBlocks[i];
      if (!block.inUse && block.size >= size) {
        block.inUse = true;
        block.taskId = taskId;
        if (block.size > size) {
          this.memoryBlocks.splice(i + 1, 0, {
            offset: block.offset + size,
            size: block.size - size,
            inUse: false
          });
          block.size = size;
        }
        return block.offset;
      }
    }
    return -1;
  }
  releaseSharedMemory(taskId) {
    for (const block of this.memoryBlocks) {
      if (block.taskId === taskId) {
        block.inUse = false;
        block.taskId = void 0;
      }
    }
    this.mergeAdjacentBlocks();
  }
  defragmentMemory() {
    this.mergeAdjacentBlocks();
    const totalUsed = this.memoryBlocks.filter((block) => block.inUse).reduce((sum, block) => sum + block.size, 0);
    const usageRatio = totalUsed / this.sharedMemoryConfig.memorySize;
    if (usageRatio > MEMORY_CLEANUP_THRESHOLD) {
    }
  }
  mergeAdjacentBlocks() {
    this.memoryBlocks.sort((a, b) => a.offset - b.offset);
    for (let i = 0; i < this.memoryBlocks.length - 1; i++) {
      const current = this.memoryBlocks[i];
      const next = this.memoryBlocks[i + 1];
      if (!current.inUse && !next.inUse && current.offset + current.size === next.offset) {
        current.size += next.size;
        this.memoryBlocks.splice(i + 1, 1);
        i--;
      }
    }
  }
  getPoolStatus() {
    const status = {
      totalWorkers: this.workers.length,
      availableWorkers: this.availableWorkers.length,
      pendingTasks: this.pendingTasks.length,
      activeTasks: this.activeTasks.size,
      maxConcurrentTasks: this.maxConcurrentTasks,
      sharedMemoryEnabled: this.sharedMemoryConfig.enabled
    };
    if (this.sharedMemoryConfig.enabled && this.sharedMemory) {
      const totalUsed = this.memoryBlocks.filter((block) => block.inUse).reduce((sum, block) => sum + block.size, 0);
      const freeBlocks = this.memoryBlocks.filter((block) => !block.inUse).length;
      status.sharedMemoryUsage = {
        total: this.sharedMemoryConfig.memorySize,
        used: totalUsed,
        available: this.sharedMemoryConfig.memorySize - totalUsed,
        fragmentation: freeBlocks
      };
    }
    return status;
  }
  enableSharedMemory(memorySize = DEFAULT_SHARED_MEMORY_SIZE, chunkSize = DEFAULT_CHUNK_SIZE) {
    if (!this.isSharedArrayBufferSupported()) {
      return false;
    }
    this.sharedMemoryConfig = {
      enabled: true,
      memorySize,
      chunkSize
    };
    this.initializeSharedMemory();
    this.destroy();
    this.initializeWorkers();
    return this.sharedMemoryConfig.enabled;
  }
  disableSharedMemory() {
    this.sharedMemoryConfig.enabled = false;
    this.sharedMemory = null;
    this.sharedMemoryView = null;
    this.memoryBlocks = [];
  }
  // 添加批量处理方法
  async calculateMd5Batch(files, md5Length = 32, onProgress) {
    const results = [];
    let completed = 0;
    const promises = files.map(
      (file, index) => this.calculateMd5(file, md5Length, 6e4, void 0, files.length - index).then((result) => {
        results[index] = result;
        completed++;
        if (onProgress) {
          onProgress(completed, files.length);
        }
        return result;
      })
    );
    await Promise.all(promises);
    return results;
  }
  // 添加取消任务的方法
  cancelTask(taskId) {
    const task = this.activeTasks.get(taskId);
    if (task) {
      this.activeTasks.delete(taskId);
      this.taskCallbacks.delete(taskId);
      this.releaseSharedMemory(taskId);
      task.reject(new Error("Task cancelled"));
      return true;
    }
    const pendingIndex = this.pendingTasks.findIndex((t) => t.id === taskId);
    if (pendingIndex !== -1) {
      const task2 = this.pendingTasks.splice(pendingIndex, 1)[0];
      task2.reject(new Error("Task cancelled"));
      return true;
    }
    return false;
  }
};
export {
  Md5Calculator,
  Md5CalculatorPool,
  WasmInit
};