Load Balancing

import http from "node:http";

// --- Types ---
interface Backend {
  url: string;
  healthy: boolean;
  activeConnections: number;
  totalRequests: number;
  weight: number;
}

interface LBConfig {
  backends: { url: string; weight?: number }[];
  healthCheckInterval: number;  // ms
  healthCheckPath: string;
  healthCheckTimeout: number;   // ms
  algorithm: "round-robin" | "least-connections" | "ip-hash";
}

// --- Load Balancer ---
class LoadBalancer {
  private backends: Backend[];
  private currentIndex = 0;
  private healthCheckTimer: ReturnType<typeof setInterval> | null = null;
  private config: LBConfig;

  constructor(config: LBConfig) {
    this.config = config;
    this.backends = config.backends.map((b) => ({
      url: b.url,
      healthy: true,
      activeConnections: 0,
      totalRequests: 0,
      weight: b.weight || 1,
    }));
  }

  // Pick a backend based on the configured algorithm
  private selectBackend(clientIP: string): Backend | null {
    const healthy = this.backends.filter((b) => b.healthy);
    if (healthy.length === 0) return null;

    switch (this.config.algorithm) {
      case "round-robin": {
        const backend = healthy[this.currentIndex % healthy.length];
        this.currentIndex++;
        return backend;
      }
      case "least-connections": {
        return healthy.reduce((min, b) =>
          b.activeConnections < min.activeConnections ? b : min
        );
      }
      case "ip-hash": {
        let hash = 0;
        for (let i = 0; i < clientIP.length; i++) {
          hash = (hash * 31 + clientIP.charCodeAt(i)) >>> 0;
        }
        return healthy[hash % healthy.length];
      }
    }
  }

  // Proxy a request to a backend
  async handleRequest(
    req: http.IncomingMessage,
    res: http.ServerResponse
  ): Promise<void> {
    const clientIP = req.socket.remoteAddress || "unknown";
    const backend = this.selectBackend(clientIP);

    if (!backend) {
      res.writeHead(503, { "Content-Type": "application/json" });
      res.end(JSON.stringify({ error: "No healthy backends available" }));
      return;
    }

    backend.activeConnections++;
    backend.totalRequests++;

    const targetUrl = new URL(req.url || "/", backend.url);
    const startTime = Date.now();

    try {
      const proxyRes = await fetch(targetUrl.toString(), {
        method: req.method,
        headers: {
          ...Object.fromEntries(
            Object.entries(req.headers).filter(([, v]) => v !== undefined) as [string, string][]
          ),
          "X-Forwarded-For": clientIP,
          "X-Forwarded-Proto": "http",
          "X-Real-IP": clientIP,
        },
        body: ["GET", "HEAD"].includes(req.method || "GET")
          ? undefined
          : await streamToBuffer(req),
        redirect: "manual",
      });

      // Forward response headers
      res.writeHead(proxyRes.status, Object.fromEntries(proxyRes.headers));

      // Stream response body
      if (proxyRes.body) {
        const reader = proxyRes.body.getReader();
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          res.write(value);
        }
      }
      res.end();

      const duration = Date.now() - startTime;
      console.log(
        `${req.method} ${req.url} -> ${backend.url} [${proxyRes.status}] ${duration}ms`
      );
    } catch (err) {
      console.error(`Backend ${backend.url} error:`, err);
      backend.healthy = false;
      res.writeHead(502, { "Content-Type": "application/json" });
      res.end(JSON.stringify({ error: "Bad gateway" }));
    } finally {
      backend.activeConnections--;
    }
  }

  // Health checking
  startHealthChecks(): void {
    this.healthCheckTimer = setInterval(async () => {
      const checks = this.backends.map(async (backend) => {
        try {
          const controller = new AbortController();
          const timeout = setTimeout(
            () => controller.abort(),
            this.config.healthCheckTimeout
          );

          const res = await fetch(
            `${backend.url}${this.config.healthCheckPath}`,
            { signal: controller.signal }
          );
          clearTimeout(timeout);

          const wasHealthy = backend.healthy;
          backend.healthy = res.ok;

          if (!wasHealthy && backend.healthy) {
            console.log(`Backend ${backend.url} is now HEALTHY`);
          } else if (wasHealthy && !backend.healthy) {
            console.log(`Backend ${backend.url} is now UNHEALTHY`);
          }
        } catch {
          if (backend.healthy) {
            console.log(`Backend ${backend.url} is now UNHEALTHY`);
          }
          backend.healthy = false;
        }
      });

      await Promise.all(checks);
    }, this.config.healthCheckInterval);
  }

  stopHealthChecks(): void {
    if (this.healthCheckTimer) {
      clearInterval(this.healthCheckTimer);
    }
  }

  getStatus(): object {
    return {
      backends: this.backends.map((b) => ({
        url: b.url,
        healthy: b.healthy,
        activeConnections: b.activeConnections,
        totalRequests: b.totalRequests,
      })),
    };
  }
}

function streamToBuffer(stream: http.IncomingMessage): Promise<Buffer> {
  return new Promise((resolve, reject) => {
    const chunks: Buffer[] = [];
    stream.on("data", (c) => chunks.push(c));
    stream.on("end", () => resolve(Buffer.concat(chunks)));
    stream.on("error", reject);
  });
}

// --- Start ---
const lb = new LoadBalancer({
  backends: [
    { url: "http://localhost:3001", weight: 2 },
    { url: "http://localhost:3002", weight: 1 },
    { url: "http://localhost:3003", weight: 1 },
  ],
  healthCheckInterval: 5000,
  healthCheckPath: "/health",
  healthCheckTimeout: 3000,
  algorithm: "round-robin",
});

lb.startHealthChecks();

const server = http.createServer((req, res) => {
  if (req.url === "/lb/status") {
    res.writeHead(200, { "Content-Type": "application/json" });
    res.end(JSON.stringify(lb.getStatus()));
    return;
  }
  lb.handleRequest(req, res);
});

server.listen(8080, () => {
  console.log("Load balancer listening on http://localhost:8080");
});

process.on("SIGTERM", () => {
  lb.stopHealthChecks();
  server.close();
});