@mui/internal-docs-infra/pipeline/loadServerTypesMeta/socketClient.mjs

MUI Infra - internal documentation creation tools.

/**
 * Socket client for communicating with the TypeScript language service worker
 *
 * When a worker starts, it checks if another worker is already running via a socket.
 * If so, it forwards requests to that worker instead of processing them locally.
 *
 * On Unix systems, this uses Unix domain sockets.
 * On Windows, this uses named pipes (which Node.js net module supports transparently).
 */

import { connect } from 'node:net';
import { mkdir, stat } from 'node:fs/promises';
import { createHash } from 'node:crypto';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import lockfile from 'proper-lockfile';
const isWindows = process.platform === 'win32';

/**
 * Short, stable hash of the current project directory. Used to scope shared
 * temp directories (CI runners, system tmp) so concurrent docs-infra processes
 * from different projects don't collide on the same socket/lock files.
 */
const projectHash = createHash('sha256').update(process.cwd()).digest('hex').slice(0, 8);

/**
 * Get the default socket directory.
 * On Unix: Prefers CI-specific temp directories, then falls back to system temp.
 * On Windows: Not used for the socket path itself (named pipes don't need directories).
 */
function getDefaultSocketDir() {
  // CI environments often have dedicated temp directories that work better
  return process.env.RUNNER_TEMP ??
  // GitHub Actions
  process.env.AGENT_TEMPDIRECTORY ??
  // Azure Pipelines
  tmpdir();
}

/**
 * Get the effective socket directory for Unix sockets and lock files.
 * An explicit `socketDir` is always used as-is (assumed to be project-scoped,
 * e.g. inside `.next/`). When no `socketDir` is given, shared temp directories
 * (CI runner temp or system tmp) are namespaced with a short hash of the project
 * directory so concurrent docs-infra processes from different projects don't
 * collide on the same socket/lock files.
 * @param socketDir - Optional custom directory for socket files
 */
function getEffectiveSocketDir(socketDir) {
  if (socketDir) {
    return socketDir;
  }
  return `${getDefaultSocketDir()}/mui-docs-infra-${projectHash}`;
}

/**
 * Get the path to the IPC endpoint (Unix socket or Windows named pipe)
 * @param socketDir - Optional custom directory for socket files (Unix only)
 */
export function getSocketPath(socketDir) {
  if (isWindows) {
    // Windows named pipe using extended-length path format
    // Uses effective socket dir to ensure uniqueness per project (prevents conflicts between parallel builds)
    return join('\\\\?\\pipe', getEffectiveSocketDir(socketDir), 'types');
  }
  const dir = getEffectiveSocketDir(socketDir);
  return join(dir, 'types.sock');
}

/**
 * Get the path to the lock file used for server election
 * @param socketDir - Optional custom directory for socket files
 */
export function getLockPath(socketDir) {
  const dir = getEffectiveSocketDir(socketDir);
  return join(dir, 'types.lock');
}

/**
 * Ensure the socket directory exists
 * @param socketDir - Optional custom directory for socket files
 */
export async function ensureSocketDir(socketDir) {
  const dir = getEffectiveSocketDir(socketDir);
  await mkdir(dir, {
    recursive: true
  });
}

/**
 * Check if a file exists
 */
async function fileExists(path) {
  try {
    await stat(path);
    return true;
  } catch {
    return false;
  }
}

/**
 * Try to connect to a named pipe (Windows)
 * @returns true if connection succeeded, false otherwise
 */
function tryConnectToPipe(socketPath) {
  return new Promise(resolve => {
    const socket = connect(socketPath);
    socket.on('connect', () => {
      socket.destroy();
      resolve(true);
    });
    socket.on('error', () => {
      resolve(false);
    });
  });
}

/**
 * Sleep for a given duration
 */
function sleep(ms) {
  return new Promise(resolve => {
    setTimeout(resolve, ms);
  });
}

/**
 * Wait for the IPC endpoint to become available.
 * On Unix: Polls the filesystem for the socket file to appear. We avoid
 * `fs.watch` here because on macOS it does not reliably fire events when a
 * unix domain socket file is created.
 * On Windows: Polls by attempting to connect to the named pipe.
 * @param socketDir - Optional custom directory for socket files (Unix only)
 * @param timeoutMs - Timeout in milliseconds (default: 5000)
 */
export async function waitForSocketFile(socketDir, timeoutMs = 5000) {
  const socketPath = getSocketPath(socketDir);
  const pollInterval = 50;
  const startTime = Date.now();
  if (isWindows) {
    while (Date.now() - startTime < timeoutMs) {
      // eslint-disable-next-line no-await-in-loop
      if (await tryConnectToPipe(socketPath)) {
        return;
      }
      // eslint-disable-next-line no-await-in-loop
      await sleep(pollInterval);
    }
    throw new Error(`Named pipe did not become available within ${timeoutMs}ms`);
  }

  // Ensure the directory exists so the first stat doesn't fail spuriously
  await mkdir(getEffectiveSocketDir(socketDir), {
    recursive: true
  });
  while (Date.now() - startTime < timeoutMs) {
    // eslint-disable-next-line no-await-in-loop
    if (await fileExists(socketPath)) {
      return;
    }
    // eslint-disable-next-line no-await-in-loop
    await sleep(pollInterval);
  }
  throw new Error(`Socket file did not appear within ${timeoutMs}ms`);
}

