Implement locking stream pull to avoid multiple catalyst pulling the stream

packages/api/src/controllers/stream.test.ts

@@ -524,6 +524,77 @@ describe("controllers/stream", () => {
         expect(server.db.stream.addDefaultFields(document)).toEqual(stream);
       });
 
+      it("should lock pull for a new stream", async () => {
+        // Create stream pull
+        const now = Date.now();
+        const res = await client.put("/stream/pull", postMockPullStream);
+        expect(res.status).toBe(201);
+        const stream = await res.json();
+
+        // Request pull lock
+        const resLockPull = await client.post(`/stream/${stream.id}/lockPull`);
+        expect(resLockPull.status).toBe(204);
+
+        // Check that the pullLockedAt is marked with the correct date
+        const res2 = await client.get(`/stream/${stream.id}`);
+        expect(res2.status).toBe(200);
+        const stream2 = await res2.json();
+        expect(stream2.pullLockedAt).toBeGreaterThan(now);
+      });
+
+      it("should not lock pull for an active stream", async () => {
+        // Create stream pull
+        const res = await client.put("/stream/pull", postMockPullStream);
+        expect(res.status).toBe(201);
+        const stream = await res.json();
+
+        // Mark stream as active
+        await db.stream.update(stream.id, { isActive: true });
+
+        // Requesting pull lock should fail, because the stream is active (so it should be replicated instead of being pulled)
+        const reslockPull = await client.post(`/stream/${stream.id}/lockPull`);
+        expect(reslockPull.status).toBe(423);
+      });
+
+      it("should not lock pull for already locked pull", async () => {
+        // Create stream pull
+        const res = await client.put("/stream/pull", postMockPullStream);
+        expect(res.status).toBe(201);
+        const stream = await res.json();
+
+        // Request pull lock by many processes at the same time, only one should acquire a lock
+        const promises = [];
+        for (let i = 0; i < 10; i++) {
+          promises.push(client.post(`/stream/${stream.id}/lockPull`));
+        }
+        const resPulls = await Promise.all(promises);
+        expect(resPulls.filter((r) => r.status === 204).length).toBe(1);
+        expect(resPulls.filter((r) => r.status === 423).length).toBe(9);
+      });
+
+      it("should lock pull for already locked pull if lease has expired", async () => {
+        // Create stream pull
+        const res = await client.put("/stream/pull", postMockPullStream);
+        expect(res.status).toBe(201);
+        const stream = await res.json();
+
+        // Request pull lock
+        const resLockPull = await client.post(`/stream/${stream.id}/lockPull`, {
+          leaseTimeout: 1,
+        });
+        expect(resLockPull.status).toBe(204);
+
+        // Wait until lease has expired
+        await sleep(1);
+
+        // Request pull lock should succeed, because the lock lease has expired (so we assume the stream is not being pulled at the moment)
+        const resLockPull2 = await client.post(
+          `/stream/${stream.id}/lockPull`,
+          { leaseTimeout: 1 }
+        );
+        expect(resLockPull2.status).toBe(204);
+      });
+
       it("should update a stream if it has the same pull source", async () => {
         let res = await client.put("/stream/pull", postMockPullStream);
         expect(res.status).toBe(201);
@@ -760,6 +831,18 @@ describe("controllers/stream", () => {
         expect(updatedStream.lastSeen).toBeGreaterThan(timeBeforeBump);
       });
 
+      it("start pull should update lastPullAt", async () => {
+        const stream = await createAndActivateStream();
+        const timeBeforeBump = Date.now();
+        expect(stream.lastSeen).toBeLessThan(timeBeforeBump);
+
+        const res = await client.post(`/stream/${stream.id}/heartbeat`);
+
+        expect(res.status).toBe(204);
+        const updatedStream = await server.db.stream.get(stream.id);
+        expect(updatedStream.lastSeen).toBeGreaterThan(timeBeforeBump);
+      });
+
       it("should allow changing the mist host as well", async () => {
         const stream = await createAndActivateStream();
 

packages/api/src/controllers/stream.ts

@@ -1086,6 +1086,41 @@ app.put(
   }
 );
 
+app.post("/:id/lockPull", authorizer({ anyAdmin: true }), async (req, res) => {
+  const { id } = req.params;
+  let { leaseTimeout } = req.body;
+  if (!leaseTimeout) {
+    // Sets the default lock lease to 60s
+    leaseTimeout = 60 * 1000;
+  }
+  logger.info(`got /lockPull for stream=${id}`);
+
+  const stream = await db.stream.get(id, { useReplica: false });
+  if (!stream || (stream.deleted && !req.user.admin)) {
+    res.status(404);
+    return res.json({ errors: ["not found"] });
+  }
+  if (stream.isActive) {
+    return res.status(423).end();
+  }
+
+  const updateRes = await db.stream.update(
+    [
+      sql`id = ${stream.id}`,
+      sql`COALESCE((data->>'pullLockedAt')::bigint,0) < ${
+        Date.now() - leaseTimeout
+      }`,
+    ],
+    { pullLockedAt: Date.now() },
+    { throwIfEmpty: false }
+  );
+
+  if (updateRes.rowCount > 0) {
+    res.status(204).end();
+  }
+  res.status(423).end();
+});
+
 function terminateDelay(stream: DBStream) {
   if (!stream.lastTerminatedAt) {
     return 0;

packages/api/src/schema/db-schema.yaml

@@ -702,6 +702,9 @@ components:
           properties:
             source:
               index: true
+        pullLockedAt:
+          type: number
+          example: 1587667174725
         playbackId:
           unique: true
         mistHost:

packages/api/src/store/stream-table.ts

@@ -408,7 +408,12 @@ export default class StreamTable extends Table<DBStream> {
   }
 }
 
-const adminOnlyFields = ["mistHost", "broadcasterHost", "createdByTokenId"];
+const adminOnlyFields = [
+  "mistHost",
+  "broadcasterHost",
+  "createdByTokenId",
+  "pullLockedAt",
+];
 
 const privateFields = [
   "recordObjectStoreId",