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

This reverts commit 8adf8e1.

Author: NDStrahilevitz

buf-perf.go

@@ -9,9 +9,7 @@ import "C"
 import (
 	"fmt"
 	"sync"
-	"sync/atomic"
 	"syscall"
-	"unsafe"
 )
 
 //
@@ -24,14 +22,15 @@ 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)
 }
@@ -42,12 +41,12 @@ func (pb *PerfBuffer) Start() {
 }
 
 func (pb *PerfBuffer) Stop() {
-	if atomic.LoadUint32(&pb.stopFlag) == 1 {
+	if pb.stop == nil {
 		return
 	}
 
 	// Signal the poll goroutine to exit
-	atomic.StoreUint32(&pb.stopFlag, 1)
+	close(pb.stop)
 
 	// The event and lost channels should be drained here since the consumer
 	// may have stopped at this point. Failure to drain it will
@@ -74,6 +73,9 @@ 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() {
@@ -91,11 +93,20 @@ func (pb *PerfBuffer) Close() {
 func (pb *PerfBuffer) poll(timeout int) error {
 	defer pb.wg.Done()
 
-	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))
+	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)
+			}
+		}
 	}
-
-	return nil
 }

buf-ring.go

@@ -9,27 +9,26 @@ import "C"
 import (
 	"fmt"
 	"sync"
-	"sync/atomic"
 	"syscall"
-	"unsafe"
 )
 
 //
 // RingBuffer
 //
 
 type RingBuffer struct {
-	rb       *C.struct_ring_buffer
-	bpfMap   *BPFMap
-	slots    []uint
-	closed   bool
-	wg       sync.WaitGroup
-	stopFlag uint32 // use with atomic operations
+	rb     *C.struct_ring_buffer
+	bpfMap *BPFMap
+	slots  []uint
+	stop   chan struct{}
+	closed bool
+	wg     sync.WaitGroup
 }
 
 // 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)
 }
@@ -40,12 +39,12 @@ func (rb *RingBuffer) Start() {
 }
 
 func (rb *RingBuffer) Stop() {
-	if atomic.LoadUint32(&rb.stopFlag) == 1 {
+	if rb.stop == nil {
 		return
 	}
 
 	// Signal the poll goroutine to exit
-	atomic.StoreUint32(&rb.stopFlag, 1)
+	close(rb.stop)
 
 	// The event channel should be drained here since the consumer
 	// may have stopped at this point. Failure to drain it will
@@ -70,6 +69,9 @@ 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() {
@@ -85,13 +87,32 @@ 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()
 
-	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))
+	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)
+		}
 	}
 
 	return nil

libbpfgo.c

@@ -89,26 +89,6 @@ 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 = {};
@@ -128,26 +108,6 @@ 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,10 +24,7 @@ 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);