go的fasthttp学习～stackless的writer

代码

package stackless

import (
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/valyala/bytebufferpool"
)

// Writer is an interface stackless writer must conform to.
// Writer 是 Stackless writer 必须遵守的接口
// The interface contains common subset for Writers from compress/* packages.
type Writer interface {
	Write(p []byte) (int, error)
	Flush() error
	Close() error
	Reset(w io.Writer)
}

// NewWriterFunc must return new writer that will be wrapped into
// stackless writer.
type NewWriterFunc func(w io.Writer) Writer

// NewWriter creates a stackless writer around a writer returned
// from newWriter.
//
// The returned writer writes data to dstW.
// 返回的Writer将数据写入到dstW
// 使用大量栈空间的Writers被封装为stackless writer，这样可以支持更高的并发
// Writers that use a lot of stack space may be wrapped into stackless writer,
// thus saving stack space for high number of concurrently running goroutines.
//
// newWriter的意思是把内部结构体writer的私有变量 xw xWriter 通过外部的方法转换为一个stackless的成员
// 然后让内部结构体writer的zw指向它
// 后续都是通过w.zw来操作写入
func NewWriter(dstW io.Writer, newWriter NewWriterFunc) Writer {
	w := &writer{
		dstW: dstW,
	}
	w.zw = newWriter(&w.xw)
	return w
}

type writer struct {
	dstW io.Writer
	zw   Writer
	xw   xWriter

	err error
	n   int

	p  []byte
	op op
}

type op int

const (
	opWrite op = iota
	opFlush
	opClose
	opReset
)

func (w *writer) Write(p []byte) (int, error) {
	w.p = p
	err := w.do(opWrite)
	w.p = nil
	return w.n, err
}

func (w *writer) Flush() error {
	return w.do(opFlush)
}

func (w *writer) Close() error {
	return w.do(opClose)
}

func (w *writer) Reset(dstW io.Writer) {
	// 归还writer中的bytebufferpool.ByteBuffer
	w.xw.Reset()
	// 初始化writer中的zw指向的stackless的内容
	w.do(opReset) //nolint:errcheck
	w.dstW = dstW
}

func (w *writer) do(op op) error {
	w.op = op
	if !stacklessWriterFunc(w) {
		return errHighLoad
	}
	err := w.err
	if err != nil {
		return err
	}
	// 如果 w.xw 的 bb 中还有数据，则将其写入dstW
	if w.xw.bb != nil && len(w.xw.bb.B) > 0 {
		_, err = w.dstW.Write(w.xw.bb.B)
	}
	// 更新xw
	w.xw.Reset()

	return err
}

var errHighLoad = errors.New("cannot compress data due to high load")

var (
	stacklessWriterFuncOnce sync.Once
	stacklessWriterFuncFunc func(ctx any) bool
)

func stacklessWriterFunc(ctx any) bool {
	stacklessWriterFuncOnce.Do(func() {
		stacklessWriterFuncFunc = NewFunc(writerFunc)
	})
	// 第一次执行stacklessWriterFuncFunc时
	// 此时macProc个子协程已经就绪，等待writerFunc执行代码
	return stacklessWriterFuncFunc(ctx)
}

// 多协程等待执行，err存储到 w 中，外部会读取
func writerFunc(ctx any) {
	w := ctx.(*writer)
	switch w.op {
	case opWrite:
		w.n, w.err = w.zw.Write(w.p)
	case opFlush:
		w.err = w.zw.Flush()
	case opClose:
		w.err = w.zw.Close()
	case opReset:
		w.zw.Reset(&w.xw) // zw内部reset的时候，需要将w.xw的位置也重新赋值上
		w.err = nil
	default:
		panic(fmt.Sprintf("BUG: unexpected op: %d", w.op))
	}
}

type xWriter struct {
	bb *bytebufferpool.ByteBuffer
}

func (w *xWriter) Write(p []byte) (int, error) {
	if w.bb == nil {
		w.bb = bufferPool.Get()
	}
	return w.bb.Write(p)
}

