cabletest/harness/stream.go

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"net"
	"net/http"
	"os"
	"sync"

	"g.fc.run/theater/cabletest/internal/drm"
)

const streamAddr = ":8080"

const page = `<!doctype html>
<meta name="viewport" content="width=device-width">
<title>cabletest</title>
<body style="margin:0;background:#000;min-height:100vh;display:grid;place-items:center">
<canvas style="max-width:100vw;max-height:100vh"></canvas>
<script>
(async () => {
  const cv = document.querySelector('canvas'), g = cv.getContext('2d');
  for (;;) {
    try {
      const res = await fetch('frames');
      if (res.ok) {
        const rd = res.body.getReader();
        let buf = new Uint8Array(0);
        for (;;) {
          const {value, done} = await rd.read();
          if (done) break;
          const b = new Uint8Array(buf.length + value.length);
          b.set(buf); b.set(value, buf.length); buf = b;
          for (;;) {
            if (buf.length < 12) break;
            const dv = new DataView(buf.buffer, buf.byteOffset);
            const x = dv.getUint32(0, true), y = dv.getUint32(4, true), n = dv.getUint32(8, true);
            if (buf.length < 12 + n) break;
            const bmp = await createImageBitmap(new Blob([buf.slice(12, 12 + n)], {type: 'image/png'}));
            if (x === 0 && y === 0 && cv.width !== bmp.width) { cv.width = bmp.width; cv.height = bmp.height; }
            g.drawImage(bmp, x, y);
            buf = buf.slice(12 + n);
          }
        }
      }
    } catch (e) {}
    await new Promise(r => setTimeout(r, 1000));
  }
})();
</script>
`

type raw struct {
	buf   []byte
	pitch int
	pw    int
	ph    int
}

var (
	captured = make(chan raw, 8)

	latestMu sync.RWMutex
	latest   raw

	clientsMu sync.RWMutex
	clients   = make(map[chan []byte]bool)
)

func serveStream() error {
	g, err := openGrabber()
	if err != nil {
		return err
	}
	ln, err := net.Listen("tcp", streamAddr)
	if err != nil {
		return err
	}
	host, err := os.Hostname()
	if err != nil {
		host = "localhost"
	}
	fmt.Printf("panel stream at http://%s%s/\n", host, streamAddr)
	go captureLoop(g)
	go encodeLoop()
	mux := http.NewServeMux()
	mux.HandleFunc("/{$}", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, page)
	})
	mux.HandleFunc("/frames", frames)
	go http.Serve(ln, mux)
	return nil
}

// Between runs the crtc scans out the console framebuffer, which is not a
// dumb buffer and fails to map on every frame; each error reads once, not at
// frame rate.
func captureLoop(g *grabber) {
	lastErr := ""
	for {
		buf, pitch, pw, ph, err := g.rawFrame()
		if err != nil {
			if err.Error() != lastErr {
				lastErr = err.Error()
				fmt.Fprintf(os.Stderr, "capture: %v\n", err)
			}
			continue
		}
		lastErr = ""
		captured <- raw{buf, pitch, pw, ph}
	}
}

// latest advances before the delta fans out, so a handler that registers and
// then snapshots latest can only ever replay a delta its keyframe already
// includes, never miss one.
func encodeLoop() {
	for cur := range captured {
		prev := latest
		if prev.buf != nil && (prev.pitch != cur.pitch || prev.pw != cur.pw || prev.ph != cur.ph) {
			panic(fmt.Sprintf("framebuffer geometry changed: %dx%d pitch %d to %dx%d pitch %d",
				prev.pw, prev.ph, prev.pitch, cur.pw, cur.ph, cur.pitch))
		}
		latestMu.Lock()
		latest = cur
		latestMu.Unlock()
		if prev.buf != nil {
			broadcast(prev, cur)
		}
	}
}

// A client that stops draining its channel is closed so its page reconnects
// into a fresh keyframe.
func broadcast(prev, cur raw) {
	var delta []byte
	for _, rect := range diffRects(prev.buf, cur.buf, cur.pitch, cur.pw, cur.ph) {
		p, err := pack(rect.Min, maskRect(prev.buf, cur.buf, cur.pitch, cur.ph, rect))
		if err != nil {
			fmt.Fprintf(os.Stderr, "encode: %v\n", err)
			return
		}
		delta = append(delta, p...)
	}
	var dead []chan []byte
	clientsMu.RLock()
	for ch := range clients {
		select {
		case ch <- delta:
		default:
			dead = append(dead, ch)
		}
	}
	clientsMu.RUnlock()
	if len(dead) == 0 {
		return
	}
	clientsMu.Lock()
	for _, ch := range dead {
		if clients[ch] {
			delete(clients, ch)
			close(ch)
		}
	}
	clientsMu.Unlock()
}

