perf add timeout argument

internal/epoll/poller.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"runtime"
 	"sync"
+	"time"
 
 	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/unix"
@@ -118,8 +119,8 @@ func (p *Poller) Add(fd int, id int) error {
 // Wait for events.
 //
 // Returns the number of pending events or an error wrapping os.ErrClosed if
-// Close is called.
-func (p *Poller) Wait(events []unix.EpollEvent) (int, error) {
+// Close is called, or os.ErrDeadlineExceeded if EpollWait timeout.
+func (p *Poller) Wait(events []unix.EpollEvent, deadline time.Time) (int, error) {
 	p.epollMu.Lock()
 	defer p.epollMu.Unlock()
 
@@ -128,7 +129,12 @@ func (p *Poller) Wait(events []unix.EpollEvent) (int, error) {
 	}
 
 	for {
-		n, err := unix.EpollWait(p.epollFd, events, -1)
+		msec := int(-1)
+		if !deadline.IsZero() {
+			msec = int(time.Until(deadline).Milliseconds())
+		}
+
+		n, err := unix.EpollWait(p.epollFd, events, msec)
 		if temp, ok := err.(temporaryError); ok && temp.Temporary() {
 			// Retry the syscall if we were interrupted, see https://github.com/golang/go/issues/20400
 			continue
@@ -138,6 +144,10 @@ func (p *Poller) Wait(events []unix.EpollEvent) (int, error) {
 			return 0, err
 		}
 
+		if n == 0 && msec != -1 {
+			return 0, fmt.Errorf("poll: %w", os.ErrDeadlineExceeded)
+		}
+
 		for _, event := range events[:n] {
 			if int(event.Fd) == p.event.raw {
 				// Since we don't read p.event the event is never cleared and

internal/epoll/poller_test.go

@@ -27,15 +27,15 @@ func TestPoller(t *testing.T) {
 	}
 
 	done := make(chan struct{}, 1)
-	read := func() {
+	read := func(timeOut time.Time) {
 		defer func() {
 			done <- struct{}{}
 		}()
 
 		events := make([]unix.EpollEvent, 1)
 
-		n, err := poller.Wait(events)
-		if errors.Is(err, os.ErrClosed) {
+		n, err := poller.Wait(events, timeOut)
+		if errors.Is(err, os.ErrClosed) || errors.Is(err, os.ErrDeadlineExceeded) {
 			return
 		}
 
@@ -57,7 +57,7 @@ func TestPoller(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	go read()
+	go read(time.Time{})
 	select {
 	case <-done:
 	case <-time.After(time.Second):
@@ -68,7 +68,14 @@ func TestPoller(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	go read()
+	go read(time.Now().Add(200 * time.Millisecond))
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("Timed out")
+	}
+
+	go read(time.Time{})
 	select {
 	case <-done:
 		t.Fatal("Wait doesn't block")

perf/reader.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"runtime"
 	"sync"
+	"time"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/internal"
@@ -133,7 +134,8 @@ func readRawSample(rd io.Reader, buf, sampleBuf []byte) ([]byte, error) {
 // Reader allows reading bpf_perf_event_output
 // from user space.
 type Reader struct {
-	poller *epoll.Poller
+	poller   *epoll.Poller
+	deadline time.Time
 
 	// mu protects read/write access to the Reader structure with the
 	// exception of 'pauseFds', which is protected by 'pauseMu'.
@@ -237,6 +239,7 @@ func NewReaderWithOptions(array *ebpf.Map, perCPUBuffer int, opts ReaderOptions)
 		array:       array,
 		rings:       rings,
 		poller:      poller,
+		deadline:    time.Time{},
 		epollEvents: make([]unix.EpollEvent, len(rings)),
 		epollRings:  make([]*perfEventRing, 0, len(rings)),
 		eventHeader: make([]byte, perfEventHeaderSize),
@@ -280,6 +283,16 @@ func (pr *Reader) Close() error {
 	return nil
 }
 
+// SetDeadline controls how long Read and ReadInto will block waiting for samples.
+//
+// Passing a zero time.Time will remove the deadline.
+func (pr *Reader) SetDeadline(t time.Time) {
+	pr.mu.Lock()
+	defer pr.mu.Unlock()
+
+	pr.deadline = t
+}
+
 // Read the next record from the perf ring buffer.
 //
 // The function blocks until there are at least Watermark bytes in one
@@ -290,6 +303,8 @@ func (pr *Reader) Close() error {
 // depending on the input sample's length.
 //
 // Calling Close interrupts the function.
+//
+// If SetDeadline the Read deadline, Read might return os.ErrDeadlineExceeded.
 func (pr *Reader) Read() (Record, error) {
 	var r Record
 	return r, pr.ReadInto(&r)
@@ -306,7 +321,7 @@ func (pr *Reader) ReadInto(rec *Record) error {
 
 	for {
 		if len(pr.epollRings) == 0 {
-			nEvents, err := pr.poller.Wait(pr.epollEvents)
+			nEvents, err := pr.poller.Wait(pr.epollEvents, pr.deadline)
 			if err != nil {
 				return err
 			}

perf/reader_test.go

@@ -58,6 +58,10 @@ func TestPerfReader(t *testing.T) {
 	if record.CPU < 0 {
 		t.Error("Record has invalid CPU number")
 	}
+
+	rd.SetDeadline(time.Now().Add(4 * time.Millisecond))
+	_, err = rd.Read()
+	qt.Assert(t, errors.Is(err, os.ErrDeadlineExceeded), qt.IsTrue, qt.Commentf("expected os.ErrDeadlineExceeded"))
 }
 
 func outputSamplesProg(sampleSizes ...int) (*ebpf.Program, *ebpf.Map, error) {

ringbuf/reader.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"os"
 	"sync"
+	"time"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/internal"
@@ -190,7 +191,7 @@ func (r *Reader) ReadInto(rec *Record) error {
 
 	for {
 		if !r.haveData {
-			_, err := r.poller.Wait(r.epollEvents[:cap(r.epollEvents)])
+			_, err := r.poller.Wait(r.epollEvents[:cap(r.epollEvents)], time.Time{})
 			if err != nil {
 				return err
 			}