🤖

昔コードのりファクターに大活躍のGoのIteratorの話

2024/09/23に公開

はじめに

最近、昔作ったGoプログラムを整理しました。その中で初めてGo言語を触ったときに作られたSNSスクレイピングツールがあります。Issueを確認すると以下が書いてありました。

スクレイピングした投稿データを一括にJSONファイルに保存するのでなく一件一件JSONファイルに追加するようにする。

すいぶん前に書いたものなので自分はIssueの存在さえ忘れました。最近Goバージョン1.23に追加されたIteratorで簡単に実装できるんじゃないかと思って実装してみました。その実装内容を今回の記事に書かせて頂きました。

Issueの背景

そもそもなんで上記のIssueを記載したかというとすべて対象となら投稿をスクレイピングできるまでJSONファイルへの保存が行われないからです。途中で何らかの理由でプログラムが終了になった場合、スクレイピングした投稿データが消えてしまいます。

またスクレイピングされたデータが多いほどメモリの消費がどんどん上がってしまいます。この2つの問題を避けるために投稿をスクレイピングした後すぐにファイルに保存することが目的でした。

実装

改善前

まず改善する前のソースコードを紹介します。

func main() {
    posts, err := ScrapePosts()
    if err != nil {
        log.Fatalf(err)
    }
    if ExportToJSON(path, posts); err != nil {
        log.Fatalf(err)
    }
}

// ScrapePosts ユーザー投稿をスクレイピングする
func ScrapePosts() ([]Post, error) {
    var posts []Post
    // ...
    // スクレイピングした投稿データをPost配列に追加
    posts = append(posts, post)
    // ...
    return posts, nil
}

// ExportToJSON スクレイピングしたユーザー投稿をJSONファイルに書き込む
func ExportToJSON(path string, posts []Post) error {
    postjson, err := json.Marshal(posts)
    if err != nil {
        return err
    }
    return os.WriteFile(path, postjson, 0644)
}

改善後

GoのIteratorを使って上記ソースコードを以下のように書き換えました。また投稿データはJSONファイルではなく、JSONLファイルに書き込むようにしました。

func main() {
    if ExportToJSON(path, ScrapePosts()); err != nil {
        log.Fatalf(err)
    }
}

// ScrapePosts ユーザー投稿をスクレイピングする
func ScrapePosts() iter.Seq2[Post, error] {
    return func(yield func(Post, error) bool) {
        // ...
        // スクレイピングした投稿データをyieldのパラメーターに渡す
        if !yield(post, nil) {
          return
        }
        // ...
    }
}

// ExportToJSONL スクレイピングしたユーザー投稿をJSONLファイルに書き込む
func ExportToJSONL(path string, posts iter.Seq2[Post, error]) error {
	  f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
	  if err != nil {
	  	  return err
	  }
	  defer f.Close()
	  enc := json.NewEncoder(f)
	  // スクレイピング処理が続く限り
	  // 投稿データはどんどんJSONLファイルに追加する
	  for post, err := range posts {
	      if err != nil {
	          return err
	      }
	  	  if err := enc.Encode(post); err != nil {
	  		    return err
	  	  }
	  }
	  return nil
}

JSONLとは

ChatGPTによるとJSON Lines(または JSONL、newline-delimited JSON、NDJSON とも呼ばれる)は、各行に1つの有効なJSONオブジェクトを含む形式で、 構造化データをエンコードするためのフォーマットです。

{"id": 1, "title": "title1", "msg": "msg1"}
{"id": 2, "title": "title2", "msg": "msg2"}
{"id": 3, "title": "title3", "msg": "msg3"}

改善前後のメモリ費用

Goのベンチマークツールを使って改善前後のヒープアロケーションを比較しました。改善後にヒープアロケーションが大幅に減少したことが確認できました。ヒープアロケーションは投稿データの取得とJSON型への変換処理で発生したものです。

BenchmarkNoIter-8         421894              2468 ns/op            1752 B/op          7 allocs/op
BenchmarkWithIter-8       340784              3402 ns/op             656 B/op         15 allocs/op

メモリプロフィールで詳細を確認してみましょう。

改善前

改善する前のヒープアロケーションは1GB超えました。

改善後

改善した後に200MB程度に抑えることができました。

まとめ

IteratorがGoに導入されること発表されたとき批判した人々多いですが、個人的にIteratorが導入されることはよかったと思います。

皆さんはいかがでしたか?

どの場面でGoのIteratorを使用していますか?

Discussion