Golang实现io.Reader的重复读取

代码一天不写我浑森蓝廋

已于 2022-10-21 03:16:25 修改

1179

分类专栏： Golang 文章标签： golang 开发语言后端

于 2022-10-19 14:56:42 首次发布

本文链接：https://blog.csdn.net/zhang197093/article/details/127407838

版权

Golang 专栏收录该内容

14 篇文章 0 订阅

订阅专栏

io.Reader 如果是实现了io.Seeker接口，想要重复读取则比较简单，直接调用Seek()方法把offset移到起始处即可，比如本地文件的IO，如果不支持Seek想要重复读取，则相对麻烦，这里我们可以借助一个buffer，从io.Reader读取数据的时候，同时把读出来的数据写到这个buffer中缓存起来，当要重新读取时，先从buffer里读，buffer读完再接着从原始的io.Reader里读，代码如下：

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"time"
)

type MultiReadable struct {
	originReader io.Reader
	reader       io.Reader
	cache        *bytes.Buffer
}

func NewMultiReadable(reader io.Reader) *MultiReadable {
	return &MultiReadable{
		originReader: reader,
		reader:       reader,
	}
}

func (mr *MultiReadable) Read(p []byte) (int, error) {
	n, err := mr.reader.Read(p)
	// 如果 reader 不支持Seek，则把读取出来的内容同时写入到一个buffer中
	if _, ok := mr.reader.(io.Seeker); !ok && n > 0 {
		if mr.cache == nil {
			mr.cache = &bytes.Buffer{}
		}
		mr.cache.Write(p[:n])
	}
	return n, err
}

func (mr *MultiReadable) Reset() error {
	// 如果reader支持Seek，直接使用Seek即可
	if seeker, ok := mr.reader.(io.Seeker); ok {
		_, err := seeker.Seek(0, io.SeekStart)
		return err
	}
	if mr.cache != nil && mr.cache.Len() > 0 {
		mr.reader = io.MultiReader(mr.cache, mr.reader)
		mr.cache = nil
	}
	return nil
}

func (mr *MultiReadable) Close() error {
	if closer, ok := mr.originReader.(io.Closer); ok {
		return closer.Close()
	}
	return nil
}

测试代码：

// getLongestValidUtf8Prefix get the longest valid utf8 string bytes from p
func getLongestValidUtf8Prefix(p []byte) []byte {
	lastValid := 0
	for lastValid < len(p) {
		r, size := utf8.DecodeRune(p[lastValid:])
		if r == utf8.RuneError {
			return p[:lastValid]
		}
		lastValid += size
	}
	return p[:lastValid]
}

func main() {

	resp, err := http.Get("http://www.baidu.com/")
	if err != nil {
		log.Fatal(err)
	}

	mr := NewMultiReadable(resp.Body)
	defer func() {
		if err := mr.Close(); err != nil {
			log.Println(err)
		}
	}()

	for i := 0; i < 10; i++ {
		if err := mr.Reset(); err != nil {
			log.Fatal(err)
		}

		buf, err := io.ReadAll(io.LimitReader(mr, int64(i)*60+100))
		if err != nil {
			log.Fatal(err)
		}
		
		fmt.Println(string(getLongestValidUtf8Prefix(buf)))
		fmt.Println("------------------------------------")
	}
}

输出：

<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charse
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name="theme-color" content="#ffffff"><meta name="description" co
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name="theme-color" content="#ffffff"><meta name="description" content="全球领先的中文搜索引擎、致力于让网
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name="theme-color" content="#ffffff"><meta name="description" content="全球领先的中文搜索引擎、致力于让网民更便捷地获取信息，找到所求。百度超过千
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name="theme-color" content="#ffffff"><meta name="description" content="全球领先的中文搜索引擎、致力于让网民更便捷地获取信息，找到所求。百度超过千亿的中文网页数据库，可以瞬间找到相关的搜
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name="theme-color" content="#ffffff"><meta name="description" content="全球领先的中文搜索引擎、致力于让网民更便捷地获取信息，找到所求。百度超过千亿的中文网页数据库，可以瞬间找到相关的搜索结果。"><link rel="shortcut icon" href="/favicon.ico" ty
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name="theme-color" content="#ffffff"><meta name="description" content="全球领先的中文搜索引擎、致力于让网民更便捷地获取信息，找到所求。百度超过千亿的中文网页数据库，可以瞬间找到相关的搜索结果。"><link rel="shortcut icon" href="/favicon.ico" type="image/x-icon" /><link rel="search" type="application/ope
------------------------------------
<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="Content-Type" content="text/html;charset=utf-8"><meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1"><meta content="always" name="referrer"><meta name="theme-color" content="#ffffff"><meta name="description" content="全球领先的中文搜索引擎、致力于让网民更便捷地获取信息，找到所求。百度超过千亿的中文网页数据库，可以瞬间找到相关的搜索结果。"><link rel="shortcut icon" href="/favicon.ico" type="image/x-icon" /><link rel="search" type="application/opensearchdescription+xml" href="/content-search.xml" title="
------------------------------------