worker-lb 0.0.2 → 0.1.3

package/README.md

@@ -5,13 +5,65 @@
 
 worker-lb is a small and extensible load balancer built on Cloudflare Workers. It's basically a poor mans Cloudflare Load Balancer, but running on Cloudflare Workers, because why not.
 
+Note that in the current state it actually works more like a failover/high availability solution rather than a fully fledged load balancer. You can see the load balancing status bits as coming soon below.
+
 ## Features
 
 - Failover between multiple HTTP endpoints
 - Health checks with customizable intervals and timeouts
-- Recovery function to dump requests that did not reach any healthy endpoints.
+- Recovery function to dump requests that did not reach any healthy endpoints
+- Geo steering (route by continent, country, region, or Cloudflare colo)
+- Response time based steering (coming soon)
+- Number of connections based steering (coming soon)
+
+## Installation
+
+```bash
+bun add worker-lb
+```
+
+## Quick Start
+
+```typescript
+import { Endpoint, LoadBalancer } from "worker-lb";
+
+const lb = new LoadBalancer({
+  endpoints: [
+    new Endpoint("https://api1.example.com"),
+    new Endpoint("https://api2.example.com"),
+    new Endpoint("https://api3.example.com"),
+  ],
+});
+
+export default {
+  async fetch(request: Request<unknown, IncomingRequestCfProperties>): Promise<Response> {
+    return lb.handleRequest(request);
+  },
+};
+```
 
 ## Use Cases
 
 - Simple load balancing for dirty small projects
 - Latency is not a concern but reliability is
+- Geographic routing to regional backends
+
+## Documentation
+
+- [Endpoints](./docs/endpoints.md) - Configure backend endpoints
+- [Availability Methods](./docs/availability.md) - Failover strategies
+- [Geo Steering](./docs/geo-steering.md) - Route by geographic location
+
+## Headers Added to Responses
+
+| Header | Description |
+|--------|-------------|
+| `X-Load-Balancer-Endpoint` | The endpoint URL that served the request |
+| `X-Load-Balancer-Latency` | Total request latency in milliseconds |
+| `X-Load-Balancer-Endpoint-Gather-Latency` | Time to select the endpoint in milliseconds |
+| `X-Load-Balancer-Tried-Count` | Number of endpoints tried (only on failover) |
+| `X-Load-Balancer-Tried-Endpoints` | Comma-separated endpoint URLs tried (only on failover) |
+
+## License
+
+MIT

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "worker-lb",
-  "version": "0.0.2",
+  "version": "0.1.3",
   "type": "module",
   "main": "src/index.ts",
   "module": "src/index.ts",

package/src/index.ts

@@ -1,36 +1,29 @@
+import { Endpoint, GeoEndpoint } from "./endpoints/index.ts";
+
+export { Endpoint, GeoEndpoint };
+export type {
+  EndpointOptions,
+  GeoConfig,
+  GeoEndpointOptions,
+} from "./endpoints/index.ts";
+
 /**
- * The process in which fall back is handled.
+ * The process in which failover is handled.
  *
  * `async-block` – One by one check and select the first available endpoint
  *
  * `promise.any` – Use Promise.any to select the first available endpoint that responds successfully
  *
  * `fail-forward` – Always use the first endpoint, fail over to the next only on specified failure conditions
- *
  */
 export type AvailabilityMethodType =
   | "async-block"
   | "promise.any"
   | "fail-forward";
 
