Go 語言 bytes-Buffer 源碼詳解 2

前言

前面一篇文章 Go 語言 bytes.Buffer 源碼詳解之 1，我們學習了 bytes.Buffer 的結構和基礎方法，瞭解了緩衝區的運行機制，其中最重要的是要理解整個結構被分爲了三段：已讀數據、未讀數據、空閒部分，忘記的小夥伴再複習下哦。緩衝區的存在，就是爲讀寫服務的，那麼本篇文章我們就一起來學習下讀寫方法是如何利用緩衝區實現的吧！

Buffer 結構示意圖

源碼分析

Write()

Write 方法將字節切片 p 中的數據寫入到緩衝切片中，返回寫入的字節長度和產生的 error

由於 Write 調用了 grow 方法，如果底層的緩衝切片太大無法重新分配，會產生 ErrTooLarge 的 panic

func (b *Buffer) Write(p []byte) (n int, err error) {
 b.lastRead = opInvalid

 // 通過調用 tryGrowByReslice 和 grow 兩個方法，確保底層的緩衝切片的長度可以寫入len(p)個字節。
 m, ok := b.tryGrowByReslice(len(p))
 if !ok {
  m = b.grow(len(p))
 }

 // 走到這裏，說明 buf 從m位置開始已經有了 len(p)個空閒字節，調用copy方法，將p中的數據複製過去
 return copy(b.buf[m:], p), nil
}

WriteString()

WriteString 和 Write 方法類似，將傳入的字符串 s 寫入到底層的緩衝切片 buf 中，返回成功寫入的字節數 n 和產生的 error

func (b *Buffer) WriteString(s string) (n int, err error) {
 b.lastRead = opInvalid

 // 通過調用 tryGrowByReslice 和 grow 兩個方法，確保底層的緩衝切片的長度可以寫入 len(s) 個字節。
 m, ok := b.tryGrowByReslice(len(s))
 if !ok {
  m = b.grow(len(s))
 }
 // 到這裏說明buf 從m位置開始已經有了 len(s) 個空閒字節，調用copy方法，將 s 複製到底層緩衝切片中
 return copy(b.buf[m:], s), nil
}

WriteByte()

和 Write 方法類似，寫入單個字節，而非字節切片

func (b *Buffer) WriteByte(c byte) error {
 b.lastRead = opInvalid

 // 通過調用 tryGrowByReslice 和 grow 兩個方法，確保底層的緩衝切片的長度可以寫入 1 個字節
 m, ok := b.tryGrowByReslice(1)
 if !ok {
  m = b.grow(1)
 }
 // m 表示擴容後寫入的開始位置，直接賦值爲要寫入的字節
 b.buf[m] = c
 return nil
}

WriteRune()

和 Write 方法類似，區別是寫入 rune，而非字節切片

func (b *Buffer) WriteRune(r rune) (n int, err error) {

 // 如果r < utf8.RuneSelf，說明 r 就是一個字節，那麼直接調用 WriteByte 方法
 if r < utf8.RuneSelf {
  b.WriteByte(byte(r))
  return 1, nil
 }
 b.lastRead = opInvalid

 // 通過調用 tryGrowByReslice 和 grow 兩個方法，確保底層的緩衝切片的長度可以寫入 utf8.UTFMax 個字節
 m, ok := b.tryGrowByReslice(utf8.UTFMax)
 if !ok {
  m = b.grow(utf8.UTFMax)
 }
  
  // 此時的buf 長度，變成了 len(buf)+ utf8.UTFMax，utf8.UTFMax 是rune 可能的最大長度，但是當前 rune 的大小可能小於這個值
 // 調用 utf8.EncodeRune() 方法，將 rune r 寫入到 buf 中，返回寫入的字節數
 n = utf8.EncodeRune(b.buf[m:m+utf8.UTFMax], r)

 // 更新 buf 的長度，因爲 n<= utf8.UTFMax
 b.buf = b.buf[:m+n]
 return n, nil
}

ReadFrom()

ReadFrom 方法從 Reader r 讀取數據，寫入底層的緩衝切片 buf 中，返回寫入的字節數和產生的 error
在讀取數據並寫入緩衝切片過程中，如果緩衝切片容量不足，會調用 grow 方法增大緩衝切片大小
讀取寫入這個過程一直循環，直至產生 error，如果最終產生的是 EOF error，即 reader r 讀取數據到了文件結尾，方法最終返回的 error 爲 nil，因爲任務已經完成了

