手把手教你寫個自己的 protoc-gen-go

• 你是否用過 protobuf 或 gRPC？

• 你們公司項目的 API 有沒有用到 proto 文件?

本文將帶你一步一步寫個類似protoc-gen-go-grpc的 proto 文件生成工具，從 proto 文件生成兼容 Go 標準庫的 HTTP 框架代碼。掌握它之後，你就能隨心所欲的從 proto 文件生成gin、echo、net/http代碼，甚至生成你自己的框架代碼。

別擔心，生成的內容不侷限於 Go 語言，別的語言也沒問題，甚至不是編程語言都可以！

你可以從 proto 文件生成任何它能描述的東西。😉

我們正在做什麼?

在學習 gRPC 時，你執行了一段命令，就將 proto 文件變成了 gRPC 代碼:

protoc -I api \
  --go_out=internal/genproto/$service \
  --go_opt=paths=source_relative \
  --go-grpc_out=internal/genproto/$service \
  --go-grpc_opt=paths=source_relative \
  $service.proto

這多虧了protoc-gen-go和protoc-gen-go-grpc這兩個可執行文件。

我們在執行protoc -xxx_out=. -xxx_opt=.命令時，protoc 會從你操作系統的 $PATH 目錄下尋找protoc-gen-xxx這個可執行文件進行執行，並把後面的參數傳給它。

你可以執行ls $GOPATH/bin | grep protoc命令來查看自己電腦上安裝了哪些 protoc 生成工具:

筆者電腦上安裝的 protoc-gen 系列生成工具

本文會帶你實現一個名叫protoc-gen-go-example的工具，在使用時，我們只需要執行以下命令，即可調用我們自己寫的生成工具來生成代碼:

protoc --go-example_out=. # 此處省略其他參數

好啦，看到這兒你也應該明白我們在做什麼事情啦，我們直接進入正題吧！

站在巨人的肩膀上

如果你學過編譯原理，你一定很清楚我們要做些什麼 (沒學過的同學先別跑😢):

1. 解析. proto 文件，構建 proto 文件的 AST(抽象語法樹)

2. 遍歷 AST，將其轉換爲想要生成的內容。

天哪，這要是從零實現，需要多大的工程量啊！更別提一些沒學過編譯原理的同學們了。我要是從零開始教，那能寫一本書了...

幸運的是，我們有一些可以利用的東西！不需要我們自己去實現 proto 文件的 Parser 啦！

protocolbuffers/protobuf-go[1] 這個庫 (也就是 protoc-gen-go) 已經幫我們實現了工作量最大的 parser 部分。

這下我們可以繼續一起愉快的玩耍了!

創建項目

我創建的項目叫protoc-gen-go-example，這就是我們最終生成的二進制文件名稱。

我們在執行go install命令時，默認會以main.go的上一級目錄名作爲可執行文件的名稱。

假設我們的main.go文件放在根目錄下，我們執行go install github.com/bootun/protoc-gen-test-name後，就會在你的 $GOPATH/bin 目錄下安裝一個名爲protoc-gen-test-name可執行文件。

你可能會覺得：“那這樣我的項目名豈不是很醜😢”。如果你不想讓項目名作爲最終的文件名稱，你可以參考 protocolbuffers/protobuf-go 的做法。protobuf-go 把main.go放在了項目的cmd/protoc-gen-go目錄下，這樣在執行go install時，生成的文件就不會與項目名一樣了，但代價就是 go install 的路徑也會變長:

go install google.golang.org/protobuf/cmd/protoc-gen-go

想了解更多可以去看看 go 的官方文檔 [2]

執行以下命令來初始化項目:

mkdir protoc-gen-go-example
cd protoc-gen-go-example
touch main.go
go mod init github.com/bootun/protoc-gen-go-example

現在你的項目看起來像下面這樣:

protoc-gen-go-example
├── main.go
└── go.mod

