refactor(transactioncache): unify transaction data in cache

parser/preinstrumentationtracing.go

@@ -3,6 +3,7 @@ package parser
 import (
 	"github.com/dave/dst"
 	"github.com/dave/dst/dstutil"
+	"github.com/newrelic/go-easy-instrumentation/parser/transactioncache"
 )
 
 // DetectTransactions analyzes the AST to identify and track transactions within function declarations.
@@ -54,14 +55,16 @@ func DetectTransactions(manager *InstrumentationManager, c *dstutil.Cursor) {
 					// Check if the transaction is passed to another function, if so track its calls
 					for _, arg := range callExpr.Args {
 						ident, ok := arg.(*dst.Ident)
-						if ok && ident.Name == currentTransaction.Name {
-							ident, ok := callExpr.Fun.(*dst.Ident)
-							if ok {
-								funcDecl, exists := manager.transactionCache.Functions[ident.Name]
-								if exists {
-									trackFunctionCalls(manager, funcDecl, currentTransaction)
-								}
-							}
+						if !ok || ident.Name != currentTransaction.Name {
+							continue
+						}
+						ident, ok = callExpr.Fun.(*dst.Ident)
+						if !ok {
+							continue
+						}
+						funcDecl, exists := manager.transactionCache.Functions[ident.Name]
+						if exists {
+							trackFunctionCalls(manager, funcDecl, currentTransaction)
 						}
 					}
 				}
@@ -76,21 +79,30 @@ func DetectTransactions(manager *InstrumentationManager, c *dstutil.Cursor) {
 func trackFunctionCalls(manager *InstrumentationManager, funcDecl *dst.FuncDecl, txn *dst.Ident) {
 	// Traverse the function body to track calls
 	dstutil.Apply(funcDecl.Body, func(c *dstutil.Cursor) bool {
-		if callExpr, ok := c.Node().(*dst.CallExpr); ok {
-			manager.transactionCache.AddCall(txn, callExpr)
+		callExpr, ok := c.Node().(*dst.CallExpr)
+		if !ok {
+			return true
+		}
 
-			// Check if the call is an End method directly on the transaction
-			if selExpr, ok := callExpr.Fun.(*dst.SelectorExpr); ok {
-				if ident, ok := selExpr.X.(*dst.Ident); ok && selExpr.Sel.Name == "End" && ident == txn {
-					manager.transactionCache.TransactionState[txn] = true // Mark transaction as closed
-					return false                                          // Stop further traversal
-				}
+		// Validate that we are able to add calls to the cache. Fail and bail if we are not.
+		if !manager.transactionCache.AddCall(txn, callExpr) {
+			return false
+		}
+
+		// Check if the call is an End method directly on the transaction
+		if transactioncache.IsTxnEnd(txn, callExpr) {
+			txnData, ok := manager.transactionCache.Transactions[txn]
+			if !ok {
+				return false
 			}
-			// Recursively track calls within functions that are called with the transaction
-			if ident, ok := callExpr.Fun.(*dst.Ident); ok {
-				if funcDecl, exists := manager.transactionCache.Functions[ident.Name]; exists {
-					trackFunctionCalls(manager, funcDecl, txn)
-				}
+			txnData.SetClosed(true)
+			return false // Stop further traversal
+		}
+
+		// Recursively track calls within functions that are called with the transaction
+		if ident, ok := callExpr.Fun.(*dst.Ident); ok {
+			if funcDecl, exists := manager.transactionCache.Functions[ident.Name]; exists {
+				trackFunctionCalls(manager, funcDecl, txn)
 			}
 		}
 		return true

parser/transactioncache/transactioncache.go

@@ -6,75 +6,160 @@ import (
 	"github.com/dave/dst"
 )
 
+// TransactionData is responsible for maintaining metadata related to an individual transaction
+// It contains the following components:
+//
+//   - Expressions: A list of dst.Expr statements active within the lifespan of the transaction
+//
+//   - IsClosed: A boolean to indicate whether a transaction is open or has ended, ensures that no
+//     further expressions are added once a transaction is marked as ended
+type TransactionData struct {
+	Expressions []dst.Expr
+	IsClosed    bool
+}
+
 // TransactionCache is responsible for tracking existing transactions within a Go application.
 // It maintains the following components:
 //
-//   - Transactions: A map where each key is a transaction name, and the value is a list of expressions
-//     that are active within the lifespan of the transaction.
+//   - Transactions: A map where each key is a transaction name, and the value is a pointer to a
+//     TransactionData struct
 //
 //   - Functions: A map that stores already seen functions alongside their declarations. This is useful
 //     for tracking transactions that span multiple function calls.
-//
-//   - TransactionState: A map that tracks whether a transaction is open or has ended, ensuring that
-//     no further expressions are added once a transaction is marked as ended.
 type TransactionCache struct {
-	Transactions     map[*dst.Ident][]dst.Expr
-	Functions        map[string]*dst.FuncDecl
-	TransactionState map[*dst.Ident]bool // Track whether a transaction is closed
+	Transactions map[*dst.Ident]*TransactionData
+	Functions    map[string]*dst.FuncDecl
 }
 
 func NewTransactionCache() *TransactionCache {
 	return &TransactionCache{
-		Transactions:     make(map[*dst.Ident][]dst.Expr),
-		Functions:        make(map[string]*dst.FuncDecl),
-		TransactionState: make(map[*dst.Ident]bool),
+		Transactions: make(map[*dst.Ident]*TransactionData),
+		Functions:    make(map[string]*dst.FuncDecl),
+	}
+}
+
+func NewTxnData() *TransactionData {
+	return &TransactionData{
+		Expressions: nil,
+		IsClosed:    false,
 	}
 }
 
+// IsTxnEnd returns true if a given dst.Expr is an `End()` operation for a given
+// *dst.Ident transaction name, else false
+func IsTxnEnd(txn *dst.Ident, expr dst.Expr) bool {
+	if txn == nil || expr == nil {
+		return false
+	}
+
+	callExpr, ok := expr.(*dst.CallExpr)
+	if !ok {
+		return false
+	}
+
+	selExpr, ok := callExpr.Fun.(*dst.SelectorExpr)
+	if !ok {
+		return false
+	}
+
+	if selExpr.Sel.Name != "End" {
+		return false
+	}
+
+	ident, ok := selExpr.X.(*dst.Ident)
+	if !ok || ident.Name != txn.Name {
+		return false
+	}
+
+	return true
+}
+
+// SetClosed is a setter function for TransactionData to control the value of
+// IsClosed. Returns true on success.
+func (td *TransactionData) SetClosed(closed bool) bool {
+	if td == nil {
+		return false
+	}
+	td.IsClosed = closed
+	return true
+}
+
+// AddExpr is a setter function for TransactionData to add an expression to the
+// list of expressions. Returns true on success.
+func (td *TransactionData) AddExpr(expr dst.Expr) bool {
+	if td == nil {
+		return false
+	}
+	td.Expressions = append(td.Expressions, expr)
+	return true
+}
+
+// AddTxnToCache is a setter function for TransactionCache to add or update a
+// TransactionData entry based on *dst.Ident key. Returns true on success.
+func (tc *TransactionCache) AddTxnToCache(txnKey *dst.Ident, txnData *TransactionData) bool {
+	if tc == nil || txnKey == nil || txnData == nil {
+		return false
+	}
+
+	tc.Transactions[txnKey] = txnData
+	return true
+}
+
 // AddCall adds an expression to the list of expressions associated with a transaction.
 // It first checks if the transaction is closed, and if so, it does not add the expression.
 // If the expression is an 'End' call directly on the transaction, it marks the transaction as closed.
 // If the 'End' call is part of a segment [ex: defer txn.StartSegment.End()], it does not mark the transaction as closed.
-func (tc *TransactionCache) AddCall(transaction *dst.Ident, expr dst.Expr) {
-	if tc.Transactions == nil {
-		tc.Transactions = make(map[*dst.Ident][]dst.Expr)
+func (tc *TransactionCache) AddCall(transaction *dst.Ident, expr dst.Expr) bool {
+	if tc == nil || tc.Transactions == nil || transaction == nil || expr == nil {
+		return false // Enforce initialization of TransactionCache
 	}
+
 	// Check if the transaction is closed
-	if closed, exists := tc.TransactionState[transaction]; exists && closed {
-		return // Do not add calls to a closed transaction
+	txn, ok := tc.Transactions[transaction]
+	if ok && txn.IsClosed {
+		return false // Do not add calls to a closed transaction
 	}
+
+	var txnData *TransactionData
+	if !ok {
+		txnData = NewTxnData()
+	} else {
+		txnData = tc.Transactions[transaction]
+	}
+
 	// Check if the call is an End method directly on the transaction
-	if callExpr, ok := expr.(*dst.CallExpr); ok {
-		if selExpr, ok := callExpr.Fun.(*dst.SelectorExpr); ok {
-			if selExpr.Sel.Name == "End" {
-				// Check if the End method is called directly on the transaction
-				if ident, ok := selExpr.X.(*dst.Ident); ok && ident.Name == transaction.Name {
-					tc.TransactionState[transaction] = true // Mark transaction as closed
-				} else if selExpr.X.(*dst.CallExpr) != nil {
-					// This is likely part of a segment operation, do not mark transaction as closed
-					return
-				}
-			}
-		}
+	if IsTxnEnd(transaction, expr) {
+		txnData.IsClosed = true
 	}
-	tc.Transactions[transaction] = append(tc.Transactions[transaction], expr)
 
+	txnData.AddExpr(expr)
+
+	return tc.AddTxnToCache(transaction, txnData)
 }
 
 // IsFunctionInTransactionScope checks if a given function name is present within any transaction.
 // It iterates over all transactions and their expressions, returning true if the function name is found.
 func (tc *TransactionCache) IsFunctionInTransactionScope(functionName string) bool {
-	for _, exprs := range tc.Transactions {
-		for _, expr := range exprs {
-			if callExpr, ok := expr.(*dst.CallExpr); ok {
-				if ident, ok := callExpr.Fun.(*dst.Ident); ok && ident.Name == functionName {
-					return true
-				}
-				if selExpr, ok := callExpr.Fun.(*dst.SelectorExpr); ok {
-					if ident, ok := selExpr.X.(*dst.Ident); ok && ident.Name == functionName {
-						return true
-					}
-				}
+	for _, txnData := range tc.Transactions {
+		for _, expr := range txnData.Expressions {
+			callExpr, ok := expr.(*dst.CallExpr)
+			if !ok {
+				continue
+			}
+
+			ident, ok := callExpr.Fun.(*dst.Ident)
+			if ok && ident.Name == functionName {
+				return true
+			}
+
+			selExpr, ok := callExpr.Fun.(*dst.SelectorExpr)
+			if !ok {
+				continue
+			}
+
+			ident, ok = selExpr.X.(*dst.Ident)
+			if ok && ident.Name == functionName {
+				return true
 			}
 		}
 	}
@@ -86,22 +171,22 @@ func (tc *TransactionCache) ExtractNames() (transactionNames []string, expressio
 	expressionNames = make(map[string][]string)
 
 	// Iterate over transactions and gather names
-	for txn, exprs := range tc.Transactions {
-		transactionNames = append(transactionNames, txn.Name)
+	for txnKey, txnData := range tc.Transactions {
+		transactionNames = append(transactionNames, txnKey.Name)
 
-		for _, expr := range exprs {
+		for _, expr := range txnData.Expressions {
 			switch e := expr.(type) {
 			case *dst.CallExpr:
 				if selExpr, ok := e.Fun.(*dst.SelectorExpr); ok {
 					if ident, identOk := selExpr.X.(*dst.Ident); identOk {
-						expressionNames[txn.Name] = append(expressionNames[txn.Name], fmt.Sprintf("%s.%s", ident.Name, selExpr.Sel.Name))
+						expressionNames[txnKey.Name] = append(expressionNames[txnKey.Name], fmt.Sprintf("%s.%s", ident.Name, selExpr.Sel.Name))
 					} else {
-						expressionNames[txn.Name] = append(expressionNames[txn.Name], selExpr.Sel.Name)
+						expressionNames[txnKey.Name] = append(expressionNames[txnKey.Name], selExpr.Sel.Name)
 					}
 				} else if ident, identOk := e.Fun.(*dst.Ident); identOk {
-					expressionNames[txn.Name] = append(expressionNames[txn.Name], ident.Name)
+					expressionNames[txnKey.Name] = append(expressionNames[txnKey.Name], ident.Name)
 				} else {
-					expressionNames[txn.Name] = append(expressionNames[txn.Name], "Unknown")
+					expressionNames[txnKey.Name] = append(expressionNames[txnKey.Name], "Unknown")
 				}
 			default:
 				continue
@@ -111,20 +196,19 @@ func (tc *TransactionCache) ExtractNames() (transactionNames []string, expressio
 
 	return transactionNames, expressionNames
 }
+
+// CheckTransactionExists returns true if the *dst.Ident transaction is already
+// recorded in the cache, otherwise false.
 func (tc *TransactionCache) CheckTransactionExists(transaction *dst.Ident) bool {
-	for txn := range tc.Transactions {
-		if txn.Name == transaction.Name {
-			return true
-		}
-	}
-	return false
+	_, ok := tc.Transactions[transaction]
+	return ok
 }
 
-// Debug printing of cache
+// Print outputs Debug printing of cache
 func (tc *TransactionCache) Print() {
-	for txn, exprs := range tc.Transactions {
-		fmt.Printf("Transaction: %s\n", txn.Name)
-		for _, expr := range exprs {
+	for txnKey, txnData := range tc.Transactions {
+		fmt.Printf("Transaction: %s\n", txnKey.Name)
+		for _, expr := range txnData.Expressions {
 			switch e := expr.(type) {
 			case *dst.CallExpr:
 				selExpr, ok := e.Fun.(*dst.SelectorExpr)