PDF.js 深入浅出 4: MessageHandler

// Source: display/api.js PDFWorker.#initialize
const worker = new Worker(workerSrc, { type: 'module' });
const messageHandler = new MessageHandler('main', 'worker', worker);

// Source: core/worker.js WorkMessageHandler.initializeFromPort
static initializeFromPort(port) {
  const handler = new MessageHandler("worker", "main", port);
  this.setup(handler, port);
  handler.send("ready", null);
}

// Source: display/api.js getDocument
const messageHandler = new MessageHandler(docId, workerId, worker.port);
const transport = new WorkerTransport(
  messageHandler,
  task,
  networkStream,
  transportParams,
  transportFactory,
  pagesMapper,
);

// Source: core/worker.js WorkMessageHandler.createDocumentHandler
const workerHandlerName = docId + '_worker';
let handler = new MessageHandler(workerHandlerName, docId, port);

// Source: shared/message_handler.js MessageHandler
constructor(sourceName, targetName, comObj) {
  this.sourceName = sourceName;
  this.targetName = targetName;
  this.comObj = comObj;
  this.callbackId = 1;
  this.streamId = 1;
  this.streamSinks = Object.create(null);
  this.streamControllers = Object.create(null);
  this.callbackCapabilities = Object.create(null);
  this.actionHandler = Object.create(null);
 
  comObj.addEventListener("message", this.#onMessage.bind(this), {
    signal: this.#messageAC.signal,
  });
}

#onMessage({ data }) {
  if (data.targetName !== this.sourceName) {
    return;
  }
  if (data.stream) {
    this.#processStreamMessage(data);
    return;
  }
  if (data.callback) {
    const callbackId = data.callbackId;
    const capability = this.callbackCapabilities[callbackId];
    if (!capability) {
      throw new Error(`Cannot resolve callback ${callbackId}`);
    }
    delete this.callbackCapabilities[callbackId];
 
    if (data.callback === CallbackKind.DATA) {
      capability.resolve(data.data);
    } else if (data.callback === CallbackKind.ERROR) {
      capability.reject(wrapReason(data.reason));
    } else {
      throw new Error("Unexpected callback case");
    }
    return;
  }
  const action = this.actionHandler[data.action];
  if (!action) {
    throw new Error(`Unknown action from worker: ${data.action}`);
  }
  if (data.callbackId) {
    const sourceName = this.sourceName,
      targetName = data.sourceName,
      comObj = this.comObj;
 
    Promise.try(action, data.data).then(
      function (result) {
        comObj.postMessage({
          sourceName,
          targetName,
          callback: CallbackKind.DATA,
          callbackId: data.callbackId,
          data: result,
        });
      },
      function (reason) {
        comObj.postMessage({
          sourceName,
          targetName,
          callback: CallbackKind.ERROR,
          callbackId: data.callbackId,
          reason: wrapReason(reason),
        });
      }
    );
    return;
  }
  if (data.streamId) {
    this.#createStreamSink(data);
    return;
  }
  action(data.data);
}

on(actionName, handler) {
  if (typeof PDFJSDev === "undefined" || PDFJSDev.test("TESTING")) {
    assert(
      typeof handler === "function",
      'MessageHandler.on: Expected "handler" to be a function.'
    );
  }
  const ah = this.actionHandler;
  if (ah[actionName]) {
    throw new Error(`There is already an actionName called "${actionName}"`);
  }
  ah[actionName] = handler;
}

send(actionName, data, transfers) {
  this.comObj.postMessage(
    {
      sourceName: this.sourceName,
      targetName: this.targetName,
      action: actionName,
      data,
    },
    transfers
  );
}

// foo
const messageHandlerFoo = new MessageHandler('foo', 'bar', comObj);
 
// bar
const messageHandlerBar = new MessageHandler('bar', 'foo', comObj);
messageHandlerBar.on('test', data => {
  console.log(data);
});
 
messageHandlerFoo.send('test', { a: 1 });
 
// result: log { a: 1 }

sendWithPromise(actionName, data, transfers) {
  const callbackId = this.callbackId++;
  const capability = Promise.withResolvers();
  this.callbackCapabilities[callbackId] = capability;
  try {
    this.comObj.postMessage(
      {
        sourceName: this.sourceName,
        targetName: this.targetName,
        action: actionName,
        callbackId,
        data,
      },
      transfers
    );
  } catch (ex) {
    capability.reject(ex);
  }
  return capability.promise;
}

