1package hints
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "encoding/binary"
8 "errors"
9 "fmt"
10 "io"
11 "math"
12 "os"
13 "os/exec"
14 "strconv"
15 "sync"
16 "syscall"
17 "time"
18
19 "github.com/bnema/neferwl/internal/adapters/logging"
20)
21
22// detFrame is a captured output as the detector receives it: BGRA pixels in
23// m, valid only during the Detect call. It holds the mapping, not the bytes,
24// so that recording a call (mocks print their arguments) never reads memory
25// that was unmapped since.
26type detFrame struct {
27 width, height, stride int
28 m *mapping
29 // regions are the crops detected independently; none means the whole
30 // frame.
31 regions []region
32}
33
34// detection is one box found by the detector, in frame (physical) pixels.
35type detection struct {
36 x1, y1, x2, y2, score float32
37 // region is the index of the request region the box comes from.
38 region int
39}
40
41// detector finds clickable elements in a frame. The worker goroutine is its
42// only caller.
43type detector interface {
44 // Detect returns the boxes found in f. An error is for this call only:
45 // the detector recovers by itself on the next one.
46 Detect(ctx context.Context, f detFrame) ([]detection, error)
47 // Close releases the detector. It is called once, by the worker
48 // goroutine, when the worker returns.
49 Close() error
50}
51
52const (
53 protocolVersion = 1
54 helloTimeout = 10 * time.Second
55 requestTimeout = 2 * time.Second
56 imgsz = 1280
57 // maxString and maxBoxes bound what a misbehaving child can make us
58 // allocate.
59 maxString = 1 << 16
60 maxBoxes = 64 * 1000 // runner default max-det per region, at most 64 regions
61)
62
63var (
64 magicHello = [4]byte{'N', 'H', 'N', 'T'}
65 magicRequest = [4]byte{'N', 'H', 'R', 'Q'}
66 magicResponse = [4]byte{'N', 'H', 'R', 'S'}
67
68 // errProtocol is a framing error: the stream cannot be trusted any more.
69 errProtocol = errors.New("hints: detector protocol error")
70 // errDetector is an error the detector reported itself; it keeps serving.
71 errDetector = errors.New("hints: detector error")
72)
73
74var le = binary.LittleEndian
75
76// writeRequest writes one request (the regions of f, default confidence)
77// followed by the pixels.
78func writeRequest(w io.Writer, seq uint64, f detFrame) error {
79 var h [4 + 4 + 8 + 5*4]byte
80 copy(h[0:], magicRequest[:])
81 le.PutUint32(h[4:], protocolVersion)
82 le.PutUint64(h[8:], seq)
83 le.PutUint32(h[16:], uint32(f.width))
84 le.PutUint32(h[20:], uint32(f.height))
85 le.PutUint32(h[24:], uint32(f.stride))
86 le.PutUint32(h[28:], formatXRGB8888)
87 le.PutUint32(h[32:], uint32(len(f.regions)))
88 if _, err := w.Write(h[:]); err != nil {
89 return err
90 }
91 var rec [16]byte
92 for _, r := range f.regions {
93 le.PutUint32(rec[0:], uint32(int32(r.x)))
94 le.PutUint32(rec[4:], uint32(int32(r.y)))
95 le.PutUint32(rec[8:], uint32(r.w))
96 le.PutUint32(rec[12:], uint32(r.h))
97 if _, err := w.Write(rec[:]); err != nil {
98 return err
99 }
100 }
101 var conf [4]byte // confidence: startup default
102 le.PutUint32(conf[:], math.Float32bits(0))
103 if _, err := w.Write(conf[:]); err != nil {
104 return err
105 }
106 _, err := w.Write(f.m.b[:f.height*f.stride])
107 return err
108}
109
110func readMagic(r io.Reader, want [4]byte) error {
111 var m [4]byte
112 if _, err := io.ReadFull(r, m[:]); err != nil {
113 return err
114 }
115 if m != want {
116 return fmt.Errorf("%w: bad magic %q", errProtocol, m[:])
117 }
118 return nil
119}
120
121func readU32(r io.Reader) (uint32, error) {
122 var b [4]byte
123 _, err := io.ReadFull(r, b[:])
124 return le.Uint32(b[:]), err
125}
126
127func readString(r io.Reader) (string, error) {
128 n, err := readU32(r)
129 if err != nil {
130 return "", err
131 }
132 if n > maxString {
133 return "", fmt.Errorf("%w: string of %d bytes", errProtocol, n)
134 }
135 b := make([]byte, n)
136 _, err = io.ReadFull(r, b)
137 return string(b), err
138}
139
140// readHello reads the hello: the device name, or the error of a failed start.
141func readHello(r io.Reader) (device string, err error) {
142 if err := readMagic(r, magicHello); err != nil {
143 return "", err
144 }
145 version, err := readU32(r)
146 if err != nil {
147 return "", err
148 }
149 if version != protocolVersion {
150 return "", fmt.Errorf("%w: version %d", errProtocol, version)
151 }
152 status, err := readU32(r)
153 if err != nil {
154 return "", err
155 }
156 msg, err := readString(r)
157 if err != nil {
158 return "", err
159 }
160 if status != 0 {
161 return "", fmt.Errorf("%w: start failed: %s", errDetector, msg)
162 }
163 if _, err := readU32(r); err != nil { // imgsz
164 return "", err
165 }
166 return msg, nil
167}
168
169// readResponse reads the response to request seq.
170func readResponse(r io.Reader, seq uint64) ([]detection, error) {
171 if err := readMagic(r, magicResponse); err != nil {
172 return nil, err
173 }
174 var h [16]byte
175 if _, err := io.ReadFull(r, h[:]); err != nil {
176 return nil, err
177 }
178 if v := le.Uint32(h[0:]); v != protocolVersion {
179 return nil, fmt.Errorf("%w: version %d", errProtocol, v)
180 }
181 if got := le.Uint64(h[4:]); got != seq {
182 return nil, fmt.Errorf("%w: response %d to request %d", errProtocol, got, seq)
183 }
184 if status := le.Uint32(h[12:]); status != 0 {
185 msg, err := readString(r)
186 if err != nil {
187 return nil, err
188 }
189 return nil, fmt.Errorf("%w: status %d: %s", errDetector, status, msg)
190 }
191 var t [8]byte // inference ms, nboxes
192 if _, err := io.ReadFull(r, t[:]); err != nil {
193 return nil, err
194 }
195 n := le.Uint32(t[4:])
196 if n > maxBoxes {
197 return nil, fmt.Errorf("%w: %d boxes", errProtocol, n)
198 }
199 out := make([]detection, 0, min(n, 4096))
200 var rec [24]byte
201 for range n {
202 if _, err := io.ReadFull(r, rec[:]); err != nil {
203 return nil, err
204 }
205 f := func(i int) float32 { return math.Float32frombits(le.Uint32(rec[i*4:])) }
206 out = append(out, detection{x1: f(0), y1: f(1), x2: f(2), y2: f(3), score: f(4), region: int(le.Uint32(rec[20:]))})
207 }
208 return out, nil
209}
210
211// procDetector runs the nefer-hints binary as a persistent child. It starts
212// on the first Detect, and again after the child died or misbehaved.
213type procDetector struct {
214 command, model string
215 conf float64
216 helloTimeout time.Duration
217 reqTimeout time.Duration
218
219 proc *process
220 seq uint64
221}
222
223func newProcDetector(command, model string, threshold float64) *procDetector {
224 return &procDetector{command: command, model: model, conf: threshold, helloTimeout: helloTimeout, reqTimeout: requestTimeout}
225}
226
227// waitDelay bounds cmd.Wait once the child is dead: a descendant that kept
228// stderr open must not block the reaping.
229const (
230 waitDelay = time.Second
231 closeGrace = time.Second
232)
233
234// process is one running child, alone in its process group. It keeps both
235// pipe ends so that closing them unblocks a pending read or write.
236type process struct {
237 cmd *exec.Cmd
238 in *os.File // write end of the child's stdin
239 out *os.File // read end of the child's stdout
240 rd *bufio.Reader
241
242 mu sync.Mutex
243 reaped bool
244 once sync.Once
245}
246
247// interrupt kills the whole process group and closes our pipe ends, without
248// waiting: it unblocks a pending read or write from a timer or a context
249// callback. kill reaps the child afterwards.
250func (p *process) interrupt() {
251 p.mu.Lock()
252 if !p.reaped {
253 _ = syscall.Kill(-p.cmd.Process.Pid, syscall.SIGKILL)
254 }
255 p.mu.Unlock()
256 _ = p.out.Close()
257 _ = p.in.Close()
258}
259
260// kill stops the child and its descendants and reaps it; safe to call more
261// than once. Wait is bounded by cmd.WaitDelay.
262func (p *process) kill() {
263 p.interrupt()
264 p.wait()
265}
266
267// wait reaps the child (once); bounded by cmd.WaitDelay after its death.
268func (p *process) wait() {
269 p.once.Do(func() {
270 _ = p.cmd.Wait()
271 p.mu.Lock()
272 p.reaped = true
273 p.mu.Unlock()
274 })
275}
276
277// cmd leaves confidence selection to the runner unless explicitly configured.
278func (d *procDetector) cmd() *exec.Cmd {
279 args := []string{"--model", d.model, "--imgsz", strconv.Itoa(imgsz), "--gpu", "auto"}
280 if d.conf != 0 {
281 args = append(args, "--conf", strconv.FormatFloat(d.conf, 'f', -1, 64))
282 }
283 return exec.Command(d.command, args...)
284}
285
286// interruptOnDeadline binds both callbacks to this process only. Stopping joins
287// any callback that already started, so its result can never be accepted.
288func interruptOnDeadline(ctx context.Context, timeout time.Duration, p *process) func() (timedOut, canceled bool) {
289 timerDone, contextDone := make(chan struct{}), make(chan struct{})
290 timer := time.AfterFunc(timeout, func() { defer close(timerDone); p.interrupt() })
291 stop := context.AfterFunc(ctx, func() { defer close(contextDone); p.interrupt() })
292 return func() (bool, bool) {
293 timedOut, canceled := !timer.Stop(), !stop()
294 if timedOut {
295 <-timerDone
296 }
297 if canceled {
298 <-contextDone
299 }
300 return timedOut, canceled
301 }
302}
303
304func (d *procDetector) start(ctx context.Context) error {
305 log := logging.For(ctx, "hints")
306 cmd := d.cmd()
307 cmd.Stderr = &lineLog{log: func(line string) { log.Debug().Str("detector", d.command).Msg(line) }}
308 cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
309 cmd.WaitDelay = waitDelay
310 inR, inW, err := os.Pipe()
311 if err != nil {
312 return err
313 }
314 outR, outW, err := os.Pipe()
315 if err != nil {
316 _ = inR.Close()
317 _ = inW.Close()
318 return err
319 }
320 cmd.Stdin, cmd.Stdout = inR, outW
321 err = cmd.Start()
322 // The child has its ends; ours are not needed any more.
323 _ = inR.Close()
324 _ = outW.Close()
325 if err != nil {
326 _ = inW.Close()
327 _ = outR.Close()
328 return fmt.Errorf("hints: start detector: %w", err)
329 }
330 p := &process{cmd: cmd, in: inW, out: outR, rd: bufio.NewReaderSize(outR, 1<<16)}
331 t0 := time.Now()
332 stop := interruptOnDeadline(ctx, d.helloTimeout, p)
333 device, err := readHello(p.rd)
334 timedOut, canceled := stop()
335 if err != nil || timedOut || canceled {
336 p.kill()
337 if ctx.Err() != nil {
338 err = ctx.Err()
339 } else if timedOut {
340 err = errors.New("hello timed out")
341 }
342 return fmt.Errorf("hints: detector hello: %w", err)
343 }
344 d.proc = p
345 log.Info().Str("device", device).Float64("start_ms", ms(time.Since(t0))).Msg("hint detector ready")
346 return nil
347}
348
349func (d *procDetector) Detect(ctx context.Context, f detFrame) ([]detection, error) {
350 if d.proc == nil {
351 if err := d.start(ctx); err != nil {
352 return nil, err
353 }
354 }
355 p := d.proc
356 d.seq++
357 stop := interruptOnDeadline(ctx, d.reqTimeout, p)
358 dets, err := d.roundTrip(p, f)
359 timedOut, canceled := stop()
360 if timedOut || canceled {
361 d.drop()
362 if ctx.Err() != nil {
363 return nil, ctx.Err()
364 }
365 return nil, fmt.Errorf("hints: detector timed out after %s", d.reqTimeout)
366 }
367 if err == nil {
368 return dets, nil
369 }
370 if errors.Is(err, errDetector) {
371 return nil, err // the child answered and keeps serving
372 }
373 d.drop()
374 switch {
375 case timedOut:
376 return nil, fmt.Errorf("hints: detector timed out after %s", d.reqTimeout)
377 case ctx.Err() != nil:
378 return nil, ctx.Err()
379 }
380 return nil, fmt.Errorf("hints: detector failed: %w", err)
381}
382
383func (d *procDetector) roundTrip(p *process, f detFrame) ([]detection, error) {
384 bw := bufio.NewWriterSize(p.in, 1<<16)
385 if err := writeRequest(bw, d.seq, f); err != nil {
386 return nil, err
387 }
388 if err := bw.Flush(); err != nil {
389 return nil, err
390 }
391 return readResponse(p.rd, d.seq)
392}
393
394// drop kills the current child: the next Detect starts a new one.
395func (d *procDetector) drop() {
396 if d.proc != nil {
397 d.proc.kill()
398 d.proc = nil
399 }
400}
401
402// Close asks the child to stop by closing its stdin (it exits 0 on EOF), then
403// kills it if it does not.
404func (d *procDetector) Close() error {
405 p := d.proc
406 d.proc = nil
407 if p == nil {
408 return nil
409 }
410 defer func() { _ = p.out.Close() }()
411 _ = p.in.Close()
412 done := make(chan struct{})
413 go func() { p.wait(); close(done) }()
414 select {
415 case <-done:
416 case <-time.After(closeGrace):
417 p.interrupt()
418 select {
419 case <-done:
420 case <-time.After(2 * waitDelay):
421 }
422 }
423 return nil
424}
425
426// lineLog is an io.Writer that hands each complete line to log.
427type lineLog struct {
428 log func(string)
429 buf []byte
430}
431
432func (l *lineLog) Write(b []byte) (int, error) {
433 l.buf = append(l.buf, b...)
434 for {
435 i := bytes.IndexByte(l.buf, '\n')
436 if i < 0 {
437 break
438 }
439 if line := string(l.buf[:i]); line != "" {
440 l.log(line)
441 }
442 l.buf = l.buf[i+1:]
443 }
444 if len(l.buf) > maxString {
445 l.buf = l.buf[:0]
446 }
447 return len(b), nil
448}