fix: Improve retry mechanisms for split pdf logic (#108)

Author: jordan-homan

Makefile

@@ -34,12 +34,12 @@ check:
 ## test-unit:	run unit tests
 .PHONY: test-unit
 test-unit:
-npx jest --detectOpenHandles --config jest.config.js test/unit
+npx jest --verbose --detectOpenHandles --config jest.config.js test/unit
 
 ## test-integration:	run integration tests
 .PHONY: test-integration
 test-integration:
-npx jest --detectOpenHandles --config jest.config.js test/integration
+npx jest --verbose --detectOpenHandles --config jest.config.js test/integration --forceExit
 
 ## test:	run all tests
 .PHONY: test

src/hooks/custom/SplitPdfHook.ts

@@ -1,6 +1,5 @@
 import async from "async";
 
-import { HTTPClient } from "../../lib/http.js";
 import {
  AfterErrorContext,
  AfterErrorHook,
@@ -25,10 +24,12 @@ import {
  stringToBoolean,
 } from "./utils/index.js";
 import {
+ HTTPClientExtension,
  MIN_PAGES_PER_THREAD,
  PARTITION_FORM_FILES_KEY,
  PARTITION_FORM_SPLIT_PDF_PAGE_KEY,
 } from "./common.js";
+import {retry, RetryConfig} from "../../lib/retries";
 
 /**
  * Represents a hook for splitting and sending PDF files as per page requests.
@@ -39,8 +40,7 @@ export class SplitPdfHook
  /**
  * The HTTP client used for making requests.
  */
- client: HTTPClient | undefined;
-
+ client: HTTPClientExtension | undefined;
 
  /**
  * Keeps the strict-mode setting for splitPdfPage feature.
@@ -68,9 +68,16 @@ export class SplitPdfHook
  * @returns The initialized SDK options.
  */
  sdkInit(opts: SDKInitOptions): SDKInitOptions {
- const { baseURL, client } = opts;
- this.client = client;
- return { baseURL: baseURL, client: client };
+ const { baseURL } = opts;
+ this.client = new HTTPClientExtension();
+
+ this.client.addHook("response", (res) => {
+ if (res.status != 200) {
+ console.error("Request failed with status code", `${res.status}`);
+ }
+ });
+
+ return { baseURL: baseURL, client: this.client };
  }
 
  /**
@@ -178,23 +185,48 @@ export class SplitPdfHook
  file.name,
  firstPageNumber
  );
+ const timeoutInMs = 60 * 10 * 1000;
  const req = new Request(requestClone, {
  headers,
  body,
+ signal: AbortSignal.timeout(timeoutInMs)
  });
  requests.push(req);
  setIndex+=1;
  }
 
  this.partitionSuccessfulResponses[operationID] = new Array(requests.length);
+ this.partitionFailedResponses[operationID] = new Array(requests.length);
 
  const allowFailed = this.allowFailed;
 
+ // These are the retry values from our api spec
+ // We need to hardcode them here until we're able to reuse the SDK
+ // from within this hook
+ const oneSecond = 1000;
+ const oneMinute = 1000 * 60;
+ const retryConfig = {
+ strategy: "backoff",
+ backoff: {
+ initialInterval: oneSecond * 3,
+ maxInterval: oneMinute * 12,
+ exponent: 1.88,
+ maxElapsedTime: oneMinute * 30,
+ },
+ } as RetryConfig;
+
+ const retryCodes = ["502", "503", "504"];
+
  this.partitionRequests[operationID] = async.parallelLimit(
- requests.slice(0, -1).map((req, pageIndex) => async () => {
+ requests.map((req, pageIndex) => async () => {
  const pageNumber = pageIndex + startingPageNumber;
  try {
- const response = await this.client!.request(req);
+ const response = await retry(
+ async () => {
+ return await this.client!.request(req.clone());
+ },
+ { config: retryConfig, statusCodes: retryCodes }
+ );
  if (response.status === 200) {
  (this.partitionSuccessfulResponses[operationID] as Response[])[pageIndex] =
  response.clone();
@@ -206,7 +238,7 @@ export class SplitPdfHook
  }
  }
  } catch (e) {
- console.error(`Failed to send request for page ${pageNumber}.`);
+ console.error(`Failed to send request for page ${pageNumber}.`, e);
  if (!allowFailed) {
  throw e;
  }
@@ -215,7 +247,8 @@ export class SplitPdfHook
  concurrencyLevel
  );
 
- return requests.at(-1) as Request;
+ const dummyRequest = new Request("https://no-op/");
+ return dummyRequest;
  }
 
  /**
@@ -230,28 +263,39 @@ export class SplitPdfHook
  successfulResponses: Response[],
  failedResponses: Response[]
  ): Promise<Response> {
+ let realResponse = response.clone();
+ const firstSuccessfulResponse = successfulResponses.at(0);
+ const isFakeResponse = response.headers.has("fake-response");
+ if (firstSuccessfulResponse !== undefined && isFakeResponse) {
+ realResponse = firstSuccessfulResponse.clone();
+ }
+
  let responseBody, responseStatus, responseStatusText;
  const numFailedResponses = failedResponses?.length ?? 0;
- const headers = prepareResponseHeaders(response);
+ const headers = prepareResponseHeaders(realResponse);
 
  if (!this.allowFailed && failedResponses && failedResponses.length > 0) {
  const failedResponse = failedResponses[0]?.clone();
  if (failedResponse) {
  responseBody = await failedResponse.text();
- responseStatus = failedResponse.status;
  responseStatusText = failedResponse.statusText;
  } else {
  responseBody = JSON.stringify({"details:": "Unknown error"});
- responseStatus = 503
  responseStatusText = "Unknown error"
  }
+ // if the response status is unknown or was 502, 503, 504, set back to 500 to ensure we don't cause more retries
+ responseStatus = 500;
  console.warn(
  `${numFailedResponses} requests failed. The partition operation is cancelled.`
  );
  } else {
- responseBody = await prepareResponseBody([...successfulResponses, response]);
- responseStatus = response.status
- responseStatusText = response.statusText
+ if (isFakeResponse) {
+ responseBody = await prepareResponseBody([...successfulResponses]);
+ } else {
+ responseBody = await prepareResponseBody([...successfulResponses, response]);
+ }
+ responseStatus = realResponse.status
+ responseStatusText = realResponse.statusText
  if (numFailedResponses > 0) {
  console.warn(
  `${numFailedResponses} requests failed. The results might miss some pages.`

src/hooks/custom/common.ts

@@ -1,3 +1,5 @@
+import {HTTPClient} from "../../lib/http";
+
 /**
  * Regular expression pattern for matching base hostnames in the form of "*.unstructuredapp.io".
  */
@@ -23,3 +25,22 @@ export const MAX_NUMBER_OF_PARALLEL_REQUESTS = 15;
 
 export const MIN_PAGES_PER_THREAD = 2;
 export const MAX_PAGES_PER_THREAD = 20;
+
+export class HTTPClientExtension extends HTTPClient {
+ constructor() {
+ super();
+ }
+
+ override async request(request: Request): Promise<Response> {
+ if (request.url === "https://no-op/") {
+ return new Response('{}', {
+ headers: [
+ ["fake-response", "fake-response"]
+ ],
+ status: 200,
+ statusText: 'OK_NO_OP'
+ });
+ }
+ return super.request(request);
+ }
+}

test/integration/SplitPdfHook.test.ts

@@ -4,7 +4,7 @@ import { UnstructuredClient } from "../../src";
 import { PartitionResponse } from "../../src/sdk/models/operations";
 import { PartitionParameters, Strategy } from "../../src/sdk/models/shared";
 
-const localServer = "http://localhost:8000"
+const localServer = "http://localhost:8000";
 
 describe("SplitPdfHook integration tests check splitted file is same as not splitted", () => {
  const FAKE_API_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
@@ -354,3 +354,42 @@ describe("SplitPdfHook integration tests page range parameter", () => {
  300000
  );
 });
+
+
+describe("SplitPDF succeeds for large PDF with high concurrency", () => {
+ const FAKE_API_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
+
+ it.each([
+ `${localServer}/general/v0/general`,
+ ])("succeed", async (serverURL) => {
+ const client = new UnstructuredClient({
+ serverURL: serverURL,
+ security: {
+ apiKeyAuth: FAKE_API_KEY,
+ },
+ });
+
+ const file = {
+ content: readFileSync("test/data/layout-parser-paper.pdf"),
+ fileName: "test/data/layout-parser-paper.pdf"
+ };
+
+ const requestParams: PartitionParameters = {
+ files: file,
+ splitPdfPage: true,
+ strategy: Strategy.HiRes,
+ splitPdfAllowFailed: false,
+ splitPdfConcurrencyLevel: 15
+ };
+
+ const res: PartitionResponse = await client.general.partition({
+ partitionParameters: {
+ ...requestParams
+ },
+ });
+
+ expect(res.statusCode).toEqual(200);
+ expect(res.elements?.length).toBeGreaterThan(0);
+ },
+ 300000);
+});