// Frames are never written again once captured, so the lock only guards the
// swap of latest itself; deltas queue in the channel while the keyframe
// encodes, and replaying one the keyframe already includes is idempotent.
func frames(w http.ResponseWriter, r *http.Request) {
	ch := make(chan []byte, 16)
	clientsMu.Lock()
	clients[ch] = true
	clientsMu.Unlock()
	defer func() {
		clientsMu.Lock()
		delete(clients, ch)
		clientsMu.Unlock()
	}()
	latestMu.RLock()
	k := latest
	latestMu.RUnlock()
	if k.buf == nil {
		http.Error(w, "no frame captured yet", http.StatusServiceUnavailable)
		return
	}
	key, err := pack(image.Pt(0, 0), decode(k.buf, k.pitch, k.pw, k.ph))
	if err != nil {
		fmt.Fprintf(os.Stderr, "keyframe %s: %v\n", r.RemoteAddr, err)
		return
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	if _, err := w.Write(key); err != nil {
		return
	}
	w.(http.Flusher).Flush()
	ctx := r.Context()
	for {
		select {
		case <-ctx.Done():
			return
		case payload, ok := <-ch:
			if !ok {
				return
			}
			if _, err := w.Write(payload); err != nil {
				return
			}
			w.(http.Flusher).Flush()
		}
	}
}

// On rect-sized parts the dense level trades 3-6x the cpu for a few percent
// of size — enough to fall behind the capture loop on the once-per-second
// bursts and stutter the stream.
var encoder = png.Encoder{CompressionLevel: png.DefaultCompression}

func pack(at image.Point, img *image.NRGBA) ([]byte, error) {
	var buf bytes.Buffer
	var hdr [12]byte
	binary.LittleEndian.PutUint32(hdr[0:], uint32(at.X))
	binary.LittleEndian.PutUint32(hdr[4:], uint32(at.Y))
	buf.Write(hdr[:])
	if err := encoder.Encode(&buf, img); err != nil {
		return nil, err
	}
	b := buf.Bytes()
	binary.LittleEndian.PutUint32(b[8:], uint32(len(b)-12))
	return b, nil
}

// Distant regions changing together (a spinner up top, counters at the foot)
// would union into one near-full-height box, so changed portrait rows are
// merged into bands only across small gaps and each band gets its own box.
const bandGap = 32

// Compares in panel space, where a portrait column is one contiguous row, and
// reports the changed bounding boxes in portrait coordinates.
func diffRects(prev, cur []byte, pitch, pw, ph int) []image.Rectangle {
	var rows []int
	for row := 0; row < ph; row++ {
		if !bytes.Equal(prev[row*pitch:row*pitch+pw*4], cur[row*pitch:row*pitch+pw*4]) {
			rows = append(rows, row)
		}
	}
	if len(rows) == 0 {
		return nil
	}

	changed := make([]bool, pw)
	for _, row := range rows {
		a := prev[row*pitch : row*pitch+pw*4]
		b := cur[row*pitch : row*pitch+pw*4]
		for y := 0; y < pw; y++ {
			if !changed[y] && !bytes.Equal(a[y*4:y*4+4], b[y*4:y*4+4]) {
				changed[y] = true
			}
		}
	}

	var rects []image.Rectangle
	for y := 0; y < pw; {
		if !changed[y] {
			y++
			continue
		}
		y0, y1 := y, y+1
		for next := y1; next < pw && next < y1+bandGap; next++ {
			if changed[next] {
				y1 = next + 1
			}
		}
		rMin, rMax := -1, -1
		for _, row := range rows {
			if !bytes.Equal(prev[row*pitch+y0*4:row*pitch+y1*4], cur[row*pitch+y0*4:row*pitch+y1*4]) {
				if rMin < 0 {
					rMin = row
				}
				rMax = row
			}
		}
		rects = append(rects, image.Rect(ph-1-rMax, y0, ph-rMin, y1))
		y = y1
	}
	return rects
}

// The box holds only the pixels that changed, transparent everywhere else, so
// the client composites it straight over the last frame.
func maskRect(prev, cur []byte, pitch, ph int, rect image.Rectangle) *image.NRGBA {
	img := image.NewNRGBA(image.Rect(0, 0, rect.Dx(), rect.Dy()))
	for x := rect.Min.X; x < rect.Max.X; x++ {
		row := (ph - 1 - x) * pitch
		for y := rect.Min.Y; y < rect.Max.Y; y++ {
			o := row + y*4
			if bytes.Equal(prev[o:o+4], cur[o:o+4]) {
				continue
			}
			v := uint32(cur[o]) | uint32(cur[o+1])<<8 |
				uint32(cur[o+2])<<16 | uint32(cur[o+3])<<24
			cr, cg, cb := drm.Unpack(v)
			img.SetNRGBA(x-rect.Min.X, y-rect.Min.Y, color.NRGBA{R: cr, G: cg, B: cb, A: 255})
		}
	}
	return img
}