AI-tianyan/internal/infer/client.go

package infer

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"tianyan-edge/internal/config"
	"tianyan-edge/internal/event"
	"tianyan-edge/internal/stream"
)

type frameMsg struct {
	StreamID int     `json:"stream_id"`
	DeviceID string  `json:"device_id"`
	URL      string  `json:"url"`
	TS       float64 `json:"ts"`
	JPEGB64  string  `json:"jpeg_b64"`
}

type detection struct {
	Class string    `json:"class"`
	Conf  float64   `json:"conf"`
	BBox  []float64 `json:"bbox"`
}

type resultMsg struct {
	StreamID   int         `json:"stream_id"`
	DeviceID   string      `json:"device_id"`
	TS         float64     `json:"ts"`
	Detections []detection `json:"detections"`
}

type dedupState struct {
	mu   sync.Mutex
	seen map[string]float64
}

type Client struct {
	cfg    *config.Config
	frames <-chan stream.Frame
	events chan<- event.SuspectedEvent
	dedup  *dedupState
}

func NewClient(cfg *config.Config, frames <-chan stream.Frame, events chan<- event.SuspectedEvent) *Client {
	return &Client{
		cfg:    cfg,
		frames: frames,
		events: events,
		dedup:  &dedupState{seen: make(map[string]float64)},
	}
}

func (c *Client) Run(ctx context.Context) {
	var wg sync.WaitGroup
	for i := 0; i < c.cfg.InferWorkers; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			c.workerLoop(ctx, workerID)
		}(i)
	}
	wg.Wait()
}

func (c *Client) workerLoop(ctx context.Context, workerID int) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		conn, err := net.Dial("unix", c.cfg.InferSocket)
		if err != nil {
			log.Printf("infer[w%d]: socket connect failed: %v, retry 2s", workerID, err)
			time.Sleep(2 * time.Second)
			continue
		}
		log.Printf("infer[w%d]: socket connected", workerID)
		if !c.loop(ctx, conn, workerID) {
			conn.Close()
			return
		}
		conn.Close()
		time.Sleep(500 * time.Millisecond)
	}
}

func (c *Client) loop(ctx context.Context, conn net.Conn, workerID int) bool {
	for {
		select {
		case <-ctx.Done():
			return false
		case f := <-c.frames:
			if err := c.send(conn, f); err != nil {
				log.Printf("infer[w%d]: send error: %v", workerID, err)
				return true
			}
			res, err := c.recv(conn)
			if err != nil {
				log.Printf("infer[w%d]: recv error: %v", workerID, err)
				return true
			}
			c.emitEvents(f, res)
		}
	}
}

func (c *Client) emitEvents(f stream.Frame, res *resultMsg) {
	for _, d := range res.Detections {
		if c.isDuplicate(f.DeviceID, d.Class, f.TS) {
			continue
		}
		ev := event.SuspectedEvent{
			EventID:   mustUUID(),
			EdgeID:    c.cfg.EdgeID,
			DeviceID:  f.DeviceID,
			StreamURL: f.URL,
			TS:        f.TS,
			Class:     d.Class,
			Conf:      d.Conf,
			BBox:      d.BBox,
			ImageB64:  base64.StdEncoding.EncodeToString(f.JPEG),
		}
		select {
		case c.events <- ev:
		default:
			log.Println("infer: event queue full, dropping")
		}
	}
}

func (c *Client) isDuplicate(deviceID, className string, ts float64) bool {
	window := float64(c.cfg.DedupWindowSec)
	key := fmt.Sprintf("%s|%s", deviceID, className)
	c.dedup.mu.Lock()
	defer c.dedup.mu.Unlock()
	last, ok := c.dedup.seen[key]
	if ok && ts-last < window {
		return true
	}
	c.dedup.seen[key] = ts
	if len(c.dedup.seen) > 5000 {
		for k, v := range c.dedup.seen {
			if ts-v > 2*window {
				delete(c.dedup.seen, k)
			}
		}
	}
	return false
}

func (c *Client) send(conn net.Conn, f stream.Frame) error {
	msg := frameMsg{StreamID: f.StreamID, DeviceID: f.DeviceID, URL: f.URL,
		TS: f.TS, JPEGB64: base64.StdEncoding.EncodeToString(f.JPEG)}
	data, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	return writeMsg(conn, data)
}

func (c *Client) recv(conn net.Conn) (*resultMsg, error) {
	data, err := readMsg(conn)
	if err != nil {
		return nil, err
	}
	var r resultMsg
	return &r, json.Unmarshal(data, &r)
}

func writeMsg(conn net.Conn, data []byte) error {
	hdr := make([]byte, 4)
	binary.BigEndian.PutUint32(hdr, uint32(len(data)))
	if _, err := conn.Write(hdr); err != nil {
		return err
	}
	_, err := conn.Write(data)
	return err
}

func readMsg(conn net.Conn) ([]byte, error) {
	hdr := make([]byte, 4)
	if _, err := io.ReadFull(conn, hdr); err != nil {
		return nil, err
	}
	size := binary.BigEndian.Uint32(hdr)
	buf := make([]byte, size)
	_, err := io.ReadFull(conn, buf)
	return buf, err
}

func mustUUID() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b[:4]) + "-" + hex.EncodeToString(b[4:6]) + "-" +
		hex.EncodeToString(b[6:8]) + "-" + hex.EncodeToString(b[8:10]) + "-" +
		hex.EncodeToString(b[10:])
}