// Store the release function globally so we can call it when needed
let lockReleaseFunction = null;

/**
 * Try to acquire the server lock using proper-lockfile
 * Returns true if successfully acquired (this worker should be server)
 * @param socketDir - Optional custom directory for socket files
 */
export async function tryAcquireServerLock(socketDir) {
  const lockPath = getLockPath(socketDir);

  // Ensure the directory exists
  await ensureSocketDir(socketDir);
  try {
    // Try to acquire the lock with no retries (immediate check)
    // Stale locks will be detected after 3 seconds (server should start quickly)
    lockReleaseFunction = await lockfile.lock(lockPath, {
      retries: 0,
      // Don't retry, just check once
      stale: 3000,
      // Consider lock stale after 3 seconds
      realpath: false // Don't resolve symlinks (file doesn't need to exist)
    });
    return true;
  } catch (error) {
    // Lock is already held by another worker
    if (error.code === 'ELOCKED') {
      return false;
    }
    // Other errors should be logged but still return false
    return false;
  }
}

/**
 * Release the server lock
 */
export async function releaseServerLock() {
  if (lockReleaseFunction) {
    try {
      await lockReleaseFunction();
      lockReleaseFunction = null;
    } catch (error) {
      // Ignore errors during cleanup
    }
  }
}

/**
 * Check if there's an existing worker socket file
 * Note: The socket server will clean up stale sockets on startup
 * @param socketDir - Optional custom directory for socket files
 */
export async function hasExistingWorker(socketDir) {
  return fileExists(getSocketPath(socketDir));
}

/**
 * Client for communicating with an existing worker via socket
 */
export class SocketClient {
  socket = null;
  messageId = 0;
  pendingRequests = new Map();
  buffer = '';
  constructor(socketDir) {
    this.socketDir = socketDir;
  }

  /**
   * Connect to the worker socket with retry logic
   */
  async connect(retryCount = 0, maxRetries = 10, retryDelay = 50) {
    const socketPath = getSocketPath(this.socketDir);
    try {
      await this.attemptConnect(socketPath);
    } catch (error) {
      // If we've exhausted retries, throw the error
      if (retryCount >= maxRetries - 1) {
        throw error;
      }

      // Wait before retrying
      await new Promise(resolve => {
        setTimeout(resolve, retryDelay);
      });

      // Recursive retry
      await this.connect(retryCount + 1, maxRetries, retryDelay);
    }
  }

  /**
   * Attempt to connect to the socket
   */
  attemptConnect(socketPath) {
    return new Promise((resolve, reject) => {
      this.socket = connect(socketPath);
      this.socket.on('connect', () => {
        // Remove error listener after successful connection
        this.socket?.removeAllListeners('error');
        resolve();
      });
      this.socket.on('error', error => {
        // Clean up and reject
        this.socket?.destroy();
        this.socket = null;
        reject(error);
      });
      this.socket.on('data', data => {
        this.handleData(data);
      });
      this.socket.on('end', () => {
        this.socket = null;
      });
    });
  }

  /**
   * Handle incoming data from socket.
   * Optimized to avoid O(n²) behavior on large messages: only split the buffer
   * when the incoming chunk actually contains a newline delimiter.
   */
  handleData(data) {
    const chunk = data.toString();
    this.buffer += chunk;

    // Fast path: skip expensive split if this chunk has no message boundary
    if (!chunk.includes('\n')) {
      return;
    }

    // Process complete messages (delimited by newlines)
    const messages = this.buffer.split('\n');
    this.buffer = messages.pop() || '';
    for (const messageStr of messages) {
      if (!messageStr.trim()) {
        continue;
      }
      try {
        const message = JSON.parse(messageStr);
        const pending = this.pendingRequests.get(message.id);
        if (pending) {
          this.pendingRequests.delete(message.id);
          if (message.type === 'success') {
            pending.resolve(message.data);
          } else {
            pending.reject(new Error(message.data?.error || 'Unknown error'));
          }
        }
      } catch (error) {
        console.error('[SocketClient] Failed to parse message:', error);
      }
    }
  }

  /**
   * Send a request to the worker
   */
  async sendRequest(request) {
    if (!this.socket) {
      throw new Error('Not connected to worker socket');
    }
    const id = `req-${this.messageId}`;
    this.messageId += 1;
    return new Promise((resolve, reject) => {
      this.pendingRequests.set(id, {
        resolve,
        reject
      });
      const message = {
        id,
        type: 'process-types',
        data: request
      };
      this.socket.write(`${JSON.stringify(message)}\n`);

      // Timeout after 5 minutes
      setTimeout(() => {
        if (this.pendingRequests.has(id)) {
          this.pendingRequests.delete(id);
          reject(new Error('Request timeout'));
        }
      }, 5 * 60 * 1000);
    });
  }

  /**
   * Close the connection
   */
  close() {
    if (this.socket) {
      this.socket.end();
      this.socket = null;
    }
  }
}