Go语言之 sync.pool 源码分析

package main

import (
	"fmt"
	"sync"
)

type Gopher struct {
	Name   string
	Remark [1024]byte
}

func (s *Gopher) Reset() {
	s.Name = ""
	s.Remark = [1024]byte{}
}

var gopherPool = sync.Pool{
	New: func() interface{} {
		return new(Gopher)
	},
}

func main() {
	g := gopherPool.Get().(*Gopher)
	fmt.Println("首次从 pool 里获取：", g.Name)

	g.Name = "first"
	fmt.Printf("设置 p.Name = %s\n", g.Name)
	gopherPool.Put(g)

	fmt.Println("Pool 里已有一个对象：&{first}，调用 Get: ", gopherPool.Get().(*Gopher).Name)
	fmt.Println("Pool 没有对象了，调用 Get: ", gopherPool.Get().(*Gopher).Name)
}

首次从 pool 里获取： 
设置 p.Name = first
Pool 里已有一个对象：&{first}，调用 Get:  first
Pool 没有对象了，调用 Get:  

var defaultGopher, _ = json.Marshal(Gopher{Name: "17x"})

func BenchmarkUnmarshal(b *testing.B) {
	var g *Gopher
	for n := 0; n < b.N; n++ {
		g = new(Gopher)
		json.Unmarshal(defaultGopher, g)
	}
}

func BenchmarkUnmarshalWithPool(b *testing.B) {
	var g *Gopher
	for n := 0; n < b.N; n++ {
		g = gopherPool.Get().(*Gopher)
		json.Unmarshal(defaultGopher, g)
		g.Reset() // 重置后在放进去
		gopherPool.Put(g)
	}
}

BenchmarkUnmarshal-6                9518            124806 ns/op            1280 B/op          6 allocs/op
BenchmarkUnmarshalWithPool-6       10000            124350 ns/op             128 B/op          5 allocs/op

func Printf(format string, a ...interface{}) (n int, err error) {
	return Fprintf(os.Stdout, format, a...)
}

func Fprintf(w io.Writer, format string, a ...interface{}) (n int, err error) {
	p := newPrinter()
	p.doPrintf(format, a)
	n, err = w.Write(p.buf)
	p.free()
	return
}

var ppFree = sync.Pool{
	New: func() interface{} { return new(pp) },
}

func newPrinter() *pp {
	p := ppFree.Get().(*pp)
	p.panicking = false
	p.erroring = false
	p.wrapErrs = false
	p.fmt.init(&p.buf)
	return p
}

func (p *pp) free() {
	if cap(p.buf) > 64<<10 {
		return
	}

	p.buf = p.buf[:0]
	p.arg = nil
	p.value = reflect.Value{}
	p.wrappedErr = nil
	ppFree.Put(p)
}

func New() *Engine {
	engine := &Engine{
		
	}
	engine.pool.New = func() interface{} {
		return engine.allocateContext()
	}
	return engine
}

func (engine *Engine) allocateContext() *Context {
	return &Context{engine: engine}
}

func (engine *Engine) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	c := engine.pool.Get().(*Context)
	c.writermem.reset(w)
	c.Request = req
	c.reset()

	engine.handleHTTPRequest(c)

	engine.pool.Put(c)
}

type Pool struct {
	noCopy noCopy

	local     unsafe.Pointer // local fixed-size per-P pool, actual type is [P]poolLocal
	localSize uintptr        // size of the local array

	victim     unsafe.Pointer // local from previous cycle
	victimSize uintptr        // size of victims array

	New func() interface{}
}

type poolLocal struct {
	poolLocalInternal
    // 将 poolLocal 补齐至128字节(即两个cache line)的倍数，防止 false sharing,
    // 伪共享，仅占位用，防止在 cache line 上分配多个 poolLocalInternal
	pad [128 - unsafe.Sizeof(poolLocalInternal{})%128]byte
}
type poolLocalInternal struct {
	private interface{} // Can be used only by the respective P.
	shared  poolChain   // Local P can pushHead/popHead; any P can popTail.
}

