douyin-archive/app/api/fetcher/browser.ts

export const runtime = "nodejs";

import { type Browser, type BrowserContext } from "playwright";
import { chromium } from "playwright-extra"; // A simple singleton manager for a persistent Chromium context with ref-counting. // Prevents concurrent tasks from closing the shared window prematurely.
import stealth from "puppeteer-extra-plugin-stealth";
chromium.use(stealth());

let contextPromise: Promise<BrowserContext> | null = null;
let context: BrowserContext | null = null;
let refCount = 0;
let idleCloseTimer: NodeJS.Timeout | null = null;

const USER_DATA_DIR = process.env.USER_DATA_DIR ?? "chrome-profile/douyin";

async function launchContext(): Promise<BrowserContext> {
  const ctx = await chromium.launchPersistentContext(USER_DATA_DIR, {
    headless: process.env.CHROMIUM_HEADLESS === "true",
    viewport: {
      width: Number(process.env.CHROMIUM_VIEWPORT_WIDTH ?? 1280),
      height: Number(process.env.CHROMIUM_VIEWPORT_HEIGHT ?? 1080),
    },
  });
  // When the context is closed externally, reset manager state
  ctx.on("close", () => {
    context = null;
    contextPromise = null;
    refCount = 0;
    if (idleCloseTimer) {
      clearTimeout(idleCloseTimer);
      idleCloseTimer = null;
    }
  });
  return ctx;
}

export async function acquireBrowserContext(): Promise<BrowserContext> {
  // Cancel any pending idle close if a new consumer arrives
  if (idleCloseTimer) {
    clearTimeout(idleCloseTimer);
    idleCloseTimer = null;
  }

  if (context) {
    refCount += 1;
    return context;
  }

  if (!contextPromise) {
    contextPromise = launchContext();
  }
  context = await contextPromise;
  refCount += 1;
  return context;
}

export async function releaseBrowserContext(options?: {
  idleMillis?: number;
}): Promise<void> {
  const idleMillis = options?.idleMillis ?? 15_000;
  refCount = Math.max(0, refCount - 1);

  if (refCount > 0 || !context) return;

  // Delay the close to allow bursty workloads to reuse the context
  if (idleCloseTimer) {
    clearTimeout(idleCloseTimer);
    idleCloseTimer = null;
  }
  idleCloseTimer = setTimeout(async () => {
    try {
      if (context && refCount === 0) {
        await context.close();
      }
    } finally {
      context = null;
      contextPromise = null;
      idleCloseTimer = null;
    }
  }, idleMillis);
}

// --- Isolated context support for per-scrape independence ---

const isolatedMap = new WeakMap<BrowserContext, Browser>();

async function getSharedStorageState(): Promise<any | undefined> {
  try {
    const shared = await acquireBrowserContext();
    const state = await shared.storageState();
    // Do not force-close immediately; keep ref-counting behavior
    await releaseBrowserContext();
    return state;
  } catch {
    // If shared context not available, proceed without storageState
    return undefined;
  }
}

/**
 * Create a brand-new, independent Browser + Context for a single scrape.
 * Attempts to import cookies/localStorage from the shared persistent profile
 * so you remain logged-in, but isolates network events, cache and listeners.
 */
export async function acquireIsolatedContext(): Promise<BrowserContext> {
  const storageState = await getSharedStorageState();
  const browser = await chromium.launch({
    headless: process.env.CHROMIUM_HEADLESS === "true",
  });
  const ctx = await browser.newContext({
    storageState,
    viewport: {
      width: Number(process.env.CHROMIUM_VIEWPORT_WIDTH ?? 1280),
      height: Number(process.env.CHROMIUM_VIEWPORT_HEIGHT ?? 1080),
    },
  });
  isolatedMap.set(ctx, browser);
  ctx.on("close", () => {
    const b = isolatedMap.get(ctx);
    if (b) {
      b.close().catch(() => {});
      isolatedMap.delete(ctx);
    }
  });
  return ctx;
}

export async function releaseIsolatedContext(
  ctx: BrowserContext | null | undefined,
): Promise<void> {
  if (!ctx) return;
  try {
    await ctx.close();
  } finally {
    const b = isolatedMap.get(ctx);
    if (b) {
      await b.close().catch(() => {});
      isolatedMap.delete(ctx);
    }
  }
}