Goでループの中で重い処理を検知する linter を作った話

初めまして、n9te9 です。

Goのループ内で特定パッケージのメソッドや関数の実行を検知する Linter を作りました。
この記事では、自作したlinterがどのようにプログラムを評価しているのか、linterを作る中での学びがあったのでまとめていきたいと思います。

linterを実装するにあたって、tenntennさんのforcetypeassertを参考にしました。
linterを作る際のパッケージについての説明やgo vetを介してのlinterの実行については、こちらの記事にまとめられていたので、これを参考にしました。

自作したlinter npluscheck

https://github.com/lkeix/npluscheck

このlinterは、ループの中で database/sql などの外部APIの呼び出し箇所を検知します。
例えば、以下のような実装があった場合、ループの中でSQLを発行するAPIの呼び出しがあったとしてlinterが警告を出します。

func insert1000records(db *sql.DB) error {
  for i := 0; i < 1000; i++ {
    db.Exec("insert into testlist (id) values (?)", i)
  }
  return nil
}

上記の実装は、insertのSQLが1000回発行されるので非効率です。
1回のSQLで1000回のSQLと同じ結果が得られればパフォーマンスが向上します。
そのため、代替の実装としては以下のようになります。

func insert1000records(db *sql.DB) error {
  db.Exec("insert into testlist (id) values (?), (?), ... ,(?)", []int{1, 2, ..., 1000}...)
  return nil
}

このようにnpluscheckは、ループの中で外部APIの呼び出し箇所を検知します。

npluscheckの中身

npluscheckが、どのようにループの中で外部APIの呼び出し箇所を検知しているかを説明していきます。
npluscheckの処理のおおまかな流れは以下の通りです。

1. 実行されている関数を抽出する
2. 1で抽出された情報をもとにループの中で外部APIが呼ばれているかチェックする

まず、外部APIの呼び出し箇所を検知する前に、定義された関数の中で実行されている関数を抽出しています。
この抽出処理は、npluscheckで実装した独自の構造体のスライスを返します。
この構造体は、実行されている関数の情報を持ったexprのフィールドとループの中で実行されているかどうかのフィールドを持っています。
抽出処理によって、関数内で実行されている関数の情報、ループの中で実行されているどうかの情報が取得できます。

抽出処理の実装は以下のようになります。

func extractCalledFuncs(stmts []ast.Stmt, calledInFor bool) []customFunc {
  funcs := []customFunc{}
  for _, stmt := range stmts {
    switch stmt := stmt.(type) {
    case *ast.ExprStmt:
      if callexpr, ok := stmt.X.(*ast.CallExpr); ok {
        if sel := extractSelectorExprFun(callexpr); sel != nil {
          funcs = append(funcs, customFunc{
            calledInFor: calledInFor,
            expr:        sel,
          })
          continue
        }
        funcs = append(funcs, customFunc{
          calledInFor: calledInFor,
          expr:        callexpr,
        })
      }
    case *ast.ForStmt:
      funcs = append(funcs, extractCalledFuncs(stmt.Body.List, true)...)
    case *ast.IfStmt:
      funcs = append(funcs, extractCalledFuncs(stmt.Body.List, calledInFor)...)
    case *ast.AssignStmt:
      for _, expr := range stmt.Rhs {
        if rhs, ok := expr.(*ast.CallExpr); ok {
          if sel := extractSelectorExprFun(rhs); sel != nil {
            funcs = append(funcs, customFunc{
              calledInFor: calledInFor,
              expr:        sel,
            })
            continue
          }
          funcs = append(funcs, customFunc{
            calledInFor: calledInFor,
            expr:        rhs,
          })
        }
      }
    }
  }
  return funcs
}

ループの中で外部APIを実行しているかどうかを検知するなら、ast.ExprStmtでcalledInForを見るだけで良いのですが、その場合、外部APIを呼び出す関数がラップされている場合において検知できません。

そのため、以下のような実装に対しては検知ができません。

var db *sql.DB

func main() {
  for i := 0; i < 100; i++ {
    insertID(i)
  }
}

func insertID(i int) error {
  db.Exec("insert into testlist (id) values (?)", i)
  return nil
}

現状、上記のような処理をnpluscheckで検知ができません。
ただ、今後このような実装もnpluscheckで検知できるように実行されている関数を構造体のスライスとして抽出しています。

次に、抽出した構造体のスライスをもとにループの中で外部APIを呼び出す関数が実行されていないかをチェックします。
構造体には、ループの中で関数が実行されているかどうかのフィールドがあるので、それがtrueになっているかを確認します。
その構造体の関数のexprのフィールドで外部APIを呼び出している関数/メソッドかどうかを確認します。
ループの中で実行されているかつ、外部APIを呼び出すパッケージの関数/メソッドである場合は、ループの中で外部APIを実行しているとしています。

