mem: Move reference-counted buffer from zodb[1] here

Freelist of buffers is frequently needed to avoid GC/zeroing overhead in
e.g. networked servers. Making buffer reference-counted is then required
if buffer is shared between several users so that it is clear when it
can go back to its pool. In ZODB this situation was arising on returning
the same buffer from cached loads if several loads for same data are
issued in close to parallel.

[1] https://lab.nexedi.com/kirr/neo/tree/75a71514/go/zodb

mem/buffer.go

+// Copyright (C) 2017  Nexedi SA and Contributors.
+//                     Kirill Smelkov <kirr@nexedi.com>
+//
+// This program is free software: you can Use, Study, Modify and Redistribute
+// it under the terms of the GNU General Public License version 3, or (at your
+// option) any later version, as published by the Free Software Foundation.
+//
+// You can also Link and Combine this program with other software covered by
+// the terms of any of the Free Software licenses or any of the Open Source
+// Initiative approved licenses and Convey the resulting work. Corresponding
+// source of such a combination shall include the source code for all other
+// software used.
+//
+// This program is distributed WITHOUT ANY WARRANTY; without even the implied
+// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+//
+// See COPYING file for full licensing terms.
+// See https://www.nexedi.com/licensing for rationale and options.
+
+package mem
+// data buffers management
+
+import (
+	"sync"
+	"sync/atomic"
+
+	"lab.nexedi.com/kirr/go123/xmath"
+)
+
+// Buf is reference-counted memory buffer.
+//
+// To lower pressure on Go garbage-collector allocate buffers with BufAlloc and
+// free them with Buf.Release.
+//
+// Custom allocation functions affect only performance, not correctness -
+// everything should work if data buffer is allocated and/or free'ed
+// via regular Go/GC-way.
+type Buf struct {
+	Data []byte
+
+	// reference counter.
+	//
+	// NOTE to allow both Bufs created via BufAlloc and via std new, Buf is
+	// created with refcnt=0. The real number of references to Buf is thus .refcnt+1
+	refcnt int32
+}
+
+const order0 = 4 // buf sizes start from 2^4 (=16)
+
+var bufPoolv = [19]sync.Pool{} // buf size stop at 2^(4+19-1) (=4M)
+
+
+func init() {
+	for i := 0; i < len(bufPoolv); i++ {
+		i := i
+		bufPoolv[i].New = func() interface{} {
+			// NOTE *Buf, not just buf, to avoid allocation when
+			// making interface{} from it (interface{} wants to always point to heap)
+			return &Buf{Data: make([]byte, 1<<(order0+uint(i)))}
+		}
+	}
+}
+
+// BufAlloc allocates buffer of requested size from freelist.
+//
+// buffer memory is not initialized.
+func BufAlloc(size int) *Buf {
+	return BufAlloc64(int64(size))
+}
+
+// BufAlloc64 is same as BufAlloc but accepts int64 for size.
+func BufAlloc64(size int64) *Buf {
+	if size < 0 {
+		panic("invalid size")
+	}
+
+	// order = min i: 2^i >= size
+	order := xmath.CeilLog2(uint64(size))
+
+	order -= order0
+	if order < 0 {
+		order = 0
+	}
+
+	// if too big - allocate straightly from heap
+	if order >= len(bufPoolv) {
+		return &Buf{Data: make([]byte, size)}
+	}
+
+	buf := bufPoolv[order].Get().(*Buf)
+	buf.Data = buf.Data[:size] // leaving cap as is = 2^i
+	buf.refcnt = 0
+	return buf
+}
+
+// Release marks buf as no longer used by caller.
+//
+// It decrements buf reference-counter and if it reaches zero returns buf to
+// freelist.
+//
+// The caller must not use buf after call to Release.
+func (buf *Buf) Release() {
+	rc := atomic.AddInt32(&buf.refcnt, -1)
+	if rc < 0 - 1 {
+		panic("Buf.Release: refcnt < 0")
+	}
+	if rc > 0 - 1 {
+		return
+	}
+
+	// order = max i: 2^i <= cap
+	order := xmath.FloorLog2(uint64(cap(buf.Data)))
+
+	order -= order0
+	if order < 0 {
+		return // too small
+	}
+
+	if order >= len(bufPoolv) {
+		return // too big
+	}
+
+	bufPoolv[order].Put(buf)
+}
+
+// Incref increments buf's reference counter by 1.
+//
+// buf must already have reference-counter > 0 before Incref call.
+func (buf *Buf) Incref() {
+	rc := atomic.AddInt32(&buf.refcnt, +1)
+	if rc <= 1 - 1 {
+		panic("Buf.Incref: refcnt was < 1")
+	}
+}
+
+// XRelease releases buf it is != nil.
+func (buf *Buf) XRelease() {
+	if buf != nil {
+		buf.Release()
+	}
+}
+
+// XIncref increments buf's reference counter by 1 if buf != nil.
+func (buf *Buf) XIncref() {
+	if buf != nil {
+		buf.Incref()
+	}
+}
+
+
+// Len returns buf's len.
+//
+// it works even if buf=nil similarly to len() on nil []byte slice.
+func (buf *Buf) Len() int {
+	if buf != nil {
+		return len(buf.Data)
+	}
+	return 0
+}
+
+// Cap returns buf's cap.
+//
+// it works even if buf=nil similarly to len() on nil []byte slice.
+func (buf *Buf) Cap() int {
+	if buf != nil {
+		return cap(buf.Data)
+	}
+	return 0
+}
+
+// XData return's buf.Data or nil if buf == nil.
+func (buf *Buf) XData() []byte {
+	if buf != nil {
+		return buf.Data
+	}
+	return nil
+}