func (w *xWriter) Reset() {
	if w.bb != nil {
		bufferPool.Put(w.bb)
		w.bb = nil
	}
}

var bufferPool bytebufferpool.Pool

单测

package stackless

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"fmt"
	"io"
	"testing"
	"time"
)

func TestCompressFlateSerial(t *testing.T) {
	t.Parallel()

	if err := testCompressFlate(); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestCompressFlateConcurrent(t *testing.T) {
	t.Parallel()

	if err := testConcurrent(testCompressFlate, 10); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func testCompressFlate() error {
	// 第一个参数是一个函数，将入参 w io.Writer 转换为 flate.Writer
	// 第二个参数是一个函数，将入参 r io.Reader 转换为 flate.Reader
	return testWriter(func(w io.Writer) Writer {
		// 这里是将 stackless的 xwriter 转换为 flate.Writer
		zw, err := flate.NewWriter(w, flate.DefaultCompression)
		if err != nil {
			panic(fmt.Sprintf("BUG: unexpected error: %v", err))
		}
		return zw
	}, func(r io.Reader) io.Reader {
		return flate.NewReader(r)
	})
}

func TestCompressGzipSerial(t *testing.T) {
	t.Parallel()

	if err := testCompressGzip(); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestCompressGzipConcurrent(t *testing.T) {
	t.Parallel()

	if err := testConcurrent(testCompressGzip, 10); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func testCompressGzip() error {
	return testWriter(func(w io.Writer) Writer {
		return gzip.NewWriter(w)
	}, func(r io.Reader) io.Reader {
		zr, err := gzip.NewReader(r)
		if err != nil {
			panic(fmt.Sprintf("BUG: cannot create gzip reader: %v", err))
		}
		return zr
	})
}

// 第一个参数是一个函数，将入参 w io.Writer 包裹为 flate.Writer
// 第二个参数是一个函数，将入参 r io.Reader 包裹为 flate.Reader
func testWriter(newWriter NewWriterFunc, newReader func(io.Reader) io.Reader) error {
	dstW := &bytes.Buffer{}
	// 生成一个stackless的Writer w
	// 并调用newWriter进行初始化：w.zw = newWriter(&w.xw) w.xw是一个xWriter
	// type xWriter struct {
	//	  bb *bytebufferpool.ByteBuffer
	// }
	//
	// 说白了就是将 w.xw 作为flate.Writer的底层writer来制作一个新的flate.Writer，并让 w.zw 指向这个 flate.Writer
	// 利用stackless的Writer去操作自己的w.zw，进而实现flate.Writer操作w.xw的目的
	w := NewWriter(dstW, newWriter)

	for i := 0; i < 5; i++ {
		// 这里明明已经将dstW赋值给stackless的writer了，又将其作为参数传入
		// 是为了通过dstW读取值，判断写入writer的内容是对的
		if err := testWriterReuse(w, dstW, newReader); err != nil {
			return fmt.Errorf("unexpected error when re-using writer on iteration %d: %w", i, err)
		}
		dstW = &bytes.Buffer{}
		w.Reset(dstW)
	}

	return nil
}

func testWriterReuse(w Writer, r io.Reader, newReader func(io.Reader) io.Reader) error {
	wantW := &bytes.Buffer{}
	// creates a writer that duplicates its writes to all the provided writers
	// similar to the Unix tee(1) command
	// 写入多个writer，如果一个失败，则全部返回失败
	mw := io.MultiWriter(w, wantW)
	for i := 0; i < 30; i++ {
		fmt.Fprintf(mw, "foobar %d\n", i)
		// 屏蔽掉也没问题
		if i%13 == 0 {
			// 实质是对 w.zw的Flush()
			// if err := w.Flush(); err != nil {
			// 	return fmt.Errorf("error on flush: %w", err)
			// }
		}
	}
	// 实质是对 w.zw的Close()
	// 没有Close和Flush的话，数据不会被刷入stackless的xWriter中
	w.Close()

	zr := newReader(r)
	data, err := io.ReadAll(zr)
	if err != nil {
		return fmt.Errorf("unexpected error: %w, data=%q", err, data)
	}

	// 从两个writer中读出来的结果相同
	wantData := wantW.Bytes()
	if !bytes.Equal(data, wantData) {
		return fmt.Errorf("unexpected data: %q. Expecting %q", data, wantData)
	}

	return nil
}

func testConcurrent(testFunc func() error, concurrency int) error {
	ch := make(chan error, concurrency)
	for i := 0; i < concurrency; i++ {
		go func() {
			ch <- testFunc()
		}()
	}
	for i := 0; i < concurrency; i++ {
		select {
		case err := <-ch:
			if err != nil {
				return fmt.Errorf("unexpected error on goroutine %d: %w", i, err)
			}
		case <-time.After(time.Second):
			return fmt.Errorf("timeout on goroutine %d", i)
		}
	}
	return nil
}

// 到底什么是flush
// Flush flushes any pending data to the underlying writer.
// 将任何挂起的数据写入底层的writer (underlying writer)
// It is useful mainly in compressed network protocols, to ensure that
// a remote reader has enough data to reconstruct a packet.
// 在压缩网络协议（compressed network protocols）中很有用
// 确保一个远端的reader有足够数据去重构一个包
//
// Flush does not return until the data has been written.
// Calling Flush when there is no pending data still causes the Writer
// to emit a sync marker of at least 4 bytes.
// 发出至少4个字节的同步标记（emit a sync marker）
// If the underlying writer returns an error, Flush returns that error.
//
// In the terminology of the zlib library, Flush is equivalent to Z_SYNC_FLUSH.

// flushes and closes the writer.
// close前先flushes writer
//
// func (w *Writer) Close() error {
//     return w.d.close()
// }

运用

var flateReaderPool sync.Pool

func acquireStacklessGzipWriter(w io.Writer, level int) stackless.Writer {
	nLevel := normalizeCompressLevel(level)
	p := stacklessGzipWriterPoolMap[nLevel]
	v := p.Get()
	if v == nil {
		// stackless.NewWriter 第二个参数是一个函数
		// 用于将一个具体的 *gzip.Writer 赋值给  stackless.Writer 接口
		//
		// newWriter的意思是把内部结构体writer的私有变量 xw xWriter 通过外部的方法转换为一个stackless的成员
		// 然后让内部结构体writer的zw指向它
		// 后续都是通过w.zw来操作写入
		return stackless.NewWriter(w, func(w io.Writer) stackless.Writer {
			// acquireRealGzipWriter：从资源池取一个/new一个gzipWriter并reset
			return acquireRealGzipWriter(w, level)
		})
	}
	sw := v.(stackless.Writer)
	sw.Reset(w)
	return sw
}

func releaseStacklessGzipWriter(sw stackless.Writer, level int) {
	sw.Close()
	nLevel := normalizeCompressLevel(level)
	p := stacklessGzipWriterPoolMap[nLevel]
	p.Put(sw)
}

// acquireRealGzipWriter
// 从资源池取一个/new一个gzipWriter并reset
func acquireRealGzipWriter(w io.Writer, level int) *gzip.Writer {
	nLevel := normalizeCompressLevel(level)
	p := realGzipWriterPoolMap[nLevel]
	v := p.Get()
	if v == nil {
		zw, err := gzip.NewWriterLevel(w, level)
		if err != nil {
			// gzip.NewWriterLevel only errors for invalid
			// compression levels. Clamp it to be min or max.
			if level < gzip.HuffmanOnly {
				level = gzip.HuffmanOnly
			} else {
				level = gzip.BestCompression
			}
			zw, _ = gzip.NewWriterLevel(w, level)
		}
		return zw
	}
	zw := v.(*gzip.Writer)
	zw.Reset(w)
	return zw
}

总结

如此精妙的设计，真的太厉害了
利用stackless的NewWriter，来实现接管一个真正的复杂的高层Writer
复杂的高层Writer利用stackless的writer提供的xWriter来做它的底层io，将编码后的数据存入里面
stackless的Flush和Close其实就是接管的复杂的高层Writer的这两，
最终stackless的worker将上述结果存入 dstWriter 中