現在讓我們來編輯 main.go 文件:

package main

import (
    "github.com/bootun/protoc-gen-go-example/parser"
    "google.golang.org/protobuf/compiler/protogen"
)

func main() {
    protogen.Options{}.Run(func(gen *protogen.Plugin) error {
        // 這個循環遍歷所有要生成的proto文件
        for _, f := range gen.Files {
            if !f.Generate {
                // 如果該文件不需要生成，則跳過
                continue
            }
            // 如果需要生成，就把文件的相關信息傳遞給生成器
            if err := parser.GenerateFile(gen, f); err != nil {
                return err
            }
        }
        return nil
    })
}

還記得我之前說過的嗎？我們在 main.go 中使用了protobuf-go中的組件，這樣我們就不需要從零開始解析 proto 文件中的內容了。

protogen.Options{}.Run()的參數是一個回調函數，回調函數的gen參數裏包括了所有已經解析好的信息。gen.Files表示所有 proto 文件的集合，我們需要遍歷這些 proto 文件，來爲它們生成代碼。

我們把 gen 和 file 向下傳遞，以便下面的組件能夠獲得足夠多的信息。現在parser.GenerateFile還在報錯，我們來實現GenerateFile這個函數:

GenerateFile 函數

GenerateFile 函數還是比較清晰明瞭的:

func GenerateFile(gen *protogen.Plugin, file *protogen.File) error {
    // 如果這個proto文件裏沒寫service
    // 我們就不需要爲它生成代碼
    if len(file.Services) == 0 {
        return nil
    }

    // 要生成的文件名稱
    filename := file.GeneratedFilenamePrefix + ".example.pb.go"
    g := gen.NewGeneratedFile(filename, file.GoImportPath)

    return NewFile(g, file).Generate()
}

代碼裏有註釋的部分我就不額外解釋了，很容易理解。我們需要額外關注下NewGeneratedFile這個函數。

NewGeneratedFile 使用給定的文件名和 ImportPath 創建一個新的生成文件實體，我們將它命名爲g。那這 ImportPath 又是個啥東西呢？

看過我上一篇文章的小夥伴們應該比較清楚，ImportPath 就是我們在 proto 文件中寫的option go_package裏的內容。protobuf-go幫我們做了許多處理，使得我們不需要過度關注像--xxx_opt=paths=source_relative等這種與代碼生成邏輯無關的內容，感興趣的話可以去看我的上篇文章——徹底搞清 protobuf-go 的文件生成位置。

有了g之後，我們只需要調用g.P("xxx")方法，即可在文件中寫入對應的字符串。

看到這裏你可能就明白了，我們只需要創建一套模板，將file參數中的信息套到這個模板上，然後傳給g.P(模板字符串)就行了。沒錯，就是這麼簡單！

在 NewGeneratedFile 函數的最後，我們調用 NewFile 創建了一個文件結構，並執行了該結構體上的 Generate 方法，整個代碼生成工作就完成了。讓我們看看 NewFile 裏做了什麼工作吧。

NewFile 函數

func NewFile(gen *protogen.GeneratedFile, protoFile *protogen.File) *File {
    f := &File{
        // 保存example.pb.go的文件實體
        // 以便後面操作
        gen: gen,
    }

    f.PackageName = string(protoFile.GoPackageName)

    for _, s := range protoFile.Services {
        f.ParseService(s)
    }

    return f
}

我在 NewFile 中創建了一個File結構體，這是我們自定義的一個結構體，用來表示一個 proto 文件的內容。

這個結構體不是必要的，甚至可能是多餘的，它只是把protogen.File參數裏的內容給轉成了我們的內部表示，所有的信息protogen.File裏都有，如果你想的話，你可以直接使用protogen.File+text/template來生成文件。這裏我出於教學目的，希望你能更容易明白這個過程，同時也爲了日後做些更騷的操作，就留下這個結構了。

proto 文件的內部表示

