mongo-go-driver を使った go の iterator 実装⌗

先日 go 1.23 がリリースされた。このバージョンの目玉機能の1つに range-over function iterators がある。簡潔に言えば、ユーザー定義のイテレーターを言語機能の仕組みを使って簡単に実装できる。ある制約を満たして関数を実装すれば、for 文の range 構文を使ってイテレーターとして扱える。ユーザー定義のイテレーターを実装するのが、ほとんど関数を実装するだけで出来てしまうので実装の難易度が減ったと言える。どんなコードかをみるには次のチュートリアルがわかりやすいと思う。

昨日で私が抱えていたプロジェクトのボトルネックを解消できた。ここからはボーナスステージで相対的に自由に開発できる。手始めにプロダクトの依存バージョンを go 1.23 へアップグレードして mongodb からの fetch 処理をイテレーターに置き換えるリファクタリングをした。アプリケーションからみたら mongodb は Store という generic なインターフェースを定義していて、そこに Iter() メソッドを追加した。mongodb のコレクションは必ず Iter() メソッドを実装しなければいけないという制約を課す。

type Store[T any] interface {
    ...
    Iter(ctx context.Context, query query.Query, opt *sort.Option) iter.Seq2[*T, error]
}

mongo-go-driver はもともと cursor でイテレーター機能を提供しているため、これを range-over function iterators を使ってアプリケーションから使えるように間をつなげてあげればよい。具体的には generics を使って次のような汎用の iter() メソッドを mongodb client に実装する。

func (c *client[R, E]) iter(
	ctx context.Context, name string, filters queryFilters, sortOpt *sort.Option,
) iter.Seq2[*E, error] {
	if c.toEntity == nil {
		panic(fmt.Errorf("require toEntity() function"))
	}
	return func(yield func(*E, error) bool) {
		collection := c.raw.Database(c.db).Collection(name)
		opts := []*options.FindOptions{}
		if sortOpt != nil {
			opt := options.Find().SetSort(makeSortSet(*sortOpt))
			opts = append(opts, opt)
		}
		cursor, err := collection.Find(ctx, filters, opts...)
		if err != nil {
			if !yield(nil, err) {
				return
			}
		}
		defer cursor.Close(ctx)
		for cursor.Next(ctx) {
			var result R
			if err := cursor.Decode(&result); err != nil {
				if !yield(nil, err) {
					return
				}
			}
			if !yield(c.toEntity(&result), nil) {
				return
			}
		}
		if err := cursor.Err(); err != nil {
			if !yield(nil, err) {
				return
			}
		}
	}
}

Store インターフェースを満たすコレクションの実装は次のように型チェックのためのメソッド実装をもてばよい。

func (c *MyCollection) Iter(
	ctx context.Context, query query.Query, opts *sort.Option,
) iter.Seq2[*entry.MyEntity, error] {
	return c.iter(ctx, c.name, makeFilters(query), opts)
}

朝から iterator 実装の試行錯誤や使い方の勘どころを確かめながら、これも1日でほとんどリファクタリングして移行できた。python のジェネレーターもそうだけど、関数をイテレーターにできると直感的にわかりやすく簡潔に実装できる。すごくよい機能だと思う。1.23 以降の go のコードはイテレーターを使うよう、大きく変わっていくと思われる。