Add counts and add_new_item methods, without tests

go/WorkQueue.go

@@ -66,6 +66,7 @@ package workqueue
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"github.com/google/uuid"
@@ -120,8 +121,117 @@ func (workQueue *WorkQueue) AddItem(ctx context.Context, db *redis.Client, item
 	return err
 }
 
+// AddNewItem adds an item to the work queue only if an item with the same ID doesn't already exist.
+//
+// This method uses WATCH to add the item atomically. The db client passed must not be used by
+// anything else while this method is running.
+//
+// Returns a boolean indicating if the item was added or not. The item is only not added if it
+// already exists or if an error (other than a transaction error, which triggers a retry) occurs.
+func (workQueue *WorkQueue) AddNewItem(ctx context.Context, db *redis.Client, item Item) (bool, error) {
+	txf := func(tx *redis.Tx) error {
+
+		processingItemsInQueueCmd := tx.LPos(ctx, workQueue.processingKey, item.ID, redis.LPosArgs{
+			Rank:   0,
+			MaxLen: 0,
+		})
+		workingItemsInQueueCmd := tx.LPos(ctx, workQueue.mainQueueKey, item.ID, redis.LPosArgs{
+			Rank:   0,
+			MaxLen: 0,
+		})
+
+		_, ProcessingQueueCheck := processingItemsInQueueCmd.Result()
+		_, WorkingQueueCheck := workingItemsInQueueCmd.Result()
+
+		if ProcessingQueueCheck == redis.Nil && WorkingQueueCheck == redis.Nil {
+
+			_, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
+				workQueue.AddItemToPipeline(ctx, pipe, item)
+				_, err := pipe.Exec(ctx)
+				return err
+			})
+			return err
+		} else {
+			return nil
+		}
+	}
+
+	for i := 0; i < 100; i++ {
+		err := db.Watch(ctx, txf, workQueue.processingKey, workQueue.mainQueueKey)
+		if err == nil {
+			return true, nil
+		}
+		if err == redis.TxFailedErr {
+			continue
+		}
+		return false, err
+	}
+
+	return false, errors.New("increment reached maximum number of retries")
+}
+
+func (workQueue *WorkQueue) AddNewItemWithSleep(ctx context.Context, db *redis.Client, item Item) (bool, error) {
+	txf := func(tx *redis.Tx) error {
+
+		processingItemsInQueueCmd := tx.LPos(ctx, workQueue.processingKey, item.ID, redis.LPosArgs{
+			Rank:   0,
+			MaxLen: 0,
+		})
+		workingItemsInQueueCmd := tx.LPos(ctx, workQueue.mainQueueKey, item.ID, redis.LPosArgs{
+			Rank:   0,
+			MaxLen: 0,
+		})
+
+		_, ProcessingQueueCheck := processingItemsInQueueCmd.Result()
+		_, WorkingQueueCheck := workingItemsInQueueCmd.Result()
+
+		if ProcessingQueueCheck == redis.Nil && WorkingQueueCheck == redis.Nil {
+			time.Sleep(100 * time.Microsecond)
+			_, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
+				workQueue.AddItemToPipeline(ctx, pipe, item)
+				_, err := pipe.Exec(ctx)
+				return err
+			})
+			return err
+		} else {
+			return nil
+		}
+	}
+
+	for i := 0; i < 100; i++ {
+		err := db.Watch(ctx, txf, workQueue.processingKey, workQueue.mainQueueKey)
+		if err == nil {
+			return true, nil
+		}
+		if err == redis.TxFailedErr {
+			continue
+		}
+		return false, err
+	}
+
+	return false, errors.New("increment reached maximum number of retries")
+}
+
+// Counts returns the queue length, and number of items currently being processed, atomically.
+func (workQueue *WorkQueue) Counts(ctx context.Context, db *redis.Client) (queueLen, processingLen int64, err error) {
+	tx := db.TxPipeline()
+
+	queueLenPipe := tx.LLen(ctx, workQueue.mainQueueKey)
+	processingLenPipe := tx.LLen(ctx, workQueue.processingKey)
+
+	_, err = tx.Exec(ctx)
+	if err == nil {
+		queueLen, err = queueLenPipe.Result()
+	}
+	if err == nil {
+		processingLen, err = processingLenPipe.Result()
+	}
+
+	return
+}
+
 // Return the length of the work queue (not including items being processed, see
