1// Package fdwait parks descriptor waits in Go's netpoller.
2package fdwait
3
4import (
5 "context"
6 "errors"
7 "fmt"
8 "os"
9 "slices"
10 "syscall"
11 "time"
12
13 "golang.org/x/sys/unix"
14)
15
16// Poller owns an epoll descriptor and is used by one goroutine; Wake and
17// Close may run from any goroutine. Watched descriptors remain owned by the
18// caller, who must not close them while a Poll waits.
19//
20// Descriptors that live as long as the poller are registered once with Pin;
21// others are registered for one Poll call. A Poll returns on a watched fd,
22// on Wake, once the context given to StopOn ends, or on Close: a loop needs
23// no polling timeout to notice work or shutdown.
24type Poller struct {
25 file *os.File
26 raw syscall.RawConn
27 // wake is an eventfd pinned in the epoll set: Wake writes it.
28 wake *os.File
29 wakeRaw syscall.RawConn
30 wakeFD int32
31 ctx context.Context
32 stop func() bool
33 pinned []int32 // besides wakeFD
34 // broken is set when a registration could not be removed: it may keep
35 // the epoll descriptor readable forever, so every later Poll fails
36 // instead of spinning.
37 broken error
38 // noDeadline: the read deadline is already cleared.
39 noDeadline bool
40
41 // State of the running Poll, used by the callbacks below. They are
42 // bound once in New, so a Poll allocates nothing.
43 fds []unix.PollFd
44 check []unix.PollFd
45 added int
46 ctlErr error
47 n int
48 pollErr error
49 buf [8]byte
50 addFn func(uintptr)
51 delFn func(uintptr)
52 checkFn func(uintptr)
53 readFn func(uintptr) bool
54 drainFn func(uintptr)
55 kickFn func(uintptr)
56}
57
58func New() (*Poller, error) {
59 fd, err := unix.EpollCreate1(unix.EPOLL_CLOEXEC)
60 if err != nil {
61 return nil, fmt.Errorf("fdwait epoll create: %w", err)
62 }
63 if err = unix.SetNonblock(fd, true); err != nil {
64 unix.Close(fd)
65 return nil, fmt.Errorf("fdwait epoll nonblock: %w", err)
66 }
67 p := &Poller{file: os.NewFile(uintptr(fd), "fdwait")}
68 if p.raw, err = p.file.SyscallConn(); err != nil {
69 p.file.Close()
70 return nil, fmt.Errorf("fdwait epoll conn: %w", err)
71 }
72 efd, err := unix.Eventfd(0, unix.EFD_CLOEXEC|unix.EFD_NONBLOCK)
73 if err != nil {
74 p.file.Close()
75 return nil, fmt.Errorf("fdwait wake eventfd: %w", err)
76 }
77 p.wake, p.wakeFD = os.NewFile(uintptr(efd), "fdwait-wake"), int32(efd)
78 p.addFn, p.delFn, p.readFn, p.checkFn = p.watch, p.unwatchAll, p.ready, p.pollNow
79 p.drainFn, p.kickFn = p.drain, p.kick
80 if p.wakeRaw, err = p.wake.SyscallConn(); err == nil {
81 err = p.control(unix.EPOLL_CTL_ADD, p.wakeFD)
82 }
83 if err != nil {
84 p.Close()
85 return nil, fmt.Errorf("fdwait wake: %w", err)
86 }
87 return p, nil
88}
89
90// Close ends a running Poll and releases the epoll and wake descriptors.
91func (p *Poller) Close() error {
92 if p.stop != nil {
93 p.stop()
94 }
95 err := p.file.Close()
96 p.wake.Close()
97 return err
98}
99
100func (p *Poller) control(op int, fd int32) error {
101 var err error
102 if cerr := p.raw.Control(func(ep uintptr) {
103 ev := unix.EpollEvent{Events: unix.EPOLLIN, Fd: fd}
104 err = unix.EpollCtl(int(ep), op, int(fd), &ev)
105 }); cerr != nil {
106 return cerr
107 }
108 return err
109}
110
111// Pin registers fd for every later Poll, which then costs it no epoll call.
112// fd must stay open until Close, and every Poll must list it: a ready pinned
113// fd that a Poll does not list fails that Poll.
114func (p *Poller) Pin(fd int32) error {
115 if err := p.control(unix.EPOLL_CTL_ADD, fd); err != nil {
116 return fmt.Errorf("fdwait pin fd %d: %w", fd, err)
117 }
118 p.pinned = append(p.pinned, fd)
119 return nil
120}
121
122// StopOn makes every Poll return (0, nil) at once after ctx ends. Call it
123// once, before polling.
124func (p *Poller) StopOn(ctx context.Context) {
125 p.ctx = ctx
126 p.stop = context.AfterFunc(ctx, p.Wake)
127}
128
129// Wake makes the running or next Poll return. Wakes before that Poll
130// coalesce. Safe from any goroutine, also after Close.
131func (p *Poller) Wake() { _ = p.wakeRaw.Control(p.kickFn) }
132
133func (p *Poller) kick(fd uintptr) {
134 one := [8]byte{1}
135 _, _ = unix.Write(int(fd), one[:])
136}
137
138func (p *Poller) drain(fd uintptr) { _, _ = unix.Read(int(fd), p.buf[:]) }
139
140func (p *Poller) stopped() bool { return p.ctx != nil && p.ctx.Err() != nil }
141
142// Poll waits like poll(2) for readable/error events on fds and fills their
143// Revents. It never consumes caller events. Negative fds are ignored and a
144// descriptor listed twice is watched once. timeout is in milliseconds; a
145// negative one waits until an event, Wake, the end of the StopOn context or
146// Close. A Wake or stop returns the events seen, possibly none.
147//
148// Each call removes the registrations it made before returning. Relying on
149// close to remove them is unsafe: another fd or a kernel fence may retain
150// the open file description and keep the epoll fd readable. When nothing is
151// ready yet, a descriptor epoll cannot watch (EPERM, EINVAL, ELOOP for epoll
152// nested too deep) fails the call.
153func (p *Poller) Poll(fds []unix.PollFd, timeout int) (int, error) {
154 if p.broken != nil {
155 return 0, p.broken
156 }
157 p.fds = fds
158 defer func() { p.fds = nil }()
159 transient := false
160 for i := range fds {
161 transient = transient || !p.skip(i)
162 }
163 // An expired deadline skips RawConn.Read's own check, and a ready
164 // transient fd needs no registration. Pinned fds are checked by Read.
165 if (timeout == 0 || transient) && p.ready(0) || timeout == 0 {
166 return p.n, p.pollErr
167 }
168 p.added, p.ctlErr = 0, nil
169 if transient {
170 err := p.raw.Control(p.addFn)
171 defer p.unwatch()
172 if err != nil {
173 return 0, err
174 }
175 if p.ctlErr != nil {
176 return 0, p.ctlErr
177 }
178 }
179 var err error
180 switch {
181 case timeout > 0:
182 err, p.noDeadline = p.file.SetReadDeadline(time.Now().Add(time.Duration(timeout)*time.Millisecond)), false
183 case !p.noDeadline:
184 err, p.noDeadline = p.file.SetReadDeadline(time.Time{}), true
185 }
186 if err != nil {
187 return 0, err
188 }
189 err = p.raw.Read(p.readFn)
190 if errors.Is(err, os.ErrDeadlineExceeded) {
191 return 0, nil
192 }
193 if err != nil {
194 return 0, err
195 }
196 return p.n, p.pollErr
197}
198
199// ready checks the caller's fds, the wake fd and unlisted pinned fds without
200// blocking, retrying an interrupted check. It reports whether Poll returns.
201func (p *Poller) ready(uintptr) bool {
202 p.check = append(append(p.check[:0], p.fds...), unix.PollFd{Fd: p.wakeFD, Events: unix.POLLIN})
203 for _, fd := range p.pinned {
204 if !slices.ContainsFunc(p.fds, func(f unix.PollFd) bool { return f.Fd == fd }) {
205 p.check = append(p.check, unix.PollFd{Fd: fd, Events: unix.POLLIN})
206 }
207 }
208 // Under the wake file's reference, so a concurrent Close cannot free
209 // its number for reuse during the check.
210 if p.wakeRaw.Control(p.checkFn) != nil {
211 p.n, p.pollErr = 0, os.ErrClosed
212 }
213 if p.pollErr != nil {
214 p.n = 0
215 return true
216 }
217 for i := range p.fds {
218 p.fds[i].Revents = p.check[i].Revents
219 }
220 for _, f := range p.check[len(p.fds)+1:] {
221 if f.Revents != 0 {
222 // Its readiness keeps the epoll fd readable: waiting would spin.
223 p.n, p.pollErr = 0, fmt.Errorf("fdwait: pinned fd %d is ready but not polled", f.Fd)
224 return true
225 }
226 }
227 woke := p.check[len(p.fds)].Revents != 0
228 if woke {
229 p.n--
230 // A stopped poller stays awake: every later Poll returns at once.
231 if !p.stopped() {
232 _ = p.wakeRaw.Control(p.drainFn)
233 if p.stopped() { // the stop's write may have been drained
234 p.Wake()
235 }
236 }
237 }
238 return p.n != 0 || woke
239}
240
241func (p *Poller) pollNow(uintptr) {
242 for {
243 p.n, p.pollErr = unix.Poll(p.check, 0)
244 if !errors.Is(p.pollErr, unix.EINTR) {
245 return
246 }
247 }
248}
249
250// skip reports whether fds[i] needs no registration of its own.
251func (p *Poller) skip(i int) bool {
252 fd := p.fds[i].Fd
253 return fd < 0 || slices.Contains(p.pinned, fd) || listedBefore(p.fds, i)
254}
255
256func (p *Poller) watch(ep uintptr) {
257 for i, f := range p.fds {
258 if !p.skip(i) {
259 ev := unix.EpollEvent{Events: unix.EPOLLIN, Fd: f.Fd}
260 if e := unix.EpollCtl(int(ep), unix.EPOLL_CTL_ADD, int(f.Fd), &ev); e != nil {
261 p.ctlErr = fmt.Errorf("fdwait watch fd %d: %w", f.Fd, e)
262 return
263 }
264 }
265 p.added++
266 }
267}
268
269// unwatch removes the registrations Poll added. A failure means a watched
270// fd was closed during the wait: its registration may outlive it.
271func (p *Poller) unwatch() {
272 if p.added > 0 {
273 _ = p.raw.Control(p.delFn)
274 }
275}
276
277func (p *Poller) unwatchAll(ep uintptr) {
278 for i, f := range p.fds[:p.added] {
279 if p.skip(i) {
280 continue
281 }
282 if e := unix.EpollCtl(int(ep), unix.EPOLL_CTL_DEL, int(f.Fd), nil); e != nil && p.broken == nil {
283 p.broken = fmt.Errorf("fdwait unwatch fd %d (closed during a wait?): %w", f.Fd, e)
284 }
285 }
286}
287
288// listedBefore reports whether fds[i].Fd also appears earlier in fds.
289func listedBefore(fds []unix.PollFd, i int) bool {
290 fd := fds[i].Fd
291 return slices.ContainsFunc(fds[:i], func(f unix.PollFd) bool { return f.Fd == fd })
292}