sendWithStream(actionName, data, queueingStrategy, transfers) {
  const streamId = this.streamId++,
    sourceName = this.sourceName,
    targetName = this.targetName,
    comObj = this.comObj;
 
  return new ReadableStream(
    {
      start: controller => {
        const startCapability = Promise.withResolvers();
        this.streamControllers[streamId] = {
          controller,
          startCall: startCapability,
          pullCall: null,
          cancelCall: null,
          isClosed: false,
        };
        comObj.postMessage(
          {
            sourceName,
            targetName,
            action: actionName,
            streamId,
            data,
            desiredSize: controller.desiredSize,
          },
          transfers
        );
        // Return Promise for Async process, to signal success/failure.
        return startCapability.promise;
      },
 
      pull: controller => {
        const pullCapability = Promise.withResolvers();
        this.streamControllers[streamId].pullCall = pullCapability;
        comObj.postMessage({
          sourceName,
          targetName,
          stream: StreamKind.PULL,
          streamId,
          desiredSize: controller.desiredSize,
        });
        // Returning Promise will not call "pull"
        // again until current pull is resolved.
        return pullCapability.promise;
      },
 
      cancel: reason => {
        assert(reason instanceof Error, "cancel must have a valid reason");
        const cancelCapability = Promise.withResolvers();
        this.streamControllers[streamId].cancelCall = cancelCapability;
        this.streamControllers[streamId].isClosed = true;
        comObj.postMessage({
          sourceName,
          targetName,
          stream: StreamKind.CANCEL,
          streamId,
          reason: wrapReason(reason),
        });
        // Return Promise to signal success or failure.
        return cancelCapability.promise;
      },
    },
    queueingStrategy
  );
}

#createStreamSink(data) {
  const streamId = data.streamId,
    sourceName = this.sourceName,
    targetName = data.sourceName,
    comObj = this.comObj;
  const self = this,
    action = this.actionHandler[data.action];
 
  const streamSink = {
    enqueue(chunk, size = 1, transfers) {
      if (this.isCancelled) {
        return;
      }
      const lastDesiredSize = this.desiredSize;
      this.desiredSize -= size;
      // Enqueue decreases the desiredSize property of sink,
      // so when it changes from positive to negative,
      // set ready as unresolved promise.
      if (lastDesiredSize > 0 && this.desiredSize <= 0) {
        this.sinkCapability = Promise.withResolvers();
        this.ready = this.sinkCapability.promise;
      }
      comObj.postMessage(
        {
          sourceName,
          targetName,
          stream: StreamKind.ENQUEUE,
          streamId,
          chunk,
        },
        transfers
      );
    },
 
    close() {
      if (this.isCancelled) {
        return;
      }
      this.isCancelled = true;
      comObj.postMessage({
        sourceName,
        targetName,
        stream: StreamKind.CLOSE,
        streamId,
      });
      delete self.streamSinks[streamId];
    },
 
    error(reason) {
      assert(reason instanceof Error, "error must have a valid reason");
      if (this.isCancelled) {
        return;
      }
      this.isCancelled = true;
      comObj.postMessage({
        sourceName,
        targetName,
        stream: StreamKind.ERROR,
        streamId,
        reason: wrapReason(reason),
      });
    },
 
    sinkCapability: Promise.withResolvers(),
    onPull: null,
    onCancel: null,
    isCancelled: false,
    desiredSize: data.desiredSize,
    ready: null,
  };
 
  streamSink.sinkCapability.resolve();
  streamSink.ready = streamSink.sinkCapability.promise;
  this.streamSinks[streamId] = streamSink;
 
  Promise.try(action, data.data, streamSink).then(
    function () {
      comObj.postMessage({
        sourceName,
        targetName,
        stream: StreamKind.START_COMPLETE,
        streamId,
        success: true,
      });
    },
    function (reason) {
      comObj.postMessage({
        sourceName,
        targetName,
        stream: StreamKind.START_COMPLETE,
        streamId,
        reason: wrapReason(reason),
      });
    }
  );
}

const readableStream = this._transport.messageHandler.sendWithStream(
  'GetOperatorList',
  {
    pageId: this.#pagesMapper.getPageId(this._pageIndex + 1) - 1,
    pageIndex: this._pageIndex,
    intent: renderingIntent,
    cacheKey,
    annotationStorage: map,
    modifiedIds,
  },
  /* queueingStrategy = */ undefined,
  transfer,
);
const reader = readableStream.getReader();
 
const intentState = this._intentStates.get(cacheKey);
intentState.streamReader = reader;
 