-// [WorkQueue.Processing]).
+// [WorkQueue.Processing] or [WorkQueue.Counts] to get both).
 func (workQueue *WorkQueue) QueueLen(ctx context.Context, db *redis.Client) (int64, error) {
 	return db.LLen(ctx, workQueue.mainQueueKey).Result()
 }

node/src/WorkQueue.ts

@@ -3,12 +3,12 @@
  * @description A work queue backed by a redis database.
  */
 
-import Redis, {ChainableCommander} from 'ioredis';
-import {v4 as uuidv4} from 'uuid';
-import {Item} from './Item';
-import {KeyPrefix} from './KeyPrefix';
+import Redis, { ChainableCommander } from 'ioredis';
+import { v4 as uuidv4 } from 'uuid';
+import { Item } from './Item';
+import { KeyPrefix } from './KeyPrefix';
 
-export {KeyPrefix, Item};
+export { KeyPrefix, Item };
 
 /**
  * A work queue backed by a redis database.
@@ -35,10 +35,11 @@ export class WorkQueue {
 
   /**
    * Add an item to the work queue. This adds the redis commands onto the pipeline passed.
+   *
    * Use `WorkQueue.addItem` if you don't want to pass a pipeline directly.
-   * Add the item data.
-   * @param {Pipeline} pipeline The pipeline that the data will be executed.
-   * @param {Item} item The Item which will be set in the Redis with the key of this.itemDataKey.of(item.id).
+   *
+   * @param {Pipeline} pipeline The pipeline that the commands to add the item will be pushed to.
+   * @param {Item} item The Item to be added.
    */
   addItemToPipeline(pipeline: ChainableCommander, item: Item) {
     // NOTE: it's important that the data is added first, otherwise someone before the data is ready.
@@ -48,11 +49,12 @@ export class WorkQueue {
   }
 
   /**
-   * Add an item to the work queue.
+   * Add an item to the work queue. See `addNewItem` to avoid adding duplicate items.
+   *
    * This creates a pipeline and executes it on the database.
    *
    * @param {Redis} db The Redis Connection.
-   * @param item The item that will be executed using the method addItemToPipeline.
+   * @param item The item to be added.
    */
   async addItem(db: Redis, item: Item): Promise<void> {
     const pipeline = db.pipeline();
@@ -61,7 +63,100 @@ export class WorkQueue {
   }
 
   /**
-   * This is used to get the length of the Main Queue.
+   * Adds an item to the work queue only if an item with the same ID doesn't already exist.
+   *
+   * This method uses WATCH to add the item atomically. The db client passed must not be used by
+   * anything else while this method is running.
+   *
+   * Returns a boolean indicating if the item was added or not. The item is only not added if it
+   * already exists or if an error (other than a transaction error, which triggers a retry) occurs.
+   *
+   * @param {Redis} db The Redis Connection, this must not be used by anything else while this method is running.
+   * @param item The item that will be added, only if an item doesn't already exist with the same ID.
+   * @returns {boolean} returns false if already in queue or true if the item is successfully added.
+   */
+  async addNewItem(db: Redis, item: Item): Promise<boolean> {
+    while (true) {
+      try {
+        await db.watch(this.mainQueueKey, this.processingKey);
+
+        const isItemInProcessingKey = await db.lpos(this.processingKey, item.id);
+        const isItemInMainQueueKey = await db.lpos(this.mainQueueKey, item.id);
+        if (isItemInProcessingKey !== null || isItemInMainQueueKey !== null) {
+          console.log("Item already exists, not added", item.id);
+          await db.unwatch();
+          return false;
+        }
+        const transaction = db.multi();
+        this.addItemToPipeline(transaction, item);
+        const results = await transaction.exec();
+
+        if (!results) {
+          console.log("Transaction failed, item not added", item.id);
+          await db.unwatch();
+        }
+        console.log("Item added successfully", item.id);
+        return true
+      } catch (e) {
+        console.log("Error", e);
+      } finally {
+        await db.unwatch();
+      }
+    }
+    return false;
+  }
+
+  /**
+   * This is a dummy function used for testing!
+   * 
+   * Adds an item to the work queue only if an item with the same ID doesn't already exist with a 50 ms sleep.
+   *
+   * This method uses WATCH to add the item atomically. The db client passed must not be used by
+   * anything else while this method is running.
+   *
+   * Returns a boolean indicating if the item was added or not. The item is only not added if it
+   * already exists or if an error (other than a transaction error, which triggers a retry) occurs.
+   *
+   * @param {Redis} db The Redis Connection, this must not be used by anything else while this method is running.
+   * @param item The item that will be added, only if an item doesn't already exist with the same ID.
+   * @returns {boolean} returns false if already in queue or true if the item is successfully added.
+   */
+  async addNewItemWithSeep(db: Redis, item: Item): Promise<boolean> {
+    while (true) {
+      try {
+        await db.watch(this.mainQueueKey, this.processingKey);
+
+        const isItemInProcessingKey = await db.lpos(this.processingKey, item.id);
+        const isItemInMainQueueKey = await db.lpos(this.mainQueueKey, item.id);
+        if (isItemInProcessingKey !== null || isItemInMainQueueKey !== null) {
+          console.log("Item already exists, not added", item.id);
+          await db.unwatch();
+          return false;
+        }
+        await new Promise<void>((resolve) => setTimeout(resolve, 50));;
+        const transaction = db.multi();
+        this.addItemToPipeline(transaction, item);
+        const results = await transaction.exec();
+
+        if (!results) {
+          console.log("Transaction failed, item not added", item.id);
+          await db.unwatch();
+        }
+        console.log("Item added successfully", item.id);
+        return true
+      } catch (e) {
+        console.log("Error", e);
+      } finally {
+        await db.unwatch();
+      }
+    }
+    return false;
+  }
+
+
+  /**
+   * Return the length of the work queue (not including items being processed, see
+   * `WorkQueue.processing` or `WorkQueue.counts` to get both).
    *
    * @param {Redis} db The Redis Connection.
    * @returns {Promise<number>} Return the length of the work queue (not including items being processed, see `WorkQueue.processing()`).
@@ -71,7 +166,7 @@ export class WorkQueue {
   }
 
   /**
-   * This is used to get the lenght of the Processing Queue.
+   * This is used to get the number of items currently being processed.
    *
    * @param {Redis} db The Redis Connection.
    * @returns {Promise<number>} The number of items being processed.
@@ -80,6 +175,22 @@ export class WorkQueue {
     return db.llen(this.processingKey);
   }
 
+  /**
+   * Returns the queue length, and number of items currently being processed, atomically.
+   * 
+   * @param {Redis} db The Redis Connection.
+   * @returns {Promise<[number, number]>} Return the length of main queue and processing queue, respectively.
+   */
+  async counts(db: Redis): Promise<[number, number]> {
+    const multi = db.multi();
+    multi.llen(this.mainQueueKey);
+    multi.llen(this.processingKey);
+    const result = (await multi.exec()) as Array<[Error | null, number]>;
+    const queueLength = result[0][1];
+    const processingLength = result[1][1];
+    return [queueLength, processingLength];
+  }
+
   /**
    * This method can be used to check if a Lease Exists or not for a itemId.
    *
@@ -116,8 +227,8 @@ export class WorkQueue {
   async lease(
     db: Redis,
     leaseSecs: number,
-    block = true,
-    timeout = 1
+    block: boolean = true,
+    timeout: number = 1
   ): Promise<Item | null> {
     let maybeItemId: string | null = null;
 
@@ -131,7 +242,7 @@ export class WorkQueue {
     } else {
       maybeItemId = await db.rpoplpush(this.mainQueueKey, this.processingKey);
     }
-
+    console.log(`Leased ${maybeItemId}`)
     if (maybeItemId == null) {
       return null;
     }
@@ -207,7 +318,7 @@ export class WorkQueue {
     if (removed === 0) {
       return false;
     }
-
+    console.log(`Completed ${item.id}`);
     const pipeline = db.pipeline();
     pipeline.del(this.itemDataKey.of(item.id));
     pipeline.del(this.leaseKey.of(item.id));