type poolChain struct {

   head *poolChainElt

   tail *poolChainElt
}

type poolChainElt struct {
	poolDequeue

	next, prev *poolChainElt
}

type poolDequeue struct {

	headTail uint64

	vals []eface
}

func (p *Pool) Get() interface{} {
	// ...
	l, pid := p.pin()
	x := l.private
	l.private = nil
	if x == nil {
		x, _ = l.shared.popHead()
		if x == nil {
			x = p.getSlow(pid)
		}
	}
	runtime_procUnpin()
	// ...
	if x == nil && p.New != nil {
		x = p.New()
	}
	return x
}

func (p *Pool) pin() (*poolLocal, int) {
   // pin 具体逻辑由 runtime_procPin 实现
   pid :=  runtime_procPin()
   // 原子操作取出 p.localSize 和 p.local
   s := runtime_LoadAcquintptr(&p.localSize) // load-acquire
   l := p.local                              // load-consume
   // 因为是把 pid 做下标从 pool.local 中取得 p 对应的 local 的，
   // 所以如果 pid 小于 pool.local size 的时候才有可能取到对应的 local
   if uintptr(pid) < s {
      return indexLocal(l, pid), pid
   }
   // 正常情况下会一直满足该条件，
   // 只有刚开始  pool.local 还没创建或者动态调整了 P 的数量这两种情况
   // 会进入到下面的逻辑 去创建 pool.local
   return p.pinSlow()
}

func (p *Pool) pinSlow() (*poolLocal, int) {
    // 这里先取消绑定，然后加锁，最后有绑定上
	runtime_procUnpin()
	allPoolsMu.Lock()
	defer allPoolsMu.Unlock()
	pid := runtime_procPin()
	// doubleCheck 因为在执行上述命令过程中 pinSlow 可能已经被其他的线程调用，因此这时候需要再次对 pid 进行检查
	s := p.localSize
	l := p.local
	if uintptr(pid) < s {
		return indexLocal(l, pid), pid
	}
	if p.local == nil {
		allPools = append(allPools, p)
	}
	// 根据当前 P 的数量创建 pool.local并更新pool.localSize
	size := runtime.GOMAXPROCS(0)
	local := make([]poolLocal, size)
	atomic.StorePointer(&p.local, unsafe.Pointer(&local[0])) // store-release
	runtime_StoreReluintptr(&p.localSize, uintptr(size))     // store-release
    // 最后根据 pid 返回当前 P 对应的 local
	return &local[pid], pid
}

x := l.private
	l.private = nil
	if x == nil {
		x, _ = l.shared.popHead()
		if x == nil {
			x = p.getSlow(pid)
		}
	}

func (c *poolChain) popHead() (interface{}, bool) {
	d := c.head
	for d != nil {
		if val, ok := d.popHead(); ok {
			return val, ok
		}
		d = loadPoolChainElt(&d.prev)
	}
	return nil, false
}

func (d *poolDequeue) popHead() (interface{}, bool) {
	var slot *eface
	for {
		ptrs := atomic.LoadUint64(&d.headTail)
		head, tail := d.unpack(ptrs)
        // 收尾相连则说明队列是空的
		if tail == head {
			return nil, false
		}

		// head 位置是队头的前一个位置，所以此处要先退一位。
        // 在读出 slot 的 value 之前就把 head 值减 1，取消对这个 slot 的控制
		head--
		ptrs2 := d.pack(head, tail)
        // 通过 CAS 操作更新头部的位置 这样当前头部第一个元素就算是被移除了
		if atomic.CompareAndSwapUint64(&d.headTail, ptrs, ptrs2) {
			slot = &d.vals[head&uint32(len(d.vals)-1)]
			break
		}
	}
    // 类型转换与 nil 判断
	val := *(*interface{})(unsafe.Pointer(slot))
	if val == dequeueNil(nil) {
		val = nil
	}
    // 然后把这个 slot 置空，因为现在这个 slot 已经是队列的 head 了，置空便于前一个 head 被回收。
	*slot = eface{}
	return val, true
}

type dequeueNil *struct{}