// 緩衝切片留出的最小空閒空間
// ReadFrom 方法會用到該參數，即從一個 Reader 寫入數據到底層緩衝切片 buf 時，buf 留出的最小空閒空間
const MinRead = 512


func (b *Buffer) ReadFrom(r io.Reader) (n int64, err error) {
 b.lastRead = opInvalid

 // for 循環，不斷讀取寫入數據，直至遇到Reader 讀取數據完畢產生 EOF error 或者 其他 error
 for {
    
  // 保證至少留出 MinRead 個空閒字節空間，並返回寫入開始位置 i
  i := b.grow(MinRead)

  // grow 方法將長度變爲了 i+MinRead，改回來
  b.buf = b.buf[:i]

  // i 位置開始，到 cap(buf) 結束的空間，即從 i位置 開始的底層數組所有空間都供 reader r 讀取數據寫入
  // Read 方法會返回讀取的字節數和產生的error，根據Read方法的定義，應該先處理m，再處理e
  m, e := r.Read(b.buf[i:cap(b.buf)])
  if m < 0 {
   panic(errNegativeRead)
  }

  // 如果 m 大於等於0，說明讀取並寫入數據到buf 中了，修改buf 的長度
  b.buf = b.buf[:i+m]
  // 已讀字節數n 加 m
  n += int64(m)

  // 如果Reader r 讀取過程中遇到了EOF error，說明讀取數據完畢了，返回 error=nil
  if e == io.EOF {
   return n, nil // e is EOF, so return nil explicitly
  }

  // 遇到了其他error，返回error
  if e != nil {
   return n, e
  }
 }
}

上面介紹的是寫入緩衝區的相關操作，接下來我們來看讀取相關的操作。

WriteTo()

WriteTo 方法，讀取字節緩衝切片中的數據，交由 Writer w 去消費使用，最終返回 Writer w 消費的字節量和產生的 error

func (b *Buffer) WriteTo(w io.Writer) (n int64, err error) {
 b.lastRead = opInvalid

 // nBytes:未讀數據的長度
 if nBytes := b.Len(); nBytes > 0 {
  // 如果未讀數據的長度大於0，將從已讀計數 off 到len()部分的未讀數據，寫入到Writer w 中
  // Write 返回消費的字節數，以及產生的error
  m, e := w.Write(b.buf[b.off:])
  // 如果消費的長度，大於可用長度，不符合邏輯，panic
  if m > nBytes {
   panic("bytes.Buffer.WriteTo: invalid Write count")
  }
  // 被消費了 m 個字節，已讀計數相應增加 m
  b.off += m

  // 消費的字節量 n = int64(m)
  n = int64(m)

  // 如果產生了 error，返回
  if e != nil {
   return n, e
  }

  // 根據 io.Writer 接口對 Write 方法的定義，如果寫入的數量 m != nBytes，一定會返回error!=nil
  // 因此上一步的 e!=nil 一定成立，會直接返回，導致到不了這一步，這一步相當於做了個double check
  if m != nBytes {
   return n, io.ErrShortWrite
  }
 }

 // 到這一步，說明緩衝切片中的未讀數據被讀完了，直接調用Reset()方法重置
 b.Reset()
 return n, nil
}

Read()

Read 方法，讀取底層緩衝字節切片 buf 中的數據，寫入到字節切片 p 中

// 方法讀取的字節數，和產生的 error
func (b *Buffer) Read(p []byte) (n int, err error) {
 b.lastRead = opInvalid

  // 如果 buf 中無數據可讀，且len(p)=0，返回 error=nil，否則返回 error=EOF
 // 如果未讀數據部分爲空，沒有數據可讀
 if b.empty() {
  // 首先將字節緩衝切片重置
  b.Reset()
  // len(p)=0，返回的 error=nil
  if len(p) == 0 {
   return 0, nil
  }
  // 讀取的數據小於 len(p)，返回 EOF error
  return 0, io.EOF
 }

 // 存在未讀數據，調用 copy 方法，從 off 位置開始複製數據到 p 中，返回複製的字節數
 n = copy(p, b.buf[b.off:])

 // 更新已讀計數
 b.off += n

 // 讀取到數據了，此次是一次合法的讀操作，更新 lastRead 爲 opRead
 if n > 0 {
  b.lastRead = opRead
 }
 // 返回讀取的字節數 n，error=nil
 return n, nil
}

