From 959ec438f5752f2f646f8ccd22e64e55f73b45a2 Mon Sep 17 00:00:00 2001
From: Prabhjot Singh Sethi <prabhjot.sethi@gmail.com>
Date: Thu, 17 Apr 2025 11:05:29 +0000
Subject: [PATCH] Add wrapped context for lock tables allowing separate cancel

Signed-off-by: Prabhjot Singh Sethi <prabhjot.sethi@gmail.com>
---
 lock/lock.go | 35 +++++++++++++++++++++++++++++------
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/lock/lock.go b/lock/lock.go
index 43c57f3..7be325d 100644
--- a/lock/lock.go
+++ b/lock/lock.go
@@ -47,6 +47,15 @@ type LockTable struct {
 
 	// context in which this lock table is being working on
 	ctx context.Context
+
+	// Context cancel function
+	cancelFn context.CancelFunc
+}
+
+func (t *LockTable) Callback(op string, wKey interface{}) {
+	// handle callback as and when needed
+	// we may need notification of release of locks
+	// allowing others to start working on it
 }
 
 func (t *LockTable) handleOwnerRelease(op string, wKey interface{}) {
@@ -94,14 +103,17 @@ func LocateLockTable(store db.Store, name string) (*LockTable, error) {
 		if ownerTable == nil {
 			return nil, errors.Wrap(errors.InvalidArgument, "Mandatory! owner table infra not initialized")
 		}
+
+		ctx, cancelFn := context.WithCancel(ownerTable.ctx)
+
 		// no existing table found, allocate a new one
 		col := store.GetCollection(name)
 		table = &LockTable{
-			colName: name,
-			col:     col,
-			ctx:     ownerTable.ctx,
+			colName:  name,
+			col:      col,
+			ctx:      ctx,
+			cancelFn: cancelFn,
 		}
-		lockTables[name] = table
 
 		matchDeleteStage := mongo.Pipeline{
 			bson.D{{
@@ -113,11 +125,22 @@ func LocateLockTable(store db.Store, name string) (*LockTable, error) {
 			}},
 		}
 
-		// watch only for delete notification
-		err := ownerTable.col.Watch(table.ctx, matchDeleteStage, table.handleOwnerRelease)
+		// watch only for delete notification of lock owner
+		err := ownerTable.col.Watch(ctx, matchDeleteStage, table.handleOwnerRelease)
 		if err != nil {
+			cancelFn()
 			return nil, err
 		}
+
+		// register to watch for locks, this is relevant for external
+		// notification and cleanup as part of handling of release of owners
+		err = table.col.Watch(ctx, nil, table.Callback)
+		if err != nil {
+			cancelFn()
+			return nil, err
+		}
+
+		lockTables[name] = table
 	}
 
 	return table, nil