遅延評価を使ってOSSのバグを修正した話

　以前は殆ど実社会に応用されておらず、ガウスをして「数学は科学の女王であり、数論は数学の女王である。」と言わしめた数論が暗号という活躍の場を得て花開いた様に、どんな技術が役に立つかは実に分からないものです。

　それと比べると幾分スケールは落ちますが、Ruby の gem を開発しているにも関わらず関数型言語で使われている遅延評価の知識が役立つといった、少し興味深い事が起きたので記事にまとめておきましょう。ウォンテッドリーで開発している OSS の pb-serializer のバグを修正した際のエピソードの紹介です。

pb-serializer とは

　pb-serializer は主にウォンテッドリーで開発している gem で、ActiveModelSerializers と似た DSL によって Protocol Buffers メッセージのシリアライザを記述出来るようにする機能を提供しています。例えば、次の User の様なモデルに対応する同名のメッセージが定義されていた場合、

# Schema: [id(integer), name(string)]
class User < ActiveRecord::Base
end

syntax = "proto3";

package example;

option ruby_package = "ExamplesPb";

message User {
  uint64 id = 1;
  string name = 2;
}

次の UserPbSerializer の様にシリアライザを実装できます。変換先のメッセージとそのフィールドの情報を pb-serializer の DSL で指定する感じですね。

class UserPbSerializer < Pb::Serializer::Base
  message ExamplesPb::User

  attribute :id
  attribute :name
end

実際にこの UserPbSerializer を使って User モデルのインスタンスをシリアライズすると、以下の様になります。

user = User.find(123)
UserPbSerializer.new(user).to_pb
# => <ExamplesPb::User: id: 123, name: "someuser">

　この pb-serializer は特定のフィールドだけメッセージに含める様なユースケースにも対応しており、オーバーフェッチ問題の解決にも便利です。例えば上に挙げた User モデルの id フィールドのみを読み出してメッセージとして返したい場合、to_pb の引数に FieldMask ないしそれに対応するハッシュを渡すことでこれを実現できます。

read_mask
# => <Google::Protobuf::FieldMask: paths: ['id']>
UserPbSerializer.new(user).to_pb(with: read_mask)
# => <ExamplesPb::User: id: 123, name: "">
UserPbSerializer.new(user).to_pb(with: [:id])
# => <ExamplesPb::User: id: 123, name: "">

　さらに、入れ子になったメッセージを簡単に取り扱うための機能も用意されています。例えば、初めの例で挙げた User モデルに関連付く Post モデルと、それに対応する同名のメッセージが定義されていた場合、

# Schema: [id(integer), title(string), author_id(integer)]
class Post < ActiveRecord::Base
  belongs_to :author, class_name: 'User'
end

message Post {
  uint64 id = 1;
  string title = 2;
  User author = 3;
}

User モデルのシリアライザ UserPbSerializer を attribute で指定するだけで、関連付け author も同時に変換するようなシリアライザ PostPbSerializer を定義できます。

class PostPbSerializer < Pb::Serializer::Base
  message ExamplesPb::Post

  attribute :id
  attribute :title
  attribute :author, serializer: UserPbSerializer
end

　本題から逸れるのでこれ以上は解説しませんが、pb-serializer には他にも様々な機能が用意されているので、もし興味があれば README を参照して下さい。

以前存在した pb-serializer のバグ

　この様に Protocol Buffers メッセージの生成に有用な pb-serializer ですが、以前のバージョンでは再帰的なスキーマを持つメッセージが（FieldMask 等を指定しない場合に）シリアライズできないバグを抱えていました。例えば、次の StringList の様なメッセージに対してシリアライザ StringListPbSerializer を定義し、

message StringList {
  string car = 1;
  StringList cdr = 2;
}

class StringListPbSerializer < Pb::Serializer::Base
  message ExamplePb::StringList

  attribute :car
  attribute :cdr, allow_nil: true, serializer: StringListSerializer
end

適当なオブジェクトのシリアライズを試みると無限ループしてしまうのです。

StringList = Struct.new(:car, :cdr)

l = StringList.new("Alpha", StringList.new("Bravo", nil))

StringListSerializer.new(l).to_pb
# /usr/local/lib/ruby/3.1.0/set.rb:609:in `merge': stack level too deep (SystemStackError)

　この問題は FieldMask 等で指定した特定のフィールドだけ Protocol Buffers メッセージに含める場合と、FieldMask 等を渡さず全てのフィールドをメッセージに含める場合の実装を共通化した事に由来していました。 pb-serializer の実装では to_pb の引数に FieldMask 等を与えなかった場合は、シリアライズ先のメッセージに含まれうる全てのフィールドを指定するハッシュ build_default_mask(descriptor) を生成して与える事で、FieldMask 等で指定した特定のフィールドだけ Protocol Buffers メッセージに含める場合と実装を共通化しています。

# https://github.com/wantedly/pb-serializer/blob/dfefdefce234e599cc98176454ea880694ade80d/lib/pb/serializer.rb#L67-L91
def build_default_mask(descriptor)
  set =
    descriptor.each_with_object(Set[]) do |fd, m|
      case fd.type
      when :message
        case fd.submsg_name
        when "google.protobuf.Timestamp", 
          "google.protobuf.StringValue",
