この記事はスタンバイAdvent Calendar 2024の11日目の記事です。
昨日の記事は@ysaki_sapさんの「MLチームのチーム運営を改善しようとしている」でした。

はじめに

株式会社スタンバイで主にBFFをはじめとする各種バックエンドシステムの開発、保守運用を行っている熊本です。
スタンバイには2024年6月に入社し、約半年が経ったところです。最近、容量が大きいCSVファイルの読み込み処理を実装して、メモリ効率を意識した開発を行うことができたので、その振り返りとメモリ使用量改善まで至ったアプローチの共有を行いたいと思います。

GoでのCSVファイルの読み込み

まず、Goでは標準パッケージのencoding/csvのReaderを使ってCSVファイルを読み込むことができます。当初は以下のようにReaderのReadAllメソッドを使ってCSVファイルの読み込み処理を行っていました。

※今回CSVの各レコードに完全一致するデータの検索を行いたかったので、mapを使用してkeyにCSVの各レコードのデータを保持するstructを持たせています。mapのvalueは何でも良いですが、今回は空のstructを保持することにします。

// CSVの各カラムのデータを保持する構造体
type CSVRecord struct {
    Data1 string
	Data2 string
	Data3 string
	Data4 string
	Data5 string
}

func readCSV(path string) (map[CSVRecord]struct{}, error) {
    file, err := os.Open(path)
	if err != nil {
        return nil, err
	}
	defer file.Close()
 
    reader := csv.NewReader(file)
    records, err := reader.ReadAll() // ReadAllメソッドを使って読み取り
    if err != nil {
    	return nil, err
    }
    for lineNum, record := range records {
    	if lineNum != 0 {
    		records[CSVRecord{Data1: record[0], Data2: record[1], Data3: record[2], Data4: record[3], Data5: record[4]}] = struct{}{}
    	}
    }

    return records, nil
}

上記実装でのメモリ負荷増加とその調査

前述の処理はサーバー起動時に呼び出していました。実際に開発環境のリソースを使ってサーバーを立ち上げると、起動時約4GBまでメモリ使用量が上がり、OOMが発生しました。そのため、Goのpprofという標準パッケージのプロファイラを使ってどこでメモリ消費しているのか調査しました。

pprofでのメモリ調査を行う方法としては、以下のような関数を作成して測定対象の処理の後に呼び出します。

func GetMemProfile() {
	f, err := os.Create("mem.pprof")
	if err != nil {
		log.Fatal("Failed to create mem profile. ", err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			log.Fatal("Failed to close mem profile. ", err)
		}
	}()
	if err := pprof.WriteHeapProfile(f); err != nil {
		log.Fatal("Failed to write mem profile. ", err)
	}
}

その後、go tool pprofでどの関数がメモリ消費しているか確認します。

go tool pprof mem.pprof
File: main
Type: inuse_space
Time: Dec 10, 2024 at 8:39am (JST)
Entering interactive mode (type "help" for commands, "o" for options)
(pprof) top10
Showing nodes accounting for 972.95MB, 99.46% of 978.27MB total
Dropped 16 nodes (cum <= 4.89MB)
Showing top 10 nodes out of 11
      flat  flat%   sum%        cum   cum%
  550.54MB 56.28% 56.28%   550.54MB 56.28%  encoding/csv.(*Reader).readRecord
  295.90MB 30.25% 86.52%   944.41MB 96.54%  readCSV
  117.48MB 12.01% 98.53%   668.02MB 68.29%  encoding/csv.(*Reader).ReadAll

上記の結果から、ReadAll()、ReadAll()が呼び出すreadRecord()、readCSV関数で直接的にメモリ容量を要していることがわかります。

CSVファイル読み込み処理改善

CSVデータ読み込み時の改善

今回のCSVファイルのサイズは約240MBと比較的大きなファイルの読み込みを行っていました。ReadAll()の実装を見てわかる通り、データサイズ分だけメモリを確保して全てのレコード返す実装なので、ファイル内のデータ量に比例して必要とするメモリ容量が大きくなります。
readRecord()の実装も見てみると、readRecord()内のreadBufferにCSVの各レコードを保持して、バッファの容量が足りなくなったら都度メモリ割り当てを行っていることがわかります。そのため、readRecord()内で利用するメモリ容量が大きくなっていると考えられます。

上記の改善として、CSVのレコードの1行ずつ読み取り、さらに各レコードを割り当てるメモリを再利用する方法を取りました。以下はサンプルの実装です。

