數據流的藝術：Go 語言中的 io-Pipe

在 Go 語言中，使用 io.Pipe() 來流式處理數據可以避免將所有數據一次性讀入內存。

io.Reader 和 io.Writer 接口堪稱 Go 語言的藝術品，相關討論已有很多。它們簡潔而強大——正如 Go 本身。本文介紹 Go 標準庫中另一個同樣簡潔強大的存在：io.Pipe。

// src/io/pipe.go

func Pipe() (*PipeReader, *PipeWriter) {
 pw := &PipeWriter{r: PipeReader{pipe: pipe{
  wrCh: make(chan []byte),
  rdCh: make(chan int),
  done: make(chan struct{}),
 }}}
 return &pw.r, pw
}

什麼是 io.Pipe？

io.Pipe 在內存中創建一個同步管道。它用於連接需要 io.Reader 的代碼和需要 io.Writer 的代碼。一端的讀取對應另一端的寫入，數據直接在兩者之間傳輸，不進行中間緩衝。

查看官方文檔的代碼示例：

 pr, pw := io.Pipe()

 go func() {
  fmt.Fprint(pw, "some io.Reader stream to be read\n")
  pw.Close()
 }()

 if _, err := io.Copy(os.Stdout, pr); err != nil {
  log.Fatal(err)
 }

代碼流程如下：

  1. io.Pipe 返回 io.PipeReader (pr) 和 io.PipeWriter (pw)。

  2. 在另一個 Go 協程中向 pw 寫入數據。pw.Write 會阻塞，直到 pr.Read 被調用。

  3. pr.Read 調用後，管道內部緩衝區中的數據被複制到 pr（即讀取到 pr 傳入的緩衝區），然後管道內部的相應緩衝區空間被釋放。pr.Read 會阻塞，直到 pw 再次寫入數據或 pw.Close 被調用。

  4. pw.Close 調用後，pr.Read 返回 EOF，處理結束。

使用 io.Pipe 時務必警惕死鎖——當讀寫操作同處一個 goroutine 時，二者將陷入相互等待的僵局，導致流程徹底停滯。避免這種情況的關鍵是啓動獨立協程執行讀寫，確保兩者互不阻塞地同步進行。

示例 1：減少文件上傳的內存使用量

思考這樣一段文件上傳的代碼，考慮其性能表現：

func uploadFile(file io.Reader, url, fieldName, fileName string) error {
 buf := &bytes.Buffer{}

 mw := multipart.NewWriter(buf)
 fw, _ := mw.CreateFormFile(fieldName, fileName)
 io.Copy(fw, file)
 mw.Close()

 resp, err := http.Post(url, mw.FormDataContentType(), buf)
 if err != nil {
  return err
 }
 defer resp.Body.Close()

 return nil
}

思考數據流，主要有兩個步驟：

  1. 在將文件內容編碼的同時，通過 io.Copy 寫入緩衝區 (bytes.Buffer)。

  2. 在 http.Post 內部：通過 io.Copy 將緩衝區的內容寫入 HTTP 連接 (net.Conn)。

這些步驟是順序執行的，意味着在第一個 io.Copy 完成之前，無法開始第二個 io.Copy。因此，緩衝區需要一次性容納整個文件內容。文件越大，程序的內存使用量就越高。讓我們用 benchmem 的基準測試來驗證一下：

┌───────┬────────┐
│ File  │ Memory │
├───────┼────────┤
│ 6b    │ 39KB   │
│ 5KB   │ 44KB   │
│ 29KB  │ 71KB   │
│ 51KB  │ 169KB  │
│ 500KB │ 1.94MB │
│ 1MB   │ 1.96MB │
│ 10MB  │ 33.5MB │
└───────┴────────┘

可以清楚地看到內存使用量隨文件大小成比例增長，甚至常常遠超文件本身的大小。在這種實現方式下，隨着文件變大，內存不足的風險顯著增加。

multipart.Writer 接收 io.Writer，而 http.Post 接收 io.Reader，所以它們無法直接連接。如果能夠並行進行讀取和寫入，那麼理論上只需要保留當前正在讀取的數據量即可。io.Pipe 正可以實現這一點。

那麼，讓我們利用 io.Pipe 實現高效的文件上傳：

func uploadFileAsync(file io.Reader, url, fieldName, fileName string) error {
 pr, pw := io.Pipe()
 mw := multipart.NewWriter(pw)

 go func() {
  defer pw.Close()
  defer mw.Close()
  fw, _ := mw.CreateFormFile(fieldName, fileName)
  io.Copy(fw, file)
 }()

 resp, err := http.Post(url, mw.FormDataContentType(), pr)
 if err != nil {
  return err
 }
 defer resp.Body.Close()

 return nil
}

這次，我們將 io.PipeWriter (pw) 傳遞給 multipart.NewWriter。在另一個 Go 協程中，我們異步地將文件內容寫入由 mw.CreateFormFile 創建的 fw（它實際上是在寫入 pw），然後關閉 mw 和 pw。同時，我們將 io.PipeReader (pr) 傳遞給 http.Post 作爲請求體。這樣，就並行地執行了上文中提到的兩個步驟。

