waitgroup.go

     1  // Copyright 2025 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package modernize
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"go/ast"
    11  	"go/printer"
    12  	"slices"
    13  
    14  	"golang.org/x/tools/go/analysis"
    15  	"golang.org/x/tools/go/analysis/passes/inspect"
    16  	"golang.org/x/tools/go/types/typeutil"
    17  	"golang.org/x/tools/internal/analysis/analyzerutil"
    18  	typeindexanalyzer "golang.org/x/tools/internal/analysis/typeindex"
    19  	"golang.org/x/tools/internal/astutil"
    20  	"golang.org/x/tools/internal/refactor"
    21  	"golang.org/x/tools/internal/typesinternal/typeindex"
    22  	"golang.org/x/tools/internal/versions"
    23  )
    24  
    25  var WaitGroupAnalyzer = &analysis.Analyzer{
    26  	Name: "waitgroup",
    27  	Doc:  analyzerutil.MustExtractDoc(doc, "waitgroup"),
    28  	Requires: []*analysis.Analyzer{
    29  		inspect.Analyzer,
    30  		typeindexanalyzer.Analyzer,
    31  	},
    32  	Run: waitgroup,
    33  	URL: "https://pkg.go.dev/golang.org/x/tools/go/analysis/passes/modernize#waitgroup",
    34  }
    35  
    36  // The waitgroup pass replaces old more complex code with
    37  // go1.25 added API WaitGroup.Go.
    38  //
    39  // Patterns:
    40  //
    41  //  1. wg.Add(1); go func() { defer wg.Done(); ... }()
    42  //     =>
    43  //     wg.Go(go func() { ... })
    44  //
    45  //  2. wg.Add(1); go func() { ...; wg.Done() }()
    46  //     =>
    47  //     wg.Go(go func() { ... })
    48  //
    49  // The wg.Done must occur within the first statement of the block in a
    50  // defer format or last statement of the block, and the offered fix
    51  // only removes the first/last wg.Done call. It doesn't fix existing
    52  // wrong usage of sync.WaitGroup.
    53  //
    54  // The use of WaitGroup.Go in pattern 1 implicitly introduces a
    55  // 'defer', which may change the behavior in the case of panic from
    56  // the "..." logic. In this instance, the change is safe: before and
    57  // after the transformation, an unhandled panic inevitably results in
    58  // a fatal crash. The fact that the transformed code calls wg.Done()
    59  // before the crash doesn't materially change anything. (If Done had
    60  // other effects, or blocked, or if WaitGroup.Go propagated panics
    61  // from child to parent goroutine, the argument would be different.)
    62  func waitgroup(pass *analysis.Pass) (any, error) {
    63  	var (
    64  		index             = pass.ResultOf[typeindexanalyzer.Analyzer].(*typeindex.Index)
    65  		info              = pass.TypesInfo
    66  		syncWaitGroupAdd  = index.Selection("sync", "WaitGroup", "Add")
    67  		syncWaitGroupDone = index.Selection("sync", "WaitGroup", "Done")
    68  	)
    69  	if !index.Used(syncWaitGroupDone) {
    70  		return nil, nil
    71  	}
    72  
    73  	for curAddCall := range index.Calls(syncWaitGroupAdd) {
    74  		// Extract receiver from wg.Add call.
    75  		addCall := curAddCall.Node().(*ast.CallExpr)
    76  		if !isIntLiteral(info, addCall.Args[0], 1) {
    77  			continue // not a call to wg.Add(1)
    78  		}
    79  		// Inv: the Args[0] check ensures addCall is not of
    80  		// the form sync.WaitGroup.Add(&wg, 1).
    81  		addCallRecv := ast.Unparen(addCall.Fun).(*ast.SelectorExpr).X
    82  
    83  		// Following statement must be go func() { ... } ().
    84  		curAddStmt := curAddCall.Parent()
    85  		if !is[*ast.ExprStmt](curAddStmt.Node()) {
    86  			continue // unnecessary parens?
    87  		}
    88  		curNext, ok := curAddCall.Parent().NextSibling()
    89  		if !ok {
    90  			continue // no successor
    91  		}
    92  		goStmt, ok := curNext.Node().(*ast.GoStmt)
    93  		if !ok {
    94  			continue // not a go stmt
    95  		}
    96  		lit, ok := goStmt.Call.Fun.(*ast.FuncLit)
    97  		if !ok || len(goStmt.Call.Args) != 0 {
    98  			continue // go argument is not func(){...}()
    99  		}
   100  		if lit.Type.Results != nil && len(lit.Type.Results.List) > 0 {
   101  			continue // function literal has return values; wg.Go requires func()
   102  		}
   103  		list := lit.Body.List
   104  		if len(list) == 0 {
   105  			continue
   106  		}
   107  
   108  		// Body must start with "defer wg.Done()" or end with "wg.Done()".
   109  		var doneStmt ast.Stmt
   110  		if deferStmt, ok := list[0].(*ast.DeferStmt); ok &&
   111  			typeutil.Callee(info, deferStmt.Call) == syncWaitGroupDone &&
   112  			astutil.EqualSyntax(ast.Unparen(deferStmt.Call.Fun).(*ast.SelectorExpr).X, addCallRecv) {
   113  			doneStmt = deferStmt // "defer wg.Done()"
   114  
   115  		} else if lastStmt, ok := list[len(list)-1].(*ast.ExprStmt); ok {
   116  			if doneCall, ok := lastStmt.X.(*ast.CallExpr); ok &&
   117  				typeutil.Callee(info, doneCall) == syncWaitGroupDone &&
   118  				astutil.EqualSyntax(ast.Unparen(doneCall.Fun).(*ast.SelectorExpr).X, addCallRecv) {
   119  				doneStmt = lastStmt // "wg.Done()"
   120  			}
   121  		}
   122  		if doneStmt == nil {
   123  			continue
   124  		}
   125  		curDoneStmt, ok := curNext.FindNode(doneStmt)
   126  		if !ok {
   127  			panic("can't find Cursor for 'done' statement")
   128  		}
   129  
   130  		file := astutil.EnclosingFile(curAddCall)
   131  		if !analyzerutil.FileUsesGoVersion(pass, file, versions.Go1_25) {
   132  			continue
   133  		}
   134  		tokFile := pass.Fset.File(file.Pos())
   135  
   136  		var addCallRecvText bytes.Buffer
   137  		err := printer.Fprint(&addCallRecvText, pass.Fset, addCallRecv)
   138  		if err != nil {
   139  			continue // error getting text for the edit
   140  		}
   141  
   142  		pass.Report(analysis.Diagnostic{
   143  			// go func() {
   144  			// ~~~~~~~~~
   145  			Pos:     goStmt.Pos(),
   146  			End:     lit.Type.End(),
   147  			Message: "Goroutine creation can be simplified using WaitGroup.Go",
   148  			SuggestedFixes: []analysis.SuggestedFix{{
   149  				Message: "Simplify by using WaitGroup.Go",
   150  				TextEdits: slices.Concat(
   151  					// delete "wg.Add(1)"
   152  					refactor.DeleteStmt(tokFile, curAddStmt),
   153  					// delete "wg.Done()" or "defer wg.Done()"
   154  					refactor.DeleteStmt(tokFile, curDoneStmt),
   155  					[]analysis.TextEdit{
   156  						// go    func()
   157  						// ------
   158  						// wg.Go(func()
   159  						{
   160  							Pos:     goStmt.Pos(),
   161  							End:     goStmt.Call.Pos(),
   162  							NewText: fmt.Appendf(nil, "%s.Go(", addCallRecvText.String()),
   163  						},
   164  						// ... }()
   165  						//      -
   166  						// ... } )
   167  						{
   168  							Pos: goStmt.Call.Lparen,
   169  							End: goStmt.Call.Rparen,
   170  						},
   171  					},
   172  				),
   173  			}},
   174  		})
   175  	}
   176  	return nil, nil
   177  }
   178
View as plain text