const pump = () => {
  reader.read().then(
    ({ value, done }) => {
      if (done) {
        intentState.streamReader = null;
        return;
      }
      if (this._transport.destroyed) {
        return; // Ignore any pending requests if the worker was terminated.
      }
      this._renderPageChunk(value, intentState);
      pump();
    },
    reason => {
      intentState.streamReader = null;
 
      if (this._transport.destroyed) {
        return; // Ignore any pending requests if the worker was terminated.
      }
      if (intentState.operatorList) {
        // Mark operator list as complete.
        intentState.operatorList.lastChunk = true;
 
        for (const internalRenderTask of intentState.renderTasks) {
          internalRenderTask.operatorListChanged();
        }
        this.#tryCleanup();
      }
 
      if (intentState.displayReadyCapability) {
        intentState.displayReadyCapability.reject(reason);
      } else if (intentState.opListReadCapability) {
        intentState.opListReadCapability.reject(reason);
      } else {
        throw reason;
      }
    },
  );
};
pump();

handler.on(
  'GetOperatorList',
  function ({ pageId, pageIndex, intent, cacheKey, annotationStorage, modifiedIds }, sink) {
    pdfManager.getPage(pageId).then(function (page) {
      const task = new WorkerTask(`GetOperatorList: page ${pageIndex}`);
      startWorkerTask(task);
 
      // NOTE: Keep this condition in sync with the `info` helper function.
      const start = verbosity >= VerbosityLevel.INFOS ? Date.now() : 0;
 
      // Pre compile the pdf page and fetch the fonts/images.
      page
        .getOperatorList({
          handler,
          sink,
          task,
          intent,
          cacheKey,
          annotationStorage,
          modifiedIds,
          pageIndex,
        })
        .then(
          opListInfo => {
            if (start) {
              info(`${task.name}; time=${Date.now() - start}ms, len=${opListInfo.length}`);
            }
            sink.close();
          },
          reason => {
            if (task.terminated) {
              return; // ignoring errors from the terminated thread
            }
            sink.error(reason);
          },
        )
        .finally(() => {
          finishWorkerTask(task);
        });
    });
  },
);

Promise.try(action, data.data, streamSink).then(
  function () {
    comObj.postMessage({
      sourceName,
      targetName,
      stream: StreamKind.START_COMPLETE,
      streamId,
      success: true,
    });
  },
  function (reason) {
    comObj.postMessage({
      sourceName,
      targetName,
      stream: StreamKind.START_COMPLETE,
      streamId,
      reason: wrapReason(reason),
    });
  },
);

#processStreamMessage(data) {
  const streamId = data.streamId,
    sourceName = this.sourceName,
    targetName = data.sourceName,
    comObj = this.comObj;
  const streamController = this.streamControllers[streamId],
    streamSink = this.streamSinks[streamId];
 
  switch (data.stream) {
    case StreamKind.START_COMPLETE:
      if (data.success) {
        streamController.startCall.resolve();
      } else {
        streamController.startCall.reject(wrapReason(data.reason));
      }
      break;
 
    case StreamKind.PULL_COMPLETE:
      if (data.success) {
        streamController.pullCall.resolve();
      } else {
        streamController.pullCall.reject(wrapReason(data.reason));
      }
      break;
 
    case StreamKind.PULL:
      // 主线程 pull 请求 — 交给 Worker 的 onPull 处理
      if (!streamSink) {
        comObj.postMessage({ stream: StreamKind.PULL_COMPLETE, success: true });
        break;
      }
      // desiredSize 从负变正 → resolve sink.ready，解除 Worker 侧背压
      if (streamSink.desiredSize <= 0 && data.desiredSize > 0) {
        streamSink.sinkCapability.resolve();
      }
      streamSink.desiredSize = data.desiredSize;
      Promise.try(streamSink.onPull).then(
        () => comObj.postMessage({ stream: StreamKind.PULL_COMPLETE, success: true }),
        reason => comObj.postMessage({ stream: StreamKind.PULL_COMPLETE, reason })
      );
      break;
 
    case StreamKind.ENQUEUE:
      // Worker 生产了数据 — 直接塞进 ReadableStream 内部队列
      if (streamController.isClosed) break;
      streamController.controller.enqueue(data.chunk);
      break;
 
    case StreamKind.CLOSE:
      streamController.isClosed = true;
      streamController.controller.close();
      this.#deleteStreamController(streamController, streamId);
      break;
 
    case StreamKind.ERROR:
      streamController.controller.error(wrapReason(data.reason));
      this.#deleteStreamController(streamController, streamId);
      break;
 
    default:
      throw new Error("Unexpected stream case");
  }
}