From 949242c4006d4f0f3fb85b2550bae692effa2fb3 Mon Sep 17 00:00:00 2001
From: Baha Aiman <bahaaiman@google.com>
Date: Tue, 7 May 2024 00:07:57 -0700
Subject: [PATCH] refactor(datastore): Rename function (#10063)

---
 datastore/query.go       |  6 +++---
 datastore/query_test.go  |  2 +-
 datastore/transaction.go | 25 +++++++++++++++++++++----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/datastore/query.go b/datastore/query.go
index 6acc6821e84f..929c133e6189 100644
--- a/datastore/query.go
+++ b/datastore/query.go
@@ -800,7 +800,7 @@ func (c *Client) RunAggregationQuery(ctx context.Context, aq *AggregationQuery)
 	// Parse the read options.
 	txn := aq.query.trans
 	if txn != nil {
-		defer txn.acquireLock()()
+		defer txn.stateLockDeferUnlock()()
 	}
 
 	req.ReadOptions, err = parseQueryReadOptions(aq.query.eventual, txn)
@@ -834,7 +834,7 @@ func validateReadOptions(eventual bool, t *Transaction) error {
 		return nil
 	}
 	if eventual {
-		return errors.New("datastore: cannot use EventualConsistency query in a transaction")
+		return errEventualConsistencyTransaction
 	}
 	if t.state == transactionStateExpired {
 		return errExpiredTransaction
@@ -963,7 +963,7 @@ func (t *Iterator) nextBatch() error {
 
 	txn := t.trans
 	if txn != nil {
-		defer txn.acquireLock()()
+		defer txn.stateLockDeferUnlock()()
 	}
 
 	var err error
diff --git a/datastore/query_test.go b/datastore/query_test.go
index 97e6fbe57ccf..4883e847acc4 100644
--- a/datastore/query_test.go
+++ b/datastore/query_test.go
@@ -886,7 +886,7 @@ func TestAggregationQueryIsNil(t *testing.T) {
 }
 
 func TestValidateReadOptions(t *testing.T) {
-	eventualInTxnErr := errors.New("datastore: cannot use EventualConsistency query in a transaction")
+	eventualInTxnErr := errEventualConsistencyTransaction
 
 	for _, test := range []struct {
 		desc     string
diff --git a/datastore/transaction.go b/datastore/transaction.go
index a8dd8149e1b6..deeaf16d057a 100644
--- a/datastore/transaction.go
+++ b/datastore/transaction.go
@@ -31,7 +31,10 @@ import (
 // to a conflict with a concurrent transaction.
 var ErrConcurrentTransaction = errors.New("datastore: concurrent transaction")
 
-var errExpiredTransaction = errors.New("datastore: transaction expired")
+var (
+	errExpiredTransaction             = errors.New("datastore: transaction expired")
+	errEventualConsistencyTransaction = errors.New("datastore: cannot use EventualConsistency query in a transaction")
+)
 
 type transactionSettings struct {
 	attempts int
@@ -228,8 +231,22 @@ func (t *Transaction) beginLaterTransaction() (err error) {
 	return nil
 }
 
-// Acquire state lock if transaction has not started
-func (t *Transaction) acquireLock() func() {
+// Acquires state lock if transaction has not started. No-op otherwise
+// The returned function unlocks the state if it was locked.
+//
+// Usage:
+//
+//	func (t *Transaction) someFunction() {
+//		...
+//		if t != nil {
+//			defer t.stateLockDeferUnlock()()
+//		}
+//		....
+//	}
+//
+// This ensures that state is locked before any of the following lines are exexcuted
+// The lock will be released after 'someFunction' ends
+func (t *Transaction) stateLockDeferUnlock() func() {
 	if t.state == transactionStateNotStarted {
 		t.stateLock.Lock()
 		// Check whether state changed while waiting to acquire lock
@@ -416,7 +433,7 @@ func (t *Transaction) get(spanName string, keys []*Key, dst interface{}) (err er
 	defer func() { trace.EndSpan(t.ctx, err) }()
 
 	if t != nil {
-		defer t.acquireLock()()
+		defer t.stateLockDeferUnlock()()
 	}
 
 	opts, err := t.parseReadOptions()