Next()

Next 方法返回未讀數據的前 n 個字節，如果未讀數據長度小於 n 個字節，那麼就返回所有的未讀數據。由於方法返回的數據是基於 buf 的切片，存在數據泄露的風險，且數據的有效期在下次調用 read 或 write 方法前，因爲調用 read、write 方法會修改底層數據。

func (b *Buffer) Next(n int) []byte {
 b.lastRead = opInvalid

 // m：未讀數據長度
 m := b.Len()

 // 如果需要的字節數 n，大於未讀數據長度m，那麼n=m
 if n > m {
  n = m
 }

 // 賦值 data 爲所需的n個字節
 data := b.buf[b.off : b.off+n]

 // 已讀計數增加 n
 b.off += n
 if n > 0 {
  b.lastRead = opRead
 }
 return data
}

ReadByte()

類似 Read 方法，ReadByte 讀取一個字節，返回讀取的字節和產生的 error

func (b *Buffer) ReadByte() (byte, error) {

 // 如果沒有未讀數據，重置，返回 EOF error
 if b.empty() {
  // Buffer is empty, reset to recover space.
  b.Reset()
  return 0, io.EOF
 }

 // 讀取一個字節，然後修改已讀計數
 c := b.buf[b.off]
 b.off++
 b.lastRead = opRead
 return c, nil
}

ReadRune()

類似 Read 方法，ReadRune 讀取一個 utf-8 編碼的 rune，返回 rune 的值、大小以及產生的 error

func (b *Buffer) ReadRune() (r rune, size int, err error) {
 // 如果沒有數據可讀，重置，返回 EOF error
 if b.empty() {
  b.Reset()
  return 0, 0, io.EOF
 }

 // c 代表開始讀取的第一個字節
 c := b.buf[b.off]
 // 如果 c < utf8.RuneSelf，表示 c 是一個單字節的 rune，直接返回這個 rune，已讀計數加一即可
 if c < utf8.RuneSelf {
  b.off++
  b.lastRead = opReadRune1
  return rune(c), 1, nil
 }

 // 從已讀計數位置開始，調用utf8.DecodeRune，方法返回從已讀計數位置開始的rune，以及對應的字節數
 r, n := utf8.DecodeRune(b.buf[b.off:])

 // 修改已讀計數
 b.off += n

 // 修改lastRead
 b.lastRead = readOp(n)
 return r, n, nil
}

UnreadRune()

回退一個 rune，只能在 ReadRune 後調用該方法纔有效，其他 read 方法之後後調用該方法非法，因爲其他相關的 read 方法記錄的 lastRead = opRead，而不是 opReadRune*

func (b *Buffer) UnreadRune() error {

 // lastRead <= opInvalid，表示上一次調用爲非ReadRune 方法，不能進行回退
 if b.lastRead <= opInvalid {
  return errors.New("bytes.Buffer: UnreadRune: previous operation was not a successful ReadRune")
 }
 // 回退
 if b.off >= int(b.lastRead) {
  b.off -= int(b.lastRead)
 }
 // 只能回退一次，不能再次回退
 b.lastRead = opInvalid
 return nil
}

UnreadByte()

回退一個字節，該方法的要求比 UnreadRune 方法要低，只要是 read 相關的方法都能回退一個字節

var errUnreadByte = errors.New("bytes.Buffer: UnreadByte: previous operation was not a successful read")

func (b *Buffer) UnreadByte() error {
 // 只有 lastRead == opInvalid 纔不能回退（ ReadRune 也可以回退 ）
 if b.lastRead == opInvalid {
  return errUnreadByte
 }

 // 只能回退一次
 b.lastRead = opInvalid

 // 已讀計數減一
 if b.off > 0 {
  b.off--
 }
 return nil
}

readSlice()

私有方法，readSlice 讀取未讀數據，直至找到 delim 這個字符停止，然後返回遍歷到的數據。返回的數據是基於底層緩衝切片的引用，存在數據泄露的風險。