剛剛我提到了File結構體，說它只是把protogen.File裏的一部分信息複製出來，轉爲我們自己的內部結構體了，事實上除了File之外，還有幾個同樣表示 proto 信息的結構體，他們都是File結構的下屬結構:

// 一個File表示一個proto文件的信息
type File struct {
    // File內部同時保存了example.pb.go的文件句柄
    // 方便我們直接調用gen.P向pb文件寫入內容
    gen *protogen.GeneratedFile

    // 內嵌了一個FileDescription結構
    // 更多信息可以繼續往下看
    FileDescription
}

// FileDescription 描述了一個解析過後的proto文件的信息
// 爲我們後邊的代碼生成做準備
type FileDescription struct {
    // PackageName 代表我們生成後的example.pb.go文件的包名
    // 也就是go文件中的 package xxx
    PackageName string

    // Services 代表我們生成後的example.pb.go文件中的所有服務
    // 我們在proto文件中寫的每個server都會轉化爲一個 Service 實體
    Services []*Service
}

type Service struct {
    // Service 的名稱
    Name string

    // Service 裏具有哪些方法
    Methods []*Method
}

type Method struct {
    // 方法名稱
    Name string
    // 請求類型
    RequestType string
    // 響應類型
    ResponseType string
}

這些結構結合起來描述了一個簡單的 proto 文件信息：

proto 源文件對應的結構體信息

因爲是教學的緣故，所以各種類型的信息都很簡單，幾乎都用字符串存儲，只保留了最核心的內容。接下來，我們需要把信息從protogen.File裏複製到我們自己的結構體裏。

複製 proto 信息到內部表示中

還記得上面的NewFile函數嗎？裏面有這樣一段代碼:

for _, s := range protoFile.Services {
    f.ParseService(s)
}

這段代碼遍歷protoFile中所有的 Service, 並調用f.ParseService()方法來處理 proto 中的每個 service:

func (f *File) ParseService(protoSvc *protogen.Service) {
    s := &Service{
        Name:    protoSvc.GoName,
        Methods: make([]*Method, 0, len(protoSvc.Methods)),
    }

    for _, m := range protoSvc.Methods {
        // 遍歷並處理Service中的所有Method
        s.Methods = append(s.Methods, f.ParseMethod(m))
    }

    f.FileDescription.Services = append(f.FileDescription.Services, s)
}

func (f *File) ParseMethod(m *protogen.Method) *Method {
    return &Method{
        Name:         m.GoName,
        RequestType:  m.Input.GoIdent.GoName,
        ResponseType: m.Output.GoIdent.GoName,
    }
}

ParseService又會調用ParseMethod方法來遍歷處理 service 中的每個 method，我將它們兩個的代碼一起貼上來了，裏面的邏輯很簡單，就是從protogen的對應結構裏找到我們需要的屬性複製過來，解析工作就完成了。

現在，我們的File結構體被 "填滿了", 裏面保存了一個 proto 文件 (比較粗略) 的信息。接下來讓我們來創建一套模板，這將是代碼生成的最後一步。

模板代碼

在給你代碼之前，我想先明確一下，我在例子中生成的是 “基於 Go 標準庫net/http的框架代碼”。當然，你可以生成 gin 或其他框架的代碼，這全看你自己。但在寫模板之前，我們要先想想，我們要生成什麼樣的代碼? 使用者又希望你能幫他做哪些事?

要知道，proto 文件不是爲 gRPC 而生的，除了 gRPC, Transport 層的框架多到數不清，gin/echo/chi 等都算 Transport 層的框架。

因此，站在業務工程師的角度上，我希望能將關注點放在業務代碼上，業務代碼中不能包含任何傳輸層的細節，這樣我就可以隨時以很低的成本更換傳輸層的框架。

疊個甲: 這裏的 Transport 層和傳輸層指的不是網絡協議中的傳輸層，別噴。