# （中略）
          "google.protobuf.BytesValue" then m << fd.name.to_sym
        else
          m << { fd.name.to_sym => build_default_mask(fd.subtype) }
        end
      else
        m << fd.name.to_sym
      end
    end
  set.to_a
end

# https://github.com/wantedly/pb-serializer/blob/b761f7f519199f413b711424d4dfe98b01f5c3fa/lib/pb/serializable.rb#L25C1-L25C1
def to_pb(with: nil)
  with ||= ::Pb::Serializer.build_default_mask(self.class.__pb_serializer_message_class.descriptor)
  with = ::Pb::Serializer.normalize_mask(with)

ところが、StringList の様にスキーマが再帰的に定義されていると StringList.new(car: "Alpha", cdr: StringList.new(car: "Bravo", cdr: StringList.new(car: "Charlie", cdr: ...))) と、（入れ子になったメッセージの分も含めて）メモリの続く限り多くのフィールドを持つメッセージが作れてしまうので、build_default_mask(descriptor) の戻り値も [:car, cdr: [:car, cdr: [:car, cdr: ...]]] と再現なく肥大化し、再帰呼び出しが止まらなくなってしまうのです。

irb(main):001:0> ::Pb::Serializer.build_default_mask(StringListSerializer.__pb_serializer_message_class.descriptor)
/usr/local/lib/ruby/3.1.0/set.rb:609:in `merge': stack level too deep (SystemStackError)

バグの修正

　再帰的なスキーマを持つメッセージのシリアライザを pb-serializer で定義した際の問題は、メッセージに含めるフィールドを指定した場合とそうでない場合の実装を共通化するために全てのフィールドをメッセージに含めるよう指定するハッシュを構成していたはずが、そのハッシュが無限に大きくなってしまう事にありました。素朴な解決法としてはメッセージに含めるフィールドを指定しなかった場合の実装を分ければ良いのですが、殆ど同じコードを繰り返し記述することはソフトウェア工学的に好ましい事ではありません。最小限の変更で修正を行うために、無限に大きなハッシュを何とか扱えないものでしょうか？

　実は、Haskell の様な関数型言語では、無限に大きなデータ構造を扱うのはそう難しい事ではありません。例えば、

ones = 1 : ones

と書くだけで 1 が無限に続くリストを定義できますし、こうして定義した無限リストも無限ループに陥る事なく使うことができます。

take 5 ones
-- [1, 1, 1, 1, 1]

　この様に Haskell が無限に大きなデータ構造を容易に扱うことができるのは、遅延評価をサポートしているからに他なりません。先程の ones の例では、ones = 1 : ones と定義した段階では ones の中身は評価されず、take 関数の中でリストの要素が必要になる度に必要な分だけ評価が行われるため、無限ループに陥る事なく無限に大きなリストを取り扱えるのです。

　それでは pb-serializer の内部実装で、シリアライズ先のメッセージに含まれうる全てのフィールドを指定するハッシュ build_default_mask(descriptor) を生成している箇所で遅延評価を使い、無限ループが起こらない様に修正しましょう。一度にハッシュを生成するのではなく必要になった分だけ返すには、入れ子になったメッセージについてのハッシュを再帰的に生成している箇所で匿名関数を用い、呼び出し元で必要になるまで実行を中断すると良さそうです。

def build_default_mask(descriptor)
  set =
    descriptor.each_with_object(Set[]) do |fd, m|
      case fd.type
      when :message
        case fd.submsg_name
        when "google.protobuf.Timestamp", 
          "google.protobuf.StringValue",
# （中略）
          "google.protobuf.BytesValue" then m << fd.name.to_sym
        else
-         m << { fd.name.to_sym => build_default_mask(fd.subtype) }
+         m << { fd.name.to_sym => -> { build_default_mask(fd.subtype) } }
        end
      else
        m << fd.name.to_sym
      end
    end
  set.to_a
end

こうして build_default_mask(descriptor) の側で再帰呼び出しを中断するよう変更したので、それを用いているコードの側で必要に応じて実行を再開しなければなりませんね。

def normalize_mask(input)
  if input.kind_of?(Google::Protobuf::FieldMask)
    input = parse_field_mask(input)
  end

  normalized = {}

+ input = input.call if input.kind_of?(Proc)
  input = [input] if input.kind_of?(Hash)

  Array(input).each do |el|
    case el
    when Symbol
      normalized[el] ||= []
    when Hash
      el.each do |k, v|
+       v = v.call if v.kind_of?(Proc)
        v = [v] if v.kind_of?(Hash)
        normalized[k] ||= []
        normalized[k].push(*Array(v))

　以上の簡単な修正によって、build_default_mask(descriptor) の再帰呼び出しが無限に続く心配は無くなり、再帰的なスキーマを持つメッセージであっても正常にシリアライズが行える様になりました。

l = StringList.new("Alpha", StringList.new("Bravo", nil))

StringListSerializer.new(l).to_pb
# => <ExamplesPb::StringList: car: "Alpha", cdr: <ExamplesPb::StringList: car: "Bravo", cdr: nil>>