go/analysis/passes/loopclosure: recursively check last statements in statements like if, switch, and for

In golang/go#16520, there was a suggestion to extend the
current loopclosure check to check more statements.

The current loopclosure flags patterns like:

	for k, v := range seq {
		go/defer func() {
			... k, v ...
		}()
	}

For this CL, the motivating example from golang/go#16520 is:

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		for j := 0; j < 1; j++ {
			wg.Add(1)
			go func() {
				fmt.Printf("%d ", i)
				wg.Done()
			}()
		}
	}
	wg.Wait()

The current loopclosure check does not flag this because of the
inner for loop, and the checker looks only at the last statement
in the outer loop body.

The suggestion is we redefine "last" recursively. For example,
if the last statement is an if, then we examine the last statements
in both of its branches. Or if the last statement is a nested loop,
then we examine the last statement of that loop's body, and so on.

A few years ago, Alan Donovan sent a sketch in CL 184537.

This CL attempts to complete Alan's sketch, as well as integrates
with the ensuing changes from golang/go#55972 to check
errgroup.Group.Go, which with this CL can now be recursively "last".

Updates golang/go#16520
Updates golang/go#55972
Fixes golang/go#30649
Fixes golang/go#32876

Change-Id: If66c6707025c20f32a2a781f6d11c4901f15742a
GitHub-Last-Rev: 04980e0
GitHub-Pull-Request: #415
Reviewed-on: https://go-review.googlesource.com/c/tools/+/452155
Reviewed-by: Tim King <[email protected]>
Run-TryBot: Tim King <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>
gopls-CI: kokoro <[email protected]>
Reviewed-by: Robert Findley <[email protected]>
Run-TryBot: Alan Donovan <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>

Author: thepudds

go/analysis/passes/loopclosure/loopclosure.go