func (b *Buffer) readSlice(delim byte) (line []byte, err error) {

 // 調用 IndexByte()，從 off 位置開始查找，找到第一個出現 delim 的索引，如果沒找到會返回 -1
 i := IndexByte(b.buf[b.off:], delim)

 // 因爲上一步索引從0開始，所以要再加 1
 end := b.off + i + 1

 // 沒有找到，需要返回所有未讀數據，因此 end 賦值爲 緩衝數組長度，err 爲 EOF
 if i < 0 {
  end = len(b.buf)
  err = io.EOF
 }

 // 返回遍歷過的數據，更新已讀計數
 line = b.buf[b.off:end]
 b.off = end

 // 此次操作也是 opRead
 b.lastRead = opRead

 // 返回數據和error
 return line, err
}

ReadBytes()

ReadBytes 遍歷未讀數據，直至遇到一個字節值爲 delim 的分隔符，然後返回遍歷過的數據（包括該分隔符）和產生的 error。

ReadBytes 直接調用的 readSlice，因此只有一種情況會返回 error!=nil，即在未讀數據中，遍歷完所有數據但沒有找到該分隔符時，此時會返回所有未讀數據和 EOF error。

func (b *Buffer) ReadBytes(delim byte) (line []byte, err error) {
 // 直接調用 readSlice() 方法，但是該方法返回的是底層切片的引用
 // readSlice() 方法已經修改了 已讀計數和 lastRead
 slice, err := b.readSlice(delim)

 // 由於readSlice() 返回的是引用，數據可能因爲其他方法調用被修改，因此拷貝一份數據
 line = append(line, slice...)
 return line, err
}

ReadString()

ReadString 和 ReadBytes 類似，區別是該方法返回的是字符串形式。

func (b *Buffer) ReadString(delim byte) (line string, err error) {
 // 直接調用 readSlice()方法，但是該方法返回的是底層切片的引用
 // readSlice() 方法已經修改了 已讀計數和 lastRead
 slice, err := b.readSlice(delim)
 // 轉爲字符串返回
 return string(slice), err
}

NewBuffer()

實例化一個 Buffer，使用傳入的字節切片 buf 作爲底層的緩衝字節切片，也可以傳入 nil
在初始化完成後，調用者不能再操作傳入的字節切片 buf 了，否則會影響數據正確性
傳入的切片數組作用：供 read 方法讀取切片中的已有數據，或者供 write 方法寫入數據，因此傳入的字節切片的容量儘量避免爲 0
bytes.Buffer 是開箱即用的，大多數情況下，直接 new(Buffer) 或者聲明一個變量就可以了，沒必要調用 NewBuffer() 方法

func NewBuffer(buf []byte) *Buffer { return &Buffer{buf: buf} }

NewBufferString()

NewBufferString 傳入一個字符串用於初始化 bytes.Buffer，因此底層的字節緩衝切片就有了初始值，就有了數據用於讀取。
bytes.Buffer 是開箱即用的，大多數情況下，直接 new(Buffer) 或者聲明一個變量就可以了，沒必要調用 NewBufferString() 方法

func NewBufferString(s string) *Buffer {
 return &Buffer{buf: []byte(s)}
}

使用示例

下面的示例，使用了 bytes.Buffer 作爲緩衝區，完成了一個文件複製的操作。

func main() {
 var buffer bytes.Buffer
 srcFile, _ := os.OpenFile("test.txt", os.O_RDWR, 0666)
 n, err := buffer.ReadFrom(srcFile)
 fmt.Println(n, err)                     // 303190 <nil>
 fmt.Println(buffer.Len(), buffer.Cap()) // 303190 523776

 targetFile, _ := os.OpenFile("target.txt", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
 n, err = buffer.WriteTo(targetFile)
 fmt.Println(n, err)                // 303190 <nil>
 fmt.Println(buffer.Len(), buffer.Cap()) // 0 523776
}

總結

本篇文章我們學習了 bytes.Buffer 中讀寫相關的源碼實現。針對寫操作，都會先確保有足夠的可用空間，然後再將數據複製到緩衝區的未讀數據部分；針對讀操作，就是將未讀部分的數據拷貝出去，然後更新已讀計數。

到這裏我們就把 bytes.Buffer 源碼給過完了，俗話說知己知彼，百戰不殆，瞭解了 bytes.Buffer 的原理後，相信你之後使用起來應該會更得心應手！

個人博客: https://lifelmy.github.io/

微信公衆號：漫漫 Coding 路

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

前言