func readCSV(path string) (map[CSVRecord]struct{}, error) {
    file, err := os.Open(path)
	if err != nil {
        return nil, err
	}
	defer file.Close()
 
    reader := csv.NewReader(file)
    reader.ReuseRecord = true // 各レコードを割り当てるメモリを再利用する
    records := map[CSVRecord]struct{}{}
    
    for i := 0; ; i++ {
		record, err := reader.Read() // 一行ずつの読み取り
		if i == 0 {
			continue
		}
		if err == io.EOF {
			break
		} else if err != nil {
			util.Logger.Error(err, "ファイルの読み込みに失敗しました")
			return nil, errors.Wrapf(err, "ファイルの読み込みに失敗しました。path: %s", path)
		}

		records[CSVRecord{Data1: record[0], Data2: record[1], Data3: record[2], Data4: record[3], Data5: record[4]}] = struct{}{}
	}
    return records, nil
}

上記では、ReadAll()ではなく、csv.ReaderのRead()を使ってCSVレコードの読み込みを行っています。またRead()ではReuseRecordというcsv.Readerが持つフラグで条件分岐しています。このフラグを利用することで、各レコードが割り当てられるメモリの再利用を行うことができます。ただメモリの再利用を行うため、複数のgoroutineで同じcsv.Readerを利用する場合はデータ競合が発生する可能性があるため注意が必要です。今回はgoroutineでの呼び出しは行わなかったため利用しました。

readCSV関数の改善

pprofの結果からreadCSV関数内でもメモリ容量を要していることがわかります。まず改善策としてmapの初期定義の際にサイズ指定して定義するようにしました。前述の実装では初期定義の際にサイズを指定していなかったので、都度メモリ割り当てを行われて不要なメモリまで確保される可能性がありました。
またCSVのレコードを保持する構造体CSVRecordの各要素がstring型であったため、ハッシュ化してint型で持つ対応も行いました。
Goのstring型のメモリ消費については以下の記事で詳しく説明されています。

上記の改善を行った実装は以下です。

func readCSV(path string, size int) (map[CSVRecord]struct{}, error) {
    file, err := os.Open(path)
	if err != nil {
        return nil, err
	}
	defer file.Close()
 
    reader := csv.NewReader(file)
    reader.ReuseRecord = true
    expandAdQueries := make(map[uint64]struct{}, size) // 引数で受け取ったサイズを指定してmapのキーの型をuint64にする
    
    for i := 0; ; i++ {
		record, err := reader.Read()
		if i == 0 {
			continue
		}
		if err == io.EOF {
			break
		} else if err != nil {
			util.Logger.Error(err, "ファイルの読み込みに失敗しました")
			return nil, errors.Wrapf(err, "ファイルの読み込みに失敗しました。path: %s", path)
		}

		records[CSVRecord{Data1: record[0], Data2: record[1], Data3: record[2], Data4: record[3], Data5: record[4]}.ToUint()] = struct{}{}
	}
    return records, nil
}

func (r CSVRecord) ToUint() uint64 {
	hasher := fnv.New64a()
    // NOTE: 以下のようなデータをハッシュ化したときに異なるデータとして扱うためハッシュ化する際に各カラムごとに接頭辞を付けています。
    // r1 := CSVRecord{Data1: "A", Data2: "", Data3: "", Data4: "", Data5: ""}
    // r2 := CSVRecord{Data1: "", Data2: "A", Data3: "", Data4: "", Data5: ""}
	hasher.Write([]byte("data1" + r.Data1))
	hasher.Write([]byte("data2" + r.Data2))
	hasher.Write([]byte("data3" + r.Data3))
	hasher.Write([]byte("data4" + r.Data4))
	hasher.Write([]byte("data5" + r.Data5))
	return hasher.Sum64()
}

改善結果

以下は前述の改善後にpprofでメモリ使用量を確認した結果です。

go tool pprof mem.pprof
File: main
Type: inuse_space
Time: Dec 10, 2024 at 8:14am (JST)
Entering interactive mode (type "help" for commands, "o" for options)
(pprof) top10
Showing nodes accounting for 115.83MB, 99.57% of 116.33MB total
Dropped 17 nodes (cum <= 0.58MB)
Showing top 10 nodes out of 36
      flat  flat%   sum%        cum   cum%
      85MB 73.07% 73.07%       85MB 73.07%  readCSV
   14.50MB 12.47% 85.53%    14.50MB 12.47%  encoding/csv.(*Reader).readRecord

上記の結果から、約240MBのCSVファイルの読み込み処理において944.41MB消費していたものが85MBまで下がり、大幅にメモリ使用量の改善ができたことがわかります。
また開発環境で再度起動時のメモリ使用量を確認したところ、4GB消費していたものが700MBまで下がったことも確認できました。

まとめ

今回のメモリ使用量改善を行ったことで、メモリリソースを必要以上に確保する必要がなくなったので、対応できて嬉しく思います。やはりファイルのI/O周りは各種リソースを意識しながら開発することが大切だと身にしみて感じました。同じようにメモリ消費量のチューニングを行う際に参考にしていただければ嬉しいです。

最後までお読みいただき、誠にありがとうございました。

参考資料