チェックの実装は以下のようになっています。

func checkNplus(pass *analysis.Pass, info *types.Info, fd *ast.FuncDecl) map[token.Pos]bool {
  funcs := extractCalledFuncs(fd.Body.List, false)
  callinLoop := make(map[token.Pos]bool)
  for _, f := range funcs {
    switch expr := f.expr.(type) {
    case *ast.SelectorExpr:
      if typ, ok := info.Selections[expr]; ok {
        if contain(UsuallyDataBasePackages, typ.Obj().Pkg().Path()) && f.calledInFor {
          callinLoop[expr.Pos()] = true
          pass.Reportf(expr.Pos(), "may call DB API in loop")
        }
      }
    }
  }
  return callinLoop
}

少々ネストが深くなりましたが、このようにしてlinterを実装しました。
外部APIを呼び出す関数がラップされている場合において検知できないなど、まだまだ課題はあります。

今回作成したlinterのnpluscheckの実装については以上になります。
次に、npluscheckを作る上でASTの扱い方を学んだので、それについて説明していきます。

関数の中身をASTで扱う

GoのASTは適切に抽象化されているため、とても扱いやすいデータ構造になっています。
その中でも、今回npluscheckを作る上で適切に抽象化されていると感じた箇所は制御文(Statement)、式(Expr)に関する抽象化です。まず、Statementの抽象化から見ていきます。
文(Statement)のデータ構造は以下のような図になっています。

具体的なStatementは、上記以外にもたくさんありますが、収まらないため3つにしました。
StatementをIfStatement, ForStatement, ExprStatement などに具体化する場合は、下記のようにswitch文で型アサーションすることで具体化したStatementに応じた処理が書けます。

func Xxx(stmts []ast.Stmt) {
  for _, stmt := range stmts {
    switch stmt := stmt.(type) {
    case *ast.ExprStatement:
    // 代入処理 :=, = での処理
    case *ast.IfStatement:
    // if文での処理
    case *ast.ForStatement:
    // for文での処理
    ...
    }
  }
}

次に、式のExprの抽象化についてです。
まず、式のExprを取得するためには、ExprStatementやAssignStatementから式を取得する必要があります
ExprStatmentとAssignStatementは、以下のように分類されます。

func main() {
  // ExprStatement
  var hoge int
  // ExprStatement
  a()
  // AssignStatement
  one := b()
}

func a() {
  // ExprStatement
  fmt.Println("called a()")
}

func b() int {
  return 1
}

ExprStatementやAssignStatementから関数を取り出す場合は、ExprStatementならば、Xフィールドから、AssignStatementならば、RhsフィールドからExprとして取り出します。
a()、b()、fmt.Println()の式はフィールドから読み取った時点では、Exprで抽象化されています。a()、b()、fmt.Println()はExprで抽象化されていますが、ここからStatementのように具体化します。
npluscheckは、ExprからCallExprとSelectionExprに具体化して関数を取得しています。CallExprやSelectionExpr以外にも具体的なExprは存在するので、Exprから取得したい情報に応じて具体化する構造体が変わってきます。
パッケージ関数や構造体のメソッドではない単純な関数実行の際にCallExprで具体化でき、パッケージ関数や構造体のメソッドなどはSelectionExprで具体化できます。
このように、StatementとExprの抽象化と具体化を利用することでlinterを作ることができます。
関数の中身の実装を解析する際は、ast.Stmtのスライスをforで回しながら、switch文で命令に応じて解析しています。
また、forやifの中身もast.Stmtのスライスで取得ができるので、関数の中身同様に処理を書けます。ただ、ASTは、木構造で再帰的なデータ構造をしているので再帰で関数を呼び出す構成にした方がよりシンプルにコードを解析できます。
実際に、npluscheckでも関数の中で呼ばれている関数を抽出する処理では、StatementからExprから抽出し、再帰処理によってスライスを作成しています。

まとめ

ループの中で外部APIを呼び出している箇所を検知するlinterのnpluscheckを作成しました。
ただ、npluscheckは、外部APIを呼び出している関数/メソッドをラップした関数がループの中で実行されてる場合には、検知ができません。
今後は、このような箇所も改修していきたいと思います。

関数の中身を解析する際は、以下のようにコードから式や制御文を分解されるのかを頭に入れておくだけでもlinterを作るハードルはグッと低くなるのではないかなと思いました。

初めてlinterを作ってみて、Goのlinter作成をサポートするパッケージやライブラリの強力さ、GoのASTの扱いやすさを実感しました。
チームで開発する際のルールをnpluscheckのような自作linterなどでカバーできる部分が生産性向上の要因にもなっているのかなと思いました。