Adds context handling to concurrent core

Author: reederc42

core/concurrent_cache/concurrent_cache.go

@@ -7,6 +7,7 @@
 package concurrent_cache
 
 import (
+	"context"
 	"fmt"
 	"slices"
 	"sync"
@@ -196,7 +197,7 @@ func Query(query ...Subquery) []Subquery {
 //  5. Merge all queries
 //     Returns sorted list of subqueries
 //  6. Acquire locks and fill in Results
-func Lock(queries ...[]Subquery) ([]Result, func(), error) {
+func Lock(ctx context.Context, queries ...[]Subquery) ([]Result, func(), error) {
 	roots := make([][]int, len(queries))
 	cachesPresent := make(map[resource]struct{}, resourceCount)
 	// phase 1, requires no locks, check errors, dedupe, and build trees
@@ -211,7 +212,11 @@ func Lock(queries ...[]Subquery) ([]Result, func(), error) {
 
 	// phase 3, takes per-resource locks
 	merged := mergeQueries(queries)
-	return lockQuery(merged, len(queries))
+	results, unlocker, err := lockQuery(merged, len(queries))
+	if err == nil && ctx.Err() != nil {
+		err = ctx.Err()
+	}
+	return results, unlocker, err
 }
 
 func assembleQueries(queries [][]Subquery, roots [][]int, cachesPresent map[resource]struct{}) error {

core/concurrent_cache/concurrent_cache_test.go

@@ -1,6 +1,7 @@
 package concurrent_cache
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"reflect"
@@ -266,7 +267,7 @@ func TestLock(t *testing.T) {
 	for i := 0; i < n; i++ {
 		go func() {
 			defer wg.Done()
-			results, unlocker, err := Lock(makeQueries()...)
+			results, unlocker, err := Lock(context.Background(), makeQueries()...)
 			defer unlocker()
 			assert.NoError(t, err)
 			assert.NotNil(t, results[2].VolumePublication.Upsert)
@@ -479,7 +480,7 @@ func TestLockNeverDeadlocks(t *testing.T) {
 					for i := 0; i < tc.numQueries; i++ {
 						queries[i] = generateRandomQuery(queryGenerators[rand.Intn(len(queryGenerators))], tc.maxSubqueriesPerQuery)
 					}
-					results, unlock, err := Lock(queries...)
+					results, unlock, err := Lock(context.Background(), queries...)
 					defer unlock()
 					time.Sleep(time.Duration(10+rand.Intn(40)) * time.Millisecond) // Simulate some processing time
 					for _, operation := range []string{"Upsert", "Delete"} {
@@ -616,7 +617,7 @@ func TestLockWithDependencyChains(t *testing.T) {
 			var err error
 
 			go func() {
-				results, unlock, err = Lock(tc.querySets...)
+				results, unlock, err = Lock(context.Background(), tc.querySets...)
 				time.Sleep(10 * time.Millisecond) // Simulate some processing time
 				for _, operation := range []string{"Upsert", "Delete"} {
 					doOperation(t, operation, results)
@@ -776,3 +777,11 @@ func cleanupTestData() {
 	volumePublications.data = make(map[string]SmartCopier)
 	volumePublications.unlock()
 }
+
+func TestLockCancel(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), -1*time.Second)
+	defer cancel()
+	_, unlocker, err := Lock(ctx, Query(UpsertNode("node1")))
+	defer unlocker()
+	assert.ErrorContains(t, err, "context deadline exceeded")
+}