snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/profiling/buffer/buffer.go

@@ -0,0 +1,272 @@
+/*
+ *
+ * Copyright 2019 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package buffer provides a high-performant lock free implementation of a
+// circular buffer used by the profiling code.
+package buffer
+
+import (
+	"errors"
+	"math/bits"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"unsafe"
+)
+
+type queue struct {
+	// An array of pointers as references to the items stored in this queue.
+	arr []unsafe.Pointer
+	// The maximum number of elements this queue may store before it wraps around
+	// and overwrites older values. Must be an exponent of 2.
+	size uint32
+	// Always size - 1. A bitwise AND is performed with this mask in place of a
+	// modulo operation by the Push operation.
+	mask uint32
+	// Each Push operation into this queue increments the acquired counter before
+	// proceeding forwarding with the actual write to arr. This counter is also
+	// used by the Drain operation's drainWait subroutine to wait for all pushes
+	// to complete.
+	acquired uint32 // Accessed atomically.
+	// After the completion of a Push operation, the written counter is
+	// incremented. Also used by drainWait to wait for all pushes to complete.
+	written uint32
+}
+
+// Allocates and returns a new *queue. size needs to be an exponent of two.
+func newQueue(size uint32) *queue {
+	return &queue{
+		arr:  make([]unsafe.Pointer, size),
+		size: size,
+		mask: size - 1,
+	}
+}
+
+// drainWait blocks the caller until all Pushes on this queue are complete.
+func (q *queue) drainWait() {
+	for atomic.LoadUint32(&q.acquired) != atomic.LoadUint32(&q.written) {
+		runtime.Gosched()
+	}
+}
+
+// A queuePair has two queues. At any given time, Pushes go into the queue
+// referenced by queuePair.q. The active queue gets switched when there's a
+// drain operation on the circular buffer.
+type queuePair struct {
+	q0 unsafe.Pointer
+	q1 unsafe.Pointer
+	q  unsafe.Pointer
+}
+
+// Allocates and returns a new *queuePair with its internal queues allocated.
+func newQueuePair(size uint32) *queuePair {
+	qp := &queuePair{}
+	qp.q0 = unsafe.Pointer(newQueue(size))
+	qp.q1 = unsafe.Pointer(newQueue(size))
+	qp.q = qp.q0
+	return qp
+}
+
+// Switches the current queue for future Pushes to proceed to the other queue
+// so that there's no blocking in Push. Returns a pointer to the old queue that
+// was in place before the switch.
+func (qp *queuePair) switchQueues() *queue {
+	// Even though we have mutual exclusion across drainers (thanks to mu.Lock in
+	// drain), Push operations may access qp.q whilst we're writing to it.
+	if atomic.CompareAndSwapPointer(&qp.q, qp.q0, qp.q1) {
+		return (*queue)(qp.q0)
+	}
+
+	atomic.CompareAndSwapPointer(&qp.q, qp.q1, qp.q0)
+	return (*queue)(qp.q1)
+}
+
+// In order to not have expensive modulo operations, we require the maximum
+// number of elements in the circular buffer (N) to be an exponent of two to
+// use a bitwise AND mask. Since a CircularBuffer is a collection of queuePairs
+// (see below), we need to divide N; since exponents of two are only divisible
+// by other exponents of two, we use floorCPUCount number of queuePairs within
+// each CircularBuffer.
+//
+// Floor of the number of CPUs (and not the ceiling) was found to be the
+// optimal number through experiments.
+func floorCPUCount() uint32 {
+	floorExponent := bits.Len32(uint32(runtime.NumCPU())) - 1
+	if floorExponent < 0 {
+		floorExponent = 0
+	}
+	return 1 << uint32(floorExponent)
+}
+
+var numCircularBufferPairs = floorCPUCount()
+
+// CircularBuffer is a lock-free data structure that supports Push and Drain
+// operations.
+//
+// Note that CircularBuffer is built for performance more than reliability.
+// That is, some Push operations may fail without retries in some situations
+// (such as during a Drain operation). Order of pushes is not maintained
+// either; that is, if A was pushed before B, the Drain operation may return an
+// array with B before A. These restrictions are acceptable within gRPC's
+// profiling, but if your use-case does not permit these relaxed constraints
+// or if performance is not a primary concern, you should probably use a
+// lock-based data structure such as internal/buffer.UnboundedBuffer.
+type CircularBuffer struct {
+	drainMutex sync.Mutex
+	qp         []*queuePair
+	// qpn is a monotonically incrementing counter that's used to determine
+	// which queuePair a Push operation should write to. This approach's
+	// performance was found to be better than writing to a random queue.
+	qpn    uint32
+	qpMask uint32
+}
+
+var errInvalidCircularBufferSize = errors.New("buffer size is not an exponent of two")
+
+// NewCircularBuffer allocates a circular buffer of size size and returns a
+// reference to the struct. Only circular buffers of size 2^k are allowed
+// (saves us from having to do expensive modulo operations).
+func NewCircularBuffer(size uint32) (*CircularBuffer, error) {
+	if size&(size-1) != 0 {
+		return nil, errInvalidCircularBufferSize
+	}
+
+	n := numCircularBufferPairs
+	if size/numCircularBufferPairs < 8 {
+		// If each circular buffer is going to hold less than a very small number
+		// of items (let's say 8), using multiple circular buffers is very likely
+		// wasteful. Instead, fallback to one circular buffer holding everything.
+		n = 1
+	}
+
+	cb := &CircularBuffer{
+		qp:     make([]*queuePair, n),
+		qpMask: n - 1,
+	}
+
+	for i := uint32(0); i < n; i++ {
+		cb.qp[i] = newQueuePair(size / n)
+	}
+
+	return cb, nil
+}
+
+// Push pushes an element in to the circular buffer. Guaranteed to complete in
+// a finite number of steps (also lock-free). Does not guarantee that push
+// order will be retained. Does not guarantee that the operation will succeed
+// if a Drain operation concurrently begins execution.
+func (cb *CircularBuffer) Push(x any) {
+	n := atomic.AddUint32(&cb.qpn, 1) & cb.qpMask
+	qptr := atomic.LoadPointer(&cb.qp[n].q)
+	q := (*queue)(qptr)
+
+	acquired := atomic.AddUint32(&q.acquired, 1) - 1
+
+	// If true, it means that we have incremented acquired before any queuePair
+	// was switched, and therefore before any drainWait completion. Therefore, it
+	// is safe to proceed with the Push operation on this queue. Otherwise, it
+	// means that a Drain operation has begun execution, but we don't know how
+	// far along the process it is. If it is past the drainWait check, it is not
+	// safe to proceed with the Push operation. We choose to drop this sample
+	// entirely instead of retrying, as retrying may potentially send the Push
+	// operation into a spin loop (we want to guarantee completion of the Push
+	// operation within a finite time). Before exiting, we increment written so
+	// that any existing drainWaits can proceed.
+	if atomic.LoadPointer(&cb.qp[n].q) != qptr {
+		atomic.AddUint32(&q.written, 1)
+		return
+	}
+
+	// At this point, we're definitely writing to the right queue. That is, one
+	// of the following is true:
+	//   1. No drainer is in execution on this queue.
+	//   2. A drainer is in execution on this queue and it is waiting at the
+	//      acquired == written barrier.
+	//
+	// Let's say two Pushes A and B happen on the same queue. Say A and B are
+	// q.size apart; i.e. they get the same index. That is,
+	//
+	//   index_A = index_B
+	//   acquired_A + q.size = acquired_B
+	//
+	// We say "B has wrapped around A" when this happens. In this case, since A
+	// occurred before B, B's Push should be the final value. However, we
+	// accommodate A being the final value because wrap-arounds are extremely
+	// rare and accounting for them requires an additional counter and a
+	// significant performance penalty. Note that the below approach never leads
+	// to any data corruption.
+	index := acquired & q.mask
+	atomic.StorePointer(&q.arr[index], unsafe.Pointer(&x))
+
+	// Allows any drainWait checks to proceed.
+	atomic.AddUint32(&q.written, 1)
+}
+
+// Dereferences non-nil pointers from arr into result. Range of elements from
+// arr that are copied is [from, to). Assumes that the result slice is already
+// allocated and is large enough to hold all the elements that might be copied.
+// Also assumes mutual exclusion on the array of pointers.
+func dereferenceAppend(result []any, arr []unsafe.Pointer, from, to uint32) []any {
+	for i := from; i < to; i++ {
+		// We have mutual exclusion on arr, there's no need for atomics.
+		x := (*any)(arr[i])
+		if x != nil {
+			result = append(result, *x)
+		}
+	}
+	return result
+}
+
+// Drain allocates and returns an array of things Pushed in to the circular
+// buffer. Push order is not maintained; that is, if B was Pushed after A,
+// drain may return B at a lower index than A in the returned array.
+func (cb *CircularBuffer) Drain() []any {
+	cb.drainMutex.Lock()
+
+	qs := make([]*queue, len(cb.qp))
+	for i := 0; i < len(cb.qp); i++ {
+		qs[i] = cb.qp[i].switchQueues()
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(len(qs))
+	for i := 0; i < len(qs); i++ {
+		go func(qi int) {
+			qs[qi].drainWait()
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+
+	result := make([]any, 0)
+	for i := 0; i < len(qs); i++ {
+		if acquired := atomic.LoadUint32(&qs[i].acquired); acquired < qs[i].size {
+			result = dereferenceAppend(result, qs[i].arr, 0, acquired)
+		} else {
+			result = dereferenceAppend(result, qs[i].arr, 0, qs[i].size)
+		}
+	}
+
+	for i := 0; i < len(qs); i++ {
+		atomic.StoreUint32(&qs[i].acquired, 0)
+		atomic.StoreUint32(&qs[i].written, 0)
+	}
+
+	cb.drainMutex.Unlock()
+	return result
+}

internal/profiling/buffer/buffer_test.go

@@ -0,0 +1,186 @@
+/*
+ *
+ * Copyright 2019 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package buffer
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/grpctest"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestCircularBufferSerial(t *testing.T) {
+	var size, i uint32
+	var result []any
+
+	size = 1 << 15
+	cb, err := NewCircularBuffer(size)
+	if err != nil {
+		t.Fatalf("error allocating CircularBuffer: %v", err)
+	}
+
+	for i = 0; i < size/2; i++ {
+		cb.Push(i)
+	}
+
+	result = cb.Drain()
+	if uint32(len(result)) != size/2 {
+		t.Fatalf("len(result) = %d; want %d", len(result), size/2)
+	}
+
+	// The returned result isn't necessarily sorted.
+	seen := make(map[uint32]bool)
+	for _, r := range result {
+		seen[r.(uint32)] = true
+	}
+
+	for i = 0; i < uint32(len(result)); i++ {
+		if !seen[i] {
+			t.Fatalf("seen[%d] = false; want true", i)
+		}
+	}
+
+	for i = 0; i < size; i++ {
+		cb.Push(i)
+	}
+
+	result = cb.Drain()
+	if uint32(len(result)) != size {
+		t.Fatalf("len(result) = %d; want %d", len(result), size/2)
+	}
+}
+
+func (s) TestCircularBufferOverflow(t *testing.T) {
+	var size, i uint32
+	var result []any
+
+	size = 1 << 10
+	cb, err := NewCircularBuffer(size)
+	if err != nil {
+		t.Fatalf("error allocating CircularBuffer: %v", err)
+	}
+
+	for i = 0; i < 10*size; i++ {
+		cb.Push(i)
+	}
+
+	result = cb.Drain()
+
+	if uint32(len(result)) != size {
+		t.Fatalf("len(result) = %d; want %d", len(result), size)
+	}
+
+	for idx, x := range result {
+		if x.(uint32) < size {
+			t.Fatalf("result[%d] = %d; want it to be >= %d", idx, x, size)
+		}
+	}
+}
+
+func (s) TestCircularBufferConcurrent(t *testing.T) {
+	for tn := 0; tn < 2; tn++ {
+		var size uint32
+		var result []any
+
+		size = 1 << 6
+		cb, err := NewCircularBuffer(size)
+		if err != nil {
+			t.Fatalf("error allocating CircularBuffer: %v", err)
+		}
+
+		type item struct {
+			R uint32
+			N uint32
+			T time.Time
+		}
+
+		var wg sync.WaitGroup
+		for r := uint32(0); r < 1024; r++ {
+			wg.Add(1)
+			go func(r uint32) {
+				for n := uint32(0); n < size; n++ {
+					cb.Push(item{R: r, N: n, T: time.Now()})
+				}
+				wg.Done()
+			}(r)
+		}
+
+		// Wait for all goroutines to finish only in one test. Draining
+		// concurrently while Pushes are still happening will test for races in the
+		// Draining lock.
+		if tn == 0 {
+			wg.Wait()
+		}
+
+		result = cb.Drain()
+
+		// Can't expect the buffer to be full if the Pushes aren't necessarily done.
+		if tn == 0 {
+			if uint32(len(result)) != size {
+				t.Fatalf("len(result) = %d; want %d", len(result), size)
+			}
+		}
+
+		// There can be absolutely no expectation on the order of the data returned
+		// by Drain because: (a) everything is happening concurrently (b) a
+		// round-robin is used to write to different queues (and therefore
+		// different cachelines) for less write contention.
+
+		// Wait for all goroutines to complete before moving on to other tests. If
+		// the benchmarks run after this, it might affect performance unfairly.
+		wg.Wait()
+	}
+}
+
+func BenchmarkCircularBuffer(b *testing.B) {
+	x := 1
+	for size := 1 << 16; size <= 1<<20; size <<= 1 {
+		for routines := 1; routines <= 1<<8; routines <<= 1 {
+			b.Run(fmt.Sprintf("goroutines:%d/size:%d", routines, size), func(b *testing.B) {
+				cb, err := NewCircularBuffer(uint32(size))
+				if err != nil {
+					b.Fatalf("error allocating CircularBuffer: %v", err)
+				}
+
+				perRoutine := b.N / routines
+				var wg sync.WaitGroup
+				for r := 0; r < routines; r++ {
+					wg.Add(1)
+					go func() {
+						for i := 0; i < perRoutine; i++ {
+							cb.Push(&x)
+						}
+						wg.Done()
+					}()
+				}
+				wg.Wait()
+			})
+		}
+	}
+}