*slot = eface{}

func (p *Pool) getSlow(pid int) interface{} {
	
	size := runtime_LoadAcquintptr(&p.localSize) // load-acquire
	locals := p.local                            // load-consume
    // 尝试从其他 p 中窃取一个对象
	for i := 0; i < int(size); i++ {
        // (pid+i+1)%int(size) 保证每次都可以从 当前pid+1 这个位置开始尝试窃取。
		l := indexLocal(locals, (pid+i+1)%int(size))
        // 如果能取到就直接返回
		if x, _ := l.shared.popTail(); x != nil {
			return x
		}
	}

    // 尝试从其他 P 的 poolLocal 窃取失败后，再尝试从victim cache中取对象
    // 这样可以使 victim 中的对象更容易被回收。
	size = atomic.LoadUintptr(&p.victimSize)
	if uintptr(pid) >= size {
		return nil
	}
	locals = p.victim
	l := indexLocal(locals, pid)
	if x := l.private; x != nil {
		l.private = nil
		return x
	}
	for i := 0; i < int(size); i++ {
		l := indexLocal(locals, (pid+i)%int(size))
		if x, _ := l.shared.popTail(); x != nil {
			return x
		}
	}

    // 清空 victim cache。下次就不用再从这里找了
	atomic.StoreUintptr(&p.victimSize, 0)

	return nil
}

func (c *poolChain) popTail() (interface{}, bool) {
    // tail 指针为空直接返回 这里的 tail 是一个双端队列
	d := loadPoolChainElt(&c.tail)
	if d == nil {
		return nil, false
	}

	for {
		// It's important that we load the next pointer
		// *before* popping the tail. In general, d may be
		// transiently empty, but if next is non-nil before
		// the pop and the pop fails, then d is permanently
		// empty, which is the only condition under which it's
		// safe to drop d from the chain.
        // 在 for 循环的一开始，就把 d.next 加载到了 d2。因为 d 可能会短暂为空，但如果 d2 在 pop 或者 pop fails 之前就不为空的话，说明 d 就会永久为空了。在这种情况下，可以安全地将 d 这个结点“甩掉”。
		d2 := loadPoolChainElt(&d.next)

		if val, ok := d.popTail(); ok {
			return val, ok
		}

		if d2 == nil {
			// This is the only dequeue. It's empty right
			// now, but could be pushed to in the future.
			return nil, false
		}

		// 同样的通过 CAS 来更新 tail 的值为 d2
		if atomic.CompareAndSwapPointer((*unsafe.Pointer)(unsafe.Pointer(&c.tail)), unsafe.Pointer(d), unsafe.Pointer(d2)) {
			storePoolChainElt(&d2.prev, nil)
		}
        // 最后将d2赋值给d便于进行下一轮循环。
		d = d2
	}
}

func (d *poolDequeue) popTail() (interface{}, bool) {
	var slot *eface
	for {
		ptrs := atomic.LoadUint64(&d.headTail)
		head, tail := d.unpack(ptrs)
        // 同样的 head=tail 说明队列为空
		if tail == head {
			return nil, false
		}

		// 同样是 CAS 更新 headTail 以移除 tail 元素
		ptrs2 := d.pack(head, tail+1)
		if atomic.CompareAndSwapUint64(&d.headTail, ptrs, ptrs2) {
			// Success.
			slot = &d.vals[tail&uint32(len(d.vals)-1)]
			break
		}
	}

	// 类型转换和 nil 判断
	val := *(*interface{})(unsafe.Pointer(slot))
	if val == dequeueNil(nil) {
		val = nil
	}

    // 最后也是将这个 slot 置空
    // 先清空 val 再清空 typ 操作顺序和 pushHead 正好相反
	slot.val = nil
	atomic.StorePointer(&slot.typ, nil)


	return val, true
}

func (p *Pool) Put(x interface{}) {
	if x == nil {
		return
	}
	// ...
	l, _ := p.pin()
	if l.private == nil {
		l.private = x
		x = nil
	}
	if x != nil {
		l.shared.pushHead(x)
	}
	runtime_procUnpin()
	// ...
}