所以站在使用者的角度上，我們可能會寫出以下代碼:

func main() {
    // 初始化Transport
    mux := http.NewServeMux()
    // 初始化業務依賴
    svc := UserService{
        store: make([]User, 0),
    }
    // 將業務Service註冊到Transport框架中
    user_pb.RegisterUserServiceHTTPServeMux(mux, &svc)
    // 啓動Transport框架
    if err := http.ListenAndServe(":8080", mux); err != nil {
        panic(err)
    }
}

// 業務Handler
type UserService struct {
    store []User
}

func (u *UserService) GetUser(ctx context.Context, req *user_pb.GetUserRequest) (resp *user_pb.GetUserResponse, err error) {
    // 這裏寫GetUser的業務代碼
}

func (u *UserService) CreateUser(ctx context.Context, req *user_pb.CreateUserRequest) (resp *user_pb.CreateUserResponse, err error) {
    // 這裏寫CreateUser的業務代碼
}

上面這段代碼中，業務代碼的GetUser和CreateUser中沒有任何 Transport 層的內容，業務代碼不知道上層使用的是 HTTP 還是 gRPC，又或者是 gin 等其他框架。

那我們就按照這個格式，來抽象出一個接口，作爲和業務之間的契約。

如果你用過 gRPC，你會發現: gRPC 也是這套 “契約”，這意味着未來我們要從net/http遷移到gRPC時，業務代碼不需要進行任何的改造，天然適配！

那爲了能讓上面那段業務代碼能夠正常運行，我們先來手寫一遍框架代碼，來 “適配” 上面的業務代碼。

這有點類似 TDD(Test-Driven Development) 的味道，從使用者的角度上開始，來定義代碼應該 “長什麼樣”。

我們很容易就能寫出下面的適配代碼, 這將是我們模板的雛形:

// 這個接口就是業務和框架的“契約”
// 實現這個接口的結構都能夠註冊進我們的框架中
// 這個Service對應着proto文件的service
type ServiceNameService interface {
    // 這裏就是service的方法列表，對應着proto文件中service的方法列表
    ServiceMethodName(ctx context.Context, req *MethodRequestName) (resp *MethodResponseName, err error)
}

// 業務代碼通過下面這段代碼將服務註冊到我們的框架中
func RegisterServerNameServiceHTTPServeMux(mux *http.ServeMux, svc ServiceNameService) {
    // 這裏用到了依賴注入的思想
    // 此時業務代碼是依賴，通過接口的形式注入進來
    s := ServiceName{
        svc: svc,
    }
    // 將對應的方法綁定到相應的路由上
    mux.HandleFunc("/UserCode", s.Name)
}

// 框架service具體實現，裏面通過接口保存了業務結構體
type ServiceName struct {
    svc ServiceNameService
}

// 每個service下都會有數個method
// 每個method也都對應着proto文件裏service的method
// 這裏用到了適配器(Adaptor)的設計思想
// 將業務代碼(通過接口)與 net/http 轉換，把它們“打通”
func (s *ServiceName) Name(rw http.ResponseWriter, r *http.Request) {
    // 我們在這個函數中要做的就是把HTTP請求中的內容解析出來
    // 嘗試將其轉換成業務需要的參數
    _ = r.ParseForm()
    var req MethodRequestName // 這個結構是protobuf生成的，和框架無關
    switch r.Method {
    // 出於教學目的，這裏只支持了POST請求
    case http.MethodPost:
        if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
            rw.WriteHeader(http.StatusBadRequest)
            rw.Write([]byte(err.Error()))
            return
        }
    default:
        rw.WriteHeader(http.StatusMethodNotAllowed)
        rw.Write([]byte(`method not allowed`))
        return
    }
    // 到這裏就順利的把HTTP請求轉爲了業務所需要的Request類型了
    // 接下來我們把控制權交給業務代碼吧
    resp, err := s.svc.ServiceMethodName(r.Context(), &req)
    if err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        rw.Write([]byte(err.Error()))
        return
    }
    // 將業務代碼返回的Response類型再轉爲HTTP請求返回給客戶端
    if err := json.NewEncoder(rw).Encode(resp); err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        rw.Write([]byte(err.Error()))
        return
    }
}

