@flink-app/flink 2.0.0-alpha.72 → 2.0.0-alpha.74

package/src/FlinkApp.ts

@@ -17,7 +17,8 @@ import { FlinkAuthPlugin } from "./auth/FlinkAuthPlugin";
 import { FlinkContext } from "./FlinkContext";
 import { FlinkError, internalServerError, notFound, unauthorized } from "./FlinkErrors";
 import { FlinkRequest, Handler, HandlerFile, HttpMethod, QueryParamMetadata, RouteProps, ValidationMode } from "./FlinkHttpHandler";
-import { FlinkJobFile } from "./FlinkJob";
+import { FlinkJobFile, FlinkJobProps } from "./FlinkJob";
+import { LeaderElection, LeaderElectionOptions } from "./LeaderElection";
 import { log } from "./FlinkLog";
 import { FlinkLogFactory } from "./FlinkLogFactory";
 import { FlinkPlugin } from "./FlinkPlugin";
@@ -30,6 +31,7 @@ import { formatValidationErrors, getPathParams, isError } from "./utils";
 
 const initLog = FlinkLogFactory.createLogger("flink.init");
 const perfLog = FlinkLogFactory.createLogger("flink.perf");
+const schedulerLog = FlinkLogFactory.createLogger("flink.scheduler");
 
 const ajv = new Ajv();
 addFormats(ajv);