func (c *poolChain) pushHead(val interface{}) {
	d := c.head
	if d == nil {
		// 第一次写入，队列为空则进行初始化 默认长度为8
		const initSize = 8 // Must be a power of 2
		d = new(poolChainElt)
		d.vals = make([]eface, initSize)
		c.head = d
		storePoolChainElt(&c.tail, d)
	}
	// 存储元素 存储成功直接返回
	if d.pushHead(val) {
		return
	}

	// The current dequeue is full. Allocate a new one of twice
	// the size.
    // 存储失败说明队列满了 进行扩容
	newSize := len(d.vals) * 2
	if newSize >= dequeueLimit { // 限制一下，不能无限扩容
		newSize = dequeueLimit
	}
	// 扩容逻辑也比较简单，就是首尾相连，构成链表
	d2 := &poolChainElt{prev: d}
	d2.vals = make([]eface, newSize)
	c.head = d2
	storePoolChainElt(&d.next, d2)
	d2.pushHead(val)
}

const dequeueBits = 32

const dequeueLimit = (1 << dequeueBits) / 4

func (d *poolDequeue) pushHead(val interface{}) bool {
	ptrs := atomic.LoadUint64(&d.headTail)
	head, tail := d.unpack(ptrs)
    // 首先判断队列是否已满： 也就是将尾部指针加上 d.vals 的长度，再取低 31 位，看它是否和 head 相等
	if (tail+uint32(len(d.vals)))&(1<<dequeueBits-1) == head {
		return false
	}
	slot := &d.vals[head&uint32(len(d.vals)-1)]

	// Check if the head slot has been released by popTail.
	typ := atomic.LoadPointer(&slot.typ)
	if typ != nil {
		// Another goroutine is still cleaning up the tail, so
		// the queue is actually still full.
		return false
	}

	// slot占位，将val存入vals中
	if val == nil {
		val = dequeueNil(nil)
	}
	*(*interface{})(unsafe.Pointer(slot)) = val

	// head 增加 1
	atomic.AddUint64(&d.headTail, 1<<dequeueBits)
	return true
}

	if (tail+uint32(len(d.vals)))&(1<<dequeueBits-1) == head {
		return false
	}

	slot := &d.vals[head&uint32(len(d.vals)-1)]

	// Check if the head slot has been released by popTail.
	typ := atomic.LoadPointer(&slot.typ)
	if typ != nil {
		// Another goroutine is still cleaning up the tail, so
		// the queue is actually still full.
		return false
	}

    // 最后也是将这个 slot 置空
    // 先清空 val 再清空 typ 操作顺序和 pushHead 正好相反
	slot.val = nil
	atomic.StorePointer(&slot.typ, nil)

// slot占位，将val存入vals中
*(*interface{})(unsafe.Pointer(slot)) = val

func (d *poolDequeue) pack(head, tail uint32) uint64 {
	const mask = 1<<dequeueBits - 1
	return (uint64(head) << dequeueBits) |
		uint64(tail&mask)
}

func (d *poolDequeue) unpack(ptrs uint64) (head, tail uint32) {
	const mask = 1<<dequeueBits - 1
	head = uint32((ptrs >> dequeueBits) & mask)
	tail = uint32(ptrs & mask)
	return
}

// sync/pool.go
func init() {
	runtime_registerPoolCleanup(poolCleanup)
}

// src/runtime/mgc.go
 
 
// Hooks for other packages
 
 
var poolcleanup func()
 
 
// 利用编译器标志将 sync 包中的清理注册到运行时
//go:linkname sync_runtime_registerPoolCleanup sync.runtime_registerPoolCleanup
func sync_runtime_registerPoolCleanup(f func()) {
  poolcleanup = f
}

func poolCleanup() {
	for _, p := range oldPools {
		p.victim = nil
		p.victimSize = 0
	}


	for _, p := range allPools {
		p.victim = p.local
		p.victimSize = p.localSize
		p.local = nil
		p.localSize = 0
	}

	oldPools, allPools = allPools, nil
}