你可以看到，上面的適配代碼其實挺簡陋的，它只支持 POST 請求，甚至 HTTP 路由都是 proto 文件裏 method 的名字。但這對於我們學習它的核心原理已經夠用了。

即使是個玩具級別的 demo，它依舊用到了許多設計模式。

PS: 如果這篇文章反響還不錯的話，可能會考慮後續繼續加點東西。這篇文章我從晚上 8 點開始寫，寫到這裏已經凌晨 1:15 了😢

知道了我們的模板大概長什麼樣子後，剩下的就簡單了，替換上面代碼中的 Name 等各個部分，就得到了我們最終的模板代碼:

package template

const HTTP = `// Code generated by github.com/bootun/protoc-gen-go-example. DO NOT EDIT.
package {{.PackageName}}

import (
    "context"
    "encoding/json"
    "net/http"
)

{{range $service := .Services}}
type {{$service.Name}}Service interface {
{{range $method := .Methods}}
    {{$method.Name}}(ctx context.Context, req *{{$method.RequestType}}) (resp *{{$method.ResponseType}}, err error){{end}}
}

type {{$service.Name}} struct {
    svc {{$service.Name}}Service
}

func Register{{$service.Name}}HTTPServeMux(mux *http.ServeMux, svc {{$service.Name}}Service) {
    s := {{$service.Name}}{
        svc: svc,
    }
    {{range $method := .Methods}}
    mux.HandleFunc("/{{$method.Name}}", s.{{$method.Name}}){{end}}
}

{{range $method := .Methods}}
func (s *{{$service.Name}}) {{$method.Name}}(rw http.ResponseWriter, r *http.Request) {
    _ = r.ParseForm()
    var req {{$method.RequestType}}
    switch r.Method {
    case http.MethodPost:
        if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
            rw.WriteHeader(http.StatusBadRequest)
            return
        }
    default:
        rw.WriteHeader(http.StatusMethodNotAllowed)
        return
    }
    resp, err := s.svc.{{$method.Name}}(r.Context(), &req)
    if err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        return
    }
    if err := json.NewEncoder(rw).Encode(resp); err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        return
    }
}
{{end}}

{{end}}
`

如果你看不懂上面的語法，你需要去看下 go 的text/template，或者你有其他的辦法能拼湊渲染出這段字符串也可以。

我們只需要將我們內部結構中的數據 “填充” 到模板裏，交給前文提到的g.P()進行打印就可以啦:

func (f *File) Generate() error {
    tmpl, err := template.New("example-template").Parse(example_tmpl.HTTP)
    if err != nil {
        return fmt.Errorf("failed to parse example template: %w", err)
    }
    buf := &bytes.Buffer{}
    if err := tmpl.Execute(buf, f.FileDescription); err != nil {
        return fmt.Errorf("failed to execute example template: %w", err)
    }
    f.gen.P(buf.String())
    return nil
}

至此，我們就完成了所有的代碼編寫。

我將完整代碼上傳到了 GitHub 上: https://github.com/bootun/protoc-gen-go-example

你也可以使用以下命令來直接安裝

go install github.com/bootun/protoc-gen-go-example@latest

然後使用--go-example_out來生成代碼:

protoc -I ./api \
   --go_out=./user \
   --go_opt=paths=source_relative \
   --go-example_out=./user \
   --go-example_opt=paths=source_relative \
    api/user.proto

參考資料

[1]

protobuf-go: https://github.com/protocolbuffers/protobuf-go

[2]

go-install: https://golang.google.cn/ref/mod#go-install

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/Azulcj2kvxGiXzkAwKlI3g