Goのreflectで任意の構造体のフィールド変数を1つ増やしちゃう

こちらはGo 5 Advent Calendar 2020の穴埋め記事です。

こんにちは、うえちょこ(uechoco)です。
DeNAから事業統合で株式会社Mobility Technologiesに転籍しまして、バックエンドエンジニアをしております。
前回も書きましたが、ええ、あの「GO」という名前のタクシー配車アプリを「Go」で作っている会社です 。

今回も業務で遭遇したネタです。

お題

任意の構造体を受け取って、 json.Marshal() してログに出力する既存実装がありました。

func outputLog(eventID string, payload interface{}) error {
    b, err := json.Marshal(payload)
    if err != nil {
        return err
    }

    // ログ出力する(例)
    fmt.Printf("event:%s payload:%s", eventID, string(b))

    return nil
}

引数の payload は任意の構造体を指定できるので、 User{...} とか Product{...} とか Order{...} とかなんでも指定することができました。

ところが、新しい機能が追加されるときに、今後はすべてのPayloadのJSONに一律で属性情報を含めてほしいと言われてしまいました。こんなイメージです：

• Before:
  • {"name":"田中","age":16}
  • {"name":"特製味噌ラーメン","price":780}
  • {"name":"さっきもらった指輪","paid_at":"2020-12-24T13:31:28Z"}
• After:
  • {"name":"田中","age":16,"_attribute":"v2"}
  • {"name":"特製味噌ラーメン","price":780,"_attribute":"v2"}
  • {"name":"さっきもらった指輪","paid_at":"2020-12-24T13:31:28Z","_attribute":"v2"}

最初は AttributedUser AttributedProduct といった構造体を新たに定義して値を入れようか、、、と思っていたのですが、「reflect使えば元の構造体にフィールド変数が1つ増えた無名の構造体をその場で作って値を詰めて返せるのでは？」と思い立って、作ってみました。

package main

import (
    "encoding/json"
    "fmt"
    "reflect"
)

type EmptyLogPayload struct {
    UnderscoreAttribute string `json:"_attribute"`
}

func NewEmptyLogPayload() EmptyLogPayload {
    return EmptyLogPayload{
        UnderscoreAttribute: "v2",
    }
}

func WrapLogPayload(src interface{}) interface{} {
    if src == nil {
        return NewEmptyLogPayload()
    }
    fv := reflect.ValueOf(src)
    ft := fv.Type()
    if fv.Kind() == reflect.Ptr {
        if fv.IsNil() {
            return NewEmptyLogPayload()
        }
        ft = ft.Elem()
        fv = fv.Elem()
    }
    switch ft.Kind() {
    case reflect.Struct:
        return forceInsertField(fv, ft)
    }
    panic(fmt.Sprintf("unsupported reflect Kind:%s", ft.Kind().String()))
}

func forceInsertField(fv reflect.Value, ft reflect.Type) interface{} {
    // 構造体のフィールドの収集
    ep := NewEmptyLogPayload()
    epRt := reflect.TypeOf(ep)
    epNum := epRt.NumField()
    num := ft.NumField()
    structFields := make([]reflect.StructField, 0, num+epNum)
    for i := 0; i < num; i++ {
        structFields = append(structFields, ft.Field(i))
    }
    for i := 0; i < epNum; i++ {
        structFields = append(structFields, epRt.Field(i))
    }
    // 構造体型の生成
    newType := reflect.StructOf(structFields)
    // 構造体の生成
    rv := reflect.New(newType).Elem()
    // 値の移植
    transplantFields := func(srcType reflect.Type, fieldNum int, srcValue reflect.Value, destValue reflect.Value) {
        for i := 0; i < fieldNum; i++ {
            field := srcType.Field(i)
            if !field.Anonymous {
                name := field.Name
                srcField := srcValue.FieldByName(name)
                dstField := destValue.FieldByName(name)
                if srcField.IsValid() && dstField.IsValid() {
                    if srcField.Type() == dstField.Type() {
                        dstField.Set(srcField)
                    }
                }
            }
        }
    }
    transplantFields(ft, num, fv, rv)
    transplantFields(epRt, epNum, reflect.ValueOf(ep), rv)
    return rv.Interface()
}

処理の流れ

WrapLogPayload 関数では、元のpayloadが構造体であることをチェックします（今回は構造体以外を扱わない）。ポインタの場合はポインタの中身が構造体であることをチェックします。nilだけは例外的に空の構造体という扱いでOKしました。

フィールド変数を1つ増やす処理は forceInsertField 関数にあります。

ここでは元の構造体のreflect.StructFieldと増やしたいフィールド変数を内包した構造体のreflect.StructFieldを合体させます。その名の通り構造体のフィールド変数の情報がここに含まれています。タグ情報も含まれているので、まるごとコピーすることでJSON変換時のキー名も引き継がれます。

そして reflect.StructOf(structFields) を呼び出すと新しい構造体型が出来上がります。その新しい構造体型の変数を作るのが reflect.New() です。

このままではゼロ値のままですので、元の構造体から値をコピーします。 transplantFields 内部関数の部分です。フィールド変数名経由でフィールドを特定して値をコピーしていますが、もしかしたらもっといいやり方があるかもしれません。

実行サンプル

これを使ったサンプル実装がこちらです: https://play.golang.org/p/X3spg0AsPZk

func main() {

    type User struct {
        Name string `json:"name"`
        Age  int    `json:"age"`
    }
    fmt.Printf("User: %+v\n", WrapLogPayload(User{Name: "田中", Age: 16}))

    type Product struct {
        Name  string `json:"name"`
        Price int    `json:"price"`
    }
    fmt.Printf("Product: %+v\n", WrapLogPayload(&Product{Name: "特製味噌ラーメン", Price: 780}))

    fmt.Printf("nil: %+v\n", WrapLogPayload(nil))

    b, _ := json.Marshal(WrapLogPayload(User{Name: "田中", Age: 16}))
    fmt.Printf("json: %s\n", string(b))

    b, _ = json.Marshal(WrapLogPayload(&Product{Name: "特製味噌ラーメン", Price: 780}))
    fmt.Printf("json: %s\n", string(b))

    b, _ = json.Marshal(WrapLogPayload(nil))
    fmt.Printf("json: %s\n", string(b))
}

出力結果

User: {Name:田中 Age:16 UnderscoreAttribute:v2}
Product: {Name:特製味噌ラーメン Price:780 UnderscoreAttribute:v2}
nil: {UnderscoreAttribute:v2}
json: {"name":"田中","age":16,"_attribute":"v2"}
json: {"name":"特製味噌ラーメン","price":780,"_attribute":"v2"}
json: {"_attribute":"v2"}

どの異なる構造体を渡しても、ちゃんとフィールド変数が1つ増えていますね。