讓我們再次通過基準測試來驗證：

┌───────┬─────────┐
│ File  │ Memory  │
├───────┼─────────┤
│ 6b    │ 71KB    │
│ 5KB   │ 71KB    │
│ 29KB  │ 71KB    │
│ 51KB  │ 71KB    │
│ 500KB │ 71KB    │
│ 1MB   │ 71KB    │
│ 10MB  │ 76KB    │
└───────┴─────────┘

這次，文件較小時似乎有一些開銷，但當文件大小超過幾十 KB 時，內存使用量大幅下降。這樣一來，即使處理大文件，內存不足的風險也幾乎消失了。

示例 2：使用 io.Pipe + io.MultiWriter 拆分數據

將一些數據，同時上傳到兩個不同的 HTTP 接口。

一般實現：

func upload(r io.Reader) error {
 contentForUpload, err := io.ReadAll(r)
 if err != nil {
  return err
 }

 if err := firstUpload(bytes.NewReader(contentForUpload)); err != nil {
  return err
 }

 if err := secondUpload(bytes.NewReader(contentForUpload)); err != nil {
  return err
 }

 return nil
}

使用 io.Pipe 對代碼進行重構：

func upload(r io.Reader) error {
 var g errgroup.Group

 fr, fw := io.Pipe()
 g.Go(func() error {
  return firstUpload(fr)
 })

 sr, sw := io.Pipe()
 g.Go(func() error {
  return secondUpload(sr)
 })

 g.Go(func() error {
  var err error
  defer func() {
   fw.CloseWithError(err)
   sw.CloseWithError(err)
  }()

  _, err = io.Copy(io.MultiWriter(fw, sw), r)
  return err
 })

 return g.Wait()
}

示例 3：文件流壓縮

考慮這樣一種場景：你需要實時壓縮數據，從數據源讀取內容，並直接將壓縮後的版本寫入目標位置，無需中間存儲：

func compressData(src io.Reader, dest io.Writer) error {
 r, w := io.Pipe()

 // Compress data and write to pipe in a separate goroutine
 go func() {
  gz := gzip.NewWriter(w)
  _, err := io.Copy(gz, src)
  gz.Close()
  w.CloseWithError(err)
 }()

 // Read from pipe and write to destination
 _, err := io.Copy(dest, r)
 return err
}

示例 4：管道傳輸 Shell 命令的輸出

這個 gist (https://gist.github.com/ifels/10392762)，它以一種很棒的方式結合了 io.Pipe 和 os.Exec。基本上，它做的就是 Jenkins 或 Travis CI 等持續集成 (CI) 服務中大多數任務運行器所做的事情：執行一些 shell 命令並在某個網站上顯示其輸出。

嘗試用以下代碼片段封裝其通用模式：

func command() {
 pr, pw := io.Pipe()
 defer pw.Close()

 // tell the command to write to our pipe
 cmd := exec.Command("cat", "test.txt")
 cmd.Stdout = pw

 go func() {
  defer pr.Close()
  // copy the data written to the PipeReader via the cmd to stdout
  if _, err := io.Copy(os.Stdout, pr); err != nil {
   log.Fatal(err)
  }
 }()

 // run the command, which writes all output to the PipeWriter
 // which then ends up in the PipeReader
 if err := cmd.Run(); err != nil {
  log.Fatal(err)
 }
}

首先，我們定義命令——在這個例子中，我們只是 cat 一個名爲 test.txt 的文件，它會將文件內容輸出到標準輸出 (stdout)。然後，我們將命令的 stdout 設置爲我們的 PipeWriter。

這樣，我們就將 Command 的輸出重定向到我們的管道，可以在另一個點通過 PipeReader 讀取它。在這個有點刻意的例子中，雖然當前示例直接輸出內容（與原始命令行爲一致），實際可擴展爲：將命令結果導出到存儲系統；實時推送至 Web 頁面（如 gist 所示）；對接需要 io.Writer 接口的任何處理模塊。

總結

io.Pipe 在 Go 中具有着實現高效、簡潔數據流處理的能力。它使開發者能夠構建處理海量數據的應用，同時保持低開銷和高靈活性。

請記住，io.Pipe 的強大之處不僅在於傳輸數據，更在於它能流暢連接應用程序的不同部分，實現字節數據的動態流轉。

References
https://www.reddit.com/r/golang/comments/1ejo7rx/issue_with_s3_iopipe_and_gzip_help_appreciated/
https://medium.com/xendit-engineering/streaming-an-avalanche-of-data-with-gos-io-package-d319226f645b
https://vishnubharathi.codes/blog/a-silly-mistake-that-i-made-with-io.teereader/
https://medium.com/@0xgotznit/mastering-io-pipe-in-go-ca8686150b5e
https://zupzup.org/io-pipe-go/

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

猜你喜歡