@@ -18,24 +18,60 @@ import (
 
 const Doc = `check references to loop variables from within nested functions
 
-This analyzer checks for references to loop variables from within a function
-literal inside the loop body. It checks for patterns where access to a loop
-variable is known to escape the current loop iteration:
- 1. a call to go or defer at the end of the loop body
- 2. a call to golang.org/x/sync/errgroup.Group.Go at the end of the loop body
- 3. a call testing.T.Run where the subtest body invokes t.Parallel()
-
-In the case of (1) and (2), the analyzer only considers references in the last
-statement of the loop body as it is not deep enough to understand the effects
-of subsequent statements which might render the reference benign.
-
-For example:
-
-	for i, v := range s {
-		go func() {
-			println(i, v) // not what you might expect
-		}()
-	}
+This analyzer reports places where a function literal references the
+iteration variable of an enclosing loop, and the loop calls the function
+in such a way (e.g. with go or defer) that it may outlive the loop
+iteration and possibly observe the wrong value of the variable.
+
+In this example, all the deferred functions run after the loop has
+completed, so all observe the final value of v.
+
+    for _, v := range list {
+        defer func() {
+            use(v) // incorrect
+        }()
+    }
+
+One fix is to create a new variable for each iteration of the loop:
+
+    for _, v := range list {
+        v := v // new var per iteration
+        defer func() {
+            use(v) // ok
+        }()
+    }
+
+The next example uses a go statement and has a similar problem.
+In addition, it has a data race because the loop updates v
+concurrent with the goroutines accessing it.
+
+    for _, v := range elem {
+        go func() {
+            use(v)  // incorrect, and a data race
+        }()
+    }
+
+A fix is the same as before. The checker also reports problems
+in goroutines started by golang.org/x/sync/errgroup.Group.
+A hard-to-spot variant of this form is common in parallel tests:
+
+    func Test(t *testing.T) {
+        for _, test := range tests {
+            t.Run(test.name, func(t *testing.T) {
+                t.Parallel()
+                use(test) // incorrect, and a data race
+            })
+        }
+    }
+
+The t.Parallel() call causes the rest of the function to execute
+concurrent with the loop.
+
+The analyzer reports references only in the last statement,
+as it is not deep enough to understand the effects of subsequent
+statements that might render the reference benign.
+("Last statement" is defined recursively in compound
+statements such as if, switch, and select.)
 
 See: https://golang.org/doc/go_faq.html#closures_and_goroutines`
 
@@ -91,59 +127,121 @@ func run(pass *analysis.Pass) (interface{}, error) {
 		//
 		// For go, defer, and errgroup.Group.Go, we ignore all but the last
 		// statement, because it's hard to prove go isn't followed by wait, or
-		// defer by return.
+		// defer by return. "Last" is defined recursively.
 		//
+		// TODO: consider allowing the "last" go/defer/Go statement to be followed by
+		// N "trivial" statements, possibly under a recursive definition of "trivial"
+		// so that that checker could, for example, conclude that a go statement is
+		// followed by an if statement made of only trivial statements and trivial expressions,
+		// and hence the go statement could still be checked.
+		forEachLastStmt(body.List, func(last ast.Stmt) {
+			var stmts []ast.Stmt
+			switch s := last.(type) {
+			case *ast.GoStmt:
+				stmts = litStmts(s.Call.Fun)
+			case *ast.DeferStmt:
+				stmts = litStmts(s.Call.Fun)
+			case *ast.ExprStmt: // check for errgroup.Group.Go
+				if call, ok := s.X.(*ast.CallExpr); ok {
+					stmts = litStmts(goInvoke(pass.TypesInfo, call))
+				}
+			}
+			for _, stmt := range stmts {
+				reportCaptured(pass, vars, stmt)
+			}
+		})
+
+		// Also check for testing.T.Run (with T.Parallel).
 		// We consider every t.Run statement in the loop body, because there is
-		// no such commonly used mechanism for synchronizing parallel subtests.
+		// no commonly used mechanism for synchronizing parallel subtests.
 		// It is of course theoretically possible to synchronize parallel subtests,
 		// though such a pattern is likely to be exceedingly rare as it would be
 		// fighting against the test runner.
-		lastStmt := len(body.List) - 1
-		for i, s := range body.List {
-			var stmts []ast.Stmt // statements that must be checked for escaping references
+		for _, s := range body.List {
 			switch s := s.(type) {
-			case *ast.GoStmt:
-				if i == lastStmt {
-					stmts = litStmts(s.Call.Fun)
-				}
-
-			case *ast.DeferStmt:
-				if i == lastStmt {
-					stmts = litStmts(s.Call.Fun)
-				}
-
-			case *ast.ExprStmt: // check for errgroup.Group.Go and testing.T.Run (with T.Parallel)
+			case *ast.ExprStmt:
 				if call, ok := s.X.(*ast.CallExpr); ok {
-					if i == lastStmt {
-						stmts = litStmts(goInvoke(pass.TypesInfo, call))
-					}
-					if stmts == nil {
-						stmts = parallelSubtest(pass.TypesInfo, call)
+					for _, stmt := range parallelSubtest(pass.TypesInfo, call) {
+						reportCaptured(pass, vars, stmt)
 					}
+
 				}
 			}
+		}
+	})
+	return nil, nil
+}
 
-			for _, stmt := range stmts {
-				ast.Inspect(stmt, func(n ast.Node) bool {
-					id, ok := n.(*ast.Ident)
-					if !ok {
-						return true
-					}
-					obj := pass.TypesInfo.Uses[id]
-					if obj == nil {
-						return true
-					}
-					for _, v := range vars {
-						if v == obj {
-							pass.ReportRangef(id, "loop variable %s captured by func literal", id.Name)
-						}
-					}
-					return true
-				})
+// reportCaptured reports a diagnostic stating a loop variable
+// has been captured by a func literal if checkStmt has escaping
+// references to vars. vars is expected to be variables updated by a loop statement,
+// and checkStmt is expected to be a statements from the body of a func literal in the loop.
+func reportCaptured(pass *analysis.Pass, vars []types.Object, checkStmt ast.Stmt) {
+	ast.Inspect(checkStmt, func(n ast.Node) bool {
+		id, ok := n.(*ast.Ident)
+		if !ok {
+			return true
+		}
+		obj := pass.TypesInfo.Uses[id]
+		if obj == nil {
+			return true
+		}
+		for _, v := range vars {
+			if v == obj {
+				pass.ReportRangef(id, "loop variable %s captured by func literal", id.Name)
 			}
 		}
+		return true
 	})
-	return nil, nil
+}
+
+// forEachLastStmt calls onLast on each "last" statement in a list of statements.
+// "Last" is defined recursively so, for example, if the last statement is
+// a switch statement, then each switch case is also visited to examine
+// its last statements.
+func forEachLastStmt(stmts []ast.Stmt, onLast func(last ast.Stmt)) {
+	if len(stmts) == 0 {
+		return
+	}
+
+	s := stmts[len(stmts)-1]
+	switch s := s.(type) {
+	case *ast.IfStmt:
+	loop:
+		for {
+			forEachLastStmt(s.Body.List, onLast)
+			switch e := s.Else.(type) {
+			case *ast.BlockStmt:
+				forEachLastStmt(e.List, onLast)
+				break loop
+			case *ast.IfStmt:
+				s = e
+			case nil:
+				break loop
+			}
+		}
+	case *ast.ForStmt:
+		forEachLastStmt(s.Body.List, onLast)
+	case *ast.RangeStmt:
+		forEachLastStmt(s.Body.List, onLast)
+	case *ast.SwitchStmt:
+		for _, c := range s.Body.List {
+			cc := c.(*ast.CaseClause)
+			forEachLastStmt(cc.Body, onLast)
+		}
+	case *ast.TypeSwitchStmt:
+		for _, c := range s.Body.List {
+			cc := c.(*ast.CaseClause)
+			forEachLastStmt(cc.Body, onLast)
+		}
+	case *ast.SelectStmt:
+		for _, c := range s.Body.List {
+			cc := c.(*ast.CommClause)
+			forEachLastStmt(cc.Body, onLast)
+		}
+	default:
+		onLast(s)
+	}
 }
 
 // litStmts returns all statements from the function body of a function

go/analysis/passes/loopclosure/testdata/src/a/a.go

@@ -7,6 +7,8 @@
 package testdata
 
 import (
+	"sync"
+
 	"golang.org/x/sync/errgroup"
 )
 
@@ -108,6 +110,70 @@ func _() {
 	}
 }
 
+// Cases that rely on recursively checking for last statements.
+func _() {
+
+	for i := range "outer" {
+		for j := range "inner" {
+			if j < 1 {
+				defer func() {
+					print(i) // want "loop variable i captured by func literal"
+				}()
+			} else if j < 2 {
+				go func() {
+					print(i) // want "loop variable i captured by func literal"
+				}()
+			} else {
+				go func() {
+					print(i)
+				}()
+				println("we don't catch the error above because of this statement")
+			}
+		}
+	}
+
+	for i := 0; i < 10; i++ {
+		for j := 0; j < 10; j++ {
+			if j < 1 {
+				switch j {
+				case 0:
+					defer func() {
+						print(i) // want "loop variable i captured by func literal"
+					}()
+				default:
+					go func() {
+						print(i) // want "loop variable i captured by func literal"
+					}()
+				}
+			} else if j < 2 {
+				var a interface{} = j
+				switch a.(type) {
+				case int:
+					defer func() {
+						print(i) // want "loop variable i captured by func literal"
+					}()
+				default:
+					go func() {
+						print(i) // want "loop variable i captured by func literal"
+					}()
+				}
+			} else {
+				ch := make(chan string)
+				select {
+				case <-ch:
+					defer func() {
+						print(i) // want "loop variable i captured by func literal"
+					}()
+				default:
+					go func() {
+						print(i) // want "loop variable i captured by func literal"
+					}()
+				}
+			}
+		}
+	}
+}
+
 // Group is used to test that loopclosure only matches Group.Go when Group is
 // from the golang.org/x/sync/errgroup package.
 type Group struct{}
@@ -125,6 +191,21 @@ func _() {
 			return nil
 		})
 	}
+
+	for i, v := range s {
+		if i > 0 {
+			g.Go(func() error {
+				print(i) // want "loop variable i captured by func literal"
+				return nil
+			})
+		} else {
+			g.Go(func() error {
+				print(v) // want "loop variable v captured by func literal"
+				return nil
+			})
+		}
+	}
+
 	// Do not match other Group.Go cases
 	g1 := new(Group)
 	for i, v := range s {
@@ -135,3 +216,28 @@ func _() {
 		})
 	}
 }
+
+// Real-world example from #16520, slightly simplified
+func _() {
+	var nodes []interface{}
+
+	critical := new(errgroup.Group)
+	others := sync.WaitGroup{}
+
+	isCritical := func(node interface{}) bool { return false }
+	run := func(node interface{}) error { return nil }
+
+	for _, node := range nodes {
+		if isCritical(node) {
+			critical.Go(func() error {
+				return run(node) // want "loop variable node captured by func literal"
+			})
+		} else {
+			others.Add(1)
+			go func() {
+				_ = run(node) // want "loop variable node captured by func literal"
+				others.Done()
+			}()
+		}
+	}
+}