-export interface AvailabilityMethodOptions {
-  "async-block": {
-    /**
-     * async-block requires a health endpoint to check availability of the service
-     *
-     * @example "/health"
-     */
-    healthCheckPathname: `/${string}`;
-  };
-  "promise.any": {
-    /**
-     * promise.any requires a health endpoint to check availability of the service
-     *
-     * @example "/health"
-     *
-     */
-    healthEndpoint: `/${string}`;
-  };
+type AvailabilityMethodOptions = {
+  "async-block": Record<string, never>;
+  "promise.any": Record<string, never>;
   "fail-forward": {
     /**
      * Statuses that should result in a failover
@@ -39,7 +32,7 @@ export interface AvailabilityMethodOptions {
      */
     failoverOnStatuses?: number[];
   };
-}
+};
 
 export type AvailabilityMethod = {
   [K in AvailabilityMethodType]: {
@@ -50,34 +43,75 @@ export type AvailabilityMethod = {
 
 export const DEFAULT_FAILOVER_STATUSES = [502, 503, 504];
 
-type RecoveryFn = (request: Request) => Promise<Response | undefined>;
+export type RecoveryContext = {
+  /**
+   * The endpoints that were tried before all failed
+   */
+  triedEndpoints: Endpoint[];
+};
+
+type RecoveryFn = (
+  request: Request,
+  context: RecoveryContext,
+) => Promise<Response | undefined | void>;
 
-export interface LoadBalancerOptions {
+type GeoSteering = {
+  type: "geo";
   /**
-   * The endpoints that you want to load balance against
-   *
-   * @example ["https://us-east-1.api.capy.lol/v1", "https://us-west-1.api.capy.lol/v1"]
+   * Default endpoints to use when no geo match is found.
+   * If not provided and no match is found, the first endpoint in the list is used.
    */
-  endpoints: string[];
+  defaultEndpoints?: Endpoint[];
+};
+
+type Steering = GeoSteering;
+
+type LoadBalancerOptionsBase = {
   /**
    * The availability method when failure happens.
    *
-   * @default { type: "fail-forward"  }
+   * @default { type: "fail-forward" }
    */
   availability?: AvailabilityMethod;
   /**
-   * A recovery function that you can customize to dump data or notify on failures. This only runs when all endpoints are unavailable.
+   * A recovery function that you can customize to dump data or notify on failures.
+   * This only runs when all endpoints are unavailable.
    *
-   * Some ideas for ya bud:
+   * Some ideas:
    * - Send a notification to a monitoring service
    * - Dump the request data to an S3/R2 bucket for replaying later.
    */
   recoveryFn?: RecoveryFn;
-}
+};
+
+type LoadBalancerOptionsWithoutSteering = LoadBalancerOptionsBase & {
+  /**
+   * The endpoints to load balance against
+   */
+  endpoints: Endpoint[];
+  steering?: never;
+};
+
+type LoadBalancerOptionsWithGeoSteering = LoadBalancerOptionsBase & {
+  /**
+   * The geo endpoints to load balance against.
+   * Each endpoint specifies which geographic regions it serves.
+   */
+  endpoints: GeoEndpoint[];
+  /**
+   * Geo steering configuration
+   */
+  steering: GeoSteering;
+};
+
+export type LoadBalancerOptions =
+  | LoadBalancerOptionsWithoutSteering
+  | LoadBalancerOptionsWithGeoSteering;
 
 export class LoadBalancer {
-  private endpoints: string[];
+  private endpoints: Endpoint[] | GeoEndpoint[];
   private availability: AvailabilityMethod;
+  private steering?: Steering;
   private recoveryFn?: RecoveryFn;
 
   constructor(options: LoadBalancerOptions) {
@@ -85,7 +119,8 @@ export class LoadBalancer {
       throw new Error("At least one endpoint is required");
     }
 
-    this.endpoints = options.endpoints.map((url) => url.replace(/\/$/, "")); // remove trailing slash
+    this.endpoints = options.endpoints;
+    this.steering = options.steering;
     this.availability = options.availability ?? {
       type: "fail-forward",
     };
@@ -102,7 +137,10 @@ export class LoadBalancer {
     this.recoveryFn = options.recoveryFn;
   }
 
-  async getAvailableEndpoint(): Promise<string> {
+  /**
+   * Get the first available endpoint based on the availability method.
+   */
+  async getAvailableEndpoint(): Promise<Endpoint> {
     switch (this.availability.type) {
       case "async-block":
         return this.getEndpointAsyncBlock();
@@ -113,20 +151,11 @@ export class LoadBalancer {
     }
   }
 
-  private async getEndpointAsyncBlock(): Promise<string> {
-    if (this.availability.type !== "async-block") {
-      throw new Error("Availability method is not async-block");
-    }
-
+  private async getEndpointAsyncBlock(): Promise<Endpoint> {
     for (const endpoint of this.endpoints) {
       try {
-        const response = await fetch(
-          endpoint + this.availability.options?.healthCheckPathname,
-          {
-            method: "GET",
-          },
-        );
-        if (response.ok) {
+        const isHealthy = await endpoint.healthCheck();
+        if (isHealthy) {
           return endpoint;
         }
       } catch {
@@ -136,19 +165,11 @@ export class LoadBalancer {
     throw new Error("No available endpoints");
   }
 
-  private async getEndpointPromiseAny(): Promise<string> {
-    if (this.availability.type !== "promise.any") {
-      throw new Error("Availability method is not promise.any");
-    }
-
-    const healthPathname = this.availability.options?.healthEndpoint;
-
+  private async getEndpointPromiseAny(): Promise<Endpoint> {
     const healthChecks = this.endpoints.map(async (endpoint) => {
-      const response = await fetch(endpoint + healthPathname, {
-        method: "GET",
-      });
-      if (!response.ok) {
-        throw new Error(`Endpoint ${endpoint} returned ${response.status}`);
+      const isHealthy = await endpoint.healthCheck();
+      if (!isHealthy) {
+        throw new Error(`Endpoint ${endpoint.url} is not healthy`);
       }
       return endpoint;
     });
@@ -160,29 +181,87 @@ export class LoadBalancer {
     }
   }
 
-  private async getEndpointsToTry(): Promise<string[]> {
+  /**
+   * Get endpoints to try based on steering and availability configuration.
+   * Returns all endpoints ordered by priority for failover:
+   * 1. Geo-matched endpoints (if geo steering enabled)
+   * 2. Default endpoints (if configured and different from matched)
+   * 3. All remaining endpoints
+   */
+  private async getEndpointsToTry(
+    request: Request<unknown, IncomingRequestCfProperties>,
+  ): Promise<Endpoint[]> {
+    if (this.steering?.type === "geo") {
+      const geoEndpoints = this.endpoints as GeoEndpoint[];
+      const matchingEndpoints = geoEndpoints.filter((endpoint) =>
+        endpoint.matchesRequest(request),
+      );
+
+      // Build ordered list: matched -> defaults -> remaining
+      const orderedEndpoints: Endpoint[] = [];
+      const seen = new Set<Endpoint>();
+
+      // Add geo-matched endpoints first
+      for (const endpoint of matchingEndpoints) {
+        orderedEndpoints.push(endpoint);
+        seen.add(endpoint);
+      }
+
+      // Add default endpoints if they're not already included
+      if (this.steering.defaultEndpoints) {
+        for (const endpoint of this.steering.defaultEndpoints) {
+          if (!seen.has(endpoint)) {
+            orderedEndpoints.push(endpoint);
+            seen.add(endpoint);
+          }
+        }
+      }
+
+      // Add all remaining endpoints for failover
+      for (const endpoint of this.endpoints) {
+        if (!seen.has(endpoint)) {
+          orderedEndpoints.push(endpoint);
+          seen.add(endpoint);
+        }
+      }
+
+      if (this.availability.type === "fail-forward") {
+        return orderedEndpoints;
+      }
+
+      // For async-block and promise.any, get the first healthy endpoint
+      const originalEndpoints = this.endpoints;
+      this.endpoints = orderedEndpoints;
+      try {
+        const healthyEndpoint = await this.getAvailableEndpoint();
+        return [healthyEndpoint];
+      } finally {
+        this.endpoints = originalEndpoints;
+      }
+    }
+
+    // No geo steering - use all endpoints
     if (this.availability.type === "fail-forward") {
       return this.endpoints;
     }
 
-    return [await this.getAvailableEndpoint()];
+    // For async-block and promise.any, get the first healthy endpoint
+    const healthyEndpoint = await this.getAvailableEndpoint();
+    return [healthyEndpoint];
   }
 
-  async handleRequest(request: Request): Promise<Response> {
+  async handleRequest(
+    request: Request<unknown, IncomingRequestCfProperties>,
+  ): Promise<Response> {
     const startTime = Date.now();
-    const endpointsToTry = await this.getEndpointsToTry();
-    const gatherEndpointLatency = Date.now();
-    const url = new URL(request.url);
+    const endpointsToTry = await this.getEndpointsToTry(request);
+    const triedEndpoints: Endpoint[] = [];
 
     for (const endpoint of endpointsToTry) {
+      triedEndpoints.push(endpoint);
+      const attemptStart = Date.now();
       try {
-        const targetUrl = endpoint + url.pathname + url.search;
-        const response = await fetch(targetUrl, {
-          method: request.method,
-          headers: request.headers,
-          body: request.body,
-          redirect: "follow",
-        });
+        const response = await endpoint.commitRequest(request);
 
         const shouldFailover =
           this.availability.type === "fail-forward" &&
@@ -193,30 +272,49 @@ export class LoadBalancer {
         if (!shouldFailover) {
           const endTime = Date.now();
           const headers = new Headers(response.headers);
-          headers.set("X-Load-Balancer-Endpoint", endpoint);
+          headers.set("X-Load-Balancer-Endpoint", endpoint.url);
           headers.set(
             "X-Load-Balancer-Latency",
             (endTime - startTime).toString(),
           );
           headers.set(
             "X-Load-Balancer-Endpoint-Gather-Latency",
-            (gatherEndpointLatency - startTime).toString(),
+            (attemptStart - startTime).toString(),
           );
 
+          // Add failover trace headers if we tried more than one endpoint
+          if (triedEndpoints.length > 1) {
+            headers.set(
+              "X-Load-Balancer-Tried-Count",
+              triedEndpoints.length.toString(),
+            );
+            headers.set(
+              "X-Load-Balancer-Tried-Endpoints",
+              triedEndpoints.map((e) => e.url).join(", "),
+            );
+          }
+
           return new Response(response.body, {
             status: response.status,
             statusText: response.statusText,
             headers,
           });
         }
-      } catch {
+      } catch (e) {
+        console.log(e);
         // Network error, try next endpoint
         continue;
       }
     }
 
     if (this.recoveryFn) {
-      await this.recoveryFn(request);
+      const recoveryResponse = await this.recoveryFn(request, {
+        triedEndpoints,
+      });
+
+      if (recoveryResponse) {
+        return recoveryResponse;
+      }
     }
 
     throw new Error("No available endpoints");