From e27085da6646f74e390627668af7f234061fab78 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@molecula.com>
Date: Wed, 8 Sep 2021 13:29:55 -0500
Subject: [PATCH] check context occasionally while processing mutex check
 results

It's hard to do this remotely sanely for the fragments, but the
translation and collation process itself could be fairly slow on
large data sets, so we should check occasionally for canceled
context and return early if no one needs the result anyway.

Also, take out no-longer-correct comments from the test case.
---
 api.go      | 35 ++++++++++++++++++++++++++++++++---
 api_test.go |  6 ------
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/api.go b/api.go
index 648696601..fe5723365 100644
--- a/api.go
+++ b/api.go
@@ -2831,6 +2831,10 @@ func (api *API) MutexCheck(ctx context.Context, qcx *Qcx, indexName string, fiel
 				}
 			}
 		}
+		// if context is done, return early.
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
 		untranslatedKeys := 0
 		// Obtain translation tables for the keys.
 		if useIndexKeys {
@@ -2849,6 +2853,10 @@ func (api *API) MutexCheck(ctx context.Context, qcx *Qcx, indexName string, fiel
 				}
 			}
 		}
+		// if context is done, return early.
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
 		if useFieldKeys {
 			fieldKeyList, err := api.cluster.translateFieldListIDs(field, fieldIDs)
 			if err != nil {
@@ -2870,6 +2878,11 @@ func (api *API) MutexCheck(ctx context.Context, qcx *Qcx, indexName string, fiel
 		}
 	}
 
+	// if context is done, return early.
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
 	// define the process functions. separated from above code just to make
 	// it easier to follow/compare them.
 	if useIndexKeys {
@@ -2879,7 +2892,9 @@ func (api *API) MutexCheck(ctx context.Context, qcx *Qcx, indexName string, fiel
 			// unlike a map, the slice won't get updated-in-place, so we have
 			// to assign to result after we're done
 			defer func() {
-				result = outStrings
+				if err == nil {
+					result = outStrings
+				}
 			}()
 			process = func(recordID uint64, valueIDs []uint64) bool {
 				if _, ok := outMap[recordID]; ok {
@@ -2930,7 +2945,9 @@ func (api *API) MutexCheck(ctx context.Context, qcx *Qcx, indexName string, fiel
 			// unlike a map, the slice won't get updated-in-place, so we have
 			// to assign to result after we're done
 			defer func() {
-				result = outIDs
+				if err == nil {
+					result = outIDs
+				}
 			}()
 			process = func(recordID uint64, valueIDs []uint64) bool {
 				if _, ok := outMap[recordID]; ok {
@@ -2982,18 +2999,30 @@ processing:
 		if len(nodeResults) == 0 {
 			continue
 		}
+		// if context is done, return early.
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
 		for _, v := range nodeResults {
 			if len(v) == 0 {
 				continue
 			}
+			counter := 0
 			for record, values := range v {
+				counter++
 				if process(record, values) {
 					break processing
 				}
+				// every 65k items or so, check the context for done-ness
+				if counter%(1<<16) == 0 {
+					if err := ctx.Err(); err != nil {
+						return nil, err
+					}
+				}
 			}
 		}
 	}
-	return result, nil
+	return result, ctx.Err()
 }
 
 type serverInfo struct {
diff --git a/api_test.go b/api_test.go
index 146924b0b..e4826374f 100644
--- a/api_test.go
+++ b/api_test.go
@@ -1138,9 +1138,6 @@ func TestAPI_MutexCheck(t *testing.T) {
 				t.Fatalf("checking mutexes: %v", err)
 			}
 			if keyedField {
-				// this just sorta comes out this way with our hashing; these are
-				// the things which were in position 1 of their shards, and did
-				// not have a value which happens to map to 3.
 				mapped, ok := results.(map[string][]string)
 				if !ok {
 					t.Fatalf("expected map[string][]string, got %T", results)
@@ -1182,9 +1179,6 @@ func TestAPI_MutexCheck(t *testing.T) {
 			if err != nil {
 				t.Fatalf("checking mutexes: %v", err)
 			}
-			// this just sorta comes out this way with our hashing; these are
-			// the things which were in position 1 of their shards, and did
-			// not have a value which happens to map to 3.
 			mapped, ok := results.([]string)
 			if !ok {
 				t.Fatalf("expected []string, got %T", results)