common/utils/io2/stats.go

package io2

import "io"

type counter struct {
	inner io.Reader
	count int64
}

func (c *counter) Read(buf []byte) (n int, err error) {
	n, err = c.inner.Read(buf)
	c.count += int64(n)
	return
}

func (c *counter) Count() int64 {
	return c.count
}

func Counter(inner io.Reader) *counter {
	return &counter{inner: inner, count: 0}
}

type counterCloser struct {
	inner    io.ReadCloser
	count    int64
	callback func(cnt int64, err error)
}

func (c *counterCloser) Read(buf []byte) (n int, err error) {
	n, err = c.inner.Read(buf)
	c.count += int64(n)
	if err != nil && c.callback != nil {
		c.callback(c.count, err)
		c.callback = nil
	}
	return
}

func (c *counterCloser) Close() error {
	// 读取方主动Close，视为正常结束
	err := c.inner.Close()
	if c.callback != nil {
		c.callback(c.count, nil)
	}
	return err
}

// 统计一个io.ReadCloser的读取字节数，在读取结束后调用callback函数。
// 仅在读取方主动调用Close时，callback的err参数才会为nil。
func CounterCloser(inner io.ReadCloser, callback func(cnt int64, err error)) io.ReadCloser {
	return &counterCloser{inner: inner, count: 0, callback: callback}
}