mem/buffer_test.go

+// Copyright (C) 2017  Nexedi SA and Contributors.
+//                     Kirill Smelkov <kirr@nexedi.com>
+//
+// This program is free software: you can Use, Study, Modify and Redistribute
+// it under the terms of the GNU General Public License version 3, or (at your
+// option) any later version, as published by the Free Software Foundation.
+//
+// You can also Link and Combine this program with other software covered by
+// the terms of any of the Free Software licenses or any of the Open Source
+// Initiative approved licenses and Convey the resulting work. Corresponding
+// source of such a combination shall include the source code for all other
+// software used.
+//
+// This program is distributed WITHOUT ANY WARRANTY; without even the implied
+// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+//
+// See COPYING file for full licensing terms.
+// See https://www.nexedi.com/licensing for rationale and options.
+
+// As of go19 sync.Pool under race-detector randomly drops items on the floor
+// https://github.com/golang/go/blob/ca360c39/src/sync/pool.go#L92
+// so it is not possible to verify we will get what we've just put there.
+// +build !race
+
+package mem
+
+import (
+	"reflect"
+	"testing"
+	"unsafe"
+)
+
+//go:linkname runtime_procPin runtime.procPin
+//go:linkname runtime_procUnpin runtime.procUnpin
+func runtime_procPin() int
+func runtime_procUnpin()
+
+
+func sliceDataPtr(b []byte) unsafe.Pointer {
+	return unsafe.Pointer((*reflect.SliceHeader)(unsafe.Pointer(&b)).Data)
+}
+
+func TestBufAllocFree(t *testing.T) {
+	// sync.Pool uses per-P free-lists. We check that after release we will
+	// allocate released object. This works only if P is not changed underneath.
+	runtime_procPin()
+	defer runtime_procUnpin()
+
+	for i := uint(0); i < 25; i++ {
+		size := 1<<i - 1
+		xcap := 1<<i
+		buf := BufAlloc(size)
+		if i < order0 {
+			xcap = 1 << order0
+		}
+		if int(i) >= order0+len(bufPoolv) {
+			xcap = size
+		}
+
+		if len(buf.Data) != size {
+			t.Fatalf("%v: len=%v  ; want %v", i, len(buf.Data), size)
+		}
+		if cap(buf.Data) != xcap {
+			t.Fatalf("%v: cap=%v  ; want %v", i, cap(buf.Data), xcap)
+		}
+
+		checkref := func(rc int32) {
+			t.Helper()
+			if buf.refcnt != rc {
+				t.Fatalf("%v: refcnt=%v  ; want %v", i, buf.refcnt, rc)
+			}
+		}
+
+		checkref(0)
+
+		// free and allocate another buf -> it must be it
+		data := buf.Data
+		buf.Release()
+		checkref(-1)
+		buf2 := BufAlloc(size)
+
+		// not from pool - memory won't be reused
+		if int(i) >= order0+len(bufPoolv) {
+			if buf2 == buf || sliceDataPtr(buf2.Data) == sliceDataPtr(data) {
+				t.Fatalf("%v: buffer reused but should not", i)
+			}
+			continue
+		}
+
+		// from pool -> it must be the same
+		if !(buf2 == buf && sliceDataPtr(buf2.Data) == sliceDataPtr(data)) {
+			t.Fatalf("%v: buffer not reused on free/realloc", i)
+		}
+		checkref(0)
+
+		// add more ref and release original buf - it must stay alive
+		buf.Incref()
+		checkref(1)
+		buf.Release()
+		checkref(0)
+
+		// another alloc must be different
+		buf2 = BufAlloc(size)
+		checkref(0)
+
+		if buf2 == buf || sliceDataPtr(buf2.Data) == sliceDataPtr(data) {
+			t.Fatalf("%v: buffer reused but should not", i)
+		}
+
+		// release buf again -> should go to pool
+		buf.Release()
+		checkref(-1)
+		buf2 = BufAlloc(size)
+		if !(buf2 == buf && sliceDataPtr(buf2.Data) == sliceDataPtr(data)) {
+			t.Fatalf("%v: buffer not reused on free/realloc", i)
+		}
+		checkref(0)
+	}
+}

mem/buffer_test.s

+// empty .s so `go build` does not use -complete for go:linkname to work

mem/mem.go

@@ -17,7 +17,8 @@
 // See COPYING file for full licensing terms.
 // See https://www.nexedi.com/licensing for rationale and options.
 
-// Package mem allows to work with memory as either string or []byte without copying
+// Package mem provides reference-counted buffer and ways to work with memory
+// as either string or []byte without copying.
 package mem
 
 import (