@@ -179,23 +181,33 @@ export interface FlinkOptions {
          */
         enabled?: boolean;
 
-        // TODO: Implement master auto assignment
-        //     /**
-        //      * If true, the master (the instance if flink app that will run jobs) will be
-        //      * automatically assigned to the first node that starts.
-        //      *
-        //      * Is persisted in database.
-        //      *
-        //      * Will throw and exception if true but no database is configured.
-        //      */
-        //     autoAssignMaster?: boolean;
-
-        //     /**
-        //      * Name of collection to be used for storing master assignment.
-        //      *
-        //      * Defaults to `flink-scheduling`
-        //      */
-        //     autoAssignCollection?: string;
+        /**
+         * Enable leader election for horizontally scaled deployments.
+         *
+         * When enabled, only one instance (the leader) will run scheduled jobs.
+         * If the leader goes down, another instance automatically takes over.
+         *
+         * Requires a database connection (`db` option) since leader election
+         * state is persisted in MongoDB. If no database is configured, a warning
+         * will be logged and jobs will run on all instances (no leader election).
+         *
+         * Set to `true` for default settings, or pass an options object to customize.
+         *
+         * @example
+         * ```ts
+         * // Use defaults (15s lease, 5s heartbeat)
+         * scheduling: { leaderElection: true }
+         *
+         * // Custom settings
+         * scheduling: {
+         *   leaderElection: {
+         *     leaseDurationMs: 30000,
+         *     heartbeatIntervalMs: 10000,
+         *   }
+         * }
+         * ```
+         */
+        leaderElection?: boolean | LeaderElectionOptions;
     };
 
     /**
@@ -337,6 +349,8 @@ export class FlinkApp<C extends FlinkContext> {
     private handlerRouteCache = new Map<string, string>();
 
     public scheduler?: ToadScheduler;
+    private allInstanceScheduler?: ToadScheduler;
+    private leaderElection?: LeaderElection;
 
     private accessLog: { enabled: boolean; format: string };
 
@@ -406,10 +420,10 @@ export class FlinkApp<C extends FlinkContext> {
         await this.initializeAgents();
         perfLog.debug(`Initialize agents took ${Date.now() - agentInitStartTime}ms`);
 
-        if (this.isSchedulingEnabled) {
+        if (this.isSchedulingEnabled && !this.leaderElectionConfig) {
             this.scheduler = new ToadScheduler();
-        } else {
-            initLog.info("🚫 Scheduling is disabled");
+        } else if (!this.isSchedulingEnabled) {
+            schedulerLog.info("Scheduling is disabled");
         }
 
         if (!this.disableHttpServer) {
@@ -459,9 +473,13 @@ export class FlinkApp<C extends FlinkContext> {
         perfLog.debug(`Register handlers took ${Date.now() - handlersStartTime}ms`);
 
         if (this.isSchedulingEnabled) {
-            const jobsStartTime = Date.now();
-            await this.registerAutoRegisterableJobs();
-            perfLog.debug(`Register jobs took ${Date.now() - jobsStartTime}ms`);
+            if (this.leaderElectionConfig) {
+                await this.startLeaderElection();
+            } else {
+                const jobsStartTime = Date.now();
+                await this.registerAutoRegisterableJobs();
+                perfLog.debug(`Register jobs took ${Date.now() - jobsStartTime}ms`);
+            }
         }
 
         // Register 404 with slight delay to allow all manually added routes to be added
@@ -495,10 +513,18 @@ export class FlinkApp<C extends FlinkContext> {
     async stop() {
         log.info("🛑 Stopping Flink app...");
 
+        if (this.leaderElection) {
+            await this.leaderElection.stop();
+        }
+
         if (this.scheduler) {
             await this.scheduler.stop();
         }
 
+        if (this.allInstanceScheduler) {
+            await this.allInstanceScheduler.stop();
+        }
+
         if (this.expressServer) {
             return new Promise<void>((resolve, reject) => {
                 const int = setTimeout(() => {
@@ -1068,9 +1094,15 @@ export class FlinkApp<C extends FlinkContext> {
 
         log.debug(`Registering ${schemaCount} schemas with AJV (manifest version: ${schemaManifest.version || "1.0"})`);
 
-        for (const { handler, assumedHttpMethod, __file } of autoRegisteredHandlers.sort(
-            (a, b) => (a.handler.Route?.order || 0) - (b.handler.Route?.order || 0)
-        )) {
+        for (const { handler, assumedHttpMethod, __file } of autoRegisteredHandlers.sort((a, b) => {
+            const orderDiff = (a.handler.Route?.order || 0) - (b.handler.Route?.order || 0);
+            if (orderDiff !== 0) return orderDiff;
+            // Static segments must be registered before parameterized ones to avoid
+            // Express matching e.g. GET /jobs/by-tags with the /jobs/:id route.
+            const aHasParam = a.handler.Route?.path?.includes("/:") ? 1 : 0;
+            const bHasParam = b.handler.Route?.path?.includes("/:") ? 1 : 0;
+            return aHasParam - bHasParam;
+        })) {
             if (!handler.Route) {
                 log.error(`Missing Props in handler ${__file}`);
                 continue;
@@ -1142,40 +1174,43 @@ export class FlinkApp<C extends FlinkContext> {
         }
     }
 
-    private async registerAutoRegisterableJobs() {
+    private async registerAutoRegisterableJobs(filter?: (jobProps: FlinkJobProps) => boolean) {
         if (!this.scheduler) {
             throw new Error("Scheduler not initialized"); // should never happen
         }
 
         for (const { Job: jobProps, default: jobFn, __file } of autoRegisteredJobs) {
+            if (filter && !filter(jobProps)) {
+                continue;
+            }
             if (jobProps.cron && jobProps.interval) {
-                log.error(`Cannot register job ${jobProps.id} - both cron and interval are set in ${__file}`);
+                schedulerLog.error(`Cannot register job ${jobProps.id} - both cron and interval are set in ${__file}`);
                 continue;
             }
 
             if (jobProps.cron && jobProps.afterDelay) {
-                log.error(`Cannot register job ${jobProps.id} - both cron and afterDelay are set in ${__file}`);
+                schedulerLog.error(`Cannot register job ${jobProps.id} - both cron and afterDelay are set in ${__file}`);
                 continue;
             }
 
             if (jobProps.interval && jobProps.afterDelay) {
-                log.error(`Cannot register job ${jobProps.id} - both interval and afterDelay are set in ${__file}`);
+                schedulerLog.error(`Cannot register job ${jobProps.id} - both interval and afterDelay are set in ${__file}`);
                 continue;
             }
 
             if (this.scheduler.existsById(jobProps.id)) {
-                log.error(`Job with id ${jobProps.id} is already registered, found duplicate in ${__file}`);
+                schedulerLog.error(`Job with id ${jobProps.id} is already registered, found duplicate in ${__file}`);
                 continue;
             }
 
-            log.debug(`Registering job ${jobProps.id}: ${JSON.stringify(jobProps)} from ${__file}`);
+            schedulerLog.debug(`Registering job ${jobProps.id}: ${JSON.stringify(jobProps)} from ${__file}`);
 
             const task = new AsyncTask(
                 jobProps.id,
                 async () => {
                     await jobFn({ ctx: this.ctx });
 
-                    log.debug(`Job ${jobProps.id} completed`);
+                    schedulerLog.debug(`Job ${jobProps.id} completed`);
 
                     if (jobProps.afterDelay) {
                         // afterDelay runs only once, so we remove the job
@@ -1183,7 +1218,7 @@ export class FlinkApp<C extends FlinkContext> {
                     }
                 },
                 (err) => {
-                    log.error(`Job ${jobProps.id} threw unhandled exception ${err}`);
+                    schedulerLog.error(`Job ${jobProps.id} threw unhandled exception ${err}`);
                     console.error(err);
                 }
             );
@@ -1216,7 +1251,7 @@ export class FlinkApp<C extends FlinkContext> {
                         try {
                             await jobFn({ ctx: this.ctx });
                         } catch (err) {
-                            log.error(`Job ${jobProps.id} threw unhandled exception ${err}`);
+                            schedulerLog.error(`Job ${jobProps.id} threw unhandled exception ${err}`);
                             console.error(err);
                         }
                     });
@@ -1235,7 +1270,7 @@ export class FlinkApp<C extends FlinkContext> {
                     this.scheduler.addSimpleIntervalJob(job);
                 }
             } else {
-                log.error(`Cannot register job ${jobProps.id} - no cron, interval or once set in ${__file}`);
+                schedulerLog.error(`Cannot register job ${jobProps.id} - no cron, interval or once set in ${__file}`);
                 continue;
             }
         }
@@ -1496,6 +1531,56 @@ export class FlinkApp<C extends FlinkContext> {
         return this.schedulingOptions?.enabled !== false;
     }
 
+    private get leaderElectionConfig(): LeaderElectionOptions | undefined {
+        const opt = this.schedulingOptions?.leaderElection;
+        if (!opt) return undefined;
+        return opt === true ? {} : opt;
+    }
+
+    private async startLeaderElection() {
+        if (!this.db) {
+            schedulerLog.warn(
+                "Leader election is enabled but no database is configured. " +
+                    "Leader election requires a MongoDB connection to coordinate between instances. " +
+                    "Either add a database connection via the `db` option, or remove `scheduling.leaderElection` from your config. " +
+                    "Jobs will run on ALL instances without leader election."
+            );
+            // Fall back to running jobs on all instances
+            this.scheduler = new ToadScheduler();
+            await this.registerAutoRegisterableJobs();
+            return;
+        }
+
+        // Register runOnAllInstances jobs immediately on a separate scheduler
+        const hasAllInstanceJobs = autoRegisteredJobs.some((j) => j.Job.runOnAllInstances);
+        if (hasAllInstanceJobs) {
+            this.allInstanceScheduler = new ToadScheduler();
+            this.scheduler = this.allInstanceScheduler;
+            await this.registerAutoRegisterableJobs((job) => !!job.runOnAllInstances);
+            this.scheduler = undefined;
+        }
+
+        const opts = this.leaderElectionConfig;
+        this.leaderElection = new LeaderElection(this.db, opts);
+
+        await this.leaderElection.start(
+            // onBecameLeader
+            async () => {
+                schedulerLog.info("This instance is now the leader - starting scheduled jobs");
+                this.scheduler = new ToadScheduler();
+                await this.registerAutoRegisterableJobs((job) => !job.runOnAllInstances);
+            },
+            // onLostLeadership
+            () => {
+                schedulerLog.info("This instance lost leadership - stopping scheduled jobs");
+                if (this.scheduler) {
+                    this.scheduler.stop();
+                    this.scheduler = undefined;
+                }
+            }
+        );
+    }
+
     private getMongoConnectionOptions() {
         if (!this.dbOpts) {
             throw new Error("No db configured");

package/src/FlinkJob.ts

@@ -37,6 +37,17 @@ export type FlinkJobProps = {
      * retried after the next interval.
      */
     singleton?: boolean;
