go io 包的使用（TeeReader, MultiReader, MultiWriter, Pipe）

背景

今天在做文件的上传，在文件上传结束后，需对比上传文件和本地文件的哈希值是否一致。
对上传客户端来说，可以提前单独计算好文件的哈希值，然后在上传成功后，与服务端返回的哈希值进行对比。但显然，这并不是一个理想的方案，边上传边计算，这样是比较合理的。
我这里一开始使用的是 io 包的 Pipe() 方法来做的，之前一直都没怎么用过，比较冷门，顺带着就看了下源码实现，再带着就看了下 io 包的其他方法的使用，最后引申出了本篇。

TeeReader

func TeeReader(r Reader, w Writer) Reader {
	return &teeReader{r, w}
}

边读边写，函数返回的 reader 会在从接收参数 r 中读取内容的同时，将内容写入 w。

源码

type teeReader struct {
	r Reader
	w Writer
}

func (t *teeReader) Read(p []byte) (n int, err error) {
	n, err = t.r.Read(p)
	if n > 0 {
		if n, err := t.w.Write(p[:n]); err != nil {
			return n, err
		}
	}
	return
}

简单示例

可以简化代码，如：计算文件哈希 (这里使用不一定合理，仅打个比方)

func TeeGetFileMD5(path string) (string, error) {
	file, err := os.Open(path)
	if err != nil {
		return "", err
	}
	h := md5.New()
	tr := io.TeeReader(file, h)
	_, err = io.ReadAll(tr) // 优化
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

MultiReader

func MultiReader(readers ...Reader) Reader {
	r := make([]Reader, len(readers))
	copy(r, readers)
	return &multiReader{r}
}

顺序读取所有的 reader，直到出错或者所有的 reader 都读取完成后返回 EOF。

源码

type multiReader struct {
	readers []Reader
}

func (mr *multiReader) Read(p []byte) (n int, err error) {
	for len(mr.readers) > 0 {
		// Optimization to flatten nested multiReaders (Issue 13558).
		if len(mr.readers) == 1 { // 适配multiReader的嵌套
			if r, ok := mr.readers[0].(*multiReader); ok {
				mr.readers = r.readers
				continue
			}
		}
		n, err = mr.readers[0].Read(p)
		if err == EOF { // 读完一个，从数组剔除，换下一个
			// Use eofReader instead of nil to avoid nil panic
			// after performing flatten (Issue 18232).
			mr.readers[0] = eofReader{} // permit earlier GC
			mr.readers = mr.readers[1:]
		}
		if n > 0 || err != EOF {
			if err == EOF && len(mr.readers) > 0 { // 是EOF，读到最后一个再返回EOF
				// Don't return EOF yet. More readers remain.
				err = nil
			}
			return 
		}
	}
	return 0, EOF
}

简单示例

func tMultiReader() {
	r1 := bytes.NewReader([]byte("ABC"))
	r2 := bytes.NewReader([]byte("DEF"))
	reader := io.MultiReader(r1, r2)
	var buf = make([]byte, 1)
	for {
		n, err := reader.Read(buf)
		if err != nil {
			if err == io.EOF {
				return 
			}
			fmt.Println(err)
			return
		}
		fmt.Println(string(buf[:n])) // ABCDEF
	}
}

MultiWriter

func MultiWriter(writers ...Writer) Writer {
	allWriters := make([]Writer, 0, len(writers))
	for _, w := range writers {
		if mw, ok := w.(*multiWriter); ok {
			allWriters = append(allWriters, mw.writers...)
		} else {
			allWriters = append(allWriters, w)
		}
	}
	return &multiWriter{allWriters}
}

写一即多，函数返回的 writer 进行写操作时，会对所有的入参 writer 都进行写操作（copy）。当有多个输出点的时候，直接使用它会简化不少代码。

源码

type multiWriter struct {
	writers []Writer
}

func (t *multiWriter) Write(p []byte) (n int, err error) {
	for _, w := range t.writers {
		n, err = w.Write(p)
		if err != nil {
			return
		}
		if n != len(p) {
			err = ErrShortWrite
			return
		}
	}
	return len(p), nil
}

简单示例

func tMultiWriter() {
	var buf []byte
	w1 := bytes.NewBuffer(buf)
	w2 := bytes.NewBuffer(buf)
	writer := io.MultiWriter(w1, w2)
	_, err := writer.Write([]byte("123"))
	if err != nil {
		fmt.Println(err)
		return
	}

	w1Res, err := ioutil.ReadAll(w1)
	fmt.Println(string(w1Res), err) // 123 <nil>
	w2Res, err := ioutil.ReadAll(w2)
	fmt.Println(string(w2Res), err) // 123 <nil>
}

Pipe

func Pipe() (*PipeReader, *PipeWriter) {
	p := &pipe{
		wrCh: make(chan []byte),
		rdCh: make(chan int),
		done: make(chan struct{}),
	}
	return &PipeReader{p}, &PipeWriter{p}
}
type pipe struct {
	wrMu sync.Mutex // Serializes Write operations
	wrCh chan []byte
	rdCh chan int

	once sync.Once // Protects closing done
	done chan struct{}
	rerr onceError
	werr onceError
}

函数返回的 PipeReader 和 PipeWriter 均有 Close 和 CloseWithError 方法，用于停止读写（done）。

• wrCh 写入的数据
• rdCh 读了多少了
• once 只close一次done
• done 结束标志
• rerr 读错
• werr 写错

源码

func (p *pipe) Write(b []byte) (n int, err error) {
	select {
	case <-p.done:
		return 0, p.writeCloseError()
	default:
		p.wrMu.Lock()
		defer p.wrMu.Unlock()
	}

	for once := true; once || len(b) > 0; once = false {
		select {
		case p.wrCh <- b:
			nw := <-p.rdCh
			b = b[nw:]
			n += nw
		case <-p.done:
			return n, p.writeCloseError()
		}
	}
	return n, nil
}

func (p *pipe) Read(b []byte) (n int, err error) {
	select {
	case <-p.done:
		return 0, p.readCloseError()
	default:
	}

	select {
	case bw := <-p.wrCh:
		nr := copy(b, bw)
		p.rdCh <- nr
		return nr, nil
	case <-p.done:
		return 0, p.readCloseError()
	}
}

Read 和 Write 方法的开头都有如下一段代码。由于 channel 的 case 选择是随机的，需要确保没有结束再进行 Read 或 Write 操作。

	select {
	case <-p.done:
		return 0, p.writeCloseError()
	default:
		p.wrMu.Lock()
		defer p.wrMu.Unlock()
	}

简单示例

func tPipe() {
	r, w := io.Pipe()
	go func  () {
		for i := 0; i < 3 ; i++ {
			fmt.Println("write now!")
			n, err := w.Write([]byte("hello"))
			if err != nil {
				fmt.Println("write err:", err.Error())
			} else {
				fmt.Println("write end n:", n)
			}
		}
		w.Close()
	}()
	//time.Sleep(time.Second)
	b := make([]byte, 100)
	for {
		n, err := r.Read(b)
		if err != nil {
			if err != io.EOF {
				fmt.Println("read err:", err.Error())
			}
			break
		} else {
			fmt.Println("read:", string(b[:n]))
		}
	}
}
// write now!
// read: hello
// write end n: 5
// write now!
// read: hello
// write end n: 5
// write now!
// read: hello
// write end n: 5

TeeReader 和 MultiWriter 的结合使用

例子：复制文件，并计算文件的哈希值

func copyFileWithHash() {
	f, dstF, hashW, err := getTestRW()
	if err != nil {
		fmt.Println(err)
		return
	}

	now := time.Now()
	defer func() {
		fmt.Println("耗时：", time.Now().Sub(now))
	}()

	multiW := io.MultiWriter(dstF, hashW)

	teeR := io.TeeReader(f, multiW)


	buf := make([]byte, 512)
	for {
		_, err := teeR.Read(buf)
		if err == io.EOF {
			break
		}
		utils.CheckErr(err)
	}

	fmt.Printf("文件sha256：%x\n",  hashW.Sum(nil))
	// 文件大小： 1840640
	// 文件sha256：b61ec80071fc414c44ff1a05f323679f9fc3e7caa2a68363019663fc16677568
	// 耗时： 15.881ms
}

func getTestRW() (f, dstF *os.File, shaW hash.Hash, err error) {
	f, err = os.Open(`E:\test\html报告-1628662433.tar`)
	if err != nil {
		return
	}
	fInfo, err := f.Stat()
	if err != nil {
		return
	}
	fmt.Println("文件大小：", fInfo.Size())

	dstF, err = os.Create(`E:\test\1.tar`)
	if err != nil {
		return
	}

	shaW = sha256.New()
	return
}

总结

io 包的一些方法还是挺好用的，实现都并不是很复杂，感兴趣的可以看下各个方法的具体实现。本篇抛砖引玉，只是做了简单的介绍，实际开发过程中，不使用这些方法也是能够完全达成目的的，可能就是稍微繁琐点。

• TeeReader 的一心二用（边读编写）
• MultiReader 的先来后到（顺序读取）
• MultiWriter 的同甘共苦（写一即多）
• Pipe 的绝不先做（写->读->写->读->…）