feat(buffer): move polling to cgo loop (#493)

Polling loops are cgo hotpaths in consuming applications. Since cgo calls
are rather expensive, moving the polling loop to cgo should significantly
reduce the number of cgo calls.

Now the go part only handles orchestration of the internal c loop.
An atomic flag is used to signal the c loop to exit.

Author: NDStrahilevitz

buf-perf.go

@@ -9,7 +9,9 @@ import "C"
 import (
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"syscall"
+	"unsafe"
 )
 
 //
@@ -22,15 +24,14 @@ type PerfBuffer struct {
 	slot       uint
 	eventsChan chan []byte
 	lostChan   chan uint64
-	stop       chan struct{}
 	closed     bool
 	wg         sync.WaitGroup
+	stopFlag   uint32 // use with atomic operations
 }
 
 // Poll will wait until timeout in milliseconds to gather
 // data from the perf buffer.
 func (pb *PerfBuffer) Poll(timeout int) {
-	pb.stop = make(chan struct{})
 	pb.wg.Add(1)
 	go pb.poll(timeout)
 }
@@ -41,12 +42,12 @@ func (pb *PerfBuffer) Start() {
 }
 
 func (pb *PerfBuffer) Stop() {
-	if pb.stop == nil {
+	if atomic.LoadUint32(&pb.stopFlag) == 1 {
 		return
 	}
 
 	// Signal the poll goroutine to exit
-	close(pb.stop)
+	atomic.StoreUint32(&pb.stopFlag, 1)
 
 	// The event and lost channels should be drained here since the consumer
 	// may have stopped at this point. Failure to drain it will
@@ -73,9 +74,6 @@ func (pb *PerfBuffer) Stop() {
 	if pb.lostChan != nil {
 		close(pb.lostChan)
 	}
-
-	// Reset pb.stop to allow multiple safe calls to Stop()
-	pb.stop = nil
 }
 
 func (pb *PerfBuffer) Close() {
@@ -93,20 +91,11 @@ func (pb *PerfBuffer) Close() {
 func (pb *PerfBuffer) poll(timeout int) error {
 	defer pb.wg.Done()
 
-	for {
-		select {
-		case <-pb.stop:
-			return nil
-		default:
-			retC := C.perf_buffer__poll(pb.pb, C.int(timeout))
-			if retC < 0 {
-				errno := syscall.Errno(-retC)
-				if errno == syscall.EINTR {
-					continue
-				}
-
-				return fmt.Errorf("error polling perf buffer: %w", errno)
-			}
-		}
+	stopFlag := (*C.uint32_t)(unsafe.Pointer(&pb.stopFlag))
+	ret := C.cgo_perf_buffer__poll(pb.pb, C.int(timeout), stopFlag)
+	if ret < 0 {
+		return fmt.Errorf("error polling perf buffer: %w", syscall.Errno(-ret))
 	}
+
+	return nil
 }

buf-ring.go

@@ -9,26 +9,27 @@ import "C"
 import (
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"syscall"
+	"unsafe"
 )
 
 //
 // RingBuffer
 //
 
 type RingBuffer struct {
-	rb     *C.struct_ring_buffer
-	bpfMap *BPFMap
-	slots  []uint
-	stop   chan struct{}
-	closed bool
-	wg     sync.WaitGroup
+	rb       *C.struct_ring_buffer
+	bpfMap   *BPFMap
+	slots    []uint
+	closed   bool
+	wg       sync.WaitGroup
+	stopFlag uint32 // use with atomic operations
 }
 
 // Poll will wait until timeout in milliseconds to gather
 // data from the ring buffer.
 func (rb *RingBuffer) Poll(timeout int) {
-	rb.stop = make(chan struct{})
 	rb.wg.Add(1)
 	go rb.poll(timeout)
 }
@@ -39,12 +40,12 @@ func (rb *RingBuffer) Start() {
 }
 
 func (rb *RingBuffer) Stop() {
-	if rb.stop == nil {
+	if atomic.LoadUint32(&rb.stopFlag) == 1 {
 		return
 	}
 
 	// Signal the poll goroutine to exit
-	close(rb.stop)
+	atomic.StoreUint32(&rb.stopFlag, 1)
 
 	// The event channel should be drained here since the consumer
 	// may have stopped at this point. Failure to drain it will
@@ -69,9 +70,6 @@ func (rb *RingBuffer) Stop() {
 		eventChan := eventChannels.get(slot).(chan []byte)
 		close(eventChan)
 	}
-
-	// Reset pb.stop to allow multiple safe calls to Stop()
-	rb.stop = nil
 }
 
 func (rb *RingBuffer) Close() {
@@ -87,32 +85,13 @@ func (rb *RingBuffer) Close() {
 	rb.closed = true
 }
 
-func (rb *RingBuffer) isStopped() bool {
-	select {
-	case <-rb.stop:
-		return true
-	default:
-		return false
-	}
-}
-
 func (rb *RingBuffer) poll(timeout int) error {
 	defer rb.wg.Done()
 
-	for {
-		retC := C.ring_buffer__poll(rb.rb, C.int(timeout))
-		if rb.isStopped() {
-			break
-		}
-
-		if retC < 0 {
-			errno := syscall.Errno(-retC)
-			if errno == syscall.EINTR {
-				continue
-			}
-
-			return fmt.Errorf("error polling ring buffer: %w", errno)
-		}
+	stopFlag := (*C.uint32_t)(unsafe.Pointer(&rb.stopFlag))
+	ret := C.cgo_ring_buffer__poll(rb.rb, C.int(timeout), stopFlag)
+	if ret < 0 {
+		return fmt.Errorf("error polling perf buffer: %w", syscall.Errno(-ret))
 	}
 
 	return nil

libbpfgo.c

@@ -89,6 +89,26 @@ int cgo_add_ring_buf(struct ring_buffer *rb, int map_fd, uintptr_t ctx)
     return ret;
 }
 
+int cgo_ring_buffer__poll(struct ring_buffer *rb, int timeout, volatile uint32_t *stop_flag)
+{
+    while (*stop_flag == 0) {
+        int ret = ring_buffer__poll(rb, timeout);
+
+        if (ret < 0) {
+            // Handle EINTR like the original Go code - continue polling
+            if (errno == EINTR) {
+                continue;
+            }
+            // Return other errors to Go
+            return ret;
+        }
+
+        // Continue polling until stop_flag is set
+        // Events are still processed in real-time via callbacks
+    }
+    return 0;
+}
+
 struct perf_buffer *cgo_init_perf_buf(int map_fd, int page_cnt, uintptr_t ctx)
 {
     struct perf_buffer_opts pb_opts = {};
@@ -108,6 +128,26 @@ struct perf_buffer *cgo_init_perf_buf(int map_fd, int page_cnt, uintptr_t ctx)
     return pb;
 }
 
+int cgo_perf_buffer__poll(struct perf_buffer *pb, int timeout_ms, volatile uint32_t *stop_flag)
+{
+    while (*stop_flag == 0) {
+        int ret = perf_buffer__poll(pb, timeout_ms);
+
+        if (ret < 0) {
+            // Handle EINTR like the original Go code - continue polling
+            if (errno == EINTR) {
+                continue;
+            }
+            // Return other errors to Go
+            return ret;
+        }
+
+        // Continue polling until stop_flag is set
+        // Events are still processed in real-time via callbacks
+    }
+    return 0;
+}
+
 void cgo_bpf_map__initial_value(struct bpf_map *map, void *value)
 {
     size_t psize;

libbpfgo.h

@@ -24,7 +24,10 @@ void cgo_libbpf_set_print_fn();
 struct ring_buffer *cgo_init_ring_buf(int map_fd, uintptr_t ctx);
 struct user_ring_buffer *cgo_init_user_ring_buf(int map_fd);
 int cgo_add_ring_buf(struct ring_buffer *rb, int map_fd, uintptr_t ctx);
+int cgo_ring_buffer__poll(struct ring_buffer *rb, int timeout, volatile uint32_t *stop_flag);
+
 struct perf_buffer *cgo_init_perf_buf(int map_fd, int page_cnt, uintptr_t ctx);
+int cgo_perf_buffer__poll(struct perf_buffer *pb, int timeout, volatile uint32_t *stop_flag);
 
 void cgo_bpf_map__initial_value(struct bpf_map *map, void *value);