+
+    /**
+     * If true, this job will run on all instances regardless of leader election.
+     *
+     * By default, when leader election is enabled, jobs only run on the leader instance.
+     * Set this to true for jobs that should run on every instance, such as
+     * local cache cleanup or instance-specific health checks.
+     *
+     * Has no effect when leader election is not enabled.
+     */
+    runOnAllInstances?: boolean;
 };
 
 /**

package/src/FlinkRepo.ts

@@ -51,7 +51,7 @@ export abstract class FlinkRepo<C extends FlinkContext, Model extends Document>
         return { ...model, _id: result.insertedId.toString() };
     }
 
-    async updateOne(id: string | ObjectId, model: PartialModel<Model>): Promise<Model | null> {
+    async updateById(id: string | ObjectId, model: PartialModel<Model>): Promise<Model | null> {
         const oid = this.buildId(id);
 
         const { _id, ...modelWithoutId } = model;
@@ -66,6 +66,13 @@ export abstract class FlinkRepo<C extends FlinkContext, Model extends Document>
         return null;
     }
 
+    /**
+     * @deprecated Use `updateById` instead. This will be removed in a future major version.
+     */
+    async updateOne(id: string | ObjectId, model: PartialModel<Model>): Promise<Model | null> {
+        return this.updateById(id, model);
+    }
+
     async updateMany<U = PartialModel<Model>>(query: any, model: U): Promise<number> {
         const { _id, ...modelWithoutId } = model as any;
 

package/src/LeaderElection.ts

@@ -0,0 +1,203 @@
+import { Collection, Db } from "mongodb";
+import { v4 } from "uuid";
+import { FlinkLogFactory } from "./FlinkLogFactory";
+
+const log = FlinkLogFactory.createLogger("flink.scheduler");
+
+export interface LeaderElectionOptions {
+    /**
+     * Duration in milliseconds before a leader's lease expires.
+     * If the leader fails to heartbeat within this time, another instance can take over.
+     * @default 15000
+     */
+    leaseDurationMs?: number;
+
+    /**
+     * Interval in milliseconds between heartbeats sent by the leader.
+     * Should be significantly less than leaseDurationMs (typically 1/3).
+     * @default 5000
+     */
+    heartbeatIntervalMs?: number;
+
+    /**
+     * Name of the MongoDB collection used for leader election.
+     * @default "_flink_leader"
+     */
+    collectionName?: string;
+}
+
+interface LeaderRecord {
+    _id: string;
+    instanceId: string;
+    lastHeartbeat: Date;
+    claimedAt: Date;
+}
+
+const LOCK_NAME = "job-scheduler";
+
+export class LeaderElection {
+    private instanceId = v4();
+    private _isLeader = false;
+    private timer: ReturnType<typeof setTimeout> | null = null;
+    private collection: Collection<LeaderRecord>;
+    private leaseDurationMs: number;
+    private heartbeatIntervalMs: number;
+    private onBecameLeader?: () => void | Promise<void>;
+    private onLostLeadership?: () => void | Promise<void>;
+    private stopped = false;
+    private transitioning = false;
+
+    constructor(db: Db, opts?: LeaderElectionOptions) {
+        const collectionName = opts?.collectionName || "_flink_leader";
+        this.leaseDurationMs = opts?.leaseDurationMs || 15000;
+        this.heartbeatIntervalMs = opts?.heartbeatIntervalMs || 5000;
+        this.collection = db.collection<LeaderRecord>(collectionName);
+    }
+
+    get isLeader() {
+        return this._isLeader;
+    }
+
+    /**
+     * Start the leader election process.
+     * @param onBecameLeader Called when this instance becomes the leader
+     * @param onLostLeadership Called when this instance loses leadership
+     */
+    async start(onBecameLeader: () => void | Promise<void>, onLostLeadership: () => void | Promise<void>) {
+        this.onBecameLeader = onBecameLeader;
+        this.onLostLeadership = onLostLeadership;
+        this.stopped = false;
+
+        // Ensure TTL index exists for cleanup
+        const ttlSeconds = Math.ceil((this.leaseDurationMs * 2) / 1000);
+        try {
+            await this.collection.createIndex({ lastHeartbeat: 1 }, { expireAfterSeconds: ttlSeconds });
+        } catch (err: any) {
+            if (err.codeName === "IndexOptionsConflict" || err.code === 85) {
+                log.debug("TTL index options changed, recreating index");
+                await this.collection.dropIndex("lastHeartbeat_1");
+                await this.collection.createIndex({ lastHeartbeat: 1 }, { expireAfterSeconds: ttlSeconds });
+            } else {
+                throw err;
+            }
+        }
+
+        log.info(`Leader election started (instance: ${this.instanceId.substring(0, 8)})`);
+
+        // Run first election attempt immediately
+        await this.tryClaimLeadership();
+    }
+
+    /**
+     * Stop the leader election and release leadership if held.
+     */
+    async stop() {
+        this.stopped = true;
+
+        if (this.timer) {
+            clearTimeout(this.timer);
+            this.timer = null;
+        }
+
+        if (this._isLeader) {
+            try {
+                await this.collection.deleteOne({
+                    _id: LOCK_NAME as any,
+                    instanceId: this.instanceId,
+                });
+                log.info("Leadership released on shutdown");
+            } catch (err) {
+                log.error(`Failed to release leadership on shutdown: ${err}`);
+            }
+            this._isLeader = false;
+        }
+    }
+
+    private async tryClaimLeadership() {
+        if (this.stopped || this.transitioning) return;
+
+        const now = new Date();
+        const leaseExpiry = new Date(now.getTime() - this.leaseDurationMs);
+
+        try {
+            const result = await this.collection.findOneAndUpdate(
+                {
+                    _id: LOCK_NAME as any,
+                    $or: [
+                        { instanceId: this.instanceId },
+                        { lastHeartbeat: { $lt: leaseExpiry } },
+                    ],
+                },
+                {
+                    $set: {
+                        instanceId: this.instanceId,
+                        lastHeartbeat: now,
+                    },
+                    $setOnInsert: {
+                        claimedAt: now,
+                    },
+                },
+                { upsert: true, returnDocument: "after" }
+            );
+
+            const gotLock = result && (result as any).instanceId === this.instanceId;
+
+            if (gotLock && !this._isLeader) {
+                log.info(`This instance became the leader (instance: ${this.instanceId.substring(0, 8)})`);
+                this._isLeader = true;
+                this.transitioning = true;
+                try {
+                    await this.onBecameLeader?.();
+                } catch (err) {
+                    log.error(`Error in onBecameLeader callback: ${err}`);
+                } finally {
+                    this.transitioning = false;
+                }
+            } else if (!gotLock && this._isLeader) {
+                log.warn(`This instance lost leadership (instance: ${this.instanceId.substring(0, 8)})`);
+                this._isLeader = false;
+                this.transitioning = true;
+                try {
+                    await this.onLostLeadership?.();
+                } catch (err) {
+                    log.error(`Error in onLostLeadership callback: ${err}`);
+                } finally {
+                    this.transitioning = false;
+                }
+            }
+        } catch (err: any) {
+            if (err.code === 11000) {
+                // Duplicate key - another instance claimed first
+                if (this._isLeader) {
+                    log.warn(`This instance lost leadership (instance: ${this.instanceId.substring(0, 8)})`);
+                    this._isLeader = false;
+                    try {
+                        await this.onLostLeadership?.();
+                    } catch (cbErr) {
+                        log.error(`Error in onLostLeadership callback: ${cbErr}`);
+                    }
+                }
+            } else {
+                log.error(`Leader election error: ${err}`);
+                // On error, assume we lost leadership to be safe
+                if (this._isLeader) {
+                    this._isLeader = false;
+                    try {
+                        await this.onLostLeadership?.();
+                    } catch (cbErr) {
+                        log.error(`Error in onLostLeadership callback: ${cbErr}`);
+                    }
+                }
+            }
+        }
+
+        // Schedule next attempt
+        if (!this.stopped) {
+            const nextInterval = this._isLeader
+                ? this.heartbeatIntervalMs
+                : this.heartbeatIntervalMs * 2;
+
+            this.timer = setTimeout(() => this.tryClaimLeadership(), nextInterval);
+        }
+    }
+}

package/src/index.ts

@@ -10,6 +10,8 @@ export * from "./FlinkRequestContext";
 export * from "./FlinkErrors";
 export * from "./FlinkPlugin";
 export * from "./FlinkJob";
+export { LeaderElection } from "./LeaderElection";
+export type { LeaderElectionOptions } from "./LeaderElection";
 export * from "./auth/FlinkAuthUser";
 export * from "./auth/FlinkAuthPlugin";
 export * from "./ai/FlinkTool";